[Kernel] Add CDNA SageAttention kernel

[LiuYinfeng01]

Summary

Add a standalone FlyDSL SageAttention V1 kernel for AMD Instinct MI308X (gfx942).

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:

• INT8 QK MFMA with FP32 accumulation.
• Per-Q-block and per-K-block descale plus softmax scaling.
• Optional lane/column-exact delta_s addition before softmax.
• Online FP32 softmax.
• Native FP8 PV MFMA with FP32 accumulation.
• 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 an MI308X/gfx942 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 the validated dense MHA path with causal/non-causal execution 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). Throughput is reported in TFLOPS using the dense attention FLOP count 4 * B * H * S^2 * D (QK^T and PV, FP32-accumulated), with B=1 and D=128; higher TFLOPS is faster. speedup = FlyDSL TFLOPS / Triton TFLOPS, 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	27.9 TFLOPS	55.1 TFLOPS	1.96x
S=4096, H=8	168.8 TFLOPS	194.1 TFLOPS	1.15x
S=4608, H=24	208.6 TFLOPS	223.4 TFLOPS	1.07x
S=8400, H=24	216.9 TFLOPS	254.0 TFLOPS	1.17x
S=8786, H=24	218.3 TFLOPS	242.1 TFLOPS	1.11x
S=8832, H=24	221.1 TFLOPS	245.3 TFLOPS	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. The timed region includes SmoothQ preprocessing, quantization, attention, and output handling, so the values below are effective attention TFLOPS over the full wrapper time (4 * B * H * S^2 * D / total_time), not attention-kernel-only throughput. 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	81.8 TFLOPS	103.4 TFLOPS	1.264x
S=8400, H=24	106.0 TFLOPS	158.1 TFLOPS	1.491x
S=8786, H=24	101.6 TFLOPS	153.2 TFLOPS	1.507x
S=8832, H=24	100.8 TFLOPS	151.8 TFLOPS	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

Validated on AMD Instinct MI308X (gfx942); the test module skips before importing the CDNA kernel on Navi and all other GPUs.

• Causal and non-causal BF16 correctness: 2 passed.
• Minimum cosine similarity requirement: >0.99; observed minimum: 0.998324 across the three validation shapes.
• Output remained bit-exact against the stored baseline for B1/S256/H8/D128, B1/S4608/H24/D128, and B1/S8448/H24/D128.
• Kernel-only non-causal performance after cleanup: 220.78 TFLOPS at S=4608/H=24 and 253.30 TFLOPS at S=8448/H=24; no regression against the 217.34 and 250.68 TFLOPS baselines.
• ruff check kernels/attention/sage_attn_cdna.py kernels/attention/sage_attn_utils.py tests/kernels/test_sage_attn.py
• black --check --line-length 120 kernels/attention/sage_attn_cdna.py kernels/attention/sage_attn_utils.py tests/kernels/test_sage_attn.py
• python3 -m py_compile kernels/attention/sage_attn_cdna.py kernels/attention/sage_attn_utils.py tests/kernels/test_sage_attn.py

GQA/MQA and arbitrary tail lengths are not claimed as validated capabilities by this PR.

Dependencies

• No new third-party dependencies added.

[coderfeli]

Clean codes using