> This page is for version Nightly (default).
> 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

MTP (multi-token prediction) acceptance-length evaluation.

Scores a model's existing native MTP head (`nemo_automodel.components.models.common.mtp`)
against its own main-head greedy prediction under teacher forcing, to report
the self-speculative accept length a finetuned MTP head would achieve, without
adding a separate draft model or recipe.

## Submodules

* **[`nemo_automodel.components.speculative.mtp.core`](/nemo/automodel/nightly/nemo-automodel/nemo_automodel/components/speculative/mtp/core)**

## Package 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                                                             |
| ------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------- |
| [`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.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.