[AMD][gfx95] Fill the chunked-prefill compute budget exactly - #32888
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
A prefill batch never reaches `chunked_prefill_size`. Two independent page roundings leak tokens out of it: `_update_prefill_budget` charges the page-ceiled extend length against `rem_chunk_tokens` and `rem_input_tokens`, and `add_one_req` floors the last request's truncation back down to a page. Every request admitted whole leaks up to page_size-1 tokens, so with page_size=64 a 16384-token chunk runs a few dozen tokens short. That shortfall is not free on gfx95. The MLA absorb bmm is an aiter Triton kernel whose `EVEN_MN` heuristic is a compile-time constexpr keyed on `M % BLOCK_SIZE_M == 0`, so a misaligned M compiles to a different specialization -- 3819 instructions against 1562 for identical MFMA work, because the masked store degenerates into 64 separate 16-bit stores and `% M` expands to a full division sequence. Measured on that kernel: 311.2 us at M=16368 vs 155.8 us at M=16384. Add `compute_charge` to `_update_prefill_budget` so the compute budgets (counted in forward-pass tokens) move independently of the KV budgets (counted in pages, and left page-ceiled). Both compute budgets have to switch together: they are typically configured to the same value, so leaving either one ceiled makes it hit zero first and stop admission while the other still has the rounding slack unspent. The last request of a batch then takes the remainder without flooring it to a page. Neither floor is a correctness requirement: `alloc_extend` handles a partial leading/trailing page explicitly, DSA builds its page tables per token, and `add_chunked_req` already commits non-page-aligned continuation chunks. The page granularity only bought conservative KV accounting, which is unchanged. An SWA-capped chunk stays capped -- the exact-fill path takes `chunk_tokens_limit`, not `rem_chunk_tokens`. Gated on gfx95 and on SGLANG_EXACT_CHUNK_FILL, so CUDA and other AMD parts keep upstream behaviour. The dLLM admission path is excluded. Side effect on the gated path: the `#new-token` metric now reports the tokens the forward actually runs, which is the TODO(lsyin) next to `log_input_tokens`. Measured on GLM-5.2-MXFP4, MI355X TP4, in8192/out1024: dense GEMM 163.5 -> 136.6 ms, of which 25.8 ms is the two absorb bmm calls. Tensile/hipBLASLt GEMMs are nearly indifferent to the alignment (1%); they have no all-or-nothing constexpr branch, they just run one partial tile. Co-authored-by: Cursor <cursoragent@cursor.com>
ef9d38f to
3f3ebf3
Compare
Except for base-c-test-8-gpu-b300, which is still in the queue, all other PRs Test Base have passed the CI tests. |
|
@amd-bot ci-status |
HaiShaw
left a comment
HaiShaw
left a comment
There was a problem hiding this comment.
@Jacob0226 please extend this to cover gfx942 next, gfx1250 later.
|
@merrymercy @hnyls2002 please have a review. |
Will cover gfx942 in a follow-up PR, gfx1250 after that. |
…hunk-fill Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # python/sglang/srt/managers/schedule_policy.py
…hunk-fill # Conflicts: # python/sglang/srt/managers/schedule_policy.py
|
/rerun-failed-ci |
|
@amd-bot ci-status |
CI Status for PR #32888Merge verdict: No PR-caused failures found. This PR's new exact-chunk-fill logic is gated behind Note The changed code path ( Caution NVIDIA "PR Test Base" is incomplete. Run 34426366171 was cancelled by fast-fail after base-b failures; Changed files: Executed CI failure attribution: AMD: 10 failures (0 related) · Others: 6 failures (0 related) · Fast-fail/gate cascade jobs collapsed into their root causes. AMD Executed Failures
Other Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
…hunk-fill Upstream independently fixed the TODO(lsyin) next to log_input_tokens by introducing raw_extend_input_len, so the metric now reports the tokens the forward actually runs on every platform. Take upstream's version and drop this branch's gated-path-only equivalent; on every exact-fill call site raw_extend_input_len equals compute_charge, so the gfx95 path is unchanged. Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # python/sglang/srt/managers/schedule_policy.py
|
@merrymercy @hnyls2002 @xiezhq-hermann Please have a review. |
|
/rerun-failed-ci |
|
@hnyls2002 please review, no functional change made to non-gfx/hip. |
CUDA Base CI all green
|
|
/rerun-failed-ci |
CUDA Extra CI all green
|
…ject#32888) Co-authored-by: Thomas Wang <thomawan@amd.com>
6 participants
Summary
A prefill batch never reaches
--chunked-prefill-size:_update_prefill_budgetcharges the page-ceiled extend length againstrem_chunk_tokensandrem_input_tokens, andadd_one_reqfloors the last truncation back to a page. On gfx95 that shortfall is expensive — the aiter MLA absorb bmm picks itsEVEN_MNspecialization from a compile-timeconstexpronM % BLOCK_SIZE_M, so a misaligned M runs a different compiled binary, not the same one with a partial tile.toks=16350toks=16384compute_chargebills the compute budgets in raw tokens; the KV budgets stay page-ceiled. Below, the ledger spends 16384 while the forward runs 16328 (page_size=64, 8100-token prompts) — only whole admissions leak.Scope
gfx95 (MI355X) chunked prefill, gated on
SGLANG_EXACT_CHUNK_FILLandis_gfx95_supported(). CUDA and other AMD parts keep upstream behaviour on every path, and the dLLM path is excluded. ROCm/aiter#4453 deepens the kernel win but does not gate this PR.schedule_policy.pyis common code, so the CUDA no-op is measured rather than asserted. Replaying 48 admission scenarios (page_size1/8/64 × chunk 2048/16384 × four prompt mixes, 114 admissions) againstupstream/mainwith the gate off reproduces every budget counter byte-for-byte; forcing it on moves 73 lines, andpage_size=1is identical either way.Test plan
Base —
rocm/sgl-dev:v0.5.15.post1-rocm720-mi35x-20260714plus six AMD PRs (#30519, #30575, #30715, #30808, #31323, #31324). This PR — that base plus this PR and ROCm/aiter#4453; nothing else moves.Kernel level, same config:
End to end, GLM-5.2-MXFP4, MI355X TP4, in8192 / out1024:
Mean TTFT improves monotonically with concurrency: the saving is per prefill forward, so a deeper queue waits behind more of them.
Prefill kernel attribution, one forward at in8192 / concurrency 64:
Only dense/linear GEMM improves, where the mechanism predicts; everything else shifts within noise.
Accuracy: GSM8K 0.933 → 0.927, single scoring pass each — within what one pass resolves.
CI States
Latest PR Test (Base): ✅ Run #34549992078
Latest PR Test (Extra): ✅ Run #34793772179
Latest PR Test (AMD ROCm 10): ❌ Run #34549992056