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`	Non-causal K3 MLA whose keys/values are `[context \| noise-block]`.
`KimiK3DFlashDecoderLayer`	Pre-norm K3 MLA block over `[context \| noise]` followed by a dense SiTU MLP.
`KimiK3DFlashDraftModel`	DFlash draft model: a small dense non-causal K3 MLA stack over `[context \| noise]`.

Functions

Name	Description
`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`	Extra `from_pretrained` kwargs for a frozen Kimi K3 DFlash target.

Data

__all__

API

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

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.

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 | NoneDefaults to None

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

**kwargs

AnyDefaults to {}

Ignored; accepted so the decoder layer can forward extras.

Returns: torch.Tensor

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

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)

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 | NoneDefaults to 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].

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

hidden_norm

layers

norm

target_layer_ids

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 | NoneDefaults to 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 | NoneDefaults to None

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

noise_embedding

torch.Tensor | NoneDefaults to None

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

target_hidden

torch.Tensor | NoneDefaults to None

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

**kwargs

AnyDefaults to {}

Ignored.

Returns: torch.Tensor

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

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.

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.

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