> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

EAGLE-3 training components.

Only the EAGLE-3-native model and trainer classes live here. Generic
helpers were moved to the project's canonical locations:

* `masked_soft_cross_entropy` -> `nemo_automodel.components.loss.soft_ce`
* `build_eagle3_dataloader` / `build_eagle3_token_mapping`
  -> `nemo_automodel.components.datasets.llm.eagle3`

Import them directly from those modules.

## Subpackages

* **[`nemo_automodel.components.speculative.eagle.remote`](/nemo-automodel/nemo_automodel/components/speculative/eagle/remote)**

## Submodules

* **[`nemo_automodel.components.speculative.eagle.backend`](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend)**
* **[`nemo_automodel.components.speculative.eagle.core`](/nemo-automodel/nemo_automodel/components/speculative/eagle/core)**
* **[`nemo_automodel.components.speculative.eagle.core_v12`](/nemo-automodel/nemo_automodel/components/speculative/eagle/core_v12)**
* **[`nemo_automodel.components.speculative.eagle.draft_deepseek`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_deepseek)**
* **[`nemo_automodel.components.speculative.eagle.draft_gemma`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_gemma)**
* **[`nemo_automodel.components.speculative.eagle.draft_gpt_oss`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_gpt_oss)**
* **[`nemo_automodel.components.speculative.eagle.draft_kimi_k3`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_kimi_k3)**
* **[`nemo_automodel.components.speculative.eagle.draft_llama`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama)**
* **[`nemo_automodel.components.speculative.eagle.draft_llama_v12`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12)**
* **[`nemo_automodel.components.speculative.eagle.msd`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd)**
* **[`nemo_automodel.components.speculative.eagle.msd_curriculum`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_curriculum)**
* **[`nemo_automodel.components.speculative.eagle.msd_decode`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_decode)**
* **[`nemo_automodel.components.speculative.eagle.msd_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_target)**
* **[`nemo_automodel.components.speculative.eagle.peagle_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_attention)**
* **[`nemo_automodel.components.speculative.eagle.peagle_data`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_data)**
* **[`nemo_automodel.components.speculative.eagle.peagle_draft`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_draft)**
* **[`nemo_automodel.components.speculative.eagle.peagle_trainer`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_trainer)**
* **[`nemo_automodel.components.speculative.eagle.registry`](/nemo-automodel/nemo_automodel/components/speculative/eagle/registry)**
* **[`nemo_automodel.components.speculative.eagle.ring_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/ring_attention)**
* **[`nemo_automodel.components.speculative.eagle.sglang_runner`](/nemo-automodel/nemo_automodel/components/speculative/eagle/sglang_runner)**
* **[`nemo_automodel.components.speculative.eagle.sglang_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/sglang_target)**
* **[`nemo_automodel.components.speculative.eagle.target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/target)**
* **[`nemo_automodel.components.speculative.eagle.target_runner`](/nemo-automodel/nemo_automodel/components/speculative/eagle/target_runner)**
* **[`nemo_automodel.components.speculative.eagle.target_v12`](/nemo-automodel/nemo_automodel/components/speculative/eagle/target_v12)**
* **[`nemo_automodel.components.speculative.eagle.ulysses_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/ulysses_attention)**
* **[`nemo_automodel.components.speculative.eagle.vispec_core`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_core)**
* **[`nemo_automodel.components.speculative.eagle.vispec_draft`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_draft)**
* **[`nemo_automodel.components.speculative.eagle.vispec_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_target)**
* **[`nemo_automodel.components.speculative.eagle.vllm_runner`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vllm_runner)**
* **[`nemo_automodel.components.speculative.eagle.vllm_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vllm_target)**
* **[`nemo_automodel.components.speculative.eagle.zigzag_ring_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/zigzag_ring_attention)**

## Package Contents

### Classes

| Name                                                                                                            | Description                                                                      |
| --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`DraftSpec`](#nemo_automodel-components-speculative-eagle-registry-DraftSpec)                                  | How to build an EAGLE draft model for a particular target architecture.          |
| [`Eagle3TargetBackend`](#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend)               | Abstract contract every EAGLE-3 target-model backend implements.                 |
| [`Eagle3TrainerModule`](#nemo_automodel-components-speculative-eagle-core-Eagle3TrainerModule)                  | Draft-side EAGLE-3 trainer module with test-time-training unroll.                |
| [`EagleTrainerModule`](#nemo_automodel-components-speculative-eagle-core_v12-EagleTrainerModule)                | Draft-side trainer for EAGLE-1 / EAGLE-2 hidden-state prediction.                |
| [`GptOssEagle3DraftModel`](#nemo_automodel-components-speculative-eagle-draft_gpt_oss-GptOssEagle3DraftModel)   | EAGLE-3 draft model for gpt-oss targets.                                         |
| [`HFEagle3TargetModel`](#nemo_automodel-components-speculative-eagle-target-HFEagle3TargetModel)                | Co-located backend that captures three auxiliary hidden states from a causal LM. |
| [`HFEagleTargetModel`](#nemo_automodel-components-speculative-eagle-target_v12-HFEagleTargetModel)              | Thin wrapper that exposes hidden-state supervision from a causal LM.             |
| [`HFMSDTargetModel`](#nemo_automodel-components-speculative-eagle-msd_target-HFMSDTargetModel)                  | Expose image-aware feature supervision from a frozen Hugging Face VLM.           |
| [`LlamaEagle3DraftModel`](#nemo_automodel-components-speculative-eagle-draft_llama-LlamaEagle3DraftModel)       | Llama-style dense EAGLE-3 draft model (Llama, Phi-3, Qwen3).                     |
| [`LlamaEagleDraftModel`](#nemo_automodel-components-speculative-eagle-draft_llama_v12-LlamaEagleDraftModel)     | Llama-style dense draft that predicts next-step hidden states.                   |
| [`MSDCurriculum`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDCurriculum)                    | Stateless-by-step two-stage curriculum with checkpointable progress.             |
| [`MSDCurriculumConfig`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDCurriculumConfig)        | Configuration for the two-stage MSD data curriculum.                             |
| [`MSDCurriculumPhase`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDCurriculumPhase)          | High-level phase of the two-stage MSD training schedule.                         |
| [`MSDDataSource`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDDataSource)                    | Data source selected for one MSD optimizer step.                                 |
| [`MSDGreedyDecoder`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDGreedyDecoder)                  | Batch-one reference MSD round with recursive drafting and tree verification.     |
| [`MSDTrainerModule`](#nemo_automodel-components-speculative-eagle-msd-MSDTrainerModule)                         | Draft-side training module for multimodal speculative decoding.                  |
| [`MSDTreeDraftGenerator`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeDraftGenerator)        | Recursively draft a top-k MSD feature tree from target prefix features.          |
| [`MSDTreeLayout`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeLayout)                        | Tree attention and retrieval metadata for an MSD proposal.                       |
| [`MSDTreeNode`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeNode)                            | One drafted text-token node in an MSD candidate tree.                            |
| [`MSDTreeProposal`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeProposal)                    | A target root token plus recursively drafted candidate-tree nodes.               |
| [`MSDVerificationResult`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDVerificationResult)        | Greedy posterior verification result for one candidate tree.                     |
| [`MultimodalEagleDraftModel`](#nemo_automodel-components-speculative-eagle-msd-MultimodalEagleDraftModel)       | EAGLE-1/2 draft model with modality-aware input feature construction.            |
| [`PEagleTrainerModule`](#nemo_automodel-components-speculative-eagle-peagle_trainer-PEagleTrainerModule)        | Draft-side P-EAGLE (parallel-drafting EAGLE-3) trainer module.                   |
| [`RunnerEagle3TargetModel`](#nemo_automodel-components-speculative-eagle-target_runner-RunnerEagle3TargetModel) | EAGLE-3 target backend that runs the frozen target through a runner.             |
| [`TargetRunner`](#nemo_automodel-components-speculative-eagle-target_runner-TargetRunner)                       | Minimal engine surface `RunnerEagle3TargetModel` depends on.                     |

### Functions

| Name                                                                                                           | Description                                                                           |
| -------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`accept_or_resample`](#nemo_automodel-components-speculative-eagle-msd_decode-accept_or_resample)             | Apply the lossless speculative acceptance rule for one candidate token.               |
| [`build_msd_tree_layout`](#nemo_automodel-components-speculative-eagle-msd_decode-build_msd_tree_layout)       | Build engine-ready tree attention and root-to-leaf retrieval metadata.                |
| [`resolve_eagle1_draft_spec`](#nemo_automodel-components-speculative-eagle-registry-resolve_eagle1_draft_spec) | Resolve the EAGLE-1 / EAGLE-2 draft spec for a target's `config.architectures` field. |
| [`resolve_eagle3_draft_spec`](#nemo_automodel-components-speculative-eagle-registry-resolve_eagle3_draft_spec) | Resolve the EAGLE-3 draft spec for a target's `config.architectures` field.           |
| [`simulated_accept_length`](#nemo_automodel-components-speculative-eagle-core-simulated_accept_length)         | Expected accepted tokens per speculative round, from prefix-hit counts.               |
| [`verify_greedy_tree`](#nemo_automodel-components-speculative-eagle-msd_decode-verify_greedy_tree)             | Select the longest target-greedy accepted path and its bonus token.                   |
| [`verify_hf_greedy_tree`](#nemo_automodel-components-speculative-eagle-msd_decode-verify_hf_greedy_tree)       | Reference target verification for a batch-one Hugging Face VLM.                       |
| [`verify_stochastic_chain`](#nemo_automodel-components-speculative-eagle-msd_decode-verify_stochastic_chain)   | Losslessly verify a linear draft chain and emit the correction token.                 |

### Data

[`EAGLE1_DRAFT_REGISTRY`](#nemo_automodel-components-speculative-eagle-registry-EAGLE1_DRAFT_REGISTRY)

[`EAGLE3_DRAFT_REGISTRY`](#nemo_automodel-components-speculative-eagle-registry-EAGLE3_DRAFT_REGISTRY)

### API

```python
class nemo_automodel.components.speculative.eagle.registry.DraftSpec(
    draft_cls: type[transformers.PreTrainedModel]
)
```

Dataclass

How to build an EAGLE draft model for a particular target architecture.

**`draft_cls`** `type[PreTrainedModel]`

---

```python
class nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend()
```

Abstract

Abstract contract every EAGLE-3 target-model backend implements.

Two supervision encodings are allowed, both consumed directly by
`Eagle3TrainerModule.forward`:

* **full logits** -- `Eagle3TargetBatch.logits` carries the target's
  full-vocab logits and the draft-vocab projection happens trainer-side.
  Cheap when co-located (the tensor never leaves the GPU); impractical to
  ship over a wire because it is full-vocab sized.
* **precomputed** -- `Eagle3TargetBatch.target_probs` and
  `Eagle3TargetBatch.position_mask` carry the already-projected
  draft-vocab distribution, so a backend that computes them itself (e.g. a
  remote server) only has to transfer draft-vocab-sized tensors.

A backend returns exactly one of the two encodings from
`generate_batch`; the recipe forwards whichever is present.

**`supports_async`** `bool`

Whether `generate_batch_async` is implemented (prefetch-capable).

---

```python
nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.close() -> None
```

Release backend resources (remote connections, server handles).

```python
nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor
) -> 'Eagle3TargetBatch'
```

abstract

Run the target and return the supervision for one training batch.

```python
nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.generate_batch_async(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor
)
```

Submit an asynchronous `generate_batch` for prefetch pipelining.

Only backends that overlap target inference with draft training
implement this; the default signals a synchronous backend so callers
fall back to `generate_batch`.

```python
nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.get_input_embeddings() -> torch.nn.Module
```

abstract

Return the target input-embedding module (used to seed the draft).

```python
nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.set_vocab_mapping(
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor
) -> None
```

Provide the draft-vocab mapping needed to precompute supervision.

Co-located backends keep the mapping on the trainer module and derive
the distribution there, so the default is a no-op. A backend that
computes `target_probs` itself overrides this to receive the mapping.

```python
class nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule(
    draft_model: torch.nn.Module,
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor,
    ttt_steps: int,
    cp_group = None,
    cp_zigzag: bool = False,
    lk_loss_type: str | None = None,
    lk_kl_scale: float = 1.0,
    lk_kl_decay: float = 3.0
)
```

**Bases:** `Module`

Draft-side EAGLE-3 trainer module with test-time-training unroll.

**`lk_kl_decay`** `= float(lk_kl_decay)`

---

**`lk_kl_scale`** `= float(lk_kl_scale)`

---

```python
nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule._lk_step_loss(
    logits: torch.Tensor,
    target_probs: torch.Tensor,
    position_mask: torch.Tensor
) -> torch.Tensor
```

One TTT step of the LK acceptance-rate loss (arXiv:2602.23881).

The per-token expected acceptance under speculative sampling is the
min-overlap of the two distributions, `alpha = sum_v min(p_target, p_draft)`,
computed over the draft vocabulary (both inputs already live there).

* `"alpha"`: the pure acceptance likelihood `-mean(log alpha)`; the log
  is taken per token before the masked mean, with zero-acceptance tokens
  contributing zero (they carry no usable gradient direction).
* `"lambda"`: the adaptive hybrid
  `kl_weight * soft_ce + (1 - kl_weight) * (1 - mean(alpha))` with
  `kl_weight = lk_kl_scale * exp(-lk_kl_decay * mean(alpha).detach())`;
  the weight is detached so gradients flow only through the two terms, and
  training shifts from distillation toward direct acceptance as the draft
  improves. A step with zero supervised positions returns 0, matching the
  soft-CE path.

Under CP every masked mean is renormalized over the full sequence via
`_cp_global_step_loss`, so the mixing weight and the loss match the
single-rank run.

**Parameters:**

**`logits`** `torch.Tensor`

Tensor of shape `[batch, sequence, draft_vocab]`; draft logits.

---

**`target_probs`** `torch.Tensor`

Tensor of shape `[batch, sequence, draft_vocab]`; target
distribution restricted to the draft vocabulary.

---

**`position_mask`** `torch.Tensor`

Bool tensor of shape `[batch, sequence, 1]`; True where
the position is supervised this step.

---

**Returns:** `torch.Tensor`

Scalar loss tensor for this TTT step.

```python
nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    aux_hidden_states: torch.Tensor,
    target_logits: torch.Tensor | None = None,
    target_probs: torch.Tensor | None = None,
    position_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics
