From ea8e585de9d82fb56ac127adeeaf1628c0cb23f8 Mon Sep 17 00:00:00 2001 From: Brett Rosequist Date: Wed, 29 Apr 2026 15:50:20 -0400 Subject: [PATCH] fix: use num_stages=1 on ROCm to avoid Triton pipelining crash on RDNA4 (gfx12xx) Co-Authored-By: Claude Sonnet 4.6 --- sageattention/triton/attn_qk_int8_block_varlen.py | 2 +- sageattention/triton/attn_qk_int8_per_block.py | 2 +- sageattention/triton/attn_qk_int8_per_block_causal.py | 2 +- sageattention/triton/attn_qk_int8_per_block_causal_varlen.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sageattention/triton/attn_qk_int8_block_varlen.py b/sageattention/triton/attn_qk_int8_block_varlen.py index 724c411a..cfe53c73 100644 --- a/sageattention/triton/attn_qk_int8_block_varlen.py +++ b/sageattention/triton/attn_qk_int8_block_varlen.py @@ -147,5 +147,5 @@ def forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, q_scale, k_scale, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM_K, STAGE=stage, num_warps=4 if head_dim == 64 else 8, - num_stages=3 if head_dim == 64 else 4) + num_stages=(1 if torch.version.hip else (3 if head_dim == 64 else 4))) return o \ No newline at end of file diff --git a/sageattention/triton/attn_qk_int8_per_block.py b/sageattention/triton/attn_qk_int8_per_block.py index 401ac3c0..5af4fce7 100644 --- a/sageattention/triton/attn_qk_int8_per_block.py +++ b/sageattention/triton/attn_qk_int8_per_block.py @@ -179,6 +179,6 @@ def forward(q, k, v, q_scale, k_scale, tensor_layout="HND", attn_mask=None, outp BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM_K, STAGE=stage, RETURN_LSE=return_lse, num_warps=4 if head_dim == 64 else 8, - num_stages=3 if head_dim == 64 else 4) + num_stages=(1 if torch.version.hip else (3 if head_dim == 64 else 4))) return o, lse \ No newline at end of file diff --git a/sageattention/triton/attn_qk_int8_per_block_causal.py b/sageattention/triton/attn_qk_int8_per_block_causal.py index 06621f0c..b403bdb0 100644 --- a/sageattention/triton/attn_qk_int8_per_block_causal.py +++ b/sageattention/triton/attn_qk_int8_per_block_causal.py @@ -170,6 +170,6 @@ def forward(q, k, v, q_scale, k_scale, tensor_layout="HND", output_dtype=torch.f STAGE=stage, RETURN_LSE=return_lse, num_warps=4 if head_dim == 64 else 8, - num_stages=4) + num_stages=(1 if torch.version.hip else 4)) return o, lse \ No newline at end of file diff --git a/sageattention/triton/attn_qk_int8_per_block_causal_varlen.py b/sageattention/triton/attn_qk_int8_per_block_causal_varlen.py index 79eec97e..619fb16f 100644 --- a/sageattention/triton/attn_qk_int8_per_block_causal_varlen.py +++ b/sageattention/triton/attn_qk_int8_per_block_causal_varlen.py @@ -162,5 +162,5 @@ def forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, q_scale, k_scale, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM_K, STAGE=stage, num_warps=4 if head_dim == 64 else 8, - num_stages=4) + num_stages=(1 if torch.version.hip else 4)) return o \ No newline at end of file