> 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.speculative.eagle.draft_kimi_k3

Kimi K3 (NoPE MLA backbone) EAGLE-3 draft model.

Kimi K3 interleaves Kimi Delta Attention (KDA) layers with Multi-head Latent
Attention (MLA) layers. The draft mirrors only the MLA layer, which differs from
the DeepSeek MLA the shared MLA draft (`draft_deepseek`) implements in three
ways:

* **NoPE.** K3 fixes `mla_use_nope=True`: the `qk_rope_head_dim` slice is
  concatenated to the query and key heads unrotated, and no rotary embedding is
  applied anywhere in the MLA layer (positional information reaches the full
  attention layers through the interleaved KDA recurrence). The draft keeps that
  contract, so it registers no rotary cache and applies no per-TTT-step rotary
  phase offset.
* **Output gate.** `mla_use_output_gate=True` multiplies the attention output
  by a sigmoid gate projected from the layer input.
* **SiTU MLP.** The feed-forward block is K3's `SituAndMul` gated MLP
  (`KimiK3MLP`), not SwiGLU, and every norm is K3's fp32-variance
  `KimiRMSNorm`.

Everything EAGLE-3 specific is inherited from the DeepSeek MLA draft: the
`[embed, hidden]` fused first layer, the `cache_hidden = [K_list, V_list]`
TTT recurrence with diagonal-extension attention, `project_hidden_states` /
`compute_logits` / `set_vocab_mapping` and the `d2t` / `t2d` buffers. The
SiTU MLP and the norm are imported from the onboarded K3 target
(`components/models/kimi_k3.model`) and the MLA projection layout and the gate
follow its `KimiMLAAttention` (widened to the fused `[embed, hidden]` input),
so the draft's math matches the target's MLA.

Scope: EAGLE-3 single fused draft layer, eager attention, dense draft (KDA,
routed experts and the attention-residual mixer stay in the target only). The
draft itself consumes the shared `[B, 1, T, T]` block-causal packed mask like
the DeepSeek MLA draft, but the recipe rejects sequence packing for a K3 target:
K3 owns its own packed-attention path, which the EAGLE-3 target wrapper does not
drive (see `examples/speculative/eagle3/README_kimi_k3.md`).

## Module Contents

### Classes

