[Rust Frontend] Take the raw buffer in mm tensor lowering when possible (#53528)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: roachsinai <roachsinai@qq.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
RoachZhao committed
Aug 29, 2026 at 09:05 UTC
4c6c9d569f83d4dc5509e77e3458734a1fda2c2d
2 files changed
+75
-10
rust/src/chat/Cargo.toml
+1
-1
@@ -14,6 +14,7 @@ itertools.workspace = true
14
llm-multimodal.workspace = true
15
minijinja.workspace = true
16
minijinja-contrib.workspace = true
17
+ndarray.workspace = true
18
openai-harmony.workspace = true
19
reqwest-0-13.workspace = true
20
serde.workspace = true
@@ -41,7 +42,6 @@ anyhow.workspace = true
42
bytes.workspace = true
43
clap.workspace = true
44
expect-test.workspace = true
44
-ndarray.workspace = true
45
paste.workspace = true
46
rmp-serde.workspace = true
47
serial_test.workspace = true
rust/src/chat/src/multimodal/tensor.rs
+74
-9
@@ -6,6 +6,7 @@ use std::mem::size_of;
6
7
use half::{bf16, f16};
8
use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs};
9
+use ndarray::ArrayD;
10
use vllm_engine_core_client::protocol::dtype::ModelDtype;
11
use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
12
use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor};
@@ -75,11 +76,7 @@ pub(super) fn collect_tensors(
76
..
77
} = preprocessed;
78
78
- let primary_value = {
79
- let shape = encoder_input.shape().to_vec();
80
- let data = encoder_input.into_iter().collect();
81
- KwargValue::from_f32_tensor(data, shape, float_dtype)?
82
- };
79
+ let primary_value = KwargValue::from_f32_array(encoder_input, float_dtype)?;
80
81
let mut tensors = HashMap::new();
82
tensors.insert(primary_key.to_string(), primary_value);
@@ -95,7 +92,7 @@ impl KwargValue {
92
93
Ok(match value {
94
ModelSpecificValue::Tensor { data, shape } => {
98
- Self::from_f32_tensor(data, shape, float_dtype)?
95
+ Self::from_f32_parts(data, shape, float_dtype)?
96
}
97
ModelSpecificValue::IntTensor { data, shape } => {
98
let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?;
@@ -126,9 +123,31 @@ impl KwargValue {
123
})
124
}
125
126
+ fn from_f32_array(array: ArrayD<f32>, float_dtype: ModelDtype) -> Result<Self> {
127
+ let shape = array.shape().to_vec();
128
+ // `into_iter()` on dynamically-dimensioned arrays walks a slow
129
+ // per-element path; take the raw buffer instead when the layout
130
+ // allows it.
131
+ let data = if array.is_standard_layout() {
132
+ let len = array.len();
133
+ let (data, offset) = array.into_raw_vec_and_offset();
134
+ let start = offset.unwrap_or(0);
135
+ if start == 0 && data.len() == len {
136
+ data
137
+ } else {
138
+ // Buffer with unused head/tail: copy the used range, which
139
+ // standard strides place contiguously in logical order.
140
+ data[start..start + len].to_vec()
141
+ }
142
+ } else {
143
+ array.into_iter().collect()
144
+ };
145
+ Self::from_f32_parts(data, shape, float_dtype)
146
+ }
147
+
148
/// Convert a float tensor to the target float dtype if needed, keeping the
149
/// same shape.
131
- fn from_f32_tensor(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
150
+ fn from_f32_parts(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
151
let (kind, wire) = match float_dtype {
152
ModelDtype::Float16 => (
153
TensorKind::F16,
@@ -316,6 +335,52 @@ fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> {
335
#[cfg(test)]
336
mod tests {
337
use super::*;
338
+ use ndarray::{Array2, s};
339
+
340
+ #[test]
341
+ fn collect_tensors_lowering_matches_logical_order_across_layouts() {
342
+ let standard =
343
+ Array2::from_shape_vec((2, 3), (1..=6).map(|v| v as f32).collect::<Vec<f32>>())
344
+ .unwrap();
345
+ let standard_with_offset =
346
+ Array2::from_shape_vec((3, 3), (1..=9).map(|v| v as f32).collect::<Vec<f32>>())
347
+ .unwrap()
348
+ .slice_move(s![1.., ..]);
349
+ let strided =
350
+ Array2::from_shape_vec((2, 3), (1..=6).map(|v| v as f32).collect::<Vec<f32>>())
351
+ .unwrap()
352
+ .reversed_axes();
353
+
354
+ assert!(standard.is_standard_layout());
355
+ assert!(standard_with_offset.is_standard_layout());
356
+ assert!(!strided.is_standard_layout());
357
+
358
+ for (name, array) in [
359
+ ("standard", standard),
360
+ ("standard_with_offset", standard_with_offset),
361
+ ("strided", strided),
362
+ ] {
363
+ let expected: Vec<u8> = array.iter().flat_map(|v| v.to_ne_bytes()).collect();
364
+ let shape = array.shape().to_vec();
365
+ let preprocessed = PreprocessedEncoderInputs::new(array, vec![6], vec![(3, 2)]);
366
+ let tensors =
367
+ collect_tensors(preprocessed, "pixel_values", ModelDtype::Float32).unwrap();
368
+
369
+ let ProtocolKwargValue::Tensor(tensor) =
370
+ ProtocolKwargValue::try_from(&tensors["pixel_values"]).unwrap()
371
+ else {
372
+ panic!("expected tensor for {name} layout");
373
+ };
374
+
375
+ assert_eq!(tensor.shape, shape, "{name} layout shape");
376
+ assert_eq!(tensor.dtype, "float32", "{name} layout dtype");
377
+ assert_eq!(
378
+ tensor.data.into_raw_view().unwrap(),
379
+ expected,
380
+ "{name} layout data"
381
+ );
382
+ }
383
+ }
384
385
#[test]
386
fn batched_wire_value_at_drops_first_axis() {
@@ -383,7 +448,7 @@ mod tests {
448
#[test]
449
fn bfloat16_tensor_wire_uses_bfloat16_dtype() {
450
let value =
386
- KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
451
+ KwargValue::from_f32_parts(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
452
453
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
454
else {
@@ -398,7 +463,7 @@ mod tests {
463
#[test]
464
fn float16_tensor_wire_uses_float16_dtype() {
465
let value =
401
- KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
466
+ KwargValue::from_f32_parts(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
467
468
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
469
else {