[Kernel] Add CDNA SageAttention kernel

[LiuYinfeng01]

Summary

Add a standalone FlyDSL SageAttention V1 kernel for CDNA GPUs.

Motivation

SageAttention uses INT8 Q/K and FP8 V to reduce attention bandwidth and improve dense diffusion-model attention performance. This contribution makes the production attention kernel available directly in the FlyDSL kernel repository.

SageAttention principle and this kernel's role

SageAttention is a mixed-precision attention recipe rather than only an INT8 matrix multiplication. The complete path has two layers:

Layer	Operation	Purpose
Preprocessing / quantization	Optional Hadamard rotation of Q/K	Spreads channel outliers while preserving the floating-point QK product because the normalized Hadamard matrix is orthogonal. This is not implemented by this PR's low-level kernel.
Preprocessing / quantization	K smoothing	Centers K over the sequence dimension, reducing its dynamic range before INT8 quantization.
Preprocessing / quantization	Optional blockwise Q smoothing	Subtracts a Q-block mean before INT8 quantization, reducing Q outliers.
Preprocessing / quantization	Score correction (delta_s)	Restores the score term removed by Q smoothing: delta_s = softmax_scale * q_mean * K_smooth^T.
Preprocessing / quantization	Quantize Q/K to INT8 and V to FP8	Produces per-block Q/K scales, per-channel V scales, and blocked FP8 V storage.
This PR's FlyDSL kernel	INT8 QK, online softmax, FP8 PV	Computes corrected attention with FP32 score/softmax/PV accumulation and writes BF16 output.

For one Q block, the data flow is:

Step	Value
K smoothing	K_smooth = K - mean_sequence(K)
Q smoothing	Q_smooth = Q - mean_block(Q)
Quantization	Q_int8 ~= Q_smooth / q_scale, K_int8 ~= K_smooth / k_scale, V_fp8 ~= V / v_scale
Approximate score	score ~= softmax_scale * (Q_int8 * K_int8^T) * q_scale * k_scale + delta_s
Output	O ~= softmax(score) * V_fp8 * v_scale

The Q correction is exact before quantization:

softmax_scale * Q_smooth * K_smooth^T + delta_s
= softmax_scale * Q * K_smooth^T

Therefore Q smoothing lowers INT8 quantization error without changing the K-smoothed floating-point target. K smoothing deliberately changes every score row only by a key-independent constant when the mean is taken over keys; softmax is invariant to that common row shift. A normalized Hadamard rotation, when enabled by a higher-level wrapper, similarly preserves floating-point QK while making channel magnitudes more uniform before quantization.

This PR starts after preprocessing. Its inputs are already-quantized INT8 Q/K, blocked FP8 V, Q/K/V descale tensors, and an optional score-bias tensor carrying delta_s. The kernel then performs:

1. INT8 QK MFMA with FP32 accumulation.
2. Per-Q-block and per-K-block descale plus softmax scaling.
3. Optional lane/column-exact delta_s addition before softmax.
4. Online FP32 softmax.
5. Native FP8 PV MFMA with FP32 accumulation.
6. V descale and BF16 output conversion.

The production Qwen path uses Q-smoothing blocks of 128 query tokens (BLKQ=128) independently of this kernel's attention scheduling tile (BLOCK_M=256). A 256-row FlyDSL workgroup therefore consumes the correction data for two 128-row Q-smoothing blocks. BLOCK_M and Hadamard BLOCK_R are unrelated: the former tiles sequence rows; the latter, when used, tiles head-dimension rotation.

Changes

• Add a gfx942/gfx950 SageAttention builder under kernels/attention/.
• Use the production BLOCK_M=256, BLOCK_N=64 FlyDSL configuration on gfx942.
• Implement INT8 QK MFMA and native FP8 PV MFMA with FP32 accumulation.
• Double-buffer K/V in LDS and consume blocked-V storage.
• Support dense causal/non-causal attention, MHA/GQA/MQA, tail handling, and optional score bias.
• Keep score-bias loads lane/column exact using cooperative scalar loads and ds_bpermute.
• Add direct BF16 correctness coverage for causal and non-causal execution.

Performance

Measured on AMD Instinct MI308X (gfx942). In both tables, speedup = baseline latency / FlyDSL latency, so values above 1 mean FlyDSL is faster.

Production configurations: attention kernel only

Quantization is outside the timed region. Each implementation uses its independently tuned production tile: Triton A3 uses BLOCK_M=128, BLOCK_N=64, 4 warps; FlyDSL uses BLOCK_M=256, BLOCK_N=64.

Shape	Triton A3	FlyDSL	Speedup
S=1024, H=8	0.154 ms	0.078 ms	1.96x
S=4096, H=8	0.407 ms	0.354 ms	1.15x
S=4608, H=24	1.251 ms	1.168 ms	1.07x
S=8400, H=24	3.997 ms	3.414 ms	1.17x
S=8786, H=24	4.346 ms	3.918 ms	1.11x
S=8832, H=24	4.335 ms	3.908 ms	1.11x

Iso-tile integration benchmark: complete SmoothQ wrapper

This historical integration benchmark forces both Triton PR4033 and FlyDSL attention to BLOCK_M=256, BLOCK_N=64. Timing includes SmoothQ preprocessing, quantization, attention, and output handling. It isolates implementation efficiency at the same problem tile, but it is not the production comparison because production Triton A3 is faster at 128x64.

Shape	Triton PR4033 256x64	FlyDSL 256x64	Speedup
S=4608, H=24	3.189 ms	2.523 ms	1.264x
S=8400, H=24	8.177 ms	5.483 ms	1.491x
S=8786, H=24	9.336 ms	6.193 ms	1.507x
S=8832, H=24	9.507 ms	6.315 ms	1.505x

Qwen-Image-Edit integration

This PR intentionally contributes the standalone low-level kernel builder. It accepts pre-quantized INT8 Q/K, blocked FP8 V, descale tensors, and an optional score-bias tensor. It does not add the AITER high-level flydsl_sage_attn_func, SmoothQ preprocessing/quantization wrapper, or Qwen-Image-Edit attention dispatcher to this repository.

The same kernel source was exercised in the external AITER/Qwen-Image-Edit integration with 1920 FlyDSL attention calls and 0 fallbacks. An 8-step, 1024x1024, seed-0 run completed and produced the validated FlyDSL PNG. However, checking out this FlyDSL PR alone is not sufficient to rerun that model: the AITER wrapper/quantization integration is also required.

Testing

• Added causal and non-causal correctness tests.
• 2 passed on AMD Instinct MI308X (gfx942).
• ruff check kernels/attention/sage_attn_cdna.py tests/kernels/test_sage_attn.py
• ruff format --check kernels/attention/sage_attn_cdna.py tests/kernels/test_sage_attn.py
• python3 -m py_compile kernels/attention/sage_attn_cdna.py tests/kernels/test_sage_attn.py

Dependencies

• No new third-party dependencies added.

Breaking Changes

None.