fix(qwen3): cap PerToken CUDA Graph buckets

docs/models/qwen3/decode-attention.md

@@ -1,8 +1,8 @@
 # Qwen3-4B Decode Attention Path Selection
 
-**TL;DR:** Decode picks between two paged-attention kernels — `NonPartition` (1 CTA per request×kv-head) and `SplitKv` (KV split into fixed-size chunks across SMs). The choice is driven by **batch (CTA count vs SM count), not context length**. The old `max_seq_len >= 1024` gate was a tuning artifact that left bs=1 mid-context decode on the SM-starved `NonPartition` kernel, producing a tpot hump that peaked ~ctx800 and dropped off a cliff exactly at ctx1024. Removing it flattens bs=1 tpot across the whole context range (5090 −16% @ctx800, 5070 Ti −7.5%) with no accuracy regression. Kept the `padded_bs <= 32` cap. Under `--batch-invariant` (`Pin`/`PerToken`) this batch-driven choice is overridden — the path is pinned to SplitKv for **every** bucket so it stops being a batch-composition axis, at a Pin-only cost in memory and long-context throughput; see *The fix*.
+**TL;DR:** Decode picks between two paged-attention kernels — `NonPartition` (1 CTA per request×kv-head) and `SplitKv` (KV split into fixed-size chunks across SMs). The choice is driven by **batch (CTA count vs SM count), not context length**. The old `max_seq_len >= 1024` gate was a tuning artifact that left bs=1 mid-context decode on the SM-starved `NonPartition` kernel, producing a tpot hump that peaked ~ctx800 and dropped off a cliff exactly at ctx1024. Removing it flattens bs=1 tpot across the whole context range (5090 −16% @ctx800, 5070 Ti −7.5%) with no accuracy regression. Kept the `padded_bs <= 32` cap. Under `--batch-invariant` (`Pin`/`PerToken`) this batch-driven choice is overridden — the path is pinned to SplitKv for **every** bucket so it stops being a batch-composition axis, at a Pin-only cost in memory and long-context throughput; see **The fix**. PerToken CUDA-Graph capture is capped at bucket 32; larger buckets retain the PerToken GEMM arithmetic but execute eagerly to bound graph executable memory.
 
-Last touched: 2026-07
+Last touched: 2026-08
 
 ## The two kernels
 