| Name                                                                                                              | Description                                                                                 |
| ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- |
| [`Eagle3KimiK3Model`](#nemo_automodel-components-speculative-eagle-draft_kimi_k3-Eagle3KimiK3Model)               | Inner backbone: `embed_tokens`, the `fc` aux-projection, the fused draft layer, and `norm`. |
| [`KimiK3Eagle3DecoderLayer`](#nemo_automodel-components-speculative-eagle-draft_kimi_k3-KimiK3Eagle3DecoderLayer) | Fused EAGLE-3 first layer: `[embed, hidden]` -> NoPE MLA -> SiTU MLP.                       |
| [`KimiK3Eagle3DraftModel`](#nemo_automodel-components-speculative-eagle-draft_kimi_k3-KimiK3Eagle3DraftModel)     | Kimi K3 (NoPE MLA) EAGLE-3 draft model.                                                     |
| [`KimiK3Eagle3MLAAttention`](#nemo_automodel-components-speculative-eagle-draft_kimi_k3-KimiK3Eagle3MLAAttention) | NoPE MLA self-attention with an output gate for the Kimi K3 EAGLE-3 draft.                  |

### Data

[`_INIT_DTYPE`](#nemo_automodel-components-speculative-eagle-draft_kimi_k3-_INIT_DTYPE)

### API

```python
class nemo_automodel.components.speculative.eagle.draft_kimi_k3.Eagle3KimiK3Model(
    config: transformers.PretrainedConfig
)
```

**Bases:** `Module`

Inner backbone: `embed_tokens`, the `fc` aux-projection, the fused draft layer, and `norm`.

**`embed_tokens`**

---

**`fc`**

---

**`fc_norm`**

---

**`layers`**

---

**`norm`**

---

```python
class nemo_automodel.components.speculative.eagle.draft_kimi_k3.KimiK3Eagle3DecoderLayer(
    config: transformers.PretrainedConfig,
    layer_id: int = 0
)
```

**Bases:** `Module`

Fused EAGLE-3 first layer: `[embed, hidden]` -> NoPE MLA -> SiTU MLP.

**`hidden_norm`**

---

**`input_layernorm`**

---

**`mlp`** `= KimiK3MLP(config, dtype=_INIT_DTYPE)`

---

**`post_attention_layernorm`**

---

**`self_attn`** `= KimiK3Eagle3MLAAttention(config)`

---

```python
nemo_automodel.components.speculative.eagle.draft_kimi_k3.KimiK3Eagle3DecoderLayer.forward(
    input_embeds: torch.Tensor,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    cache_hidden: list[list[torch.Tensor]]
) -> torch.Tensor
```

Fuse the embedding and hidden streams, then run attention and the MLP.

**Parameters:**

**`input_embeds`** `torch.Tensor`

Tensor of shape `[batch, sequence, hidden]`, the draft
token embeddings.

---

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

Tensor of shape `[batch, sequence, hidden]`, the
projected target hidden states; also the residual stream.

---

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

Additive mask of shape `[batch, 1, sequence, sequence]`.

---

**`position_ids`** `torch.Tensor`

Tensor of shape `[batch, sequence]`; unused (NoPE MLA).

---

**`cache_hidden`** `list[list[torch.Tensor]]`

TTT cache `[K_list, V_list]` forwarded to the attention.

---

**Returns:** `torch.Tensor`

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

```python
class nemo_automodel.components.speculative.eagle.draft_kimi_k3.KimiK3Eagle3DraftModel()
```

**Bases:** [DeepseekV3Eagle3DraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_deepseek#nemo_automodel-components-speculative-eagle-draft_deepseek-DeepseekV3Eagle3DraftModel)

Kimi K3 (NoPE MLA) EAGLE-3 draft model.

Same public training API as every other EAGLE-3 draft
(`project_hidden_states` / `embed_input_ids` / `compute_logits` /
`set_vocab_mapping` / `freeze_embeddings` / `forward`) and the same
`d2t` / `t2d` vocab-remap buffers, so the EAGLE-3 trainer and
checkpointing are reused unchanged; only the backbone is K3's.

```python
nemo_automodel.components.speculative.eagle.draft_kimi_k3.KimiK3Eagle3DraftModel._build_inner_model(
    config: transformers.PretrainedConfig
) -> torch.nn.Module
```

```python
class nemo_automodel.components.speculative.eagle.draft_kimi_k3.KimiK3Eagle3MLAAttention(
    config: transformers.PretrainedConfig,
    fuse_input: bool = True
)
```

**Bases:** [Eagle3DeepseekMLAAttention](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_deepseek#nemo_automodel-components-speculative-eagle-draft_deepseek-Eagle3DeepseekMLAAttention)

NoPE MLA self-attention with an output gate for the Kimi K3 EAGLE-3 draft.

Reuses the DeepSeek draft's TTT recurrence verbatim (`_eager_attention_forward`:
full `T x T` causal attention against the step-0 keys plus one diagonal column
per cached later step). Only the projections change: no rotary cache is built,
the rope slice rides along unrotated, and the attention output is gated. The
parent `__init__` is bypassed deliberately -- it registers a rotary buffer
that a NoPE layer must not carry.

**`g_proj`**

---

**`kv_a_layernorm`**

---

**`kv_a_proj_with_mqa`**

---

**`kv_b_proj`**

---

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

---

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

---

**`o_proj`**

---

**`q_a_layernorm`**

---

**`q_a_proj`**

---

**`q_b_proj`**

---

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

---

**`q_proj`**

---

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

---

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

---

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

---

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

---

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

---

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

---

```python
nemo_automodel.components.speculative.eagle.draft_kimi_k3.KimiK3Eagle3MLAAttention.forward(
    combined_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    cache_hidden: list[list[torch.Tensor]]
) -> torch.Tensor
```

Run one NoPE MLA step over the fused EAGLE-3 input.

**Parameters:**

**`combined_states`** `torch.Tensor`

Tensor of shape `[batch, sequence, 2 * hidden]`, the
fused `[normed embed, normed hidden]` EAGLE-3 first-layer input.

---

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

Additive mask of shape `[batch, 1, sequence, sequence]`.

---

**`position_ids`** `torch.Tensor`

Tensor of shape `[batch, sequence]`; unused (NoPE MLA).

---

**`cache_hidden`** `list[list[torch.Tensor]]`

TTT cache `[K_list, V_list]`; each list holds one entry
per TTT step, shaped `[batch, heads, sequence, qk_head_dim]` for
keys and `[batch, heads, sequence, v_head_dim]` for values.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.speculative.eagle.draft_kimi_k3._INIT_DTYPE = torch.float32
```