[ROCm] Add FlyDSL paged-attention decode backend (dense + fp8) with benchmark - #463
[ROCm] Add FlyDSL paged-attention decode backend (dense + fp8) with benchmark#463avbokovoy wants to merge 2 commits into
Conversation
Replace the CK paged-attention decode operators with FlyDSL kernels and add a native-fp8 decode path, wired into the fmha op registry (flydsl_decoder / flydsl_splitk). Kernels (mslk/attention/fmha/flydsl/): - pa_decode_gfx950: primary head-packed MFMA fast path (double-buffered wide V load, ds_read_tr16_b64), gfx950, GQA ratio in [1,16]. - pa_decode_gfx950_coop: gfx950 cooperative-DMA kernel for ratios the primary path can't head-pack. - pa_decode_generic: arch-generic per-warp fallback (gfx942 + off-gfx950). - pa_decode_fp8: native-fp8 (e4m3fn + symmetric per-token scale) paged decode, with a per-call quantizing adapter and a guarded public dispatcher. MQA/GQA, D=128/256, any context length. Stream threaded through compute + reduce launches so the kernel is CUDA-graph capturable. fp8-KV is opt-in per call via Inputs.quantize_kv_to_fp8. - pa_decode_reduce: split-K partial combine. - utils: shared low-level FlyDSL helpers (WARP_SIZE, dpp/wave-reduce, exp2/rcp/max). AOT: dense (generic/gfx950/coop), split-K reduce, and native-fp8 kernels are all registered in mslk/flydsl/aot.py and precompiled into the bundled cache. Also: - Fix the Triton fp8 decode to use OCP e4m3fn on gfx950 (not fnuz). - Import FlyDSL helpers from the mslk.flydsl package (common/jit), matching the upstream flash-attention layout. - bench/attn/decoder_bench.py: FlyDSL + Triton dense/fp8 backends, eager timing via the shared do_bench, real CUDA-graph timing with empty-graph detection, and subprocess isolation per shape. - ROCm CI path triggers for the decode backend + tests.
|
@q10 has imported this pull request. If you are a Meta employee, you can view this in D114929079. |
cthi
left a comment
cthi
left a comment
There was a problem hiding this comment.
Thanks, I left some high level questions. Mainly wondering about the comments arounds the benchmarking.
| * eager (``--no-cuda-graph``, default): shared ``mslk.bench.common.utils.do_bench``. | ||
| * graph (``--cuda-graph``): HIP graph capture + replay (removes launch overhead). | ||
|
|
||
| gfx950 gotchas (see the runners / _bench_ms_graph for detail): |
There was a problem hiding this comment.
Is these "gotchas" for graph benchmark specific stuff? Or more like limitations we should follow up on? Graph is actually more important than eager, so we should always default our benchmarking/optimization/enablement to be graph-enabled first.
|
|
||
| """Shared low-level FlyDSL helpers for attention kernels. | ||
|
|
||
| Only pip `flydsl==0.2.2` is imported (no ~/FlyDSL/kernels imports). |
There was a problem hiding this comment.
I don't understand this comment, is it left over from a debugging?
|
|
||
| # pyre-strict | ||
|
|
||
| """Shared low-level FlyDSL helpers for attention kernels. |
There was a problem hiding this comment.
These don't look attention specific, should they be moved into some more common flydsl/utils.py?
| * fp8 runners cache the ``flyc.compile`` CompiledFunction so timing is kernel-only | ||
| and scales with KV (calling the public dispatcher directly pays ~0.38ms/call of | ||
| JIT dispatch that hides the ~0.02ms kernel — flat, meaningless numbers). | ||
| * flydsl_fp8 IS graph-capturable (its kernels thread the capture stream); dense |
There was a problem hiding this comment.
Can you explain this "empty graph" issue? It seems we modify a lot of the benchmarking to handle this.
| - 'test/attention/flydsl/**' | ||
| - 'test/flydsl/**' | ||
| # FlyDSL paged-attention decode backend (fmha op layer) and tests | ||
| - 'mslk/attention/fmha/flydsl/**' |
There was a problem hiding this comment.
Btw I have fixed the flakiness in the grouped gemm tests, so if you rebase those issues should go away, and we should be able to get some clean CI signal.
2 participants
Summary
Replaces the CK paged-attention decode operators with a FlyDSL backend and adds a native-fp8 decode path, wired into the existing fmha op registry (flydsl_decoder / flydsl_splitk) with backward-compatible op aliases. FlyDSL becomes the sole decode backend across gfx942 and gfx950; the CK operator path is removed.
Kernels
Dense f16/bf16 decode uses three kernels selected by shape: pa_decode_gfx950 (head-packed MFMA + double-buffered wide V load, ds_read_tr16_b64; gfx950, GQA ratio 1..16), pa_decode_gfx950_coop (gfx950 cooperative-DMA for ratios that cannot head-pack), and pa_decode_generic (arch-generic per-warp fallback for gfx942 and off-gfx950). pa_decode_fp8 adds native fp8 (e4m3fn + symmetric per-token scale) paged decode covering MQA/GQA, D=128/256, and any context length; it is opt-in per call via Inputs.quantize_kv_to_fp8 and is CUDA-graph capturable. All kernels (dense, split-K reduce, and fp8) are AOT-registered and precompiled into the bundled cache.
Performance
Measured on gfx950 (MI350X). Dense f16/bf16: FlyDSL is 1.6-10x faster than the old CK decoder (CK uses no matrix cores) and within 1.05-1.37x of Triton (near parity; Triton is marginally faster on raw kernel time). Native fp8: FlyDSL is 1.8-2.5x faster than Triton fp8 and also beats Triton dense on most shapes. Correctness is verified against a torch reference for all backends.
Other changes
Fixes the Triton fp8 decode to use OCP e4m3fn on gfx950 (was hardcoded fnuz, which mis-decoded gfx950 caches). Aligns FlyDSL helper imports with the mslk.flydsl package layout. Adds bench/attn/decoder_bench.py (FlyDSL vs Triton, dense and fp8) with eager timing via the shared do_bench, real CUDA-graph timing with empty-graph detection, and per-shape subprocess isolation. Adds ROCm CI path triggers for the decode backend and its tests.
Test plan