AGENTS.md

1# AGENTS.md
2 
3## Project Overview
4 
5Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime.
6 
7Member 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` (async hf-hub API bridged to sync) instead of hf-hub's ureq backend.
8 
9## Build & Test
10 
11Run from the `rust/` workspace root:
12 
13```bash
14# Build release binary (rust/target/release/vllm-bench)
15cargo build -p vllm-bench --release
16 
17# Run all tests
18cargo test -p vllm-bench
19 
20# Run ignored integration tests (requires network for tokenizer download)
21cargo 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)
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/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback
36- `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source)
37- `src/hub.rs` — `HubRepo`: sync facade over hf-hub's async (reqwest/native-tls) API — per-download thread with its own runtime; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans
38- `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up
39- `src/ready_checker.rs` — Endpoint readiness with retry
40- `src/backends/` — Backend implementations (enum dispatch, not trait objects)
41    - `mod.rs` — `Backend` enum, `RequestFuncInput`/`RequestFuncOutput` (includes `messages` field for multi-turn)
42    - `streaming.rs` — SSE parser (`StreamedResponseHandler`) with speculative JSON parse for split TCP segments
43    - `openai_completions.rs` — `/v1/completions` backend
44    - `openai_chat.rs` — `/v1/chat/completions` backend (uses `input.messages` when set; zero-copy raw JSON payload for multimodal)
45    - `pooling.rs` — Non-streaming pooling/embedding backends: `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank`
46- `src/datasets/random.rs` — Random dataset generation with rayon parallelism
47- `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`)
48- `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching
49- `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`)
50- `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache)
51- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection)
52- `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)
53- `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`)
54- `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`)
55- `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`)
56- `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs
57- `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics`
58- `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)
59- `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown
60- `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics`
61 
62## Key Design Decisions
63 
64- **Enum dispatch** for backends (avoids async trait object issues with `dyn`)
65- **reqwest http1_only()** to match Python aiohttp behavior
66- **rayon** for parallel dataset generation (key perf win over Python)
67- **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)
68- **Arc\<str\> prompts** zero-copy sharing across tokio tasks (~3GB savings at 100k prompts with 8k-token inputs)
69- **Spawn-per-request** `tokio::spawn` + `Semaphore` (matches Python asyncio pattern)
70- **Speculative JSON parse** in SSE handler — detects complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split
71- **Tokenizer fallback chain**: Local HF → Tiktoken (`.tiktoken`/`.model` + built-in encodings) → Server-side `/tokenize`+`/detokenize`. Blocking HTTP in rayon threads for server fallback.
72- **hf-hub** for downloading tokenizers and datasets from HuggingFace Hub
73- **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
74- **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`
75- **`--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)
76- **`--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.)
77- **`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
78- JSON output schema must match Python `vllm bench serve` exactly
79 
80## Common Issues
81 
82- **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.
83- **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`.
84- **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`).
85 
86## Typical Usage
87 
88```bash
89# Embedding benchmark (openai-embeddings, 8 inputs batched per request)
90./target/release/vllm-bench \
91  --backend openai-embeddings \
92  --base-url http://gb200-10:30000 \
93  --model BAAI/bge-large-en-v1.5 \
94  --dataset-name random \
95  --random-input-len 512 \
96  --random-batch-size 8 \
97  --num-prompts 1000 \
98  --save-result
99 
100# vLLM rerank benchmark (one query + 8 documents per request)
101./target/release/vllm-bench \
102  --backend vllm-rerank \
103  --base-url http://gb200-10:30000 \
104  --model BAAI/bge-reranker-v2-m3 \
105  --dataset-name random-rerank \
106  --random-input-len 512 \
107  --random-batch-size 8 \
108  --num-prompts 500 \
109  --save-result
110 
111# Prefix-cache stress (10 shared prefixes, 256+256 tokens)
112./target/release/vllm-bench \
113  --backend vllm \
114  --base-url http://gb200-10:30000 \
115  --model nvidia/Kimi-K2.5-NVFP4 \
116  --dataset-name prefix_repetition \
117  --prefix-repetition-prefix-len 256 \
118  --prefix-repetition-suffix-len 256 \
119  --prefix-repetition-num-prefixes 10 \
120  --num-prompts 1000
121 
122# Custom JSONL workload ({"prompt": ..., "output_tokens": ...} per line)
123./target/release/vllm-bench \
124  --backend openai-chat \
125  --base-url http://gb200-10:30000 \
126  --model nvidia/Kimi-K2.5-NVFP4 \
127  --dataset-name custom \
128  --dataset-path workload.jsonl \
129  --custom-output-len -1 \
130  --num-prompts 1000
131 
132# Random dataset
133./target/release/vllm-bench \
134  --backend vllm \
135  --base-url http://gb200-10:30000 \
136  --model nvidia/Kimi-K2.5-NVFP4 \
137  --dataset-name random \
138  --random-input-len 8192 \
139  --random-output-len 1024 \
140  --ignore-eos \
141  --num-prompts 4096 \
142  --percentile-metrics &quot;ttft,tpot,itl,e2el&quot; \
143  --save-result \
144  --max-concurrency 1400
145 
146# Random multimodal dataset (VLM benchmark)
147./target/release/vllm-bench \
148  --backend openai-chat \
149  --base-url http://gb200-10:30000 \
150  --model Qwen/Qwen2.5-VL-7B-Instruct \
151  --dataset-name random-mm \
152  --random-input-len 512 \
153  --random-output-len 128 \
154  --num-prompts 100 \
155  --random-mm-base-items-per-request 1 \
156  --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
157  --random-mm-bucket-config '{(1024, 800, 1): 1.0}'
158 
159# HuggingFace dataset (WildChat)
160./target/release/vllm-bench \
161  --backend openai-chat \
162  --base-url http://gb200-10:30000 \
163  --model nvidia/Kimi-K2.5-NVFP4 \
164  --dataset-name hf \
165  --dataset-path allenai/WildChat-4.8M \
166  --hf-split train \
167  --num-prompts 1000 \
168  --save-result
169 
170# HuggingFace dataset (LongBench with subset)
171./target/release/vllm-bench \
172  --backend openai-chat \
173  --base-url http://gb200-10:30000 \
174  --model nvidia/Kimi-K2.5-NVFP4 \
175  --dataset-name hf \
176  --dataset-path THUDM/LongBench \
177  --hf-subset narrativeqa \
178  --hf-split test \
179  --hf-output-len 512 \
180  --num-prompts 200
181```
182