perf(qwen35): align the logits GEMM width so the output projection stops landing on an align-1 sm_75 kernel #1046
There are no files selected for viewing
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,79 @@ | ||
| # Qwen3.5 decode kernel attribution vs vLLM 0.27 (A100) | ||
|
|
||
| > **TL;DR:** nsys kernel-level attribution of the remaining serving gap (1×A100-40GB, upstream/main September 2026, vLLM 0.27.0, 1024-token prompts, `--cuda-graph-trace=node`, steady-decode session capture). Four findings: (1) at c16 total kernel-busy is 16.5 vs ~13.2 ms/step — GEMM family +3.5 ms/step (sm_80 decode buckets never tuned vs torch.compile-selected kernels); (2) at bs1 the GPU is 92% busy — kernel-busy 10.1 vs 8.2 ms/step, a genuine ~1.9 ms kernel-time deficit of which ~1.0 ms is FlashInfer `BatchDecodeWithPagedKVCache` (156 µs/layer-step vs `flash_fwd_splitkv` 26 µs, 6×); (3) per layer-step the GDN decode kernel is 97.0 µs vs vLLM's FLA `fused_recurrent` 43.6 µs (2.2×) and the FlashInfer paged decode 153.9 µs vs 64.6 µs (2.4×); (4) the once-per-step output-projection GEMM lands on a `cutlass_75_tensorop` align-1 kernel costing 1.67 ms/step (~12% of c16 TPOT) because the selection width is odd — fixed by aligning the selection width to the tile multiple (see the `perf(qwen35)` output-projection alignment change). | ||
| > | ||
| > **Last touched:** 2026-09 | ||
|
|
||
| ## Contract | ||
|
|
||
| - GPU 1× A100-SXM4-40GB (sm_80); model Qwen3.5-4B (Qwen/Qwen3.5-4B checkpoint); greedy, seed 42, random dataset (1024-token prompts). | ||
| - PegaInfer upstream/main (September 2026), default serve flags (serial, policy off — kernel composition is the subject, not scheduling). | ||
| - vLLM 0.27.0: FLASH_ATTN full attention + FLA Triton GDN + FlashInfer sampler, piecewise CUDA graphs. | ||
| - Capture: server under `nsys launch`/`start`/`stop` sessions with `--cuda-graph-trace=node`; capture armed after readiness + warmup, stopped after the bench; `--export=sqlite`; `nsys stats --report cuda_gpu_kern_sum`. | ||
| - Absolute times are node-trace inflated — used for **composition and cross-engine ratios at the same workload**, not as TPOT claims. Wall TPOT from the HTTP runs: PegaInfer bs1@1024 `9.82 ms`, vLLM `8.32`; c16 `14.26` vs `11.00`. | ||
|
|
||
| ## Findings | ||
|
|
||
| ### 1. c16 per-step kernel deficits: busy 16.5 vs ~13.2 ms/step (matches the 3.3 ms wall TPOT gap) | ||
|
|
||
| The decode shape itself is healthy: an `PEGAINFER_ITL_DEBUG` run shows 271 scheduler steps where `decode_n` ramps 0→16 over the first second (chunked prefill, one 1024-token request per unified step) and then stays **16/16 for the entire decode**; the clean capture agrees (270 implied steps, bucket-16 grids throughout). | ||
|
|
||
| The per-step kernel deficits at c16 (270 steps, bucket 16): | ||
|
|
||
| | family | PegaInfer | vLLM 0.27 | delta | | ||
| | --- | --- | --- | --- | | ||
| | GEMM family | 13.27 ms/step | ~9.78 | **+3.5** | | ||
| | GDN decode (24/step) | 2.33 (97.0 µs/layer-step) | 1.05 (43.6 µs) | **+1.3** | | ||
| | full-attn decode (8/step) | 1.23 (153.9 µs/layer-step) | ~0.53 (64.6 µs flash splitkv) | +0.7 | | ||
|
|
||
| vLLM numbers from its capture (251 steps × batch 16, `fused_recurrent_gated_delta_rule_packed` one launch per layer-step). | ||
|
|
||
| ### 2. bs1: GPU 92% busy — the deficit is kernel time, concentrated in paged decode attention | ||
|
|
||
| Clean single-request capture (256 steps, kernel-busy 92% of wall): PegaInfer **10.1 ms/step** kernel-busy vs wall TPOT `9.82`; vLLM `8.2 ms/step` busy on a full-coverage capture. Composition: GEMM/GEMV `7.42 ms/step` (73%, ≈ vLLM's ~7.1 — both near the weight-read bandwidth floor, not the gap), **full-attn paged decode `1.25 ms/step` vs vLLM ~`0.21` — FlashInfer `BatchDecodeWithPagedKVCacheKernel` at 156 µs/layer-step vs `flash::flash_fwd_splitkv_kernel` 26 µs (6×)**, GDN decode 0.53 vs 0.21, rest small. The ~1.9 ms/step deficit ≈ attention (+1.0) + GDN (+0.3) + GEMM (+0.3). | ||
|
|
||
| ### 3. Decode attention + GDN kernels: 2.1–2.4× per layer-step at batch 16 (bs1: up to 6×) | ||
|
|
||
| | kernel family | PegaInfer (batch 16) | vLLM (batch 16) | per layer-step | | ||
| | --- | --- | --- | --- | | ||
| | GDN decode | `gated_delta_rule_decode_batch_kernel` 97.0 µs ×24/step | `fused_recurrent_gated_delta_rule_packed` 43.6 µs ×24/step | **2.2×** | | ||
| | full-attn decode | FlashInfer `BatchDecodeWithPagedKVCacheKernel` 153.9 µs ×8/step | `flash::flash_fwd_splitkv_kernel` 64.6 µs ×8/step | **2.4×** | | ||
|
|
||
| At bs1@1024 the attention gap widens to 156 vs 26 µs/layer-step (**6×**): FlashInfer's paged decode is especially weak at tiny batch on sm_80. Our GDN kernel is one block per value head (grid = value_heads × batch, sequential k-head slices inside); FLA packs the work into one launch with a different tiling. | ||
|
|
||
| Neither kernel has had an sm_80 tuning pass (decode tuning history is sm_120/RTX 5090). | ||
|
|
||
| ### 4. The once-per-step output-projection GEMM runs on an sm_75-era align-1 kernel (1.67 ms/step) | ||
|
|
||
| `cutlass_75_tensorop_bf16_s1688gemm_bf16_128x64_tn_align1` runs **once per decode step at 1.67 ms** (280 instances ≈ 270 steps) — an sm_75-era kernel at alignment 1, ~12% of c16 TPOT for a single launch. One GEMM per step points at the output projection (`selection_vocab × hidden` over the bounded vocab): 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**. | ||
|
|
||
| Fixed by aligning the selection width to the 128-token tile multiple (`248077 → 248192`, still inside the 248,320-row checkpoint weight) at the `bound_selection_vocab` boundary, so every downstream buffer and the sampler stay consistent. The align-1 kernel disappears; measured on A100-40GB: c16 TPOT `14.26 → 13.51 ms` (−5.3%), c8 `11.98 → 11.23`, QPS16 `23.77 → 22.68` (vLLM `23.60`). | ||
|
|
||
| The widened width is a GEMM decision, not a token-space one, so two invariants keep it from leaking into selection. Both are silent when broken — the pad rows are trained embeddings with plausible logits: | ||
|
|
||
| - **The pad rows must be unselectable.** `248077..248192` are real checkpoint rows but not decodable tokens; `Qwen35Model::output_logits_into` forces them to `-inf` right after the output-projection GEMM (uploaded once at load, runs inside the decode CUDA Graph). Without it a permissive sample can emit an undecodable id and feed it back into later decode steps. The GEMM and the mask are one method rather than two statements at each site, because a site that ran the GEMM alone fails silently — the pad rows carry plausible logits. | ||
| - **The pad rows must not move the routing threshold.** `effectively_greedy`'s `top_p <= 1/vocab` nucleus has to be measured against the decodable width, not the arena: a `top_p` between `1/248077` and `1/248192` is effectively greedy and must keep taking the deterministic argmax path instead of the rejection sampler, which picks an arbitrary member of a bf16-tied top. `SampleScratch` carries that width separately (`with_selection_width`) so a padded arena routes exactly the rows an unpadded one would. | ||
|
|
||
| Neither invariant is covered by the HF gate above. `hf_golden_gate` replays fixed teacher-forced sequences and compares logprobs for the tokens it is given — it never selects a token, and the row logsumexp it scores against is mask-invariant (`exp(-inf) = 0`), so a missing mask is invisible to it by construction. The two are pinned elsewhere: | ||
|
|
||
| - The mask is asserted on the model itself by `pad_columns_are_suppressed_on_the_model_logits_paths`, which drives the real checkpoint through both logits paths (`batch_prefill_logits`, `batch_decode_graph`) and requires every column in `decodable_vocab..selection_vocab` to come out `-inf` while the decodable prefix does not. Commenting the mask out of `output_logits_into` fails it. | ||
| - Routing is pinned by `pegainfer-sample`'s `select_batch` tests; `padded_arena_width_does_not_suppress_the_argmax_routing` places `top_p` exactly on the `1/decodable` boundary with a padded arena. | ||
|
|
||
| Both are silent when broken — the pad rows are trained embeddings with plausible logits — which is why the GEMM and the mask are one method rather than two statements at each site. | ||
|
|
||
| ### 5. What is NOT the problem | ||
|
|
||
| - bs1 GEMM/GEMV: ours 7.42 vs vLLM ~7.1 ms/step — both near the weight-read bandwidth floor. | ||
| - Host/launch overhead: the clean captures show 91–92% GPU busy with only ~0.75 ms/step of inter-kernel gaps. | ||
| - Admission/batching shape at c16: decode_n stays 16/16 after the ramp. | ||
| - Prefill: not captured in this pass (decode attribution only). | ||
|
|
||
| ## Improvement queue (ordered by expected value) | ||
|
|
||
| 1. Retune the decode GEMM family on sm_80 (+3.5 ms/step at c16 — cublasLt algo selection vs torch.compile's kernel choices; possible qkv/z/b/a projection fusion). | ||
| 2. Replace or retune the full-attn paged decode path (FlashInfer `BatchDecodeWithPagedKVCache` → `flash_fwd_splitkv`-class: ~1.0 ms/step at bs1, 2.4× per layer-step at c16). | ||
| 3. Retune the GDN decode kernel vs FLA `fused_recurrent` (2.2× per layer-step). | ||
|
|
||
| ## Claim boundary | ||
|
|
||
| Single runs per capture, one GPU, node-trace-inflated absolute times (composition/ratio claims only), kernel-family grouping with template args stripped. Aggregate-step counts are inferred from kernel instance counts (24 GDN layers/step, 8 full-attn layers/step). The vLLM bs1 capture covers its full run (span 4.61 s ≈ 512 steps). Two earlier readings from polluted or partial captures — "admission waves at c16" and "a 5.9 ms/step bs1 host gap" — were retracted during review; every number here comes from the clean captures listed above. |
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -103,8 +103,13 @@ pub(crate) struct Config35 { | |
| /// `false` requires a top-level `lm_head.weight`; `true` reuses `embed_tokens`. | ||
| pub(crate) tie_word_embeddings: bool, | ||
|
|
||
| /// Token-selection width: `vocab_size` bounded to the frontend-decodable | ||
| /// vocab, then rounded up to the logits GEMM's tile multiple. Buffers and | ||
| /// the sampler arena span this width. | ||
| pub(crate) selection_vocab: usize, | ||
| /// Tokenizer-decodable width: the rows past it are suppressed to -inf and | ||
| /// are what the argmax-vs-sample routing is measured against. | ||
| pub(crate) decodable_vocab: usize, | ||
| } | ||
|
|
||
| impl Config35 { | ||
|
|
@@ -146,7 +151,8 @@ impl Config35 { | |
| self.linear_num_value_heads * self.linear_value_head_dim | ||
| } | ||
|
|
||
| /// Bound the output-selection width to the frontend-decodable vocab, which | ||
| /// also records the semantic width the argmax-vs-sample routing uses. | ||
| /// | ||
| /// The frontend decodes a dense prefix of the vocab; the checkpoint may pad | ||
| /// beyond it. Refusing a tokenizer wider than the checkpoint is the | ||
|
|
@@ -162,7 +168,15 @@ impl Config35 { | |
| vocab_size: self.vocab_size, | ||
| }); | ||
| } | ||
| // The logits GEMM's M dimension and the output matrix's leading | ||
| // dimension both take this width; an unaligned value (Qwen3.5-4B's | ||
| // tokenizer decodes 248077) forces cublasLt onto an align-1 sm_75-era | ||
| // kernel that costs ~1.7 ms/decode-step on sm_80. Round the GEMM width | ||
| // 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); | ||
| self.decodable_vocab = effective_vocab; | ||
| self.selection_vocab = aligned.min(self.vocab_size); | ||
|
Comment on lines
+177
to
+179
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 AGENTS.md reference: AGENTS.md:L122-L124 Useful? React with 👍 / 👎. |
||
| Ok(()) | ||
| } | ||
| } | ||
|
|
@@ -278,6 +292,7 @@ impl TryFrom<RawConfig> for Config35 { | |
| layer_types, | ||
| tie_word_embeddings, | ||
| selection_vocab: t.vocab_size, | ||
| decodable_vocab: t.vocab_size, | ||
| }) | ||
| } | ||
| } | ||
|
|
@@ -324,6 +339,37 @@ mod tests { | |
| assert_eq!(config.vocab_size, 1000); | ||
| } | ||
|
|
||
| #[test] | ||
| fn selection_vocab_bound_aligns_to_tile_multiple() { | ||
| let mut config = config(VALID_RAW); | ||
| config | ||
| .bound_selection_vocab(977) | ||
| .expect("977 decodes within vocab 1000"); | ||
| // 977 -> next 128 multiple (1024) clamped to the checkpoint's rows. | ||
| assert_eq!(config.selection_vocab, 1000); | ||
| assert_eq!(config.decodable_vocab, 977); | ||
|
|
||
| config | ||
| .bound_selection_vocab(769) | ||
| .expect("769 decodes within vocab 1000"); | ||
| // 769 -> 896: aligned, below the checkpoint rows, so the pad region is | ||
| // 769..896. | ||
| assert_eq!(config.selection_vocab, 896); | ||
| assert_eq!(config.decodable_vocab, 769); | ||
| } | ||
|
|
||
| #[test] | ||
| fn selection_vocab_bound_still_rejects_wider_tokenizers() { | ||
| let mut config = config(VALID_RAW); | ||
| let err = config | ||
| .bound_selection_vocab(1001) | ||
| .expect_err("a tokenizer wider than the checkpoint is rejected"); | ||
| assert!(matches!( | ||
| err, | ||
| ConfigError::EffectiveVocabExceedsCheckpoint { .. } | ||
| )); | ||
| } | ||
|
|
||
| #[test] | ||
| fn missing_tie_word_embeddings_is_typed() { | ||
| let json = VALID_RAW.replace("\"tie_word_embeddings\": true,", ""); | ||
|
|
||
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
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
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could we keep using the decodable vocabulary size when deciding whether a request takes the argmax path? This changes
selection_vocabfrom 248077 to 248192, andselect_batchpasses that padded width fromlogits.hidden_dimintoeffectively_greedy.For example, with
temperature=1,top_k=-1, andtop_p=4.03e-6, the old1 / vocabthreshold 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 existingtiny_top_p_routes_to_argmax_even_under_bf16_tiestest covers why this distinction matters.The
-infmask keeps padded IDs out of the result, but it doesn't preserve this routing decision. Could we usedecodable_vocabfor that decision while keeping the aligned width for the GEMM and buffer layout?