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
`_block_mask_cache_generation`	Drop cached block masks when a new batch (new document map) arrives.
`_compiled_flex_attention`	-
`_document_causal_block_mask`	Build (and cache for the step) the FlexAttention document-causal block mask.
`_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`	-
`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_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.

Data

_BLOCK_MASK_CACHE

_BLOCK_MASK_GENERATION

_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

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

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

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._block_mask_cache_generation(
    doc_ids: torch.Tensor
) -> None

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

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
)

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

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.

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

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

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

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._BLOCK_MASK_CACHE: dict[tuple, Any] = {}

nemo_automodel.components.models.kimi_k3.cp._BLOCK_MASK_GENERATION: list[Any] = [None, 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