Diff Coverage

                    matched_converter = converter
                    break

            if matched_converter is not None:
                transform = conversion_mapping.setdefault(
                    renamed_key, deepcopy(matched_converter)
                )
            else:
                transform = conversion_mapping.setdefault(
                    renamed_key, WeightRenaming(original_key, renamed_key)
                )
                source_pattern = original_key
            transform.add_tensor(renamed_key, original_key, source_pattern, tensor)
        return conversion_mapping


    def _materialize_reverse_conversions(
        model: Any,

        model: Any,
        state_dict: dict[str, torch.Tensor],
    ) -> dict[str, torch.Tensor]:
        """Restore checkpoint names and layouts used before weight conversion."""
        weight_conversions = getattr(model, "_weight_conversions", None)
        if weight_conversions is None:
            weight_conversions = get_model_conversion_mapping(model, add_legacy=False)
        if not weight_conversions:
            return state_dict

        reverse_conversions = [
            conversion.reverse_transform() for conversion in weight_conversions
        ]
        renamings = [
            conversion
            for conversion in reverse_conversions
            if isinstance(conversion, WeightRenaming)
        ]
        converters = [
            conversion
            for conversion in reverse_conversions
            if isinstance(conversion, WeightConverter)
        ]
        conversion_mapping = _collect_reverse_conversions(
            state_dict, renamings, converters
        )
        return _materialize_reverse_conversions(model, conversion_mapping)

    sink_sum = sink_softmax_sum * torch.exp(sink_softmax_max - combined_max)
    combined_sum = output_sum + sink_sum
    output_scale = (output_sum / combined_sum).unsqueeze(3)
    sink_scale = (sink_sum / combined_sum).unsqueeze(3)
    return (
        (output * output_scale + sink_output * sink_scale).to(dtype=output.dtype),
        output_scale,
        sink_scale,
    )

        sink_scale,
    )


def _normal_attention_backward(ctx, grad_rescaled_output):
    """Compute the normal branch's Q/K/V gradients."""
    (
        query, key, value, sink_key, _, attention_mask, softmax_max, softmax_sum,
        _, _, rescaled_output, output_scale, _,
    ) = ctx.saved_tensors
    grad_output = rearrange(output_scale * grad_rescaled_output, "s b n d -> (b s) n d")
    return torch_npu.npu_fusion_attention_grad(
        query, key, value, grad_output.to(sink_key.dtype), ctx.num_heads, "TND",
        pse=None,
        padding_mask=None,
        atten_mask=attention_mask,

        softmax_layout="TND",
    )[:3]


def _sink_attention_backward(ctx, grad_rescaled_output):
    """Compute sink gradients and restore the query gradient's TND layout."""
    (
        query, _, _, sink_key, sink_value, _, _, _,
        sink_softmax_max, sink_softmax_sum, rescaled_output, _, sink_scale,
    ) = ctx.saved_tensors
    sink_query = rearrange(query, "(b s) n d -> s b (n d)", b=ctx.batch_size, s=ctx.sequence_length)
    grad_output = rearrange(sink_scale * grad_rescaled_output, "s b n d -> s b (n d)")
    grad_query, grad_key, grad_value = torch_npu.npu_fusion_attention_grad(
        sink_query, sink_key, sink_value, grad_output.to(sink_key.dtype), ctx.num_heads, "SBH",
        pse=None,
        padding_mask=None,
        atten_mask=None,

        actual_seq_qlen=None,
        actual_seq_kvlen=None,
        sparse_mode=0,
    )[:3]
    grad_query = rearrange(
        grad_query, "s b (n d) -> (b s) n d", b=ctx.batch_size, s=ctx.sequence_length, n=ctx.num_heads,
    )
    return grad_query, grad_key, grad_value


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

            rescaled_output,
            output_scale,
            sink_scale,
        )
        ctx.batch_size = batch_size
        ctx.sequence_length = sequence_length
        ctx.num_heads = num_heads
        ctx.scale = scale
        ctx.pre_tokens = pre_tokens
        ctx.next_tokens = next_tokens
        ctx.keep_prob = keep_prob
        ctx.sparse_mode = sparse_mode
        ctx.actual_seq_qlen = actual_seq_qlen
        ctx.actual_seq_kvlen = actual_seq_kvlen
        return rescaled_output, softmax_max

    @staticmethod
    def backward(ctx: Any, grad_rescaled_output: torch.Tensor, grad_softmax_max: torch.Tensor) -> tuple:
        """Run the explicit fusion-attention backward operators."""
        del grad_softmax_max
        grad_query, grad_key, grad_value = _normal_attention_backward(ctx, grad_rescaled_output)
        sink_grad_query, sink_grad_key, sink_grad_value = _sink_attention_backward(ctx, grad_rescaled_output)
        return (
            grad_query + sink_grad_query,
            grad_key,
            grad_value,

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


def _sparse_attention_forward(
    query_nope, compressed_kv, query_rope, key_rope, topk_indices,
    batch_size, sequence_length, scale, actual_seq_qlen, actual_seq_kvlen,
):
    """Run sparse attention and restore the padded output's BSND layout."""
    query_nope, compressed_kv, query_rope, key_rope = [
        rearrange(tensor, "b s n d -> (b s) n d")
        for tensor in (query_nope, compressed_kv, query_rope, key_rope)
    ]
    output, softmax_max, softmax_sum = torch.ops.custom.npu_sparse_flash_attention_enhance(
        query_nope, compressed_kv, compressed_kv, topk_indices, scale,
        block_table=None,
        actual_seq_lengths_query=actual_seq_qlen,
        actual_seq_lengths_kv=actual_seq_kvlen,

        sparse_mode=3,
        attention_mode=2,
        return_softmax_lse=True,
    )
    if query_rope.size(-1) > 0:
        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)
    return output, softmax_max, softmax_sum


