> 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_linear.model

Native Automodel support for Moonshot Kimi Linear causal LM checkpoints.

## Module Contents

### Classes

| Name                                                                                                       | Description                                                              |
| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`KimiDecoderLayer`](#nemo_automodel-components-models-kimi_linear-model-KimiDecoderLayer)                 | Kimi decoder block with KDA/MLA attention and dense or MoE MLP.          |
| [`KimiDeltaAttention`](#nemo_automodel-components-models-kimi_linear-model-KimiDeltaAttention)             | Kimi Delta Attention backed by FLA KDA kernels.                          |
| [`KimiKDAFp32Params`](#nemo_automodel-components-models-kimi_linear-model-KimiKDAFp32Params)               | Owns Kimi KDA fp32 recurrent-decay parameters and computes the gate.     |
| [`KimiLinear48BForCausalLM`](#nemo_automodel-components-models-kimi_linear-model-KimiLinear48BForCausalLM) | Kimi Linear causal LM with native trainable MoE layers.                  |
| [`KimiLinear48BModel`](#nemo_automodel-components-models-kimi_linear-model-KimiLinear48BModel)             | Kimi Linear decoder backbone with trainable Automodel MoE layers.        |
| [`KimiMLAAttention`](#nemo_automodel-components-models-kimi_linear-model-KimiMLAAttention)                 | Kimi MLA full-attention layer copied from the HF reference math.         |
| [`KimiRMSNorm`](#nemo_automodel-components-models-kimi_linear-model-KimiRMSNorm)                           | Kimi RMSNorm with fp32 variance computation.                             |
| [`_KimiKDAFp32Param`](#nemo_automodel-components-models-kimi_linear-model-_KimiKDAFp32Param)               | Descriptor exposing a KDA fp32 parameter from the `_fp32_params` holder. |

### Functions

| Name                                                                                                             | Description                                                                         |
| ---------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`_build_moe_config`](#nemo_automodel-components-models-kimi_linear-model-_build_moe_config)                     | -                                                                                   |
| [`_fused_kda_gate`](#nemo_automodel-components-models-kimi_linear-model-_fused_kda_gate)                         | Call FLA fused KDA gate across FLA versions.                                        |
| [`_get_unpad_data`](#nemo_automodel-components-models-kimi_linear-model-_get_unpad_data)                         | Build metadata for converting padded batches to flattened valid tokens.             |
| [`_index_first_axis`](#nemo_automodel-components-models-kimi_linear-model-_index_first_axis)                     | Gather rows from the first axis while preserving trailing tensor layout.            |
| [`_index_put_first_axis`](#nemo_automodel-components-models-kimi_linear-model-_index_put_first_axis)             | Scatter rows into the first axis while preserving trailing tensor layout.           |
| [`_make_causal_mask`](#nemo_automodel-components-models-kimi_linear-model-_make_causal_mask)                     | Create the additive causal attention mask for full-attention layers.                |
| [`_packed_context_from_inputs`](#nemo_automodel-components-models-kimi_linear-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_linear-model-_pad_input)                                   | Restore flattened valid tokens to padded batch layout.                              |
| [`_require_fla`](#nemo_automodel-components-models-kimi_linear-model-_require_fla)                               | -                                                                                   |
| [`_torch_kda_gate`](#nemo_automodel-components-models-kimi_linear-model-_torch_kda_gate)                         | Torch equivalent of FLA's KDA gate.                                                 |

### Data

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

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

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

### API

```python
class nemo_automodel.components.models.kimi_linear.model.KimiDecoderLayer(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    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.

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

---

**`input_layernorm`**

---

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

---

**`is_moe_layer`**

---

**`mlp`** `= MoE(moe_config, backend)`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.models.kimi_linear.model.KimiDecoderLayer._has_dtensor_expert_params(
    experts: nemo_automodel.components.moe.experts.GroupedExperts
) -> bool
```

staticmethod

```python
nemo_automodel.components.models.kimi_linear.model.KimiDecoderLayer._moe(
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None
) -> torch.Tensor
```

Run a Kimi MoE layer.

**Parameters:**

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

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

---

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

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

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.models.kimi_linear.model.KimiDecoderLayer._moe_infer_hf_order(
    moe: nemo_automodel.components.moe.layers.MoE,
    experts: nemo_automodel.components.moe.experts.GroupedExperts,
    hidden_states: torch.Tensor
) -> torch.Tensor
```

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`** `MoE`

MoE module containing the router and optional shared experts.

---

**`experts`** `GroupedExperts`

Grouped routed experts for the layer.

---

**`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_linear.model.KimiDecoderLayer.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> 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.

---

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

Extra attention kwargs forwarded to KDA/MLA.

---

**Returns:** `torch.Tensor`

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

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

```python
class nemo_automodel.components.models.kimi_linear.model.KimiDeltaAttention(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    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`**

---

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

---

**`v_conv1d`**

---

**`v_proj`**

---

```python
nemo_automodel.components.models.kimi_linear.model.KimiDeltaAttention._forward_with_cp(
    hidden_states: torch.Tensor,
    packed_context: nemo_automodel.components.models.kimi_linear.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_linear.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_linear.model.KimiDeltaAttention._resolve_mode(
    seq_len: int
) -> str
```

Return the KDA kernel to use for a sequence of `seq_len` tokens.

```python
nemo_automodel.components.models.kimi_linear.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_linear.model.KimiDeltaAttention.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None
```

```python
nemo_automodel.components.models.kimi_linear.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_linear.model.KimiKDAFp32Params(
    num_heads: int,
    projection_size: int
)
```

**Bases:** `Module`

Owns Kimi KDA fp32 recurrent-decay parameters and computes the gate.

**`A_log`**

---

**`dt_bias`**

---

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

Compute KDA decay gate while holder params are unsharded by FSDP.

**Parameters:**

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

Tensor of shape \[batch, sequence, heads \* head\_dim].

---

**`head_dim`** `int`

Per-head KDA dimension.

---

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

Whether to use FLA's fused KDA gate kernel.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, heads, head\_dim].

```python
class nemo_automodel.components.models.kimi_linear.model.KimiLinear48BForCausalLM(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | 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', 'e_score_correction_bias']`

---

**`backend`**

---

**`lm_head`**

---

**`model`**

---

**`state_dict_adapter`**

---

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

---

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

---

```python
nemo_automodel.components.models.kimi_linear.model.KimiLinear48BForCausalLM.forward(
    input_ids: 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
```

Run Kimi Linear causal LM.

**Parameters:**

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

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

---

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

Causal LM output whose logits have shape \[batch, sequence, vocab] unless `logits_to_keep` trims sequence.

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

classmethod

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

classmethod

```python
nemo_automodel.components.models.kimi_linear.model.KimiLinear48BForCausalLM.get_input_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.kimi_linear.model.KimiLinear48BForCausalLM.get_output_embeddings() -> torch.nn.Module
```

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

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

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

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

The full-sequence batch; left intact, the sharder shards it.

---

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

Accepted for hook-signature parity; unused, because Kimi
Linear shards the `[batch, sequence]` layout directly.

---

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

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

```python
nemo_automodel.components.models.kimi_linear.model.KimiLinear48BForCausalLM.update_moe_gate_bias() -> None
```

```python
class nemo_automodel.components.models.kimi_linear.model.KimiLinear48BModel(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    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`**

---

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

---

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

---

```python
nemo_automodel.components.models.kimi_linear.model.KimiLinear48BModel._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_linear.model.KimiLinear48BModel.forward(
    input_ids: 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_linear.cp.KimiPackedContext | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
```

Run the Kimi Linear decoder.

**Parameters:**

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

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

---

**`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_linear.cp.shard_batch_for_kimi_cp`;
required under context parallelism and otherwise derived here.

---

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

Additional attention kwargs used by packed or THD execution.

---

**Returns:** `torch.Tensor`

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

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

```python
nemo_automodel.components.models.kimi_linear.model.KimiLinear48BModel.update_moe_gate_bias() -> None
```

```python
class nemo_automodel.components.models.kimi_linear.model.KimiMLAAttention(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    layer_idx: int
)
```

**Bases:** `Module`

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

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

---

**`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_head_dim`** `= self.qk_nope_head_dim + self.qk_rope_head_dim`

---

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

---

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

---

```python
nemo_automodel.components.models.kimi_linear.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_linear.model.KimiMLAAttention._forward_with_cp(
    hidden_states: torch.Tensor,
    packed_context: nemo_automodel.components.models.kimi_linear.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_linear.model.KimiMLAAttention.forward(
    hidden_states: torch.Tensor,
    attention_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].

---

**`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_linear.model.KimiMLAAttention.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None
```

```python
nemo_automodel.components.models.kimi_linear.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_linear.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_linear.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_linear.model.KimiRMSNorm.reset_parameters() -> None
```

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

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

```python
nemo_automodel.components.models.kimi_linear.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_linear.model._build_moe_config(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    model_dtype: torch.dtype,
    moe_overrides: dict[str, typing.Any] | None
) -> nemo_automodel.components.moe.config.MoEConfig
```

```python
nemo_automodel.components.models.kimi_linear.model._fused_kda_gate(
    g: torch.Tensor,
    a_log: torch.Tensor,
    head_dim: int,
    dt_bias: torch.Tensor
) -> torch.Tensor
```

Call FLA fused KDA gate across FLA versions.

**Parameters:**

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

Tensor of shape \[batch, sequence, heads \* head\_dim].

---

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

Tensor of shape \[1, 1, heads, 1].

---

**`head_dim`** `int`

Per-head KDA dimension.

---

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

Tensor of shape \[heads \* head\_dim].

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, heads, head\_dim].

```python
nemo_automodel.components.models.kimi_linear.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_linear.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_linear.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_linear.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_linear.model._packed_context_from_inputs(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | None,
    cu_seqlens: torch.Tensor | None
) -> nemo_automodel.components.models.kimi_linear.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_linear.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_linear.model._require_fla() -> None
```

```python
nemo_automodel.components.models.kimi_linear.model._torch_kda_gate(
    g: torch.Tensor,
    a_log: torch.Tensor,
    head_dim: int,
    dt_bias: torch.Tensor
) -> torch.Tensor
```

Torch equivalent of FLA's KDA gate.

**Parameters:**

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

Tensor of shape \[batch, sequence, heads \* head\_dim] or \[batch, sequence, heads, head\_dim].

---

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

Tensor of shape \[1, 1, heads, 1].

---

**`head_dim`** `int`

Per-head KDA dimension.

---

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

Tensor of shape \[heads \* head\_dim].

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, heads, head\_dim].

```python
nemo_automodel.components.models.kimi_linear.model.ModelClass = KimiLinear48BForCausalLM
```

```python
nemo_automodel.components.models.kimi_linear.model._FLA_MSG = 'Kimi Linear requires the flash-linear-attention/fla extra. Install with `uv syn...
```

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