Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsFlintnemo_automodel.components.models.flint.magi

nemo_automodel.components.models.flint.magi

View as Markdown

Flint learned-sink Magi CP using public mask keys and model-owned token plans.

Module Contents

Classes

NameDescription
FlintMagiBatchPlansStep-owned plans replicated by PP; a chunked index selects each plan.
FlintMagiPlanPer-microbatch keys; all masks retain the same dispatched token order.
_FlintMagiModelFlint model state needed to build model-owned Magi CP inputs.
_MagiAttentionKeyMagi key fields consumed by the Flint integration.

Functions

NameDescription
_build_magi_planBuild masks with one shared token partition for a packed microbatch.
_prepare_magi_inputsPrepare model-owned packed CP, including independent pipeline microbatches.

Data

FlintMagiBatch

FlintMagiBatchValue

logger

API

class nemo_automodel.components.models.flint.magi.FlintMagiBatchPlans(
    plans: tuple[nemo_automodel.components.models.flint.magi.FlintMagiPlan, ...]
)
Dataclass

Step-owned plans replicated by PP; a chunked index selects each plan.

plans
tuple[FlintMagiPlan, ...]
nemo_automodel.components.models.flint.magi.FlintMagiBatchPlans.select(
    microbatch_index: torch.Tensor
) -> nemo_automodel.components.models.flint.magi.FlintMagiPlan

Select without a mutable cursor, before entering checkpointed layers.

Parameters:

microbatch_index
torch.Tensor

Int64 tensor [1], split by PP and possibly moved by FSDP.

Returns: FlintMagiPlan

The immutable microbatch plan retained by that forward’s AC graph.

class nemo_automodel.components.models.flint.magi.FlintMagiPlan(
    keys: dict[int, nemo_automodel.components.models.flint.magi._MagiAttentionKey]
)
Dataclass

Per-microbatch keys; all masks retain the same dispatched token order.

keys
dict[int, _MagiAttentionKey]
nemo_automodel.components.models.flint.magi.FlintMagiPlan.attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    sink: torch.Tensor,
    window: int,
    scale: float
) -> torch.Tensor

Compute local THD output with Magi’s own KV communication.

Parameters:

query
torch.Tensor

CP-local queries [tokens, query_heads, head_dim].

key
torch.Tensor

CP-local keys [tokens, kv_heads, head_dim].

value
torch.Tensor

CP-local values [tokens, kv_heads, head_dim].

sink
torch.Tensor

Protected FP32 learned logits [query_heads], replicated over CP.

window
int

Sliding-window length including the current token, or zero for full causal attention.

scale
float

Attention score scale.

Returns: torch.Tensor

Attention output [tokens, query_heads, head_dim] in the local token order.

class nemo_automodel.components.models.flint.magi._FlintMagiModel()
Protocol

Flint model state needed to build model-owned Magi CP inputs.

config
FlintConfig
cp_mesh
DeviceMesh | None
nemo_automodel.components.models.flint.magi._FlintMagiModel._is_pipeline_parallel_stage() -> bool

Return whether this module is a physical pipeline stage.

class nemo_automodel.components.models.flint.magi._MagiAttentionKey()
Protocol

Magi key fields consumed by the Flint integration.

total_seqlen_q
int
nemo_automodel.components.models.flint.magi._build_magi_plan(
    cu_seqlens: torch.Tensor,
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    group: torch.distributed.ProcessGroup
) -> nemo_automodel.components.models.flint.magi.FlintMagiPlan

Build masks with one shared token partition for a packed microbatch.

Parameters:

cu_seqlens
torch.Tensor

CPU int32 padded document boundaries [documents + 1].

config
FlintConfig

Flint attention configuration, shared across pipeline stages.

group
torch.distributed.ProcessGroup

This pipeline stage’s context-parallel process group.

Returns: FlintMagiPlan

Full-causal and sliding-window keys with identical local token ownership.

Prepare model-owned packed CP, including independent pipeline microbatches.

Parameters:

model
_FlintMagiModel

Full Flint model or a local pipeline stage.

batch
FlintMagiBatch

Global tokens/labels/positions [batch, sequence] and packed document lengths [batch, documents]. Each PP microbatch holds one pack.

num_chunks
int

Number of independent pipeline microbatches.

Returns: dict[str, ContextParallelSharder]

A model-owned CP sharder. In PP, token ids stay global [chunks, sequence];

nemo_automodel.components.models.flint.magi.FlintMagiBatch = dict[str, FlintMagiBatchValue]
 
nemo_automodel.components.models.flint.magi.FlintMagiBatchValue = torch.Tensor | str | int | float | bool | FlintMagiPlan | FlintMagiBatchPlans | ...
 
nemo_automodel.components.models.flint.magi.logger = logging.getLogger(__name__)