[envy-ai]

Serial top-K logprobs currently copy the full vocabulary to the CPU and scan it for each emitted token. This change computes the chosen-token score and top-K on the GPU for ordinary serial decode, including accepted MTP rows, and transfers only those results. For vocabulary 248,320 and K=20, the result download changes from 993,280 to 168 bytes per row. This is a data-size comparison, not a measured throughput result.

This targets the logprobs branch behind Niko1221/Strata#3200. It preserves the T/LP protocol and existing server behavior. Named allowed_token_ids, pipeline windows, gfx906 and unsupported dimensions retain CPU scoring. The reduction uses shifted log-sum-exp, deterministic token-ID ordering for ties, and at most 20 alternatives. GPU rounding and tie ordering can differ from the CPU path. Requests without logprobs launch no new scoring work. No RAM-standby changes are included.

The optional startup-allocation follow-up is included: Verifier::init() allocates the three separate logprobs buffers, so their allocation failures are reported at startup rather than during the first GPU-scored request. Existing !ptr guards preserve those buffers on repeated initialization. Zero scratch size allocates none of them and retains the decoder's CPU fallback; named IDs remain on the existing CPU path. This reserves a small fixed amount even when logprobs are never requested: 83,360 bytes for max_t=8, vocabulary 248,320 and K=20.

Validation on 7265fc75dfbbf714cc03d740e95b41f9ea0586fb: Linux, RTX 4090, CUDA 12.8, SM89, GCC 14.3.1. Commands below normalize the isolated build location to $BUILD; $PYTHON is an isolated CPython 3.13.9 environment with repository-pinned dependencies.

Passed:

• Full Release build: cmake --build "$BUILD" --target strata logprobs_parity -j2. Fresh pinned ggml dependency 3cf03257f219afbe7334045ff7c6a06ac68c627d; changed translation units compiled and linked. Build took 7m52s, with 1.5 GiB reported peak host memory and zero swap.
• "$BUILD/logprobs_parity": six numeric cases, including vocabularies 1/4097/8193/248320, K=0/1/20 and K=vocabulary, multiple rows, ties, a chosen token outside top-K, and finite extreme/equal logits. GPU scores match the double-precision CPU reference within 3e-5.
• A temporary local harness linked the actual candidate libraries and called the complete Verifier::init() with synthetic metadata and real CUDA allocations. Supported and unsupported-vocabulary cases each initialized twice; the three optional buffers were reused or absent as appropriate. It checked all three injected allocation-error returns, immediate failure, actual destructor cleanup, and Verifier::logprobs() versus CPU reference for rows 1/8 and K=0/20. This verifies the three logprobs buffers, not general repeated-init leak freedom or natural hardware OOM.
• All 33 existing serve.test_logprobs cases passed across compatible runs: 15 non-API cases in $PYTHON -m unittest -v serve.test_logprobs, then all 18 API cases in $PYTHON -m unittest -v serve.test_logprobs.Api after resolving the missing dependency described below.
• Real-model API checks passed for ordinary decode, GPU top-K 0, GPU top-K 20, named-ID CPU renormalization, and ordinary decode after logprobs. Model: Qwen3.8-Flash-Next GSQ-RCO IQ3_S, plain tokenizer, MTP/spec 4, INT8 KV, context limit 131072, serial requests, vision disabled. The five requests used temperature 0, a fixed short prompt and max_tokens=24. All returned Paris.; this is one observed sequence match, not a general bit-exact guarantee. K0/K20 scores can vary across requests.
• git diff --check and final tracked-tree checks.

Builds and standalone tests used bounded transient systemd units with a 4 GiB host-memory cap and no swap. The isolated full-model unit used an 88 GiB cap and no swap; observed host peak was 61.9 GiB. Its corrected --expert-cache 6452 sets a nominal pool budget below 16 GiB for this pack's largest aligned blob of 2,662,400 bytes; profile/free-room limits produced a 13.27 GiB cache. VRAM reserve was 2048 MiB. Sampled candidate GPU peak was 21,900 MiB, with at least 1,479 MiB physically free. These are correctness-test observations, not benchmark results. The original serving deployment was restored and its text/logprobs health check passed afterward.

Failed attempts, recovered:

• System Python initially lacked jinja2; the first isolated full-suite attempt then had 18 API import errors because regex was missing. Repository-pinned dependencies were installed only in the isolated test environment; the 18 affected cases were rerun and passed.
• The first model startup was stopped by the VRAM guard before API tests: 22,416 MiB candidate use and 967 MiB free. Its --expert-cache 8192 was mistakenly treated as a MiB value; the parameter is a count-based byte budget and produced a 14.27 GiB cache. The corrected-count/reserve retry passed all five API cases within the original total resource limits. Failed-run evidence was retained.

