> 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.dflash.draft_kimi_k3

DFlash draft model with a dense Kimi K3 MLA backbone, plus its build helpers.

The Qwen3 draft (`draft_qwen3.py`) documents the DFlash contract: the draft
predicts a whole `block_size` block in one non-causal forward whose keys and
values are `[target-hidden context | noise block]`, with the block structure
supplied entirely by the attention mask built in
`nemo_automodel.components.speculative.dflash.core`.

This module keeps that contract and swaps the backbone for Kimi K3's:

* **MLA with Q-LoRA and a compressed KV latent.** The attention subclasses the
  target's own :class:`KimiMLAAttention` so the draft's projection layout stays
  in lockstep with the layers whose hidden states it consumes; only the forward
  is replaced, because the draft's keys and values span the context as well as
  the queried noise block.
* **NoPE.** K3 requires `mla_use_nope=True`: no rotary is applied anywhere in
  its full-attention layers (the `qk_rope_head_dim` slice is a plain
  head-shared extension of the key). The draft therefore has no rotary
  embedding and ignores `position_ids`.
* **SiTU feed-forward and fp32-variance RMSNorm**, reusing the target's
  `KimiK3MLP` / `KimiRMSNorm` modules.

The draft is always dense: K3's KDA linear-attention layers, routed and shared
experts, MTP heads, and the learned attention-residual mixer live in the target
only. Because the draft has no FlexAttention path, it consumes the dense
additive DFlash mask (`attention_backend='sdpa'`).

## Module Contents

### Classes

