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

Core EAGLE-3 training logic for the minimal Llama MVP.

## Module Contents

### Classes

| Name                                                                                           | Description                                                       |
| ---------------------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`Eagle3StepMetrics`](#nemo_automodel-components-speculative-eagle-core-Eagle3StepMetrics)     | Aggregated metrics from one EAGLE-3 training step.                |
| [`Eagle3TrainerModule`](#nemo_automodel-components-speculative-eagle-core-Eagle3TrainerModule) | Draft-side EAGLE-3 trainer module with test-time-training unroll. |
| [`_CpAllReduceSum`](#nemo_automodel-components-speculative-eagle-core-_CpAllReduceSum)         | Differentiable sum-all-reduce across the cp group.                |

### Functions

| Name                                                                                                             | Description                                                                   |
| ---------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_compute_target_distribution`](#nemo_automodel-components-speculative-eagle-core-_compute_target_distribution) | Project target logits into draft vocabulary space and build supervision mask. |
| [`_cp_global_step_loss`](#nemo_automodel-components-speculative-eagle-core-_cp_global_step_loss)                 | Renormalize a per-shard masked-mean loss over the full cp sequence.           |
| [`_cp_shift_left`](#nemo_automodel-components-speculative-eagle-core-_cp_shift_left)                             | Left-shift a context-parallel-sharded sequence tensor by one position.        |
| [`_cp_shift_left_zigzag`](#nemo_automodel-components-speculative-eagle-core-_cp_shift_left_zigzag)               | Left-shift a ZIG-ZAG-sharded sequence tensor by one position.                 |
| [`_masked_mean`](#nemo_automodel-components-speculative-eagle-core-_masked_mean)                                 | Mean of `values` over the supervised positions.                               |
| [`_shift_left_with_zero`](#nemo_automodel-components-speculative-eagle-core-_shift_left_with_zero)               | Shift a batched sequence tensor left and zero-fill the tail.                  |
| [`simulated_accept_length`](#nemo_automodel-components-speculative-eagle-core-simulated_accept_length)           | Expected accepted tokens per speculative round, from prefix-hit counts.       |

### Data

[`_LK_LOSS_TYPES`](#nemo_automodel-components-speculative-eagle-core-_LK_LOSS_TYPES)

### API

```python
class nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics(
    loss: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor,
    step_prefix_hits: torch.Tensor | None = None,
    step_valid: torch.Tensor | None = None
)
```

Dataclass

Aggregated metrics from one EAGLE-3 training step.

`step_prefix_hits` / `step_valid` are `[ttt_steps]` per-TTT-step
counts for the simulated accept length (:func:`simulated_accept_length`):
`step_prefix_hits[k]` counts positions whose greedy draft chain is still
fully correct through depth `k + 1` (top-1 hit at every step `<= k`),
and `step_valid[k]` counts positions supervised at every step `<= k`
under the shifted loss/document masks, so numerator and denominator cover
the same chain population (a chain with an unsupervised earlier depth,
e.g. from a gappy multi-turn loss mask, can never be a hit and must not
count as a miss). `step_valid` deliberately does NOT exclude positions
whose target token falls outside the compressed draft vocabulary: serving
must reject those tokens (the draft cannot emit them), so they count as
chain breaks instead of dropping out of the denominator. Both are kept
unreduced so the recipe can accumulate them over a logging window. The
P-EAGLE trainer, whose depths are drafted in parallel from one anchor
rather than as a TTT chain, leaves them `None`.

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

```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
:func:`_cp_global_step_loss`, so the mixing weight and the loss match the
single-rank run.

**Parameters:**

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

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

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._CpAllReduceSum()
```

**Bases:** `Function`

Differentiable sum-all-reduce across the cp group.

Forward sums per-rank inputs into a replicated total. Each rank uses that total
identically, so the loss gradient w\.r.t. this rank's input is just the incoming
grad (coefficient 1) -- hence the identity backward.

```python
nemo_automodel.components.speculative.eagle.core._CpAllReduceSum.backward(
    ctx,
    grad
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.core._CpAllReduceSum.forward(
    ctx,
    x,
    cp_group
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.core._compute_target_distribution(
    target_logits: torch.Tensor,
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor,
    loss_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Project target logits into draft vocabulary space and build supervision mask.

```python
nemo_automodel.components.speculative.eagle.core._cp_global_step_loss(
    step_loss: torch.Tensor,
    position_mask: torch.Tensor,
    cp_group
) -> torch.Tensor
```

Renormalize a per-shard masked-mean loss over the full cp sequence.

`step_loss` is the mean over this rank's LOCAL supervised positions. Recover the
local sum (mean \* count), sum both across cp (the sum differentiably), and divide
by the global count -- so the loss value equals the full-sequence loss and the
backprop'd gradient is the global gradient.

```python
nemo_automodel.components.speculative.eagle.core._cp_shift_left(
    tensor: torch.Tensor,
    cp_group
) -> torch.Tensor
```

Left-shift a context-parallel-sharded sequence tensor by one position.

Each rank holds a contiguous `S/cp` shard. A plain left-shift would zero-fill
the boundary, but the token that rolls into rank `r`'s tail lives at rank
`r+1`'s head. Shift locally, then P2P the neighbour's (current) head into the
tail; the last rank keeps the zero fill. Applied in lockstep each TTT step this
reproduces the global shift exactly. Used for labels / input\_ids only (no grad).

```python
nemo_automodel.components.speculative.eagle.core._cp_shift_left_zigzag(
    tensor: torch.Tensor,
    cp_group
) -> torch.Tensor
```

Left-shift a ZIG-ZAG-sharded sequence tensor by one position.

Under zig-zag sharding rank `r` holds two non-contiguous global chunks --
the early chunk `r` and the late chunk `2*cp-1-r` -- laid out locally as
`[early | late]` (each of length `c = local_len/2`). A global left-shift
therefore rolls each half locally and needs two boundary tokens:

* the early half's tail (global end of chunk `r`) is the head of chunk
  `r+1` -- rank `r+1`'s early head -- except on the last rank, where chunk
  `r+1` is that rank's OWN late chunk, so the fill is local.
* the late half's tail (global end of chunk `2*cp-1-r`) is the head of chunk
  `2*cp-r` -- rank `r-1`'s late head -- except on rank 0, whose late tail is
  the global last position and stays zero-filled.

Two ring P2P exchanges (early head -> `r-1`, late head -> `r+1`) cover both.
Used for labels / input\_ids only (no grad). Applied in lockstep each TTT step
this reproduces the global shift exactly. See :func:`_cp_shift_left` for the
contiguous-shard analogue.

```python
nemo_automodel.components.speculative.eagle.core._masked_mean(
    values: torch.Tensor,
    position_mask: torch.Tensor
) -> torch.Tensor
```

Mean of `values` over the supervised positions.

**Parameters:**

Tensor of shape `[batch, sequence]`; a per-position quantity.

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

**Returns:** `torch.Tensor`

Scalar tensor; the masked mean (zero when no position is supervised).

```python
nemo_automodel.components.speculative.eagle.core._shift_left_with_zero(
    tensor: torch.Tensor
) -> torch.Tensor
```

Shift a batched sequence tensor left and zero-fill the tail.

```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.core._LK_LOSS_TYPES = ('alpha', 'lambda')
```