> This page is for version 0.7.0 · 26.10.
> 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.flint.layers

Flint's sink-aware attention and shared dense/MoE decoder blocks.

Eager, SDPA and TE attention have distinct numerical validation scopes. Magi
uses model-owned packed token plans and requires an upstream learned-sink
training kernel. FP32 component parity does not establish full-model parity
with the serving kernels' BF16 rounding order.

## Module Contents

### Classes

| Name                                                                                          | Description                                                              |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`FlintAttention`](#nemo_automodel-components-models-flint-layers-FlintAttention)             | Separate Q/K/V, per-head RMSNorm, selective Neox RoPE and learned sinks. |
| [`FlintDecoderLayer`](#nemo_automodel-components-models-flint-layers-FlintDecoderLayer)       | Pre-norm attention and a dense or shared NeMo MoE, each with a residual. |
| [`_FlintSinkParameters`](#nemo_automodel-components-models-flint-layers-_FlintSinkParameters) | Own the trainable FP32 denominator logits in a separate dtype unit.      |

### Functions

| Name                                                                                    | Description                                                             |
| --------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_validate_backend`](#nemo_automodel-components-models-flint-layers-_validate_backend) | Reject backend choices outside the supported native component contract. |

### API

```python
class nemo_automodel.components.models.flint.layers.FlintAttention(
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Separate Q/K/V, per-head RMSNorm, selective Neox RoPE and learned sinks.

**`_keep_in_fp32_modules_strict`** `= ['_fp32_params']`

---

**`_te_cp_size`** `= 1`

---

**`attention_compute_dtype`**

---

**`attn_module`** `Module | None`

Expose shared TE attention for CP discovery without a second state path.

---

**`eps`** `= config.rms_norm_eps`

---

**`head_dim`** `= config.head_dim`

---

**`hidden_size`** `= config.hidden_size`

---

**`initializer_range`** `= config.initializer_range`

---

**`k_norm`**

---

**`k_proj`**

---

**`num_heads`** `= config.num_attention_heads`

---

**`num_kv_heads`** `= config.num_key_value_heads`

---

**`o_proj`**

---

**`q_norm`**

---

**`q_proj`**

---

**`rope_theta`** `= float(config.rope_parameters['rope_theta'])`

---

**`sliding_window`**

---

**`use_rope`** `= layer_idx not in config.no_rope_layer_indices`

---

**`v_proj`**

---

```python
nemo_automodel.components.models.flint.layers.FlintAttention._apply_rope(
    values: torch.Tensor,
    position_ids: torch.Tensor
) -> torch.Tensor
```

Rotate corresponding channels in the two halves of each head.

**Parameters:**

**`values`** `torch.Tensor`

Tensor \[batch, sequence, heads, head_dim] or THD \[tokens, heads, head_dim].

---

**`position_ids`** `torch.Tensor`

Integer tensor \[batch, sequence] or \[tokens] on values' device.

---

**Returns:** `torch.Tensor`

Rotated tensor of shape \[batch, sequence, heads, head_dim], in values'

```python
nemo_automodel.components.models.flint.layers.FlintAttention.forward(
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
```

Compute full-sequence causal attention with a zero-value sink category.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Floating projection-dtype tensor \[batch, sequence,
hidden_size], or \[tokens, hidden_size] for explicit THD packing.

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional integer tensor of shape \[batch, sequence] or
document-local \[tokens] for THD;
contiguous zero-based positions are used by default.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary tensor of shape \[batch, sequence],
with one/True for valid tokens. Both padded queries and keys are
excluded from real attention; the sink remains valid. TE BSHD
supports trailing padding; arbitrary padding uses THD packing.

---

**`**attn_kwargs`** `Any` — default: \{}

Shared TE THD boundaries, padding metadata and formats,
model-owned Magi plan, or integer document_ids \[batch, sequence]
for the eager/SDPA block-diagonal mask.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden_size] in the output

```python
nemo_automodel.components.models.flint.layers.FlintAttention.init_weights(
    buffer_device: torch.device,
    init_std: float | None = None
) -> None
```

Initialize projections, head-normalization scales and trainable sinks.

**Parameters:**

**`buffer_device`** `torch.device`

Device context after any meta materialization.

---

**`init_std`** `float | None` — default: None

Optional normal-initialization deviation for projections.

---

```python
class nemo_automodel.components.models.flint.layers.FlintDecoderLayer(
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Pre-norm attention and a dense or shared NeMo MoE, each with a residual.

**`initializer_range`** `= config.initializer_range`

---

**`input_layernorm`**

---

**`mlp`** `MLP | MoE`

---

**`moe_config`** `MoEConfig | None = None`

---

**`post_attention_layernorm`**

---

**`residual_dtype`** `= config.residual_dtype`

---

**`self_attn`** `= FlintAttention(config, layer_idx, backend)`

---

```python
nemo_automodel.components.models.flint.layers.FlintDecoderLayer.forward(
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
```

Apply the block while preserving the shared MoE's padding convention.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Residual tensor \[batch, sequence, hidden_size] or
THD \[tokens, hidden_size]. Explicit residual_dtype preserves its
chosen dtype while normalized branches use projection dtype.

---

**`position_ids`** `torch.Tensor | None` — default: None

Integer \[batch, sequence] or document-local \[tokens].

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary \[batch, sequence], with one for valid
tokens; shared MoE instead receives True for padding.

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional True-for-padding expert mask \[batch, sequence]
or THD \[tokens], derived from real/padded document lengths.

---

**`**attn_kwargs`** `Any` — default: \{}

Shared TE metadata or an immutable model-owned Magi
plan retained across activation-checkpoint recomputation.

---

**Returns:** `torch.Tensor`

Residual tensor in the input layout and configured stream dtype. Padded-query outputs

```python
nemo_automodel.components.models.flint.layers.FlintDecoderLayer.init_weights(
    buffer_device: torch.device,
    init_std: float | None = None
) -> None
```

Initialize every component without depending on a base checkpoint.

**Parameters:**

**`buffer_device`** `torch.device`

Device context after any meta materialization.

---

**`init_std`** `float | None` — default: None

Optional normal-initialization deviation for projections/experts.

---

```python
class nemo_automodel.components.models.flint.layers._FlintSinkParameters(
    num_heads: int
)
```

**Bases:** `Module`

Own the trainable FP32 denominator logits in a separate dtype unit.

**`softmax_offset`**

---

```python
nemo_automodel.components.models.flint.layers._FlintSinkParameters.forward() -> torch.Tensor
```

Return unscaled sink logits through the module's normal call path.

**Returns:** `torch.Tensor`

FP32 tensor of shape \[query_heads], with independent storage and

```python
nemo_automodel.components.models.flint.layers._validate_backend(
    backend: nemo_automodel.components.models.common.BackendConfig
) -> None
```

Reject backend choices outside the supported native component contract.

**Parameters:**

**`backend`** `BackendConfig`

Shared backend settings, preserved without mutation.

---