feat(logprobs): batched device-side logprobs reduction for Qwen3#723
Open
n-WN wants to merge 2 commits into
Open
feat(logprobs): batched device-side logprobs reduction for Qwen3#723n-WN wants to merge 2 commits into
n-WN wants to merge 2 commits into
Conversation
…ninfer-project#720) The engine contract collapsed `logprobs: 0` and "logprobs disabled" into one value, dropped `prompt_logprobs` at the wire boundary, and the bridge silently discarded `TokenEvent::PromptTokens`. As a result `echo=true`+`logprobs` and explicit `prompt_logprobs` requests hit HTTP 500 in response assembly, and `logprobs=-1` was silently mapped to 0 instead of failing fast. Engine contract (openinfer-engine): - `GenerateRequest.logprobs: usize` -> `Option<usize>` (None = disabled, Some(0) = chosen-token logprob only, Some(k) = +k alts) - `GenerateRequest.echo: bool` -> `prompt_logprobs: Option<usize>` vLLM frontend: - wire.rs: `requested_logprobs` preserves None/Some(0)/Some(k); `prompt_logprobs` is now read; `-1` (full-vocab sentinel) rejected early via `unsupported_sampling` for both fields, before any GPU work. - bridge.rs: submit maps `prompt_logprobs=Some(k)` to the engine prompt-logprobs request; `TokenEvent::PromptTokens` is encoded into `new_prompt_logprobs_tensors` with vLLM semantics (scored positions only, leading None position omitted; single-token prompt sends no payload). Qwen3 (only crate that computes prompt logprobs): - scheduler/plan/resolve gate on `prompt_logprobs.is_some()`; all echo invariants kept (whole-prompt single forward, no chunking, no prefix cache match, dedicated Prefill routing, all-position logits). Honor-or-reject for crates without prompt-logprob support: - Kimi-K2: reject prompt logprobs at admission (lm_head runs on last position only); row options take `Option<usize>`. - Qwen3.5: reject at admission instead of emitting all-None PromptTokens. - GLM5.2: admission message now names completion/prompt logprobs. - DeepSeek-V2-Lite: reject prompt logprobs (was silent all-None echo stub). Tests: - wire: None/0/k/-1 matrix for both fields. - bridge: PromptTokens -> wire payload (multi-token + single-token cases). - sim frontend_e2e: HTTP regression for echo+logprobs (was HTTP 500), explicit prompt_logprobs, logprobs=0, -1 early rejection, single-token prompt. - Model crates: scheduler/plan/executor tests migrated to Option fields; batch-invariance gates updated for Option step-item logprobs. CPU logprob computation is unchanged; GPU-side reduction is openinfer-project#719. Verified: - cargo check + clippy --all-targets: 9 crates clean - openinfer-vllm-frontend: 27 unit tests - openinfer-sim frontend_e2e: 19 HTTP regression tests - model-crate lib tests: 150 passed - qwen3 dflash_speculative_gate on H100: 4 passed - qwen3 hf_golden_gate (logprobs vs HF golden, bf16 tolerance): 1 passed
Decode rows requesting logprobs previously each paid a full-vocab
extract_vec + D2H (297 KiB) + stream sync + three O(V) host passes,
serialized in the shared step loop — ~0.85 ms per requesting row per
step on top of TPOT, plus a batch-wide stall for unrelated rows. Prompt
logprobs repeated the same per scored position.
This moves the whole reduction onto the GPU:
* New `logprobs.cu` kernels: per-row log-sum-exp (bf16 max pass + f64
exp-sum block reduction, mirroring the host reference's f64
accumulation) with the picked token's logprob; a row-index gather
feeding FlashInfer FilteredTopK (deterministic, smallest-index
tie-break), followed by SortTopKByIndex + StableSortTopKByValue so the
emitted order is exactly (value desc, token id asc) like the host
insertion pass.
* One batched call per step scores every logprobs row; D2H is
O(rows * (k + 1)) with a single sync, replacing O(rows * V) with a
sync per row. Rows sharing a step with different k ride one max-k
launch and truncate on the host (prefix property preserves exact
selection).
* Prompt logprobs reuse the same device reduction over
`all_position_logits` rows instead of per-position full-row extracts.
* Host path kept as per-row fallback for k > 2048 (FilteredTopK smem
cap) and pre-Hopper GPUs without FilteredTopK support.
* Frontend wire conversion now emits the chosen token plus all top-k
entries per position instead of skipping an entry when the sampled
token is already in the top-k set: the skip produced ragged rows that
vLLM's msgpack WireLogprobs width validation rejects, killing the
whole server on ordinary greedy + logprobs requests. Covered by a new
uniform-width wire test.
Semantics match `openinfer_sample::token_logprob_from_row`: top-k token
ids match exactly (tie order included); logprobs agree within 1e-4
(sequential f64 sum vs tree-reduced f64 partials), covered by new
model-free golden tests over random, tie-heavy, and chosen-only rows.
The hf_golden_gate logprobs gate passes unchanged.
Measured on Qwen3-4B, H100 (sm_90), `--no-prefix-cache`, greedy
~170-in/256-out, 3 reps, medians (baseline is the parent commit on the
same box):
| concurrency | path | logprobs off | logprobs=1 | logprobs=5 |
|---|---|---|---|---|
| 1 | before | 4.57 ms / 648 tok/s | 5.65 ms / 530 | 5.53 ms / 541 |
| 1 | after | 4.57 ms / 647 | 4.75 ms / 630 | 4.76 ms / 629 |
| 8 | before | 4.77 ms / 4851 | 12.39 ms / 1918 | 12.47 ms / 1910 |
| 8 | after | 4.77 ms / 4942 | 4.95 ms / 4754 | 4.96 ms / 4745 |
| 32 | before | 5.55 ms / 16544 | 36.04 ms / 2630 | 36.23 ms / 2621 |
| 32 | after | 5.77 ms / 15659 | 5.71 ms / 16085 | 5.75 ms / 15971 |
Mixed c8 batch (7 plain + 1 logprobs=1): plain-row TPOT 5.73 ms before
(+20.0% over the no-logprobs batch) vs 4.95 ms after (+3.8%, within
run-to-run noise of the 4.77 ms no-logprobs batch).
D2H/sync accounting via an LD_PRELOAD shim interposing the driver's
cuMemcpyDtoH{,Async}_v2 + cuStreamSynchronize/cuCtxSynchronize entry
points (nsys is broken on this box: CUDA 13.3 nsys reports only memsets
then TargetProfilingFailed). One c8 logprobs=5 wave, 8 reqs x 256
tokens, engine process totals: D2H drops from 2306 calls / 593.7 MiB
(one 297 KiB row per token per requesting row) to 1290 calls / 334 KiB
(~304 B of logprobs payload per step; the rest is fixed per-step infra
traffic present in both builds), and host syncs from 2308 to 518.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 01da9a30af
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| token_tx: req.token_tx, | ||
| logprobs: req.logprobs, | ||
| echo: req.echo, | ||
| logprobs: req.logprobs.unwrap_or(0), |
There was a problem hiding this comment.
When a client sends logprobs=0, the new engine contract treats it as a real request for the sampled token's logprob, but this unwrap_or(0) collapses it to the same value as None. Since admission_decision still only rejects req.logprobs > 0, DeepSeek-V2-Lite will accept logprobs=0 requests even though this path does not emit logprobs, so the frontend can return missing logprob data instead of honoring or rejecting the request. Keep the Option through admission or reject Some(0) before converting.
Useful? React with 👍 / 👎.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resolves #719.
What
Decode rows requesting logprobs previously each paid a full-vocab
extract_vec+ D2H (297 KiB) + stream sync + three O(V) host passes, serialized in the shared step loop — ~0.85 ms per requesting row per step on top of TPOT, plus a batch-wide stall for unrelated rows. Prompt logprobs repeated the same per scored position.This moves the whole reduction onto the GPU:
logprobs.cukernels: per-row log-sum-exp (bf16 max pass + f64 exp-sum block reduction, mirroring the host reference's f64 accumulation) with the picked token's logprob; a row-index gather feeding FlashInferFilteredTopK(deterministic, smallest-index tie-break), followed bySortTopKByIndex+StableSortTopKByValueso the emitted order is exactly (value desc, token id asc) like the host insertion pass.all_position_logitsrows instead of per-position full-row extracts.to_wire_position_logprobsskipped an entry when the sampled token was already in the top-k set, producing ragged rows that vLLM's msgpackWireLogprobswidth validation rejects — killing the whole server on ordinary greedy + logprobs requests. Now emits chosen + all top-k entries (uniform width), with a new regression test.Semantics
Matches
openinfer_sample::token_logprob_from_row: top-k token ids match exactly (tie order included); logprobs agree within 1e-4 (sequential f64 sum vs tree-reduced f64 partials). Covered by new model-free golden tests (random / tie-heavy / chosen-only rows).hf_golden_gatelogprobs gate passes unchanged. Frontend 28/28, qwen3 lib 75/75, clippy-D warningsclean.Bench — Qwen3-4B, H100 (sm_90),
--no-prefix-cache, greedy ~170-in/256-out, 3 reps, medians (baseline = parent commit, same box)Mixed c8 batch (7 plain + 1 logprobs=1): plain-row TPOT 5.73 ms before (+20.0% over the no-logprobs batch) vs 4.95 ms after (+3.8%, within run-to-run noise).
D2H / sync accounting
Via an LD_PRELOAD shim interposing
cuMemcpyDtoH{,Async}_v2+cuStreamSynchronize/cuCtxSynchronize(nsys is broken on this box — CUDA 13.3 nsys reports only memsets thenTargetProfilingFailed). One c8 logprobs=5 wave (8 reqs × 256 tokens), engine-process totals:The remaining D2H/sync traffic is fixed per-step infra (sampled-token readback etc.) present in both builds; logprobs payload is now ~304 B per step.
E2E: streaming completions with
logprobsandprompt_logprobsverified against the live server, including the greedy chosen-token-in-top-k case that previously crashed it.