From 1107626da248455bb2214959b64a3e544deb6054 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Thu, 15 Jan 2026 14:08:44 +0000 Subject: [PATCH] Revert "non TMA [T4]" This reverts commit 56a72c677ab8eecb3f66d75ac8750af608639363. --- unsloth/kernels/moe/grouped_gemm/interface.py | 25 ++++++---------- .../moe/grouped_gemm/kernels/backward.py | 29 +++++-------------- .../moe/grouped_gemm/kernels/forward.py | 21 ++------------ 3 files changed, 20 insertions(+), 55 deletions(-) diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index 10a97d1d71..572b02702b 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/unsloth/kernels/moe/grouped_gemm/interface.py @@ -191,15 +191,12 @@ 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" - # 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 = 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") 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 @@ -401,14 +398,12 @@ 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" - # Hard disable TMA on unsupported setups (e.g. T4 / SM75). - if not supports_tma(): + 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") 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 @@ -571,14 +566,12 @@ 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" - # Hard disable TMA on unsupported setups (e.g. T4 / SM75). - if not supports_tma(): + 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") 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 e976d9ef31..5e07056b52 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/backward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/backward.py @@ -12,19 +12,6 @@ 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 @@ -95,7 +82,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 = _make_tensor_descriptor( + dY_desc = tl.make_tensor_descriptor( dY_ptr, shape = [TOTAL_TOKENS, N], strides = [N, 1], @@ -104,7 +91,7 @@ def _grouped_gemm_dX_kernel( if USE_TMA_LOAD_W: expert_stride = N * K - w_desc = _make_tensor_descriptor( + w_desc = tl.make_tensor_descriptor( w_ptr, shape = [NUM_EXPERTS, N, K], strides = [expert_stride, K, 1], @@ -136,7 +123,7 @@ def _grouped_gemm_dX_kernel( tl.static_assert( K % BLOCK_SIZE_K == 0, "K must be divisible by BLOCK_SIZE_K" ) - dX_desc = _make_tensor_descriptor( + dX_desc = tl.make_tensor_descriptor( dX_ptr, shape = [m_end, K], strides = [K, 1], @@ -337,7 +324,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 = _make_tensor_descriptor( + dY_desc = tl.make_tensor_descriptor( dY_ptr, shape = [TOTAL_TOKENS, N], strides = [N, 1], @@ -345,7 +332,7 @@ def _grouped_gemm_dW_kernel( ) if USE_TMA_LOAD_X and not TMA_LOAD_BOTH: - x_desc = _make_tensor_descriptor( + x_desc = tl.make_tensor_descriptor( x_ptr, shape = [TOTAL_TOKENS, K], strides = [K, 1], @@ -364,7 +351,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 = _make_tensor_descriptor( + dW_desc = tl.make_tensor_descriptor( dW_ptr, shape = [NUM_EXPERTS, N, K], strides = [N * K, K, 1], @@ -405,14 +392,14 @@ def _grouped_gemm_dW_kernel( if m_size > 0: if TMA_LOAD_BOTH: - dY_desc = _make_tensor_descriptor( + dY_desc = tl.make_tensor_descriptor( dY_ptr, shape = [m_end, N], strides = [N, 1], block_shape = [BLOCK_SIZE_M, BLOCK_SIZE_N], ) - x_desc = _make_tensor_descriptor( + x_desc = tl.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 d038d55688..a42ec5ffe9 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py @@ -10,21 +10,6 @@ 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 @@ -81,7 +66,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 = _make_tensor_descriptor( + x_desc = tl.make_tensor_descriptor( x_ptr, shape = [TOTAL_TOKENS, K], strides = [K, 1], @@ -90,7 +75,7 @@ def _grouped_gemm_forward_kernel( if USE_TMA_LOAD_W: expert_stride = N * K - w_desc = _make_tensor_descriptor( + w_desc = tl.make_tensor_descriptor( w_ptr, shape = [NUM_EXPERTS, N, K], strides = [expert_stride, K, 1], @@ -115,7 +100,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 = _make_tensor_descriptor( + y_desc = tl.make_tensor_descriptor( y_ptr, # + m_start * N, shape = [m_end, N], strides = [N, 1],