@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b721a4c709

[Feature] Add fixed-token prefill scoring (#54335)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

aoshen02 committed Sep 28, 2026 at 22:50 UTC b721a4c709da44551b614e2562fcfa7f4cbd1352
36 files changed +768 -8
docs/training/prompt_token_id_logprobs.md new
+64
@@ -0,0 +1,64 @@
1 +# Prompt Token ID Logprobs (Teacher Scoring)
2 +
3 +In on-policy distillation, the student records its top-K candidate tokens at
4 +every response position, and the teacher provides the log probabilities of
5 +exactly those candidates as the distillation target. The caller passes the
6 +candidates as `prompt_logprob_token_ids`; the teacher runs a single prefill
7 +over the prompt plus response and gathers their log probabilities on the GPU at
8 +every scored row, starting from `prompt_logprob_start` (typically the prompt
9 +length minus one, so only response positions are scored). The result is a
10 +`[rows, K]` matrix whose row `i` scores prompt token
11 +`prompt_logprob_start + i + 1` and whose columns follow the candidate order, so
12 +it lines up position by position with the student's top-K.
13 +
14 +## Quick start
15 +
16 +```python
17 +from vllm import LLM, SamplingParams
18 +from vllm.inputs import TokensPrompt
19 +
20 +llm = LLM(model)
21 +output = llm.generate(
22 + TokensPrompt(prompt_token_ids=prompt_ids + response_ids),
23 + SamplingParams(
24 + max_tokens=1,
25 + prompt_logprob_token_ids=candidate_ids,
26 + prompt_logprob_start=len(prompt_ids) - 1,
27 + ),
28 +)
29 +scores = output[0].prompt_token_id_logprobs
30 +# scores.shape == (len(response_ids), len(candidate_ids))
31 +# scores[i, j] = log p(candidate_ids[j] | prompt_ids + response_ids[:i])
32 +```
33 +
34 +The same parameters are accepted by the `/inference/v1/generate` HTTP endpoint
35 +(Python and Rust frontends), which returns the matrix as a base64-encoded
36 +`.npy` float32 array:
37 +
38 +```python
39 +import base64, io
40 +import numpy as np
41 +
42 +scores = np.load(io.BytesIO(base64.b64decode(response["prompt_token_id_logprobs"])))
43 +```
44 +
45 +The matrix has `max(prompt_len - 1 - prompt_logprob_start, 0)` rows. With a
46 +logits `--logprobs-mode`, it holds logits instead of log probabilities.
47 +
48 +## Requirements
49 +
50 +- The V2 model runner (`VLLM_USE_V2_MODEL_RUNNER=1`).
51 +- `len(prompt_logprob_token_ids)` at most `--max-logprobs`.
52 +- No `--kv-sharing-fast-prefill`.
53 +
54 +The request skips reading the prefix cache (local and KV connector), since
55 +cached rows have no logits to score.
56 +
57 +## Limitations
58 +
59 +- One candidate set per request, shared by every scored row; per-position
60 + candidates are planned ([#56860](https://github.com/vllm-project/vllm/issues/56860)).
61 +- Non-streaming only; `stream=true` requests are rejected.
62 +- Not exposed through the OpenAI-compatible endpoints.
63 +- If a prefill starts past the first scored row (e.g. `skip_reading_prefix_cache=False`
64 + with a cache hit), the result is `None` rather than a partial matrix.
rust/Cargo.lock
+1
@@ -5968,6 +5968,7 @@ dependencies = [
5968 "asynk-strim-attr",
5969 "auto_enums",
5970 "axum",
5971 + "base64 0.22.1",
5972 "bytes",
5973 "clap",
5974 "educe",
rust/src/engine-core-client/src/protocol/output.rs
+11
@@ -17,6 +17,7 @@ use crate::error::{Error, Result, ext_value_decode};
17 use crate::protocol::logprobs::MaybeWireLogprobs;
18 use crate::protocol::sampling_mask::MaybeWireSamplingMask;
19 use crate::protocol::stats::{PrefillStats, SchedulerStats};
20 +use crate::protocol::tensor::WireNdArray;
21 use crate::protocol::{OpaqueValue, decode_msgpack};
22
23 /// The stop reason associated with a finished output.
@@ -133,6 +134,10 @@ pub struct EngineCoreOutput {
134 /// output when `--per-request-spec-decode-metrics` is enabled.
135 #[serde(default)]
136 pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
137 + /// Log probabilities of `SamplingParams.prompt_logprob_token_ids`, set on
138 + /// the first output of a request that asked for them.
139 + #[serde(default)]
140 + pub prompt_token_id_logprobs: Option<WireNdArray>,
141 }
142
143 /// Raw per-sequence speculative-decoding accumulator.
@@ -180,6 +185,11 @@ impl EngineCoreOutput {
185 self.new_sampling_mask = (self.new_sampling_mask.take())
186 .map(|value| value.resolve(frames, "new_sampling_mask"))
187 .transpose()?;
188 + if let Some(value) = self.prompt_token_id_logprobs.as_mut() {
189 + value
190 + .resolve_aux_frame(frames)
191 + .map_err(|message| ext_value_decode!("prompt_token_id_logprobs: {message}"))?;
192 + }
193 Ok(())
194 }
195 }
@@ -531,6 +541,7 @@ mod tests {
541 mm_cache_miss_hashes: None,
542 new_sampling_mask: None,
543 spec_decode_metrics: None,
544 + prompt_token_id_logprobs: None,
545 },
546 ],
547 scheduler_stats: None,
rust/src/engine-core-client/src/protocol/sampling.rs
+7
@@ -101,6 +101,11 @@ pub struct EngineCoreSamplingParams {
101 ///
102 /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
103 pub prompt_logprobs: Option<i32>,
104 + /// Candidate token IDs scored at every scored causal prompt row, where row
105 + /// `i` scores them as predictions of prompt token `i + 1`.
106 + pub prompt_logprob_token_ids: Option<Vec<u32>>,
107 + /// First causal prompt row to score; `None` scores from the first row.
108 + pub prompt_logprob_start: Option<u32>,
109 /// Minimum probability threshold for token sampling.
110 pub min_p: f32,
111 /// Frequency penalty applied by the sampler.
@@ -171,6 +176,8 @@ impl EngineCoreSamplingParams {
176 thinking_token_budget: None,
177 logprobs: None,
178 prompt_logprobs: None,
179 + prompt_logprob_token_ids: None,
180 + prompt_logprob_start: None,
181 min_p: 0.0,
182 frequency_penalty: 0.0,
183 presence_penalty: 0.0,
rust/src/engine-core-client/src/tests/client.rs
+3
@@ -2684,6 +2684,8 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2684 thinking_token_budget: None,
2685 logprobs: None,
2686 prompt_logprobs: None,
2687 + prompt_logprob_token_ids: None,
2688 + prompt_logprob_start: None,
2689 min_p: 0.0,
2690 frequency_penalty: 0.0,
2691 presence_penalty: 0.0,
@@ -2770,6 +2772,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2772 mm_cache_miss_hashes: None,
2773 new_sampling_mask: None,
2774 spec_decode_metrics: None,
2775 + prompt_token_id_logprobs: None,
2776 },
2777 ],
2778 scheduler_stats: None,
rust/src/engine-core-client/src/tests/python_compat.py
+1
@@ -116,6 +116,7 @@ class EngineCoreOutput(
116 mm_cache_miss_hashes: list[str] | None = None
117 new_sampling_mask: object | None = None
118 spec_decode_metrics: object | None = None
119 + prompt_token_id_logprobs: object | None = None
120
121
122 class ExtendedEngineCoreOutput(EngineCoreOutput):
rust/src/llm/src/output.rs
+16
@@ -15,6 +15,7 @@ use vllm_engine_core_client::protocol::output::{
15 EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason,
16 };
17 use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
18 +use vllm_engine_core_client::protocol::tensor::WireNdArray;
19 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
20
21 use crate::error::Result;
@@ -38,6 +39,7 @@ pub struct CollectedGenerateOutput {
39 pub request_id: String,
40 pub prompt_token_ids: Vec<u32>,
41 pub prompt_logprobs: Option<Logprobs>,
42 + pub prompt_token_id_logprobs: Option<WireNdArray>,
43 pub token_ids: Vec<u32>,
44 pub logprobs: Option<Logprobs>,
45 pub finish_reason: FinishReason,
@@ -62,6 +64,8 @@ pub struct GeneratePromptInfo {
64 /// Prompt logprobs returned by engine-core for scored prompt positions,
65 /// when requested.
66 pub prompt_logprobs: Option<Logprobs>,
67 + /// Log probabilities of `prompt_logprob_token_ids`, when requested.
68 + pub prompt_token_id_logprobs: Option<WireNdArray>,
69 }
70
71 /// The reason a request finished.
@@ -206,6 +210,7 @@ impl GenerateOutput {
210 prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
211 prompt_token_ids: ids,
212 prompt_logprobs: None,
213 + prompt_token_id_logprobs: None,
214 }),
215 token_ids,
216 logprobs: None,
@@ -247,6 +252,7 @@ impl GenerateOutputStream {
252 pending_prompt_info: Some(GeneratePromptInfo {
253 prompt_token_ids,
254 prompt_logprobs: None,
255 + prompt_token_id_logprobs: None,
256 }),
257 raw_stream,
258 request_metrics,
@@ -282,6 +288,11 @@ impl Stream for GenerateOutputStream {
288 info.prompt_logprobs =
289 raw.new_prompt_logprobs_tensors.map(|value| value.into_direct().unwrap());
290 }
291 + if let Some(info) = &mut self.pending_prompt_info
292 + && info.prompt_token_id_logprobs.is_none()
293 + {
294 + info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs;
295 + }
296
297 let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
298 let sampling_mask = raw.new_sampling_mask.map(|value| value.into_direct().unwrap());
@@ -361,6 +372,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
372 pin_mut!(stream);
373 let mut prompt_token_ids = None;
374 let mut prompt_logprobs = None;
375 + let mut prompt_token_id_logprobs = None;
376 let mut cached_token_count = 0;
377 let mut collected: Option<CollectedGenerateOutput> = None;
378
@@ -374,6 +386,9 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
386 if prompt_logprobs.is_none() {
387 prompt_logprobs = info.prompt_logprobs;
388 }
389 + if prompt_token_id_logprobs.is_none() {
390 + prompt_token_id_logprobs = info.prompt_token_id_logprobs;
391 + }
392 }
393
394 if let Some(existing) = collected.as_mut() {
@@ -393,6 +408,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
408 request_id: output.request_id,
409 prompt_token_ids: prompt_token_ids.take().unwrap_or_default(),
410 prompt_logprobs: prompt_logprobs.take(),
411 + prompt_token_id_logprobs: prompt_token_id_logprobs.take(),
412 token_ids: output.token_ids,
413 logprobs: output.logprobs,
414 finish_reason: FinishReason::Error,
rust/src/llm/tests/generate.rs
+2
@@ -277,6 +277,7 @@ async fn generate_streams_outputs() {
277 Some(GeneratePromptInfo {
278 prompt_token_ids: vec![11, 22].into(),
279 prompt_logprobs: Some(prompt_logprobs()),
280 + prompt_token_id_logprobs: None,
281 })
282 );
283 assert_eq!(first.token_ids, vec![1, 2]);
@@ -405,6 +406,7 @@ async fn collect_output_rejects_partial_sampling_mask() {
406 prompt_info: Some(GeneratePromptInfo {
407 prompt_token_ids: Arc::from([11_u32, 22]),
408 prompt_logprobs: None,
409 + prompt_token_id_logprobs: None,
410 }),
411 token_ids: vec![33, 44],
412 logprobs: None,
rust/src/server/Cargo.toml
+1
@@ -9,6 +9,7 @@ anyhow.workspace = true
9 asynk-strim-attr.workspace = true
10 auto_enums.workspace = true
11 axum.workspace = true
12 +base64.workspace = true
13 educe.workspace = true
14 futures.workspace = true
15 http-body.workspace = true
rust/src/server/src/routes/inference/generate.rs
+48
@@ -16,12 +16,15 @@ use axum::extract::State;
16 use axum::http::HeaderMap;
17 use axum::response::sse::Event;
18 use axum::response::{IntoResponse, Response};
19 +use base64::Engine as _;
20 use futures::{Stream, StreamExt as _, pin_mut};
21 use thiserror_ext::AsReport as _;
22 use tracing::{error, info, trace};
23 use tracing_futures::Instrument as _;
24 +use vllm_engine_core_client::protocol::dtype::TensorDtype;
25 use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
26 use vllm_engine_core_client::protocol::multimodal::MmFeatureSpec;
27 +use vllm_engine_core_client::protocol::tensor::WireNdArray;
28 use vllm_llm::{
29 CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage,
30 };
@@ -330,6 +333,11 @@ fn collect_generate(
333 token_ids: collected.token_ids,
334 }],
335 prompt_logprobs,
336 + prompt_token_id_logprobs: collected
337 + .prompt_token_id_logprobs
338 + .as_ref()
339 + .map(npy_base64)
340 + .transpose()?,
341 prompt_token_ids: return_token_ids.then_some(collected.prompt_token_ids),
342 mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(),
343 kv_transfer_params: collected.kv_transfer_params,
@@ -432,6 +440,29 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
440 .collect()
441 }
442
443 +/// Encode a float32 matrix as Python `numpy2base64` does.
444 +fn npy_base64(array: &WireNdArray) -> Result<String, ApiError> {
445 + let (TensorDtype::F32, Some(data), [rows, cols]) = (
446 + array.dtype.scalar,
447 + array.data.as_raw_view(),
448 + array.shape.as_slice(),
449 + ) else {
450 + bail_server_error!("prompt_token_id_logprobs: expected a float32 matrix");
451 + };
452 + let mut header =
453 + format!("{{'descr': '<f4', 'fortran_order': False, 'shape': ({rows}, {cols}), }}");
454 + header.push_str(&" ".repeat(21 - rows.to_string().len()));
455 + header.push_str(&" ".repeat(64 - (11 + header.len()) % 64));
456 + header.push('\n');
457 +
458 + let mut npy = Vec::with_capacity(10 + header.len() + data.len());
459 + npy.extend_from_slice(b"\x93NUMPY\x01\x00");
460 + npy.extend_from_slice(&(header.len() as u16).to_le_bytes());
461 + npy.extend_from_slice(header.as_bytes());
462 + npy.extend_from_slice(data);
463 + Ok(base64::engine::general_purpose::STANDARD.encode(npy))
464 +}
465 +
466 fn format_token_id(token_id: u32) -> String {
467 format!("token_id:{token_id}")
468 }
@@ -509,6 +540,7 @@ mod tests {
540 prompt_info: Some(GeneratePromptInfo {
541 prompt_token_ids: Arc::from([11_u32, 22_u32]),
542 prompt_logprobs: None,
543 + prompt_token_id_logprobs: None,
544 }),
545 token_ids: vec![33],
546 logprobs: None,
@@ -592,6 +624,7 @@ mod tests {
624 prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
625 prompt_token_ids: Arc::from(ids),
626 prompt_logprobs: None,
627 + prompt_token_id_logprobs: None,
628 }),
629 token_ids,
630 logprobs: None,
@@ -747,6 +780,7 @@ mod tests {
780 let output = CollectedGenerateOutput {
781 request_id: "raw-1".to_string(),
782 prompt_logprobs: None,
783 + prompt_token_id_logprobs: None,
784 token_ids: vec![30],
785 logprobs: None,
786 finish_reason: FinishReason::stop_eos(),
@@ -781,6 +815,17 @@ mod tests {
815 assert!(response.mm_placeholders.is_none());
816 }
817
818 + #[test]
819 + fn prompt_token_id_logprobs_match_numpy2base64() {
820 + let scores =
821 + WireNdArray::from_f32(vec![2, 3], vec![-0.5, -1.5, -2.5, -3.5, -4.5, -5.5]).unwrap();
822 + // numpy2base64(np.array([[-0.5, -1.5, -2.5], [-3.5, -4.5, -5.5]], np.float32))
823 + assert_eq!(
824 + npy_base64(&scores).unwrap(),
825 + "k05VTVBZAQB2AHsnZGVzY3InOiAnPGY0JywgJ2ZvcnRyYW5fb3JkZXInOiBGYWxzZSwgJ3NoYXBlJzogKDIsIDMpLCB9ICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIAoAAAC/AADAvwAAIMAAAGDAAACQwAAAsMA="
826 + );
827 + }
828 +
829 #[test]
830 fn collect_generate_returns_spec_decode_metrics() {
831 let metrics = RequestSpecDecodeMetrics {
@@ -792,6 +837,7 @@ mod tests {
837 let output = CollectedGenerateOutput {
838 request_id: "raw-metrics".to_string(),
839 prompt_logprobs: None,
840 + prompt_token_id_logprobs: None,
841 token_ids: vec![30],
842 logprobs: None,
843 finish_reason: FinishReason::stop_eos(),
@@ -879,6 +925,7 @@ mod tests {
925 let output = CollectedGenerateOutput {
926 request_id: "raw-1".to_string(),
927 prompt_logprobs: None,
928 + prompt_token_id_logprobs: None,
929 token_ids: vec![30],
930 logprobs: None,
931 finish_reason: FinishReason::stop_eos(),
@@ -914,6 +961,7 @@ mod tests {
961 let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
962 request_id: "raw-1".to_string(),
963 prompt_logprobs: None,
964 + prompt_token_id_logprobs: None,
965 token_ids: vec![3],
966 logprobs: None,
967 finish_reason: FinishReason::stop_eos(),
rust/src/server/src/routes/inference/generate/types.rs
+2
@@ -94,6 +94,8 @@ pub(super) struct GenerateResponse {
94 pub request_id: String,
95 pub choices: Vec<GenerateResponseChoice>,
96 pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
97 + /// Base64 `.npy` float32 array, as Python `numpy2base64` encodes it.
98 + pub prompt_token_id_logprobs: Option<String>,
99 pub prompt_token_ids: Option<Vec<u32>>,
100 pub mm_placeholders: Option<MultiModalPlaceholders>,
101 pub kv_transfer_params: Option<Value>,
rust/src/server/src/routes/inference/generate/validate.rs
+18
@@ -57,6 +57,13 @@ pub(crate) fn validate_request_compat(
57 }
58 }
59
60 + if request.stream && request.sampling_params.inner.prompt_logprob_token_ids.is_some() {
61 + bail_invalid_request!(
62 + param = "sampling_params",
63 + "`prompt_logprob_token_ids` are not available when `stream=true`."
64 + );
65 + }
66 +
67 Ok(())
68 }
69
@@ -169,6 +176,17 @@ mod tests {
176 assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
177 }
178
179 + #[test]
180 + fn validate_request_compat_rejects_streaming_prompt_logprob_token_ids() {
181 + let request: GenerateRequest = serde_json::from_value(json!({
182 + "token_ids": [11, 22],
183 + "stream": true,
184 + "sampling_params": {"prompt_logprob_token_ids": [1, 2]}
185 + }))
186 + .expect("parse request");
187 + assert!(validate_request_compat(&request, &served(&["test-model"])).is_err());
188 + }
189 +
190 #[test]
191 fn validate_request_compat_accepts_non_stream_prompt_logprobs() {
192 let request: GenerateRequest = serde_json::from_value(json!({
rust/src/server/src/routes/openai/chat_completions/convert.rs
+2
@@ -159,6 +159,8 @@ pub(super) fn prepare_chat_request(
159 thinking_token_budget: request.thinking_token_budget,
160 logprobs: request.logprobs.then_some(top_logprobs),
161 prompt_logprobs,
162 + prompt_logprob_token_ids: None,
163 + prompt_logprob_start: None,
164 min_p: request.min_p,
165 frequency_penalty: request.frequency_penalty,
166 presence_penalty: request.presence_penalty,
rust/src/server/src/routes/openai/completions/convert.rs
+2
@@ -124,6 +124,8 @@ pub(super) fn prepare_completion_request(
124 thinking_token_budget: request.thinking_token_budget,
125 logprobs,
126 prompt_logprobs,
127 + prompt_logprob_token_ids: None,
128 + prompt_logprob_start: None,
129 min_p: request.min_p,
130 frequency_penalty: request.frequency_penalty,
131 presence_penalty: request.presence_penalty,
rust/src/text/src/lower.rs
+56
@@ -103,6 +103,8 @@ pub fn lower_sampling_params(
103 thinking_token_budget,
104 logprobs,
105 prompt_logprobs,
106 + prompt_logprob_token_ids,
107 + prompt_logprob_start,
108 min_p,
109 frequency_penalty,
110 presence_penalty,
@@ -123,6 +125,8 @@ pub fn lower_sampling_params(
125 logprobs,
126 prompt_logprobs,
127 logprob_token_ids.as_deref(),
128 + prompt_logprob_token_ids.as_deref(),
129 + prompt_logprob_start,
130 sampling_limits,
131 )?;
132 validate_repetition_detection(repetition_detection.as_ref())?;
@@ -176,6 +180,8 @@ pub fn lower_sampling_params(
180 thinking_token_budget,
181 logprobs,
182 prompt_logprobs,
183 + prompt_logprob_token_ids,
184 + prompt_logprob_start,
185 min_p,
186 frequency_penalty,
187 presence_penalty,
@@ -646,6 +652,8 @@ mod tests {
652 thinking_token_budget: None,
653 logprobs: None,
654 prompt_logprobs: None,
655 + prompt_logprob_token_ids: None,
656 + prompt_logprob_start: None,
657 min_p: 0.0,
658 frequency_penalty: 0.0,
659 presence_penalty: 0.0,
@@ -701,6 +709,8 @@ mod tests {
709 thinking_token_budget: None,
710 logprobs: None,
711 prompt_logprobs: None,
712 + prompt_logprob_token_ids: None,
713 + prompt_logprob_start: None,
714 min_p: 0.0,
715 frequency_penalty: 0.0,
716 presence_penalty: 0.0,
@@ -870,6 +880,8 @@ mod tests {
880 thinking_token_budget: None,
881 logprobs: None,
882 prompt_logprobs: None,
883 + prompt_logprob_token_ids: None,
884 + prompt_logprob_start: None,
885 min_p: 0.0,
886 frequency_penalty: 0.0,
887 presence_penalty: 0.0,
@@ -935,6 +947,8 @@ mod tests {
947 thinking_token_budget: None,
948 logprobs: None,
949 prompt_logprobs: None,
950 + prompt_logprob_token_ids: None,
951 + prompt_logprob_start: None,
952 min_p: 0.0,
953 frequency_penalty: 0.0,
954 presence_penalty: 0.0,
@@ -1008,6 +1022,8 @@ mod tests {
1022 thinking_token_budget: None,
1023 logprobs: None,
1024 prompt_logprobs: None,
1025 + prompt_logprob_token_ids: None,
1026 + prompt_logprob_start: None,
1027 min_p: 0.1,
1028 frequency_penalty: 0.0,
1029 presence_penalty: 0.0,
@@ -1122,6 +1138,44 @@ mod tests {
1138 ));
1139 }
1140
1141 + #[test]
1142 + fn lower_sampling_params_validates_prompt_logprob_token_ids() {
1143 + let lower = |ids: Option<Vec<u32>>, start: Option<u32>| {
1144 + lower_sampling_params_with_limits(
1145 + SamplingParams {
1146 + prompt_logprob_token_ids: ids,
1147 + prompt_logprob_start: start,
1148 + ..Default::default()
1149 + },
1150 + sample_sampling_limits(),
1151 + )
1152 + };
1153 + let rejects = |ids, start| match lower(ids, start) {
1154 + Err(error) => error.to_string(),
1155 + Ok(_) => panic!("expected rejection"),
1156 + };
1157 +
1158 + let params = lower(Some(vec![1, 2]), Some(1)).unwrap();
1159 + assert_eq!(params.prompt_logprob_token_ids, Some(vec![1, 2]));
1160 + assert_eq!(params.prompt_logprob_start, Some(1));
1161 + expect_test::expect![[r#"
1162 + [
1163 + "prompt_logprob_token_ids must not be empty.",
1164 + "requested prompt_logprob_token_ids of 21, which is greater than max allowed: 20",
1165 + "token_id(s) [1000] in prompt_logprob_token_ids are out of vocabulary. Vocabulary size: 1000",
1166 + "prompt_logprob_token_ids must not contain duplicates.",
1167 + "prompt_logprob_start requires prompt_logprob_token_ids.",
1168 + ]
1169 + "#]]
1170 + .assert_debug_eq(&[
1171 + rejects(Some(vec![]), None),
1172 + rejects(Some((0..21).collect()), None),
1173 + rejects(Some(vec![1000]), None),
1174 + rejects(Some(vec![1, 1]), None),
1175 + rejects(None, Some(0)),
1176 + ]);
1177 + }
1178 +
1179 #[test]
1180 fn lower_sampling_params_rejects_out_of_vocab_stop_token_ids() {
1181 let error = lower_sampling_params_with_limits(
@@ -1259,6 +1313,8 @@ mod tests {
1313 thinking_token_budget: None,
1314 logprobs: None,
1315 prompt_logprobs: None,
1316 + prompt_logprob_token_ids: None,
1317 + prompt_logprob_start: None,
1318 min_p: 0.1,
1319 frequency_penalty: 0.0,
1320 presence_penalty: 0.0,
rust/src/text/src/lower/logprobs.rs
+28
@@ -6,6 +6,7 @@
6 //! `-1` is expanded only for bounds checks. The original request values are
7 //! passed through to engine-core.
8
9 +use itertools::Itertools as _;
10 use thiserror::Error;
11
12 use crate::backend::SamplingLimits;
@@ -35,6 +36,12 @@ pub enum LogprobsError {
36 len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
37 )]
38 TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
39 + #[error("prompt_logprob_token_ids must not be empty.")]
40 + EmptyPromptLogprobTokenIds,
41 + #[error("prompt_logprob_token_ids must not contain duplicates.")]
42 + DuplicatePromptLogprobTokenIds,
43 + #[error("prompt_logprob_start requires prompt_logprob_token_ids.")]
44 + PromptLogprobStartWithoutTokenIds,
45 }
46
47 /// Validate logprobs count sampling parameters.
@@ -42,6 +49,8 @@ pub(super) fn validate_logprobs(
49 logprobs: Option<i32>,
50 prompt_logprobs: Option<i32>,
51 logprob_token_ids: Option<&[u32]>,
52 + prompt_logprob_token_ids: Option<&[u32]>,
53 + prompt_logprob_start: Option<u32>,
54 sampling_limits: SamplingLimits,
55 ) -> Result<(), LogprobsError> {
56 let vocab_size = sampling_limits.model_vocab_size;
@@ -50,9 +59,28 @@ pub(super) fn validate_logprobs(
59
60 validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
61 validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
62 + validate_logprobs_count(
63 + prompt_logprob_token_ids.map(|ids| ids.len() as i32),
64 + max_logprobs,
65 + vocab_size,
66 + "prompt_logprob_token_ids",
67 + )?;
68 + validate_prompt_logprob_token_ids(prompt_logprob_token_ids, prompt_logprob_start)?;
69 validate_logprob_token_ids(logprobs, logprob_token_ids)
70 }
71
72 +fn validate_prompt_logprob_token_ids(
73 + token_ids: Option<&[u32]>,
74 + start: Option<u32>,
75 +) -> Result<(), LogprobsError> {
76 + match token_ids {
77 + Some([]) => Err(LogprobsError::EmptyPromptLogprobTokenIds),
78 + Some(ids) if !ids.iter().all_unique() => Err(LogprobsError::DuplicatePromptLogprobTokenIds),
79 + None if start.is_some() => Err(LogprobsError::PromptLogprobStartWithoutTokenIds),
80 + _ => Ok(()),
81 + }
82 +}
83 +
84 fn validate_logprobs_count(
85 requested: Option<i32>,
86 max_logprobs: usize,
rust/src/text/src/lower/token_ids.rs
+8
@@ -95,6 +95,14 @@ pub(crate) fn validate_vocab_range(
95 )?;
96 }
97
98 + if let Some(token_ids) = params.prompt_logprob_token_ids.as_deref() {
99 + validate_param(
100 + "prompt_logprob_token_ids",
101 + token_ids.iter().copied(),
102 + limits.model_vocab_size,
103 + )?;
104 + }
105 +
106 if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() {
107 validate_param(
108 "bad_words",
rust/src/text/src/request.rs
+7
@@ -97,6 +97,11 @@ pub struct SamplingParams {
97 ///
98 /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
99 pub prompt_logprobs: Option<i32>,
100 + /// Candidate token IDs scored at every scored causal prompt row, where row
101 + /// `i` scores them as predictions of prompt token `i + 1`.
102 + pub prompt_logprob_token_ids: Option<Vec<u32>>,
103 + /// First causal prompt row to score; `None` scores from the first row.
104 + pub prompt_logprob_start: Option<u32>,
105 /// Minimum probability threshold for token sampling. `None` means no
106 /// explicit user override.
107 pub min_p: Option<f32>,
@@ -156,6 +161,8 @@ impl Default for SamplingParams {
161 thinking_token_budget: None,
162 logprobs: None,
163 prompt_logprobs: None,
164 + prompt_logprob_token_ids: None,
165 + prompt_logprob_start: None,
166 min_p: None,
167 frequency_penalty: None,
168 presence_penalty: None,
tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+49
@@ -1,9 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import io
5 import json
6
7 import httpx
8 +import numpy as np
9 +import pybase64 as base64
10 import pytest
11 import pytest_asyncio
12 from transformers import AutoTokenizer
@@ -311,6 +314,52 @@ async def test_generate_logprobs(client, logprobs_value):
314 assert len(entry["top_logprobs"]) == max(logprobs_value, 1)
315
316
317 +@pytest.mark.asyncio
318 +async def test_generate_prompt_token_id_logprobs(client):
319 + token_ids = [11, 22, 33, 44, 55]
320 + candidates = [22, 33, 44, 55]
321 + sampling_params = {
322 + "max_tokens": 1,
323 + "prompt_logprobs": 0,
324 + "prompt_logprob_token_ids": candidates,
325 + "prompt_logprob_start": 1,
326 + }
327 + payload = {
328 + "model": MODEL_NAME,
329 + "token_ids": token_ids,
330 + "sampling_params": sampling_params,
331 + }
332 + resp = await client.post(GEN_ENDPOINT, json=payload)
333 + resp.raise_for_status()
334 + data = resp.json()
335 +
336 + scores = np.load(io.BytesIO(base64.b64decode(data["prompt_token_id_logprobs"])))
337 + assert scores.shape == (len(token_ids) - 2, len(candidates))
338 + for row, target in enumerate(token_ids[2:]):
339 + expected = data["prompt_logprobs"][row + 2][str(target)]["logprob"]
340 + assert scores[row, candidates.index(target)] == pytest.approx(
341 + expected, abs=1e-3
342 + )
343 +
344 + stream_only_ids = {"prompt_logprob_token_ids": candidates}
345 + resp = await client.post(
346 + GEN_ENDPOINT,
347 + json={**payload, "sampling_params": stream_only_ids, "stream": True},
348 + )
349 + assert resp.status_code == 400
350 + assert "prompt_logprob_token_ids" in resp.text
351 +
352 + past_end = {**sampling_params, "prompt_logprob_start": len(token_ids)}
353 + resp = await client.post(
354 + GEN_ENDPOINT, json={**payload, "sampling_params": past_end}
355 + )
356 + resp.raise_for_status()
357 + empty = np.load(
358 + io.BytesIO(base64.b64decode(resp.json()["prompt_token_id_logprobs"]))
359 + )
360 + assert empty.shape == (0, len(candidates))
361 +
362 +
363 @pytest.mark.asyncio
364 async def test_same_response_as_chat_completions(client, tokenizer, messages):
365 token_ids = tokenizer.apply_chat_template(
tests/v1/engine/test_logprobs_processor.py
+23
@@ -11,8 +11,10 @@ row), the trailing positions are populated with sentinel values
11 """
12
13 import numpy as np
14 +import torch
15
16 from vllm.logprobs import create_sample_logprobs
17 +from vllm.v1.engine import EngineCoreOutput
18 from vllm.v1.engine.logprobs import LogprobsProcessor
19 from vllm.v1.outputs import LogprobsLists
20
@@ -64,3 +66,24 @@ def test_accepts_exactly_sized_row():
66
67 pos = processor.logprobs[0]
68 assert set(pos.keys()) == {7, 11, 13}
69 +
70 +
71 +def test_prompt_token_id_logprobs_are_popped_once():
72 + """DELTA outputs carry fixed-ID prompt scores exactly once."""
73 + processor = _make_processor(num_logprobs=1)
74 + assert processor.pop_prompt_token_id_logprobs() is None
75 +
76 + processor.update_from_output(
77 + EngineCoreOutput(
78 + request_id="req-0",
79 + new_token_ids=[],
80 + prompt_token_id_logprobs=torch.tensor(
81 + [[-0.5, -1.5], [-2.5, -3.5]], dtype=torch.float32
82 + ),
83 + )
84 + )
85 +
86 + scores = processor.pop_prompt_token_id_logprobs()
87 + assert scores is not None
88 + assert scores.tolist() == [[-0.5, -1.5], [-2.5, -3.5]]
89 + assert processor.pop_prompt_token_id_logprobs() is None
tests/v1/sample/test_logprobs.py
+177 -2
@@ -6,7 +6,9 @@ import math
6 from collections.abc import Generator
7 from types import SimpleNamespace
8 from typing import get_args
9 +from unittest.mock import patch
10
11 +import numpy as np
12 import pytest
13 import torch
14
@@ -23,6 +25,7 @@ from vllm.config.model import LogprobsMode
25 from vllm.distributed import cleanup_dist_env_and_memory
26 from vllm.exceptions import VLLMValidationError
27 from vllm.platforms import current_platform
28 +from vllm.v1.engine.input_processor import InputProcessor
29
30 from ...conftest import HfRunner, VllmRunner
31
@@ -81,9 +84,9 @@ def hf_model(hf_runner) -> Generator[HfRunner, None, None]:
84 yield hf_model
85
86
84 -def _model_config(vocab_size: int = 10):
87 +def _model_config(vocab_size: int = 10, max_logprobs: int = 20):
88 return SimpleNamespace(
86 - max_logprobs=20,
89 + max_logprobs=max_logprobs,
90 logits_processors=None,
91 is_diffusion=False,
92 get_vocab_size=lambda: vocab_size,
@@ -432,6 +435,57 @@ def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]):
435 )
436
437
438 +def test_prompt_logprob_token_ids_bounded_by_max_logprobs():
439 + """The candidate count is bounded by max_logprobs."""
440 + model_config = _model_config(vocab_size=100, max_logprobs=4)
441 +
442 + def verify(**kwargs):
443 + SamplingParams(**kwargs).verify(
444 + model_config,
445 + speculative_config=None,
446 + structured_outputs_config=None,
447 + tokenizer=None,
448 + )
449 +
450 + verify(prompt_logprob_token_ids=[1, 2, 3, 4])
451 +
452 + # The error names the value to raise max_logprobs to.
453 + with pytest.raises(VLLMValidationError, match=r"max_logprobs.*at least 5"):
454 + verify(prompt_logprob_token_ids=[1, 2, 3, 4, 5])
455 +
456 + # prompt_logprob_start alone is a caller mistake, not a silent no-op.
457 + with pytest.raises(VLLMValidationError, match="requires prompt_logprob_token_ids"):
458 + verify(prompt_logprob_start=3)
459 +
460 +
461 +def test_prompt_logprob_token_ids_require_v2_model_runner():
462 + """Only the V2 runner scores them; the V1 runner would return None, and
463 + kv-sharing fast prefill would score rows the cross-decoder never ran."""
464 + processor = SimpleNamespace(
465 + model_config=SimpleNamespace(
466 + return_sampling_mask=False, enable_trace_replay=False, is_diffusion=False
467 + ),
468 + vllm_config=SimpleNamespace(
469 + reasoning_config=None,
470 + use_v2_model_runner=False,
471 + cache_config=SimpleNamespace(kv_sharing_fast_prefill=False),
472 + ),
473 + speculative_config=None,
474 + structured_outputs_config=None,
475 + tokenizer=None,
476 + validate_logits_processors_params=lambda params: None,
477 + )
478 + params = SamplingParams(prompt_logprob_token_ids=[1, 2])
479 + with patch.object(SamplingParams, "verify"):
480 + with pytest.raises(VLLMValidationError, match="V2 model runner"):
481 + InputProcessor._validate_params(processor, params, ("generate",))
482 + processor.vllm_config.use_v2_model_runner = True
483 + InputProcessor._validate_params(processor, params, ("generate",))
484 + processor.vllm_config.cache_config.kv_sharing_fast_prefill = True
485 + with pytest.raises(VLLMValidationError, match="fast-prefill"):
486 + InputProcessor._validate_params(processor, params, ("generate",))
487 +
488 +
489 def test_none_logprobs(vllm_model, example_prompts):
490 """Engine should return `logprobs` and `prompt_logprobs` as `None`.
491
@@ -1298,6 +1352,127 @@ def test_prompt_logprobs_with_chunking_and_preemption():
1352 print(f"Test passed with {preemptions} preemptions")
1353
1354
1355 +def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch):
1356 + """Fixed-ID scores stay row-aligned across chunked prefill and preemption."""
1357 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
1358 +
1359 + prompts = [
1360 + "The following numbers of the sequence "
1361 + + ", ".join(str(i) for i in range(10))
1362 + + " are:",
1363 + "In one word, the capital of France is ",
1364 + ] + [f"Tell me about the number {i}: " for i in range(32)]
1365 +
1366 + start = 2
1367 + candidate_ids = [10, 100, 1000, 10000]
1368 + sampling_params = SamplingParams(
1369 + temperature=0.0,
1370 + max_tokens=40,
1371 + min_tokens=20,
1372 + prompt_logprob_token_ids=candidate_ids,
1373 + prompt_logprob_start=start,
1374 + )
1375 +
1376 + with VllmRunner(
1377 + "Qwen/Qwen3-0.6B",
1378 + max_model_len=512,
1379 + enable_chunked_prefill=True,
1380 + max_num_batched_tokens=48, # Force prefill chunking
1381 + num_gpu_blocks_override=33, # Force preemptions (32 usable + null block)
1382 + disable_log_stats=False,
1383 + gpu_memory_utilization=0.25,
1384 + ) as vllm_model:
1385 + metrics_before = vllm_model.llm.get_metrics()
1386 + outputs = vllm_model.llm.generate(prompts, sampling_params)
1387 +
1388 + for i, output in enumerate(outputs):
1389 + scores = output.prompt_token_id_logprobs
1390 + assert scores is not None, f"Output {i} missing fixed-ID scores"
1391 + expected_shape = (
1392 + len(output.prompt_token_ids) - 1 - start,
1393 + len(candidate_ids),
1394 + )
1395 + assert scores.shape == expected_shape, (
1396 + f"Output {i} scored {scores.shape}, expected {expected_shape}"
1397 + )
1398 + assert math.isfinite(float(scores.min()))
1399 + assert float(scores.max()) <= 1e-3, "logprobs must be <= 0"
1400 +
1401 + metrics_after = vllm_model.llm.get_metrics()
1402 + preemptions_before = next(
1403 + (m.value for m in metrics_before if m.name == "vllm:num_preemptions"), 0
1404 + )
1405 + preemptions_after = next(
1406 + (m.value for m in metrics_after if m.name == "vllm:num_preemptions"), 0
1407 + )
1408 + assert preemptions_after - preemptions_before > 0, (
1409 + "Test did not trigger any preemptions"
1410 + )
1411 +
1412 + # Row alignment is numerical: chunked and preempted scores must match an
1413 + # unchunked, unpreempted run of the same requests. Batch composition moves
1414 + # bf16 tail logprobs by a few percent; a misaligned row differs by nats.
1415 + with VllmRunner(
1416 + "Qwen/Qwen3-0.6B",
1417 + max_model_len=512,
1418 + max_logprobs=64,
1419 + gpu_memory_utilization=0.25,
1420 + ) as reference:
1421 + reference_outputs = reference.llm.generate(prompts, sampling_params)
1422 + own_ids = [list(dict.fromkeys(o.prompt_token_ids)) for o in reference_outputs]
1423 + self_scored = reference.llm.generate(
1424 + prompts,
1425 + [
1426 + SamplingParams(
1427 + max_tokens=1, prompt_logprobs=0, prompt_logprob_token_ids=ids
1428 + )
1429 + for ids in own_ids
1430 + ],
1431 + )
1432 + for output, ref in zip(outputs, reference_outputs):
1433 + np.testing.assert_allclose(
1434 + output.prompt_token_id_logprobs, ref.prompt_token_id_logprobs, rtol=0.1
1435 + )
1436 + for output, ids in zip(self_scored, own_ids):
1437 + scores = output.prompt_token_id_logprobs
1438 + for row, target in enumerate(output.prompt_token_ids[1:]):
1439 + expected = output.prompt_logprobs[row + 1][target].logprob
1440 + assert scores[row, ids.index(target)] == pytest.approx(expected, abs=1e-3)
1441 +
1442 +
1443 +def test_prompt_logprob_token_ids_drop_partially_scored_prefills(monkeypatch):
1444 + """A prefill that starts past the first scored row returns no scores.
1445 +
1446 + Such a prefill (here a prefix-cache hit the caller opted back into) never
1447 + computes the leading rows, so emitting the buffer would return values the
1448 + model never produced for them.
1449 + """
1450 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
1451 +
1452 + prompt = "The capital of France is Paris. " * 20
1453 + sampling_params = SamplingParams(
1454 + temperature=0.0,
1455 + max_tokens=1,
1456 + prompt_logprob_token_ids=[10, 100, 1000],
1457 + skip_reading_prefix_cache=False,
1458 + )
1459 +
1460 + with VllmRunner(
1461 + "Qwen/Qwen3-0.6B",
1462 + max_model_len=512,
1463 + enable_prefix_caching=True,
1464 + gpu_memory_utilization=0.25,
1465 + ) as vllm_model:
1466 + first = vllm_model.llm.generate([prompt], sampling_params)[0]
1467 + assert first.prompt_token_id_logprobs is not None
1468 + assert math.isfinite(float(first.prompt_token_id_logprobs.min()))
1469 +
1470 + # The re-send is served from the prefix cache, so its prefill starts
1471 + # past row 0 and the leading rows are never computed.
1472 + second = vllm_model.llm.generate([prompt], sampling_params)[0]
1473 + assert second.prompt_token_id_logprobs is None
1474 +
1475 +
1476 @large_gpu_mark(min_gb=24)
1477 def test_token_logprobs_large_batch_int64_row_offset():
1478 """Regression: logprob kernel row offset (row * vocab_size) must use int64.
tests/v1/sample/test_trace_replay_params.py
+8
@@ -41,6 +41,14 @@ def test_sampling_params_trace_field_rejects_empty_list():
41 )
42
43
44 +def test_sampling_params_trace_field_rejects_prompt_scoring():
45 + params = SamplingParams(trace_decode_token_ids=[1], prompt_logprob_token_ids=[2, 3])
46 + with pytest.raises(ValueError, match="prompt_logprob_token_ids"):
47 + params._validate_trace_replay(
48 + _make_model_config(vocab_size=100), speculative_config=None
49 + )
50 +
51 +
52 def test_sampling_params_trace_field_requires_single_output():
53 params = SamplingParams(n=2, trace_decode_token_ids=[1])
54 with pytest.raises(ValueError, match="requires n=1"):
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+1
@@ -381,6 +381,7 @@ class GenerateResponse(BaseModel):
381 choices: list[GenerateResponseChoice]
382 usage: UsageInfo | None = Field(default=None)
383 prompt_logprobs: list[dict[int, Logprob] | None] | None = None
384 + prompt_token_id_logprobs: str | None = None
385 prompt_token_ids: list[int] | None = None
386 mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
387 metrics: PerRequestMetrics | None = None
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+10
@@ -138,6 +138,11 @@ class ServingTokens(GenerateBaseServing):
138 f"sampling_params.n must be at most the server's max_num_seqs "
139 f"({max_num_seqs}), got {sampling_params.n}."
140 )
141 + # The stream schema has no field for the scores.
142 + if request.stream and sampling_params.prompt_logprob_token_ids is not None:
143 + return self.create_error_response(
144 + "prompt_logprob_token_ids are not available when stream=true."
145 + )
146 if self.force_no_detokenize and sampling_params.stop:
147 # SamplingParams rejects stop with detokenize=False at request
148 # validation, but this server forces detokenize=False afterwards,
@@ -389,6 +394,11 @@ class ServingTokens(GenerateBaseServing):
394 choices=choices,
395 usage=usage,
396 prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
397 + prompt_token_id_logprobs=(
398 + numpy2base64(final_res.prompt_token_id_logprobs)
399 + if final_res.prompt_token_id_logprobs is not None
400 + else None
401 + ),
402 prompt_token_ids=(
403 final_res.prompt_token_ids if request.return_token_ids else None
404 ),
vllm/outputs.py
+7
@@ -117,6 +117,10 @@ class RequestOutput:
117 For encoder/decoder models, this is the
118 decoder input prompt token ids.
119 prompt_logprobs: The log probabilities to return per prompt token.
120 + prompt_token_id_logprobs: Logprobs of prompt_logprob_token_ids, shaped
121 + [max(prompt_len - 1 - prompt_logprob_start, 0),
122 + len(prompt_logprob_token_ids)]; row i scores them as predictions of
123 + prompt token prompt_logprob_start + i + 1.
124 outputs: The output sequences of the request.
125 finished: Whether the whole request is finished.
126 metrics: Metrics associated with the request.
@@ -148,6 +152,7 @@ class RequestOutput:
152 num_cached_tokens: int | None = None,
153 num_cache_creation_tokens: int | None = None,
154 *,
155 + prompt_token_id_logprobs: np.ndarray | None = None,
156 kv_transfer_params: dict[str, Any] | None = None,
157 ec_transfer_params: dict[str, Any] | None = None,
158 # Forward compatibility, code that uses args added in new release can
@@ -162,6 +167,7 @@ class RequestOutput:
167 self.prompt = prompt
168 self.prompt_token_ids = prompt_token_ids
169 self.prompt_logprobs = prompt_logprobs
170 + self.prompt_token_id_logprobs = prompt_token_id_logprobs
171 self.outputs = outputs
172 self.finished = finished
173 self.metrics = metrics
@@ -215,6 +221,7 @@ class RequestOutput:
221 f"encoder_prompt={self.encoder_prompt!r}, "
222 f"encoder_prompt_token_ids={self.encoder_prompt_token_ids}, "
223 f"prompt_logprobs={self.prompt_logprobs}, "
224 + f"prompt_token_id_logprobs={self.prompt_token_id_logprobs}, "
225 f"outputs={self.outputs}, "
226 f"finished={self.finished}, "
227 f"metrics={self.metrics}, "
vllm/sampling_params.py
+68 -1
@@ -289,6 +289,12 @@ class SamplingParams(
289 prompt_logprobs: int | None = None
290 """Number of log probabilities to return per prompt token.
291 When set to -1, return all `vocab_size` log probabilities."""
292 + prompt_logprob_token_ids: list[int] | None = None
293 + """Token IDs to score at each scored causal prompt row, where row i scores
294 + them as predictions of prompt token i + 1. The last prompt row is excluded,
295 + so the result has `max(prompt_len - 1 - prompt_logprob_start, 0)` rows."""
296 + prompt_logprob_start: int | None = None
297 + """First causal prompt row to score; defaults to the first row."""
298 logprob_token_ids: list[int] | None = None
299 """Specific token IDs to return logprobs for. More efficient than
300 logprobs=-1 when you only need logprobs for a small set of tokens.
@@ -398,6 +404,8 @@ class SamplingParams(
404 min_tokens: int = 0,
405 logprobs: int | None = None,
406 prompt_logprobs: int | None = None,
407 + prompt_logprob_token_ids: list[int] | None = None,
408 + prompt_logprob_start: int | None = None,
409 detokenize: bool = True,
410 skip_special_tokens: bool = True,
411 spaces_between_special_tokens: bool = True,
@@ -464,6 +472,8 @@ class SamplingParams(
472 min_tokens=min_tokens,
473 logprobs=logprobs,
474 prompt_logprobs=prompt_logprobs,
475 + prompt_logprob_token_ids=prompt_logprob_token_ids,
476 + prompt_logprob_start=prompt_logprob_start,
477 logprob_token_ids=logprob_token_ids,
478 detokenize=detokenize,
479 skip_special_tokens=skip_special_tokens,
@@ -550,7 +560,10 @@ class SamplingParams(
560 # If prefix caching is enabled,
561 # the output of prompt logprobs may less than n_prompt_tokens,
562 # we need to skip reading cache at this request.
553 - self.skip_reading_prefix_cache = self.prompt_logprobs is not None
563 + self.skip_reading_prefix_cache = (
564 + self.prompt_logprobs is not None
565 + or self.prompt_logprob_token_ids is not None
566 + )
567
568 def _verify_args(self) -> None:
569 _verify_num_sequences(self.n, "n")
@@ -902,6 +915,56 @@ class SamplingParams(
915 value=num_prompt_logprobs,
916 )
917
918 + # Validate prompt_logprob_token_ids.
919 + if self.prompt_logprob_token_ids is not None:
920 + n = len(self.prompt_logprob_token_ids)
921 + if n == 0:
922 + raise VLLMValidationError(
923 + "prompt_logprob_token_ids must not be empty.",
924 + parameter="prompt_logprob_token_ids",
925 + value=n,
926 + )
927 + if n > max_logprobs:
928 + raise VLLMValidationError(
929 + f"Requested prompt_logprob_token_ids of length {n}, "
930 + f"which is greater than max allowed: {max_logprobs}. "
931 + f"Set max_logprobs (--max-logprobs) to at least {n}.",
932 + parameter="prompt_logprob_token_ids",
933 + value=n,
934 + )
935 + vocab_size = model_config.get_vocab_size()
936 + invalid_token_ids = [
937 + token_id
938 + for token_id in self.prompt_logprob_token_ids
939 + if token_id < 0 or token_id >= vocab_size
940 + ]
941 + if invalid_token_ids:
942 + raise VLLMValidationError(
943 + f"token_id(s) {invalid_token_ids} in "
944 + f"prompt_logprob_token_ids contain out-of-vocab token ids. "
945 + f"Vocabulary size: {vocab_size}",
946 + parameter="prompt_logprob_token_ids",
947 + value=invalid_token_ids,
948 + )
949 + if len(set(self.prompt_logprob_token_ids)) != n:
950 + raise VLLMValidationError(
951 + "prompt_logprob_token_ids must not contain duplicates.",
952 + parameter="prompt_logprob_token_ids",
953 + value=self.prompt_logprob_token_ids,
954 + )
955 + if self.prompt_logprob_start is not None and self.prompt_logprob_start < 0:
956 + raise VLLMValidationError(
957 + "prompt_logprob_start must be non-negative.",
958 + parameter="prompt_logprob_start",
959 + value=self.prompt_logprob_start,
960 + )
961 + elif self.prompt_logprob_start is not None:
962 + raise VLLMValidationError(
963 + "prompt_logprob_start requires prompt_logprob_token_ids.",
964 + parameter="prompt_logprob_start",
965 + value=self.prompt_logprob_start,
966 + )
967 +
968 def _validate_stop_token_ids(self, model_config: ModelConfig) -> None:
969 """Validate stop_token_ids are within vocabulary range."""
970 if not self.stop_token_ids:
@@ -969,6 +1032,10 @@ class SamplingParams(
1032 raise ValueError(
1033 "trace_decode_token_ids is not supported with prompt_logprobs."
1034 )
1035 + if self.prompt_logprob_token_ids is not None:
1036 + raise ValueError(
1037 + "trace_decode_token_ids is not supported with prompt_logprob_token_ids."
1038 + )
1039 if speculative_config is not None:
1040 raise ValueError(
1041 "trace_decode_token_ids is not supported with speculative decoding."
vllm/v1/core/sched/scheduler.py
+6
@@ -1972,6 +1972,9 @@ class Scheduler(SchedulerInterface):
1972 sampled_token_ids = model_runner_output.sampled_token_ids
1973 logprobs = model_runner_output.logprobs
1974 prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
1975 + prompt_token_id_logprobs_dict = (
1976 + model_runner_output.prompt_token_id_logprobs_dict
1977 + )
1978 num_scheduled_tokens = scheduler_output.num_scheduled_tokens
1979 pooler_outputs = model_runner_output.pooler_output
1980 num_nans_in_logits = model_runner_output.num_nans_in_logits
@@ -2194,6 +2197,7 @@ class Scheduler(SchedulerInterface):
2197
2198 # Get prompt logprobs for this request.
2199 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
2200 + prompt_token_id_logprobs = prompt_token_id_logprobs_dict.get(req_id)
2201 if should_emit_output:
2202 # Add EngineCoreOutput for this Request.
2203 outputs[request.client_index].append(
@@ -2204,6 +2208,7 @@ class Scheduler(SchedulerInterface):
2208 new_logprobs=new_logprobs,
2209 new_sampling_mask=new_sampling_mask,
2210 new_prompt_logprobs_tensors=prompt_logprobs_tensors,
2211 + prompt_token_id_logprobs=prompt_token_id_logprobs,
2212 pooling_output=pooler_output,
2213 stop_reason=request.stop_reason,
2214 events=request.take_events(),
@@ -2223,6 +2228,7 @@ class Scheduler(SchedulerInterface):
2228 else:
2229 # Invariant: EngineCore returns no partial prefill outputs.
2230 assert not prompt_logprobs_tensors
2231 + assert prompt_token_id_logprobs is None
2232
2233 # Remove the stopped requests from the running and waiting queues.
2234 if stopped_running_reqs:
vllm/v1/engine/__init__.py
+3
@@ -236,6 +236,9 @@ class EngineCoreOutput(
236 # Appended last so `array_like` positional serialization stays compatible.
237 spec_decode_metrics: RequestSpecDecodeMetrics | None = None
238
239 + # [num_scored_rows, num_token_ids]; appended last for array_like compatibility.
240 + prompt_token_id_logprobs: torch.Tensor | None = None
241 +
242 @property
243 def finished(self) -> bool:
244 return self.finish_reason is not None
vllm/v1/engine/input_processor.py
+13
@@ -131,6 +131,19 @@ class InputProcessor:
131 self.structured_outputs_config,
132 self.tokenizer,
133 )
134 + if params.prompt_logprob_token_ids is not None:
135 + if not self.vllm_config.use_v2_model_runner:
136 + raise VLLMValidationError(
137 + "prompt_logprob_token_ids requires the V2 model runner "
138 + "(VLLM_USE_V2_MODEL_RUNNER=1).",
139 + parameter="prompt_logprob_token_ids",
140 + )
141 + if self.vllm_config.cache_config.kv_sharing_fast_prefill:
142 + raise VLLMValidationError(
143 + "prompt_logprob_token_ids is incorrect with "
144 + "--kv-sharing-fast-prefill; disable it for scoring.",
145 + parameter="prompt_logprob_token_ids",
146 + )
147
148 self.validate_logits_processors_params(params)
149
vllm/v1/engine/logprobs.py
+12
@@ -5,6 +5,8 @@ import itertools
5 from collections.abc import Iterable
6 from dataclasses import dataclass
7
8 +import numpy as np
9 +
10 from vllm.logger import init_logger
11 from vllm.logprobs import (
12 FlatLogprobs,
@@ -38,6 +40,8 @@ class LogprobsProcessor:
40 cumulative_logprob: float | None
41 num_logprobs: int | None
42 num_prompt_logprobs: int | None
43 + # [num_scored_rows, num_token_ids], set once on the final prefill chunk.
44 + prompt_token_id_logprobs: np.ndarray | None = None
45
46 @classmethod
47 def from_new_request(
@@ -204,6 +208,12 @@ class LogprobsProcessor:
208 self.prompt_logprobs = []
209 return plp
210
211 + def pop_prompt_token_id_logprobs(self) -> np.ndarray | None:
212 + """Pop and return the fixed-ID prompt scores."""
213 + scores = self.prompt_token_id_logprobs
214 + self.prompt_token_id_logprobs = None
215 + return scores
216 +
217 @staticmethod
218 def _get_sampled_context_ids(
219 logprobs_source: SampleLogprobs | PromptLogprobs | None,
@@ -352,3 +362,5 @@ class LogprobsProcessor:
362 self._update_sample_logprobs(output.new_logprobs)
363 if output.new_prompt_logprobs_tensors is not None:
364 self._update_prompt_logprobs(output.new_prompt_logprobs_tensors)
365 + if output.prompt_token_id_logprobs is not None:
366 + self.prompt_token_id_logprobs = output.prompt_token_id_logprobs.numpy()
vllm/v1/engine/output_processor.py
+5
@@ -390,8 +390,12 @@ class RequestState:
390 if self.output_kind == RequestOutputKind.DELTA:
391 # Side effect: logprobs processor forgets prompt logprobs
392 prompt_logprobs = self.logprobs_processor.pop_prompt_logprobs()
393 + prompt_token_id_logprobs = (
394 + self.logprobs_processor.pop_prompt_token_id_logprobs()
395 + )
396 else:
397 prompt_logprobs = self.logprobs_processor.prompt_logprobs
398 + prompt_token_id_logprobs = self.logprobs_processor.prompt_token_id_logprobs
399
400 return RequestOutput(
401 request_id=external_req_id, # request_id is what was provided externally
@@ -399,6 +403,7 @@ class RequestState:
403 prompt=self.prompt,
404 prompt_token_ids=prompt_token_ids,
405 prompt_logprobs=prompt_logprobs,
406 + prompt_token_id_logprobs=prompt_token_id_logprobs,
407 outputs=cast(list[CompletionOutput], outputs),
408 finished=finished,
409 kv_transfer_params=kv_transfer_params,
vllm/v1/outputs.py
+3
@@ -283,6 +283,9 @@ class ModelRunnerOutput:
283 default_factory=dict
284 )
285
286 + # req_id -> [num_scored_rows, num_token_ids] prompt_logprob_token_ids scores.
287 + prompt_token_id_logprobs_dict: dict[str, torch.Tensor] = field(default_factory=dict)
288 +
289 # [num_reqs, hidden_size]
290 pooler_output: list[torch.Tensor | None] | None = None
291
vllm/v1/worker/gpu/async_utils.py
+7
@@ -156,6 +156,10 @@ class AsyncOutput(AsyncModelRunnerOutput):
156 k: v.to_cpu_nonblocking() if v is not None else None
157 for k, v in self.model_runner_output.prompt_logprobs_dict.items()
158 }
159 + token_id_logprobs = self.model_runner_output.prompt_token_id_logprobs_dict
160 + self.prompt_token_id_logprobs_dict = {
161 + k: v.to("cpu", non_blocking=True) for k, v in token_id_logprobs.items()
162 + }
163 if self.pending_aux_output is not None:
164 self.pending_aux_output.enqueue_cpu_copy(
165 num_sampled=self.num_sampled_tokens_np,
@@ -194,6 +198,9 @@ class AsyncOutput(AsyncModelRunnerOutput):
198 if self.logprobs_tensors is not None:
199 self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
200 self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
201 + self.model_runner_output.prompt_token_id_logprobs_dict = (
202 + self.prompt_token_id_logprobs_dict
203 + )
204 if self.pending_aux_output is not None:
205 self.model_runner_output.aux_output_connector_output = (
206 self.pending_aux_output.process_output()
vllm/v1/worker/gpu/model_runner.py
+12 -1
@@ -494,7 +494,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
494 ),
495 )
496 self.prompt_logprobs_worker = PromptLogprobsWorker(
497 - self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
497 + self.max_num_reqs,
498 + self.device,
499 + logprobs_mode=self.model_config.logprobs_mode,
500 )
501 self.structured_outputs_worker = StructuredOutputsWorker(
502 max_num_logits=self.max_num_reqs * self.decode_query_len,
@@ -2063,6 +2065,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2065 self.req_states.num_computed_tokens.gpu,
2066 self.req_states.prompt_len.np,
2067 )
2068 + prompt_token_id_logprobs_dict = (
2069 + self.prompt_logprobs_worker.compute_prompt_token_id_logprobs(
2070 + self.model.compute_logits,
2071 + hidden_states,
2072 + input_batch,
2073 + self.req_states.prompt_len.np,
2074 + )
2075 + )
2076
2077 # Prepare the model runner output.
2078 model_runner_output = ModelRunnerOutput(
@@ -2072,6 +2082,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2082 req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
2083 sampled_token_ids=None, # type: ignore
2084 prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]
2085 + prompt_token_id_logprobs_dict=prompt_token_id_logprobs_dict,
2086 cudagraph_stats=cudagraph_stats,
2087 )
2088 pending_aux_output = None
vllm/v1/worker/gpu/sample/logprob.py
+3 -1
@@ -19,6 +19,7 @@ def _topk_log_softmax_kernel(
19 logits_ptr,
20 logits_stride,
21 topk_ids_ptr,
22 + topk_ids_stride,
23 topk,
24 vocab_size,
25 BLOCK_SIZE: tl.constexpr,
@@ -49,7 +50,7 @@ def _topk_log_softmax_kernel(
50 k_offset = j + tl.arange(0, TOPK_BLOCK_SIZE)
51 k_mask = k_offset < topk
52 topk_ids = tl.load(
52 - topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0
53 + topk_ids_ptr + req_idx * topk_ids_stride + k_offset, mask=k_mask, other=0
54 )
55 logits = tl.load(row_ptr + topk_ids, mask=k_mask)
56 logits = logits.to(tl.float32)
@@ -98,6 +99,7 @@ def compute_token_logprobs(
99 logits,
100 logits.stride(0),
101 token_ids,
102 + token_ids.stride(0),
103 num_logprobs,
104 vocab_size,
105 BLOCK_SIZE=1024, # type: ignore
vllm/v1/worker/gpu/sample/prompt_logprob.py
+84 -3
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Callable
4 +from dataclasses import dataclass
5
6 import numpy as np
7 import torch
@@ -8,20 +9,41 @@ import torch
9 from vllm.config.model import LogprobsMode
10 from vllm.sampling_params import SamplingParams
11 from vllm.triton_utils import tl, triton
12 +from vllm.utils.torch_utils import async_tensor_h2d
13 from vllm.v1.outputs import LogprobsTensors
14 from vllm.v1.worker.gpu.input_batch import InputBatch
13 -from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
15 +from vllm.v1.worker.gpu.sample.logprob import (
16 + compute_token_logprobs,
17 + compute_topk_scores,
18 +)
19 +
20 +CHUNK_SIZE = 1024
21 +
22 +
23 +@dataclass
24 +class _TokenIdScores:
25 + token_ids: torch.Tensor # [num_ids], on device
26 + start: int # first prompt row to score
27 + scores: torch.Tensor | None = None # [num_rows, num_ids], filled per chunk
28
29
30 class PromptLogprobsWorker:
17 - def __init__(self, max_num_reqs: int, logprobs_mode: LogprobsMode = "raw_logprobs"):
31 + def __init__(
32 + self,
33 + max_num_reqs: int,
34 + device: torch.device,
35 + logprobs_mode: LogprobsMode = "raw_logprobs",
36 + ):
37 self.max_num_reqs = max_num_reqs
38 + self.device = device
39 self.logprobs_mode = logprobs_mode
40
41 self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
42 self.num_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=np.int32)
43 # req_idx -> list of in-progress LogprobsTensors
44 self.in_progress_prompt_logprobs: dict[str, list[LogprobsTensors]] = {}
45 + # req_id -> fixed-ID scoring state
46 + self.token_id_scores: dict[str, _TokenIdScores] = {}
47
48 def add_request(self, req_id: str, req_idx: int, sampling_params: SamplingParams):
49 uses_prompt_logprobs = sampling_params.prompt_logprobs is not None
@@ -29,9 +51,69 @@ class PromptLogprobsWorker:
51 self.num_prompt_logprobs[req_idx] = sampling_params.prompt_logprobs or 0
52 if uses_prompt_logprobs:
53 self.in_progress_prompt_logprobs[req_id] = []
54 + if sampling_params.prompt_logprob_token_ids is not None:
55 + self.token_id_scores[req_id] = _TokenIdScores(
56 + async_tensor_h2d(
57 + sampling_params.prompt_logprob_token_ids, self.device, torch.int64
58 + ),
59 + sampling_params.prompt_logprob_start or 0,
60 + )
61
62 def remove_request(self, req_id: str) -> None:
63 self.in_progress_prompt_logprobs.pop(req_id, None)
64 + self.token_id_scores.pop(req_id, None)
65 +
66 + def compute_prompt_token_id_logprobs(
67 + self,
68 + logits_fn: Callable[[torch.Tensor], torch.Tensor],
69 + hidden_states: torch.Tensor,
70 + input_batch: InputBatch,
71 + prompt_lens: np.ndarray,
72 + ) -> dict[str, torch.Tensor]:
73 + """Fill scores row (prompt row - req.start) per chunk; emit on the last."""
74 + if not self.token_id_scores:
75 + return {}
76 + logits_mode = self.logprobs_mode in ("raw_logits", "processed_logits")
77 + out: dict[str, torch.Tensor] = {}
78 + for i, req_id in enumerate(input_batch.req_ids):
79 + req = self.token_id_scores.get(req_id)
80 + if req is None:
81 + continue
82 + prompt_len = int(prompt_lens[input_batch.idx_mapping_np[i]])
83 + chunk_start = int(input_batch.num_computed_prefill_tokens_np[i])
84 + chunk_end = chunk_start + int(input_batch.num_scheduled_tokens[i])
85 + # Skip decode steps and, as compute_prompt_logprobs does, requests
86 + # resumed after preemption: their scores were already emitted.
87 + if chunk_start >= prompt_len or prompt_len < input_batch.prefill_len_np[i]:
88 + continue
89 + if req.scores is None:
90 + if chunk_start > req.start:
91 + # This prefill starts past the first scored row, so those
92 + # rows will never be written; drop the request instead of
93 + # emitting a buffer with unwritten rows.
94 + self.token_id_scores.pop(req_id, None)
95 + continue
96 + req.scores = hidden_states.new_empty(
97 + (max(prompt_len - 1 - req.start, 0), len(req.token_ids)),
98 + dtype=torch.float32,
99 + )
100 + # The last prompt row predicts the first decode token; skip it.
101 + lo = max(req.start, chunk_start)
102 + hi = min(chunk_end, prompt_len - 1)
103 + base = int(input_batch.query_start_loc_np[i]) - chunk_start
104 + for a in range(lo, hi, CHUNK_SIZE):
105 + b = min(a + CHUNK_SIZE, hi)
106 + logits = logits_fn(hidden_states[base + a : base + b])
107 + ids = req.token_ids.expand(b - a, -1)
108 + req.scores[a - req.start : b - req.start] = (
109 + logits.gather(-1, ids)
110 + if logits_mode
111 + else compute_token_logprobs(logits, ids)
112 + )
113 + if chunk_end >= prompt_len:
114 + out[req_id] = req.scores
115 + req.scores = None
116 + return out
117
118 def compute_prompt_logprobs(
119 self,
@@ -205,7 +287,6 @@ def compute_prompt_logprobs_with_chunking(
287 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
288 # Since materializing the full prompt logits can take too much memory,
289 # we compute it in chunks.
208 - CHUNK_SIZE = 1024
290 token_ids = []
291 scores = []
292 ranks = []