> This page is for version Latest.
> 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.magi

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

## Module Contents

### Classes

| Name                                                                                      | Description                                                            |
| ----------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`FlintMagiBatchPlans`](#nemo_automodel-components-models-flint-magi-FlintMagiBatchPlans) | Step-owned plans replicated by PP; a chunked index selects each plan.  |
| [`FlintMagiPlan`](#nemo_automodel-components-models-flint-magi-FlintMagiPlan)             | Per-microbatch keys; all masks retain the same dispatched token order. |
| [`_FlintMagiModel`](#nemo_automodel-components-models-flint-magi-_FlintMagiModel)         | Flint model state needed to build model-owned Magi CP inputs.          |
| [`_MagiAttentionKey`](#nemo_automodel-components-models-flint-magi-_MagiAttentionKey)     | Magi key fields consumed by the Flint integration.                     |

### Functions

| Name                                                                                        | Description                                                                 |
| ------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_build_magi_plan`](#nemo_automodel-components-models-flint-magi-_build_magi_plan)         | Build masks with one shared token partition for a packed microbatch.        |
| [`_prepare_magi_inputs`](#nemo_automodel-components-models-flint-magi-_prepare_magi_inputs) | Prepare model-owned packed CP, including independent pipeline microbatches. |

### Data

[`FlintMagiBatch`](#nemo_automodel-components-models-flint-magi-FlintMagiBatch)

[`FlintMagiBatchValue`](#nemo_automodel-components-models-flint-magi-FlintMagiBatchValue)

[`logger`](#nemo_automodel-components-models-flint-magi-logger)

### API

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

---

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

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

---

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

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

---

```python
nemo_automodel.components.models.flint.magi._FlintMagiModel._is_pipeline_parallel_stage() -> bool
```

Return whether this module is a physical pipeline stage.

```python
class nemo_automodel.components.models.flint.magi._MagiAttentionKey()
```

Protocol

Magi key fields consumed by the Flint integration.

**`total_seqlen_q`** `int`

---

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

```python
nemo_automodel.components.models.flint.magi._prepare_magi_inputs(
    model: nemo_automodel.components.models.flint.magi._FlintMagiModel,
    batch: nemo_automodel.components.models.flint.magi.FlintMagiBatch,
    num_chunks: int
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]
```

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];

```python
nemo_automodel.components.models.flint.magi.FlintMagiBatch = dict[str, FlintMagiBatchValue]
```

```python
nemo_automodel.components.models.flint.magi.FlintMagiBatchValue = torch.Tensor | str | int | float | bool | FlintMagiPlan | FlintMagiBatchPlans | ...
```

```python
nemo_automodel.components.models.flint.magi.logger = logging.getLogger(__name__)
```