@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f84325c48c

[Frontend][Rust] Add mm-processor benchmark for Rust frontend (#51922)

Signed-off-by: Karthik Gangula <gangula-karthik@users.noreply.github.com> Signed-off-by: gangula-karthik <gkarthik923@gmail.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Karthik Gangula <gangula-karthik@users.noreply.github.com> Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>

karthik committed Sep 22, 2026 at 13:08 UTC f84325c48c0acc1e3703103788c5f2976e719762

14 files changed +980 -33

rust/Cargo.lock

@@ -5735,6 +5735,11 @@ dependencies = [

5735 "url",

5736 "uuid",

5737 "vllm-build-info",

5738 + "vllm-chat",

5739 + "vllm-engine-core-client",

5740 + "vllm-llm",

5741 + "vllm-managed-engine",

5742 + "vllm-text",

5743 "vllm-tracing",

5744 ]

5745

@@ -5784,6 +5789,7 @@ dependencies = [

5789 "vllm-parser",

5790 "vllm-text",

5791 "vllm-tokenizer",

5792 + "vllm-tracing",

5793 "xgrammar-structural-tag",

5794 "zeromq",

5795 ]

rust/src/bench/AGENTS.md

+2 -1

@@ -23,7 +23,7 @@ cargo test -p vllm-bench -- --ignored

24 ## Architecture

26 -- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single)

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)

@@ -32,6 +32,7 @@ cargo test -p vllm-bench -- --ignored

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

rust/src/bench/Cargo.toml

@@ -36,6 +36,11 @@ tracing.workspace = true

36 url.workspace = true

37 uuid.workspace = true

38 vllm-build-info.workspace = true

39 +vllm-chat.workspace = true

40 +vllm-engine-core-client.workspace = true

41 +vllm-llm.workspace = true

42 +vllm-managed-engine.workspace = true

43 +vllm-text.workspace = true

44 vllm-tracing.workspace = true

46 [lints]

rust/src/bench/README.md

+59 -1

@@ -18,6 +18,7 @@ vllm-bench --backend vllm --base-url http://127.0.0.1:8000 \

18 - **Beyond a single run** — concurrency/rate **sweeps**, **multi-run** stats, **multi-turn** conversations, **LoRA** multi-adapter, and result **comparison**.

19 - **Steady-state metrics** — throughput/latency measured over the saturated plateau, excluding ramp-up and drain.

20 - **Parity** — JSON output schema and timing semantics match Python `vllm bench serve` exactly.

21 +- **Multimodal preprocessing** — `mm-processor` subcommand measures per-stage VLM preprocessing latency against a managed headless engine (mirrors `vllm bench mm-processor`).

23 ### Performance vs. Python

@@ -404,6 +405,32 @@ vllm-bench \

405

406 </details>

407

408 +<details>

409 +<summary><b>Multimodal preprocessing (offline)</b></summary>

410 +

411 +The `mm-processor` subcommand benchmarks the multimodal preprocessing pipeline

412 +itself (media fetch → per-modality preprocess → prompt expansion), mirroring

413 +`vllm bench mm-processor`. It spawns a managed headless Python engine (real

414 +weights, real encoder) — no serving endpoint needed — and reports per-stage

415 +timing (mean/median/std/percentiles) plus end-to-end latency.

416 +

417 +```bash

418 +vllm-bench mm-processor \

419 + --model Qwen/Qwen2.5-VL-7B-Instruct \

420 + --num-prompts 100 --num-warmups 5 \

421 + --random-input-len 512 --random-output-len 128 \

422 + --random-mm-base-items-per-request 1 \

423 + --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \

424 + --random-mm-bucket-config '{(1024, 800, 1): 1.0}' \

425 + --metric-percentiles "50,90,99" \

426 + --output-json mm_processor_stats.json

427 +```

428 +

429 +`--max-concurrency` defaults to `1`, matching Python's serial driver; raise it

430 +to measure a concurrent serving-style workload.

431 +

432 +</details>

433 +

434 ## Supported Backends

435

436 ### Generation

@@ -705,6 +732,36 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n

732

733 </details>

734

735 +<details>

736 +<summary><b><code>mm-processor</code> subcommand</b></summary>

737 +

738 +Offline multimodal preprocessing latency benchmark (mirrors `vllm bench mm-processor`). Spawns a managed headless Python engine and drives the full chat preprocessing pipeline.

739 +

740 +| Flag | Default | Description |

741 +| ------ | --------- | ------------- |

742 +| `--model` | (required) | Model to serve and benchmark (HF id or local path) |

743 +| `--num-prompts` | `10` | Prompts to process, excluding warmups |

744 +| `--num-warmups` | `1` | Warmup prompts processed and discarded before timing |

745 +| `--metric-percentiles` | `99` | Comma-separated percentiles to report |

746 +| `--output-json` | — | Path to write aggregate stats as JSON (same schema as Python) |

747 +| `--random-input-len` | `1024` | Text input length per prompt |

748 +| `--random-output-len` | `128` | Expected output length per prompt |

749 +| `--random-prefix-len` | `0` | Prefix token length per prompt |

750 +| `--random-range-ratio` | `0.0` | Input/output length range ratio (float or `{"input": i, "output": o}`) |

751 +| `--random-mm-base-items-per-request` | `1` | Base number of multimodal items per request |

752 +| `--random-mm-num-mm-items-range-ratio` | `0.0` | Range ratio (in [0, 1]) for the number of mm items per request |

753 +| `--random-mm-limit-mm-per-prompt` | `{"image": 255, "video": 1}` | Per-modality item caps |

754 +| `--random-mm-bucket-config` | `{(256,256,1): 0.5, (720,1280,1): 0.5}` | Bucket config |

755 +| `--seed` | `0` | Seed for dataset generation |

756 +| `--request-id-prefix` | `mm-proc-` | Request-id prefix for per-request timing keys |

