nemo_automodel.components.speculative.eagle.msd_decode

Reference multimodal speculative decoding primitives for MSD.

The module implements the feature-level EAGLE/MSD decoding contract in plain PyTorch. It intentionally keeps the proposal tree and posterior verification separate from serving-engine kernels: engines can use :class:MSDTreeLayout to execute the tree in one attention pass, while the Hugging Face reference path verifies each leaf independently for correctness and integration testing.

Module Contents

Classes

Name	Description
`MSDGreedyDecoder`	Batch-one reference MSD round with recursive drafting and tree verification.
`MSDStochasticStep`	Result of one lossless speculative acceptance decision.
`MSDStochasticVerificationResult`	Lossless stochastic verification result for a linear candidate chain.
`MSDTreeDraftGenerator`	Recursively draft a top-k MSD feature tree from target prefix features.
`MSDTreeLayout`	Tree attention and retrieval metadata for an MSD proposal.
`MSDTreeNode`	One drafted text-token node in an MSD candidate tree.
`MSDTreeProposal`	A target root token plus recursively drafted candidate-tree nodes.
`MSDVerificationResult`	Greedy posterior verification result for one candidate tree.
`_DraftState`	Internal recursive feature-draft state for one tree node.

Functions

Name	Description
`_right_padded_active_length`	Return the prefix length of a batch-one right-padded attention mask.
`_validate_generator_device`	Reject a random generator whose device type differs from the sampling device.
`accept_or_resample`	Apply the lossless speculative acceptance rule for one candidate token.
`build_msd_tree_layout`	Build engine-ready tree attention and root-to-leaf retrieval metadata.
`verify_greedy_tree`	Select the longest target-greedy accepted path and its bonus token.
`verify_hf_greedy_tree`	Reference target verification for a batch-one Hugging Face VLM.
`verify_stochastic_chain`	Losslessly verify a linear draft chain and emit the correction token.

API

class nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder(
    target: nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel,
    draft_model: torch.nn.Module
)

Batch-one reference MSD round with recursive drafting and tree verification.

generator

target_embeddings

= self._resolve_target_embeddings(target.model)

nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder._resolve_target_embeddings(
    model: torch.nn.Module
) -> torch.nn.Module

staticmethod

Find the language token embedding module across common HF VLM layouts.

nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder.decode_round(
    model_inputs: dict[str, torch.Tensor],
    draft_steps: int,
    top_k: int,
    beam_width: int
) -> tuple[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal, nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult]

Draft and greedily verify one MSD candidate tree for a VLM prompt.

class nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticStep(
    token_id: int,
    accepted: bool
)

Dataclass

Result of one lossless speculative acceptance decision.

accepted

bool

token_id

int

class nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticVerificationResult(
    emitted_token_ids: tuple[int, ...],
    accepted_draft_tokens: int
)

Dataclass

Lossless stochastic verification result for a linear candidate chain.

accepted_draft_tokens

int

emitted_token_ids

tuple[int, ...]

class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeDraftGenerator(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    target_embeddings: torch.nn.Module
)

Recursively draft a top-k MSD feature tree from target prefix features.

nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeDraftGenerator.propose(
    shifted_inputs_embeds: torch.Tensor,
    input_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    shifted_image_mask: torch.Tensor,
    root_token_id: int,
    draft_steps: int,
    top_k: int,
    beam_width: int
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal

Create a feature-level candidate tree after a target-sampled root.

shifted_inputs_embeds and shifted_image_mask use the same next-token alignment as :class:HFMSDTargetModel. The final shifted embedding is replaced with the target root token embedding, then each recursive draft prediction is appended as the feature context for its child nodes. attention_mask must be right-padded, since the prefix is taken as its leading active slice.

class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout(
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    retrieve_indices: torch.Tensor
)

Dataclass

Tree attention and retrieval metadata for an MSD proposal.

Node zero represents the target-sampled root token. attention_mask is boolean and contains the root plus every ancestor for each node. Serving backends turn its false entries into additive -inf attention bias.

attention_mask

Tensor

position_ids

Tensor

retrieve_indices

Tensor

class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode(
    index: int,
    parent_index: int,
    token_id: int,
    depth: int,
    log_probability: float
)

Dataclass

One drafted text-token node in an MSD candidate tree.

depth

int

index

int

log_probability

float

parent_index

int

token_id

int

class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal(
    root_token_id: int,
    nodes: tuple[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode, ...],
    leaf_indices: tuple[int, ...],
    layout: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout
)

Dataclass

A target root token plus recursively drafted candidate-tree nodes.

layout

MSDTreeLayout

leaf_indices

tuple[int, ...]

nodes

tuple[MSDTreeNode, ...]

root_token_id

int

nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal.candidate_path(
    leaf_index: int
) -> tuple[int, ...]

Return the root-to-leaf candidate-token path.

nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal.candidate_paths() -> tuple[tuple[int, ...], ...]

Return all leaf paths in deterministic tree order.

class nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult(
    accepted_token_ids: tuple[int, ...],
    bonus_token_id: int,
    accepted_draft_tokens: int,
    leaf_index: int | None
)

Dataclass

Greedy posterior verification result for one candidate tree.

accepted_draft_tokens

int

accepted_token_ids

tuple[int, ...]

bonus_token_id

int

emitted_token_ids

tuple[int, ...]

Return accepted candidate tokens followed by the target bonus token.

leaf_index

int | None

class nemo_automodel.components.speculative.eagle.msd_decode._DraftState(
    node_index: int,
    inputs_embeds: torch.Tensor,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    image_mask: torch.Tensor,
    log_probability: float
)

Dataclass

Internal recursive feature-draft state for one tree node.

attention_mask

Tensor

hidden_states

Tensor

image_mask

Tensor

inputs_embeds

Tensor

log_probability

float

node_index

int

nemo_automodel.components.speculative.eagle.msd_decode._right_padded_active_length(
    attention_mask: torch.Tensor
) -> int

Return the prefix length of a batch-one right-padded attention mask.

The reference decoder slices every tensor as [:, :active_length], which is only the real prefix when padding sits at the end. Left padding would silently select pad positions, so reject it instead.

nemo_automodel.components.speculative.eagle.msd_decode._validate_generator_device(
    generator: torch.Generator | None,
    device: torch.device
) -> None

Reject a random generator whose device type differs from the sampling device.

torch.rand and torch.multinomial both refuse such a generator, so a mismatch here would otherwise surface as an opaque failure part-way through verification. The check is deliberately type-level, matching PyTorch’s own generator check, which likewise ignores the device index.

nemo_automodel.components.speculative.eagle.msd_decode.accept_or_resample(
    target_logits: torch.Tensor,
    draft_logits: torch.Tensor,
    candidate_token_id: int,
    generator: torch.Generator | None = None,
    random_value: float | None = None
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticStep

Apply the lossless speculative acceptance rule for one candidate token.

generator must live on the same device as the logits, because both the acceptance draw and the residual resample run there. random_value comes from the half-open interval [0, 1) that torch.rand samples, and the candidate is accepted on a strict draw < acceptance so a token the target assigns zero probability can never be accepted.

nemo_automodel.components.speculative.eagle.msd_decode.build_msd_tree_layout(
    nodes: typing.Iterable[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode],
    leaf_indices: typing.Iterable[int],
    device: torch.device
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout

Build engine-ready tree attention and root-to-leaf retrieval metadata.

nemo_automodel.components.speculative.eagle.msd_decode.verify_greedy_tree(
    proposal: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal,
    target_logits_by_leaf: typing.Iterable[torch.Tensor]
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult

Select the longest target-greedy accepted path and its bonus token.

Each logits tensor has one row for every candidate token in the matching leaf path plus one final row for the target bonus token. This is the greedy form of EAGLE/MSD posterior verification.

nemo_automodel.components.speculative.eagle.msd_decode.verify_hf_greedy_tree(
    model: torch.nn.Module,
    model_inputs: dict[str, torch.Tensor],
    proposal: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult

Reference target verification for a batch-one Hugging Face VLM.

The path-wise implementation is intentionally engine neutral and works for VLMs whose visual tensors are prompt-only. A serving backend can instead consume proposal.layout to verify the same tree in one tree-attention target forward.

nemo_automodel.components.speculative.eagle.msd_decode.verify_stochastic_chain(
    candidate_token_ids: typing.Iterable[int],
    target_logits: torch.Tensor,
    draft_logits: torch.Tensor,
    generator: torch.Generator | None = None
) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticVerificationResult

Losslessly verify a linear draft chain and emit the correction token.