fix(qwen35): chunk dense-TP prefill at the rank's prefill batch scaled by the TP degree - #725
Open
alpineQ wants to merge 1 commit into
Open
fix(qwen35): chunk dense-TP prefill at the rank's prefill batch scaled by the TP degree#725alpineQ wants to merge 1 commit into
alpineQ wants to merge 1 commit into
Conversation
…d by the TP degree
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
Dense-TP serve fed the prefill 32 tokens at a time, so every kernel below ran at M=32 (under the MMQ batch floor) and the 128 per-layer collectives fired once per 32 tokens instead of once per chunk; the chunk is now the rank's prefill batch scaled by the TP degree, and prefill on 2× gfx1100 goes from 31 % slower than one card to faster than it (8k prompt: 72.1 → 54.7 s wall including load; 33k prompt tp=2: 649 → 773 tok/s, byte-identical output).
What was wrong
ep_serve_qwen35_dense_tp(crates/hipfire-generate/src/qwen.rs) iteratedprompt_ids.chunks(32).forward_prefill_dense_tpre-chunks internally atprefill_max_batch(512 on gfx1100), so the inner chunker was fine — it just never saw more than 32 tokens per call. Two consequences visible in a rocprof kernel trace of the same 8k prompt (tp=2vstp=1):attention_q8_0_flash_prefill_wmmacalls / grid_mmq_full_set)_residual_wmmashould_use_mmqneedsbatch_size >= 128on RDNA3 (arch_caps.rs), so TP silently lost the i8-WMMA MMQ prefill path that single-GPU gets (#60), and issued 32 768 collectives for an 8k prompt instead of 2 048.What changes
qwen35::prefill_max_batch_tp(gpu, tp): the arch default scaled bytp(cap 2048). The prefill attention grid islocal_heads × chunk / M_TILEand TP divides the heads, so a rank running the single-card chunk launchestp×fewer workgroups than one card and starves an already latency-bound kernel; scaling the chunk restores the single-card workgroup count. An explicitHIPFIRE_PREFILL_MAX_BATCHstill wins.ep_serve_qwen35_dense_tpchunks by it instead of 32; a missing EP state is a validation error instead of a silent32fallback.forward_prefill_dense_tpsizes its per-rank prefill scratch to the call (min(cap, tokens.len())) instead of the arch ceiling.PR #662's own table already shows TP2 raising prefill on 5× gfx1201 (462 → 549 tok/s); on 2× gfx1100 it fell 465 → 153 before this change.
Which surface(s) does this touch?
crates/hipfire-generateserve path for dense TPcrates/hipfire-generate/src/qwen.rshipfire-arch-qwen35(qwen35/prefill.rs,qwen35/forward.rs)crates/hipfire-quantizeTest plan
./scripts/no-gpu-ci.sh: Rust check + no-GPU unit tests + env/docs drift check pass; the Python stage reports 838 passed / 6 failed intests/test_mq4c_repack.py, and those 6 fail identically on untouchedmasterin this environment (NixOS: no/bin/bash, and the test module resolves a differentmq4c_repackthan the script) — not touched by this changecargo build --releaseclean (cargo build --release --workspace --all-targets --locked)cargo test --lib --workspacepassesserve_harness.py --mode batteryand--mode chainonqwen3.6:27b(registry fixture, sha86a5f80f…) at--tp 2, plus battery attp=1for the unchanged path — JSON attached belowqwen_dense_tp2_parity(32-step greedy parity vs single GPU) PASSdocs/methodology/perf-benchmarking.md(byte-identical prompt, fresh process per arm); speed-gate baselines are single-GPU and not affectedlocal serve_harness JSON (tp=2 battery + chain, tp=1 battery)
Hardware validation request (optional)
{ "routes": [ {"mode": "battery", "tag": "qwen3.6:27b"}, {"mode": "chain", "tag": "qwen3.6:27b"} ], "claim": "dense TP2 prefill is chunked at the rank prefill batch; tp=1 serve is unchanged. If the gate can drive `--tp 2` on the multi-GPU host, that is the route this change is about." }Architecture-trait change?
No.