diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index 572b02702b..10a97d1d71 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/unsloth/kernels/moe/grouped_gemm/interface.py @@ -191,12 +191,15 @@ def grouped_gemm_forward( # TMA load for activations, TMA gather only supported on Blackwell+ assert not permute_x, "Cannot use both use_tma_load_x and permute_x" - use_tma = use_tma_load_w or use_tma_load_x or use_tma_store - if not supports_tma() and use_tma: - warnings.warn("TMA not supported, tma_load will be set to False") + # Hard disable TMA on unsupported setups (e.g. T4 / SM75). + # This guarantees we never compile or run a kernel with any USE_TMA_* flags. + if not supports_tma(): use_tma_load_w = False use_tma_load_x = False use_tma_store = False + use_tma = False + else: + use_tma = use_tma_load_w or use_tma_load_x or use_tma_store if use_tma or autotune: # Respect global persistent allocator if set @@ -398,12 +401,14 @@ def grouped_gemm_dX( assert not (permute_y and use_tma_load_dy), "Cannot use both TMA load and permute_y" assert not (permute_x and use_tma_store), "Cannot use both TMA store and permute_x" - use_tma = use_tma_load_dy or use_tma_load_w or use_tma_store - if not supports_tma() and use_tma: - warnings.warn("TMA not supported, tma_load will be set to False") + # Hard disable TMA on unsupported setups (e.g. T4 / SM75). + if not supports_tma(): use_tma_load_w = False use_tma_load_dy = False use_tma_store = False + use_tma = False + else: + use_tma = use_tma_load_dy or use_tma_load_w or use_tma_store if use_tma or autotune: # Respect global persistent allocator if set @@ -566,12 +571,14 @@ def grouped_gemm_dW( assert not (permute_y and use_tma_load_dy), "Cannot use both TMA load and permute_y" assert not (permute_x and use_tma_load_x), "Cannot use both TMA load and permute_x" - use_tma = use_tma_load_dy or use_tma_load_x or use_tma_store - if not supports_tma() and use_tma: - warnings.warn("TMA not supported, tma_load will be set to False") + # Hard disable TMA on unsupported setups (e.g. T4 / SM75). + if not supports_tma(): use_tma_load_x = False use_tma_load_dy = False use_tma_store = False + use_tma = False + else: + use_tma = use_tma_load_dy or use_tma_load_x or use_tma_store if use_tma or autotune: # Respect global persistent allocator if set diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/backward.py b/unsloth/kernels/moe/grouped_gemm/kernels/backward.py index 5e07056b52..e976d9ef31 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/backward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/backward.py @@ -12,6 +12,19 @@ from .autotuning import ( prune_kernel_configs_backward_dW, ) +# ----------------------------------------------------------------------------- +# Triton compatibility: +# Some Triton versions (e.g. 3.2.0) do NOT expose `triton.language.make_tensor_descriptor`. +# Triton's JIT dependency hashing will try to resolve attribute accesses like +# `tl.make_tensor_descriptor` even if they are guarded by compile-time constexpr +# branches, leading to an AttributeError before compilation. +# +# Avoid direct attribute access; use a module-level alias created via getattr. +# ----------------------------------------------------------------------------- +_make_tensor_descriptor = getattr(tl, "make_tensor_descriptor", None) +if _make_tensor_descriptor is None: + _make_tensor_descriptor = getattr(tl, "_experimental_make_tensor_descriptor", None) + """ dX backward kernel @@ -82,7 +95,7 @@ def _grouped_gemm_dX_kernel( # Also, we are defining a single global descriptor with single block shape # Need to check that this does not result in errors when crossing expert boundaries if USE_TMA_LOAD_dY: - dY_desc = tl.make_tensor_descriptor( + dY_desc = _make_tensor_descriptor( dY_ptr, shape = [TOTAL_TOKENS, N], strides = [N, 1], @@ -91,7 +104,7 @@ def _grouped_gemm_dX_kernel( if USE_TMA_LOAD_W: expert_stride = N * K - w_desc = tl.make_tensor_descriptor( + w_desc = _make_tensor_descriptor( w_ptr, shape = [NUM_EXPERTS, N, K], strides = [expert_stride, K, 1], @@ -123,7 +136,7 @@ def _grouped_gemm_dX_kernel( tl.static_assert( K % BLOCK_SIZE_K == 0, "K must be divisible by BLOCK_SIZE_K" ) - dX_desc = tl.make_tensor_descriptor( + dX_desc = _make_tensor_descriptor( dX_ptr, shape = [m_end, K], strides = [K, 1], @@ -324,7 +337,7 @@ def _grouped_gemm_dW_kernel( output_dtype = dW_ptr.dtype.element_ty if USE_TMA_LOAD_dY and not TMA_LOAD_BOTH: - dY_desc = tl.make_tensor_descriptor( + dY_desc = _make_tensor_descriptor( dY_ptr, shape = [TOTAL_TOKENS, N], strides = [N, 1], @@ -332,7 +345,7 @@ def _grouped_gemm_dW_kernel( ) if USE_TMA_LOAD_X and not TMA_LOAD_BOTH: - x_desc = tl.make_tensor_descriptor( + x_desc = _make_tensor_descriptor( x_ptr, shape = [TOTAL_TOKENS, K], strides = [K, 1], @@ -351,7 +364,7 @@ def _grouped_gemm_dW_kernel( if USE_TMA_STORE: tl.static_assert(N % BLOCK_SIZE_N == 0, "N must be divisible by BLOCK_SIZE_N") tl.static_assert(K % BLOCK_SIZE_K == 0, "K must be divisible by BLOCK_SIZE_K") - dW_desc = tl.make_tensor_descriptor( + dW_desc = _make_tensor_descriptor( dW_ptr, shape = [NUM_EXPERTS, N, K], strides = [N * K, K, 1], @@ -392,14 +405,14 @@ def _grouped_gemm_dW_kernel( if m_size > 0: if TMA_LOAD_BOTH: - dY_desc = tl.make_tensor_descriptor( + dY_desc = _make_tensor_descriptor( dY_ptr, shape = [m_end, N], strides = [N, 1], block_shape = [BLOCK_SIZE_M, BLOCK_SIZE_N], ) - x_desc = tl.make_tensor_descriptor( + x_desc = _make_tensor_descriptor( x_ptr, shape = [m_end, K], strides = [K, 1], diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py index a42ec5ffe9..d038d55688 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py @@ -10,6 +10,21 @@ from .autotuning import ( prune_kernel_configs_fwd, ) +# ----------------------------------------------------------------------------- +# Triton compatibility: +# Some Triton versions (e.g. 3.2.0) do NOT expose `triton.language.make_tensor_descriptor`. +# Importantly, Triton's JIT dependency hashing will attempt to resolve attribute +# accesses like `tl.make_tensor_descriptor` even if the code is behind a +# compile-time `tl.constexpr` branch (e.g. USE_TMA_* is False), causing an +# AttributeError at compile time. +# +# To avoid this, never reference `tl.make_tensor_descriptor` directly in kernel +# source. Instead, alias it via getattr at module import time. +# ----------------------------------------------------------------------------- +_make_tensor_descriptor = getattr(tl, "make_tensor_descriptor", None) +if _make_tensor_descriptor is None: + _make_tensor_descriptor = getattr(tl, "_experimental_make_tensor_descriptor", None) + # # PERMUTE_X -> permute tokens so that they are ordered by expert @@ -66,7 +81,7 @@ def _grouped_gemm_forward_kernel( # Also, we are defining a single global descriptor with single block shape # Need to check that this does not result in errors when crossing expert boundaries if USE_TMA_LOAD_X: - x_desc = tl.make_tensor_descriptor( + x_desc = _make_tensor_descriptor( x_ptr, shape = [TOTAL_TOKENS, K], strides = [K, 1], @@ -75,7 +90,7 @@ def _grouped_gemm_forward_kernel( if USE_TMA_LOAD_W: expert_stride = N * K - w_desc = tl.make_tensor_descriptor( + w_desc = _make_tensor_descriptor( w_ptr, shape = [NUM_EXPERTS, N, K], strides = [expert_stride, K, 1], @@ -100,7 +115,7 @@ def _grouped_gemm_forward_kernel( # Need to create tma_store within loop since we need to predicate stores based on m_size if USE_TMA_STORE: - y_desc = tl.make_tensor_descriptor( + y_desc = _make_tensor_descriptor( y_ptr, # + m_start * N, shape = [m_end, N], strides = [N, 1],