[Bugfix][Rust Frontend] Map missing prompt logprobs for single-token prompts in chat and raw generate (#49111)
Signed-off-by: Feathbow <feathbow@gmail.com>
Salt Sato committed
Jul 20, 2026 at 11:00 UTC
d835ad572ca76d8ed6308e81750813c814fb3343
4 files changed
+94
-53
rust/src/server/src/routes/inference/generate.rs
+56
-7
@@ -238,13 +238,18 @@ fn collect_generate(
238
None
239
};
240
let prompt_logprobs = if include_prompt_logprobs {
241
- let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
242
- ApiError::server_error(
243
- "raw generate response requested prompt_logprobs but generation returned none"
244
- .to_string(),
245
- )
246
- })?;
247
- Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
241
+ match collected.prompt_logprobs.as_ref() {
242
+ Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
243
+ // A single-token prompt has no scored positions; same mapping
244
+ // as /v1/completions.
245
+ None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
246
+ None => {
247
+ return Err(ApiError::server_error(
248
+ "raw generate response requested prompt_logprobs but generation returned none"
249
+ .to_string(),
250
+ ));
251
+ }
252
+ }
253
} else {
254
None
255
};
@@ -472,4 +477,48 @@ mod tests {
477
Some(2)
478
);
479
}
480
+
481
+ #[test]
482
+ fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
483
+ let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
484
+ request_id: "raw-1".to_string(),
485
+ prompt_logprobs: None,
486
+ token_ids: vec![3],
487
+ logprobs: None,
488
+ finish_reason: FinishReason::stop_eos(),
489
+ usage: vllm_llm::TokenUsage {
490
+ prompt_token_count: prompt_token_ids.len(),
491
+ output_token_count: 1,
492
+ cached_token_count: 0,
493
+ },
494
+ kv_transfer_params: None,
495
+ ec_transfer_params: None,
496
+ prompt_token_ids,
497
+ };
498
+
499
+ let response = collect_generate(
500
+ output_without_payload(vec![9707]),
501
+ "raw-1".to_string(),
502
+ ApiServerOptions::default(),
503
+ ResponseOptions {
504
+ include_prompt_logprobs: true,
505
+ ..Default::default()
506
+ },
507
+ )
508
+ .expect("single-token prompt without payload maps to [None]");
509
+ let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
510
+ assert_eq!(prompt_logprobs.len(), 1);
511
+ assert!(prompt_logprobs[0].is_none());
512
+
513
+ collect_generate(
514
+ output_without_payload(vec![9707, 11]),
515
+ "raw-2".to_string(),
516
+ ApiServerOptions::default(),
517
+ ResponseOptions {
518
+ include_prompt_logprobs: true,
519
+ ..Default::default()
520
+ },
521
+ )
522
+ .expect_err("multi-token prompt without payload is an engine failure");
523
+ }
524
}
rust/src/server/src/routes/openai/chat_completions.rs
+5
-8
@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
35
ChatMessageDelta,
36
};
37
use crate::routes::openai::utils::logprobs::{
38
- decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
38
+ decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
39
};
40
use crate::routes::openai::utils::types::{
41
ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
181
None
182
};
183
let prompt_logprobs = if include_prompt_logprobs {
184
- Some(decoded_prompt_logprobs_to_maps(
185
- prompt_logprobs.as_ref().ok_or_else(|| {
186
- server_error!(
187
- "chat response requested prompt_logprobs but generation returned none"
188
- )
189
- })?,
184
+ Some(prompt_logprobs_to_maps(
185
+ prompt_logprobs.as_ref(),
186
+ &prompt_token_ids,
187
return_tokens_as_token_ids,
191
- ))
188
+ )?)
189
} else {
190
None
191
};
rust/src/server/src/routes/openai/completions.rs
+2
-24
@@ -5,7 +5,6 @@ mod convert;
5
mod types;
6
mod validate;
7
8
-use std::collections::HashMap;
8
use std::convert::Infallible;
9
use std::result::Result;
10
use std::sync::Arc;
@@ -29,8 +28,8 @@ use vllm_text::{
28
29
use self::convert::{ResponseOptions, prepare_completion_request};
30
use super::utils::logprobs::{
32
- collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
33
- decoded_prompt_logprobs_to_openai, text_len,
31
+ collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
32
+ prompt_logprobs_to_maps, text_len,
33
};
34
use super::utils::types::Usage;
35
use crate::config::ApiServerOptions;
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
504
))
505
}
506
508
-fn prompt_logprobs_to_maps(
509
- prompt_logprobs: Option<&DecodedPromptLogprobs>,
510
- prompt_token_ids: &[u32],
511
- return_tokens_as_token_ids: bool,
512
-) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
513
- if let Some(prompt_logprobs) = prompt_logprobs {
514
- return Ok(decoded_prompt_logprobs_to_maps(
515
- prompt_logprobs,
516
- return_tokens_as_token_ids,
517
- ));
518
- }
519
-
520
- if let [_token_id] = prompt_token_ids {
521
- return Ok(vec![None]);
522
- }
523
-
524
- Err(server_error!(
525
- "completion response requested prompt_logprobs but generation returned none"
526
- ))
527
-}
528
-
507
fn usage_chunk(
508
request_id: &str,
509
response_model: &str,
rust/src/server/src/routes/openai/utils/logprobs.rs
+31
-14
@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
100
})
101
}
102
103
-/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
104
-/// shape.
105
-pub fn decoded_prompt_logprobs_to_maps(
106
- prompt_logprobs: &DecodedPromptLogprobs,
103
+/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
104
+/// missing single-token payload as `[None]`.
105
+pub fn prompt_logprobs_to_maps(
106
+ prompt_logprobs: Option<&DecodedPromptLogprobs>,
107
+ prompt_token_ids: &[u32],
108
return_tokens_as_token_ids: bool,
108
-) -> Vec<Option<HashMap<String, f32>>> {
109
- std::iter::once(None)
110
- .chain(prompt_logprobs.scored_positions.iter().map(|position| {
111
- Some(position_top_logprobs_map(
112
- position,
113
- return_tokens_as_token_ids,
114
- ))
115
- }))
116
- .collect()
109
+) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
110
+ if let Some(prompt_logprobs) = prompt_logprobs {
111
+ return Ok(std::iter::once(None)
112
+ .chain(prompt_logprobs.scored_positions.iter().map(|position| {
113
+ Some(position_top_logprobs_map(
114
+ position,
115
+ return_tokens_as_token_ids,
116
+ ))
117
+ }))
118
+ .collect());
119
+ }
120
+
121
+ if let [_token_id] = prompt_token_ids {
122
+ return Ok(vec![None]);
123
+ }
124
+
125
+ Err(server_error!(
126
+ "prompt_logprobs were requested but generation returned none"
127
+ ))
128
}
129
130
/// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
286
mod tests {
287
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
288
278
- use super::decoded_logprobs_to_openai_chat;
289
+ use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
290
+
291
+ #[test]
292
+ fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
293
+ prompt_logprobs_to_maps(None, &[9707, 11], false)
294
+ .expect_err("multi-token prompt without payload is an engine failure");
295
+ }
296
297
fn sample_logprobs() -> DecodedLogprobs {
298
DecodedLogprobs {