[Bugfix][Rust Frontend] Preserve selected-token logprob mode (#56406)
Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Alec committed
Sep 14, 2026 at 00:05 UTC
e6b4e47d2d894a1a0815fbf03669adb1b3405083
1 file changed
+84
-14
rust/src/server/src/grpc/convert.rs
+84
-14
@@ -193,7 +193,7 @@ fn build_sampling_params(
193
if let Some(r) = response {
194
if r.output_logprobs {
195
let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref());
196
- params.logprobs = Some(count);
196
+ params.logprobs = count;
197
params.logprob_token_ids = token_ids;
198
}
199
if r.prompt_logprobs {
@@ -210,7 +210,7 @@ fn build_sampling_params(
210
));
211
}
212
let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref());
213
- params.prompt_logprobs = Some(count);
213
+ params.prompt_logprobs = count;
214
}
215
}
216
@@ -220,18 +220,22 @@ fn build_sampling_params(
220
/// Map the proto `CandidateTokens` selector to a `(logprobs_count,
221
/// logprob_token_ids)` pair.
222
///
223
-/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability
224
-/// - `all` → `(-1, None)` — return the full vocabulary
225
-/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the
226
-/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs
227
-/// are carried via `logprob_token_ids` on `SamplingParams`)
228
-/// - absent → `(1, None)` — just the sampled/scored token
229
-fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option<Vec<u32>>) {
223
+/// - `top_n(k)` → `(Some(k), None)` — return top-k candidates by probability
224
+/// - `all` → `(Some(-1), None)` — return the full vocabulary
225
+/// - nonempty `token_ids` → `(None, Some(ids))` — let the engine derive the count
226
+/// - absent or empty `token_ids` → `(Some(0), None)` — return only the sampled/scored token
227
+fn candidate_logprob_spec(
228
+ candidates: Option<&pb::CandidateTokens>,
229
+) -> (Option<i32>, Option<Vec<u32>>) {
230
match candidates.and_then(|c| c.select.as_ref()) {
231
- Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None),
232
- Some(pb::candidate_tokens::Select::All(true)) => (-1, None),
233
- Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())),
234
- _ => (1, None),
231
+ Some(pb::candidate_tokens::Select::TopN(n)) => (Some(*n as i32), None),
232
+ Some(pb::candidate_tokens::Select::All(true)) => (Some(-1), None),
233
+ Some(pb::candidate_tokens::Select::TokenIds(ids)) if !ids.ids.is_empty() => {
234
+ // Match Python HTTP: selected IDs have their own limit, independent
235
+ // of the numeric top-logprobs count and its max_logprobs cap.
236
+ (None, Some(ids.ids.clone()))
237
+ }
238
+ _ => (Some(0), None),
239
}
240
}
241
@@ -518,7 +522,10 @@ impl ResponseOpts {
522
#[cfg(test)]
523
mod tests {
524
use vllm_engine_core_client::protocol::output::StopReason;
521
- use vllm_text::{FinishReason, Finished, Prompt};
525
+ use vllm_text::{
526
+ FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params,
527
+ };
528
+ use vllm_tokenizer::test_utils::TestTokenizer;
529
530
use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason};
531
use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request};
@@ -590,6 +597,69 @@ mod tests {
597
assert_eq!(text.sampling_params.seed, Some(0));
598
}
599
600
+ #[test]
601
+ fn output_logprob_selectors_survive_request_lowering() {
602
+ use pb::candidate_tokens::Select;
603
+
604
+ let selected_ids: Vec<u32> = (100..121).collect();
605
+ let cases = [
606
+ (
607
+ Some(Select::TokenIds(pb::TokenIds {
608
+ ids: selected_ids.clone(),
609
+ })),
610
+ None,
611
+ Some(selected_ids),
612
+ ),
613
+ (
614
+ Some(Select::TokenIds(pb::TokenIds {
615
+ ids: vec![198, 198],
616
+ })),
617
+ None,
618
+ Some(vec![198, 198]),
619
+ ),
620
+ (
621
+ Some(Select::TokenIds(pb::TokenIds { ids: vec![] })),
622
+ Some(0),
623
+ None,
624
+ ),
625
+ (None, Some(0), None),
626
+ (Some(Select::TopN(0)), Some(0), None),
627
+ (Some(Select::TopN(2)), Some(2), None),
628
+ ];
629
+ for (select, expected_count, expected_ids) in cases {
630
+ let req = pb::GenerateRequest {
631
+ response: Some(pb::ResponseOptions {
632
+ output_logprobs: true,
633
+ output_candidates: select.map(|select| pb::CandidateTokens {
634
+ select: Some(select),
635
+ }),
636
+ ..Default::default()
637
+ }),
638
+ ..base_request()
639
+ };
640
+ let text =
641
+ to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
642
+ let params = lower_sampling_params(
643
+ text.sampling_params,
644
+ SamplingHints::default(),
645
+ SamplingLimits {
646
+ max_model_len: 32,
647
+ max_logprobs: 20,
648
+ model_vocab_size: 512,
649
+ tokenizer_vocab_size: 512,
650
+ },
651
+ 1,
652
+ &TestTokenizer::new(),
653
+ )
654
+ .expect("logprob selector should lower to an engine request");
655
+
656
+ assert_eq!(
657
+ (params.logprobs, params.logprob_token_ids),
658
+ (expected_count, expected_ids)
659
+ );
660
+ }
661
+ }
662
+
663
#[test]
664
fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() {
665
let req = pb::GenerateRequest {