@@ -60,12 +60,11 @@ The transition lands exactly where CTA count crosses SM count: bs≤8 (≤64 CTA
 
 The grid must be fixed for graph replay, but SplitKv's chunk count varies with context. Resolved by **fixed-upper-bound grid + out-of-graph metadata**:
 
-- One captured graph per `(batch bucket, attention_path)` — `graph_index = bucket_idx × 2 + path.graph_slot()`. bs=1 has a `NonPartition` slot and a `SplitKv` slot; first use of a combination captures, later steps replay.
+- For `Tuned` and `Pin`, one captured graph is maintained per `(batch bucket, attention_path)` — `graph_index = bucket_idx × 2 + path.graph_slot()`. The first use of a combination captures it and later steps replay it.
+- `PerToken` uses CUDA Graph only for buckets up to `PERTOKEN_GRAPH_MAX_BUCKET` (currently 32). Larger buckets retain the PerToken GEMM arithmetic but run the decode eagerly, so they do not allocate or replay a CUDA Graph executable.
 - SplitKv workspace is pre-allocated to the active policy's grid (`bs × max_split_chunks()`: default `Tuned` `bs × 64` = `SPLIT_KV_TUNED_MAX_CHUNKS`, `Pin`/`PerToken` `bs × 256`), so buffer pointers stay stable across replay — the policy is fixed before construction, so the size never shifts under a live executor. The grid is fixed per `(bucket, policy)` (see *Chunk size and batch-invariance*).
 - Per-step context differences go through `memcpy_htod` in `sync_split_kv_meta` **before** `run_or_capture` (outside the graph): chunk_size, valid-chunk count, `valid_mask`, `o_indptr`. Chunks beyond the real count are masked off (`valid_mask = 0`, those CTAs early-exit).
 
-So under the `Tuned` 64-token floor, ctx=300 (5 chunks) and ctx=1024 (16 chunks) share the same SplitKv graph — only the metadata buffer contents differ (`Pin`/`PerToken` give different per-context counts, still within the same fixed grid). This is why dropping the context gate is safe: capture-time context never determined the grid. The accuracy gate's CUDA-graph replay over small-context sequences passes, confirming it.
-
 ## Chunk size and batch-invariance
 
 The chunk *size* sets a request's chunk count, hence its online-softmax merge order; bf16 non-associativity makes the decoded logits depend on that count. `split_chunk_size()` picks it by `NumericPolicy`:
@@ -75,11 +74,11 @@ The chunk *size* sets a request's chunk count, hence its online-softmax merge or
 
 `SPLIT_KV_MAX_CHUNKS_PER_REQUEST` (256) is the absolute upper bound — both the `Pin`/`PerToken` chunk cap and the `pin_chunk_size` divisor, which must match so a request yields ≤ cap chunks and the guard stays tight. Both the workspace and the grid are sized to the active policy. The default `Tuned` path caps the split batch at `min(bs, 32)`, so at `MAX_BATCH` its workspace is `32 × 64` slots ≈ 16 MiB. Under `Pin`/`PerToken` the split batch is the full `bs` and the chunk cap is 256 (pinning SplitKv at every bucket lifted the `Tuned` cap of 32), so at `MAX_BATCH = 256` the workspace is `256 × 256` slots ≈ 512 MiB (`split_tmp_v` dominates) — **+448 MiB** over the pre-fix Pin cap (`32 × 256` ≈ 64 MiB) it replaces, and ≈ **+496 MiB** over the default `Tuned` sizing, pre-subtracted from the KV cache. TP shrinks it (per-GPU `local_q_dim`; TP8 ~64 MiB). `--batch-invariant` also pins the decode GEMM-N reduction order (an orthogonal axis); chunk size alone does not make decode fully batch-invariant.
 
-The `batch_invariance_decode_splitkv_graph` gate covers this: co-batching a request with a longer neighbour drifts its `Tuned` chunk count (the decoded top-K changes) while `Pin`/`PerToken` replay the requested top-K logprobs bit-identically across the SplitKv CUDA-graph (the gate compares A's prefill first token and its `LOGPROBS=64` decode top-K, not full logits).
+The `batch_invariance_decode_splitkv_graph` gate covers this: co-batching a request with a longer neighbour drifts its `Tuned` chunk count (the decoded top-K changes), while `Pin` and graph-backed `PerToken` buckets preserve the requested top-K logprobs across SplitKv graph replay. For PerToken buckets above the graph cap, the same PerToken arithmetic is executed eagerly instead of through CUDA Graph. The gate compares A's prefill first token and its `LOGPROBS=64` decode top-K, not full logits.
 
 ## The fix
 
-`attention_path()` gates on `NumericPolicy` first, then `padded_bs`. Under the default `Tuned` it is the batch-driven choice this doc describes: SplitKv when `padded_bs <= SPLIT_KV_MAX_BATCH_SIZE`, else NonPartition (the `SPLIT_KV_MIN_SEQ_LEN` constant is gone — small-context sequences run entirely on SplitKv; `tests/hf_golden_gate.rs` head delta at bf16 noise, mean ~0.03 p99 ~0.11, no regression). Under `Pin`/`PerToken` (`--batch-invariant`) it pins SplitKv for **every** bucket: otherwise a request crossing the 32 cap under co-scheduled load switches kernel, and the two paths are different reductions, so its logits move with its batch-mates — the batch-invariance residual this closes. The chunk size is already pinned (above), so the constant path is batch-invariant by construction.
+`attention_path()` gates on `NumericPolicy` first, then `padded_bs`. Under the default `Tuned` it is the batch-driven choice this doc describes: SplitKv when `padded_bs <= SPLIT_KV_MAX_BATCH_SIZE`, else NonPartition (the `SPLIT_KV_MIN_SEQ_LEN` constant is gone; small-context sequences run entirely on SplitKv; `tests/hf_golden_gate.rs` head delta at bf16 noise, mean ~0.03 p99 ~0.11, no regression). Under `Pin`/`PerToken` (`--batch-invariant`) it pins SplitKv for every bucket. Without this override, a request crossing `SPLIT_KV_MAX_BATCH_SIZE` under co-scheduled load would switch attention kernels, and the two paths are different reductions, so its logits would move with its batch-mates. The chunk size is already pinned (above), so the constant path is batch-invariant by construction. This attention-path decision is independent of the PerToken CUDA Graph cap: PerToken buckets above 32 still use SplitKv attention, but their decode execution is eager.
 
 Pinning the path above the cap is not free — past `SPLIT_KV_MAX_BATCH_SIZE` NonPartition already fills the SMs, so SplitKv's merge is the *bs>1* overhead described above, now paid deliberately for invariance. Measured (`attention_path_perf`, sm_89, Pin decode tok/s, SplitKv vs the NonPartition it replaces): short context wins at moderate batch (CTX=256 net +6.6% @bs=64), narrowing to ≈ −1% net by `MAX_BATCH`; long context goes negative (CTX=1024 −3.7% @bs=64 — the worst *measured* point, not a floor: bs ≥ 128 at that context exceeds a single 48 GB card's KV pool) as chunk count (∝ context) grows the merge. sm_90 untested. `Pin`/`PerToken` opt-in only; default `Tuned` keeps the batch-driven choice and is byte-identical. Verified by the `batch_invariance_attention_path` gate.
 

pegainfer-qwen3/src/batch_decode.rs

@@ -5,6 +5,7 @@ use cudarc::driver::CudaSlice;
 use half::bf16;
 use pegainfer_core::kv_pool::KvLayout;
 use pegainfer_core::ops;
+use pegainfer_kernels::ops::NumericPolicy;
 use pegainfer_kernels::tensor::KvDim;
 use pegainfer_kernels::tensor::QDim;
 use pegainfer_kv_cache::KvView;
@@ -59,6 +60,81 @@ pub(crate) enum DecodeGraphUse {
     Replay,
 }
 
+/// Largest batch bucket for which PerToken retains a CUDA Graph. PerToken
+/// records one GEMM node per row, so each additional bucket has a growing,
+/// independently resident graph executable.
+const PERTOKEN_GRAPH_MAX_BUCKET: usize = 32;
+
+/// Canonical CUDA Graph coverage and dispatch policy for decode
+///
+/// Serving and memory profiling both consum this plan so the profiler reserves
+/// every retained PerToken full-SM Graph covered by this policy
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(crate) struct DecodeGraphPlan {
+    policy: NumericPolicy,
+}
+
+impl DecodeGraphPlan {
+    pub(crate) const fn new(policy: NumericPolicy) -> Self {
+        Self { policy }
+    }
+
+    const fn allows_bucket(self, bucket: usize) -> bool {
+        !matches!(self.policy, NumericPolicy::PerToken) || bucket <= PERTOKEN_GRAPH_MAX_BUCKET
+    }
+
+    pub(crate) const fn requires_cumulative_profile(self) -> bool {
+        matches!(self.policy, NumericPolicy::PerToken)
+    }
+
+    pub(crate) fn retained_buckets(self) -> impl Iterator<Item = usize> {
+        BATCH_BUCKETS
+            .iter()
+            .copied()
+            .filter(move |&bucket| self.allows_bucket(bucket))
+    }
+
+    fn resolve(
+        self,
+        graph_use: DecodeGraphUse,
+        graphs_available: bool,
+        batch_size: usize,
+    ) -> Result<Option<usize>> {
+        if !graphs_available {
+            anyhow::ensure!(
+                matches!(graph_use, DecodeGraphUse::Serve | DecodeGraphUse::Eager),
+                "batch_decode {graph_use:?} requires CUDA graph enabled and no LoRA rows"
+            );
+            return Ok(None);
+        }
+        if graph_use == DecodeGraphUse::Eager {
+            return Ok(None);
+        }
+
+        let bucket = bucket_for(batch_size);
+        if self.allows_bucket(bucket) {
+            return Ok(Some(bucket));
+        }
+
+        match graph_use {
+            DecodeGraphUse::Serve => {
+                log::debug!(
+                    "PerToken decode bucket {bucket} exceeds graph cap \
+                    {PERTOKEN_GRAPH_MAX_BUCKET}; using eager decode"
+                );
+                Ok(None)
+            }
+            DecodeGraphUse::CaptureOnly | DecodeGraphUse::Replay => {
+                anyhow::bail!(
+                    "batch_decode {graph_use:?} requested PerToken CUDA Graph bucket {bucket} \
+                    above cap {PERTOKEN_GRAPH_MAX_BUCKET}"
+                )
+            }
+            DecodeGraphUse::Eager => unreachable!("handled above"),
+        }
+    }
+}
+
 impl Qwen3Model {
     /// Batch decode step: N requests, 1 new token each, one forward pass.
     ///
@@ -90,17 +166,14 @@ impl Qwen3Model {
         let lora_slots = self.decode_lora_slots(lora_adapters)?;
         let use_lora = lora_slots.is_some();
         let graphs_available = self.enable_cuda_graph && lora_slots.is_none();
-        anyhow::ensure!(
-            graphs_available || matches!(graph_use, DecodeGraphUse::Serve | DecodeGraphUse::Eager),
-            "batch_decode {graph_use:?} requires CUDA graphs enabled and no LoRA rows"
-        );
-        let use_cuda_graph = graphs_available && graph_use != DecodeGraphUse::Eager;
+        let graph_plan = DecodeGraphPlan::new(bufs.policy_at_construction);
+        let graph_bucket = graph_plan.resolve(graph_use, graphs_available, bs)?;
 
         // Derive positions from views (seq_len - 1 = position of the new token)
         let mut positions: Vec<i32> = kv_views.iter().map(|v| (v.seq_len() - 1) as i32).collect();
 
         // Pad to bucket size for CUDA Graph stability
-        let padded_bs = if use_cuda_graph { bucket_for(bs) } else { bs };
+        let padded_bs = graph_bucket.unwrap_or(bs);
 
         // Set batch size on all buffers (padded — kernels run at bucket width)
         bufs.set_batch_size(padded_bs);
@@ -129,8 +202,11 @@ impl Qwen3Model {
             BatchDecodeBuffers::attention_path(padded_bs, bufs.policy_at_construction);
         #[cfg(feature = "kernel-call-trace")]
         let trace_kv_len = kv_views.iter().map(|v| v.seq_len()).max().unwrap_or(0);
-        if use_cuda_graph {
-            let bucket_idx = BATCH_BUCKETS.iter().position(|&b| b == padded_bs).unwrap();
+        if let Some(bucket) = graph_bucket {
+            let bucket_idx = BATCH_BUCKETS
+                .iter()
+                .position(|&candidate| candidate == bucket)
+                .expect("resolved graph bucket must exist in BATCH_BUCKETS");
             let graph_idx = BatchDecodeBuffers::graph_index(bucket_idx, attention_path);
             // Override-stream captures use the split cache; full-SM captures use the normal cache.
             let on_split_stream = pegainfer_kernels::tensor::has_stream_override();
@@ -157,7 +233,7 @@ impl Qwen3Model {
                 DecodeGraphUse::Serve => graph.run_or_capture(&self.ctx, kernels),
                 DecodeGraphUse::CaptureOnly => graph.capture_only(&self.ctx, kernels),
                 DecodeGraphUse::Replay => graph.launch_captured(&self.ctx),
-                DecodeGraphUse::Eager => unreachable!("use_cuda_graph excludes Eager"),
+                DecodeGraphUse::Eager => unreachable!("graph_bucket excludes Eager"),
             };
             if on_split_stream {
                 bufs.graphs_split = graphs;
@@ -503,3 +579,51 @@ fn grouped_projection_from_packed<'a>(
         out_dim: packed.out_dim,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use pegainfer_kernels::ops::NumericPolicy;
+
+    use super::DecodeGraphPlan;
+    use super::DecodeGraphUse;
+    use super::PERTOKEN_GRAPH_MAX_BUCKET;
+
+    #[test]
+    fn decode_graph_plan_matches_runtime_and_profile_coverage() {
+        let per_token = DecodeGraphPlan::new(NumericPolicy::PerToken);
+        assert!(per_token.requires_cumulative_profile());
+        assert_eq!(
+            per_token.retained_buckets().collect::<Vec<_>>(),
+            vec![1, 2, 4, 8, 16, 20, 24, 32]
+        );
+
+        assert_eq!(
+            per_token.resolve(DecodeGraphUse::Serve, true, 32).unwrap(),
+            Some(PERTOKEN_GRAPH_MAX_BUCKET)
+        );
+        assert_eq!(
+            per_token.resolve(DecodeGraphUse::Serve, true, 33).unwrap(),
+            None
+        );
+
+        for graph_use in [DecodeGraphUse::CaptureOnly, DecodeGraphUse::Replay] {
+            let error = per_token
+                .resolve(graph_use, true, 33)
+                .expect_err("explicit graph use above the PerToken cap must fail");
+            let message = error.to_string();
+            assert!(
+                message.contains("PerToken CUDA Graph bucket 40 above cap 32"),
+                "{graph_use:?}: {message}"
+            );
+        }
+
+        for policy in [NumericPolicy::Pin, NumericPolicy::Tuned] {
+            let plan = DecodeGraphPlan::new(policy);
+            assert!(!plan.requires_cumulative_profile());
+            assert_eq!(
+                plan.resolve(DecodeGraphUse::Serve, true, 256).unwrap(),
+                Some(256)
+            );
+        }
+    }
+}

pegainfer-qwen3/src/batch_decode_buffers.rs

@@ -26,6 +26,7 @@ pub(crate) const BATCH_BUCKETS: &[usize] = &[
     1, 2, 4, 8, 16, 20, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152,
     160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256,
 ];
+
 const DECODE_ATTENTION_PATH_COUNT: usize = 2;
 // Split-KV decode attention: the non-partitioned kernel issues one CTA per
 // (request x kv-head), starving SMs at small batch. The path is therefore

pegainfer-qwen3/src/unified_forward.rs

@@ -17,6 +17,8 @@ use pegainfer_kernels::ops::numeric_policy;
 use pegainfer_kv_cache::KvBuffer;
 use pegainfer_kv_cache::KvView;
 
+use super::batch_decode::DecodeGraphPlan;
+use super::batch_decode::DecodeGraphUse;
 use super::batch_decode_buffers::BatchDecodeBuffers;
 use super::config::PREFILL_ATTENTION_CTA_TILE_Q;
 use super::prefill::PrefillBuffers;
@@ -61,17 +63,41 @@ impl Qwen3Model {
         let decode_tokens = vec![0u32; profile_decode_rows];
         let decode_adapters = vec![None; profile_decode_rows];
 
-        // Force the decode CUDA-Graph/buffer path before the unified peak
-        // sample. The synthetic views are short, but the pre-allocated decode
-        // arena and graph state are the same serving objects used later. Skip it
-        // for uncompiled-group models: the unified sample below bounds their KV.
-        // Eager under TP: ranks profile uncoordinated, so an in-profile capture
-        // would hit the same deadlock the sweep avoids (see `PrecapturePhase`).
+        // Exercise decode before the unified peak sample. PerToken first captures
+        // every retained graph into this one buffer set so their cumulative
+        // residency remains live under the later eager/unified probes. Skip this
+        // for uncompiled-group modles; the unified sample below bounds their KV.
+        // TP stays eager because its ranks profile independently.
         if self.config.decode_group_is_compiled() {
+            let graph_plan = DecodeGraphPlan::new(decode_bufs.policy_at_construction);
+            if self.enable_cuda_graph
+                && self.tensor_parallel.world_size == 1
+                && graph_plan.requires_cumulative_profile()
+            {
+                for graph_rows in graph_plan.retained_buckets() {
+                    anyhow::ensure!(
+                        graph_rows <= profile_decode_rows,
+                        "retained decode Graph bucket {graph_rows} exceeds profile row capacity \
+                        {profile_decode_rows}"
+                    );
+                    self.batch_decode(
+                        &decode_tokens[..graph_rows],
+                        &decode_views[..graph_rows],
+                        &decode_adapters[..graph_rows],
+                        kv_buffer.buffer(),
+                        &layout,
+                        decode_bufs,
+                        DecodeGraphUse::Serve,
+                    )?;
+                    self.ctx.sync()?;
+                    mark_peak()?;
+                }
+            }
+
             let graph_use = if self.tensor_parallel.world_size > 1 {
-                crate::batch_decode::DecodeGraphUse::Eager
+                DecodeGraphUse::Eager
             } else {
-                crate::batch_decode::DecodeGraphUse::Serve
+                DecodeGraphUse::Serve
             };
             self.batch_decode(
                 &decode_tokens,