[Rust Frontend] Avoid extra copies for multimodal tensors (#47581)

Signed-off-by: reidliu41 <reid201711@gmail.com>

rust/src/engine-core-client/src/protocol/tensor.rs

@@ -11,6 +11,21 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 /// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L41-L43>
 const CUSTOM_TYPE_RAW_VIEW: i8 = 3;
 
+#[derive(Serialize)]
+#[serde(rename = "_ExtStruct")]
+struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));
+
+struct ByteSlice<'a>(&'a [u8]);
+
+impl Serialize for ByteSlice<'_> {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_bytes(self.0)
+    }
+}
+
 #[easy_ext::ext(ShapeExt)]
 impl [usize] {
     /// Returned the total number of elements implied by this shape, or `None`
@@ -184,7 +199,7 @@ impl Serialize for WireArrayData {
         match self {
             Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
             Self::RawView(bytes) => {
-                Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()).serialize(serializer)
+                MsgpackExtRef((CUSTOM_TYPE_RAW_VIEW, ByteSlice(bytes))).serialize(serializer)
             }
         }
     }
@@ -194,6 +209,21 @@ impl Serialize for WireArrayData {
 mod tests {
     use super::*;
 
+    #[test]
+    fn raw_view_serializes_as_msgpack_ext() {
+        let bytes = vec![1, 2, 3, 4];
+        let encoded =
+            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
+        let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
+            .expect("encode expected");
+
+        assert_eq!(encoded, expected);
+        assert_eq!(
+            rmpv::decode::read_value(&mut std::io::Cursor::new(encoded)).expect("decode"),
+            Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes)
+        );
+    }
+
     #[test]
     fn constructors_build_raw_view_tensors() {
         let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();

rust/src/text/src/lower.rs

@@ -18,7 +18,7 @@ use crate::request::{SamplingParams, TextRequest};
 /// One text request after it has been lowered into the raw generate boundary.
 #[derive(Debug)]
 pub struct PreparedTextRequest {
-    /// The original high-level request, preserved for response-side metadata
+    /// The high-level request fields still needed for response-side metadata
     /// and decoding options.
     pub text_request: TextRequest,
     /// The southbound request ready to be sent to `vllm-llm`.
@@ -28,7 +28,7 @@ pub struct PreparedTextRequest {
 /// Convert a high-level [`TextRequest`] into one lower-level
 /// [`GenerateRequest`] ready for the `llm` crate.
 pub fn lower_text_request(
-    request: TextRequest,
+    mut request: TextRequest,
     prompt_token_ids: Vec<u32>,
     sampling_hints: SamplingHints,
     sampling_limits: SamplingLimits,
@@ -40,7 +40,10 @@ pub fn lower_text_request(
     let generate_request = GenerateRequest {
         request_id: request.request_id.clone(),
         prompt_token_ids,
-        mm_features: request.mm_features.clone(),
+        // Align with Python's response path: decoded output state does not retain
+        // `mm_features`; move them to the engine request to avoid cloning large
+        // multimodal tensor payloads.
+        mm_features: request.mm_features.take(),
         sampling_params: lower_sampling_params(
             request.sampling_params.clone(),
             sampling_hints,
@@ -307,6 +310,7 @@ mod tests {
     use std::collections::{BTreeSet, HashMap};
 
     use serial_test::file_serial;
+    use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, PlaceholderRange};
     use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::*;
@@ -574,6 +578,35 @@ mod tests {
         .assert_debug_eq(&params);
     }
 
+    #[test]
+    fn lower_text_request_moves_multimodal_features_to_generate_request() {
+        let features = vec![MmFeatureSpec {
+            data: None,
+            modality: "image".to_string(),
+            identifier: "image-1".to_string(),
+            mm_position: PlaceholderRange {
+                offset: 2,
+                length: 4,
+                is_embed: None,
+            },
+            mm_hash: Some("hash-1".to_string()),
+        }];
+        let mut request = sample_request();
+        request.mm_features = Some(features.clone());
+
+        let prepared = lower_text_request(
+            request,
+            vec![1, 2, 3],
+            sample_sampling_hints(),
+            sample_sampling_limits(),
+            &stub_tokenizer(),
+        )
+        .unwrap();
+
+        assert_eq!(prepared.generate_request.mm_features, Some(features));
+        assert_eq!(prepared.text_request.mm_features, None);
+    }
+
     #[test]
     fn lower_text_request_uses_union_vocab_for_prompt_token_ids() {
         lower_text_request(