def _sink_attention_forward(query_nope, query_rope, sink_key, sink_value, num_heads, scale, keep_prob):
    """Run the sink branch in SBH layout."""
    query = torch.cat([query_nope, query_rope], dim=-1)
    return torch_npu.npu_fusion_attention(
        rearrange(query, "b s n d -> s b (n d)"),
        rearrange(sink_key, "b s n d -> s b (n d)"),
        rearrange(sink_value, "b s n d -> s b (n d)"),
        num_heads, "SBH",

        actual_seq_kvlen=None,
    )[:3]


def _rescale_sparse_outputs(
    output, softmax_max, softmax_sum, sink_output, sink_softmax_max, sink_softmax_sum,
    batch_size, sequence_length, num_heads,
):
    """Combine the two attention outputs using stable softmax weights."""
    sink_output = rearrange(sink_output, "s b (n d) -> b s n d", n=num_heads)
    softmax_max = softmax_max.squeeze(0).view(batch_size, sequence_length, num_heads)
    softmax_sum = softmax_sum.squeeze(0).view(batch_size, sequence_length, num_heads)
    sink_softmax_max = sink_softmax_max[:, :, :, 0].transpose(1, 2)
    sink_softmax_sum = sink_softmax_sum[:, :, :, 0].transpose(1, 2)
    combined_max = torch.maximum(softmax_max, sink_softmax_max)
    output_sum = softmax_sum * torch.exp(softmax_max - combined_max)
    sink_sum = sink_softmax_sum * torch.exp(sink_softmax_max - combined_max)
    combined_sum = output_sum + sink_sum
    output_scale = (output_sum / combined_sum).unsqueeze(-1)
    sink_scale = (sink_sum / combined_sum).unsqueeze(-1)
    return (output * output_scale + sink_output * sink_scale).to(dtype=output.dtype), output_scale, sink_scale


def _sparse_attention_backward(ctx, grad_rescaled_output):
    """Compute sparse gradients and restore their BSND layout."""
    (
        query_nope, compressed_kv, query_rope, key_rope, sink_key, _, topk_indices,
        softmax_max, softmax_sum, _, _, rescaled_output, output_scale, _,
    ) = ctx.saved_tensors
    query_nope, compressed_kv, query_rope, key_rope = [
        rearrange(tensor, "b s n d -> (b s) n d")
        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 = 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 = rescaled_output[:, :, :-query_rope.size(-1)]
    gradients = torch.ops.custom.npu_sparse_flash_attention_grad_enhance(
        query_nope, compressed_kv, compressed_kv, topk_indices,
        grad_output.to(sink_key.dtype), rescaled_output, softmax_max, softmax_sum, ctx.scale,
        sparse_block_size=1,
        actual_seq_qlen=ctx.actual_seq_qlen,

        sparse_mode=3,
        attention_mode=2,
        deterministic=torch.are_deterministic_algorithms_enabled(),
    )
    return tuple(
        rearrange(tensor, "(b s) n d -> b s n d", b=ctx.batch_size, s=ctx.sequence_length)
        for tensor in gradients
    )

        for tensor in gradients
    )


