> 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.vispec_draft

ViSpec (vision-aware) draft model for speculative decoding on VLM targets.

ViSpec (Kang et al., NeurIPS 2025, arXiv:2509.15235) extends the EAGLE-1/2
draft with two vision-specific modules, keeping the rest of the draft
(`embed_tokens` / `fc` / decoder layers / `norm`) byte-identical to
:class:`~nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel`
so a text-only EAGLE-1/2 checkpoint can initialize stage-2 ViSpec training:

* :class:`VispecImageAdaptor` -- `num_query_tokens` learnable queries
  cross-attend over the target's image-token features and compress a whole
  image span (hundreds to thousands of tokens) into `num_query_tokens`
  vectors. `num_query_tokens - 1` of them are spliced back into the draft
  sequence at the *original* trailing positions of the image span, so the
  positional layout of the surrounding text is untouched.
* `img_fc` -- the remaining ("global") image vector is broadcast onto every
  subsequent text position and mixed in with a `[2*hidden -> hidden]`
  projection, giving each text token vision context without paying for image
  tokens in the draft's KV cache.

Reference implementation: `vispec/model/cnets_ours.py` in
[https://github.com/KangJialiang/ViSpec](https://github.com/KangJialiang/ViSpec).

## Module Contents

### Classes

| Name                                                                                                 | Description                                                                        |
| ---------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- |
| [`VispecDraftModel`](#nemo_automodel-components-speculative-eagle-vispec_draft-VispecDraftModel)     | EAGLE-1/2 draft extended with ViSpec's image compression and global image feature. |
| [`VispecImageAdaptor`](#nemo_automodel-components-speculative-eagle-vispec_draft-VispecImageAdaptor) | Compress an image-token span into a small set of learnable-query vectors.          |

### Functions

| Name                                                                                                                           | Description                                                                |
| ------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------- |
| [`apply_vispec_draft_architecture`](#nemo_automodel-components-speculative-eagle-vispec_draft-apply_vispec_draft_architecture) | Pin a target-derived draft config to ViSpec's released draft architecture. |

### API

```python
class nemo_automodel.components.speculative.eagle.vispec_draft.VispecDraftModel(
    config: transformers.PretrainedConfig
)
```

**Bases:** [LlamaEagleDraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12#nemo_automodel-components-speculative-eagle-draft_llama_v12-LlamaEagleDraftModel)

EAGLE-1/2 draft extended with ViSpec's image compression and global image feature.

Adds two config fields on top of the EAGLE-1/2 draft config:

* `vispec_num_query_tokens` (default 2) -- queries per image span.
* `draft_num_hidden_layers` -- inherited, unchanged.

The base draft's parameters (`embed_tokens`, `fc`, `layers`, `norm`)
keep their names, so `load_state_dict(..., strict=False)` restores a
stage-1 EAGLE-1/2 checkpoint and leaves only `img_adaptor` / `img_fc`
freshly initialized.

**`img_adaptor`** `= VispecImageAdaptor(config, self.num_query_tokens)`

---

**`img_fc`**

---

**`num_query_tokens`** `= int(getattr(config, 'vispec_num_query_tokens', 2))`

---

```python
nemo_automodel.components.speculative.eagle.vispec_draft.VispecDraftModel._compress_sequence(
    inputs_embeds: torch.Tensor,
    target_hidden_states: torch.Tensor,
    image_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Build the compressed draft sequence for one (batch-size-1) sample.

Walks the sample image span by image span. Each span contributes its
leading text positions (fused with the *previous* span's global image
vector) followed by `num_query_tokens - 1` compressed image tokens;
the trailing text after the last span is fused with the last span's
global vector. Because every span ends with its image run, this
reordering preserves the original left-to-right token order.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[sequence, hidden].

---

**`target_hidden_states`** `torch.Tensor`

Tensor of shape \[sequence, hidden].

---

**`image_mask`** `torch.Tensor`

Bool tensor of shape \[sequence]; True at image positions.

---

**Returns:** `torch.Tensor`

Tuple of `(hidden_states, source_index)`:

```python
nemo_automodel.components.speculative.eagle.vispec_draft.VispecDraftModel._fuse(
    inputs_embeds: torch.Tensor,
    target_hidden_states: torch.Tensor,
    global_image_feature: torch.Tensor
) -> torch.Tensor
```

Mix the global image vector into the target hidden state, then run EAGLE's `fc`.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[tokens, hidden].

---

**`target_hidden_states`** `torch.Tensor`

Tensor of shape \[tokens, hidden].

---

**`global_image_feature`** `torch.Tensor`

Tensor of shape \[1, hidden], broadcast over `tokens`.

---

**Returns:** `torch.Tensor`

Tensor of shape \[tokens, hidden].

```python
nemo_automodel.components.speculative.eagle.vispec_draft.VispecDraftModel.forward(
    inputs_embeds: torch.Tensor,
    target_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    image_mask: torch.Tensor
) -> torch.Tensor
```

Predict the next-position target hidden states from vision-aware features.

Batch size must be 1: image spans have per-sample lengths, so the
compressed sequences of two samples would not share a length. This
matches the reference implementation, which raises on batch size > 1.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[1, sequence, hidden] -- the target's
embedding-layer output (vision features already spliced in),
shifted left by one position so index `i` holds the embedding
of token `i + 1`.

---

**`target_hidden_states`** `torch.Tensor`

Tensor of shape \[1, sequence, hidden] -- the
target's last hidden state, *not* shifted.

---

**`attention_mask`** `torch.Tensor`

Tensor of shape \[1, sequence]; 1 for real tokens,
0 for padding.

---

**`image_mask`** `torch.Tensor`

Bool tensor of shape \[1, sequence], aligned with
`inputs_embeds` (True where token `i + 1` is an image token).

---

**Returns:** `torch.Tensor`

Tensor of shape \[1, sequence, hidden]. Positions consumed by the

```python
nemo_automodel.components.speculative.eagle.vispec_draft.VispecDraftModel.reset_vispec_parameters() -> None
```

Initialize the ViSpec-only modules (identity-start for `img_fc`).

`img_fc` starts as `[I | 0]`: it copies the target hidden state
through and ignores the global image vector, so a stage-2 run
initialized from a stage-1 EAGLE-1/2 checkpoint starts numerically
equal to that checkpoint and only then learns to use vision context.

```python
class nemo_automodel.components.speculative.eagle.vispec_draft.VispecImageAdaptor(
    config: transformers.PretrainedConfig,
    num_query_tokens: int
)
```

**Bases:** `Module`

Compress an image-token span into a small set of learnable-query vectors.

A single non-causal cross-attention: `num_query_tokens` learnable queries
attend over the image-token features, so the output length is independent
of how many image tokens the target emitted.

**Parameters:**

**`config`** `PretrainedConfig`

Draft config supplying `hidden_size` and `num_attention_heads`.

---

**`num_query_tokens`** `int`

Number of learnable queries (ViSpec's `num_q`).

---

**`head_dim`** `= self.hidden_size // self.num_heads`

---

**`hidden_size`** `= config.hidden_size`

---

**`k_proj`**

---

**`num_heads`** `= config.num_attention_heads`

---

**`o_proj`**

---

**`query`**

---

**`v_proj`**

---

```python
nemo_automodel.components.speculative.eagle.vispec_draft.VispecImageAdaptor.forward(
    image_features: torch.Tensor
) -> torch.Tensor
```

Compress image features into `num_query_tokens` vectors.

**Parameters:**

**`image_features`** `torch.Tensor`

Tensor of shape \[batch, image\_tokens, hidden] holding
the target's embedding-layer output at the image positions of one
image span.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, num\_query\_tokens, hidden].

```python
nemo_automodel.components.speculative.eagle.vispec_draft.VispecImageAdaptor.reset_query() -> None
```

Re-draw the learnable queries from `N(0, head_dim ** -0.5)` (ViSpec's init).

```python
nemo_automodel.components.speculative.eagle.vispec_draft.apply_vispec_draft_architecture(
    config: transformers.PretrainedConfig
) -> None
```

Pin a target-derived draft config to ViSpec's released draft architecture.

Both ViSpec stages derive the draft config from the target's text config, so
without this the draft silently inherits whatever the target's language tower
happens to use. The released draft is not that: `qwen2.5_vl_7B_config.json`
in the reference repository sets `num_attention_heads: 28`,
`num_key_value_heads: 28` and `qkv_bias: true`, whereas the matching
Qwen2.5-VL-7B text tower is 28-head GQA over 4 KV heads. Copying the target
therefore produced a materially smaller draft than the paper's, which is not
a configuration difference a reader of the recipe would notice.

The three settings cannot be read off the target: HF's `Qwen2_5_VLTextConfig`
exposes neither `attention_bias` nor `qkv_bias` (its attention hard-codes
the qkv bias inside the module), so there is nothing to inherit. They are
properties of the ViSpec draft rather than of any one target, and are applied
unconditionally for every ViSpec target.

**Parameters:**

**`config`** `PretrainedConfig`

The draft config, already derived from the target's text config.
Modified in place.

---