fix: use num_stages=1 on ROCm to avoid Triton pipelining crash on RDNA4 (gfx12xx)

[brosequist]

Problem

On AMD RDNA4 GPUs (gfx1200 / gfx1201, e.g. RX 9060 XT / RX 9070 XT) with ROCm 7.2.2 and Triton 3.6.0, the sageattention Triton kernels fail immediately on first use with:

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

Full traceback snippet:

File ".../sageattention/triton/attn_qk_int8_per_block.py", line 182, in forward
    num_stages=3 if head_dim == 64 else 4)
triton.compiler.errors.CompilationError: ...
  File "<triton-generated>", line ...: error: 'tt.load' op operation destroyed but still has uses
RuntimeError: PassManager::run failed

Root Cause

The Triton AMD software pipelining pass (tritonamdgpu-pipeline) has a use-after-free bug triggered when num_stages >= 2 on gfx12xx targets. RDNA4 (gfx1200/gfx1201) is not in Triton 3.6.0's verified target list, and the pipelining pass does not handle these architectures correctly.

Fix

Use num_stages=1 when running on ROCm. torch.version.hip is a non-None string on ROCm builds and None on CUDA builds, making it a reliable runtime discriminator with no import overhead.

num_stages=1 disables software pipelining entirely, which avoids the bug. CUDA behavior is completely unchanged.

Testing

Tested on:

• AMD RX 9070 XT (gfx1201) + RX 9060 XT (gfx1200)
• ROCm 7.2.2 / PyTorch 2.10.0+rocm7.2.2 / Triton 3.6.0
• sageattention 1.0.6

With this fix, all 4 Triton kernel files compile and run correctly on RDNA4. Attention quality is unaffected; num_stages=1 only disables prefetching, not correctness.

Notes

• The two *_h96.py variant files (head_dim=96) contain the same num_stages=4 pattern and may need the same fix if present in other installed versions — they were absent from the upstream source but appeared in the pip-installed 1.0.6 package.
• A longer-term fix could detect the specific Triton/GPU combination more selectively, but num_stages=1 on all ROCm targets is safe and simple.

[arsic3]

Hello @jt-zhang @whx1003 @Weiqi-Feng, I am from AMD and we verified that OP's commit fixes the issue on gfx12xx GPUs. Is there a possibility to merge this PR? Also, is there a possibility to have this included in the v1.0.6?