From 07dabf2e6c89cb74f6b75bce675f959ce65f3540 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Sun, 13 Sep 2026 13:31:18 +0800
Subject: [PATCH 01/12] docs(qwen35): nsys decode kernel attribution vs vLLM
 0.27 on A100

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>
---
 docs/index.md                                 |  1 +
 .../qwen35/decode-kernel-attribution.md       | 68 +++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 docs/models/qwen35/decode-kernel-attribution.md

diff --git a/docs/index.md b/docs/index.md
index 0f11d9e49..635409bfb 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -57,6 +57,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 | `models/qwen35/tp-implementation.md` | Qwen3.5 TP landed through P2c (2026-08, 2× RTX 4090): Phase 1/P2A lifecycle and ID contracts kept; GDR state sharded per rank (27B TP2 fits 48 GB pairs); batched eager TP decode; TP decode CUDA Graphs for 4B/9B (27B group-6 stays eager). 9B/27B TP2 HF + e2e gates pass; the throughput A/B on the merged stack is tracked in #1001. |
 | `models/qwen35/mixed-load-itl-470.md` | Issue #470: full cold `--max-batch 8/bg=4` matrix on RTX 4090 (24/24 valid) + starvation negative control. Qwen3.5 is not immune; chunking bounds max/per-step stall but raises p99 at low QPS (~14→~80–92ms) and pulls p99/max back from the prefill wall to the chunk wall at high load; `qps·prefill_s≳1` is a throughput wall (chunking can't fix it, and ON's +15% TTFT can trip it earlier). The old "p99 immunity" was a slot-starvation artifact. |
 | `models/qwen35/adaptive-scheduler-policy.md` | Issue #727 adaptive scheduler policy record: default `off`, opt-in `auto`, hard `--max-prefill-tokens` cap, TP `auto` rejection, and pre-review whole-prefill benchmark tradeoff retained as non-default evidence. |
+| `models/qwen35/decode-kernel-attribution.md` | nsys kernel attribution of the remaining serving gap vs vLLM 0.27 on A100: c16 kernel-busy 16.5 vs ~13.2 ms/step (GEMM family +3.5), bs1 92% busy with the deficit in FlashInfer paged decode (6x at bs1), GDN decode 2.2x per layer-step, and the once-per-step output-projection GEMM on an align-1 sm_75 kernel (1.67 ms/step, fixed by tile-aligning the selection width). |
 | `models/qwen35/unified-prefill-overlap.md` | Issue #715 implementation record: opt-in single-GPU shared-SM overlap keeps one prefill chunk in flight while active decode continues; default serial policy and unsupported-combination guards remain explicit. |
 
 ## models / gemma4
diff --git a/docs/models/qwen35/decode-kernel-attribution.md b/docs/models/qwen35/decode-kernel-attribution.md
new file mode 100644
index 000000000..bd0d073f6
--- /dev/null
+++ b/docs/models/qwen35/decode-kernel-attribution.md
@@ -0,0 +1,68 @@
+# 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**.
+
+Follow-up: round the selection width to the 128-token tile multiple at the `bound_selection_vocab` boundary (still inside the 248,320-row checkpoint weight), so every downstream buffer and the sampler stay consistent.
+
+### 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).
+4. Fix the output-projection align-1 kernel (1.67 ms/step at c16): round the selection width to the tile multiple or pin a cublasLt algo.
+
+## 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.

From 10cea8f5f595a5fab66fcbb5786769ba1a4e9b8e Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Sun, 13 Sep 2026 13:31:18 +0800
Subject: [PATCH 02/12] perf(qwen35): align the logits GEMM width so the output
 projection stops landing on an align-1 sm_75 kernel
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 .../qwen35/decode-kernel-attribution.md       |  3 +-
 pegainfer-qwen35/src/config/model.rs          | 38 ++++++++++++++++++-
 pegainfer-qwen35/src/weights.rs               |  6 +--
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/docs/models/qwen35/decode-kernel-attribution.md b/docs/models/qwen35/decode-kernel-attribution.md
index bd0d073f6..2600732cc 100644
--- a/docs/models/qwen35/decode-kernel-attribution.md
+++ b/docs/models/qwen35/decode-kernel-attribution.md
@@ -47,7 +47,7 @@ Neither kernel has had an sm_80 tuning pass (decode tuning history is sm_120/RTX
 
 `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**.
 
-Follow-up: round the selection width to the 128-token tile multiple at the `bound_selection_vocab` boundary (still inside the 248,320-row checkpoint weight), so every downstream buffer and the sampler stay consistent.
+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`); `hf_golden_gate` TP1/TP2 validate output equivalence against HF, which computes logits over the full checkpoint vocab itself.
 
 ### 5. What is NOT the problem
 
