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`	Per-step document layout shared by the KDA and MLA layers.
`_AllGatherSequence`	Autograd-aware all-gather of equal-sized shards along the sequence axis.

Functions

Name	Description
`_compiled_flex_attention`	-
`_document_causal_block_mask`	Build the FlexAttention mask owned by one microbatch’s packed context.
`_global_doc_ids_from_batch`	Resolve the global document-id map for a batch about to be CP-sharded.
`_pad_position_ids`	-
`_pad_sequence_dim`	-
`_valid_first_token_order`	Return a stable per-row order that moves masked tokens behind valid tokens.
`all_gather_sequence`	All-gather a sequence-sharded tensor while keeping autograd connected.
`build_document_causal_mask`	Build the additive causal mask that also blocks cross-document attention.
`build_fla_cp_context`	Build FLA’s per-row context-parallel context for a KDA layer.
`doc_ids_from_attention_mask`	Build document ids from a binary or indexed attention mask.
`doc_ids_from_cu_seqlens`	Build single-row document ids from cumulative sequence lengths.
`doc_ids_from_seq_lens`	Build document ids from the packed-sequence collater’s `seq_lens`.
`document_causal_allowed`	Mark the keys each query may attend to: causal and within its own document.
`document_causal_flex_attention`	Run causal, per-document attention of local queries against global keys.
`segment_cu_seqlens`	Return segment boundaries for one row of document ids.
`shard_batch_for_kimi_cp`	Shard a batch contiguously across the context-parallel mesh for Kimi K3.
`shard_batch_for_kimi_vlm_cp`	Keep complete media token streams until embedding inside the VLM forward.
`sparse_keep_block_mask`	Build a FlexAttention mask from a dense DSA keep-mask.

Data

KimiVlmBatch

_COMPILED_FLEX_ATTENTION

_PAD_DOC_ID

API

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

Number of context-parallel shards the global sequence was split into.

doc_ids

Tensor

Global document ids of shape [batch, sequence] covering the full (unsharded) sequence, where 0 marks padding. In indexed packed masks, each positive document ID must occupy one contiguous run per row once masked (0) gaps are ignored: [1, 1, 0, 1, 1, 2, 2] is one document with a masked hole, [1, 2, 1] is rejected.

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.

mla_block_mask

BlockMask

Return the document-causal mask retained through this microbatch’s backward.

seq_start

int = 0

Global sequence offset of the local shard; 0 without CP.

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

nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext.cudnn_dsa_metadata(
    row: int
) -> nemo_automodel.components.models.glm_moe_dsa.kernels.cudnn_dsa.CudnnDsaPackedMetadata

Cache cuDNN metadata for one row’s local queries and global keys.

Parameters:

row

int

Batch row in the document map of shape [batch, global_sequence].

Returns: CudnnDsaPackedMetadata

Segmented cuDNN metadata with row-local global key coordinates.

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

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

Bases: Function

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

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

staticmethod

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

staticmethod

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

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
) -> torch.nn.attention.flex_attention.BlockMask

Build the FlexAttention mask owned by one microbatch’s packed context.

Parameters:

q_doc_ids

torch.Tensor

Local query document IDs of shape [batch, local_sequence].

kv_doc_ids

torch.Tensor

Global key document IDs of shape [batch, global_sequence].

q_global_start

int

Offset of the contiguous local query shard.

Returns: BlockMask

BlockMask for local queries against global keys, causal within documents.

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.

Parameters:

batch

dict

Model inputs. _packed_seq_ids or a rank-two attention_mask carries document IDs of shape [batch, sequence]; seq_lens and seq_lens_padded have shape [batch, documents]; padding_mask and input_ids have shape [batch, sequence].

seq_len

int

Global sequence length before context-parallel sharding.

device

torch.device

Device of the global token tensors.

Returns: torch.Tensor

Int32 document IDs of shape [batch, sequence], with zero marking padding.

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

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

nemo_automodel.components.models.kimi_k3.cp._valid_first_token_order(
    doc_ids: torch.Tensor
) -> torch.Tensor | None

Return a stable per-row order that moves masked tokens behind valid tokens.

Without CP, KDA compacts masked tokens away, so a masked span inside a document, such as an injected fake image, does not interrupt that document’s convolution or recurrent state. FLA’s CP kernels instead start a new segment at every document-id change. Moving masked tokens to the row’s trailing padding gives every CP rank the non-CP valid-token sequence; K3 attention applies no rotary positions, so only the relative order of valid tokens matters.

Parameters:

doc_ids

torch.Tensor

Document IDs of shape [batch, sequence], with zero marking masked tokens.

Returns: torch.Tensor | None

Source indices of shape [batch, sequence], or None when no masked

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

intDefaults to 1

Sequence axis.

Returns: torch.Tensor

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

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

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: torch.distributed.ProcessGroup,
    conv_kernel_size: int
) -> fla.ops.cp.FLACPContext

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

dist.ProcessGroup

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: FLACPContext

The FLA FLACPContext for this row.

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

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

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

intDefaults to -1000

Sentinel marking unused pack slots.

Returns: torch.Tensor

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

nemo_automodel.components.models.kimi_k3.cp.document_causal_allowed(
    q_doc_ids: torch.Tensor,
    kv_doc_ids: torch.Tensor,
    q_global_start: int
) -> torch.Tensor

Mark the keys each query may attend to: causal and within its own document.

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.

Returns: torch.Tensor

Boolean tensor of shape [batch, query_sequence, key_sequence], true where attention is allowed.

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,
    block_mask: torch.nn.attention.flex_attention.BlockMask,
    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.

block_mask

BlockMask

Document-causal mask prepared outside checkpointed decoder layers.

scale

float

Softmax scale applied to the query-key product.

Returns: torch.Tensor

Tensor of shape [batch, heads, query_sequence, v_head_dim].

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.

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 ContextParallelSharder returned by 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

Defaults to None

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

padding_token_id

intDefaults to 0

Token id used when padding input_ids.

Returns:

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

nemo_automodel.components.models.kimi_k3.cp.shard_batch_for_kimi_vlm_cp(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    batch: nemo_automodel.components.models.kimi_k3.cp.KimiVlmBatch,
    loss_mask: torch.Tensor | None = None,
    padding_token_id: int = 0
) -> tuple[collections.abc.Callable[[], contextlib.AbstractContextManager[None]], nemo_automodel.components.models.kimi_k3.cp.KimiVlmBatch, nemo_automodel.components.distributed.context_parallel.sharder.ShardLayout]

Keep complete media token streams until embedding inside the VLM forward.

Labels, positions, masks, and global document metadata use the text model’s contiguous CP layout. The full primary stream is batch-chunkable for PP.

Parameters:

cp_mesh

DeviceMesh | None

Context-parallel mesh, or None for a single context rank.

tp_mesh

DeviceMesh | None

Tensor-parallel mesh passed to the native text sharder.

batch

KimiVlmBatch

Mutable VLM batch. input_ids, labels, position_ids and document masks have shape [batch, global_sequence]; pixel_values has shape [patches, channels, patch_height, patch_width] and image grids have shape [images, 3]. Media tensors remain replicated over CP.

loss_mask

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, global_sequence].

padding_token_id

intDefaults to 0

Token ID used for sequence alignment.

Returns: Callable[[], contextlib.AbstractContextManager[None]]

Context factory, mutated batch, and original/padded sequence lengths.

nemo_automodel.components.models.kimi_k3.cp.sparse_keep_block_mask(
    keep: torch.Tensor
) -> torch.nn.attention.flex_attention.BlockMask

Build a FlexAttention mask from a dense DSA keep-mask.

Parameters:

keep

torch.Tensor

Boolean tensor of shape [batch, query_sequence, key_sequence], true where a query attends. It already folds in the causal and document constraints.

Returns: BlockMask

BlockMask for local queries against global keys that reads keep directly.

nemo_automodel.components.models.kimi_k3.cp.KimiVlmBatch = dict[str, torch.Tensor | bool | int | str | None]

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

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