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

1. 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.
2. `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`.
3. 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`](#nemo_automodel-components-speculative-eagle-vllm_runner-VLLMTargetRunner) | Offline vLLM engine that returns EAGLE-3 supervision tensors. |
| [`_VLLMModelShim`](#nemo_automodel-components-speculative-eagle-vllm_runner-_VLLMModelShim)     | Lightweight stand-in exposing `.config` + `.parameters()`.    |

### Functions

| Name                                                                                                              | Description                                                                |
| ----------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_load_target_head_weights`](#nemo_automodel-components-speculative-eagle-vllm_runner-_load_target_head_weights) | Load `(embed, final_norm, lm_head)` weights straight from the safetensors. |
| [`vllm_dtype_str`](#nemo_automodel-components-speculative-eagle-vllm_runner-vllm_dtype_str)                       | Map a torch dtype to the string form vLLM's `dtype` argument expects.      |

### Data

[`_VLLM_DTYPE_STRINGS`](#nemo_automodel-components-speculative-eagle-vllm_runner-_VLLM_DTYPE_STRINGS)

[`logger`](#nemo_automodel-components-speculative-eagle-vllm_runner-logger)

### API

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

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

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

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

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

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

Release the vLLM engine (best effort).

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

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

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

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

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

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

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

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

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

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