@@ -61,7 +61,6 @@ Follow-up: round the selection width to the 128-token tile multiple at the `boun
 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).
-4. Fix the output-projection align-1 kernel (1.67 ms/step at c16): round the selection width to the tile multiple or pin a cublasLt algo.
 
 ## Claim boundary
 
diff --git a/pegainfer-qwen35/src/config/model.rs b/pegainfer-qwen35/src/config/model.rs
index 566c25459..06a74ceb1 100644
--- a/pegainfer-qwen35/src/config/model.rs
+++ b/pegainfer-qwen35/src/config/model.rs
@@ -162,7 +162,14 @@ impl Config35 {
                 vocab_size: self.vocab_size,
             });
         }
-        self.selection_vocab = effective_vocab;
+        // 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.selection_vocab = aligned.min(self.vocab_size);
         Ok(())
     }
 }
@@ -324,6 +331,35 @@ 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);
+
+        config
+            .bound_selection_vocab(769)
+            .expect("769 decodes within vocab 1000");
+        // 769 -> 896: aligned, below the checkpoint rows, pad rows ride the
+        // mapped weight.
+        assert_eq!(config.selection_vocab, 896);
+    }
+
+    #[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,", "");
diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs
index 30527c48e..7d9d2772a 100644
--- a/pegainfer-qwen35/src/weights.rs
+++ b/pegainfer-qwen35/src/weights.rs
@@ -167,10 +167,10 @@ impl Qwen35Model {
         config
             .bound_selection_vocab(effective_vocab)
             .map_err(anyhow::Error::from)?;
-        if config.selection_vocab < config.vocab_size {
+        if config.selection_vocab != effective_vocab {
             info!(
-                "output projection: selection bounded to decodable vocab {} (checkpoint pads to {})",
-                config.selection_vocab, config.vocab_size
+                "output projection: selection width {} = decodable vocab {} + tile-alignment pad (checkpoint has {})",
+                config.selection_vocab, effective_vocab, config.vocab_size
             );
         }
 

From 66493f3fa29d349d3564d6c2f7de89e680766dbf Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Sun, 13 Sep 2026 14:06:08 +0800
Subject: [PATCH 03/12] fix(qwen35): suppress the selection-width pad rows from
 token selection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review follow-up on #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>
---
 pegainfer-qwen35/src/batch_decode.rs |  2 ++
 pegainfer-qwen35/src/config/model.rs | 10 ++++++++-
 pegainfer-qwen35/src/ops.rs          |  2 ++
 pegainfer-qwen35/src/prefill.rs      |  1 +
 pegainfer-qwen35/src/weights.rs      | 31 ++++++++++++++++++++++++++++
 5 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs
index b41da41d4..9ac625a78 100644
--- a/pegainfer-qwen35/src/batch_decode.rs
+++ b/pegainfer-qwen35/src/batch_decode.rs
@@ -702,6 +702,7 @@ impl Qwen35Model {
             &bufs.normed,
             &mut bufs.logits,
         )?;
+        self.suppress_pad_logits(&mut bufs.logits)?;
         debug_assert_eq!(bufs.logits.seq_len, padded_bs);
 
         Ok(())
@@ -780,6 +781,7 @@ impl Qwen35Model {
             &bufs.normed,
             &mut bufs.logits,
         )?;
+        self.suppress_pad_logits(&mut bufs.logits)?;
         debug_assert_eq!(bufs.logits.seq_len, bs);
         Ok(())
     }
diff --git a/pegainfer-qwen35/src/config/model.rs b/pegainfer-qwen35/src/config/model.rs
index 06a74ceb1..43f451d66 100644
--- a/pegainfer-qwen35/src/config/model.rs
+++ b/pegainfer-qwen35/src/config/model.rs
@@ -105,6 +105,10 @@ pub(crate) struct Config35 {
 
     /// Token-selection width: `vocab_size` bounded to the frontend-decodable vocab.
     pub(crate) selection_vocab: usize,
+    /// Tokenizer-decodable width. The selection width may be tile-aligned
+    /// past this (see [`Config35::bound_selection_vocab`]); logits rows beyond
+    /// it are suppressed to -inf before selection.
+    pub(crate) decodable_vocab: usize,
 }
 
 impl Config35 {
@@ -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);
+        self.decodable_vocab = effective_vocab;
         self.selection_vocab = aligned.min(self.vocab_size);
         Ok(())
     }
@@ -285,6 +290,7 @@ impl TryFrom<RawConfig> for Config35 {
             layer_types,
             tie_word_embeddings,
             selection_vocab: t.vocab_size,
+            decodable_vocab: t.vocab_size,
         })
     }
 }
@@ -339,13 +345,15 @@ mod tests {
             .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, pad rows ride the
-        // mapped weight.
+        // mapped weight and are suppressed to -inf before selection.
         assert_eq!(config.selection_vocab, 896);
+        assert_eq!(config.decodable_vocab, 769);
     }
 
     #[test]
diff --git a/pegainfer-qwen35/src/ops.rs b/pegainfer-qwen35/src/ops.rs
index 3d288b074..c50423c34 100644
--- a/pegainfer-qwen35/src/ops.rs
+++ b/pegainfer-qwen35/src/ops.rs
@@ -2,6 +2,7 @@
 
 pub(crate) use pegainfer_core::ops::GEMM_LT_MAX_N;
 pub(crate) use pegainfer_core::ops::PrefillPagedPlan;
+pub(crate) use pegainfer_core::ops::SuppressIds;
 pub(crate) use pegainfer_core::ops::add_batch;
 pub(crate) use pegainfer_core::ops::add_batch_into;
 pub(crate) use pegainfer_core::ops::embedding_batch;
@@ -17,6 +18,7 @@ pub use pegainfer_core::ops::rms_norm_batch_offset_into;
 pub(crate) use pegainfer_core::ops::rms_norm_gated_batch_into;
 pub use pegainfer_core::ops::rms_norm_offset_into;
 pub(crate) use pegainfer_core::ops::silu_mul_fused_batch_into;
+pub(crate) use pegainfer_core::ops::suppress_logits_bf16_in_place;
 pub(crate) use pegainfer_core::ops::write_vec_into;
 pub(crate) use recurrent::conv1d_decode_batch_into;
 pub(crate) use recurrent::conv1d_prefill_batch_into;
diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs
index 6fa2555d7..f44f4e8f8 100644
--- a/pegainfer-qwen35/src/prefill.rs
+++ b/pegainfer-qwen35/src/prefill.rs
@@ -125,6 +125,7 @@ impl Qwen35Model {
             &normed,
             &mut logits,
         )?;
+        self.suppress_pad_logits(&mut logits)?;
         debug_assert_eq!(logits.seq_len, n);
         Ok(logits)
     }
diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs
index 7d9d2772a..9df8a1d6b 100644
--- a/pegainfer-qwen35/src/weights.rs
+++ b/pegainfer-qwen35/src/weights.rs
@@ -69,6 +69,9 @@ pub struct Qwen35Model {
     /// (e.g. `--max-batch 5` allocates bucket 8 but admits at most 5). See #470.
     pub(super) decode_admission_batch: usize,
     tp_comm: Option<Comm>,
+    /// -inf suppression for the tile-alignment pad rows of the logits GEMM
+    /// (absent when the selection width needs no padding).
+    pad_logit_suppress: Option<crate::ops::SuppressIds>,
 }
 
 // SAFETY: A Qwen3.5 model instance is bound to one CUDA device and driven from
@@ -308,6 +311,21 @@ impl Qwen35Model {
             page_size,
             num_pages,
         )?;
+        // Rows the selection width adds past the decodable vocab are real
+        // checkpoint embeddings but not decodable tokens; force their logits
+        // to -inf so they can never win selection.
+        let pad_logit_suppress = if config.selection_vocab > config.decodable_vocab {
+            let ids: Vec<u32> = (config.decodable_vocab..config.selection_vocab)
+                .map(|id| id as u32)
+                .collect();
+            Some(crate::ops::SuppressIds::upload(
+                &ctx,
+                &ids,
+                config.selection_vocab,
+            )?)
+        } else {
+            None
+        };
 
         Ok(Self {
             ctx,
@@ -323,6 +341,7 @@ impl Qwen35Model {
             reserved_decode_slots: max_batch,
             decode_admission_batch,
             tp_comm: None,
+            pad_logit_suppress,
         })
     }
 
@@ -334,6 +353,18 @@ impl Qwen35Model {
         self.lm_head.as_ref().unwrap_or(&self.embed_tokens)
     }
 
+    /// Suppress the tile-alignment pad rows of a logits buffer: rows past the
+    /// decodable vocab are forced to -inf so they can never win greedy or
+    /// sampled selection (a no-op when the width needs no padding).
+    pub(crate) fn suppress_pad_logits(&self, logits: &mut HiddenStates) -> Result<()> {
+        match &self.pad_logit_suppress {
+            Some(suppress) => {
+                crate::ops::suppress_logits_bf16_in_place(&self.ctx, logits, suppress)
+            }
+            None => Ok(()),
+        }
+    }
+
     pub(crate) fn ensure_rope_cache_covers(&self, positions: usize) -> Result<()> {
         let cache_positions = self.cos_cache.len / self.config.rotary_dim;
         anyhow::ensure!(

From 95c733342606994664946592fb160a432b41fd6b Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Tue, 15 Sep 2026 13:37:57 +0800
Subject: [PATCH 04/12] fix(qwen35): measure argmax routing against the
 decodable vocab
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review follow-up on #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>
---
 pegainfer-qwen35/src/config/model.rs    | 28 +++++++++--
 pegainfer-qwen35/src/decode_buffers.rs  | 10 +++-
 pegainfer-qwen35/src/tp_executor.rs     |  3 +-
 pegainfer-qwen35/src/unified_forward.rs | 12 +++--
 pegainfer-sample/src/lib.rs             | 46 +++++++++++++++++-
 pegainfer-sample/tests/select_batch.rs  | 63 +++++++++++++++++++++++++
 6 files changed, 151 insertions(+), 11 deletions(-)

diff --git a/pegainfer-qwen35/src/config/model.rs b/pegainfer-qwen35/src/config/model.rs
index 43f451d66..7a54b17ee 100644
--- a/pegainfer-qwen35/src/config/model.rs
+++ b/pegainfer-qwen35/src/config/model.rs
@@ -103,11 +103,14 @@ 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.
+    /// 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 selection width may be tile-aligned
-    /// past this (see [`Config35::bound_selection_vocab`]); logits rows beyond
-    /// it are suppressed to -inf before selection.
+    /// Tokenizer-decodable width. The selection width may be tile-aligned past
+    /// this (see [`Config35::bound_selection_vocab`]); logits rows beyond it are
+    /// suppressed to -inf before selection, and the argmax-vs-sample routing
+    /// decision is measured against this width rather than the aligned one.
     pub(crate) decodable_vocab: usize,
 }
 
@@ -150,12 +153,27 @@ impl Config35 {
         self.linear_num_value_heads * self.linear_value_head_dim
     }
 
-    /// Bound the output-selection width to the frontend-decodable vocab.
+    /// Bound the output-selection width to the frontend-decodable vocab and
+    /// record the two widths that follow from it.
     ///
     /// The frontend decodes a dense prefix of the vocab; the checkpoint may pad
     /// beyond it. Refusing a tokenizer wider than the checkpoint is the
     /// fail-closed rule, and it is checked here at the validation boundary
     /// rather than scattered through the loader.
+    ///
+    /// The two widths the rest of the crate keys off:
+    ///
+    /// * [`Config35::selection_vocab`] — the alignable one: logits buffers, the
+    ///   sampler arena and the output-projection GEMM all span it, and the pad
+    ///   rows past the decodable vocab are suppressed to -inf before selection.
+    /// * [`Config35::decodable_vocab`] — the semantic one: the tokens that can
+    ///   actually be emitted, hence the width the argmax-vs-sample routing
+    ///   decision is measured against.
+    ///
+    /// Keeping them separate is what lets the GEMM be widened for throughput
+    /// without moving `top_p <= 1/vocab` routing: pad columns are not tokens, so
+    /// they must not decide whether an effectively-greedy request takes the
+    /// deterministic argmax path or the rejection sampler.
     pub(crate) fn bound_selection_vocab(
         &mut self,
         effective_vocab: usize,
diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs
index cea677445..260aabb9f 100644
--- a/pegainfer-qwen35/src/decode_buffers.rs
+++ b/pegainfer-qwen35/src/decode_buffers.rs
@@ -116,7 +116,15 @@ impl BatchDecodeBuffers35 {
             kv_tile_indices_d: ctx.stream.alloc_zeros(bs)?,
             kv_chunk_size_d: ctx.stream.alloc_zeros(bs)?,
 
-            sample: pegainfer_sample::SampleScratch::new(ctx, config.selection_vocab, bs)?,
+            // The arena spans the tile-aligned selection width; the
+            // argmax-vs-sample routing keys off the decodable vocab instead, so
+            // pad columns cannot move the `top_p <= 1/vocab` routing decision.
+            sample: pegainfer_sample::SampleScratch::with_selection_width(
+                ctx,
+                config.selection_vocab,
+                config.decodable_vocab,
+                bs,
+            )?,
             steps: Vec::new(),
 
             padding_page_id,
diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs
index e02ace2c3..2320dfebd 100644
--- a/pegainfer-qwen35/src/tp_executor.rs
+++ b/pegainfer-qwen35/src/tp_executor.rs
@@ -1395,9 +1395,10 @@ impl TpWorkerPrepared {
             prefill_scratch_bytes as f64 / 1024.0 / 1024.0,
         );
         let decode_buffers = model.create_batch_decode_buffers_with_capacity(max_batch)?;
-        let sample_scratch = pegainfer_sample::SampleScratch::new(
+        let sample_scratch = pegainfer_sample::SampleScratch::with_selection_width(
             model.device_ctx(),
             model.config().selection_vocab,
+            model.config().decodable_vocab,
             max_batch,
         )?;
         Ok((
diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs
index e4dfc4560..659b4aba4 100644
--- a/pegainfer-qwen35/src/unified_forward.rs
+++ b/pegainfer-qwen35/src/unified_forward.rs
@@ -153,9 +153,15 @@ mod tests {
         let params = vec![pegainfer_frontend::sampler::SamplingParams::default(); rows];
         let params_refs: Vec<&pegainfer_frontend::sampler::SamplingParams> =
             params.iter().collect();
-        let mut scratch =
-            pegainfer_sample::SampleScratch::new(&model.ctx, model.config.selection_vocab, rows)
-                .unwrap();
+        // Mirrors the production scratch: the arena spans the tile-aligned
+        // selection width, the routing decision keys off the decodable vocab.
+        let mut scratch = pegainfer_sample::SampleScratch::with_selection_width(
+            &model.ctx,
+            model.config.selection_vocab,
+            model.config.decodable_vocab,
+            rows,
+        )
+        .unwrap();
         let steps = vec![0u64; params_refs.len()];
         pegainfer_sample::select_batch(&model.ctx, logits, &params_refs, &steps, 0, &mut scratch)
             .unwrap()
diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs
index b96e8eee2..4ae6a66de 100644
--- a/pegainfer-sample/src/lib.rs
+++ b/pegainfer-sample/src/lib.rs
@@ -83,15 +83,46 @@ pub struct SampleScratch {
     /// Vocab width every buffer above was sized for; `select_batch` rejects a
     /// logits arena whose `hidden_dim` differs, since the sizes are baked in.
     vocab: usize,
+    /// Width the argmax-vs-sample routing decision is measured against. Equal
+    /// to `vocab` unless a model tile-aligns its logits width past the tokens
+    /// it can actually emit (see [`SampleScratch::with_selection_width`]).
+    selection_width: usize,
     max_rows: usize,
 }
 
 impl SampleScratch {
     pub fn new(ctx: &DeviceContext, vocab: usize, max_rows: usize) -> Result<Self> {
+        Self::with_selection_width(ctx, vocab, vocab, max_rows)
+    }
+
+    /// Build scratch for a logits arena that spans `vocab` columns while the
+    /// model can only emit the first `selection_width` ids.
+    ///
+    /// A model may widen its logits past the decodable vocab to reach a GEMM
+    /// tile multiple (qwen35's output projection: an odd decodable vocab drops
+    /// cublasLt onto an align-1 kernel), suppressing the extra rows to `-inf`
+    /// before selection. The arena still spans the widened `vocab`, but those
+    /// pad columns are not tokens, so they must not widen the
+    /// `top_p <= 1/vocab` nucleus [`effectively_greedy`] keys off: a `top_p` at
+    /// or below `1/selection_width` stays an effectively-greedy request and has
+    /// to keep taking the deterministic argmax path, rather than falling to the
+    /// rejection sampler over bf16-tied maxima because the arena was widened.
+    /// Equivalently, `selection_width` is the width at which a padded arena
+    /// routes exactly the rows an unpadded arena would.
+    pub fn with_selection_width(
+        ctx: &DeviceContext,
+        vocab: usize,
+        selection_width: usize,
+        max_rows: usize,
+    ) -> Result<Self> {
         ensure!(
             vocab > 0 && max_rows > 0,
             "SampleScratch requires vocab > 0 and max_rows > 0"
         );
+        ensure!(
+            selection_width > 0 && selection_width <= vocab,
+            "SampleScratch selection width {selection_width} must be in 1..={vocab}"
+        );
         let partials = argmax_batch_bf16_split_partials_len(max_rows, vocab);
         let alloc_i32 = |n: usize| -> Result<CudaSlice<i32>> {
             ctx.stream
@@ -124,6 +155,7 @@ impl SampleScratch {
                 .map_err(|e| anyhow!("SampleScratch identity upload failed: {e}"))?,
             sampling: BatchSamplingScratch::new(ctx, max_rows, vocab)?,
             vocab,
+            selection_width,
             max_rows,
         })
     }
@@ -135,6 +167,12 @@ impl SampleScratch {
     pub fn vocab(&self) -> usize {
         self.vocab
     }
+
+    /// Width the argmax-vs-sample routing decision is measured against — the
+    /// semantic vocab, not the (possibly tile-aligned) arena width.
+    pub fn selection_width(&self) -> usize {
+        self.selection_width
+    }
 }
 
 /// Pick the next token for every row of a logits arena.
@@ -151,6 +189,9 @@ impl SampleScratch {
 /// argmax survives. Routing those through argmax keeps an effectively-greedy
 /// request deterministic — the rejection sampler would otherwise pick an
 /// arbitrary member of a bf16-tied top — and skips a softmax it does not need.
+/// `vocab` here is `scratch`'s semantic selection width, which is narrower than
+/// the arena when a model aligned its logits GEMM (see
+/// [`SampleScratch::with_selection_width`]).
 ///
 /// `seed` must be fresh per decode step (one engine seed at startup, advanced
 /// per step); unseeded rows decorrelate through the philox subsequence.
@@ -196,7 +237,10 @@ pub fn select_batch(
         "select_batch: logits vocab {vocab} != scratch vocab {}",
         scratch.vocab
     );
-    let is_argmax = |p: &&SamplingParams| effectively_greedy(p, vocab);
+    // Route on the semantic width, not the arena width: `scratch`'s buffers
+    // span `vocab`, but pad columns a model aligned its GEMM to are not
+    // emittable tokens, so they must not move the `top_p <= 1/vocab` nucleus.
+    let is_argmax = |p: &&SamplingParams| effectively_greedy(p, scratch.selection_width);
     let mut tokens = vec![0u32; n];
 
     // Argmax rows -> one batched indexed argmax.
diff --git a/pegainfer-sample/tests/select_batch.rs b/pegainfer-sample/tests/select_batch.rs
index 802aa1a86..604f05b03 100644
--- a/pegainfer-sample/tests/select_batch.rs
+++ b/pegainfer-sample/tests/select_batch.rs
@@ -177,6 +177,69 @@ fn tiny_top_p_routes_to_argmax_even_under_bf16_ties() {
     }
 }
 
+#[test]
+fn padded_arena_width_does_not_suppress_the_argmax_routing() {
+    // Regression for the qwen35 tile-aligned logits width (#1046 review). A
+    // model may widen its logits past the decodable vocab to reach a GEMM tile
+    // multiple and suppress the pad columns to -inf. The routing decision must
+    // still be measured against the decodable vocab: here the arena spans 512
+    // columns while only the first 256 can be emitted, and top_p = 1/256 is
+    // exactly the effectively-greedy boundary. Keyed off the arena width
+    // (1/512) it would fall to the rejection sampler and pick the bf16-tied
+    // peer; keyed off the decodable vocab it keeps the deterministic argmax.
+    let ctx = DeviceContext::new().unwrap();
+    let decodable = 256usize;
+    // The arena spans twice the decodable width, so 1/decodable and 1/vocab
+    // land in different routing buckets (top_p = 1/256 is above 1/512).
+    let vocab = 2 * decodable;
+    let lo = 128usize;
+    let hi = 200usize;
+    let mut row = vec![-1.0f32; vocab];
+    row[lo] = 8.0; // bf16-exact, identical to `hi` -> a true top tie
+    row[hi] = 8.0;
+    let arena = make_arena(&ctx, &[row]);
+
+    let tiny = sampling(1.0, -1, 1.0 / decodable as f32);
+
+    // Without the semantic bound the row is not effectively greedy: the
+    // sampler may return either tied peer, which is what the bound prevents.
+    let mut wide = SampleScratch::new(&ctx, vocab, 1).unwrap();
+    let mut sampled_the_peer = false;
+    for s in 0..64u64 {
+        let picked = select_batch(&ctx, &arena, &[&tiny], &[0], s, &mut wide).unwrap()[0];
+        assert!(picked == lo as u32 || picked == hi as u32);
+        sampled_the_peer |= picked == hi as u32;
+    }
+    assert!(
+        sampled_the_peer,
+        "control: an arena-width routing decision puts this row on the sampler"
+    );
+
+    let mut bounded = SampleScratch::with_selection_width(&ctx, vocab, decodable, 1).unwrap();
+    assert_eq!(bounded.vocab(), vocab);
+    assert_eq!(bounded.selection_width(), decodable);
+    for s in 0..64u64 {
+        assert_eq!(
+            select_batch(&ctx, &arena, &[&tiny], &[0], s, &mut bounded).unwrap(),
+            vec![lo as u32],
+            "seed {s}: the decodable vocab must keep this row on the argmax path"
+        );
+    }
+}
+
+#[test]
+fn selection_width_above_the_arena_is_rejected() {
+    let ctx = DeviceContext::new().unwrap();
+    assert!(
+        SampleScratch::with_selection_width(&ctx, 256, 257, 1).is_err(),
+        "a selection width wider than the arena must be refused"
+    );
+    assert!(
+        SampleScratch::with_selection_width(&ctx, 256, 0, 1).is_err(),
+        "a zero selection width must be refused"
+    );
+}
+
 #[test]
 fn batch_larger_than_scratch_is_rejected() {
     let ctx = DeviceContext::new().unwrap();

From f9aa4f8f156ead498bf41dcbdf4868b3ef42d81a Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Tue, 15 Sep 2026 14:59:46 +0800
Subject: [PATCH 05/12] style(qwen35): trim the selection-width comments and
 test-only surface

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>
---
 pegainfer-qwen35/src/config/model.rs    | 28 ++++-------------
 pegainfer-qwen35/src/decode_buffers.rs  |  3 --
 pegainfer-qwen35/src/unified_forward.rs |  2 --
 pegainfer-qwen35/src/weights.rs         |  9 ++----
 pegainfer-sample/src/lib.rs             | 40 +++++++------------------
 pegainfer-sample/tests/select_batch.rs  | 32 ++++----------------
 6 files changed, 25 insertions(+), 89 deletions(-)

diff --git a/pegainfer-qwen35/src/config/model.rs b/pegainfer-qwen35/src/config/model.rs
index 7a54b17ee..31683eba7 100644
--- a/pegainfer-qwen35/src/config/model.rs
+++ b/pegainfer-qwen35/src/config/model.rs
@@ -107,10 +107,8 @@ pub(crate) struct Config35 {
     /// 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 selection width may be tile-aligned past
-    /// this (see [`Config35::bound_selection_vocab`]); logits rows beyond it are
-    /// suppressed to -inf before selection, and the argmax-vs-sample routing
-    /// decision is measured against this width rather than the aligned one.
+    /// 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,
 }
 
@@ -153,27 +151,13 @@ impl Config35 {
         self.linear_num_value_heads * self.linear_value_head_dim
     }
 
-    /// Bound the output-selection width to the frontend-decodable vocab and
-    /// record the two widths that follow from it.
+    /// 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
     /// fail-closed rule, and it is checked here at the validation boundary
     /// rather than scattered through the loader.
-    ///
-    /// The two widths the rest of the crate keys off:
-    ///
-    /// * [`Config35::selection_vocab`] — the alignable one: logits buffers, the
-    ///   sampler arena and the output-projection GEMM all span it, and the pad
-    ///   rows past the decodable vocab are suppressed to -inf before selection.
-    /// * [`Config35::decodable_vocab`] — the semantic one: the tokens that can
-    ///   actually be emitted, hence the width the argmax-vs-sample routing
-    ///   decision is measured against.
-    ///
-    /// Keeping them separate is what lets the GEMM be widened for throughput
-    /// without moving `top_p <= 1/vocab` routing: pad columns are not tokens, so
-    /// they must not decide whether an effectively-greedy request takes the
-    /// deterministic argmax path or the rejection sampler.
     pub(crate) fn bound_selection_vocab(
         &mut self,
         effective_vocab: usize,
@@ -368,8 +352,8 @@ mod tests {
         config
             .bound_selection_vocab(769)
             .expect("769 decodes within vocab 1000");
-        // 769 -> 896: aligned, below the checkpoint rows, pad rows ride the
-        // mapped weight and are suppressed to -inf before selection.
+        // 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);
     }
diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs
index 260aabb9f..5e56273ad 100644
--- a/pegainfer-qwen35/src/decode_buffers.rs
+++ b/pegainfer-qwen35/src/decode_buffers.rs
@@ -116,9 +116,6 @@ impl BatchDecodeBuffers35 {
             kv_tile_indices_d: ctx.stream.alloc_zeros(bs)?,
             kv_chunk_size_d: ctx.stream.alloc_zeros(bs)?,
 
-            // The arena spans the tile-aligned selection width; the
-            // argmax-vs-sample routing keys off the decodable vocab instead, so
-            // pad columns cannot move the `top_p <= 1/vocab` routing decision.
             sample: pegainfer_sample::SampleScratch::with_selection_width(
                 ctx,
                 config.selection_vocab,
diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs
index 659b4aba4..4c313a63d 100644
--- a/pegainfer-qwen35/src/unified_forward.rs
+++ b/pegainfer-qwen35/src/unified_forward.rs
@@ -153,8 +153,6 @@ mod tests {
         let params = vec![pegainfer_frontend::sampler::SamplingParams::default(); rows];
         let params_refs: Vec<&pegainfer_frontend::sampler::SamplingParams> =
             params.iter().collect();
-        // Mirrors the production scratch: the arena spans the tile-aligned
-        // selection width, the routing decision keys off the decodable vocab.
         let mut scratch = pegainfer_sample::SampleScratch::with_selection_width(
             &model.ctx,
             model.config.selection_vocab,
diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs
index 9df8a1d6b..71b2d14a1 100644
--- a/pegainfer-qwen35/src/weights.rs
+++ b/pegainfer-qwen35/src/weights.rs
@@ -311,9 +311,8 @@ impl Qwen35Model {
             page_size,
             num_pages,
         )?;
-        // Rows the selection width adds past the decodable vocab are real
-        // checkpoint embeddings but not decodable tokens; force their logits
-        // to -inf so they can never win selection.
+        // The alignment pad rows are real checkpoint embeddings but not
+        // decodable tokens, so they must never win selection.
         let pad_logit_suppress = if config.selection_vocab > config.decodable_vocab {
             let ids: Vec<u32> = (config.decodable_vocab..config.selection_vocab)
                 .map(|id| id as u32)
@@ -353,9 +352,7 @@ impl Qwen35Model {
         self.lm_head.as_ref().unwrap_or(&self.embed_tokens)
     }
 
-    /// Suppress the tile-alignment pad rows of a logits buffer: rows past the
-    /// decodable vocab are forced to -inf so they can never win greedy or
-    /// sampled selection (a no-op when the width needs no padding).
+    /// Force the tile-alignment pad rows of a logits buffer to -inf.
     pub(crate) fn suppress_pad_logits(&self, logits: &mut HiddenStates) -> Result<()> {
         match &self.pad_logit_suppress {
             Some(suppress) => {
diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs
index 4ae6a66de..a5206c72b 100644
--- a/pegainfer-sample/src/lib.rs
+++ b/pegainfer-sample/src/lib.rs
@@ -83,9 +83,8 @@ pub struct SampleScratch {
     /// Vocab width every buffer above was sized for; `select_batch` rejects a
     /// logits arena whose `hidden_dim` differs, since the sizes are baked in.
     vocab: usize,
-    /// Width the argmax-vs-sample routing decision is measured against. Equal
-    /// to `vocab` unless a model tile-aligns its logits width past the tokens
-    /// it can actually emit (see [`SampleScratch::with_selection_width`]).
+    /// Width the argmax-vs-sample routing decision is measured against: the
+    /// emittable tokens, not the arena a model may have tile-aligned wider.
     selection_width: usize,
     max_rows: usize,
 }
@@ -95,20 +94,11 @@ impl SampleScratch {
         Self::with_selection_width(ctx, vocab, vocab, max_rows)
     }
 
-    /// Build scratch for a logits arena that spans `vocab` columns while the
-    /// model can only emit the first `selection_width` ids.
-    ///
-    /// A model may widen its logits past the decodable vocab to reach a GEMM
-    /// tile multiple (qwen35's output projection: an odd decodable vocab drops
-    /// cublasLt onto an align-1 kernel), suppressing the extra rows to `-inf`
-    /// before selection. The arena still spans the widened `vocab`, but those
-    /// pad columns are not tokens, so they must not widen the
-    /// `top_p <= 1/vocab` nucleus [`effectively_greedy`] keys off: a `top_p` at
-    /// or below `1/selection_width` stays an effectively-greedy request and has
-    /// to keep taking the deterministic argmax path, rather than falling to the
-    /// rejection sampler over bf16-tied maxima because the arena was widened.
-    /// Equivalently, `selection_width` is the width at which a padded arena
-    /// routes exactly the rows an unpadded arena would.
+    /// For an arena spanning `vocab` columns whose last `vocab - selection_width`
+    /// columns the model cannot emit (it widened them to reach a GEMM tile
+    /// multiple). Those pad columns must not widen the `top_p <= 1/vocab`
+    /// nucleus [`effectively_greedy`] keys off, or a request that is effectively
+    /// greedy would drop to the rejection sampler over bf16-tied maxima.
     pub fn with_selection_width(
         ctx: &DeviceContext,
         vocab: usize,
@@ -167,12 +157,6 @@ impl SampleScratch {
     pub fn vocab(&self) -> usize {
         self.vocab
     }
-
-    /// Width the argmax-vs-sample routing decision is measured against — the
-    /// semantic vocab, not the (possibly tile-aligned) arena width.
-    pub fn selection_width(&self) -> usize {
-        self.selection_width
-    }
 }
 
 /// Pick the next token for every row of a logits arena.
@@ -189,9 +173,8 @@ impl SampleScratch {
 /// argmax survives. Routing those through argmax keeps an effectively-greedy
 /// request deterministic — the rejection sampler would otherwise pick an
 /// arbitrary member of a bf16-tied top — and skips a softmax it does not need.
-/// `vocab` here is `scratch`'s semantic selection width, which is narrower than
-/// the arena when a model aligned its logits GEMM (see
-/// [`SampleScratch::with_selection_width`]).
+/// `vocab` here is `scratch`'s selection width, narrower than the arena when a
+/// model aligned its logits GEMM.
 ///
 /// `seed` must be fresh per decode step (one engine seed at startup, advanced
 /// per step); unseeded rows decorrelate through the philox subsequence.
@@ -237,9 +220,8 @@ pub fn select_batch(
         "select_batch: logits vocab {vocab} != scratch vocab {}",
         scratch.vocab
     );
-    // Route on the semantic width, not the arena width: `scratch`'s buffers
-    // span `vocab`, but pad columns a model aligned its GEMM to are not
-    // emittable tokens, so they must not move the `top_p <= 1/vocab` nucleus.
+    // Pad columns a model aligned its GEMM to are not emittable tokens, so they
+    // must not move the `top_p <= 1/vocab` nucleus.
     let is_argmax = |p: &&SamplingParams| effectively_greedy(p, scratch.selection_width);
     let mut tokens = vec![0u32; n];
 
diff --git a/pegainfer-sample/tests/select_batch.rs b/pegainfer-sample/tests/select_batch.rs
index 604f05b03..065c082c1 100644
--- a/pegainfer-sample/tests/select_batch.rs
+++ b/pegainfer-sample/tests/select_batch.rs
@@ -179,18 +179,11 @@ fn tiny_top_p_routes_to_argmax_even_under_bf16_ties() {
 
 #[test]
 fn padded_arena_width_does_not_suppress_the_argmax_routing() {
-    // Regression for the qwen35 tile-aligned logits width (#1046 review). A
-    // model may widen its logits past the decodable vocab to reach a GEMM tile
-    // multiple and suppress the pad columns to -inf. The routing decision must
-    // still be measured against the decodable vocab: here the arena spans 512
-    // columns while only the first 256 can be emitted, and top_p = 1/256 is
-    // exactly the effectively-greedy boundary. Keyed off the arena width
-    // (1/512) it would fall to the rejection sampler and pick the bf16-tied
-    // peer; keyed off the decodable vocab it keeps the deterministic argmax.
+    // qwen35 tile-aligns its logits width past the decodable vocab (#1046). A
+    // top_p at exactly 1/decodable collapses to the argmax on the unpadded
+    // arena, so it must still do so at twice the arena width.
     let ctx = DeviceContext::new().unwrap();
     let decodable = 256usize;
-    // The arena spans twice the decodable width, so 1/decodable and 1/vocab
-    // land in different routing buckets (top_p = 1/256 is above 1/512).
     let vocab = 2 * decodable;
     let lo = 128usize;
     let hi = 200usize;
@@ -201,8 +194,8 @@ fn padded_arena_width_does_not_suppress_the_argmax_routing() {
 
     let tiny = sampling(1.0, -1, 1.0 / decodable as f32);
 
-    // Without the semantic bound the row is not effectively greedy: the
-    // sampler may return either tied peer, which is what the bound prevents.
+    // Control: at the arena width 1/256 is no longer a single-token nucleus,
+    // and the sampler may return the tied peer.
     let mut wide = SampleScratch::new(&ctx, vocab, 1).unwrap();
     let mut sampled_the_peer = false;
     for s in 0..64u64 {
@@ -216,8 +209,6 @@ fn padded_arena_width_does_not_suppress_the_argmax_routing() {
     );
 
     let mut bounded = SampleScratch::with_selection_width(&ctx, vocab, decodable, 1).unwrap();
-    assert_eq!(bounded.vocab(), vocab);
-    assert_eq!(bounded.selection_width(), decodable);
     for s in 0..64u64 {
         assert_eq!(
             select_batch(&ctx, &arena, &[&tiny], &[0], s, &mut bounded).unwrap(),
@@ -227,19 +218,6 @@ fn padded_arena_width_does_not_suppress_the_argmax_routing() {
     }
 }
 
-#[test]
-fn selection_width_above_the_arena_is_rejected() {
-    let ctx = DeviceContext::new().unwrap();
-    assert!(
-        SampleScratch::with_selection_width(&ctx, 256, 257, 1).is_err(),
-        "a selection width wider than the arena must be refused"
-    );
-    assert!(
-        SampleScratch::with_selection_width(&ctx, 256, 0, 1).is_err(),
-        "a zero selection width must be refused"
-    );
-}
-
 #[test]
 fn batch_larger_than_scratch_is_rejected() {
     let ctx = DeviceContext::new().unwrap();

From 80e3c9584953661647bba950107e67cd28ac0aee Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Tue, 15 Sep 2026 15:05:14 +0800
Subject: [PATCH 06/12] docs(qwen35): record the two invariants the widened
 logits width depends 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>
---
 docs/models/qwen35/decode-kernel-attribution.md | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/docs/models/qwen35/decode-kernel-attribution.md b/docs/models/qwen35/decode-kernel-attribution.md
index 2600732cc..1887ab105 100644
--- a/docs/models/qwen35/decode-kernel-attribution.md
+++ b/docs/models/qwen35/decode-kernel-attribution.md
@@ -49,6 +49,11 @@ Neither kernel has had an sm_80 tuning pass (decode tuning history is sm_120/RTX
 
 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`); `hf_golden_gate` TP1/TP2 validate output equivalence against HF, which computes logits over the full checkpoint vocab itself.
 
