Diff Coverage

from einops import rearrange
import omni_training_custom_ops  # noqa: F401  # pylint: disable=unused-import


def _sink_attention_forward(
    query: torch.Tensor,
    sink_key: torch.Tensor,
    sink_value: torch.Tensor,
    num_heads: int,

    Returns:
        Sink-attention output and softmax statistics.
    """
    sink_query = rearrange(query, "b s n d -> s b (n d)")
    sink_key_sbh = rearrange(sink_key, "b s n d -> s b (n d)")
    sink_value_sbh = rearrange(sink_value, "b s n d -> s b (n d)")
    sink_output, sink_softmax_max, sink_softmax_sum = torch_npu.npu_fusion_attention(
        sink_query,
        sink_key_sbh,
        sink_value_sbh,
        num_heads,

        sparse_mode=0,
        actual_seq_qlen=None,
        actual_seq_kvlen=None,
    )[:3]
    sink_output = rearrange(sink_output, "s b (n d) -> b s n d", n=num_heads)
    return sink_output, sink_softmax_max, sink_softmax_sum


def _sink_attention_backward(
    query: torch.Tensor,
    sink_key: torch.Tensor,
    sink_value: torch.Tensor,
    sink_grad_output: torch.Tensor,

    Returns:
        Gradients of the sink query, key, and value.
    """
    sink_query = rearrange(query, "b s n d -> s b (n d)")
    sink_key_sbh = rearrange(sink_key, "b s n d -> s b (n d)")
    sink_value_sbh = rearrange(sink_value, "b s n d -> s b (n d)")
    sink_grad_query, sink_grad_key, sink_grad_value, *_ = torch_npu.npu_fusion_attention_grad(
        sink_query,
        sink_key_sbh,
        sink_value_sbh,
        sink_grad_output.to(sink_key.dtype),

        actual_seq_qlen=None,
        actual_seq_kvlen=None,
        sparse_mode=0,
    )
    sink_grad_query = rearrange(sink_grad_query, "s b (n d) -> b s n d", n=num_heads)
    sink_grad_key = rearrange(sink_grad_key, "s b (n d) -> b s n d", n=sink_key.size(2))
    sink_grad_value = rearrange(sink_grad_value, "s b (n d) -> b s n d", n=sink_value.size(2))
    return sink_grad_query, sink_grad_key, sink_grad_value


class _SparseAttentionRescale(torch.autograd.Function):
    """Autograd bridge for sparse attention with separate sink parameters."""

            output = F.pad(output, [0, query_rope.size(-1)])
        output = rearrange(output, "(b s) n d -> b s n d", b=batch_size, s=sequence_length)

        query = torch.cat([query_nope, query_rope], dim=-1)
        sink_output, sink_softmax_max, sink_softmax_sum = _sink_attention_forward(
            query, sink_key, sink_value, num_heads, scale, keep_prob
        )
        softmax_max_rescale = softmax_max.squeeze(0).view(batch_size, sequence_length, num_heads)
        softmax_sum_rescale = softmax_sum.squeeze(0).view(batch_size, sequence_length, num_heads)

        Returns:
            Gradients corresponding to the forward inputs.
        """
        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,

            )
        )
        sink_grad_output = rearrange(sink_scale * grad_rescaled_output, "b s n d -> s b (n d)")
        query = torch.cat([query_nope, query_rope], dim=-1)
        sink_grad_query, sink_grad_key, sink_grad_value = _sink_attention_backward(
            query,
            sink_key,
            sink_value,
            sink_grad_output,