nemo_automodel.components.speculative.eagle.draft_deepseek

DeepSeek (MLA-backbone) EAGLE-3 draft model.

The Llama-style draft (draft_llama.LlamaEagle3DraftModel) covers dense multi-head-attention targets (Llama / Phi-3 / Qwen3). DeepSeek-V3 targets use Multi-head Latent Attention (MLA): queries and keys are produced through low-rank q_lora / kv_lora projections, the rotary part of the head is a separate qk_rope_head_dim slice rotated with DeepSeek’s interleaved RoPE, and the value head dimension (v_head_dim) differs from the query/key one (qk_nope_head_dim + qk_rope_head_dim). A Llama-shaped draft cannot represent that, so DeepSeek targets get this dedicated draft.

This mirrors draft_llama one-to-one for everything EAGLE-3 specific (the [embed, hidden] fused first layer, the cache_hidden = [K_list, V_list] TTT recurrence with per-step rotary phase offset and diagonal-extension attention, project_hidden_states / compute_logits / set_vocab_mapping and the d2t / t2d buffers). Only the attention block is replaced with MLA.

To guarantee the draft’s rotary math matches the target’s exactly, the MLA projection layout and the interleaved RoPE are taken from the onboarded DeepSeek target (components/models/deepseek_v3: the MLA projection structure and rope_utils), not reimplemented.

Scope: EAGLE-3 single fused draft layer, eager attention. Sequence packing is supported via the shared [B, 1, T, T] block-causal mask (see forward); the MLA attention is eager-only, so packing reuses the eager mask path rather than a FlashAttention varlen kernel. P-EAGLE parallel-drafting and flash-attention remain follow-ups (orthogonal to the MLA attention this file adds).

Module Contents

Classes

Name	Description
`DeepseekV3Eagle3DraftModel`	DeepSeek-V3 (MLA) EAGLE-3 draft model.
`Eagle3DeepseekDecoderLayer`	Fused EAGLE-3 first layer: `[embed, hidden]` -> MLA -> MLP (mirrors Eagle3LlamaDecoderLayer).
`Eagle3DeepseekMLAAttention`	MLA self-attention for the DeepSeek EAGLE-3 draft.
`Eagle3DeepseekMLP`	Plain SwiGLU MLP for the draft (the draft is small; it does not replicate the target MoE).
`Eagle3DeepseekModel`	Inner backbone: `embed_tokens`, the `fc` aux-projection, the fused draft layer, and `norm`.

Functions

Name	Description
`_build_causal_mask`	Build a standard `[B, 1, T, T]` additive causal + padding mask for eager attention.
`_resolve_rope_theta`	Read the RoPE base the way the DeepSeek target does.

API

class nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel(
    config: transformers.PretrainedConfig
)

Bases: PreTrainedModel

DeepSeek-V3 (MLA) EAGLE-3 draft model.

The MLA counterpart to LlamaEagle3DraftModel: same public training API (project_hidden_states / embed_input_ids / compute_logits / set_vocab_mapping / forward) and the same d2t / t2d vocab-remap buffers, so the EAGLE-3 trainer and checkpointing are reused unchanged.

base_model_prefix

= 'model'

draft_vocab_size

has_vocab_compression

= self.draft_vocab_size < config.vocab_size

lm_head

model

= Eagle3DeepseekModel(config)

target_hidden_size

nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

Draft logits over the draft vocabulary (applies model.norm unless already normed).

nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel.copy_embeddings_from_target(
    target_embedding: torch.nn.Embedding
) -> None

Seed the draft embedding table from the (possibly FSDP-sharded) target embeddings.

nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel.embed_input_ids(
    input_ids: torch.Tensor
) -> torch.Tensor

Embed input ids with the draft embedding table.

nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel.forward(
    input_ids: torch.Tensor,
    projected_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: typing.Optional[torch.Tensor] = None,
    cache_hidden: typing.Optional[list[list[torch.Tensor]]] = None,
    seq_lens: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

Run one EAGLE-3 TTT draft step (eager attention).

Args (batch B, sequence length T, draft hidden size H): input_ids: [B, T] long token ids. projected_hidden_states: [B, T, H] draft-hidden features (fc output). attention_mask: [B, T] (1 = real, 0 = pad); used on the unpacked path only. position_ids: [B, T] long, or None to default to arange. Packing requires per-document positions (reset to range(doc_len)). cache_hidden: EAGLE-3 TTT cache [K_list, V_list] — [[], []] on the first step, the same list on later steps, or None for a fresh cache. seq_lens: [B, max_docs] long (0-padded per-document lengths that sum to T per row) turns on sequence packing — Block-1 attention becomes document-level block-causal via :func:build_block_causal_additive_mask (eager mask path; the MLA attention has no FA2 varlen kernel). The TTT diagonal block is position-wise and stays document-safe; cross-document supervision is masked by the trainer’s doc_remaining gate, not here.

nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel.project_hidden_states(
    aux_hidden_states: torch.Tensor
) -> torch.Tensor

Project concatenated target aux hidden states to draft hidden size (with optional fc_norm).

nemo_automodel.components.speculative.eagle.draft_deepseek.DeepseekV3Eagle3DraftModel.set_vocab_mapping(
    selected_token_ids: torch.Tensor
) -> None

Populate d2t / t2d from the draft->target id map (offset form vLLM/SGLang expect).

class nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekDecoderLayer(
    config: transformers.PretrainedConfig,
    layer_id: int = 0
)

Bases: Module

Fused EAGLE-3 first layer: [embed, hidden] -> MLA -> MLP (mirrors Eagle3LlamaDecoderLayer).

hidden_norm

input_layernorm

mlp

= Eagle3DeepseekMLP(config)

post_attention_layernorm

self_attn

= Eagle3DeepseekMLAAttention(config)

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekDecoderLayer.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

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

Bases: Module

MLA self-attention for the DeepSeek EAGLE-3 draft.

Driven through the shared cache_hidden = [K_list, V_list] recurrence, exactly like Eagle3LlamaAttention: step_idx = len(K_list) is both the TTT step index and the rotary phase offset; each step appends its rotated K/V and attends over the step-0 keys (full T x T causal) plus one diagonal column per later step. The only difference from the Llama attention is how q / k / v are produced (MLA low-rank projections + interleaved RoPE on the qk_rope_head_dim slice) and that the value head dim differs from the q/k head dim.

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

= getattr(config, 'q_lora_rank', None)

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

v_head_dim

= config.v_head_dim

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLAAttention._apply(
    fn,
    recurse: bool = True
)

Keep the RoPE rope_freqs buffer in fp32 across dtype casts.

precompute_freqs_cis returns fp32 base frequencies, but the EAGLE-3 training build casts the whole draft with a raw .to(dtype=compute_dtype) (train_eagle3.py), and nn.Module.to rounds floating-point buffers — so rope_freqs is downcast to bf16. The rounded frequencies dephase with absolute position (worse with longer context, and a bf16 round-trip cannot be undone by a later fp32 upcast), so the draft’s RoPE drifts from the fp32 target and draft acceptance erodes. Recompute fresh fp32 frequencies after any cast that would round them, keeping the same fp32 guarantee the Llama draft and the DeepSeek target get by recomputing their RoPE cache from config in fp32.

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLAAttention._assemble_qkv(
    q_nope: torch.Tensor,
    q_pe: torch.Tensor,
    kv_latent: torch.Tensor,
    k_pe: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Up-project the KV latent and assemble per-head q[..,qk], k[..,qk], v[..,v] (heads first).

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLAAttention._eager_attention_forward(
    q: torch.Tensor,
    cache_k: list[torch.Tensor],
    cache_v: list[torch.Tensor],
    attention_mask: torch.Tensor,
    step_idx: int,
    batch_size: int,
    seq_len: int
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLAAttention._project_qkv(
    combined_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Return (q_nope, q_pe, kv_latent, k_pe) from the MLA down/up projections.

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLAAttention.forward(
    combined_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    cache_hidden: list[list[torch.Tensor]]
) -> torch.Tensor

class nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLP(
    config: transformers.PretrainedConfig
)

Bases: Module

Plain SwiGLU MLP for the draft (the draft is small; it does not replicate the target MoE).

act_fn

= nn.SiLU()

down_proj

gate_proj

up_proj

nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.speculative.eagle.draft_deepseek.Eagle3DeepseekModel(
    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

nemo_automodel.components.speculative.eagle.draft_deepseek._build_causal_mask(
    attention_mask: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor

Build a standard [B, 1, T, T] additive causal + padding mask for eager attention.

nemo_automodel.components.speculative.eagle.draft_deepseek._resolve_rope_theta(
    config: transformers.PretrainedConfig
) -> float

Read the RoPE base the way the DeepSeek target does.

transformers 5.x moved rope_theta under config.rope_parameters and dropped the top-level attribute (DeepseekV3Config(rope_theta=...).rope_theta raises AttributeError), so a bare getattr(config, "rope_theta", 10000.0) always returns the fallback and the draft’s rotary phase silently diverges from any target whose base is not 10000. The top-level read is kept only as a fallback for older configs.