Repository navigation
[Triton/Gluon] [Perf] [KDA] Optimize the gfx950 FlashKDA prefill for GLM-5.3-Flash - #7360
waqahmed-amd-fi wants to merge 3 commits into
Conversation
vLLM passes q, k and v as column slices of one conv output, so flash_kda_fwd packed them on gfx950 before K1 and passes A and C, about 45 us per layer at 8192 tokens. The gfx950 Gluon kernels now read them at their own token stride, as the gfx1250 ones do. - Take `stride_q`, `stride_k` and `stride_g` in the gfx950 K1, and `stride_v` in passes A and C - Decide each operand's packing after the routes are known; pack rows that are not 16-byte aligned or need more than 32-bit offsets - Check gfx950 in-place reads in `test_strided_operands_match_packed`
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision. Questions to answer in the PR description
Edit the Missing or invalid answers
Shadow mode reports what would block but keeps the status successful. Set |
Summary
On gfx950,
flash_kda_fwdnow reads strided q, k, v and g in place in its Gluon kernels. Before, itpacked these operands into contiguous copies. The results do not change: they are bit for bit
the same as with packed inputs.
This PR is stacked on #6912, because the change uses the
segment plan that #6912 adds. Review only commit
02689fad. vLLM uses
it in vllm-project/vllm#65224.
Motivation
vLLM's GLM-5.3-Flash KDA layer gives FlashKDA q, k and v as column slices of one conv output. Their
token stride is 3 × H × D. g is also a strided view. On gfx950,
flash_kda_fwdpacked the fouroperands before K1 and before passes A and C. The gfx1250 Gluon kernels already read them at their
own stride.
On MI355X, at TP4, the copies take 53.3 us in each of the 34 KDA layers of an 8192-token prefill
step. That is 1.8 ms per step.
Changes
flash_kda_k1.py) takesstride_q,stride_kandstride_g.flash_kda_k2.py) takestride_v.flash_kda_fwddecides how to pass each operand after it knows the kernel routes. It still packsan operand that a Triton kernel reads. It also packs an operand whose rows are not 16-byte aligned,
or whose span needs offsets larger than 32 bits.
chunk_delta_attndocstring says that gfx950 also reads these operands in place.test_strided_operands_match_packedchecks the in-place reads on gfx950 too.The gfx1250 kernels, the Triton kernels and the other archs do not change. Contiguous inputs take
the same path as before.
Performance (if applicable)
GLM-5.3-Flash MXFP4 (
amd/GLM-5.3-Flash-Quark-MXFP4) served by vLLM, rank 0, torch profiler over60000/600 at concurrency 16. The step is one 8192-token prefill with 8-15 decodes, the most common
prefill step there. Both arms use the same vLLM, with #6912 in AITER; only this commit changes.
The FlashKDA kernels take about the same time with strided inputs: K1 57.3 → 56.1 us, pass A
41.1 → 42.3 us and pass C 40.5 → 41.0 us per layer. The copies left are the
betapack and othersoutside FlashKDA. Over nine serving points, the
change alone is within run-to-run noise: output throughput moves -1.8% to +1.4% (geomean +0.2%),
and median TPOT -1.3% to +0.5% (geomean -0.4%). The long-context points gain most.
Together with vllm#65224, this branch serves GLM-5.3-Flash faster than vLLM and AITER
main. Overthe same nine serving points, output throughput goes up by 11.5% and median TPOT goes down by 10.8%
(geomeans). These numbers include vLLM's #63568 and #63612, which vllm#65224 carries. vllm#65224
gives the full tables.
Testing
The change is for gfx950 only, so it was not run on MI250X or MI300X. See "Validation and unverified
scope" below.
Documentation
Dependencies
Breaking Changes
None.
PR Self-Review
Reviewed revision
1373434cfcaaBehavior and intent
On gfx950, FlashKDA's Gluon kernels read strided q, k, v and g where they are. Before, the wrapper
copied them to contiguous tensors first. The copies cost 1.8 ms in each 8192-token prefill step of
GLM-5.3-Flash. The output does not change.
Validation and unverified scope
On 4x MI355X (gfx950), in
vllm/vllm-openai-rocm:nightly-76c6f3f2a95f6049657db1e513133b824a5c8175(torch 2.13.0+rocm10.0.0, Triton 3.8.0), from the AITER checkout:
Result: 157 passed, 1 skipped.
The performance numbers above come from vLLM serving GLM-5.3-Flash on MI355X at TP4. GSM8K, GPQA
Diamond and NIAH at 65536 do not change; vllm#65224 gives the scores.
Not verified: gfx1250, MI300X and MI250X. The gfx1250 kernels and the Triton kernels do not change.
PR-specific question
Before this PR,
test_strided_operands_match_packedchecked in-place reads only on gfx1250. Ongfx950 it checked only that the output was equal. Its new check on gfx950 fails before this PR:
the Gluon launchers got packed copies in all layouts, also in the aligned
bandslayout. No othertest checks the storage that the gfx950 launchers get.
Changed op tests
op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py: this PR changestest_strided_operands_match_packed. On gfx950 it now checks that the K1 and pass A and Claunchers get the storage of the strided views for the aligned layout, and packed copies for the
others. The helper that records the storage now reads the current arch. [Triton/Gluon] KDA: accept cu_seqlens_cpu, upload varlen metadata async #6912 adds the
cu_seqlens_cputests in this file; [Triton/Gluon] KDA: accept cu_seqlens_cpu, upload varlen metadata async #6912 explains them.op_tests/op_benchmarks/triton/bench_flash_kda.py: [Triton/Gluon] KDA: accept cu_seqlens_cpu, upload varlen metadata async #6912 adds its varlen mode. This PR's owncommit does not change the file.