757 +| `--max-concurrency` | `1` | Maximum concurrent requests (1 matches Python's serial driver) |

758 +| `--trust-remote-code` | `false` | Trust remote code for the tokenizer and the managed engine |

759 +| `--chat-template` | — | Chat-template override (inline text or path) |

760 +

761 +Managed headless-engine options (e.g. `--python`, `--data-parallel-address`, `--data-parallel-rpc-port`, `--data-parallel-size`, `--max-model-len`, `--python-args`) are shared with other managed-engine commands; see `vllm-bench mm-processor --help`.

762 +

763 +</details>

764 +

765 ## Tokenizer Support

766

767 Tokenizers are loaded with a three-tier fallback chain:

@@ -734,10 +791,11 @@ Use `--append-result` to append multiple runs to the same file in JSONL format.

791

792 ```text

793 src/

737 -├── main.rs # Entry point, mimalloc, tokio runtime, mode dispatch

794 +├── main.rs # Entry point, mimalloc, tokio runtime, mode/subcommand dispatch

795 ├── cli.rs # clap CLI argument definitions

796 ├── config.rs # Validated config, goodput/ramp-up parsing

797 ├── benchmark.rs # Core orchestrator (schedule, spawn, collect, verify, profile)

798 +├── mm_processor.rs # Offline mm preprocessing benchmark (mm-processor subcommand)

799 ├── multi_turn.rs # Multi-turn conversation orchestrator (channel workers)

800 ├── compare.rs # Result diff (--compare file_a.json file_b.json)

801 ├── sweep.rs # Parameter sweep (--sweep-max-concurrency, --sweep-request-rate)

rust/src/bench/src/lib.rs

@@ -10,6 +10,7 @@ mod datasets;

10 mod error;

11 mod hub;

12 mod metrics;

13 +mod mm_processor;

14 mod multi_run;

15 mod multi_turn;

16 mod output;

@@ -25,6 +26,7 @@ pub use cli::{

26 BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,

27 };

28 use config::BenchConfig;

29 +pub use mm_processor::{MmProcessorArgs, run_mm_processor};

31 /// Prepare process-wide resources for a benchmark run.

32 pub fn prepare_process() {

rust/src/bench/src/main.rs

+23 -5

@@ -6,22 +6,30 @@

6 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

8 use anyhow::Context;

9 -use clap::Parser;

9 +use clap::{Parser, Subcommand};

11 #[derive(Parser)]

12 #[command(

13 name = "vllm-bench",

14 - about = "Benchmark online serving throughput",

14 + about = "Benchmark online serving throughput and offline multimodal preprocessing",

15 version = vllm_build_info::VERSION

16 )]

17 struct Cli {

18 #[command(flatten)]

19 args: vllm_bench::BenchServeArgs,

20 +

21 + /// Optional subcommand; absent by default for the online serving benchmark.

22 + #[command(subcommand)]

23 + command: Option<Command>,

24 }

22 -fn main() -> anyhow::Result<()> {

23 - vllm_tracing::init_tracing("Bench");

26 +#[derive(Subcommand)]

27 +enum Command {

28 + /// Offline multimodal preprocessing latency benchmark.

29 + MmProcessor(vllm_bench::MmProcessorArgs),

30 +}

32 +fn main() -> anyhow::Result<()> {

33 let cli = Cli::parse();

34 vllm_bench::prepare_process();

@@ -30,5 +38,15 @@ fn main() -> anyhow::Result<()> {

38 .build()

39 .context("Failed to build tokio runtime")?;

33 - runtime.block_on(vllm_bench::run(cli.args))

41 + match cli.command {

42 + Some(Command::MmProcessor(args)) => {

43 + let (timing_layer, timing_stats) = vllm_chat::mm_timing_layer();

44 + vllm_tracing::init_tracing_with("Bench", timing_layer);

45 + runtime.block_on(vllm_bench::run_mm_processor(args, timing_stats))

46 + }

47 + None => {

48 + vllm_tracing::init_tracing("Bench");

49 + runtime.block_on(vllm_bench::run(cli.args))

50 + }

51 + }

52 }

rust/src/bench/src/metrics/calculator.rs

+5 -5

@@ -216,20 +216,20 @@ pub fn calculate_metrics(

216 (metrics, actual_output_lens)

217 }

218

219 -fn mean(data: &[f64]) -> f64 {

219 +pub(crate) fn mean(data: &[f64]) -> f64 {

220 if data.is_empty() {

221 return 0.0;

222 }

223 data.iter().sum::<f64>() / data.len() as f64

224 }

225

226 -fn sort_clone(data: &[f64]) -> Vec<f64> {

226 +pub(crate) fn sort_clone(data: &[f64]) -> Vec<f64> {

227 let mut sorted = data.to_vec();

228 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

229 sorted

230 }

231

232 -fn median_sorted(sorted: &[f64]) -> f64 {

232 +pub(crate) fn median_sorted(sorted: &[f64]) -> f64 {

233 if sorted.is_empty() {

234 return 0.0;

235 }

@@ -241,7 +241,7 @@ fn median_sorted(sorted: &[f64]) -> f64 {

241 }

242 }

243

244 -fn std_dev(data: &[f64]) -> f64 {

244 +pub(crate) fn std_dev(data: &[f64]) -> f64 {

245 if data.is_empty() {

246 return 0.0;

247 }

@@ -250,7 +250,7 @@ fn std_dev(data: &[f64]) -> f64 {

250 variance.sqrt()

251 }

252

253 -fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {

253 +pub(crate) fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {

254 if sorted_data.is_empty() {

255 return 0.0;

256 }

rust/src/bench/src/mm_processor.rs new

+572

@@ -0,0 +1,572 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +//! Multimodal preprocessing benchmark (`vllm-bench mm-processor`), mirroring

5 +//! `vllm bench mm-processor`.

6 +

7 +use std::collections::{BTreeMap, HashMap};

8 +use std::sync::Arc;

9 +use std::time::{Duration, Instant};

10 +

11 +use anyhow::{Context, Result};

12 +use clap::Args;

13 +use futures::future::join_all;

14 +use serde::Serialize;

15 +use tokio::sync::Semaphore;

16 +

17 +use vllm_chat::{

18 + ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRequest, LoadModelBackendsOptions,

19 + MultiModalTimingStats, load_model_backends,

20 +};

21 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};

22 +use vllm_llm::Llm;

23 +use vllm_managed_engine::cli::ManagedEngineArgs;

24 +use vllm_managed_engine::{ManagedEngineConfig, ManagedEngineHandle};

25 +use vllm_text::TextLlm;

26 +

27 +use crate::config::RangeRatio;

28 +use crate::datasets::SampleRequest;

29 +use crate::datasets::random_mm::{

30 + MmBucketKey, MmLimitPerPrompt, generate_random_mm_dataset, parse_bucket_config,

31 + parse_limit_mm_per_prompt,

32 +};

33 +use crate::error::BenchError;

34 +use crate::metrics::calculator::{mean, median_sorted, percentile_sorted, sort_clone, std_dev};

35 +use crate::tokenizer::{TokenizerKind, load_tokenizer};

36 +

37 +/// Seconds -> milliseconds.

38 +const SEC_TO_MS: f64 = 1000.0;

39 +

40 +/// Preferred column order; other stages appended alphabetically.

41 +const STAGE_ORDER: &[&str] = &[

42 + "preprocessor_total_ms",

43 + "media_fetch_ms",

44 + "preprocess_image_ms",

45 + "preprocess_video_ms",

46 + "preprocess_audio_ms",

47 + "prompt_expansion_ms",

48 +];

49 +

50 +/// Time to wait for managed engines to register.

51 +const ENGINE_READY_TIMEOUT: Duration = Duration::from_secs(600);

52 +

53 +/// Maximum time to wait for the managed engine to drain on shutdown.

54 +const ENGINE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(120);

55 +

56 +#[derive(Args, Debug, Clone)]

57 +pub struct MmProcessorArgs {

58 + /// Model to serve and benchmark (Hugging Face id or local path).

59 + #[arg(long)]

60 + pub model: String,

61 +

62 + /// Number of prompts to process, excluding warmups.

63 + #[arg(long, default_value_t = 10)]

64 + pub num_prompts: usize,

65 +

66 + /// Number of warmup prompts processed and discarded before timing.

67 + #[arg(long, default_value_t = 1)]

68 + pub num_warmups: usize,

69 +

70 + /// Comma-separated percentiles to report, e.g. "50,90,99".

71 + #[arg(long, default_value = "99")]

72 + pub metric_percentiles: String,

73 +

74 + /// Path to write the aggregate stats as JSON.

75 + #[arg(long)]

76 + pub output_json: Option<String>,

77 +

78 + /// Text input length per prompt (before range scaling).

79 + #[arg(long, default_value_t = 1024)]

80 + pub random_input_len: usize,

81 +

82 + /// Expected output length per prompt.

83 + #[arg(long, default_value_t = 128)]

84 + pub random_output_len: usize,

85 +

86 + /// Prefix token length per prompt.

87 + #[arg(long, default_value_t = 0)]

88 + pub random_prefix_len: usize,

89 +

90 + /// Input/output length range ratio: a float or `{"input": i, "output": o}`.

91 + #[arg(long, default_value = "0.0")]

92 + pub random_range_ratio: String,

93 +

94 + /// Base number of multimodal items per request.

95 + #[arg(long, default_value_t = 1)]

96 + pub random_mm_base_items_per_request: usize,

97 +

98 + /// Range ratio (in [0, 1]) for the number of mm items per request.

99 + #[arg(long, default_value_t = 0.0)]

100 + pub random_mm_num_mm_items_range_ratio: f64,

101 +

102 + /// Per-modality item caps, e.g. `{"image": 255, "video": 1}`.

103 + #[arg(long, default_value = "{\"image\": 255, \"video\": 1}")]

104 + pub random_mm_limit_mm_per_prompt: String,

105 +

106 + /// Bucket config, e.g. `{(256,256,1): 0.5, (720,1280,1): 0.5}`.

107 + #[arg(long, default_value = "{(256,256,1): 0.5, (720,1280,1): 0.5}")]

108 + pub random_mm_bucket_config: String,

109 +

110 + /// Seed for dataset generation.

111 + #[arg(long)]

112 + pub seed: Option<u64>,

113 +

114 + /// Request-id prefix for per-request timing keys.

115 + #[arg(long, default_value = "mm-proc-")]

116 + pub request_id_prefix: String,

117 +

118 + /// Maximum concurrent requests (default 1, matching Python's serial driver).

119 + #[arg(long, default_value_t = 1)]

120 + pub max_concurrency: usize,

121 +

122 + /// Trust remote code for the tokenizer and the managed engine.

123 + #[arg(long, default_value_t = false)]

124 + pub trust_remote_code: bool,

125 +

126 + /// Optional chat-template override (inline text or path).

127 + #[arg(long)]

128 + pub chat_template: Option<String>,

129 +

130 + /// Managed headless-engine options.

131 + #[command(flatten)]

132 + pub engine: ManagedEngineArgs,

133 +}

134 +

135 +/// Entry point for `vllm-bench mm-processor`.

136 +pub async fn run_mm_processor(args: MmProcessorArgs, timing: MultiModalTimingStats) -> Result<()> {

137 + if args.engine.data_parallel_size_local == Some(0) {

138 + anyhow::bail!(

139 + "`--data-parallel-size-local 0` is not supported; the benchmark requires \

140 + a local managed engine"

141 + );

142 + }

143 +

144 + // Spawn the managed headless Python engine (real weights, real encoder).

145 + let handshake_port = args

146 + .engine

147 + .resolve_handshake_port()

148 + .context("failed to allocate handshake port")?;

149 + let mut python_args = args.engine.python_args.clone();

150 + if args.trust_remote_code {

151 + python_args.push("--trust-remote-code".to_string());

152 + }

153 + let engine_config = ManagedEngineConfig {

154 + python: args.engine.python.clone(),

155 + model: args.model.clone(),

156 + handshake_host: args.engine.handshake_host.clone(),

157 + handshake_port,

158 + data_parallel_size: args.engine.data_parallel_size,

159 + python_args,

160 + };

161 + eprintln!(

162 + "Spawning managed headless engine: python={} model={} data_parallel_size={}",

163 + engine_config.python, engine_config.model, engine_config.data_parallel_size

164 + );

165 + let engine = ManagedEngineHandle::spawn(engine_config)

166 + .await

167 + .context("failed to start managed Python headless engine")?;

168 +

169 + let result = run_mm_processor_with_engine(args, handshake_port, &timing).await;

170 + let shutdown_result = engine

171 + .shutdown(ENGINE_SHUTDOWN_TIMEOUT)

172 + .await

173 + .context("failed to shut down managed engine");

174 + result?;

175 + shutdown_result?;

176 + Ok(())

177 +}

178 +

179 +async fn run_mm_processor_with_engine(

180 + args: MmProcessorArgs,

181 + handshake_port: u16,

182 + timing: &MultiModalTimingStats,

183 +) -> Result<()> {

184 + let range_ratio =

185 + RangeRatio::parse(&args.random_range_ratio).context("invalid --random-range-ratio")?;

186 + let limit: MmLimitPerPrompt = parse_limit_mm_per_prompt(&args.random_mm_limit_mm_per_prompt)

187 + .context("invalid --random-mm-limit-mm-per-prompt")?;

188 + let buckets: Vec<(MmBucketKey, f64)> = parse_bucket_config(&args.random_mm_bucket_config)

189 + .context("invalid --random-mm-bucket-config")?;

190 + let percentiles = parse_percentiles(&args.metric_percentiles)?;

191 + let seed = args.seed.unwrap_or(0);

192 +

193 + // Offline tokenizer: local Hugging Face tokenizer or built-in tiktoken.

194 + let tokenizer = load_tokenizer(&args.model, args.trust_remote_code, None)

195 + .await

196 + .context("failed to load tokenizer")?;

197 +

198 + // Frontend backends: model configs + multimodal processor metadata.

199 + let backends = load_model_backends(

200 + &args.model,

201 + LoadModelBackendsOptions {

202 + chat_template: args.chat_template.clone(),

203 + ..LoadModelBackendsOptions::default()

204 + },

205 + )

206 + .await

207 + .context("failed to load model backends")?;

208 +

209 + if backends.chat_backend.multimodal_model_info().is_none() {

210 + anyhow::bail!(

211 + "model `{}` has no multimodal processor registered",

212 + args.model

213 + );

214 + }

215 +

216 + // Connect to the engine over the same transport the serving frontend uses.

217 + let client = EngineCoreClient::connect(EngineCoreClientConfig {

218 + transport_mode: TransportMode::HandshakeOwner {

219 + handshake_address: format!("tcp://{}:{}", args.engine.handshake_host, handshake_port),

220 + advertised_host: args.engine.handshake_host.clone(),

221 + engine_count: args

222 + .engine

223 + .data_parallel_size_local

224 + .unwrap_or(args.engine.data_parallel_size),

225 + ready_timeout: ENGINE_READY_TIMEOUT,

226 + local_input_address: None,

227 + local_output_address: None,

228 + },

229 + coordinator_mode: None,

230 + model_name: args.model.clone(),

231 + client_index: 0,

232 + })

233 + .await

234 + .context("failed to connect to engine core")?;

235 +

236 + let chat = ChatLlm::new(

237 + TextLlm::new(Llm::new(client), backends.text_backend),

238 + backends.chat_backend,

239 + );

240 +

241 + let semaphore = Arc::new(Semaphore::new(args.max_concurrency));

242 +

243 + // Warmup: generate a separate dataset, process it, then drain and discard

244 + // the collected timings so cold caches do not skew the results.

245 + if args.num_warmups > 0 {

246 + let warmup = generate_dataset(

247 + &tokenizer,

248 + args.num_warmups,

249 + seed.wrapping_add(1),

250 + &args,

251 + range_ratio,

252 + &limit,

253 + &buckets,

254 + )?;

255 + eprintln!("Processing {} warmup requests...", warmup.len());

256 + run_batch(&chat, &warmup, &semaphore).await;

257 + let _ = timing.stat();

258 + }

259 +

260 + let samples = generate_dataset(

261 + &tokenizer,

262 + args.num_prompts,

263 + seed,

264 + &args,

265 + range_ratio,

266 + &limit,

267 + &buckets,

268 + )?;

269 + eprintln!("Processing {} requests...", samples.len());

270 + let start_time = Instant::now();

271 + let outcomes = run_batch(&chat, &samples, &semaphore).await;

272 + let total_time = start_time.elapsed();

273 +

274 + let completed = outcomes.iter().filter(|outcome| outcome.is_ok()).count();

275 + let failed = outcomes.len() - completed;

276 + let e2el_times: Vec<f64> = outcomes

277 + .iter()

278 + .filter_map(|outcome| outcome.as_ref().ok().map(|elapsed| elapsed.as_secs_f64()))

279 + .collect();

280 +

281 + let per_request = timing.stat();

282 + let stats = aggregate_stats(per_request, &percentiles);

283 +

284 + eprintln!(

285 + "Processed {} requests (+{} warmup) for model `{}`",

286 + samples.len(),

287 + args.num_warmups,

288 + args.model

289 + );

290 + print_report(&stats, &percentiles);

291 + print_e2el_summary(completed, failed, total_time, &e2el_times, &percentiles);

292 +

293 + if let Some(path) = &args.output_json {

294 + write_result_json(path, completed, failed, &e2el_times, &percentiles, stats)?;

295 + eprintln!("Wrote stats JSON to {path}");

296 + }

297 +

298 + chat.shutdown().await.context("failed to shut down chat facade")?;

299 + Ok(())

300 +}

301 +

302 +/// Submit every sample through the chat pipeline and drive each stream to

303 +/// completion, returning per-request wall-clock latencies.

304 +async fn run_batch(

305 + chat: &ChatLlm,

306 + samples: &[SampleRequest],

307 + semaphore: &Arc<Semaphore>,

308 +) -> Vec<Result<Duration>> {

309 + let submissions = samples.iter().map(|sample| {

310 + let request = build_chat_request(sample);

311 + let semaphore = semaphore.clone();

312 + async move {

313 + let _permit = semaphore.acquire().await.context("semaphore closed")?;

314 + let request = request?;

315 + let start = Instant::now();

316 + chat.chat(request)

317 + .await

318 + .context("failed to submit request")?

319 + .collect_message()

320 + .await

321 + .context("request failed")?;

322 + Ok(start.elapsed())

323 + }

324 + });

325 + join_all(submissions).await

326 +}

327 +

328 +fn generate_dataset(

329 + tokenizer: &TokenizerKind,

330 + num_requests: usize,

331 + seed: u64,

332 + args: &MmProcessorArgs,

333 + range_ratio: RangeRatio,

334 + limit: &MmLimitPerPrompt,

335 + buckets: &[(MmBucketKey, f64)],

336 +) -> Result<Vec<SampleRequest>> {

337 + generate_random_mm_dataset(

338 + tokenizer,

339 + num_requests,

340 + args.random_input_len,

341 + args.random_output_len,

342 + args.random_prefix_len,

343 + range_ratio,

344 + seed,

345 + &args.request_id_prefix,

346 + args.random_mm_base_items_per_request,

347 + args.random_mm_num_mm_items_range_ratio,

348 + limit,

349 + buckets,

350 + false,

351 + )

352 + .map_err(|e| anyhow::anyhow!("{e}"))

353 +}

354 +

355 +/// Build a chat request from a sample (text prompt + `image_url` parts).

356 +fn build_chat_request(sample: &SampleRequest) -> Result<ChatRequest> {

357 + let mut request = ChatRequest::for_test();

358 + request.request_id = sample.request_id.clone().unwrap_or_else(|| "mm-proc".to_string());

359 + request.sampling_params.temperature = Some(0.0);

360 + request.sampling_params.max_tokens = Some(u32::try_from(sample.expected_output_len)?);

361 +

362 + let mut parts = vec![ChatContentPart::text(sample.prompt.as_ref())];

363 + if let Some(mm_items) = &sample.multi_modal_content {

364 + for fragment in mm_items.iter() {

365 + let url = extract_image_url(fragment)?;

366 + parts.push(ChatContentPart::image_url(url));

367 + }

368 + }

369 + request.messages = vec![ChatMessage::user(ChatContent::from(parts))];

370 + Ok(request)

371 +}

372 +

373 +/// Extract the `image_url.url` string from a pre-serialized content fragment

374 +/// (`{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,..."}}`).

375 +fn extract_image_url(fragment: &str) -> Result<String> {

376 + let value: serde_json::Value = serde_json::from_str(fragment)

377 + .with_context(|| format!("invalid multimodal fragment: {fragment}"))?;

378 + let url = value

379 + .get("image_url")

380 + .and_then(|v| v.get("url"))

381 + .and_then(|v| v.as_str())

382 + .ok_or_else(|| BenchError::Config("fragment missing image_url.url".to_string()))?;

383 + Ok(url.to_string())

384 +}

385 +

386 +fn parse_percentiles(raw: &str) -> Result<Vec<f64>> {

387 + let mut values = Vec::new();

388 + for part in raw.split(',') {

389 + let part = part.trim();

390 + if part.is_empty() {

391 + continue;

392 + }

393 + let value: f64 = part.parse().with_context(|| format!("invalid percentile `{part}`"))?;

394 + if !(0.0..=100.0).contains(&value) {

395 + anyhow::bail!("percentile `{part}` out of range [0, 100]");

396 + }

397 + values.push(value);

398 + }

399 + Ok(values)

400 +}

401 +

402 +/// Aggregate per-request stage timings into `{stage_ms: {statistic: ms}}`,

403 +/// matching the Python `mm_processor_stats` schema orientation.

404 +fn aggregate_stats(

405 + per_request: HashMap<String, HashMap<String, f64>>,

406 + percentiles: &[f64],

407 +) -> BTreeMap<String, BTreeMap<String, f64>> {

408 + let mut stage_values: BTreeMap<String, Vec<f64>> = BTreeMap::new();

409 + for stage_map in per_request.into_values() {

410 + for (stage, seconds) in stage_map {

411 + stage_values.entry(stage).or_default().push(seconds);

412 + }

413 + }

414 +

415 + stage_values

416 + .iter()

417 + .map(|(stage, values)| {

418 + let sorted = sort_clone(values);

419 + let mut stat_map = BTreeMap::new();

420 + stat_map.insert("mean".to_string(), mean(values) * SEC_TO_MS);

421 + stat_map.insert("median".to_string(), median_sorted(&sorted) * SEC_TO_MS);

422 + stat_map.insert("std".to_string(), std_dev(values) * SEC_TO_MS);

423 + for p in percentiles {

424 + stat_map.insert(format!("p{p}"), percentile_sorted(&sorted, *p) * SEC_TO_MS);

425 + }

426 + (stage.replace("_secs", "_ms"), stat_map)

427 + })

428 + .collect()

429 +}

430 +

431 +/// Print a right-aligned plain-text table (pandas `to_string(index=False)` style).

432 +fn print_table(rows: &[Vec<String>]) {

433 + let ncols = rows.first().map_or(0, Vec::len);

434 + let mut widths = vec![0usize; ncols];

435 + for row in rows {

436 + for (i, cell) in row.iter().enumerate() {

437 + widths[i] = widths[i].max(cell.len());

438 + }

439 + }

440 + for row in rows {

441 + let line = row

442 + .iter()

443 + .enumerate()

444 + .map(|(i, cell)| format!("{cell:>width$}", width = widths[i]))

445 + .collect::<Vec<_>>()

446 + .join(" ");

447 + println!("{line}");

448 + }

449 +}

450 +

451 +/// Print a plain-text table with one row per stage and one column per

452 +/// statistic (mirrors the Python pandas report layout).

453 +fn print_report(stats: &BTreeMap<String, BTreeMap<String, f64>>, percentiles: &[f64]) {

454 + // Column order: mean/median/std then the requested percentiles, like Python.

455 + let mut columns: Vec<(String, String)> =

456 + [("mean", "Mean"), ("median", "Median"), ("std", "Std")]

457 + .into_iter()

458 + .map(|(key, label)| (key.to_string(), label.to_string()))

459 + .collect();

460 + columns.extend(percentiles.iter().map(|p| (format!("p{p}"), format!("P{p:?}"))));

461 + for key in stats.values().flat_map(|map| map.keys()) {

462 + if !columns.iter().any(|(k, _)| k == key) {

463 + columns.push((key.clone(), key.clone()));

464 + }

465 + }

466 +

467 + let mut stages: Vec<&String> = stats.keys().collect();

468 + stages.sort_by_key(|stage| {

469 + STAGE_ORDER

470 + .iter()

471 + .position(|s| s == stage)

472 + .map(|i| (i as isize, (*stage).clone()))

473 + .unwrap_or((isize::MAX, (*stage).clone()))

474 + });

475 +

476 + let mut rows: Vec<Vec<String>> = vec![

477 + std::iter::once("Stage".to_string())

478 + .chain(columns.iter().map(|(_, label)| label.clone()))

479 + .collect(),

480 + ];

481 + for stage in &stages {

482 + let mut row = vec![(*stage).clone()];

483 + for (key, _) in &columns {

484 + row.push(

485 + stats[*stage]

486 + .get(key)

487 + .map(|v| format!("{v:.2}"))

488 + .unwrap_or_else(|| "-".to_string()),

489 + );

490 + }

491 + rows.push(row);

492 + }

493 + println!("MM Processor Metrics:");

494 + print_table(&rows);

495 +}

496 +

497 +/// Print the end-to-end latency summary (mirrors the Python report layout).

498 +fn print_e2el_summary(

499 + completed: usize,

500 + failed: usize,

501 + total_time: Duration,

502 + e2el_times_secs: &[f64],

503 + percentiles: &[f64],

504 +) {

505 + println!("Completed: {completed}, Failed: {failed}");

506 + if e2el_times_secs.is_empty() {

507 + println!("End-to-End Latency (ms): no completed requests");

508 + return;

509 + }

510 +

511 + let e2el_ms: Vec<f64> = e2el_times_secs.iter().map(|secs| secs * SEC_TO_MS).collect();

512 + let sorted = sort_clone(&e2el_ms);

513 + let mut rows = vec![

514 + vec!["Metric".to_string(), "Value (ms)".to_string()],

515 + vec!["Mean".to_string(), format!("{:.2}", mean(&e2el_ms))],

516 + vec![

517 + "Median".to_string(),

518 + format!("{:.2}", median_sorted(&sorted)),

519 + ],

520 + vec!["Std".to_string(), format!("{:.2}", std_dev(&e2el_ms))],

521 + ];

522 + for p in percentiles {

523 + rows.push(vec![

524 + format!("P{p:?}"),

525 + format!("{:.2}", percentile_sorted(&sorted, *p)),

526 + ]);

527 + }

528 + println!("End-to-End Latency (ms):");

529 + print_table(&rows);

530 + let throughput = completed as f64 / total_time.as_secs_f64();

531 + println!("Request throughput: {throughput:.3} req/s");

532 +}

533 +

534 +/// JSON result schema mirroring `vllm bench mm-processor`.

535 +#[derive(Serialize)]

536 +struct BenchmarkResultJson {

537 + completed: usize,

538 + failed: usize,

539 + mean_e2el_ms: f64,

540 + median_e2el_ms: f64,

541 + std_e2el_ms: f64,

542 + percentiles_e2el_ms: Vec<(f64, f64)>,

543 + /// `{stage_ms: {statistic: ms}}`, keyed by stage like the Python schema.

544 + mm_processor_stats: BTreeMap<String, BTreeMap<String, f64>>,

545 +}

546 +

547 +fn write_result_json(

548 + path: &str,

549 + completed: usize,

550 + failed: usize,

551 + e2el_times_secs: &[f64],

552 + percentiles: &[f64],

553 + mm_processor_stats: BTreeMap<String, BTreeMap<String, f64>>,

554 +) -> Result<()> {

555 + let e2el_ms: Vec<f64> = e2el_times_secs.iter().map(|secs| secs * SEC_TO_MS).collect();

556 + let sorted = sort_clone(&e2el_ms);

557 + let result = BenchmarkResultJson {

558 + completed,

559 + failed,

560 + mean_e2el_ms: mean(&e2el_ms),

561 + median_e2el_ms: median_sorted(&sorted),

562 + std_e2el_ms: std_dev(&e2el_ms),

563 + percentiles_e2el_ms: percentiles

564 + .iter()

565 + .map(|p| (*p, percentile_sorted(&sorted, *p)))

566 + .collect(),

567 + mm_processor_stats,

568 + };

569 + let json = serde_json::to_string_pretty(&result).context("failed to serialize result")?;

570 + std::fs::write(path, format!("{json}\n")).with_context(|| format!("failed to write {path}"))?;

571 + Ok(())

572 +}

rust/src/chat/Cargo.toml

@@ -35,6 +35,7 @@ vllm-llm.workspace = true

35 vllm-parser.workspace = true

36 vllm-text.workspace = true

37 vllm-tokenizer.workspace = true

38 +vllm-tracing.workspace = true

39 xgrammar-structural-tag.workspace = true

41 [dev-dependencies]

rust/src/chat/src/lib.rs

@@ -46,6 +46,9 @@ pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessag

46 pub use vllm_engine_core_client::protocol::multimodal::MmFeatures;

47 pub use vllm_llm::FinishReason;

48 pub use vllm_text::GenerationConfigMode;

49 +pub use vllm_tracing::timing::{

50 + RequestTimingStats as MultiModalTimingStats, mm_request_span, mm_timing_layer,

51 +};

53 mod backend;

54 mod error;

rust/src/chat/src/multimodal.rs

+40 -11

@@ -26,7 +26,7 @@ use llm_multimodal::{

26 };

27 use serde::{Deserialize, Serialize};

28 use thiserror_ext::AsReport as _;

29 -use tracing::warn;

29 +use tracing::{Instrument as _, warn};

30 use vllm_engine_core_client::protocol::dtype::ModelDtype;

31 use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem};

32 use vllm_text::Prompt;

@@ -47,6 +47,7 @@ mod video;

48 use self::expand::expand_prompt_token_ids;

49 pub use self::input::MultimodalInput;

50 +use vllm_tracing::timing::{mm_request_span, mm_stage_span};

52 /// Resolved multimodal support for one loaded model.

53 #[derive(Clone)]

@@ -562,7 +563,10 @@ pub(crate) async fn finalize_rendered_prompt(

563 .map_err(|error| multimodal!("{error}"))?,

564 Prompt::TokenIds(token_ids) => token_ids,

565 };

565 - let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;

566 + let prepared = info

567 + .prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype)