```

Run the EAGLE-3 unrolled draft loss for one batch.

The attention layer is driven through a shared `cache_hidden`
list so each TTT step can attend to the K/V branches produced by
every previous step at the same position. This matches the
SpecForge `llama3_eagle.py` recurrence; without it, multi-step
TTT would degenerate into `ttt_steps` independent single-step
passes and the draft would never learn the multi-step
distribution it sees at deployment time.

`attention_mask` is held constant across TTT steps -- only
`input_ids` / `loss_mask` / `position_mask` /
`target_probs` roll forward by one position per step.

Packing: `position_ids` / `seq_lens` make the draft's Block-1 attention
document-level block-causal, and `doc_remaining` gates supervision per
step (slot `t` valid at step `k` only while `k < doc_remaining[t]`),
masking every cross-document TTT prediction.

Two supervision sources are accepted: the live path passes the
target's full-vocab `target_logits` and the draft distribution is
derived here; the offline-cache path (`precompute_eagle3`) passes the
already-derived `target_probs` (over the draft vocab) and
`position_mask` directly, so the full-vocab logits never have to be
stored. Provide exactly one of the two.

```python
class nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    hidden_loss_weight: float = 1.0,
    token_loss_weight: float = 0.1,
    feature_noise_config: nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig | None = None,
    rank_loss_weight: float = 0.0,
    rank_loss_topk: int = 10
)
```

**Bases:** `Module`

Draft-side trainer for EAGLE-1 / EAGLE-2 hidden-state prediction.

**`hidden_loss_fn`** `= nn.SmoothL1Loss(reduction='none')`

---

```python
nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Project predicted hidden states through the frozen target lm\_head.

