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`	Per-depth joint-prefix-hit counts for `simulated_accept_length`.

Functions

Name	Description
`_roll_per_depth`	Cumulative per-depth left-roll of `base`, mirroring `calculate_mtp_loss`’s label roll.
`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._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.

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.