568 + .instrument(mm_request_span(&request.request_id))

569 + .await?;

570

571 Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))

572 }

@@ -753,27 +757,52 @@ impl MultimodalModelInfo {

757 prompt_token_ids: &mut Vec<u32>,

758 model_dtype: ModelDtype,

759 ) -> Result<MmFeatures> {

756 - let media_parts_len = media_parts.len();

757 - if media_parts_len == 0 {

760 + if media_parts.is_empty() {

761 return Ok(Vec::new());

762 }

763 + self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype)

764 + .instrument(mm_stage_span("preprocessor_total"))

765 + .await

766 + }

767 +

768 + /// Timed body of [`Self::prepare_multimodal`]; each stage is a `tracing`

769 + /// span aggregated per request by the `vllm-tracing` timing layer.

770 + async fn prepare_multimodal_timed(

771 + &self,

772 + media_parts: Vec<MediaContentPart>,

773 + prompt_token_ids: &mut Vec<u32>,

774 + model_dtype: ModelDtype,

775 + ) -> Result<MmFeatures> {

776 + let media_parts_len = media_parts.len();

777 self.validate_mm_limits(&media_parts)?;

761 - let fetched = self.fetch_media(media_parts).await?;

778 + let fetched =

779 + self.fetch_media(media_parts).instrument(mm_stage_span("media_fetch")).await?;

780

781 let mut prepared = Vec::new();

782 if !fetched.images.is_empty() {

765 - prepared

766 - .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?);

783 + prepared.push(

784 + self.prepare_images(fetched.images, fetched.image_uuids, model_dtype)

785 + .instrument(mm_stage_span("preprocess_image"))

786 + .await?,

787 + );

788 }