```python
nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.eagle.core_v12.EagleStepMetrics
```

Run one EAGLE-1 / EAGLE-2 training step.

Per-token tensors are `[B, T]` (`position_ids` / `doc_remaining`
included) except the `[B, T, H]` hidden states and `[B, T, V]`
`target_logits`; `seq_lens` is `[B, max_docs]`. When packing is on,
`position_ids` / `seq_lens` make the draft block-causal and per-document,
and `doc_remaining` (real tokens after each slot within its document) gates
supervision: a document's last real token (`doc_remaining == 0`) is dropped
because the wrapper's global left-shift makes its target the next document's
first token.

```python
class nemo_automodel.components.speculative.eagle.draft_gpt_oss.GptOssEagle3DraftModel(
    config: transformers.PretrainedConfig
)
```

**Bases:** [LlamaEagle3DraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama#nemo_automodel-components-speculative-eagle-draft_llama-LlamaEagle3DraftModel)

EAGLE-3 draft model for gpt-oss targets.

Identical to `LlamaEagle3DraftModel` except that the single draft
layer's rotary embedding is replaced with `GptOssDraftRotaryEmbedding`
so the draft reproduces gpt-oss's YaRN RoPE instead of the YaRN-incapable
`LlamaRotaryEmbedding`.

```python
class nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel(
    model: torch.nn.Module,
    aux_layer_ids: typing.Sequence[int] | None = None,
    cp_mesh = None
)
```

**Bases:** [Eagle3TargetBackend](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend)

Co-located backend that captures three auxiliary hidden states from a causal LM.

**`_cp_size`** `= cp_mesh.size() if cp_mesh is not None else 1`

---

**`aux_layer_ids`** `= self._validate_aux_layer_ids(candidate_ids)`

---

**`model`** `= model.eval()`

---

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._check_captured(
    captured: dict[int, torch.Tensor]
) -> None
```

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._default_aux_layer_ids() -> list[int]
```

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._get_transformer_layers() -> list[torch.nn.Module]
```

Return decoder layers as an ordered list indexable by integer.

Supports both the HuggingFace layouts (where `layers` is a
`ModuleList`) and AutoModel's custom-impl layouts (where
`layers` is a `ModuleDict` keyed by `str(i)`). Returning a
plain list normalizes the access pattern for downstream
`register_forward_hook` calls.

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._num_hidden_layers() -> int
```

Return the target's decoder depth.

Multimodal targets (e.g. `Gemma4ForConditionalGeneration`) carry no
top-level `num_hidden_layers`; the text backbone's depth lives on
`config.text_config`. Text-only targets (and the lightweight config
stubs used in tests) expose `num_hidden_layers` directly.

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._validate_aux_layer_ids(
    aux_layer_ids: typing.Sequence[int]
) -> list[int]
```

Validate aux-layer selection before any forward hooks are registered.

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch
```

Run the target model and capture aux hidden states plus logits.

With `seq_lens` (packing), the target runs with a `[B, 1, T, T]`
block-causal mask and per-document `position_ids` so its outputs respect
document boundaries; the packing metadata is forwarded unshifted to the
trainer. `seq_lens=None` keeps the original 2D-mask path.

```python
nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel.get_input_embeddings() -> torch.nn.Embedding
```

Return the target model input embeddings.

```python
class nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel(
    model: torch.nn.Module
)
```

Thin wrapper that exposes hidden-state supervision from a causal LM.

**`model`** `= model.eval()`

---

```python
nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.eagle.target_v12.EagleTargetBatch
```

Run the target transformer and prepare shifted supervision tensors.

All per-token inputs are `[B, T]`. With `seq_lens` (`[B, max_docs]`
long, per-document lengths summing to `T`) the target runs with a
document-level block-causal mask and per-document `position_ids` so its
hidden states do not leak across document boundaries; SDPA/eager targets
consume the `[B, 1, T, T]` block-causal additive mask, FlashAttention
targets infer document boundaries from `position_ids` and are passed
`attention_mask=None` (batch size 1 only). `position_ids` / `seq_lens`
/ `doc_remaining` are carried through (unshifted) so the trainer can build
the draft's block-causal mask and drop cross-document supervision.

```python
nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.get_input_embeddings() -> torch.nn.Embedding
```

Return the target model input embeddings.

```python
nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.get_lm_head() -> torch.nn.Module
```

Return the target model lm\_head.

```python
class nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel(
    model: torch.nn.Module
)
```

Expose image-aware feature supervision from a frozen Hugging Face VLM.

**`image_token_id`** `= int(image_token_id)`

---

**`language_backbone`**

---

**`model`** `= model.eval()`

---

```python
nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel.generate_batch(
    loss_mask: torch.Tensor,
    model_inputs: dict[str, torch.Tensor]
) -> nemo_automodel.components.speculative.eagle.msd_target.MSDTargetBatch
```

Run a VLM and construct shifted multimodal draft supervision.

`input_ids` and `attention_mask` are read from `model_inputs` rather
than taken separately, so the tensors the target actually runs on are the
same ones the image mask and the returned attention mask are derived from.
A second copy would misalign the draft's modality routing without raising.

**Parameters:**

**`loss_mask`** `torch.Tensor`

Bool tensor of shape \[batch, sequence], already aligned to
next-token logits. VLM collators produce this from `labels != -100`;
they emit `labels[:, 1:]` against `input_ids[:, :-1]`, so no
further shift is applied here. Its final position is always
dropped because the shifted supervision tensors zero-fill there.

---

**`model_inputs`** `dict[str, torch.Tensor]`

Mapping of VLM forward inputs. It must include
`input_ids` of shape \[batch, sequence], with image-token
placeholders at positions that receive projected vision features,
and `attention_mask` of the same shape, where one denotes a real
token and zero denotes padding. It also carries any vision tensors
such as `pixel_values` and `image_grid_thw`.

---

**Returns:** `MSDTargetBatch`

MSDTargetBatch with shifted image-aware embeddings, target hidden

```python
nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel.get_lm_head() -> torch.nn.Module
```

Return the frozen VLM language-model head used by the draft loss.

```python
class nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel(
    config: transformers.PretrainedConfig
)
```

**Bases:** [\_PeagleDraftMixin](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_draft#nemo_automodel-components-speculative-eagle-peagle_draft-_PeagleDraftMixin), `PreTrainedModel`

Llama-style dense EAGLE-3 draft model (Llama, Phi-3, Qwen3).

State dict keys match SGLang's `LlamaForCausalLMEagle3` so the saved
checkpoint can be loaded by SGLang's inference engine without any
remapping (SGLang's `load_weights` fuses `q/k/v_proj` into
`qkv_proj` and `gate/up_proj` into `gate_up_proj` via its
standard `stacked_params_mapping`).

The class name is retained for checkpoint-architectures compatibility; the
implementation is config-driven and works for any HF dense decoder-only
config that exposes `hidden_size`, `num_attention_heads`,
`num_key_value_heads`, `attention_bias`, `mlp_bias`, `rope_theta`,
and `rms_norm_eps`. A decoupled `head_dim` is read via
`getattr(config, "head_dim", ...)` in the attention layer.

Scope:

* single draft decoder layer
* no KV-cache optimization
* no speculative runtime integration

**`base_model_prefix`** `= 'model'`

---

**`draft_vocab_size`**

---

**`has_vocab_compression`** `= self.draft_vocab_size < config.vocab_size`

---

**`lm_head`**

---

**`model`** `= Eagle3LlamaModel(config)`

---

**`target_hidden_size`**

---

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Compute draft logits on the configured draft vocabulary.

With `config.norm_output` unset (EAGLE-3 default) the input is the
raw decoder-layer output and the final `model.norm` is applied
here. With `config.norm_output` set (EAGLE-3.1) `forward` has
already returned the post-norm state, so `lm_head` is applied
directly to avoid a double normalisation.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.copy_embeddings_from_target(
    target_embedding: torch.nn.Embedding
) -> None
```

Initialize draft embeddings from the target model embeddings.

When the target model is wrapped with FSDP2, `target_embedding.weight`
is a `DTensor` sharded across ranks.  The draft embedding is a plain
`nn.Parameter` (the draft is not FSDP-wrapped), so a direct
`copy_` of a DTensor into a regular tensor raises a mixed-type
distributed-operator error.  Gather to a full local tensor first.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.embed_input_ids(
    input_ids: torch.Tensor
) -> torch.Tensor
```

Embed input ids with the draft embedding table.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.forward(
    input_ids: torch.Tensor,
    projected_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    cache_hidden: list[list[torch.Tensor]] | None = None,
    seq_lens: torch.Tensor | None = None
) -> torch.Tensor
```

Run one full-sequence draft update step.

`cache_hidden` is the EAGLE-3 TTT cache. Pass `[[], []]` on
the first step of a TTT unroll and the same list object on each
subsequent step; the attention layer appends the per-step K and V
to it. If `None` is passed (e.g. from a one-shot evaluation
call) a fresh `[[], []]` is allocated locally -- step 0 of TTT
is mathematically equivalent to a plain causal forward.

`seq_lens` (packing) makes Block-1 attention document-level block-causal
(eager mask / FA2 varlen); callers must pass per-document `position_ids`.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.freeze_embeddings() -> None
```

Freeze draft input embeddings.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.gradient_checkpointing_disable() -> None
```

Disable activation checkpointing for the P-EAGLE draft layers.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs = None
) -> None
```

Enable activation checkpointing for the P-EAGLE draft layers.

Training-only memory knob: recomputes each `forward_peagle` layer in the
backward instead of storing its activations (the EAGLE-3 TTT `forward`
path is unaffected). `gradient_checkpointing_kwargs` is accepted for
HF-API parity but ignored -- recompute is always non-reentrant, the only
mode compatible with the non-tensor `block_mask`.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.project_hidden_states(
    aux_hidden_states: torch.Tensor
) -> torch.Tensor
```

Project concatenated target aux states from `num_aux * H_target` to draft hidden size.

When `config.fc_norm` is set (EAGLE-3.1), the input is split into
`num_aux_hidden_states` equal chunks along the last dim and each
chunk is passed through its own RMSNorm in `model.fc_norm` (the
modules are independent, matching vLLM's upstream implementation).
The normalized chunks are then re-concatenated and fed to `fc`,
stabilising the per-aux-state scale before the projection mixes them
and removing the speculation-depth drift observed with raw inputs.

```python
nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.set_vocab_mapping(
    selected_token_ids: torch.Tensor
) -> None
```

Populate the `d2t` / `t2d` vocab-remap buffers from the draft->target id map.

`selected_token_ids` has shape `[draft_vocab_size]`; entry `i` is
the *target* vocab id of draft id `i` (the frequency-pruned mapping
built by `build_eagle3_token_mapping`). This writes the two buffers
inference engines consume:

* `d2t[i] = selected_token_ids[i] - i` -- the offset form vLLM expects
  (`target_id = draft_id + d2t[draft_id]`);
* `t2d[target_id] = True` for every selected target id -- the boolean
  presence mask SGLang consumes.

These must be in the saved checkpoint: without them vLLM/SGLang find no
mapping, silently align draft ids to the first `draft_vocab_size`
target ids, and acceptance rate collapses.

No-op when the draft vocab is not compressed (the buffers do not exist
and the draft logits are already in target space).

```python
class nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel(
    config: transformers.PretrainedConfig
)
```

**Bases:** `PreTrainedModel`

Llama-style dense draft that predicts next-step hidden states.

Works with Llama, Phi-3, and Qwen3 dense configs. The class name is
retained for backward compatibility with already-trained checkpoints.

**`embed_tokens`**

---

**`fc`**

---

**`layers`**

---

**`main_input_name`** `= 'input_ids'`

---

**`norm`**

---

```python
nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel.copy_embeddings_from_target(
    target_embeddings: torch.nn.Embedding
) -> None
```

Copy the target model token embeddings into the draft embeddings.

When the target is wrapped with FSDP2, its `embed_tokens.weight` is
a `DTensor` sharded across ranks.  Gather to a local full tensor
before copying into the (unsharded) draft parameter -- otherwise
`aten.copy_` raises a mixed Tensor/DTensor error.

```python
nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel.forward(
    input_ids: torch.Tensor,
    target_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None
) -> torch.Tensor
```

Predict the next-step target hidden state for each position.

Args (batch `B`, sequence length `T`, hidden size `H`):
input\_ids: `[B, T]` long token ids.
target\_hidden\_states: `[B, T, H]` target hidden features fed to the draft.
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)`).
seq\_lens: `[B, max_docs]` long (0-padded per-document lengths that sum to
`T` per row) turns on sequence packing -- attention becomes
document-level block-causal via `build_block_causal_additive_mask`
instead of the plain causal + padding mask, and RoPE uses the per-document
`position_ids` so each document is rotated from its own origin.

```python
nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel.freeze_embeddings() -> None
```

Freeze draft token embeddings.

```python
class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum(
    config: nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig,
    global_step: int = 0
)
```

Stateless-by-step two-stage curriculum with checkpointable progress.

Call `source_for_step` from a training loop that owns separate text
and multimodal dataloaders, then draw the selected dataloader's next batch.
`next_source` is a convenience for single-process loops and stores
the next global step in its checkpoint state.

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.load_state_dict(
    state_dict: dict[str, typing.Any]
) -> None
```

Restore curriculum progress from a checkpoint state dictionary.

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.multimodal_probability(
    global_step: int
) -> float
```

Return the multimodal batch probability at `global_step`.

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.next_source() -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource
```

Return the current source selection and advance one optimizer step.

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.phase_for_step(
    global_step: int
) -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumPhase
```

Return the curriculum phase active at `global_step`.

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.source_for_step(
    global_step: int
) -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource
```

Deterministically choose the data source for `global_step`.

The random draw is seeded from the global step rather than mutable RNG
state. Resumed and distributed runs therefore make the same sampling
decision for a given optimizer step.

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.state_dict() -> dict[str, typing.Any]
```

Return progress required to resume `next_source`.

```python
class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig(
    text_only_steps: int,
    multimodal_ramp_steps: int,
    final_multimodal_ratio: float = 1.0,
    seed: int = 42
)
```

Dataclass

Configuration for the two-stage MSD data curriculum.

The first stage consumes text-only instruction data. In the second stage,
the probability of selecting a multimodal batch rises linearly to
`final_multimodal_ratio`. A final ratio of one makes the post-ramp phase
entirely multimodal, matching the curriculum described in the MSD paper.

**Parameters:**

**`text_only_steps`** `int`

Number of initial optimizer steps using only text
batches.

---

**`multimodal_ramp_steps`** `int`

Number of stage-two steps over which the
multimodal sampling probability rises to its final value. Set to
zero to switch immediately after the text-only stage.

---

**`final_multimodal_ratio`** `float` — default: 1.0

Multimodal batch probability after the ramp.

---

**`seed`** `int` — default: 42

Seed used for stateless, step-indexed source selection. All
distributed ranks therefore select the same source for a global
optimizer step without communicating RNG state.

---

**`final_multimodal_ratio`** `float = 1.0`

---

**`multimodal_ramp_steps`** `int`

---

**`seed`** `int = 42`

---

**`text_only_steps`** `int`

---

```python
nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig.__post_init__() -> None
```

```python
class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumPhase
```

**Bases:** `enum.Enum`

High-level phase of the two-stage MSD training schedule.

**`MULTIMODAL`** `= 'multimodal'`

---

**`MULTIMODAL_RAMP`** `= 'multimodal_ramp'`

---

**`TEXT_ONLY`** `= 'text_only'`

---

```python
class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource
```

**Bases:** `enum.Enum`

Data source selected for one MSD optimizer step.

**`MULTIMODAL`** `= 'multimodal'`

---

**`TEXT`** `= 'text'`

---

```python
class nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder(
    target: nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel,
    draft_model: torch.nn.Module
)
```

Batch-one reference MSD round with recursive drafting and tree verification.

**`generator`**

---

**`target_embeddings`** `= self._resolve_target_embeddings(target.model)`

---

```python
nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder._resolve_target_embeddings(
    model: torch.nn.Module
) -> torch.nn.Module
```

staticmethod

Find the language token embedding module across common HF VLM layouts.

```python
nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder.decode_round(
    model_inputs: dict[str, torch.Tensor],
    draft_steps: int,
    top_k: int,
    beam_width: int
) -> tuple[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal, nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult]
```

Draft and greedily verify one MSD candidate tree for a VLM prompt.

```python
class nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    hidden_loss_weight: float = 1.0,
    token_loss_weight: float = 0.1
)
```

**Bases:** `Module`

Draft-side training module for multimodal speculative decoding.

**`hidden_loss_fn`** `= nn.SmoothL1Loss(reduction='none')`

---

```python
nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Project predicted hidden states through the frozen target language head.

**Parameters:**

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

Tensor of shape \[..., hidden], with arbitrary leading
dimensions.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., vocab], containing target-vocabulary logits.

```python
nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule.forward(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    image_mask: torch.Tensor
) -> nemo_automodel.components.speculative.eagle.msd.MSDStepMetrics
```

Run one multimodal speculative-decoding training step.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], the target
VLM input embeddings shifted left by one position.

---

**`attention_mask`** `torch.Tensor`

Tensor of shape \[batch, sequence], where one denotes
a real token and zero denotes padding.

---

**`loss_mask`** `torch.Tensor`

Bool tensor of shape \[batch, sequence], true where the
shifted next-token target is supervised.

---

**`input_hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], the
target's unshifted final hidden states used as draft features.

