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

Package Contents

Classes

Name	Description
`MTPAcceptanceMetrics`	Per-depth joint-prefix-hit counts for `simulated_accept_length`.

Functions

Name	Description
`compute_mtp_prefix_hits`	Score a model’s MTP depths against its own main-head greedy prediction.
`simulated_accept_length`	Expected accepted tokens per speculative round, from prefix-hit counts.

API

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

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] | NoneDefaults to None

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

mtp_per_depth_logits

list[torch.Tensor] | NoneDefaults to 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 | NoneDefaults to 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 | NoneDefaults to 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.

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.