[Rust Frontend][gRPC] Preserve multimodal metadata for remote-prefill decode (#54814)
Connor Carpenter committed
Sep 4, 2026 at 12:37 UTC
761c5861e34edb52c290e3171feb3763732ba8ca
3 files changed
+154
-25
rust/src/server/src/grpc/convert.rs
+9
-1
@@ -516,6 +516,11 @@ fn positions_to_proto(
516
// KV transfer params conversion (serde_json::Value ↔ prost_types::Struct)
517
// ========================================================================================
518
519
+/// Largest integer exactly representable as `f64` (2^53 - 1). Integral
520
+/// numbers within this bound are emitted as JSON integers so consumers
521
+/// expecting ints (e.g. `image_grid_thw`) don't see `16.0`.
522
+const MAX_SAFE_INTEGER_F64: f64 = ((1u64 << 53) - 1) as f64;
523
+
524
fn proto_struct_to_json(s: &prost_types::Struct) -> serde_json::Value {
525
serde_json::Value::Object(
526
s.fields.iter().map(|(k, v)| (k.clone(), proto_value_to_json(v))).collect(),
@@ -527,6 +532,9 @@ fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value {
532
match v.kind.as_ref() {
533
None | Some(Kind::NullValue(_)) => serde_json::Value::Null,
534
Some(Kind::BoolValue(b)) => serde_json::Value::Bool(*b),
535
+ Some(Kind::NumberValue(n)) if n.fract() == 0.0 && n.abs() <= MAX_SAFE_INTEGER_F64 => {
536
+ serde_json::Value::Number(serde_json::Number::from(*n as i64))
537
+ }
538
Some(Kind::NumberValue(n)) => serde_json::json!(*n),
539
Some(Kind::StringValue(s)) => serde_json::Value::String(s.clone()),
540
Some(Kind::ListValue(list)) => {
@@ -536,7 +544,7 @@ fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value {
544
}
545
}
546
539
-fn json_to_proto_struct(value: &serde_json::Value) -> Option<prost_types::Struct> {
547
+pub(super) fn json_to_proto_struct(value: &serde_json::Value) -> Option<prost_types::Struct> {
548
match value {
549
serde_json::Value::Object(map) => Some(prost_types::Struct {
550
fields: map.iter().map(|(k, v)| (k.clone(), json_to_proto_value(v))).collect(),
rust/src/server/src/grpc/inference.rs
+60
@@ -35,6 +35,65 @@ struct PreparedGrpcRequest {
35
started_at: Instant,
36
}
37
38
+/// Keep producer metadata for matching EC items; leave unmatched inputs intact.
39
+fn apply_encoder_cache_placeholders(text_request: &mut TextRequest) {
40
+ let is_decode_kv_consumer = text_request
41
+ .sampling_params
42
+ .vllm_xargs
43
+ .as_ref()
44
+ .and_then(|args| args.get("kv_transfer_params"))
45
+ .and_then(|params| params.get("do_remote_prefill"))
46
+ .and_then(serde_json::Value::as_bool)
47
+ .unwrap_or(false);
48
+ let ec_items = text_request
49
+ .sampling_params
50
+ .vllm_xargs
51
+ .as_ref()
52
+ .and_then(|args| args.get("ec_transfer_params"))
53
+ .and_then(|params| params.get("ec_items"))
54
+ .and_then(serde_json::Value::as_array)
55
+ .cloned();
56
+
57
+ // Decode uses EC metadata only to prepare the prompt; EngineCore consumes KV.
58
+ if is_decode_kv_consumer && let Some(args) = text_request.sampling_params.vllm_xargs.as_mut() {
59
+ args.remove("ec_transfer_params");
60
+ }
61
+
62
+ let Some(ec_items) = ec_items else {
63
+ return;
64
+ };
65
+ let Some(features) = text_request.mm_features.as_mut() else {
66
+ return;
67
+ };
68
+ let ec_items_by_hash: std::collections::HashMap<_, _> = ec_items
69
+ .iter()
70
+ .filter_map(serde_json::Value::as_object)
71
+ .filter_map(|item| {
72
+ item.get("mm_hash")
73
+ .and_then(serde_json::Value::as_str)
74
+ .map(|mm_hash| (mm_hash, item))
75
+ })
76
+ .collect();
77
+
78
+ for feature in features.iter_mut() {
79
+ let Some(item) = ec_items_by_hash.get(feature.identifier.as_str()) else {
80
+ continue;
81
+ };
82
+ let Some(data) = feature.data.as_mut() else {
83
+ continue;
84
+ };
85
+ let metadata_keys: Vec<_> = data
86
+ .keys()
87
+ .filter(|key| key.as_str() != "mm_hash" && item.contains_key(key.as_str()))
88
+ .cloned()
89
+ .collect();
90
+ if metadata_keys.is_empty() {
91
+ continue;
92
+ }
93
+ data.retain(|key, _| metadata_keys.contains(key));
94
+ }
95
+}
96
+
97
impl InferenceServiceImpl {
98
pub fn new(state: Arc<AppState>) -> Self {
99
Self { state }
@@ -106,6 +165,7 @@ impl InferenceServiceImpl {
165
text_request.prompt = Prompt::TokenIds(token_ids);
166
text_request.mm_features = mm_features;
167
}
168
+ apply_encoder_cache_placeholders(&mut text_request);
169
170
Ok(text_request)
171
}
rust/src/server/src/grpc/tests.rs
+85
-24
@@ -49,6 +49,7 @@ use zeromq::prelude::{SocketRecv, SocketSend};
49
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
50
51
use super::control::kv_event_source;
52
+use super::convert::json_to_proto_struct;
53
use super::pb::control_client::ControlClient;
54
use super::pb::inference_client::InferenceClient;
55
use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb};
@@ -143,6 +144,29 @@ fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason
144
]
145
}
146
147
+fn ec_proto_struct(mm_hashes: &[&str]) -> prost_types::Struct {
148
+ let ec_items: Vec<_> = mm_hashes
149
+ .iter()
150
+ .map(|mm_hash| {
151
+ serde_json::json!({
152
+ "image_grid_thw": [[1, 16, 16]],
153
+ "mm_hash": mm_hash,
154
+ })
155
+ })
156
+ .collect();
157
+ json_to_proto_struct(&serde_json::json!({ "ec_items": ec_items }))
158
+ .expect("valid EC proto struct")
159
+}
160
+
161
+fn decode_kv_proto_struct() -> prost_types::Struct {
162
+ json_to_proto_struct(&serde_json::json!({
163
+ "do_remote_prefill": true,
164
+ "pp_size": 1,
165
+ "remote_block_ids": [[7]],
166
+ }))
167
+ .expect("valid KV proto struct")
168
+}
169
+
170
async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
171
push.send(ZmqMessage::from(
172
rmp_serde::to_vec_named(&outputs).expect("encode outputs"),
@@ -701,22 +725,51 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() {
725
|request| {
726
let token_ids = request.prompt_token_ids.as_ref().expect("prompt token ids");
727
let features = request.mm_features.as_ref().expect("multimodal features");
704
- assert_eq!(features.len(), 1);
705
-
706
- let feature = &features[0];
707
- assert_eq!(feature.modality, "image");
708
- assert_eq!(feature.identifier, "image-1");
709
- assert_eq!(feature.mm_position.offset, 1);
710
- assert!(feature.mm_position.length > 1);
711
- assert_eq!(token_ids.len(), feature.mm_position.length + 2);
712
- assert_eq!(token_ids[0], 11);
713
- assert_eq!(token_ids.last(), Some(&12));
714
- assert!(
715
- token_ids[feature.mm_position.offset
716
- ..feature.mm_position.offset + feature.mm_position.length]
717
- .iter()
718
- .all(|token_id| *token_id == QWEN_IMAGE_TOKEN_ID)
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");
732
+ assert_eq!(feature.identifier, identifier);
733
+ assert!(feature.mm_position.length > 1);
734
+ assert_eq!(
735
+ feature
736
+ .data
737
+ .as_ref()
738
+ .expect("multimodal feature data")
739
+ .keys()
740
+ .map(String::as_str)
741
+ .collect::<Vec<_>>(),
742
+ vec!["image_grid_thw"]
743
+ );
744
+ }
745
+ assert_eq!(features[0].mm_position.offset, 1);
746
+ let xargs = request
747
+ .sampling_params
748
+ .as_ref()
749
+ .and_then(|params| params.extra_args.as_ref())
750
+ .expect("KV transfer args");
751
+ let kv_transfer_params =
752
+ xargs.get("kv_transfer_params").expect("KV transfer params");
753
+ assert_eq!(kv_transfer_params["pp_size"].as_i64(), Some(1));
754
+ assert_eq!(
755
+ kv_transfer_params["remote_block_ids"][0][0].as_i64(),
756
+ Some(7)
757
);
758
+ assert!(!xargs.contains_key("ec_transfer_params"));
759
+ assert_eq!(
760
+ token_ids.len(),
761
+ features.iter().map(|feature| feature.mm_position.length).sum::<usize>() + 3
762
+ );
763
+ assert_eq!(token_ids[0], 11);
764
+ assert_eq!(token_ids.last(), Some(&13));
765
+ for feature in features {
766
+ assert!(
767
+ token_ids[feature.mm_position.offset
768
+ ..feature.mm_position.offset + feature.mm_position.length]
769
+ .iter()
770
+ .all(|token_id| *token_id == QWEN_IMAGE_TOKEN_ID)
771
+ );
772
+ }
773
},
774
)
775
.await;
@@ -734,16 +787,24 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() {
787
request_id: "test-multimodal".to_string(),
788
model: "test-model".to_string(),
789
prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds {
737
- ids: vec![11, QWEN_IMAGE_TOKEN_ID, 12],
790
+ ids: vec![11, QWEN_IMAGE_TOKEN_ID, 12, QWEN_IMAGE_TOKEN_ID, 13],
791
})),
739
- media: vec![pb::MediaItem {
740
- modality: pb::Modality::Image as i32,
741
- source: Some(pb::media_item::Source::DataUri(
742
- TINY_PNG_DATA_URI.to_string(),
743
- )),
744
- mime_type: String::new(),
745
- uuid: "image-1".to_string(),
746
- }],
792
+ media: ["image-1", "image-2"]
793
+ .into_iter()
794
+ .map(|uuid| pb::MediaItem {
795
+ modality: pb::Modality::Image as i32,
796
+ source: Some(pb::media_item::Source::DataUri(
797
+ TINY_PNG_DATA_URI.to_string(),
798
+ )),
799
+ mime_type: String::new(),
800
+ uuid: uuid.to_string(),
801
+ })
802
+ .collect(),
803
+ kv: Some(pb::KvCacheParameters {
804
+ kv_transfer_params: Some(decode_kv_proto_struct()),
805
+ ec_transfer_params: Some(ec_proto_struct(&["image-2", "image-1"])),
806
+ ..Default::default()
807
+ }),
808
stopping: Some(pb::StoppingCriteria {
809
max_new_tokens: 10,
810
..Default::default()