nemo_automodel.components.models.glm_moe_dsa.layers

Bases: Module

Indexer for top-k sparse attention selection.

Based on the official GLM-5.2 training implementation. Computes attention scores between queries and keys with per-head weights, applies ReLU activation, then selects the top-k positions to attend to.

Key features:

• Uses LayerNorm (not RMSNorm) for key normalization
• Has a weights_proj that learns per-head importance weights
• Optional Hadamard transform (rotate_activation) on Q and K
• ReLU activation on attention scores before weighting

nemo_automodel.components.models.glm_moe_dsa.layers.GlmMoeDsaIndexer.forward(
    x: torch.Tensor,
    q_resid: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

nemo_automodel.components.models.glm_moe_dsa.layers.GlmMoeDsaIndexer.init_weights(
    init_std: float = 0.02
)

Bases: Module

Multi-head Latent Attention with Indexer for sparse attention.

This extends the V3 MLA with an Indexer module that performs top-k selection for sparse attention. The indexer uses the q_lora residual and hidden states to compute which positions to attend to.

nemo_automodel.components.models.glm_moe_dsa.layers.GlmMoeDsaMLA._build_sparse_mask(
    topk_indices: torch.Tensor,
    seq_len: int,
    qkv_format: str,
    bsz: int = 1,
    n_heads: int = 1,
    dtype: torch.dtype = torch.bfloat16,
    attention_mask: torch.Tensor | None = None,
    union_across_batches: bool = False
) -> torch.Tensor

nemo_automodel.components.models.glm_moe_dsa.layers.GlmMoeDsaMLA.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    prev_topk_indices: torch.Tensor | None = None,
    return_topk_indices: bool = False,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.glm_moe_dsa.layers.GlmMoeDsaMLA.init_weights(
    _buffer_device: torch.device,
    init_std: float = 0.02
)

Apply NON-interleaved (half-split) RoPE to the indexer’s rope slice.

The DSA indexer uses half-split RoPE (rotate_half: pair dim j with j + d/2), unlike the main MLA attention which uses interleaved RoPE. freqs_cis is the same complex tensor used by the MLA (exp(i * theta_j * pos) for j in [0, d/2)); we read its real/imag parts as cos/sin so the angles match exactly.

Parameters:

Resolve the DSA kernel path from the configured attention backend.

Apply Hadamard rotation activation.

Parameters:

Returns: torch.Tensor

Rotated tensor.

Convert a {0,1} keep-mask (1=attend, 0=mask) to an ADDITIVE key mask (0 / finfo.min).

Masked positions use finfo.min rather than -inf: F.scaled_dot_product_attention mishandles -inf float masks (its fused kernels corrupt the softmax). HF builds the attention bias with create_causal_mask, which likewise masks padding to finfo.min. The recipe, however, hands the model a 2D {0,1} padding mask; adding it to the scores raw (the previous behaviour) both fails to mask padding (0 -> +0 instead of finfo.min) AND adds +1.0 to every kept key, which is only softmax-invariant in fp32 — in bf16 the +1.0 swamps the (scaled) score differences and collapses attention toward uniform. A mask that is already additive (values <= 0) is returned unchanged.

Fallback hadamard_transform when fast_hadamard_transform is not available.

Multiply H_n @ u where H_n is the Hadamard matrix of dimension n x n. n must be a power of 2. Parameters: u: Tensor of shape (…, n) normalize: if True, divide the result by 2^{m/2} where m = log_2(n). Returns: product: Tensor of shape (…, n)