Repository navigation
Conversation
Add gfx1201_sage_prepare, the input stage of a SageAttention-style attention on gfx1201 (RDNA4): BF16 [B, S, H, 128] Q/K are quantized to INT8 with one FP32 scale per (batch, head, 32-row block), Q carrying softmax_scale * log2(e); V is quantized to FP8 e4m3, transposed to [B, H, 128, S_pad], with one FP32 scale per channel. Optionally Q/K go through RMSNorm and rotate-half RoPE (rope_dim 32/64/96/128) first, rounded like the eager PyTorch BF16 sequence. The module_gfx1201_sage_prepare JIT module is excluded from prebuild unless gfx1201 is in GPU_ARCHS, and the op test exits early on other architectures. Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>
Add gfx1201_sage_attention, a non-causal SageAttention forward for gfx1201 (RDNA4), registered as torch.ops.aiter.gfx1201_sage_attention: BF16 [B, S, H, 128] Q/K/V go through gfx1201_sage_prepare, then a HIP kernel computes Q K^T with INT8 WMMA and P V with FP8 WMMA, FP32 accumulation and online softmax, writing BF16. The optional QK RMSNorm + RoPE arguments are forwarded to the prepare stage. gfx1201_sage_attention_fwd runs the kernel alone on prepared tensors. The op is lossy by design. op_tests/test_gfx1201_sage_attention.py checks it against attention_ref with the tolerance band used by the Triton SageAttention test; op_tests/op_benchmarks/hip/bench_gfx1201_sage_attention.py times it against BF16 SDPA and Triton fav3_sage and can report accuracy against an FP32 reference. Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
github-actions
Bot
changed the title
This was referenced Oct 10, 2026
This was referenced Oct 12, 2026
Draft
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 1 of 4 of the restructured #6292 (see the series list there);
Motivation
aiter has no fast attention path for gfx1201 (RDNA4, Radeon AI PRO R9700). BF16 SDPA reaches ~68 TFLOPS there and the
Triton
fav3_sagepath (#1818) is untuned for RDNA4 (~15-18 TFLOPS). This PR adds a SageAttention-style forward forgfx1201 written in HIP: INT8 Q K^T and FP8 P V on the RDNA4 WMMA units with FP32 accumulation.
This is the first PR of a series; later PRs add a hand-scheduled code object for the same kernel and a sparse variant.
Technical Details
Two commits:
gfx1201_sage_prepare- BF16[B, S, H, 128]Q/K -> INT8 with one FP32 scale per (batch, head, 32-row block),Q carrying
softmax_scale * log2(e); V -> FP8 e4m3, transposed to[B, H, 128, S_pad], one FP32 scale per channel.Optional fused QK RMSNorm + rotate-half RoPE (rope_dim 32/64/96/128), rounded like the eager PyTorch BF16 sequence.
gfx1201_sage_attention(torch.ops.aiter.gfx1201_sage_attention) - prepare + HIP kernel: online softmax, INT8 WMMAfor Q K^T, FP8 WMMA for P V, FP32 accumulation, BF16 output.
gfx1201_sage_attention_fwdruns the kernel alone.Scope: gfx1201 only, head dim 128, non-causal, no varlen/mask, BF16 in/out, forward only. Both JIT modules are excluded
from prebuild unless
gfx1201is inGPU_ARCHS; the tests/benchmark exit early on other architectures.Test Plan
python op_tests/test_gfx1201_sage_prepare.py- bitwise against a PyTorch reference of the quantization (+ norm/RoPE).python op_tests/test_gfx1201_sage_attention.py- againstattention_refwith the tolerance band of the TritonSageAttention test (atol 0.3, rtol 0.25, <= 0.5 % failing elements) plus cosine >= 0.998 and median row cosine >= 0.99;
torch.library.opcheck.python op_tests/op_benchmarks/hip/bench_gfx1201_sage_attention.py --compare_to_ref- speed + accuracy vs BF16 SDPAand Triton
fav3_sage.Test Result
Radeon AI PRO R9700 (gfx1201, 32 GB), public image above, this branch.
Op tests:
test_gfx1201_sage_prepare.pyandtest_gfx1201_sage_attention.pyALL PASS (64 checks, incl.opcheck,non-default stream, fused == unfused bitwise, input validation).
Kernel benchmark (
randninputs;sageincludes prepare; accuracy vs the FP32 reference):fav3_sagems (TFLOPS)Cosine vs FP32 reference 0.9992-0.9993 (Triton
fav3_sage0.9991-0.9993, BF16 SDPA 1.0000); relative L2 error0.038-0.039 (Triton 0.039-0.042, BF16 SDPA 0.0023); 0 % elements outside the test band for all providers.
Real MiniMax-H3 tensors (
DUMP_INPUTS, 4 inputs x steps 0/2 x layers below, heads 0/9/18/27, S = 117k-128k),relative L2 error vs FP32 reference, mean / max over the 8 dumps per layer:
fav3_sageAll outputs finite. Layers 45/49 have the largest error for both quantized kernels; this PR is below Triton
fav3_sageon every layer.
End-to-end (
repro.sh, TP2, street_dance, 3 steps, block offload):The DiT self-attention went through the new op for all 150 calls (50 layers x 3 steps), 0 SDPA fallbacks. A second
run of each backend gives bitwise identical frames and audio.
NOISE_FLOOR=1)Reversing the K/V token order leaves attention mathematically unchanged and only perturbs the FP32 summation order,
so the second row shows how far rounding-level attention differences already move the 3-step output.
Frames 0/120/241/361 (left BF16 SDPA, middle this PR, right |difference| x 8):
End-to-end at the stock step count (
STEPS=29 NOISE_FLOOR=1 repro.sh; 29 is theinfer_stepsofminimax_h3_tp.json; same input, seed and config as above):Step 1 of this PR (125.1 s) includes the first-call JIT build. All 1450 DiT self-attention calls went through the
op, 0 SDPA fallbacks.
The reference row is the upstream SageAttention implementation as LightX2V uses it:
sage_attn2(theattn_typeofminimax_h3_tp.json) with thu-ml/SageAttentiond1a57a5, throughsageattn_qk_int8_pv_fp16_triton(INT8 QK^T,FP16 PV), the kernel
sage_attn2selects on SM89; the CUDA kernels of the package do not build on ROCm. It is thereto separate the effect of SageAttention itself from this kernel, not as a performance baseline (the Triton kernel is
not tuned for gfx1201). 1450 calls per rank, 0 fallbacks.
Over 29 steps the sampling trajectory is chaotic: reordering the FP32 sums of exact BF16 attention alone already
produces a visibly different video (14.6 dB, SSIM 0.63 against the SDPA run), so pixel metrics against a single
reference run no longer measure kernel accuracy. LightX2V's own MiniMax-H3 scheduler
notes the same sensitivity: an FP32 result that "differs by an ulp on CUDA" is "amplified by the 48-layer
denoiser" (scheduler.py). The reference SageAttention implementation also ends far
from the SDPA run (14.09 dB, SSIM 0.597, close to the exact-reorder floor), so a low PSNR against a single BF16 run at
29 steps is expected for SageAttention in general and is not specific to this kernel. This PR ends 1.0 dB / 0.05
SSIM below the reference implementation and 1.6 dB / 0.08 SSIM below the exact-reorder floor. Global statistics of
all outputs stay within the spread of the two exact runs:
29-step frames 0/120/241/361 (columns: BF16 SDPA, BF16 SDPA K/V reversed, this PR):
How to reproduce
Everything below uses public artifacts only: the
rocm/pytorch:rocm10.0_ubuntu24.04_py3.13_pytorch_release_2.11.0image, this branch, upstream LightX2V 0.5.0 (
95e9b86b) and the publicMiniMaxAI/MiniMax-H3weights(revision
42ed227e). Inputs are LightX2V's own example images (assets/inputs/imgs) with our prompts.Kernel test / benchmark (1 GPU):
The benchmark times each provider with CUDA events (2 warm-up + 5 timed calls), repeats this for 3 rounds with the
providers interleaved, and reports the median. Inputs are
randnBF16 Q/K/V.sageincludes the prepare stage,sage_coreis the attention kernel alone on prepared inputs; TFLOPS = 4 * B * H * S^2 * 128 / time.End-to-end (2 GPUs, TP2) with the attached gfx1201_sage_repro.zip (
repro.sh,h3_e2e.py,real_dumps.py, prompts),run from the unzipped directory:
h3_e2e.pyruns stock LightX2V MiniMax-H3 ref2av (configs/minimax_h3/minimax_h3_tp.json, 768x1344, 362 frames,seed 42) and only swaps the attention backend: it registers
gfx1201_sagein LightX2V's attention registry, which callstorch.ops.aiter.gfx1201_sage_attentionfor the DiT self-attention and falls back to SDPA for anything else (thenumber of calls of each kind is printed at exit). Kernels that differ between the two runs only through attention:
RMSNorm/RoPE use LightX2V's
torch_native/torch_real_rope, and FP32 GEMMs/convolutions go through rocBLAS so thatrepeated runs are bitwise identical (hipBLASLt picks a per-process kernel for some FP32 GEMMs,
ModelTC/LightX2V#1594, opt-in fix proposed in ModelTC/LightX2V#1596). The pre-encode frames and audio are saved as
*_raw.ptand compared directly (no codec inthe loop).
Notes for running on a 32 GB card with the public image (all handled by
repro.sh):repro.shpins the installedversions through
PIP_CONSTRAINT, and installs aiter withAITER_USE_SYSTEM_TRITON=1._rocm_sdk_core) with onlylibamdhip64.so.7; aiter's JIT linkswith
-lamdhip64, sorepro.shadds an unversioned symlink toLIBRARY_PATH.repro.shuses LightX2V's AMD platform (PLATFORM=amd_rocm) and exportsPYTORCH_CUDA_ALLOC_CONF=expandable_segments:True: withamd_rocm, LightX2V 0.5.0 initializes the allocator before itsets this itself and the run OOMs at DiT step 1 ([Bug] PLATFORM=amd_rocm ignores expandable_segments:True, MiniMax-H3 OOMs on 32 GB GPUs ModelTC/LightX2V#1593, fixed in Set PYTORCH_CUDA_ALLOC_CONF before platform init ModelTC/LightX2V#1595).
which for MiniMax-H3 needs the persistent AdaLN cache;
h3_e2e.py adalnbuilds it once with LightX2V'stools/cache_minimax_h3_adaln. Offload makes the absolute s/step slower than a resident setup; both backends usethe same config.
DUMP_INPUTS="street_dance calm_interview mountain_flyover snake_jungle"additionally dumps the real DiT Q/K/V(rank 0, steps 0 and last, selected layers; ~5 GB per input) and runs
real_dumps.py, which compares this kernel,Triton
fav3_sageand BF16 SDPA against an FP32 reference on them.NOISE_FLOOR=1additionally reruns BF16 SDPA with the K/V token order reversed (h3_e2e.py run --sdpa_kv_flip)and writes its comparison against the SDPA run to
compare_floor.json.STEPS=29runs the stock step count instead of the default 3 (the AdaLN cache is built per step count; about1.8 h for the SDPA run and 0.8-0.9 h for the op run on R9700 with block offload).
OFFICIAL_SAGE=1additionally runs LightX2V's ownsage_attn2with thu-ml/SageAttentiond1a57a5throughsageattn_qk_int8_pv_fp16_triton(the reference row of the 29-step tables) and writes its comparison against theSDPA run to
compare_official_sage.json(about 2.9 h at 29 steps).