sageattention crashes on RDNA4 (gfx1200/gfx1201) with ROCm 7.2.2 — Triton num_stages bug + fix

[brosequist]

Summary

When attention_mode: sageattn is selected on AMD RDNA4 GPUs (gfx1200 / gfx1201) running ROCm 7.2.2, the Triton JIT compiler crashes on first use with:

RuntimeError: PassManager::run failed
'tt.load' op operation destroyed but still has uses

This makes sageattn completely unavailable on RDNA4 hardware without a manual patch to the sageattention package.

Root Cause

sageattention's Triton kernels hardcode num_stages=4 (or num_stages=3 if head_dim == 64 else 4). On RDNA4 (gfx12xx), Triton 3.6.0's AMD software pipelining pass (tritonamdgpu-pipeline) has a use-after-free bug when num_stages >= 2. The gfx12xx architecture is not in Triton 3.6.0's verified target list.

Affected hardware: RX 9060 XT (gfx1200), RX 9070 XT (gfx1201) — all RDNA4 GPUs.

Fix (sageattention upstream PR)

A PR has been submitted to sageattention to use num_stages=1 on ROCm:
thu-ml/SageAttention#365

torch.version.hip is a non-None string on ROCm and None on CUDA, making it a zero-cost runtime discriminator. num_stages=1 disables pipelining (not correctness).

Suggested Workaround for This Repo

Until the fix lands in a sageattention release, WanVideoWrapper could apply a runtime patch in attention.py before the first kernel invocation:

def _patch_sageattention_for_rocm():
    """Patch sageattention Triton kernels to use num_stages=1 on ROCm (gfx12xx bug)."""
    import torch, pathlib, re
    if not torch.version.hip:
        return
    try:
        import sageattention
        sa_dir = pathlib.Path(sageattention.__file__).parent / "triton"
        for f in sa_dir.glob("attn_qk_int8*.py"):
            src = f.read_text()
            orig = src
            src = src.replace(
                "num_stages=3 if head_dim == 64 else 4",
                "num_stages=(1 if torch.version.hip else (3 if head_dim == 64 else 4))"
            )
            src = re.sub(r"num_stages=4\)", "num_stages=(1 if torch.version.hip else 4))", src)
            if src != orig:
                f.write_text(src)
                pyc = f.parent / "__pycache__" / (f.stem + ".cpython-312.pyc")
                pyc.unlink(missing_ok=True)
    except Exception:
        pass

Call this once at module import time before sageattention kernels are first invoked.

Environment

• GPU: RDNA4 (gfx1200 / gfx1201)
• ROCm: 7.2.2
• PyTorch: 2.10.0+rocm7.2.2
• Triton: 3.6.0
• sageattention: 1.0.6