+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; `suppress_pad_logits` forces them to `-inf` after every logits 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 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.
+
 ### 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.

From 05ffa360aac220fbab09572a869cdcd4b4c6c805 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Tue, 15 Sep 2026 19:24:24 +0800
Subject: [PATCH 07/12] refactor(qwen35): make the pad-row mask unskippable at
 a logits site
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review follow-up on #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>
---
 .../qwen35/decode-kernel-attribution.md       |  2 +-
 pegainfer-qwen35/src/batch_decode.rs          | 20 +---------
 pegainfer-qwen35/src/ops.rs                   |  2 -
 pegainfer-qwen35/src/prefill.rs               | 10 +----
 pegainfer-qwen35/src/weights.rs               | 40 ++++++++++++++-----
 5 files changed, 35 insertions(+), 39 deletions(-)

diff --git a/docs/models/qwen35/decode-kernel-attribution.md b/docs/models/qwen35/decode-kernel-attribution.md
index 1887ab105..837e5641b 100644
--- a/docs/models/qwen35/decode-kernel-attribution.md
+++ b/docs/models/qwen35/decode-kernel-attribution.md
@@ -51,7 +51,7 @@ Fixed by aligning the selection width to the 128-token tile multiple (`248077 
 
 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; `suppress_pad_logits` forces them to `-inf` after every logits 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 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.
 
 ### 5. What is NOT the problem
diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs
index 9ac625a78..a558a5054 100644
--- a/pegainfer-qwen35/src/batch_decode.rs
+++ b/pegainfer-qwen35/src/batch_decode.rs
@@ -694,15 +694,7 @@ impl Qwen35Model {
             eps,
             &mut bufs.normed,
         )?;
-        ops::gemm_rows_into_checked(
-            &self.ctx,
-            self.output_projection(),
-            0,
-            self.config.selection_vocab,
-            &bufs.normed,
-            &mut bufs.logits,
-        )?;
-        self.suppress_pad_logits(&mut bufs.logits)?;
+        self.output_logits_into(&bufs.normed, &mut bufs.logits)?;
         debug_assert_eq!(bufs.logits.seq_len, padded_bs);
 
         Ok(())
@@ -773,15 +765,7 @@ impl Qwen35Model {
             eps,
             &mut bufs.normed,
         )?;
-        ops::gemm_rows_into_checked(
-            &self.ctx,
-            self.output_projection(),
-            0,
-            self.config.selection_vocab,
-            &bufs.normed,
-            &mut bufs.logits,
-        )?;
-        self.suppress_pad_logits(&mut bufs.logits)?;
+        self.output_logits_into(&bufs.normed, &mut bufs.logits)?;
         debug_assert_eq!(bufs.logits.seq_len, bs);
         Ok(())
     }
diff --git a/pegainfer-qwen35/src/ops.rs b/pegainfer-qwen35/src/ops.rs
index c50423c34..af08deb6b 100644
--- a/pegainfer-qwen35/src/ops.rs
+++ b/pegainfer-qwen35/src/ops.rs
@@ -10,7 +10,6 @@ pub(crate) use pegainfer_core::ops::extract_vec;
 pub(crate) use pegainfer_core::ops::gemm;
 pub(crate) use pegainfer_core::ops::gemm_into;
 pub(crate) use pegainfer_core::ops::gemm_lt_tune;
-pub(crate) use pegainfer_core::ops::gemm_rows_into_checked;
 pub(crate) use pegainfer_core::ops::paged_attention_batch_decode_hd256_into;
 pub(crate) use pegainfer_core::ops::paged_attention_batch_decode_via_prefill_hd256_into;
 pub(crate) use pegainfer_core::ops::qk_norm_partial_rope_batched_decode_hd256_into;
@@ -18,7 +17,6 @@ pub use pegainfer_core::ops::rms_norm_batch_offset_into;
 pub(crate) use pegainfer_core::ops::rms_norm_gated_batch_into;
 pub use pegainfer_core::ops::rms_norm_offset_into;
 pub(crate) use pegainfer_core::ops::silu_mul_fused_batch_into;
-pub(crate) use pegainfer_core::ops::suppress_logits_bf16_in_place;
 pub(crate) use pegainfer_core::ops::write_vec_into;
 pub(crate) use recurrent::conv1d_decode_batch_into;
 pub(crate) use recurrent::conv1d_prefill_batch_into;
diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs
index f44f4e8f8..1537e458c 100644
--- a/pegainfer-qwen35/src/prefill.rs
+++ b/pegainfer-qwen35/src/prefill.rs
@@ -117,15 +117,7 @@ impl Qwen35Model {
             &mut normed,
         )?;
         let mut logits = HiddenStates::zeros(&self.ctx, self.config.selection_vocab, n)?;
-        ops::gemm_rows_into_checked(
-            &self.ctx,
-            self.output_projection(),
-            0,
-            self.config.selection_vocab,
-            &normed,
-            &mut logits,
-        )?;
-        self.suppress_pad_logits(&mut logits)?;
+        self.output_logits_into(&normed, &mut logits)?;
         debug_assert_eq!(logits.seq_len, n);
         Ok(logits)
     }
diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs
index 71b2d14a1..d500bec96 100644
--- a/pegainfer-qwen35/src/weights.rs
+++ b/pegainfer-qwen35/src/weights.rs
@@ -7,6 +7,8 @@ use cudarc::nccl::safe::Comm;
 use cudarc::nccl::safe::ReduceOp;
 use log::debug;
 use log::info;
+use pegainfer_core::ops::gemm_rows_into_checked;
+use pegainfer_core::ops::suppress_logits_bf16_in_place;
 use pegainfer_core::rope::RopeTableSpec;
 use pegainfer_core::rope::precompute_rope;
 use pegainfer_core::tensor::DeviceContext;
@@ -69,8 +71,9 @@ pub struct Qwen35Model {
     /// (e.g. `--max-batch 5` allocates bucket 8 but admits at most 5). See #470.
     pub(super) decode_admission_batch: usize,
     tp_comm: Option<Comm>,
-    /// -inf suppression for the tile-alignment pad rows of the logits GEMM
-    /// (absent when the selection width needs no padding).
+    /// -inf suppression for the tile-alignment pad rows, applied by
+    /// [`Qwen35Model::output_logits_into`] (absent when the selection width
+    /// needs no padding).
     pad_logit_suppress: Option<crate::ops::SuppressIds>,
 }
 
@@ -348,18 +351,37 @@ impl Qwen35Model {
         &self.config
     }
 
+    /// Only the GEMM tuning helper samples this directly; logits go through
+    /// [`Qwen35Model::output_logits_into`] so the pad-row mask cannot be skipped.
     pub(super) fn output_projection(&self) -> &DeviceMatrix {
         self.lm_head.as_ref().unwrap_or(&self.embed_tokens)
     }
 
-    /// Force the tile-alignment pad rows of a logits buffer to -inf.
-    pub(crate) fn suppress_pad_logits(&self, logits: &mut HiddenStates) -> Result<()> {
-        match &self.pad_logit_suppress {
-            Some(suppress) => {
-                crate::ops::suppress_logits_bf16_in_place(&self.ctx, logits, suppress)
-            }
-            None => Ok(()),
+    /// Write selectable logits for `normed` rows: the output-projection GEMM
+    /// followed by the tile-alignment pad-row mask.
+    ///
+    /// The two belong together — a site that ran the GEMM alone would leave the
+    /// pad rows selectable, which on the wire is an undecodable id fed back into
+    /// later decode steps rather than a failure — so this is the only way the
+    /// model produces logits.
+    pub(crate) fn output_logits_into(
+        &self,
+        normed: &HiddenStates,
+        logits: &mut HiddenStates,
+    ) -> Result<()> {
+        let vocab = self.config.selection_vocab;
+        gemm_rows_into_checked(
+            &self.ctx,
+            self.output_projection(),
+            0,
+            vocab,
+            normed,
+            logits,
+        )?;
+        if let Some(suppress) = &self.pad_logit_suppress {
+            suppress_logits_bf16_in_place(&self.ctx, logits, suppress)?;
         }
+        Ok(())
     }
 
     pub(crate) fn ensure_rope_cache_covers(&self, positions: usize) -> Result<()> {

From 423d908a863e209ff240e693a587f6b60471ce46 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Tue, 15 Sep 2026 19:51:14 +0800
Subject: [PATCH 08/12] refactor(qwen35): make output_projection private

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>
---
 pegainfer-qwen35/src/weights.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs
index d500bec96..ca2b158d1 100644
--- a/pegainfer-qwen35/src/weights.rs
+++ b/pegainfer-qwen35/src/weights.rs
@@ -353,7 +353,7 @@ impl Qwen35Model {
 
     /// Only the GEMM tuning helper samples this directly; logits go through
     /// [`Qwen35Model::output_logits_into`] so the pad-row mask cannot be skipped.
-    pub(super) fn output_projection(&self) -> &DeviceMatrix {
+    fn output_projection(&self) -> &DeviceMatrix {
         self.lm_head.as_ref().unwrap_or(&self.embed_tokens)
     }
 

From 87ad0be70055fa482b22f34d0ee5cdb627b1b062 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Wed, 16 Sep 2026 11:46:20 +0800
Subject: [PATCH 09/12] docs(qwen35): correct the alignment doc's verification
 claims

Review follow-up on #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>
---
 docs/models/qwen35/decode-kernel-attribution.md | 4 +++-
 docs/models/qwen35/optimization.md              | 6 ++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/docs/models/qwen35/decode-kernel-attribution.md b/docs/models/qwen35/decode-kernel-attribution.md
index 837e5641b..1fd8ef029 100644
--- a/docs/models/qwen35/decode-kernel-attribution.md
+++ b/docs/models/qwen35/decode-kernel-attribution.md
@@ -47,13 +47,15 @@ Neither kernel has had an sm_80 tuning pass (decode tuning history is sm_120/RTX
 
 `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`); `hf_golden_gate` TP1/TP2 validate output equivalence against HF, which computes logits over the full checkpoint vocab itself.
+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 regression coverage for both invariants is in `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, and the suite's other tests pin argmax/sampler routing. Nothing in the suite would fail if the mask stopped being applied on the real model path, 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.
diff --git a/docs/models/qwen35/optimization.md b/docs/models/qwen35/optimization.md
index 2bc19fe90..69584f8d7 100644
--- a/docs/models/qwen35/optimization.md
+++ b/docs/models/qwen35/optimization.md
@@ -2,7 +2,7 @@
 
 > **TL;DR:** Hybrid 24 linear + 8 full attn. This file is the historical optimization ledger; the current RTX 5090 comparison is in [Qwen3.5 serving: pegainfer vs vLLM on RTX 5090](../../benchmarks/qwen35-4b-serving-vllm-rtx5090.md). The decode-tuning refresh fuses MLP gate/up and tunes decode cuBLASLt buckets, improving direct TPOT by `2.1-3.2%`; vLLM still leads 1024/256 HTTP decode and high-concurrency throughput.
 >
-> **Last touched:** 2026-08. Qwen3.5 runtime code lives in top-level `pegainfer-qwen35`. Current accuracy coverage is `PEGAINFER_CUDA_SM=120 PEGAINFER_TEST_MODEL_PATH=<absolute Qwen3.5-4B path> cargo test --release -p pegainfer-qwen35 --test hf_golden_gate -- --nocapture`; run `e2e_scheduler` when scheduler request-flow behavior changes. The old exact-text e2e/regen baseline was retired by the HF logits gate in `docs/models/qwen35/accuracy.md`.
+> **Last touched:** 2026-09. Qwen3.5 runtime code lives in top-level `pegainfer-qwen35`. Current accuracy coverage is `PEGAINFER_CUDA_SM=120 PEGAINFER_TEST_MODEL_PATH=<absolute Qwen3.5-4B path> cargo test --release -p pegainfer-qwen35 --test hf_golden_gate -- --nocapture`; run `e2e_scheduler` when scheduler request-flow behavior changes. The old exact-text e2e/regen baseline was retired by the HF logits gate in `docs/models/qwen35/accuracy.md`.
 
 Historical command logs below keep the command paths that were actually run at the time (the in-process `bench_serving` bin they reference is retired as of 2026-08). For new Qwen3.5 accuracy tests, use `-p pegainfer-qwen35 --test hf_golden_gate`; for serving benchmarks, use HTTP-level benching (`scripts/bench_http_serving.py` / vllm-bench, see [profiling-guide](../../playbooks/profiling-guide.md)).
 
@@ -254,10 +254,12 @@ Trace note: the `nsys` capture includes one warmup run plus one measured run, so
 
 The GEMV family can be split further by launch shape (`gridX`) because each output width maps to a distinct projection class in the current Qwen3.5 decode path:
 
+The LM-head rows below were measured while selection still spanned the full 248,320-row checkpoint. The output projection now runs at the tile-aligned selection width (248,192 on 4B) — see [decode-kernel-attribution](decode-kernel-attribution.md) — so read those rows as a shape class, not the current cost.
+
 | GEMV subfamily | Time/step | % | Count/step | Avg each | Mapping |
 |----------------|-----------|---|------------|----------|---------|
 | Q / QKV (8192-dim) | 1.65ms | 13.1% | 32 | 51.5μs | 8 full-attn `q_proj` + 24 linear-attn `in_proj_qkv` |
-| LM head (248320-dim) | 1.50ms | 12.0% | 1 | 1.50ms | final logits projection |
+| LM head (248320-dim, pre-alignment) | 1.50ms | 12.0% | 1 | 1.50ms | final logits projection |
 | O projection (2560-dim) | 0.84ms | 6.7% | 32 | 26.2μs | 8 full-attn `o_proj` + 24 linear-attn `out_proj` |
 | Z projection (4096-dim) | 0.65ms | 5.2% | 24 | 27.1μs | 24 linear-attn `in_proj_z` |
 | B / A projection (32-dim) | 0.20ms | 1.6% | 48 | 4.2μs | 24 linear-attn `in_proj_b` + 24 `in_proj_a` |

From 026bca982d6bd6a8cf905db1f1471258e45b8c24 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Wed, 16 Sep 2026 11:54:00 +0800
Subject: [PATCH 10/12] test(qwen35): assert the pad columns are masked on the
 real logits paths
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review follow-up on #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>
---
 pegainfer-qwen35/src/unified_forward.rs | 67 +++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs
index 4c313a63d..fdce4eeea 100644
--- a/pegainfer-qwen35/src/unified_forward.rs
+++ b/pegainfer-qwen35/src/unified_forward.rs
@@ -165,6 +165,73 @@ mod tests {
             .unwrap()
     }
 
+    /// The alignment pad columns must come out of the model already masked, on
+    /// both the prefill and the decode logits path. A logits site that ran the
+    /// output projection without the mask leaves them selectable, and nothing
+    /// else in the suite would notice: the pad rows are trained embeddings with
+    /// plausible logits, so the failure shows up as an undecodable id on the
+    /// wire rather than a crash.
+    #[test]
+    fn pad_columns_are_suppressed_on_the_model_logits_paths() {
+        let Some(model_path) = crate::test_fixture::model_path_or_skip(
+            "pad_columns_are_suppressed_on_the_model_logits_paths",
+        ) else {
+            return;
+        };
+        let model = Qwen35Model::from_safetensors(&model_path, 0, 2).unwrap();
+
+        let decodable = model.config.decodable_vocab;
+        let selection = model.config.selection_vocab;
+        if decodable == selection {
+            eprintln!("checkpoint needs no alignment pad; nothing to assert");
+            return;
+        }
+
+        let pad_tail_is_suppressed = |logits: &HiddenStates, label: &str| {
+            assert_eq!(logits.hidden_dim, selection, "{label}: arena width");
+            let row = crate::ops::extract_vec(&model.ctx, logits, 0).unwrap();
+            let row = row.to_host(&model.ctx).unwrap();
+            for id in decodable..selection {
+                assert_eq!(
+                    f32::from(row[id]),
+                    f32::NEG_INFINITY,
+                    "{label}: pad id {id} survived selection (want -inf)"
+                );
+            }
+            let live = (0..decodable)
+                .filter(|&id| f32::from(row[id]).is_finite())
+                .count();
+            assert!(live > 0, "{label}: decodable vocab came out entirely -inf");
+        };
+
+        let prompts: Vec<Vec<u32>> = vec![vec![9707, 374, 220, 17]];
+        let prompt_refs: Vec<&[u32]> = prompts.iter().map(|p| p.as_slice()).collect();
+        let mut kv_states = vec![model.alloc_kv()];
+        let mut rec_states =
+            vec![RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap()];
+        let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect();
+        let prefill_logits = model
+            .batch_prefill_logits(&prompt_refs, &mut kv_states, &mut rec_refs)
+            .unwrap();
+        pad_tail_is_suppressed(&prefill_logits, "prefill");
+
+        let next = greedy_sample_batch(&model, &prefill_logits, 1)[0];
+        let mut graph_state = model.create_batch_decode_graph_state().unwrap();
+        graph_state
+            .copy_state_to_slot(&model.ctx, &rec_states[0], 0)
+            .unwrap();
+        let mut kv_refs: Vec<&mut KvState> = kv_states.iter_mut().collect();
+        model
+            .batch_decode_graph(
+                &[next],
+                &mut kv_refs,
+                &mut graph_state,
+                DecodeGraphUse::Serve,
+            )
+            .unwrap();
+        pad_tail_is_suppressed(&graph_state.buffers.logits, "decode");
+    }
+
     /// Verify that unified_step decode output matches batch_decode_graph standalone.
     #[test]
     fn unified_step_decode_matches_graph_decode() {

From fe48b666501a6c6182ff519d065ed1fffdbd7c93 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Wed, 16 Sep 2026 11:54:09 +0800
Subject: [PATCH 11/12] docs(qwen35): the mask invariant now has a model-path
 test

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>
---
 docs/models/qwen35/decode-kernel-attribution.md | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/docs/models/qwen35/decode-kernel-attribution.md b/docs/models/qwen35/decode-kernel-attribution.md
index 1fd8ef029..45a1e553d 100644
--- a/docs/models/qwen35/decode-kernel-attribution.md
+++ b/docs/models/qwen35/decode-kernel-attribution.md
@@ -54,7 +54,12 @@ The widened width is a GEMM decision, not a token-space one, so two invariants k
 - **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 regression coverage for both invariants is in `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, and the suite's other tests pin argmax/sampler routing. Nothing in the suite would fail if the mask stopped being applied on the real model path, which is why the GEMM and the mask are one method rather than two statements at each site.
+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
 

From 302b46b0b4e03fa3b5b967a97aab49aa2bcfa237 Mon Sep 17 00:00:00 2001
From: CAICAIIs <3360776475@qq.com>
Date: Wed, 16 Sep 2026 12:10:44 +0800
Subject: [PATCH 12/12] test(qwen35): satisfy clippy on the pad-column test

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>
---
 pegainfer-qwen35/src/unified_forward.rs | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs
index fdce4eeea..a831500a0 100644
--- a/pegainfer-qwen35/src/unified_forward.rs
+++ b/pegainfer-qwen35/src/unified_forward.rs
@@ -191,24 +191,22 @@ mod tests {
             assert_eq!(logits.hidden_dim, selection, "{label}: arena width");
             let row = crate::ops::extract_vec(&model.ctx, logits, 0).unwrap();
             let row = row.to_host(&model.ctx).unwrap();
-            for id in decodable..selection {
-                assert_eq!(
-                    f32::from(row[id]),
-                    f32::NEG_INFINITY,
-                    "{label}: pad id {id} survived selection (want -inf)"
+            for (id, value) in row.iter().enumerate().take(selection).skip(decodable) {
+                assert!(
+                    value.is_sign_negative() && value.is_infinite(),
+                    "{label}: pad id {id} survived selection (want -inf, got {value})"
                 );
             }
-            let live = (0..decodable)
-                .filter(|&id| f32::from(row[id]).is_finite())
-                .count();
-            assert!(live > 0, "{label}: decodable vocab came out entirely -inf");
+            assert!(
+                row[..decodable].iter().any(|value| value.is_finite()),
+                "{label}: decodable vocab came out entirely -inf"
+            );
         };
 
-        let prompts: Vec<Vec<u32>> = vec![vec![9707, 374, 220, 17]];
-        let prompt_refs: Vec<&[u32]> = prompts.iter().map(|p| p.as_slice()).collect();
+        let prompt_refs: Vec<&[u32]> = vec![&[9707, 374, 220, 17]];
         let mut kv_states = vec![model.alloc_kv()];
         let mut rec_states =
-            vec![RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap()];
+            [RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap()];
         let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect();
         let prefill_logits = model
             .batch_prefill_logits(&prompt_refs, &mut kv_states, &mut rec_refs)