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

fix for tma (#4023)

Browse source
This commit is contained in:
Lei Zhenyuan 2026-02-11 09:50:33 +08:00 • committed by GitHub
commit cdc9dc1fb1
1 changed files with 3 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -4,6 +4,7 @@
import logging
import warnings
from dataclasses import asdict
from unsloth import DEVICE_TYPE
import torch
import triton
@ -41,6 +42,8 @@ logger.addHandler(ch)
# 1. GPU 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":
return False
import triton.language as tl
gpu_supports_tma = torch.cuda.get_device_capability()[0] >= 9