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.

Module Contents

Classes

Name	Description
`VispecDraftModel`	EAGLE-1/2 draft extended with ViSpec’s image compression and global image feature.
`VispecImageAdaptor`	Compress an image-token span into a small set of learnable-query vectors.

Functions

Name	Description
`apply_vispec_draft_architecture`	Pin a target-derived draft config to ViSpec’s released draft architecture.

API

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

Bases: 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))

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):

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

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

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.

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

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

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

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.