[Bugfix][Rust Frontend] Tolerate NaN-corrupted logprobs in engine-core (#51026)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Islam <islam.almersawi@openinnovation.ai> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Islam committed
Sep 16, 2026 at 07:10 UTC
c16bb6068f70878fb8a2f7c4d6cda95cd03a778b
3 files changed
+168
-4
rust/src/engine-core-client/src/protocol/logprobs.rs
+3
-4
@@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray};
25
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
26
pub struct TokenLogprob {
27
pub token_id: u32,
28
+ /// Preserves the engine's value, including NaN and infinities.
29
pub logprob: f32,
30
/// The sampled/selected token uses its actual vocab rank. Remaining entries
31
/// use 1-based top-k ranks matching the engine's returned candidate
32
/// order.
33
+ /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's
34
+ /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values.
35
pub rank: u32,
36
}
37
@@ -53,10 +56,6 @@ impl PositionLogprobs {
56
logprobs.len()
57
);
58
}
56
- if sampled_rank == 0 {
57
- bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs");
58
- }
59
-
59
let mut entries = Vec::with_capacity(token_ids.len());
60
for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() {
61
let rank = if index == 0 {
rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+118
@@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() {
396
"new_logprobs: zero-column logprobs payload with 2 rows"
397
);
398
}
399
+
400
+fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value {
401
+ let rows = ranks.len();
402
+ let cols = ids.len() / rows;
403
+ Value::Array(vec![
404
+ ndarray_value(
405
+ "<i8",
406
+ &[rows, cols],
407
+ Value::Ext(3, ids.iter().flat_map(|v| v.to_le_bytes()).collect()),
408
+ ),
409
+ ndarray_value(
410
+ "<f4",
411
+ &[rows, cols],
412
+ Value::Ext(3, probs.iter().flat_map(|v| v.to_le_bytes()).collect()),
413
+ ),
414
+ ndarray_value(
415
+ "<i8",
416
+ &[rows],
417
+ Value::Ext(3, ranks.iter().flat_map(|v| v.to_le_bytes()).collect()),
418
+ ),
419
+ Value::Nil,
420
+ ])
421
+}
422
+
423
+#[test]
424
+fn zero_sampled_rank_does_not_fail_frame_mates() {
425
+ // A batched EngineCoreOutputs frame where one request's row has rank 0
426
+ // must still resolve; the other requests' logprobs stay intact.
427
+ let frames = vec![Bytes::from(encode_value(&Value::Array(vec![
428
+ Value::from(0),
429
+ Value::Array(vec![
430
+ Value::Array(vec![
431
+ Value::from("req-nan"),
432
+ Value::Array(vec![Value::from(7), Value::from(8)]),
433
+ logprobs_value(&[1, 2], &[f32::NAN, f32::NEG_INFINITY], &[0]),
434
+ logprobs_value(&[7, 8], &[f32::NAN, f32::NEG_INFINITY], &[0]),
435
+ Value::Nil,
436
+ Value::from(EngineCoreFinishReason::Length as u8),
437
+ ]),
438
+ Value::Array(vec![
439
+ Value::from("req-ok"),
440
+ Value::Array(vec![Value::from(9)]),
441
+ logprobs_value(&[3, 4], &[-3.0, -0.5], &[7]),
442
+ Value::Nil,
443
+ Value::Nil,
444
+ Value::from(EngineCoreFinishReason::Length as u8),
445
+ ]),
446
+ ]),
447
+ Value::Nil,
448
+ Value::from(0.0),
449
+ Value::Nil,
450
+ Value::Nil,
451
+ ])))];
452
+ let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
453
+
454
+ let bad = decoded.outputs[0].new_logprobs.as_ref().unwrap();
455
+ let prompt = decoded.outputs[0].new_prompt_logprobs_tensors.as_ref().unwrap();
456
+ let good = decoded.outputs[1].new_logprobs.as_ref().unwrap();
457
+ expect_test::expect![[r#"
458
+ (
459
+ Logprobs {
460
+ positions: [
461
+ PositionLogprobs {
462
+ entries: [
463
+ TokenLogprob {
464
+ token_id: 1,
465
+ logprob: NaN,
466
+ rank: 0,
467
+ },
468
+ TokenLogprob {
469
+ token_id: 2,
470
+ logprob: -inf,
471
+ rank: 1,
472
+ },
473
+ ],
474
+ },
475
+ ],
476
+ },
477
+ Logprobs {
478
+ positions: [
479
+ PositionLogprobs {
480
+ entries: [
481
+ TokenLogprob {
482
+ token_id: 7,
483
+ logprob: NaN,
484
+ rank: 0,
485
+ },
486
+ TokenLogprob {
487
+ token_id: 8,
488
+ logprob: -inf,
489
+ rank: 1,
490
+ },
491
+ ],
492
+ },
493
+ ],
494
+ },
495
+ Logprobs {
496
+ positions: [
497
+ PositionLogprobs {
498
+ entries: [
499
+ TokenLogprob {
500
+ token_id: 3,
501
+ logprob: -3.0,
502
+ rank: 7,
503
+ },
504
+ TokenLogprob {
505
+ token_id: 4,
506
+ logprob: -0.5,
507
+ rank: 1,
508
+ },
509
+ ],
510
+ },
511
+ ],
512
+ },
513
+ )
514
+ "#]]
515
+ .assert_debug_eq(&(&**bad, &**prompt, &**good));
516
+}
rust/src/server/src/routes/openai/utils/logprobs.rs
+47
@@ -369,4 +369,51 @@ mod tests {
369
assert_eq!(chat_top_logprobs_len(1), 1);
370
assert_eq!(chat_top_logprobs_len(-1), 3);
371
}
372
+
373
+ #[test]
374
+ fn chat_logprobs_clamps_nan_and_negative_infinity_for_json() {
375
+ let mut logprobs = sample_logprobs();
376
+ logprobs.positions[0].entries[0].logprob = f32::NAN;
377
+ logprobs.positions[0].entries[0].rank = 0;
378
+ logprobs.positions[0].entries[1].logprob = f32::NEG_INFINITY;
379
+
380
+ let response = decoded_logprobs_to_openai_chat(&logprobs, -1, false).unwrap();
381
+ let json = serde_json::to_string_pretty(&response).unwrap();
382
+ expect_test::expect![[r#"
383
+ {
384
+ "content": [
385
+ {
386
+ "token": "A",
387
+ "logprob": -9999.0,
388
+ "bytes": [
389
+ 65
390
+ ],
391
+ "top_logprobs": [
392
+ {
393
+ "token": "A",
394
+ "logprob": -9999.0,
395
+ "bytes": [
396
+ 65
397
+ ]
398
+ },
399
+ {
400
+ "token": "B",
401
+ "logprob": -9999.0,
402
+ "bytes": [
403
+ 66
404
+ ]
405
+ },
406
+ {
407
+ "token": "C",
408
+ "logprob": -2.0,
409
+ "bytes": [
410
+ 67
411
+ ]
412
+ }
413
+ ]
414
+ }
415
+ ]
416
+ }"#]]
417
+ .assert_eq(&json);
418
+ }
419
}