From 0cc694158f99be78bb00d11718b01aa6cac01691 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=91=E9=BB=84=E8=89=B2=E8=91=A1=E8=90=84=E7=90=83?= =?UTF-8?q?=E5=90=9B=E5=90=9B?= Date: Sun, 1 Mar 2026 15:59:27 +0800 Subject: [PATCH] fix(ROCm): prevent false TMA support detection on AMD GPUs (#4126) TMA (Tensor Memory Accelerator) is an NVIDIA Hopper+ feature that does not exist on AMD GPUs. However, _check_tma_support() incorrectly returns True on ROCm because: 1. torch.cuda.get_device_capability() returns (11, 0) for gfx1100, satisfying the >= 9 check intended for Hopper (sm_90). 2. ROCm Triton exports tl.make_tensor_descriptor (the symbol exists even though the hardware does not support TMA). This would cause MoE grouped_gemm to attempt TMA operations on AMD GPUs, leading to runtime failures. Fix: early-return False for HIP devices, matching the existing XPU guard. --- unsloth/kernels/moe/grouped_gemm/interface.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index 554e5fcc03..5588458973 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/unsloth/kernels/moe/grouped_gemm/interface.py @@ -39,10 +39,10 @@ logger.addHandler(ch) # Precompute TMA support to avoid graph breaks # TMA requires both: -# 1. GPU capability >= 9 (Hopper+) +# 1. NVIDIA GPU with capability >= 9 (Hopper+) # 2. Triton version with TMA API (make_tensor_descriptor or _experimental_make_tensor_descriptor) def _check_tma_support(): - if DEVICE_TYPE == "xpu": + if DEVICE_TYPE in ("xpu", "hip"): return False import triton.language as tl