> 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.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`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDGreedyDecoder)                               | Batch-one reference MSD round with recursive drafting and tree verification. |
| [`MSDStochasticStep`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDStochasticStep)                             | Result of one lossless speculative acceptance decision.                      |
| [`MSDStochasticVerificationResult`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDStochasticVerificationResult) | Lossless stochastic verification result for a linear candidate chain.        |
| [`MSDTreeDraftGenerator`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeDraftGenerator)                     | Recursively draft a top-k MSD feature tree from target prefix features.      |
| [`MSDTreeLayout`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeLayout)                                     | Tree attention and retrieval metadata for an MSD proposal.                   |
| [`MSDTreeNode`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeNode)                                         | One drafted text-token node in an MSD candidate tree.                        |
| [`MSDTreeProposal`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeProposal)                                 | A target root token plus recursively drafted candidate-tree nodes.           |
| [`MSDVerificationResult`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDVerificationResult)                     | Greedy posterior verification result for one candidate tree.                 |
| [`_DraftState`](#nemo_automodel-components-speculative-eagle-msd_decode-_DraftState)                                         | Internal recursive feature-draft state for one tree node.                    |

### Functions

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

### API

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

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

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

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

Dataclass

Result of one lossless speculative acceptance decision.

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

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

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

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

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

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

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

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

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

Return all leaf paths in deterministic tree order.

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

Return accepted candidate tokens followed by the target bonus token.

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

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

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

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

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

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

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

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