> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.mistral_large4.attention

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

## Module Contents

### Classes

| Name                                                                                                          | Description                                                                     |
| ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`MistralLarge4Attention`](#nemo_automodel-components-models-mistral_large4-attention-MistralLarge4Attention) | Direct-query MLA; KV consumers retain their source's autograd graph.            |
| [`MistralLarge4Indexer`](#nemo_automodel-components-models-mistral_large4-attention-MistralLarge4Indexer)     | Frozen hard DSA selector; an indexer training objective has not been specified. |
| [`SequenceLayout`](#nemo_automodel-components-models-mistral_large4-attention-SequenceLayout)                 | Local token metadata plus optional CP-global keys, all flattened in rank order. |

### Functions

| Name                                                                                  | Description                                                    |
| ------------------------------------------------------------------------------------- | -------------------------------------------------------------- |
| [`sparse_mla`](#nemo_automodel-components-models-mistral_large4-attention-sparse_mla) | Run sparse MLA on local packed queries and gathered latent KV. |

### Data

[`_INDEXER_NORM_EPS`](#nemo_automodel-components-models-mistral_large4-attention-_INDEXER_NORM_EPS)

### API

```python
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]`

---

```python
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

```python
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,

```python
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`**

---

```python
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

```python
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`

---

```python
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.

```python
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

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