def _sink_attention_backward(ctx, grad_rescaled_output):
    """Compute sink gradients and restore their BSND layout."""
    (
        query_nope, _, query_rope, _, sink_key, sink_value, _, _, _,
        sink_softmax_max, sink_softmax_sum, rescaled_output, _, sink_scale,
    ) = ctx.saved_tensors
    query = rearrange(torch.cat([query_nope, query_rope], dim=-1), "b s n d -> s b (n d)")
    grad_output = rearrange(sink_scale * grad_rescaled_output, "b s n d -> s b (n d)")
    grad_query, grad_key, grad_value = torch_npu.npu_fusion_attention_grad(
        query,
        rearrange(sink_key, "b s n d -> s b (n d)"),
        rearrange(sink_value, "b s n d -> s b (n d)"),
        grad_output.to(sink_key.dtype), ctx.num_heads, "SBH",

        actual_seq_qlen=None,
        actual_seq_kvlen=None,
        sparse_mode=0,
    )[:3]
    return (
        rearrange(grad_query, "s b (n d) -> b s n d", n=ctx.num_heads),
        rearrange(grad_key, "s b (n d) -> b s n d", n=sink_key.size(2)),
        rearrange(grad_value, "s b (n d) -> b s n d", n=sink_value.size(2)),
    )

        actual_seq_qlen: torch.Tensor,
        actual_seq_kvlen: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Run sparse and sink attention and rescale their outputs."""
        output, softmax_max, softmax_sum = _sparse_attention_forward(
            query_nope, compressed_kv, query_rope, key_rope, topk_indices,
            batch_size, sequence_length, scale, actual_seq_qlen, actual_seq_kvlen,
        )
        sink_output, sink_softmax_max, sink_softmax_sum = _sink_attention_forward(
            query_nope, query_rope, sink_key, sink_value, num_heads, scale, keep_prob,
        )
        rescaled_output, output_scale, sink_scale = _rescale_sparse_outputs(
            output, softmax_max, softmax_sum, sink_output, sink_softmax_max, sink_softmax_sum,
            batch_size, sequence_length, num_heads,
        )

            rescaled_output,
            output_scale,
            sink_scale,
        )
        ctx.batch_size = batch_size
        ctx.sequence_length = sequence_length
        ctx.num_heads = num_heads
        ctx.scale = scale
        ctx.keep_prob = keep_prob
        ctx.actual_seq_qlen = actual_seq_qlen
        ctx.actual_seq_kvlen = actual_seq_kvlen
        return rescaled_output, softmax_max, softmax_sum

    @staticmethod
    def backward(

        grad_softmax_sum: torch.Tensor,
    ) -> tuple:
        """Run the explicit sparse- and fusion-attention backward operators."""
        del grad_softmax_max, grad_softmax_sum
        grad_query_nope, grad_key, grad_value, grad_query_rope, grad_key_rope = _sparse_attention_backward(
            ctx, grad_rescaled_output,
        )
        sink_grad_query, sink_grad_key, sink_grad_value = _sink_attention_backward(ctx, grad_rescaled_output)
        sink_grad_query_nope, sink_grad_query_rope = torch.split(
            sink_grad_query,
            [ctx.saved_tensors[0].size(-1), ctx.saved_tensors[2].size(-1)],
            dim=-1,

from hyper_parallel.components.functional.sinkhorn import sinkhorn


def _mhc_pre_backward(ctx, dh_in, dh_post, dh_res, dh_x):
    """Invoke the MHC gradient operator with the saved forward state."""
    x, phi, alpha, gamma, h_post, inv_rms, h_mix, h_pre = ctx.saved_tensors
    return torch.ops.custom.npu_manifold_constrained_hyper_connection_pre_grad(
        x, phi, alpha, dh_in, dh_post, dh_res, inv_rms, h_mix, h_pre, h_post,
        gamma=gamma,
        hc_eps=ctx.hc_eps,
        grad_x_post=dh_x,

        dh_x: torch.Tensor,
    ) -> tuple:
        """Run the NPU MHC pre backward operator."""
        del dh_pre
        dx, dphi, dalpha, dbias, dgamma = _mhc_pre_backward(ctx, dh_in, dh_post, dh_res, dh_x)
        if not ctx.has_gamma:
            dgamma = None
        grads = [dx, dphi, dalpha, dbias, dgamma, None, None, None]
        return tuple(grads)

from __future__ import annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Optional

import torch  # pylint: disable=forbidden-backend-import
import torch_npu

import torch  # pylint: disable=forbidden-backend-import
import torch_npu


@dataclass
class _FusionAttentionContext:
    """Normalized arguments required by the NPU fusion-attention operator."""

    query: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor
    input_layout: str
    attention_mask: Optional[torch.Tensor]
    sparse_mode: int
    query_lengths: Optional[Sequence[int]]
    key_lengths: Optional[Sequence[int]]
    batch_size: int
    query_length: int
    head_dim: int
    pre_tokens: int
    next_tokens: int
    is_packed: bool


def _npu_attention_mask(attention_mask: torch.Tensor) -> torch.Tensor:
    """Convert the Transformers attention mask to the NPU mask convention."""

    npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
    return query, key, value, "BNSD", npu_mask, sparse_mode


def _prepare_fusion_attention_context(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,

    attention_mask: Optional[torch.Tensor],
    kwargs: dict[str, Any],
) -> _FusionAttentionContext:
    """Normalize model inputs into one NPU fusion-attention invocation."""
    query_lengths, key_lengths = resolve_packed_sequence_lengths(
        kwargs,
        query.shape[0] * query.shape[2],
        key.shape[0] * key.shape[2],
    )
    options = _attention_options(module, kwargs)
    prepared = _prepare_attention_inputs(
        query,
        key,
        value,
        attention_mask,

        is_causal=options[4],
        sliding_window=options[3],
        sparse_mode=options[2],
    )
    return _FusionAttentionContext(
        *prepared[:5],
        prepared[5],
        query_lengths,
        key_lengths,

        raise ValueError(
            "npu_fusion_attention_forward does not consume sparse attention indices; "
            "select a DSA sparse-attention implementation instead."
        )
    context = _prepare_fusion_attention_context(
        module,
        query,
        key,
        value,

        sparse_mode=context.sparse_mode,
        actual_seq_qlen=context.query_lengths,
        actual_seq_kvlen=context.key_lengths,
    )[0]
    if context.is_packed:
        output = output.reshape(
            context.batch_size,
            context.query_length,
            output.shape[1],
            output.shape[2],

    k_embed = torch_npu.npu_rotary_mul(k.clone(), cos, sin, rotary_mode="half")
    return q_embed, k_embed


def _apply_interleaved_rope(
    tensor: torch.Tensor,
    pass_through: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,

    sin: torch.Tensor,
    unsqueeze_dim: int,
) -> torch.Tensor:
    """Apply interleaved RoPE and append dimensions that bypass rotation."""
    sequence_first = (
        tensor.permute(2, 0, 1, 3)
        if unsqueeze_dim == 1
        else tensor.permute(1, 0, 2, 3)
    )
    seq_length, batch_size, num_heads, head_dim = sequence_first.shape
    if batch_size > 1 and seq_length > 1:
        rotated = torch_npu.npu_rotary_mul(
            sequence_first.reshape(batch_size * seq_length, 1, num_heads, head_dim),
            cos.reshape(batch_size * seq_length, 1, 1, head_dim),
            sin.reshape(batch_size * seq_length, 1, 1, head_dim),
            rotary_mode="interleave",

            sin.reshape(batch_size * seq_length, 1, 1, head_dim),
            rotary_mode="interleave",
        ).reshape(seq_length, batch_size, num_heads, head_dim)
    else:
        rotated = torch_npu.npu_rotary_mul(
            sequence_first.clone(),
            cos.unsqueeze(-2),
            sin.unsqueeze(-2),
            rotary_mode="interleave",

            cos.unsqueeze(-2),
            sin.unsqueeze(-2),
            rotary_mode="interleave",
        )
    rotated = rotated.permute(1, 2, 0, 3)
    rotated = torch.cat((rotated[..., 0::2], rotated[..., 1::2]), dim=-1)
    if unsqueeze_dim == 2:
        rotated = rotated.transpose(1, 2)
    return torch.cat((rotated, pass_through), dim=-1)


def apply_rotary_pos_emb_interleave(
    q: torch.Tensor,

        sin = sin.expand(q.shape[0], -1, -1)
    cos = cos.transpose(0, 1)
    sin = sin.transpose(0, 1)

    return (
        _apply_interleaved_rope(q_rot, q_pass, cos, sin, unsqueeze_dim),
        _apply_interleaved_rope(k_rot, k_pass, cos, sin, unsqueeze_dim),
    )

    padded_states = F.pad(hidden_states, (0, 0, padding, 0))
    if fused:
        padded_mask = F.pad(mome_mask, (padding, 0), value=False)
        weight = convolution.weight.squeeze(1).transpose(0, 1)
        mixed_states = aggregate_hidden(  # pylint: disable=not-callable
            padded_states.transpose(0, 1).contiguous(),
            weight,
            padded_mask,
        )

            device=device,
        )

    @staticmethod
    def _invalid_position_ids(
        position_ids: torch.Tensor | None,
        batch_size: int,
        seq_length: int,
    ) -> bool:

        batch_size: int,
        seq_length: int,
    ) -> bool:
        """Return whether position IDs violate the supported dense layout."""
        if position_ids is None:
            return False
        if position_ids.ndim != 2:
            return True
        return position_ids.shape[0] not in (1, batch_size) or position_ids.shape[1] != seq_length

    @staticmethod
    def _validate_forward_inputs(
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None,
        past_key_values: Any | None,

        """Validate inputs supported by the DeepSeek sparse-attention path."""
        batch_size, seq_length = hidden_states.shape[:-1]
        if past_key_values is not None:
            raise NotImplementedError("DeepseekV32DSAAttention does not support KV cache")
        if DeepseekV32DSAAttention._invalid_position_ids(position_ids, batch_size, seq_length):
            raise ValueError("position_ids must have shape [1 or batch_size, sequence_length]")
        if attention_mask is not None:
            expected_shape = (batch_size, 1, seq_length, seq_length)
            if tuple(attention_mask.shape) != expected_shape:

        query_tnd, key_tnd, q_rot_tnd, k_rot_tnd = (
            tensor.reshape(-1, tensor.shape[2], tensor.shape[3])
            for tensor in attention_states
        )
        aux_loss = dsa_kl_loss(  # pylint: disable=not-callable
            *index_states, query_tnd, key_tnd,
            topk_indices, *softmax_stats, q_rot_tnd, k_rot_tnd,
            actual_q_len, actual_kv_len, self.scaling, self.dsa_loss_coeff,
        )

            actual_q_len, actual_kv_len, self.scaling, self.dsa_loss_coeff,
        )
        return aux_loss_auto_scale(attn_output, aux_loss)

    def _prepare_attention_states(self, hidden_states, position_embeddings):
        """Project attention states and apply their rotary embeddings."""
        q_resid, absorbed_query_states, kv_nope, q_rot, k_rot, kv_weight = self._project_attention_states(
            hidden_states
        )

        )
        q_rot, k_rot = _apply_attention_rope(
            q_rot, k_rot, position_embeddings, interleaved=self.rotary_interleaved
        )
        return q_resid, (absorbed_query_states, kv_nope, q_rot, k_rot), kv_weight

    def _run_indexer(self, hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs):
        """Resolve packed boundaries and select sparse attention indices."""
        index_states = self._project_index_states(hidden_states, q_resid, position_embeddings)
        packed_kwargs = dict(kwargs)
        packed_kwargs["actual_seq_len"] = actual_seq_len
        actual_q_len, actual_kv_len = resolve_packed_sequence_lengths(
            packed_kwargs,

        )
        actual_kv_len = self._get_actual_seq_len(
            actual_kv_len, *hidden_states.shape[:-1], hidden_states.device,
        )
        indexed = dsa_indexer(  # pylint: disable=not-callable
            *index_states, actual_q_len, actual_kv_len, self.index_topk
        )
        return indexed[0], indexed[1:], actual_q_len, actual_kv_len

    def _compute_sparse_attention(
        self, hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs,
    ):
        """Run sparse attention and attach its optional auxiliary loss."""
        topk_indices, index_states, actual_q_len, actual_kv_len = self._run_indexer(
            hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs,
        )
        attn_output, softmax_max, softmax_sum = dsa_sparse_attention(  # pylint: disable=not-callable
            *attention_states, topk_indices, self.scaling, actual_q_len, actual_kv_len,
        )
        return self._apply_auxiliary_loss(
            attn_output, index_states, attention_states, topk_indices,
            (softmax_max, softmax_sum), actual_q_len, actual_kv_len,
        )

    def _compute_attention_output(
        self, hidden_states, position_embeddings, actual_seq_len, kwargs,
    ):
        """Compute sparse attention and restore its value projection."""
        batch_size, seq_length = hidden_states.shape[:-1]
        q_resid, attention_states, kv_weight = self._prepare_attention_states(
            hidden_states, position_embeddings,
        )
        attn_output = self._compute_sparse_attention(
            hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs,
        )
        value_up_weight = kv_weight[:, self.qk_nope_head_dim :].transpose(1, 2)
        return _restore_attention_projection(
            attn_output,
            value_up_weight,
            num_heads=self.num_heads,
            batch_size=batch_size,

            kv_lora_rank=self.kv_lora_rank,
            value_head_dim=self.v_head_dim,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,

        actual_seq_len: torch.Tensor | Sequence[int] | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, None]:
        """Run causal or packed DSA with the NPU sparse-attention kernels."""
        self._validate_forward_inputs(hidden_states, attention_mask, past_key_values, position_ids, kwargs)
        attn_output = self._compute_attention_output(
            hidden_states, position_embeddings, actual_seq_len, kwargs,
        )
        attn_output = self.o_proj(attn_output)
        return attn_output, None

        self.use_mome = bool(getattr(module, "use_mome", False))
        self.use_fused_mome = bool(getattr(module, "use_fused_mome", False))
        self.param_sink_number = int(getattr(module, "param_sink_number", 0))
        self.param_sink_scalar = getattr(module, "param_sink_scalar", None)
        self.apply_fa_rescale = bool(
            getattr(module, "apply_fa_rescale", False)
        )
        self.attention_dropout = getattr(module, "attention_dropout", nn.Dropout(0.0))
        if self.param_sink_number > 0:

            if self.param_sink_scalar:
                raise NotImplementedError(
                    "DSAAttention does not support scalar parameter sink"
                )
            if not self.apply_fa_rescale:
                raise NotImplementedError(
                    "DSAAttention supports parameter sink through "
                    "apply_fa_rescale only"
                )

        """Run sparse attention with or without parameter-sink rescaling."""
        absorbed_query, kv_nope, q_rot, k_rot = attention_states
        sparse_scale = self.qk_head_dim**-0.5
        if self.param_sink_number <= 0:
            return dsa_sparse_attention(  # pylint: disable=not-callable
                absorbed_query, kv_nope, q_rot, k_rot, topk_indices,
                sparse_scale, actual_q_len, actual_kv_len,
            )
        sink_key, sink_value = self._prepare_param_sink(batch_size)
        return dsa_sparse_attention_rescale(  # pylint: disable=not-callable
            absorbed_query, kv_nope, q_rot, k_rot, sink_key, sink_value,
            topk_indices, batch_size, seq_length, self.num_heads, sparse_scale,
            1 - self.attention_dropout.p, actual_q_len, actual_kv_len,
        )

        query_tnd, key_tnd, q_rot_tnd, k_rot_tnd = (
            tensor.reshape(-1, tensor.shape[2], tensor.shape[3])
            for tensor in attention_states
        )
        aux_loss = dsa_kl_loss(  # pylint: disable=not-callable
            *index_states, query_tnd, key_tnd,
            topk_indices, *softmax_stats, q_rot_tnd, k_rot_tnd,
            actual_q_len, actual_kv_len, self.qk_head_dim**-0.5, self.dsa_loss_coeff,
        )

            actual_q_len, actual_kv_len, self.qk_head_dim**-0.5, self.dsa_loss_coeff,
        )
        return aux_loss_auto_scale(attn_output, aux_loss)

    def _prepare_attention_states(self, hidden_states, position_embeddings, mome_mask):
        """Project attention states and apply their rotary embeddings."""
        q_resid, absorbed_query_states, kv_nope, q_rot, k_rot, kv_weight = self._project_attention_states(
            hidden_states, mome_mask
        )

        )
        q_rot, k_rot = _apply_attention_rope(
            q_rot, k_rot, position_embeddings, interleaved=self.rotary_interleaved
        )
        return q_resid, (absorbed_query_states, kv_nope, q_rot, k_rot), kv_weight

    def _run_indexer(self, hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs):
        """Resolve packed boundaries and select sparse attention indices."""
        index_states = self._project_index_states(hidden_states, q_resid, position_embeddings)
        packed_kwargs = dict(kwargs)
        packed_kwargs["actual_seq_len"] = actual_seq_len
        actual_q_len, actual_kv_len = resolve_packed_sequence_lengths(
            packed_kwargs,

        )
        actual_kv_len = self._get_actual_seq_len(
            actual_kv_len, *hidden_states.shape[:-1], hidden_states.device,
        )
        indexed = dsa_indexer(  # pylint: disable=not-callable
            *index_states, actual_q_len, actual_kv_len, self.index_topk
        )
        return indexed[0], indexed[1:], actual_q_len, actual_kv_len

    def _compute_sparse_attention(
        self, hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs,
    ):
        """Run sparse attention and attach its optional auxiliary loss."""
        topk_indices, index_states, actual_q_len, actual_kv_len = self._run_indexer(
            hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs,
        )
        attn_output, softmax_max, softmax_sum = self._run_sparse_attention(
            attention_states, topk_indices, actual_q_len, actual_kv_len, *hidden_states.shape[:-1],

        )
        attn_output, softmax_max, softmax_sum = self._run_sparse_attention(
            attention_states, topk_indices, actual_q_len, actual_kv_len, *hidden_states.shape[:-1],
        )
        return self._apply_auxiliary_loss(
            attn_output, index_states, attention_states, topk_indices,
            (softmax_max, softmax_sum), actual_q_len, actual_kv_len,
        )

    def _compute_attention_output(
        self, hidden_states, position_embeddings, actual_seq_len, kwargs, mome_mask,
    ):
        """Compute sparse attention and restore its value projection."""
        batch_size, seq_length = hidden_states.shape[:-1]
        q_resid, attention_states, kv_weight = self._prepare_attention_states(
            hidden_states, position_embeddings, mome_mask,
        )
        attn_output = self._compute_sparse_attention(
            hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs,
        )
        value_up_weight = kv_weight[:, self.qk_nope_head_dim :].transpose(1, 2)
        return _restore_attention_projection(
            attn_output,
            value_up_weight,
            num_heads=self.num_heads,
            batch_size=batch_size,

            kv_lora_rank=self.kv_lora_rank,
            value_head_dim=self.v_head_dim,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,

        mome_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Run DSA with its configured MOME and parameter-sink paths."""
        self._validate_forward_inputs(
            attention_mask, kv_reuse_states, past_key_values, cache_position, output_attentions
        )
        attn_output = self._compute_attention_output(
            hidden_states, position_embeddings, actual_seq_len, kwargs, mome_mask,
        )
        if self.use_mome:
            attn_output = apply_mome(

        and (biases[0] is None or len({bias.requires_grad for bias in biases}) == 1)
    )


