@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 242c591d5a

[Rust Frontend] Send multimodal tensors in auxiliary frames (#49341)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Reid committed Jul 29, 2026 at 20:21 UTC 242c591d5a38be6c80ff5197e83e4fa34c5f74b5
11 files changed +530 -159
rust/src/chat/src/multimodal/expand.rs
+3 -1
@@ -222,7 +222,9 @@ mod tests {
222 assert_eq!(tensor.shape, vec![expected.len()]);
223 assert_eq!(
224 tensor.data,
225 - WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect())
225 + WireArrayData::RawView(
226 + expected.iter().map(|value| u8::from(*value)).collect::<Vec<_>>().into(),
227 + )
228 );
229 }
230
rust/src/chat/src/multimodal/tensor.rs
+157 -95
@@ -2,28 +2,58 @@
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::collections::HashMap;
5 +use std::mem::size_of;
6
7 use half::{bf16, f16};
8 use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs};
9 use vllm_engine_core_client::protocol::dtype::ModelDtype;
10 use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
10 -use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
11 +use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor};
12
13 use crate::error::{Error, Result, bail_multimodal, multimodal};
14
15 +/// Element type retained alongside an encoded tensor during multimodal lowering.
16 +#[derive(Debug, Clone, Copy)]
17 +pub(super) enum TensorKind {
18 + /// 32-bit floating point.
19 + F32,
20 + /// IEEE 16-bit floating point.
21 + F16,
22 + /// Brain floating point.
23 + Bf16,
24 + /// Signed 64-bit integer.
25 + I64,
26 + /// Unsigned 32-bit integer.
27 + U32,
28 +}
29 +
30 +impl TensorKind {
31 + const fn element_size(self) -> usize {
32 + match self {
33 + Self::F32 => size_of::<f32>(),
34 + Self::F16 => size_of::<f16>(),
35 + Self::Bf16 => size_of::<bf16>(),
36 + Self::I64 => size_of::<i64>(),
37 + Self::U32 => size_of::<u32>(),
38 + }
39 + }
40 +
41 + const fn wire_dtype(self) -> &'static str {
42 + match self {
43 + Self::F32 => "float32",
44 + Self::F16 => "float16",
45 + Self::Bf16 => "bfloat16",
46 + Self::I64 => "int64",
47 + Self::U32 => "uint32",
48 + }
49 + }
50 +}
51 +
52 /// Representation for multimodal kwarg values for transformation.
53 #[derive(Debug)]
54 pub(super) enum KwargValue {
17 - /// Float tensor with row-major flat data and shape.
18 - F32Tensor { data: Vec<f32>, shape: Vec<usize> },
19 - /// Float16 tensor with row-major flat data and shape.
20 - F16Tensor { data: Vec<f16>, shape: Vec<usize> },
21 - /// BFloat16 tensor with row-major flat data and shape.
22 - Bf16Tensor { data: Vec<bf16>, shape: Vec<usize> },
23 - /// Signed integer tensor with row-major flat data and shape.
24 - I64Tensor { data: Vec<i64>, shape: Vec<usize> },
25 - /// Unsigned integer tensor with row-major flat data and shape.
26 - U32Tensor { data: Vec<u32>, shape: Vec<usize> },
55 + /// Tensor with row-major flat data and shape.
56 + Tensor { kind: TensorKind, wire: WireTensor },
57 /// Non-tensor kwarg value that is shared or copied as-is.
58 Passthrough(ProtocolKwargValue),
59 }
@@ -67,8 +97,14 @@ impl KwargValue {
97 ModelSpecificValue::Tensor { data, shape } => {
98 Self::from_f32_tensor(data, shape, float_dtype)?
99 }
70 - ModelSpecificValue::IntTensor { data, shape } => Self::I64Tensor { data, shape },
71 - ModelSpecificValue::UintTensor { data, shape } => Self::U32Tensor { data, shape },
100 + ModelSpecificValue::IntTensor { data, shape } => {
101 + let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?;
102 + Self::tensor(TensorKind::I64, wire)
103 + }
104 + ModelSpecificValue::UintTensor { data, shape } => {
105 + let wire = WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?;
106 + Self::tensor(TensorKind::U32, wire)
107 + }
108 ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)),
109 ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)),
110 ModelSpecificValue::IntVec(values) => {
@@ -93,17 +129,23 @@ impl KwargValue {
129 /// Convert a float tensor to the target float dtype if needed, keeping the
130 /// same shape.
131 fn from_f32_tensor(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
96 - match float_dtype {
97 - ModelDtype::Float16 => Ok(Self::F16Tensor {
98 - data: data.into_iter().map(f16::from_f32).collect(),
99 - shape,
100 - }),
101 - ModelDtype::BFloat16 => Ok(Self::Bf16Tensor {
102 - data: data.into_iter().map(bf16::from_f32).collect(),
103 - shape,
104 - }),
105 - ModelDtype::Float32 => Ok(Self::F32Tensor { data, shape }),
106 - }
132 + let (kind, wire) = match float_dtype {
133 + ModelDtype::Float16 => (
134 + TensorKind::F16,
135 + WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect()),
136 + ),
137 + ModelDtype::BFloat16 => (
138 + TensorKind::Bf16,
139 + WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect()),
140 + ),
141 + ModelDtype::Float32 => (TensorKind::F32, WireTensor::from_f32(shape, data)),
142 + };
143 + wire.map(|wire| Self::tensor(kind, wire)).map_err(Error::Multimodal)
144 + }
145 +
146 + fn tensor(kind: TensorKind, wire: WireTensor) -> Self {
147 + debug_assert_eq!(wire.dtype, kind.wire_dtype());
148 + Self::Tensor { kind, wire }
149 }
150 }
151
@@ -111,15 +153,11 @@ impl TryFrom<&KwargValue> for ProtocolKwargValue {
153 type Error = Error;
154
155 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),
156 + let wire = match value {
157 + KwargValue::Tensor { wire, .. } => wire.clone(),
158 KwargValue::Passthrough(value) => return Ok(value.clone()),
159 };
122 - tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
160 + Ok(ProtocolKwargValue::Tensor(wire))
161 }
162 }
163
@@ -127,11 +165,7 @@ impl KwargValue {
165 /// First-axis length for tensor values; `None` for passthrough kwargs.
166 pub(super) fn first_dim(&self) -> Option<usize> {
167 match self {
130 - Self::F32Tensor { shape, .. }
131 - | Self::F16Tensor { shape, .. }
132 - | Self::Bf16Tensor { shape, .. }
133 - | Self::I64Tensor { shape, .. }
134 - | Self::U32Tensor { shape, .. } => shape.first().copied(),
168 + Self::Tensor { wire, .. } => wire.shape.first().copied(),
169 Self::Passthrough(_) => None,
170 }
171 }
@@ -161,30 +195,13 @@ impl KwargValue {
195 end: usize,
196 drop_axis: bool,
197 ) -> Result<ProtocolKwargValue> {
164 - let tensor = match self {
165 - 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 } => {
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 } => {
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 } => {
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 } => {
182 - let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
183 - WireTensor::from_u32(shape, data)
198 + let wire = match self {
199 + Self::Tensor { kind, wire } => {
200 + slice_first_axis_range(wire, kind.element_size(), start, end, drop_axis)
201 }
202 Self::Passthrough(value) => return Ok(value.clone()),
203 };
187 - tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
204 + wire.map(ProtocolKwargValue::Tensor)
205 }
206 }
207
@@ -208,44 +225,53 @@ pub(super) fn flat_range_for_index(
225 /// Read a tensor value as per-image sizes for flat slicing.
226 fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
227 match tensor {
211 - KwargValue::I64Tensor { data, .. } => data
212 - .iter()
228 + KwargValue::Tensor {
229 + kind: TensorKind::I64,
230 + wire,
231 + } => raw_tensor_bytes(wire, size_of::<i64>())?
232 + .chunks_exact(size_of::<i64>())
233 + .map(|bytes| i64::from_ne_bytes(bytes.try_into().expect("exact int64 chunk")))
234 .map(|value| {
214 - usize::try_from(*value)
235 + usize::try_from(value)
236 .map_err(|_| multimodal!("negative flat tensor size `{value}`"))
237 })
238 .collect(),
218 - KwargValue::U32Tensor { data, .. } => {
219 - Ok(data.iter().map(|value| *value as usize).collect())
220 - }
239 + KwargValue::Tensor {
240 + kind: TensorKind::U32,
241 + wire,
242 + } => Ok(raw_tensor_bytes(wire, size_of::<u32>())?
243 + .chunks_exact(size_of::<u32>())
244 + .map(|bytes| u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize)
245 + .collect()),
246 _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")),
247 }
248 }
249
250 /// Slice a flat row-major tensor along its first axis.
226 -fn slice_first_axis_range<'a, T>(
227 - shape: &[usize],
228 - data: &'a [T],
251 +fn slice_first_axis_range(
252 + tensor: &WireTensor,
253 + element_size: usize,
254 start: usize,
255 end: usize,
256 drop_axis: bool,
232 -) -> Result<(Vec<usize>, &'a [T])> {
257 +) -> Result<WireTensor> {
258 + let shape = tensor.shape.as_slice();
259 + raw_tensor_bytes(tensor, element_size)?;
260 let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
261 if start > end || end > first_dim {
262 bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
263 }
237 - let expected_len = shape
238 - .checked_numel()
239 - .ok_or_else(|| multimodal!("tensor shape {shape:?} has too many elements"))?;
240 - if expected_len != data.len() {
241 - bail_multimodal!(
242 - "tensor shape {shape:?} expects {expected_len} elements, got {}",
243 - data.len()
244 - );
245 - }
246 - let stride = shape[1..].iter().product::<usize>();
247 - let data_start = start * stride;
248 - let data_end = end * stride;
264 + let stride = shape[1..]
265 + .iter()
266 + .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))
267 + .and_then(|stride| stride.checked_mul(element_size))
268 + .ok_or_else(|| multimodal!("tensor shape {shape:?} byte stride overflowed usize"))?;
269 + let data_start = start
270 + .checked_mul(stride)
271 + .ok_or_else(|| multimodal!("tensor slice start byte offset overflowed usize"))?;
272 + let data_end = end
273 + .checked_mul(stride)
274 + .ok_or_else(|| multimodal!("tensor slice end byte offset overflowed usize"))?;
275 let out_shape = if drop_axis {
276 shape[1..].to_vec()
277 } else {
@@ -253,7 +279,38 @@ fn slice_first_axis_range<'a, T>(
279 shape[0] = end - start;
280 shape
281 };
256 - Ok((out_shape, &data[data_start..data_end]))
282 + let WireArrayData::RawView(data) = &tensor.data else {
283 + return Err(multimodal!("cannot slice an aux tensor buffer"));
284 + };
285 + Ok(WireTensor::from_raw_bytes(
286 + tensor.dtype.clone(),
287 + out_shape,
288 + data.slice(data_start..data_end),
289 + ))
290 +}
291 +
292 +fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> {
293 + let WireArrayData::RawView(data) = &tensor.data else {
294 + return Err(multimodal!("expected an inline tensor buffer"));
295 + };
296 + let expected_bytes = tensor
297 + .shape
298 + .checked_numel()
299 + .and_then(|numel| numel.checked_mul(element_size))
300 + .ok_or_else(|| {
301 + multimodal!(
302 + "tensor shape {:?} byte length overflowed usize",
303 + tensor.shape
304 + )
305 + })?;
306 + if expected_bytes != data.len() {
307 + bail_multimodal!(
308 + "tensor shape {:?} expects {expected_bytes} bytes, got {}",
309 + tensor.shape,
310 + data.len()
311 + );
312 + }
313 + Ok(data)
314 }
315
316 #[cfg(test)]
@@ -262,28 +319,32 @@ mod tests {
319
320 #[test]
321 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 - };
322 + let data = vec![1.0_f32, 2.0, 3.0, 4.0];
323 + let expected_ptr = data.as_ptr().cast::<u8>().wrapping_add(2 * size_of::<f32>());
324 + let value = KwargValue::tensor(
325 + TensorKind::F32,
326 + WireTensor::from_f32(vec![2, 2], data).unwrap(),
327 + );
328
329 let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
330 panic!("expected tensor");
331 };
332
333 assert_eq!(tensor.shape, vec![2]);
334 + let raw_view = tensor.data.into_raw_view().unwrap();
335 + assert_eq!(raw_view.as_ptr(), expected_ptr);
336 assert_eq!(
276 - tensor.data.into_raw_view().unwrap(),
337 + raw_view,
338 [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
339 );
340 }
341
342 #[test]
343 fn flat_wire_value_range_keeps_first_axis() {
283 - let value = KwargValue::U32Tensor {
284 - data: (0..10).collect(),
285 - shape: vec![5, 2],
286 - };
344 + let value = KwargValue::tensor(
345 + TensorKind::U32,
346 + WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(),
347 + );
348
349 let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
350 panic!("expected tensor");
@@ -298,10 +359,10 @@ mod tests {
359
360 #[test]
361 fn flat_range_for_index_uses_size_tensor() {
301 - let sizes = KwargValue::I64Tensor {
302 - data: vec![2, 3, 4],
303 - shape: vec![3],
304 - };
362 + let sizes = KwargValue::tensor(
363 + TensorKind::I64,
364 + WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(),
365 + );
366
367 assert_eq!(
368 flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(),
@@ -311,10 +372,11 @@ mod tests {
372
373 #[test]
374 fn slice_first_axis_range_errors_on_shape_data_mismatch() {
314 - let error = slice_first_axis_range(&[2, 2], &[1.0_f32, 2.0, 3.0], 0, 1, true).unwrap_err();
375 + let tensor = WireTensor::from_raw("float32", vec![2, 2], vec![0; 3 * size_of::<f32>()]);
376 + let error = slice_first_axis_range(&tensor, size_of::<f32>(), 0, 1, true).unwrap_err();
377
378 assert!(
317 - matches!(error, Error::Multimodal(message) if message.contains("expects 4 elements"))
379 + matches!(error, Error::Multimodal(message) if message.contains("expects 16 bytes"))
380 );
381 }
382
rust/src/engine-core-client/src/client.rs
+1 -1
@@ -519,7 +519,7 @@ impl EngineCoreClient {
519 "registered request to engine"
520 );
521
522 - self.inner.send_to_engine(&engine_id, EngineCoreRequestType::Add, &req).await?;
522 + self.inner.send_request_to_engine(&engine_id, req).await?;
523 Ok(())
524 }
525 .await;
rust/src/engine-core-client/src/client/imp.rs
+40 -4
@@ -6,6 +6,7 @@ use std::sync::Arc;
6 use std::time::{SystemTime, UNIX_EPOCH};
7
8 use arc_swap::ArcSwapOption;
9 +use bytes::Bytes;
10 use parking_lot::Mutex;
11 use thiserror_ext::AsReport as _;
12 use tokio::runtime::Handle;
@@ -21,17 +22,30 @@ use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_outpu
22 use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
23 use crate::protocol::encode_msgpack;
24 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
24 -use crate::protocol::request::EngineCoreRequestType;
25 +use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
26 use crate::protocol::stats::SchedulerStats;
27 use crate::protocol::utility::UtilityOutput;
28 use crate::transport::{ConnectedEngine, EngineId};
29 use crate::{Error, Result, transport};
30
31 +const MSGPACK_ZERO_COPY_THRESHOLD_ENV: &str = "VLLM_MSGPACK_ZERO_COPY_THRESHOLD";
32 +const DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD: usize = 256;
33 +
34 +fn msgpack_zero_copy_threshold() -> usize {
35 + std::env::var(MSGPACK_ZERO_COPY_THRESHOLD_ENV)
36 + .ok()
37 + .and_then(|value| value.parse().ok())
38 + .unwrap_or(DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD)
39 +}
40 +
41 pub(crate) struct ClientInner {
42 input_send: RouterSendHalf,
43 /// The runtime handle used for sending messages to the engine.
44 handle: Handle,
45 model_name: String,
46 + /// Per-tensor byte threshold loaded from env variable
47 + /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
48 + msgpack_zero_copy_threshold: usize,
49 scheduler_stats_recorder: SchedulerStatsRecorder,
50 request_reg: Mutex<RequestRegistry>,
51 utility_reg: Mutex<UtilityRegistry>,
@@ -54,6 +68,7 @@ impl ClientInner {
68 input_send,
69 handle,
70 model_name,
71 + msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
72 scheduler_stats_recorder,
73 request_reg: Mutex::new(RequestRegistry::new(engines)),
74 utility_reg: Mutex::new(UtilityRegistry::default()),
@@ -229,9 +244,29 @@ impl ClientInner {
244 where
245 T: serde::Serialize + std::fmt::Debug,
246 {
232 - // TODO: for `EngineCoreRequest`, split outbound tensor raw views into aux
233 - // frames instead of always producing a single msgpack frame.
234 - let payload = encode_msgpack(payload)?;
247 + let payload = Bytes::from(encode_msgpack(payload)?);
248 + self.send_encoded_to_engine(engine_id, request_type, payload, Vec::new()).await
249 + }
250 +
251 + /// Send an add request, moving large tensor buffers into auxiliary frames.
252 + pub async fn send_request_to_engine(
253 + &self,
254 + engine_id: &EngineId,
255 + mut payload: EngineCoreRequest,
256 + ) -> Result<()> {
257 + let aux_frames = payload.extract_aux_frames(self.msgpack_zero_copy_threshold);
258 + let payload = Bytes::from(encode_msgpack(&payload)?);
259 + self.send_encoded_to_engine(engine_id, EngineCoreRequestType::Add, payload, aux_frames)
260 + .await
261 + }
262 +
263 + async fn send_encoded_to_engine(
264 + &self,
265 + engine_id: &EngineId,
266 + request_type: EngineCoreRequestType,
267 + payload: Bytes,
268 + aux_frames: Vec<Bytes>,
269 + ) -> Result<()> {
270 let mut input_send = self.input_send.clone();
271 let engine_id = engine_id.clone();
272
@@ -242,6 +277,7 @@ impl ClientInner {
277 &engine_id,
278 request_type.to_frame(),
279 payload,
280 + aux_frames,
281 )
282 .await
283 })
rust/src/engine-core-client/src/protocol/logprobs.rs
+3 -3
@@ -209,17 +209,17 @@ impl WireLogprobs {
209 logprob_token_ids: WireNdArray {
210 dtype: "<i8".to_string(),
211 shape: vec![rows, cols],
212 - data: WireArrayData::RawView(token_ids),
212 + data: WireArrayData::RawView(token_ids.into()),
213 },
214 logprobs: WireNdArray {
215 dtype: "<f4".to_string(),
216 shape: vec![rows, cols],
217 - data: WireArrayData::RawView(logprobs),
217 + data: WireArrayData::RawView(logprobs.into()),
218 },
219 token_ranks: WireNdArray {
220 dtype: "<i8".to_string(),
221 shape: vec![rows],
222 - data: WireArrayData::RawView(token_ranks),
222 + data: WireArrayData::RawView(token_ranks.into()),
223 },
224 cu_num_generated_tokens: None,
225 })
rust/src/engine-core-client/src/protocol/logprobs/array.rs
+4 -3
@@ -4,6 +4,7 @@
4 use std::io::Cursor;
5
6 use byteorder::{BigEndian, LittleEndian, NativeEndian, ReadBytesExt};
7 +use bytes::Bytes;
8 use itertools::Itertools as _;
9
10 use crate::error::{Error, Result, ext_value_decode};
@@ -126,7 +127,7 @@ pub(super) fn decode_array_metadata<Frame>(
127 field: &str,
128 frames: &[Frame],
129 expected_scalars: &[ScalarType],
129 -) -> Result<(Vec<usize>, Vec<u8>, ScalarType, Endianness)>
130 +) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)>
131 where
132 Frame: AsRef<[u8]>,
133 {
@@ -171,7 +172,7 @@ pub(super) fn resolve_array_bytes<Frame>(
172 value: WireArrayData,
173 field: &str,
174 frames: &[Frame],
174 -) -> Result<Vec<u8>>
175 +) -> Result<Bytes>
176 where
177 Frame: AsRef<[u8]>,
178 {
@@ -187,7 +188,7 @@ where
188 ),
189 )
190 })?;
190 - Ok(frame.as_ref().to_vec())
191 + Ok(Bytes::copy_from_slice(frame.as_ref()))
192 }
193 }
194 }
rust/src/engine-core-client/src/protocol/multimodal.rs
+39
@@ -3,6 +3,7 @@
3
4 use std::collections::BTreeMap;
5
6 +use bytes::Bytes;
7 use serde::{Deserialize, Serialize};
8 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
9
@@ -105,6 +106,44 @@ pub enum MmKwargValue {
106 List(Vec<MmKwargValue>),
107 }
108
109 +impl MmFeatureSpec {
110 + /// Extract large tensor buffers from this feature in serialized field order.
111 + pub(crate) fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
112 + if let Some(data) = &mut self.data {
113 + for elem in data.values_mut() {
114 + elem.extract_aux_frames(aux_frames, threshold);
115 + }
116 + }
117 + if let Some(is_embed) = &mut self.mm_position.is_embed {
118 + is_embed.extract_aux_frame(aux_frames, threshold);
119 + }
120 + }
121 +}
122 +
123 +impl MmFieldElem {
124 + /// Extract large tensor buffers from this field element.
125 + fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
126 + if let Some(data) = &mut self.data {
127 + data.extract_aux_frames(aux_frames, threshold);
128 + }
129 + }
130 +}
131 +
132 +impl MmKwargValue {
133 + /// Recursively extract large tensor buffers from this nested value.
134 + fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
135 + match self {
136 + Self::Tensor(tensor) => tensor.extract_aux_frame(aux_frames, threshold),
137 + Self::List(values) => {
138 + for value in values {
139 + value.extract_aux_frames(aux_frames, threshold);
140 + }
141 + }
142 + Self::Int(_) | Self::Float(_) => {}
143 + }
144 + }
145 +}
146 +
147 /// Defines how to interpret tensor data belonging to a keyword argument for
148 /// `MultiModalKwargsItems`, and vice versa.
149 ///
rust/src/engine-core-client/src/protocol/request.rs
+97
@@ -137,6 +137,17 @@ impl EngineCoreRequest {
137 }
138 Ok(())
139 }
140 +
141 + /// Extract large request tensors into ordered auxiliary frames.
142 + pub(crate) fn extract_aux_frames(&mut self, threshold: usize) -> Vec<Bytes> {
143 + let mut aux_frames = Vec::new();
144 + if let Some(features) = &mut self.mm_features {
145 + for feature in features {
146 + feature.extract_aux_frames(&mut aux_frames, threshold);
147 + }
148 + }
149 + aux_frames
150 + }
151 }
152
153 #[cfg(test)]
@@ -144,9 +155,15 @@ mod tests {
155 use rmpv::Value;
156
157 use super::*;
158 + use crate::protocol::multimodal::{
159 + MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, PlaceholderRange,
160 + };
161 use crate::protocol::sampling::EngineCoreSamplingParams;
162 + use crate::protocol::tensor::{WireArrayData, WireTensor};
163 use crate::protocol::{decode_value, encode_msgpack};
164
165 + const AUX_FRAME_THRESHOLD: usize = 256;
166 +
167 #[test]
168 fn engine_core_request_serializes_as_full_array() {
169 let request = EngineCoreRequest {
@@ -175,4 +192,84 @@ mod tests {
192 assert_eq!(array[10], Value::Nil);
193 assert_eq!(array[11], Value::from(7));
194 }
195 +
196 + #[test]
197 + fn engine_core_request_extracts_large_nested_tensors_in_wire_order() {
198 + let inline = vec![1_u8; AUX_FRAME_THRESHOLD - 1];
199 + let first_aux = vec![2_u8; AUX_FRAME_THRESHOLD];
200 + let second_aux = vec![3_u8; AUX_FRAME_THRESHOLD + 1];
201 + let first_aux_ptr = first_aux.as_ptr();
202 + let second_aux_ptr = second_aux.as_ptr();
203 + let mut request = EngineCoreRequest {
204 + mm_features: Some(vec![MmFeatureSpec {
205 + data: Some(BTreeMap::from([
206 + (
207 + "inline".to_string(),
208 + MmFieldElem {
209 + data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
210 + "uint8",
211 + vec![inline.len()],
212 + inline,
213 + ))),
214 + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
215 + },
216 + ),
217 + (
218 + "nested".to_string(),
219 + MmFieldElem {
220 + data: Some(MmKwargValue::List(vec![
221 + MmKwargValue::Int(7),
222 + MmKwargValue::Tensor(WireTensor::from_raw(
223 + "uint8",
224 + vec![first_aux.len()],
225 + first_aux,
226 + )),
227 + ])),
228 + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
229 + },
230 + ),
231 + ])),
232 + modality: "image".to_string(),
233 + identifier: "id".to_string(),
234 + mm_position: PlaceholderRange {
235 + offset: 0,
236 + length: second_aux.len(),
237 + is_embed: Some(WireTensor::from_raw(
238 + "bool",
239 + vec![second_aux.len()],
240 + second_aux,
241 + )),
242 + },
243 + mm_hash: None,
244 + }]),
245 + ..EngineCoreRequest::default()
246 + };
247 +
248 + let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD);
249 +
250 + assert_eq!(aux_frames.len(), 2);
251 + assert_eq!(aux_frames[0].as_ptr(), first_aux_ptr);
252 + assert_eq!(aux_frames[1].as_ptr(), second_aux_ptr);
253 + let feature = &request.mm_features.as_ref().unwrap()[0];
254 + let MmKwargValue::Tensor(inline) =
255 + feature.data.as_ref().unwrap()["inline"].data.as_ref().unwrap()
256 + else {
257 + panic!("expected inline tensor");
258 + };
259 + assert!(matches!(inline.data, WireArrayData::RawView(_)));
260 + let MmKwargValue::List(nested) =
261 + feature.data.as_ref().unwrap()["nested"].data.as_ref().unwrap()
262 + else {
263 + panic!("expected nested tensor list");
264 + };
265 + let MmKwargValue::Tensor(nested_tensor) = &nested[1] else {
266 + panic!("expected nested tensor");
267 + };
268 + assert_eq!(nested_tensor.data, WireArrayData::AuxIndex(1));
269 + assert_eq!(
270 + feature.mm_position.is_embed.as_ref().unwrap().data,
271 + WireArrayData::AuxIndex(2)
272 + );
273 + assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty());
274 + }
275 }
rust/src/engine-core-client/src/protocol/tensor.rs
+98 -44
@@ -1,7 +1,8 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -use bytemuck::allocation::pod_collect_to_vec;
4 +use bytemuck::{Pod, cast_slice};
5 +use bytes::Bytes;
6 use enum_as_inner::EnumAsInner;
7 use half::{bf16, f16};
8 use rmpv::Value;
@@ -20,6 +21,21 @@ struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));
21
22 struct ByteSlice<'a>(&'a [u8]);
23
24 +struct PodVec<T: Pod>(Vec<T>);
25 +
26 +impl<T: Pod> AsRef<[u8]> for PodVec<T> {
27 + fn as_ref(&self) -> &[u8] {
28 + cast_slice(&self.0)
29 + }
30 +}
31 +
32 +fn bytes_from_pod_vec<T>(data: Vec<T>) -> Bytes
33 +where
34 + T: Pod + Send + 'static,
35 +{
36 + Bytes::from_owner(PodVec(data))
37 +}
38 +
39 impl Serialize for ByteSlice<'_> {
40 fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
41 where
@@ -55,58 +71,63 @@ pub struct WireNdArray {
71
72 impl WireNdArray {
73 /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
58 - pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
59 - let data = data.as_ref();
74 + ///
75 + /// Takes ownership of the backing buffer without copying its data.
76 + pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
77 validate_element_count(&shape, data.len())?;
61 - Ok(Self {
62 - dtype: "float32".to_string(),
78 + Ok(Self::from_raw_bytes(
79 + "float32",
80 shape,
64 - data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
65 - })
81 + bytes_from_pod_vec(data),
82 + ))
83 }
84
85 /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
69 - pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
70 - let data = data.as_ref();
86 + ///
87 + /// Takes ownership of the backing buffer without copying its data.
88 + pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
89 validate_element_count(&shape, data.len())?;
72 - Ok(Self {
73 - dtype: "float16".to_string(),
90 + Ok(Self::from_raw_bytes(
91 + "float16",
92 shape,
75 - data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
76 - })
93 + bytes_from_pod_vec(data),
94 + ))
95 }
96
97 /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
80 - pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
81 - let data = data.as_ref();
98 + ///
99 + /// Takes ownership of the backing buffer without copying its data.
100 + pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
101 validate_element_count(&shape, data.len())?;
83 - Ok(Self {
84 - dtype: "bfloat16".to_string(),
102 + Ok(Self::from_raw_bytes(
103 + "bfloat16",
104 shape,
86 - data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
87 - })
105 + bytes_from_pod_vec(data),
106 + ))
107 }
108
109 /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
91 - pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
92 - let data = data.as_ref();
110 + ///
111 + /// Takes ownership of the backing buffer without copying its data.
112 + pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
113 validate_element_count(&shape, data.len())?;
94 - Ok(Self {
95 - dtype: "int64".to_string(),
114 + Ok(Self::from_raw_bytes(
115 + "int64",
116 shape,
97 - data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
98 - })
117 + bytes_from_pod_vec(data),
118 + ))
119 }
120
121 /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
102 - pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
103 - let data = data.as_ref();
122 + ///
123 + /// Takes ownership of the backing buffer without copying its data.
124 + pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
125 validate_element_count(&shape, data.len())?;
105 - Ok(Self {
106 - dtype: "uint32".to_string(),
126 + Ok(Self::from_raw_bytes(
127 + "uint32",
128 shape,
108 - data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
109 - })
129 + bytes_from_pod_vec(data),
130 + ))
131 }
132
133 /// Build a bool tensor/ndarray backed by raw-view bytes.
@@ -118,7 +139,9 @@ impl WireNdArray {
139 Ok(Self {
140 dtype: "bool".to_string(),
141 shape,
121 - data: WireArrayData::RawView(data.iter().map(|value| u8::from(*value)).collect()),
142 + data: WireArrayData::RawView(Bytes::from(
143 + data.into_iter().map(u8::from).collect::<Vec<_>>(),
144 + )),
145 })
146 }
147
@@ -127,12 +150,22 @@ impl WireNdArray {
150 /// Use this as an escape hatch when the caller already owns bytes that
151 /// match the requested `dtype` and `shape`.
152 pub fn from_raw(dtype: impl Into<String>, shape: Vec<usize>, data: Vec<u8>) -> Self {
153 + Self::from_raw_bytes(dtype, shape, Bytes::from(data))
154 + }
155 +
156 + /// Build a tensor/ndarray from an owned immutable raw-view buffer.
157 + pub fn from_raw_bytes(dtype: impl Into<String>, shape: Vec<usize>, data: Bytes) -> Self {
158 Self {
159 dtype: dtype.into(),
160 shape,
161 data: WireArrayData::RawView(data),
162 }
163 }
164 +
165 + /// Move a sufficiently large inline buffer into the ordered auxiliary-frame list.
166 + pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
167 + self.data.extract_aux_frame(aux_frames, threshold);
168 + }
169 }
170
171 /// Validate that the number of elements implied by the shape matches the length
@@ -170,7 +203,25 @@ pub enum WireArrayData {
203 /// stored.
204 AuxIndex(usize),
205 /// The raw bytes of this array/tensor.
173 - RawView(Vec<u8>),
206 + RawView(Bytes),
207 +}
208 +
209 +impl WireArrayData {
210 + /// Replace a sufficiently large raw view with its one-based auxiliary-frame index.
211 + fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
212 + let Self::RawView(bytes) = self else {
213 + return;
214 + };
215 + if bytes.len() < threshold {
216 + return;
217 + }
218 +
219 + let index = aux_frames.len() + 1;
220 + let bytes = std::mem::replace(self, Self::AuxIndex(index))
221 + .into_raw_view()
222 + .expect("raw view was matched above");
223 + aux_frames.push(bytes);
224 + }
225 }
226
227 impl<'de> Deserialize<'de> for WireArrayData {
@@ -180,7 +231,9 @@ impl<'de> Deserialize<'de> for WireArrayData {
231 {
232 let value = Value::deserialize(deserializer)?;
233 match value {
183 - Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => Ok(Self::RawView(bytes)),
234 + Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => {
235 + Ok(Self::RawView(Bytes::from(bytes)))
236 + }
237 Value::Ext(tag, _) => Err(serde::de::Error::custom(format!(
238 "unsupported extension type code {tag}"
239 ))),
@@ -201,9 +254,6 @@ impl Serialize for WireArrayData {
254 where
255 S: Serializer,
256 {
204 - // TODO: outbound request serialization currently only supports inline
205 - // raw-view bytes. Emitting aux frames needs transport-level plumbing;
206 - // serializing `AuxIndex` here only preserves an already-built reference.
257 match self {
258 Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
259 Self::RawView(bytes) => {
@@ -221,7 +271,7 @@ mod tests {
271 fn raw_view_serializes_as_msgpack_ext() {
272 let bytes = vec![1, 2, 3, 4];
273 let encoded =
224 - rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
274 + rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone().into())).expect("encode");
275 let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
276 .expect("encode expected");
277
@@ -234,11 +284,15 @@ mod tests {
284
285 #[test]
286 fn constructors_build_raw_view_tensors() {
237 - let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
287 + let f32_data = vec![1.0, 2.5];
288 + let f32_data_ptr = f32_data.as_ptr().cast::<u8>();
289 + let f32_tensor = WireNdArray::from_f32(vec![2], f32_data).unwrap();
290 assert_eq!(f32_tensor.dtype, "float32");
291 assert_eq!(f32_tensor.shape, vec![2]);
292 + let f32_raw_view = f32_tensor.data.into_raw_view().expect("raw view");
293 + assert_eq!(f32_raw_view.as_ptr(), f32_data_ptr);
294 assert_eq!(
241 - f32_tensor.data.into_raw_view().expect("raw view"),
295 + f32_raw_view,
296 [1.0_f32, 2.5].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
297 );
298
@@ -258,15 +312,15 @@ mod tests {
312 let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap();
313 assert_eq!(i64_tensor.dtype, "int64");
314 assert_eq!(
261 - i64_tensor.data.into_raw_view().expect("raw view"),
262 - (-7_i64).to_ne_bytes()
315 + i64_tensor.data.into_raw_view().expect("raw view").as_ref(),
316 + (-7_i64).to_ne_bytes().as_ref()
317 );
318
319 let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap();
320 assert_eq!(u32_tensor.dtype, "uint32");
321 assert_eq!(
268 - u32_tensor.data.into_raw_view().expect("raw view"),
269 - 42_u32.to_ne_bytes()
322 + u32_tensor.data.into_raw_view().expect("raw view").as_ref(),
323 + 42_u32.to_ne_bytes().as_ref()
324 );
325
326 let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap();
rust/src/engine-core-client/src/tests/client.rs
+81 -1
@@ -32,7 +32,7 @@ use crate::protocol::output::{
32 use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
33 use crate::protocol::sampling::EngineCoreSamplingParams;
34 use crate::protocol::stats::SchedulerStats;
35 -use crate::protocol::tensor::WireTensor;
35 +use crate::protocol::tensor::{WireArrayData, WireTensor};
36 use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
37 use crate::test_utils::{
38 IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
@@ -1728,6 +1728,86 @@ async fn client_decodes_multipart_logprob_outputs() {
1728 client.shutdown().await.unwrap();
1729 }
1730
1731 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1732 +async fn client_sends_large_multimodal_tensor_as_aux_frame() {
1733 + init_tracing();
1734 + let ipc = IpcNamespace::new().unwrap();
1735 + let handshake_address = ipc.handshake_endpoint();
1736 + let engine_id = b"engine-multimodal-aux".to_vec();
1737 + let tensor_data = (0..64).map(|value| value as f32).collect::<Vec<_>>();
1738 + let expected_bytes =
1739 + tensor_data.iter().flat_map(|value| value.to_ne_bytes()).collect::<Vec<_>>();
1740 + let mut request = sample_multimodal_request();
1741 + request.mm_features.as_mut().unwrap()[0]
1742 + .data
1743 + .as_mut()
1744 + .unwrap()
1745 + .get_mut("pixel_values")
1746 + .unwrap()
1747 + .data = Some(MmKwargValue::Tensor(
1748 + WireTensor::from_f32(vec![64], tensor_data).unwrap(),
1749 + ));
1750 +
1751 + let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1752 + handshake_address.clone(),
1753 + engine_id.clone(),
1754 + move |dealer, push| {
1755 + Box::pin(async move {
1756 + let add = recv_engine_message(dealer).await;
1757 + assert_eq!(add.len(), 3);
1758 + assert_eq!(add[0].as_ref(), &[0x00]);
1759 + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
1760 + let MmKwargValue::Tensor(tensor) =
1761 + request.mm_features.as_ref().unwrap()[0].data.as_ref().unwrap()["pixel_values"]
1762 + .data
1763 + .as_ref()
1764 + .unwrap()
1765 + else {
1766 + panic!("expected tensor");
1767 + };
1768 + assert_eq!(tensor.data, WireArrayData::AuxIndex(1));
1769 + assert_eq!(add[2].as_ref(), expected_bytes);
1770 +
1771 + send_outputs(
1772 + push,
1773 + RequestBatchOutputs {
1774 + outputs: vec![request_output(
1775 + "req-mm",
1776 + vec![],
1777 + Some(EngineCoreFinishReason::Length),
1778 + )],
1779 + finished_requests: Some(BTreeSet::from(["req-mm".to_string()])),
1780 + ..Default::default()
1781 + }
1782 + .into(),
1783 + )
1784 + .await;
1785 + })
1786 + },
1787 + );
1788 +
1789 + let client = connect_client_with_ipc(
1790 + handshake_test_config(
1791 + handshake_address,
1792 + 1,
1793 + "test-model",
1794 + Duration::from_secs(2),
1795 + 0,
1796 + None,
1797 + ),
1798 + &ipc,
1799 + )
1800 + .await;
1801 +
1802 + let outputs = client.call(request).await.unwrap().collect::<Vec<_>>().await;
1803 + assert_eq!(outputs.len(), 1);
1804 + assert!(outputs[0].is_ok());
1805 +
1806 + let _ = shutdown_tx.send(());
1807 + engine_task.await.unwrap();
1808 + client.shutdown().await.unwrap();
1809 +}
1810 +
1811 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1812 async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
1813 init_tracing();
rust/src/engine-core-client/src/transport.rs
+7 -7
@@ -512,14 +512,14 @@ pub async fn send_message(
512 input_send: &mut RouterSendHalf,
513 engine_id: &EngineId,
514 request_type: Bytes,
515 - payload: Vec<u8>,
515 + payload: Bytes,
516 + aux_frames: Vec<Bytes>,
517 ) -> Result<()> {
517 - let message = ZmqMessage::try_from(vec![
518 - engine_id.to_frame(),
519 - request_type,
520 - Bytes::from(payload),
521 - ])
522 - .expect("router messages must contain identity and payload");
518 + let mut frames = Vec::with_capacity(3 + aux_frames.len());
519 + frames.extend([engine_id.to_frame(), request_type, payload]);
520 + frames.extend(aux_frames);
521 + let message =
522 + ZmqMessage::try_from(frames).expect("router messages must contain identity and payload");
523
524 trace!(
525 ?engine_id,