[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>

rust/src/engine-core-client/src/protocol/logprobs.rs

@@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray};
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct TokenLogprob {
     pub token_id: u32,
+    /// Preserves the engine's value, including NaN and infinities.
     pub logprob: f32,
     /// The sampled/selected token uses its actual vocab rank. Remaining entries
     /// use 1-based top-k ranks matching the engine's returned candidate
     /// order.
+    /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's
+    /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values.
     pub rank: u32,
 }
 
@@ -53,10 +56,6 @@ impl PositionLogprobs {
                 logprobs.len()
             );
         }
-        if sampled_rank == 0 {
-            bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs");
-        }
-
         let mut entries = Vec::with_capacity(token_ids.len());
         for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() {
             let rank = if index == 0 {

rust/src/engine-core-client/src/protocol/logprobs/tests.rs

@@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() {
         "new_logprobs: zero-column logprobs payload with 2 rows"
     );
 }
+
+fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value {
+    let rows = ranks.len();
+    let cols = ids.len() / rows;
+    Value::Array(vec![
+        ndarray_value(
+            "<i8",
+            &[rows, cols],
+            Value::Ext(3, ids.iter().flat_map(|v| v.to_le_bytes()).collect()),
+        ),
+        ndarray_value(
+            "<f4",
+            &[rows, cols],
+            Value::Ext(3, probs.iter().flat_map(|v| v.to_le_bytes()).collect()),
+        ),
+        ndarray_value(
+            "<i8",
+            &[rows],
+            Value::Ext(3, ranks.iter().flat_map(|v| v.to_le_bytes()).collect()),
+        ),
+        Value::Nil,
+    ])
+}
+
+#[test]
+fn zero_sampled_rank_does_not_fail_frame_mates() {
+    // A batched EngineCoreOutputs frame where one request's row has rank 0
+    // must still resolve; the other requests' logprobs stay intact.
+    let frames = vec![Bytes::from(encode_value(&Value::Array(vec![
+        Value::from(0),
+        Value::Array(vec![
+            Value::Array(vec![
+                Value::from("req-nan"),
+                Value::Array(vec![Value::from(7), Value::from(8)]),
+                logprobs_value(&[1, 2], &[f32::NAN, f32::NEG_INFINITY], &[0]),
+                logprobs_value(&[7, 8], &[f32::NAN, f32::NEG_INFINITY], &[0]),
+                Value::Nil,
+                Value::from(EngineCoreFinishReason::Length as u8),
+            ]),
+            Value::Array(vec![
+                Value::from("req-ok"),
+                Value::Array(vec![Value::from(9)]),
+                logprobs_value(&[3, 4], &[-3.0, -0.5], &[7]),
+                Value::Nil,
+                Value::Nil,
+                Value::from(EngineCoreFinishReason::Length as u8),
+            ]),
+        ]),
+        Value::Nil,
+        Value::from(0.0),
+        Value::Nil,
+        Value::Nil,
+    ])))];
+    let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
+
+    let bad = decoded.outputs[0].new_logprobs.as_ref().unwrap();
+    let prompt = decoded.outputs[0].new_prompt_logprobs_tensors.as_ref().unwrap();
+    let good = decoded.outputs[1].new_logprobs.as_ref().unwrap();
+    expect_test::expect![[r#"
+        (
+            Logprobs {
+                positions: [
+                    PositionLogprobs {
+                        entries: [
+                            TokenLogprob {
+                                token_id: 1,
+                                logprob: NaN,
+                                rank: 0,
+                            },
+                            TokenLogprob {
+                                token_id: 2,
+                                logprob: -inf,
+                                rank: 1,
+                            },
+                        ],
+                    },
+                ],
+            },
+            Logprobs {
+                positions: [
+                    PositionLogprobs {
+                        entries: [
+                            TokenLogprob {
+                                token_id: 7,
+                                logprob: NaN,
+                                rank: 0,
+                            },
+                            TokenLogprob {
+                                token_id: 8,
+                                logprob: -inf,
+                                rank: 1,
+                            },
+                        ],
+                    },
+                ],
+            },
+            Logprobs {
+                positions: [
+                    PositionLogprobs {
+                        entries: [
+                            TokenLogprob {
+                                token_id: 3,
+                                logprob: -3.0,
+                                rank: 7,
+                            },
+                            TokenLogprob {
+                                token_id: 4,
+                                logprob: -0.5,
+                                rank: 1,
+                            },
+                        ],
+                    },
+                ],
+            },
+        )
+    "#]]
+    .assert_debug_eq(&(&**bad, &**prompt, &**good));
+}

rust/src/server/src/routes/openai/utils/logprobs.rs

@@ -369,4 +369,51 @@ mod tests {
         assert_eq!(chat_top_logprobs_len(1), 1);
         assert_eq!(chat_top_logprobs_len(-1), 3);
     }
+
+    #[test]
+    fn chat_logprobs_clamps_nan_and_negative_infinity_for_json() {
+        let mut logprobs = sample_logprobs();
+        logprobs.positions[0].entries[0].logprob = f32::NAN;
+        logprobs.positions[0].entries[0].rank = 0;
+        logprobs.positions[0].entries[1].logprob = f32::NEG_INFINITY;
+
+        let response = decoded_logprobs_to_openai_chat(&logprobs, -1, false).unwrap();
+        let json = serde_json::to_string_pretty(&response).unwrap();
+        expect_test::expect![[r#"
+            {
+              "content": [
+                {
+                  "token": "A",
+                  "logprob": -9999.0,
+                  "bytes": [
+                    65
+                  ],
+                  "top_logprobs": [
+                    {
+                      "token": "A",
+                      "logprob": -9999.0,
+                      "bytes": [
+                        65
+                      ]
+                    },
+                    {
+                      "token": "B",
+                      "logprob": -9999.0,
+                      "bytes": [
+                        66
+                      ]
+                    },
+                    {
+                      "token": "C",
+                      "logprob": -2.0,
+                      "bytes": [
+                        67
+                      ]
+                    }
+                  ]
+                }
+              ]
+            }"#]]
+        .assert_eq(&json);
+    }
 }