nemo_automodel.components.models.kimi_k3.model
nemo_automodel.components.models.kimi_k3.model
Native AutoModel implementation of the Moonshot Kimi K3 architecture.
Module Contents
Classes
Functions
Data
_FUSED_KDA_GATE_HAS_LOWER_BOUND
API
Bases: Module
Kimi decoder block with KDA/MLA attention and dense or MoE MLP.
Run one decoder layer using K3’s learned block-residual mixing.
Parameters:
Tensor of shape [batch, sequence, hidden].
Tensor of shape [batch * sequence, blocks, hidden].
KDA padding mask or MLA additive causal mask.
Optional boolean tensor of shape [batch, sequence], where true marks padding tokens.
Extra attention arguments.
Returns: tuple[torch.Tensor, torch.Tensor]
Updated hidden states and block residuals.
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.
Prior block starts with shape [batch * sequence, blocks, hidden].
Extra attention kwargs forwarded to KDA/MLA.
Returns: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
Tensor of shape [batch, sequence, hidden], plus updated block residuals when enabled.
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 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: HFCheckpointingMixin, Module, MoEFSDPSyncMixin
Kimi Linear causal LM with native trainable MoE layers.
Keep every K3 attention-residual block within one pipeline stage.
Run Kimi Linear causal LM.
Parameters:
Optional token ids of shape [batch, sequence].
Prior K3 attention-residual block starts. Pipeline stages after the first receive this as their second activation.
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 | torch.Tensor | tuple[torch.Tensor, torch.Tensor]
Causal LM output whose logits have shape [batch, sequence, vocab] unless logits_to_keep trims sequence.
Return static PP metadata for hidden states and block residuals.
Hand the recipe Kimi K3’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 K3 replaces the default
load-balanced context-parallel sharding with
:func:~nemo_automodel.components.models.kimi_k3.cp.shard_batch_for_kimi_cp.
Parameters:
Full-sequence batch; left untouched until the returned sharder runs.
Accepted for CP hook signature parity; K3 uses one contiguous shard.
Returns: dict[str, Any]
Batch updates carrying the model-owned context-parallel sharder.
Bases: Gate
K3’s fp32 sigmoid router with correction-bias-only expert selection.
Route local token states and return fp32 top-k weights.
Parameters:
Tensor of shape [tokens, hidden] containing this rank’s local token states.
Boolean tensor of shape [tokens]. Kimi K3 currently routes every supplied token, so this mask is unused.
Optional context-parallel mesh. Kimi K3 currently operates on already-local token states, so this mesh is unused.
Returns: torch.Tensor
Tuple containing fp32 routing weights of shape
Bases: Module
Dense or shared K3 SiTU MLP.
Transform hidden_states of shape [..., hidden].
Bases: MoE
K3 routed experts with latent projections and a SiTU shared expert.
Match the checkpoint implementation’s expert-ordered inference loop.
Whether grouped expert parameters are DTensors.
Run K3 MoE on [batch, sequence, hidden] states.
Bases: Module
Kimi Linear decoder backbone with trainable Automodel MoE layers.
Mix final [batch, sequence, hidden] states with block starts.
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].
Prior attention-residual block starts with shape [batch * sequence, blocks, hidden]. Pipeline stages after the first receive this as their second positional activation.
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_k3.cp.shard_batch_for_kimi_cp;
required under context parallelism and otherwise derived here.
Pipeline-safe global document map used to
reconstruct kimi_packed_context after microbatch chunking.
Global offset of this CP rank’s sequence shard.
Number of context-parallel sequence shards.
Additional attention kwargs used by packed or THD execution.
Returns: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
Tensor of shape [batch, sequence, hidden], or the hidden states and
Bases: Module
Own KDA recurrent-decay parameters and compute the FP32 decay gate.
Compute the KDA decay while this holder’s FSDP unit is unsharded.
Parameters:
Raw gate tensor of shape [batch, sequence, heads * head_dim].
Per-head KDA dimension.
Optional lower bound for K3’s bounded decay function.
Whether to use FLA’s fused gate kernel.
Returns: torch.Tensor
FP32 decay tensor of shape [batch, sequence, heads, head_dim].
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 boolean mask of shape [batch, sequence], where true marks padding.
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].
Bases: Module
K3 SiTU gated activation with fp32 nonlinearities.
Apply SiTU to [... , 2 * intermediate] gate/up projections.
Bases: Module
Keep a callable FLA operator in its own fp32 FSDP unit.
Expose the wrapped weight under the reference module API.
Run the operator while its fp32 FSDP unit is unsharded.
Reset the wrapped operator.
Descriptor exposing a KDA fp32 parameter from the _fp32_params holder.
Mix [tokens, hidden] with prior [tokens, blocks, hidden] residuals.
Call FLA’s fused KDA gate across supported FLA APIs.
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.
Partition decoder layers without splitting an attention-residual block.
Initialize DTensor RNG without a world broadcast during PP weight init.
Compute K3’s KDA decay gate with torch FP32 operations.
Parameters:
Raw gate tensor of shape [batch, sequence, heads * head_dim] or [batch, sequence, heads, head_dim].
Log decay tensor of shape [heads].
Per-head KDA dimension.
Gate bias tensor of shape [heads * head_dim].
Optional lower bound for K3’s bounded decay function.
Returns: torch.Tensor
FP32 decay tensor of shape [batch, sequence, heads, head_dim].
Apply SiTU and routing weights to [tokens, 2 * intermediate] projections.