fix(qwen3): re-enable CUDA Graph under TP via startup decode-graph pre-capture

[FeathBow]

Description

Fixes #498.

Re-enables CUDA Graph decode under --tp-size > 1 for qwen3. #496 guarded CUDA graph off under TP after the #481 hang; this removes both guards and makes graph-on TP serving the working path.

The issue proposed three runtime mechanisms: a dedicated capture stream, a CPU barrier before capture, and a per-bucket all-ranks-captured replay gate. Each defends against divergence the TP step protocol already excludes — the graph key is a pure function of the broadcast StepCommand (the controller picks the bucket and clones one command to every rank) and the step lock already prevents any rank from starting step N+1 before all ranks finish N, so under lazy capture every rank hits a bucket's first capture on the same global step. The dedicated capture stream would also conflict with the stream-binding invariant documented in cuda_graph.rs: a graph is tied to the stream it was captured on, so capturing on a non-decode stream silently strips the Green Context SM partition from graphs_split.

What actually broke in #481 is narrower: mid-serving capture overlaps the previous step's eager NCCL tail. Removing the mid-serving capture window removes the bug. So this takes the shape glm52 has shipped since #570 — pre-capture every reachable (bucket, attention-path) decode graph at startup, per rank, right after all rank workers spawn, while the ranks are trivially in lock-step and no eager collectives are in flight. Uniform captured-ness across ranks becomes an invariant established at construction rather than a property maintained at runtime. Serving decode-only steps (batch_decode) are then pure replay; mixed prefill+decode steps keep running the eager unified path, which never captures. The single-GPU lazy-capture path keeps its behaviour: openinfer-core gains one additive method (capture_only, which records + instantiates + uploads without launching), and run_or_capture is unchanged.

Supporting pieces: an eager warm-up collective per rank before the first capture (NCCL lazy connect must stay outside cuStreamBeginCapture, and since NCCL >= 2.22 connects per size-selected algorithm the warm-up runs one all-reduce per bucket message size); a fail-loud check in the serving decode arm so a missed shape can never lazily capture under TP; decode-overlap rejected under TP (the graphs_split cache has no cross-rank alignment story); and lane buffers dropping before the model so captured graphs are destroyed before the NCCL communicator (NCCL >= 2.26 comm teardown polls until referencing graphs are gone).

The first sweep fused capture + instantiate + first-launch per bucket (the single-GPU run_or_capture shape) and deadlocked on two GPUs: the launched rank's captured all-reduce spins on its peer while the peer is still inside capture/instantiate driver calls that hold the process-wide driver lock — and the first launch of a never-uploaded exec performs the implicit upload, which is more of the same. The sweep is therefore two-phase per bucket under a controller barrier: every rank records + instantiates + cuGraphUploads the bucket's graph before any rank launches it, and the launch phase is a pure enqueue + sync. NCCL has no device-side timeout, so a startup watchdog turns a wedged sweep into a loud abort() (not exit(), which would re-enter the same driver lock in cudart teardown); separately, a rank that errors mid-sweep fails startup and logs which peers were not yet collected and may still be wedged. Applying the same reasoning backward, the TP memory-profile decode exercise now runs eager instead of lazily capturing per rank uncoordinated.

Measured — Dual-GH200 (aarch64, sm_90), Qwen3-4B, tp-size 2

Startup pre-capture covers all 36 (bucket, attention-path) graphs per rank in 0.74s; graph-exec residency costs ~500 MiB per GPU over the eager server. Because the TP memory profile now runs the decode exercise eager (above), that residency lands after profiling rather than inside the profiled KV budget — measured, graph-on used 88002 MiB against a profiled 87552 MiB request, so it draws ~500 MiB from the 0.90-utilization headroom rather than OOMing. Budgeting the TP graph-exec into the profile (and correcting the --gpu-memory-utilization help text, which still says profiling accounts for CUDA-graph capture — true single-GPU, not under TP) is a small follow-up, not a blocker here.

Graph-on vs eager serving (vllm bench serve, random 128 in / 256 out):

concurrency	mean TPOT (ms) graph / eager	output tok/s graph / eager
1	5.77 / 6.35	173 / 156
4	5.79 / 6.38	642 / 595
8	5.92 / 6.48	1206 / 1092

Decode TPOT improves ~9% and throughput 8-11% across concurrencies.

Testing

• tests/tp_concurrent_decode.rs passes graph-on with assertions unchanged (the fix(qwen3): disable cuda graph under tensor parallelism #496 acceptance gate), 13.3s on two GH200s — the first exercise of NCCL collectives under cuStreamBeginCapture on this stack.
• hf_golden_gate gains a graph-on TP=2 pass (from_runtime(true, [0, 1]) — the startup sweep captures each rank's decode graphs, serving replays them) checked against the HF golden at the same tolerances as the eager TP suite. The graph-on TP logits match: worst head delta 0.375 / 0.162 at the straddle buckets, inside the eager TP suite's 0.375 noise floor. This closes the "TP CUDA-graph deferred" note the golden gate carried.
• Single-GPU regression: hf_golden_gate (single-GPU eager + graph passes) and the workspace lib suite green on sm_89; clippy clean (no new warnings vs base).
• Both fix(qwen3): disable cuda graph under tensor parallelism #496 guards deleted.
• Everything above ran on this branch rebased onto main, so the numbers are the merge target, not a stale base.

Type of Change

• Bug fix (non-breaking change which fixes an issue)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 347a72b848

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".