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..45a1e553d --- /dev/null +++ b/docs/models/qwen35/decode-kernel-attribution.md @@ -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. 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= 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= 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` | diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index b41da41d4..a558a5054 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -694,14 +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.output_logits_into(&bufs.normed, &mut bufs.logits)?; debug_assert_eq!(bufs.logits.seq_len, padded_bs); Ok(()) @@ -772,14 +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.output_logits_into(&bufs.normed, &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 566c25459..31683eba7 100644 --- a/pegainfer-qwen35/src/config/model.rs +++ b/pegainfer-qwen35/src/config/model.rs @@ -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. + /// 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. + /// 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, }); } - 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.decodable_vocab = effective_vocab; + self.selection_vocab = aligned.min(self.vocab_size); Ok(()) } } @@ -278,6 +292,7 @@ impl TryFrom 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,", ""); diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs index cea677445..5e56273ad 100644 --- a/pegainfer-qwen35/src/decode_buffers.rs +++ b/pegainfer-qwen35/src/decode_buffers.rs @@ -116,7 +116,12 @@ 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)?, + 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/ops.rs b/pegainfer-qwen35/src/ops.rs index 3d288b074..af08deb6b 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; @@ -9,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; diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index 6fa2555d7..1537e458c 100644 --- a/pegainfer-qwen35/src/prefill.rs +++ b/pegainfer-qwen35/src/prefill.rs @@ -117,14 +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.output_logits_into(&normed, &mut logits)?; debug_assert_eq!(logits.seq_len, n); Ok(logits) } 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..a831500a0 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -153,14 +153,83 @@ 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(); + 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, ¶ms_refs, &steps, 0, &mut scratch) .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, 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})" + ); + } + assert!( + row[..decodable].iter().any(|value| value.is_finite()), + "{label}: decodable vocab came out entirely -inf" + ); + }; + + let prompt_refs: Vec<&[u32]> = vec![&[9707, 374, 220, 17]]; + let mut kv_states = vec![model.alloc_kv()]; + let mut rec_states = + [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() { diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index 30527c48e..ca2b158d1 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,6 +71,10 @@ 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, + /// -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, } // SAFETY: A Qwen3.5 model instance is bound to one CUDA device and driven from @@ -167,10 +173,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 ); } @@ -308,6 +314,20 @@ impl Qwen35Model { page_size, num_pages, )?; + // 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 = (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 +343,7 @@ impl Qwen35Model { reserved_decode_slots: max_batch, decode_admission_batch, tp_comm: None, + pad_logit_suppress, }) } @@ -330,10 +351,39 @@ impl Qwen35Model { &self.config } - pub(super) fn output_projection(&self) -> &DeviceMatrix { + /// Only the GEMM tuning helper samples this directly; logits go through + /// [`Qwen35Model::output_logits_into`] so the pad-row mask cannot be skipped. + fn output_projection(&self) -> &DeviceMatrix { self.lm_head.as_ref().unwrap_or(&self.embed_tokens) } + /// 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<()> { let cache_positions = self.cos_cache.len / self.config.rotary_dim; anyhow::ensure!( diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs index b96e8eee2..a5206c72b 100644 --- a/pegainfer-sample/src/lib.rs +++ b/pegainfer-sample/src/lib.rs @@ -83,15 +83,36 @@ 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: the + /// emittable tokens, not the arena a model may have tile-aligned wider. + selection_width: usize, max_rows: usize, } impl SampleScratch { pub fn new(ctx: &DeviceContext, vocab: usize, max_rows: usize) -> Result { + Self::with_selection_width(ctx, vocab, vocab, max_rows) + } + + /// 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, + selection_width: usize, + max_rows: usize, + ) -> Result { 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> { ctx.stream @@ -124,6 +145,7 @@ impl SampleScratch { .map_err(|e| anyhow!("SampleScratch identity upload failed: {e}"))?, sampling: BatchSamplingScratch::new(ctx, max_rows, vocab)?, vocab, + selection_width, max_rows, }) } @@ -151,6 +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 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. @@ -196,7 +220,9 @@ pub fn select_batch( "select_batch: logits vocab {vocab} != scratch vocab {}", scratch.vocab ); - let is_argmax = |p: &&SamplingParams| effectively_greedy(p, vocab); + // 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..065c082c1 100644 --- a/pegainfer-sample/tests/select_batch.rs +++ b/pegainfer-sample/tests/select_batch.rs @@ -177,6 +177,47 @@ fn tiny_top_p_routes_to_argmax_even_under_bf16_ties() { } } +#[test] +fn padded_arena_width_does_not_suppress_the_argmax_routing() { + // 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; + 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); + + // 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 { + 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(); + 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 batch_larger_than_scratch_is_rejected() { let ctx = DeviceContext::new().unwrap();