feat(qwen35): cuda graph decode under TP (split 3/4 of #946)

docs/index.md

@@ -53,8 +53,8 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 | `models/qwen35/accuracy.md` | Qwen3.5 HF bf16 logits goldens, size-keyed (0.8b/2b/4b/9b/27b all committed), through `past_key_values`: short replay covers sequential graph, bucket-straddling batched graph, and slot-compaction; long replay covers 4097/8192-token prompts; full GSM8K 8-shot now matches the HF baseline within 0.15 percentage points. |
 | `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. |
 | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. |
-| `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 is eager dense TP on Qwen3's controller/worker runtime; validate TP2 first, fail closed for indivisible degrees and TP+CUDA Graph, shard dense full-attention/MLP, and leave sharded linear/GDR state to follow-up. |
-| `models/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1 and P2A are complete: TP2 has start-gated eager unified prefill+decode, strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load ordinal validation; P2B GDR state sharding is next. |
+| `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 eager dense TP on Qwen3's controller/worker runtime, P2a mixed-step protocol, P2b rank-local GDR sharding, P2c CUDA Graph under TP gated on the compiled decode GQA group (27B group-6 stays eager). |
+| `models/qwen35/tp-implementation.md` | Qwen3.5 TP landed through P2c on #870 (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 decode (27B: 292 tok/s ×16); TP decode CUDA Graphs for 4B/9B (9B: 767 vs 706 tok/s ×16 eager). 9B/27B TP2 HF + e2e gates pass. |
 | `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/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. |

docs/models/qwen35/tp-design.md

@@ -269,6 +269,30 @@ Validation scope:
 - recurrent-state cleanup on finish/drop/cancellation
 - no stale local recurrent state after a new `RequestId` is admitted
 
+## P2c: CUDA Graph under TP
+
+Status: landed (2026-08-20) on `feat/qwen35-tp2-rebased`, gated on
+`local_decode_group_is_compiled` — 4B/9B TP2 capture and replay decode graphs;
+27B TP2 (group 6) stays on the batched eager path byte-for-byte until group-6
+batch-decode kernels are compiled. Execution record: `tp-implementation.md`
+section "P2c — CUDA Graph under TP".
+
+**Gate**: graph mode active iff `enable_cuda_graph && config.local_decode_group_is_compiled(tp)`. 27B TP2 is group-6 (`SUPPORTED_GQA_GROUP_SIZES = [1,2,3,4,8]`, group ratio is TP-invariant), so 27B TP2 keeps the batched eager path byte-for-byte until group-6 batch-decode kernels are compiled; 4B/9B TP2 capture graphs. Startup logs once when graph was requested but the group gate keeps decode eager.
+
+**State model**: scheduler owns slot semantics (TP1 mirror); workers execute slot copies on command, never infer slots worker-side.
+
+- KV paged state unchanged (pool stable; page tables are per-step H2D via `sync_paged_meta`).
+- Per rank: `BatchDecodeGraphState`-equivalent at `bucket_for(effective_max_batch)` slots — fixed-address `slot_states: Vec<RecurrentState>` + one persistent `LinearStatePointerTables` built once over slots (contents stable → replay-safe).
+- Admission: decode command rows carry explicit `slot_idx` (`slot_for_new_request`); first decode row D2D-copies prefill `RecurrentState` into the slot (`copy_state_to_slot`), drops the per-request allocation.
+- Retirement: `DropRequest` gains `compaction: Option<(RequestId, from, to)>`; worker D2D-moves slot state (`move_slot_within`), asserts occupancy, poisons on mismatch.
+- Decode rows arrive dense slot order `0..bs`; padding rows clobber free slots (benign — admission overwrites).
+
+**Capture/replay**: startup pre-capture sweep ported from qwen3 (`executor.rs:1424`): `Warmup` (port `warmup_tp_collective`, one all-reduce per bucket message size — lazy NCCL connect inside capture wedges), `Capture`/`Launch` per bucket `[1,2,4,8,16,32,64]` with synthetic rows, `Finalize` asserts all captured; dedicated 600 s abort watchdog (60 s startup timeout too small). New `TpWorkerCommand::Precapture { phase }` via existing exact-rank dispatch. Serve time: replay-only (`ensure is_captured` + `launch_captured`), never capture mid-serving. Sampling/logprobs stay rank-0 host-side outside the graph. Mixed ticks: prefill eager + decode replay; collective order canonical per plan. `TpWorkerState` declares graph state before `model` so graphs drop before the NCCL comm (teardown hang precedent qwen3 `executor.rs:3076`).
+
+**Memory** (27B TP2/rank): weights ~17.5 GB + KV pool ~5.9 GiB + slot state reserve ~6.1 GiB + buffers/graphs ~0.3 + scratch/NCCL ~2.5 ≈ 32 GiB → fits 48 GB. 9B TP2 slot state ~1.6 GiB. Loader already reserves `2 × max_batch × bytes_per_request` before sizing KV.
+
+**Validation ladder**: CPU lib suite → TP2 graph HF gate (9B: sequential + bucket-straddling + post-compaction replay vs eager stats) → e2e scheduler graph variant → serving_tp2 graph smoke → 27B TP2 regression unchanged (group-6 stays eager) → per-bucket eager-vs-graph decode benchmark recorded in `bench_snapshots/`.
+
 ## References
 
 - `docs/models/qwen3/tp-design.md`

docs/models/qwen35/tp-implementation.md

@@ -1,6 +1,6 @@
 # Qwen3.5 TP Implementation Record
 
-> **TL;DR:** Qwen3.5 TP Phase 1 and P2A are complete: TP2 now supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation; P2B GDR state sharding is next.
+> **TL;DR:** Qwen3.5 TP is complete through Phase 2b plus batched eager TP decode and P2c CUDA Graph under TP: TP2 supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation (#870); linear-attention/GDR state is sharded per rank (27B TP2 fits on 2×48 GB); TP decode rows run as one batched forward per step; and 4B/9B TP2 decode replays pre-captured CUDA Graphs (27B group-6 stays eager by gate). Remaining TP work: group-6 batch-decode kernels for 27B graphs, perf gates.
 >
 > **Last touched:** 2026-08
 
@@ -460,9 +460,160 @@ Non-negotiable invariant:
 
 - Never all-reduce GDR recurrent state or conv state. These states are owned by rank-local request state.
 
+## Rebase onto #870 (2026-08-20)
+
+#870 landed its own Phase 1/2a upstream while the parallel line on
+`feat/qwen35-tp2-batched-decode` had implemented its own Phase 1/2a/2b
+plus a batched-decode fix on the old main. The rebase onto #870
+ported only the two deltas #870 lacks, with #870's code as the base:
+
+1. **Phase 2b — sharded linear-attention/GDR state** (commit
+   `feat(qwen35): shard linear-attention/GDR state per TP rank`). #870's
+   `recurrent_state.rs` had no rank sharding and its linear-attention
+   weights loaded replicated; the port adds rank-local slices end to end
+   (`weight_loader` stitch/shard loaders, `config.rs` `local_linear_*`
+   accessors, `weights.rs` per-rank stitched qkv/conv1d + row/col shards,
+   `recurrent_state`/`decode_buffers`/`prefill_buffers` at local sizes,
+   `batch_decode`/`prefill` local head counts + all-reduce after linear
+   `out_proj`, TP-local `batch_decode_full_attention_via_prefill` so 27B
+   TP2 group-6 eager decode routes through prefill). `tp_executor.rs`
+   only took the capacity-math and `RecurrentState::new` signature
+   changes; #870's worker protocol untouched.
+2. **Step 3 — batched eager TP decode** (commit
+   `perf(qwen35): batch eager decode rows under TP`). #870's
+   `execute_decode_rows` looped per request with bs=1 forwards and
+   capacity-1 per-request pointer tables. The port adds `run_decode_batch`
+   (one `batch_decode_eager_logits` over all decode rows, step-scoped
+   `LinearStatePointerTables::from_recurrent_refs(..., bs, ...)`, one
+   batched rank-0 `select_batch`, per-row fan-out in command order) inside
+   #870's `execute_decode_rows`, keeping its validation and response
+   contracts; `TpRequestState.linear_pointer_tables` removed.
+
+What #870 already covered (not ported): Phase 1 dense TP, the Phase 2a
+unified command/scheduler surface (`TpUnifiedPlan`, command start gates,
+dispatch/response validators, drop-expectation lifecycle proofs), and the
+scheduler planner-gate/test updates — ours' `scheduler.rs`,
+`scheduler/tests.rs`, and `e2e_scheduler.rs` deltas were subsumed
+upstream, so those files resolved to #870's versions except the
+`alloc_recurrent` signature change.
+
+Validation on 2× RTX 4090:
+
+- `cargo check --release -p pegainfer-qwen35 --features qwen35` clean;
+  `cargo fmt --check -p pegainfer-qwen35` clean.
+- Lib unit suite 101/101 (includes the four sharding layout tests:
+  segment tables, conv kernel-dim scaling, TP1 identity, synthetic
+  safetensors stitch contract).
+- 9B TP2 HF short+long gates PASS (24.7 s); 9B TP2 scheduler e2e
+  (`test_e2e_qwen35_scheduler_tp2`) PASS (27.1 s).
+- 27B TP2 HF short+long gates PASS (64.7 s) — 27B TP2 fits on 2×48 GB
+  only because of the Phase-2b sharding (the acceptance criterion for the
+  port). 27B TP2 scheduler e2e PASS (68.8 s).
+
+## P2c — CUDA Graph under TP (2026-08-20)
+
+Implemented the locked P2c design from `tp-design.md`: decode CUDA Graphs
+under TP, gated on the TP-local decode GQA group.
+
+**Gate.** Graph mode is active iff `enable_cuda_graph &&
+config.local_decode_group_is_compiled(tp)`. The group ratio is TP-invariant,
+so 27B TP2 (group 6, not in `SUPPORTED_GQA_GROUP_SIZES`) keeps the batched
+eager path byte-for-byte while 4B/9B TP2 capture. The old fail-closed
+rejections (`config.rs` `validate_for`, `tp_executor.rs` startup ensure,
+`lib.rs` TP branch) were replaced by the gate; startup logs once when graph
+was requested but the group gate keeps decode eager.
+
+**State model.** Scheduler owns slots, workers execute:
+
+- `ActiveBackendState::Tp` gains `slot_idx` (dense `active` position, assigned
+  at promote via `slot_for_new_request`, updated on compaction);
+  `tp_decode_items` emits rows with explicit `slot_idx` and workers assert
+  `slot_idx == row`.
+- Graph workers hold a `BatchDecodeGraphState` at
+  `bucket_for(effective_max_batch)` fixed-address slots plus a
+  `slot_map: Vec<Option<RequestId>>`. On a request's first decode row the
+  worker D2D-copies its prefill recurrent state into the slot
+  (`copy_state_to_slot`) and drops the per-request allocation.
+- `DropRequest` carries `compaction: Option<TpSlotCompaction>`; the worker
+  validates occupancy against `slot_map` (`slot_compact`), applies the D2D
+  move (`BatchDecodeGraphState::move_slot_within`), and poisons on mismatch.
+  Requests retired between promotion and their first decode row legitimately
+  have no materialized slot; `slot_compact` tolerates exactly that case and
+  skips the GPU move.
+- Convenience `execute_prefill`/`execute_decode`/`drop_request` (model-local
+  tests) keep a Mutex-guarded slot tracker mirroring the single-GPU
+  `Qwen35Executor`; scheduler flows pass explicit slots via
+  `execute_decode_items`/`drop_request_with_compaction` and never touch the
+  tracker.
+
+**Capture/replay.** Startup pre-capture sweep ported from qwen3:
+`TpWorkerCommand::Precapture { phase }` over Warmup (new
+`Qwen35Model::warmup_tp_collective` — one all-reduce per bucket message size;
+without it the lazy NCCL connect wedges inside capture), Capture + Launch per
+bucket `[1,2,4,8,16,32,64]` up to `bucket_for(max_batch)` with synthetic rows,
+Finalize asserting all buckets captured. Dedicated 600 s abort watchdog (the
+60 s NCCL startup timeout is too small for the sweep).
+`batch_decode_graph` gained `DecodeGraphUse` (Serve lazy / CaptureOnly /
+Replay); TP serve time is Replay-only. Workers tune decode GEMM algos on the
+worker thread before capture (cuBLASLt plans are thread-local). Graph state
+is declared before `model` in `TpWorkerState` so graphs drop before the NCCL
+comm. Sampling/logprobs stay rank-0 host-side outside the graph. Eager
+workers ignore `slot_idx`/`compaction`, keeping 27B TP2 byte-identical.
+
+**Gotcha fixed during validation:** `track_retired_slot` used
+`bool::then_some`, which evaluates eagerly and indexed past the tail when the
+retired request was the last slot — use `then` for the lazy closure.
+
+**Validation (2× RTX 4090):**
+
+- `cargo check --release -p pegainfer-qwen35 --features qwen35` clean; lib
+  suite 105/105 (new: group-gate acceptance incl. group-6 stays eager, slot
+  map admit/compact/mismatch CPU tests); `cargo fmt --check` clean.
+- 9B TP2 HF gates (`--test-threads=1`): eager sequential+batched PASS;
+  graph sequential replay (identical fingerprints across reruns),
+  bucket-straddling batched replay (5→bucket 8, 3→bucket 4), and
+  post-compaction replay after a mid-batch drop all within the existing TP2
+  tolerances (worst graph arm: mean 0.0228, p99 0.1002 against MEAN_TOL 0.06
+  / P99_TOL 0.20; eager arm mean 0.0227/0.0230).
+- 9B TP2 scheduler e2e eager + graph variants PASS; 9B TP2 serving smoke now
+  launches with `--cuda-graph true` (graph acceptance replaced the old
+  rejection assertion).
+- 27B TP2 HF short+long and scheduler e2e PASS unchanged — the gate log
+  confirms group 6 keeps decode eager (the graph test variant skips itself
+  via `graph_enabled()`).
+- 9B TP2 serving benchmark (`pegainfer-server --tp-size 2 --port 18093`,
+  vllm-bench `openai` backend, random 128-in/256-out, 64 prompts at
+  concurrency 16, greedy, ignore_eos, seed 42):
+
+  | arm | steady output tok/s | mean TPOT (ms) | total tok/s |
+  |-----|--------------------:|---------------:|------------:|
+  | CUDA Graph on  | 767.15 | 20.04 | 1146.65 |
+  | CUDA Graph off | 705.86 | 21.99 | 1054.56 |
+
+  Graph decode is +8.7% steady output tok/s (-8.8% TPOT) at 16 concurrent.
+  The design's "record in `bench_snapshots/`" step was skipped: the
+  in-process snapshot gate is retired (`docs/conventions/bench-regression.md`),
+  so the HTTP bench numbers live here instead.
+
+**Test-isolation note:** TP2 GPU tests must run with `--test-threads=1`. Two
+TP executors sharing the GPUs perturb cuBLASLt algorithm selection (workspace
+pressure), which once flipped a sequential-replay fingerprint comparison in
+the *eager* test while the graph test ran concurrently.
+
 ## Follow-Ups
 
-- Design and implement P2B sharded linear-attention/GDR state without weakening the completed P2A lifecycle and ID contracts.
+- P2c CUDA Graph under TP landed for compiled decode GQA groups (4B/9B); 27B
+  TP2 graphs stay gated off until group-6 batch-decode kernels are compiled
+  (`SUPPORTED_GQA_GROUP_SIZES`). The eager path is the 27B fallback and must
+  not regress.
+- 27B TP2 knowledge-benchmark parity (2026-08-20, validated pre-rebase on
+  the parallel TP line; `docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md`):
+  MMLU-Redux 94.09 vs official 93.2 (full 5330), C-Eval 88.11 vs 90.5
+  (full 1346, thinking-cap truncation rerun-merged) — inside the
+  cross-harness band, no TP-induced accuracy regression. MMLU-Pro /
+  SuperGPQA sampled runs remain outstanding; rerun on this rebased branch
+  before citing parity.
+- P2B sharded linear-attention/GDR state landed (see "Rebase onto #870"); keep the completed P2A lifecycle and ID contracts unweakened.
 - Promote any stable contract changes discovered here back into `tp-design.md` through the design-doc branch.
 - Decide whether Qwen3.5 server CLI should accept arbitrary TP device ordinals instead of only `0..tp_size`.
 - Consider lifting the per-device Triton AOT handle lesson into a kernels or runtime subsystem doc if another model hits the same issue.