A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 16:30 UTC
[Rust Frontend] Optimize multimodal prompt expansion (#43670)
Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com>
Chao-Ju Chen committed
May 29, 2026 at 00:46 UTC
099024762cc7622890f4887d27e5f547a8e73225
1 file changed
+119
-47
rust/src/chat/src/multimodal.rs
+119
-47
@@ -392,53 +392,13 @@ impl MultimodalModelInfo {
392
prompt_token_ids: &mut Vec<u32>,
393
replacements: Vec<PromptReplacement>,
394
) -> Result<Vec<PlaceholderRange>> {
395
- let mut cursor = 0;
396
- let mut ranges = Vec::with_capacity(replacements.len());
397
- for replacement in replacements {
398
- if replacement.modality != Modality::Image {
399
- bail_multimodal!(
400
- "unsupported prompt replacement modality `{}`",
401
- replacement.modality
402
- );
403
- }
404
- let offset = find_next_token(
405
- prompt_token_ids,
406
- self.spec.placeholder_marker_token_id,
407
- cursor,
408
- )
409
- .ok_or_else(|| {
410
- multimodal!(
411
- "placeholder token `{}` was not found in tokenized prompt",
412
- self.spec.placeholder_token
413
- )
414
- })?;
415
-
416
- if replacement.tokens.is_empty() {
417
- bail_multimodal!(
418
- "placeholder token `{}` expanded to no tokens",
419
- self.spec.placeholder_token
420
- );
421
- }
422
- let replacement_len = replacement.tokens.len();
423
- let replacement_tokens =
424
- replacement.tokens.iter().map(|&token| token as u32).collect::<Vec<_>>();
425
- let is_embed = {
426
- let mask = replacement_tokens
427
- .iter()
428
- .map(|&token| token == self.spec.placeholder_embed_token_id)
429
- .collect::<Vec<_>>();
430
- WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
431
- };
432
-
433
- prompt_token_ids.splice(offset..offset + 1, replacement_tokens);
434
- ranges.push(PlaceholderRange {
435
- offset,
436
- length: replacement_len,
437
- is_embed: Some(is_embed),
438
- });
439
- cursor = offset + replacement_len;
440
- }
441
- Ok(ranges)
395
+ expand_prompt_token_ids(
396
+ prompt_token_ids,
397
+ replacements,
398
+ self.spec.placeholder_marker_token_id,
399
+ self.spec.placeholder_embed_token_id,
400
+ &self.spec.placeholder_token,
401
+ )
402
}
403
404
/// Convert preprocessed image tensors into engine-core multimodal features.
@@ -516,6 +476,71 @@ impl MultimodalModelInfo {
476
}
477
}
478
479
+fn expand_prompt_token_ids(
480
+ prompt_token_ids: &mut Vec<u32>,
481
+ replacements: Vec<PromptReplacement>,
482
+ placeholder_marker_token_id: u32,
483
+ placeholder_embed_token_id: u32,
484
+ placeholder_token: &str,
485
+) -> Result<Vec<PlaceholderRange>> {
486
+ if replacements.is_empty() {
487
+ return Ok(Vec::new());
488
+ }
489
+
490
+ let replacement_growth = replacements.iter().fold(0usize, |total, replacement| {
491
+ total.saturating_add(replacement.tokens.len().saturating_sub(1))
492
+ });
493
+ let mut expanded =
494
+ Vec::with_capacity(prompt_token_ids.len().saturating_add(replacement_growth));
495
+ let mut ranges = Vec::with_capacity(replacements.len());
496
+ let mut cursor = 0usize;
497
+
498
+ for replacement in replacements {
499
+ if replacement.modality != Modality::Image {
500
+ bail_multimodal!(
501
+ "unsupported prompt replacement modality `{}`",
502
+ replacement.modality
503
+ );
504
+ }
505
+
506
+ let offset = find_next_token(prompt_token_ids, placeholder_marker_token_id, cursor)
507
+ .ok_or_else(|| {
508
+ multimodal!(
509
+ "placeholder token `{placeholder_token}` was not found in tokenized prompt"
510
+ )
511
+ })?;
512
+
513
+ if replacement.tokens.is_empty() {
514
+ bail_multimodal!("placeholder token `{placeholder_token}` expanded to no tokens");
515
+ }
516
+
517
+ let replacement_len = replacement.tokens.len();
518
+ let is_embed = {
519
+ let mask = replacement
520
+ .tokens
521
+ .iter()
522
+ .map(|&token| token as u32 == placeholder_embed_token_id)
523
+ .collect::<Vec<_>>();
524
+ WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
525
+ };
526
+
527
+ expanded.extend_from_slice(&prompt_token_ids[cursor..offset]);
528
+ let expanded_offset = expanded.len();
529
+ expanded.extend(replacement.tokens.into_iter().map(|token| token as u32));
530
+ ranges.push(PlaceholderRange {
531
+ offset: expanded_offset,
532
+ length: replacement_len,
533
+ is_embed: Some(is_embed),
534
+ });
535
+ cursor = offset + 1;
536
+ }
537
+
538
+ expanded.extend_from_slice(&prompt_token_ids[cursor..]);
539
+ *prompt_token_ids = expanded;
540
+
541
+ Ok(ranges)
542
+}
543
+
544
/// Find `needle` in `haystack`, starting at `start`.
545
///
546
/// This is intentionally order-preserving rather than a global replace: each
@@ -736,6 +761,53 @@ mod tests {
761
assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
762
}
763
764
+ #[test]
765
+ fn expand_prompt_tokens_ignores_empty_replacements() {
766
+ let info = llama4_info();
767
+ let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
768
+ let original_prompt_token_ids = prompt_token_ids.clone();
769
+
770
+ let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, Vec::new()).unwrap();
771
+
772
+ assert!(ranges.is_empty());
773
+ assert_eq!(prompt_token_ids, original_prompt_token_ids);
774
+ }
775
+
776
+ #[test]
777
+ fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() {
778
+ let info = llama4_info();
779
+ let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
780
+ let original_prompt_token_ids = prompt_token_ids.clone();
781
+ let replacements = vec![
782
+ llama4_single_tile_replacement(),
783
+ llama4_single_tile_replacement(),
784
+ ];
785
+
786
+ let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
787
+
788
+ assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
789
+ assert_eq!(prompt_token_ids, original_prompt_token_ids);
790
+ }
791
+
792
+ #[test]
793
+ fn expand_prompt_tokens_errors_when_replacement_is_empty() {
794
+ let info = llama4_info();
795
+ let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
796
+ let original_prompt_token_ids = prompt_token_ids.clone();
797
+ let replacements = vec![PromptReplacement::sequence(
798
+ Modality::Image,
799
+ "<|image|>",
800
+ Vec::new(),
801
+ )];
802
+
803
+ let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
804
+
805
+ assert!(
806
+ matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens"))
807
+ );
808
+ assert_eq!(prompt_token_ids, original_prompt_token_ids);
809
+ }
810
+
811
#[test]
812
fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() {
813
let info = llama4_info();