nemo_automodel.components.speculative.bench_vllm

View as Markdown

Offline acceptance / speedup benchmark for a trained drafter served by vLLM.

This is the vLLM companion to bench_sglang: the same chat-completions workload machinery (which it imports), with the acceptance statistics read from vLLM’s Prometheus /metrics endpoint instead of SGLang’s /server_info. It covers every drafter serve_vllm can launch — EAGLE-3, P-EAGLE, and the DFlash family — and reports:

  • accept_length — mean tokens emitted per target verification step, including the one guaranteed bonus token: 1 + num_accepted_tokens / num_drafts.
  • acceptance_rate — fraction of proposed draft tokens accepted: num_accepted_tokens / num_draft_tokens.
  • output_throughput_tok_s — measured decode throughput (output tokens per wall-clock second).
  • speedup — optional: throughput divided by the same workload’s throughput against a --baseline-server running without speculation.

The spec-decode counters (vllm:spec_decode_num_drafts_total, vllm:spec_decode_num_draft_tokens_total, vllm:spec_decode_num_accepted_tokens_total) are snapshotted before and after the workload and differenced, so — unlike SGLang’s server-cumulative avg_spec_accept_length — the numbers cover exactly this benchmark’s requests even on a server that has already handled other traffic.

Typical usage (after serve_vllm launches the drafter on port 8000):

python -m nemo_automodel.components.speculative.bench_vllm
—server http://localhost:8000
—model Qwen/Qwen3-8B
—input-data Aeala/ShareGPT_Vicuna_unfiltered
—num-prompts 64 —concurrency 16 —max-new-tokens 256

Add --baseline-server http://localhost:8001 (a second server started without --speculative-config) to also report the end-to-end speedup.

vLLM is intentionally NOT a dependency of this script — it talks to the server over HTTP, so only aiohttp is required (already pulled in by the project). The server itself must be running separately; see serve_vllm.

Module Contents

Classes

NameDescription
SpecMetricsOne snapshot of vLLM’s cumulative spec-decode counters.

Functions

NameDescription
_accept_lengthMean tokens per verify step incl. the bonus token: 1 + accepted / drafts.
_acceptance_rateFraction of proposed draft tokens accepted: accepted / draft_tokens.
_build_parser-
_fetch_spec_metricsGET <server>/metrics and parse the spec-decode counters; None on failure.
_parse_spec_metricsParse the spec-decode counters out of a Prometheus exposition payload.
_runAsync driver: compute the benchmark summary and report it. Returns an exit code.
_run_summaryValidate args, run the workload(s), and return the metrics dict.
_summarizeAssemble the metrics dict reported to stdout / --output-json.
mainCLI entry point. Parses argv and returns the process exit code.

Data

_METRIC_NUM_ACCEPTED_TOKENS

_METRIC_NUM_DRAFTS

_METRIC_NUM_DRAFT_TOKENS

_SPEC_METRIC_NAMES

logger

API

class nemo_automodel.components.speculative.bench_vllm.SpecMetrics(
num_drafts: float,
num_draft_tokens: float,
num_accepted_tokens: float
)
Dataclass

One snapshot of vLLM’s cumulative spec-decode counters.

num_accepted_tokens
float
num_draft_tokens
float
num_drafts
float

Return the counter increase since before (clamped at zero per counter).

nemo_automodel.components.speculative.bench_vllm._accept_length(
delta: nemo_automodel.components.speculative.bench_vllm.SpecMetrics | None
) -> float | None

Mean tokens per verify step incl. the bonus token: 1 + accepted / drafts.

nemo_automodel.components.speculative.bench_vllm._acceptance_rate(
delta: nemo_automodel.components.speculative.bench_vllm.SpecMetrics | None
) -> float | None

Fraction of proposed draft tokens accepted: accepted / draft_tokens.

nemo_automodel.components.speculative.bench_vllm._build_parser() -> argparse.ArgumentParser
nemo_automodel.components.speculative.bench_vllm._fetch_spec_metrics(
server: str,
timeout_s: float
) -> nemo_automodel.components.speculative.bench_vllm.SpecMetrics | None
async

GET <server>/metrics and parse the spec-decode counters; None on failure.

nemo_automodel.components.speculative.bench_vllm._parse_spec_metrics(
metrics_text: str
) -> nemo_automodel.components.speculative.bench_vllm.SpecMetrics | None

Parse the spec-decode counters out of a Prometheus exposition payload.

Values are summed across label sets (multiple engines / model labels), and None is returned when the drafts counter is absent — i.e. the server is not running with speculative decoding.

nemo_automodel.components.speculative.bench_vllm._run(
args: argparse.Namespace
) -> int
async

Async driver: compute the benchmark summary and report it. Returns an exit code.

nemo_automodel.components.speculative.bench_vllm._run_summary(
args: argparse.Namespace
) -> dict[str, typing.Any] | None
async

Validate args, run the workload(s), and return the metrics dict.

Returns None when no usable prompts were loaded (the caller’s cue to report a failure without raising — a bad --num-prompts/etc. value is a real programming error and still raises via _validate_workload_args). Split out of _run so bench_sweep can drive one dataset at a time without the printing / --output-json side effects below.

Assemble the metrics dict reported to stdout / --output-json.

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

CLI entry point. Parses argv and returns the process exit code.

nemo_automodel.components.speculative.bench_vllm._METRIC_NUM_ACCEPTED_TOKENS = 'vllm:spec_decode_num_accepted_tokens_total'
nemo_automodel.components.speculative.bench_vllm._METRIC_NUM_DRAFTS = 'vllm:spec_decode_num_drafts_total'
nemo_automodel.components.speculative.bench_vllm._METRIC_NUM_DRAFT_TOKENS = 'vllm:spec_decode_num_draft_tokens_total'
nemo_automodel.components.speculative.bench_vllm._SPEC_METRIC_NAMES = (_METRIC_NUM_DRAFTS, _METRIC_NUM_DRAFT_TOKENS, _METRIC_NUM_ACCEPTED_TOKENS)
nemo_automodel.components.speculative.bench_vllm.logger = logging.getLogger(__name__)