---

**`target_hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], the
target's final hidden states shifted left by one position.

---

**`target_logits`** `torch.Tensor`

Tensor of shape \[batch, sequence, vocab], the frozen
target logits shifted left by one position.

---

**`image_mask`** `torch.Tensor`

Bool tensor of shape \[batch, sequence], true where
`inputs_embeds` contains a projected image token.

---

**Returns:** `MSDStepMetrics`

MSDStepMetrics containing scalar loss terms, first-token accuracy, and

```python
class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeDraftGenerator(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    target_embeddings: torch.nn.Module
)
```

Recursively draft a top-k MSD feature tree from target prefix features.

```python
nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeDraftGenerator.propose(
    shifted_inputs_embeds: torch.Tensor,
    input_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    shifted_image_mask: torch.Tensor,
    root_token_id: int,
    draft_steps: int,
    top_k: int,
    beam_width: int
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal
```

Create a feature-level candidate tree after a target-sampled root.

`shifted_inputs_embeds` and `shifted_image_mask` use the same
next-token alignment as `HFMSDTargetModel`. The final shifted
embedding is replaced with the target root token embedding, then each
recursive draft prediction is appended as the feature context for its
child nodes. `attention_mask` must be right-padded, since the prefix is
taken as its leading active slice.

```python
class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout(
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    retrieve_indices: torch.Tensor
)
```

Dataclass

Tree attention and retrieval metadata for an MSD proposal.

Node zero represents the target-sampled root token. `attention_mask` is
boolean and contains the root plus every ancestor for each node. Serving
backends turn its false entries into additive `-inf` attention bias.

**`attention_mask`** `Tensor`

---

**`position_ids`** `Tensor`

---

**`retrieve_indices`** `Tensor`

---

```python
class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode(
    index: int,
    parent_index: int,
    token_id: int,
    depth: int,
    log_probability: float
)
```

Dataclass

One drafted text-token node in an MSD candidate tree.

**`depth`** `int`

---

**`index`** `int`

---

**`log_probability`** `float`

---

**`parent_index`** `int`

---

**`token_id`** `int`

---

```python
class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal(
    root_token_id: int,
    nodes: tuple[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode, ...],
    leaf_indices: tuple[int, ...],
    layout: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout
)
```

Dataclass

A target root token plus recursively drafted candidate-tree nodes.

**`layout`** `MSDTreeLayout`

---

**`leaf_indices`** `tuple[int, ...]`

---

**`nodes`** `tuple[MSDTreeNode, ...]`

---

**`root_token_id`** `int`

---

```python
nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal.candidate_path(
    leaf_index: int
) -> tuple[int, ...]
```

Return the root-to-leaf candidate-token path.

```python
nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal.candidate_paths() -> tuple[tuple[int, ...], ...]
```

Return all leaf paths in deterministic tree order.

```python
class nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult(
    accepted_token_ids: tuple[int, ...],
    bonus_token_id: int,
    accepted_draft_tokens: int,
    leaf_index: int | None
)
```

Dataclass

Greedy posterior verification result for one candidate tree.

**`accepted_draft_tokens`** `int`

---

**`accepted_token_ids`** `tuple[int, ...]`

---

**`bonus_token_id`** `int`

---

**`emitted_token_ids`** `tuple[int, ...]`

Return accepted candidate tokens followed by the target bonus token.

---

**`leaf_index`** `int | None`

---

```python
class nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel()
```

**Bases:** [LlamaEagleDraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12#nemo_automodel-components-speculative-eagle-draft_llama_v12-LlamaEagleDraftModel)

EAGLE-1/2 draft model with modality-aware input feature construction.

```python
nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.copy_embeddings_from_target(
    target_embeddings: torch.nn.Module
) -> None
```

Raise: MSD drafts consume target embeddings instead of holding their own.

```python
nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.forward(
    inputs_embeds: torch.Tensor,
    target_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    image_mask: torch.Tensor
) -> torch.Tensor
```

Predict next-position target hidden states from multimodal features.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], containing
the target VLM input embeddings shifted one position to the left.
Image positions already contain projected vision features.

---

**`target_hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], the
frozen target language model's unshifted final hidden states.

