Diff Coverage

        grad_softmax_max: torch.Tensor,
        grad_softmax_sum: torch.Tensor,
    ) -> tuple:
        """Run the explicit sparse- and fusion-attention backward operators."""
        del grad_softmax_max, grad_softmax_sum
        (
            query_nope,
            compressed_kv,
            query_rope,

            for tensor in (query_nope, compressed_kv, query_rope, key_rope)
        ]
        grad_output = rearrange(output_scale * grad_rescaled_output, "b s n d -> (b s) n d")
        rescaled_output_tnd = rearrange(rescaled_output, "b s n d -> (b s) n d")
        if query_rope.size(-1) > 0:
            grad_output = grad_output[:, :, :-query_rope.size(-1)]
            rescaled_output_tnd = rescaled_output_tnd[:, :, :-query_rope.size(-1)]

        grad_query_nope, grad_key, grad_value, grad_query_rope, grad_key_rope = (
            torch.ops.custom.npu_sparse_flash_attention_grad_enhance(
                query_nope_tnd,