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; _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].
Run KDA linear attention.
Parameters:
Tensor of shape [batch, sequence, hidden]; the sequence axis holds this rank’s contiguous shard under context parallelism.
Optional tensor of shape [batch, sequence]. Without
packed metadata this is a binary validity mask (1 means valid).
With packed metadata it may hold 1-based document IDs and zero
for padding. Non-2D masks are not consumed directly; the legacy
fallback reads padding_mask instead.
Optional document layout of the batch, required under context parallelism and used to reset the recurrent state at every packed-document boundary.
Optional KDA metadata. packed_token_indices contains
row-local valid-token positions of shape [batch, sequence] with
-1 padding, or flat positions of shape [tokens]. cu_seqlens
contains row-local boundaries of shape [batch, max_documents + 1]
with -1 unused entries, or flat boundaries of shape [documents + 1].
With token indices, boundaries address the unpadded valid-token
stream; alone, flat boundaries address one physical row only when
neither a document mask nor _packed_seq_ids is supplied.
_packed_seq_ids optionally carries document IDs of shape
[batch, sequence]; padding_mask is boolean with the same
shape, where True marks padding.
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.
Packed or THD metadata: _packed_seq_ids has shape
[batch, sequence] with 1-based document IDs and zero padding.
packed_token_indices has shape [batch, sequence] with row-local
positions and -1 padding, or [tokens] indexing flattened batch
and sequence axes. cu_seqlens has shape [batch, max_documents + 1]
with row-local boundaries and -1 padding, or [documents + 1] for
flattened/THD inputs. max_seqlen is an integer document length.
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
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
shard_batch_for_kimi_cp;
required under context parallelism and otherwise derived here.
Packed or THD metadata: _packed_seq_ids has shape
[batch, sequence] with 1-based document IDs and zero padding.
packed_token_indices has shape [batch, sequence] with row-local
positions and -1 padding, or [tokens] indexing flattened batch
and sequence axes. cu_seqlens has shape [batch, max_documents + 1]
with row-local boundaries and -1 padding, or [documents + 1] for
flattened/THD inputs. max_seqlen is an integer document length.
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].