non TMA [T4]
This commit is contained in:
Datta Nimmaturi
parent
ee669f3f30
commit
9bb0db2a66
3 changed files with 55 additions and 20 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue