Repository navigation
[FlyDSL] Optimize PA partition reduction (split from #5949) - #6619
Conversation
Use host-bounded register ladders and a dual register/LDS reducer, with paired loads, launch-shape partition groups and XCD-aware CTA order. Preserve the existing main attention implementation and reusable planner scratch. Add FP32-oracle tests for mixed and empty partition counts, sinks, poisoned scratch and graph replay across the supported head sizes and output dtypes. Co-authored-by: Sami Remes <samremes@amd.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
The rewritten reducer now lives in #6619. This branch keeps the gfx950 head-256 pipe, including per-token scales. Co-authored-by: Cursor <cursoragent@cursor.com>
Main's FlyDSL PA reducer spends unnecessary work on a large partition cap, and reduction becomes expensive for long contexts. Extract the rewritten partition reducer from #5949 (
179193716f3e7b6136e9f98f85c8f3c0292a963b; reducer introduced in9bcb97c9d) into a separate change based on main42c2b1427.The reducer uses exact register specializations through 64 partitions, a runtime register ladder for up to 32 live partitions under larger bounds, and LDS partition groups for larger live counts. BF16/FP16 paired loads, launch-shape-based group selection, and XCD-aware CTA order reduce work and scheduling overhead. The host bounds specialization by context tiles, plan capacity and sliding-window tiles while retaining reusable plan/scratch buffers.
Scope is the reducer kernel and its host dispatch/call site. Attention tile compilation/launch, attention pipeline, planner, tuning defaults and NHD support are unchanged. The reducer kernel matches #5949 verbatim. Credits to @samremes for the reducer implementation.
Performance
MI355X/gfx950, 256 CUs, FlyDSL 0.4.0, PyTorch 2.12.0+rocm7.14.0, HIP 7.14.60850. These are fresh measurements of this extracted branch against main, using warm CUDA graphs and medians across three independent processes. All times are microseconds.
Reducer-only: D256, BF16 partials/output, QL4, GQA8, Hkv1. Both columns retain compile-time partition cap 256, including the low-live-count rows. Same packed scratch and input tensors; 100 reducer calls per graph, seven timing rounds with rotating implementation order.
These high-cap cases also cover the earlier regression reported in #5949: none regresses versus main. A tighter valid host bound can further reduce short-count latency.
Full PA uses main's existing attention implementation in both columns. Each graph contains 50 steps, including GPU plan refresh, attention and reduce, with planner defaults
max_partitions=256,workgroup_budget=512. FP8 KV with per-token scales, BF16 Q/output, D256, GQA8, Hkv1; host context bound equals live context. The reduction timing is measured separately on attention-produced scratch.The full PA improvement comes from reduction and its specialization bound; attention code is unchanged. Isolated reducer numbers are warm-scratch measurements; each full step writes scratch again in attention. In particular, main's long-context reducer has a substantially higher latency inside the full step than when repeatedly reducing the same scratch. Do not subtract isolated reducer time from full-step time to infer attention latency.
Profiling the long-context 50-step graphs corroborates that attribution. Attention kernel time stays essentially equal while the reducer accounts for the full-step improvement. Per-process medians of three profiler captures, then medians across three processes:
Validation
op_tests/test_flydsl_pa_decode.pysuite: 106 passed on MI355X/gfx950 and 106 passed on MI308X/gfx942 (80 CUs).git diff --checkpass.GPU suites run with
AITER_AOT_IMPORT=1and equivalent pure-Torch quantization helpers injected into the existing test harness to avoid unrelated C++/CK builds. PA implementation and test assertions are unchanged by that harness.