diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 61eb21ef70..aea04495eb 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/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>) { +/// - `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, Option>) { 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 = (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 {