perf(qwen35): align the logits GEMM width so the output projection stops landing on an align-1 sm_75 kernel - #1046
Conversation
Steady-decode kernel captures of both engines on the same GPU and workload (node-trace, composition and ratio claims only). Four findings for the remaining serving gap: the c16 GEMM family runs +3.5 ms/step (sm_80 decode buckets never tuned), the GDN decode kernel is 2.2x and the FlashInfer paged decode 2.4x per layer-step (6x at bs1) against FLA/flash references, the once-per-step output-projection GEMM lands on an align-1 sm_75 kernel at 1.67 ms/step because the selection width is odd, and bs1 is 92% GPU-busy (no host-gap problem). Improvement queue ordered by expected value. Signed-off-by: CAICAIIs <3360776475@qq.com>
…ops landing on an align-1 sm_75 kernel Kernel attribution (previous commit) found the once-per-step output-projection GEMM picking cutlass_75_tensorop_bf16_128x64_tn_align1 at 1.67 ms/decode-step (~12% of c16 TPOT): the selection width is Qwen3.5-4B's tokenizer-decodable vocab 248077 — odd — so both the GEMM M and the logits leading dimension defeat cublasLt's vectorized ampere kernels. bound_selection_vocab now rounds the selection width up to the 128-token tile multiple (248077 -> 248192), clamped by the checkpoint's 248,320 rows so the GEMM still reads inside the mapped weight; the extra rows are real trained embeddings, and sampling over them matches HF, which computes logits over the full checkpoint vocab. Downstream buffers and the sampler key off the same config width, so nothing else moves. Measured (A100-40GB, vLLM 0.27.0 baseline, 1024-token prompts, zero failed requests): c16 TPOT 14.26 -> 13.51 ms (-5.3%), c8 11.98 -> 11.23, QPS16 23.77 -> 22.68 (vLLM 23.60); throughputs +4-6%. hf_golden_gate TP1 2/2 and the non-graph TP2 tests green; the tp2_graph NCCL precapture failure reproduces on clean upstream main without this change (host-environmental, differential-tested). e2e_scheduler passed. Signed-off-by: CAICAIIs <3360776475@qq.com>
![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: 10cea8f5f5
ℹ️ 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".
| let aligned = effective_vocab.next_multiple_of(128); | ||
| self.selection_vocab = aligned.min(self.vocab_size); |
There was a problem hiding this comment.
When the checkpoint has more rows than the tokenizer (the documented 4B case is 248,320 versus 248,077), this makes IDs 248077–248191 eligible for greedy and sampled selection even though tokenizer_effective_vocab explicitly defines only the frontend-decodable prefix. The bridge forwards sampled IDs without a range check, so permissive sampling can emit an undecodable/blank token and then feed that ID back into subsequent decode steps. Keep the GEMM width padded, but mask the extra logits to -inf or retain effective_vocab as a separate sampling bound; the HF top-logprob replay cited here does not exercise selection from these rows, and the repository requires model-eval evidence for output changes.
AGENTS.md reference: AGENTS.md:L122-L124
Useful? React with 👍 / 👎.
Review follow-up on pegainfer-project#1046 (Codex P1): tile-aligning the selection width makes IDs past the tokenizer-decodable vocab (248077..248192 on 4B) eligible for greedy and sampled selection, and the bridge forwards sampled IDs without a range check — a permissive sample could emit an undecodable token and feed it back into later decode steps. The logits GEMM keeps the aligned width; each logits GEMM is now followed by the existing SuppressIds kernel forcing the pad rows to -inf (a no-op when the width needs no padding). Greedy and sampled selection therefore scan exactly the pre-fix decodable vocab — bit-identical behavior, pure GEMM perf. The suppress list is uploaded once at model load and the suppression runs inside the decode CUDA Graph, so there is no per-step host cost. A100-40GB re-run with the mask: c8/c16/QPS16 TPOT 11.50/14.09/35.48 ms — within run noise of the unmasked fix; hf_golden_gate TP1 2/2, non-graph TP2 2/2 (tp2_graph NCCL precapture failure is the pre-existing host-environmental one, differential-tested), e2e_scheduler passed. Signed-off-by: CAICAIIs <3360776475@qq.com>
| @@ -169,6 +173,7 @@ impl Config35 { | |||
| // up to the tile-aligned multiple — the checkpoint's remaining rows are | |||
| // real trained embeddings and stay inside the mapped weight. | |||
| let aligned = effective_vocab.next_multiple_of(128); | |||
There was a problem hiding this comment.
Could we keep using the decodable vocabulary size when deciding whether a request takes the argmax path? This changes selection_vocab from 248077 to 248192, and select_batch passes that padded width from logits.hidden_dim into effectively_greedy.
For example, with temperature=1, top_k=-1, and top_p=4.03e-6, the old 1 / vocab threshold routes the request to argmax, while the new threshold sends it to the rejection sampler. That can change which token wins when the largest BF16 logits are tied. The existing tiny_top_p_routes_to_argmax_even_under_bf16_ties test covers why this distinction matters.
The -inf mask keeps padded IDs out of the result, but it doesn't preserve this routing decision. Could we use decodable_vocab for that decision while keeping the aligned width for the GEMM and buffer layout?
Review follow-up on pegainfer-project#1046 (FeathBow): tile-aligning the selection width moved `select_batch`'s `effectively_greedy` threshold from 1/248077 to 1/248192, because the routing read `logits.hidden_dim`. A request whose `top_p` sits between those two values (e.g. 4.03e-6) used to take the deterministic argmax path and would now fall to the rejection sampler, which can pick either member of a bf16-tied top — the behavior `tiny_top_p_routes_to_argmax_even_under_bf16_ties` exists to pin. The GEMM width stays aligned; only the routing width moves. SampleScratch now carries both: `vocab` is the (possibly aligned) arena width the buffers are sized for, `selection_width` is the semantic bound the argmax-vs-sample decision is measured against. `new` keeps its meaning (selection_width == vocab), and qwen35 builds its scratch through `with_selection_width(selection_vocab, decodable_vocab)` so a padded arena routes exactly the rows an unpadded one would. Since every pad row is already suppressed to -inf, this restores the pre-fix routing with no change to what can be selected, and a width wider than the arena is refused at construction. Tests: `padded_arena_width_does_not_suppress_the_argmax_routing` (512-column arena, 256 decodable, top_p = 1/256 exactly on the boundary) asserts both the sampler-path control and the argmax result; `selection_width_above_the_arena_is_rejected` covers the fail-closed bound. Not run here: this checkout has no CUDA toolkit or GPU, so the GPU test targets still need a run on the A100 box. Signed-off-by: CAICAIIs <3360776475@qq.com>
Review pass over the two follow-up commits. Comments here were restating the diff instead of recording a constraint a reader cannot recover from the code; cut them to the one line that cannot be. Drops the `selection_width()` getter and its assertions (the behaviour assertions already pin the routing), and the construction-guard test, which covered a bounds check that is an assertion on a programmer error rather than an input boundary. Signed-off-by: CAICAIIs <3360776475@qq.com>
…nds on The attribution doc recorded the alignment but not what keeps it from leaking into token selection. Widening the GEMM is a throughput decision, so the pad rows have to stay both unselectable (suppressed to -inf) and outside the argmax-vs-sample routing width. Both failure modes are silent, because the pad rows are real trained embeddings with plausible logits. Signed-off-by: CAICAIIs <3360776475@qq.com>
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Somethings I'd like to settle before merging. The suppression invariant is a call-site convention today and nothing in the suite would catch a fourth logits site that skips it, which matters because this PR's own doc says the failure is silent, btw some parts annotations outdated.
| self.lm_head.as_ref().unwrap_or(&self.embed_tokens) | ||
| } | ||
|
|
||
| /// Force the tile-alignment pad rows of a logits buffer to -inf. |
There was a problem hiding this comment.
Right now the aligned GEMM and the pad-row mask are two separate statements repeated at three sites (prefill.rs:120-128, batch_decode.rs:697-705, batch_decode.rs:776-784). A fourth logits site can call the GEMM and skip the mask, and nothing in the suite would catch it. The config unit tests only assert config values, the new sampler test builds its own arena, and hf_golden_gate compares HF top-K, which is only sensitive to the mask if HF's stored top-K lands inside 248077..248192.
That matters more than usual here because this commit's own doc says the failure is silent: the pad rows are trained embeddings with plausible logits, so a missed mask surfaces as an undecodable id on the wire and fed back into later decode steps, not as a crash.
Would you consider folding the two statements into one method so the invariant cannot be skipped?
Review follow-up on pegainfer-project#1046 (FeathBow): the aligned GEMM and the pad-row mask were two statements repeated at three sites, so a fourth logits site could run the GEMM and skip the mask with nothing in the suite catching it — and this PR's own doc says that failure is silent, because the pad rows are trained embeddings with plausible logits. output_logits_into now owns both statements, and the three sites call it, so the GEMM is no longer reachable without the mask. The re-exports the old split needed are gone with it. Signed-off-by: CAICAIIs <3360776475@qq.com>
The output-projection GEMM is now reached only through output_logits_into, so the raw matrix accessor has no callers outside this module (the GEMM tuning helper being the remaining one). Private stops a future logits site from reaching the GEMM without the pad-row mask. Verified by compiling with the tightened visibility. Signed-off-by: CAICAIIs <3360776475@qq.com>
|
@codex review please |
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
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". |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
non-blocking but docs outdated.
Review follow-up on pegainfer-project#1046 (FeathBow: annotations outdated). Three claims did not match the code or the evidence: - The doc said hf_golden_gate validates output equivalence at TP1/TP2. What actually ran on this branch is TP1 2/2; the gate replays fixed teacher-forced sequences and compares logprobs, so it never selects a token, and the row logsumexp it scores is mask-invariant (exp(-inf) = 0). It cannot see the mask by construction, which is now stated instead of implied. - The optimization ledger's LM-head rows were measured while selection spanned the full 248,320 rows; they are marked pre-alignment and pointed at the attribution doc. Signed-off-by: CAICAIIs <3360776475@qq.com>
Review follow-up on pegainfer-project#1046 (FeathBow): the doc admitted nothing in the suite would catch a logits site that ran the output projection without the mask. The structural fix removes the sites, but there was still no test of the invariant on the model itself. This drives the real checkpoint through both logits paths — batch_prefill_logits and batch_decode_graph — and asserts every column in decodable_vocab..selection_vocab came out -inf while the decodable prefix did not. Mutation-checked: commenting the mask out of output_logits_into fails it with "prefill: pad id 248077 survived selection (want -inf)". Signed-off-by: CAICAIIs <3360776475@qq.com>
The paragraph said nothing in the suite would catch a logits site that skipped the mask. That is no longer true, so point at the test instead of the gap. Signed-off-by: CAICAIIs <3360776475@qq.com>
CI's `-D warnings` rejected six lints in the test added by the previous commit: manual indexing in the pad loop, two useless f32 conversions (to_host already returns f32), a float equality compare, a redundant closure, and a one-element vec. The comparison now checks the value is negative infinity, which is what the invariant means, and reports the offending value when it fails. Signed-off-by: CAICAIIs <3360776475@qq.com>
2 participants
Summary
The once-per-step output-projection GEMM (logits over the selection vocab) was landing on
cutlass_75_tensorop_bf16_s1688gemm_bf16_128x64_tn_align1— an sm_75-era kernel at alignment 1 — costing 1.67 ms/decode-step, ~12% of c16 TPOT. Root cause: Qwen3.5-4B's tokenizer-decodable vocab is 248077 — odd — so both the GEMM M dimension and the logits output leading dimension defeat cublasLt's vectorized ampere kernels.bound_selection_vocabnow rounds the selection width up to the 128-token tile multiple (248077 → 248192), clamped by the checkpoint's 248,320 rows so the GEMM still reads inside the mapped weight. The extra rows are real trained embeddings, and sampling over them matches HF, which computes logits over the full checkpoint vocab. Downstream buffers and the sampler key off the same config width, so nothing else moves. Found via nsys kernel attribution; the full doc ships in the second commit (docs/models/qwen35/decode-kernel-attribution.md).Measured (1x A100-40GB, upstream/main
139d925e+ this change, vLLM 0.27.0 baseline)1024-token prompts, greedy,
vllm bench serve, zero failed requests; single run per cell:Measured on top of the merged
auto+streamposture (#1033/#1034): c16 TPOT −5.3%, c8 −6.3%, QPS16 −4.6% (back below vLLM), throughputs +4–6%, TTFT/ITL unchanged or better.Review follow-ups
Three commits sit on top of the perf commit:
248077..248192eligible for greedy/sampled selection. Each logits GEMM is now followed by the existingSuppressIdskernel forcing those rows to-inf, so selection scans exactly the pre-fix decodable vocab. The suppress list is uploaded once at load and runs inside the decode CUDA Graph (no per-step host cost).select_batchread the routing width fromlogits.hidden_dim, soeffectively_greedy'stop_p <= 1/vocabthreshold moved from1/248077to1/248192— a request withtop_pbetween those values would drop off the deterministic argmax path onto the rejection sampler, which can pick either member of a bf16-tied top.SampleScratchnow carries the arena width and the semantic selection width separately (newkeepsselection_width == vocab; qwen35 builds throughwith_selection_width(selection_vocab, decodable_vocab)), so a padded arena routes exactly the rows an unpadded one would.Qwen35Model::output_logits_intonow owns both statements and the three sites call it, andoutput_projection()is private (the GEMM-tuning helper is its only remaining caller), so the GEMM is not reachable without the mask. The ops re-exports the old split needed are gone with it.Testing
Verified on 1x A100-SXM4-40GB (
dgx, CUDA 11.5, sm_80), branch atfe48b666, modelQwen3.5-4B, built with--features qwen35andPEGAINFER_TRITON_PYTHONon a Triton 3.7.1 venv.lib tests for the touched packages (
pegainfer-sample,pegainfer-qwen35) 102 passed / 0 failed / 7 ignored, including the twobound_selection_vocabalignment/clamp unit tests.pad_columns_are_suppressed_on_the_model_logits_paths— new, and the one test that covers the mask invariant on the model itself rather than on a synthetic arena. It drives the real checkpoint through both logits paths (batch_prefill_logitsandbatch_decode_graph) and requires every column indecodable_vocab..selection_vocab(248077..248192) to come out-infwhile the decodable prefix does not. Mutation-checked: commenting the mask out ofoutput_logits_intofails it withprefill: pad id 248077 survived selection (want -inf), so it is a test with teeth rather than a tautology.pegainfer-sampleintegration suite 11 passed / 0 failed (was 9), including the newpadded_arena_width_does_not_suppress_the_argmax_routing: a 512-column arena with 256 decodable andtop_p = 1/256exactly on the boundary, where the arena-width scratch routes to the sampler and can return the bf16-tied peer, and the decodable-width scratch keeps the deterministic argmax for all 64 seeds.hf_golden_gate2/2 passed (TP1; the TP2 graph variant is the pre-existing host-environmental NCCL precapture issue noted above). It replays fixed teacher-forced sequences and compares logprobs against HF fixtures — it never selects a token, and the row logsumexp it scores is mask-invariant (exp(-inf) = 0), so it is evidence that the aligned width leaves prefill/decode numerics alone, not evidence about the mask. That distinction is now written down in the attribution doc instead of being implied by a blanket "output equivalence" claim.e2e_scheduler2/2 passed.Production A/B, same session on the same GPU (base
139d925vs this branch, real server on GPU 7,vllm bench serverandom 1024/256, greedy, seed 42, zero failed requests in every cell):Both cells improve ~5.4%, so the speed-up reproduces against a same-session baseline rather than only across sessions, and the review follow-ups do not give any of it back. Re-measured after the
output_logits_intorefactor (at05ffa360; the only later commit makesoutput_projectionprivate): c1614.12 ms/ TTFT 708.5, c811.48 ms/ TTFT 377.2, 0 failed.HTTP invariant probe against the branch server, which keys the observed contract rather than raw text equality (greedy decode on this host is not bit-reproducible, see below): across 4 prompts with
logprobs=1, all 96 emitted token strings are real tokenizer tokens with zero empty/undecodable strings, so no id at or above the decodable248077ever reaches the wire;temperature=1, top_p=4.03e-6(the case FeathBow raised) returns byte-identical text to greedy on every prompt; andtop_p=0.9differs from greedy on 3 of 4 prompts, confirming the wide-nucleus sampling path is still live rather than everything being pinned to argmax.Widths confirmed against the checkpoint: decodable
248077, weight rows248320, aligned selection248192, so the suppressed region is exactly248077..248192(115 rows) and the GEMM stays inside the mapped weight.Kernel check on a clean c16 trace: the align-1 kernel no longer appears; the replacement ampere GEMM family drops the step cost by ~0.7–1.3 ms.
What was NOT verifiable, and why
hf_golden_gate, the sampler tests) are the sound evidence for output behaviour.--test-threads=1; the flake below is why. The fullcargo test --workspace --libcould not run: the workspace build pulls themoefeature, which needs NCCL >= 2.30.4, and every NCCL on this host is 2.29.7. The touched-package lib tests were run instead.139d925, always asbatch sampling kernel failed with error 1from whichever test happened to overlap; with--test-threads=1the suite passed 4/4 on both. Earlier in this PR that flake was misattributed tomin_p_row_takes_the_sampler_path_and_filtersspecifically — it moves between tests, which is what pointed at cross-test GPU contention rather than any one test. Pre-existing, not caused by this PR, and not fixed here. Every serial run quoted above is--test-threads=1.Claim boundary
Single-run cells on one GPU; the GDN/full-attn decode kernel gaps (2.2×/2.4× per layer-step, 6× at bs1) documented in the attribution doc remain open and are the next optimization targets. The perf table's "before/after" cells predate the review follow-ups; the re-run above is within run noise of the original fix.