From 3fd72ffd7b78bc80f249cdfc28d67177dbab2b56 Mon Sep 17 00:00:00 2001 From: Jung Jiyu <142137932+jungjiyu@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:16:43 +0900 Subject: [PATCH] [Bugfix][Rust Frontend] Fix LogprobsTensors wire schema mismatch (#53939) Signed-off-by: jungjiyu Signed-off-by: Tyler Michael Smith Co-authored-by: Tyler Michael Smith Co-authored-by: OpenAI Codex --- .../src/protocol/logprobs.rs | 10 ++++++++ .../src/protocol/logprobs/tests.rs | 24 +++++++++++++++++++ .../src/protocol/logprobs/wire.rs | 19 ++++++++++++++- .../src/tests/python_compat.py | 6 +++-- 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index 28f6327d9b..6949dc0330 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -222,6 +222,7 @@ impl WireLogprobs { data: WireArrayData::RawView(token_ranks.into()), }, cu_num_generated_tokens: None, + cu_num_generated_tokens_tensor: None, }) } @@ -239,6 +240,15 @@ impl WireLogprobs { ); } + // Unlike the sibling check above, don't Debug-print the payload: + // an opaque non-None value here may embed a full tensor blob. + if self.cu_num_generated_tokens_tensor.is_some() { + bail_ext_value_decode!( + "{field_prefix}.cu_num_generated_tokens_tensor: \ + expected None for per-request engine-core logprobs payload" + ); + } + let token_ids = array::decode_array2_u32( self.logprob_token_ids, &format!("{field_prefix}.logprob_token_ids"), diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index 42f7ec78b5..2708637ebc 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -87,6 +87,7 @@ fn inline_prompt_logprobs_value() -> Value { ndarray_value("float32", &[2, 3], probs), ndarray_value("int64", &[2], ranks), Value::Nil, + Value::Nil, ]) } @@ -240,6 +241,29 @@ fn decodes_inline_prompt_logprobs() { assert_eq!(logprobs, expected_prompt_logprobs()); } +#[test] +fn rejects_non_none_cu_num_generated_tokens_tensor() { + let Value::Array(mut fields) = inline_prompt_logprobs_value() else { + panic!("inline_prompt_logprobs_value must be an array"); + }; + fields[4] = Value::from(42); + + let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields( + None, + Some(Value::Array(fields)), + )))]; + + let error = decode_engine_core_outputs(&frames).unwrap_err(); + let crate::error::Error::ExtValueDecode { message } = &error else { + panic!("expected ValueDecodeExt"); + }; + assert_eq!( + message, + "new_prompt_logprobs_tensors.cu_num_generated_tokens_tensor: \ + expected None for per-request engine-core logprobs payload" + ); +} + #[test] fn decodes_big_endian_payloads() { let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields( diff --git a/rust/src/engine-core-client/src/protocol/logprobs/wire.rs b/rust/src/engine-core-client/src/protocol/logprobs/wire.rs index 2f37d7ee46..375cf466f4 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/wire.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/wire.rs @@ -3,6 +3,7 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; +use crate::protocol::OpaqueValue; use crate::protocol::tensor::WireNdArray; /// Python wire representation of `LogprobsLists` / `LogprobsTensors` before @@ -13,7 +14,15 @@ use crate::protocol::tensor::WireNdArray; /// into semantic per-position logprobs records. /// /// Original Python definition: -/// +/// +/// +/// # Wire compatibility +/// +/// Python's `LogprobsLists` / `LogprobsTensors` are `NamedTuple`s: every +/// field is serialized positionally, and `Deserialize_tuple` here hard-fails +/// on payloads longer than this struct. To preserve positional compatibility, +/// append new Python fields and add matching trailing `#[serde(default)]` +/// fields here in the same change. #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)] pub struct WireLogprobs { /// Wire array with shape `[num_positions, max_num_logprobs + 1]`. @@ -31,4 +40,12 @@ pub struct WireLogprobs { /// the semantic Rust decoder rejects any other value. #[serde(default)] pub cu_num_generated_tokens: Option>, + /// Device-only generation boundaries added to `LogprobsTensors` in + /// vllm-project/vllm#52242. Preserved only for wire compatibility: + /// NamedTuple encoding always emits the slot even when it is `None`, + /// so this field must exist for the 5-element payload to decode. + /// Scheduler-sliced per-request outputs should emit `None` here, and + /// the semantic Rust decoder rejects any other value. + #[serde(default)] + pub cu_num_generated_tokens_tensor: Option, } diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 0cbd2f4a59..d3705cc391 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -353,7 +353,8 @@ inline_prompt_logprobs = engine_outputs_wire( ).tobytes(), ), ("int64", [2], np.array([3, 4], dtype=np.int64).tobytes()), - None, + None, # cu_num_generated_tokens + None, # cu_num_generated_tokens_tensor ), ) ) @@ -376,7 +377,8 @@ multipart_prompt_logprobs = engine_outputs_wire( ).tobytes(), ), ("int64", [2], np.array([3, 4], dtype=np.int64).tobytes()), - None, + None, # cu_num_generated_tokens + None, # cu_num_generated_tokens_tensor ), ) )