Not run: HIP compilation/hardware execution; full SYCL/Metal builds or execution; a controlled before/after performance benchmark; long-prefix/long-generation stress. The 128K context value above is a configured limit, not a 128K-prompt test. SYCL/Metal scoring paths are unchanged. Keeping this PR draft under the repository's backend-validation rules; targeted HIP validation is still needed. No throughput improvement is claimed from these tests.

Developed with AI assistance; the diff, actual command results and limits were reviewed before publication.

[xiezhipeng-git]

I built and ran both branches end-to-end (real model, real logprobs requests) before approving the workflow run.

Setup

• Local RTX 4090, Swift-Qwen3.8-Flash-Next-IQ3_XXS (~40 GB), served through the stock OpenAI-compatible server (serve/server.py --engine strata).
• base feat/openai-logprobs @ a82a9b4d vs PR gpu-token-logprobs @ a2ea8f95; identical CMake feature toggles, CUDA arch 89, MSVC, CUDA 13.3.
• Request: fixed prompt, temperature 0, max_tokens 200, logprobs=true, top_logprobs=20, thinking off. Two runs per configuration.
• Kernel microbenchmark rows were measured on Kaggle T4×2.
• One methodology note that matters here: the expert arena needs ~40 GiB of pinned host RAM. If something else is building or holding memory at the same time, the engine logs the expert arena needs 39.97 GiB of RAM but 39.39 GiB is available and decode collapses to ~50 tok/s. All numbers below were taken with the box otherwise idle.

Disclosure: the builds, the server runs and the log/result parsing here were executed with AI assistance (an assistant driving the build + benchmark scripts). Every number below comes from an actual run on this machine — nothing is estimated — but treat the write-up as machine-assisted analysis, and the code review as mine.

1. Build and test pass rates

Check	base feat/openai-logprobs	PR gpu-token-logprobs
Full engine build (strata)	ok (prebuilt at HEAD)	ok, no errors — logprobs.cu + generate.cpp + verify.cpp compiled and linked
logprobs_parity (GPU kernel vs double-precision CPU ref, 6 cases, tol 3e-5)	n/a — no GPU scoring path	PASS
serve.test_logprobs	32/32 PASS (0 fail / 0 err)	33/33 PASS (0 fail / 0 err)
End-to-end: 200 tokens with full logprobs	✅	✅

The extra PR test is test_engine_wire_keeps_chosen_score_separate_from_raw_top_k; it passes.

2. Speed — this is where the PR pays off

Decode throughput, 200 tokens, engine's own log line (N generated in T ms), two runs each:

Configuration	run 1	run 2	mean
base + logprobs (CPU scoring)	137.0 tok/s	143.2 tok/s	140.1 tok/s
PR + logprobs (GPU scoring)	174.0 tok/s	168.0 tok/s	171.0 tok/s
base, no logprobs (control)	177.0 tok/s	175.4 tok/s	176.2 tok/s

Derived
What logprobs costs on the CPU path	176.2 → 140.1 tok/s = −20.5% (−36 tok/s)
PR with logprobs vs the no-logprobs ceiling	171.0 vs 176.2 = −3.0%
base → PR, both with logprobs	140.1 → 171.0 = +22.1%

So the honest framing is: requesting logprobs used to cost a fifth of decode throughput, and with this PR it is essentially free (−3% off the ceiling). That is a bigger deal than the microbenchmark alone suggests, and it is well outside run-to-run spread (the two runs of each configuration differ by ≤6 tok/s).

Kernel-level numbers (Kaggle T4×2):

Metric	Old path (CPU full-vocab copy)	New path (GPU extract)	Gain
Per-call latency (vocab 248,320, top_k 20, 200 iters)	16.5343 ms	0.6418 ms	25.8×
Data returned per call	3,973,120 B (~3.79 MB)	672 B	5912×
Per-token payload	993,280 B	168 B	5912×

The old path copies the whole 248,320-way logit row back to the CPU and reduces it there every step; the new path reduces on the GPU and returns only chosen + top-K (672 B = 4 accepted tokens × 168 B). That removed per-step cost is exactly the 36 tok/s the control run shows.

3. Logprob values

I ran each branch twice to get a noise floor before attributing any value delta to this PR:

Comparison	tokens	max Δ chosen	max Δ top-K	top-K set diff
base vs base2 (same engine)	identical (200/200)	0.124	2.09	180
PR vs PR2 (same engine)	diverged at token 55	1.689	18.83	4966
base vs PR	identical (200/200)	0.123	2.33	171

