[Rust][Benchmark] Align speed-bench CLI flags with Python and add flag parity test (#51592)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>

.buildkite/test_areas/benchmarks.yaml

@@ -9,6 +9,7 @@ steps:
   source_file_dependencies:
   - vllm/
   - "!vllm/distributed/kv_transfer/"
+  - rust/src/bench/tests/python_serve_flags.txt
   - tests/benchmarks/
   commands:
   - pytest -v -s benchmarks/

rust/src/bench/README.md

@@ -483,7 +483,7 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below.
 | `--model` | Auto-detect | Model name (fetched from `/v1/models` if omitted) |
 | `--served-model-name` | — | Model name used in API requests |
 | `--tokenizer` | Same as model | Tokenizer name or path (supports HF, tiktoken, server fallback) |
-| `--tokenizer-mode` | `auto` | Tokenizer mode (`auto`, `hf`, `slow`, `mistral`) |
+| `--tokenizer-mode` | `auto` | Accepted for Python CLI compatibility but ignored (warns on non-`auto`): resolution always follows the HF `tokenizer.json` → tiktoken → server-side `/tokenize` chain, so `mistral_common` (tekken) tokenizers fall back to server-side tokenization |
 | `--trust-remote-code` | `false` | Trust remote code for tokenizer |
 | `--skip-tokenizer-init` | `false` | Skip tokenizer initialization |
 
@@ -521,7 +521,8 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below.
 | `--sonnet-output-len` | `150` | Output tokens per request |
 | `--sonnet-prefix-len` | `200` | Prefix tokens shared across requests |
 | **SPEED-Bench** | | |
-| `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`) |
+| `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`); alias: `--speed-bench-dataset-subset` (Python name) |
+| `--speed-bench-output-len` | `4096` | Output tokens per request (matches Python default) |
 | `--speed-bench-category` | — | Filter by category (`low_entropy`, `high_entropy`, `mixed_entropy`, `coding`, `math`, …) |
 | `--speed-bench-max-input-len` | — | Truncate prompts to at most N tokens |
 | **HuggingFace** | | |
@@ -794,6 +795,11 @@ The Rust implementation matches Python `vllm bench serve` in:
 - Rate control (Gamma distribution, normalization, burstiness, linear/exponential ramp-up)
 - Metrics (TTFT/TPOT/ITL/E2EL percentiles, peak tokens/sec, peak concurrency, goodput)
 - Sampling parameters merged into the request body via `extra_body` (same precedence rules)
+- CLI flags: a superset of Python's, enforced by `tests/cli_parity.rs` against a
+  snapshot of the Python parser (`tests/python_serve_flags.txt`, regenerated by
+  `tests/benchmarks/test_rust_bench_cli_parity.py` in the repo root) with an
+  explicit allowlist for the Python-only remainder — so `VLLM_USE_RUST_BENCH=1`
+  delegation cannot silently reject documented flags
 
 ## Environment Variables
 

rust/src/bench/src/benchmark.rs

@@ -134,6 +134,17 @@ pub(crate) async fn fetch_spec_decode_metrics(
         Err(_) => return None,
     };
 
