diff --git a/tests/jax/test_mxfp8_cutedsl_backend.py b/tests/jax/test_mxfp8_cutedsl_backend.py index 2ee82f8c5d5..c44b93b5614 100644 --- a/tests/jax/test_mxfp8_cutedsl_backend.py +++ b/tests/jax/test_mxfp8_cutedsl_backend.py @@ -23,6 +23,7 @@ from transformer_engine.common import _get_shared_object_file from transformer_engine.common.CuTeDSL.utils import device_compute_capability from transformer_engine.jax import cpp_extensions as tex +from transformer_engine.jax.cpp_extensions.activation import ActivationParams from transformer_engine.jax.quantize import ( QuantizerFactory, QuantizeLayout, @@ -68,8 +69,8 @@ # Every activation the CuTeDSL backend implements, as (JAX activation_type, name it goes by # in the config key). The two differ for qgelu/srelu: JAX spells them quick_gelu/squared_relu # (ActivationEnum in jax/cpp_extensions/activation.py), the backend uses the C++ spelling -# (Activation in common/util/cutedsl_utils.h). Gated variants are absent because they go to a -# separate TE/common kernel that the backend does not cover. +# (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel and are +# covered by GATED_ACTIVATIONS below. ACTIVATIONS = [ (("gelu",), "gelu"), (("relu",), "relu"), @@ -316,3 +317,119 @@ def test_dtypes(method, act_type, act_desc, fp8_dtype, in_dtype): ) def test_specialized_kernels(shape, q_layout, in_dtype): run_test_case("CAST_ONLY", *NO_ACT, shape, q_layout, in_dtype, jnp.float8_e4m3fn) + + +# Gated activations (gated_mxfp8.cuh), as (JAX activation_type, name in the config key +# (GatedActivation in common/util/cutedsl_utils.h), activation parameters). The clamped SwiGLU +# parameters are chosen so that its clamps are active on the [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + (("gelu", "linear"), "geglu", {}), + (("quick_gelu", "linear"), "qgeglu", {}), + (("relu", "linear"), "reglu", {}), + (("squared_relu", "linear"), "sreglu", {}), + (("silu", "linear"), "swiglu", {}), + (("sigmoid", "linear"), "glu", {}), + ( + ("clamped_silu", "clamped_linear"), + "clamped_swiglu", + {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + ), +] +GATED_IDS = [desc for _, desc, _ in GATED_ACTIVATIONS] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + # JAX quantize emits scales in the linear (non-swizzled) layout + flags = (True, q_layout.has_colwise, False, is_bwd) + major, minor = device_compute_capability() + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act_desc}" + ) + + +def run_gated_test_case( + act_type, act_desc, act_kwargs, is_bwd, shape, q_layout, in_dtype, fp8_dtype +): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is (M, 2, N), [act | gate] along axis -2. + """ + M, N = shape + # generate_inputs returns (M, N) and (M, 1, N) draws; stack them into the (M, 2, N) input + act, gate = generate_inputs(M, N, in_dtype) + x = jnp.concatenate([jnp.expand_dims(act, axis=-2), gate], axis=-2) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + act_params = ActivationParams.create(act_type, **act_kwargs) + + def run(): + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=fp8_dtype, q_layout=q_layout + ) + if is_bwd: + out, _ = tex.quantize_dact_dbias( + grad, x, act_type, is_dbias=False, quantizer=quantizer, act_params=act_params + ) + else: + out = tex.act_lu(x, act_type, quantizer=quantizer, act_params=act_params) + return extract_quantized_output(out, None)[0] + + set_cutedsl_backend(False) + cuda_output = run() + set_cutedsl_backend(True) + try: + cutedsl_output = run() + finally: + set_cutedsl_backend(False) + + key = get_gated_cfg_key(act_desc, in_dtype, fp8_dtype, q_layout, is_bwd) + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + tag = ( + f"{act_desc}/{get_direction_id(is_bwd)}/{get_layout_id(q_layout)}/{M}x{N}/" + f"{DTYPE_TO_STR[in_dtype]}/{FP8_TO_STR[fp8_dtype]}" + ) + for name, cuda_bytes in cuda_output.items(): + assert np.array_equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("q_layout", Q_LAYOUTS, ids=get_layout_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_sizes(is_bwd, act_type, act_desc, act_kwargs, q_layout, shape): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + shape, + q_layout, + jnp.bfloat16, + jnp.float8_e4m3fn, + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act_type,act_desc,act_kwargs", GATED_ACTIVATIONS, ids=GATED_IDS) +@pytest.mark.parametrize("is_bwd", [False, True], ids=get_direction_id) +def test_gated_dtypes(is_bwd, act_type, act_desc, act_kwargs, fp8_dtype, in_dtype): + run_gated_test_case( + act_type, + act_desc, + act_kwargs, + is_bwd, + (256, 384), + QuantizeLayout.ROWWISE_COLWISE, + in_dtype, + fp8_dtype, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py index 5d7feeb6ce6..fa51d292cb7 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py +++ b/tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py @@ -63,8 +63,8 @@ BLOCK_SIZES = [ROW, COL, BOTH] # Every activation the CuTeDSL backend implements; "desc" is the name it goes by in the -# config key (Activation in common/util/cutedsl_utils.h). Gated variants are absent because -# they go to a separate TE/common kernel that the backend does not cover. +# config key (Activation in common/util/cutedsl_utils.h). Gated variants go to a separate kernel +# and are covered by GATED_ACTIVATIONS below. IDENTITY = {"name": "Identity", "act": None, "dact": None, "dbias_dact": None, "desc": "none"} ACTIVATIONS = [ { @@ -431,3 +431,126 @@ def test_specialized_kernels(shape, block_size, in_dtype, swizzled, input_offset swizzled, input_offset, ) + + +# Gated activations (gated_mxfp8.cuh). Input is [M, 2N] = [act | gate]; the forward writes +# [M, N] and the backward, given grad [M, N], writes [M, 2N]. "desc" is the name the activation +# goes by in the config key (GatedActivation in common/util/cutedsl_utils.h). The parameters of +# the parameterized ones are chosen so that their clamps and tanh terms are active on the +# [-2, 2] inputs generate_inputs draws. +GATED_ACTIVATIONS = [ + {"name": "GeGLU", "act": tex.geglu, "dact": tex.dgeglu, "desc": "geglu", "kwargs": {}}, + {"name": "QGeGLU", "act": tex.qgeglu, "dact": tex.dqgeglu, "desc": "qgeglu", "kwargs": {}}, + {"name": "ReGLU", "act": tex.reglu, "dact": tex.dreglu, "desc": "reglu", "kwargs": {}}, + {"name": "SReGLU", "act": tex.sreglu, "dact": tex.dsreglu, "desc": "sreglu", "kwargs": {}}, + {"name": "SwiGLU", "act": tex.swiglu, "dact": tex.dswiglu, "desc": "swiglu", "kwargs": {}}, + {"name": "GLU", "act": tex.glu, "dact": tex.dglu, "desc": "glu", "kwargs": {}}, + { + "name": "ClampedSwiGLU", + "act": tex.clamped_swiglu, + "dact": tex.clamped_dswiglu, + "desc": "clamped_swiglu", + "kwargs": {"limit": 1.0, "alpha": 1.702, "glu_linear_offset": 0.5}, + }, + { + "name": "SiTUGLU", + "act": tex.situglu, + "dact": tex.dsituglu, + "desc": "situglu", + "kwargs": {"beta1": 1.0, "beta2": 2.0}, + }, +] +get_gated_id = lambda a: a["name"] +GATED_DIRECTIONS = [False, True] +get_direction_id = lambda is_bwd: "bwd" if is_bwd else "fwd" + + +def run_on_both_backends(run): + """Return (CUDA result, CuTeDSL result) of run(), with run() responsible for materializing + whatever it returns before the backend is switched.""" + set_cutedsl_backend(False) + cuda_result = run() + set_cutedsl_backend(True) + try: + cutedsl_result = run() + finally: + set_cutedsl_backend(False) + return cuda_result, cutedsl_result + + +def assert_cutedsl_kernel_registered(key): + """Guard against a silent CUDA fallback, which would make a cross-backend comparison compare + CUDA against itself.""" + assert tvm_ffi.get_global_func(key, allow_missing=True) is not None, ( + f"CuTeDSL kernel not registered for {key}; the CuTeDSL backend fell back " + "to CUDA and this case compared CUDA against itself" + ) + + +def get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, colwise, swizzled, is_bwd): + """Mirror of MXFP8GatedQuantConfig::to_key (gated_mxfp8_cutedsl.cuh).""" + major, minor = device_compute_capability() + flags = (rowwise, colwise, swizzled, is_bwd) + return ( + f"cutedsl_mxfp8_gated_sm{major * 10 + minor}_{DTYPE_TO_STR[in_dtype]}_" + f"{FP8_TO_KEY[fp8_dtype]}_" + + "_".join("1" if f else "0" for f in flags) + + f"_{act['desc']}" + ) + + +def run_gated_test_case(act, is_bwd, shape, block_size, in_dtype, fp8_dtype, swizzled=False): + """Assert the CuTeDSL and CUDA backends produce bit-identical gated activation outputs. + shape is (M, N): the gated input is [M, 2N]. + """ + M, N = shape + if is_bwd and swizzled: + pytest.skip("CuTeDSL hands backward gated activations with GEMM-swizzled scales to CUDA") + rowwise = block_size[1] != 1 + columnwise = block_size[0] != 1 + x, _ = generate_inputs(M, 2 * N, in_dtype) + grad, _ = generate_inputs(M, N, in_dtype, seed=1) + + def run(): + q = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=rowwise, columnwise=columnwise) + q.optimize_for_gemm = swizzled + if is_bwd: + out = act["dact"](grad, x, q, **act["kwargs"]) + else: + out = act["act"](x, q, **act["kwargs"]) + return extract_quantized_output(out, rowwise, columnwise, swizzled) + + cuda_output, cutedsl_output = run_on_both_backends(run) + assert_cutedsl_kernel_registered( + get_gated_cfg_key(act, in_dtype, fp8_dtype, rowwise, columnwise, swizzled, is_bwd) + ) + + layout = "swizzled" if swizzled else "linear" + tag = ( + f"{act['name']}/{get_direction_id(is_bwd)}/{M}x{N}/{DTYPE_TO_STR[in_dtype]}/" + f"{FP8_TO_STR[fp8_dtype]}/{layout}" + ) + for name, cuda_bytes in cuda_output.items(): + assert torch.equal( + cutedsl_output[name], cuda_bytes + ), f"{tag}: {name} differ between backends" + + +@pytest.mark.parametrize("shape", MATRIX_SIZES, ids=get_shape_id) +@pytest.mark.parametrize("block_size", BLOCK_SIZES, ids=get_block_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_sizes(swizzled, is_bwd, act, block_size, shape): + run_gated_test_case( + act, is_bwd, shape, block_size, torch.bfloat16, tex.DType.kFloat8E4M3, swizzled + ) + + +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=get_dtype_id) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES, ids=get_fp8_id) +@pytest.mark.parametrize("act", GATED_ACTIVATIONS, ids=get_gated_id) +@pytest.mark.parametrize("is_bwd", GATED_DIRECTIONS, ids=get_direction_id) +@pytest.mark.parametrize("swizzled", SWIZZLE_MODES, ids=get_swizzle_id) +def test_gated_dtypes(swizzled, is_bwd, act, fp8_dtype, in_dtype): + run_gated_test_case(act, is_bwd, (256, 384), BOTH, in_dtype, fp8_dtype, swizzled) diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 15aae1aba2d..aac7f232ab0 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -12,6 +12,7 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8_common import ( get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, ) from transformer_engine.common.CuTeDSL.cast.mxfp8.requantize_mxfp8 import ( get_mxfp8_requantization_function, @@ -22,9 +23,9 @@ def register_cutedsl_backends(): """Register all available CuTeDSL backends for on-demand compilation via TVM-FFI. The C++ dispatcher retrieves them by the names defined here. """ - tvm_ffi.register_global_func( - "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True - ) - tvm_ffi.register_global_func( - "get_mxfp8_requantization_function", get_mxfp8_requantization_function, override=True - ) + for entrypoint in ( + get_mxfp8_quantization_function, + get_mxfp8_gated_quantization_function, + get_mxfp8_requantization_function, + ): + tvm_ffi.register_global_func(entrypoint.__name__, entrypoint, override=True) diff --git a/transformer_engine/common/CuTeDSL/activations.py b/transformer_engine/common/CuTeDSL/activations.py index 34772eb7257..d65fbca6f49 100644 --- a/transformer_engine/common/CuTeDSL/activations.py +++ b/transformer_engine/common/CuTeDSL/activations.py @@ -8,6 +8,7 @@ # Each function is a 1:1 port of its math.h counterpart (see module docstring). # pylint: disable=missing-function-docstring import os +from typing import Optional from cutlass import cute from cutlass import Float32 @@ -21,8 +22,9 @@ def tanh(x: Float32) -> Float32: return cute.math.tanh(x, fastmath=USE_FAST_MATH) -def exp(x: Float32) -> Float32: - return cute.math.exp(x, fastmath=USE_FAST_MATH) +def exp(x: Float32, fast_math: bool = False) -> Float32: + """exp(x); fast_math, or NVTE_USE_FAST_MATH, selects the fast approximate exp.""" + return cute.math.exp(x, fastmath=fast_math or USE_FAST_MATH) def act_relu(x: Float32) -> Float32: @@ -35,8 +37,11 @@ def act_gelu(x: Float32) -> Float32: return x * (Float32(0.5) + Float32(0.5) * tanh(x * (A + B * x * x))) -def act_silu(x: Float32) -> Float32: - return x / (Float32(1.0) + exp(-x)) +def act_silu(x: Float32, sig: Optional[Float32] = None, fast_math: bool = False) -> Float32: + # Recompute the sigmoid if not provided one + if sig is None: + sig = sigmoid(x, fast_math) + return x * sig def act_qgelu(x: Float32) -> Float32: @@ -57,17 +62,20 @@ def dact_dsrelu(x: Float32) -> Float32: return cute.arch.fmax(Float32(2.0) * x, Float32(0.0)) -def sigmoid(x: Float32) -> Float32: - return Float32(1.0) / (Float32(1.0) + exp(-x)) +def sigmoid(x: Float32, fast_math: bool = False) -> Float32: + return Float32(1.0) / (Float32(1.0) + exp(-x, fast_math)) -def dsigmoid(x: Float32) -> Float32: - s = sigmoid(x) +def dsigmoid(x: Float32, fast_math: bool = False) -> Float32: + s = sigmoid(x, fast_math) return s * (Float32(1.0) - s) -def dact_dsilu(x: Float32) -> Float32: - return fma_f32(x, dsigmoid(x), sigmoid(x)) +def dact_dsilu(x: Float32, sig: Optional[Float32] = None, fast_math: bool = False) -> Float32: + # Recompute the sigmoid if not provided one + if sig is None: + return fma_f32(x, dsigmoid(x, fast_math), sigmoid(x, fast_math)) + return x * sig * (Float32(1.0) - sig) + sig def dact_dqgelu(x: Float32) -> Float32: @@ -80,3 +88,54 @@ def dact_dgelu(x: Float32) -> Float32: return Float32(0.5) * x * ( (Float32(1.0) - t * t) * (Float32(0.79788456) + Float32(0.1070322243) * x * x) ) + Float32(0.5) * (Float32(1.0) + t) + + +# ---- Gated activations (gated_mxfp8.cuh) ---- +# The parameterized ones take their ClampedSwiGLUParam / SiTUGLUParam fields as +# trailing runtime arguments. + + +def act_sigmoid(x: Float32) -> Float32: + return sigmoid(x) + + +def dact_dsigmoid(x: Float32) -> Float32: + return dsigmoid(x) + + +def act_clamped_silu(x: Float32, limit: Float32, alpha: Float32) -> Float32: + c = cute.arch.fmin(limit, x) + return c * sigmoid(alpha * c) + + +def act_situ_gate(x: Float32, beta1: Float32) -> Float32: + """SiTU-GLU's activation, applied to the act half.""" + return beta1 * tanh(x / beta1) * sigmoid(x) + + +def dact_dsitu_gate(x: Float32, beta1: Float32) -> Float32: + """Derivative of act_situ_gate.""" + t = tanh(x / beta1) + s = sigmoid(x) + return (Float32(1.0) - t * t) * s + beta1 * t * s * (Float32(1.0) - s) + + +def identity(x: Float32) -> Float32: + """The default gate: gating is ideally the identity.""" + return x + + +def didentity(x: Float32) -> Float32: # pylint: disable=unused-argument + """Derivative of identity.""" + return Float32(1.0) + + +def situ_up(x: Float32, beta2: Float32) -> Float32: + """SiTU-GLU's gate transform, applied to the gate half.""" + return beta2 * tanh(x / beta2) + + +def dsitu_up(x: Float32, beta2: Float32) -> Float32: + """Derivative of situ_up.""" + t = tanh(x / beta2) + return Float32(1.0) - t * t diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py index 8a66b55e48f..164a93dabff 100644 --- a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py @@ -30,6 +30,8 @@ dact_dsilu, dact_dqgelu, dact_dgelu, + act_sigmoid, + dact_dsigmoid, ) CUTEDSL_DEBUG_LOGGING = os.environ.get("CUTEDSL_DEBUG_LOGGING", "0") == "1" @@ -67,6 +69,17 @@ "dsrelu": dact_dsrelu, } +GATED_ACTIVATIONS = { + "geglu": (act_gelu, dact_dgelu), + "qgeglu": (act_qgelu, dact_dqgelu), + "reglu": (act_relu, dact_drelu), + "sreglu": (act_srelu, dact_dsrelu), + "glu": (act_sigmoid, dact_dsigmoid), + "swiglu": (act_silu, dact_dsilu), + "clamped_swiglu": (None, None), + "situglu": (None, None), +} + @cute.jit def derive_swizzled_scale_layout( @@ -212,6 +225,233 @@ def is_plain_cast(self) -> bool: """Whether cfg is a plain cast: no amax, dbias or (d)activation fused in.""" return not (self.WITH_AMAX or self.WITH_DBIAS or self.WITH_DACT or self.WITH_ACT) + def compile(self): + """ + Return the compiled CuTeDSL function object for the given MXFP8 quantization config. + """ + + kernel_obj = MXFP8QuantizeEntry(self) + sym_M = cute.sym_int32() + sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) + in_shape = out_shape = (sym_M, sym_N) + # TE allocates scale tensors at a padded shape (see + # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): + # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) + # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) + # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the + # scales their own fresh syms carrying the divisibility the padding + # guarantees (rowwise: 128 x 4; colwise: 4 x 128). + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N + # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ + # bridge sends). + out_dtype = self.FP8_DTYPE + scale_dtype = cutlass.Float8E8M0FNU + + in_fake = cute.runtime.make_fake_compact_tensor( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + out_row_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.ROWWISE + else None + ) + scale_row_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_rowwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.ROWWISE + else None + ) + out_col_fake = ( + cute.runtime.make_fake_compact_tensor( + out_dtype, + out_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.COLWISE + else None + ) + scale_col_fake = ( + cute.runtime.make_fake_compact_tensor( + scale_dtype, + scale_colwise_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.COLWISE + else None + ) + amax_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 + ) + if self.WITH_AMAX + else None + ) + # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that + # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). + # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher + # passes at launch. + noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) + act_input_fake = ( + cute.runtime.make_fake_compact_tensor( + self.DTYPE, + in_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + if self.WITH_DACT + else None + ) + workspace_fake = ( + cute.runtime.make_fake_compact_tensor( + Float32, + ws_shape, + stride_order=(1, 0), + memspace=cute.AddressSpace.gmem, + assumed_align=4, + ) + if self.WITH_DBIAS + else None + ) + + compiled = cute.compile( + kernel_obj, + in_fake, # mX + out_row_fake, + scale_row_fake, # mO_row, mS_row + out_col_fake, + scale_col_fake, # mO_col, mS_col + amax_fake, # mAmax + noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) + act_input_fake, # mDActInput (activation input, dact only) + workspace_fake, # mWorkspace(backward slot, unused) + cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi + # "handle" arg; C++ passes the CUDA stream + # as void*) + options="--enable-tvm-ffi", + ) + return compiled + + +def _fake_gmem_tensor(dtype, shape, align): + """Row-major fake gmem tensor for cute.compile.""" + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + memspace=cute.AddressSpace.gmem, + assumed_align=align, + ) + + +class MXFP8GatedQuantizeConfig(MXFP8QuantizeConfig): + """Compile-time config of gated_mxfp8.cuh's port: a plain MXFP8 cast config plus the gated op.""" + + def __init__( + self, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, + ): + super().__init__( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=False, + activation="none", + ) + if activation not in GATED_ACTIVATIONS: + raise ValueError( + f"unknown gated activation {activation!r}; expected one of" + f" {sorted(GATED_ACTIVATIONS)}" + ) + if is_bwd and with_gemm_swizzled_scales: + # TODO(kainingz): technically this is supported but I need to come up with a less disgusting way + # to implement it in CuTe + raise ValueError("backward gated activations with GEMM-swizzled scales") + self.IS_BWD = is_bwd + self.GATED_ACTIVATION = activation + + def __str__(self): + return ( + f"MXFP8GatedQuantizeConfig(dtype={self.DTYPE_STR}, fp8_dtype={self.FP8_DTYPE_STR}, " + f"rowwise={self.ROWWISE}, colwise={self.COLWISE}, " + f"swizzled={self.WITH_GEMM_SWIZZLED_SCALES}, is_bwd={self.IS_BWD}, " + f"activation={self.GATED_ACTIVATION})" + ) + + __repr__ = __str__ + + def compile(self): + """Return the compiled CuTeDSL function object for this config.""" + # The kernel module imports this one + # pylint: disable-next=import-outside-toplevel + from .quantize_mxfp8_gated_g2s import MXFP8GatedQuantizeG2SKernel + + sym_M = cute.sym_int32() + # The C++ dispatcher requires N % 32 == 0, like the CUDA kernel. + sym_N = cute.sym_int32(divisibility=32) + sym_2N = cute.sym_int32(divisibility=64) + out_cols = sym_2N if self.IS_BWD else sym_N + # Padded scale shapes, see MXFP8QuantizeConfig.compile + scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) + scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) + # Only the activation's own params are runtime arguments; the others are compiled as None + # and the C++ launcher passes None for them. + clamped = self.GATED_ACTIVATION == "clamped_swiglu" + situ = self.GATED_ACTIVATION == "situglu" + params = [Float32(0.0) if clamped else None] * 3 # limit, alpha, glu_linear_offset + params += [Float32(0.0) if situ else None] * 2 # beta1, beta2 + return cute.compile( + MXFP8GatedQuantizeG2SKernel(self), + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_2N), 16), # mGatedInput + _fake_gmem_tensor(self.DTYPE, (sym_M, sym_N), 16) if self.IS_BWD else None, # mYGrad + # mO_row, mS_row + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.ROWWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_rowwise_shape, 4) + if self.ROWWISE + else None + ), + # mO_col, mS_col + _fake_gmem_tensor(self.FP8_DTYPE, (sym_M, out_cols), 16) if self.COLWISE else None, + ( + _fake_gmem_tensor(cutlass.Float8E8M0FNU, scale_colwise_shape, 4) + if self.COLWISE + else None + ), + *params, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + class MXFP8QuantizeKernelBase(abc.ABC): """Base class for MXFP8 quantize kernels.""" @@ -404,142 +644,8 @@ def __call__( ) -def compile_cutedsl_function_from_cfg(cfg): - """ - Return the compiled CuTeDSL function object for the given MXFP8 quantization config. - """ - - kernel_obj = MXFP8QuantizeEntry(cfg) - sym_M = cute.sym_int32() - sym_N = cute.sym_int32(divisibility=SYM_N_DIVISIBILITY) - in_shape = out_shape = (sym_M, sym_N) - # TE allocates scale tensors at a padded shape (see - # MXFP8Quantizer::get_scale_shape in transformer_engine/pytorch/csrc): - # rowwise: (roundup(M, 128), roundup(ceildiv(N, 32), 4)) - # columnwise: (roundup(ceildiv(M, 32), 4), roundup(N, 128)) - # These padded extents are NOT M/N (and SymInt has no `//`/`+`), so give the - # scales their own fresh syms carrying the divisibility the padding - # guarantees (rowwise: 128 x 4; colwise: 4 x 128). - scale_rowwise_shape = (cute.sym_int32(divisibility=128), cute.sym_int32(divisibility=4)) - scale_colwise_shape = (cute.sym_int32(divisibility=4), cute.sym_int32(divisibility=128)) - ws_shape = (cute.sym_int32(), sym_N) # (blocks_Y, N); N ties to input N - # Native FP8/E8M0 dtypes at the FFI boundary (matches the DLPack dtype the C++ - # bridge sends). - out_dtype = cfg.FP8_DTYPE - scale_dtype = cutlass.Float8E8M0FNU - - in_fake = cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, in_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=16 - ) - out_row_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.ROWWISE - else None - ) - scale_row_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_rowwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.ROWWISE - else None - ) - out_col_fake = ( - cute.runtime.make_fake_compact_tensor( - out_dtype, - out_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.COLWISE - else None - ) - scale_col_fake = ( - cute.runtime.make_fake_compact_tensor( - scale_dtype, - scale_colwise_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=4, - ) - if cfg.COLWISE - else None - ) - amax_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, (1,), stride_order=(0,), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_AMAX - else None - ) - # The cast-noop flag is an always-present f32 pointer instead of an optional tensor, so that - # the same compiled kernel serves both an absent and a present flag (see noop_flag_is_set). - # The null address here is only a compile-time placeholder for the pointer the C++ dispatcher - # passes at launch. - noop_fake = cute.runtime.nullptr(Float32, mem_space=cute.AddressSpace.gmem, assumed_align=4) - act_input_fake = ( - cute.runtime.make_fake_compact_tensor( - cfg.DTYPE, - in_shape, - stride_order=(1, 0), - memspace=cute.AddressSpace.gmem, - assumed_align=16, - ) - if cfg.WITH_DACT - else None - ) - workspace_fake = ( - cute.runtime.make_fake_compact_tensor( - Float32, ws_shape, stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4 - ) - if cfg.WITH_DBIAS - else None - ) - - compiled = cute.compile( - kernel_obj, - in_fake, # mX - out_row_fake, - scale_row_fake, # mO_row, mS_row - out_col_fake, - scale_col_fake, # mO_col, mS_col - amax_fake, # mAmax - noop_fake, # mNoop (pointer to the cast_noop flag, may be null at launch) - act_input_fake, # mDActInput (activation input, dact only) - workspace_fake, # mWorkspace(backward slot, unused) - cute.runtime.make_fake_stream(), # stream (compiled as an explicit tvm-ffi - # "handle" arg; C++ passes the CUDA stream - # as void*) - options="--enable-tvm-ffi", - ) - return compiled - - -def get_mxfp8_quantization_function( - fn_name: str, - dtype: str, - fp8_dtype: str, - rowwise: bool, - colwise: bool, - with_gemm_swizzled_scales: bool, - with_amax: bool, - with_dbias: bool, - with_dact: bool, - with_act: bool, - use_2d_quantization: bool, - activation: str, -) -> bool: - """Compile the MXFP8 quantize kernel for this config and register it in the TVM-FFI global registry +def compile_and_register_function(fn_name: str, cfg: MXFP8QuantizeConfig) -> bool: + """Compile the CuTeDSL kernel for this config and register it in the TVM-FFI global registry under EXACTLY `fn_name` (the key the C++ dispatcher built; Python treats it as an opaque name). Returns True if a kernel is successfully registered under `fn_name` (the C++ side then fetches it with GetGlobal(fn_name)); False if the config is unsupported, so the caller caches the negative result and falls back to the CUDA C++ kernel. @@ -559,33 +665,8 @@ def get_mxfp8_quantization_function( ) return False - try: - cfg = MXFP8QuantizeConfig( - dtype=dtype, - fp8_dtype=fp8_dtype, - rowwise=rowwise, - colwise=colwise, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - with_amax=with_amax, - with_dbias=with_dbias, - with_dact=with_dact, - with_act=with_act, - use_2d_quantization=use_2d_quantization, - activation=activation, - ) - except ValueError as e: - # The exception message states exactly why the config is unsupported - # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a - # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. - logger.warning( - "CuTeDSL MXFP8 backend does not support this config, " - "falling back to the CUDA C++ kernel: %s", - e, - ) - return False - logger.debug("Compiling CuTeDSL MXFP8 quantization kernel for %s", cfg) - compiled = compile_cutedsl_function_from_cfg(cfg) + compiled = cfg.compile() # The returned compiled object is not neccessarily the compiled function itself. It could be a Python # wrapper that parses arguments and calls the underlying function. This is needed in the general case # as TVM-FFI is positional-only, however, for us this is unnecessary launch overhead as we pass the @@ -607,3 +688,79 @@ def get_mxfp8_quantization_function( ) # Unconditionally fallback to CUDA path because we can't tell if this exception is transient or permanent. return False + + +def get_mxfp8_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + with_amax: bool, + with_dbias: bool, + with_dact: bool, + with_act: bool, + use_2d_quantization: bool, + activation: str, +) -> bool: + """Compile and register the MXFP8 quantize kernel for this config under `fn_name`.""" + try: + cfg = MXFP8QuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + with_amax=with_amax, + with_dbias=with_dbias, + with_dact=with_dact, + with_act=with_act, + use_2d_quantization=use_2d_quantization, + activation=activation, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + return compile_and_register_function(fn_name, cfg) + + +def get_mxfp8_gated_quantization_function( + fn_name: str, + dtype: str, + fp8_dtype: str, + rowwise: bool, + colwise: bool, + with_gemm_swizzled_scales: bool, + is_bwd: bool, + activation: str, +) -> bool: + """Compile and register the gated MXFP8 quantize kernel for this config under `fn_name`.""" + try: + cfg = MXFP8GatedQuantizeConfig( + dtype=dtype, + fp8_dtype=fp8_dtype, + rowwise=rowwise, + colwise=colwise, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + is_bwd=is_bwd, + activation=activation, + ) + except ValueError as e: + # The exception message states exactly why the config is unsupported + # (unknown dtype/activation, dbias not implemented, ...). Surfacing it as a + # warning lets the C++ dispatcher's CUDA fallback be recognized as expected. + logger.warning( + "CuTeDSL MXFP8 backend does not support this config, " + "falling back to the CUDA C++ kernel: %s", + e, + ) + return False + return compile_and_register_function(fn_name, cfg) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py new file mode 100644 index 00000000000..dea96563611 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2s.py @@ -0,0 +1,852 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Gated-activation MXFP8 casts in CuTeDSL, ported from gated_mxfp8.cuh. + +The kernel reuses MXFP8QuantizeGeneralG2SKernel's TMA pipeline and its plain-cast SMEM tile +quantizer. What it adds is an elementwise pass that runs on each staged tile before it is +quantized: it combines the act and gate halves of the input (and the incoming gradient y_grad, +in backward) and writes the result back into the act (and gate) input buffer in the input dtype. +That is exactly the truncation the CUDA kernel applies before computing block amaxes, so the +plain cast that follows reproduces its output. + +Gated input is [M, 2N] with [act | gate] along the last dim. Every operand is addressed through +(M, N) views (stride 2N, column offset 0 or N), mirroring the CUDA kernel's tensor maps: + forward : y[M, N] = act(x[:, :N]) * gate(x[:, N:]) + backward: [act_grad | gate_grad][M, 2N], the gradients of the act and gate halves given + y_grad[M, N] +""" + +from typing import Optional + +import cutlass +from cutlass import cute +from cutlass import pipeline +from cutlass import Float32, Int32 +from cuda.bindings.driver import CUstream # pylint: disable=no-name-in-module + +from transformer_engine.common.CuTeDSL.utils import ( + fabs_f32, + is_packed16, + select_f32, + truncate_f32_bf16, + truncate_f32_f16, +) +from transformer_engine.common.CuTeDSL.activations import ( + act_clamped_silu, + act_situ_gate, + dact_dsitu_gate, + identity, + didentity, + situ_up, + dsitu_up, + act_silu, + dact_dsilu, + sigmoid, +) +from .quantize_mxfp8_common import ( + CUTEDSL_DEBUG_LOGGING, + GATED_ACTIVATIONS, + MXFP8GatedQuantizeConfig, + MXFP8_BLOCK_SCALING_SIZE, + derive_swizzled_scale_layout, +) +from .quantize_mxfp8_g2s_general import ( + MXFP8QuantizeGeneralG2SKernel, + quantize_mxfp8_smem_tile_general, +) + + +@cute.jit +def gated_elementwise_op( + act_element: Float32, + gate_element: Float32, + y_grad_element: Optional[Float32], + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTU-GLU only + beta2: Optional[Float32], # SiTU-GLU only + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], +): + """ + FWD: (y, None) y = act(a) * gate(g) + BWD: (act_grad, gate_grad) act_grad = y_grad * gate(g) * dact(a), + gate_grad = y_grad * act(a) * dgate(g) + """ + + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + # FWD gate function for clamped_swiglu: + # gate(g) = clamp(g, -limit, limit) + glu_linear_offset + gate = cute.arch.fmin(cute.arch.fmax(-limit, gate_element), limit) + glu_linear_offset + # BWD dgate function for clamped_swiglu: + # dgate(g) = 0 if g < -limit or g > limit + # = 1 if -limit <= g <= limit + dgate = select_f32(fabs_f32(gate_element) > limit, Float32(0.0), Float32(1.0)) + elif cutlass.const_expr(ACTIVATION == "situglu"): + gate = situ_up(gate_element, beta2) + dgate = dsitu_up(gate_element, beta2) + else: + gate = identity(gate_element) + dgate = didentity(gate_element) + + act_fn, dact_fn = GATED_ACTIVATIONS[ACTIVATION] + + if cutlass.const_expr(not IS_BWD): + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + act = act_clamped_silu(act_element, limit, alpha) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act = act_situ_gate(act_element, beta1) + else: + act = act_fn(act_element) + y = act * gate + outputs = (y, None) + else: + if cutlass.const_expr(ACTIVATION == "clamped_swiglu"): + x = cute.arch.fmin(act_element, limit) + # Somehow CUDA always uses fast math for this and we should too to match its output + s = sigmoid(alpha * x, fast_math=True) + act = x * s + dact = select_f32( + act_element <= limit, s + s * (Float32(1.0) - s) * alpha * x, Float32(0.0) + ) + elif cutlass.const_expr(ACTIVATION == "situglu"): + act = act_situ_gate(act_element, beta1) + dact = dact_dsitu_gate(act_element, beta1) + elif cutlass.const_expr(ACTIVATION == "swiglu"): + # Match CUDA implementation's behavior too + s = sigmoid(act_element, fast_math=True) + act = act_silu(act_element, s) + dact = dact_dsilu(act_element, s) + else: + act = act_fn(act_element) + dact = dact_fn(act_element) + act_grad = dact * y_grad_element * gate + gate_grad = dgate * act * y_grad_element + outputs = (act_grad, gate_grad) + return outputs + + +@cute.jit +def gated_elementwise_smem_tile( + sAct_tile: cute.Tensor, # also the y tile in FWD / the act_grad tile in BWD + sGate_tile: cute.Tensor, # also the gate_grad tile in BWD + sYGrad_tile: Optional[cute.Tensor], # BWD only + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTU-GLU only + beta2: Optional[Float32], # SiTU-GLU only + ACTIVATION: cutlass.Constexpr[str], + IS_BWD: cutlass.Constexpr[bool], + ROWWISE: cutlass.Constexpr[bool], + DTYPE: cutlass.Constexpr[type], + TILE_X: cutlass.Constexpr[int], +): + """Compute the gated op on one staged tile, in place, rounded to the input dtype.""" + tidx, _, _ = cute.arch.thread_idx() + + _, tv_layout = cute.make_layout_tv( + thr_layout=cute.make_layout((1, TILE_X), stride=(TILE_X, 1)), + val_layout=cute.make_layout((MXFP8_BLOCK_SCALING_SIZE, 1), stride=(1, 1)), + ) + + sAct_tile = cute.composition(sAct_tile, tv_layout)[tidx, None] + sGate_tile = cute.composition(sGate_tile, tv_layout)[tidx, None] + if cutlass.const_expr(IS_BWD): + sYGrad_tile = cute.composition(sYGrad_tile, tv_layout)[tidx, None] + + if cutlass.const_expr(is_packed16(DTYPE)): + truncate_f32 = truncate_f32_f16 if DTYPE is cutlass.Float16 else truncate_f32_bf16 + + for i in cutlass.range_constexpr(MXFP8_BLOCK_SCALING_SIZE): + act_element = Float32(sAct_tile[i]) + gate_element = Float32(sGate_tile[i]) + # Round to the input dtype, like the CUDA kernels do before taking the amax. The cast to + # DTYPE below is then exact. + if cutlass.const_expr(IS_BWD): + y_grad_element = Float32(sYGrad_tile[i]) + act_grad_element, gate_grad_element = gated_elementwise_op( + act_element, + gate_element, + y_grad_element, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + if cutlass.const_expr(is_packed16(DTYPE)): + act_grad_element = truncate_f32(act_grad_element) + gate_grad_element = truncate_f32(gate_grad_element) + sAct_tile[i] = DTYPE(act_grad_element) + sGate_tile[i] = DTYPE(gate_grad_element) + else: + y_element, _ = gated_elementwise_op( + act_element, + gate_element, + None, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=ACTIVATION, + IS_BWD=IS_BWD, + ) + if cutlass.const_expr(is_packed16(DTYPE)): + y_element = truncate_f32(y_element) + sAct_tile[i] = DTYPE(y_element) + + # The rowwise quantization pass reads elements other threads wrote so we need to sync here + if cutlass.const_expr(ROWWISE): + cute.arch.sync_threads() + + +@cute.jit +def quantize_mxfp8_smem_tile_general_partial( + sX_tile: cute.Tensor, + sO_row_tile: Optional[cute.Tensor], + mS_row_stage: Optional[cute.Tensor], + sO_col_tile: Optional[cute.Tensor], + mS_col_stage: Optional[cute.Tensor], + tile_row_start: Int32, + tile_col_start: Int32, + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + MAX_NORM_RCP: cutlass.Constexpr[float], + DTYPE: cutlass.Constexpr[type], + FP8_DTYPE: cutlass.Constexpr[type], + WITH_GEMM_SWIZZLED_SCALES: cutlass.Constexpr[bool], + TILE_X: cutlass.Constexpr[int], + TILE_Y: cutlass.Constexpr[int], +): + """A partial quantize_mxfp8_smem_tile_general function with some args fixed.""" + quantize_mxfp8_smem_tile_general( + sX_tile, + None, + sO_row_tile, + mS_row_stage, + sO_col_tile, + mS_col_stage, + None, + tile_row_start, + tile_col_start, + M, + N, + ROWWISE=ROWWISE, + COLWISE=COLWISE, + WITH_AMAX=False, + WITH_ACT=False, + WITH_DACT=False, + ACTIVATION=None, + MAX_NORM_RCP=MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_X, + TILE_Y=TILE_Y, + # The gated op maps the zeros TMA fills out-of-bounds elements with to zero. + SKIP_INPUT_MASKING=True, + SKIP_SCALE_BOUNDS=False, + CACHE_ACTIVATION=False, + DBIAS_REDUCTION_IN_COLWISE=False, + COLWISE_DBIAS_REDUCTION_ONLY=False, + DBIAS_REDUCTION_IN_ROWWISE=False, + ) + + +@cute.jit +def _split_gated_scales( + mS_row: Optional[cute.Tensor], + mS_col: Optional[cute.Tensor], + M: Int32, + N: Int32, + ROWWISE: cutlass.Constexpr[bool], + COLWISE: cutlass.Constexpr[bool], + SWIZZLED: cutlass.Constexpr[bool], +): + """Split the scales of a gated input with logical shape [M, 2N] into the act and gate halves, each with logical shape [M, N].""" + mS_row_act = mS_row_gate = mS_col_act = mS_col_gate = None + if cutlass.const_expr(SWIZZLED): + mS_row, mS_col = derive_swizzled_scale_layout(M, 2 * N, ROWWISE, COLWISE, mS_row, mS_col) + # To slice the left half (act) and right half (gate) of the swizzled scales, we need to tile using the tiler + # with the same mode pattern as the swizzled layout's shape: + # mS_row: ((32, 4, ceil(M/128)), (4, ceil(N/64))) : ((16, 4, 512·ceil(N/64)), (1, 512)) + # mS_col: ((4, ceil(M/128)), (32, 4, ceil(N/64))) : ((1, 512), (16, 4, 512·ceil(M/128))) + if cutlass.const_expr(ROWWISE): + row_tiler = ( + cute.make_layout((32, 4, cute.ceil_div(M, 128))), # slice all rows + cute.make_layout( + (4, cute.ceil_div(N, 128)) + ), # slice half columns so cdiv uses 128 instead of 64 + ) + mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) + mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = ( + cute.make_layout((4, cute.ceil_div(M, 128))), # slice all rows + cute.make_layout( + (32, 4, cute.ceil_div(N, 128)) + ), # slice half columns so cdiv uses 128 instead of 64 + ) + mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) + mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) + else: + # Compact scales are row-major, so each half is a block of scale columns + if cutlass.const_expr(ROWWISE): + row_tiler = (M, cute.ceil_div(N, MXFP8_BLOCK_SCALING_SIZE)) + mS_row_act = cute.local_tile(mS_row, row_tiler, (0, 0)) + mS_row_gate = cute.local_tile(mS_row, row_tiler, (0, 1)) + if cutlass.const_expr(COLWISE): + col_tiler = (cute.ceil_div(M, MXFP8_BLOCK_SCALING_SIZE), N) + mS_col_act = cute.local_tile(mS_col, col_tiler, (0, 0)) + mS_col_gate = cute.local_tile(mS_col, col_tiler, (0, 1)) + return mS_row_act, mS_row_gate, mS_col_act, mS_col_gate + + +class MXFP8GatedQuantizeG2SKernel(MXFP8QuantizeGeneralG2SKernel): + """gated_mxfp8.cuh in CuTeDSL. The general kernel's pipeline, loading the act and gate halves + (and y_grad in BWD) each stage and storing one (forward) or two (backward) quantized halves. + + Inherits the plain-cast tile config, which is gated_mxfp8.cuh's: 64 threads per CTA, each CTA + covering two vertically stacked 32x64 tiles. + """ + + def __init__(self, cfg: MXFP8GatedQuantizeConfig): + super().__init__(cfg) + # Same occupancy hint as the general kernel with a fused activation + self._MIN_BLOCKS_PER_MP = 1 + # Every pipeline stage has its own SMEM buffers, so a CTA with fewer tiles than stages + # would allocate buffers it never fills + assert ( + self._NUM_STAGES <= self._NUM_TILES + ), f"{self._NUM_STAGES} pipeline stages for {self._NUM_TILES} tiles per CTA" + + @cute.jit + def __call__( + self, + mGatedInput: cute.Tensor, # [M, 2N] = [act | gate] + mYGrad: Optional[cute.Tensor], # [M, N], BWD only + # Outputs, with C = N (FWD) or 2N (BWD) + mO_row: Optional[cute.Tensor], # [M, C] + mS_row: Optional[cute.Tensor], # [roundup(M, 128), roundup(C / 32, 4)] + mO_col: Optional[cute.Tensor], # [M, C], not transposed + mS_col: Optional[cute.Tensor], # [roundup(ceil(M / 32), 4), roundup(C, 128)] + # Special activation params, None if not used by the activation + limit: Optional[Float32], # Clamped SwiGLU only + alpha: Optional[Float32], # Clamped SwiGLU only + glu_linear_offset: Optional[Float32], # Clamped SwiGLU only + beta1: Optional[Float32], # SiTUGLU only + beta2: Optional[Float32], # SiTUGLU only + stream: CUstream, + ): + """ + mGatedInput: [a | g] with shape M x 2N, where a is the activation half and g is the gate half. + mYGrad: None for FWD, y_grad with shape [M, N] for BWD + FWD mode: + y = quantize(act(a) * gate(g)) with shape [M, N] + BWD mode: + [act_grad | gate_grad] = quantize([ y_grad * gate(g) * dact(a) | y_grad * act(a) * dgate(g) ]) with shape [M, 2N] + """ + if cutlass.const_expr(CUTEDSL_DEBUG_LOGGING): + cute.printf( + f"[CuTeDSL] MXFP8GatedQuantizeG2SKernel.__call__() with config: {self.cfg}\n" + ) + + M = mGatedInput.shape[0] + # Pick any mO that is used just to get N + mO = mO_row if mO_row is not None else mO_col + # FWD: mO is [M, N]; BWD: mYGrad is [M, N] + N = mO.shape[1] if not self.cfg.IS_BWD else mYGrad.shape[1] + + # The act and gate halves of the [M, 2N] input + mAct = cute.local_tile(mGatedInput, (M, N), (0, 0)) + mGate = cute.local_tile(mGatedInput, (M, N), (0, 1)) + + mO_row_act_grad = mO_row_gate_grad = mO_col_act_grad = mO_col_gate_grad = None # BWD only + mS_row_act_grad = mS_row_gate_grad = mS_col_act_grad = mS_col_gate_grad = None # BWD only + if cutlass.const_expr(self.cfg.IS_BWD): + # The [M, 2N] output is [act_grad | gate_grad]; the kernel takes its two halves + mS_row_act_grad, mS_row_gate_grad, mS_col_act_grad, mS_col_gate_grad = ( + _split_gated_scales( + mS_row, + mS_col, + M, + N, + self.cfg.ROWWISE, + self.cfg.COLWISE, + self.cfg.WITH_GEMM_SWIZZLED_SCALES, + ) + ) + if cutlass.const_expr(self.cfg.ROWWISE): + mO_row_act_grad = cute.local_tile(mO_row, (M, N), (0, 0)) + mO_row_gate_grad = cute.local_tile(mO_row, (M, N), (0, 1)) + if cutlass.const_expr(self.cfg.COLWISE): + mO_col_act_grad = cute.local_tile(mO_col, (M, N), (0, 0)) + mO_col_gate_grad = cute.local_tile(mO_col, (M, N), (0, 1)) + # In BWD mode we won't use these, and we should set them to None so we don't build TMA atoms for them + mO_row = mO_col = mS_row = mS_col = None + elif cutlass.const_expr(self.cfg.WITH_GEMM_SWIZZLED_SCALES): + mS_row, mS_col = derive_swizzled_scale_layout( + M, N, self.cfg.ROWWISE, self.cfg.COLWISE, mS_row, mS_col + ) + + atoms = self._make_tma_atoms( + mYGrad, + mAct, + mGate, + mO_row, + mO_col, + mO_row_act_grad, + mO_row_gate_grad, + mO_col_act_grad, + mO_col_gate_grad, + ) + # The TMA atoms are spread from a list, which pylint cannot count + self.kernel( # pylint: disable=no-value-for-parameter + M, + N, + mS_row, + mS_col, + mS_row_act_grad, + mS_row_gate_grad, + mS_col_act_grad, + mS_col_gate_grad, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + *atoms, + ).launch( + grid=[ + cute.ceil_div(N, self._TILE_COLS), + cute.ceil_div(M, self._TILE_ROWS * self._NUM_TILES), + ], + block=[self._THREADS_PER_CTA], + min_blocks_per_mp=self._MIN_BLOCKS_PER_MP, + stream=stream, + ) + + @cute.jit + def _make_tma_atoms( + self, + mYGrad: Optional[cute.Tensor], + mAct: cute.Tensor, + mGate: cute.Tensor, + mO_row: Optional[cute.Tensor], # FWD only + mO_col: Optional[cute.Tensor], # FWD only + mO_row_act_grad: Optional[cute.Tensor], # BWD only + mO_row_gate_grad: Optional[cute.Tensor], # BWD only + mO_col_act_grad: Optional[cute.Tensor], # BWD only + mO_col_gate_grad: Optional[cute.Tensor], # BWD only + ): + """Build TMA atoms for input / quantized output; scales will not be loaded / stored by TMA.""" + smem_tile_layout = cute.make_ordered_layout( + (self._TILE_ROWS, self._TILE_COLS), order=(1, 0) + ) + tiler = (self._TILE_ROWS, self._TILE_COLS) + op_load = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + op_store = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + atoms = [] + for op, mT in ( + (op_load, mYGrad), + (op_load, mAct), + (op_load, mGate), + (op_store, mO_row), + (op_store, mO_col), + (op_store, mO_row_act_grad), + (op_store, mO_row_gate_grad), + (op_store, mO_col_act_grad), + (op_store, mO_col_gate_grad), + ): + if cutlass.const_expr(mT is None): + atoms.extend([None, None]) + else: + atoms.extend( + cute.nvgpu.cpasync.make_tiled_tma_atom( + op, mT, smem_tile_layout, tiler, num_multicast=1 + ) + ) + return atoms + + @cute.kernel + def kernel( + self, + M: Int32, + N: Int32, + mS_row: Optional[cute.Tensor], # FWD only + mS_col: Optional[cute.Tensor], # FWD only + mS_row_act_grad: Optional[cute.Tensor], # BWD only + mS_row_gate_grad: Optional[cute.Tensor], # BWD only + mS_col_act_grad: Optional[cute.Tensor], # BWD only + mS_col_gate_grad: Optional[cute.Tensor], # BWD only + limit: Optional[Float32], + alpha: Optional[Float32], + glu_linear_offset: Optional[Float32], + beta1: Optional[Float32], + beta2: Optional[Float32], + # TMA for mYGrad, only used in BWD + tma_atom_y_grad: Optional[cute.CopyAtom], + tma_src_y_grad: Optional[cute.Tensor], + # TMA for mAct and mGate, used in both FWD and BWD + tma_atom_act: cute.CopyAtom, + tma_src_act: cute.Tensor, + tma_atom_gate: cute.CopyAtom, + tma_src_gate: cute.Tensor, + # TMA for mO_row/col, only used in FWD + tma_atom_row: Optional[cute.CopyAtom], + tma_dst_row: Optional[cute.Tensor], + tma_atom_col: Optional[cute.CopyAtom], + tma_dst_col: Optional[cute.Tensor], + # TMA for mO_row/col_act_grad/gate_grad, only used in BWD + tma_atom_row_act_grad: Optional[cute.CopyAtom], + tma_dst_row_act_grad: Optional[cute.Tensor], + tma_atom_row_gate_grad: Optional[cute.CopyAtom], + tma_dst_row_gate_grad: Optional[cute.Tensor], + tma_atom_col_act_grad: Optional[cute.CopyAtom], + tma_dst_col_act_grad: Optional[cute.Tensor], + tma_atom_col_gate_grad: Optional[cute.CopyAtom], + tma_dst_col_gate_grad: Optional[cute.Tensor], + ): + """ + FWD mode: + mAct: activation half input + mGate: gate half input + mO_row/col: quantized output for row / column directions + mS_row/col: output scales for row / column directions + + BWD mode: + mYGrad: y_grad, the incoming gradient + mAct: activation half input from FWD pass + mGate: gate half input from FWD pass + mO_row/col_act_grad: quantized activation gradient for row / column directions + mO_row/col_gate_grad: quantized gate gradient for row / column directions + mS_row/col_act_grad: output scales for activation gradient row / column directions + mS_row/col_gate_grad: output scales for gate gradient row / column directions + mO_row/col: quantized gradient (both activation and gate) for row / column directions + mS_row/col: output scales for gradient (both activation and gate) row / column directions + + mAct: both FWD and BWD + mGate: both FWD and BWD + mYGrad: BWD only + mO_row/col: both FWD (N) and BWD (2N) + mS_row/col: both FWD (N) and BWD (2N) + mO_row/col_act_grad: BWD only + mO_row/col_gate_grad: BWD only + mS_row/col_act_grad: BWD only + mS_row/col_gate_grad: BWD only + """ + + cfg = self.cfg + DTYPE = cfg.DTYPE + TILE_ROWS = self._TILE_ROWS + TILE_COLS = self._TILE_COLS + NUM_STAGES = self._NUM_STAGES + NUM_TILES = self._NUM_TILES + + # The scales of one (TILE_ROWS, TILE_COLS) tile + row_scale_tile = (TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + col_scale_tile = (TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + + # Shared memory: one ((TILE_ROWS, TILE_COLS), NUM_STAGES) ring buffer per operand + smem = cutlass.memory.SmemAllocator() + + @cute.struct + class BarrierStorage: # pylint: disable=missing-class-docstring + mbar_storage: cute.struct.MemRange[cute.Int64, 2 * NUM_STAGES] + + storage = smem.allocate(BarrierStorage) + stages_layout = cute.make_layout( + ((TILE_ROWS, TILE_COLS), NUM_STAGES), + stride=((TILE_COLS, 1), TILE_ROWS * TILE_COLS), + ) + + # Allocate SMEM buffers for inputs + sAct = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sGate = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + sYGrad = None + if cutlass.const_expr(cfg.IS_BWD): + sYGrad = smem.allocate_tensor(DTYPE, stages_layout, byte_alignment=128) + + # Allocate SMEM buffers for outputs + sO_row = sO_col = None + sO_row_act_grad = sO_row_gate_grad = sO_col_act_grad = sO_col_gate_grad = None + if cutlass.const_expr(cfg.ROWWISE): + if cutlass.const_expr(cfg.IS_BWD): + sO_row_act_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + sO_row_gate_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + else: + sO_row = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + if cutlass.const_expr(cfg.COLWISE): + if cutlass.const_expr(cfg.IS_BWD): + sO_col_act_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + sO_col_gate_grad = smem.allocate_tensor( + cfg.FP8_DTYPE, stages_layout, byte_alignment=128 + ) + else: + sO_col = smem.allocate_tensor(cfg.FP8_DTYPE, stages_layout, byte_alignment=128) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, _ = cute.arch.block_idx() + + # Issue TMA descriptor prefetch using the TMA warp (warp 0) + if warp_idx == 0: + for load_atom in (tma_atom_act, tma_atom_gate, tma_atom_y_grad): + if cutlass.const_expr(load_atom is not None): + cute.nvgpu.cpasync.prefetch_descriptor(load_atom) + + # TMA warp (warp 0)'s elected thread produces (issues TMA), every warp consumes + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self._NUM_WARPS) + # Input has 2 tiles (act + gate) in FWD, 3 tiles (y_grad + act + gate) in BWD + num_inputs = 3 if cfg.IS_BWD else 2 + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_storage.data_ptr(), + num_stages=NUM_STAGES, + producer_group=producer_group, + consumer_group=consumer_group, + # Every input tile of a stage arrives on the same barrier + tx_count=num_inputs * TILE_ROWS * TILE_COLS * DTYPE.width // 8, + cta_layout_vmnk=None, + ) + prod_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, NUM_STAGES) + cons_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, NUM_STAGES) + + cta_coord = (bidy, bidx) + # Shape of each individual tile. A CTA processes one tile at a time + tiler = (TILE_ROWS, TILE_COLS) + # The block processed by a CTA, consisting of NUM_TILES vertically stacked tiles + cta_block = (NUM_TILES * TILE_ROWS, TILE_COLS) + + # Prepare parameters to pass to the SMEM tile quantize function per output (FWD: y; BWD: + # act_grad, gate_grad), where input and quantized output are SMEM tiles and scales are GMEM + # tiles + # Prepare parameters to pass to the SMEM tile quantize function for all outputs + # (FWD: y; BWD: act_grad, gate_grad) + # Note: input, quantized output are SMEM tiles, and scales are GMEM tiles + cta_block_row_scales = (NUM_TILES * TILE_ROWS, TILE_COLS // MXFP8_BLOCK_SCALING_SIZE) + cta_block_col_scales = (NUM_TILES * TILE_ROWS // MXFP8_BLOCK_SCALING_SIZE, TILE_COLS) + if cutlass.const_expr(cfg.IS_BWD): + if cutlass.const_expr(cfg.ROWWISE): + mS_row_act_grad = cute.local_tile(mS_row_act_grad, cta_block_row_scales, cta_coord) + mS_row_gate_grad = cute.local_tile( + mS_row_gate_grad, cta_block_row_scales, cta_coord + ) + if cutlass.const_expr(cfg.COLWISE): + mS_col_act_grad = cute.local_tile(mS_col_act_grad, cta_block_col_scales, cta_coord) + mS_col_gate_grad = cute.local_tile( + mS_col_gate_grad, cta_block_col_scales, cta_coord + ) + # In BWD mode we output gradient for both the act and gate halves, which are stored in sAct and sGate respectively + tile_quantizer_params = [ + ( + sAct, + sO_row_act_grad, + mS_row_act_grad, + sO_col_act_grad, + mS_col_act_grad, + ), # We will write unquantized act_grad to sAct in-place + ( + sGate, + sO_row_gate_grad, + mS_row_gate_grad, + sO_col_gate_grad, + mS_col_gate_grad, + ), # We will write unquantized gate_grad to sGate in-place + ] + else: + if cutlass.const_expr(cfg.ROWWISE): + mS_row = cute.local_tile(mS_row, cta_block_row_scales, cta_coord) + if cutlass.const_expr(cfg.COLWISE): + mS_col = cute.local_tile(mS_col, cta_block_col_scales, cta_coord) + # In FWD mode we only output the result after the gated op, which is stored in sAct + tile_quantizer_params = [ + (sAct, sO_row, mS_row, sO_col, mS_col) + ] # We will write unquantized output to sAct in-place + + # Prepare input TMA (atom, smem partition, gmem partition) triples: act and gate, and BWD's y_grad + TMA_load_atom_and_partitions = [] + for atom_in, src_in, sX_in in ( + (tma_atom_act, tma_src_act, sAct), + (tma_atom_gate, tma_src_gate, sGate), + (tma_atom_y_grad, tma_src_y_grad, sYGrad), + ): + if cutlass.const_expr(atom_in is not None): + tXsX_in, tXgX_in = cute.nvgpu.cpasync.tma_partition( + atom_in, + 0, + cute.make_layout(1), + sX_in, + cute.zipped_divide(cute.local_tile(src_in, cta_block, cta_coord), tiler), + ) + TMA_load_atom_and_partitions.append((atom_in, tXsX_in, tXgX_in)) + + # Prepare output TMA (atom, smem partition, gmem partition) triples of the outputs this config has + TMA_store_atom_and_partitions = [] + for atom_out, dst_out, sO_out in ( + (tma_atom_row, tma_dst_row, sO_row), # FWD only + (tma_atom_col, tma_dst_col, sO_col), # FWD only + (tma_atom_row_act_grad, tma_dst_row_act_grad, sO_row_act_grad), # BWD only + (tma_atom_row_gate_grad, tma_dst_row_gate_grad, sO_row_gate_grad), # BWD only + (tma_atom_col_act_grad, tma_dst_col_act_grad, sO_col_act_grad), # BWD only + (tma_atom_col_gate_grad, tma_dst_col_gate_grad, sO_col_gate_grad), # BWD only + ): + if cutlass.const_expr(atom_out is not None): + tXsO_out, tXgO_out = cute.nvgpu.cpasync.tma_partition( + atom_out, + 0, + cute.make_layout(1), + sO_out, + cute.zipped_divide(cute.local_tile(dst_out, cta_block, cta_coord), tiler), + ) + TMA_store_atom_and_partitions.append((atom_out, tXsO_out, tXgO_out)) + + # Barrier init must be visible to every thread + cute.arch.sync_threads() + + # Prologue: fill the pipeline + if warp_idx == 0: + for s in cutlass.range_constexpr(NUM_STAGES): + mainloop_pipeline.producer_acquire(prod_state) + for load_atom, load_sX, load_gX in TMA_load_atom_and_partitions: + cute.copy( + load_atom, + load_gX[(None, (s, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + for tile_idx in cutlass.range(NUM_TILES, unroll=1): + mainloop_pipeline.consumer_wait(cons_state) + # We only reuse previously used SMEM buffers when we have more tiles to process than pipeline stages, and only + # after we've already finished one round and wrap around to the start of the ring buffer + if cutlass.const_expr(NUM_TILES > NUM_STAGES): + if tile_idx >= NUM_STAGES: + # Only allow at most NUM_STAGES - 1 store groups reading SMEM on the fly, so the tile NUM_STAGES + # ago must finish before we reuse its SMEM buffer + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(NUM_STAGES - 1, read=True) + cute.arch.sync_threads() + # The tile should be ready to use now + + stage_idx = cons_state.index + tile_coord = (tile_idx, 0) # within this CTA's block + + # Select current stage's SMEM tile + sAct_tile = sAct[(None, stage_idx)] + sGate_tile = sGate[(None, stage_idx)] + sYGrad_tile = None + if cutlass.const_expr(cfg.IS_BWD): + sYGrad_tile = sYGrad[(None, stage_idx)] + + # Compute the gated op on this tile in fp32, cast back to the original dtype, + # and store the result in-place in the SMEM tile without quantization. + # FWD: write the output into sAct_tile + # BWD: write the act half's gradient to sAct_tile, and write the gate half's gradient to sGate_tile + gated_elementwise_smem_tile( + sAct_tile, # we also store the FWD output / BWD act half's gradient here + sGate_tile, # we also store the BWD gate half's gradient here + sYGrad_tile, + limit, + alpha, + glu_linear_offset, + beta1, + beta2, + ACTIVATION=cfg.GATED_ACTIVATION, + IS_BWD=cfg.IS_BWD, + ROWWISE=cfg.ROWWISE, + DTYPE=DTYPE, + TILE_X=TILE_COLS, + ) + + # Quantize every output of this tile into its SMEM buffers and scales + for sX_out, sO_row_out, mS_row_cta, sO_col_out, mS_col_cta in tile_quantizer_params: + # Slice the SMEM / GMEM tile for the current iteration + sO_row_tile = None + mS_row_stage = None + if cutlass.const_expr(cfg.ROWWISE): + sO_row_tile = sO_row_out[(None, stage_idx)] + mS_row_stage = cute.flatten( + cute.local_tile(mS_row_cta, row_scale_tile, tile_coord) + ) + sO_col_tile = None + mS_col_stage = None + if cutlass.const_expr(cfg.COLWISE): + sO_col_tile = sO_col_out[(None, stage_idx)] + mS_col_stage = cute.flatten( + cute.local_tile(mS_col_cta, col_scale_tile, tile_coord) + ) + # Use the general MXFP8 SMEM tile quantize function to quantize it + quantize_mxfp8_smem_tile_general_partial( + sX_out[(None, stage_idx)], # The SMEM tile to quantize + sO_row_tile, # SMEM tile for rowwise quantized output + mS_row_stage, # GMEM tile for rowwise quantized scales + sO_col_tile, # SMEM tile for columnwise quantized output + mS_col_stage, # GMEM tile for columnwise quantized scales + # The tile's first row and column in the whole tensor, for bounds checks + (bidy * NUM_TILES + tile_idx) * TILE_ROWS, + bidx * TILE_COLS, + M, + N, + ROWWISE=cfg.ROWWISE, + COLWISE=cfg.COLWISE, + MAX_NORM_RCP=cfg.MAX_NORM_RCP, + DTYPE=DTYPE, + FP8_DTYPE=cfg.FP8_DTYPE, + WITH_GEMM_SWIZZLED_SCALES=cfg.WITH_GEMM_SWIZZLED_SCALES, + TILE_X=TILE_COLS, + TILE_Y=TILE_ROWS, + ) + + # Make the SMEM writes visible to the TMA's async proxy before it reads them + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + mainloop_pipeline.consumer_release(cons_state) + + if warp_idx == 0: + # The tile quantize function has already written the quantized outputs to SMEM, + # so we can issue TMA stores from SMEM to GMEM + for store_atom, store_sO, store_gO in TMA_store_atom_and_partitions: + cute.copy( + store_atom, + store_sO[(None, stage_idx)], + store_gO[(None, (tile_idx, 0))], + ) + cute.arch.cp_async_bulk_commit_group() + + cons_state.advance() + + # Refill the stage just released with the tile NUM_STAGES ahead + if warp_idx == 0: + next_tile_idx = tile_idx + NUM_STAGES + if next_tile_idx < NUM_TILES: + mainloop_pipeline.producer_acquire(prod_state) + for load_atom, load_sX, load_gX in TMA_load_atom_and_partitions: + cute.copy( + load_atom, + load_gX[(None, (next_tile_idx, 0))], + load_sX[(None, prod_state.index)], + tma_bar_ptr=mainloop_pipeline.producer_get_barrier(prod_state), + ) + mainloop_pipeline.producer_commit(prod_state) + prod_state.advance() + + # Only wait until the in-flight TMA stores have read SMEM, so the CTA retires early + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/transformer_engine/common/CuTeDSL/utils.py b/transformer_engine/common/CuTeDSL/utils.py index 2200823d058..ed45324dcd4 100644 --- a/transformer_engine/common/CuTeDSL/utils.py +++ b/transformer_engine/common/CuTeDSL/utils.py @@ -9,7 +9,7 @@ import cutlass from cutlass import cute -from cutlass import Float32, Int64, Int32, Uint16, Uint32 +from cutlass import Boolean, Float32, Int64, Int32, Uint16, Uint32 # This must match `inline std::string to_string(const DType type)`` in transformer_engine/common/common.h exactly # which is how we encode types to strings in C++ @@ -83,6 +83,15 @@ def fabs_f32(val: Float32) -> Float32: return (val.bitcast(Int32) & Int32(0x7FFFFFFF)).bitcast(Float32) +@cute.jit +def select_f32(pred: Boolean, if_true: Float32, if_false: Float32) -> Float32: + """Return if_true if the runtime predicate pred holds, else if_false.""" + result = if_false + if pred: + result = if_true + return result + + @cute.jit def fma_f32(a: Float32, b: Float32, c: Float32) -> Float32: """Compute the fused multiply-add of three float32 values: a * b + c.""" diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index 0c1761264a1..02a0740f635 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -19,6 +19,10 @@ #include "../fp8/gated_fp8.cuh" #include "../mxfp8/gated_mxfp8.cuh" +#ifdef NVTE_WITH_CUTEDSL +#include "../mxfp8/gated_mxfp8_cutedsl.cuh" +#endif + namespace transformer_engine { namespace dispatch { @@ -89,8 +93,16 @@ void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_outp NVTE_CHECK(is_supported_by_CC_100(), "Gated FWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); Tensor dummy_grad_tensor; - mxfp8::quantize_gated(input, dummy_grad_tensor, - output, p, stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + input, dummy_grad_tensor, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(input, dummy_grad_tensor, + output, p, stream); + } break; } default: @@ -182,8 +194,16 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte NVTE_CHECK(is_supported_by_CC_100(), "Gated BWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); - mxfp8::quantize_gated(gated_input, grad, output, p, - stream); + bool quantized_with_cutedsl = false; +#ifdef NVTE_WITH_CUTEDSL + quantized_with_cutedsl = + cutedsl_backend::mxfp8_quantize_gated_cutedsl( + gated_input, grad, output, p, stream); +#endif + if (!quantized_with_cutedsl) { + mxfp8::quantize_gated(gated_input, grad, output, p, + stream); + } break; } default: diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh new file mode 100644 index 00000000000..adb3ced0f78 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8_cutedsl.cuh @@ -0,0 +1,204 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated_mxfp8_cutedsl.cuh + * \brief CuTeDSL backend of gated_mxfp8.cuh. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" +#include "../../util/cutedsl_utils.h" + +namespace transformer_engine { +namespace cutedsl_backend { + +using namespace tvm_ffi_bridge; + +struct MXFP8GatedQuantConfig { + static constexpr const char *kEntrypointName = "get_mxfp8_gated_quantization_function"; + + DType dtype; // The input format + DType fp8_dtype; // The fp8 output format + bool rowwise; // If quantize rowwisely + bool colwise; // If quantize columnwisely + bool swizzled; // If the scale output is used for cudnn's swizzled layout + bool is_bwd; // Backward (dgated) rather than forward + GatedActivation activation; + uint32_t sm_arch = static_cast(cuda::sm_arch()); + + // Bit layout: dtype [3:0], fp8_dtype [7:4], flags [11:8], activation [17:12] and + // SM architecture [26:18]. + uint32_t to_id() const { + static_assert(static_cast(DType::kNumTypes) <= 16, + "DType no longer fits in the 4 bits to_id() gives it."); + static_assert(static_cast(GatedActivation::kNumTypes) <= 64, + "GatedActivation no longer fits in the 6 bits to_id() gives it."); + NVTE_CHECK(sm_arch < 512, "SM architecture no longer fits in the 9 bits to_id() gives it."); + return static_cast(dtype) | (static_cast(fp8_dtype) << 4) | + (static_cast(rowwise) << 8) | (static_cast(colwise) << 9) | + (static_cast(swizzled) << 10) | (static_cast(is_bwd) << 11) | + (static_cast(activation) << 12) | (sm_arch << 18); + } + + std::optional get_kernel() const { + static TVMFFIConfigCache &cache = TVMFFIConfigCache::create(); + return cache.get_or_load(*this); + } + + std::string to_key() const { + std::string key; + key.reserve(80); + key.append("cutedsl_mxfp8_gated_sm") + .append(std::to_string(sm_arch)) + .append("_") + .append(to_string(dtype)) + .append("_") + .append(to_string(fp8_dtype)) + .append("_") + .append(rowwise ? "1" : "0") + .append("_") + .append(colwise ? "1" : "0") + .append("_") + .append(swizzled ? "1" : "0") + .append("_") + .append(is_bwd ? "1" : "0") + .append("_") + .append(gated_activation_to_str(activation)); + return key; + } + + bool retrieve_func_from_python(const std::string &fn_name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal(kEntrypointName); + if (!entrypoint.has_value()) { + return false; + } + tvm::ffi::Any result = + (*entrypoint)(tvm::ffi::String(fn_name), tvm::ffi::String(to_string(dtype)), + tvm::ffi::String(to_string(fp8_dtype)), rowwise, colwise, swizzled, is_bwd, + tvm::ffi::String(gated_activation_to_str(activation))); + return result.try_cast().value_or(false); + } +}; + +// Signature mirrors mxfp8::quantize_gated. Returns false to fall back to the CUDA kernel. +template +bool mxfp8_quantize_gated_cutedsl(const Tensor &gated_input, const Tensor &grad, Tensor *output, + const ParamOP &p, cudaStream_t stream) { + if (!TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { + maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); + return false; + } + constexpr GatedActivation activation = gated_activation_func_to_enum(); + if constexpr (activation == GatedActivation::kUnsupported) { + maybe_warn_cutedsl_not_chosen("the gated activation is not supported."); + return false; + } else { + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + const MXFP8GatedQuantConfig config{/*dtype=*/gated_input.dtype(), + /*fp8_dtype=*/output->dtype(), + /*rowwise=*/output->has_data(), + /*colwise=*/output->has_columnwise_data(), + /*swizzled=*/output->with_gemm_swizzled_scales, + /*is_bwd=*/IS_BWD, + /*activation=*/activation}; + // The kernel offsets the gate half's swizzled scales by a fixed amount, which only matches + // the swizzled layout when the half spans whole 128-column scale tiles. + if (IS_BWD && config.swizzled && cols % 128 != 0) { + maybe_warn_cutedsl_not_chosen( + "backward gated activations with GEMM-swizzled scales need the last dimension of the " + "output halves to be divisible by 128."); + return false; + } + // TMA needs aligned global addresses + bool aligned = + is_aligned_ptr(gated_input.data.dptr, TMA_GMEM_ALIGNMENT) && + (!config.rowwise || is_aligned_ptr(output->data.dptr, TMA_GMEM_ALIGNMENT)) && + (!config.colwise || is_aligned_ptr(output->columnwise_data.dptr, TMA_GMEM_ALIGNMENT)); + if constexpr (IS_BWD) { + aligned = aligned && is_aligned_ptr(grad.data.dptr, TMA_GMEM_ALIGNMENT); + } + if (!aligned) { + maybe_warn_cutedsl_not_chosen("an input or output is not 16-byte aligned."); + return false; + } + + std::optional fn = config.get_kernel(); + if (!fn.has_value()) { + return false; + } + checkCuDriverContext(stream); + + // Sanity checks, mirroring mxfp8::quantize_gated + if (config.rowwise) { + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + if (config.colwise) { + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + } + + // Zero out swizzled scales if padding is needed. + if (config.swizzled && (cols % 128 != 0 || rows % 128 != 0)) { + if (config.rowwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, + output->scale_inv.buffer_size_bytes(), stream)); + } + if (config.colwise) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->columnwise_scale_inv.dptr, 0, + output->columnwise_scale_inv.buffer_size_bytes(), stream)); + } + } + + // The activation's own parameters are runtime arguments; the kernel is compiled with None in + // the slots of the others, so they are passed as None (nullptr). + tvm::ffi::Any limit = nullptr, alpha = nullptr, glu_linear_offset = nullptr, beta1 = nullptr, + beta2 = nullptr; + if constexpr (std::is_same_v) { + limit = p.limit; + alpha = p.alpha; + glu_linear_offset = p.glu_linear_offset; + } else if constexpr (std::is_same_v) { + beta1 = p.beta1; + beta2 = p.beta2; + } + + const int32_t device_index = transformer_engine::cuda::current_device(); + DLTensorWrapper mGatedInput(gated_input.data, true, device_index); + DLTensorWrapper mGrad, mO_row, mS_row, mO_col, mS_col; + if constexpr (IS_BWD) { + mGrad = DLTensorWrapper(grad.data, true, device_index); + } + if (config.rowwise) { + mO_row = DLTensorWrapper(output->data, true, device_index); + mS_row = DLTensorWrapper(output->scale_inv, true, device_index); + } + if (config.colwise) { + mO_col = DLTensorWrapper(output->columnwise_data, true, device_index); + mS_col = DLTensorWrapper(output->columnwise_scale_inv, true, device_index); + } + (*fn)(&mGatedInput, &mGrad, &mO_row, &mS_row, &mO_col, &mS_col, limit, alpha, glu_linear_offset, + beta1, beta2, static_cast(stream)); + return true; + } +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_GATED_MXFP8_CUTEDSL_CUH_ diff --git a/transformer_engine/common/util/cutedsl_utils.h b/transformer_engine/common/util/cutedsl_utils.h index 4fd4ee56f69..4af0d78e906 100644 --- a/transformer_engine/common/util/cutedsl_utils.h +++ b/transformer_engine/common/util/cutedsl_utils.h @@ -90,6 +90,76 @@ NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM(dsrelu, kDSReLU) #undef NVTE_SPECIALIZE_CUTEDSL_ACTIVATION_FUNC_TO_ENUM +// Encodes the gated activation selected by the gated dispatcher (gated_mxfp8.cuh), identified by +// its forward ActOP and, for backward, its DActOP. +enum class GatedActivation { + kGeGLU, + kQGeGLU, + kReGLU, + kSReGLU, + kSwiGLU, + kGLU, + kSiTUGLU, + kClampedSwiGLU, + kUnsupported, + kNumTypes +}; + +inline const char *gated_activation_to_str(GatedActivation act) { + switch (act) { + case GatedActivation::kGeGLU: + return "geglu"; + case GatedActivation::kQGeGLU: + return "qgeglu"; + case GatedActivation::kReGLU: + return "reglu"; + case GatedActivation::kSReGLU: + return "sreglu"; + case GatedActivation::kSwiGLU: + return "swiglu"; + case GatedActivation::kGLU: + return "glu"; + case GatedActivation::kSiTUGLU: + return "situglu"; + case GatedActivation::kClampedSwiGLU: + return "clamped_swiglu"; + case GatedActivation::kUnsupported: + case GatedActivation::kNumTypes: + return "unsupported"; + } + return "unsupported"; +} + +// DActOP is nullptr for the forward. +template +constexpr GatedActivation gated_activation_func_to_enum() { + return GatedActivation::kUnsupported; +} + +#define NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(PARAM, ACT, DACT, ACTIVATION) \ + template <> \ + constexpr GatedActivation gated_activation_func_to_enum, nullptr>() { \ + return GatedActivation::ACTIVATION; \ + } \ + template <> \ + constexpr GatedActivation \ + gated_activation_func_to_enum, DACT>() { \ + return GatedActivation::ACTIVATION; \ + } + +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, gelu, dgelu, kGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, qgelu, dqgelu, kQGeGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, relu, drelu, kReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, srelu, dsrelu, kSReGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, silu, dsilu, kSwiGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(Empty, sigmoid, dsigmoid, kGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(SiTUGLUParam, situ_gate, dsitu_gate, kSiTUGLU) +NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM(ClampedSwiGLUParam, clamped_silu, + clamped_dsilu, kClampedSwiGLU) + +#undef NVTE_SPECIALIZE_CUTEDSL_GATED_ACTIVATION_FUNC_TO_ENUM + } // namespace cutedsl_backend } // namespace transformer_engine