From ee669f3f30becfd78b81fd2990726c08329f3432 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Thu, 15 Jan 2026 12:10:27 +0000 Subject: [PATCH] fixup derp --- unsloth/kernels/moe/autotune_cache.py | 3 +++ unsloth/kernels/moe/grouped_gemm/interface.py | 9 ++++++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/moe/autotune_cache.py b/unsloth/kernels/moe/autotune_cache.py index 8f03af75db..5483e74c03 100644 --- a/unsloth/kernels/moe/autotune_cache.py +++ b/unsloth/kernels/moe/autotune_cache.py @@ -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, diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index d69af9ebee..572b02702b 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/unsloth/kernels/moe/grouped_gemm/interface.py @@ -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)