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`	Separate Q/K/V, per-head RMSNorm, selective Neox RoPE and learned sinks.
`FlintDecoderLayer`	Pre-norm attention and a dense or shared NeMo MoE, each with a residual.
`_FlintSinkParameters`	Own the trainable FP32 denominator logits in a separate dtype unit.

Functions

Name	Description
`_validate_backend`	Reject backend choices outside the supported native component contract.

API

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

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’

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 | NoneDefaults to 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 | NoneDefaults to 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

AnyDefaults to {}

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

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 | NoneDefaults to None

Optional normal-initialization deviation for projections.

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)

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 | NoneDefaults to None

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

attention_mask

torch.Tensor | NoneDefaults to None

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

padding_mask

torch.Tensor | NoneDefaults to None

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

**attn_kwargs

AnyDefaults to {}

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

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 | NoneDefaults to None

Optional normal-initialization deviation for projections/experts.

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

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

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.