[Bugfix][Rust Frontend] Fix LogprobsTensors wire schema mismatch (#53939)
Signed-off-by: jungjiyu <libraryofjiyu@gmail.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Jung Jiyu committed
Aug 27, 2026 at 08:16 UTC
3fd72ffd7b78bc80f249cdfc28d67177dbab2b56
4 files changed
+56
-3
rust/src/engine-core-client/src/protocol/logprobs.rs
+10
@@ -222,6 +222,7 @@ impl WireLogprobs {
222
data: WireArrayData::RawView(token_ranks.into()),
223
},
224
cu_num_generated_tokens: None,
225
+ cu_num_generated_tokens_tensor: None,
226
})
227
}
228
@@ -239,6 +240,15 @@ impl WireLogprobs {
240
);
241
}
242
243
+ // Unlike the sibling check above, don't Debug-print the payload:
244
+ // an opaque non-None value here may embed a full tensor blob.
245
+ if self.cu_num_generated_tokens_tensor.is_some() {
246
+ bail_ext_value_decode!(
247
+ "{field_prefix}.cu_num_generated_tokens_tensor: \
248
+ expected None for per-request engine-core logprobs payload"
249
+ );
250
+ }
251
+
252
let token_ids = array::decode_array2_u32(
253
self.logprob_token_ids,
254
&format!("{field_prefix}.logprob_token_ids"),
rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+24
@@ -87,6 +87,7 @@ fn inline_prompt_logprobs_value() -> Value {
87
ndarray_value("float32", &[2, 3], probs),
88
ndarray_value("int64", &[2], ranks),
89
Value::Nil,
90
+ Value::Nil,
91
])
92
}
93
@@ -240,6 +241,29 @@ fn decodes_inline_prompt_logprobs() {
241
assert_eq!(logprobs, expected_prompt_logprobs());
242
}
243
244
+#[test]
245
+fn rejects_non_none_cu_num_generated_tokens_tensor() {
246
+ let Value::Array(mut fields) = inline_prompt_logprobs_value() else {
247
+ panic!("inline_prompt_logprobs_value must be an array");
248
+ };
249
+ fields[4] = Value::from(42);
250
+
251
+ let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
252
+ None,
253
+ Some(Value::Array(fields)),
254
+ )))];
255
+
256
+ let error = decode_engine_core_outputs(&frames).unwrap_err();
257
+ let crate::error::Error::ExtValueDecode { message } = &error else {
258
+ panic!("expected ValueDecodeExt");
259
+ };
260
+ assert_eq!(
261
+ message,
262
+ "new_prompt_logprobs_tensors.cu_num_generated_tokens_tensor: \
263
+ expected None for per-request engine-core logprobs payload"
264
+ );
265
+}
266
+
267
#[test]
268
fn decodes_big_endian_payloads() {
269
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
rust/src/engine-core-client/src/protocol/logprobs/wire.rs
+18
-1
@@ -3,6 +3,7 @@
3
4
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
5
6
+use crate::protocol::OpaqueValue;
7
use crate::protocol::tensor::WireNdArray;
8
9
/// Python wire representation of `LogprobsLists` / `LogprobsTensors` before
@@ -13,7 +14,15 @@ use crate::protocol::tensor::WireNdArray;
14
/// into semantic per-position logprobs records.
15
///
16
/// Original Python definition:
16
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/outputs.py#L23-L56>
17
+/// <https://github.com/vllm-project/vllm/blob/d5cadcee8641a9fcec15facb5a9157d157daa207/vllm/v1/outputs.py#L30-L106>
18
+///
19
+/// # Wire compatibility
20
+///
21
+/// Python's `LogprobsLists` / `LogprobsTensors` are `NamedTuple`s: every
22
+/// field is serialized positionally, and `Deserialize_tuple` here hard-fails
23
+/// on payloads longer than this struct. To preserve positional compatibility,
24
+/// append new Python fields and add matching trailing `#[serde(default)]`
25
+/// fields here in the same change.
26
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
27
pub struct WireLogprobs {
28
/// Wire array with shape `[num_positions, max_num_logprobs + 1]`.
@@ -31,4 +40,12 @@ pub struct WireLogprobs {
40
/// the semantic Rust decoder rejects any other value.
41
#[serde(default)]
42
pub cu_num_generated_tokens: Option<Vec<usize>>,
43
+ /// Device-only generation boundaries added to `LogprobsTensors` in
44
+ /// vllm-project/vllm#52242. Preserved only for wire compatibility:
45
+ /// NamedTuple encoding always emits the slot even when it is `None`,
46
+ /// so this field must exist for the 5-element payload to decode.
47
+ /// Scheduler-sliced per-request outputs should emit `None` here, and
48
+ /// the semantic Rust decoder rejects any other value.
49
+ #[serde(default)]
50
+ pub cu_num_generated_tokens_tensor: Option<OpaqueValue>,
51
}
rust/src/engine-core-client/src/tests/python_compat.py
+4
-2
@@ -353,7 +353,8 @@ inline_prompt_logprobs = engine_outputs_wire(
353
).tobytes(),
354
),
355
("int64", [2], np.array([3, 4], dtype=np.int64).tobytes()),
356
- None,
356
+ None, # cu_num_generated_tokens
357
+ None, # cu_num_generated_tokens_tensor
358
),
359
)
360
)
@@ -376,7 +377,8 @@ multipart_prompt_logprobs = engine_outputs_wire(
377
).tobytes(),
378
),
379
("int64", [2], np.array([3, 4], dtype=np.int64).tobytes()),
379
- None,
380
+ None, # cu_num_generated_tokens
381
+ None, # cu_num_generated_tokens_tensor
382
),
383
)
384
)