> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.kimi_k3.model

Native AutoModel implementation of the Moonshot Kimi K3 architecture.

## Module Contents

### Classes

| Name                                                                                       | Description                                                              |
| ------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------ |
| [`KimiDecoderLayer`](#nemo_automodel-components-models-kimi_k3-model-KimiDecoderLayer)     | Kimi decoder block with KDA/MLA attention and dense or MoE MLP.          |
| [`KimiDeltaAttention`](#nemo_automodel-components-models-kimi_k3-model-KimiDeltaAttention) | Kimi Delta Attention backed by FLA KDA kernels.                          |
| [`KimiK3ForCausalLM`](#nemo_automodel-components-models-kimi_k3-model-KimiK3ForCausalLM)   | Kimi Linear causal LM with native trainable MoE layers.                  |
| [`KimiK3Gate`](#nemo_automodel-components-models-kimi_k3-model-KimiK3Gate)                 | K3's fp32 sigmoid router with correction-bias-only expert selection.     |
| [`KimiK3MLP`](#nemo_automodel-components-models-kimi_k3-model-KimiK3MLP)                   | Dense or shared K3 SiTU MLP.                                             |
| [`KimiK3MoE`](#nemo_automodel-components-models-kimi_k3-model-KimiK3MoE)                   | K3 routed experts with latent projections and a SiTU shared expert.      |
| [`KimiK3TextModel`](#nemo_automodel-components-models-kimi_k3-model-KimiK3TextModel)       | Kimi Linear decoder backbone with trainable Automodel MoE layers.        |
| [`KimiKDAFp32Params`](#nemo_automodel-components-models-kimi_k3-model-KimiKDAFp32Params)   | Own KDA recurrent-decay parameters and compute the FP32 decay gate.      |
| [`KimiMLAAttention`](#nemo_automodel-components-models-kimi_k3-model-KimiMLAAttention)     | Kimi MLA full-attention layer copied from the HF reference math.         |
| [`KimiRMSNorm`](#nemo_automodel-components-models-kimi_k3-model-KimiRMSNorm)               | Kimi RMSNorm with fp32 variance computation.                             |
| [`SituAndMul`](#nemo_automodel-components-models-kimi_k3-model-SituAndMul)                 | K3 SiTU gated activation with fp32 nonlinearities.                       |
| [`_KimiFp32Module`](#nemo_automodel-components-models-kimi_k3-model-_KimiFp32Module)       | Keep a callable FLA operator in its own fp32 FSDP unit.                  |
| [`_KimiKDAFp32Param`](#nemo_automodel-components-models-kimi_k3-model-_KimiKDAFp32Param)   | Descriptor exposing a KDA fp32 parameter from the `_fp32_params` holder. |

### Functions

| Name                                                                                                                           | Description                                                                         |
| ------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------- |
| [`_apply_attn_res`](#nemo_automodel-components-models-kimi_k3-model-_apply_attn_res)                                           | Mix `[tokens, hidden]` with prior `[tokens, blocks, hidden]` residuals.             |
| [`_build_moe_config`](#nemo_automodel-components-models-kimi_k3-model-_build_moe_config)                                       | -                                                                                   |
| [`_fused_kda_gate`](#nemo_automodel-components-models-kimi_k3-model-_fused_kda_gate)                                           | Call FLA's fused KDA gate across supported FLA APIs.                                |
| [`_get_unpad_data`](#nemo_automodel-components-models-kimi_k3-model-_get_unpad_data)                                           | Build metadata for converting padded batches to flattened valid tokens.             |
| [`_index_first_axis`](#nemo_automodel-components-models-kimi_k3-model-_index_first_axis)                                       | Gather rows from the first axis while preserving trailing tensor layout.            |
| [`_index_put_first_axis`](#nemo_automodel-components-models-kimi_k3-model-_index_put_first_axis)                               | Scatter rows into the first axis while preserving trailing tensor layout.           |
| [`_make_causal_mask`](#nemo_automodel-components-models-kimi_k3-model-_make_causal_mask)                                       | Create the additive causal attention mask for full-attention layers.                |
| [`_packed_context_from_inputs`](#nemo_automodel-components-models-kimi_k3-model-_packed_context_from_inputs)                   | Derive the document layout of a batch that was not sharded for context parallelism. |
| [`_pad_input`](#nemo_automodel-components-models-kimi_k3-model-_pad_input)                                                     | Restore flattened valid tokens to padded batch layout.                              |
| [`_partition_attn_residual_blocks`](#nemo_automodel-components-models-kimi_k3-model-_partition_attn_residual_blocks)           | Partition decoder layers without splitting an attention-residual block.             |
| [`_require_fla`](#nemo_automodel-components-models-kimi_k3-model-_require_fla)                                                 | -                                                                                   |
| [`_seed_dtensor_rng_for_pipeline_stage`](#nemo_automodel-components-models-kimi_k3-model-_seed_dtensor_rng_for_pipeline_stage) | Initialize DTensor RNG without a world broadcast during PP weight init.             |
| [`_torch_kda_gate`](#nemo_automodel-components-models-kimi_k3-model-_torch_kda_gate)                                           | Compute K3's KDA decay gate with torch FP32 operations.                             |
| [`_weighted_situ`](#nemo_automodel-components-models-kimi_k3-model-_weighted_situ)                                             | Apply SiTU and routing weights to `[tokens, 2 * intermediate]` projections.         |

### Data

[`ModelClass`](#nemo_automodel-components-models-kimi_k3-model-ModelClass)

[`_FLA_MSG`](#nemo_automodel-components-models-kimi_k3-model-_FLA_MSG)

[`_FUSED_KDA_GATE_HAS_G_BIAS`](#nemo_automodel-components-models-kimi_k3-model-_FUSED_KDA_GATE_HAS_G_BIAS)

[`_FUSED_KDA_GATE_HAS_LOWER_BOUND`](#nemo_automodel-components-models-kimi_k3-model-_FUSED_KDA_GATE_HAS_LOWER_BOUND)

### API

```python
class nemo_automodel.components.models.kimi_k3.model.KimiDecoderLayer(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Kimi decoder block with KDA/MLA attention and dense or MoE MLP.

**`attn_res_block_size`** `= config.attn_res_block_size`

---

**`hidden_size`** `= config.hidden_size`

---

**`input_layernorm`**

---

**`is_linear_attn`** `= config.is_kda_layer(layer_idx)`

---

**`is_moe_layer`**

---

**`mlp`** `= KimiK3MoE(config, moe_config, backend)`

---

**`mlp_res_norm`**

---

**`mlp_res_proj`**

---

**`post_attention_layernorm`**

---

**`self_attention_res_norm`**

---

**`self_attention_res_proj`**

---

**`self_attn`**

---

**`use_attn_residuals`** `= config.attn_res_block_size is not None`

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiDecoderLayer._forward_attn_residual(
    hidden_states: torch.Tensor,
    block_residual: torch.Tensor,
    attention_mask: torch.Tensor | None,
    padding_mask: torch.Tensor | None,
    attn_kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor]
```

Run one decoder layer using K3's learned block-residual mixing.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**`block_residual`** `torch.Tensor`

Tensor of shape \[batch \* sequence, blocks, hidden].

---

**`attention_mask`** `torch.Tensor | None`

KDA padding mask or MLA additive causal mask.

---

**`padding_mask`** `torch.Tensor | None`

Optional boolean tensor of shape \[batch, sequence], where true marks padding tokens.

---

**`**attn_kwargs`** `Any` — default: \{}

Extra attention arguments.

---

**Returns:** `tuple[torch.Tensor, torch.Tensor]`

Updated hidden states and block residuals.

```python
nemo_automodel.components.models.kimi_k3.model.KimiDecoderLayer.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    block_residual: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]
```

Run one Kimi decoder layer.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**`attention_mask`** `torch.Tensor | None` — default: None

KDA layers receive a binary mask \[batch, sequence]; MLA layers receive an additive
causal mask \[batch, 1, sequence, sequence].

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional boolean tensor of shape \[batch, sequence], where true marks padding tokens.

---

**`block_residual`** `torch.Tensor | None` — default: None

Prior block starts with shape \[batch \* sequence, blocks, hidden].

---

**`**attn_kwargs`** `Any` — default: \{}

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.

```python
nemo_automodel.components.models.kimi_k3.model.KimiDecoderLayer.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None
```

```python
class nemo_automodel.components.models.kimi_k3.model.KimiDeltaAttention(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    layer_idx: int
)
```

**Bases:** `Module`

Kimi Delta Attention backed by FLA KDA kernels.

**`A_log`** `= _KimiKDAFp32Param('A_log')`

---

**`_fp32_params`** `= KimiKDAFp32Params(self.num_heads, projection_size)`

---

**`b_proj`**

---

**`conv_size`**

---

**`dt_bias`** `= _KimiKDAFp32Param('dt_bias')`

---

**`f_a_proj`**

---

**`f_b_proj`**

---

**`g_a_proj`**

---

**`g_b_proj`**

---

**`g_proj`**

---

**`gate_lower_bound`** `= config.linear_attn_config.get('gate_lower_bound')`

---

**`head_dim`** `= config.linear_attn_config['head_dim']`

---

**`head_k_dim`** `= self.head_dim`

---

**`hidden_size`** `= config.hidden_size`

---

**`k_conv1d`**

---

**`k_proj`**

---

**`mode`** `= getattr(config, 'kda_mode', 'chunk')`

---

**`num_heads`** `= config.linear_attn_config['num_heads']`

---

**`num_k_heads`** `= self.num_heads`

---

**`o_norm`**

---

**`o_proj`**

---

**`q_conv1d`**

---

**`q_proj`**

---

**`use_full_rank_gate`**

---

**`v_conv1d`**

---

**`v_proj`**

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiDeltaAttention._forward_with_cp(
    hidden_states: torch.Tensor,
    packed_context: nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext
) -> torch.Tensor
```

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

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, local\_sequence, hidden].

---

**`packed_context`** `KimiPackedContext`

Document layout of the batch.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, local\_sequence, hidden].

```python
nemo_automodel.components.models.kimi_k3.model.KimiDeltaAttention._kda_core(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None,
    cp_context: typing.Any = None
) -> torch.Tensor
```

Run the KDA projections, convolutions and delta-rule kernel.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; the batch must be
one whenever `cu_seqlens` or `cp_context` is given.

---

**`cu_seqlens`** `torch.Tensor | None` — default: None

Optional cumulative document lengths of shape \[documents + 1].

---

**`cp_context`** `Any` — default: None

Optional FLA context-parallel context, which supersedes
`cu_seqlens` with its per-rank local segments.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

```python
nemo_automodel.components.models.kimi_k3.model.KimiDeltaAttention.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    packed_context: 'KimiPackedContext | None' = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run KDA linear attention.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; the sequence axis
holds this rank's contiguous shard under context parallelism.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary padding mask of shape \[batch, sequence] where 1 marks valid tokens.

---

**`packed_context`** `'KimiPackedContext | None'` — default: None

Optional document layout of the batch, required under
context parallelism and used to reset the recurrent state at every
packed-document boundary.

---

**`**kwargs`** `Any` — default: \{}

Optional KDA kwargs, including `cu_seqlens` for packed sequences.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

```python
nemo_automodel.components.models.kimi_k3.model.KimiDeltaAttention.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiDeltaAttention.setup_cp_attention(
    cp_mesh
) -> None
```

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

**`cp_mesh`**

One-dimensional context-parallel device mesh.

---

```python
class nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3Config | nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    num_hidden_layers: int | None = None,
    kda_mode: str | None = None,
    kwargs: typing.Any = {}
)
```

**Bases:** [HFCheckpointingMixin](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Kimi Linear causal LM with native trainable MoE layers.

**`_keep_in_fp32_modules`** `= ['_fp32_params', 'e_score_correction_bias']`

---

**`_keep_in_fp32_modules_strict`** `= ['_fp32_params']`

---

**`_pp_keep_self_forward`** `bool = True`

---

**`_pp_return_hidden_states_supported`** `bool = True`

---

**`backend`**

---

**`lm_head`**

---

**`model`**

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.UNTIED_ONLY`

---

**`vocab_size`** `= text_config.vocab_size`

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]
```

Keep every K3 attention-residual block within one pipeline stage.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    block_residual: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, torch.Tensor]
```

Run Kimi Linear causal LM.

**Parameters:**

**`input_ids`** `torch.Tensor | None` — default: None

Optional token ids of shape \[batch, sequence].

---

**`block_residual`** `torch.Tensor | None` — default: None

Prior K3 attention-residual block starts. Pipeline
stages after the first receive this as their second activation.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary padding mask of shape \[batch, sequence].

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional positions of shape \[batch, sequence].

---

**`inputs_embeds`** `torch.Tensor | None` — default: None

Optional embeddings of shape \[batch, sequence, hidden].

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional boolean tensor of shape \[batch, sequence], where true marks padding tokens.

---

**`logits_to_keep`** `int | torch.Tensor` — default: 0

Number of trailing sequence logits to compute, or tensor indices.

---

**`output_hidden_states`** `bool | None` — default: None

Whether to include hidden states in the output.

---

**`**attn_kwargs`** `Any` — default: \{}

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.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.from_config(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> 'KimiK3ForCausalLM'
```

classmethod

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> 'KimiK3ForCausalLM'
```

classmethod

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.get_input_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.get_output_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]
```

Return static PP metadata for hidden states and block residuals.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]
```

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

**`batch`** `dict[str, Any]`

Full-sequence batch; left untouched until the returned sharder runs.

---

**`num_chunks`** `int` — default: 1

Accepted for CP hook signature parity; K3 uses one contiguous shard.

---

**Returns:** `dict[str, Any]`

Batch updates carrying the model-owned context-parallel sharder.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.set_input_embeddings(
    value: torch.nn.Module
) -> None
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.set_output_embeddings(
    new_embeddings: torch.nn.Module
) -> None
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3ForCausalLM.update_moe_gate_bias() -> None
```

```python
class nemo_automodel.components.models.kimi_k3.model.KimiK3Gate()
```

**Bases:** [Gate](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-Gate)

K3's fp32 sigmoid router with correction-bias-only expert selection.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3Gate.forward(
    hidden_states: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: typing.Any = None
) -> tuple[torch.Tensor, torch.Tensor, None]
```

Route local token states and return fp32 top-k weights.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[tokens, hidden] containing this rank's
local token states.

---

**`token_mask`** `torch.Tensor`

Boolean tensor of shape \[tokens]. Kimi K3 currently routes
every supplied token, so this mask is unused.

---

**`cp_mesh`** `Any` — default: None

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

```python
class nemo_automodel.components.models.kimi_k3.model.KimiK3MLP(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    hidden_size: int | None = None,
    intermediate_size: int | None = None,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** `Module`

Dense or shared K3 SiTU MLP.

**`act_fn`**

---

**`down_proj`**

---

**`gate_proj`**

---

**`hidden_size`** `= hidden_size or config.hidden_size`

---

**`intermediate_size`** `= intermediate_size or config.intermediate_size`

---

**`up_proj`**

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3MLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Transform `hidden_states` of shape `[..., hidden]`.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3MLP.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None
```

```python
class nemo_automodel.components.models.kimi_k3.model.KimiK3MoE(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [MoE](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-MoE)

K3 routed experts with latent projections and a SiTU shared expert.

**`_situ`**

---

**`dim`** `= moe_config.dim`

---

**`experts`**

---

**`gate`**

---

**`n_activated_experts`** `= moe_config.n_activated_experts`

---

**`n_routed_experts`** `= moe_config.n_routed_experts`

---

**`routed_expert_down_proj`**

---

**`routed_expert_norm`**

---

**`routed_expert_up_proj`**

---

**`shared_experts`**

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3MoE._forward_reference_order(
    hidden_states: torch.Tensor,
    topk_ids: torch.Tensor,
    topk_weight: torch.Tensor
) -> torch.Tensor
```

Match the checkpoint implementation's expert-ordered inference loop.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3MoE._has_distributed_experts() -> bool
```

Whether grouped expert parameters are DTensors.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3MoE.forward(
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
    cp_mesh: typing.Any = None
) -> torch.Tensor
```

Run K3 MoE on `[batch, sequence, hidden]` states.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3MoE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.models.kimi_k3.model.KimiK3TextModel(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict[str, typing.Any] | None = None
)
```

**Bases:** `Module`

Kimi Linear decoder backbone with trainable Automodel MoE layers.

**`embed_tokens`**

---

**`layers`**

---

**`moe_config`**

---

**`norm`**

---

**`output_attn_res_norm`**

---

**`output_attn_res_proj`**

---

**`padding_idx`** `= config.pad_token_id`

---

**`use_attn_residuals`** `= config.attn_res_block_size is not None`

---

**`vocab_size`** `= config.vocab_size`

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3TextModel._apply_output_attn_res(
    hidden_states: torch.Tensor,
    block_residual: torch.Tensor
) -> torch.Tensor
```

Mix final `[batch, sequence, hidden]` states with block starts.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3TextModel._update_linear_attn_mask(
    attention_mask: torch.Tensor | None,
    cache_position: torch.Tensor
) -> torch.Tensor | None
```

Select the padding mask passed to KDA layers.

**Parameters:**

**`attention_mask`** `torch.Tensor | None`

Optional binary padding mask tensor of shape \[batch, sequence].

---

**`cache_position`** `torch.Tensor`

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.

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3TextModel.forward(
    input_ids: torch.Tensor | None = None,
    block_residual: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    cache_position: torch.Tensor | None = None,
    kimi_packed_context: nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext | None = None,
    kimi_packed_doc_ids: torch.Tensor | None = None,
    kimi_packed_seq_start: int = 0,
    kimi_packed_cp_size: int = 1,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]
```

Run the Kimi Linear decoder.

**Parameters:**

**`input_ids`** `torch.Tensor | None` — default: None

Optional token ids of shape \[batch, sequence].

---

**`block_residual`** `torch.Tensor | None` — default: None

Prior attention-residual block starts with shape
\[batch \* sequence, blocks, hidden]. Pipeline stages after the
first receive this as their second positional activation.

---

**`inputs_embeds`** `torch.Tensor | None` — default: None

Optional embeddings of shape \[batch, sequence, hidden].

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary or indexed packing mask of shape \[batch, sequence].

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional positions of shape \[batch, sequence]; accepted for HF compatibility.

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional boolean tensor of shape \[batch, sequence], where true marks padding tokens.

---

**`cache_position`** `torch.Tensor | None` — default: None

Optional position vector of shape \[sequence].

---

**`kimi_packed_context`** `KimiPackedContext | None` — default: None

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.

---

**`kimi_packed_doc_ids`** `torch.Tensor | None` — default: None

Pipeline-safe global document map used to
reconstruct `kimi_packed_context` after microbatch chunking.

---

**`kimi_packed_seq_start`** `int` — default: 0

Global offset of this CP rank's sequence shard.

---

**`kimi_packed_cp_size`** `int` — default: 1

Number of context-parallel sequence shards.

---

**`**attn_kwargs`** `Any` — default: \{}

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

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3TextModel.init_weights(
    buffer_device: torch.device | None = None
) -> None
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiK3TextModel.update_moe_gate_bias() -> None
```

```python
class nemo_automodel.components.models.kimi_k3.model.KimiKDAFp32Params(
    num_heads: int,
    projection_size: int
)
```

**Bases:** `Module`

Own KDA recurrent-decay parameters and compute the FP32 decay gate.

**`A_log`**

---

**`dt_bias`**

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiKDAFp32Params.forward(
    g: torch.Tensor,
    head_dim: int,
    lower_bound: float | None,
    use_fused_gate: bool = True
) -> torch.Tensor
```

Compute the KDA decay while this holder's FSDP unit is unsharded.

**Parameters:**

**`g`** `torch.Tensor`

Raw gate tensor of shape \[batch, sequence, heads \* head\_dim].

---

**`head_dim`** `int`

Per-head KDA dimension.

---

**`lower_bound`** `float | None`

Optional lower bound for K3's bounded decay function.

---

**`use_fused_gate`** `bool` — default: True

Whether to use FLA's fused gate kernel.

---

**Returns:** `torch.Tensor`

FP32 decay tensor of shape \[batch, sequence, heads, head\_dim].

```python
class nemo_automodel.components.models.kimi_k3.model.KimiMLAAttention(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Kimi MLA full-attention layer copied from the HF reference math.

**`attention_dropout`** `= getattr(config, 'attention_dropout', 0.0)`

---

**`g_proj`**

---

**`hidden_size`** `= config.hidden_size`

---

**`kv_a_layernorm`** `= KimiRMSNorm(self.kv_lora_rank, dtype=dtype)`

---

**`kv_a_proj_with_mqa`**

---

**`kv_b_proj`**

---

**`kv_lora_rank`** `= config.kv_lora_rank`

---

**`num_heads`** `= config.num_attention_heads`

---

**`num_key_value_groups`** `= self.num_heads // self.num_key_value_heads`

---

**`num_key_value_heads`** `= config.num_key_value_heads`

---

**`o_proj`**

---

**`q_a_layernorm`** `= KimiRMSNorm(self.q_lora_rank, dtype=dtype)`

---

**`q_a_proj`**

---

**`q_b_proj`**

---

**`q_head_dim`** `= self.qk_nope_head_dim + self.qk_rope_head_dim`

---

**`q_lora_rank`** `= config.q_lora_rank`

---

**`q_proj`**

---

**`qk_nope_head_dim`** `= config.qk_nope_head_dim`

---

**`qk_rope_head_dim`** `= config.qk_rope_head_dim`

---

**`scaling`** `= self.q_head_dim ** -0.5`

---

**`use_output_gate`** `= config.mla_use_output_gate`

---

**`v_head_dim`** `= config.v_head_dim`

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiMLAAttention._expand_key_value_groups(
    key_states: torch.Tensor,
    value_states: torch.Tensor,
    seq_length: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Repeat key/value heads to match the query heads.

**Parameters:**

**`key_states`** `torch.Tensor`

Tensor of shape \[batch, key\_value\_heads, sequence, qk\_head\_dim].

---

**`value_states`** `torch.Tensor`

Tensor of shape \[batch, key\_value\_heads, sequence, v\_head\_dim].

---

**`seq_length`** `int`

Sequence length of the key/value tensors.

---

**Returns:** `tuple[torch.Tensor, torch.Tensor]`

Key and value tensors expanded to \[batch, heads, sequence, head\_dim].

```python
nemo_automodel.components.models.kimi_k3.model.KimiMLAAttention._forward_with_cp(
    hidden_states: torch.Tensor,
    packed_context: nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext
) -> torch.Tensor
```

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

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, local\_sequence, hidden].

---

**`packed_context`** `KimiPackedContext`

Document layout of the batch.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, local\_sequence, hidden].

```python
nemo_automodel.components.models.kimi_k3.model.KimiMLAAttention.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    packed_context: 'KimiPackedContext | None' = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run MLA full attention.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; the sequence axis
holds this rank's contiguous shard under context parallelism.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional additive attention mask of shape \[batch, 1, sequence, sequence].

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional boolean mask of shape \[batch, sequence], where true marks padding.

---

**`packed_context`** `'KimiPackedContext | None'` — default: None

Optional document layout of the batch, required under
context parallelism.

---

**`**kwargs`** `Any` — default: \{}

Extra attention options accepted for HF compatibility.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

```python
nemo_automodel.components.models.kimi_k3.model.KimiMLAAttention.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None
```

```python
nemo_automodel.components.models.kimi_k3.model.KimiMLAAttention.setup_cp_attention(
    cp_mesh
) -> None
```

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

**`cp_mesh`**

One-dimensional context-parallel device mesh.

---

```python
class nemo_automodel.components.models.kimi_k3.model.KimiRMSNorm(
    hidden_size: int,
    eps: float = 1e-06,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** `Module`

Kimi RMSNorm with fp32 variance computation.

**`weight`** `= nn.Parameter(torch.ones(hidden_size, dtype=dtype))`

---

```python
nemo_automodel.components.models.kimi_k3.model.KimiRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Normalize hidden states.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

```python
nemo_automodel.components.models.kimi_k3.model.KimiRMSNorm.reset_parameters() -> None
```

```python
class nemo_automodel.components.models.kimi_k3.model.SituAndMul(
    beta: float = 1.0,
    linear_beta: float | None = None
)
```

**Bases:** `Module`

K3 SiTU gated activation with fp32 nonlinearities.

```python
nemo_automodel.components.models.kimi_k3.model.SituAndMul.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Apply SiTU to `[... , 2 * intermediate]` gate/up projections.

```python
class nemo_automodel.components.models.kimi_k3.model._KimiFp32Module(
    module: torch.nn.Module
)
```

**Bases:** `Module`

Keep a callable FLA operator in its own fp32 FSDP unit.

**`weight`** `Parameter`

Expose the wrapped weight under the reference module API.

---

```python
nemo_automodel.components.models.kimi_k3.model._KimiFp32Module.forward(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any
```

Run the operator while its fp32 FSDP unit is unsharded.

```python
nemo_automodel.components.models.kimi_k3.model._KimiFp32Module.reset_parameters() -> None
```

Reset the wrapped operator.

```python
class nemo_automodel.components.models.kimi_k3.model._KimiKDAFp32Param(
    name: str
)
```

Descriptor exposing a KDA fp32 parameter from the `_fp32_params` holder.

```python
nemo_automodel.components.models.kimi_k3.model._KimiKDAFp32Param.__get__(
    obj: torch.nn.Module | None,
    owner: type[torch.nn.Module] | None = None
) -> torch.nn.Parameter | '_KimiKDAFp32Param'
```

```python
nemo_automodel.components.models.kimi_k3.model._apply_attn_res(
    prefix_sum: torch.Tensor,
    block_residual: torch.Tensor,
    projection: torch.nn.Linear,
    norm: nemo_automodel.components.models.kimi_k3.model.KimiRMSNorm
) -> torch.Tensor
```

Mix `[tokens, hidden]` with prior `[tokens, blocks, hidden]` residuals.

```python
nemo_automodel.components.models.kimi_k3.model._build_moe_config(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    model_dtype: torch.dtype,
    moe_overrides: dict[str, typing.Any] | None
) -> nemo_automodel.components.moe.config.MoEConfig
```

```python
nemo_automodel.components.models.kimi_k3.model._fused_kda_gate(
    g: torch.Tensor,
    a_log: torch.Tensor,
    head_dim: int,
    dt_bias: torch.Tensor,
    lower_bound: float | None
) -> torch.Tensor
```

Call FLA's fused KDA gate across supported FLA APIs.

```python
nemo_automodel.components.models.kimi_k3.model._get_unpad_data(
    attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, int]
```

Build metadata for converting padded batches to flattened valid tokens.

**Parameters:**

**`attention_mask`** `torch.Tensor`

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

```python
nemo_automodel.components.models.kimi_k3.model._index_first_axis(
    x: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor
```

Gather rows from the first axis while preserving trailing tensor layout.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[tokens, ...], with arbitrary trailing axes.

---

**`indices`** `torch.Tensor`

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

```python
nemo_automodel.components.models.kimi_k3.model._index_put_first_axis(
    x: torch.Tensor,
    indices: torch.Tensor,
    first_axis_dim: int
) -> torch.Tensor
```

Scatter rows into the first axis while preserving trailing tensor layout.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[selected\_tokens, ...], with arbitrary trailing axes.

---

**`indices`** `torch.Tensor`

Tensor of shape \[selected\_tokens] containing destination row indices.

---

**`first_axis_dim`** `int`

Size of the output first axis.

---

**Returns:** `torch.Tensor`

Tensor of shape \[first\_axis\_dim, ...], with the same trailing axes as `x`.

```python
nemo_automodel.components.models.kimi_k3.model._make_causal_mask(
    inputs_embeds: torch.Tensor,
    packed_context: 'KimiPackedContext | None',
    dtype: torch.dtype
) -> torch.Tensor | None
```

Create the additive causal attention mask for full-attention layers.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**`packed_context`** `'KimiPackedContext | None'`

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.

---

**`dtype`** `torch.dtype`

Floating-point dtype used for the additive mask values.

---

**Returns:** `torch.Tensor | None`

Additive causal mask tensor of shape \[batch, 1, sequence, sequence].

```python
nemo_automodel.components.models.kimi_k3.model._packed_context_from_inputs(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | None,
    cu_seqlens: torch.Tensor | None
) -> nemo_automodel.components.models.kimi_k3.cp.KimiPackedContext | None
```

Derive the document layout of a batch that was not sharded for context parallelism.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**`attention_mask`** `torch.Tensor | None`

Optional binary or indexed packing mask of shape \[batch, sequence].

---

**`cu_seqlens`** `torch.Tensor | None`

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

```python
nemo_automodel.components.models.kimi_k3.model._pad_input(
    hidden_states: torch.Tensor,
    indices: torch.Tensor,
    batch_size: int,
    seq_len: int
) -> torch.Tensor
```

Restore flattened valid tokens to padded batch layout.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[total\_valid\_tokens, ...], with arbitrary trailing axes.

---

**`indices`** `torch.Tensor`

Tensor of shape \[total\_valid\_tokens] containing flattened padded-batch row indices.

---

**`batch_size`** `int`

Number of sequences in the padded output batch.

---

**`seq_len`** `int`

Sequence length in the padded output batch.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, ...], with the same trailing axes as `hidden_states`.

```python
nemo_automodel.components.models.kimi_k3.model._partition_attn_residual_blocks(
    num_layers: int,
    block_size: int,
    num_stages: int,
    allow_output_only_stage: bool = False
) -> list[range]
```

Partition decoder layers without splitting an attention-residual block.

```python
nemo_automodel.components.models.kimi_k3.model._require_fla() -> None
```

```python
nemo_automodel.components.models.kimi_k3.model._seed_dtensor_rng_for_pipeline_stage(
    model: torch.nn.Module
) -> None
```

Initialize DTensor RNG without a world broadcast during PP weight init.

```python
nemo_automodel.components.models.kimi_k3.model._torch_kda_gate(
    g: torch.Tensor,
    a_log: torch.Tensor,
    head_dim: int,
    dt_bias: torch.Tensor,
    lower_bound: float | None
) -> torch.Tensor
```

Compute K3's KDA decay gate with torch FP32 operations.

**Parameters:**

**`g`** `torch.Tensor`

Raw gate tensor of shape \[batch, sequence, heads \* head\_dim] or
\[batch, sequence, heads, head\_dim].

---

**`a_log`** `torch.Tensor`

Log decay tensor of shape \[heads].

---

**`head_dim`** `int`

Per-head KDA dimension.

---

**`dt_bias`** `torch.Tensor`

Gate bias tensor of shape \[heads \* head\_dim].

---

**`lower_bound`** `float | None`

Optional lower bound for K3's bounded decay function.

---

**Returns:** `torch.Tensor`

FP32 decay tensor of shape \[batch, sequence, heads, head\_dim].

```python
nemo_automodel.components.models.kimi_k3.model._weighted_situ(
    gate_up: torch.Tensor,
    routing_weights: torch.Tensor,
    beta: float,
    linear_beta: float | None
) -> torch.Tensor
```

Apply SiTU and routing weights to `[tokens, 2 * intermediate]` projections.

```python
nemo_automodel.components.models.kimi_k3.model.ModelClass = KimiK3ForCausalLM
```

```python
nemo_automodel.components.models.kimi_k3.model._FLA_MSG = 'Kimi K3 requires the flash-linear-attention/fla extra. Install with `uv sync --...
```

```python
nemo_automodel.components.models.kimi_k3.model._FUSED_KDA_GATE_HAS_G_BIAS = _KDA_GATE_OK and 'g_bias' in inspect.signature(fused_kda_gate).parameters
```

```python
nemo_automodel.components.models.kimi_k3.model._FUSED_KDA_GATE_HAS_LOWER_BOUND = _KDA_GATE_OK and 'lower_bound' in inspect.signature(fused_kda_gate).parameters
```