| Name                                                                                                               | Description                                                                          |
| ------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------ |
| [`KimiK3DFlashAttention`](#nemo_automodel-components-speculative-dflash-draft_kimi_k3-KimiK3DFlashAttention)       | Non-causal K3 MLA whose keys/values are `[context \| noise-block]`.                  |
| [`KimiK3DFlashDecoderLayer`](#nemo_automodel-components-speculative-dflash-draft_kimi_k3-KimiK3DFlashDecoderLayer) | Pre-norm K3 MLA block over `[context \| noise]` followed by a dense SiTU MLP.        |
| [`KimiK3DFlashDraftModel`](#nemo_automodel-components-speculative-dflash-draft_kimi_k3-KimiK3DFlashDraftModel)     | DFlash draft model: a small dense non-causal K3 MLA stack over `[context \| noise]`. |

### Functions

| Name                                                                                                                                   | Description                                                                |
| -------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`build_kimi_k3_dflash_draft_config`](#nemo_automodel-components-speculative-dflash-draft_kimi_k3-build_kimi_k3_dflash_draft_config)   | Build a dense MLA DFlash draft config from a Kimi K3 target's text config. |
| [`build_kimi_k3_dflash_target_kwargs`](#nemo_automodel-components-speculative-dflash-draft_kimi_k3-build_kimi_k3_dflash_target_kwargs) | Extra `from_pretrained` kwargs for a frozen Kimi K3 DFlash target.         |

### Data

[`__all__`](#nemo_automodel-components-speculative-dflash-draft_kimi_k3-__all__)

### API

```python
class nemo_automodel.components.speculative.dflash.draft_kimi_k3.KimiK3DFlashAttention(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [KimiMLAAttention](/nemo-automodel/nemo_automodel/components/models/kimi_k3/model#nemo_automodel-components-models-kimi_k3-model-KimiMLAAttention)

Non-causal K3 MLA whose keys/values are `[context | noise-block]`.

Inherits the target's MLA projections unchanged and replaces only the
forward. Queries come from the draft (noise) tokens only. K3's MLA emits one
K/V per query head (`kv_b_proj` produces `num_attention_heads` slices of
`qk_nope_head_dim + v_head_dim`), so there is no GQA group repeat and the
target's `_expand_key_value_groups` is not needed here.

```python
nemo_automodel.components.speculative.dflash.draft_kimi_k3.KimiK3DFlashAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Attend the draft block over `[context | noise]`.

Unlike the target's eager forward, this runs
`F.scaled_dot_product_attention`: a DFlash block batch makes both the
query and the key axis sequence-scale (`num_anchors * block_size` by
`sequence + num_anchors * block_size`), so materializing the score
matrix is not affordable. The dense additive mask is passed through
unchanged, which keeps the visibility semantics identical.

**Parameters:**

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

Tensor of shape `[batch, blocks * block_size, hidden]`;
the draft (noise) tokens, which are the queries.

---

**`target_hidden`** `torch.Tensor`

Tensor of shape `[batch, sequence, hidden]`; the
projected target-hidden context prepended to the keys and values.

---

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

Additive mask of shape
`[batch, 1, blocks * block_size, sequence + blocks * block_size]`,
or None.

---

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

Ignored; accepted so the decoder layer can forward extras.

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, blocks * block_size, hidden]`.

```python
class nemo_automodel.components.speculative.dflash.draft_kimi_k3.KimiK3DFlashDecoderLayer(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Pre-norm K3 MLA block over `[context | noise]` followed by a dense SiTU MLP.

**`input_layernorm`**

---

**`mlp`** `= KimiK3MLP(config)`

---

**`post_attention_layernorm`**

---

**`self_attn`** `= KimiK3DFlashAttention(config, layer_idx, backend)`

---

```python
nemo_automodel.components.speculative.dflash.draft_kimi_k3.KimiK3DFlashDecoderLayer.forward(
    hidden_states: torch.Tensor,
    target_hidden: torch.Tensor,
    attention_mask: torch.Tensor | None = None
) -> torch.Tensor
```

Run one draft layer.

**Parameters:**

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

Tensor of shape `[batch, blocks * block_size, hidden]`.

---

**`target_hidden`** `torch.Tensor`

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

---

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

Additive mask of shape
`[batch, 1, blocks * block_size, sequence + blocks * block_size]`,
or None.

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, blocks * block_size, hidden]`.

```python
class nemo_automodel.components.speculative.dflash.draft_kimi_k3.KimiK3DFlashDraftModel(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig
)
```

**Bases:** `Module`

DFlash draft model: a small dense non-causal K3 MLA stack over `[context | noise]`.

The draft owns no embedding table and no LM head: the DFlash trainer embeds
the `[anchor, MASK, ...]` blocks with the frozen target's `embed_tokens`
and decodes this model's output with the frozen target's `lm_head`.

**`_no_split_modules`** `= ['KimiK3DFlashDecoderLayer']`

---

**`fc`**

---

**`hidden_norm`**

---

**`layers`**

---

**`norm`**

---

**`target_layer_ids`**

---

```python
nemo_automodel.components.speculative.dflash.draft_kimi_k3.KimiK3DFlashDraftModel.forward(
    position_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    noise_embedding: torch.Tensor | None = None,
    target_hidden: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Predict the draft blocks' hidden states.

**Parameters:**

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

Unused. K3's MLA is NoPE, so the draft has no rotary
embedding; the argument is accepted to keep the trainer's call
signature identical across DFlash drafts.

---

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

Additive DFlash mask of shape
`[batch, 1, blocks * block_size, sequence + blocks * block_size]`.

---

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

Tensor of shape `[batch, blocks * block_size, hidden]`.

---

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

Tensor of shape
`[batch, sequence, len(target_layer_ids) * hidden]`.

---

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

Ignored.

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, blocks * block_size, hidden]`.

```python
nemo_automodel.components.speculative.dflash.draft_kimi_k3.build_kimi_k3_dflash_draft_config(
    target_config,
    num_draft_layers: int,
    num_target_layers: int,
    block_size: int,
    dflash_config: dict,
    attention_backend: str
) -> nemo_automodel.components.models.kimi_k3.config.KimiK3TextConfig
```

Build a dense MLA DFlash draft config from a Kimi K3 target's text config.

The draft consumes the target's frozen `embed_tokens` / `lm_head` and
fuses its hidden states, so it keeps the target's MLA dims, hidden size, and
vocabulary and only shrinks the depth. Everything the draft does not build is
switched off explicitly rather than left at the target's value, so the
serialized draft config describes the draft and not the target: KDA linear
attention, routed and shared experts, MTP layers, and the learned
attention-residual mixer all stay in the target.

**Parameters:**

**`target_config`**

The Kimi K3 target's text config (`kimi_linear`).

---

**`num_draft_layers`** `int`

Number of draft decoder layers.

---

**`num_target_layers`** `int`

Depth of the target's text backbone; recorded so a
reloaded draft config still describes which target it was trained on.

---

**`block_size`** `int`

DFlash block size.

---

**`dflash_config`** `dict`

The recipe's DFlash block, carrying `mask_token_id` and
`target_layer_ids` (the target layers whose hidden states the draft
consumes, which set the `fc` input width).

---

**`attention_backend`** `str`

The draft's attention implementation. Recorded on the
config for the serving runtime; the draft itself always attends over
the dense additive mask.

---

**Returns:** `KimiK3TextConfig`

A `KimiK3TextConfig` describing the draft.

**Raises:**

* `ValueError`: If `target_config` is not a Kimi K3 text config.

```python
nemo_automodel.components.speculative.dflash.draft_kimi_k3.build_kimi_k3_dflash_target_kwargs(
    recipe_cfg
) -> dict
```

Extra `from_pretrained` kwargs for a frozen Kimi K3 DFlash target.

Two things a Qwen3-shaped target does not need:

* `config` pins the architecture to the text-only `KimiK3ForCausalLM`. A K3
  checkpoint declares the multimodal `KimiK3ForConditionalGeneration`, which
  would additionally build the vision tower that DFlash never reads.
* `backend` selects the expert-parallel token dispatcher and the HF
  state-dict adapter (which also dequantizes an FP8 base checkpoint on load),
  mirroring the frozen large-MoE target backends the DSpark recipe builds.
  `experts` defaults to `torch_mm` rather than `gmm` because the latter
  needs the optional `grouped_gemm` package. `attn` is left at `eager`
  and is inert -- `KimiK3ForCausalLM` never reads `backend.attn`, since
  its MLA and KDA layers each have a fixed attention path -- and
  `gate_precision` is left unset because K3 already defaults it to fp32.

**Parameters:**

**`recipe_cfg`**

The recipe's `recipe_args` mapping.

---

**Returns:** `dict`

Keyword arguments to merge into the target's `from_pretrained` call.

```python
nemo_automodel.components.speculative.dflash.draft_kimi_k3.__all__ = ['KimiK3DFlashDraftModel', 'build_kimi_k3_dflash_draft_config', 'build_kimi_k3_d...
```