nemo_automodel.components.models.kimi_linear.model
nemo_automodel.components.models.kimi_linear.model
Native Automodel support for Moonshot Kimi Linear causal LM checkpoints.
Module Contents
Classes
Functions
Data
API
Bases: Module
Kimi decoder block with KDA/MLA attention and dense or MoE MLP.
Run a Kimi MoE layer.
Parameters:
Tensor of shape [batch, sequence, hidden].
Optional boolean tensor of shape [batch, sequence], where true marks padding tokens.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Run Kimi inference MoE in the same expert-ordered loop as the HF reference.
Transcribed from upstream KimiLinearMoE.moe_infer so eval-time output
matches HF’s expert ordering and accumulation; :meth:_moe above decides when
it applies. Parity with the canonical MoE/GroupedExperts path is pinned
by test_hf_order_eval_moe_matches_standard_grouped_experts_path.
Parameters:
MoE module containing the router and optional shared experts.
Grouped routed experts for the layer.
Tensor of shape [batch, sequence, hidden].
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Run one Kimi decoder layer.
Parameters:
Tensor of shape [batch, sequence, hidden].
KDA layers receive a binary mask [batch, sequence]; MLA layers receive an additive causal mask [batch, 1, sequence, sequence].
Optional boolean tensor of shape [batch, sequence], where true marks padding tokens.
Extra attention kwargs forwarded to KDA/MLA.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Bases: Module
Kimi Delta Attention backed by FLA KDA kernels.
Run KDA over a contiguous context-parallel shard.
FLA’s context-parallel kernels take the global cu_seqlens and derive
each rank’s local segments, passing the recurrent state (and the short
convolution’s boundary tokens) rank to rank. Batch rows are processed one
at a time because FLA’s variable-length path expects a single flattened
sequence per call.
Parameters:
Tensor of shape [batch, local_sequence, hidden].
Document layout of the batch.
Returns: torch.Tensor
Tensor of shape [batch, local_sequence, hidden].
Run the KDA projections, convolutions and delta-rule kernel.
Parameters:
Tensor of shape [batch, sequence, hidden]; the batch must be
one whenever cu_seqlens or cp_context is given.
Optional cumulative document lengths of shape [documents + 1].
Optional FLA context-parallel context, which supersedes
cu_seqlens with its per-rank local segments.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Return the KDA kernel to use for a sequence of seq_len tokens.
Run KDA linear attention.
Parameters:
Tensor of shape [batch, sequence, hidden]; the sequence axis holds this rank’s contiguous shard under context parallelism.
Optional binary padding mask of shape [batch, sequence] where 1 marks valid tokens.
Optional document layout of the batch, required under context parallelism and used to reset the recurrent state at every packed-document boundary.
Optional KDA kwargs, including cu_seqlens for packed sequences.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Attach the context-parallel mesh used to build FLA’s CP context.
Called by the MoE parallelizer’s apply_cp for every attention block.
Parameters:
One-dimensional context-parallel device mesh.
Bases: Module
Owns Kimi KDA fp32 recurrent-decay parameters and computes the gate.
Compute KDA decay gate while holder params are unsharded by FSDP.
Parameters:
Tensor of shape [batch, sequence, heads * head_dim].
Per-head KDA dimension.
Whether to use FLA’s fused KDA gate kernel.
Returns: torch.Tensor
Tensor of shape [batch, sequence, heads, head_dim].
Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin
Kimi Linear causal LM with native trainable MoE layers.
Run Kimi Linear causal LM.
Parameters:
Optional token ids of shape [batch, sequence].
Optional binary padding mask of shape [batch, sequence].
Optional positions of shape [batch, sequence].
Optional embeddings of shape [batch, sequence, hidden].
Optional boolean tensor of shape [batch, sequence], where true marks padding tokens.
Number of trailing sequence logits to compute, or tensor indices.
Whether to include hidden states in the output.
Additional attention kwargs used by packed or THD execution.
Returns: CausalLMOutputWithPast
Causal LM output whose logits have shape [batch, sequence, vocab] unless logits_to_keep trims sequence.
Hand the recipe Kimi Linear’s own context-parallel batch sharding.
KDA’s recurrent state (and FLA’s CP kernels) require every rank to own one
contiguous slice of the sequence, so Kimi Linear replaces the default
load-balanced context_parallel sharding with
:func:~nemo_automodel.components.models.kimi_linear.cp.shard_batch_for_kimi_cp.
The returned sharder is handed back to the CP dispatch under the
"cp_sharder" key.
Parameters:
The full-sequence batch; left intact, the sharder shards it.
Accepted for hook-signature parity; unused, because Kimi
Linear shards the [batch, sequence] layout directly.
Bases: Module
Kimi Linear decoder backbone with trainable Automodel MoE layers.
Select the padding mask passed to KDA layers.
Parameters:
Optional binary padding mask tensor of shape [batch, sequence].
Tensor of shape [sequence] containing current token positions.
Returns: torch.Tensor | None
Binary padding mask tensor of shape [batch, sequence], or None when no KDA mask is needed.
Run the Kimi Linear decoder.
Parameters:
Optional token ids of shape [batch, sequence].
Optional embeddings of shape [batch, sequence, hidden].
Optional binary or indexed packing mask of shape [batch, sequence].
Optional positions of shape [batch, sequence]; accepted for HF compatibility.
Optional boolean tensor of shape [batch, sequence], where true marks padding tokens.
Optional position vector of shape [sequence].
Optional document layout attached by
:func:~nemo_automodel.components.models.kimi_linear.cp.shard_batch_for_kimi_cp;
required under context parallelism and otherwise derived here.
Additional attention kwargs used by packed or THD execution.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Bases: Module
Kimi MLA full-attention layer copied from the HF reference math.
Repeat key/value heads to match the query heads.
Parameters:
Tensor of shape [batch, key_value_heads, sequence, qk_head_dim].
Tensor of shape [batch, key_value_heads, sequence, v_head_dim].
Sequence length of the key/value tensors.
Returns: tuple[torch.Tensor, torch.Tensor]
Key and value tensors expanded to [batch, heads, sequence, head_dim].
Run MLA attention over a contiguous context-parallel shard.
Queries stay local while the compressed KV latent — kv_lora_rank + qk_rope_head_dim values per token, far smaller than the expanded per-head
keys and values — is all-gathered across the context-parallel group and
expanded locally. Attention then runs as FlexAttention with a causal,
per-document block mask over the full sequence.
Parameters:
Tensor of shape [batch, local_sequence, hidden].
Document layout of the batch.
Returns: torch.Tensor
Tensor of shape [batch, local_sequence, hidden].
Run MLA full attention.
Parameters:
Tensor of shape [batch, sequence, hidden]; the sequence axis holds this rank’s contiguous shard under context parallelism.
Optional additive attention mask of shape [batch, 1, sequence, sequence].
Optional document layout of the batch, required under context parallelism.
Extra attention options accepted for HF compatibility.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Attach the context-parallel mesh used to gather full-sequence keys and values.
Called by the MoE parallelizer’s apply_cp for every attention block.
Parameters:
One-dimensional context-parallel device mesh.
Bases: Module
Kimi RMSNorm with fp32 variance computation.
Normalize hidden states.
Parameters:
Tensor of shape [batch, sequence, hidden].
Returns: torch.Tensor
Tensor of shape [batch, sequence, hidden].
Descriptor exposing a KDA fp32 parameter from the _fp32_params holder.
Call FLA fused KDA gate across FLA versions.
Parameters:
Tensor of shape [batch, sequence, heads * head_dim].
Tensor of shape [1, 1, heads, 1].
Per-head KDA dimension.
Tensor of shape [heads * head_dim].
Returns: torch.Tensor
Tensor of shape [batch, sequence, heads, head_dim].
Build metadata for converting padded batches to flattened valid tokens.
Parameters:
Binary mask tensor of shape [batch, sequence] where 1 marks valid tokens.
Returns: torch.Tensor
Tuple containing indices of shape [total_valid_tokens], cu_seqlens of shape [batch + 1],
Gather rows from the first axis while preserving trailing tensor layout.
Parameters:
Tensor of shape [tokens, …], with arbitrary trailing axes.
Tensor of shape [selected_tokens] containing first-axis row indices.
Returns: torch.Tensor
Tensor of shape [selected_tokens, …], with the same trailing axes as x.
Scatter rows into the first axis while preserving trailing tensor layout.
Parameters:
Tensor of shape [selected_tokens, …], with arbitrary trailing axes.
Tensor of shape [selected_tokens] containing destination row indices.
Size of the output first axis.
Returns: torch.Tensor
Tensor of shape [first_axis_dim, …], with the same trailing axes as x.
Create the additive causal attention mask for full-attention layers.
Parameters:
Tensor of shape [batch, sequence, hidden].
Optional document layout of the batch. When it marks more than one document per row, the mask is block-diagonal so tokens never attend across packed documents.
Floating-point dtype used for the additive mask values.
Returns: torch.Tensor | None
Additive causal mask tensor of shape [batch, 1, sequence, sequence].
Derive the document layout of a batch that was not sharded for context parallelism.
Parameters:
Tensor of shape [batch, sequence, hidden].
Optional binary or indexed packing mask of shape [batch, sequence].
Optional cumulative document lengths of shape [documents + 1] from the THD packed path.
Returns: KimiPackedContext | None
The document layout, or None when the batch is a single unpadded document per
Restore flattened valid tokens to padded batch layout.
Parameters:
Tensor of shape [total_valid_tokens, …], with arbitrary trailing axes.
Tensor of shape [total_valid_tokens] containing flattened padded-batch row indices.
Number of sequences in the padded output batch.
Sequence length in the padded output batch.
Returns: torch.Tensor
Tensor of shape [batch, sequence, …], with the same trailing axes as hidden_states.
Torch equivalent of FLA’s KDA gate.
Parameters:
Tensor of shape [batch, sequence, heads * head_dim] or [batch, sequence, heads, head_dim].
Tensor of shape [1, 1, heads, 1].
Per-head KDA dimension.
Tensor of shape [heads * head_dim].
Returns: torch.Tensor
Tensor of shape [batch, sequence, heads, head_dim].