[Bugfix][Rust Frontend] Handle zero-column logprobs payloads without panicking (#49113)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Feathbow <feathbow@gmail.com>
Salt Sato committed
Jul 21, 2026 at 05:30 UTC
f25953cc59f9b4ba9b04b16228d2b86dcfbcbdb1
2 files changed
+59
rust/src/engine-core-client/src/protocol/logprobs.rs
+13
@@ -269,6 +269,19 @@ impl WireLogprobs {
269
);
270
}
271
272
+ // Empty position lists may be encoded as either [0, 0] or [0, k + 1].
273
+ if token_ids.rows == 0 {
274
+ return Ok(Logprobs {
275
+ positions: Vec::new(),
276
+ });
277
+ }
278
+ if token_ids.cols == 0 {
279
+ bail_ext_value_decode!(
280
+ "{field_prefix}: zero-column logprobs payload with {} rows",
281
+ token_ids.rows
282
+ );
283
+ }
284
+
285
let mut positions = Vec::with_capacity(token_ids.rows);
286
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
287
.data
rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+46
@@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() {
303
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
304
);
305
}
306
+
307
+#[test]
308
+fn decodes_zero_row_logprobs_as_empty() {
309
+ for shape in [[0usize, 0], [0, 3]] {
310
+ let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
311
+ None,
312
+ Some(Value::Array(vec![
313
+ ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
314
+ ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
315
+ ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
316
+ Value::Nil,
317
+ ])),
318
+ )))];
319
+ let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
320
+ let logprobs = decoded.outputs[0]
321
+ .new_prompt_logprobs_tensors
322
+ .clone()
323
+ .unwrap()
324
+ .into_direct()
325
+ .unwrap();
326
+ assert!(logprobs.is_empty());
327
+ }
328
+}
329
+
330
+#[test]
331
+fn rejects_zero_column_logprobs_with_rows() {
332
+ let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
333
+ let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
334
+ Some(Value::Array(vec![
335
+ ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
336
+ ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
337
+ ndarray_value("<i8", &[2], ranks),
338
+ Value::Nil,
339
+ ])),
340
+ None,
341
+ )))];
342
+
343
+ let error = decode_engine_core_outputs(&frames).unwrap_err();
344
+ let crate::error::Error::ExtValueDecode { message } = &error else {
345
+ panic!("expected ExtValueDecode");
346
+ };
347
+ assert_eq!(
348
+ message,
349
+ "new_logprobs: zero-column logprobs payload with 2 rows"
350
+ );
351
+}