@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 41cffa9af0

[Rust Frontend] Strongly type wire dtypes and multimodal modalities (#56174)

Co-authored-by: Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Sep 10, 2026 at 12:11 UTC 41cffa9af0b7317d71838262004186b68e976773
15 files changed +453 -270
rust/src/chat/src/multimodal.rs
+20 -41
@@ -60,30 +60,9 @@ pub struct MultimodalModelInfo {
60 /// Per-modality item-count limits configured by `--limit-mm-per-prompt`.
61 ///
62 /// Modalities absent from the map are unlimited.
63 -pub type MmLimitPerPrompt = HashMap<MmLimitModality, MmLimitSpec>;
63 +pub type MmLimitPerPrompt = HashMap<MmModality, MmLimitSpec>;
64
65 -/// Modalities that `--limit-mm-per-prompt` can be keyed by.
66 -///
67 -/// Closed on purpose: these are exactly the keys Python accepts, per
68 -/// `MultiModalDummyOptionsBuiltins` in `vllm/config/multimodal.py`.
69 -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
70 -#[serde(rename_all = "snake_case")]
71 -pub enum MmLimitModality {
72 - Image,
73 - Audio,
74 - Video,
75 -}
76 -
77 -impl MmLimitModality {
78 - /// The wire name, matching Python's modality strings.
79 - pub fn as_str(self) -> &'static str {
80 - match self {
81 - Self::Image => "image",
82 - Self::Audio => "audio",
83 - Self::Video => "video",
84 - }
85 - }
86 -}
65 +pub use vllm_engine_core_client::protocol::multimodal::MmModality;
66
67 /// One modality's limit, in either of the two shapes Python accepts.
68 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -646,17 +625,17 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result<String> {
625 /// Embedding inputs share their base modality's budget rather than getting one
626 /// of their own, matching Python's `modality.replace("_embeds", "")` in
627 /// `vllm/entrypoints/chat_utils.py`.
649 -fn media_part_limit_modality(part: &MediaContentPart) -> Option<MmLimitModality> {
628 +fn media_part_limit_modality(part: &MediaContentPart) -> Option<MmModality> {
629 match part {
630 MediaContentPart::Text { .. } => None,
631 MediaContentPart::ImageUrl { .. }
632 | MediaContentPart::ImageData { .. }
654 - | MediaContentPart::ImageEmbeds { .. } => Some(MmLimitModality::Image),
633 + | MediaContentPart::ImageEmbeds { .. } => Some(MmModality::Image),
634 MediaContentPart::AudioUrl { .. } | MediaContentPart::AudioData { .. } => {
656 - Some(MmLimitModality::Audio)
635 + Some(MmModality::Audio)
636 }
637 MediaContentPart::VideoUrl { .. } | MediaContentPart::VideoData { .. } => {
659 - Some(MmLimitModality::Video)
638 + Some(MmModality::Video)
639 }
640 }
641 }
@@ -667,7 +646,7 @@ impl MultimodalModelInfo {
646 ///
647 /// Modalities without a configured count are unlimited.
648 fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> {
670 - let mut counts: HashMap<MmLimitModality, usize> = HashMap::new();
649 + let mut counts: HashMap<MmModality, usize> = HashMap::new();
650 for part in media_parts {
651 if let Some(modality) = media_part_limit_modality(part) {
652 *counts.entry(modality).or_default() += 1;
@@ -738,7 +717,11 @@ impl MultimodalModelInfo {
717 for (item, range) in izip!(media.items, media_ranges) {
718 features.push(MmFeatureSpec {
719 data: Some(item.data),
741 - modality: media.modality.to_string(),
720 + modality: match media.modality {
721 + Modality::Image | Modality::ImageEmbeds => MmModality::Image,
722 + Modality::Video => MmModality::Video,
723 + Modality::Audio => MmModality::Audio,
724 + },
725 identifier: item.uuid.unwrap_or_else(|| item.hash.clone()),
726 mm_position: range,
727 mm_hash: Some(item.hash),
@@ -1029,10 +1012,8 @@ mod tests {
1012
1013 #[test]
1014 fn validate_mm_limits_enforces_configured_limit_at_the_boundary() {
1032 - let info = qwen3_vl_info_with_limits(HashMap::from([(
1033 - MmLimitModality::Image,
1034 - MmLimitSpec::Count(1),
1035 - )]));
1015 + let info =
1016 + qwen3_vl_info_with_limits(HashMap::from([(MmModality::Image, MmLimitSpec::Count(1))]));
1017 assert!(info.validate_mm_limits(&[image_url_part()]).is_ok());
1018
1019 let error = info.validate_mm_limits(&[image_url_part(), image_url_part()]).unwrap_err();
@@ -1047,10 +1028,8 @@ mod tests {
1028
1029 #[test]
1030 fn validate_mm_limits_counts_image_embeds_against_the_image_limit() {
1050 - let info = qwen3_vl_info_with_limits(HashMap::from([(
1051 - MmLimitModality::Image,
1052 - MmLimitSpec::Count(1),
1053 - )]));
1031 + let info =
1032 + qwen3_vl_info_with_limits(HashMap::from([(MmModality::Image, MmLimitSpec::Count(1))]));
1033 let image_embeds_part = MediaContentPart::ImageEmbeds {
1034 payload: serde_json::Value::String("AAAA".to_string()),
1035 uuid: None,
@@ -1068,7 +1047,7 @@ mod tests {
1047 #[test]
1048 fn validate_mm_limits_treats_a_count_less_options_object_as_unlimited() {
1049 let info = qwen3_vl_info_with_limits(HashMap::from([(
1071 - MmLimitModality::Image,
1050 + MmModality::Image,
1051 MmLimitSpec::Options {
1052 count: None,
1053 extra: BTreeMap::from([("width".to_string(), serde_json::json!(512))]),
@@ -1086,9 +1065,9 @@ mod tests {
1065 fn limit_map_parses_both_shapes_python_accepts() {
1066 let limits = parse_limits(r#"{"image": 16, "video": {"count": 1, "num_frames": 32}}"#);
1067
1089 - assert_eq!(limits[&MmLimitModality::Image].count(), Some(16));
1090 - assert_eq!(limits[&MmLimitModality::Video].count(), Some(1));
1091 - assert_eq!(limits.get(&MmLimitModality::Audio), None);
1068 + assert_eq!(limits[&MmModality::Image].count(), Some(16));
1069 + assert_eq!(limits[&MmModality::Video].count(), Some(1));
1070 + assert_eq!(limits.get(&MmModality::Audio), None);
1071 }
1072
1073 #[test]
rust/src/chat/src/multimodal/audio.rs
+4 -4
@@ -268,14 +268,14 @@ mod tests {
268 assert!(matches!(
269 features.data.as_ref(),
270 Some(MmKwargValue::Tensor(tensor))
271 - if tensor.dtype == "float32" && tensor.shape.first() == Some(&128)
271 + if tensor.dtype.as_str() == "float32" && tensor.shape.first() == Some(&128)
272 ));
273 let lengths = &item.data["audio_feature_lengths"];
274 assert!(matches!(&lengths.field, MmField::Batched(_)));
275 assert!(matches!(
276 lengths.data.as_ref(),
277 Some(MmKwargValue::Tensor(tensor))
278 - if tensor.dtype == "int64" && tensor.shape.is_empty()
278 + if tensor.dtype.as_str() == "int64" && tensor.shape.is_empty()
279 ));
280 }
281
@@ -314,14 +314,14 @@ mod tests {
314 assert!(matches!(
315 features.data.as_ref(),
316 Some(MmKwargValue::Tensor(tensor))
317 - if tensor.dtype == "float32" && tensor.shape.get(1) == Some(&80)
317 + if tensor.dtype.as_str() == "float32" && tensor.shape.get(1) == Some(&80)
318 ));
319 let count = &item.data["num_audio_tokens"];
320 assert!(matches!(&count.field, MmField::Batched(_)));
321 assert!(matches!(
322 count.data.as_ref(),
323 Some(MmKwargValue::Tensor(tensor))
324 - if tensor.dtype == "int64" && tensor.shape.is_empty()
324 + if tensor.dtype.as_str() == "int64" && tensor.shape.is_empty()
325 ));
326 }
327 }
rust/src/chat/src/multimodal/expand.rs
+1 -1
@@ -228,7 +228,7 @@ mod tests {
228
229 fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) {
230 let tensor = range.is_embed.as_ref().expect("is_embed mask");
231 - assert_eq!(tensor.dtype, "bool");
231 + assert_eq!(tensor.dtype.as_str(), "bool");
232 assert_eq!(tensor.shape, vec![expected.len()]);
233 assert_eq!(
234 tensor.data,
rust/src/chat/src/multimodal/tensor.rs
+46 -115
@@ -7,54 +7,17 @@ use std::mem::size_of;
7 use half::{bf16, f16};
8 use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs};
9 use ndarray::ArrayD;
10 -use vllm_engine_core_client::protocol::dtype::ModelDtype;
10 +use vllm_engine_core_client::protocol::dtype::{ModelDtype, TensorDtype};
11 use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
12 -use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor};
12 +use vllm_engine_core_client::protocol::tensor::{WireArrayData, WireTensor};
13
14 use crate::error::{Error, Result, bail_multimodal, multimodal};
15
16 -/// Element type retained alongside an encoded tensor during multimodal lowering.
17 -#[derive(Debug, Clone, Copy)]
18 -pub(super) enum TensorKind {
19 - /// 32-bit floating point.
20 - F32,
21 - /// IEEE 16-bit floating point.
22 - F16,
23 - /// Brain floating point.
24 - Bf16,
25 - /// Signed 64-bit integer.
26 - I64,
27 - /// Unsigned 32-bit integer.
28 - U32,
29 -}
30 -
31 -impl TensorKind {
32 - const fn element_size(self) -> usize {
33 - match self {
34 - Self::F32 => size_of::<f32>(),
35 - Self::F16 => size_of::<f16>(),
36 - Self::Bf16 => size_of::<bf16>(),
37 - Self::I64 => size_of::<i64>(),
38 - Self::U32 => size_of::<u32>(),
39 - }
40 - }
41 -
42 - const fn wire_dtype(self) -> &'static str {
43 - match self {
44 - Self::F32 => "float32",
45 - Self::F16 => "float16",
46 - Self::Bf16 => "bfloat16",
47 - Self::I64 => "int64",
48 - Self::U32 => "uint32",
49 - }
50 - }
51 -}
52 -
16 /// Representation for multimodal kwarg values for transformation.
17 #[derive(Debug)]
18 pub(super) enum KwargValue {
19 /// Tensor with row-major flat data and shape.
57 - Tensor { kind: TensorKind, wire: WireTensor },
20 + Tensor { wire: WireTensor },
21 /// Non-tensor kwarg value that is shared or copied as-is.
22 Passthrough(ProtocolKwargValue),
23 }
@@ -96,11 +59,11 @@ impl KwargValue {
59 }
60 ModelSpecificValue::IntTensor { data, shape } => {
61 let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?;
99 - Self::tensor(TensorKind::I64, wire)
62 + Self::Tensor { wire }
63 }
64 ModelSpecificValue::UintTensor { data, shape } => {
65 let wire = WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?;
103 - Self::tensor(TensorKind::U32, wire)
66 + Self::Tensor { wire }
67 }
68 ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)),
69 ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)),
@@ -148,23 +111,16 @@ impl KwargValue {
111 /// Convert a float tensor to the target float dtype if needed, keeping the
112 /// same shape.
113 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,
154 - WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect()),
155 - ),
156 - ModelDtype::BFloat16 => (
157 - TensorKind::Bf16,
158 - WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect()),
159 - ),
160 - ModelDtype::Float32 => (TensorKind::F32, WireTensor::from_f32(shape, data)),
114 + let wire = match float_dtype {
115 + ModelDtype::Float16 => {
116 + WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect())
117 + }
118 + ModelDtype::BFloat16 => {
119 + WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect())
120 + }
121 + ModelDtype::Float32 => WireTensor::from_f32(shape, data),
122 };
162 - wire.map(|wire| Self::tensor(kind, wire)).map_err(Error::Multimodal)
163 - }
164 -
165 - fn tensor(kind: TensorKind, wire: WireTensor) -> Self {
166 - debug_assert_eq!(wire.dtype, kind.wire_dtype());
167 - Self::Tensor { kind, wire }
123 + wire.map(|wire| Self::Tensor { wire }).map_err(Error::Multimodal)
124 }
125 }
126
@@ -215,9 +171,7 @@ impl KwargValue {
171 drop_axis: bool,
172 ) -> Result<ProtocolKwargValue> {
173 let wire = match self {
218 - Self::Tensor { kind, wire } => {
219 - slice_first_axis_range(wire, kind.element_size(), start, end, drop_axis)
220 - }
174 + Self::Tensor { wire } => slice_first_axis_range(wire, start, end, drop_axis),
175 Self::Passthrough(value) => return Ok(value.clone()),
176 };
177 wire.map(ProtocolKwargValue::Tensor)
@@ -244,10 +198,7 @@ pub(super) fn flat_range_for_index(
198 /// Read a tensor value as per-image sizes for flat slicing.
199 fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
200 match tensor {
247 - KwargValue::Tensor {
248 - kind: TensorKind::I64,
249 - wire,
250 - } => raw_tensor_bytes(wire, size_of::<i64>())?
201 + KwargValue::Tensor { wire } if wire.dtype == TensorDtype::I64 => raw_tensor_bytes(wire)?
202 .chunks_exact(size_of::<i64>())
203 .map(|bytes| i64::from_ne_bytes(bytes.try_into().expect("exact int64 chunk")))
204 .map(|value| {
@@ -255,13 +206,14 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
206 .map_err(|_| multimodal!("negative flat tensor size `{value}`"))
207 })
208 .collect(),
258 - KwargValue::Tensor {
259 - kind: TensorKind::U32,
260 - wire,
261 - } => Ok(raw_tensor_bytes(wire, size_of::<u32>())?
262 - .chunks_exact(size_of::<u32>())
263 - .map(|bytes| u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize)
264 - .collect()),
209 + KwargValue::Tensor { wire } if wire.dtype == TensorDtype::U32 => {
210 + Ok(raw_tensor_bytes(wire)?
211 + .chunks_exact(size_of::<u32>())
212 + .map(|bytes| {
213 + u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize
214 + })
215 + .collect())
216 + }
217 _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")),
218 }
219 }
@@ -269,13 +221,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
221 /// Slice a flat row-major tensor along its first axis.
222 fn slice_first_axis_range(
223 tensor: &WireTensor,
272 - element_size: usize,
224 start: usize,
225 end: usize,
226 drop_axis: bool,
227 ) -> Result<WireTensor> {
228 let shape = tensor.shape.as_slice();
278 - raw_tensor_bytes(tensor, element_size)?;
229 + raw_tensor_bytes(tensor)?;
230 + let element_size = tensor.dtype.element_size();
231 let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
232 if start > end || end > first_dim {
233 bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -302,34 +254,15 @@ fn slice_first_axis_range(
254 return Err(multimodal!("cannot slice an aux tensor buffer"));
255 };
256 Ok(WireTensor::from_raw_bytes(
305 - tensor.dtype.clone(),
257 + tensor.dtype,
258 out_shape,
259 data.slice(data_start..data_end),
260 ))
261 }
262
311 -fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> {
312 - let WireArrayData::RawView(data) = &tensor.data else {
313 - return Err(multimodal!("expected an inline tensor buffer"));
314 - };
315 - let expected_bytes = tensor
316 - .shape
317 - .checked_numel()
318 - .and_then(|numel| numel.checked_mul(element_size))
319 - .ok_or_else(|| {
320 - multimodal!(
321 - "tensor shape {:?} byte length overflowed usize",
322 - tensor.shape
323 - )
324 - })?;
325 - if expected_bytes != data.len() {
326 - bail_multimodal!(
327 - "tensor shape {:?} expects {expected_bytes} bytes, got {}",
328 - tensor.shape,
329 - data.len()
330 - );
331 - }
332 - Ok(data)
263 +fn raw_tensor_bytes(tensor: &WireTensor) -> Result<&[u8]> {
264 + tensor.validate_inline().map_err(Error::Multimodal)?;
265 + Ok(tensor.data.as_raw_view().expect("validated inline tensor"))
266 }
267
268 #[cfg(test)]
@@ -373,7 +306,7 @@ mod tests {
306 };
307
308 assert_eq!(tensor.shape, shape, "{name} layout shape");
376 - assert_eq!(tensor.dtype, "float32", "{name} layout dtype");
309 + assert_eq!(tensor.dtype.as_str(), "float32", "{name} layout dtype");
310 assert_eq!(
311 tensor.data.into_raw_view().unwrap(),
312 expected,
@@ -386,10 +319,9 @@ mod tests {
319 fn batched_wire_value_at_drops_first_axis() {
320 let data = vec![1.0_f32, 2.0, 3.0, 4.0];
321 let expected_ptr = data.as_ptr().cast::<u8>().wrapping_add(2 * size_of::<f32>());
389 - let value = KwargValue::tensor(
390 - TensorKind::F32,
391 - WireTensor::from_f32(vec![2, 2], data).unwrap(),
392 - );
322 + let value = KwargValue::Tensor {
323 + wire: WireTensor::from_f32(vec![2, 2], data).unwrap(),
324 + };
325
326 let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
327 panic!("expected tensor");
@@ -406,10 +338,9 @@ mod tests {
338
339 #[test]
340 fn flat_wire_value_range_keeps_first_axis() {
409 - let value = KwargValue::tensor(
410 - TensorKind::U32,
411 - WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(),
412 - );
341 + let value = KwargValue::Tensor {
342 + wire: WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(),
343 + };
344
345 let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
346 panic!("expected tensor");
@@ -424,10 +355,9 @@ mod tests {
355
356 #[test]
357 fn flat_range_for_index_uses_size_tensor() {
427 - let sizes = KwargValue::tensor(
428 - TensorKind::I64,
429 - WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(),
430 - );
358 + let sizes = KwargValue::Tensor {
359 + wire: WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(),
360 + };
361
362 assert_eq!(
363 flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(),
@@ -437,11 +367,12 @@ mod tests {
367
368 #[test]
369 fn slice_first_axis_range_errors_on_shape_data_mismatch() {
440 - let tensor = WireTensor::from_raw("float32", vec![2, 2], vec![0; 3 * size_of::<f32>()]);
441 - let error = slice_first_axis_range(&tensor, size_of::<f32>(), 0, 1, true).unwrap_err();
370 + let tensor =
371 + WireTensor::from_raw(TensorDtype::F32, vec![2, 2], vec![0; 3 * size_of::<f32>()]);
372 + let error = slice_first_axis_range(&tensor, 0, 1, true).unwrap_err();
373
374 assert!(
444 - matches!(error, Error::Multimodal(message) if message.contains("expects 16 bytes"))
375 + matches!(error, Error::Multimodal(message) if message.contains("does not match expected 16"))
376 );
377 }
378
@@ -455,7 +386,7 @@ mod tests {
386 panic!("expected tensor");
387 };
388
458 - assert_eq!(tensor.dtype, "bfloat16");
389 + assert_eq!(tensor.dtype.as_str(), "bfloat16");
390 assert_eq!(tensor.shape, vec![2]);
391 assert_eq!(tensor.data.into_raw_view().unwrap().len(), 4);
392 }
@@ -470,7 +401,7 @@ mod tests {
401 panic!("expected tensor");
402 };
403
473 - assert_eq!(tensor.dtype, "float16");
404 + assert_eq!(tensor.dtype.as_str(), "float16");
405 assert_eq!(tensor.shape, vec![2]);
406 assert_eq!(tensor.data.into_raw_view().unwrap().len(), 4);
407 }
rust/src/engine-core-client/src/protocol/dtype.rs
+231 -1
@@ -1,7 +1,191 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +use std::fmt;
5 +use std::str::FromStr;
6 +
7 use serde::{Deserialize, Serialize};
8 +use serde_with::{DeserializeFromStr, SerializeDisplay};
9 +
10 +/// Element types accepted by Python's tensor decoder.
11 +#[derive(Debug, Clone, Copy, PartialEq, Eq, SerializeDisplay, DeserializeFromStr)]
12 +pub enum TensorDtype {
13 + Bool,
14 + U8,
15 + I8,
16 + U16,
17 + I16,
18 + U32,
19 + I32,
20 + U64,
21 + I64,
22 + F16,
23 + Bf16,
24 + F32,
25 + F64,
26 +}
27 +
28 +impl TensorDtype {
29 + pub const fn as_str(self) -> &'static str {
30 + match self {
31 + Self::Bool => "bool",
32 + Self::U8 => "uint8",
33 + Self::I8 => "int8",
34 + Self::U16 => "uint16",
35 + Self::I16 => "int16",
36 + Self::U32 => "uint32",
37 + Self::I32 => "int32",
38 + Self::U64 => "uint64",
39 + Self::I64 => "int64",
40 + Self::F16 => "float16",
41 + Self::Bf16 => "bfloat16",
42 + Self::F32 => "float32",
43 + Self::F64 => "float64",
44 + }
45 + }
46 +
47 + pub const fn element_size(self) -> usize {
48 + match self {
49 + Self::Bool | Self::U8 | Self::I8 => 1,
50 + Self::U16 | Self::I16 | Self::F16 | Self::Bf16 => 2,
51 + Self::U32 | Self::I32 | Self::F32 => 4,
52 + Self::U64 | Self::I64 | Self::F64 => 8,
53 + }
54 + }
55 +}
56 +
57 +impl fmt::Display for TensorDtype {
58 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59 + f.write_str(self.as_str())
60 + }
61 +}
62 +
63 +impl FromStr for TensorDtype {
64 + type Err = String;
65 +
66 + fn from_str(value: &str) -> Result<Self, Self::Err> {
67 + match value {
68 + "bool" => Ok(Self::Bool),
69 + "uint8" => Ok(Self::U8),
70 + "int8" => Ok(Self::I8),
71 + "uint16" => Ok(Self::U16),
72 + "int16" => Ok(Self::I16),
73 + "uint32" => Ok(Self::U32),
74 + "int32" => Ok(Self::I32),
75 + "uint64" => Ok(Self::U64),
76 + "int64" => Ok(Self::I64),
77 + "float16" => Ok(Self::F16),
78 + "bfloat16" => Ok(Self::Bf16),
79 + "float32" => Ok(Self::F32),
80 + "float64" => Ok(Self::F64),
81 + _ => Err(format!("unsupported tensor dtype {value:?}")),
82 + }
83 + }
84 +}
85 +
86 +/// Byte order carried by a NumPy dtype descriptor.
87 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88 +pub enum Endianness {
89 + Little,
90 + Big,
91 + Native,
92 +}
93 +
94 +/// Numeric array dtype with an explicit byte-order contract.
95 +///
96 +/// Also accepts Torch type names as native-endian dtypes because prompt
97 +/// logprobs use tensor tuples with the same wire shape as NumPy arrays.
98 +/// Serializes byte order explicitly, including `=` for native byte order.
99 +/// Bfloat16 uses the Torch name after the prefix as a nonstandard descriptor.
100 +#[derive(Debug, Clone, Copy, PartialEq, Eq, SerializeDisplay, DeserializeFromStr)]
101 +pub struct NumpyDtype {
102 + pub scalar: TensorDtype,
103 + pub endianness: Endianness,
104 +}
105 +
106 +impl NumpyDtype {
107 + pub const fn little(scalar: TensorDtype) -> Self {
108 + Self {
109 + scalar,
110 + endianness: Endianness::Little,
111 + }
112 + }
113 +
114 + pub const fn native(scalar: TensorDtype) -> Self {
115 + Self {
116 + scalar,
117 + endianness: Endianness::Native,
118 + }
119 + }
120 +
121 + pub const fn big(scalar: TensorDtype) -> Self {
122 + Self {
123 + scalar,
124 + endianness: Endianness::Big,
125 + }
126 + }
127 +}
128 +
129 +impl From<TensorDtype> for NumpyDtype {
130 + fn from(scalar: TensorDtype) -> Self {
131 + Self::native(scalar)
132 + }
133 +}
134 +
135 +impl fmt::Display for NumpyDtype {
136 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137 + let prefix = match self.endianness {
138 + Endianness::Native => '=',
139 + Endianness::Little => '<',
140 + Endianness::Big => '>',
141 + };
142 + let code = match self.scalar {
143 + TensorDtype::Bool => "b1",
144 + TensorDtype::U8 => "u1",
145 + TensorDtype::I8 => "i1",
146 + TensorDtype::U16 => "u2",
147 + TensorDtype::I16 => "i2",
148 + TensorDtype::U32 => "u4",
149 + TensorDtype::I32 => "i4",
150 + TensorDtype::U64 => "u8",
151 + TensorDtype::I64 => "i8",
152 + TensorDtype::F16 => "f2",
153 + TensorDtype::F32 => "f4",
154 + TensorDtype::F64 => "f8",
155 + TensorDtype::Bf16 => "bfloat16",
156 + };
157 + write!(f, "{prefix}{code}")
158 + }
159 +}
160 +
161 +impl FromStr for NumpyDtype {
162 + type Err = String;
163 +
164 + fn from_str(dtype: &str) -> Result<Self, Self::Err> {
165 + let (endianness, body) = match dtype.as_bytes().first().copied() {
166 + Some(b'<') => (Endianness::Little, &dtype[1..]),
167 + Some(b'>') => (Endianness::Big, &dtype[1..]),
168 + Some(b'=' | b'|') => (Endianness::Native, &dtype[1..]),
169 + _ => (Endianness::Native, dtype),
170 + };
171 + let scalar = match body {
172 + "b1" => TensorDtype::Bool,
173 + "u1" => TensorDtype::U8,
174 + "i1" => TensorDtype::I8,
175 + "u2" => TensorDtype::U16,
176 + "i2" => TensorDtype::I16,
177 + "u4" => TensorDtype::U32,
178 + "i4" => TensorDtype::I32,
179 + "u8" => TensorDtype::U64,
180 + "i8" => TensorDtype::I64,
181 + "f2" => TensorDtype::F16,
182 + "f4" => TensorDtype::F32,
183 + "f8" => TensorDtype::F64,
184 + _ => body.parse().map_err(|_| format!("unsupported dtype string {dtype:?}"))?,
185 + };
186 + Ok(Self { scalar, endianness })
187 + }
188 +}
189
190 /// Effective model dtype reported by the engine after config resolution.
191 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -26,7 +210,53 @@ impl ModelDtype {
210
211 #[cfg(test)]
212 mod tests {
29 - use super::ModelDtype;
213 + use super::{Endianness, ModelDtype, NumpyDtype, TensorDtype};
214 +
215 + #[test]
216 + fn numpy_descriptors_preserve_explicit_byte_order() {
217 + for (code, scalar) in [
218 + ("b1", TensorDtype::Bool),
219 + ("u1", TensorDtype::U8),
220 + ("i1", TensorDtype::I8),
221 + ("u2", TensorDtype::U16),
222 + ("i2", TensorDtype::I16),
223 + ("u4", TensorDtype::U32),
224 + ("i4", TensorDtype::I32),
225 + ("u8", TensorDtype::U64),
226 + ("i8", TensorDtype::I64),
227 + ("f2", TensorDtype::F16),
228 + ("bfloat16", TensorDtype::Bf16),
229 + ("f4", TensorDtype::F32),
230 + ("f8", TensorDtype::F64),
231 + ] {
232 + for (prefix, endianness) in [("<", Endianness::Little), (">", Endianness::Big)] {
233 + let wire = serde_json::json!(format!("{prefix}{code}"));
234 + let dtype: NumpyDtype = serde_json::from_value(wire.clone()).unwrap();
235 + assert_eq!(dtype, NumpyDtype { scalar, endianness });
236 + assert_eq!(serde_json::to_value(dtype).unwrap(), wire);
237 + }
238 + for wire in [
239 + format!("={code}"),
240 + format!("|{code}"),
241 + code.into(),
242 + scalar.as_str().into(),
243 + ] {
244 + let dtype: NumpyDtype = serde_json::from_value(serde_json::json!(wire)).unwrap();
245 + assert_eq!(dtype, NumpyDtype::from(scalar));
246 + assert_eq!(
247 + serde_json::to_value(dtype).unwrap(),
248 + serde_json::json!(format!("={code}"))
249 + );
250 + }
251 + }
252 + }
253 +
254 + #[test]
255 + fn rejects_unsupported_numpy_dtypes() {
256 + for wire in ["", "object", "<c8", "<bf2"] {
257 + assert!(serde_json::from_value::<NumpyDtype>(serde_json::json!(wire)).is_err());
258 + }
259 + }
260
261 #[test]
262 fn serde_uses_protocol_dtype_strings() {
rust/src/engine-core-client/src/protocol/logprobs.rs
+4 -3
@@ -14,6 +14,7 @@ use serde::{Deserialize, Deserializer, Serialize};
14
15 use self::wire::*;
16 use crate::error::{Error, Result, bail_ext_value_decode};
17 +use crate::protocol::dtype::{NumpyDtype, TensorDtype};
18 use crate::protocol::tensor::{WireArrayData, WireNdArray};
19
20 /// One token candidate and its logprob metadata for a single sequence position.
@@ -205,17 +206,17 @@ impl WireLogprobs {
206
207 Ok(Self {
208 logprob_token_ids: WireNdArray {
208 - dtype: "<i8".to_string(),
209 + dtype: NumpyDtype::little(TensorDtype::I64),
210 shape: vec![rows, cols],
211 data: WireArrayData::RawView(token_ids.into()),
212 },
213 logprobs: WireNdArray {
213 - dtype: "<f4".to_string(),
214 + dtype: NumpyDtype::little(TensorDtype::F32),
215 shape: vec![rows, cols],
216 data: WireArrayData::RawView(logprobs.into()),
217 },
218 token_ranks: WireNdArray {
218 - dtype: "<i8".to_string(),
219 + dtype: NumpyDtype::little(TensorDtype::I64),
220 shape: vec![rows],
221 data: WireArrayData::RawView(token_ranks.into()),
222 },
rust/src/engine-core-client/src/protocol/logprobs/array.rs
+17 -55
@@ -8,22 +8,9 @@ use bytes::Bytes;
8 use itertools::Itertools as _;
9
10 use crate::error::{Error, Result, ext_value_decode};
11 +use crate::protocol::dtype::{Endianness, TensorDtype};
12 use crate::protocol::tensor::{ShapeExt as _, WireNdArray};
13
13 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14 -pub(super) enum ScalarType {
15 - I32,
16 - I64,
17 - F32,
18 -}
19 -
20 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21 -pub(super) enum Endianness {
22 - Little,
23 - Big,
24 - Native,
25 -}
26 -
14 #[derive(Debug, Clone, PartialEq)]
15 pub(super) struct DecodedArray2<T> {
16 pub rows: usize,
@@ -37,7 +24,7 @@ pub(super) fn decode_array2_u32(
24 frames: &[Bytes],
25 ) -> Result<DecodedArray2<u32>> {
26 let (shape, bytes, scalar, endianness) =
40 - decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
27 + decode_array_metadata(value, field, frames, &[TensorDtype::I32, TensorDtype::I64])?;
28 if shape.len() != 2 {
29 return Err(decode_error(
30 field,
@@ -46,15 +33,15 @@ pub(super) fn decode_array2_u32(
33 }
34
35 let data = match scalar {
49 - ScalarType::I32 => decode_i32_vec(&bytes, endianness, field)?
36 + TensorDtype::I32 => decode_i32_vec(&bytes, endianness, field)?
37 .into_iter()
38 .map(|value| convert_to_u32(value, field))
39 .try_collect()?,
53 - ScalarType::I64 => decode_i64_vec(&bytes, endianness, field)?
40 + TensorDtype::I64 => decode_i64_vec(&bytes, endianness, field)?
41 .into_iter()
42 .map(|value| convert_to_u32(value, field))
43 .try_collect()?,
57 - ScalarType::F32 => unreachable!("scalar validation should reject f32"),
44 + _ => unreachable!("scalar validation should accept only i32 and i64"),
45 };
46 Ok(DecodedArray2 {
47 rows: shape[0],
@@ -69,7 +56,7 @@ pub(super) fn decode_array1_u32(
56 frames: &[Bytes],
57 ) -> Result<Vec<u32>> {
58 let (shape, bytes, scalar, endianness) =
72 - decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
59 + decode_array_metadata(value, field, frames, &[TensorDtype::I32, TensorDtype::I64])?;
60 if shape.len() != 1 {
61 return Err(decode_error(
62 field,
@@ -78,15 +65,15 @@ pub(super) fn decode_array1_u32(
65 }
66
67 let data = match scalar {
81 - ScalarType::I32 => decode_i32_vec(&bytes, endianness, field)?
68 + TensorDtype::I32 => decode_i32_vec(&bytes, endianness, field)?
69 .into_iter()
70 .map(|value| convert_to_u32(value, field))
71 .try_collect()?,
85 - ScalarType::I64 => decode_i64_vec(&bytes, endianness, field)?
72 + TensorDtype::I64 => decode_i64_vec(&bytes, endianness, field)?
73 .into_iter()
74 .map(|value| convert_to_u32(value, field))
75 .try_collect()?,
89 - ScalarType::F32 => unreachable!("scalar validation should reject f32"),
76 + _ => unreachable!("scalar validation should accept only i32 and i64"),
77 };
78 Ok(data)
79 }
@@ -97,7 +84,7 @@ pub(super) fn decode_array2_f32(
84 frames: &[Bytes],
85 ) -> Result<DecodedArray2<f32>> {
86 let (shape, bytes, _, endianness) =
100 - decode_array_metadata(value, field, frames, &[ScalarType::F32])?;
87 + decode_array_metadata(value, field, frames, &[TensorDtype::F32])?;
88 if shape.len() != 2 {
89 return Err(decode_error(
90 field,
@@ -117,15 +104,16 @@ pub(super) fn decode_array_metadata(
104 value: WireNdArray,
105 field: &str,
106 frames: &[Bytes],
120 - expected_scalars: &[ScalarType],
121 -) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)> {
107 + expected_scalars: &[TensorDtype],
108 +) -> Result<(Vec<usize>, Bytes, TensorDtype, Endianness)> {
109 let mut value = value;
123 - let (scalar, endianness) = parse_dtype(&value.dtype, field)?;
110 + let scalar = value.dtype.scalar;
111 + let endianness = value.dtype.endianness;
112 if !expected_scalars.contains(&scalar) {
113 return Err(decode_error(
114 field,
115 &format!(
128 - "expected dtype in {:?}, got {}",
116 + "expected dtype in {:?}, got {:?}",
117 expected_scalars, value.dtype
118 ),
119 ));
@@ -140,42 +128,16 @@ pub(super) fn decode_array_metadata(
128 Ok((shape, bytes, scalar, endianness))
129 }
130
143 -pub(super) fn parse_dtype(dtype: &str, field: &str) -> Result<(ScalarType, Endianness)> {
144 - let (endianness, body) = match dtype.as_bytes().first().copied() {
145 - Some(b'<') => (Endianness::Little, &dtype[1..]),
146 - Some(b'>') => (Endianness::Big, &dtype[1..]),
147 - Some(b'=') => (Endianness::Native, &dtype[1..]),
148 - Some(b'|') => (Endianness::Native, &dtype[1..]),
149 - _ => (Endianness::Native, dtype),
150 - };
151 -
152 - let scalar = match body {
153 - "i4" | "int32" => ScalarType::I32,
154 - "i8" | "int64" => ScalarType::I64,
155 - "f4" | "float32" => ScalarType::F32,
156 - _ => {
157 - return Err(decode_error(
158 - field,
159 - &format!("unsupported dtype string {dtype:?}"),
160 - ));
161 - }
162 - };
163 - Ok((scalar, endianness))
164 -}
165 -
131 pub(super) fn validate_byte_length(
132 shape: &[usize],
133 byte_len: usize,
134 field: &str,
170 - scalar: ScalarType,
135 + scalar: TensorDtype,
136 ) -> Result<()> {
137 let element_count = shape
138 .checked_numel()
139 .ok_or_else(|| decode_error(field, "shape element count overflowed usize"))?;
175 - let element_size = match scalar {
176 - ScalarType::I32 | ScalarType::F32 => 4,
177 - ScalarType::I64 => 8,
178 - };
140 + let element_size = scalar.element_size();
141 let expected = element_count
142 .checked_mul(element_size)
143 .ok_or_else(|| decode_error(field, "byte length overflowed usize"))?;
rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+23
@@ -302,6 +302,29 @@ fn decodes_big_endian_payloads() {
302 );
303 }
304
305 +#[test]
306 +fn rejects_supported_array_dtypes_in_incompatible_logprobs_fields() {
307 + for (ids_dtype, probs_dtype, field) in [
308 + ("<u4", "<f4", "logprob_token_ids"),
309 + ("<i4", "<i4", "logprobs"),
310 + ] {
311 + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
312 + Some(Value::Array(vec![
313 + ndarray_value(ids_dtype, &[1, 1], Value::Ext(3, vec![0; 4])),
314 + ndarray_value(probs_dtype, &[1, 1], Value::Ext(3, vec![0; 4])),
315 + ndarray_value("<i4", &[1], Value::Ext(3, vec![1, 0, 0, 0])),
316 + Value::Nil,
317 + ])),
318 + None,
319 + )))];
320 + let error = decode_engine_core_outputs(&frames).unwrap_err();
321 + let crate::error::Error::ExtValueDecode { message } = error else {
322 + panic!("expected ExtValueDecode");
323 + };
324 + assert!(message.starts_with(&format!("new_logprobs.{field}: expected dtype")));
325 + }
326 +}
327 +
328 #[test]
329 fn rejects_non_none_cu_num_generated_tokens() {
330 let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
rust/src/engine-core-client/src/protocol/multimodal.rs
+21 -1
@@ -9,6 +9,26 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple};
9
10 use super::tensor::WireTensor;
11
12 +/// Modalities represented by engine-ready multimodal features.
13 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
14 +#[serde(rename_all = "snake_case")]
15 +pub enum MmModality {
16 + Image,
17 + Audio,
18 + Video,
19 +}
20 +
21 +impl MmModality {
22 + /// The wire name, matching Python's modality strings.
23 + pub fn as_str(self) -> &'static str {
24 + match self {
25 + Self::Image => "image",
26 + Self::Audio => "audio",
27 + Self::Video => "video",
28 + }
29 + }
30 +}
31 +
32 /// Multimodal feature payload accepted from higher-level frontend code.
33 ///
34 /// Original Python definition:
@@ -32,7 +52,7 @@ pub struct MmFeatureSpec {
52 pub data: Option<MmKwargsItem>,
53
54 /// The input modality, e.g., `"image"`, `"audio"`, `"video"`.
35 - pub modality: String,
55 + pub modality: MmModality,
56
57 /// The hash for caching encoder outputs (with LoRA prefix if applicable).
58 pub identifier: String,
rust/src/engine-core-client/src/protocol/request.rs
+7 -5
@@ -158,8 +158,10 @@ mod tests {
158 use rmpv::Value;
159
160 use super::*;
161 + use crate::protocol::dtype::TensorDtype;
162 use crate::protocol::multimodal::{
162 - MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, PlaceholderRange,
163 + MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, MmModality,
164 + PlaceholderRange,
165 };
166 use crate::protocol::sampling::EngineCoreSamplingParams;
167 use crate::protocol::tensor::{WireArrayData, WireTensor};
@@ -212,7 +214,7 @@ mod tests {
214 "inline".to_string(),
215 MmFieldElem {
216 data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
215 - "uint8",
217 + TensorDtype::U8,
218 vec![inline.len()],
219 inline,
220 ))),
@@ -225,7 +227,7 @@ mod tests {
227 data: Some(MmKwargValue::List(vec![
228 MmKwargValue::Int(7),
229 MmKwargValue::Tensor(WireTensor::from_raw(
228 - "uint8",
230 + TensorDtype::U8,
231 vec![first_aux.len()],
232 first_aux,
233 )),
@@ -234,13 +236,13 @@ mod tests {
236 },
237 ),
238 ])),
237 - modality: "image".to_string(),
239 + modality: MmModality::Image,
240 identifier: "id".to_string(),
241 mm_position: PlaceholderRange {
242 offset: 0,
243 length: second_aux.len(),
244 is_embed: Some(WireTensor::from_raw(
243 - "bool",
245 + TensorDtype::Bool,
246 vec![second_aux.len()],
247 second_aux,
248 )),
rust/src/engine-core-client/src/protocol/tensor.rs
+69 -36
@@ -6,9 +6,12 @@ use bytes::Bytes;
6 use enum_as_inner::EnumAsInner;
7 use half::{bf16, f16};
8 use rmpv::Value;
9 +use serde::de::DeserializeOwned;
10 use serde::{Deserialize, Deserializer, Serialize, Serializer};
11 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
12
13 +use super::dtype::{NumpyDtype, TensorDtype};
14 +
15 /// Tensors and ndarrays are encoded with this extension type in Python.
16 ///
17 /// Original Python definition:
@@ -63,20 +66,20 @@ impl [usize] {
66 /// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L237-L273>
67 /// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L389-L425>
68 #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
66 -pub struct WireNdArray {
67 - pub dtype: String,
69 +pub struct WireArray<Dtype: Serialize + DeserializeOwned> {
70 + pub dtype: Dtype,
71 pub shape: Vec<usize>,
72 pub data: WireArrayData,
73 }
74
72 -impl WireNdArray {
75 +impl<Dtype: From<TensorDtype> + Serialize + DeserializeOwned> WireArray<Dtype> {
76 /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
77 ///
78 /// Takes ownership of the backing buffer without copying its data.
79 pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
80 validate_element_count(&shape, data.len())?;
81 Ok(Self::from_raw_bytes(
79 - "float32",
82 + TensorDtype::F32,
83 shape,
84 bytes_from_pod_vec(data),
85 ))
@@ -88,7 +91,7 @@ impl WireNdArray {
91 pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
92 validate_element_count(&shape, data.len())?;
93 Ok(Self::from_raw_bytes(
91 - "float16",
94 + TensorDtype::F16,
95 shape,
96 bytes_from_pod_vec(data),
97 ))
@@ -100,7 +103,7 @@ impl WireNdArray {
103 pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
104 validate_element_count(&shape, data.len())?;
105 Ok(Self::from_raw_bytes(
103 - "bfloat16",
106 + TensorDtype::Bf16,
107 shape,
108 bytes_from_pod_vec(data),
109 ))
@@ -112,7 +115,7 @@ impl WireNdArray {
115 pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
116 validate_element_count(&shape, data.len())?;
117 Ok(Self::from_raw_bytes(
115 - "int64",
118 + TensorDtype::I64,
119 shape,
120 bytes_from_pod_vec(data),
121 ))
@@ -124,7 +127,7 @@ impl WireNdArray {
127 pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
128 validate_element_count(&shape, data.len())?;
129 Ok(Self::from_raw_bytes(
127 - "uint32",
130 + TensorDtype::U32,
131 shape,
132 bytes_from_pod_vec(data),
133 ))
@@ -137,7 +140,7 @@ impl WireNdArray {
140 pub fn from_bool(shape: Vec<usize>, data: Vec<bool>) -> Result<Self, String> {
141 validate_element_count(&shape, data.len())?;
142 Ok(Self {
140 - dtype: "bool".to_string(),
143 + dtype: TensorDtype::Bool.into(),
144 shape,
145 data: WireArrayData::RawView(Bytes::from(
146 data.into_iter().map(u8::from).collect::<Vec<_>>(),
@@ -149,12 +152,12 @@ impl WireNdArray {
152 ///
153 /// Use this as an escape hatch when the caller already owns bytes that
154 /// match the requested `dtype` and `shape`.
152 - pub fn from_raw(dtype: impl Into<String>, shape: Vec<usize>, data: Vec<u8>) -> Self {
155 + pub fn from_raw(dtype: impl Into<Dtype>, shape: Vec<usize>, data: Vec<u8>) -> Self {
156 Self::from_raw_bytes(dtype, shape, Bytes::from(data))
157 }
158
159 /// 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 {
160 + pub fn from_raw_bytes(dtype: impl Into<Dtype>, shape: Vec<usize>, data: Bytes) -> Self {
161 Self {
162 dtype: dtype.into(),
163 shape,
@@ -192,7 +195,10 @@ fn validate_element_count(shape: &[usize], len: usize) -> Result<(), String> {
195 ///
196 /// This is the same wire shape as [`WireNdArray`]; multimodal request payloads
197 /// use it for `torch.Tensor` values.
195 -pub type WireTensor = WireNdArray;
198 +pub type WireTensor = WireArray<TensorDtype>;
199 +
200 +/// Numeric array with a typed NumPy descriptor or native-endian Torch dtype.
201 +pub type WireNdArray = WireArray<NumpyDtype>;
202
203 /// Python array/tensor payload reference inside [`WireNdArray`].
204 ///
@@ -211,6 +217,27 @@ pub enum WireArrayData {
217 RawView(Bytes),
218 }
219
220 +impl WireTensor {
221 + /// Validate the shape and storage of an inline tensor.
222 + pub fn validate_inline(&self) -> Result<(), String> {
223 + let expected = self
224 + .shape
225 + .checked_numel()
226 + .and_then(|numel| numel.checked_mul(self.dtype.element_size()))
227 + .ok_or_else(|| format!("tensor shape or byte length overflows: {:?}", self.shape))?;
228 + match &self.data {
229 + WireArrayData::RawView(bytes) if bytes.len() == expected => Ok(()),
230 + WireArrayData::RawView(bytes) => Err(format!(
231 + "tensor byte length {} does not match expected {expected}",
232 + bytes.len()
233 + )),
234 + WireArrayData::AuxIndex(index) => {
235 + Err(format!("inline tensor references auxiliary frame {index}"))
236 + }
237 + }
238 + }
239 +}
240 +
241 impl WireArrayData {
242 /// Resolve an auxiliary-frame reference into owned raw bytes.
243 pub(crate) fn resolve_aux_frame(&mut self, frames: &[Bytes]) -> Result<(), String> {
@@ -288,6 +315,21 @@ impl Serialize for WireArrayData {
315 mod tests {
316 use super::*;
317
318 + #[test]
319 + fn tensor_and_ndarray_preserve_their_dtype_wire_contracts() {
320 + let tensor = WireTensor::from_f32(vec![1], vec![1.0]).unwrap();
321 + let bytes = rmp_serde::to_vec_named(&tensor).unwrap();
322 + let wire: Value = rmp_serde::from_slice(&bytes).unwrap();
323 + assert_eq!(wire.as_array().unwrap()[0], Value::from("float32"));
324 +
325 + let array: WireNdArray = rmp_serde::from_slice(&bytes).unwrap();
326 + assert_eq!(array.dtype, NumpyDtype::native(TensorDtype::F32));
327 + let bytes = rmp_serde::to_vec_named(&array).unwrap();
328 + let wire: Value = rmp_serde::from_slice(&bytes).unwrap();
329 + assert_eq!(wire.as_array().unwrap()[0], Value::from("=f4"));
330 + assert!(rmp_serde::from_slice::<WireTensor>(&bytes).is_err());
331 + }
332 +
333 #[test]
334 fn raw_view_serializes_as_msgpack_ext() {
335 let bytes = vec![1, 2, 3, 4];
@@ -307,8 +349,8 @@ mod tests {
349 fn constructors_build_raw_view_tensors() {
350 let f32_data = vec![1.0, 2.5];
351 let f32_data_ptr = f32_data.as_ptr().cast::<u8>();
310 - let f32_tensor = WireNdArray::from_f32(vec![2], f32_data).unwrap();
311 - assert_eq!(f32_tensor.dtype, "float32");
352 + let f32_tensor = WireTensor::from_f32(vec![2], f32_data).unwrap();
353 + assert_eq!(f32_tensor.dtype, TensorDtype::F32);
354 assert_eq!(f32_tensor.shape, vec![2]);
355 let f32_raw_view = f32_tensor.data.into_raw_view().expect("raw view");
356 assert_eq!(f32_raw_view.as_ptr(), f32_data_ptr);
@@ -318,58 +360,49 @@ mod tests {
360 );
361
362 let f16_tensor =
321 - WireNdArray::from_f16(vec![2], vec![f16::from_f32(1.0), f16::from_f32(2.5)]).unwrap();
322 - assert_eq!(f16_tensor.dtype, "float16");
363 + WireTensor::from_f16(vec![2], vec![f16::from_f32(1.0), f16::from_f32(2.5)]).unwrap();
364 + assert_eq!(f16_tensor.dtype, TensorDtype::F16);
365 assert_eq!(f16_tensor.shape, vec![2]);
366 assert_eq!(f16_tensor.data.into_raw_view().expect("raw view").len(), 4);
367
368 let bf16_tensor =
327 - WireNdArray::from_bf16(vec![2], vec![bf16::from_f32(1.0), bf16::from_f32(2.5)])
328 - .unwrap();
329 - assert_eq!(bf16_tensor.dtype, "bfloat16");
369 + WireTensor::from_bf16(vec![2], vec![bf16::from_f32(1.0), bf16::from_f32(2.5)]).unwrap();
370 + assert_eq!(bf16_tensor.dtype, TensorDtype::Bf16);
371 assert_eq!(bf16_tensor.shape, vec![2]);
372 assert_eq!(bf16_tensor.data.into_raw_view().expect("raw view").len(), 4);
373
333 - let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap();
334 - assert_eq!(i64_tensor.dtype, "int64");
374 + let i64_tensor = WireTensor::from_i64(vec![1], vec![-7]).unwrap();
375 + assert_eq!(i64_tensor.dtype, TensorDtype::I64);
376 assert_eq!(
377 i64_tensor.data.into_raw_view().expect("raw view").as_ref(),
378 (-7_i64).to_ne_bytes().as_ref()
379 );
380
340 - let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap();
341 - assert_eq!(u32_tensor.dtype, "uint32");
381 + let u32_tensor = WireTensor::from_u32(vec![1], vec![42]).unwrap();
382 + assert_eq!(u32_tensor.dtype, TensorDtype::U32);
383 assert_eq!(
384 u32_tensor.data.into_raw_view().expect("raw view").as_ref(),
385 42_u32.to_ne_bytes().as_ref()
386 );
387
347 - let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap();
348 - assert_eq!(bool_tensor.dtype, "bool");
388 + let bool_tensor = WireTensor::from_bool(vec![2], vec![false, true]).unwrap();
389 + assert_eq!(bool_tensor.dtype, TensorDtype::Bool);
390 assert_eq!(
391 bool_tensor.data.into_raw_view().expect("raw view"),
392 vec![0, 1]
393 );
353 -
354 - let raw_tensor = WireNdArray::from_raw("custom", vec![3], vec![1, 2, 3]);
355 - assert_eq!(raw_tensor.dtype, "custom");
356 - assert_eq!(raw_tensor.shape, vec![3]);
357 - assert_eq!(
358 - raw_tensor.data.into_raw_view().expect("raw view"),
359 - vec![1, 2, 3]
360 - );
394 }
395
396 #[test]
397 fn constructors_validate_shape_product() {
365 - let err = WireNdArray::from_f32(vec![2, 2], vec![1.0, 2.0]).unwrap_err();
398 + let err = WireTensor::from_f32(vec![2, 2], vec![1.0, 2.0]).unwrap_err();
399 assert!(err.contains("does not match shape"));
400 }
401
402 #[test]
403 fn resolve_aux_frame_shares_bytes_using_one_based_message_index() {
371 - let mut tensor = WireNdArray {
372 - dtype: "float32".to_string(),
404 + let mut tensor = WireTensor {
405 + dtype: TensorDtype::F32,
406 shape: vec![2],
407 data: WireArrayData::AuxIndex(1),
408 };
rust/src/engine-core-client/src/tests/client.rs
+3 -3
@@ -22,8 +22,8 @@ use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, Zmq
22 use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
23 use crate::protocol::logprobs::MaybeWireLogprobs;
24 use crate::protocol::multimodal::{
25 - MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange,
26 - SliceSpec,
25 + MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmModality, MmSlice,
26 + PlaceholderRange, SliceSpec,
27 };
28 use crate::protocol::output::{
29 DpControlMessage, DpControlOutput, EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
@@ -191,7 +191,7 @@ fn sample_multimodal_request() -> EngineCoreRequest {
191 }),
192 },
193 )])),
194 - modality: "image".to_string(),
194 + modality: MmModality::Image,
195 identifier: "mm-cache-key".to_string(),
196 mm_position: PlaceholderRange {
197 offset: 1,
rust/src/server/src/grpc/tests.rs
+1 -1
@@ -728,7 +728,7 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() {
728 assert_eq!(features.len(), 2);
729
730 for (feature, identifier) in features.iter().zip(["image-1", "image-2"]) {
731 - assert_eq!(feature.modality, "image");
731 + assert_eq!(feature.modality.as_str(), "image");
732 assert_eq!(feature.identifier, identifier);
733 assert!(feature.mm_position.length > 1);
734 assert_eq!(
rust/src/server/src/routes/tests.rs
+2 -2
@@ -2736,7 +2736,7 @@ async fn non_stream_chat_image_url_reaches_engine_mm_features() {
2736
2737 let features = request.mm_features.as_ref().expect("multimodal features");
2738 assert_eq!(features.len(), 1);
2739 - assert_eq!(features[0].modality, "image");
2739 + assert_eq!(features[0].modality.as_str(), "image");
2740 assert_eq!(features[0].identifier, "image-1");
2741 assert!(features[0].mm_position.length > 0);
2742 assert!(features[0].mm_position.is_embed.is_some());
@@ -2799,7 +2799,7 @@ async fn non_stream_chat_rejects_when_image_count_exceeds_limit_mm_per_prompt()
2799 default_stream_output_specs(),
2800 Arc::new(FakeChatBackend::with_multimodal_model_info(
2801 qwen_multimodal_model_info_with_limits(std::collections::HashMap::from([(
2802 - vllm_chat::multimodal::MmLimitModality::Image,
2802 + vllm_chat::multimodal::MmModality::Image,
2803 vllm_chat::multimodal::MmLimitSpec::Count(1),
2804 )])),
2805 )),
rust/src/text/src/lower.rs
+4 -2
@@ -318,7 +318,9 @@ mod tests {
318 use std::collections::{BTreeSet, HashMap};
319
320 use serial_test::file_serial;
321 - use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, PlaceholderRange};
321 + use vllm_engine_core_client::protocol::multimodal::{
322 + MmFeatureSpec, MmModality, PlaceholderRange,
323 + };
324 use vllm_tokenizer::test_utils::TestTokenizer;
325
326 use super::*;
@@ -720,7 +722,7 @@ mod tests {
722 fn lower_text_request_moves_multimodal_features_to_generate_request() {
723 let features = vec![MmFeatureSpec {
724 data: None,
723 - modality: "image".to_string(),
725 + modality: MmModality::Image,
726 identifier: "image-1".to_string(),
727 mm_position: PlaceholderRange {
728 offset: 2,