feat(qwen35): shard GDR state per TP rank + batched eager/graph decode under TP (+27B benchmark validation) - #946
Open
Ma1oneZhang wants to merge 5 commits into
Conversation
Phase 2b ported onto pegainfer-project#870: recurrent/conv state, GDR scratch, and the linear-attention weight surface are allocated and addressed per rank instead of replicated, which is what makes 27B TP2 fit on 2x48 GB cards. - weight_loader: additive stitch/shard loaders (2D row stitch for the fused qkv [q|k|v] segments, 1D element stitch for conv1d channels, bf16/f32 1D shards for dt_bias/A_log) - config: local_linear_* accessors mirroring the full-attn TP style; linear head divisibility fails closed in TensorParallelConfig - weights: per-rank stitched shards for in_proj_qkv/conv1d (head-local slices per segment), row shards for z/b/a, col shard for out_proj, dt_bias/A_log sliced, norm_weight kept replicated (head-shared); loader reserve uses TP-aware estimates - recurrent_state/decode_buffers/prefill_buffers: state and GDR scratch at local value-head/qkv sizes; capacity math derives from local allocation_bytes - batch_decode/prefill: local head counts into the GDR decode/conv/ Triton-AOT prefill chains, gated RMSNorm at local v heads, all-reduce after linear out_proj; batch_decode_full_attention_via_prefill is now TP-local so eager decode routes 27B TP2 group-6 full attention through prefill (was FlashInfer Unsupported group_size: 6) - tp_executor: worker capacity math and per-request state use the rank-local sizes; decode rows still run as a per-request bs=1 loop (batched in a follow-up) Recurrent/conv state is never all-reduced. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Port of the batched eager TP decode step onto pegainfer-project#870's worker structure: decode rows in one command now run as ONE batched forward per step on every rank plus one batched rank-0 sampling pass, instead of a per-request bs=1 loop. - run_decode_batch resolves every decode row's worker state in command order, builds a step-scoped LinearStatePointerTables over the whole batch (from_recurrent_refs(..., bs, ...)), runs one batch_decode_eager_logits forward, then rank 0 snapshots all requested logprob rows before one batched select_batch over per-row params - execute_decode_rows (used by both decode-only and unified steps) calls run_decode_batch once; per-row results fan out in command order - TpRequestState.linear_pointer_tables (capacity-1, decode-only) removed; ensure_prefill_state no longer builds it. The step-scoped table is rebuilt every step, so swap_remove retirement can't stale it Seeded rows keep per-row semantics: select_batch isolates each seeded row into its own single-row philox call keyed on (request seed, step 0), so seeded output stays independent of batch composition. Unseeded rows decorrelate via the per-step command seed, same as the single-GPU batched path. Reference (27B TP2, 2x RTX 4090, eager): 16 concurrent 256-token completions aggregate 24.9 -> 292.3 tok/s; single-request unchanged. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
…eval - tp-implementation.md: Phase 2a/2b + Step 3 batched-decode landing record - benchmarks/qwen35-27b-tp2-knowledge-eval.md: MMLU-Redux 94.09 vs 93.2, C-Eval 88.11 vs 90.5 (full runs, truncation rerun-merged); in band - scripts/eval_mc.py + eval_rerun_truncated.py: benchmark runner (OpenCompass/lm-eval recipe replicas over the chat completions API) - developer-onboarding: NCCL libnccl.so dlopen note Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Ma1oneZhang
force-pushed
the
feat/qwen35-tp2-batched-decode
branch
from
August 20, 2026 16:00
6b44a25 to
f56ce51
Compare
Ma1oneZhang
changed the title
Ma1oneZhang
marked this pull request as ready for review
August 20, 2026 16:00
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>
Ma1oneZhang
changed the title
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
1 participant
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.
Summary
Rebased onto current main (#870 landed its own Phase 1/2a). This PR carries the deltas #870 lacks, plus 27B benchmark validation:
2530dfc2): rank-local slices end to end (weight_loader stitch/shard loaders, per-rank stitched qkv/conv1d + row/col shards, state/buffers at local sizes, all-reduce after local linearout_proj, TP-local group-6 decode reroute). Acceptance: 27B TP2 now fits on 2×48 GB cards (Phase-1 replicated state OOMed).b3aa5e11): feat(qwen35): tp2 phase 2a #870'sexecute_decode_rowsran one bs=1 forward per request per tick; this addsrun_decode_batch— onebatch_decode_eager_logitsover all decode rows with a step-scoped pointer table, one batched rank-0select_batch, per-row fan-out. Measured: 16-concurrency aggregate 24.9 → 292.3 tok/s.0570313f): startup pre-capture sweep ported from qwen3 (warmup collectives per bucket message size, Capture/Launch per bucket, Finalize assert, 600 s watchdog), slot-indexed per-rank recurrent/conv state with scheduler-ownedslot_idx+ D2D compaction on retire, serve-time replay-only. Gate:enable_cuda_graph && local_decode_group_is_compiled(tp)— 4B/9B TP2 capture; 27B TP2 (group-6, TP-invariant) stays batched eager until group-6 batch-decode kernels exist. 9B TP2 serving bench: steady output 706 → 767 tok/s (+8.7%), TPOT −8.8%.b376772a,f56ce511): rebase + P2c landing records intp-implementation.md; 27B knowledge-benchmark parity snapshot indocs/benchmarks/qwen35-27b-tp2-knowledge-eval.md; runner scriptsscripts/eval_mc.py+scripts/eval_rerun_truncated.py.Validation (2× RTX 4090, --features qwen35, on this branch)
cargo check/fmtclean; lib suite 105/105Knowledge-benchmark parity (27B, chat completions, thinking on)
Scores were measured on the pre-rebase f4c6678 line; both lines pass identical logits gates, but rerun on this branch before citing parity externally (noted in the doc).
Follow-ups
Group-6 batch-decode kernels (unlocks 27B TP2 graphs), arbitrary TP device ordinals.