| 1 | # AGENTS.md |
| 2 | |
| 3 | ## Project Overview |
| 4 | |
| 5 | Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime. |
| 6 | |
| 7 | Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (hf-hub's async reqwest/native-tls API) instead of hf-hub's ureq backend. |
| 8 | |
| 9 | ## Build & Test |
| 10 | |
| 11 | Run from the `rust/` workspace root: |
| 12 | |
| 13 | ```bash |
| 14 | # Build release binary (rust/target/release/vllm-bench) |
| 15 | cargo build -p vllm-bench --release |
| 16 | |
| 17 | # Run all tests |
| 18 | cargo test -p vllm-bench |
| 19 | |
| 20 | # Run ignored integration tests (requires network for tokenizer download) |
| 21 | cargo test -p vllm-bench -- --ignored |
| 22 | ``` |
| 23 | |
| 24 | ## Architecture |
| 25 | |
| 26 | - `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single, `mm-processor` subcommand) |
| 27 | - `src/cli.rs` — clap derive CLI args (~50+ flags) |
| 28 | - `src/config.rs` — Validated config from CLI; `GoodputConfig`, `RampUpConfig`, sampling param merging |
| 29 | - `src/error.rs` — `BenchError` enum (Http, Json, Tokenizer, Config, EndpointTimeout, Backend, Io) |
| 30 | - `src/benchmark.rs` — Core benchmark orchestrator (spawn-per-request with tokio + Semaphore; fetches speculative decoding metrics from `/metrics`) |
| 31 | - `src/multi_turn.rs` — Multi-turn conversation orchestrator (channel-based worker pool, sequential turns per conversation) |
| 32 | - `src/sweep.rs` — Concurrency/rate parameter sweep (`--sweep-max-concurrency`, `--sweep-request-rate`) |
| 33 | - `src/multi_run.rs` — N-run aggregation with mean/std/min/max/CV (`--num-runs`) |
| 34 | - `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`) |
| 35 | - `src/mm_processor.rs` — `vllm-bench mm-processor` subcommand: offline multimodal preprocessing latency benchmark against a managed headless engine (mirrors `vllm bench mm-processor`); per-stage timings come from `vllm_tracing::timing::RequestTimingLayer` (via `vllm_chat::mm_timing_layer`) |
| 36 | - `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback |
| 37 | - `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source) |
| 38 | - `src/hub.rs` — `HubRepo`: thin wrapper over hf-hub's async (reqwest/native-tls) API — file download and repo listing, optional revision; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans |
| 39 | - `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up |
| 40 | - `src/ready_checker.rs` — Endpoint readiness with retry |
| 41 | - `src/backends/` — Backend implementations (enum dispatch, not trait objects) |
| 42 | - `mod.rs` — `Backend` enum, `RequestFuncInput`/`RequestFuncOutput` (includes `messages` field for multi-turn) |
| 43 | - `streaming.rs` — SSE parser (`StreamedResponseHandler`) with speculative JSON parse for split TCP segments |
| 44 | - `openai_completions.rs` — `/v1/completions` backend |
| 45 | - `openai_chat.rs` — `/v1/chat/completions` backend (uses `input.messages` when set; zero-copy raw JSON payload for multimodal) |
| 46 | - `pooling.rs` — Non-streaming pooling/embedding backends: `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank` |
| 47 | - `src/datasets/random.rs` — Random dataset generation with rayon parallelism |
| 48 | - `src/datasets/random_mm.rs` — Random multimodal dataset (synthetic JPEG images, bucket config sampling, pre-serialized JSON fragments); `--enable-multimodal-chat` pre-builds the chat `messages` array at dataset time (mirrors Python's `apply_multimodal_chat_transformation`) |
| 49 | - `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching |
| 50 | - `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`) |
| 51 | - `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache) |
| 52 | - `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (parquet shards via hf-hub into the standard HF hub cache; datasets-server for config/split discovery; native main-branch parquet/JSON fallback for private datasets; column auto-detection) |
| 53 | - `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template) |
| 54 | - `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`) |
| 55 | - `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`) |
| 56 | - `src/datasets/multi_turn.rs` — Multi-turn synthetic generator + ShareGPT multi-turn loader (3-tier prefix sharing: global/conversation/unique-suffix; `per_turn_input_len`) |
| 57 | - `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs |
| 58 | - `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics` |
| 59 | - `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` (closed-loop) |
| 60 | - `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown |
| 61 | - `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics` |
| 62 | |
| 63 | ## Key Design Decisions |
| 64 | |
| 65 | - **Enum dispatch** for backends (avoids async trait object issues with `dyn`) |
| 66 | - **reqwest http1_only()** to match Python aiohttp behavior |
| 67 | - **rayon** for parallel dataset generation (key perf win over Python) |
| 68 | - **mimalloc** global allocator to reduce contention at 1400+ concurrency (page-agnostic; works on aarch64 64K-page kernels where jemalloc aborts with `LG_PAGE=12` builds) |
| 69 | - **Arc\<str\> prompts** zero-copy sharing across tokio tasks (~3GB savings at 100k prompts with 8k-token inputs) |
| 70 | - **Spawn-per-request** `tokio::spawn` + `Semaphore` (matches Python asyncio pattern) |
| 71 | - **Speculative JSON parse** in SSE handler — detects complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split |
| 72 | - **Tokenizer fallback chain**: Local HF → Tiktoken (`.tiktoken`/`.model` + built-in encodings) → Server-side `/tokenize`+`/detokenize`. Blocking HTTP in rayon threads for server fallback. |
| 73 | - **hf-hub** for downloading tokenizers and datasets from HuggingFace Hub |
| 74 | - **Pre-serialized mm fragments** (`Arc<str>`) for multimodal: image content stored as JSON strings, zero-copy concatenated into payload — avoids deep-cloning ~200KB+ base64 per request |
| 75 | - **Steady-state metrics** (default-on in closed-loop): measure throughput/TTFT/TPOT only over the saturated plateau to cut run-to-run variance at high concurrency; `steady_state` is an `Option` in JSON (`#[serde(default)]` for backward compat), null when the scope gate fails or `--no-steady-state` |
| 76 | - **`--prompt-token-ids`** (random dataset only): send token-ID arrays instead of text to skip server-side tokenization; also skips the token-length verification pass (counts exact by construction) |
| 77 | - **`--random-range-ratio`** follows Python semantics: lengths sampled uniformly from `[len*(1-r), len*(1+r)]`, default `0.0` = fixed; accepts a float in `[0,1)` or `'{"input": r1, "output": r2}'`. (The pre-2026-07 Rust-only form `[len*r, len]` with default 1.0 is rejected with a migration hint.) |
| 78 | - **`prompt_list`** (`Arc<[Arc<str>]>` on `SampleRequest`/`RequestFuncInput`): multiple inputs per request for pooling backends — embeddings batches (`--random-batch-size`) send `"input": [...]`, rerank sends `[0]` as query + `[1..]` as documents |
| 79 | - JSON output schema must match Python `vllm bench serve` exactly |
| 80 | |
| 81 | ## Common Issues |
| 82 | |
| 83 | - **localhost vs 127.0.0.1**: Some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only. Use `127.0.0.1` or the actual hostname. |
| 84 | - **Models without tokenizer.json** (e.g., `nvidia/Kimi-K2.5-NVFP4`): Automatically falls back to server-side tokenization. Can also use `--tokenizer` to point to a model with `tokenizer.json`. |
| 85 | - **usage.completion_tokens parsing**: vLLM sends final usage chunk with `"choices":[]` (empty array). The usage `if` must be separate from the choices `if` (not `else if`). |
| 86 | |
| 87 | ## Typical Usage |
| 88 | |
| 89 | ```bash |
| 90 | # Embedding benchmark (openai-embeddings, 8 inputs batched per request) |
| 91 | ./target/release/vllm-bench \ |
| 92 | --backend openai-embeddings \ |
| 93 | --base-url http://gb200-10:30000 \ |
| 94 | --model BAAI/bge-large-en-v1.5 \ |
| 95 | --dataset-name random \ |
| 96 | --random-input-len 512 \ |
| 97 | --random-batch-size 8 \ |
| 98 | --num-prompts 1000 \ |
| 99 | --save-result |
| 100 | |
| 101 | # vLLM rerank benchmark (one query + 8 documents per request) |
| 102 | ./target/release/vllm-bench \ |
| 103 | --backend vllm-rerank \ |
| 104 | --base-url http://gb200-10:30000 \ |
| 105 | --model BAAI/bge-reranker-v2-m3 \ |
| 106 | --dataset-name random-rerank \ |
| 107 | --random-input-len 512 \ |
| 108 | --random-batch-size 8 \ |
| 109 | --num-prompts 500 \ |
| 110 | --save-result |
| 111 | |
| 112 | # Prefix-cache stress (10 shared prefixes, 256+256 tokens) |
| 113 | ./target/release/vllm-bench \ |
| 114 | --backend vllm \ |
| 115 | --base-url http://gb200-10:30000 \ |
| 116 | --model nvidia/Kimi-K2.5-NVFP4 \ |
| 117 | --dataset-name prefix_repetition \ |
| 118 | --prefix-repetition-prefix-len 256 \ |
| 119 | --prefix-repetition-suffix-len 256 \ |
| 120 | --prefix-repetition-num-prefixes 10 \ |
| 121 | --num-prompts 1000 |
| 122 | |
| 123 | # Custom JSONL workload ({"prompt": ..., "output_tokens": ...} per line) |
| 124 | ./target/release/vllm-bench \ |
| 125 | --backend openai-chat \ |
| 126 | --base-url http://gb200-10:30000 \ |
| 127 | --model nvidia/Kimi-K2.5-NVFP4 \ |
| 128 | --dataset-name custom \ |
| 129 | --dataset-path workload.jsonl \ |
| 130 | --custom-output-len -1 \ |
| 131 | --num-prompts 1000 |
| 132 | |
| 133 | # Random dataset |
| 134 | ./target/release/vllm-bench \ |
| 135 | --backend vllm \ |
| 136 | --base-url http://gb200-10:30000 \ |
| 137 | --model nvidia/Kimi-K2.5-NVFP4 \ |
| 138 | --dataset-name random \ |
| 139 | --random-input-len 8192 \ |
| 140 | --random-output-len 1024 \ |
| 141 | --ignore-eos \ |
| 142 | --num-prompts 4096 \ |
| 143 | --percentile-metrics "ttft,tpot,itl,e2el" \ |
| 144 | --save-result \ |
| 145 | --max-concurrency 1400 |
| 146 | |
| 147 | # Random multimodal dataset (VLM benchmark) |
| 148 | ./target/release/vllm-bench \ |
| 149 | --backend openai-chat \ |
| 150 | --base-url http://gb200-10:30000 \ |
| 151 | --model Qwen/Qwen2.5-VL-7B-Instruct \ |
| 152 | --dataset-name random-mm \ |
| 153 | --random-input-len 512 \ |
| 154 | --random-output-len 128 \ |
| 155 | --num-prompts 100 \ |
| 156 | --random-mm-base-items-per-request 1 \ |
| 157 | --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \ |
| 158 | --random-mm-bucket-config '{(1024, 800, 1): 1.0}' |
| 159 | |
| 160 | # HuggingFace dataset (WildChat) |
| 161 | ./target/release/vllm-bench \ |
| 162 | --backend openai-chat \ |
| 163 | --base-url http://gb200-10:30000 \ |
| 164 | --model nvidia/Kimi-K2.5-NVFP4 \ |
| 165 | --dataset-name hf \ |
| 166 | --dataset-path allenai/WildChat-4.8M \ |
| 167 | --hf-split train \ |
| 168 | --num-prompts 1000 \ |
| 169 | --save-result |
| 170 | |
| 171 | # HuggingFace dataset (LongBench with subset) |
| 172 | ./target/release/vllm-bench \ |
| 173 | --backend openai-chat \ |
| 174 | --base-url http://gb200-10:30000 \ |
| 175 | --model nvidia/Kimi-K2.5-NVFP4 \ |
| 176 | --dataset-name hf \ |
| 177 | --dataset-path THUDM/LongBench \ |
| 178 | --hf-subset narrativeqa \ |
| 179 | --hf-split test \ |
| 180 | --hf-output-len 512 \ |
| 181 | --num-prompts 200 |
| 182 | ``` |