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

fixup derp

Browse source
This commit is contained in:
Datta Nimmaturi 2026-01-15 12:10:27 +00:00
commit ee669f3f30
2 changed files with 9 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/kernels/moe/autotune_cache.py
View file

@ -426,6 +426,9 @@ def _get_heuristic_configs() -> Tuple[Any, Any, Any]:
)
return config_fwd, config_bwd_dx, config_bwd_dw
def _get_default_configs() -> Tuple[Any, Any, Any]:
"""Get default kernel configurations as fallback."""
from .grouped_gemm.kernels.tuning import (
KernelConfigForward,

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

@ -53,6 +53,9 @@ def _check_tma_support():
_SUPPORTS_TMA = _check_tma_support()
# Check if triton.set_allocator is available (Triton 3.0+)
_HAS_SET_ALLOCATOR = hasattr(triton, "set_allocator")
def supports_tma():
return _SUPPORTS_TMA
@ -197,7 +200,7 @@ def grouped_gemm_forward(
if use_tma or autotune:
# Respect global persistent allocator if set
if not getattr(triton, "_unsloth_allocator_set", False):
if _HAS_SET_ALLOCATOR and not getattr(triton, "_unsloth_allocator_set", False):
def alloc_fn(size: int, alignment: int, stream: int):
return torch.empty(size, device = "cuda", dtype = torch.int8)
@ -404,7 +407,7 @@ def grouped_gemm_dX(
if use_tma or autotune:
# Respect global persistent allocator if set
if not getattr(triton, "_unsloth_allocator_set", False):
if _HAS_SET_ALLOCATOR and not getattr(triton, "_unsloth_allocator_set", False):
def alloc_fn(size: int, alignment: int, stream: int):
# print(f"DEBUG::GROUPED_GEMM alloc_fn {size=} {alignment=} {stream=}")
@ -572,7 +575,7 @@ def grouped_gemm_dW(
if use_tma or autotune:
# Respect global persistent allocator if set
if not getattr(triton, "_unsloth_allocator_set", False):
if _HAS_SET_ALLOCATOR and not getattr(triton, "_unsloth_allocator_set", False):
def alloc_fn(size: int, alignment: int, stream: int):
return torch.empty(size, device = "cuda", dtype = torch.int8)