---

**`attention_mask`** `torch.Tensor`

Tensor of shape \[batch, sequence], where one denotes
a real token and zero denotes padding.

---

**`image_mask`** `torch.Tensor`

Bool tensor of shape \[batch, sequence], aligned with
`inputs_embeds` and true for projected image-token positions.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], the draft predictions for

```python
nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.freeze_embeddings() -> None
```

Raise: MSD drafts have no token embedding table to freeze.

```python
class nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule(
    draft_model: torch.nn.Module,
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor,
    num_depths: int,
    mask_token_id: int,
    down_sample_ratio: float = 0.7,
    down_sample_ratio_min: float = 0.2,
    sequence_partitions: int = 1
)
```

**Bases:** `Module`

Draft-side P-EAGLE (parallel-drafting EAGLE-3) trainer module.

Faithful port of speculators' P-EAGLE
([https://github.com/vllm-project/speculators/pull/480](https://github.com/vllm-project/speculators/pull/480)): the draft predicts
all `num_depths` tokens in a *single* parallel forward over a flat,
COD-subsampled sequence -- it does NOT run EAGLE-3's autoregressive TTT
recurrence.

Per training sequence:

1. **COD sampling** (`generate_cod_sample_indices`) draws
   `(anchor_pos, depth)`: depth 0 keeps every position, depth `d` keeps a
   geometrically decaying `down_sample_ratio**d` fraction. The reference
   position of each element is `anchor_pos + depth`.
2. **Flat input assembly.** All depths are concatenated into one
   `[1, total_sampled]` sequence. Depth-0 slots take the real token id and
   the `fc`-projected target aux hidden state; depth >= 1 slots take the
   masked `mask_token_id` and the single learnable `mask_hidden`
   placeholder (projected through the same `fc`).
3. **COD flex attention.** A single `flex_attention` forward with the
   `create_peagle_mask_mod` block mask: each element attends to the
   causal depth-0 context of its document plus earlier-or-equal depths of
   its own rollout. This is exactly what vLLM's parallel-drafting runtime
   sees at inference.
4. **Count-normalized KL loss.** `KL(target || draft)` over the draft vocab
   at every supervised sampled position, normalized by a single total token
   count -- deeper depths (fewer COD positions) naturally contribute less
   gradient. No `0.8**d` schedule.

Batches with `batch_size &gt; 1` are processed row-by-row (speculators is
batch-size-1); per-row losses are accumulated with a shared denominator so
the normalization stays count-based across the whole batch.

**Sequence partitioning (`sequence_partitions &gt; 1`).** The flat COD
forward attends over `n * sum(r**d)` positions, so its peak attention /
activation memory grows with the context length and OOMs on long sequences.
P-EAGLE's Algorithm 1 (arXiv:2602.01469) splits each sequence into `S`
segments by dependency lineage (`assign_cod_segments`) and runs a
*separate* forward+backward per segment so only one segment's activations are
resident at a time. The partition is exact: each segment additionally reads
every depth-0 position up to its right boundary as key/value context (causal
completion), so a segment's queries see exactly the key/value set they would
in the single flat forward -- the gradients accumulated across segments equal
the single-forward gradient.

The split is *caller-driven* so the gradient sync stays correct under DDP:
`build_peagle_plan` (no-grad) assigns COD elements to segments, then the
recipe runs one `forward(..., peagle_segment=(plan, i))` per segment and
owns the `backward()`. Doing the per-segment backward here (inside a single
`forward`) would bypass `DistributedDataParallel`'s reducer -- its grad
all-reduce hooks only fire for backwards over the tensor `DDP.forward`
returns -- and silently desynchronize ranks. `sequence_partitions == 1` and
eval take the single flat `forward` unchanged.

**`down_sample_ratio`** `= float(down_sample_ratio)`

---

**`down_sample_ratio_min`** `= float(down_sample_ratio_min)`

---

**`mask_token_id`** `= int(mask_token_id)`

---

**`sequence_partitions`** `= int(sequence_partitions)`

---

```python
nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule._forward_peagle_segment(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    aux_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    peagle_segment: tuple
) -> nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics
```

Loss for one segment of a `build_peagle_plan` plan.

The recipe drives this once per `plan.units` entry and back-propagates
each result, so each segment owns a self-contained autograd graph that is
freed before the next -- and the backward flows through `DDP.forward` so
gradients all-reduce correctly. `metrics.loss` is the segment's share of
the count-normalized batch loss (`loss / total_den`); summing it over the
plan reproduces the single flat forward's loss.

```python
nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule._peagle_position_loss(
    input_ids_row: torch.Tensor,
    aux_row: torch.Tensor,
    target_logits_row: torch.Tensor,
    anchor_pos: torch.Tensor,
    depth: torch.Tensor,
    orig_positions: torch.Tensor,
    loss_positions: torch.Tensor,
    row_length: torch.Tensor,
    seq_len: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
```

Draft forward + count-normalized KL for one row's COD elements.

Shared by the single flat forward (one call per row, all sampled
positions charged) and the partitioned forward (one call per segment,
only the segment's owned/supervised positions charged -- the rest ride
along as key/value context). Returns `(loss_num, loss_den, correct,
valid)` as float scalars, where the loss is `Σ KL` over `loss_positions`
and `loss_den` is their count; the caller normalizes.

```python
nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule.build_peagle_plan(
    loss_mask: torch.Tensor
) -> '_PeaglePlan'
```

Assign COD elements to segments for the sequence-partitioning path.

Samples COD once per row (the indices must be reused across the segment
forwards), runs Algorithm 1 assignment (`assign_cod_segments`) plus
causal completion, and emits one `unit` per non-empty `(row, segment)`
as `(b, anchor, depth, orig_positions, loss_positions)`. `loss_positions`
marks the segment's *owned* supervised elements (charged loss); the other
elements are depth-0 causal-completion context (key/value only). The shared
`total_den` is the batch's total supervised-token count, so each segment's
`loss / total_den` sums to the single-forward loss.

```python
nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    aux_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    peagle_segment: tuple | None = None
) -> nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics
```

Run the P-EAGLE parallel-drafting loss for one batch.

`attention_mask` supplies the per-row valid length so padded positions
are excluded from attention (document mask) and from supervision.

`peagle_segment` selects the sequence-partitioning path: when it is a
`(plan, index)` pair (built by `build_peagle_plan`) this computes
the loss for that one segment only -- the recipe calls this once per
segment and owns the `backward()` so DDP's gradient sync stays correct.
When `None` (`sequence_partitions == 1` and eval) a single flat
forward over the whole COD sequence returns a grad-carrying loss.

```python
class nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel(
    runner: nemo_automodel.components.speculative.eagle.target_runner.TargetRunner,
    aux_layer_ids: typing.Sequence[int] | None = None
)
```

**Bases:** [Eagle3TargetBackend](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend)

EAGLE-3 target backend that runs the frozen target through a runner.

Engine-agnostic: it owns only the supervision *contract* and delegates the
actual forward to a `TargetRunner`. Concrete engines subclass this to
add their own `from_pretrained` (which builds the runner); everything else
is inherited.

**Parameters:**

**`runner`** `TargetRunner`

A loaded runner implementing `TargetRunner`.

---

**`aux_layer_ids`** `Sequence[int] | None` — default: None

The three decoder layers to capture (low / mid / high). When `None`
the shared EAGLE-3 default recipe is used, matching every other backend.

---

**`aux_layer_ids`**

---

**`model`** `= runner.model`

---

```python
nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel.close() -> None
```

Release the runner (frees GPU memory / engine handles).

```python
nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch
```

Run the runner's target and capture aux hidden states plus logits.

Produces an `Eagle3TargetBatch` byte-for-byte compatible with
`HFEagle3TargetModel.generate_batch`: the logits, `input_ids`
and `loss_mask` are shifted left by one (next-token alignment) while
the aux hidden states are kept position-aligned. The draft-vocab
projection happens trainer-side / server-side from `logits`, exactly
as in the co-located path.

Note: sequences are assumed right-padded (loss\_mask zeros the pad). With
causal attention, trailing pad tokens do not affect earlier positions,
so the captured supervision matches a masked HuggingFace forward.

```python
nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel.get_input_embeddings() -> types.SimpleNamespace
```

Return an object exposing `.weight` (the target input embeddings).

Matches the offline-cache / remote path: the draft's
`copy_embeddings_from_target` only reads `.weight`.

```python
class nemo_automodel.components.speculative.eagle.target_runner.TargetRunner()
```

Protocol

Minimal engine surface `RunnerEagle3TargetModel` depends on.

Implemented for real by an inference-engine runner (`SGLangTargetRunner`
today, a vLLM runner later) and faked in unit tests, which is why the
backend depends on this protocol rather than on any engine directly.

**`model`** `'nn.Module'`

---

```python
nemo_automodel.components.speculative.eagle.target_runner.TargetRunner.forward_eagle3(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Run the target once and return `(logits, aux_hidden_states)`.

`logits` is `[batch, seq, vocab]` (full vocab, unshifted) and
`aux_hidden_states` is `[batch, seq, 3 * hidden]` (the three capture
layers concatenated on the last dim, unshifted).

```python
nemo_automodel.components.speculative.eagle.target_runner.TargetRunner.input_embedding_weight() -> torch.Tensor
```

Return the target input-embedding weight `[vocab, hidden]`.

```python
nemo_automodel.components.speculative.eagle.target_runner.TargetRunner.set_aux_layers(
    aux_layer_ids: typing.Sequence[int]
) -> None
```

Tell the underlying model which 3 decoder layers to capture.

```python
nemo_automodel.components.speculative.eagle.msd_decode.accept_or_resample(
    target_logits: torch.Tensor,
    draft_logits: torch.Tensor,
    candidate_token_id: int,
    generator: torch.Generator | None = None,
    random_value: float | None = None
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticStep
```

Apply the lossless speculative acceptance rule for one candidate token.

`generator` must live on the same device as the logits, because both the
acceptance draw and the residual resample run there. `random_value` comes
from the half-open interval `[0, 1)` that `torch.rand` samples, and the
candidate is accepted on a strict `draw &lt; acceptance` so a token the target
assigns zero probability can never be accepted.

```python
nemo_automodel.components.speculative.eagle.msd_decode.build_msd_tree_layout(
    nodes: typing.Iterable[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode],
    leaf_indices: typing.Iterable[int],
    device: torch.device
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout
```

Build engine-ready tree attention and root-to-leaf retrieval metadata.

```python
nemo_automodel.components.speculative.eagle.registry.resolve_eagle1_draft_spec(
    architectures: list[str]
) -> nemo_automodel.components.speculative.eagle.registry.DraftSpec
```

Resolve the EAGLE-1 / EAGLE-2 draft spec for a target's `config.architectures` field.

```python
nemo_automodel.components.speculative.eagle.registry.resolve_eagle3_draft_spec(
    architectures: list[str]
) -> nemo_automodel.components.speculative.eagle.registry.DraftSpec
```

Resolve the EAGLE-3 draft spec for a target's `config.architectures` field.

```python
nemo_automodel.components.speculative.eagle.core.simulated_accept_length(
    step_prefix_hits: torch.Tensor,
    step_valid: torch.Tensor
) -> torch.Tensor
```

Expected accepted tokens per speculative round, from prefix-hit counts.

Models greedy chain drafting: `step_prefix_hits[k] / step_valid[k]`
estimates the joint probability that the first `k + 1` drafted tokens
all match the target's greedy choices, i.e. that the chain survives depth
`k + 1`, so the expectation is `1 + sum_k P(survives depth k + 1)`.
Joint prefix counts keep the correlation between depths that a product of
per-step marginal accuracies discards (coincident and disjoint hits score
differently). The leading 1 counts the token the target itself emits on
every verification round, matching the `accept_length` convention of
the serving benchmarks (`1 + accepted/drafts`). A step with no
supervised positions contributes zero via the `clamp_min`. This is a
training-time proxy for greedy chain decoding; engine tree drafting
typically accepts more.

```python
nemo_automodel.components.speculative.eagle.msd_decode.verify_greedy_tree(
    proposal: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal,
    target_logits_by_leaf: typing.Iterable[torch.Tensor]
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult
```

Select the longest target-greedy accepted path and its bonus token.

Each logits tensor has one row for every candidate token in the matching
leaf path plus one final row for the target bonus token. This is the greedy
form of EAGLE/MSD posterior verification.

```python
nemo_automodel.components.speculative.eagle.msd_decode.verify_hf_greedy_tree(
    model: torch.nn.Module,
    model_inputs: dict[str, torch.Tensor],
    proposal: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult
```

Reference target verification for a batch-one Hugging Face VLM.

The path-wise implementation is intentionally engine neutral and works for
VLMs whose visual tensors are prompt-only. A serving backend can instead
consume `proposal.layout` to verify the same tree in one tree-attention
target forward.

```python
nemo_automodel.components.speculative.eagle.msd_decode.verify_stochastic_chain(
    candidate_token_ids: typing.Iterable[int],
    target_logits: torch.Tensor,
    draft_logits: torch.Tensor,
    generator: torch.Generator | None = None
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticVerificationResult
```

Losslessly verify a linear draft chain and emit the correction token.

```python
nemo_automodel.components.speculative.eagle.registry.EAGLE1_DRAFT_REGISTRY: dict[str, DraftSpec] = {arch: (DraftSpec(draft_cls=LlamaEagleDraftModel)) for arch in _DENSE_ARCHITECTU...
```

```python
nemo_automodel.components.speculative.eagle.registry.EAGLE3_DRAFT_REGISTRY: dict[str, DraftSpec] = {arch: (DraftSpec(draft_cls=LlamaEagle3DraftModel)) for arch in _DENSE_ARCHITECT...
```