Skip to content

[ROCm] Run Sol-Attn on gfx1201 (RDNA4) through aiter - #1599

Draft
Shan2L wants to merge 3 commits into
ModelTC:mainfrom
Shan2L:feat/rocm-aiter-sol-attn
Draft

Shan2L wants to merge 3 commits into
ModelTC:mainfrom
Shan2L:feat/rocm-aiter-sol-attn

Conversation

@Shan2L

@Shan2L Shan2L commented Oct 11, 2026 •

Copy link
Copy Markdown

Draft: stacked on #1598 (only the last commit, feat(sol_attn): ..., is new here). Depends on aiter
gfx1201_sol_attention (aiter gfx1201 series parts 3-4, upstream PRs follow ROCm/aiter#6390; previews
Shan2L/aiter#2 and Shan2L/aiter#3). This PR will be rebased and marked ready after those are merged.

This PR runs Sol-Attn on AMD Radeon RDNA4 (gfx1201: Radeon AI PRO R9700 / R9600D, RX 9070 series) through aiter.

Today sol_attn needs the CUDA NVlabs sol_attn package; on ROCm SolAttnWeight reports the call as ineligible and
falls back to SDPA (or raises with strict).

Changes

  • lightx2v/common/ops/attn/sol_attn.py, ROCm only (CUDA path unchanged):
    • the kernel is aiter.gfx1201_sol_attention, which takes the same arguments as the NVlabs sol_attn forward
      (tau, diag threshold, sink tokens, softmax scale); kv_splits does not apply to it and is ignored;
    • eligibility: gfx1201 only; other ROCm GPUs get an ineligibility reason (strict raises, otherwise the existing
      SDPA fallback);
    • set_config raises on ROCm for thresh_type other than diag and compile_mode other than default;
    • aiter_sage_attn ([ROCm] Add aiter_sage_attn attention backend for gfx1201 (RDNA4) #1598) is accepted as dense_backend;
    • the "Sol-Attn active" log prints the GPU arch on ROCm.
  • configs/platforms/amd_rocm/minimax_h3_ref2av_r9700_sol.json: the [ROCm] Add aiter_sage_attn attention backend for gfx1201 (RDNA4) #1598 MiniMax-H3 config with attn_type: sol_attn
    (tau 1.5, diag, dense_steps 6, dense_layers [0], dense backend aiter_sage_attn, strict).
  • scripts/platforms/amd_rocm/run_minimax_h3_ref2av_r9700.sh: sol option.

sol_attn is selected only by attn_type; no defaults change.

Test

2 x R9700 (gfx1201, 32 GB), ROCm 7.1, PLATFORM=amd_rocm, MiniMax-H3 ref2av TP2, input street_dance
(girl.png, img_0.jpg), seed 42, 768x1344, 362 frames, DiT block offload, AdaLN cache. Only the DiT self-attention
differs between runs.

  • Unit: SolAttnWeight output (3D/4D, softmax scale) and the dense guard (by step, by layer) bitwise identical to direct
    aiter.gfx1201_sol_attention / aiter.gfx1201_sage_attention calls; unsupported thresh_type / compile_mode and
    FP32 inputs with strict raise.
  • ruff check / ruff format --check (v0.11.0) clean.

End to end on main (4fe984c + #1595 + #1596, 3 steps, dense_steps 1, ROCM_DETERMINISTIC_FP32_BLAS=1):
this PR vs main without it, with the same aiter ops installed as the sol_attn / sageattention packages (stock
CUDA code path of SolAttnWeight, sage_attn2 dense guard):

DiT s/step (step 1 dense guard, steps 2-3 Sol) pipeline total
this PR (dense_backend aiter_sage_attn) 89.3 / 54.2 / 54.7 513.9 s
main + packages 93.1 / 53.9 / 54.4 511.8 s

Video and audio bitwise identical. Per rank 98 Sol calls and 52 dense-guard calls. For comparison, dense
aiter_sage_attn on main (#1598, other GPU pair) runs at 89.0 s/step, so a Sol step is 1.63x faster.

End to end at the stock 29 steps (LightX2V 0.5.0 unmodified, sol_attn with the sol_attn_setting of
minimax_h3_sol_block_offload.json (the settings of the config added here), the aiter ops installed as the sol_attn /
sageattention packages, i.e. the path this PR is bitwise equal to on main; dense guard on steps 1-6 and layer 0):

attention DiT s/step (steps 2-6 / 7-29 / mean 2-29) pipeline total vs BF16 SDPA (s/step / total)
BF16 SDPA 211.4 / 211.7 / 211.6 6450 s
dense gfx1201_sage_attention (code object) 89.0 / 89.1 / 89.1 2877 s 2.38x / 2.24x
Sol, gfx1201_sol_attention HIP kernel 89.4 / 58.1 / 63.7 2144 s 3.32x / 3.01x
Sol, gfx1201_sol_attention code object 89.5 / 54.5 / 60.8 2091 s 3.48x / 3.08x

Per rank 1127 Sol calls and 323 dense-guard calls, 0 fallbacks, all outputs finite.

decoded output pair video PSNR (min frame) SSIM (min frame) audio SNR / cos
Sol code object vs BF16 SDPA 12.80 dB (10.91) 0.522 (0.408) 1.51 dB / 0.661
Sol HIP kernel vs BF16 SDPA 12.82 dB (10.88) 0.523 (0.408) 1.48 dB / 0.660
dense code object vs BF16 SDPA 13.59 dB (10.93) 0.578 (0.416) 6.13 dB / 0.877
BF16 SDPA K/V reversed vs BF16 SDPA (noise floor) 14.64 dB (10.30) 0.629 (0.391) 9.16 dB / 0.938

Over 29 steps the sampling trajectory is chaotic (reordering the FP32 sums of exact attention alone gives 14.6 dB
against the SDPA run), so pixel metrics against one reference run do not measure attention accuracy. Global statistics
of the Sol run (luma 70.0 / 56.9, sharpness 40.2, audio RMS 0.249) are close to the
dense run (69.3 / 56.9, 40.6, 0.238). Sol is block-sparse and lossy by design; its accuracy against FP32 attention
on real MiniMax-H3 tensors is reported with the aiter op.

29-step frames 0/120/241/361 (columns: BF16 SDPA, dense code object, Sol HIP kernel, Sol code object):

frames_street_dance_29steps_sol_asm

Before ready for review

Registers the ROCm platform attention backend "aiter_sage_attn", which runs
non-causal BF16/FP16 self-attention through aiter.gfx1201_sage_attention
(INT8 QK^T, FP8 PV, FP32 accumulation). The backend is explicit opt-in via
attn_type and raises at construction on non-ROCm devices, on GPUs other than
gfx1201, and when the installed aiter does not provide the op. Causal
attention, masks, dropout and packed multi-sequence cu_seqlens are rejected.

Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>
…h aiter_sage_attn

Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>
On ROCm, SolAttnWeight calls aiter.gfx1201_sol_attention (same arguments as
the NVlabs sol_attn forward: tau, diag threshold, sink tokens) instead of the
CUDA sol_attn package. The ROCm path is restricted to gfx1201 and to
thresh_type='diag' with the default compile mode; other GPUs report an
ineligible call (strict mode raises, otherwise the existing SDPA fallback).
aiter_sage_attn is accepted as dense guard backend, and a MiniMax-H3 Ref2AV
config for 2 x R9700 is added.

Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant