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

[Rust Frontend] Separate multimodal instrumentation from request timing (#58084)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Sep 22, 2026 at 14:57 UTC a6c47fbbf4187e8ae5396ebdcf93234db617a4ad
9 files changed +143 -52
rust/src/bench/src/mm_processor.rs
+4 -3
@@ -16,13 +16,14 @@ use tokio::sync::Semaphore;
16
17 use vllm_chat::{
18 ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRequest, LoadModelBackendsOptions,
19 - MultiModalTimingStats, load_model_backends,
19 + 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 +use vllm_tracing::timing::RequestTimingStats;
27
28 use crate::config::RangeRatio;
29 use crate::datasets::SampleRequest;
@@ -133,7 +134,7 @@ pub struct MmProcessorArgs {
134 }
135
136 /// Entry point for `vllm-bench mm-processor`.
136 -pub async fn run_mm_processor(args: MmProcessorArgs, timing: MultiModalTimingStats) -> Result<()> {
137 +pub async fn run_mm_processor(args: MmProcessorArgs, timing: RequestTimingStats) -> Result<()> {
138 if args.engine.data_parallel_size_local == Some(0) {
139 anyhow::bail!(
140 "`--data-parallel-size-local 0` is not supported; the benchmark requires \
@@ -179,7 +180,7 @@ pub async fn run_mm_processor(args: MmProcessorArgs, timing: MultiModalTimingSta
180 async fn run_mm_processor_with_engine(
181 args: MmProcessorArgs,
182 handshake_port: u16,
182 - timing: &MultiModalTimingStats,
183 + timing: &RequestTimingStats,
184 ) -> Result<()> {
185 let range_ratio =
186 RangeRatio::parse(&args.random_range_ratio).context("invalid --random-range-ratio")?;
rust/src/chat/src/lib.rs
+1 -3
@@ -22,6 +22,7 @@ pub use event::{
22 };
23 use futures::{StreamExt, TryStreamExt as _};
24 pub use llm_multimodal::MediaContentPart;
25 +pub use multimodal::{mm_request_span, mm_timing_layer};
26 use output::apply_output_grammar;
27 pub use output::{
28 ChatOutputProcessor, DefaultChatOutputProcessor, DynChatOutputProcessor,
@@ -47,9 +48,6 @@ pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessag
48 pub use vllm_engine_core_client::protocol::multimodal::MmFeatures;
49 pub use vllm_llm::FinishReason;
50 pub use vllm_text::GenerationConfigMode;
50 -pub use vllm_tracing::timing::{
51 - RequestTimingStats as MultiModalTimingStats, mm_request_span, mm_timing_layer,
52 -};
51
52 mod backend;
53 mod error;
rust/src/chat/src/multimodal.rs
+23 -23
@@ -43,11 +43,13 @@ mod input;
43 mod item;
44 mod preprocessed;
45 mod tensor;
46 +mod timing;
47 mod video;
48
49 use self::expand::expand_prompt_token_ids;
50 pub use self::input::MultimodalInput;
50 -use vllm_tracing::timing::{mm_request_span, mm_stage_span};
51 +use self::timing::MM_STAGE_TARGET;
52 +pub use self::timing::{mm_request_span, mm_timing_layer};
53
54 /// Resolved multimodal support for one loaded model.
55 #[derive(Clone)]
@@ -760,13 +762,17 @@ impl MultimodalModelInfo {
762 if media_parts.is_empty() {
763 return Ok(Vec::new());
764 }
763 - self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype)
764 - .instrument(mm_stage_span("preprocessor_total"))
765 - .await
765 + self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype).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 + #[tracing::instrument(
771 + name = "mm_stage",
772 + target = MM_STAGE_TARGET,
773 + skip_all,
774 + fields(stage = "preprocessor_total")
775 + )]
776 async fn prepare_multimodal_timed(
777 &self,
778 media_parts: Vec<MediaContentPart>,
@@ -775,34 +781,22 @@ impl MultimodalModelInfo {
781 ) -> Result<MmFeatures> {
782 let media_parts_len = media_parts.len();
783 self.validate_mm_limits(&media_parts)?;
778 - let fetched =
779 - self.fetch_media(media_parts).instrument(mm_stage_span("media_fetch")).await?;
784 + let fetched = self.fetch_media(media_parts).await?;
785
786 let mut prepared = Vec::new();
787 if !fetched.images.is_empty() {
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 + prepared
789 + .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?);
790 }
791 if !fetched.videos.is_empty() {
790 - prepared.push(
791 - self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype)
792 - .instrument(mm_stage_span("preprocess_video"))
793 - .await?,
794 - );
792 + prepared
793 + .push(self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype).await?);
794 }
795 if !fetched.audios.is_empty() {
797 - prepared.push(
798 - self.prepare_audios(fetched.audios, fetched.audio_uuids)
799 - .instrument(mm_stage_span("preprocess_audio"))
800 - .await?,
801 - );
796 + prepared.push(self.prepare_audios(fetched.audios, fetched.audio_uuids).await?);
797 }
798
804 - let mut ranges = mm_stage_span("prompt_expansion")
805 - .in_scope(|| expand_prompt_token_ids(prompt_token_ids, &prepared))?;
799 + let mut ranges = expand_prompt_token_ids(prompt_token_ids, &prepared)?;
800
801 let mut features = Vec::with_capacity(media_parts_len);
802 for media in prepared {
@@ -845,6 +839,12 @@ impl MultimodalModelInfo {
839
840 /// Fetch all connector-backed media parts and split them per modality,
841 /// preserving their request-order UUID metadata.
842 + #[tracing::instrument(
843 + name = "mm_stage",
844 + target = MM_STAGE_TARGET,
845 + skip_all,
846 + fields(stage = "media_fetch")
847 + )]
848 async fn fetch_media(&self, media_parts: Vec<MediaContentPart>) -> Result<FetchedMedia> {
849 let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector));
850 for part in media_parts {
rust/src/chat/src/multimodal/audio.rs
+46
@@ -8,6 +8,7 @@ use std::sync::Arc;
8 use llm_multimodal::{AudioClip, Modality, PreprocessedEncoderInputs};
9 use vllm_engine_core_client::protocol::dtype::ModelDtype;
10
11 +use super::timing::MM_STAGE_TARGET;
12 use super::{AudioModalitySupport, MultimodalModelInfo, PreparedMedia, item};
13 use crate::error::{Error, Result, bail_multimodal, multimodal};
14
@@ -16,6 +17,12 @@ pub(super) const AUDIO_PRIMARY_KEY: &str = "input_audio_features";
17
18 impl MultimodalModelInfo {
19 /// Preprocess fetched audio clips as one batch and build per-item features.
20 + #[tracing::instrument(
21 + name = "mm_stage",
22 + target = MM_STAGE_TARGET,
23 + skip_all,
24 + fields(stage = "preprocess_audio")
25 + )]
26 pub(super) async fn prepare_audios(
27 &self,
28 clips: Vec<Arc<AudioClip>>,
@@ -236,6 +243,45 @@ mod tests {
243 ));
244 }
245
246 + #[tokio::test]
247 + async fn audio_pipeline_attributes_stage_timings_to_request() {
248 + use tracing::Instrument as _;
249 + use tracing::instrument::WithSubscriber as _;
250 + use tracing_subscriber::layer::SubscriberExt as _;
251 +
252 + use crate::{mm_request_span, mm_timing_layer};
253 +
254 + let info = qwen3_asr_info();
255 + let (layer, timings) = mm_timing_layer();
256 + let parts = vec![MediaContentPart::AudioData {
257 + data: wav_i16_mono(16_000, &[0; 1_600]),
258 + mime_type: Some("audio/wav".to_string()),
259 + uuid: Some("audio-1".to_string()),
260 + }];
261 + async {
262 + info.prepare_multimodal(parts, &mut vec![AUDIO_PAD_ID], ModelDtype::Float32)
263 + .instrument(mm_request_span("request-1"))
264 + .await
265 + .unwrap();
266 + }
267 + .with_subscriber(tracing_subscriber::registry().with(layer))
268 + .await;
269 +
270 + let records = timings.stat();
271 + assert_eq!(records.len(), 1);
272 + let mut stages = records["request-1"].keys().collect::<Vec<_>>();
273 + stages.sort();
274 + expect_test::expect![[r#"
275 + [
276 + "media_fetch_secs",
277 + "preprocess_audio_secs",
278 + "preprocessor_total_secs",
279 + "prompt_expansion_secs",
280 + ]
281 + "#]]
282 + .assert_debug_eq(&stages);
283 + }
284 +
285 #[tokio::test]
286 async fn tracker_processor_and_lowering_preserve_audio_contract() {
287 let info = qwen3_asr_info();
rust/src/chat/src/multimodal/expand.rs
+7
@@ -10,6 +10,7 @@ use vllm_engine_core_client::protocol::multimodal::PlaceholderRange;
10 use vllm_engine_core_client::protocol::tensor::WireTensor;
11
12 use super::PreparedMedia;
13 +use super::timing::MM_STAGE_TARGET;
14 use crate::error::{Error, Result, bail_multimodal};
15
16 /// One modality's queue of pending placeholder replacements for prompt
@@ -55,6 +56,12 @@ impl<'a> ExpansionLane<'a> {
56 ///
57 /// The returned ranges point into the already-expanded prompt, grouped per
58 /// modality in item order.
59 +#[tracing::instrument(
60 + name = "mm_stage",
61 + target = MM_STAGE_TARGET,
62 + skip_all,
63 + fields(stage = "prompt_expansion")
64 +)]
65 pub(super) fn expand_prompt_token_ids(
66 prompt_token_ids: &mut Vec<u32>,
67 prepared: &[PreparedMedia],
rust/src/chat/src/multimodal/image.rs
+7
@@ -9,6 +9,7 @@ use std::sync::Arc;
9 use llm_multimodal::{ImageFrame, Modality, PreprocessedEncoderInputs};
10 use vllm_engine_core_client::protocol::dtype::ModelDtype;
11
12 +use super::timing::MM_STAGE_TARGET;
13 use super::{ModalitySupport, MultimodalModelInfo, PreparedMedia, item};
14 use crate::error::{Error, Result, bail_multimodal, multimodal};
15
@@ -18,6 +19,12 @@ pub(super) const IMAGE_PRIMARY_KEY: &str = "pixel_values";
19 impl MultimodalModelInfo {
20 /// Preprocess all fetched image frames as one batch and build per-item
21 /// features.
22 + #[tracing::instrument(
23 + name = "mm_stage",
24 + target = MM_STAGE_TARGET,
25 + skip_all,
26 + fields(stage = "preprocess_image")
27 + )]
28 pub(super) async fn prepare_images(
29 &self,
30 frames: Vec<Arc<ImageFrame>>,
rust/src/chat/src/multimodal/timing.rs new
+22
@@ -0,0 +1,22 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Multimodal preprocessing spans and their timing collector configuration.
5 +
6 +use tracing::{Span, info_span};
7 +use vllm_tracing::timing::{RequestTimingLayer, RequestTimingStats};
8 +
9 +/// Target of the multimodal preprocessing stage spans.
10 +pub(super) const MM_STAGE_TARGET: &str = "mm_processor_timing";
11 +
12 +/// Create the timing layer and its stats handle for multimodal preprocessing
13 +/// stage spans (`vllm-bench mm-processor`), mirroring the Python
14 +/// `TimingContext` / `MultiModalTimingRegistry`.
15 +pub fn mm_timing_layer() -> (RequestTimingLayer, RequestTimingStats) {
16 + RequestTimingLayer::new(MM_STAGE_TARGET)
17 +}
18 +
19 +/// Span carrying the `request_id` used to attribute multimodal stage timings.
20 +pub fn mm_request_span(request_id: &str) -> Span {
21 + info_span!("mm_request", request_id)
22 +}
rust/src/chat/src/multimodal/video.rs
+7
@@ -16,6 +16,7 @@ use vllm_engine_core_client::protocol::multimodal::{
16 SliceSpec,
17 };
18
19 +use super::timing::MM_STAGE_TARGET;
20 use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor};
21 use crate::error::{Error, Result, bail_multimodal, multimodal};
22
@@ -32,6 +33,12 @@ impl MultimodalModelInfo {
33 /// Unlike images, each clip runs through the preprocessor independently
34 /// (a batch of one), so its tensors are complete per item and need no
35 /// cross-item slicing.
36 + #[tracing::instrument(
37 + name = "mm_stage",
38 + target = MM_STAGE_TARGET,
39 + skip_all,
40 + fields(stage = "preprocess_video")
41 + )]
42 pub(super) async fn prepare_videos(
43 &self,
44 clips: Vec<Arc<VideoClip>>,
rust/src/tracing/src/timing.rs
+26 -23
@@ -16,8 +16,8 @@ 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};
19 +use tracing::span::Attributes;
20 +use tracing::{Id, Subscriber};
21 use tracing_subscriber::layer::{Context, Layer};
22 use tracing_subscriber::registry::LookupSpan;
23
@@ -60,6 +60,26 @@ impl Visit for SpanFields {
60
61 /// `tracing` layer that aggregates stage-span timings per request, keyed by
62 /// the nearest ancestor span with a `request_id` field.
63 +///
64 +/// # Span contract
65 +///
66 +/// Timed spans must have the configured target and a `stage` field. The
67 +/// `request_id` is taken from the timed span itself or its nearest ancestor
68 +/// carrying that field, regardless of the ancestor's target. Spans without
69 +/// a request ID are ignored. Span names are unrestricted.
70 +///
71 +/// Both fields are read at span creation; subsequent [`tracing::Span::record`]
72 +/// calls are ignored. String values are used verbatim; other values use their
73 +/// debug representation. Prefer string fields for stable keys.
74 +///
75 +/// # Timing semantics
76 +///
77 +/// Each sample measures wall-clock time from span creation to final close,
78 +/// including async waits and any time the span remains alive through cloned
79 +/// handles. Nested and concurrent spans are measured independently, so their
80 +/// durations can overlap. Repeated stages accumulate under `"{stage}_secs"`
81 +/// for each request. Closing spans are recorded regardless of the operation's
82 +/// success, failure, or cancellation; callers select the samples to report.
83 pub struct RequestTimingLayer {
84 target: &'static str,
85 stats: Arc<Mutex<HashMap<String, StageStats>>>,
@@ -87,6 +107,10 @@ impl RequestTimingLayer {
107
108 impl RequestTimingStats {
109 /// Drain and return `{request_id: {stage_secs}}` records.
110 + ///
111 + /// Returns accumulated timings since the previous drain. Spans still open
112 + /// contribute to a later drain; finish the measured work before calling
113 + /// this method to collect complete request timings.
114 pub fn stat(&self) -> HashMap<String, StageStats> {
115 std::mem::take(&mut *self.stats.lock().unwrap())
116 }
@@ -139,27 +163,6 @@ where
163 }
164 }
165
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 -
166 #[cfg(test)]
167 mod tests {
168 use super::*;