> This page is for version Latest.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.mtp.core

Self-speculative acceptance-length evaluation for a model's native MTP head.

MTP finetuning trains the existing `MTPModule` depths directly (see
`nemo_automodel.components.loss.mtp.calculate_mtp_loss`), but that loss only
reports scalar per-depth CE -- it says nothing about how many tokens the MTP
head would actually get accepted under self-speculative decoding. This module
computes that acceptance-length proxy from the same per-depth outputs the loss
already consumes, without an extra model forward pass.

Acceptance is checked against the model's own main-head greedy prediction (the
actual verifier in self-speculative decoding), not against the dataset label:
a draft is "accepted" when it agrees with what the model itself would have
produced, which can diverge from the ground-truth label whenever the main head
is imperfect. The dataset label only decides which positions are supervised
(masking ignored/padded/cross-document positions), the same role it plays in
`calculate_mtp_loss`.

## Module Contents

### Classes

| Name                                                                                           | Description                                                      |
| ---------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- |
| [`MTPAcceptanceMetrics`](#nemo_automodel-components-speculative-mtp-core-MTPAcceptanceMetrics) | Per-depth joint-prefix-hit counts for `simulated_accept_length`. |

### Functions

| Name                                                                                                   | Description                                                                            |
| ------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------- |
| [`_roll_per_depth`](#nemo_automodel-components-speculative-mtp-core-_roll_per_depth)                   | Cumulative per-depth left-roll of `base`, mirroring `calculate_mtp_loss`'s label roll. |
| [`compute_mtp_prefix_hits`](#nemo_automodel-components-speculative-mtp-core-compute_mtp_prefix_hits)   | Score a model's MTP depths against its own main-head greedy prediction.                |
| [`simulated_accept_length`](#nemo_automodel-components-speculative-eagle-core-simulated_accept_length) | Expected accepted tokens per speculative round, from prefix-hit counts.                |

### API

```python
class nemo_automodel.components.speculative.mtp.core.MTPAcceptanceMetrics(
    prefix_hits: torch.Tensor,
    valid: torch.Tensor
)
```

Dataclass

Per-depth joint-prefix-hit counts for `simulated_accept_length`.

`prefix_hits[k]` counts positions whose greedy MTP chain is still fully
correct through depth `k + 1` (argmax hit at every depth `<= k`), and
`valid[k]` counts positions supervised at every depth `<= k`. Kept
unreduced (shape `[num_depths]`) so a recipe can sum them across
microbatches and the DP/CP group before dividing once.

**`prefix_hits`** `Tensor`

---

**`valid`** `Tensor`

---

```python
nemo_automodel.components.speculative.mtp.core._roll_per_depth(
    base: torch.Tensor,
    num_depths: int,
    fill_value: int,
    seq_idx: torch.Tensor | None
) -> list[torch.Tensor]
```

Cumulative per-depth left-roll of `base`, mirroring `calculate_mtp_loss`'s label roll.

`fill_value` only matters at trailing/cross-sub-sequence positions, which
the caller always excludes from both the hit and valid counts.

```python
nemo_automodel.components.speculative.mtp.core.compute_mtp_prefix_hits(
    model: torch.nn.Module,
    mtp_per_depth_h: list[torch.Tensor] | None = None,
    mtp_per_depth_logits: list[torch.Tensor] | None = None,
    verifier_logits: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int,
    seq_idx: torch.Tensor | None = None,
    cu_seqlens: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.mtp.core.MTPAcceptanceMetrics
```

Score a model's MTP depths against its own main-head greedy prediction.

**Parameters:**

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

The wrapped model; used to fetch the shared LM head when
`mtp_per_depth_h` is supplied instead of logits.

---

**`mtp_per_depth_h`** `list[torch.Tensor] | None` — default: None

Per-depth hidden-state tensors of shape
`[batch, sequence, hidden]`.

---

**`mtp_per_depth_logits`** `list[torch.Tensor] | None` — default: None

Per-depth logit tensors of shape
`[batch, sequence, vocab]`.

---

**`verifier_logits`** `torch.Tensor`

The model's own main-head logits, `[batch, sequence,
vocab]`, from the same forward pass. Its argmax is the acceptance
reference -- not `labels`, which can disagree with what the
model itself would produce.

---

**`labels`** `torch.Tensor`

Original unshifted label tensor of shape `[batch, sequence]`,
used only to mask ignored/padded positions out of the counts.

---

**`ignore_index`** `int`

Label value excluded from both the hit and valid counts.

---

**`seq_idx`** `torch.Tensor | None` — default: None

Optional per-token sub-sequence index used to mask rolls
whose source position crosses a packed-sequence boundary. Takes
precedence over `cu_seqlens`.

---

**`cu_seqlens`** `torch.Tensor | None` — default: None

Optional cumulative sequence lengths `[num_seqs+1]`
(THD-pack layout) used to derive `seq_idx` when it is not given,
the same boundary `calculate_mtp_loss` uses for this batch.

---

**Returns:** `MTPAcceptanceMetrics`

class:`MTPAcceptanceMetrics` with one raw sum per depth.

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