> 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_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`](#nemo_automodel-components-speculative-eagle-draft_deepseek-DeepseekV3Eagle3DraftModel) | DeepSeek-V3 (MLA) EAGLE-3 draft model.                                                        |
| [`Eagle3DeepseekDecoderLayer`](#nemo_automodel-components-speculative-eagle-draft_deepseek-Eagle3DeepseekDecoderLayer) | Fused EAGLE-3 first layer: `[embed, hidden]` -> MLA -> MLP (mirrors Eagle3LlamaDecoderLayer). |
| [`Eagle3DeepseekMLAAttention`](#nemo_automodel-components-speculative-eagle-draft_deepseek-Eagle3DeepseekMLAAttention) | MLA self-attention for the DeepSeek EAGLE-3 draft.                                            |
| [`Eagle3DeepseekMLP`](#nemo_automodel-components-speculative-eagle-draft_deepseek-Eagle3DeepseekMLP)                   | Plain SwiGLU MLP for the draft (the draft is small; it does not replicate the target MoE).    |
| [`Eagle3DeepseekModel`](#nemo_automodel-components-speculative-eagle-draft_deepseek-Eagle3DeepseekModel)               | Inner backbone: `embed_tokens`, the `fc` aux-projection, the fused draft layer, and `norm`.   |

### Functions

| Name                                                                                                     | Description                                                                         |
| -------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`_build_causal_mask`](#nemo_automodel-components-speculative-eagle-draft_deepseek-_build_causal_mask)   | Build a standard `[B, 1, T, T]` additive causal + padding mask for eager attention. |
| [`_resolve_rope_theta`](#nemo_automodel-components-speculative-eagle-draft_deepseek-_resolve_rope_theta) | Read the RoPE base the way the DeepSeek target does.                                |

### API

```python
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.

```python
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).

```python
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.

```python
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.

```python
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.

```python
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).

```python
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).

```python
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).

```python
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
```

```python
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.

```python
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.

```python
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).

```python
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
```

```python
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.

```python
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
```

```python
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).

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

```python
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`.

```python
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.

```python
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.