[Kernel] Migrate SWA decode to tile programming

[fsx950223]

Summary

Sliding-window paged-attention decode now uses FlyDSL layout views and copy atoms for global memory traffic while preserving the compact raw-pointer launch ABI and hardware-specific FP8 MFMA/LDS paths. Supported production configurations retain their accuracy and effectively neutral aggregate performance.

Motivation

The SWA path still bypassed FlyDSL's layout algebra through legacy buffer helpers, making its addressing and data movement inconsistent with current kernel authoring patterns. Its large-partition reduction fallback also could not compile because runtime branch values lacked valid SSA initialization.

Changes

• Route query, scale, metadata, partial-output, and reduction traffic through buffer tensors, logical divides, register fragments, and copy atoms.
• Keep packed K/V cache accesses as aligned 128-bit universal copies and retain the existing FP8 MFMA and LDS synchronization strategy.
• Fix reduction for more than 64 context partitions and reject unsupported head dimensions instead of allowing invalid results.

Performance

GPU-only rocprofv3 medians on AMD Instinct MI308X (gfx942), paired across two reversed-order A/B rounds. Path time includes the SWA kernel and reduction; lower is better.

Configuration	Original	Tile programming	Delta
Q1, per-token	23.77 us	23.84 us	+0.29%
Q2, per-token	30.02 us	29.42 us	-1.98%
Q3, per-token	47.66 us	48.21 us	+1.16%
Q4, per-token	91.68 us	91.97 us	+0.32%
Q1, per-tensor	22.31 us	22.02 us	-1.30%
Q3, per-token, non-transposed V	64.21 us	67.62 us	+5.31%
Short-window Q2	27.62 us	27.32 us	-1.09%
Head dimension 64, Q1	20.32 us	20.30 us	-0.10%

Geometric mean across the matrix is +0.31%; the transposed-V configurations used by the main regression matrix are effectively unchanged.

Testing

• FlyDSL Python style gate
• Official SWA accuracy matrix: Q1-Q4 x 8/16 query heads, batch 128, variable lengths
• Per-token and per-tensor scales; transposed and non-transposed V
• Sliding-window boundaries at 1, 255, 256, and 257 tokens
• Reduction with 1, 2, 64, and 65 partitions using bf16, f16, and f32
• rocprofv3 performance comparison on MI308X

Dependencies

• No new third-party dependencies added

Breaking Changes

None for supported configurations. SWA now fails fast for head dimensions other than 64 or 128; the previously accepted 256 path produced invalid results.

Made with Cursor

[Copilot]

Pull request overview

Migrates the sliding-window paged-attention (SWA) decode path to FlyDSL tile programming primitives (layout views + copy atoms) while keeping the existing raw-pointer launch ABI and preserving the existing FP8 MFMA/LDS compute path. The PR also tightens supported configurations and fixes the large-partition reduction fallback to compile correctly.

Changes:

• Replaces legacy global/buffer helper loads/stores with buffer tensors, logical divides, and fx.copy-based copy atoms for most global traffic.
• Keeps K/V cache reads as aligned 128-bit universal copies while routing query/scale/metadata/partial-output/reduction through tile views.
• Fixes the >64-partition reduction fallback SSA/branching and fails fast for unsupported head_dim values (now only 64 or 128).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.