diff --git a/crates/engine_zmq_client/src/codec/tensor.rs b/crates/engine_zmq_client/src/codec/tensor.rs index 638078ef6..9bb0fc7de 100644 --- a/crates/engine_zmq_client/src/codec/tensor.rs +++ b/crates/engine_zmq_client/src/codec/tensor.rs @@ -134,6 +134,42 @@ impl WireNdArray { Self::from_raw_bytes(dtype, shape, Bytes::from(data)) } + /// Build from little-endian `float32` bytes, casting each element to + /// `dtype`. Mirrors the model-dtype cast the engine's own frontend applies + /// to floating multimodal tensors before they reach the model. + pub fn from_f32_bytes_cast( + dtype: super::dtype::ModelDtype, + shape: Vec, + data: &[u8], + ) -> std::result::Result { + use super::dtype::ModelDtype; + if !data.len().is_multiple_of(4) { + return Err(format!( + "float32 buffer length {} is not a multiple of 4", + data.len() + )); + } + validate_element_count(&shape, data.len() / 4)?; + let floats = data + .as_chunks::<4>() + .0 + .iter() + .map(|c| f32::from_le_bytes(*c)); + Ok(match dtype { + ModelDtype::Float32 => Self::from_raw("float32", shape, data.to_vec()), + ModelDtype::Float16 => Self::from_raw_bytes( + "float16", + shape, + bytes_from_pod_vec(floats.map(f16::from_f32).collect::>()), + ), + ModelDtype::BFloat16 => Self::from_raw_bytes( + "bfloat16", + shape, + bytes_from_pod_vec(floats.map(bf16::from_f32).collect::>()), + ), + }) + } + /// Build from an owned immutable raw-view buffer. pub fn from_raw_bytes(dtype: impl Into, shape: Vec, data: Bytes) -> Self { Self { diff --git a/crates/engine_zmq_client/src/protocol/vllm/mod.rs b/crates/engine_zmq_client/src/protocol/vllm/mod.rs index 4edb05a41..933a0b32a 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/mod.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/mod.rs @@ -5,16 +5,16 @@ //! shapes, field order, and `array_like` positional-tuple encoding are the wire //! contract with Python `EngineCoreProc` — do not reorder fields. //! -//! Text generation and structured outputs (guided decoding) are typed fully. -//! Multimodal features and pooling params are carried as -//! [`crate::codec::OpaqueValue`] for now — they serialize as `nil` on the text -//! path and get strongly typed in the multimodal phase. +//! Text generation, structured outputs (guided decoding), and multimodal +//! features are typed fully. Pooling params and prompt embeds are carried as +//! [`crate::codec::OpaqueValue`] — they serialize as `nil` on supported paths. // The startup handshake is engine-neutral (TokenSpeed speaks the same // protocol); re-exported here so existing `vllm::handshake` paths keep working. pub use crate::protocol::handshake; pub mod logprobs; pub mod lora; +pub mod multimodal; pub mod output; pub mod request; pub mod sampling; diff --git a/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs b/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs new file mode 100644 index 000000000..68ee336d7 --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs @@ -0,0 +1,322 @@ +// Ported from the Apache-2.0 reference `vllm-engine-core-client` +// (vllm-project/vllm): protocol/multimodal.rs. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::codec::tensor::WireTensor; + +/// Multimodal feature payload carried at `EngineCoreRequest.mm_features`. +/// +/// Python: `list[MultiModalFeatureSpec] | None` (`vllm/v1/engine/__init__.py`). +pub type MmFeatures = Vec; + +/// A single multimodal input with its processed data and metadata. A request +/// containing multiple multimodal items carries one `MmFeatureSpec` per item. +/// +/// Python: `MultiModalFeatureSpec` (`vllm/multimodal/inputs.py`), a dataclass — +/// encodes as a string-keyed msgpack map. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MmFeatureSpec { + /// Processed multimodal data for this item. `None` only when the engine's + /// receiver cache already holds `identifier` — an external frontend that + /// does not mirror that cache protocol must always send the full data. + pub data: Option, + + /// The input modality, e.g. `"image"`, `"audio"`, `"video"`. + pub modality: String, + + /// The hash for caching encoder outputs (with LoRA prefix if applicable). + pub identifier: String, + + /// The location of the `modality` tokens corresponding to this item in + /// the prompt. + pub mm_position: PlaceholderRange, + + /// The hash for caching processor outputs (without LoRA prefix). + #[serde(default)] + pub mm_hash: Option, +} + +/// Placeholder location information for one multimodal item. +/// +/// Python: `PlaceholderRange` (`vllm/multimodal/inputs.py`), a dataclass — +/// encodes as a string-keyed msgpack map. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct PlaceholderRange { + /// The start index of the placeholder in the prompt. + pub offset: usize, + + /// The length of the placeholder. + pub length: usize, + + /// A boolean mask of shape `(length,)` indicating which positions between + /// `offset` and `offset + length` receive embeddings. `None` means all. + #[serde(default)] + pub is_embed: Option, +} + +/// Processed keyword arguments for a single multimodal item, keyed by model +/// kwarg name (e.g. `pixel_values`). +/// +/// Python: `MultiModalKwargsItem` (`vllm/multimodal/inputs.py`) — encoded by +/// the serializer hooks as a string-keyed map. +pub type MmKwargsItem = BTreeMap; + +/// One processed keyword argument of a `MmKwargsItem`. +/// +/// Python: `MultiModalFieldElem` (`vllm/multimodal/inputs.py`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MmFieldElem { + /// The keyword argument value passed to the model. `None` only when the + /// item is cached engine-side (see [`MmFeatureSpec::data`]). + pub data: Option, + + /// How this field's values combine with other items' for batching. + pub field: MmField, +} + +/// Processed multimodal keyword argument value (Python `NestedTensors`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MmKwargValue { + Tensor(WireTensor), + Int(i64), + Float(f64), + List(Vec), +} + +/// How to interpret tensor data belonging to a keyword argument. +/// +/// Wire form is a 2-tuple `(factory_name, kwargs_map)` with factory names +/// `"batched"`, `"flat"`, `"shared"` — the serializer's +/// `MMF_CLASS_TO_FACTORY` encoding. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(try_from = "MmFieldWire", into = "MmFieldWire")] +pub enum MmField { + Batched(MmBatchedField), + Flat(MmFlatField), + Shared(MmSharedField), +} + +/// Python `MultiModalFieldConfig.batched`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmBatchedField { + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python `MultiModalFieldConfig.flat` / `flat_from_sizes`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmFlatField { + /// For each multimodal item, a slice (`dim=0`) or a tuple of slices + /// (`dim>0`) that extracts the data corresponding to it. + pub slices: Vec, + + /// The dimension to extract data from, default 0. + pub dim: i32, + + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python `MultiModalFieldConfig.shared`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmSharedField { + pub batch_size: usize, + + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python slice encoded as `(start, stop, step)`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] +pub struct SliceSpec { + pub start: Option, + pub stop: Option, + pub step: Option, +} + +/// A single slice or a tuple of slices used by [`MmFlatField`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MmSlice { + Slice(SliceSpec), + Slices(Vec), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] +struct MmFieldWire { + name: String, + inner: MmFieldWireInner, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +enum MmFieldWireInner { + Batched(MmBatchedField), + Flat(MmFlatField), + Shared(MmSharedField), +} + +impl TryFrom for MmField { + type Error = String; + + fn try_from(value: MmFieldWire) -> Result { + match (value.name.as_str(), value.inner) { + ("batched", MmFieldWireInner::Batched(kwargs)) => Ok(Self::Batched(kwargs)), + ("flat", MmFieldWireInner::Flat(kwargs)) => Ok(Self::Flat(kwargs)), + ("shared", MmFieldWireInner::Shared(kwargs)) => Ok(Self::Shared(kwargs)), + (name, _) => Err(format!( + "mismatched or unknown multimodal field factory {name:?}" + )), + } + } +} + +impl From for MmFieldWire { + fn from(value: MmField) -> Self { + match value { + MmField::Batched(kwargs) => Self { + name: "batched".to_string(), + inner: MmFieldWireInner::Batched(kwargs), + }, + MmField::Flat(kwargs) => Self { + name: "flat".to_string(), + inner: MmFieldWireInner::Flat(kwargs), + }, + MmField::Shared(kwargs) => Self { + name: "shared".to_string(), + inner: MmFieldWireInner::Shared(kwargs), + }, + } + } +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use rmpv::Value; + + use super::*; + use crate::codec::encode_msgpack; + + fn encode_value(value: &T) -> Value { + let bytes = encode_msgpack(value).expect("encode value"); + rmpv::decode::read_value(&mut Cursor::new(bytes)).expect("decode value") + } + + #[test] + fn field_serializes_to_python_factory_tuple() { + let field = MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(1200), + step: None, + })], + dim: 0, + keep_on_cpu: false, + }); + + let value = encode_value(&field); + let Value::Array(items) = value else { + panic!("field should encode as a 2-tuple array"); + }; + assert_eq!(items.len(), 2); + assert_eq!(items[0].as_str(), Some("flat")); + let Value::Map(kwargs) = &items[1] else { + panic!("field kwargs should encode as a map"); + }; + for key in ["slices", "dim", "keep_on_cpu"] { + assert!( + kwargs.iter().any(|(k, _)| k.as_str() == Some(key)), + "missing kwarg {key}" + ); + } + } + + #[test] + fn field_round_trips_python_factory_tuple() { + for field in [ + MmField::Batched(MmBatchedField { keep_on_cpu: true }), + MmField::Shared(MmSharedField { + batch_size: 4, + keep_on_cpu: false, + }), + ] { + let encoded = encode_msgpack(&field).expect("encode field"); + let decoded: MmField = rmp_serde::from_slice(&encoded).expect("decode field"); + assert_eq!(decoded, field); + } + } + + #[test] + fn feature_spec_serializes_as_named_map_with_tensor_ext() { + let mut item = MmKwargsItem::new(); + item.insert( + "pixel_values".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + WireTensor::from_f32(vec![2, 3], vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + .expect("tensor built"), + )), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + ); + let spec = MmFeatureSpec { + data: Some(item), + modality: "image".to_string(), + identifier: "abc123".to_string(), + mm_position: PlaceholderRange { + offset: 5, + length: 6, + is_embed: None, + }, + mm_hash: Some("abc123".to_string()), + }; + + let value = encode_value(&spec); + let Value::Map(entries) = value else { + panic!("feature spec should encode as a map"); + }; + for key in ["data", "modality", "identifier", "mm_position", "mm_hash"] { + assert!( + entries.iter().any(|(k, _)| k.as_str() == Some(key)), + "missing key {key}" + ); + } + + // The tensor payload must reach the wire as the 3-tuple + // (dtype, shape, ext-3 raw view). + let data = entries + .iter() + .find(|(k, _)| k.as_str() == Some("data")) + .map(|(_, v)| v) + .expect("data present"); + let Value::Map(kwargs) = data else { + panic!("kwargs item should encode as a map"); + }; + let Value::Map(elem) = &kwargs[0].1 else { + panic!("field elem should encode as a map"); + }; + let tensor = elem + .iter() + .find(|(k, _)| k.as_str() == Some("data")) + .map(|(_, v)| v) + .expect("elem data present"); + let Value::Array(tuple) = tensor else { + panic!("tensor should encode as (dtype, shape, data)"); + }; + assert_eq!(tuple[0].as_str(), Some("float32")); + assert!(matches!(&tuple[2], Value::Ext(3, _))); + } +} diff --git a/crates/engine_zmq_client/src/protocol/vllm/request.rs b/crates/engine_zmq_client/src/protocol/vllm/request.rs index 3e2e70959..d59ea8be7 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/request.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/request.rs @@ -1,8 +1,5 @@ // Ported from the Apache-2.0 reference `vllm-engine-core-client` // (vllm-project/vllm): protocol/request.rs. -// -// `mm_features` is carried as OpaqueValue for now (typed in the multimodal -// phase); on the text path it is `nil` and carries no aux-frame tensors. use std::collections::{BTreeMap, HashMap}; @@ -13,7 +10,7 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use crate::{ codec::OpaqueValue, - protocol::vllm::{lora, sampling::EngineCoreSamplingParams}, + protocol::vllm::{lora, multimodal::MmFeatures, sampling::EngineCoreSamplingParams}, Error, Result, }; @@ -70,8 +67,8 @@ pub struct ReasoningParserKwargs { pub struct EngineCoreRequest { pub request_id: String, pub prompt_token_ids: Option>, - /// Multimodal features (untyped for now; `nil` on the text path). - pub mm_features: Option, + /// Multimodal features, one per input item, sorted by placeholder offset. + pub mm_features: Option, pub sampling_params: Option, /// Pooling parameters, preserved in the schema but not yet strongly typed. pub pooling_params: Option, @@ -129,9 +126,10 @@ impl EngineCoreRequest { Ok(()) } - // NOTE: send-side aux-frame extraction (walking `mm_features` for large - // tensors) is added with the typed multimodal module. Text requests carry - // no tensors, so the transport send path appends no aux frames for now. + // NOTE: multimodal tensors are sent as inline ext-3 raw views in the + // request frame — valid at any size (the engine's aux-frame split is an + // encoder-side optimization only). Send-side aux extraction is a perf + // follow-up. } #[cfg(test)] diff --git a/model_gateway/src/routers/grpc/backend_client.rs b/model_gateway/src/routers/grpc/backend_client.rs index c3da6c153..f7ede0bba 100644 --- a/model_gateway/src/routers/grpc/backend_client.rs +++ b/model_gateway/src/routers/grpc/backend_client.rs @@ -13,8 +13,8 @@ use openai_protocol::{ messages::CreateMessageRequest, worker::WorkerLoadResponse, }; use smg_grpc_client::{ - common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, TokenSpeedSchedulerClient, - VllmEngineClient, + common_proto, tokenizer_bundle::StreamBundle, tokenspeed_proto, vllm_proto, + SglangSchedulerClient, TokenSpeedSchedulerClient, VllmEngineClient, }; use crate::{ @@ -27,6 +27,7 @@ use crate::{ ProtoGenerateRequest, ProtoStream, }, zmq_client::ZmqEngineClient, + MultimodalData, }, worker::RuntimeType, }; @@ -205,10 +206,10 @@ impl BackendClient { // A ZMQ backend speaks vLLM EngineCore or TokenSpeed directly; build // the native request for its runtime, mirroring the gRPC per-engine // dispatch in `GrpcClient::build_chat_request`. - Self::Zmq(client) => { - reject_zmq_multimodal(&options)?; - match client.runtime() { - RuntimeType::TokenSpeed => finish_tokenspeed_request(None, |mm| { + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let tokenspeed_mm = zmq_tokenspeed_mm(options.multimodal_inputs)?; + finish_tokenspeed_request(tokenspeed_mm, |mm| { TokenSpeedSchedulerClient::build_generate_request_from_chat( request_id, body, @@ -217,8 +218,11 @@ impl BackendClient { mm, options.tool_constraints, ) - }), - _ => finish_vllm_request(None, |mm| { + }) + } + _ => { + let vllm_mm = zmq_vllm_mm(options.multimodal_inputs)?; + finish_vllm_request(vllm_mm, |mm| { VllmEngineClient::build_generate_request_from_chat( request_id, body, @@ -227,9 +231,9 @@ impl BackendClient { mm, options.tool_constraints, ) - }), + }) } - } + }, } } @@ -247,10 +251,10 @@ impl BackendClient { } // Mirrors the gRPC per-engine dispatch: build the request natively for // the ZMQ backend's runtime (vLLM EngineCore or TokenSpeed). - Self::Zmq(client) => { - reject_zmq_multimodal(&options)?; - match client.runtime() { - RuntimeType::TokenSpeed => finish_tokenspeed_request(None, |mm| { + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let tokenspeed_mm = zmq_tokenspeed_mm(options.multimodal_inputs)?; + finish_tokenspeed_request(tokenspeed_mm, |mm| { TokenSpeedSchedulerClient::build_generate_request_from_messages( request_id, body, @@ -259,8 +263,11 @@ impl BackendClient { mm, options.tool_constraints, ) - }), - _ => finish_vllm_request(None, |mm| { + }) + } + _ => { + let vllm_mm = zmq_vllm_mm(options.multimodal_inputs)?; + finish_vllm_request(vllm_mm, |mm| { VllmEngineClient::build_generate_request_from_messages( request_id, body, @@ -269,9 +276,9 @@ impl BackendClient { mm, options.tool_constraints, ) - }), + }) } - } + }, } } @@ -344,10 +351,41 @@ impl BackendClient { } } -/// ZMQ text path does not carry multimodal inputs yet. -fn reject_zmq_multimodal(options: &GenerateRequestBuildOptions) -> Result<(), String> { - if options.multimodal_inputs.is_some() { - return Err("ZMQ backend does not support multimodal inputs yet".to_string()); - } - Ok(()) +/// Convert assembled multimodal data for a vLLM ZMQ backend. A backend/variant +/// mismatch is a gateway bug (the assembly stage should produce the backend's +/// own variant), surfaced as a build error rather than a panic. +fn zmq_vllm_mm( + inputs: Option, +) -> Result, String> { + inputs + .map(|mm| match mm { + MultimodalData::Vllm(data) => Ok(data.into_proto()), + other => Err(mm_variant_mismatch("vLLM", &other)), + }) + .transpose() +} + +/// Convert assembled multimodal data for a TokenSpeed ZMQ backend. See +/// [`zmq_vllm_mm`] for the mismatch semantics. +fn zmq_tokenspeed_mm( + inputs: Option, +) -> Result, String> { + inputs + .map(|mm| match mm { + MultimodalData::TokenSpeed(data) => Ok(data.into_proto(true)), + other => Err(mm_variant_mismatch("TokenSpeed", &other)), + }) + .transpose() +} + +/// Name the variant of a mismatched `MultimodalData` without dumping its tensor +/// payloads into the error string. +fn mm_variant_mismatch(expected: &str, got: &MultimodalData) -> String { + let got = match got { + MultimodalData::Sglang(_) => "SGLang", + MultimodalData::Vllm(_) => "vLLM", + MultimodalData::Trtllm(_) => "TRT-LLM", + MultimodalData::TokenSpeed(_) => "TokenSpeed", + }; + format!("multimodal data variant mismatch: {expected} ZMQ backend got {got} data") } diff --git a/model_gateway/src/routers/grpc/mod.rs b/model_gateway/src/routers/grpc/mod.rs index 6014c7b2f..6f374ba86 100644 --- a/model_gateway/src/routers/grpc/mod.rs +++ b/model_gateway/src/routers/grpc/mod.rs @@ -18,6 +18,7 @@ pub(crate) mod regular; pub(crate) mod router; // Used by routers/factory pub mod utils; // Used by routers/http and bindings/golang pub mod zmq_client; // ZMQ backend adapter behind the vLLM client surface +pub(crate) mod zmq_multimodal; // Proto mm inputs → EngineCore mm_features // Re-export for convenience pub use proto_wrapper::{MultimodalData, TensorBytes}; diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index ff839eec7..d08abcbb8 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -28,16 +28,19 @@ use super::{ transport::{mm_encoder_input_dtype, resolve_mm_shm_enabled, resolve_mm_shm_min_bytes}, MediaBatch, MultimodalIntermediate, PrecomputedMultimodalIntermediate, PromptBinding, }; -use crate::routers::grpc::{ - backend_client::BackendClient, - client::GrpcClient, - context::WorkerSelection, - proto_wrapper::{ - cleanup_tokenspeed_items_encoder_shm, SglangMultimodalData, TensorBytes, - TokenSpeedModality, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, TokenSpeedTensor, - TrtllmMultimodalData, VllmMultimodalData, +use crate::{ + routers::grpc::{ + backend_client::BackendClient, + client::GrpcClient, + context::WorkerSelection, + proto_wrapper::{ + cleanup_tokenspeed_items_encoder_shm, SglangMultimodalData, TensorBytes, + TokenSpeedModality, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, + TokenSpeedTensor, TrtllmMultimodalData, VllmMultimodalData, + }, + MultimodalData, }, - MultimodalData, + worker::RuntimeType, }; /// Assemble backend-specific multimodal data from the intermediate. @@ -108,7 +111,20 @@ async fn assemble_multimodal_data_impl( BackendClient::Grpc(GrpcClient::Mlx(_)) => { anyhow::bail!("MLX does not support multimodal inputs") } - BackendClient::Zmq(_) => anyhow::bail!("ZMQ backend does not support multimodal inputs"), + BackendClient::Zmq(client) => match client.runtime() { + RuntimeType::Vllm | RuntimeType::Unspecified => { + let batch = into_single_batch(intermediate, "vLLM")?; + let mut data = assemble_vllm(batch, workers)?; + // The ZMQ translate reads tensor bytes inline; this wire has no + // /dev/shm or RDMA pull on the engine side. + data.shm_enabled = false; + data.rdma_enabled = false; + Ok(MultimodalData::Vllm(data)) + } + runtime => anyhow::bail!( + "multimodal inputs are not supported over the {runtime} ZMQ backend yet" + ), + }, } } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 42c1fea1a..d6165fa91 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -17,6 +17,7 @@ use std::{ }; use engine_zmq_client::{ + codec::dtype::ModelDtype, connect_handshake, connector::{EngineCoreClient, EngineCoreStream, TokenSpeedClient, TokenSpeedStream}, protocol::{ @@ -42,6 +43,7 @@ use crate::{ routers::grpc::{ client::{ModelInfo, ServerInfo}, proto_wrapper::ProtoGenerateRequest, + zmq_multimodal, }, worker::RuntimeType, }; @@ -183,10 +185,10 @@ impl ZmqEngineClient { // (which the ZMQ path bypasses) defaults an unset value to // `max_model_len - prompt_len`. The context length comes from the // engine's ready handshake, so a connected engine is required. - let max_model_len = client + let (max_model_len, model_dtype) = client .engines() .first() - .map(|e| e.ready_response.max_model_len) + .map(|e| (e.ready_response.max_model_len, e.ready_response.dtype)) .ok_or_else(|| tonic::Status::unavailable("no connected ZMQ engine"))?; let mut streams = SelectAll::new(); for (index, sub) in fan_out_requests(*req).into_iter().enumerate() { @@ -199,7 +201,7 @@ impl ZmqEngineClient { .and_then(|sp| sp.logprobs) .filter(|&n| n > 0) .map_or(0, |n| n as usize); - let request = translate_request(sub, max_model_len) + let request = translate_request(sub, max_model_len, model_dtype) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; streams.push(VllmGenerateStream::new(stream, index as u32, top_logprobs)); @@ -755,6 +757,13 @@ fn fan_out_tokenspeed_requests( fn translate_request_tokenspeed( req: tokenspeed_proto::GenerateRequest, ) -> Result { + // The TokenSpeed ZMQ wire has no multimodal slot yet; reject loudly rather + // than silently dropping pixels (assembly also refuses upstream). + if req.mm_inputs.is_some() { + return Err( + "multimodal inputs are not supported over the TokenSpeed ZMQ backend".to_string(), + ); + } let input_ids = match req.tokenized { Some(tokenized) => tokenized.input_ids, None => { @@ -850,6 +859,7 @@ fn apply_tokenspeed_constraint( fn translate_request( req: vllm::GenerateRequest, max_model_len: u64, + model_dtype: ModelDtype, ) -> Result { let prompt_token_ids = match req.input { Some(vllm::generate_request::Input::Tokenized(tokenized)) => Some(tokenized.input_ids), @@ -860,6 +870,19 @@ fn translate_request( return Err("ZMQ mode requires pre-tokenized input; no input provided".to_string()); } }; + // Per-item mm features: the split the Python servicer performs before the + // engine happens here instead (the ZMQ path bypasses it). + let mm_features = req + .mm_inputs + .map(|mm| { + zmq_multimodal::build_mm_features( + mm, + prompt_token_ids.as_deref().unwrap_or(&[]), + model_dtype, + ) + }) + .transpose()? + .filter(|features| !features.is_empty()); let data_parallel_rank = req .data_parallel_rank .map(|rank| u32::try_from(rank).map_err(|_| format!("invalid data_parallel_rank: {rank}"))) @@ -881,7 +904,7 @@ fn translate_request( Ok(EngineCoreRequest { request_id: req.request_id, prompt_token_ids, - mm_features: None, + mm_features, sampling_params: req .sampling_params .map(|sp| translate_sampling(sp, default_max_tokens)), @@ -1640,6 +1663,15 @@ mod tests { assert_eq!(req.sampling_params.stop, None); } + #[test] + fn tokenspeed_rejects_multimodal_inputs() { + // The TokenSpeed ZMQ wire has no multimodal slot yet; reject rather than + // silently drop pixels. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.mm_inputs = Some(tokenspeed_proto::MultimodalInputs::default()); + assert!(translate_request_tokenspeed(req).is_err()); + } + #[test] fn vllm_rejects_unsupported_sampling_features() { // Prompt logprobs have no renderer merge on the ZMQ path. @@ -1649,6 +1681,7 @@ mod tests { ..Default::default() }), 4096, + ModelDtype::BFloat16, ) .expect_err("prompt logprobs rejected"); assert!(err.contains("prompt logprobs"), "{err}"); @@ -1657,7 +1690,7 @@ mod tests { #[test] fn vllm_defaults_unset_max_tokens_to_remaining_context() { let max_tokens = |sampling, max_model_len| { - translate_request(tokenized_req(sampling), max_model_len) + translate_request(tokenized_req(sampling), max_model_len, ModelDtype::BFloat16) .expect("request translated") .sampling_params .expect("sampling params present") @@ -1693,6 +1726,7 @@ mod tests { ..Default::default() }), 4096, + ModelDtype::BFloat16, ) .expect("constraint translated") .sampling_params diff --git a/model_gateway/src/routers/grpc/zmq_multimodal.rs b/model_gateway/src/routers/grpc/zmq_multimodal.rs new file mode 100644 index 000000000..584a17fbc --- /dev/null +++ b/model_gateway/src/routers/grpc/zmq_multimodal.rs @@ -0,0 +1,627 @@ +//! Proto multimodal inputs → EngineCore `mm_features` for the direct-ZMQ path. +//! +//! The gRPC servicer converts the batched proto tensors into per-item engine +//! structures Python-side (`_build_preprocessed_mm_inputs` + the engine's +//! `from_hf_inputs` split). The ZMQ path bypasses that process, so the same +//! split happens here: batched keys index row `i`, flat keys slice by the +//! cumulative sizes tensor, everything else is shared (replicated per item). +//! Floating tensors are cast to the model dtype — the engine applies no cast +//! on this path. + +use std::collections::{BTreeMap, HashMap, HashSet}; + +use bytes::Bytes; +use engine_zmq_client::{ + codec::{ + dtype::ModelDtype, + tensor::{WireArrayData, WireTensor}, + }, + protocol::vllm::multimodal::{ + MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargValue, + MmKwargsItem, MmSharedField, MmSlice, PlaceholderRange, SliceSpec, + }, +}; +use smg_grpc_client::{common_proto as common, vllm_proto as vllm}; + +/// A decoded (and dtype-cast) proto tensor ready for per-item slicing. +struct Decoded { + dtype: String, + shape: Vec, + bytes: Bytes, +} + +impl Decoded { + fn elem_size(&self) -> Result { + match self.dtype.as_str() { + "bool" => Ok(1), + "float16" | "bfloat16" => Ok(2), + "float32" | "uint32" | "int32" => Ok(4), + "int64" | "float64" => Ok(8), + other => Err(format!("unsupported multimodal tensor dtype {other:?}")), + } + } + + /// Bytes per index step along dim 0. + fn row_nbytes(&self) -> Result { + let inner: usize = self.shape.iter().skip(1).product(); + Ok(inner * self.elem_size()?) + } + + /// Zero-copy view of rows `[start, stop)` along dim 0. + fn slice_rows(&self, start: usize, stop: usize) -> Result { + let row = self.row_nbytes()?; + let (lo, hi) = (start * row, stop * row); + if hi > self.bytes.len() || start > stop { + return Err(format!( + "row slice {start}..{stop} out of bounds for tensor of {} bytes", + self.bytes.len() + )); + } + let mut shape = self.shape.clone(); + shape[0] = stop - start; + Ok(WireTensor::from_raw_bytes( + self.dtype.clone(), + shape, + self.bytes.slice(lo..hi), + )) + } + + fn whole(&self) -> WireTensor { + WireTensor::from_raw_bytes(self.dtype.clone(), self.shape.clone(), self.bytes.clone()) + } + + /// Flattened values as widened i64 (sizes tensors are int64 or uint32). + fn flat_i64(&self) -> Result, String> { + match self.dtype.as_str() { + "int64" => Ok((self.bytes.as_chunks::<8>().0.iter()) + .map(|c| i64::from_le_bytes(*c)) + .collect()), + "uint32" => Ok((self.bytes.as_chunks::<4>().0.iter()) + .map(|c| i64::from(u32::from_le_bytes(*c))) + .collect()), + other => Err(format!("flat sizes tensor has unsupported dtype {other:?}")), + } + } +} + +fn decode_tensor( + name: &str, + tensor: vllm::TensorData, + model_dtype: ModelDtype, +) -> Result { + let shape: Vec = tensor.shape.iter().map(|&d| d as usize).collect(); + let data = match tensor.payload { + Some(vllm::tensor_data::Payload::Inline(data)) => data, + Some(_) => { + return Err(format!( + "multimodal tensor {name:?} uses a non-inline payload; the ZMQ wire carries \ + tensors inline" + )); + } + None => return Err(format!("multimodal tensor {name:?} has no payload")), + }; + // Floating tensors arrive as float32 and are cast to the model dtype, + // mirroring the cast the engine's own frontend applies. + if tensor.dtype == "float32" { + let cast = WireTensor::from_f32_bytes_cast(model_dtype, shape.clone(), &data)?; + let WireArrayData::RawView(bytes) = cast.data else { + return Err(format!("cast tensor {name:?} lost its raw view")); + }; + return Ok(Decoded { + dtype: cast.dtype, + shape, + bytes, + }); + } + // Non-float32 tensors are forwarded as-is (integer/bool kwargs like the + // flat sizes or grid tensors). A floating dtype other than float32 would + // reach the engine uncast — reject it rather than produce garbage. + if matches!(tensor.dtype.as_str(), "float16" | "bfloat16" | "float64") { + return Err(format!( + "multimodal tensor {name:?} has floating dtype {:?}; the ZMQ path expects float32 \ + so it can cast to the model dtype", + tensor.dtype + )); + } + let decoded = Decoded { + dtype: tensor.dtype, + shape, + bytes: Bytes::from(data), + }; + // Guard against a truncated or oversized inline payload: the engine would + // otherwise reinterpret the raw buffer against the declared shape. + let expected = decoded + .shape + .iter() + .try_fold(decoded.elem_size()?, |acc, &d| acc.checked_mul(d)); + if expected != Some(decoded.bytes.len()) { + return Err(format!( + "multimodal tensor {name:?} has {} bytes, which does not match shape {:?} of dtype {:?}", + decoded.bytes.len(), + decoded.shape, + decoded.dtype + )); + } + Ok(decoded) +} + +/// Rename generic keys for video inputs, mirroring the servicer's `mm_key`. +fn mm_key(key: &str, is_video: bool) -> String { + if is_video && key == "pixel_values" { + "pixel_values_videos".to_string() + } else { + key.to_string() + } +} + +/// Build per-item `mm_features` from batched proto multimodal inputs. +pub(crate) fn build_mm_features( + mm: vllm::MultimodalInputs, + prompt_token_ids: &[u32], + model_dtype: ModelDtype, +) -> Result { + let num_items = mm.mm_placeholders.len(); + if num_items == 0 { + // No placeholders is only valid for a genuinely empty payload. Tensors + // or hashes with nowhere to attach means malformed input — surface it + // instead of silently building a text-only request. + if mm.pixel_values.is_some() + || !mm.model_specific_tensors.is_empty() + || !mm.mm_hashes.is_empty() + { + return Err( + "multimodal inputs carry tensors or hashes but no placeholders".to_string(), + ); + } + return Ok(Vec::new()); + } + if mm.mm_hashes.len() != num_items { + return Err(format!( + "multimodal hash count {} does not match placeholder count {num_items}", + mm.mm_hashes.len() + )); + } + let is_video = mm.modality == common::Modality::Video as i32; + let modality = if is_video { "video" } else { "image" }; + + // Decode every tensor once, applying the video key rename. + let mut tensors: BTreeMap = BTreeMap::new(); + if let Some(pixel_values) = mm.pixel_values { + tensors.insert( + mm_key("pixel_values", is_video), + decode_tensor("pixel_values", pixel_values, model_dtype)?, + ); + } + for (key, tensor) in mm.model_specific_tensors { + let decoded = decode_tensor(&key, tensor, model_dtype)?; + tensors.insert(mm_key(&key, is_video), decoded); + } + + let batched: HashSet = mm + .batched_keys + .iter() + .map(|k| mm_key(k, is_video)) + .collect(); + let flat: HashMap = mm + .flat_keys + .iter() + .map(|(k, v)| (mm_key(k, is_video), mm_key(v, is_video))) + .collect(); + let keep_on_cpu: HashSet = mm + .keep_on_cpu_keys + .iter() + .map(|k| mm_key(k, is_video)) + .collect(); + + // Split every kwarg into per-item elems. + let mut items: Vec = vec![MmKwargsItem::new(); num_items]; + for (key, decoded) in &tensors { + let on_cpu = keep_on_cpu.contains(key); + if batched.contains(key) { + if decoded.shape.first() != Some(&num_items) { + return Err(format!( + "batched tensor {key:?} has leading dim {:?}, expected {num_items} items", + decoded.shape.first() + )); + } + for (i, item) in items.iter_mut().enumerate() { + let mut tensor = decoded.slice_rows(i, i + 1)?; + tensor.shape.remove(0); + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(tensor)), + field: MmField::Batched(MmBatchedField { + keep_on_cpu: on_cpu, + }), + }, + ); + } + } else if let Some(sizes_key) = flat.get(key) { + let sizes = tensors + .get(sizes_key) + .ok_or_else(|| format!("flat sizes tensor {sizes_key:?} missing for {key:?}"))? + .flat_i64()?; + if sizes.len() != num_items { + return Err(format!( + "flat sizes tensor {sizes_key:?} has {} entries, expected {num_items}", + sizes.len() + )); + } + // Cumulative row offsets, and the full per-item slice list every + // elem carries (the engine's flat field serializes all slices). + let mut bounds = Vec::with_capacity(num_items + 1); + let mut total = 0usize; + bounds.push(total); + for size in &sizes { + let size = usize::try_from(*size) + .map_err(|_| format!("negative size in flat sizes tensor {sizes_key:?}"))?; + total = total.checked_add(size).ok_or_else(|| { + format!("flat sizes tensor {sizes_key:?} sums past usize range") + })?; + bounds.push(total); + } + if decoded.shape.first() != Some(&total) { + return Err(format!( + "flat tensor {key:?} has leading dim {:?}, expected {total} total rows", + decoded.shape.first(), + )); + } + let slices: Vec = bounds + .windows(2) + .map(|w| { + MmSlice::Slice(SliceSpec { + start: Some(w[0] as isize), + stop: Some(w[1] as isize), + step: None, + }) + }) + .collect(); + for (i, item) in items.iter_mut().enumerate() { + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + decoded.slice_rows(bounds[i], bounds[i + 1])?, + )), + field: MmField::Flat(MmFlatField { + slices: slices.clone(), + dim: 0, + keep_on_cpu: on_cpu, + }), + }, + ); + } + } else { + // Shared: the full tensor replicated per item (the servicer's + // fallback for keys in neither batched nor flat sets). + for item in &mut items { + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(decoded.whole())), + field: MmField::Shared(MmSharedField { + batch_size: num_items, + keep_on_cpu: on_cpu, + }), + }, + ); + } + } + } + + // One feature per placeholder, in prompt-offset order. + let mut features: MmFeatures = Vec::with_capacity(num_items); + for ((placeholder, item), hash) in mm + .mm_placeholders + .iter() + .zip(items) + .zip(mm.mm_hashes.iter()) + { + let offset = placeholder.offset as usize; + let length = placeholder.length as usize; + features.push(MmFeatureSpec { + data: Some(item), + modality: modality.to_string(), + identifier: hash.clone(), + mm_position: PlaceholderRange { + offset, + length, + is_embed: is_embed_mask(prompt_token_ids, offset, length, mm.im_token_id)?, + }, + mm_hash: Some(hash.clone()), + }); + } + features.sort_by_key(|f| f.mm_position.offset); + Ok(features) +} + +/// Boolean embed mask over a placeholder range: `true` where the prompt token +/// is the image token, excluding structural tokens (vision start/end markers) +/// from the embedding scatter. `None` when every position is an embed slot. +fn is_embed_mask( + prompt_token_ids: &[u32], + offset: usize, + length: usize, + im_token_id: Option, +) -> Result, String> { + // Validate the range first — it must hold regardless of whether a mask is + // needed, so an absent `im_token_id` can't skip the bounds check. + let end = offset + .checked_add(length) + .filter(|&end| end <= prompt_token_ids.len()) + .ok_or_else(|| { + format!( + "placeholder range {offset}+{length} exceeds prompt of {} tokens", + prompt_token_ids.len() + ) + })?; + let Some(im_token_id) = im_token_id else { + return Ok(None); + }; + let mask: Vec = prompt_token_ids[offset..end] + .iter() + .map(|&id| id == im_token_id) + .collect(); + if mask.iter().all(|&m| m) { + return Ok(None); + } + Ok(Some(WireTensor::from_bool(vec![length], mask)?)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn inline_tensor(shape: Vec, dtype: &str, data: Vec) -> vllm::TensorData { + vllm::TensorData { + shape, + dtype: dtype.to_string(), + payload: Some(vllm::tensor_data::Payload::Inline(data)), + } + } + + fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + fn i64_bytes(values: &[i64]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + fn placeholders(ranges: &[(u32, u32)]) -> Vec { + ranges + .iter() + .map(|&(offset, length)| vllm::PlaceholderRange { offset, length }) + .collect() + } + + fn base_inputs() -> vllm::MultimodalInputs { + vllm::MultimodalInputs { + pixel_values: Some(inline_tensor( + vec![2, 4], + "float32", + f32_bytes(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]), + )), + model_specific_tensors: Default::default(), + im_token_id: None, + mm_placeholders: placeholders(&[(1, 3), (6, 3)]), + mm_hashes: vec!["h0".to_string(), "h1".to_string()], + batched_keys: vec!["pixel_values".to_string()], + flat_keys: Default::default(), + keep_on_cpu_keys: vec![], + modality: common::Modality::Image as i32, + } + } + + fn tensor_of(elem: &MmFieldElem) -> &WireTensor { + match elem.data.as_ref().expect("data present") { + MmKwargValue::Tensor(tensor) => tensor, + other => panic!("expected tensor, got {other:?}"), + } + } + + #[test] + fn batched_keys_split_per_row_and_cast_to_model_dtype() { + let features = + build_mm_features(base_inputs(), &[0; 9], ModelDtype::BFloat16).expect("built"); + assert_eq!(features.len(), 2); + + for (i, feature) in features.iter().enumerate() { + assert_eq!(feature.modality, "image"); + assert_eq!(feature.identifier, format!("h{i}")); + assert_eq!(feature.mm_hash.as_deref(), Some(format!("h{i}").as_str())); + let item = feature.data.as_ref().expect("item present"); + let tensor = tensor_of(&item["pixel_values"]); + // Row i of the [2, 4] float32 batch, cast to bfloat16. + assert_eq!(tensor.dtype, "bfloat16"); + assert_eq!(tensor.shape, vec![4]); + assert!(matches!( + item["pixel_values"].field, + MmField::Batched(MmBatchedField { keep_on_cpu: false }) + )); + } + assert_eq!(features[0].mm_position.offset, 1); + assert_eq!(features[1].mm_position.offset, 6); + } + + #[test] + fn flat_keys_slice_by_cumulative_sizes() { + let mut mm = base_inputs(); + mm.pixel_values = Some(inline_tensor(vec![5, 2], "float32", f32_bytes(&[0.0; 10]))); + mm.batched_keys = vec!["patches_per_image".to_string()]; + mm.flat_keys = [("pixel_values".to_string(), "patches_per_image".to_string())].into(); + mm.model_specific_tensors = [( + "patches_per_image".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[2, 3])), + )] + .into(); + + let features = build_mm_features(mm, &[0; 9], ModelDtype::Float32).expect("built"); + let item0 = features[0].data.as_ref().expect("item 0"); + let item1 = features[1].data.as_ref().expect("item 1"); + assert_eq!(tensor_of(&item0["pixel_values"]).shape, vec![2, 2]); + assert_eq!(tensor_of(&item1["pixel_values"]).shape, vec![3, 2]); + + // Every elem carries the full per-item slice list. + let expected_slices = vec![ + MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(2), + step: None, + }), + MmSlice::Slice(SliceSpec { + start: Some(2), + stop: Some(5), + step: None, + }), + ]; + for item in [item0, item1] { + let MmField::Flat(flat) = &item["pixel_values"].field else { + panic!("expected flat field"); + }; + assert_eq!(flat.slices, expected_slices); + assert_eq!(flat.dim, 0); + } + } + + #[test] + fn unlisted_keys_are_shared_and_replicated() { + let mut mm = base_inputs(); + mm.model_specific_tensors = [( + "video_second_per_grid".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[1, 1])), + )] + .into(); + + let features = build_mm_features(mm, &[0; 9], ModelDtype::BFloat16).expect("built"); + for feature in &features { + let item = feature.data.as_ref().expect("item present"); + let elem = &item["video_second_per_grid"]; + assert_eq!(tensor_of(elem).shape, vec![2]); + assert!(matches!( + elem.field, + MmField::Shared(MmSharedField { + batch_size: 2, + keep_on_cpu: false, + }) + )); + } + } + + #[test] + fn is_embed_masks_structural_tokens() { + let mut mm = base_inputs(); + mm.im_token_id = Some(7); + // Placeholder 0 covers tokens [7, 7, 5] (mixed); placeholder 1 covers + // [7, 7, 7] (all image tokens). + let prompt = [9, 7, 7, 5, 9, 9, 7, 7, 7]; + + let features = build_mm_features(mm, &prompt, ModelDtype::BFloat16).expect("built"); + let mask = features[0] + .mm_position + .is_embed + .as_ref() + .expect("mixed range keeps a mask"); + assert_eq!(mask.dtype, "bool"); + assert_eq!(mask.shape, vec![3]); + assert!(features[1].mm_position.is_embed.is_none()); + } + + #[test] + fn video_renames_pixel_values() { + let mut mm = base_inputs(); + mm.modality = common::Modality::Video as i32; + + let features = build_mm_features(mm, &[0; 9], ModelDtype::BFloat16).expect("built"); + let item = features[0].data.as_ref().expect("item present"); + assert!(item.contains_key("pixel_values_videos")); + assert!(!item.contains_key("pixel_values")); + assert_eq!(features[0].modality, "video"); + } + + #[test] + fn rejects_hash_mismatch_and_non_inline_payloads() { + let mut mm = base_inputs(); + mm.mm_hashes.pop(); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("hash mismatch"); + assert!(err.contains("hash count"), "{err}"); + + let mut mm = base_inputs(); + mm.pixel_values = Some(vllm::TensorData { + shape: vec![2, 4], + dtype: "float32".to_string(), + payload: Some(vllm::tensor_data::Payload::Shm(Default::default())), + }); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("shm rejected"); + assert!(err.contains("inline"), "{err}"); + } + + #[test] + fn rejects_tensors_without_placeholders() { + // A payload that carries tensors but no placeholders must not silently + // degrade to a text-only request. + let mut mm = base_inputs(); + mm.mm_placeholders = placeholders(&[]); + mm.mm_hashes = vec![]; + let err = build_mm_features(mm, &[], ModelDtype::BFloat16) + .expect_err("tensors with no placeholders"); + assert!(err.contains("no placeholders"), "{err}"); + } + + #[test] + fn rejects_non_float32_floating_dtype() { + // A bf16 pixel tensor would reach the engine uncast — reject it. + let mut mm = base_inputs(); + mm.pixel_values = Some(inline_tensor(vec![2, 4], "bfloat16", vec![0u8; 16])); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16) + .expect_err("non-float32 floating dtype"); + assert!(err.contains("float32"), "{err}"); + } + + #[test] + fn rejects_tensor_byte_length_mismatch() { + // int64 [2] needs 16 bytes; supply 8 so the buffer can't match the shape. + let mut mm = base_inputs(); + mm.batched_keys = vec!["patches_per_image".to_string()]; + mm.model_specific_tensors = [( + "patches_per_image".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[2])), + )] + .into(); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("truncated payload"); + assert!(err.contains("does not match shape"), "{err}"); + } + + #[test] + fn shared_branch_preserves_keep_on_cpu() { + let mut mm = base_inputs(); + mm.model_specific_tensors = [( + "video_second_per_grid".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[1, 1])), + )] + .into(); + mm.keep_on_cpu_keys = vec!["video_second_per_grid".to_string()]; + + let features = build_mm_features(mm, &[0; 9], ModelDtype::BFloat16).expect("built"); + let item = features[0].data.as_ref().expect("item present"); + assert!(matches!( + item["video_second_per_grid"].field, + MmField::Shared(MmSharedField { + keep_on_cpu: true, + .. + }) + )); + } + + #[test] + fn validates_placeholder_range_without_im_token() { + // With no im_token_id the range check must still run. + let mut mm = base_inputs(); + mm.im_token_id = None; + let prompt = [9, 7, 7, 5, 9]; // 5 tokens; placeholder 1 spans [6, 9) + let err = build_mm_features(mm, &prompt, ModelDtype::BFloat16) + .expect_err("out-of-range placeholder"); + assert!(err.contains("exceeds prompt"), "{err}"); + } +}