[Rust Frontend] Avoid extra copies for multimodal tensors (#47581)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jul 7, 2026 at 11:09 UTC
39a1d32b598280c4bc5b6853a4a55622e67deb60
2 files changed
+67
-4
rust/src/engine-core-client/src/protocol/tensor.rs
+31
-1
@@ -11,6 +11,21 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple};
11
/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L41-L43>
12
const CUSTOM_TYPE_RAW_VIEW: i8 = 3;
13
14
+#[derive(Serialize)]
15
+#[serde(rename = "_ExtStruct")]
16
+struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));
17
+
18
+struct ByteSlice<'a>(&'a [u8]);
19
+
20
+impl Serialize for ByteSlice<'_> {
21
+ fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
22
+ where
23
+ S: Serializer,
24
+ {
25
+ serializer.serialize_bytes(self.0)
26
+ }
27
+}
28
+
29
#[easy_ext::ext(ShapeExt)]
30
impl [usize] {
31
/// Returned the total number of elements implied by this shape, or `None`
@@ -184,7 +199,7 @@ impl Serialize for WireArrayData {
199
match self {
200
Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
201
Self::RawView(bytes) => {
187
- Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()).serialize(serializer)
202
+ MsgpackExtRef((CUSTOM_TYPE_RAW_VIEW, ByteSlice(bytes))).serialize(serializer)
203
}
204
}
205
}
@@ -194,6 +209,21 @@ impl Serialize for WireArrayData {
209
mod tests {
210
use super::*;
211
212
+ #[test]
213
+ fn raw_view_serializes_as_msgpack_ext() {
214
+ let bytes = vec![1, 2, 3, 4];
215
+ let encoded =
216
+ rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
217
+ let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
218
+ .expect("encode expected");
219
+
220
+ assert_eq!(encoded, expected);
221
+ assert_eq!(
222
+ rmpv::decode::read_value(&mut std::io::Cursor::new(encoded)).expect("decode"),
223
+ Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes)
224
+ );
225
+ }
226
+
227
#[test]
228
fn constructors_build_raw_view_tensors() {
229
let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
rust/src/text/src/lower.rs
+36
-3
@@ -18,7 +18,7 @@ use crate::request::{SamplingParams, TextRequest};
18
/// One text request after it has been lowered into the raw generate boundary.
19
#[derive(Debug)]
20
pub struct PreparedTextRequest {
21
- /// The original high-level request, preserved for response-side metadata
21
+ /// The high-level request fields still needed for response-side metadata
22
/// and decoding options.
23
pub text_request: TextRequest,
24
/// The southbound request ready to be sent to `vllm-llm`.
@@ -28,7 +28,7 @@ pub struct PreparedTextRequest {
28
/// Convert a high-level [`TextRequest`] into one lower-level
29
/// [`GenerateRequest`] ready for the `llm` crate.
30
pub fn lower_text_request(
31
- request: TextRequest,
31
+ mut request: TextRequest,
32
prompt_token_ids: Vec<u32>,
33
sampling_hints: SamplingHints,
34
sampling_limits: SamplingLimits,
@@ -40,7 +40,10 @@ pub fn lower_text_request(
40
let generate_request = GenerateRequest {
41
request_id: request.request_id.clone(),
42
prompt_token_ids,
43
- mm_features: request.mm_features.clone(),
43
+ // Align with Python's response path: decoded output state does not retain
44
+ // `mm_features`; move them to the engine request to avoid cloning large
45
+ // multimodal tensor payloads.
46
+ mm_features: request.mm_features.take(),
47
sampling_params: lower_sampling_params(
48
request.sampling_params.clone(),
49
sampling_hints,
@@ -307,6 +310,7 @@ mod tests {
310
use std::collections::{BTreeSet, HashMap};
311
312
use serial_test::file_serial;
313
+ use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, PlaceholderRange};
314
use vllm_tokenizer::test_utils::TestTokenizer;
315
316
use super::*;
@@ -574,6 +578,35 @@ mod tests {
578
.assert_debug_eq(¶ms);
579
}
580
581
+ #[test]
582
+ fn lower_text_request_moves_multimodal_features_to_generate_request() {
583
+ let features = vec![MmFeatureSpec {
584
+ data: None,
585
+ modality: "image".to_string(),
586
+ identifier: "image-1".to_string(),
587
+ mm_position: PlaceholderRange {
588
+ offset: 2,
589
+ length: 4,
590
+ is_embed: None,
591
+ },
592
+ mm_hash: Some("hash-1".to_string()),
593
+ }];
594
+ let mut request = sample_request();
595
+ request.mm_features = Some(features.clone());
596
+
597
+ let prepared = lower_text_request(
598
+ request,
599
+ vec![1, 2, 3],
600
+ sample_sampling_hints(),
601
+ sample_sampling_limits(),
602
+ &stub_tokenizer(),
603
+ )
604
+ .unwrap();
605
+
606
+ assert_eq!(prepared.generate_request.mm_features, Some(features));
607
+ assert_eq!(prepared.text_request.mm_features, None);
608
+ }
609
+
610
#[test]
611
fn lower_text_request_uses_union_vocab_for_prompt_token_ids() {
612
lower_text_request(