Description
Qwen3 currently handles every decode row that requests logprobs by extracting a full logits row, copying it to the host, and calling the shared CPU implementation token_logprob_from_row. The Qwen path performs that work synchronously for each requested row in the shared step loop (decode result construction, host readback and reduction). The CPU routine makes full-vocabulary max and f64-exp-sum passes, followed by an optional top-k pass, so its cost grows with the number of logprobs rows in the batch and stalls unrelated rows before the next decode step can start.
Measured on Qwen3-4B, one sm_89 GPU, commit f5d2bf2ad, --no-prefix-cache, greedy 128-input/256-output requests, three repetitions with medians:
| concurrency |
logprobs off |
logprobs=1 |
logprobs=5 |
| 1 |
9.82 ms TPOT / 102 tok/s |
10.69 ms / 94 tok/s |
10.71 ms / 93 tok/s |
| 8 |
10.72 ms / 730 tok/s |
17.40 ms / 456 tok/s |
17.44 ms / 455 tok/s |
| 32 |
13.08 ms / 2328 tok/s |
41.34 ms / 788 tok/s |
39.87 ms / 788 tok/s |
The added step time is approximately 0.85 ms per requested row at all three concurrency levels. In a mixed concurrency-8 batch, one logprobs request raises non-logprobs rows from 10.70 ms to 11.43 ms TPOT (+6.8%) and lowers total throughput from 733 to 686 tok/s (-6.4%). This makes logprobs a batch-wide multi-tenant latency cost rather than a cost isolated to the requesting row.
Nsight Systems counted one approximately 297 KiB D2H per requested token. CUDA copy, synchronization, allocation, and related API work total approximately 30 µs per extraction. Because the capture used --sample=none, it does not directly sample individual CPU functions; the approximately 0.82 ms residual is attributed to the host conversion and full-vocabulary scans by the end-to-end slope, process CPU-time delta, and the source path above. The near-identical logprobs=1 and logprobs=5 results show that the fixed logsumexp/full-row path dominates the incremental top-k work.
Prompt logprobs have the same scaling problem: a 1024-token prompt performs one full-row extraction per scored position and adds approximately 0.8 ms per prompt token. The device-side implementation should therefore be reusable for both decode rows and batches of prompt positions.
Suggested direction: compute logsumexp, the selected-token logprob, and the requested top-k while logits remain on the GPU, then return a compact O(k+1) result rather than an O(vocabulary) row. This is an output contract, not a requirement that the implementation be a single kernel; a numerically sound multi-stage reduction is acceptable.
Acceptance criteria:
- Decode logprobs no longer perform one full-vocabulary D2H and CPU scan per requested row.
- Prompt logprobs reuse the device reduction and do not extract every full logits row to the host position by position.
- For finite k, host readback is O(number of requested rows × (k+1)), including the selected-token result, rather than O(number of requested rows × vocabulary); full-vocabulary
-1 behavior remains governed by the frontend contract issue below.
- Selected-token logprobs, requested top-k entries, deterministic tie ordering, and error behavior match the current host reference within an explicit numerical tolerance.
- A before/after run of the concurrency and mixed-batch matrix reports TPOT/throughput, D2H count and bytes, and synchronization count. Kimi's larger whole-arena readback should be measured when it adopts the shared implementation; it is not a prerequisite for the Qwen fix.
Description
Qwen3 currently handles every decode row that requests logprobs by extracting a full logits row, copying it to the host, and calling the shared CPU implementation
token_logprob_from_row. The Qwen path performs that work synchronously for each requested row in the shared step loop (decode result construction, host readback and reduction). The CPU routine makes full-vocabulary max and f64-exp-sum passes, followed by an optional top-k pass, so its cost grows with the number of logprobs rows in the batch and stalls unrelated rows before the next decode step can start.Measured on Qwen3-4B, one sm_89 GPU, commit
f5d2bf2ad,--no-prefix-cache, greedy 128-input/256-output requests, three repetitions with medians:logprobs=1logprobs=5The added step time is approximately 0.85 ms per requested row at all three concurrency levels. In a mixed concurrency-8 batch, one logprobs request raises non-logprobs rows from 10.70 ms to 11.43 ms TPOT (+6.8%) and lowers total throughput from 733 to 686 tok/s (-6.4%). This makes logprobs a batch-wide multi-tenant latency cost rather than a cost isolated to the requesting row.
Nsight Systems counted one approximately 297 KiB D2H per requested token. CUDA copy, synchronization, allocation, and related API work total approximately 30 µs per extraction. Because the capture used
--sample=none, it does not directly sample individual CPU functions; the approximately 0.82 ms residual is attributed to the host conversion and full-vocabulary scans by the end-to-end slope, process CPU-time delta, and the source path above. The near-identicallogprobs=1andlogprobs=5results show that the fixed logsumexp/full-row path dominates the incremental top-k work.Prompt logprobs have the same scaling problem: a 1024-token prompt performs one full-row extraction per scored position and adds approximately 0.8 ms per prompt token. The device-side implementation should therefore be reusable for both decode rows and batches of prompt positions.
Suggested direction: compute logsumexp, the selected-token logprob, and the requested top-k while logits remain on the GPU, then return a compact O(k+1) result rather than an O(vocabulary) row. This is an output contract, not a requirement that the implementation be a single kernel; a numerically sound multi-stage reduction is acceptable.
Acceptance criteria:
-1behavior remains governed by the frontend contract issue below.