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

[Rust Frontend] Pass vision preprocessing context for Nemotron-H (#57634)

Pass the remaining engine context-length budget to model-owned vision processors through VisionPreprocessingContext. Preserve Nemotron batched engine fields and recognize llm_config as a text_config alias. Use the merged upstream llm-multimodal revision f0985ef65967615db2c79279aa07818499301bfd. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Biswa Panda committed Sep 23, 2026 at 23:08 UTC cf78797cb79e73bb798746c6cce7c34f56860a0a
8 files changed +240 -21
rust/Cargo.lock
+1 -1
@@ -2290,7 +2290,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
2290 [[package]]
2291 name = "llm-multimodal"
2292 version = "1.7.1"
2293 -source = "git+https://github.com/smg-project/llm-multimodal?rev=53ca0d6ea829af7477b6b63851cc70bed409bbc8#53ca0d6ea829af7477b6b63851cc70bed409bbc8"
2293 +source = "git+https://github.com/smg-project/llm-multimodal?rev=f0985ef65967615db2c79279aa07818499301bfd#f0985ef65967615db2c79279aa07818499301bfd"
2294 dependencies = [
2295 "anyhow",
2296 "base64 0.22.1",
rust/Cargo.toml
+1 -1
@@ -61,7 +61,7 @@ indexmap = "2.13.0"
61 indicatif = "0.18.4"
62 itertools = "0.14.0"
63 libc = "0.2.177"
64 -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "53ca0d6ea829af7477b6b63851cc70bed409bbc8", default-features = false, features = ["native-tls"] }
64 +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "f0985ef65967615db2c79279aa07818499301bfd", default-features = false, features = ["native-tls"] }
65 mimalloc = "0.1.52"
66 minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
67 minijinja-contrib = { version = "2.24", features = ["pycompat"] }
rust/src/chat/src/lib.rs
+11 -3
@@ -72,6 +72,9 @@ pub struct ChatRequestProcessor {
72 /// Effective model dtype reported by the engine.
73 /// Absent for text-only frontends without an engine handshake.
74 model_dtype: Option<ModelDtype>,
75 + /// Effective runtime context length reported by the engine.
76 + /// Absent for render-only frontends without an engine handshake.
77 + max_model_len: Option<usize>,
78 /// Tool-call parser selection used when preparing generation requests.
79 tool_call_parser: ParserSelection,
80 /// Reasoning parser selection used when preparing generation requests.
@@ -83,10 +86,11 @@ pub struct ChatRequestProcessor {
86 impl ChatRequestProcessor {
87 /// Create a processor with multimodal support using the effective model
88 /// dtype reported by the engine.
86 - fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
89 + fn new(backend: DynChatBackend, model_dtype: ModelDtype, max_model_len: u32) -> Self {
90 Self {
91 backend,
92 model_dtype: Some(model_dtype),
93 + max_model_len: Some(max_model_len as usize),
94 tool_call_parser: ParserSelection::Auto,
95 reasoning_parser: ParserSelection::Auto,
96 tool_strict_level: ToolStrictLevel::Auto,
@@ -98,6 +102,7 @@ impl ChatRequestProcessor {
102 Self {
103 backend,
104 model_dtype: None,
105 + max_model_len: None,
106 tool_call_parser: ParserSelection::Auto,
107 reasoning_parser: ParserSelection::Auto,
108 tool_strict_level: ToolStrictLevel::Auto,
@@ -137,6 +142,7 @@ impl ChatRequestProcessor {
142 rendered,
143 self.backend.multimodal_model_info(),
144 model_dtype,
145 + self.max_model_len,
146 )
147 .await
148 }
@@ -161,7 +167,8 @@ impl ChatRequestProcessor {
167 let features = match media {
168 multimodal::MultimodalInput::Raw(parts) => {
169 let model_dtype = self.model_dtype.ok_or(Error::UnsupportedMultimodalRenderer)?;
164 - info.prepare_multimodal(parts, token_ids, model_dtype).await?
170 + info.prepare_multimodal(parts, token_ids, model_dtype, self.max_model_len)
171 + .await?
172 }
173 multimodal::MultimodalInput::Preprocessed(features) => {
174 info.prepare_preprocessed(features, token_ids.len())?
@@ -256,10 +263,11 @@ impl ChatLlm {
263 /// backend.
264 pub fn new(text: TextLlm, backend: DynChatBackend) -> Self {
265 let model_dtype = text.engine_core_client().model_dtype();
266 + let max_model_len = text.engine_core_client().max_model_len();
267
268 Self {
269 text,
262 - processor: ChatRequestProcessor::new(backend, model_dtype),
270 + processor: ChatRequestProcessor::new(backend, model_dtype, max_model_len),
271 }
272 }
273
rust/src/chat/src/multimodal.rs
+142 -4
@@ -23,6 +23,7 @@ use llm_multimodal::{
23 ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata,
24 ModelProcessorSpec, ModelRegistry, PreProcessorConfig, PreprocessedEncoderInputs,
25 PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor,
26 + VisionPreprocessingContext,
27 };
28 use serde::{Deserialize, Serialize};
29 use thiserror_ext::AsReport as _;
@@ -55,6 +56,8 @@ pub use self::timing::{mm_request_span, mm_timing_layer};
56 #[derive(Clone)]
57 pub struct MultimodalModelInfo {
58 context: MultimodalModelContext,
59 + /// Rendered placeholder marker IDs for all resolved modalities.
60 + placeholder_token_ids: Vec<u32>,
61 image: Option<VisionModalitySupport>,
62 video: Option<VisionModalitySupport>,
63 audio: Option<AudioModalitySupport>,
@@ -420,8 +423,18 @@ impl MultimodalModelInfo {
423 MediaConnectorConfig::default(),
424 )?);
425
426 + let placeholder_token_ids = [
427 + image.as_ref().map(|support| support.placeholder.marker_token_id),
428 + video.as_ref().map(|support| support.placeholder.marker_token_id),
429 + audio.as_ref().map(|support| support.placeholder.marker_token_id),
430 + ]
431 + .into_iter()
432 + .flatten()
433 + .collect();
434 +
435 Ok(Some(Self {
436 context,
437 + placeholder_token_ids,
438 image,
439 video,
440 audio,
@@ -547,6 +560,20 @@ impl MultimodalModelInfo {
560 Modality::ImageEmbeds => None,
561 }
562 }
563 +
564 + fn vision_preprocessing_context(
565 + &self,
566 + prompt_token_ids: &[u32],
567 + max_model_len: Option<usize>,
568 + ) -> VisionPreprocessingContext {
569 + let text_prompt_length = prompt_token_ids
570 + .iter()
571 + .filter(|token_id| !self.placeholder_token_ids.contains(token_id))
572 + .count();
573 + VisionPreprocessingContext {
574 + token_budget: max_model_len.map(|limit| limit.saturating_sub(text_prompt_length)),
575 + }
576 + }
577 }
578
579 /// Finalize a rendered chat prompt into text-generation input.
@@ -559,6 +586,7 @@ pub(crate) async fn finalize_rendered_prompt(
586 rendered: RenderedPrompt,
587 info: Option<&MultimodalModelInfo>,
588 model_dtype: ModelDtype,
589 + max_model_len: Option<usize>,
590 ) -> Result<(Prompt, Option<MmFeatures>)> {
591 let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?;
592 if media_parts.is_empty() {
@@ -574,7 +602,12 @@ pub(crate) async fn finalize_rendered_prompt(
602 Prompt::TokenIds(token_ids) => token_ids,
603 };
604 let prepared = info
577 - .prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype)
605 + .prepare_multimodal(
606 + media_parts,
607 + &mut prompt_token_ids,
608 + model_dtype,
609 + max_model_len,
610 + )
611 .instrument(mm_request_span(&request.request_id))
612 .await?;
613
@@ -768,11 +801,13 @@ impl MultimodalModelInfo {
801 media_parts: Vec<MediaContentPart>,
802 prompt_token_ids: &mut Vec<u32>,
803 model_dtype: ModelDtype,
804 + max_model_len: Option<usize>,
805 ) -> Result<MmFeatures> {
806 if media_parts.is_empty() {
807 return Ok(Vec::new());
808 }
775 - self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype).await
809 + self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype, max_model_len)
810 + .await
811 }
812
813 /// Timed body of [`Self::prepare_multimodal`]; each stage is a `tracing`
@@ -788,15 +823,24 @@ impl MultimodalModelInfo {
823 media_parts: Vec<MediaContentPart>,
824 prompt_token_ids: &mut Vec<u32>,
825 model_dtype: ModelDtype,
826 + max_model_len: Option<usize>,
827 ) -> Result<MmFeatures> {
828 let media_parts_len = media_parts.len();
829 self.validate_mm_limits(&media_parts)?;
830 + let vision_context = self.vision_preprocessing_context(prompt_token_ids, max_model_len);
831 let fetched = self.fetch_media(media_parts).await?;
832
833 let mut prepared = Vec::new();
834 if !fetched.images.is_empty() {
798 - prepared
799 - .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?);
835 + prepared.push(
836 + self.prepare_images(
837 + fetched.images,
838 + fetched.image_uuids,
839 + model_dtype,
840 + vision_context,
841 + )
842 + .await?,
843 + );
844 }
845 if !fetched.videos.is_empty() {
846 prepared
@@ -1044,6 +1088,100 @@ mod tests {
1088 );
1089 }
1090
1091 + #[test]
1092 + fn vision_preprocessing_context_uses_runtime_budget_and_text_tokens() {
1093 + let info = qwen3_vl_info();
1094 + let context = info.vision_preprocessing_context(
1095 + &[11, QWEN3_IMAGE_PAD_ID, 12, QWEN3_VIDEO_PAD_ID, 13],
1096 + Some(4096),
1097 + );
1098 +
1099 + assert_eq!(context.token_budget, Some(4096 - 3));
1100 + assert_eq!(
1101 + info.vision_preprocessing_context(&[11, 12, 13], Some(2)).token_budget,
1102 + Some(0)
1103 + );
1104 + assert_eq!(
1105 + info.vision_preprocessing_context(&[11], None).token_budget,
1106 + None
1107 + );
1108 + }
1109 +
1110 + #[tokio::test]
1111 + async fn nemotron_budget_controls_preprocessing_and_engine_features() {
1112 + use vllm_engine_core_client::protocol::multimodal::MmKwargValue;
1113 +
1114 + let tokenizer = TestTokenizer::new()
1115 + .with_regular_token("<image>", 1018)
1116 + .with_regular_token("<img>", 1019)
1117 + .with_regular_token("</img>", 1020);
1118 + let config = serde_json::json!({
1119 + "model_type": "nemotron_h_omni",
1120 + "img_context_token_id": 1018,
1121 + "max_position_embeddings": 512,
1122 + "patch_size": 16,
1123 + "downsample_ratio": 0.5,
1124 + "norm_mean": [0.0, 0.0, 0.0],
1125 + "norm_std": [1.0, 1.0, 1.0],
1126 + "vision_config": {
1127 + "args": {"min_num_patches": 1024, "max_num_patches": 13312}
1128 + }
1129 + });
1130 + let info = test_info("nemotron_h_omni", config, tokenizer);
1131 + let media = MediaContentPart::ImageUrl {
1132 + url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=".to_string(),
1133 + detail: None,
1134 + uuid: None,
1135 + };
1136 + let mut prompt = vec![11; 100];
1137 + prompt.push(1018);
1138 +
1139 + let mut tokens = prompt.clone();
1140 + let features = info
1141 + .prepare_multimodal(
1142 + vec![media.clone()],
1143 + &mut tokens,
1144 + ModelDtype::Float32,
1145 + Some(512),
1146 + )
1147 + .await
1148 + .unwrap();
1149 + assert_eq!(features.len(), 1);
1150 + let feature = &features[0];
1151 + assert_eq!(
1152 + (feature.mm_position.offset, feature.mm_position.length),
1153 + (100, 258)
1154 + );
1155 + assert_eq!(tokens.len(), 358);
1156 + assert_eq!(tokens[100], 1019);
1157 + assert!(tokens[101..357].iter().all(|&token| token == 1018));
1158 + assert_eq!(tokens[357], 1020);
1159 + let data = feature.data.as_ref().unwrap();
1160 + assert!(matches!(
1161 + data["pixel_values_flat"].data.as_ref(),
1162 + Some(MmKwargValue::Tensor(tensor)) if tensor.shape == [3, 512, 512]
1163 + ));
1164 + assert_eq!(
1165 + data["num_tokens_per_image"].data,
1166 + Some(MmKwargValue::Int(256))
1167 + );
1168 + assert_eq!(
1169 + data["imgs_sizes"].data,
1170 + Some(MmKwargValue::List(vec![
1171 + MmKwargValue::Int(512),
1172 + MmKwargValue::Int(512),
1173 + ]))
1174 + );
1175 +
1176 + let error = info
1177 + .prepare_multimodal(vec![media], &mut prompt, ModelDtype::Float32, Some(100))
1178 + .await
1179 + .unwrap_err();
1180 + assert!(
1181 + matches!(error, Error::Multimodal(message) if message.contains("exceeding token_budget"))
1182 + );
1183 + }
1184 +
1185 #[test]
1186 fn qwen3_vl_without_video_token_id_disables_video_support_only() {
1187 let config = serde_json::json!({
rust/src/chat/src/multimodal/audio.rs
+1 -1
@@ -259,7 +259,7 @@ mod tests {
259 uuid: Some("audio-1".to_string()),
260 }];
261 async {
262 - info.prepare_multimodal(parts, &mut vec![AUDIO_PAD_ID], ModelDtype::Float32)
262 + info.prepare_multimodal(parts, &mut vec![AUDIO_PAD_ID], ModelDtype::Float32, None)
263 .instrument(mm_request_span("request-1"))
264 .await
265 .unwrap();
rust/src/chat/src/multimodal/image.rs
+9 -5
@@ -6,7 +6,7 @@
6
7 use std::sync::Arc;
8
9 -use llm_multimodal::{ImageFrame, Modality, PreprocessedEncoderInputs};
9 +use llm_multimodal::{ImageFrame, Modality, PreprocessedEncoderInputs, VisionPreprocessingContext};
10 use vllm_engine_core_client::protocol::dtype::ModelDtype;
11
12 use super::timing::MM_STAGE_TARGET;
@@ -30,11 +30,12 @@ impl MultimodalModelInfo {
30 frames: Vec<Arc<ImageFrame>>,
31 uuids: Vec<Option<String>>,
32 model_dtype: ModelDtype,
33 + preprocessing_context: VisionPreprocessingContext,
34 ) -> Result<PreparedMedia> {
35 let support = self.image.as_ref().ok_or_else(|| Error::UnsupportedModality {
36 modality: Modality::Image.to_string(),
37 })?;
37 - let preprocessed = self.preprocess_images(support, &frames).await?;
38 + let preprocessed = self.preprocess_images(support, &frames, preprocessing_context).await?;
39 let replacements = support.spec.prompt_replacements_for(&self.context, &preprocessed)?;
40 if replacements.len() != frames.len() {
41 bail_multimodal!(
@@ -65,13 +66,16 @@ impl MultimodalModelInfo {
66 &self,
67 support: &VisionModalitySupport,
68 image_frames: &[Arc<ImageFrame>],
69 + preprocessing_context: VisionPreprocessingContext,
70 ) -> Result<PreprocessedEncoderInputs> {
71 let processor = Arc::clone(&support.processor);
72 let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();
73
74 // TODO: is it still necessary given that we've already in a dedicated runtime?
73 - tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images)?))
74 - .await
75 - .map_err(|error| multimodal!("image preprocessing task failed: {error}"))?
75 + tokio::task::spawn_blocking(move || {
76 + Ok(processor.preprocess_with_context(&images, &preprocessing_context)?)
77 + })
78 + .await
79 + .map_err(|error| multimodal!("image preprocessing task failed: {error}"))?
80 }
81 }
rust/src/chat/src/multimodal/tensor.rs
+47 -2
@@ -18,7 +18,7 @@ use crate::error::{Error, Result, bail_multimodal, multimodal};
18 pub(super) enum KwargValue {
19 /// Tensor with row-major flat data and shape.
20 Tensor { wire: WireTensor },
21 - /// Non-tensor kwarg value that is shared or copied as-is.
21 + /// Non-tensor kwarg value preserved in protocol form.
22 Passthrough(ProtocolKwargValue),
23 }
24
@@ -150,7 +150,17 @@ impl KwargValue {
150 /// Batched fields use their first axis as media-item index and drop that
151 /// axis in the per-feature value, matching vLLM's batched-field semantics.
152 pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
153 - self.wire_value_range(index, index + 1, true)
153 + match self {
154 + Self::Passthrough(ProtocolKwargValue::List(values)) => values
155 + .get(index)
156 + .cloned()
157 + .ok_or_else(|| multimodal!("batched passthrough has no item at index {index}")),
158 + Self::Passthrough(value) if index == 0 => Ok(value.clone()),
159 + Self::Passthrough(_) => Err(multimodal!(
160 + "scalar batched passthrough has no item at index {index}"
161 + )),
162 + Self::Tensor { .. } => self.wire_value_range(index, index + 1, true),
163 + }
164 }
165
166 /// Convert one media item's flat tensor range directly to wire bytes.
@@ -336,6 +346,41 @@ mod tests {
346 );
347 }
348
349 + #[test]
350 + fn batched_passthrough_value_at_selects_media_item() {
351 + let image_sizes = KwargValue::from_model_specific(
352 + ModelSpecificValue::TupleVec(vec![(384, 736), (512, 512)]),
353 + ModelDtype::BFloat16,
354 + )
355 + .unwrap();
356 +
357 + assert_eq!(
358 + image_sizes.batched_wire_value_at(1).unwrap(),
359 + ProtocolKwargValue::List(vec![
360 + ProtocolKwargValue::Int(512),
361 + ProtocolKwargValue::Int(512),
362 + ])
363 + );
364 + assert!(image_sizes.batched_wire_value_at(2).is_err());
365 +
366 + let token_counts = KwargValue::from_model_specific(
367 + ModelSpecificValue::IntVec(vec![276, 400]),
368 + ModelDtype::BFloat16,
369 + )
370 + .unwrap();
371 + assert_eq!(
372 + token_counts.batched_wire_value_at(1).unwrap(),
373 + ProtocolKwargValue::Int(400)
374 + );
375 +
376 + let scalar = KwargValue::Passthrough(ProtocolKwargValue::Int(7));
377 + assert_eq!(
378 + scalar.batched_wire_value_at(0).unwrap(),
379 + ProtocolKwargValue::Int(7)
380 + );
381 + assert!(scalar.batched_wire_value_at(1).is_err());
382 + }
383 +
384 #[test]
385 fn flat_wire_value_range_keeps_first_axis() {
386 let value = KwargValue::Tensor {
rust/src/text/src/backend/hf/config.rs
+28 -4
@@ -84,10 +84,11 @@ impl HfSpecialTokens {
84
85 /// Minimal subset of `config.json` (the model's main HF config).
86 ///
87 -/// This intentionally supports only the two layouts we currently care about in
87 +/// This intentionally supports only the layouts we currently care about in
88 /// the Rust frontend:
89 /// - pure text models that keep text metadata at the top level
90 -/// - composite models that expose a single nested `text_config`
90 +/// - composite models that expose a single nested `text_config` or its
91 +/// `llm_config` alias (used by Nemotron)
92 ///
93 /// We do not support additional entry points such as `decoder`, `generator`, or
94 /// `text_encoder`.
@@ -102,6 +103,7 @@ pub struct ModelConfig {
103 n_routed_experts: Option<OneOrManyExpertCount>,
104 num_local_experts: Option<OneOrManyExpertCount>,
105 block_configs: Vec<BlockConfig>,
106 + #[serde(alias = "llm_config")]
107 text_config: Option<Box<ModelConfig>>,
108 }
109
@@ -188,7 +190,7 @@ impl ModelConfig {
190 /// Return the config that the Rust frontend treats as the text/LLM config.
191 ///
192 /// This is deliberately narrower than Python/transformers: we only support
191 - /// either the top-level config itself or a single nested `text_config`.
193 + /// either the top-level config itself or a single nested text config.
194 fn effective_text_config(&self) -> &Self {
195 self.text_config.as_deref().unwrap_or(self)
196 }
@@ -233,7 +235,8 @@ impl ModelConfig {
235 /// config.
236 ///
237 /// The only intentional simplification here is how we pick the text config:
236 - /// Rust only looks at the top level or `text_config`, not the broader
238 + /// Rust only looks at the top level or `text_config` (including its
239 + /// `llm_config` alias), not the broader
240 /// transformers composite-config surface.
241 fn num_experts_from_block_configs(&self) -> u32 {
242 self.effective_text_config()
@@ -405,6 +408,27 @@ mod tests {
408 assert_eq!(config.vocab_size().unwrap(), 151936);
409 }
410
411 + #[test]
412 + fn model_config_uses_llm_config_for_nemotron_composite_models() {
413 + let config: ModelConfig = serde_json::from_str(
414 + r#"{
415 + "model_type": "nemotron_h_omni",
416 + "llm_config": {
417 + "model_type": "nemotron_h",
418 + "vocab_size": 131072,
419 + "n_routed_experts": 256,
420 + "eos_token_id": 2
421 + }
422 + }"#,
423 + )
424 + .unwrap();
425 +
426 + assert_eq!(config.model_type(), Some("nemotron_h_omni"));
427 + assert_eq!(config.vocab_size().unwrap(), 131072);
428 + assert_eq!(config.eos_token_ids(), &[2]);
429 + assert_eq!(config.num_experts(), 256);
430 + }
431 +
432 #[test]
433 fn model_config_reads_top_level_eos_token_ids() {
434 let single: ModelConfig = serde_json::from_str(r#"{"eos_token_id":151645}"#).unwrap();