[Bugfix][Rust Frontend] Limit chat top_logprobs in responses (#48134)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jul 16, 2026 at 17:41 UTC
8c3393f3733fa37d253c4d7a3126cfc9ae27b50d
4 files changed
+86
-5
rust/src/server/src/routes/openai/chat_completions.rs
+10
-1
@@ -125,6 +125,7 @@ async fn collect_chat_completion(
125
// Ignored: non-streaming responses are collected before usage is attached.
126
include_continuous_usage: _,
127
requested_logprobs,
128
+ output_top_logprobs,
129
include_prompt_logprobs,
130
include_reasoning,
131
echo,
@@ -173,6 +174,7 @@ async fn collect_chat_completion(
174
logprobs.as_ref().ok_or_else(|| {
175
server_error!("chat response requested logprobs but generation returned none")
176
})?,
177
+ output_top_logprobs,
178
return_tokens_as_token_ids,
179
)?)
180
} else {
@@ -248,6 +250,7 @@ async fn chat_completion_chunk_stream(
250
include_usage,
251
include_continuous_usage,
252
requested_logprobs,
253
+ output_top_logprobs,
254
// Ignored: chat streaming prompt logprobs are rejected for Python parity.
255
include_prompt_logprobs: _,
256
include_reasoning,
@@ -337,7 +340,13 @@ async fn chat_completion_chunk_stream(
340
let openai_logprobs = if include_metadata {
341
logprobs
342
.as_ref()
340
- .map(|lp| decoded_logprobs_to_openai_chat(lp, return_tokens_as_token_ids))
343
+ .map(|lp| {
344
+ decoded_logprobs_to_openai_chat(
345
+ lp,
346
+ output_top_logprobs,
347
+ return_tokens_as_token_ids,
348
+ )
349
+ })
350
.transpose()?
351
} else {
352
None
rust/src/server/src/routes/openai/chat_completions/convert.rs
+3
@@ -42,6 +42,8 @@ pub(super) struct ResponseOptions {
42
pub include_continuous_usage: bool,
43
/// Whether the caller requested output logprobs on chat choices.
44
pub requested_logprobs: bool,
45
+ /// Number of top logprobs to include for each output token.
46
+ pub output_top_logprobs: i32,
47
/// Whether the caller requested top-level prompt logprobs.
48
pub include_prompt_logprobs: bool,
49
/// Whether to include reasoning content in OpenAI responses.
@@ -172,6 +174,7 @@ pub(super) fn prepare_chat_request(
174
include_usage,
175
include_continuous_usage,
176
requested_logprobs,
177
+ output_top_logprobs: top_logprobs,
178
include_prompt_logprobs,
179
include_reasoning,
180
echo,
rust/src/server/src/routes/openai/utils/logprobs.rs
+65
-4
@@ -120,12 +120,13 @@ pub fn decoded_prompt_logprobs_to_maps(
120
/// shape.
121
pub fn decoded_logprobs_to_openai_chat(
122
logprobs: &DecodedLogprobs,
123
+ top_logprobs: i32,
124
return_tokens_as_token_ids: bool,
125
) -> Result<ChatLogProbs, ApiError> {
126
let content = logprobs
127
.positions
128
.iter()
128
- .map(|pos| position_to_chat_logprobs_content(pos, return_tokens_as_token_ids))
129
+ .map(|pos| position_to_chat_logprobs_content(pos, top_logprobs, return_tokens_as_token_ids))
130
.try_collect()?;
131
132
Ok(ChatLogProbs {
@@ -224,6 +225,7 @@ fn position_top_logprobs_map(
225
226
fn position_to_chat_logprobs_content(
227
position: &DecodedPositionLogprobs,
228
+ top_logprobs: i32,
229
return_tokens_as_token_ids: bool,
230
) -> Result<ChatLogProbsContent, ApiError> {
231
let chosen = position.entries.first().ok_or_else(|| {
@@ -235,9 +237,7 @@ fn position_to_chat_logprobs_content(
237
token: token_str.clone(),
238
logprob: clamp_logprob(chosen.logprob),
239
bytes: Some(token_bytes(&token_str)),
238
- top_logprobs: position
239
- .entries
240
- .iter()
240
+ top_logprobs: chat_top_logprob_entries(position, top_logprobs)
241
.map(|entry| {
242
let t = format_token(entry, return_tokens_as_token_ids);
243
TopLogProb {
@@ -250,6 +250,19 @@ fn position_to_chat_logprobs_content(
250
})
251
}
252
253
+fn chat_top_logprob_entries(
254
+ position: &DecodedPositionLogprobs,
255
+ top_logprobs: i32,
256
+) -> impl Iterator<Item = &DecodedTokenLogprob> {
257
+ let limit = if top_logprobs == -1 {
258
+ position.entries.len()
259
+ } else {
260
+ usize::try_from(top_logprobs).unwrap_or(0)
261
+ };
262
+
263
+ position.entries.iter().take(limit)
264
+}
265
+
266
fn token_bytes(token: &str) -> Vec<u8> {
267
token.as_bytes().to_vec()
268
}
@@ -257,3 +270,51 @@ fn token_bytes(token: &str) -> Vec<u8> {
270
pub fn clamp_logprob(logprob: f32) -> f32 {
271
logprob.max(-9999.0)
272
}
273
+
274
+#[cfg(test)]
275
+mod tests {
276
+ use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
277
+
278
+ use super::decoded_logprobs_to_openai_chat;
279
+
280
+ fn sample_logprobs() -> DecodedLogprobs {
281
+ DecodedLogprobs {
282
+ positions: vec![DecodedPositionLogprobs {
283
+ entries: vec![
284
+ DecodedTokenLogprob {
285
+ token_id: 1,
286
+ token: "A".to_string(),
287
+ logprob: -0.1,
288
+ rank: 1,
289
+ },
290
+ DecodedTokenLogprob {
291
+ token_id: 2,
292
+ token: "B".to_string(),
293
+ logprob: -1.0,
294
+ rank: 2,
295
+ },
296
+ DecodedTokenLogprob {
297
+ token_id: 3,
298
+ token: "C".to_string(),
299
+ logprob: -2.0,
300
+ rank: 3,
301
+ },
302
+ ],
303
+ }],
304
+ }
305
+ }
306
+
307
+ fn chat_top_logprobs_len(top_logprobs: i32) -> usize {
308
+ let chat_logprobs =
309
+ decoded_logprobs_to_openai_chat(&sample_logprobs(), top_logprobs, false)
310
+ .expect("chat logprobs");
311
+ chat_logprobs.content.expect("content")[0].top_logprobs.len()
312
+ }
313
+
314
+ #[test]
315
+ fn chat_logprobs_respects_requested_top_logprobs_count() {
316
+ assert_eq!(chat_top_logprobs_len(0), 0);
317
+ assert_eq!(chat_top_logprobs_len(1), 1);
318
+ assert_eq!(chat_top_logprobs_len(-1), 3);
319
+ }
320
+}
rust/src/server/src/routes/tests.rs
+8
@@ -2380,6 +2380,14 @@ async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() {
2380
json["choices"][0]["logprobs"]["content"][1]["token"],
2381
json!("i")
2382
);
2383
+ assert_eq!(
2384
+ json["choices"][0]["logprobs"]["content"][0]["top_logprobs"],
2385
+ json!([])
2386
+ );
2387
+ assert_eq!(
2388
+ json["choices"][0]["logprobs"]["content"][1]["top_logprobs"],
2389
+ json!([])
2390
+ );
2391
assert_eq!(json["prompt_logprobs"][0], serde_json::Value::Null);
2392
assert!(json["prompt_logprobs"][1].is_object());
2393
}