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`	Inner backbone: `embed_tokens`, the `fc` aux-projection, the fused draft layer, and `norm`.
`KimiK3Eagle3DecoderLayer`	Fused EAGLE-3 first layer: `[embed, hidden]` -> NoPE MLA -> SiTU MLP.
`KimiK3Eagle3DraftModel`	Kimi K3 (NoPE MLA) EAGLE-3 draft model.
`KimiK3Eagle3MLAAttention`	NoPE MLA self-attention with an output gate for the Kimi K3 EAGLE-3 draft.

Data

_INIT_DTYPE

API

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_norm

layers

norm

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)

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

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

Bases: 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.

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

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

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

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

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