diff --git a/docs/index.md b/docs/index.md index f5325870f..31d61e9d1 100644 --- a/docs/index.md +++ b/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, P2A, and P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a post-`out_proj` hidden all-reduce; TP decode rows run as one batched eager forward per step (#1004); TP CUDA Graph (#1005) 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. | diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index 27d90b170..117214921 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -1,8 +1,8 @@ # Qwen3.5 Tensor Parallelism Design -> **TL;DR:** Qwen3.5 TP Phase 2 is two separately delivered correctness milestones: P2a adds eager `RunUnifiedStep` with a shared ordered `RequestId` plan while retaining Phase 1 replicated GDR; P2b shards the head-indexed linear-attention/GDR surface and adds only the hidden all-reduce after local `out_proj`. +> **TL;DR:** Qwen3.5 TP Phase 2 is two separately delivered correctness milestones: P2a adds eager `RunUnifiedStep` with a shared ordered `RequestId` plan while retaining Phase 1 replicated GDR; P2b shards the head-indexed linear-attention/GDR surface and adds only the hidden all-reduce after local `out_proj`. P2c adds decode CUDA Graphs under TP, gated on the compiled decode GQA group (4B/9B TP2 capture; 27B group-6 stays eager). > -> **Last touched:** 2026-08 +> **Last touched:** 2026-09 ## Goal @@ -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 + +Landed 2026-08-20 (#1005), gated on `Config35::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.decode_group_is_compiled()`. 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` + 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` diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index eb452f767..a5e449a38 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -1,6 +1,6 @@ # Qwen3.5 TP Implementation Record -> **TL;DR:** Qwen3.5 TP Phase 1, P2A, and the P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a single post-`out_proj` hidden all-reduce, and TP decode rows run as one batched eager forward per step (#1004); TP CUDA Graph under TP 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-09 @@ -502,9 +502,169 @@ Not in this step: TP CUDA Graph capture (#1005) and the same-context per-request-vs-batched throughput A/B tracked in #1001; no performance claim is made before that rerun lands. +## 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/tp.rs` `local_linear_*` + accessors, `weights/layers.rs` per-rank stitched qkv/conv1d + row/col + shards carried on `WeightSource`, + `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, one persistent + `LinearStatePointerTables` per worker refilled each step — see Step 3 + above — 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, +`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. #967/#968 have since split the +scheduler into `scheduler/{mod,backend,plan,tp}.rs`, so this branch's +slot-tracking deltas land there: the `ActiveBackendState::Tp` slot field +and dispatch arms in `mod.rs`, `TpSchedulerBackend` slot compaction in +`backend.rs`, and the decode-item/alignment helpers in `tp.rs`. + +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.decode_group_is_compiled()`. 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>`. 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`; 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()`). +- **Historical, not re-measured on this stack:** 9B TP2 serving benchmark + from the pre-rebase #946 branch (2026-08-20; `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 | + + On that branch graph decode measured +8.7% steady output tok/s (-8.8% TPOT) + at 16 concurrent. No performance claim is made for this stack until the + same-context A/B is rerun on it (tracked in #1001). + 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 -- Land TP CUDA Graph (#1005) on top of batched eager TP decode; rerun the #946 throughput A/B (per-request vs batched TP decode) on the merged stack before any performance claim. +- 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. diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index f294000f0..e8565198e 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -26,6 +26,22 @@ use crate::ops; static LOG_UNCOMPILED_DECODE_ROUTE: std::sync::Once = std::sync::Once::new(); +/// How a `batch_decode_graph` call interacts with the per-bucket CUDA graphs. +/// +/// TP serving never captures lazily: a mid-serving capture on one rank while a +/// peer replays desyncs the recorded NCCL collectives, so tensor-parallel +/// decodes are replay-only after the startup pre-capture sweep. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum DecodeGraphUse { + /// Replay if captured, lazily capture otherwise (single-GPU serving). + Serve, + /// Record + instantiate + upload, no launch (the TP sweep's Capture phase). + CaptureOnly, + /// Replay only; error if never captured (TP serving, and the TP sweep's + /// Launch phase that drains the captured collectives across ranks). + Replay, +} + impl Qwen35Model { pub(crate) fn select_tokens_from_logits_varied( &self, @@ -340,6 +356,25 @@ impl Qwen35Model { token_ids: &[u32], kv_states: &mut [&mut KvState], graph_state: &mut BatchDecodeGraphState, + graph_use: DecodeGraphUse, + ) -> Result<()> { + let padded_bs = bucket_for(token_ids.len()); + self.batch_decode_graph_padded(token_ids, kv_states, graph_state, graph_use, padded_bs) + } + + /// `batch_decode_graph` with the bucket chosen by the caller instead of + /// derived from `bs`. Rows `bs..padded_bs` are padding either way — they + /// ride the pool's reserved padding page and a free recurrent slot — so a + /// caller that only needs a *bucket* (the TP pre-capture sweep) can pass + /// one real row and still capture or launch the bucket-`padded_bs` graph + /// without holding `padded_bs` KV pages. + pub(crate) fn batch_decode_graph_padded( + &self, + token_ids: &[u32], + kv_states: &mut [&mut KvState], + graph_state: &mut BatchDecodeGraphState, + graph_use: DecodeGraphUse, + padded_bs: usize, ) -> Result<()> { let bs = token_ids.len(); anyhow::ensure!(bs > 0, "batch_decode_graph requires at least one request"); @@ -349,8 +384,16 @@ impl Qwen35Model { "batch size {bs} exceeds decode capacity {}", graph_state.slot_states.len() ); + anyhow::ensure!( + padded_bs >= bs && BATCH_BUCKETS.contains(&padded_bs), + "padded batch {padded_bs} is not a decode bucket covering bs={bs}" + ); if !self.config.decode_group_is_compiled() { + anyhow::ensure!( + graph_use == DecodeGraphUse::Serve, + "Qwen3.5 batched hybrid eager fallback only supports lazy serve-mode decode, got {graph_use:?}" + ); LOG_UNCOMPILED_DECODE_ROUTE.call_once(|| { let group = self.config.num_attention_heads / self.config.num_key_value_heads; log::info!( @@ -365,7 +408,6 @@ impl Qwen35Model { return self.batch_decode_batched_hybrid(token_ids, kv_states, graph_state); } - let padded_bs = bucket_for(bs); graph_state.linear_pointer_tables.validate_for( &self.config, padded_bs, @@ -411,17 +453,35 @@ impl Qwen35Model { let mut graphs = std::mem::take(&mut graph_state.graphs); let linear_state_ptrs = &graph_state.linear_pointer_tables.state_ptrs; let linear_conv_state_ptrs = &graph_state.linear_pointer_tables.conv_state_ptrs; - let result = graphs[bucket_idx].run_or_capture(&self.ctx, || { - self.batch_decode_kernels_graph( - kv_buffer, - &layout, - padded_bs, - None, - linear_state_ptrs, - linear_conv_state_ptrs, - &mut graph_state.buffers, - ) - }); + let result = match graph_use { + DecodeGraphUse::Serve => graphs[bucket_idx].run_or_capture(&self.ctx, || { + self.batch_decode_kernels_graph( + kv_buffer, + &layout, + padded_bs, + None, + linear_state_ptrs, + linear_conv_state_ptrs, + &mut graph_state.buffers, + ) + }), + DecodeGraphUse::CaptureOnly => graphs[bucket_idx].capture_only(&self.ctx, || { + self.batch_decode_kernels_graph( + kv_buffer, + &layout, + padded_bs, + None, + linear_state_ptrs, + linear_conv_state_ptrs, + &mut graph_state.buffers, + ) + }), + // Replay is a pure enqueue: every bucket was recorded by the + // startup pre-capture sweep, so a missing graph here means the + // sweep was skipped or incomplete — fail loudly, never capture + // mid-serving (a one-sided capture desyncs TP collectives). + DecodeGraphUse::Replay => graphs[bucket_idx].launch_captured(&self.ctx), + }; graph_state.graphs = graphs; result } diff --git a/pegainfer-qwen35/src/batch_decode_graph.rs b/pegainfer-qwen35/src/batch_decode_graph.rs index 59d5225d3..e68348ea1 100644 --- a/pegainfer-qwen35/src/batch_decode_graph.rs +++ b/pegainfer-qwen35/src/batch_decode_graph.rs @@ -128,4 +128,52 @@ impl BatchDecodeGraphState { dst.seq_len = src.seq_len; Ok(()) } + + /// D2D move slot `from` into slot `to`, leaving `to` as the canonical + /// state. Used by TP slot compaction when a mid-batch request retires: + /// the last occupied slot moves into the vacated slot so decode rows stay + /// dense (`0..bs`). Slot `from` keeps stale bytes afterwards; admission + /// overwrites it (`copy_state_to_slot`), and padding rows may clobber it — + /// both benign by design. + pub(crate) fn move_slot_within( + &mut self, + ctx: &DeviceContext, + from: usize, + to: usize, + ) -> Result<()> { + anyhow::ensure!( + from != to, + "TP slot compaction move {from} -> {to} is a no-op" + ); + anyhow::ensure!( + from < self.slot_states.len() && to < self.slot_states.len(), + "TP slot compaction move {from} -> {to} exceeds {} slots", + self.slot_states.len() + ); + let (lo, hi) = (from.min(to), from.max(to)); + let (left, right) = self.slot_states.split_at_mut(hi); + let (src, dst) = if from < to { + (&left[from], &mut right[0]) + } else { + (&right[0], &mut left[lo]) + }; + anyhow::ensure!( + src.layers.len() == dst.layers.len(), + "TP slot compaction layer count mismatch: {} vs {}", + src.layers.len(), + dst.layers.len() + ); + for (src_layer, dst_layer) in src.layers.iter().zip(dst.layers.iter_mut()) { + ctx.stream + .memcpy_dtod(&src_layer.state, &mut dst_layer.state) + .map_err(|e| anyhow::anyhow!("move recurrent state slot {from} -> {to}: {e}"))?; + ctx.stream + .memcpy_dtod(&src_layer.conv_state.data, &mut dst_layer.conv_state.data) + .map_err(|e| anyhow::anyhow!("move conv state slot {from} -> {to}: {e}"))?; + } + let seq_len = src.seq_len; + let dst = &mut self.slot_states[to]; + dst.seq_len = seq_len; + Ok(()) + } } diff --git a/pegainfer-qwen35/src/config/error.rs b/pegainfer-qwen35/src/config/error.rs index b44a5b996..3b8dbfbb1 100644 --- a/pegainfer-qwen35/src/config/error.rs +++ b/pegainfer-qwen35/src/config/error.rs @@ -68,10 +68,6 @@ pub(crate) enum ConfigError { TpZeroWorldSize, #[error("tensor_parallel.rank {rank} must be < world_size {world_size}")] TpRankOutOfRange { rank: usize, world_size: usize }, - #[error( - "Qwen3.5 tensor parallelism is eager-only; disable CUDA Graph for tp world_size={world_size}" - )] - TpRequiresEager { world_size: usize }, #[error("{field}={value} not divisible by tp world_size={world_size}")] TpIndivisible { field: &'static str, diff --git a/pegainfer-qwen35/src/config/tp.rs b/pegainfer-qwen35/src/config/tp.rs index bb58d8e79..e63b2eb21 100644 --- a/pegainfer-qwen35/src/config/tp.rs +++ b/pegainfer-qwen35/src/config/tp.rs @@ -90,26 +90,22 @@ pub(crate) struct LocalGeometry { } impl LocalGeometry { - /// Validate `config` against `tp` and the runtime execution mode, then derive - /// this rank's local dimensions. + /// Validate `config` against `tp`, then derive this rank's local dimensions. /// /// Fails on unsupported combinations before expensive loading: - /// - sharded TP demands eager execution (`enable_cuda_graph` off); /// - every sharded model dimension must divide evenly by `world_size` /// (linear-attention key heads included; the value-head count follows /// from the `Config35` key/value-head invariant); /// - `rank < world_size` and `world_size >= 1` are guaranteed by /// `TensorParallelConfig::try_from`. + /// + /// CUDA Graph under TP is gated at executor startup on + /// [`Config35::decode_group_is_compiled`] (P2c): uncompiled GQA + /// groups keep the batched eager path instead of failing validation here. pub(crate) fn try_new( config: &Config35, tp: TensorParallelConfig, - enable_cuda_graph: bool, ) -> Result { - if tp.is_sharded() && enable_cuda_graph { - return Err(ConfigError::TpRequiresEager { - world_size: tp.world_size(), - }); - } if !config.num_attention_heads.is_multiple_of(tp.world_size()) { return Err(ConfigError::TpIndivisible { field: "num_attention_heads", @@ -279,7 +275,7 @@ mod tests { fn tp2_local_geometry_matches_dense_dims() { let cfg = config(); let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); - let geom = LocalGeometry::try_new(&cfg, tp, false).unwrap(); + let geom = LocalGeometry::try_new(&cfg, tp).unwrap(); assert!(geom.is_sharded()); assert_eq!(geom.shard_range(4096), (2048, 2048)); assert_eq!(geom.local_num_attention_heads(), 8); @@ -309,7 +305,7 @@ mod tests { fn rejects_indivisible_dense_dimensions() { let tp = TensorParallelConfig::try_from((0, 3)).unwrap(); let cfg = config(); - let mut err = LocalGeometry::try_new(&cfg, tp, false).unwrap_err(); + let mut err = LocalGeometry::try_new(&cfg, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -322,7 +318,7 @@ mod tests { let mut broken = cfg; broken.num_attention_heads = 15; broken.num_key_value_heads = 4; - err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -334,7 +330,7 @@ mod tests { broken.num_key_value_heads = 3; broken.intermediate_size = 9217; - err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -345,16 +341,6 @@ mod tests { ); } - #[test] - fn rejects_tensor_parallel_with_cuda_graph() { - let cfg = config(); - let tp = TensorParallelConfig::try_from((0, 2)).unwrap(); - assert_eq!( - LocalGeometry::try_new(&cfg, tp, true), - Err(ConfigError::TpRequiresEager { world_size: 2 }) - ); - } - #[test] fn requires_linear_attention_key_head_divisibility() { let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); @@ -363,7 +349,7 @@ mod tests { // branch is the key-head TP guard itself. broken.linear_num_key_heads = 17; broken.linear_num_value_heads = 34; - let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + let err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { diff --git a/pegainfer-qwen35/src/executor.rs b/pegainfer-qwen35/src/executor.rs index 956d0f9ab..67bf5568e 100644 --- a/pegainfer-qwen35/src/executor.rs +++ b/pegainfer-qwen35/src/executor.rs @@ -232,8 +232,12 @@ impl Qwen35Executor { let token_ids: Vec = plan.requests.iter().map(|req| req.token_id).collect(); let mut kv_refs: Vec<&mut KvState> = self.active.iter_mut().map(|req| &mut req.kv).collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state)?; + self.model.batch_decode_graph( + &token_ids, + &mut kv_refs, + &mut self.graph_state, + crate::batch_decode::DecodeGraphUse::Serve, + )?; let requested_logprobs: Vec = plan.requests.iter().map(|req| req.logprobs).collect(); let cpu_logits = snapshot_requested_logprobs( diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index f46db580c..9f02fc00d 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -98,6 +98,9 @@ pub enum Qwen35DecodeOverlap { SharedSm, } +/// TP decode runs CUDA Graphs when `cuda_graph` is set AND the rank-local +/// decode GQA group has a compiled kernel (`tp-design.md` P2c gate); otherwise +/// TP keeps the batched eager path. pub fn start_engine( model_path: &Path, options: EngineLoadOptions, @@ -119,7 +122,8 @@ pub struct Qwen35LaunchOptions { device_ordinal: usize, /// Tensor-parallel world size; `> 1` uses devices `0..tp_size`. tp_size: usize, - /// TP Phase 1 supports eager-only multi-GPU execution. + /// TP decode captures CUDA Graphs when the rank-local decode GQA group has + /// a compiled kernel; uncompiled groups (27B) stay on the eager path. cuda_graph: bool, max_batch: usize, max_prefill_tokens: usize, @@ -231,11 +235,6 @@ pub fn start_engine_with_capacity_policy_and_overlap( "Qwen3.5 TP uses the fixed off scheduler policy; --qwen35-scheduler-policy=auto is single-GPU only" )); } - if enable_cuda_graph { - return Err(anyhow!( - "Qwen3.5 TP Phase 1 supports eager execution only; disable CUDA Graph" - )); - } let model_path = model_path .to_str() .ok_or_else(|| anyhow!("model path must be valid UTF-8"))?; @@ -245,6 +244,7 @@ pub fn start_engine_with_capacity_policy_and_overlap( &device_ordinals, max_batch, max_prefill_tokens, + enable_cuda_graph, ); } @@ -309,6 +309,31 @@ mod tests { assert_eq!(Qwen35SchedulerPolicy::default(), Qwen35SchedulerPolicy::Off); } + #[test] + fn tp_cuda_graph_is_no_longer_rejected_before_model_load() { + // P2c: the graph/eager decision is gated on the model's TP-local decode + // GQA group, which needs the loaded config — so TP + CUDA Graph passes + // pre-load validation and fails here only because the path is bogus. + let err = start_engine_with_capacity_and_policy( + Path::new("unused-model-path"), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0, 1], + parallel_config: None, + ep_backend: EpBackend::Nccl, + seed: 42, + }, + 1, + 1, + Qwen35SchedulerPolicy::Off, + ) + .err() + .expect("nonexistent model path should fail at load") + .to_string(); + + assert!(!err.contains("eager execution only")); + } + #[test] fn tp_rejects_auto_scheduler_policy_before_loading_model() { let err = start_engine_with_capacity_and_policy( diff --git a/pegainfer-qwen35/src/scheduler/backend.rs b/pegainfer-qwen35/src/scheduler/backend.rs index 819f1cfef..6b0709f6b 100644 --- a/pegainfer-qwen35/src/scheduler/backend.rs +++ b/pegainfer-qwen35/src/scheduler/backend.rs @@ -115,8 +115,11 @@ fn attached_logprobs( } pub(super) struct TpSchedulerBackend { - executor: Qwen35TpExecutor, + pub(super) executor: Qwen35TpExecutor, next_request_id: u64, + /// Slot move derived by the in-flight `take_active_request`; consumed by + /// the paired `drop_active_state` so the workers apply the same move. + pub(super) pending_compaction: Option, } impl SingleGpuBackend { @@ -265,8 +268,12 @@ impl SingleGpuBackend { pub(super) fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { let (tokens, mut kvs) = single_decode_views(active); - self.model - .batch_decode_graph(&tokens, &mut kvs, &mut self.graph_state) + self.model.batch_decode_graph( + &tokens, + &mut kvs, + &mut self.graph_state, + crate::batch_decode::DecodeGraphUse::Serve, + ) } pub(super) fn sample_prefill_logits( @@ -386,10 +393,11 @@ impl TpSchedulerBackend { device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + enable_cuda_graph: bool, ) -> Result { let executor = Qwen35TpExecutor::from_runtime_with_limits( model_path, - false, + enable_cuda_graph, device_ordinals, max_batch, max_prefill_tokens, @@ -397,6 +405,7 @@ impl TpSchedulerBackend { Ok(Self { executor, next_request_id: 1, + pending_compaction: None, }) } @@ -498,7 +507,40 @@ impl TpSchedulerBackend { request_id: RequestId, expectation: DropExpectation, ) -> Result<()> { - self.executor.drop_request(request_id, expectation) + self.executor + .drop_request_with_compaction(request_id, expectation, None) + } + + /// Remove the TP request at `idx` via swap_remove and stash the resulting + /// slot compaction for the paired `drop_active_state`. Mirrors + /// `compact_single_slot`: after the swap, slots `0..active.len()` stay + /// dense because the moved request's slot follows it. + pub(super) fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + let compaction = compaction_after_retire(active.len(), idx); + let removed = active.swap_remove(idx); + + self.pending_compaction = compaction.map(|compaction| { + let moved = &mut active[idx]; + let ActiveBackendState::Tp { + request_id, + slot_idx, + } = &mut moved.backend_state + else { + panic!("TP scheduler received single-GPU active state") + }; + debug_assert_eq!(*slot_idx, compaction.moved_from); + *slot_idx = compaction.moved_to; + TpSlotCompaction { + moved_request_id: *request_id, + from: compaction.moved_from, + to: compaction.moved_to, + } + }); + removed } } diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs index f8531998e..3efb497b6 100644 --- a/pegainfer-qwen35/src/scheduler/mod.rs +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -72,6 +72,7 @@ use crate::tp_executor::DropExpectation; use crate::tp_executor::Qwen35TpExecutor; use crate::tp_executor::TpDecodeStepItem; use crate::tp_executor::TpPrefillChunkItem; +use crate::tp_executor::TpSlotCompaction; use crate::tp_executor::TpUnifiedPlan; use crate::weights::Qwen35Model; @@ -112,6 +113,9 @@ enum ActiveBackendState { }, Tp { request_id: RequestId, + /// Dense decode slot (`active` position). Graph-mode workers assert + /// `slot_idx == row` on every decode command; eager workers ignore it. + slot_idx: usize, }, } @@ -421,13 +425,19 @@ pub(crate) fn start_tp_with_capacity( device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + enable_cuda_graph: bool, ) -> Result { assert!( max_prefill_tokens > 0, "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" ); - let backend = - TpSchedulerBackend::new(model_path, device_ordinals, max_batch, max_prefill_tokens)?; + let backend = TpSchedulerBackend::new( + model_path, + device_ordinals, + max_batch, + max_prefill_tokens, + enable_cuda_graph, + )?; let servable = servable_len( backend.max_position_embeddings(), backend.capacity_pages_for_requests(), @@ -1488,15 +1498,20 @@ impl DecodeDispatchBackend for SchedulerBackend { ) -> ActiveRequest35 { match self { SchedulerBackend::Single(backend) => compact_single_slot(backend, active, idx), - SchedulerBackend::Tp(_) => active.swap_remove(idx), + SchedulerBackend::Tp(backend) => backend.take_active_request(active, idx), } } fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { match (self, state) { (SchedulerBackend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), - (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id }) => { - backend.drop_request(*request_id, DropExpectation::MustExist) + (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id, .. }) => { + let compaction = backend.pending_compaction.take(); + backend.executor.drop_request_with_compaction( + *request_id, + DropExpectation::MustExist, + compaction, + ) } _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), } @@ -1854,8 +1869,13 @@ impl PrefillPromoteBackend for SchedulerBackend { (SchedulerBackend::Single(single), state @ PrefillBackendState::Single { .. }) => { single.promote_prefill_state(active_len, state) } - (SchedulerBackend::Tp(_), PrefillBackendState::Tp { request_id }) => { - ActiveBackendState::Tp { request_id } + (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + let slot_idx = slot_for_new_request(active_len, backend.max_batch()) + .expect("admission must reserve a TP decode slot"); + ActiveBackendState::Tp { + request_id, + slot_idx, + } } _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), } diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 634f108b8..0aa92ebc4 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -44,6 +44,7 @@ fn active_request(request_id: u64, label: &str, token_tx: TokenSink) -> ActiveRe token_tx, backend_state: ActiveBackendState::Tp { request_id: RequestId::new(request_id), + slot_idx: 0, }, last_token: 1, generated_count: 1, @@ -89,7 +90,7 @@ impl DecodeDispatchBackend for PruneTestBackend { } fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id } = state else { + let ActiveBackendState::Tp { request_id, .. } = state else { panic!("prune test expected TP active state"); }; self.retired_active.push(*request_id); @@ -180,7 +181,7 @@ impl DecodeDispatchBackend for LifecycleTestBackend { } fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id } = state else { + let ActiveBackendState::Tp { request_id, .. } = state else { panic!("lifecycle test expected TP active state"); }; if self.active_completion_requires_drop_ack { @@ -880,7 +881,10 @@ fn submit_parking_requires_idle_owned_work_and_no_inflight_prefill() { } #[test] -fn tp_engine_rejects_cuda_graph_before_model_load() { +fn tp_engine_cuda_graph_passes_preload_validation() { + // P2c: TP + CUDA Graph is gated on the model's TP-local decode GQA group + // after load, so startup with a bogus path fails at load, not at the old + // eager-only rejection. let err = match crate::start_engine_with_capacity( Path::new("unused"), EngineLoadOptions { @@ -893,10 +897,10 @@ fn tp_engine_rejects_cuda_graph_before_model_load() { 1, 1, ) { - Ok(_) => panic!("TP CUDA Graph startup should fail"), + Ok(_) => panic!("TP CUDA Graph startup with a nonexistent path should fail at load"), Err(err) => err.to_string(), }; - assert!(err.contains("eager execution only")); + assert!(!err.contains("eager execution only")); } #[test] @@ -908,7 +912,7 @@ fn tp2_scheduler_runs_forced_mixed_steps() { return; }; let handle = - start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1).expect("start TP2 scheduler"); + start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1, false).expect("start TP2 scheduler"); let (decode_tx, mut decode_rx) = TokenSink::standalone(); let (prefill_tx, mut prefill_rx) = TokenSink::standalone(); diff --git a/pegainfer-qwen35/src/scheduler/tp.rs b/pegainfer-qwen35/src/scheduler/tp.rs index d33b94e52..9b636256d 100644 --- a/pegainfer-qwen35/src/scheduler/tp.rs +++ b/pegainfer-qwen35/src/scheduler/tp.rs @@ -34,15 +34,25 @@ pub(super) fn tp_prefill_items(chunk: &ScheduledChunk) -> Result Result> { active .iter() - .map(|req| { - let ActiveBackendState::Tp { request_id } = &req.backend_state else { + .enumerate() + .map(|(row, req)| { + let ActiveBackendState::Tp { + request_id, + slot_idx, + } = &req.backend_state + else { anyhow::bail!("TP decode received single-GPU active state"); }; - Ok(TpDecodeStepItem::new( + debug_assert_eq!( + *slot_idx, row, + "TP decode slots must stay dense in active order" + ); + Ok(TpDecodeStepItem::new_with_slot( *request_id, req.last_token, req.logprobs, req.params, + *slot_idx, )) }) .collect() @@ -120,7 +130,7 @@ pub(super) fn align_decode_results( let expected: Vec = active .iter() .map(|active_req| { - let ActiveBackendState::Tp { request_id } = active_req.backend_state else { + let ActiveBackendState::Tp { request_id, .. } = active_req.backend_state else { anyhow::bail!("align_decode_results requires TP active state"); }; Ok(request_id) diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 7584edacf..613c67100 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -16,16 +16,22 @@ use std::sync::atomic::Ordering; use std::sync::mpsc; use std::thread::JoinHandle; use std::thread::{self}; +use std::time::Instant; use anyhow::Result; use pegainfer_core::kv_pool::KvState; use pegainfer_frontend::sampler::SamplingParams; +use crate::batch_decode::DecodeGraphUse; +use crate::batch_decode_graph::BATCH_BUCKETS; +use crate::batch_decode_graph::BatchDecodeGraphState; +use crate::batch_decode_graph::bucket_for; use crate::config::TensorParallelConfig; use crate::decode_buffers::BatchDecodeBuffers35; use crate::executor::DecodePlan; use crate::executor::DecodeRequestResult; use crate::executor::DecodeResult; +#[cfg(test)] use crate::executor::DecodeStepItem; use crate::executor::PrefillPlan; use crate::executor::PrefillRequestResult; @@ -43,9 +49,43 @@ use crate::weights::Qwen35Model; const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); const TP_RUNTIME_STEP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); const TP_WORKER_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +/// The pre-capture sweep records every decode bucket per rank; the 60 s NCCL +/// startup budget is far too small for that (qwen3 uses the same 600 s). +const TP_PRECAPTURE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600); const TP_RUNTIME_MEMORY_RESERVE_BYTES: usize = 512 * 1024 * 1024; const TRITON_AOT_DEVICE_TABLE_LEN: usize = 16; +/// One controller-barriered phase of the TP decode-graph pre-capture sweep. +/// +/// Capture and launch are separate phases because a captured collective's +/// first launch blocks on its peers: overlapping that with a peer still in +/// capture/instantiate/upload (which contend driver locks and allocate device +/// memory) deadlocks the driver. So every rank finishes capturing a bucket +/// before any rank launches it. (Ported from qwen3's TP sweep.) +#[derive(Clone, Copy, Debug)] +enum PrecapturePhase { + /// One eager all-reduce per bucket message size, so the size-selected NCCL + /// algorithm connects before any `cuStreamBeginCapture` records it. + Warmup, + /// Record + instantiate + upload one bucket; no launch, no cross-rank dependency. + Capture { bucket_idx: usize }, + /// Launch one bucket (pure enqueue after `Capture`) + sync; collectives pair across ranks. + Launch { bucket_idx: usize }, + /// Verify every reachable bucket captured. + Finalize, +} + +/// Scheduler-owned slot move for TP graph decode: when the request at slot +/// `to` retires mid-batch, the request at slot `from` (the last occupied slot) +/// takes over slot `to` so decode rows stay dense. Workers apply the move and +/// fail (poisoning the executor) if slot occupancy does not match. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct TpSlotCompaction { + pub(crate) moved_request_id: RequestId, + pub(crate) from: usize, + pub(crate) to: usize, +} + #[allow(dead_code)] enum TpWorkerCommand { Ping { @@ -70,6 +110,17 @@ enum TpWorkerCommand { }, DropRequest { request_id: RequestId, + /// Slot move the scheduler already applied to its own bookkeeping; + /// `Some` only when the dropped request held a decode slot that a + /// still-active request now takes over. Eager workers ignore it. + compaction: Option, + start: Arc, + resp: mpsc::Sender, + }, + /// Startup-only (graph-enabled TP): one phase of the decode-graph + /// pre-capture sweep, barriered across ranks by the controller. + Precapture { + phase: PrecapturePhase, start: Arc, resp: mpsc::Sender, }, @@ -195,6 +246,18 @@ pub struct Qwen35TpExecutor { capacity_pages_for_requests: usize, max_position_embeddings: usize, eos_token_id: u32, + /// Whether decode steps replay pre-captured CUDA Graphs that record NCCL + /// collectives (`enable_cuda_graph` AND a compiled TP-local decode GQA + /// group; see the P2c gate in `tp-design.md`). + graph_enabled: bool, + /// Slot tracker for the convenience `execute_prefill`/`execute_decode`/ + /// `drop_request` API (model-local tests), mirroring the single-GPU + /// `Qwen35Executor`: prefill completion appends, decode plans must cover + /// every tracked request in slot order, drop swap-removes and derives the + /// slot compaction. Scheduler-driven flows bypass it entirely — they pass + /// explicit slots via `execute_decode_items` and + /// `drop_request_with_compaction`. Never mix the two flows on one executor. + active_slots: Mutex>, } #[derive(Clone)] @@ -245,6 +308,11 @@ pub(crate) struct TpDecodeStepItem { token_id: u32, logprobs: usize, sampling_params: SamplingParams, + /// Scheduler-assigned decode slot under CUDA Graph TP. Rows must arrive in + /// dense slot order (`slot_idx == row`); on the request's first decode row + /// the worker D2D-copies its prefill recurrent state into the slot and + /// drops the per-request allocation. `None` on the slot-free eager path. + slot_idx: Option, } impl TpDecodeStepItem { @@ -259,6 +327,20 @@ impl TpDecodeStepItem { token_id, logprobs, sampling_params, + slot_idx: None, + } + } + + pub(crate) fn new_with_slot( + request_id: RequestId, + token_id: u32, + logprobs: usize, + sampling_params: SamplingParams, + slot_idx: usize, + ) -> Self { + Self { + slot_idx: Some(slot_idx), + ..Self::new(request_id, token_id, logprobs, sampling_params) } } } @@ -306,10 +388,6 @@ impl Qwen35TpExecutor { "Qwen3.5 TP executor requires at least two CUDA devices, got {}", device_ordinals.len() ); - anyhow::ensure!( - !enable_cuda_graph, - "Qwen3.5 TP Phase 1 supports eager execution only; disable CUDA Graph" - ); anyhow::ensure!( max_prefill_tokens > 0, "Qwen3.5 TP max_prefill_tokens must be positive" @@ -321,7 +399,7 @@ impl Qwen35TpExecutor { models.push(Qwen35Model::from_safetensors_with_runtime( model_path, ModelRuntimeConfig { - enable_cuda_graph: false, + enable_cuda_graph, tensor_parallel: Some(TensorParallelConfig::try_from((rank, world_size))?), device_ordinal, }, @@ -330,6 +408,23 @@ impl Qwen35TpExecutor { let first = models .first() .ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP executor loaded no models"))?; + // P2c gate: graph decode under TP requires a compiled batch-decode + // kernel for the rank-local GQA group. The group ratio is TP-invariant, + // so every rank decides identically; an uncompiled group (e.g. 27B's + // group 6) keeps the batched eager path byte-for-byte. + let geometry = first.geometry; + let graph_enabled = enable_cuda_graph && first.config().decode_group_is_compiled(); + if enable_cuda_graph && !graph_enabled { + static LOG_GRAPH_GATE: std::sync::Once = std::sync::Once::new(); + LOG_GRAPH_GATE.call_once(|| { + log::info!( + "Qwen3.5 TP decode GQA group {} ({} q heads / {} kv heads per rank) has no compiled batch-decode kernel; CUDA Graph requested but decode stays on the batched eager path", + geometry.local_num_attention_heads() / geometry.local_num_key_value_heads(), + geometry.local_num_attention_heads(), + geometry.local_num_key_value_heads(), + ); + }); + } let page_size = first.kv_pool().layout().page_size; let mut min_capacity_pages = usize::MAX; for (rank, model) in models.iter().enumerate() { @@ -359,6 +454,7 @@ impl Qwen35TpExecutor { model, max_batch, max_prefill_tokens, + graph_enabled, nccl_id, Arc::clone(&startup_gate), Arc::clone(&effective_max_batch), @@ -426,7 +522,7 @@ impl Qwen35TpExecutor { } disarm_nccl_startup_watchdog(watchdog_done, watchdog)?; - Ok(Self { + let executor = Self { workers, poison, world_size, @@ -435,7 +531,31 @@ impl Qwen35TpExecutor { capacity_pages_for_requests, max_position_embeddings, eos_token_id, - }) + graph_enabled, + active_slots: Mutex::new(Vec::new()), + }; + // Pre-capture every reachable decode graph now: after NCCL connect, + // exactly once, before serving. A mid-serving capture on one rank while + // a peer replays would desync the recorded collectives, so serve time + // is replay-only. + if graph_enabled { + executor.run_decode_graph_precapture_sweep()?; + log::info!( + "Qwen3.5 TP decode CUDA Graph enabled: {} bucket(s) up to batch {} captured per rank", + BATCH_BUCKETS + .iter() + .take_while(|&&b| b <= bucket_for(executor.max_batch)) + .count(), + bucket_for(executor.max_batch), + ); + } + Ok(executor) + } + + /// Whether decode replays pre-captured CUDA Graphs (P2c gate: requested + /// AND the TP-local decode GQA group has a compiled kernel). + pub fn graph_enabled(&self) -> bool { + self.graph_enabled } #[cfg(test)] @@ -463,6 +583,90 @@ impl Qwen35TpExecutor { token_id == self.eos_token_id } + /// Pre-capture every reachable decode bucket on every rank, phase-by-phase + /// and barriered by the controller so a captured collective's first launch + /// never overlaps a peer's capture (qwen3 sweep precedent). + fn run_decode_graph_precapture_sweep(&self) -> Result<()> { + // NCCL has no device timeout, so a desynced sweep wedges forever; this + // watchdog aborts on the deadline. abort() not exit() — exit's cudart + // atexit teardown takes the same wedged driver lock — and it disarms + // only on the explicit success send (drop-on-error stays armed). + let (sweep_done_tx, sweep_done_rx) = mpsc::sync_channel::<()>(1); + let deadline = Instant::now() + TP_PRECAPTURE_TIMEOUT; + let watchdog = thread::Builder::new() + .name("qwen35-tp-precapture-watchdog".into()) + .spawn(move || { + // Disarmed only by the explicit success send. A sender drop + // (error path) also returns Err here; stay armed to the + // deadline before deciding startup is wedged. + if sweep_done_rx.recv_timeout(TP_PRECAPTURE_TIMEOUT).is_ok() { + return; + } + std::thread::sleep(deadline.saturating_duration_since(Instant::now())); + eprintln!( + "Qwen3.5 TP decode graph pre-capture did not complete within {}s — NCCL wedge suspected, aborting", + TP_PRECAPTURE_TIMEOUT.as_secs() + ); + log::error!( + "Qwen3.5 TP decode graph pre-capture did not complete within {}s — NCCL wedge suspected, aborting", + TP_PRECAPTURE_TIMEOUT.as_secs() + ); + std::process::abort(); + }) + .map_err(|e| anyhow::anyhow!("failed to spawn Qwen3.5 TP pre-capture watchdog: {e}"))?; + + let started = Instant::now(); + let max_bucket = bucket_for(self.max_batch); + let sweep = (|| { + self.run_precapture_phase(PrecapturePhase::Warmup)?; + for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() { + if bucket > max_bucket { + break; + } + self.run_precapture_phase(PrecapturePhase::Capture { bucket_idx })?; + self.run_precapture_phase(PrecapturePhase::Launch { bucket_idx })?; + } + self.run_precapture_phase(PrecapturePhase::Finalize) + })(); + match sweep { + Ok(()) => { + // Disarm: only the explicit success send stops the watchdog. + sweep_done_tx + .send(()) + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP pre-capture watchdog exited"))?; + watchdog + .join() + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP pre-capture watchdog panicked"))?; + log::info!( + "Qwen3.5 TP decode graph pre-capture: buckets up to {max_bucket} captured per rank in {:.2}s", + started.elapsed().as_secs_f64() + ); + Ok(()) + } + // On error the watchdog stays armed: peers may be wedged in + // unpaired collectives, and the abort makes the wedge attributable. + Err(err) => Err(err), + } + } + + fn run_precapture_phase(&self, phase: PrecapturePhase) -> Result<()> { + self.poison.ensure_healthy()?; + let resp_rx = self.dispatch_mutating("decode graph precapture", |start, resp| { + TpWorkerCommand::Precapture { phase, start, resp } + })?; + let responses = recv_runtime_responses( + &resp_rx, + self.world_size, + "decode graph precapture", + &self.poison, + )?; + validate_dispatched_responses( + validate_ack_responses(responses, self.world_size, "decode graph precapture"), + "decode graph precapture", + &self.poison, + ) + } + #[cfg(test)] fn ping_all(&self) -> Result<()> { self.poison.ensure_healthy()?; @@ -495,7 +699,20 @@ impl Qwen35TpExecutor { .cloned() .map(TpPrefillChunkItem::from) .collect(); - self.execute_prefill_chunks(&chunks) + let result = self.execute_prefill_chunks(&chunks)?; + if self.graph_enabled { + // Convenience-API slot tracking: every prefill plan item finishes + // prefill (TpPrefillChunkItem::from sets finish_prefill), so each + // request takes the next dense decode slot. + let mut active = self + .active_slots + .lock() + .unwrap_or_else(PoisonError::into_inner); + for chunk in &chunks { + active.push(chunk.request_id); + } + } + Ok(result) } fn execute_prefill_chunks(&self, chunks: &[TpPrefillChunkItem]) -> Result { @@ -536,18 +753,49 @@ impl Qwen35TpExecutor { !plan.requests.is_empty(), "Qwen3.5 TP decode plan requires at least one request" ); - let requests: Vec = plan - .requests - .iter() - .map(|request| { - TpDecodeStepItem::new( - request.request_id, - request.token_id, - request.logprobs, - SamplingParams::default(), - ) - }) - .collect(); + let requests: Vec = if self.graph_enabled { + let active = self + .active_slots + .lock() + .unwrap_or_else(PoisonError::into_inner); + anyhow::ensure!( + plan.requests.len() == active.len(), + "Qwen3.5 TP graph decode must cover all {} active requests in slot order, got {}", + active.len(), + plan.requests.len() + ); + plan.requests + .iter() + .enumerate() + .map(|(slot, request)| { + anyhow::ensure!( + active[slot] == request.request_id, + "Qwen3.5 TP graph decode slot {slot} holds request {} but the plan carries {}", + active[slot].get(), + request.request_id.get() + ); + Ok(TpDecodeStepItem::new_with_slot( + request.request_id, + request.token_id, + request.logprobs, + SamplingParams::default(), + slot, + )) + }) + .collect::>()? + } else { + plan.requests + .iter() + .map(|request| { + TpDecodeStepItem::new( + request.request_id, + request.token_id, + request.logprobs, + SamplingParams::default(), + ) + }) + .collect() + }; self.execute_decode_items(&requests, 0) } @@ -611,10 +859,24 @@ impl Qwen35TpExecutor { } pub fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { + let compaction = self.track_retired_slot(request_id); + self.drop_request_with_compaction(request_id, expectation, compaction) + } + + /// Retire a request, attaching the slot compaction the caller (scheduler) + /// already applied to its own dense-slot bookkeeping. Workers apply the + /// move and poison on occupancy mismatch; eager workers ignore it. + pub(crate) fn drop_request_with_compaction( + &self, + request_id: RequestId, + expectation: DropExpectation, + compaction: Option, + ) -> Result<()> { self.poison.ensure_healthy()?; let resp_rx = self.dispatch_mutating("drop request", |start, resp| TpWorkerCommand::DropRequest { request_id, + compaction, start, resp, })?; @@ -627,6 +889,29 @@ impl Qwen35TpExecutor { ) } + /// Convenience-API tracker: swap-remove the retired request and derive the + /// slot compaction (last occupied slot moves into the vacated one). + /// Returns `None` on the eager path and for untracked requests. + fn track_retired_slot(&self, request_id: RequestId) -> Option { + if !self.graph_enabled { + return None; + } + let mut active = self + .active_slots + .lock() + .unwrap_or_else(PoisonError::into_inner); + let idx = active.iter().position(|&id| id == request_id)?; + let last = active.len() - 1; + active.swap_remove(idx); + // `then`, not `then_some`: the moved request only exists when the + // retired request was not the tail slot. + (idx < active.len()).then(|| TpSlotCompaction { + moved_request_id: active[idx], + from: last, + to: idx, + }) + } + #[cfg(test)] fn snapshot_workers(&self) -> Result> { self.poison.ensure_healthy()?; @@ -880,6 +1165,7 @@ impl TpStartupGate { } impl TpWorker { + #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] fn spawn( rank: usize, @@ -887,6 +1173,7 @@ impl TpWorker { model: Qwen35Model, max_batch: usize, max_prefill_tokens: usize, + graph_enabled: bool, nccl_id: cudarc::nccl::safe::Id, startup_gate: Arc, effective_max_batch: Arc, @@ -911,6 +1198,7 @@ impl TpWorker { model, max_batch, max_prefill_tokens, + graph_enabled, ); let prepared = match prepared { Ok((prepared, rank_max_batch)) => { @@ -926,7 +1214,7 @@ impl TpWorker { return; } let max_batch = effective_max_batch.load(Ordering::Acquire); - match prepared.connect(nccl_id, max_batch, poison) { + match prepared.connect(nccl_id, max_batch, graph_enabled, poison) { Ok(mut state) => { let _ = startup_tx.send(Ok(())); state.run(rx); @@ -984,8 +1272,17 @@ struct TpWorkerState { rank: usize, _world_size: usize, max_batch: usize, + /// Before `model` on purpose: NCCL comm teardown polls until every graph + /// that recorded its collectives is destroyed, so the decode graphs must + /// drop before `model.tp_comm` (qwen3 teardown-hang precedent). + graph_state: Option, model: Qwen35Model, requests: Vec, + /// Graph-mode slot ownership: `slot_map[i]` is the request whose recurrent + /// state lives in `graph_state.slot_states[i]`. The scheduler owns slot + /// assignment and compaction; the worker only applies and checks them. + /// Empty in eager mode. + slot_map: Vec>, decode_buffers: BatchDecodeBuffers35, /// Eager decode GDR pointer tables: allocated once at capacity, refilled /// with the live rows every step. @@ -1009,7 +1306,10 @@ struct TpRequestState { request_id: RequestId, phase: TpRequestPhase, kv: KvState, - recurrent: RecurrentState, + /// Prefill-owned recurrent state. Graph mode moves it into the decode slot + /// on the request's first decode row (`None` afterwards); the eager path + /// keeps it for the request's whole lifetime. + recurrent: Option, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -1033,6 +1333,7 @@ impl TpWorkerPrepared { model: Qwen35Model, requested_max_batch: usize, max_prefill_tokens: usize, + graph_enabled: bool, ) -> Result<(Self, usize)> { let cublas_guard = bind_worker_thread(&model)?; let (free_bytes, total_bytes) = model @@ -1049,13 +1350,42 @@ impl TpWorkerPrepared { model.geometry, prefill_scratch_tokens, ); - let max_batch = effective_recurrent_capacity( - requested_max_batch, - free_bytes, - recurrent_bytes, - TP_RUNTIME_MEMORY_RESERVE_BYTES, - prefill_scratch_bytes, - ); + // Graph mode pre-allocates one fixed-address slot state per decode + // bucket position up front; reserve that before sizing per-request + // (prefill-transient) state capacity. The reserve must track the + // bucket of the *effective* batch, not the requested one: reserving + // for `bucket_for(requested)` can starve a tight-memory rank down to + // zero capacity. Iterate the bucket downward until it stabilises — + // the bucket only shrinks, so this converges — and clamp the fitted + // batch to the reserved bucket so the later `bucket_for(effective)` + // graph allocation never exceeds the reserve. + let max_batch = if graph_enabled { + let mut slot_bucket = bucket_for(requested_max_batch); + loop { + let reserve = slot_bucket * recurrent_bytes; + let candidate = effective_recurrent_capacity( + requested_max_batch, + free_bytes.saturating_sub(reserve), + recurrent_bytes, + TP_RUNTIME_MEMORY_RESERVE_BYTES, + prefill_scratch_bytes, + ); + let fitted = candidate.min(slot_bucket); + let next = bucket_for(fitted); + if next >= slot_bucket { + break fitted; + } + slot_bucket = next; + } + } else { + effective_recurrent_capacity( + requested_max_batch, + free_bytes, + recurrent_bytes, + TP_RUNTIME_MEMORY_RESERVE_BYTES, + prefill_scratch_bytes, + ) + }; anyhow::ensure!( max_batch > 0, "Qwen3.5 TP rank {rank} has {} MiB free after fixed buffers, but one recurrent request needs {} MiB plus {} MiB runtime reserve and {} MiB prefill scratch for {} tokens", @@ -1098,6 +1428,7 @@ impl TpWorkerPrepared { self, nccl_id: cudarc::nccl::safe::Id, effective_max_batch: usize, + graph_enabled: bool, poison: Arc, ) -> Result { let Self { @@ -1127,12 +1458,25 @@ impl TpWorkerPrepared { effective_max_batch, "Qwen3.5 TP eager decode", )?; + let (graph_state, slot_map) = if graph_enabled { + // cuBLASLt plans are thread-local: tune the decode bucket GEMMs on + // this worker thread now so plan selection never runs inside + // cuStreamBeginCapture during the pre-capture sweep. + model.tune_decode_gemm_algos()?; + let slots = bucket_for(effective_max_batch); + let graph_state = model.create_batch_decode_graph_state_with_capacity(slots)?; + (Some(graph_state), vec![None; slots]) + } else { + (None, Vec::new()) + }; Ok(TpWorkerState { rank, _world_size: world_size, max_batch: effective_max_batch, + graph_state, model, requests: Vec::new(), + slot_map, decode_buffers, decode_pointer_tables, sample_scratch, @@ -1208,14 +1552,25 @@ impl TpWorkerState { } TpWorkerCommand::DropRequest { request_id, + compaction, start, resp, } => { if start.wait() == TpCommandDecision::Cancel { false } else { - let existed = self.drop_request(request_id); - self.respond(resp, "drop request", Ok(TpWorkerReply::DropAck { existed })) + let result = self + .drop_request(request_id, compaction) + .map(|existed| TpWorkerReply::DropAck { existed }); + self.respond(resp, "drop request", result) + } + } + TpWorkerCommand::Precapture { phase, start, resp } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.precapture_phase(phase).map(|()| TpWorkerReply::Ack); + self.respond(resp, "decode graph precapture", result) } } #[cfg(test)] @@ -1237,7 +1592,7 @@ impl TpWorkerState { } #[cfg(test)] TpWorkerCommand::RemoveRequestStateForTest { request_id, resp } => { - let _ = resp.send(self.drop_request(request_id)); + let _ = resp.send(self.drop_request(request_id, None).unwrap_or(false)); false } #[cfg(test)] @@ -1327,7 +1682,12 @@ impl TpWorkerState { ); let prompt = [chunk.prompt_tokens.as_slice()]; - let mut recurrent_refs = vec![&mut state.recurrent]; + let mut recurrent_refs = vec![ + state + .recurrent + .as_mut() + .expect("prefill-phase TP request owns its recurrent state"), + ]; let logits = self.model.batch_prefill_logits( &prompt, std::slice::from_mut(&mut state.kv), @@ -1371,6 +1731,9 @@ impl TpWorkerState { if bs == 0 { return Ok(Vec::new()); } + if self.graph_state.is_some() { + return self.run_decode_batch_graph(requests, sample_seed); + } // Resolve the worker state slot of every row in command order. // Decode request ids are unique within one command @@ -1398,7 +1761,12 @@ impl TpWorkerState { for (state_idx, state) in self.requests.iter_mut().enumerate() { if let Some(row) = row_of_state[state_idx] { kv_slots[row] = Some(&mut state.kv); - recurrent_slots[row] = Some(&mut state.recurrent); + recurrent_slots[row] = Some( + state + .recurrent + .as_mut() + .expect("eager TP decode request owns its recurrent state"), + ); } } let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); @@ -1478,6 +1846,144 @@ impl TpWorkerState { .collect()) } + /// CUDA Graph decode step under TP: replay-only (every bucket was recorded + /// by the startup pre-capture sweep), one forward for the whole batch on + /// every rank, then (rank 0 only) the same batched host-side sampling pass + /// as the eager path. + /// + /// Rows must arrive in the scheduler-owned dense slot order + /// (`slot_idx == row`). On a request's first decode row its prefill-owned + /// recurrent state is D2D-copied into `graph_state.slot_states[slot]` and + /// the per-request allocation is dropped; the persistent linear-state + /// pointer tables then keep every replay reading the fixed slot addresses. + fn run_decode_batch_graph( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result> { + let bs = requests.len(); + let graph_state = self + .graph_state + .as_mut() + .expect("graph decode arm requires graph state"); + let ctx = self.model.device_ctx(); + + // Resolve the worker state of every row, enforce dense slot order, and + // admit first-decode rows into their slots. Decode request ids are + // unique within one command (validate_decode_requests), so each slot + // is borrowed at most once. + let mut row_of_state: Vec> = vec![None; self.requests.len()]; + for (row, request) in requests.iter().enumerate() { + anyhow::ensure!( + request.slot_idx == Some(row), + "Qwen3.5 TP graph decode row {row} carries slot {:?}; rows must arrive in dense slot order 0..{bs}", + request.slot_idx + ); + let state_idx = self + .requests + .iter() + .position(|state| state.request_id == request.request_id) + .ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode request {} has no worker state", + request.request_id.get() + ) + })?; + anyhow::ensure!( + self.requests[state_idx].phase == TpRequestPhase::Decoding, + "Qwen3.5 TP request {} is not ready for decode", + request.request_id.get() + ); + debug_assert!(row_of_state[state_idx].is_none()); + row_of_state[state_idx] = Some(row); + + if self.slot_map.get(row).copied().flatten() == Some(request.request_id) { + anyhow::ensure!( + self.requests[state_idx].recurrent.is_none(), + "Qwen3.5 TP request {} was admitted to slot {row} but still owns prefill recurrent state", + request.request_id.get() + ); + } else { + slot_admit(&mut self.slot_map, row, request.request_id)?; + let recurrent = self.requests[state_idx].recurrent.take().ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP request {} lost its prefill recurrent state before slot admission", + request.request_id.get() + ) + })?; + graph_state.copy_state_to_slot(ctx, &recurrent, row)?; + } + } + + // KV refs in row (slot) order; page tables stay per-step H2D via + // sync_paged_meta inside batch_decode_graph. + let mut kv_slots: Vec> = + std::iter::repeat_with(|| None).take(bs).collect(); + for (state_idx, state) in self.requests.iter_mut().enumerate() { + if let Some(row) = row_of_state[state_idx] { + kv_slots[row] = Some(&mut state.kv); + } + } + let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); + for kv in kv_slots { + kv_refs.push(kv.expect("decode row state resolved above")); + } + let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); + self.model.batch_decode_graph( + &token_ids, + &mut kv_refs, + graph_state, + DecodeGraphUse::Replay, + )?; + + if self.rank != 0 { + return Ok(Vec::new()); + } + + // Snapshot requested logits rows BEFORE sampling: the sampler may + // modify bufs.logits in place. + let requested_logprobs: Vec = + requests.iter().map(|request| request.logprobs).collect(); + let cpu_logits = + snapshot_requested_logprobs(ctx, &graph_state.buffers.logits, &requested_logprobs)?; + let params_refs: Vec<&SamplingParams> = requests + .iter() + .map(|request| &request.sampling_params) + .collect(); + let steps = vec![0u64; bs]; + let tokens = pegainfer_sample::select_batch( + ctx, + &graph_state.buffers.logits, + ¶ms_refs, + &steps, + sample_seed, + &mut self.sample_scratch, + )?; + anyhow::ensure!( + tokens.len() == bs, + "Qwen3.5 TP decode sampling returned {} tokens for {bs} rows", + tokens.len() + ); + Ok(requests + .iter() + .enumerate() + .map(|(row, request)| { + let logprob = cpu_logits[row].as_ref().and_then(|logits_row| { + pegainfer_sample::token_logprob_from_row( + logits_row, + tokens[row], + request.logprobs, + ) + }); + DecodeRequestResult { + request_id: request.request_id, + token: tokens[row], + logprob, + } + }) + .collect()) + } + fn sample_final_prefill_chunk( &mut self, chunk: &TpPrefillChunkItem, @@ -1578,7 +2084,7 @@ impl TpWorkerState { request_id, phase: TpRequestPhase::Prefilling, kv: self.model.alloc_kv(), - recurrent, + recurrent: Some(recurrent), }; self.requests.push(state); Ok(self.requests.len() - 1) @@ -1590,14 +2096,175 @@ impl TpWorkerState { .position(|state| state.request_id == request_id) } - fn drop_request(&mut self, request_id: RequestId) -> bool { - if let Some(idx) = self.request_index(request_id) { - self.requests.swap_remove(idx); - true - } else { - false + /// One phase of the startup pre-capture sweep (graph mode only). + fn precapture_phase(&mut self, phase: PrecapturePhase) -> Result<()> { + match phase { + PrecapturePhase::Warmup => self.model.warmup_tp_collective(), + PrecapturePhase::Capture { bucket_idx } => { + self.precapture_bucket(bucket_idx, DecodeGraphUse::CaptureOnly) + } + PrecapturePhase::Launch { bucket_idx } => { + self.precapture_bucket(bucket_idx, DecodeGraphUse::Replay) + } + PrecapturePhase::Finalize => { + let graph_state = self.graph_state.as_ref().ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 TP pre-capture Finalize without graph state") + })?; + for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() { + if bucket > graph_state.slot_states.len() { + break; + } + anyhow::ensure!( + graph_state.graphs[bucket_idx].is_captured(), + "Qwen3.5 TP decode graph pre-capture left bucket {bucket} uncaptured" + ); + } + Ok(()) + } + } + } + + /// Capture or launch one bucket with synthetic rows. Outputs are + /// discarded; the rows exist only to give the recorded kernels valid + /// addresses. One real row (token 0 at position 0 over a freshly + /// allocated one-page KV state) selects nothing — the bucket is passed + /// explicitly — and every other row is padding on the pool's reserved + /// padding page, exactly as when serving. The sweep therefore holds one + /// KV page at a time regardless of pool size or bucket. + fn precapture_bucket(&mut self, bucket_idx: usize, graph_use: DecodeGraphUse) -> Result<()> { + let bucket = BATCH_BUCKETS[bucket_idx]; + let graph_state = self.graph_state.as_mut().ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 TP pre-capture on a worker without graph state") + })?; + anyhow::ensure!( + bucket <= graph_state.slot_states.len(), + "Qwen3.5 TP pre-capture bucket {bucket} exceeds {} slots", + graph_state.slot_states.len() + ); + let mut synthetic_kv = self.model.alloc_kv(); + let mut kv_refs = [&mut synthetic_kv]; + self.model.batch_decode_graph_padded( + &[0u32], + &mut kv_refs, + graph_state, + graph_use, + bucket, + )?; + // Capture acks only after the async cuGraphUpload lands; Launch acks + // only after the collectives drained. + self.model + .device_ctx() + .stream + .synchronize() + .map_err(|e| anyhow::anyhow!("Qwen3.5 TP pre-capture bucket {bucket} sync: {e}"))?; + Ok(()) + } + + /// Retire a request. Graph mode also applies the scheduler's slot + /// compaction (D2D move + occupancy assertions) so the slot layout stays + /// dense; any mismatch between the scheduler's claim and the worker's slot + /// map is a divergence and fails the command (poisoning the executor). + fn drop_request( + &mut self, + request_id: RequestId, + compaction: Option, + ) -> Result { + let Some(idx) = self.request_index(request_id) else { + anyhow::ensure!( + compaction.is_none(), + "Qwen3.5 TP drop of absent request {} carries a slot compaction", + request_id.get() + ); + return Ok(false); + }; + if let Some(graph_state) = self.graph_state.as_mut() { + match compaction { + Some(compaction) => { + let needs_move = slot_compact(&mut self.slot_map, request_id, compaction)?; + if needs_move { + graph_state.move_slot_within( + self.model.device_ctx(), + compaction.from, + compaction.to, + )?; + } + } + None => { + slot_release(&mut self.slot_map, request_id); + } + } } + self.requests.swap_remove(idx); + Ok(true) + } +} + +/// Admit `request_id` to decode `slot`: the slot must be free (retirement and +/// compaction keep the map dense, so an occupied slot here is a scheduler +/// divergence). +fn slot_admit(owners: &mut [Option], slot: usize, request_id: RequestId) -> Result<()> { + let slot_count = owners.len(); + let owner = owners.get_mut(slot).ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 TP decode slot {slot} exceeds worker slot map {slot_count}") + })?; + anyhow::ensure!( + owner.is_none(), + "Qwen3.5 TP decode slot {slot} still owned by request {} at admission of request {}", + owner.expect("checked").get(), + request_id.get() + ); + *owner = Some(request_id); + Ok(()) +} + +/// Clear `request_id`'s slot if it held one. Requests retired before their +/// first decode row never materialized a slot; that is not an error. +fn slot_release(owners: &mut [Option], request_id: RequestId) -> Option { + let slot = owners.iter().position(|owner| *owner == Some(request_id))?; + owners[slot] = None; + Some(slot) +} + +/// Apply the scheduler's slot compaction to the worker's slot map and report +/// whether a GPU state move is needed. Both requests may legitimately be +/// unmaterialized (retired/compacted before their first decode row), but a +/// materialized slot must hold exactly the request the scheduler claims. +fn slot_compact( + owners: &mut [Option], + dropped: RequestId, + compaction: TpSlotCompaction, +) -> Result { + let TpSlotCompaction { + moved_request_id, + from, + to, + } = compaction; + anyhow::ensure!( + from < owners.len() && to < owners.len(), + "Qwen3.5 TP slot compaction {from} -> {to} exceeds worker slot map {}", + owners.len() + ); + let dropped_owner = owners[to]; + let moved_owner = owners[from]; + if let Some(owner) = dropped_owner { + anyhow::ensure!( + owner == dropped, + "Qwen3.5 TP slot {to} holds request {} where the scheduler dropped request {}", + owner.get(), + dropped.get() + ); } + if let Some(owner) = moved_owner { + anyhow::ensure!( + owner == moved_request_id, + "Qwen3.5 TP slot {from} holds request {} where the scheduler moved request {}", + owner.get(), + moved_request_id.get() + ); + } + owners[to] = moved_owner; + owners[from] = None; + Ok(moved_owner.is_some()) } fn validate_prefill_chunks(chunks: &[TpPrefillChunkItem]) -> Result<()> { @@ -1735,17 +2402,6 @@ impl From for TpPrefillChunkItem { } } -impl From for TpDecodeStepItem { - fn from(request: DecodeStepItem) -> Self { - Self::new( - request.request_id, - request.token_id, - request.logprobs, - SamplingParams::default(), - ) - } -} - fn recv_runtime_responses( responses: &mpsc::Receiver, expected: usize, @@ -1826,7 +2482,6 @@ fn validate_exact_rank_responses( Ok(replies) } -#[cfg(test)] fn validate_ack_responses( responses: Vec, world_size: usize, @@ -2231,12 +2886,122 @@ mod tests { } #[test] - fn rejects_tensor_parallel_cuda_graph() { + fn tensor_parallel_cuda_graph_gate_defers_to_model_load() { + // P2c: TP + CUDA Graph is no longer rejected up front; the graph/eager + // decision needs the model config, so a nonexistent path fails at load. let err = match Qwen35TpExecutor::from_runtime_with_capacity("unused", true, &[0, 1], 1) { - Ok(_) => panic!("TP CUDA Graph should fail"), + Ok(_) => panic!("TP CUDA Graph with a nonexistent model path should fail at load"), Err(err) => err.to_string(), }; - assert!(err.contains("eager execution only")); + assert!(!err.contains("eager execution only")); + } + + #[test] + fn slot_map_admit_release_and_compact() { + let id = |value: u64| RequestId::new(value); + let mut owners = vec![None, None, None, None]; + + slot_admit(&mut owners, 0, id(1)).unwrap(); + slot_admit(&mut owners, 1, id(2)).unwrap(); + slot_admit(&mut owners, 2, id(3)).unwrap(); + + let err = slot_admit(&mut owners, 1, id(9)).unwrap_err().to_string(); + assert!(err.contains("still owned by request 2")); + + // Retire slot 1: last occupied slot (2, request 3) moves into it. + let needs_move = slot_compact( + &mut owners, + id(2), + TpSlotCompaction { + moved_request_id: id(3), + from: 2, + to: 1, + }, + ) + .unwrap(); + assert!(needs_move, "materialized moved request needs the GPU move"); + assert_eq!(owners, vec![Some(id(1)), Some(id(3)), None, None]); + + // Retire the tail slot: release without compaction. + assert_eq!(slot_release(&mut owners, id(3)), Some(1)); + assert_eq!(owners, vec![Some(id(1)), None, None, None]); + + // Releasing a request that never materialized a slot is not an error. + assert_eq!(slot_release(&mut owners, id(77)), None); + } + + #[test] + fn slot_map_compact_tolerates_unmaterialized_requests() { + let id = |value: u64| RequestId::new(value); + let mut owners = vec![None, None, None]; + + // Dropped request materialized, moved request not yet admitted to its + // slot (retired between promotion and its first decode row): clear + // only, no GPU move. + slot_admit(&mut owners, 0, id(1)).unwrap(); + let needs_move = slot_compact( + &mut owners, + id(1), + TpSlotCompaction { + moved_request_id: id(2), + from: 2, + to: 0, + }, + ) + .unwrap(); + assert!(!needs_move); + assert_eq!(owners, vec![None, None, None]); + + // Moved request materialized, dropped request not: the move is needed + // and the moved request takes over the vacated slot. + slot_admit(&mut owners, 2, id(3)).unwrap(); + let needs_move = slot_compact( + &mut owners, + id(4), + TpSlotCompaction { + moved_request_id: id(3), + from: 2, + to: 0, + }, + ) + .unwrap(); + assert!(needs_move); + assert_eq!(owners, vec![Some(id(3)), None, None]); + } + + #[test] + fn slot_map_compact_poisons_on_occupancy_mismatch() { + let id = |value: u64| RequestId::new(value); + let mut owners = vec![Some(id(1)), Some(id(2))]; + + let err = slot_compact( + &mut owners, + id(9), + TpSlotCompaction { + moved_request_id: id(2), + from: 1, + to: 0, + }, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("slot 0 holds request 1")); + + let err = slot_compact( + &mut owners, + id(1), + TpSlotCompaction { + moved_request_id: id(9), + from: 1, + to: 0, + }, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("slot 1 holds request 2")); + + let err = slot_admit(&mut owners, 5, id(1)).unwrap_err().to_string(); + assert!(err.contains("exceeds worker slot map")); } #[test] diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs index 32f216672..e4dfc4560 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -21,6 +21,7 @@ use pegainfer_kernels::tensor::StreamOverrideGuard; use super::batch_decode_graph::BatchDecodeGraphState; use super::recurrent_state::RecurrentState; use super::weights::Qwen35Model; +use crate::batch_decode::DecodeGraphUse; pub(crate) struct UnifiedStepOutput { pub(crate) prefill_logits: Option, @@ -125,7 +126,12 @@ impl Qwen35Model { let decoded = if decode_tokens.is_empty() { false } else { - self.batch_decode_graph(decode_tokens, decode_kv_states, graph_state)?; + self.batch_decode_graph( + decode_tokens, + decode_kv_states, + graph_state, + DecodeGraphUse::Serve, + )?; true }; @@ -199,7 +205,7 @@ mod tests { for _ in 1..num_steps { let tids = [*tokens_a.last().unwrap(), *tokens_b.last().unwrap()]; model - .batch_decode_graph(&tids, &mut kv_refs, &mut gs) + .batch_decode_graph(&tids, &mut kv_refs, &mut gs, DecodeGraphUse::Serve) .unwrap(); let next = greedy_sample_batch(&model, &gs.buffers.logits, 2); tokens_a.push(next[0]); diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index 214e9442d..30527c48e 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -151,8 +151,8 @@ impl Qwen35Model { let mut config = Config35::from_file(model_path)?; let tensor_parallel = runtime.tensor_parallel.unwrap_or_default(); - let geometry = LocalGeometry::try_new(&config, tensor_parallel, runtime.enable_cuda_graph) - .map_err(anyhow::Error::from)?; + let geometry = + LocalGeometry::try_new(&config, tensor_parallel).map_err(anyhow::Error::from)?; debug!( "Config: hidden_size={}, num_layers={}, full_attn={}, linear_attn={}, max_position_embeddings={}, tp_rank={}, tp_world_size={}", config.hidden_size, @@ -360,6 +360,34 @@ impl Qwen35Model { self.tp_comm = Some(comm); } + /// Force NCCL connect before any CUDA Graph capture records a collective + /// (lazy connect inside `cuStreamBeginCapture` wedges the capture). NCCL + /// 2.22+ connects per size-selected algorithm, so warm one all-reduce at + /// every decode bucket's message size. No-op without a TP communicator. + pub(crate) fn warmup_tp_collective(&self) -> Result<()> { + if let Some(comm) = &self.tp_comm { + let buckets = super::batch_decode_graph::BATCH_BUCKETS; + let max_elems = buckets.last().unwrap() * self.config.hidden_size; + let mut scratch = self + .ctx + .stream + .alloc_zeros::(max_elems) + .map_err(|e| anyhow::anyhow!("alloc NCCL warm-up scratch: {e}"))?; + for &bucket in buckets { + let mut view = scratch.slice_mut(0..bucket * self.config.hidden_size); + comm.all_reduce_in_place(&mut view, &ReduceOp::Sum) + .map_err(|e| { + anyhow::anyhow!("Qwen3.5 NCCL warm-up all-reduce failed: {e:?}") + })?; + } + self.ctx + .stream + .synchronize() + .map_err(|e| anyhow::anyhow!("Qwen3.5 NCCL warm-up sync failed: {e}"))?; + } + Ok(()) + } + pub(crate) fn all_reduce_hidden(&self, hidden: &mut HiddenStates) -> Result<()> { self.all_reduce_hidden_untraced(hidden) } diff --git a/pegainfer-qwen35/src/weights/layers.rs b/pegainfer-qwen35/src/weights/layers.rs index fd593956f..f346343b9 100644 --- a/pegainfer-qwen35/src/weights/layers.rs +++ b/pegainfer-qwen35/src/weights/layers.rs @@ -433,7 +433,7 @@ mod tests { fn test_geometry(rank: usize, world_size: usize) -> LocalGeometry { let config = test_config(); let tp = TensorParallelConfig::try_from((rank, world_size)).unwrap(); - LocalGeometry::try_new(&config, tp, false).unwrap() + LocalGeometry::try_new(&config, tp).unwrap() } #[test] diff --git a/pegainfer-qwen35/tests/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index 9886399bb..9933d2e22 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -829,7 +829,6 @@ fn test_e2e_qwen35_scheduler_tp2() { info!("Loading Qwen3.5 TP2 model for scheduler test..."); let start = Instant::now(); let tokenizer = common::load_tokenizer(&model_path); - // TP Phase 1 is eager-only; CUDA Graph must stay disabled for multi-device startup. let handle = pegainfer_qwen35::start_engine_with_capacity( Path::new(&model_path), EngineLoadOptions { @@ -847,3 +846,34 @@ fn test_e2e_qwen35_scheduler_tp2() { let max_context_tokens = max_position_embeddings(&model_path); run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2"); } + +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn test_e2e_qwen35_scheduler_tp2_graph() { + let Some(model_path) = common::model_path_or_skip("test_e2e_qwen35_scheduler_tp2_graph") else { + return; + }; + + info!("Loading Qwen3.5 TP2 model with CUDA Graph for scheduler test..."); + let start = Instant::now(); + let tokenizer = common::load_tokenizer(&model_path); + // P2c: decode replays pre-captured CUDA Graphs when the TP-local decode + // GQA group has a compiled kernel (4B/9B); uncompiled groups (27B group 6) + // keep the batched eager path under the same request flow. + let handle = pegainfer_qwen35::start_engine_with_capacity( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: common::tp2_device_ordinals(), + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, + ) + .expect("Failed to start Qwen3.5 TP2 graph scheduler"); + info!("TP2 graph scheduler loaded in {:.2?}", start.elapsed()); + + let max_context_tokens = max_position_embeddings(&model_path); + run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2 graph"); +} diff --git a/pegainfer-qwen35/tests/hf_golden_gate.rs b/pegainfer-qwen35/tests/hf_golden_gate.rs index 558866a6f..a2e3220b5 100644 --- a/pegainfer-qwen35/tests/hf_golden_gate.rs +++ b/pegainfer-qwen35/tests/hf_golden_gate.rs @@ -747,6 +747,117 @@ fn build_tp2_executor(model_path: &str) -> Qwen35TpExecutor { .expect("build Qwen3.5 TP2 logits executor") } +/// TP2 executor with CUDA Graph requested. Returns `None` when the loaded +/// model's TP-local decode GQA group has no compiled kernel (27B group 6) — +/// the P2c gate keeps that path eager, so there is no graph to gate on. +fn build_tp2_graph_executor(model_path: &str, label: &str) -> Option { + let devices = common::tp2_device_ordinals(); + let ex = Qwen35TpExecutor::from_runtime_with_capacity( + model_path, + true, + &devices, + MAX_EXECUTOR_BATCH, + ) + .expect("build Qwen3.5 TP2 graph logits executor"); + if !ex.graph_enabled() { + eprintln!( + "qwen35 hf_golden_gate [{label}]: CUDA Graph gated off (uncompiled TP-local decode GQA group); skipping" + ); + return None; + } + Some(ex) +} + +/// Mid-batch drop with slot compaction under TP: prefill `seqs`, decode one +/// step, retire `SLOT_COMPACTION_DROP_INDEX` (the last slot moves into the +/// gap), then keep decoding the survivors in their new dense slot order. +fn run_tp_with_slot_compaction( + g: &Golden, + ex: &Qwen35TpExecutor, + seqs: &[usize], +) -> (Stats, Vec) { + assert!( + seqs.len() > SLOT_COMPACTION_DROP_INDEX + 1, + "TP slot-compaction replay needs a non-tail request to drop" + ); + assert!( + g.decode_len >= 2, + "TP slot-compaction replay needs at least two decode tokens" + ); + + let mut stats = Stats::default(); + let mut fingerprint = Vec::new(); + let mut fold = |stats: &mut Stats, seq, pos, pega: &[(u32, f32)]| { + fingerprint.push(pega[0].1); + check_position(stats, seq, pos, pega, &g.topk(seq, pos)); + }; + + let mut live: Vec<(usize, RequestId)> = seqs + .iter() + .map(|&seq| (seq, RequestId::new(30_000 + seq as u64))) + .collect(); + let items: Vec = live + .iter() + .map(|&(seq, id)| prefill_item(id, g.prompt(seq))) + .collect(); + let pr = ex + .execute_prefill(PrefillPlan { requests: &items }) + .expect("TP2 prefill"); + for (i, &(seq, _)) in live.iter().enumerate() { + fold( + &mut stats, + seq, + 0, + &top_logprobs(pr.requests[i].first_token_logprob.as_ref()), + ); + } + + let step0: Vec = live + .iter() + .map(|&(seq, id)| decode_item(id, g.decode(seq, 0))) + .collect(); + let dr = ex + .execute_decode(DecodePlan { requests: &step0 }) + .expect("TP2 decode before compaction"); + for (i, &(seq, _)) in live.iter().enumerate() { + fold( + &mut stats, + seq, + 1, + &top_logprobs(dr.requests[i].logprob.as_ref()), + ); + } + + let (_, dropped_id) = live[SLOT_COMPACTION_DROP_INDEX]; + ex.drop_request(dropped_id, DropExpectation::MustExist) + .expect("TP2 drop request"); + live.swap_remove(SLOT_COMPACTION_DROP_INDEX); + + for step in 1..g.decode_len { + let items: Vec = live + .iter() + .map(|&(seq, id)| decode_item(id, g.decode(seq, step))) + .collect(); + let dr = ex + .execute_decode(DecodePlan { requests: &items }) + .expect("TP2 decode after compaction"); + for (i, &(seq, _)) in live.iter().enumerate() { + fold( + &mut stats, + seq, + step + 1, + &top_logprobs(dr.requests[i].logprob.as_ref()), + ); + } + } + + for (_, id) in live { + ex.drop_request(id, DropExpectation::MustExist) + .expect("TP2 drop request"); + } + (stats, fingerprint) +} + #[test] fn pega_logprobs_match_hf_golden_within_qwen35_tolerance() { let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_golden") else { @@ -888,3 +999,62 @@ fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance_tp2() { "TP2 long sequential Qwen3.5 replay must reproduce identical logprobs" ); } + +/// P2c TP2 CUDA Graph gate: sequential replay, bucket-straddling batched +/// replay, and post-compaction replay after a mid-batch drop, all compared +/// against the same HF golden within the existing TP2 tolerances. +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2_graph() { + let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_golden_tp2_graph") + else { + return; + }; + let Some(golden) = Golden::load_for(&model_path, false) else { + return; + }; + if !check_fixture_metadata(&model_path, &golden) { + return; + } + report_fixture_shape(&golden); + let all: Vec = (0..golden.num_seqs).collect(); + + let Some(ex) = build_tp2_graph_executor(&model_path, "TP2 graph") else { + return; + }; + let (stats, fp1) = run_tp(&golden, &ex, &all, false); + report_and_assert("TP2 sequential graph", &stats); + let (_, fp2) = run_tp(&golden, &ex, &all, false); + assert_eq!( + fp1, fp2, + "TP2 sequential Qwen3.5 graph replay must reproduce identical logprobs" + ); + + for n in BUCKET_STRADDLES { + if all.len() >= n { + let (batched, _) = run_tp(&golden, &ex, &all[..n], true); + report_and_assert(&format!("TP2 batched graph ({n} padded)"), &batched); + } else { + eprintln!( + "qwen35 hf_golden_gate: skipping TP2 batched graph ({n} padded); fixture has only {} sequence(s)", + all.len() + ); + } + } + + if golden.num_seqs >= SLOT_COMPACTION_BATCH && golden.decode_len >= 2 { + let (compacted, fp1) = + run_tp_with_slot_compaction(&golden, &ex, &all[..SLOT_COMPACTION_BATCH]); + report_and_assert("TP2 slot-compaction graph", &compacted); + let (_, fp2) = run_tp_with_slot_compaction(&golden, &ex, &all[..SLOT_COMPACTION_BATCH]); + assert_eq!( + fp1, fp2, + "TP2 slot-compaction Qwen3.5 graph replay must reproduce identical logprobs" + ); + } else { + eprintln!( + "qwen35 hf_golden_gate: skipping TP2 slot-compaction graph; fixture has {} sequence(s), decode_len {}", + golden.num_seqs, golden.decode_len + ); + } +} diff --git a/pegainfer-qwen35/tests/serving_tp2.rs b/pegainfer-qwen35/tests/serving_tp2.rs index 2ed223e49..7340e0d79 100644 --- a/pegainfer-qwen35/tests/serving_tp2.rs +++ b/pegainfer-qwen35/tests/serving_tp2.rs @@ -1,5 +1,4 @@ use std::net::TcpListener; -use std::path::Path; use std::path::PathBuf; use std::time::Duration; @@ -51,7 +50,10 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { return Ok(()); }; let frontend_model_path = PathBuf::from(frontend_model_path); - let invalid_graph_model_path = engine_model_path.clone(); + // P2c graph acceptance smoke: CUDA Graph requested at TP2. Models with a + // compiled TP-local decode GQA group (4B/9B) replay pre-captured decode + // graphs; uncompiled groups (27B group 6) stay on the batched eager path + // under the same serving flow. let server = spawn_ready_server(engine_model_path, frontend_model_path, 1).await?; let client = test_client()?; @@ -59,11 +61,6 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { assert_non_streaming_completion(&client, &server.base_url).await?; assert_streaming_completion(&client, &server.base_url).await?; assert_concurrent_completions(&client, &server.base_url).await?; - assert_invalid_cuda_graph_tp_startup_fails( - invalid_graph_model_path - .to_str() - .context("Qwen3.5 engine fixture path is not valid UTF-8")?, - )?; server.shutdown().await } @@ -78,7 +75,7 @@ async fn spawn_ready_server( pegainfer_qwen35::start_engine_with_capacity( &engine_model_path, EngineLoadOptions { - enable_cuda_graph: false, + enable_cuda_graph: true, device_ordinals, seed: 42, ..EngineLoadOptions::default() @@ -199,27 +196,6 @@ async fn assert_concurrent_completions(client: &Client, base_url: &str) -> Resul Ok(()) } -fn assert_invalid_cuda_graph_tp_startup_fails(model_path: &str) -> Result<()> { - let Err(error) = pegainfer_qwen35::start_engine_with_capacity( - Path::new(model_path), - EngineLoadOptions { - enable_cuda_graph: true, - device_ordinals: common::tp2_device_ordinals(), - seed: 42, - ..EngineLoadOptions::default() - }, - 8, - 1, - ) else { - bail!("TP2 + CUDA Graph must fail before serving requests"); - }; - let message = error.to_string(); - if !message.contains("eager execution only") { - bail!("unexpected TP2 + CUDA Graph startup error: {message}"); - } - Ok(()) -} - async fn post_completion(client: &Client, base_url: &str, body: Value) -> Result { client .post(format!("{base_url}/v1/completions"))