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

Serve an Automodel-trained EAGLE-3 / P-EAGLE / DFlash drafter with vLLM.

This is the vLLM companion to `serve_sglang`. It exists because the
parallel-drafting (P-EAGLE) head produced by the EAGLE-3 recipe with
`parallel_drafting: true` is *not* servable by SGLang today (`serve_sglang`
rejects it); its inference runtime is vLLM's parallel-drafting path
(vLLM >= 0.16, [https://github.com/vllm-project/speculators/pull/480](https://github.com/vllm-project/speculators/pull/480)). The same
entry point also serves a plain (non-parallel) EAGLE-3 draft under vLLM, and a
DFlash-family draft (DFlash / JetSpec; vLLM's `dflash` method landed in 0.20).

The EAGLE drafter checkpoints produced by the recipes are written by the
consolidated checkpointer as an HF-style `model/` directory
(`model.safetensors` + `config.json`), optionally nested one level deeper in
a `consolidated/` subdir. The draft already stores the `d2t` / `t2d`
vocab-remap buffers inside its weights in the exact offset form vLLM expects
(`target_id = draft_id + d2t[draft_id]`), so -- unlike the SGLang path -- no
separate speculative-token-map file is emitted.

Three fixups bridge the Automodel on-disk format to what vLLM loads:

* `architectures`: the recipe writes `["LlamaEagle3DraftModel"]` (the
  Automodel class name), which is not in vLLM's model registry. vLLM routes
  every EAGLE-3 draft on `LlamaForCausalLMEagle3` (`PEagleDraftModel` /
  `Eagle3LlamaForCausalLM` are registered aliases), so the field is rewritten.
* `pard_token` (P-EAGLE only): vLLM's parallel-drafting proposer reads the
  masked-slot token id from `pard_token` / `ptd_token_id` /
  `dflash_config.mask_token_id`, but the recipe only writes the top-level
  `mask_token_id`. The value is copied to `pard_token` (mirroring vLLM's own
  speculators `update_peagle` mapping).
* weight keys: Automodel wraps the draft as `self.model`, so its weights are
  saved with a leading `model.` (`model.embed_tokens.weight`). vLLM's loader
  re-adds `model.` to every non-top-level weight, so the prefix is stripped
  during a one-time export into a `vllm_export/` subdir (the source checkpoint
  is left untouched). A draft whose weights are already vLLM-standard skips the
  export and only gets the config fixups in place.

The script then shells out to `python -m vllm.entrypoints.openai.api_server`
with the right `--speculative-config`.

NOTE -- vLLM is NOT bundled with the NeMo-AutoModel container image and is
intentionally NOT declared in `pyproject.toml`. To use this entry point,
install it yourself into the same environment:

uv pip install "vllm>=0.16"

Refer to [https://github.com/vllm-project/vllm](https://github.com/vllm-project/vllm) for the version matching your
CUDA / PyTorch stack. If vLLM is missing this script exits with a clear install
hint rather than crashing on import.

Typical usage (after training produces a P-EAGLE checkpoint at
`./outputs/peagle/checkpoints/epoch_2_step_44326`):

python -m nemo\_automodel.components.speculative.serve\_vllm \
\--target Qwen/Qwen3-8B \
\--draft ./outputs/peagle/checkpoints/epoch\_2\_step\_44326 \
\--num-speculative-tokens 8

`--num-speculative-tokens` defaults to the draft config's `num_depths` (K),
so for a P-EAGLE head it can be omitted. Pass `--print-only` to inspect the
command without launching it; in that mode the on-disk `architectures` rewrite
is skipped and the printed paths reflect what a real launch would produce.

DFlash family: a draft trained by the DFlash / JetSpec recipes (the config's
`architectures` is `["Qwen3DFlashDraftModel"]`) is auto-detected and served
with vLLM's `dflash` method -- `--method` can be omitted. Its weights carry
no `model.` wrapper prefix and no embed/lm\_head (vLLM shares the target's), so
only the `architectures` rewrite (to vLLM's registered `DFlashDraftModel`)
applies, and K defaults to `block_size - 1` (a block is one anchor token plus
`block_size - 1` mask slots). A JetSpec draft is the same on-disk format
trained with *causal* in-block attention: the recipe stamps
`dflash_config.causal=true` so vLLM matches it at inference; for a JetSpec
checkpoint trained before that stamp existed, pass `--dflash-causal`. A Domino
draft is rejected: its GRU correction head (`prefix_gru` / `embed_proj`) has
no vLLM runtime.

## Module Contents

### Functions

| Name                                                                                                                   | Description                                                                                       |
| ---------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------- |
| [`_build_speculative_config`](#nemo_automodel-components-speculative-serve_vllm-_build_speculative_config)             | Build the vLLM `--speculative-config` dict for a resolved draft.                                  |
| [`_check_vllm_available`](#nemo_automodel-components-speculative-serve_vllm-_check_vllm_available)                     | Verify the `vllm` package can actually be imported, else exit (code 2).                           |
| [`_dflash_config_updates`](#nemo_automodel-components-speculative-serve_vllm-_dflash_config_updates)                   | Return the `config.json` key changes needed for vLLM to load a DFlash-family draft.               |
| [`_draft_weight_keys`](#nemo_automodel-components-speculative-serve_vllm-_draft_weight_keys)                           | Return the weight tensor names in `draft_dir`.                                                    |
| [`_export_for_vllm`](#nemo_automodel-components-speculative-serve_vllm-_export_for_vllm)                               | Materialize a vLLM-loadable copy of `draft_dir` in a `vllm_export/` subdir.                       |
| [`_export_is_fresh`](#nemo_automodel-components-speculative-serve_vllm-_export_is_fresh)                               | True when `export_dir` holds a complete config + weights newer than the source weights or config. |
| [`_find_draft_dir`](#nemo_automodel-components-speculative-serve_vllm-_find_draft_dir)                                 | Return the HF-style drafter directory under `draft_path`, if present.                             |
| [`_has_hf_weight_file`](#nemo_automodel-components-speculative-serve_vllm-_has_hf_weight_file)                         | Return True if `path` contains an HF-style weight artifact.                                       |
| [`_is_dflash_config`](#nemo_automodel-components-speculative-serve_vllm-_is_dflash_config)                             | True when the draft config identifies a DFlash-family draft (DFlash / Domino / JetSpec).          |
| [`_load_config`](#nemo_automodel-components-speculative-serve_vllm-_load_config)                                       | Load a drafter `config.json` into a dict.                                                         |
| [`_load_safetensors_io`](#nemo_automodel-components-speculative-serve_vllm-_load_safetensors_io)                       | Return `(load_file, save_file)` from `safetensors.torch` or exit with a hint.                     |
| [`_needs_weight_remap`](#nemo_automodel-components-speculative-serve_vllm-_needs_weight_remap)                         | True when any weight carries the Automodel `self.model` wrapper prefix.                           |
| [`_parse_args`](#nemo_automodel-components-speculative-serve_vllm-_parse_args)                                         | Parse command-line arguments for the vLLM serve helper.                                           |
| [`_remap_key_for_vllm`](#nemo_automodel-components-speculative-serve_vllm-_remap_key_for_vllm)                         | Strip the `self.model` wrapper prefix so vLLM's loader (which re-adds it) resolves it.            |
| [`_resolve_method`](#nemo_automodel-components-speculative-serve_vllm-_resolve_method)                                 | Resolve the vLLM speculative method: explicit `--method`, else detect from the draft config.      |
| [`_resolve_num_speculative_tokens`](#nemo_automodel-components-speculative-serve_vllm-_resolve_num_speculative_tokens) | Resolve K: CLI value, else `block_size - 1` (DFlash) / `num_depths` (EAGLE).                      |
| [`_rewrite_config_for_vllm`](#nemo_automodel-components-speculative-serve_vllm-_rewrite_config_for_vllm)               | Patch `config_path` in place with the keys vLLM needs (see `_vllm_config_updates`).               |
| [`_vllm_config_updates`](#nemo_automodel-components-speculative-serve_vllm-_vllm_config_updates)                       | Return the `config.json` key changes needed for vLLM to load this draft.                          |
| [`build_vllm_argv`](#nemo_automodel-components-speculative-serve_vllm-build_vllm_argv)                                 | Build the `python -m vllm.entrypoints.openai.api_server` argv for a config.                       |
| [`main`](#nemo_automodel-components-speculative-serve_vllm-main)                                                       | Validate the environment, resolve the drafter ckpt, then exec vLLM.                               |
| [`resolve_draft_artifacts`](#nemo_automodel-components-speculative-serve_vllm-resolve_draft_artifacts)                 | Resolve a user-supplied drafter path to the directory and config vLLM expects.                    |

### Data

[`_AUTOMODEL_DFLASH_ARCHITECTURE`](#nemo_automodel-components-speculative-serve_vllm-_AUTOMODEL_DFLASH_ARCHITECTURE)

[`_AUTOMODEL_DRAFT_ARCHITECTURE`](#nemo_automodel-components-speculative-serve_vllm-_AUTOMODEL_DRAFT_ARCHITECTURE)

[`_MODEL_PREFIX`](#nemo_automodel-components-speculative-serve_vllm-_MODEL_PREFIX)

[`_SAFETENSORS_INSTALL_HINT`](#nemo_automodel-components-speculative-serve_vllm-_SAFETENSORS_INSTALL_HINT)

[`_VLLM_DFLASH_ARCHITECTURE`](#nemo_automodel-components-speculative-serve_vllm-_VLLM_DFLASH_ARCHITECTURE)

[`_VLLM_DRAFT_ARCHITECTURE`](#nemo_automodel-components-speculative-serve_vllm-_VLLM_DRAFT_ARCHITECTURE)

[`_VLLM_EXPORT_DIRNAME`](#nemo_automodel-components-speculative-serve_vllm-_VLLM_EXPORT_DIRNAME)

[`_VLLM_INSTALL_HINT`](#nemo_automodel-components-speculative-serve_vllm-_VLLM_INSTALL_HINT)

[`logger`](#nemo_automodel-components-speculative-serve_vllm-logger)

### API

```python
nemo_automodel.components.speculative.serve_vllm._build_speculative_config(
    args: argparse.Namespace,
    draft_path: str,
    config: dict[str, typing.Any]
) -> dict[str, typing.Any]
```

Build the vLLM `--speculative-config` dict for a resolved draft.

`parallel_drafting` is a SpeculativeConfig field that vLLM defaults to
False; it must be set here for a P-EAGLE head (the draft's own
`config.json` flag is not auto-promoted to the speculative config). A
DFlash draft needs no such flag: vLLM forces parallel drafting for the
`dflash` method itself.

```python
nemo_automodel.components.speculative.serve_vllm._check_vllm_available() -> None
```

Verify the `vllm` package can actually be imported, else exit (code 2).

```python
nemo_automodel.components.speculative.serve_vllm._dflash_config_updates(
    config: dict[str, typing.Any],
    dflash_causal: bool
) -> dict[str, typing.Any]
```

Return the `config.json` key changes needed for vLLM to load a DFlash-family draft.

Only the `architectures` rewrite (to vLLM's registered `DFlashDraftModel`)
is required: the draft weights carry no `model.` prefix and vLLM reads
`dflash_config.mask_token_id` / `dflash_config.target_layer_ids` where the
recipe already writes them. `--dflash-causal` additionally stamps
`dflash_config.causal=true` (a JetSpec draft trained before the recipe
stamped it).

```python
nemo_automodel.components.speculative.serve_vllm._draft_weight_keys(
    draft_dir: pathlib.Path
) -> list[str]
```

Return the weight tensor names in `draft_dir`.

Reads the names from `model.safetensors.index.json` when sharded, otherwise
from the single `model.safetensors` header (no tensor data is loaded).
Returns an empty list when neither is present / readable.

```python
nemo_automodel.components.speculative.serve_vllm._export_for_vllm(
    draft_dir: pathlib.Path,
    config: dict[str, typing.Any]
) -> pathlib.Path
```

Materialize a vLLM-loadable copy of `draft_dir` in a `vllm_export/` subdir.

The safetensors keys are remapped to strip the Automodel `self.model`
wrapper prefix, the config receives the architectures / pard\_token fixups, and
the tokenizer assets are copied alongside. The source checkpoint is left
untouched. The export is cached and rebuilt only when stale.

```python
nemo_automodel.components.speculative.serve_vllm._export_is_fresh(
    draft_dir: pathlib.Path,
    export_dir: pathlib.Path
) -> bool
```

True when `export_dir` holds a complete config + weights newer than the source weights or config.

```python
nemo_automodel.components.speculative.serve_vllm._find_draft_dir(
    draft_path: pathlib.Path
) -> pathlib.Path | None
```

Return the HF-style drafter directory under `draft_path`, if present.

Accepts the outer `epoch_<E>_step_<S>` directory, its `model/` subdir, or
the consolidated checkpointer's `model/consolidated/` (and a bare
`consolidated/`) layout, as well as a directory that is already the HF
model dir. Returns the first candidate holding `config.json` + a weight
file, else `None`.

```python
nemo_automodel.components.speculative.serve_vllm._has_hf_weight_file(
    path: pathlib.Path
) -> bool
```

Return True if `path` contains an HF-style weight artifact.

```python
nemo_automodel.components.speculative.serve_vllm._is_dflash_config(
    config: dict[str, typing.Any]
) -> bool
```

True when the draft config identifies a DFlash-family draft (DFlash / Domino / JetSpec).

```python
nemo_automodel.components.speculative.serve_vllm._load_config(
    config_path: pathlib.Path
) -> dict[str, typing.Any]
```

Load a drafter `config.json` into a dict.

```python
nemo_automodel.components.speculative.serve_vllm._load_safetensors_io()
```

Return `(load_file, save_file)` from `safetensors.torch` or exit with a hint.

```python
nemo_automodel.components.speculative.serve_vllm._needs_weight_remap(
    weight_keys: list[str]
) -> bool
```

True when any weight carries the Automodel `self.model` wrapper prefix.

```python
nemo_automodel.components.speculative.serve_vllm._parse_args(
    argv: list[str] | None = None
) -> argparse.Namespace
```

Parse command-line arguments for the vLLM serve helper.

```python
nemo_automodel.components.speculative.serve_vllm._remap_key_for_vllm(
    key: str
) -> str
```

Strip the `self.model` wrapper prefix so vLLM's loader (which re-adds it) resolves it.

Top-level tensors (`lm_head.weight`, `d2t`, `t2d`, `mask_hidden`) do
not carry the prefix and are returned unchanged.

```python
nemo_automodel.components.speculative.serve_vllm._resolve_method(
    cli_method: str | None,
    config: dict[str, typing.Any]
) -> str
```

Resolve the vLLM speculative method: explicit `--method`, else detect from the draft config.

```python
nemo_automodel.components.speculative.serve_vllm._resolve_num_speculative_tokens(
    cli_value: int | None,
    config: dict[str, typing.Any],
    method: str
) -> int
```

Resolve K: CLI value, else `block_size - 1` (DFlash) / `num_depths` (EAGLE).

```python
nemo_automodel.components.speculative.serve_vllm._rewrite_config_for_vllm(
    config_path: pathlib.Path,
    config: dict[str, typing.Any],
    dflash_causal: bool = False
) -> None
```

Patch `config_path` in place with the keys vLLM needs (see `_vllm_config_updates`).

Takes the already-parsed `config` so the caller's read is not repeated.
No-op when the config already satisfies vLLM. The write is staged through a
sibling `.tmp` file and finalized with `os.replace` so an interrupted
write cannot leave the destination half-truncated.

```python
nemo_automodel.components.speculative.serve_vllm._vllm_config_updates(
    config: dict[str, typing.Any],
    dflash_causal: bool = False
) -> dict[str, typing.Any]
```

Return the `config.json` key changes needed for vLLM to load this draft.

A DFlash-family draft gets the DFlash fixups (see `_dflash_config_updates`).
For an EAGLE draft, two fixups bridge the Automodel on-disk format to what
vLLM reads:

* `architectures` -> the vLLM-canonical `LlamaForCausalLMEagle3` (the
  recipe writes the Automodel class name `LlamaEagle3DraftModel`, which
  vLLM's registry does not know).
* for a P-EAGLE head, `pard_token` (= `mask_token_id`). This mirrors
  vLLM's own speculators `update_peagle` mapping; the parallel-drafting
  proposer reads `pard_token` / `ptd_token_id` /
  `dflash_config.mask_token_id`, none of which the recipe writes -- it
  only writes the top-level `mask_token_id`.

Returns an empty dict when the config already satisfies vLLM.

```python
nemo_automodel.components.speculative.serve_vllm.build_vllm_argv(
    args: argparse.Namespace
) -> list[str]
```

Build the `python -m vllm.entrypoints.openai.api_server` argv for a config.

```python
nemo_automodel.components.speculative.serve_vllm.main(
    argv: list[str] | None = None
) -> int
```

Validate the environment, resolve the drafter ckpt, then exec vLLM.

Returns the vLLM server's exit code, or `2` if vLLM is missing.

```python
nemo_automodel.components.speculative.serve_vllm.resolve_draft_artifacts(
    draft: str,
    dry_run: bool = False,
    dflash_causal: bool = False
) -> tuple[str, dict[str, typing.Any]]
```

Resolve a user-supplied drafter path to the directory and config vLLM expects.

**Parameters:**

A local path to a recipe checkpoint (the `epoch_*` dir, its
`model/` / `model/consolidated/` subdir, or the HF model dir
itself) or a Hugging Face Hub repo id.

When True, no on-disk `architectures` rewrite is performed and
the returned paths reflect what a real launch would produce.

When True, stamp `dflash_config.causal=true` on a
DFlash-family draft (a JetSpec checkpoint trained before the recipe
stamped it).

**Returns:** `str`

`(draft_path, config)` where `draft_path` is what vLLM's

```python
nemo_automodel.components.speculative.serve_vllm._AUTOMODEL_DFLASH_ARCHITECTURE = 'Qwen3DFlashDraftModel'
```

```python
nemo_automodel.components.speculative.serve_vllm._AUTOMODEL_DRAFT_ARCHITECTURE = 'LlamaEagle3DraftModel'
```

```python
nemo_automodel.components.speculative.serve_vllm._MODEL_PREFIX = 'model.'
```

```python
nemo_automodel.components.speculative.serve_vllm._SAFETENSORS_INSTALL_HINT = 'safetensors is required to remap an Automodel drafter checkpoint for vLLM. Inst...
```

```python
nemo_automodel.components.speculative.serve_vllm._VLLM_DFLASH_ARCHITECTURE = 'DFlashDraftModel'
```

```python
nemo_automodel.components.speculative.serve_vllm._VLLM_DRAFT_ARCHITECTURE = 'LlamaForCausalLMEagle3'
```

```python
nemo_automodel.components.speculative.serve_vllm._VLLM_EXPORT_DIRNAME = 'vllm_export'
```

```python
nemo_automodel.components.speculative.serve_vllm._VLLM_INSTALL_HINT = 'vllm is not installed in this environment. Install it manually with `uv pip ins...
```

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