[Rust Frontend] Use zero-copy slicing for multimodal tensors (#48781)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
Sage committed
Jul 20, 2026 at 11:28 UTC
f1f1259692004c909f8b7f8f371d42c7a7871fa3
4 files changed
+91
-99
rust/src/chat/src/multimodal/item.rs
+4
-4
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
38
let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
39
let (value, field) = match spec.field_layout_for(key) {
40
Some(FieldLayout::Batched) => (
41
- tensor.batched_value_at(index)?,
41
+ tensor.batched_wire_value_at(index)?,
42
MmField::Batched(MmBatchedField { keep_on_cpu }),
43
),
44
Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
47
})?;
48
let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
49
(
50
- tensor.flat_value_range(start, end)?,
50
+ tensor.flat_wire_value_range(start, end)?,
51
MmField::Flat(MmFlatField {
52
slices: vec![MmSlice::Slice(SliceSpec {
53
start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
60
)
61
}
62
None => (
63
- tensor.clone(),
63
+ tensor.try_into()?,
64
MmField::Shared(MmSharedField {
65
batch_size: len,
66
keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
71
data.insert(
72
key.clone(),
73
MmFieldElem {
74
- data: Some(value.try_into()?),
74
+ data: Some(value),
75
field,
76
},
77
);
rust/src/chat/src/multimodal/tensor.rs
+68
-81
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
12
use crate::error::{Error, Result, bail_multimodal, multimodal};
13
14
/// Representation for multimodal kwarg values for transformation.
15
-#[derive(Debug, Clone)]
15
+#[derive(Debug)]
16
pub(super) enum KwargValue {
17
/// Float tensor with row-major flat data and shape.
18
F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,28 +107,19 @@ impl KwargValue {
107
}
108
}
109
110
-impl TryFrom<KwargValue> for ProtocolKwargValue {
110
+impl TryFrom<&KwargValue> for ProtocolKwargValue {
111
type Error = Error;
112
113
- fn try_from(value: KwargValue) -> Result<Self> {
114
- match value {
115
- KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
116
- WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
117
- )),
118
- KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
119
- WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
120
- )),
121
- KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
122
- WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
123
- )),
124
- KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
125
- WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
126
- )),
127
- KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
128
- WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
129
- )),
130
- KwargValue::Passthrough(value) => Ok(value),
131
- }
113
+ fn try_from(value: &KwargValue) -> Result<Self> {
114
+ let tensor = match value {
115
+ KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
116
+ KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
117
+ KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
118
+ KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
119
+ KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
120
+ KwargValue::Passthrough(value) => return Ok(value.clone()),
121
+ };
122
+ tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
123
}
124
}
125
@@ -145,63 +136,55 @@ impl KwargValue {
136
}
137
}
138
148
- /// Extract one media item from a batched tensor field.
139
+ /// Convert one media item from a batched tensor field to wire bytes.
140
///
141
/// Batched fields use their first axis as media-item index and drop that
142
/// axis in the per-feature value, matching vLLM's batched-field semantics.
152
- pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
153
- match self {
154
- Self::F32Tensor { data, shape } => {
155
- let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
156
- Ok(Self::F32Tensor { data, shape })
157
- }
158
- Self::F16Tensor { data, shape } => {
159
- let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
160
- Ok(Self::F16Tensor { data, shape })
161
- }
162
- Self::Bf16Tensor { data, shape } => {
163
- let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
164
- Ok(Self::Bf16Tensor { data, shape })
165
- }
166
- Self::I64Tensor { data, shape } => {
167
- let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
168
- Ok(Self::I64Tensor { data, shape })
169
- }
170
- Self::U32Tensor { data, shape } => {
171
- let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
172
- Ok(Self::U32Tensor { data, shape })
173
- }
174
- Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
175
- }
143
+ pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
144
+ self.wire_value_range(index, index + 1, true)
145
}
146
178
- /// Extract one media item's variable-length range from a flat tensor field.
147
+ /// Convert one media item's flat tensor range directly to wire bytes.
148
///
149
/// Flat fields keep the first axis as the sliced length for this item.
181
- pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
182
- match self {
150
+ pub(super) fn flat_wire_value_range(
151
+ &self,
152
+ start: usize,
153
+ end: usize,
154
+ ) -> Result<ProtocolKwargValue> {
155
+ self.wire_value_range(start, end, false)
156
+ }
157
+
158
+ fn wire_value_range(
159
+ &self,
160
+ start: usize,
161
+ end: usize,
162
+ drop_axis: bool,
163
+ ) -> Result<ProtocolKwargValue> {
164
+ let tensor = match self {
165
Self::F32Tensor { data, shape } => {
184
- let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
185
- Ok(Self::F32Tensor { data, shape })
166
+ let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
167
+ WireTensor::from_f32(shape, data)
168
}
169
Self::F16Tensor { data, shape } => {
188
- let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
189
- Ok(Self::F16Tensor { data, shape })
170
+ let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
171
+ WireTensor::from_f16(shape, data)
172
}
173
Self::Bf16Tensor { data, shape } => {
192
- let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
193
- Ok(Self::Bf16Tensor { data, shape })
174
+ let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
175
+ WireTensor::from_bf16(shape, data)
176
}
177
Self::I64Tensor { data, shape } => {
196
- let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
197
- Ok(Self::I64Tensor { data, shape })
178
+ let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
179
+ WireTensor::from_i64(shape, data)
180
}
181
Self::U32Tensor { data, shape } => {
200
- let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
201
- Ok(Self::U32Tensor { data, shape })
182
+ let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
183
+ WireTensor::from_u32(shape, data)
184
}
203
- Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
204
- }
185
+ Self::Passthrough(value) => return Ok(value.clone()),
186
+ };
187
+ tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
188
}
189
}
190
@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
223
}
224
225
/// Slice a flat row-major tensor along its first axis.
243
-fn slice_first_axis_range<T: Clone>(
226
+fn slice_first_axis_range<'a, T>(
227
shape: &[usize],
245
- data: &[T],
228
+ data: &'a [T],
229
start: usize,
230
end: usize,
231
drop_axis: bool,
249
-) -> Result<(Vec<usize>, Vec<T>)> {
232
+) -> Result<(Vec<usize>, &'a [T])> {
233
let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
234
if start > end || end > first_dim {
235
bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
253
shape[0] = end - start;
254
shape
255
};
273
- Ok((out_shape, data[data_start..data_end].to_vec()))
256
+ Ok((out_shape, &data[data_start..data_end]))
257
}
258
259
#[cfg(test)]
@@ -278,35 +261,39 @@ mod tests {
261
use super::*;
262
263
#[test]
281
- fn batched_value_at_drops_first_axis() {
264
+ fn batched_wire_value_at_drops_first_axis() {
265
let value = KwargValue::F32Tensor {
266
data: vec![1.0, 2.0, 3.0, 4.0],
267
shape: vec![2, 2],
268
};
269
287
- let value = value.batched_value_at(1).unwrap();
270
+ let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
271
+ panic!("expected tensor");
272
+ };
273
289
- assert!(matches!(
290
- value,
291
- KwargValue::F32Tensor { data, shape }
292
- if shape == vec![2] && data == vec![3.0, 4.0]
293
- ));
274
+ assert_eq!(tensor.shape, vec![2]);
275
+ assert_eq!(
276
+ tensor.data.into_raw_view().unwrap(),
277
+ [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
278
+ );
279
}
280
281
#[test]
297
- fn flat_value_range_keeps_first_axis() {
282
+ fn flat_wire_value_range_keeps_first_axis() {
283
let value = KwargValue::U32Tensor {
284
data: (0..10).collect(),
285
shape: vec![5, 2],
286
};
287
303
- let value = value.flat_value_range(1, 3).unwrap();
288
+ let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
289
+ panic!("expected tensor");
290
+ };
291
305
- assert!(matches!(
306
- value,
307
- KwargValue::U32Tensor { data, shape }
308
- if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
309
- ));
292
+ assert_eq!(tensor.shape, vec![2, 2]);
293
+ assert_eq!(
294
+ tensor.data.into_raw_view().unwrap(),
295
+ [2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
296
+ );
297
}
298
299
#[test]
@@ -336,7 +323,7 @@ mod tests {
323
let value =
324
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
325
339
- let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
326
+ let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
327
else {
328
panic!("expected tensor");
329
};
@@ -351,7 +338,7 @@ mod tests {
338
let value =
339
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
340
354
- let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
341
+ let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
342
else {
343
panic!("expected tensor");
344
};
rust/src/chat/src/multimodal/video.rs
+4
-4
@@ -130,7 +130,7 @@ fn build_video_item(
130
let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
131
let (value, field) = match support.spec.field_layout_for(&key) {
132
Some(FieldLayout::Batched) => (
133
- tensor.batched_value_at(0)?,
133
+ tensor.batched_wire_value_at(0)?,
134
MmField::Batched(MmBatchedField { keep_on_cpu }),
135
),
136
Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
138
.first_dim()
139
.ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
140
(
141
- tensor,
141
+ (&tensor).try_into()?,
142
MmField::Flat(MmFlatField {
143
slices: vec![MmSlice::Slice(SliceSpec {
144
start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
151
)
152
}
153
None => (
154
- tensor,
154
+ (&tensor).try_into()?,
155
MmField::Shared(MmSharedField {
156
batch_size: 1,
157
keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
162
data.insert(
163
key,
164
MmFieldElem {
165
- data: Some(value.try_into()?),
165
+ data: Some(value),
166
field,
167
},
168
);
rust/src/engine-core-client/src/protocol/tensor.rs
+15
-10
@@ -55,52 +55,57 @@ pub struct WireNdArray {
55
56
impl WireNdArray {
57
/// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
58
- pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
58
+ pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
59
+ let data = data.as_ref();
60
validate_element_count(&shape, data.len())?;
61
Ok(Self {
62
dtype: "float32".to_string(),
63
shape,
63
- data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
64
+ data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
65
})
66
}
67
68
/// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
68
- pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
69
+ pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
70
+ let data = data.as_ref();
71
validate_element_count(&shape, data.len())?;
72
Ok(Self {
73
dtype: "float16".to_string(),
74
shape,
73
- data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
75
+ data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
76
})
77
}
78
79
/// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
78
- pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
80
+ pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
81
+ let data = data.as_ref();
82
validate_element_count(&shape, data.len())?;
83
Ok(Self {
84
dtype: "bfloat16".to_string(),
85
shape,
83
- data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
86
+ data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
87
})
88
}
89
90
/// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
88
- pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
91
+ pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
92
+ let data = data.as_ref();
93
validate_element_count(&shape, data.len())?;
94
Ok(Self {
95
dtype: "int64".to_string(),
96
shape,
93
- data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
97
+ data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
98
})
99
}
100
101
/// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
98
- pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
102
+ pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
103
+ let data = data.as_ref();
104
validate_element_count(&shape, data.len())?;
105
Ok(Self {
106
dtype: "uint32".to_string(),
107
shape,
103
- data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
108
+ data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
109
})
110
}
111