> 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.kimi_k3.cp

Context-parallel and packed-sequence support for Kimi Linear.

Kimi Linear interleaves KDA linear-attention layers with MLA full-attention
layers, so context parallelism has to satisfy both at once:

* KDA carries a sequential recurrent state, so FLA's context-parallel kernels
  require every rank to own one **contiguous** slice of the global token stream
  (rank `r` owns `[r * S / cp, (r + 1) * S / cp)`) and take document
  boundaries through `cu_seqlens`. PyTorch's default load-balanced
  `context_parallel` layout (head/tail chunk swap) does not satisfy that, so
  Kimi Linear owns its batch sharding through `_cp_make_batch_fn`.
* MLA attends globally. Under the contiguous layout each rank all-gathers the
  *compressed* KV latent (`kv_lora_rank + qk_rope_head_dim` values per token,
  roughly an order of magnitude smaller than the expanded per-head K/V) and runs
  FlexAttention with a causal, per-document block mask against the full-sequence
  keys.

Everything is driven by one `[batch, sequence]` document-id map (`0` marks
padding, `1..n` are 1-based document indices), which is also what makes packed
sequences work with and without CP.

## Module Contents

### Classes

| Name                                                                                    | Description                                                              |
| --------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`KimiPackedContext`](#nemo_automodel-components-models-kimi_k3-cp-KimiPackedContext)   | Per-step document layout shared by the KDA and MLA layers.               |
| [`_AllGatherSequence`](#nemo_automodel-components-models-kimi_k3-cp-_AllGatherSequence) | Autograd-aware all-gather of equal-sized shards along the sequence axis. |

### Functions

| Name                                                                                                            | Description                                                                  |
| --------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`_block_mask_cache_generation`](#nemo_automodel-components-models-kimi_k3-cp-_block_mask_cache_generation)     | Drop cached block masks when a new batch (new document map) arrives.         |
| [`_compiled_flex_attention`](#nemo_automodel-components-models-kimi_k3-cp-_compiled_flex_attention)             | -                                                                            |
| [`_document_causal_block_mask`](#nemo_automodel-components-models-kimi_k3-cp-_document_causal_block_mask)       | Build (and cache for the step) the FlexAttention document-causal block mask. |
| [`_global_doc_ids_from_batch`](#nemo_automodel-components-models-kimi_k3-cp-_global_doc_ids_from_batch)         | Resolve the global document-id map for a batch about to be CP-sharded.       |
| [`_pad_position_ids`](#nemo_automodel-components-models-kimi_k3-cp-_pad_position_ids)                           | -                                                                            |
| [`_pad_sequence_dim`](#nemo_automodel-components-models-kimi_k3-cp-_pad_sequence_dim)                           | -                                                                            |
| [`all_gather_sequence`](#nemo_automodel-components-models-kimi_k3-cp-all_gather_sequence)                       | All-gather a sequence-sharded tensor while keeping autograd connected.       |
| [`build_document_causal_mask`](#nemo_automodel-components-models-kimi_k3-cp-build_document_causal_mask)         | Build the additive causal mask that also blocks cross-document attention.    |
| [`build_fla_cp_context`](#nemo_automodel-components-models-kimi_k3-cp-build_fla_cp_context)                     | Build FLA's per-row context-parallel context for a KDA layer.                |
| [`doc_ids_from_attention_mask`](#nemo_automodel-components-models-kimi_k3-cp-doc_ids_from_attention_mask)       | Build document ids from a binary or indexed attention mask.                  |
| [`doc_ids_from_cu_seqlens`](#nemo_automodel-components-models-kimi_k3-cp-doc_ids_from_cu_seqlens)               | Build single-row document ids from cumulative sequence lengths.              |
| [`doc_ids_from_seq_lens`](#nemo_automodel-components-models-kimi_k3-cp-doc_ids_from_seq_lens)                   | Build document ids from the packed-sequence collater's `seq_lens`.           |
| [`document_causal_flex_attention`](#nemo_automodel-components-models-kimi_k3-cp-document_causal_flex_attention) | Run causal, per-document attention of local queries against global keys.     |
| [`segment_cu_seqlens`](#nemo_automodel-components-models-kimi_k3-cp-segment_cu_seqlens)                         | Return segment boundaries for one row of document ids.                       |
| [`shard_batch_for_kimi_cp`](#nemo_automodel-components-models-kimi_k3-cp-shard_batch_for_kimi_cp)               | Shard a batch contiguously across the context-parallel mesh for Kimi K3.     |

### Data

[`_BLOCK_MASK_CACHE`](#nemo_automodel-components-models-kimi_k3-cp-_BLOCK_MASK_CACHE)

[`_BLOCK_MASK_GENERATION`](#nemo_automodel-components-models-kimi_k3-cp-_BLOCK_MASK_GENERATION)

[`_COMPILED_FLEX_ATTENTION`](#nemo_automodel-components-models-kimi_k3-cp-_COMPILED_FLEX_ATTENTION)

[`_PAD_DOC_ID`](#nemo_automodel-components-models-kimi_k3-cp-_PAD_DOC_ID)

### API

```python
class nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext(
    doc_ids: torch.Tensor,
    seq_start: int = 0,
    cp_size: int = 1
)
```

Dataclass

Per-step document layout shared by the KDA and MLA layers.

**`cp_enabled`** `bool`

Whether the batch was sharded across a context-parallel mesh.

---

**`cp_size`** `int = 1`

---

**`doc_ids`** `Tensor`

---

**`has_multiple_documents`** `bool`

Whether any batch row contains more than one non-padding document.

---

**`local_doc_ids`** `Tensor`

Document ids of shape \[batch, local\_sequence] for this rank's shard.

---

**`seq_start`** `int = 0`

---

```python
nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext.__post_init__() -> None
```

```python
nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext.row_cu_seqlens(
    row: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Return the global segment boundaries of one batch row.

Computed on first use (and cached for the step) because the device-to-host
copy is only needed by the context-parallel path.

**Parameters:**

**`row`** `int`

Batch row to describe.

---

**Returns:** `torch.Tensor`

The device and CPU copies of the row's cumulative segment lengths, each

```python
class nemo_automodel.components.models.kimi_k3.cp._AllGatherSequence()
```

**Bases:** `Function`

Autograd-aware all-gather of equal-sized shards along the sequence axis.

```python
nemo_automodel.components.models.kimi_k3.cp._AllGatherSequence.backward(
    ctx,
    grad_output: torch.Tensor
)
```

staticmethod

```python
nemo_automodel.components.models.kimi_k3.cp._AllGatherSequence.forward(
    ctx,
    local_tensor: torch.Tensor,
    group: typing.Any,
    dim: int
) -> torch.Tensor
```

staticmethod

```python
nemo_automodel.components.models.kimi_k3.cp._block_mask_cache_generation(
    doc_ids: torch.Tensor
) -> None
```

Drop cached block masks when a new batch (new document map) arrives.

```python
nemo_automodel.components.models.kimi_k3.cp._compiled_flex_attention()
```

```python
nemo_automodel.components.models.kimi_k3.cp._document_causal_block_mask(
    q_doc_ids: torch.Tensor,
    kv_doc_ids: torch.Tensor,
    q_global_start: int
)
```

Build (and cache for the step) the FlexAttention document-causal block mask.

```python
nemo_automodel.components.models.kimi_k3.cp._global_doc_ids_from_batch(
    batch: dict,
    seq_len: int,
    device: torch.device
) -> torch.Tensor
```

Resolve the global document-id map for a batch about to be CP-sharded.

```python
nemo_automodel.components.models.kimi_k3.cp._pad_position_ids(
    position_ids: torch.Tensor,
    seq_dim: int,
    pad_len: int
) -> torch.Tensor
```

```python
nemo_automodel.components.models.kimi_k3.cp._pad_sequence_dim(
    tensor: torch.Tensor,
    seq_dim: int,
    pad_len: int,
    value: float | int
) -> torch.Tensor
```

```python
nemo_automodel.components.models.kimi_k3.cp.all_gather_sequence(
    tensor: torch.Tensor,
    cp_group: typing.Any,
    dim: int = 1
) -> torch.Tensor
```

All-gather a sequence-sharded tensor while keeping autograd connected.

**Parameters:**

**`tensor`** `torch.Tensor`

Tensor of shape \[..., local\_sequence, ...] whose sequence axis is
selected by `dim`. Every rank must contribute the same shape.

---

**`cp_group`** `Any`

Context-parallel process group.

---

**`dim`** `int` — default: 1

Sequence axis.

---

**Returns:** `torch.Tensor`

Tensor with the sequence axis expanded to the full global sequence.

```python
nemo_automodel.components.models.kimi_k3.cp.build_document_causal_mask(
    q_doc_ids: torch.Tensor,
    kv_doc_ids: torch.Tensor,
    q_global_start: int,
    dtype: torch.dtype
) -> torch.Tensor
```

Build the additive causal mask that also blocks cross-document attention.

**Parameters:**

**`q_doc_ids`** `torch.Tensor`

Tensor of shape \[batch, query\_sequence] with 1-based document ids.

---

**`kv_doc_ids`** `torch.Tensor`

Tensor of shape \[batch, key\_sequence] with 1-based document ids.

---

**`q_global_start`** `int`

Global sequence offset of the first query token.

---

**`dtype`** `torch.dtype`

Floating-point dtype used for the additive mask values.

---

**Returns:** `torch.Tensor`

Additive mask tensor of shape \[batch, 1, query\_sequence, key\_sequence].

```python
nemo_automodel.components.models.kimi_k3.cp.build_fla_cp_context(
    packed_context: nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext,
    row: int,
    cp_group: typing.Any,
    conv_kernel_size: int
)
```

Build FLA's per-row context-parallel context for a KDA layer.

**Parameters:**

**`packed_context`** `KimiPackedContext`

Context describing the global document layout.

---

**`row`** `int`

Batch row the context is built for.

---

**`cp_group`** `Any`

Context-parallel process group.

---

**`conv_kernel_size`** `int`

Short-convolution kernel size, used by FLA to exchange
the conv boundary tokens between neighbouring ranks.

---

**Returns:**

The FLA `FLACPContext` for this row.

```python
nemo_automodel.components.models.kimi_k3.cp.doc_ids_from_attention_mask(
    attention_mask: torch.Tensor
) -> torch.Tensor
```

Build document ids from a binary or indexed attention mask.

**Parameters:**

**`attention_mask`** `torch.Tensor`

Tensor of shape \[batch, sequence]. A binary mask marks
valid tokens with 1; an Automodel packing mask marks document `i`
(1-based) with the value `i` and padding with 0.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence] with 1-based document ids and 0 for

```python
nemo_automodel.components.models.kimi_k3.cp.doc_ids_from_cu_seqlens(
    cu_seqlens: torch.Tensor,
    seq_len: int
) -> torch.Tensor
```

Build single-row document ids from cumulative sequence lengths.

**Parameters:**

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

Tensor of shape \[segments + 1] with cumulative token counts.
THD batches pad unused entries with a negative sentinel, which is
dropped here.

---

**`seq_len`** `int`

Sequence length of the batch's token tensors.

---

**Returns:** `torch.Tensor`

Tensor of shape \[1, sequence] with 1-based document ids and 0 for the

```python
nemo_automodel.components.models.kimi_k3.cp.doc_ids_from_seq_lens(
    seq_lens: torch.Tensor,
    seq_len: int,
    padding_value: int = -1000
) -> torch.Tensor
```

Build document ids from the packed-sequence collater's `seq_lens`.

**Parameters:**

**`seq_lens`** `torch.Tensor`

Tensor of shape \[batch, packs] with per-pack token counts, using
`padding_value` for unused pack slots.

---

**`seq_len`** `int`

Sequence length of the batch's token tensors.

---

**`padding_value`** `int` — default: -1000

Sentinel marking unused pack slots.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence] with 1-based document ids and 0 for the

```python
nemo_automodel.components.models.kimi_k3.cp.document_causal_flex_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    q_doc_ids: torch.Tensor,
    kv_doc_ids: torch.Tensor,
    q_global_start: int,
    scale: float
) -> torch.Tensor
```

Run causal, per-document attention of local queries against global keys.

**Parameters:**

**`query`** `torch.Tensor`

Tensor of shape \[batch, heads, query\_sequence, qk\_head\_dim].

---

**`key`** `torch.Tensor`

Tensor of shape \[batch, heads, key\_sequence, qk\_head\_dim].

---

**`value`** `torch.Tensor`

Tensor of shape \[batch, heads, key\_sequence, v\_head\_dim].

---

**`q_doc_ids`** `torch.Tensor`

Tensor of shape \[batch, query\_sequence] with 1-based document ids.

---

**`kv_doc_ids`** `torch.Tensor`

Tensor of shape \[batch, key\_sequence] with 1-based document ids.

---

**`q_global_start`** `int`

Global sequence offset of the first query token.

---

**`scale`** `float`

Softmax scale applied to the query-key product.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, heads, query\_sequence, v\_head\_dim].

```python
nemo_automodel.components.models.kimi_k3.cp.segment_cu_seqlens(
    doc_ids_row: torch.Tensor
) -> torch.Tensor
```

Return segment boundaries for one row of document ids.

Consecutive runs of the same id -- including runs of padding -- become their
own segment so that the boundaries always tile the full row, which is what
FLA's context-parallel partitioning expects.

**Parameters:**

**`doc_ids_row`** `torch.Tensor`

Tensor of shape \[sequence] with 1-based document ids.

---

**Returns:** `torch.Tensor`

Tensor of shape \[segments + 1] with cumulative segment lengths.

```python
nemo_automodel.components.models.kimi_k3.cp.shard_batch_for_kimi_cp(
    cp_mesh,
    tp_mesh,
    batch: dict,
    loss_mask = None,
    padding_token_id: int = 0
)
```

Shard a batch contiguously across the context-parallel mesh for Kimi K3.

Exposed through the :class:`ContextParallelSharder` returned by
:meth:`KimiK3ForCausalLM.prepare_model_inputs_for_cp`. Every rank starts
from the same full batch, keeps the `[seq_start, seq_end)` slice of each
sequence-aligned tensor, and gets the (unsharded) global document-id map
needed by the KDA and MLA layers.

**Parameters:**

**`cp_mesh`**

One-dimensional context-parallel mesh, or None.

---

**`tp_mesh`**

Tensor-parallel mesh; unused, accepted for interface parity.

---

**`batch`** `dict`

Batch mapping containing `input_ids` of shape \[batch, sequence]
plus `labels` and optional sequence-aligned tensors.

---

**`loss_mask`** — default: None

Optional tensor of shape \[batch, sequence] sharded with the labels.

---

**`padding_token_id`** `int` — default: 0

Token id used when padding `input_ids`.

---

**Returns:**

`(context_factory, batch, layout)`; the context factory is a null

```python
nemo_automodel.components.models.kimi_k3.cp._BLOCK_MASK_CACHE: dict[tuple, Any] = {}
```

```python
nemo_automodel.components.models.kimi_k3.cp._BLOCK_MASK_GENERATION: list[Any] = [None, None]
```

```python
nemo_automodel.components.models.kimi_k3.cp._COMPILED_FLEX_ATTENTION: list[Any] = [None]
```

```python
nemo_automodel.components.models.kimi_k3.cp._PAD_DOC_ID = 0
```