From 4e2849539fc2e353b0d06909a6defe807a309b1b Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 1 Sep 2026 10:30:27 -0700 Subject: [PATCH 1/4] fix(grpc): preserve multimodal metadata for KV decode Signed-off-by: Connor Carpenter --- rust/src/server/src/grpc/convert.rs | 3 + rust/src/server/src/grpc/inference.rs | 57 +++++++++++++ rust/src/server/src/grpc/tests.rs | 110 ++++++++++++++++++++++++++ 3 files changed, 170 insertions(+) diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 93669bd1b2ab..6ffd5bdd3fe3 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -527,6 +527,9 @@ fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value { match v.kind.as_ref() { None | Some(Kind::NullValue(_)) => serde_json::Value::Null, Some(Kind::BoolValue(b)) => serde_json::Value::Bool(*b), + Some(Kind::NumberValue(n)) if n.fract() == 0.0 && n.abs() <= 9_007_199_254_740_991.0 => { + serde_json::Value::Number(serde_json::Number::from(*n as i64)) + } Some(Kind::NumberValue(n)) => serde_json::json!(*n), Some(Kind::StringValue(s)) => serde_json::Value::String(s.clone()), Some(Kind::ListValue(list)) => { diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index c3f590af99f6..93d31b7821e9 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -35,6 +35,62 @@ struct PreparedGrpcRequest { started_at: Instant, } +/// Keep producer metadata for matching EC items; leave unmatched inputs intact. +fn apply_encoder_cache_placeholders(text_request: &mut TextRequest) { + let is_decode_kv_consumer = text_request + .sampling_params + .vllm_xargs + .as_ref() + .and_then(|args| args.get("kv_transfer_params")) + .and_then(|params| params.get("do_remote_prefill")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let Some(ec_items) = text_request + .sampling_params + .vllm_xargs + .as_ref() + .and_then(|args| args.get("ec_transfer_params")) + .and_then(|params| params.get("ec_items")) + .and_then(serde_json::Value::as_array) + .cloned() + else { + return; + }; + let Some(features) = text_request.mm_features.as_mut() else { + return; + }; + + for (feature, item) in features.iter_mut().zip(ec_items) { + let Some(item) = item.as_object() else { + continue; + }; + if item.get("mm_hash").and_then(serde_json::Value::as_str) + != Some(feature.identifier.as_str()) + { + continue; + } + let Some(data) = feature.data.as_mut() else { + continue; + }; + let metadata_keys: Vec<_> = data + .keys() + .filter(|key| key.as_str() != "mm_hash" && item.contains_key(key.as_str())) + .cloned() + .collect(); + if metadata_keys.is_empty() { + continue; + } + data.retain(|key, _| metadata_keys.contains(key)); + } + + // Decode uses EC metadata only to prepare the prompt; EngineCore consumes KV. + if is_decode_kv_consumer { + if let Some(args) = text_request.sampling_params.vllm_xargs.as_mut() { + args.remove("ec_transfer_params"); + } + } +} + impl InferenceServiceImpl { pub fn new(state: Arc) -> Self { Self { state } @@ -105,6 +161,7 @@ impl InferenceServiceImpl { .map_err(|error| Status::internal(error.to_report_string()))?; text_request.prompt = Prompt::TokenIds(token_ids); text_request.mm_features = mm_features; + apply_encoder_cache_placeholders(&mut text_request); } Ok(text_request) diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 34667c31ea35..a3cb740a126e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -143,6 +143,88 @@ fn default_stream_output_specs() -> Vec<(Vec, Option prost_types::Struct { + use prost_types::value::Kind; + + prost_types::Struct { + fields: std::collections::BTreeMap::from([( + "ec_items".to_string(), + prost_types::Value { + kind: Some(Kind::ListValue(prost_types::ListValue { + values: vec![prost_types::Value { + kind: Some(Kind::StructValue(prost_types::Struct { + fields: std::collections::BTreeMap::from([ + ( + "image_grid_thw".to_string(), + prost_types::Value { + kind: Some(Kind::ListValue(prost_types::ListValue { + values: vec![prost_types::Value { + kind: Some(Kind::ListValue( + prost_types::ListValue { + values: vec![1.0, 16.0, 16.0] + .into_iter() + .map(|value| prost_types::Value { + kind: Some(Kind::NumberValue( + value, + )), + }) + .collect(), + }, + )), + }], + })), + }, + ), + ( + "mm_hash".to_string(), + prost_types::Value { + kind: Some(Kind::StringValue("image-1".to_string())), + }, + ), + ]), + })), + }], + })), + }, + )]), + } +} + +fn decode_kv_proto_struct() -> prost_types::Struct { + use prost_types::value::Kind; + + prost_types::Struct { + fields: std::collections::BTreeMap::from([ + ( + "do_remote_prefill".to_string(), + prost_types::Value { + kind: Some(Kind::BoolValue(true)), + }, + ), + ( + "pp_size".to_string(), + prost_types::Value { + kind: Some(Kind::NumberValue(1.0)), + }, + ), + ( + "remote_block_ids".to_string(), + prost_types::Value { + kind: Some(Kind::ListValue(prost_types::ListValue { + values: vec![prost_types::Value { + kind: Some(Kind::ListValue(prost_types::ListValue { + values: vec![prost_types::Value { + kind: Some(Kind::NumberValue(7.0)), + }], + })), + }], + })), + }, + ), + ]), + } +} + async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) { push.send(ZmqMessage::from( rmp_serde::to_vec_named(&outputs).expect("encode outputs"), @@ -708,6 +790,29 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() { assert_eq!(feature.identifier, "image-1"); assert_eq!(feature.mm_position.offset, 1); assert!(feature.mm_position.length > 1); + assert_eq!( + feature + .data + .as_ref() + .expect("multimodal feature data") + .keys() + .map(String::as_str) + .collect::>(), + vec!["image_grid_thw"] + ); + let xargs = request + .sampling_params + .as_ref() + .and_then(|params| params.extra_args.as_ref()) + .expect("KV transfer args"); + let kv_transfer_params = + xargs.get("kv_transfer_params").expect("KV transfer params"); + assert_eq!(kv_transfer_params["pp_size"].as_i64(), Some(1)); + assert_eq!( + kv_transfer_params["remote_block_ids"][0][0].as_i64(), + Some(7) + ); + assert!(!xargs.contains_key("ec_transfer_params")); assert_eq!(token_ids.len(), feature.mm_position.length + 2); assert_eq!(token_ids[0], 11); assert_eq!(token_ids.last(), Some(&12)); @@ -744,6 +849,11 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() { mime_type: String::new(), uuid: "image-1".to_string(), }], + kv: Some(pb::KvCacheParameters { + kv_transfer_params: Some(decode_kv_proto_struct()), + ec_transfer_params: Some(ec_proto_struct()), + ..Default::default() + }), stopping: Some(pb::StoppingCriteria { max_new_tokens: 10, ..Default::default() From 905b964a31c596d8082e046dcb0e23fdf25d56e4 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 3 Sep 2026 12:32:59 -0700 Subject: [PATCH 2/4] fix(grpc): handle encoder cache handoff edge cases Signed-off-by: Connor Carpenter --- rust/src/server/src/grpc/inference.rs | 41 +++---- rust/src/server/src/grpc/tests.rs | 150 ++++++++++++++------------ 2 files changed, 104 insertions(+), 87 deletions(-) diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 93d31b7821e9..3f23cb7e04e7 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -45,30 +45,42 @@ fn apply_encoder_cache_placeholders(text_request: &mut TextRequest) { .and_then(|params| params.get("do_remote_prefill")) .and_then(serde_json::Value::as_bool) .unwrap_or(false); - let Some(ec_items) = text_request + let ec_items = text_request .sampling_params .vllm_xargs .as_ref() .and_then(|args| args.get("ec_transfer_params")) .and_then(|params| params.get("ec_items")) .and_then(serde_json::Value::as_array) - .cloned() - else { + .cloned(); + + // Decode uses EC metadata only to prepare the prompt; EngineCore consumes KV. + if is_decode_kv_consumer { + if let Some(args) = text_request.sampling_params.vllm_xargs.as_mut() { + args.remove("ec_transfer_params"); + } + } + + let Some(ec_items) = ec_items else { return; }; let Some(features) = text_request.mm_features.as_mut() else { return; }; + let ec_items_by_hash: std::collections::HashMap<_, _> = ec_items + .iter() + .filter_map(serde_json::Value::as_object) + .filter_map(|item| { + item.get("mm_hash") + .and_then(serde_json::Value::as_str) + .map(|mm_hash| (mm_hash, item)) + }) + .collect(); - for (feature, item) in features.iter_mut().zip(ec_items) { - let Some(item) = item.as_object() else { + for feature in features.iter_mut() { + let Some(item) = ec_items_by_hash.get(feature.identifier.as_str()) else { continue; }; - if item.get("mm_hash").and_then(serde_json::Value::as_str) - != Some(feature.identifier.as_str()) - { - continue; - } let Some(data) = feature.data.as_mut() else { continue; }; @@ -82,13 +94,6 @@ fn apply_encoder_cache_placeholders(text_request: &mut TextRequest) { } data.retain(|key, _| metadata_keys.contains(key)); } - - // Decode uses EC metadata only to prepare the prompt; EngineCore consumes KV. - if is_decode_kv_consumer { - if let Some(args) = text_request.sampling_params.vllm_xargs.as_mut() { - args.remove("ec_transfer_params"); - } - } } impl InferenceServiceImpl { @@ -161,8 +166,8 @@ impl InferenceServiceImpl { .map_err(|error| Status::internal(error.to_report_string()))?; text_request.prompt = Prompt::TokenIds(token_ids); text_request.mm_features = mm_features; - apply_encoder_cache_placeholders(&mut text_request); } + apply_encoder_cache_placeholders(&mut text_request); Ok(text_request) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index a3cb740a126e..cf31e9dcdcec 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -143,7 +143,7 @@ fn default_stream_output_specs() -> Vec<(Vec, Option prost_types::Struct { +fn ec_proto_struct(mm_hashes: &[&str]) -> prost_types::Struct { use prost_types::value::Kind; prost_types::Struct { @@ -151,39 +151,42 @@ fn ec_proto_struct() -> prost_types::Struct { "ec_items".to_string(), prost_types::Value { kind: Some(Kind::ListValue(prost_types::ListValue { - values: vec![prost_types::Value { - kind: Some(Kind::StructValue(prost_types::Struct { - fields: std::collections::BTreeMap::from([ - ( - "image_grid_thw".to_string(), - prost_types::Value { - kind: Some(Kind::ListValue(prost_types::ListValue { - values: vec![prost_types::Value { - kind: Some(Kind::ListValue( - prost_types::ListValue { - values: vec![1.0, 16.0, 16.0] - .into_iter() - .map(|value| prost_types::Value { - kind: Some(Kind::NumberValue( - value, - )), - }) - .collect(), - }, - )), - }], - })), - }, - ), - ( - "mm_hash".to_string(), - prost_types::Value { - kind: Some(Kind::StringValue("image-1".to_string())), - }, - ), - ]), - })), - }], + values: mm_hashes + .iter() + .map(|mm_hash| prost_types::Value { + kind: Some(Kind::StructValue(prost_types::Struct { + fields: std::collections::BTreeMap::from([ + ( + "image_grid_thw".to_string(), + prost_types::Value { + kind: Some(Kind::ListValue(prost_types::ListValue { + values: vec![prost_types::Value { + kind: Some(Kind::ListValue( + prost_types::ListValue { + values: vec![1.0, 16.0, 16.0] + .into_iter() + .map(|value| prost_types::Value { + kind: Some(Kind::NumberValue( + value, + )), + }) + .collect(), + }, + )), + }], + })), + }, + ), + ( + "mm_hash".to_string(), + prost_types::Value { + kind: Some(Kind::StringValue((*mm_hash).to_string())), + }, + ), + ]), + })), + }) + .collect(), })), }, )]), @@ -783,23 +786,24 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() { |request| { let token_ids = request.prompt_token_ids.as_ref().expect("prompt token ids"); let features = request.mm_features.as_ref().expect("multimodal features"); - assert_eq!(features.len(), 1); - - let feature = &features[0]; - assert_eq!(feature.modality, "image"); - assert_eq!(feature.identifier, "image-1"); - assert_eq!(feature.mm_position.offset, 1); - assert!(feature.mm_position.length > 1); - assert_eq!( - feature - .data - .as_ref() - .expect("multimodal feature data") - .keys() - .map(String::as_str) - .collect::>(), - vec!["image_grid_thw"] - ); + assert_eq!(features.len(), 2); + + for (feature, identifier) in features.iter().zip(["image-1", "image-2"]) { + assert_eq!(feature.modality, "image"); + assert_eq!(feature.identifier, identifier); + assert!(feature.mm_position.length > 1); + assert_eq!( + feature + .data + .as_ref() + .expect("multimodal feature data") + .keys() + .map(String::as_str) + .collect::>(), + vec!["image_grid_thw"] + ); + } + assert_eq!(features[0].mm_position.offset, 1); let xargs = request .sampling_params .as_ref() @@ -813,15 +817,20 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() { Some(7) ); assert!(!xargs.contains_key("ec_transfer_params")); - assert_eq!(token_ids.len(), feature.mm_position.length + 2); - assert_eq!(token_ids[0], 11); - assert_eq!(token_ids.last(), Some(&12)); - assert!( - token_ids[feature.mm_position.offset - ..feature.mm_position.offset + feature.mm_position.length] - .iter() - .all(|token_id| *token_id == QWEN_IMAGE_TOKEN_ID) + assert_eq!( + token_ids.len(), + features.iter().map(|feature| feature.mm_position.length).sum::() + 3 ); + assert_eq!(token_ids[0], 11); + assert_eq!(token_ids.last(), Some(&13)); + for feature in features { + assert!( + token_ids[feature.mm_position.offset + ..feature.mm_position.offset + feature.mm_position.length] + .iter() + .all(|token_id| *token_id == QWEN_IMAGE_TOKEN_ID) + ); + } }, ) .await; @@ -839,19 +848,22 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() { request_id: "test-multimodal".to_string(), model: "test-model".to_string(), prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { - ids: vec![11, QWEN_IMAGE_TOKEN_ID, 12], + ids: vec![11, QWEN_IMAGE_TOKEN_ID, 12, QWEN_IMAGE_TOKEN_ID, 13], })), - media: vec![pb::MediaItem { - modality: pb::Modality::Image as i32, - source: Some(pb::media_item::Source::DataUri( - TINY_PNG_DATA_URI.to_string(), - )), - mime_type: String::new(), - uuid: "image-1".to_string(), - }], + media: ["image-1", "image-2"] + .into_iter() + .map(|uuid| pb::MediaItem { + modality: pb::Modality::Image as i32, + source: Some(pb::media_item::Source::DataUri( + TINY_PNG_DATA_URI.to_string(), + )), + mime_type: String::new(), + uuid: uuid.to_string(), + }) + .collect(), kv: Some(pb::KvCacheParameters { kv_transfer_params: Some(decode_kv_proto_struct()), - ec_transfer_params: Some(ec_proto_struct()), + ec_transfer_params: Some(ec_proto_struct(&["image-2", "image-1"])), ..Default::default() }), stopping: Some(pb::StoppingCriteria { From 32b31c94027c5690a2449c4a158b3c0d06caec10 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 3 Sep 2026 20:51:07 +0000 Subject: [PATCH 3/4] refactor(grpc): name f64 safe-integer bound, build EC fixtures via json! Replace the 9_007_199_254_740_991.0 literal with a named MAX_SAFE_INTEGER_F64 constant, and build the ec_proto_struct / decode_kv_proto_struct test fixtures with serde_json::json! plus the existing json_to_proto_struct helper instead of hand-assembled prost trees. Co-authored-by: Kimi Signed-off-by: Bugen Zhao --- rust/src/server/src/grpc/convert.rs | 9 ++- rust/src/server/src/grpc/tests.rs | 97 ++++++----------------------- 2 files changed, 25 insertions(+), 81 deletions(-) diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 6ffd5bdd3fe3..2a3f41ffa438 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -516,6 +516,11 @@ fn positions_to_proto( // KV transfer params conversion (serde_json::Value ↔ prost_types::Struct) // ======================================================================================== +/// Largest integer exactly representable as `f64` (2^53 - 1). Integral +/// numbers within this bound are emitted as JSON integers so consumers +/// expecting ints (e.g. `image_grid_thw`) don't see `16.0`. +const MAX_SAFE_INTEGER_F64: f64 = ((1u64 << 53) - 1) as f64; + fn proto_struct_to_json(s: &prost_types::Struct) -> serde_json::Value { serde_json::Value::Object( s.fields.iter().map(|(k, v)| (k.clone(), proto_value_to_json(v))).collect(), @@ -527,7 +532,7 @@ fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value { match v.kind.as_ref() { None | Some(Kind::NullValue(_)) => serde_json::Value::Null, Some(Kind::BoolValue(b)) => serde_json::Value::Bool(*b), - Some(Kind::NumberValue(n)) if n.fract() == 0.0 && n.abs() <= 9_007_199_254_740_991.0 => { + Some(Kind::NumberValue(n)) if n.fract() == 0.0 && n.abs() <= MAX_SAFE_INTEGER_F64 => { serde_json::Value::Number(serde_json::Number::from(*n as i64)) } Some(Kind::NumberValue(n)) => serde_json::json!(*n), @@ -539,7 +544,7 @@ fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value { } } -fn json_to_proto_struct(value: &serde_json::Value) -> Option { +pub(super) fn json_to_proto_struct(value: &serde_json::Value) -> Option { match value { serde_json::Value::Object(map) => Some(prost_types::Struct { fields: map.iter().map(|(k, v)| (k.clone(), json_to_proto_value(v))).collect(), diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index cf31e9dcdcec..ed465e5c7c6c 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -49,6 +49,7 @@ use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; use super::control::kv_event_source; +use super::convert::json_to_proto_struct; use super::pb::control_client::ControlClient; use super::pb::inference_client::InferenceClient; use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb}; @@ -144,88 +145,26 @@ fn default_stream_output_specs() -> Vec<(Vec, Option prost_types::Struct { - use prost_types::value::Kind; - - prost_types::Struct { - fields: std::collections::BTreeMap::from([( - "ec_items".to_string(), - prost_types::Value { - kind: Some(Kind::ListValue(prost_types::ListValue { - values: mm_hashes - .iter() - .map(|mm_hash| prost_types::Value { - kind: Some(Kind::StructValue(prost_types::Struct { - fields: std::collections::BTreeMap::from([ - ( - "image_grid_thw".to_string(), - prost_types::Value { - kind: Some(Kind::ListValue(prost_types::ListValue { - values: vec![prost_types::Value { - kind: Some(Kind::ListValue( - prost_types::ListValue { - values: vec![1.0, 16.0, 16.0] - .into_iter() - .map(|value| prost_types::Value { - kind: Some(Kind::NumberValue( - value, - )), - }) - .collect(), - }, - )), - }], - })), - }, - ), - ( - "mm_hash".to_string(), - prost_types::Value { - kind: Some(Kind::StringValue((*mm_hash).to_string())), - }, - ), - ]), - })), - }) - .collect(), - })), - }, - )]), - } + let ec_items: Vec<_> = mm_hashes + .iter() + .map(|mm_hash| { + serde_json::json!({ + "image_grid_thw": [[1, 16, 16]], + "mm_hash": mm_hash, + }) + }) + .collect(); + json_to_proto_struct(&serde_json::json!({ "ec_items": ec_items })) + .expect("valid EC proto struct") } fn decode_kv_proto_struct() -> prost_types::Struct { - use prost_types::value::Kind; - - prost_types::Struct { - fields: std::collections::BTreeMap::from([ - ( - "do_remote_prefill".to_string(), - prost_types::Value { - kind: Some(Kind::BoolValue(true)), - }, - ), - ( - "pp_size".to_string(), - prost_types::Value { - kind: Some(Kind::NumberValue(1.0)), - }, - ), - ( - "remote_block_ids".to_string(), - prost_types::Value { - kind: Some(Kind::ListValue(prost_types::ListValue { - values: vec![prost_types::Value { - kind: Some(Kind::ListValue(prost_types::ListValue { - values: vec![prost_types::Value { - kind: Some(Kind::NumberValue(7.0)), - }], - })), - }], - })), - }, - ), - ]), - } + json_to_proto_struct(&serde_json::json!({ + "do_remote_prefill": true, + "pp_size": 1, + "remote_block_ids": [[7]], + })) + .expect("valid KV proto struct") } async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) { From 21b373efbe18d437cdaba10a819d5f9a794431e7 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 3 Sep 2026 14:47:38 -0700 Subject: [PATCH 4/4] style(grpc): collapse encoder cache condition Signed-off-by: Connor Carpenter --- rust/src/server/src/grpc/inference.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 3f23cb7e04e7..f864cbde5959 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -55,10 +55,8 @@ fn apply_encoder_cache_placeholders(text_request: &mut TextRequest) { .cloned(); // Decode uses EC metadata only to prepare the prompt; EngineCore consumes KV. - if is_decode_kv_consumer { - if let Some(args) = text_request.sampling_params.vllm_xargs.as_mut() { - args.remove("ec_transfer_params"); - } + if is_decode_kv_consumer && let Some(args) = text_request.sampling_params.vllm_xargs.as_mut() { + args.remove("ec_transfer_params"); } let Some(ec_items) = ec_items else {