789 if !fetched.videos.is_empty() {

769 - prepared

770 - .push(self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype).await?);

790 + prepared.push(

791 + self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype)

792 + .instrument(mm_stage_span("preprocess_video"))

793 + .await?,

794 + );

795 }

796 if !fetched.audios.is_empty() {

773 - prepared.push(self.prepare_audios(fetched.audios, fetched.audio_uuids).await?);

797 + prepared.push(

798 + self.prepare_audios(fetched.audios, fetched.audio_uuids)

799 + .instrument(mm_stage_span("preprocess_audio"))

800 + .await?,

801 + );

802 }

803

776 - let mut ranges = expand_prompt_token_ids(prompt_token_ids, &prepared)?;

804 + let mut ranges = mm_stage_span("prompt_expansion")

805 + .in_scope(|| expand_prompt_token_ids(prompt_token_ids, &prepared))?;

806

807 let mut features = Vec::with_capacity(media_parts_len);

808 for media in prepared {

rust/src/server/src/routes/inference/generate.rs

+6 -1

@@ -53,7 +53,12 @@ pub async fn generate(

53 let lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await;

55 let mm_features = if let Some(parts) = body.content_parts.take() {

56 - match state.chat.prepare_media(parts, &mut body.token_ids).await {

56 + match state

57 + .chat

58 + .prepare_media(parts, &mut body.token_ids)

59 + .instrument(vllm_chat::mm_request_span(&request_context.request_id))

60 + .await

61 + {

62 Ok(features) => features,

63 Err(e) => {

64 return ApiError::invalid_request(

rust/src/tracing/src/lib.rs

+27 -9

@@ -3,13 +3,14 @@

4 //! Shared tracing subscriber and log formatting for vLLM Rust binaries.

6 +pub mod timing;

7 +

8 use std::{env, fmt, process};

10 use time::UtcOffset;

11 use time::macros::format_description;

12 use tracing::level_filters::LevelFilter;

13 use tracing::{Event, Level, Subscriber};

12 -use tracing_subscriber::Layer as _;

14 use tracing_subscriber::filter::Targets;

15 use tracing_subscriber::fmt::format::{FormatEvent, FormatFields, Writer};

16 use tracing_subscriber::fmt::time::FormatTime;

@@ -17,6 +18,7 @@ use tracing_subscriber::fmt::{FmtContext, FormattedFields};

18 use tracing_subscriber::layer::SubscriberExt as _;

19 use tracing_subscriber::registry::LookupSpan;

20 use tracing_subscriber::util::SubscriberInitExt as _;

21 +use tracing_subscriber::{Layer, Registry};

23 const CYAN: &str = "\x1b[0;36m";

24 const GREY: &str = "\x1b[90m";

@@ -30,20 +32,36 @@ const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =

33 /// Install the process-wide vLLM-style tracing subscriber.

34 pub fn init_tracing(process_label: &str) {

35 + let _ = tracing_subscriber::registry().with(vllm_fmt_layer(process_label)).try_init();

36 +}

37 +

38 +/// Install the process-wide vLLM-style tracing subscriber, with an

39 +/// additional layer (e.g. the multimodal preprocessing timing collector).

40 +pub fn init_tracing_with<L>(process_label: &str, layer: L)

41 +where

42 + L: Layer<Registry> + Send + Sync + 'static,

43 +{

44 + let _ = tracing_subscriber::registry()

45 + .with(layer)

46 + .with(vllm_fmt_layer(process_label))

47 + .try_init();

48 +}

49 +

50 +/// The vLLM-style fmt layer filtered by `VLLM_LOGGING_LEVEL` / `RUST_LOG`.

51 +fn vllm_fmt_layer<S>(process_label: &str) -> impl Layer<S> + Send + Sync + 'static + use<S>

52 +where

53 + S: Subscriber + for<'lookup> LookupSpan<'lookup>,

54 +{

55 let filter = build_targets_filter(

56 env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),

57 env::var("RUST_LOG").ok().as_deref(),

58 );

59 let formatter = VllmEventFormatter::new(process_label);

39 - let _ = tracing_subscriber::registry()

40 - .with(

41 - tracing_subscriber::fmt::layer()

42 - .event_format(formatter)

43 - .with_writer(std::io::stderr)

44 - .with_filter(filter),

45 - )

46 - .try_init();

61 + tracing_subscriber::fmt::layer()

62 + .event_format(formatter)

63 + .with_writer(std::io::stderr)

64 + .with_filter(filter)

65 }

67 /// Build the CLI log filter by merging the vLLM-style default level with

rust/src/tracing/src/timing.rs new

+229

@@ -0,0 +1,229 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +//! Request-scoped span timing.

5 +//!

6 +//! Pipelines emit stage spans (target = the layer's configured target, a

7 +//! `stage` field naming the stage) inside a request span carrying a

8 +//! `request_id` field; [`RequestTimingLayer`] attributes each stage span's

9 +//! wall-clock time to the nearest ancestor `request_id` and accumulates per

10 +//! request. Callers that want timings install the layer on their subscriber;

11 +//! everyone else pays only the cost of untracked spans.

12 +

13 +use std::collections::HashMap;

14 +use std::fmt;

15 +use std::sync::{Arc, Mutex};

16 +use std::time::Instant;

17 +

18 +use tracing::field::{Field, Visit};

19 +use tracing::span::{Attributes, Span};

20 +use tracing::{Id, Subscriber, info_span};

21 +use tracing_subscriber::layer::{Context, Layer};

22 +use tracing_subscriber::registry::LookupSpan;

23 +

24 +/// `"{stage}_secs" => seconds` maps returned by [`RequestTimingStats::stat`].

25 +pub type StageStats = HashMap<String, f64>;

26 +

27 +/// `request_id` extracted from a span, propagated to stage timings.

28 +struct RequestId(String);

29 +

30 +/// Start timestamp of one in-flight stage span.

31 +struct StageStart {

32 + stage: String,

33 + start: Instant,

34 +}

35 +

36 +/// Extracts the fields the layer cares about from span attributes.

37 +#[derive(Default)]

38 +struct SpanFields {

39 + request_id: Option<String>,

40 + stage: Option<String>,

41 +}

42 +

43 +impl Visit for SpanFields {

44 + fn record_str(&mut self, field: &Field, value: &str) {

45 + match field.name() {

46 + "request_id" => self.request_id = Some(value.to_string()),

47 + "stage" => self.stage = Some(value.to_string()),

48 + _ => {}

49 + }

50 + }

51 +

52 + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {

53 + match field.name() {

54 + "request_id" => self.request_id = Some(format!("{value:?}")),

55 + "stage" => self.stage = Some(format!("{value:?}")),

56 + _ => {}

57 + }

58 + }

59 +}

60 +

61 +/// `tracing` layer that aggregates stage-span timings per request, keyed by

62 +/// the nearest ancestor span with a `request_id` field.

63 +pub struct RequestTimingLayer {

64 + target: &'static str,

65 + stats: Arc<Mutex<HashMap<String, StageStats>>>,

66 +}

67 +

68 +/// Drain handle for the timings collected by [`RequestTimingLayer`].

69 +pub struct RequestTimingStats {

70 + stats: Arc<Mutex<HashMap<String, StageStats>>>,

71 +}

72 +

73 +impl RequestTimingLayer {

74 + /// Create the layer and its stats handle. Only spans whose target equals

75 + /// `target` and that carry a `stage` field are timed.

76 + pub fn new(target: &'static str) -> (Self, RequestTimingStats) {

77 + let stats = Arc::new(Mutex::new(HashMap::new()));

78 + (

79 + Self {

80 + target,

81 + stats: Arc::clone(&stats),

82 + },

83 + RequestTimingStats { stats },

84 + )

85 + }

86 +}

87 +

88 +impl RequestTimingStats {

89 + /// Drain and return `{request_id: {stage_secs}}` records.

90 + pub fn stat(&self) -> HashMap<String, StageStats> {

91 + std::mem::take(&mut *self.stats.lock().unwrap())

92 + }

93 +}

94 +

95 +impl<S> Layer<S> for RequestTimingLayer

96 +where

97 + S: Subscriber + for<'lookup> LookupSpan<'lookup>,

98 +{

99 + fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {

100 + let Some(span) = ctx.span(id) else {

101 + return;

102 + };

103 + let mut fields = SpanFields::default();

104 + attrs.record(&mut fields);

105 + let mut extensions = span.extensions_mut();

106 + if let Some(request_id) = fields.request_id {

107 + extensions.insert(RequestId(request_id));

108 + }

109 + if span.metadata().target() == self.target

110 + && let Some(stage) = fields.stage

111 + {

112 + extensions.insert(StageStart {

113 + stage,

114 + start: Instant::now(),

115 + });

116 + }

117 + }

118 +

119 + fn on_close(&self, id: Id, ctx: Context<'_, S>) {

120 + let Some(span) = ctx.span(&id) else {

121 + return;

122 + };

123 + let (stage, start) = {

124 + let extensions = span.extensions();

125 + let Some(stage_start) = extensions.get::<StageStart>() else {

126 + return;

127 + };

128 + (stage_start.stage.clone(), stage_start.start)

129 + };

130 + let Some(request_id) = span

131 + .scope()

132 + .find_map(|ancestor| ancestor.extensions().get::<RequestId>().map(|id| id.0.clone()))

133 + else {

134 + return;

135 + };

136 + let mut stats = self.stats.lock().unwrap();

137 + *stats.entry(request_id).or_default().entry(format!("{stage}_secs")).or_default() +=

138 + start.elapsed().as_secs_f64();

139 + }

140 +}

141 +

142 +/// Target of the multimodal preprocessing stage spans.

143 +const MM_STAGE_TARGET: &str = "mm_processor_timing";

144 +

145 +/// Create the timing layer and its stats handle for multimodal preprocessing

146 +/// stage spans (`vllm-bench mm-processor`), mirroring the Python

147 +/// `TimingContext` / `MultiModalTimingRegistry`.

148 +pub fn mm_timing_layer() -> (RequestTimingLayer, RequestTimingStats) {

149 + RequestTimingLayer::new(MM_STAGE_TARGET)

150 +}

151 +

152 +/// Span carrying the `request_id` used to attribute multimodal stage timings.

153 +pub fn mm_request_span(request_id: &str) -> Span {

154 + info_span!("mm_request", request_id)

155 +}

156 +

157 +/// Span for one multimodal preprocessing stage; the elapsed time is recorded

158 +/// on close.

159 +pub fn mm_stage_span(stage: &'static str) -> Span {

160 + info_span!(target: MM_STAGE_TARGET, "mm_stage", stage)

161 +}

162 +

163 +#[cfg(test)]

164 +mod tests {

165 + use super::*;

166 + use tracing::info_span;

167 + use tracing_subscriber::layer::SubscriberExt as _;

168 +

169 + const TARGET: &str = "test_request_timing";

170 +

171 + fn stage_span(stage: &'static str) -> tracing::Span {

172 + info_span!(target: TARGET, "stage", stage)

173 + }

174 +

175 + #[test]

176 + fn records_stage_timings_per_request() {

177 + let (layer, stats) = RequestTimingLayer::new(TARGET);

178 + let subscriber = tracing_subscriber::registry().with(layer);

179 + tracing::subscriber::with_default(subscriber, || {

180 + info_span!("request", request_id = "req-1").in_scope(|| {

181 + drop(stage_span("media_fetch"));

182 + stage_span("prompt_expansion").in_scope(|| {});

183 + });

184 + });

185 + let drained = stats.stat();

186 + assert!(drained["req-1"].contains_key("media_fetch_secs"));

187 + assert!(drained["req-1"].contains_key("prompt_expansion_secs"));

188 + // `stat` drains: a second read is empty.

189 + assert!(stats.stat().is_empty());

190 + }

191 +

192 + #[test]

193 + fn attributes_to_nearest_request_id_ancestor() {

194 + let (layer, stats) = RequestTimingLayer::new(TARGET);

195 + let subscriber = tracing_subscriber::registry().with(layer);

196 + tracing::subscriber::with_default(subscriber, || {

197 + info_span!("request", request_id = "req-outer").in_scope(|| {

198 + info_span!("request", request_id = "req-inner").in_scope(|| {

199 + drop(stage_span("media_fetch"));

200 + });

201 + });

202 + });

203 + let drained = stats.stat();

204 + assert!(drained["req-inner"].contains_key("media_fetch_secs"));

205 + assert!(!drained.contains_key("req-outer"));

206 + }

207 +

208 + #[test]

209 + fn stage_without_request_id_is_dropped() {

210 + let (layer, stats) = RequestTimingLayer::new(TARGET);

211 + let subscriber = tracing_subscriber::registry().with(layer);

212 + tracing::subscriber::with_default(subscriber, || {

213 + drop(stage_span("media_fetch"));

214 + });

215 + assert!(stats.stat().is_empty());

216 + }

217 +

218 + #[test]

219 + fn other_targets_are_ignored() {

220 + let (layer, stats) = RequestTimingLayer::new(TARGET);

221 + let subscriber = tracing_subscriber::registry().with(layer);

222 + tracing::subscriber::with_default(subscriber, || {

223 + info_span!("request", request_id = "req-1").in_scope(|| {

224 + drop(info_span!("unrelated", stage = "media_fetch"));

225 + });

226 + });

227 + assert!(stats.stat().is_empty());

228 + }

229 +}