Two things follow. First, base-vs-PR (0.123 / 2.33 / 171) sits at the same-engine noise floor (0.124 / 2.09 / 180), so the differences are the engine's own run-to-run rounding — expert-cache hit rate and the prefill CPU share decide which experts are computed where, and the engine itself logs that the GPU rounds differently from the CPU. Second, that noise is large enough to flip a near-tie: the two PR runs both had ~-0.98 on the chosen token at position 55 and picked different tokens (函数 vs 当前), after which the sequences diverge. So token-exact comparison is not a usable equivalence check here — which is precisely what logprobs_parity exists for (same logits, GPU kernel vs double-precision CPU reference, tol 3e-5), and it passes.

Two non-blocking observations

1. The run logs only 385 MiB of VRAM free with this model on a 48 GB card — runtime pressure, not a correctness issue, but worth a --vram-reserve-mib note for large models.
2. Float scoring can round differently from the CPU path (DETAILS.md already says so) — fine for logprobs, just flagging it for any downstream bit-exact assumption.

[xiezhipeng-git]

One optional, non-blocking follow-up from reading the engine side of this PR.

Consider allocating the three logprobs buffers in Verifier::init() rather than on the first GPU-scored request.

verify.cpp:4183/4185/4187 currently allocate logprobs_scratch_, logprobs_chosen_, and logprobs_out_ lazily. This happens only on the first call and the buffers are reused afterwards, so this is not a steady-state throughput concern, and the sizing does not depend on the request's top_k (the scratch uses max_t_ and min(kLogprobsMaxTop, n_vocab)). For max_t = 8 (the kVerifyMaxT ceiling), vocabulary 248,320, and K = 20, the three total 83,360 bytes: 81,984 for the scratch (8 rows x 61 blocks x 168) plus 32 and 1,344.

The reason to consider moving the allocations is failure timing. A first-use cudaMalloc failure is currently discovered while a request is running; the caller then prints ERR and returns 1 from main (generate.cpp:15270-15274), ending the engine process and interrupting all of its slots. Allocating during initialization would turn that particular capacity failure into a fail-fast startup error and would also remove the three allocation branches from the per-request path. init() already works that way for a far larger sibling: head_logits_ is carved from the arena in the same function, and the arena's own cudaMalloc failure already fails startup with a message, so both the precedent and the plumbing are there.

The tradeoff is that initialization would reserve this small amount of memory for every applicable verifier even when logprobs are never requested, and an allocation failure would move from first use to startup. Given the tiny fixed size, that fail-fast tradeoff seems reasonable to me, although keeping the lazy allocation is also defensible.

If you do move it, please keep scratch_bytes == 0 as the no-allocation case so unsupported vocabulary sizes can still start and retain the existing CPU scoring path. Note that the CPU path is selected by the caller's gate (generate.cpp:15264-15266 tests logprobs_scratch_bytes(1, vocab, req_logprobs) != 0), not inside logprobs(), which returns an error for that case; the == 0 branch there would then be unreachable from the decoder, so a comment saying so would help. Keeping the existing !ptr guards would also keep a repeated init() from leaking. I would leave these buffers separate from arena_; there is no need to complicate that sizing for such a small optional change.

Disclosure: I went through the engine side with AI assistance (an assistant reading and searching the sources). The control flow and buffer sizing above were checked against a2ea8f95; the judgement call is mine.

[xiezhipeng-git]

One note on the boundary this PR draws, which I think is the right one.

Keeping named allowed_token_ids on the existing CPU path is not a gap, and the reason is not only the output shape. On the named path the cost scales with the size of the named set, not with the vocabulary copy: a single max_tokens: 1 request naming 20,000 ids takes 51.7 s on the base branch, linear in the set size. What this PR removes is a per-token constant: measured on a local RTX 4090 with Swift-Qwen3.8-Flash-Next IQ3_XXS, logprobs cost 1.46 ms per token on the base branch and 0.17 ms here (176.2 / 140.1 / 171.0 tok/s). So GPU-scoring the named path would save about 1.3 ms inside requests measured in tens of seconds, and it is not worth a second output shape and a second buffer sizing in this change.

That tradeoff reverses once the named set becomes small and per-request. If upstream Niko1221#3743 lands (typed decisions scoring a few dozen single-token numeric labels per request), the per-token cost becomes the whole cost again, and this is worth revisiting as a separate change.

Disclosure: the numbers above come from runs I made on my own hardware and branch; the engine-side reading was done with AI assistance.

[envy-ai]

Implemented the optional startup-allocation follow-up in 7265fc75. The three buffers now allocate in Verifier::init(), retaining the pointer guards, zero-scratch no-allocation case, and named-ID CPU scoring.

On this SHA, the full CUDA/SM89 build, kernel parity, real-CUDA initializer/failure/reuse checks, all 33 server cases, and five real-model API cases passed. The updated PR body records the exact coverage, corrected cache configuration, and recovered failed attempt. No throughput or general bit-exact claim is made.

HIP compilation/hardware validation remains unavailable, so this PR stays draft. A targeted HIP check would help complete backend validation.