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

Data

_AUTOMODEL_DFLASH_ARCHITECTURE

_AUTOMODEL_DRAFT_ARCHITECTURE

_MODEL_PREFIX

_SAFETENSORS_INSTALL_HINT

_VLLM_DFLASH_ARCHITECTURE

_VLLM_DRAFT_ARCHITECTURE

_VLLM_EXPORT_DIRNAME

_VLLM_INSTALL_HINT

logger

API

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.

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

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

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

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.

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.

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.

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.

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

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

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

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

Load a drafter config.json into a dict.

nemo_automodel.components.speculative.serve_vllm._load_safetensors_io()

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

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.

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

Parse command-line arguments for the vLLM serve helper.

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.

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.

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

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.

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.

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.

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.

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:

draft

str

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.

dry_run

boolDefaults to False

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

dflash_causal

boolDefaults to False

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

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

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

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

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

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

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

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

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

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