+    parse_spec_decode_metrics(&text)
+}
+
+/// Parse spec decode counters from Prometheus text exposition format.
+///
+/// Matches on the metric name (the part before any labels) rather than the
+/// whole line so label values cannot be mistaken for metric names, and only
+/// reads `_total` counter samples — skipping e.g. `_created` timestamp series,
+/// which would otherwise be summed into the counts (mirrors the Python fix
+/// in vllm#41916).
+pub(crate) fn parse_spec_decode_metrics(text: &str) -> Option<SpecDecodeMetrics> {
     let mut num_drafts: u64 = 0;
     let mut num_draft_tokens: u64 = 0;
     let mut num_accepted_tokens: u64 = 0;
@@ -145,24 +156,25 @@ pub(crate) async fn fetch_spec_decode_metrics(
         if line.is_empty() || line.starts_with('#') {
             continue;
         }
-        if !line.starts_with("vllm:spec_decode") {
+        let first_token = match line.split_whitespace().next() {
+            Some(t) => t,
+            None => continue,
+        };
+        let metric_name = first_token.split('{').next().unwrap_or(first_token);
+        if !metric_name.starts_with("vllm:spec_decode") || !metric_name.ends_with("_total") {
             continue;
         }
         found_spec_decode = true;
-        let parts: Vec<&str> = line.split_whitespace().collect();
-        if parts.is_empty() {
-            continue;
-        }
-        let val = match parts.last().and_then(|s| s.parse::<f64>().ok()) {
+        let val = match line.split_whitespace().last().and_then(|s| s.parse::<f64>().ok()) {
             Some(v) => v as u64,
             None => continue,
         };
 
-        if line.contains("num_drafts") {
+        if metric_name.contains("num_drafts") {
             num_drafts += val;
-        } else if line.contains("num_draft_tokens") {
+        } else if metric_name.contains("num_draft_tokens") {
             num_draft_tokens += val;
-        } else if line.contains("num_accepted_tokens_per_pos") {
+        } else if metric_name.contains("num_accepted_tokens_per_pos") {
             // Parse position label: position="N"
             if let Some(start) = line.find("position=\"") {
                 let start = start + "position=\"".len();
@@ -172,7 +184,7 @@ pub(crate) async fn fetch_spec_decode_metrics(
                     *accepted_per_pos.entry(pos).or_insert(0) += val;
                 }
             }
-        } else if line.contains("num_accepted_tokens") {
+        } else if metric_name.contains("num_accepted_tokens") {
             num_accepted_tokens += val;
         }
     }
@@ -502,12 +514,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
                     downloaded.as_str()
                 }
             };
-            let output_len = config.sharegpt_output_len.unwrap_or(config.random_output_len);
             crate::datasets::speed_bench::load_speed_bench_dataset(
                 tok,
                 path,
                 config.num_prompts,
-                output_len,
+                config.speed_bench_output_len,
                 config.seed,
                 &config.request_id_prefix,
                 config.speed_bench_category.as_deref(),
@@ -1875,6 +1886,52 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_parse_spec_decode_metrics_sums_counters_across_engines() {
+        let text = "\
+# TYPE vllm:spec_decode_num_drafts_total counter
+vllm:spec_decode_num_drafts_total{model_name=\"m\",engine=\"0\"} 100.0
+vllm:spec_decode_num_drafts_total{model_name=\"m\",engine=\"1\"} 50.0
+vllm:spec_decode_num_draft_tokens_total{model_name=\"m\",engine=\"0\"} 700.0
+vllm:spec_decode_num_accepted_tokens_total{model_name=\"m\",engine=\"0\"} 300.0
+vllm:spec_decode_num_accepted_tokens_per_pos_total{position=\"0\",engine=\"0\"} 80.0
+vllm:spec_decode_num_accepted_tokens_per_pos_total{position=\"1\",engine=\"0\"} 40.0
+vllm:num_requests_running{model_name=\"m\"} 0
+";
+        let m = parse_spec_decode_metrics(text).unwrap();
+        assert_eq!(m.num_drafts, 150);
+        assert_eq!(m.num_draft_tokens, 700);
+        assert_eq!(m.num_accepted_tokens, 300);
+        assert_eq!(m.accepted_per_pos, [(0, 80), (1, 40)].into_iter().collect());
+    }
+
+    #[test]
+    fn test_parse_spec_decode_metrics_skips_created_series_and_label_values() {
+        // `_created` timestamps must not be summed into counters, and metric
+        // matching must key off the metric name, not substrings inside label
+        // values (vllm#41916).
+        let text = "\
+vllm:spec_decode_num_drafts_total{model_name=\"m\"} 100.0
+vllm:spec_decode_num_drafts_created{model_name=\"m\"} 1.7863e+09
+vllm:spec_decode_num_draft_tokens_total{model_name=\"m\"} 700.0
+vllm:spec_decode_num_draft_tokens_created{model_name=\"m\"} 1.7863e+09
+vllm:spec_decode_num_accepted_tokens_total{model_name=\"num_drafts\"} 300.0
+";
+        let m = parse_spec_decode_metrics(text).unwrap();
+        assert_eq!(m.num_drafts, 100);
+        assert_eq!(m.num_draft_tokens, 700);
+        assert_eq!(m.num_accepted_tokens, 300);
+    }
+
+    #[test]
+    fn test_parse_spec_decode_metrics_none_without_spec_metrics() {
+        let text = "\
+vllm:num_requests_running{model_name=\"m\"} 0
+vllm:prefix_cache_queries_total{model_name=\"m\"} 0
+";
+        assert!(parse_spec_decode_metrics(text).is_none());
+    }
+
     #[test]
     fn test_compute_spec_decode_stats_basic() {
         let before = make_metrics(100, 300, 200, &[(0, 90), (1, 70), (2, 50)]);

rust/src/bench/src/cli.rs

@@ -176,7 +176,8 @@ pub struct BenchServeArgs {
     #[arg(long)]
     pub tokenizer: Option<String>,
 
-    /// Tokenizer mode (auto, hf, slow, mistral).
+    /// Tokenizer mode (auto, hf, slow, mistral). Accepted for Python CLI
+    /// compatibility; non-auto values are ignored with a warning.
     #[arg(long, default_value = "auto")]
     pub tokenizer_mode: String,
 
@@ -535,9 +536,17 @@ pub struct BenchServeArgs {
 
     /// SPEED-Bench config/split (qualitative, throughput_1k, throughput_2k, throughput_8k,
     /// throughput_16k, throughput_32k).
-    #[arg(long, default_value = "qualitative")]
+    #[arg(
+        long,
+        visible_alias = "speed-bench-dataset-subset",
+        default_value = "qualitative"
+    )]
     pub speed_bench_config: SpeedBenchConfig,
 
+    /// Number of output tokens per request (SPEED-Bench dataset).
+    #[arg(long, default_value_t = 4096)]
+    pub speed_bench_output_len: usize,
+
     /// Filter SPEED-Bench by category (e.g. low_entropy, high_entropy, coding, math).
     #[arg(long)]
     pub speed_bench_category: Option<String>,

rust/src/bench/src/config.rs

@@ -196,6 +196,7 @@ pub struct BenchConfig {
     pub speed_bench_config: SpeedBenchConfig,
     pub speed_bench_category: Option<String>,
     pub speed_bench_max_input_len: Option<usize>,
+    pub speed_bench_output_len: usize,
     pub hf_split: Option<String>,
     pub hf_subset: Option<String>,
     pub hf_output_len: Option<usize>,
@@ -640,6 +641,15 @@ impl BenchConfig {
             ));
         }
 
+        if args.tokenizer_mode != "auto" {
+            tracing::warn!(
+                mode = %args.tokenizer_mode,
+                "--tokenizer-mode is ignored by the Rust client; tokenizer resolution always \
+                 follows the HF tokenizer.json -> tiktoken -> server-side /tokenize fallback \
+                 chain (mistral_common tokenizers are not supported locally)"
+            );
+        }
+
         Ok(BenchConfig {
             backend: args.backend,
             base_url,
@@ -727,6 +737,7 @@ impl BenchConfig {
             speed_bench_config: args.speed_bench_config,
             speed_bench_category: args.speed_bench_category.clone(),
             speed_bench_max_input_len: args.speed_bench_max_input_len,
+            speed_bench_output_len: args.speed_bench_output_len,
             hf_split: args.hf_split.clone(),
             hf_subset: args.hf_subset.clone(),
             hf_output_len: args.hf_output_len,
@@ -881,6 +892,22 @@ mod tests {
         ]
     }
 
+    #[test]
+    fn test_speed_bench_flags_match_python() {
+        let args = parse_args(vec![
+            "vllm-bench",
+            "--model",
+            "test-model",
+            "--speed-bench-dataset-subset",
+            "throughput_8k",
+        ]);
+        assert!(matches!(
+            args.speed_bench_config,
+            crate::cli::SpeedBenchConfig::Throughput8k
+        ));
+        assert_eq!(args.speed_bench_output_len, 4096);
+    }
+
     #[test]
     fn test_prefix_sharing_defaults_to_zero() {
         let args = base_multi_turn_args();

rust/src/bench/tests/cli_parity.rs

@@ -0,0 +1,119 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Flag parity with the Python `vllm bench serve` parser.
+//!
+//! `VLLM_USE_RUST_BENCH=1` makes `vllm bench serve` exec this crate's binary
+//! with the raw argv, so every documented Python flag must either parse here
+//! or be a deliberate, allowlisted gap. The snapshot is kept current by
+//! `tests/benchmarks/test_rust_bench_cli_parity.py` in the repo root.
+
+use std::collections::HashSet;
+
+use clap::{CommandFactory, Parser};
+use vllm_bench::BenchServeArgs;
+
+#[derive(Parser)]
+struct TestCli {
+    #[command(flatten)]
+    args: BenchServeArgs,
+}
+
+const PYTHON_FLAGS: &str = include_str!("python_serve_flags.txt");
+
+/// Python-only flags: `vllm bench serve` features vllm-bench deliberately does
+/// not implement. Adding an entry is a decision to let the flag fail under
+/// Rust delegation; remove the entry once the flag gains Rust support.
+const PYTHON_ONLY: &[&str] = &[
+    // asr dataset
+    "--asr-max-audio-len-sec",
+    "--asr-min-audio-len-sec",
+    // bfcl dataset
+    "--bfcl-categories",
+    // blazedit dataset
+    "--blazedit-max-distance",
+    "--blazedit-min-distance",
+    // spec_bench dataset
+    "--spec-bench-category",
+    "--spec-bench-output-len",
+    // timed-trace dataset
+    "--timed-trace-chunk-hash-size",
+    "--timed-trace-label-hash-ids",
+    "--timed-trace-label-input-length",
+    "--timed-trace-label-output-length",
+    "--timed-trace-label-timestamp",
+    "--timed-trace-sec-multiplier",
+    // client-side chat templating / request shaping
+    "--chat-template-kwargs",
+    "--custom-ensure-client-side-data",
+    "--use-beam-search",
+    // result post-processing and plotting
+    "--plot-dataset-stats",
+    "--plot-timeline",
+    "--timeline-itl-thresholds",
+    // misc python-side controls
+    "--hf-name",
+    "--no-self-timed",
+    "--self-timed",
+    "--no-stream",
+    "--probe-request-rate",
+];
+
+fn rust_flags() -> HashSet<String> {
+    let cmd = TestCli::command();
+    let mut flags = HashSet::new();
+    for arg in cmd.get_arguments() {
+        if let Some(long) = arg.get_long() {
+            flags.insert(format!("--{long}"));
+        }
+        if let Some(aliases) = arg.get_all_aliases() {
+            for alias in aliases {
+                flags.insert(format!("--{alias}"));
+            }
+        }
+    }
+    flags
+}
+
+fn python_flags() -> Vec<&'static str> {
+    PYTHON_FLAGS
+        .lines()
+        .map(str::trim)
+        .filter(|line| !line.is_empty() && !line.starts_with('#'))
+        .collect()
+}
+
+#[test]
+fn python_serve_flags_parse_or_are_allowlisted() {
+    let known = rust_flags();
+    let missing: Vec<&str> = python_flags()
+        .into_iter()
+        .filter(|flag| !known.contains(*flag) && !PYTHON_ONLY.contains(flag))
+        .collect();
+    assert!(
+        missing.is_empty(),
+        "Python `vllm bench serve` flags unknown to vllm-bench: {missing:?}. \
+         Support them in src/cli.rs (a clap alias is enough for renames) or \
+         add them to PYTHON_ONLY above as a deliberate gap."
+    );
+}
+
+#[test]
+fn python_only_allowlist_is_not_stale() {
+    let known = rust_flags();
+    let stale: Vec<&&str> = PYTHON_ONLY.iter().filter(|f| known.contains(**f)).collect();
+    assert!(
+        stale.is_empty(),
+        "PYTHON_ONLY entries now supported by vllm-bench, remove them: {stale:?}"
+    );
+}
+
+#[test]
+fn python_only_entries_exist_in_snapshot() {
+    let snapshot: HashSet<&str> = python_flags().into_iter().collect();
+    let gone: Vec<&&str> = PYTHON_ONLY.iter().filter(|f| !snapshot.contains(**f)).collect();
+    assert!(
+        gone.is_empty(),
+        "PYTHON_ONLY entries no longer exist in the Python parser, remove them: {gone:?}"
+    );
+}

rust/src/bench/tests/python_serve_flags.txt

@@ -0,0 +1,110 @@
+# Long CLI flags of the Python `vllm bench serve` parser
+# (vllm.benchmarks.serve.add_cli_args), sorted; --help excluded.
+# Regenerate: .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py
+# Consumed by rust/src/bench/tests/cli_parity.rs, which requires every flag
+# below to be a known vllm-bench flag/alias or a PYTHON_ONLY allowlist entry.
+--append-result
+--asr-max-audio-len-sec
+--asr-min-audio-len-sec
+--backend
+--base-url
+--bfcl-categories
+--blazedit-max-distance
+--blazedit-min-distance
+--burstiness
+--chat-template-kwargs
+--custom-ensure-client-side-data
+--custom-output-len
+--dataset-name
+--dataset-path
+--disable-shuffle
+--disable-tqdm
+--enable-multimodal-chat
+--endpoint
+--extra-body
+--frequency-penalty
+--goodput
+--header
+--hf-name
+--hf-output-len
+--hf-split
+--hf-subset
+--host
+--ignore-eos
+--input-len
+--insecure
+--label
+--logprobs
+--lora-assignment
+--lora-modules
+--max-concurrency
+--metadata
+--metric-percentiles
+--min-p
+--model
+--no-oversample
+--no-reranker
+--no-self-timed
+--no-stream
+--num-prompts
+--num-warmups
+--output-len
+--percentile-metrics
+--plot-dataset-stats
+--plot-timeline
+--port
+--prefix-repetition-num-prefixes
+--prefix-repetition-output-len
+--prefix-repetition-prefix-len
+--prefix-repetition-suffix-len
+--presence-penalty
+--probe-request-rate
+--profile
+--ramp-up-end-rps
+--ramp-up-start-rps
+--ramp-up-strategy
+--random-batch-size
+--random-input-len
+--random-mm-base-items-per-request
+--random-mm-bucket-config
+--random-mm-limit-mm-per-prompt
+--random-mm-num-mm-items-range-ratio
+--random-output-len
+--random-prefix-len
+--random-range-ratio
+--ready-check-timeout-sec
+--repetition-penalty
+--request-id-prefix
+--request-rate
+--result-dir
+--result-filename
+--save-detailed
+--save-result
+--seed
+--self-timed
+--served-model-name
+--sharegpt-output-len
+--skip-chat-template
+--skip-tokenizer-init
+--sonnet-input-len
+--sonnet-output-len
+--sonnet-prefix-len
+--spec-bench-category
+--spec-bench-output-len
+--speed-bench-category
+--speed-bench-dataset-subset
+--speed-bench-output-len
+--temperature
+--timed-trace-chunk-hash-size
+--timed-trace-label-hash-ids
+--timed-trace-label-input-length
+--timed-trace-label-output-length
+--timed-trace-label-timestamp
+--timed-trace-sec-multiplier
+--timeline-itl-thresholds
+--tokenizer
+--tokenizer-mode
+--top-k
+--top-p
+--trust-remote-code
+--use-beam-search

tests/benchmarks/test_rust_bench_cli_parity.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Keep the `vllm bench serve` flag snapshot in sync for the Rust CLI parity test.
+
+`rust/src/bench/tests/cli_parity.rs` asserts that every flag in the snapshot is
+either a known `vllm-bench` clap flag/alias or explicitly allowlisted as
+Python-only, so `VLLM_USE_RUST_BENCH=1` delegation cannot silently reject
+documented flags. This test keeps the snapshot itself current.
+
+Run this file directly to regenerate the snapshot:
+
+    .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py
+"""
+
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+SNAPSHOT = REPO_ROOT / "rust" / "src" / "bench" / "tests" / "python_serve_flags.txt"
+HEADER = """\
+# Long CLI flags of the Python `vllm bench serve` parser
+# (vllm.benchmarks.serve.add_cli_args), sorted; --help excluded.
+# Regenerate: .venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py
+# Consumed by rust/src/bench/tests/cli_parity.rs, which requires every flag
+# below to be a known vllm-bench flag/alias or a PYTHON_ONLY allowlist entry.
+"""
+
+
+def _current_flags() -> list[str]:
+    from vllm.benchmarks.serve import add_cli_args
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+    parser = FlexibleArgumentParser()
+    add_cli_args(parser)
+    return sorted(
+        {
+            opt
+            for action in parser._actions
+            for opt in action.option_strings
+            if opt.startswith("--") and opt != "--help"
+        }
+    )
+
+
+def _snapshot_flags() -> list[str]:
+    lines = SNAPSHOT.read_text().splitlines()
+    return [line.strip() for line in lines if line.strip() and not line.startswith("#")]
+
+
+def test_serve_flag_snapshot_is_current():
+    assert SNAPSHOT.is_file(), f"missing snapshot {SNAPSHOT}"
+    assert _current_flags() == _snapshot_flags(), (
+        "`vllm bench serve` flags changed. Regenerate the snapshot with "
+        "`.venv/bin/python tests/benchmarks/test_rust_bench_cli_parity.py`, "
+        "then for any added flag either support it in rust/src/bench/src/cli.rs "
+        "or add it to PYTHON_ONLY in rust/src/bench/tests/cli_parity.rs"
+    )
+
+
+if __name__ == "__main__":
+    SNAPSHOT.write_text(HEADER + "\n".join(_current_flags()) + "\n")
+    print(f"wrote {SNAPSHOT}")