feat(qwen35): cuda graph decode under TP (split 3/4 of #946) - #1005
Merged
FeathBow merged 6 commits intoSep 7, 2026
Conversation
This was referenced Aug 30, 2026
Contributor
Author
|
@codex please review |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Ma1oneZhang
force-pushed
the
feat/qwen35-tp-cuda-graph
branch
6 times, most recently
from
September 7, 2026 11:36
1ceef18 to
8e7da56
Compare
Collaborator
|
@codex review please |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e7da56b30
ℹ️ 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".
Ma1oneZhang
force-pushed
the
feat/qwen35-tp-cuda-graph
branch
from
September 7, 2026 11:59
8e7da56 to
e8a2449
Compare
P2c from docs/models/qwen35/tp-design.md. TP decode captures/replays per-bucket CUDA Graphs when --cuda-graph is set AND the TP-local decode GQA group has a compiled kernel (4B/9B TP2); 27B TP2 (group 6) keeps the batched eager path byte-for-byte under the gate. - Gate: drop the fail-closed TP+graph rejections in config.rs/lib.rs/ tp_executor.rs; log once when graph was requested but the group gate keeps decode eager. - State: scheduler owns dense decode slots (slot_idx on TP decode rows, slot_for_new_request at promote, compaction_after_retire on retire); workers hold a fixed-address BatchDecodeGraphState plus slot_map, D2D copy prefill state into the slot on the first decode row, and apply DropRequest compactions via move_slot_within with occupancy assertions (poison on mismatch). - Capture/replay: startup pre-capture sweep ported from qwen3 (Warmup via new warmup_tp_collective, Capture/Launch per bucket, Finalize) behind a 600 s abort watchdog; batch_decode_graph gains DecodeGraphUse (Serve/CaptureOnly/Replay); TP serving is replay-only; graphs drop before the NCCL comm. - Convenience executor API (model-local tests) keeps a slot tracker mirroring Qwen35Executor; scheduler flows pass explicit slots. Validation (2x RTX 4090, venv NCCL): lib 105/105; 9B TP2 HF gates eager + graph (sequential replay, bucket-straddling, post-compaction) pass within existing tolerances; 9B TP2 scheduler e2e eager+graph pass; serving_tp2 now launches with graph on; 27B TP2 HF+e2e pass unchanged (group-6 gate keeps eager, graph variant self-skips). Serving benchmark 9B TP2, 16 concurrent x 256 out: 767.15 tok/s graph vs 705.86 eager (+8.7% steady output, TPOT 20.04 vs 21.99 ms). Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
The graph-slot pre-reserve used bucket_for(requested_max_batch), but the graph state is later allocated as bucket_for(effective_max_batch). On a tight-memory rank the oversized reserve could starve the effective recurrent capacity to zero and refuse startup. Iterate the reserve bucket downward against the fitted capacity until it stabilises (the bucket only shrinks, so it converges), and clamp the fitted batch to the reserved bucket so the later bucket_for(effective) allocation never exceeds the reserve. Absorbs the codex review comment on Ma1oneZhang/pegainfer PR pegainfer-project#946 (tp_executor.rs graph_slot_reserve). Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
…group predicate The rank-local GQA group equals the global group under validated head sharding, so the P2B review removed LocalGeometry::local_decode_group_is_compiled. Route the P2c graph gate through Config35::decode_group_is_compiled instead of re-adding the local duplicate; the gate decision was already identical on every rank. Evidence (2x RTX 4090, sm_89): cargo check/clippy --release --all-targets -D warnings clean; qwen35 lib tests 107 passed / 0 failed; cargo fmt clean. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
The rebase onto main lands this branch's TP deltas in the layout pegainfer-project#967/pegainfer-project#968/pegainfer-project#969 introduced, so the record names where the code actually lives: linear-attention sharding in weights/layers.rs on WeightSource, and the TP slot bookkeeping across scheduler/{mod,backend,tp}.rs. Paths only — no invariant changed. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
4 tasks
…et row precapture_bucket allocated a real KvState for every row of the bucket being captured, so a rank whose usable KV pages were fewer than bucket_for(max_batch) failed at startup even though serving would have padded those rows. The sweep now passes one real row and the bucket explicitly through batch_decode_graph_padded; every other row is padding on the pool's reserved padding page, exactly as when serving, so the sweep's page demand is a constant 1 regardless of pool size or bucket. batch_decode_graph is unchanged for its callers (it delegates with bucket_for(bs)); the padded entry point checks the bucket covers bs. Evidence: 9B TP2 HF gates on 2x RTX 4090 (eager, graph incl. 5->8 and 3->4 padded buckets and post-compaction replay, long) 3/3 pass, 39 s. Codex review follow-up on pegainfer-project#1005. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… A/B tp-design.md: the P2c section's "Status:" line becomes dated prose (the doc convention forbids status enums), the stale branch name goes, and the TL;DR gains the landed P2c gate so the header matches the body. tp-implementation.md: the 9B TP2 767-vs-706 tok/s table is marked as the pre-rebase pegainfer-project#946 measurement it is; no throughput claim is made for this stack until the same-context A/B is rerun on it (pegainfer-project#1001). Codex review follow-up on pegainfer-project#1005. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Ma1oneZhang
force-pushed
the
feat/qwen35-tp-cuda-graph
branch
from
September 7, 2026 12:35
e8a2449 to
80c9b14
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Split 3/4 of #946, tracked in #1001. Stacked on #1003 and #1004 (contains their commits). Merge last of the three.
Enables CUDA Graph decode under TP for rank-local GQA groups that have a compiled batch-decode kernel, gated at load time via
LocalGeometry::local_decode_group_is_compiled(); uncompiled groups (e.g. 27B group 6) stay on the batched eager path byte-for-byte. Graph bucket state pre-allocates one fixed-address recurrent slot per bucket position.Also absorbs a codex review fix from #946: the graph-slot pre-reserve is now sized by iterating the effective batch bucket instead of
bucket_for(requested_max_batch), so a tight-memory rank is no longer over-reserved into zero effective capacity, and the laterbucket_for(effective)allocation can never exceed the reserve.Type of Change
Evidence
cargo check -p pegainfer-qwen35 --features qwen35 --all-targets: clean.cargo test -p pegainfer-qwen35 --features qwen35 --lib: 109/109 pass (7 GPU-gated skipped), includingunified_step_decode_matches_graph_decodeand the graph-gate tests.cargo fmt --all -- --check: clean.The graph on/off A/B (+8.7% output tok/s, −8.8% TPOT @ 16 concurrent) was measured 2026-08-20 on the pre-rebase #946 branch; GPU rerun on this rebased stack is tracked in #1001 and not re-claimed here.
Checklist
docs/conventions/coding-style.md).CLAUDE.md).