nemo_automodel.components.speculative.eagle.vllm_runner

vLLM forward for the EAGLE-3 target (server-side, GPU only).

This module owns every vLLM-internal touch point so the rest of the speculative stack stays vLLM-agnostic and importable without vLLM. It is imported lazily (only from :meth:VLLMEagle3TargetModel.from_pretrained) and implements the engine-agnostic :class:~nemo_automodel.components.speculative.eagle.target_runner.TargetRunner surface, exactly like sglang_runner.SGLangTargetRunner.

Mechanism (vLLM’s native extract_hidden_states speculative method, the supported way to pull EAGLE-3 supervision out of vLLM without driving v1 worker internals):

Build an offline LLM with speculative_config={"method": "extract_hidden_states", ...} and the three EAGLE-3 capture layers (plus num_hidden_layers so the final pre-norm hidden is captured too) in draft_model_config.hf_config.eagle_aux_hidden_state_layer_ids. Chunked prefill is disabled so every prompt is captured in one prefill.
generate(max_tokens=1) over the batch; vLLM writes the captured hidden states to disk through ExampleHiddenStatesConnector (a KV connector), and each request output carries the path under kv_transfer_params.
Read the per-prompt [seq, num_capture_layers, hidden] tensor back, split the three EAGLE-3 layers (concatenated into [seq, 3 * hidden]) from the final-layer hidden, and rebuild full-vocab logits in-process by applying the target’s final RMSNorm + LM head (loaded once from the model’s safetensors).

vLLM’s extract_hidden_states API is version-coupled: the calls here track vllm==0.23.0. This forward path requires a GPU and vLLM, so it is validated on the training server, not in CPU unit tests; the CPU tests exercise the shared contract layer against a fake runner instead.

Module Contents

Classes

Name	Description
`VLLMTargetRunner`	Offline vLLM engine that returns EAGLE-3 supervision tensors.
`_VLLMModelShim`	Lightweight stand-in exposing `.config` + `.parameters()`.

Functions

Name	Description
`_load_target_head_weights`	Load `(embed, final_norm, lm_head)` weights straight from the safetensors.
`vllm_dtype_str`	Map a torch dtype to the string form vLLM’s `dtype` argument expects.

Data

_VLLM_DTYPE_STRINGS

logger

API

class nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner(
    model_path: str,
    dtype: typing.Optional[torch.dtype] = None,
    tp_size: int = 1,
    trust_remote_code: bool = False,
    gpu_memory_utilization: float = 0.5,
    shared_storage_path: typing.Optional[str] = None,
    vllm_kwargs: typing.Optional[dict] = None
)

Offline vLLM engine that returns EAGLE-3 supervision tensors.

Built via :meth:build; consumed through the engine-agnostic :class:~nemo_automodel.components.speculative.eagle.target_runner.TargetRunner surface (model / set_aux_layers / forward_eagle3 / input_embedding_weight). The vLLM LLM is constructed lazily in :meth:set_aux_layers, because the capture layers must be known up front.

_aux_layer_ids

Optional[list[int]] = None

_embed_w

Optional[Tensor] = None

_hidden

= int(hf_config.hidden_size)

_lm_head_w

Optional[Tensor] = None

_norm_w

Optional[Tensor] = None

_num_layers

= int(hf_config.num_hidden_layers)

_rms_eps

= float(getattr(hf_config, 'rms_norm_eps', 1e-06))

_shared_storage_path

_vllm_kwargs

= dict(vllm_kwargs or {})

model

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner._build_llm() -> None

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner._compute_logits(
    final_hidden: torch.Tensor
) -> torch.Tensor

Rebuild full-vocab logits from the captured pre-norm final hidden state.

Capture id num_hidden_layers yields the pre-(final-norm) residual hidden (verified for vllm==0.23.0), so apply the target’s final RMSNorm + LM head here; if a future vLLM captures post-norm, drop the RMSNorm below.

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner._ensure_weights_loaded(
    device: torch.device
) -> None

Load the final-norm + LM-head + input embeddings once and cache them.

norm is kept in fp32 and lm_head is cast to fp32 and transposed to [hidden, vocab] here so the per-microbatch forward_eagle3 does no redundant cast/transpose (the LM head is multi-GB at large vocab).

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner.build(
    model_path: str,
    dtype: typing.Optional[torch.dtype] = None,
    tp_size: int = 1,
    trust_remote_code: bool = False,
    vllm_kwargs = {}
) -> 'VLLMTargetRunner'

classmethod

Construct the runner for a standalone target server.

vllm_kwargs are forwarded to vllm.LLM (e.g. gpu_memory_utilization, max_model_len, quantization). GPU/vLLM-only.

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner.close() -> None

Release the vLLM engine (best effort).

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner.forward_eagle3(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Run one prefill per row and return (logits, aux_hidden_states).

logits is [batch, seq, vocab] (full vocab, unshifted) rebuilt from the captured final hidden via the target’s final RMSNorm + LM head; aux_hidden_states is [batch, seq, 3 * hidden] (the three capture layers concatenated, unshifted). Sequences must share a length (training batches are right-padded); attention_mask is unused because each row is a full causal prefill in vLLM (trailing pad tokens do not affect earlier positions under causal attention).

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner.input_embedding_weight() -> torch.Tensor

Return the target input-embedding weight [vocab, hidden].

nemo_automodel.components.speculative.eagle.vllm_runner.VLLMTargetRunner.set_aux_layers(
    aux_layer_ids: typing.Sequence[int]
) -> None

Record the 3 capture layers and build the vLLM engine around them.

class nemo_automodel.components.speculative.eagle.vllm_runner._VLLMModelShim(
    hf_config,
    device: torch.device
)

Lightweight stand-in exposing .config + .parameters().

The target’s real parameters live in the vLLM engine process, so this only carries the HF config (for num_hidden_layers / hidden_size / vocab_size) and a single device-marker tensor so the remote server can still infer the target’s device via next(model.parameters()).device.

_marker

= torch.empty(0, device=device)

nemo_automodel.components.speculative.eagle.vllm_runner._VLLMModelShim.parameters()

nemo_automodel.components.speculative.eagle.vllm_runner._load_target_head_weights(
    model_path: str,
    device: torch.device
)

Load (embed, final_norm, lm_head) weights straight from the safetensors.

The real weights live inside the vLLM engine (a separate process), so the final-norm + LM-head used to rebuild logits, and the input embeddings the draft copies, are read directly off disk here instead. lm_head falls back to the input embeddings for tied-embedding models (e.g. Qwen3).

nemo_automodel.components.speculative.eagle.vllm_runner.vllm_dtype_str(
    dtype: typing.Optional[torch.dtype]
) -> str

Map a torch dtype to the string form vLLM’s dtype argument expects.

None means “let vLLM pick” ("auto").

nemo_automodel.components.speculative.eagle.vllm_runner._VLLM_DTYPE_STRINGS = {torch.float32: 'float32', torch.float16: 'float16', torch.bfloat16: 'bfloat16'}

nemo_automodel.components.speculative.eagle.vllm_runner.logger = logging.getLogger(__name__)