> 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.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.muse_spark.magi_cp

Model-owned packed Magi CP with shared token ownership across Muse masks.

## Module Contents

### Classes

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

### Functions

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

### Data

[`MuseMagiBatch`](#nemo_automodel-components-models-muse_spark-magi_cp-MuseMagiBatch)

[`MuseMagiBatchValue`](#nemo_automodel-components-models-muse_spark-magi_cp-MuseMagiBatchValue)

[`logger`](#nemo_automodel-components-models-muse_spark-magi_cp-logger)

### API

```python
class nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan(
    keys: dict[int, nemo_automodel.components.models.muse_spark.magi_cp._MagiAttentionKey]
)
```

Dataclass

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

**`keys`** `dict[int, _MagiAttentionKey]`

---

```python
nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan.attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: 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].

---

**`window`** `int`

Sliding-window length, 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.muse_spark.magi_cp._MagiAttentionKey()
```

Protocol

Magi key fields consumed by the Muse integration.

**`total_seqlen_q`** `int`

---

```python
class nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiBatchPlans(
    plans: tuple[nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan, ...]
)
```

Dataclass

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

**`plans`** `tuple[MuseMagiPlan, ...]`

---

```python
nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiBatchPlans.select(
    microbatch_index: torch.Tensor
) -> nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan
```

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:** `MuseMagiPlan`

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

```python
class nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiModel()
```

Protocol

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

**`config`** `MuseSpark1xOmniConfig`

---

**`cp_mesh`** `DeviceMesh | None`

---

```python
nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiModel._is_pipeline_parallel_stage() -> bool
```

Return whether this module is a physical pipeline stage.

```python
nemo_automodel.components.models.muse_spark.magi_cp._build_magi_plan(
    cu_seqlens: torch.Tensor,
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    group: torch.distributed.ProcessGroup
) -> nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan
```

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

**Parameters:**

**`cu_seqlens`** `torch.Tensor`

CPU int32 padded document boundaries \[documents + 1].

---

**`config`** `MuseSpark1xOmniConfig`

Muse attention configuration, shared across pipeline stages.

---

**`group`** `torch.distributed.ProcessGroup`

This pipeline stage's context-parallel process group.

---

**Returns:** `MuseMagiPlan`

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

```python
nemo_automodel.components.models.muse_spark.magi_cp.prepare_magi_inputs(
    model: nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiModel,
    batch: nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiBatch,
    num_chunks: int
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]
```

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

**Parameters:**

**`model`** `_MuseMagiModel`

Full Muse model or a local pipeline stage.

---

**`batch`** `MuseMagiBatch`

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.muse_spark.magi_cp.MuseMagiBatch = dict[str, MuseMagiBatchValue]
```

```python
nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiBatchValue = torch.Tensor | str | int | float | bool | MuseMagiPlan | _MuseMagiBatchPlans | N...
```

```python
nemo_automodel.components.models.muse_spark.magi_cp.logger = logging.getLogger(__name__)
```