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

Open
  • Is something wrong?
#2,007 0 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
35/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Stale
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start by reading attention.py and how this repository imports and invokes sageattention; the issue proposes applying the workaround before the first kernel invocation. Check the upstream SageAttention PR and package version first, then verify the behavior on ROCm and CUDA; done means RDNA4 ROCm no longer crashes while CUDA behavior remains unchanged.

Written by the indexing model from the issue text.

Description

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:
https://github.com/thu-ml/SageAttention/pull/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
Dominant language
Python
Stars
6.8k
Forks
712
No merged PRs in 30d

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from kijai/ComfyUI-WanVideoWrapper

All issues in kijai/ComfyUI-WanVideoWrapper

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.