@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 64 lines 2.44 KB
Rendered Raw
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.