Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMistral Large4nemo_automodel.components.models.mistral_large4.attention

nemo_automodel.components.models.mistral_large4.attention

View as Markdown

Floating DSA/MLA baseline with explicit, differentiable per-forward KV sharing.

Module Contents

Classes

NameDescription
MistralLarge4AttentionDirect-query MLA; KV consumers retain their source’s autograd graph.
MistralLarge4IndexerFrozen hard DSA selector; an indexer training objective has not been specified.
SequenceLayoutLocal token metadata plus optional CP-global keys, all flattened in rank order.

Functions

NameDescription
sparse_mlaRun sparse MLA on local packed queries and gathered latent KV.

Data

_INDEXER_NORM_EPS

API

class nemo_automodel.components.models.mistral_large4.attention.MistralLarge4Attention(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    layer_idx: int,
    dtype: torch.dtype
)

Bases: Module

Direct-query MLA; KV consumers retain their source’s autograd graph.

attn_backend
= backend.attn
indexer
kv_a_layernorm
kv_a_proj_with_mqa
kv_b_proj
nope_dim
o_proj
q_proj
shares_kv
= config.kv_source_layers[layer_idx] != layer_idx
ssmax
= SSMax(config.num_attention_heads, dtype)
use_rope
= not config.nope_layer_flags[layer_idx]
nemo_automodel.components.models.mistral_large4.attention.MistralLarge4Attention._sparse_forward(
    q: torch.Tensor,
    latent: torch.Tensor,
    k_pe: torch.Tensor,
    indices: torch.Tensor,
    layout: nemo_automodel.components.models.mistral_large4.attention.SequenceLayout
) -> torch.Tensor

Absorb KV-B into queries/outputs and gather compressed KV for sparse MLA.

Parameters:

q
torch.Tensor

Local [tokens, heads, head_dim] rotated/scaled queries.

latent
torch.Tensor

Local [tokens, 512] normalized KV, possibly shared from another layer.

k_pe
torch.Tensor

Local [tokens, 64] rotated keys, or [tokens, 0] for NoPE.

indices
torch.Tensor

Integer [tokens, selected_keys] global KV indices with -1 padding.

layout
SequenceLayout

Local/global document, position and validity metadata as documented by SequenceLayout. Global keys use CP-rank concatenation order.

Returns: torch.Tensor

Local [tokens, heads, value_dim] attention outputs. KV-B effective weights

nemo_automodel.components.models.mistral_large4.attention.MistralLarge4Attention.forward(
    x: torch.Tensor,
    tables: tuple[torch.Tensor, torch.Tensor],
    layout: nemo_automodel.components.models.mistral_large4.attention.SequenceLayout,
    shared_kv: tuple[torch.Tensor, torch.Tensor] | None,
    indices: torch.Tensor | None
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor], torch.Tensor]

Run selected sparse MLA or the expanded-KV SDPA reference (including LoRA).

Parameters:

x
torch.Tensor

Floating tensor [tokens, hidden].

tables
tuple[torch.Tensor, torch.Tensor]

FP32 cosine/sine pair, each [tokens, rotary_dim / 2].

layout
SequenceLayout

Local positions/document IDs/validity [tokens], optional global CP key metadata.

shared_kv
tuple[torch.Tensor, torch.Tensor] | None

Optional source latent [tokens, latent_dim] and rotated keys [tokens, rotary_dim], or [tokens, 0] for NoPE. Never detached.

indices
torch.Tensor | None

Optional preceding full indexer’s int64 selected keys [tokens, min(topk, tokens)], with -1 padding.

Returns: torch.Tensor

Output [tokens, hidden], latent/rotary pair in shared_kv layouts,

class nemo_automodel.components.models.mistral_large4.attention.MistralLarge4Indexer(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    layer_idx: int,
    dtype: torch.dtype
)

Bases: Module

Frozen hard DSA selector; an indexer training objective has not been specified.

k_norm
q_input_norm
ssmax
= SSMax(config.index_n_heads, dtype)
use_rope
= not config.indexer_nope_layer_flags[layer_idx]
weights_proj
wk
wq_b
nemo_automodel.components.models.mistral_large4.attention.MistralLarge4Indexer.forward(
    x: torch.Tensor,
    tables: tuple[torch.Tensor, torch.Tensor],
    layout: nemo_automodel.components.models.mistral_large4.attention.SequenceLayout
) -> torch.Tensor

Select valid causal keys with stable ascending-key tie breaking.

Parameters:

x
torch.Tensor

Floating tensor [tokens, hidden].

tables
tuple[torch.Tensor, torch.Tensor]

FP32 cosine/sine pair, each [tokens, rotary_dim / 2].

layout
SequenceLayout

Local positions/document IDs/validity [tokens], optional global CP key metadata.

Returns: torch.Tensor

Independent int64 indices [tokens, min(topk, tokens)], with -1 for

class nemo_automodel.components.models.mistral_large4.attention.SequenceLayout(
    positions: torch.Tensor,
    sequence_ids: torch.Tensor,
    valid: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    key_sequence_ids: torch.Tensor | None = None,
    key_valid: torch.Tensor | None = None
)
Dataclass

Local token metadata plus optional CP-global keys, all flattened in rank order.

cp_mesh
DeviceMesh | None = None
key_sequence_ids
Tensor | None = None
key_valid
Tensor | None = None
positions
Tensor
sequence_ids
Tensor
valid
Tensor
nemo_automodel.components.models.mistral_large4.attention.SequenceLayout.allowed(
    start: int,
    end: int
) -> torch.Tensor

Return an independent bool [query_tokens, key_tokens] causal document mask.

nemo_automodel.components.models.mistral_large4.attention.sparse_mla(
    q: torch.Tensor,
    kv: torch.Tensor,
    indices: torch.Tensor,
    value_weight: torch.Tensor,
    scale: float
) -> torch.Tensor

Run sparse MLA on local packed queries and gathered latent KV.

Parameters:

q
torch.Tensor

Contiguous CUDA BF16 [local_tokens, heads, 576] absorbed queries. The last 64 coordinates are rotary keys, or zeros in NoPE layers.

kv
torch.Tensor

CUDA BF16 [global_tokens, 576] latent/rotary KV in CP-rank order.

indices
torch.Tensor

Integer [local_tokens, selected_keys] global KV coordinates; -1 denotes invalid slots. Caller enforces document/causal boundaries.

value_weight
torch.Tensor

BF16 [heads, value_dim, 512] differentiable KV-B value weight.

scale
float

Original Mistral head-dimension attention scale.

Returns: torch.Tensor

BF16 [local_tokens, heads, value_dim]. Empty query rows are zero and

nemo_automodel.components.models.mistral_large4.attention._INDEXER_NORM_EPS = 1e-06