def _apply_gqa_rope(
    query: torch.Tensor,
    key: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
    rotary_interleaved: bool,

    rotary_interleaved: bool,
    head_dim: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply the configured rotary convention to grouped-query attention."""
    if position_embeddings is None:
        return query, key
    cos, sin = position_embeddings
    if rotary_interleaved:
        return apply_rotary_pos_emb_interleave(query, key, cos, sin, unsqueeze_dim=2)
    if cos.shape[-1] >= head_dim:
        return apply_rotary_pos_emb(query, key, cos, sin, unsqueeze_dim=2)
    rotary_dim = cos.shape[-1]
    query_rot, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
    key_rot, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
    cos = cos.unsqueeze(2)
    sin = sin.unsqueeze(2)
    query = torch.cat((query_rot * cos + rotate_half(query_rot) * sin, query_pass), dim=-1)
    key = torch.cat((key_rot * cos + rotate_half(key_rot) * sin, key_pass), dim=-1)
    return query, key


@module_replacement
class GQAAttention(nn.Module):

            query_states = self.q_norm(query_states)
        if self.k_norm is not None:
            key_states = self.k_norm(key_states)

        query_states, key_states = _apply_gqa_rope(
            query_states,
            key_states,
            position_embeddings,
            self.rotary_interleaved,

            query_states = self.q_norm(query_states)
        if self.k_norm is not None:
            key_states = self.k_norm(key_states)

        query_states, key_states = _apply_gqa_rope(
            query_states,
            key_states,
            position_embeddings,
            self.rotary_interleaved,

        # GroupGemm state (NPU-only)
        self._group_list = None
        self._tokens_per_expert_gmm = None

        self._initialize_projection_layout(
            source_gate_up, source_down, source_is_transposed, gated_linear_unit
        )
        self._initialize_biases(module)

        self.has_gate = gated_linear_unit
        self.has_bias = self.add_bias
        self.is_concatenated = bool(getattr(module, "is_concatenated", True))
        if self.has_gate and not self.is_concatenated:
            raise ValueError("GroupedExperts requires concatenated gate/up expert weights")
        self.is_transposed = True
        self.train(module.training)

    def _initialize_projection_layout(
        self, source_gate_up, source_down, source_is_transposed, gated_linear_unit,
    ) -> None:
        """Resolve expert weight layouts and initialize their target parameters."""
        hidden_size = self.hidden_size
        intermediate_size = self.intermediate_size
        fc1_output_size = intermediate_size
        if gated_linear_unit:
            fc1_output_size *= 2

from __future__ import annotations

from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any

# This package provides PyTorch-specific high-performance modules.
# pylint: disable=forbidden-backend-import

from hyper_parallel.components.functional import apply_rotary_pos_emb, apply_rotary_pos_emb_interleave
from hyper_parallel.components.functional import npu_fusion_attention_forward


@dataclass
class _MLALatents:
    batch_size: int
    sequence_length: int
    query_pass: torch.Tensor
    query_rope: torch.Tensor
    kv_nope: torch.Tensor
    key_rope: torch.Tensor


@module_replacement
class MLAAttention(nn.Module):

                ),
            )
        return transforms

    def _project_latents(self, hidden_states: torch.Tensor) -> _MLALatents:
        """Project hidden states into query and compressed KV latent states."""
        batch_size, sequence_length = hidden_states.shape[:-1]
        latent_states = self.linear_qkv(hidden_states)
        query_latent, kv_nope, key_rope = torch.split(
            latent_states,
            (self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim),
            dim=-1,
        )
        query_states = self.q_b_proj(self.q_a_layernorm(query_latent)).view(
            batch_size,
            sequence_length,
            self.num_heads,
            self.qk_head_dim,

            sequence_length,
            self.num_heads,
            self.qk_head_dim,
        )
        query_pass, query_rope = torch.split(
            query_states,
            (self.qk_nope_head_dim, self.qk_rope_head_dim),
            dim=-1,
        )
        return _MLALatents(
            batch_size,
            sequence_length,
            query_pass,
            query_rope.transpose(1, 2),

            ),
            key_rope.view(batch_size, 1, sequence_length, self.qk_rope_head_dim),
        )

    def _project_attention_inputs(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
        past_key_values: Any | None,

        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
        past_key_values: Any | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build query, key, and value tensors from compressed latent states."""
        latents = self._project_latents(hidden_states)
        query_rope, key_rope = latents.query_rope, latents.key_rope
        if position_embeddings is not None:
            cos, sin = position_embeddings
            rope_fn = apply_rotary_pos_emb_interleave if self.rotary_interleaved else apply_rotary_pos_emb
            query_rope, key_rope = rope_fn(query_rope, key_rope, cos, sin)
        kv_nope = latents.kv_nope
        if past_key_values is not None:
            kv_nope, key_rope = past_key_values.update(kv_nope, key_rope, self.layer_idx)
        kv_states = self.kv_b_proj(kv_nope).view(
            latents.batch_size,
            kv_nope.shape[2],
            self.num_heads,

            kv_nope.shape[2],
            self.num_heads,
            self.qk_nope_head_dim + self.v_head_dim,
        ).transpose(1, 2)
        key_nope, value_states = torch.split(
            kv_states,
            (self.qk_nope_head_dim, self.v_head_dim),
            dim=-1,
        )
        return (
            torch.cat((latents.query_pass.transpose(1, 2), query_rope), dim=-1),
            torch.cat((key_nope, key_rope.expand(-1, self.num_heads, -1, -1)), dim=-1),
            value_states,
        )