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

rust/src/server/src/grpc/convert.rs

@@ -193,7 +193,7 @@ fn build_sampling_params(
     if let Some(r) = response {
         if r.output_logprobs {
             let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref());
-            params.logprobs = Some(count);
+            params.logprobs = count;
             params.logprob_token_ids = token_ids;
         }
         if r.prompt_logprobs {
@@ -210,7 +210,7 @@ fn build_sampling_params(
                 ));
             }
             let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref());
-            params.prompt_logprobs = Some(count);
+            params.prompt_logprobs = count;
         }
     }
 
@@ -220,18 +220,22 @@ fn build_sampling_params(
 /// Map the proto `CandidateTokens` selector to a `(logprobs_count,
 /// logprob_token_ids)` pair.
 ///
-/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability
-/// - `all` → `(-1, None)` — return the full vocabulary
-/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the
-///   count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs
-///   are carried via `logprob_token_ids` on `SamplingParams`)
-/// - absent → `(1, None)` — just the sampled/scored token
-fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option<Vec<u32>>) {
+/// - `top_n(k)` → `(Some(k), None)` — return top-k candidates by probability
+/// - `all` → `(Some(-1), None)` — return the full vocabulary
+/// - nonempty `token_ids` → `(None, Some(ids))` — let the engine derive the count
+/// - absent or empty `token_ids` → `(Some(0), None)` — return only the sampled/scored token
+fn candidate_logprob_spec(
+    candidates: Option<&pb::CandidateTokens>,
+) -> (Option<i32>, Option<Vec<u32>>) {
     match candidates.and_then(|c| c.select.as_ref()) {
-        Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None),
-        Some(pb::candidate_tokens::Select::All(true)) => (-1, None),
-        Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())),
-        _ => (1, None),
+        Some(pb::candidate_tokens::Select::TopN(n)) => (Some(*n as i32), None),
+        Some(pb::candidate_tokens::Select::All(true)) => (Some(-1), None),
+        Some(pb::candidate_tokens::Select::TokenIds(ids)) if !ids.ids.is_empty() => {
+            // Match Python HTTP: selected IDs have their own limit, independent
+            // of the numeric top-logprobs count and its max_logprobs cap.
+            (None, Some(ids.ids.clone()))
+        }
+        _ => (Some(0), None),
     }
 }
 
@@ -518,7 +522,10 @@ impl ResponseOpts {
 #[cfg(test)]
 mod tests {
     use vllm_engine_core_client::protocol::output::StopReason;
-    use vllm_text::{FinishReason, Finished, Prompt};
+    use vllm_text::{
+        FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params,
+    };
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason};
     use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request};
@@ -590,6 +597,69 @@ mod tests {
         assert_eq!(text.sampling_params.seed, Some(0));
     }
 
+    #[test]
+    fn output_logprob_selectors_survive_request_lowering() {
+        use pb::candidate_tokens::Select;
+
+        let selected_ids: Vec<u32> = (100..121).collect();
+        let cases = [
+            (
+                Some(Select::TokenIds(pb::TokenIds {
+                    ids: selected_ids.clone(),
+                })),
+                None,
+                Some(selected_ids),
+            ),
+            (
+                Some(Select::TokenIds(pb::TokenIds {
+                    ids: vec![198, 198],
+                })),
+                None,
+                Some(vec![198, 198]),
+            ),
+            (
+                Some(Select::TokenIds(pb::TokenIds { ids: vec![] })),
+                Some(0),
+                None,
+            ),
+            (None, Some(0), None),
+            (Some(Select::TopN(0)), Some(0), None),
+            (Some(Select::TopN(2)), Some(2), None),
+        ];
+        for (select, expected_count, expected_ids) in cases {
+            let req = pb::GenerateRequest {
+                response: Some(pb::ResponseOptions {
+                    output_logprobs: true,
+                    output_candidates: select.map(|select| pb::CandidateTokens {
+                        select: Some(select),
+                    }),
+                    ..Default::default()
+                }),
+                ..base_request()
+            };
+            let text =
+                to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+            let params = lower_sampling_params(
+                text.sampling_params,
+                SamplingHints::default(),
+                SamplingLimits {
+                    max_model_len: 32,
+                    max_logprobs: 20,
+                    model_vocab_size: 512,
+                    tokenizer_vocab_size: 512,
+                },
+                1,
+                &TestTokenizer::new(),
+            )
+            .expect("logprob selector should lower to an engine request");
+
+            assert_eq!(
+                (params.logprobs, params.logprob_token_ids),
+                (expected_count, expected_ids)
+            );
+        }
+    }
+
     #[test]
     fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() {
         let req = pb::GenerateRequest {