ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(ROCm): prevent false TMA support detection on AMD GPUs (#4126)

Browse source 
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.
This commit is contained in:
金黄色葡萄球君君 2026-03-01 15:59:27 +08:00 • committed by GitHub
commit 48e8f78042
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/moe/grouped_gemm/interface.py
View file

@ -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