From 7e08b94fd079f07fa7e9a922028b306938a57134 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Thu, 13 Aug 2026 00:44:49 +0000 Subject: [PATCH] Add a fused gated-activation (SwiGLU) + MXFP8 quantization kernel for Blackwell One CuTe DSL kernel fuses the gated activation (forward and backward) with the RCEIL MXFP8 cast: gate/up are read once and the bf16 activation never round-trips through global memory. Rowwise (1x32), colwise (32x1), or both scale modes come out of a single pass, in the same blocked tcgen05 layouts as the standalone quantizers, and bit-compatible with their special-value semantics (#4725 contract: NaN/Inf amax invalidates the block with scale byte 255 and all-NaN data; byte-0 scales descale by 2^127). Ops: torchao::gated_act_mxfp8_{forward,backward} custom ops with fake impls (torch.compile fullgraph works), public wrappers, 80-case numerics suite (bitwise forward, one-code-bounded backward), and an A/B benchmark vs the unfused Triton-activation + standalone-quantizer path. Co-Authored-By: Claude Fable 5 --- .../mxfp8/bench_cutedsl_gated_act_mxfp8.py | 350 ++++ .../test_cutedsl_gated_act_mxfp8.py | 664 +++++++ .../kernels/mxfp8/cutedsl_gated_act_mxfp8.py | 1706 +++++++++++++++++ 3 files changed, 2720 insertions(+) create mode 100644 benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py create mode 100644 test/prototype/moe_training/test_cutedsl_gated_act_mxfp8.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/cutedsl_gated_act_mxfp8.py diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py b/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py new file mode 100644 index 0000000000..2b7787b1e9 --- /dev/null +++ b/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py @@ -0,0 +1,350 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +# this benchmarking script is a modified version of the original script from: https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/utils/benchmark.py +# +# The baseline is a torch.compile-fused SwiGLU (one Triton kernel) followed by +# the standalone MXFP8 quantizers: +# python benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py [--compile] + +import argparse +import itertools +import os +from dataclasses import dataclass +from typing import List + +import torch +import torch.nn.functional as F +from tabulate import tabulate +from tqdm import tqdm + +from benchmarks.utils import benchmark_cuda_function_in_microseconds +from torchao.prototype.moe_training.kernels.mxfp8 import ( + mxfp8_quantize_2d_1x32_cutedsl, + mxfp8_quantize_2d_32x1_cutedsl, +) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_gated_act_mxfp8 import ( + gated_act_mxfp8_cutedsl_backward, + gated_act_mxfp8_cutedsl_forward, +) + +device = torch.device("cuda") +VALIDATE = os.environ.get("MXFP8_BENCH_VALIDATE", "0") == "1" + +SCALING_MODE = "rceil" + +# Backward E4M3 only; see eager_reference. Keep in sync with +# _MAX_DIFFERING_FRACTION in test/prototype/moe_training/ +# test_cutedsl_gated_act_mxfp8.py. +MAX_DIFFERING_FRACTION = 1e-5 + + +def _swiglu_fwd(gate, up): + return (F.silu(gate.float()) * up.float()).bfloat16() + + +def _swiglu_bwd(grad_h, gate, up): + # Mirror the kernel's evaluation order (silu path); see eager_reference. + gate, up, grad_h = gate.float(), up.float(), grad_h.float() + sig = torch.sigmoid(gate) + act = gate * sig + dact = act * (1.0 - sig) + sig + return torch.cat( + [ + ((dact * grad_h) * up).bfloat16(), + (act * grad_h).bfloat16(), + ], + dim=1, + ) + + +# One fused Triton kernel each, so the timing baseline's activation never +# round-trips intermediates through DRAM. dynamic=False keeps every shape on +# a static specialization (the auto-dynamic recompile is ~3x slower at large +# shapes); run_experiment resets dynamo between configs so specializations +# never accumulate toward the recompile limit. +_swiglu_fwd_c = torch.compile(_swiglu_fwd, fullgraph=True, dynamic=False) +_swiglu_bwd_c = torch.compile(_swiglu_bwd, fullgraph=True, dynamic=False) + + +@dataclass(frozen=True) +class ExperimentConfig: + input_shape: tuple[int, int] + direction: str + scales: str + + +@dataclass(frozen=True) +class ExperimentResult: + # time + baseline_us: float + fused_us: float + # mem bw + baseline_gbps: float + fused_gbps: float + + +@dataclass(frozen=True) +class Experiment: + config: ExperimentConfig + result: ExperimentResult + + +def get_configs(args: argparse.Namespace) -> List[ExperimentConfig]: + # (M, K): token counts x gate/up widths. 2048 = DSv3 expert FFN + # intermediate; 7168/8192 are DSv3/Llama3-70B model dims used as + # representative large widths. (128, 128) is the minimum legal size + # (launch-bound); (131072, 8192) backward lands just under the kernel's + # INT32 addressing bound (2*K*M = 2^31 exactly). + input_shapes = [ + (128, 128), + (4096, 2048), + (4096, 7168), + (16384, 7168), + (131072, 8192), + ] + if args.shape is not None: + input_shapes = [tuple(args.shape)] + directions = ( + ["forward", "backward"] if args.direction == "both" else [args.direction] + ) + scales_modes = ( + ["rowwise", "colwise", "both"] if args.scales == "all" else [args.scales] + ) + configs = [] + for shape, direction, scales in itertools.product( + input_shapes, directions, scales_modes + ): + configs.append( + ExperimentConfig( + input_shape=shape, + direction=direction, + scales=scales, + ) + ) + return configs + + +def _quantize_reference(reference, rowwise, colwise): + empty_qdata = reference.new_empty(0, dtype=torch.float8_e4m3fn) + empty_scales = reference.new_empty(0, dtype=torch.float8_e8m0fnu) + row = ( + mxfp8_quantize_2d_1x32_cutedsl(reference, scaling_mode=SCALING_MODE) + if rowwise + else (empty_qdata, empty_scales) + ) + col = ( + mxfp8_quantize_2d_32x1_cutedsl(reference, scaling_mode=SCALING_MODE) + if colwise + else (empty_qdata, empty_scales) + ) + return row[0], col[0], row[1], col[1] + + +def baseline(gated_input, grad_h, rowwise, colwise): + # torch.compile-fused SwiGLU, then the standalone MXFP8 quantizers: the + # activation is already fused, so the measured win is removing the + # bfloat16 round trip between it and the cast. + k = gated_input.shape[1] // 2 + gate, up = gated_input[:, :k], gated_input[:, k:] + if grad_h is None: + reference = _swiglu_fwd_c(gate, up) + else: + reference = _swiglu_bwd_c(grad_h, gate, up) + return _quantize_reference(reference, rowwise, colwise) + + +def eager_reference(gated_input, grad_h, rowwise, colwise): + # Ground truth for validation (not the timing baseline): the kernel's fast + # sigmoid and d_silu FMA contraction have no bit-exact eager equivalent, + # so exact agreement is only achievable in the forward direction. Keep in + # sync with _eager_reference in test/prototype/moe_training/ + # test_cutedsl_gated_act_mxfp8.py: both mirror the kernel's evaluation + # order. + k = gated_input.shape[1] // 2 + gate, up = gated_input[:, :k], gated_input[:, k:] + if grad_h is None: + reference = _swiglu_fwd(gate, up) + else: + reference = _swiglu_bwd(grad_h, gate, up) + return _quantize_reference(reference, rowwise, colwise) + + +def fused(gated_input, grad_h, rowwise, colwise): + if grad_h is None: + return gated_act_mxfp8_cutedsl_forward( + gated_input, rowwise=rowwise, colwise=colwise + ) + return gated_act_mxfp8_cutedsl_backward( + grad_h, gated_input, rowwise=rowwise, colwise=colwise + ) + + +def _e4m3_ordinal(u): + # Map sign-magnitude E4M3 bytes onto a signed number line so adjacent + # codes differ by 1 across the +/-0 boundary (raw byte distance jumps to + # 128 there). + s = u.to(torch.int16) + return torch.where(s >= 0x80, 0x80 - s, s) + + +def check(actual, expected, msg, exact): + # Scales and forward data are bitwise exact; backward data within one code. + assert actual.shape == expected.shape, f"{msg}: {actual.shape} vs {expected.shape}" + assert actual.stride() == expected.stride(), f"{msg}: stride mismatch" + a, e = actual.view(torch.uint8), expected.view(torch.uint8) + if exact or actual.dtype == torch.float8_e8m0fnu: + assert bool((a == e).all()), f"{msg}: not bitwise identical" + return + # A disabled direction is zero-sized: torch.max() has no empty-reduction identity. + if actual.numel() == 0: + return + gap = (_e4m3_ordinal(a) - _e4m3_ordinal(e)).abs() + assert int(gap.max()) <= 1, f"{msg}: max E4M3 code gap > 1" + count = int((gap != 0).sum()) + # Count floor mirrors the test suite's: at small shapes the fractional + # bound alone allows less than one differing code. + limit = max(8, int(MAX_DIFFERING_FRACTION * a.numel())) + assert count <= limit, f"{msg}: {count} codes differ, limit {limit}" + + +def validate_outputs(actual, gated_input, grad_h, rowwise, colwise): + M, two_k = gated_input.shape + direction = "forward" if grad_h is None else "backward" + expected = eager_reference(gated_input, grad_h, rowwise, colwise) + for i, (a, e) in enumerate(zip(actual, expected)): + check( + a, e, f"M={M} K={two_k // 2} {direction} output {i}", exact=grad_h is None + ) + + +def run_experiment( + config: ExperimentConfig, args: argparse.Namespace +) -> ExperimentResult: + M, K = config.input_shape + is_backward = config.direction == "backward" + rowwise = config.scales in ("rowwise", "both") + colwise = config.scales in ("colwise", "both") + + gated_input = torch.randn(M, 2 * K, dtype=torch.bfloat16, device=device) + grad_h = ( + torch.randn(M, K, dtype=torch.bfloat16, device=device) if is_backward else None + ) + bench_args = (gated_input, grad_h, rowwise, colwise) + if args.compile: + baseline_fn = torch.compile(baseline, fullgraph=True) + fused_fn = torch.compile(fused, fullgraph=True) + else: + baseline_fn, fused_fn = baseline, fused + + try: + outputs = fused_fn(*bench_args) + if VALIDATE: + validate_outputs(outputs, *bench_args) + baseline_time_us = benchmark_cuda_function_in_microseconds( + baseline_fn, *bench_args + ) + fused_time_us = benchmark_cuda_function_in_microseconds(fused_fn, *bench_args) + finally: + torch._dynamo.reset() + + # Memory bandwidth calculations, using the logical traffic of the fused op; + # the baseline additionally round-trips the bf16 activation through DRAM. + bytes_per_input_el = torch.finfo(torch.bfloat16).bits / 8 + bytes_per_output_el = torch.finfo(torch.float8_e4m3fn).bits / 8 + bytes_per_scale_el = torch.finfo(torch.float8_e8m0fnu).bits / 8 + + read_bytes = gated_input.numel() * bytes_per_input_el + if grad_h is not None: + read_bytes += grad_h.numel() * bytes_per_input_el + output_rowwise, output_colwise, scales_rowwise, scales_colwise = outputs + write_bytes = ( + output_rowwise.numel() + output_colwise.numel() + ) * bytes_per_output_el + ( + scales_rowwise.numel() + scales_colwise.numel() + ) * bytes_per_scale_el + + baseline_gbps = ((read_bytes + write_bytes) / 1e9) / (baseline_time_us / 1e6) + fused_gbps = ((read_bytes + write_bytes) / 1e9) / (fused_time_us / 1e6) + + return ExperimentResult( + baseline_us=baseline_time_us, + fused_us=fused_time_us, + baseline_gbps=baseline_gbps, + fused_gbps=fused_gbps, + ) + + +def print_results(experiments: List[Experiment]): + headers = [ + "input_shape", + "direction", + "scales", + "baseline_us", + "fused_us", + "speedup", + "baseline_gbps", + "fused_gbps", + ] + rows = [] + for experiment in experiments: + speedup = experiment.result.baseline_us / experiment.result.fused_us + rows.append( + [ + str(experiment.config.input_shape), + experiment.config.direction, + experiment.config.scales, + f"{experiment.result.baseline_us:.2f}", + f"{experiment.result.fused_us:.2f}", + f"{speedup:.2f}x", + f"{experiment.result.baseline_gbps:.1f}", + f"{experiment.result.fused_gbps:.1f}", + ] + ) + print(tabulate(rows, headers=headers)) + + +def main(args: argparse.Namespace): + torch.random.manual_seed(123) + configs = get_configs(args) + results = [] + for config in tqdm(configs): + result = run_experiment(config, args) + results.append(Experiment(config=config, result=result)) + torch.cuda.empty_cache() + + # Use Tabulate to print results + print(f"\nmode: {'compile' if args.compile else 'eager'}") + print_results(results) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--compile", + action="store_true", + help="benchmark torch.compile(fullgraph=True) instead of eager", + ) + parser.add_argument( + "--shape", + nargs=2, + type=int, + default=None, + metavar=("M", "K"), + help="run a single (M, K) shape instead of the sweep", + ) + parser.add_argument( + "--direction", + choices=("forward", "backward", "both"), + default="both", + ) + parser.add_argument( + "--scales", + choices=("rowwise", "colwise", "both", "all"), + default="all", + help="'both' is the single both-scales mode; 'all' sweeps all three", + ) + args = parser.parse_args() + main(args) diff --git a/test/prototype/moe_training/test_cutedsl_gated_act_mxfp8.py b/test/prototype/moe_training/test_cutedsl_gated_act_mxfp8.py new file mode 100644 index 0000000000..06023a1406 --- /dev/null +++ b/test/prototype/moe_training/test_cutedsl_gated_act_mxfp8.py @@ -0,0 +1,664 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch +import torch.nn.functional as F + + +def _is_sm_10x() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +if not _is_sm_10x(): + pytest.skip("MXFP8 requires CUDA SM 10.x", allow_module_level=True) + +from torchao.prototype.moe_training.kernels.mxfp8 import ( + mxfp8_quantize_2d_1x32_cutedsl, + mxfp8_quantize_2d_32x1_cutedsl, +) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( + _mxfp8_cutedsl_kernels_available, +) + +if not _mxfp8_cutedsl_kernels_available: + pytest.skip("MXFP8 cutedsl kernels not available", allow_module_level=True) + +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_gated_act_mxfp8 import ( + _gemm_swizzled_scale_idx, + _launch_gated_act_mxfp8, + _validate_inputs, + gated_act_mxfp8_cutedsl_backward, + gated_act_mxfp8_cutedsl_forward, +) +from torchao.prototype.mx_formats.utils import from_blocked + +# Irregular, multi-chunk shape for the edge-input pattern axis. +_PATTERN_SHAPE = (256, 384) + + +# Forward data and all forward scales must match the eager reference bitwise. +# Backward data may differ by one E4M3 code in a bounded fraction of elements: +# the kernel's fast sigmoid and d_silu FMA contraction have no bit-exact eager +# equivalent (measured rate 5.7e-7). A backward scale byte may additionally +# flip by one code where such a difference lands on a block amax at the RCEIL +# mantissa-carry boundary; the affected block's codes then shift by ~2x, so +# those elements are excluded from the data compare and budgeted by the scales +# check instead. Keep in sync with MAX_DIFFERING_FRACTION in +# benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py +# (whose opt-in validation holds scales fully bitwise). +_MAX_DIFFERING_FRACTION = 1e-5 + +# Absolute floor on the backward diff budgets: at small shapes the fractional +# bound rounds to zero, which would demand a bitwise-exact backward. Matches +# the bench's floor. +_COUNT_FLOOR = 8 + + +def _eager_reference(gated_input, grad_h, act_kind): + """Compute the bf16 tensor the kernel is expected to quantize. + + Keep in sync with eager_reference() in benchmarks/prototype/moe_training/ + mxfp8/bench_cutedsl_gated_act_mxfp8.py: both mirror the kernel's + evaluation order. + """ + assert act_kind == "silu", f"no eager reference wired up for {act_kind!r}" + K = gated_input.shape[1] // 2 + gate = gated_input[:, :K].float() + up = gated_input[:, K:].float() + if grad_h is None: + return (F.silu(gate) * up).bfloat16() + grad_h_f = grad_h.float() + # Mirror the kernel's evaluation order (silu path); the kernel contracts + # dact into a single FMA, which eager cannot reproduce bit for bit. + sigmoid_gate = torch.sigmoid(gate) + silu = gate * sigmoid_gate + dact = silu * (1.0 - sigmoid_gate) + sigmoid_gate + dgate = ((dact * grad_h_f) * up).bfloat16() + dup = (silu * grad_h_f).bfloat16() + return torch.cat([dgate, dup], dim=1) + + +def _assert_layout(actual, ref, msg): + assert actual.shape == ref.shape, f"{msg}: shape {actual.shape} vs {ref.shape}" + assert actual.stride() == ref.stride(), ( + f"{msg}: stride {actual.stride()} vs {ref.stride()}" + ) + assert actual.dtype == ref.dtype, f"{msg}: dtype {actual.dtype} vs {ref.dtype}" + + +def _assert_bitwise(actual, ref, msg): + # Raw-byte compare: NaN codes are expected content in the edge patterns. + _assert_layout(actual, ref, msg) + torch.testing.assert_close( + actual.contiguous().view(torch.uint8), + ref.contiguous().view(torch.uint8), + rtol=0, + atol=0, + msg=msg, + ) + + +def _e4m3_ordinal(codes): + # Map sign-magnitude E4M3 bytes onto a signed number line so adjacent + # codes differ by 1 across the +/-0 boundary (raw byte distance jumps to + # 128 there). Keep in sync with _e4m3_ordinal in the bench. + c = codes.int() + return torch.where(c >= 0x80, 0x80 - c, c) + + +def _assert_scales_match(actual, ref, msg, exact, logical_rows, logical_cols): + """Forward scales are bitwise. A backward scale byte may flip by one code + in a bounded count of blocks (a one-ulp amax difference at the RCEIL + mantissa-carry boundary). Returns the logical (rows, cols) mask of + differing blocks for the data compare to exclude, or None when exact.""" + if exact: + _assert_bitwise(actual, ref, msg) + return None + _assert_layout(actual, ref, msg) + a = from_blocked( + actual.contiguous().view(torch.uint8), logical_rows, logical_cols + ).int() + r = from_blocked( + ref.contiguous().view(torch.uint8), logical_rows, logical_cols + ).int() + gap = (a - r).abs() + max_gap = int(gap.max()) + assert max_gap <= 1, f"{msg}: max E8M0 code gap {max_gap} > 1" + mismatch = gap != 0 + count = int(mismatch.sum()) + limit = max(_COUNT_FLOOR, int(_MAX_DIFFERING_FRACTION * gap.numel())) + assert count <= limit, f"{msg}: {count} scale bytes differ, limit {limit}" + return mismatch + + +def _assert_qdata_matches(actual, ref, msg, exact, exclude_mask=None): + """Forward data is bitwise; backward data within one E4M3 code in a + bounded count of elements. exclude_mask marks elements of blocks whose + scale byte differs: their codes legitimately shift ~2x and are budgeted + by the scales check.""" + if exact: + assert exclude_mask is None + _assert_bitwise(actual, ref, msg) + return + _assert_layout(actual, ref, msg) + a = _e4m3_ordinal(actual.contiguous().view(torch.uint8)) + r = _e4m3_ordinal(ref.contiguous().view(torch.uint8)) + gap = (a - r).abs() + numel = gap.numel() + if exclude_mask is not None: + gap = gap[~exclude_mask] + max_gap = int(gap.max()) if gap.numel() else 0 + assert max_gap <= 1, f"{msg}: max E4M3 code gap {max_gap} > 1" + count = int((gap != 0).sum()) + limit = max(_COUNT_FLOOR, int(_MAX_DIFFERING_FRACTION * numel)) + assert count <= limit, f"{msg}: {count} codes differ, limit {limit}" + + +def _run_and_check(gated_input, grad_h, rowwise, colwise, act_kind, tag): + """Run the public op and check all four outputs against the eager reference.""" + M = gated_input.shape[0] + K = gated_input.shape[1] // 2 + if grad_h is not None: + outputs = gated_act_mxfp8_cutedsl_backward( + grad_h, gated_input, rowwise=rowwise, colwise=colwise + ) + else: + outputs = gated_act_mxfp8_cutedsl_forward( + gated_input, rowwise=rowwise, colwise=colwise + ) + + # Fixed four-output tuple regardless of which directions are enabled. + assert isinstance(outputs, tuple) and len(outputs) == 4, f"{tag}: arity" + output_rowwise, output_colwise, scales_rowwise, scales_colwise = outputs + + # Forward emits h (width K); backward emits [dGate | dUp] (width 2K). + expected_width = K if grad_h is None else 2 * K + reference = _eager_reference(gated_input, grad_h, act_kind) + assert reference.shape == (M, expected_width) + exact = grad_h is None + + if rowwise: + ref_q, ref_s = mxfp8_quantize_2d_1x32_cutedsl(reference, scaling_mode="rceil") + assert output_rowwise.shape == (M, expected_width), f"{tag}: rowwise width" + assert output_rowwise.stride() == (expected_width, 1), ( + f"{tag}: rowwise qdata must be row-major, got {output_rowwise.stride()}" + ) + mismatch = _assert_scales_match( + scales_rowwise, + ref_s, + f"{tag}: rowwise scales", + exact, + M, + expected_width // 32, + ) + exclude = None if mismatch is None else mismatch.repeat_interleave(32, dim=1) + _assert_qdata_matches( + output_rowwise, ref_q, f"{tag}: rowwise qdata", exact, exclude + ) + else: + assert output_rowwise.numel() == 0, f"{tag}: rowwise output should be empty" + assert scales_rowwise.numel() == 0, f"{tag}: rowwise scales should be empty" + assert output_rowwise.dtype == torch.float8_e4m3fn + assert scales_rowwise.dtype == torch.float8_e8m0fnu + assert output_rowwise.device == gated_input.device + + if colwise: + ref_q, ref_s = mxfp8_quantize_2d_32x1_cutedsl(reference, scaling_mode="rceil") + assert output_colwise.shape == (M, expected_width), f"{tag}: colwise width" + assert output_colwise.stride() == (1, M), ( + f"{tag}: colwise qdata must have stride (1, M), got {output_colwise.stride()}" + ) + assert scales_colwise.ndim == 1, f"{tag}: colwise scales should be flat" + # Colwise scales use transposed blocked coordinates: rows are output + # columns, columns are 32-row blocks. + mismatch = _assert_scales_match( + scales_colwise, + ref_s, + f"{tag}: colwise scales", + exact, + expected_width, + M // 32, + ) + exclude = None if mismatch is None else mismatch.repeat_interleave(32, dim=1).T + _assert_qdata_matches( + output_colwise, ref_q, f"{tag}: colwise qdata", exact, exclude + ) + else: + assert output_colwise.numel() == 0, f"{tag}: colwise output should be empty" + assert scales_colwise.numel() == 0, f"{tag}: colwise scales should be empty" + assert output_colwise.dtype == torch.float8_e4m3fn + assert scales_colwise.dtype == torch.float8_e8m0fnu + assert output_colwise.device == gated_input.device + + return outputs + + +def _make_gated_act_edge_input( + M: int, K: int, pattern: str +) -> tuple[torch.Tensor, torch.Tensor]: + """Build (gated_input, grad_h) for one input pattern. Pinned rows use + gate = 20 (sigmoid saturates to 1.0) with power-of-two `up` values so h is + bitwise stable; rowwise patterns sit in rows 0..8, colwise patterns in + columns 100..103 over rows 0..31.""" + if pattern == "normal": + torch.manual_seed(42) + gated = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + return gated, grad + + if pattern == "boundary": + # Block amaxes at and above the E4M3 max (448), plus large/tiny mixes. + torch.manual_seed(13) + gated = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + gate, up = gated[:, :K], gated[:, K:] + gate[:6, :] = 20.0 + up[0, :] = 22.375 # h = bf16(447.5) = 448: amax exactly at E4M3 max + up[1, :] = -22.375 # negative boundary + up[2, :] = 23.0 # h = 460: amax in (448, 512), RCEIL boundary region + up[3, ::2] = 16.0 # h = 320 mixed with ... + up[3, 1::2] = 2.0**-12 # ... tiny values in the same 1x32 blocks + up[4, 0] = 22.375 # block amax 448 with everything else ~448 * 2^-9, + up[4, 1:] = 2.0**-4 # landing at the E4M3 subnormal boundary + gate[:32, 100:104] = 20.0 + up[:32, 100] = 22.375 # colwise variants of the same boundaries + up[:32, 101] = -22.375 + up[:32, 102] = 23.0 + up[:32:2, 103] = 16.0 + up[1:32:2, 103] = 2.0**-12 + return gated, grad + + if pattern == "zeros": + # Zero rows/blocks: amax 0 takes the byte-0 scale and must still + # produce zero codes. + torch.manual_seed(5) + gated = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + gate, up = gated[:, :K], gated[:, K:] + gate[:2, :] = 20.0 + up[:2, :] = 0.5 + up[0, :64] = 0.0 # zero rowwise blocks amid a nonzero row + gated[1, :] = 0.0 # fully zero row + gate[:, 128:160] = 0.0 # silu(0) * up = 0: zero colwise stripes + gate[:32, 100] = 20.0 + up[:32, 100] = 0.0 # zero colwise block + grad[2, :] = 0.0 # zero gradient row (backward) + grad[:32, 7] = 0.0 # zero gradient column block + return gated, grad + + if pattern == "subnormal_tiny": + # Tiny amaxes: the byte-0 scale must pair with the 2^127 reciprocal so + # the blocks do not collapse to zero codes. + torch.manual_seed(7) + gated = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + gate, up = gated[:, :K], gated[:, K:] + gate[:4, :] = 20.0 + up[:4, :] = 0.5 + up[0, :64] = 2.0**-125 # tiny amax: scale byte clamps to 0 + up[1, :] = 2.0**-130 # bf16 subnormal inputs + up[2, ::2] = 2.0**-126 # bf16 min-normal mixed with zeros + up[2, 1::2] = 0.0 + gate[:32, 100:102] = 20.0 + up[:32, 100] = 2.0**-125 # tiny colwise block + up[:32, 101] = 2.0**-130 + grad[:32, 9] = 2.0**-120 # tiny gradients (backward) + return gated, grad + + if pattern == "nan_inf": + # NaN or Inf amax invalidates the block: scale byte 255 and every + # element quantizes to the E4M3 NaN code. + torch.manual_seed(11) + gated = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + gate, up = gated[:, :K], gated[:, K:] + gate[0, 0] = float("nan") # NaN input propagates through the activation + up[1, 5] = float("nan") + gate[2, :] = 88.0 # sigmoid saturates: silu(x) == x + up[2, :] = 3.0e38 # products overflow bf16 to Inf: Inf-amax blocks + gate[3:5, :] = 20.0 + up[3:5, :] = 0.5 + up[3, 0] = float("inf") # Inf element + up[3, 1] = -3.0e38 # h overflows f32 to -Inf: mixed-sign Inf block + gate[4, 0] = float("inf") + up[4, 0] = 0.0 # silu(inf) * 0 -> NaN element + gate[:32, 100] = float("inf") + up[:32, 100] = 0.0 # all-NaN colwise block + gate[:32, 101] = 88.0 + up[:32, 101] = 3.0e38 # Inf colwise block + grad[5, 7] = float("nan") # NaN gradient (backward) + return gated, grad + + if pattern == "mixed_extreme": + # Every special pattern in one tensor on a random background. + torch.manual_seed(17) + gated = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + grad = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + gate, up = gated[:, :K], gated[:, K:] + gate[:4, :] = 20.0 + up[:4, :] = 0.5 + up[0, :] = 22.375 # boundary row (amax 448) + up[1, ::2] = 16.0 # large/tiny mix + up[1, 1::2] = 2.0**-12 + up[2, :64] = 0.0 # zero blocks + up[3, :64] = 2.0**-125 # tiny-amax blocks + gate[4, 0] = float("nan") # NaN input + gate[5, :] = 88.0 + up[5, :] = 3.0e38 # Inf-amax row + gate[6, :] = 20.0 + up[6, :] = 0.5 + gate[6, 0] = float("inf") + up[6, 0] = 0.0 # silu(inf) * 0 -> NaN element + gate[:32, 100:103] = 20.0 + up[:32, 100] = 0.0 # zero colwise block + up[:32, 101] = 2.0**-125 # tiny colwise block + up[:32, 102] = 22.375 # boundary colwise block + gate[:32, 103] = float("inf") + up[:32, 103] = 0.0 # all-NaN colwise block + grad[7, 3] = float("nan") # NaN gradient (backward) + grad[8, :] = 0.0 # zero gradient row + return gated, grad + + raise AssertionError(f"unknown pattern: {pattern}") + + +@pytest.mark.parametrize( + "M,K", + ( + (128, 128), + (256, 256), + (640, 384), # one non-square, irregular shape + (512, 2048), + (1024, 7168), + ), +) +@pytest.mark.parametrize("is_backward", (False, True)) +@pytest.mark.parametrize( + "rowwise,colwise", ((True, False), (False, True), (True, True)) +) +# The public ops hard-wire silu today; new activation kinds slot into this +# axis once the ops expose a selector. +@pytest.mark.parametrize("act_kind", ("silu",)) +def test_gated_act_mxfp8_numerics(M, K, is_backward, rowwise, colwise, act_kind): + gated_input, grad = _make_gated_act_edge_input(M, K, "normal") + tag = ( + f"{act_kind} bwd={is_backward} rowwise={rowwise} colwise={colwise} M={M} K={K}" + ) + _run_and_check( + gated_input, grad if is_backward else None, rowwise, colwise, act_kind, tag + ) + + +@pytest.mark.parametrize( + "pattern", + ("normal", "boundary", "zeros", "subnormal_tiny", "nan_inf", "mixed_extreme"), +) +@pytest.mark.parametrize("is_backward", (False, True)) +@pytest.mark.parametrize( + "rowwise,colwise", ((True, False), (False, True), (True, True)) +) +@pytest.mark.parametrize("act_kind", ("silu",)) +def test_gated_act_mxfp8_edge_inputs(pattern, is_backward, rowwise, colwise, act_kind): + M, K = _PATTERN_SHAPE + gated_input, grad = _make_gated_act_edge_input(M, K, pattern) + tag = f"{act_kind} {pattern} bwd={is_backward} rowwise={rowwise} colwise={colwise}" + outputs = _run_and_check( + gated_input, + grad if is_backward else None, + rowwise, + colwise, + act_kind, + tag, + ) + + # Targeted probes at known coordinates (forward rowwise only: exact codes). + if is_backward or not rowwise: + return + q = outputs[0].view(torch.uint8) + scales_r = outputs[2].contiguous().view(torch.uint8).flatten() + ncb = (K // 32 + 3) // 4 # 128x4 scale-column blocks in the swizzled layout + if pattern in ("subnormal_tiny", "mixed_extreme"): + # Tiny-amax blocks must not collapse to zero codes (byte-0 scale + # descales by 2^127). + row = 0 if pattern == "subnormal_tiny" else 3 + assert bool(q[row, :64].ne(0).all()), ( + f"{tag}: tiny-amax block quantized to zero codes" + ) + assert int(scales_r[_gemm_swizzled_scale_idx(row, 0, ncb)]) == 0x00, ( + f"{tag}: tiny amax did not clamp to scale byte 0x00" + ) + if pattern in ("nan_inf", "mixed_extreme"): + # Invalidated (NaN-amax) blocks map to scale byte 0xFF and E4M3 NaN + # codes. + row = 4 if pattern == "nan_inf" else 6 + assert int(q[row, 0]) & 0x7F == 0x7F, ( + f"{tag}: NaN element did not map to the E4M3 NaN code" + ) + assert int(scales_r[_gemm_swizzled_scale_idx(row, 0, ncb)]) == 0xFF, ( + f"{tag}: NaN-amax block did not take scale byte 0xFF" + ) + + +@pytest.mark.parametrize("is_backward", (False, True)) +@pytest.mark.parametrize( + "rowwise,colwise", ((True, False), (False, True), (True, True)) +) +def test_gated_act_mxfp8_torch_compile(is_backward, rowwise, colwise): + """torch.compile(fullgraph=True) must trace the fake impls and return + bitwise the eager op's outputs (both paths run the identical kernel).""" + M, K = _PATTERN_SHAPE + gated_input, grad = _make_gated_act_edge_input(M, K, "normal") + grad = grad if is_backward else None + + def run_op(gated, grad_h): + if grad_h is None: + return gated_act_mxfp8_cutedsl_forward( + gated, rowwise=rowwise, colwise=colwise + ) + return gated_act_mxfp8_cutedsl_backward( + grad_h, gated, rowwise=rowwise, colwise=colwise + ) + + tag = f"compile bwd={is_backward} rowwise={rowwise} colwise={colwise}" + try: + eager_outputs = run_op(gated_input, grad) + compiled_outputs = torch.compile(run_op, fullgraph=True)(gated_input, grad) + for i, (c, e) in enumerate(zip(compiled_outputs, eager_outputs)): + _assert_bitwise(c, e, f"{tag}: output {i}") + finally: + torch._dynamo.reset() + + +@pytest.mark.parametrize("M,K", ((128, 128), (256, 384))) +@pytest.mark.parametrize("is_backward", (False, True)) +def test_gated_act_mxfp8_mode_consistency(M, K, is_backward): + """Both-scales mode must reproduce each single-mode run bitwise; the modes + take different kernel paths (cached-activation vs single-orientation).""" + gated_input, grad = _make_gated_act_edge_input(M, K, "normal") + + def run_op(rowwise, colwise): + if is_backward: + return gated_act_mxfp8_cutedsl_backward( + grad, gated_input, rowwise=rowwise, colwise=colwise + ) + return gated_act_mxfp8_cutedsl_forward( + gated_input, rowwise=rowwise, colwise=colwise + ) + + row_only = run_op(True, False) + col_only = run_op(False, True) + both = run_op(True, True) + tag = f"mode-consistency bwd={is_backward} M={M} K={K}" + _assert_bitwise(both[0], row_only[0], f"{tag}: rowwise qdata") + _assert_bitwise(both[2], row_only[2], f"{tag}: rowwise scales") + _assert_bitwise(both[1], col_only[1], f"{tag}: colwise qdata") + _assert_bitwise(both[3], col_only[3], f"{tag}: colwise scales") + + +def test_gated_act_mxfp8_invalid_inputs(): + ok = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + with pytest.raises(TypeError, match="bfloat16"): + gated_act_mxfp8_cutedsl_forward(ok.float()) + with pytest.raises(ValueError, match="multiples of 128"): + gated_act_mxfp8_cutedsl_forward( + torch.randn(130, 256, device="cuda", dtype=torch.bfloat16) + ) + with pytest.raises(ValueError, match="multiples of 128"): + gated_act_mxfp8_cutedsl_forward( + torch.randn(128, 260, device="cuda", dtype=torch.bfloat16) + ) + # Zero-size inputs satisfy every modulus but cannot form a launch grid + # or a TMA descriptor. + with pytest.raises(ValueError, match="nonzero"): + gated_act_mxfp8_cutedsl_forward( + torch.empty(0, 256, device="cuda", dtype=torch.bfloat16) + ) + with pytest.raises(ValueError, match="nonzero"): + gated_act_mxfp8_cutedsl_forward( + torch.empty(128, 0, device="cuda", dtype=torch.bfloat16) + ) + with pytest.raises(ValueError, match="even"): + gated_act_mxfp8_cutedsl_forward( + torch.randn(128, 257, device="cuda", dtype=torch.bfloat16) + ) + with pytest.raises(ValueError, match="contiguous"): + gated_act_mxfp8_cutedsl_forward( + torch.randn(256, 128, device="cuda", dtype=torch.bfloat16).t() + ) + with pytest.raises(ValueError, match="rowwise/colwise"): + gated_act_mxfp8_cutedsl_forward(ok, rowwise=False, colwise=False) + with pytest.raises(ValueError, match="grad_h"): + gated_act_mxfp8_cutedsl_backward( + torch.randn(128, 256, device="cuda", dtype=torch.bfloat16), ok + ) + + # 32-byte pointer alignment: contiguous storage-offset views are legal + # torch tensors but break the launcher's assumed_align contract. + base = torch.randn(128 * 256 + 16, device="cuda", dtype=torch.bfloat16) + misaligned = base[1 : 1 + 128 * 256].view(128, 256) + assert misaligned.is_contiguous() + with pytest.raises(ValueError, match="32-byte aligned"): + gated_act_mxfp8_cutedsl_forward(misaligned) + with pytest.raises(ValueError, match="32-byte aligned"): + gated_act_mxfp8_cutedsl_backward(base[1 : 1 + 128 * 128].view(128, 128), ok) + # A 32-byte-aligned storage offset must still run and match a fresh copy. + offset_ok = base[16 : 16 + 128 * 256].view(128, 256) + assert offset_ok.data_ptr() % 32 == 0 + got = gated_act_mxfp8_cutedsl_forward(offset_ok) + want = gated_act_mxfp8_cutedsl_forward(offset_ok.clone()) + _assert_bitwise(got[0], want[0], "aligned-offset view: rowwise qdata") + _assert_bitwise(got[2], want[2], "aligned-offset view: rowwise scales") + + # INT32 indexing bound, allocation-free: _validate_inputs only inspects + # metadata, so a FakeTensor exercises the check without 4 GiB of HBM. + from torch._subclasses.fake_tensor import FakeTensorMode + + with FakeTensorMode(): + big = torch.empty(8388736, 256, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="32-bit indexing limit"): + _validate_inputs(big) + + # The fakes must reject both-False too: otherwise compile/export would + # trace a call eager rejects (returning aliased zero-size outputs). + with FakeTensorMode(): + fake = torch.empty(128, 512, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="rowwise/colwise"): + torch.ops.torchao.gated_act_mxfp8_cutedsl_forward( + fake, rowwise=False, colwise=False + ) + + +def test_gated_act_mxfp8_invalid_geometry(): + """The private geometry override must reject values the kernel's grid and + bit-mask thread mapping cannot represent: the floor-division grid would + silently skip columns/rows instead of failing.""" + gated_input = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + outputs = ( + torch.empty(128, 128, device="cuda", dtype=torch.float8_e4m3fn), + torch.empty(0, device="cuda", dtype=torch.float8_e4m3fn), + torch.empty(0, device="cuda", dtype=torch.uint8), + torch.empty(0, device="cuda", dtype=torch.uint8), + ) + for bad in ((96, 64, True), (256, 64, True)): # non-pow2; K % CX != 0 + with pytest.raises(ValueError, match=f"CX={bad[0]}"): + _launch_gated_act_mxfp8( + gated_input, None, outputs, True, False, geometry=bad + ) + with pytest.raises(ValueError, match="CY=48"): + _launch_gated_act_mxfp8( + gated_input, None, outputs, True, False, geometry=(64, 48, True) + ) + # The staged path's output smem is double-buffered with no in-loop + # TMA-store drain, so stage counts past the double buffer must be + # rejected rather than silently corrupting reused buffers. + tall = torch.randn(384, 256, device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="pipeline stages"): + _launch_gated_act_mxfp8( + tall, None, outputs, True, False, geometry=(64, 96, False) + ) + # Row chunks ride CUDA grid dim y (cap 65535): a small-K shape inside the + # int32 element bound must still be rejected. The gate fires before + # compile or launch, so only the input allocation is paid. + if torch.cuda.get_device_properties(0).total_memory >= 8 * 2**30: + grid_y_input = torch.empty(65536 * 32, 256, device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="grid y-dimension"): + _launch_gated_act_mxfp8( + grid_y_input, None, outputs, True, False, geometry=(64, 32, True) + ) + del grid_y_input + # A valid non-default geometry must produce bit-identical results. + ref = gated_act_mxfp8_cutedsl_forward(gated_input, rowwise=True, colwise=False) + alt = tuple(torch.empty_like(t) for t in ref) + _launch_gated_act_mxfp8( + gated_input, None, alt, True, False, geometry=(64, 64, True) + ) + _assert_bitwise(alt[0], ref[0], "geometry override: rowwise qdata") + _assert_bitwise(alt[2], ref[2], "geometry override: rowwise scales") + + +def test_gated_act_mxfp8_wrappers_unavailable(monkeypatch): + """When the CuTeDSL runtime is unavailable the public wrappers must raise + the informative NotImplementedError, not a raw import error. The flag is + read at call time, so monkeypatching it simulates the unavailable case.""" + from torchao.prototype.moe_training.kernels.mxfp8 import quant as _quant + + monkeypatch.setattr(_quant, "_mxfp8_cutedsl_kernels_available", False) + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + with pytest.raises( + NotImplementedError, match="gated_act_mxfp8_cutedsl_forward requires" + ): + gated_act_mxfp8_cutedsl_forward(x) + with pytest.raises( + NotImplementedError, match="gated_act_mxfp8_cutedsl_backward requires" + ): + gated_act_mxfp8_cutedsl_backward( + torch.randn(128, 128, device="cuda", dtype=torch.bfloat16), x + ) + + +def test_gated_act_mxfp8_int32_boundary(): + """Largest legal shape: 2*K*M - K - 1 lands just under INT32_MAX; verify + boundary rows against the standalone quantizer without materializing a + full-size reference.""" + if torch.cuda.get_device_properties(0).total_memory < 32 * 2**30: + pytest.skip("needs >= 32 GiB of device memory") + M, K = 131072, 8192 + torch.manual_seed(3) + gated_input = torch.randn(M, 2 * K, device="cuda", dtype=torch.bfloat16) + output_rowwise = gated_act_mxfp8_cutedsl_forward( + gated_input, rowwise=True, colwise=False + )[0] + # Rowwise 1x32 blocks are row-local, so row slabs compare cleanly (the + # blocked scales' swizzle offsets are not slice-local; qdata only). + for rows in (slice(0, 128), slice(M - 128, M)): + reference = _eager_reference(gated_input[rows], None, "silu") + ref_q, _ = mxfp8_quantize_2d_1x32_cutedsl(reference, scaling_mode="rceil") + _assert_bitwise( + output_rowwise[rows].contiguous(), ref_q, f"int32-boundary rows {rows}" + ) + del gated_input, output_rowwise + torch.cuda.empty_cache() diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_gated_act_mxfp8.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_gated_act_mxfp8.py new file mode 100644 index 0000000000..5f441a3249 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_gated_act_mxfp8.py @@ -0,0 +1,1706 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +"""Fused gated activation (SwiGLU) + MXFP8 CuTe DSL kernel for Blackwell, +tuned and validated on GB200/SM100 (the inline PTX deliberately avoids the +sm_100a-only cvt so sm_103a is expected to work; see ``_mul_cvt_2x``). One +pass computes the activation and its RCEIL MXFP8 cast, so the activation +never round-trips through global memory: + + forward: h = silu(gate) * up + backward: dGate = grad_h * up * d_silu(gate), dUp = grad_h * silu(gate) + +``gated_input`` is bf16 [M, 2K] holding ``gate`` then ``up``; forward outputs +are K wide, backward outputs 2K wide (``[dGate | dUp]``). Rowwise (1x32) +scales, colwise (32x1) scales, or both come from that single read, in the +blocked tcgen05 layouts. Mode flags, chunk geometry, and the ``ACT_PAIR`` +activation policy are ``Constexpr`` — each (mode, geometry, device) +combination compiles once and is cached — while M and K are runtime +arguments, so a single specialization serves every shape. Requires M and K +multiples of 128 and +``2*K*M - K - 1 <= INT32_MAX`` (index arithmetic assumes 32-bit offsets). +""" + +import functools +from typing import Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float32, Int32 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm +from cutlass.cute import AddressSpace +from cutlass.cutlass_dsl import T, dsl_user_op + +from torchao.prototype.moe_training.kernels.mxfp8.cute_utils import ( + _missing_cutedsl_runtime_packages, +) +from torchao.utils import ceil_div + +# Kernel geometry constants. SCALE_DIM_X/Y: MXFP8 block sizes (rowwise 1x32, +# colwise 32x1). BUFFS_NUM: smem double-buffer depth. BUFF_DIM_Y: rows per +# pipeline stage buffer. PACK_SIZE: E4M3 bytes packed per b32 store word. +# WAVES: 8B vector groups covering one 1x32 block in the staged rowwise +# swizzled traversal. Chunk shapes (CX, CY) are per-mode Constexpr launch +# parameters; defaults live in _DEFAULT_GEOMETRY below. +SCALE_DIM_Y = 32 +SCALE_DIM_X = 32 +BUFFS_NUM = 2 +BUFF_DIM_Y = 32 +PACK_SIZE = 4 +WAVES = SCALE_DIM_X // PACK_SIZE + +# Chunk geometry (CX, CY, direct) per (is_bwd, rowwise, colwise), tuned on +# GB200; ``direct`` selects the single-pass path over the staged pipeline. +_DEFAULT_GEOMETRY = { + (False, True, False): (128, 64, True), + (False, False, True): (64, 32, True), + (False, True, True): (64, 32, True), + (True, True, False): (64, 64, True), + (True, False, True): (64, 32, True), + (True, True, True): (64, 64, True), +} + +_INT32_MAX = 2**31 - 1 + + +# -- Scale indexing, PTX numeric helpers, activation policy (kernel-private) -- + +EVICT_FIRST = cute.nvgpu.common.CacheEvictionPriority.EVICT_FIRST + +# All TMA-accessed shared-memory buffers (G2S destinations and S2G sources) +# must be 128-byte aligned. +TMA_SHMEM_ALIGNMENT = 128 + + +def _gemm_swizzled_scale_idx(row, scale_col, num_scale_col_blocks): + """Index into the blocked (tcgen05) scale layout expected by MXFP8 GEMMs: + the logical ``[rows, cols/32]`` scale matrix stored as 512-byte blocks of + 128 rows x 4 scale columns (cuBLAS "128x4 block scaling factors layout"); + ``num_scale_col_blocks`` = ceil(num_scale_cols / 4). For colwise scales + pass transposed coordinates. + """ + return ( + ((row >> 7) * num_scale_col_blocks + (scale_col >> 2)) * 512 + + (row & 31) * 16 + + ((row >> 5) & 3) * 4 + + (scale_col & 3) + ) + + +@cute.jit +def _scale_idx(row, scale_col, ncb, stride, SWIZ: cutlass.Constexpr): + """Scale-tensor index for one 32-block: the blocked (GEMM-swizzled) + layout, or the compact row-major ``[rows, cols/32]`` layout.""" + if cutlass.const_expr(SWIZ): + return _gemm_swizzled_scale_idx(row, scale_col, ncb) + else: + return row * stride + scale_col + + +@dsl_user_op +def _bitcast_i32_to_f32(val: Int32, *, loc=None, ip=None) -> Float32: + """Bitcast an int32 value to float32 without changing the bit pattern.""" + return Float32( + mlir_arith.bitcast(T.f32(), val.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + + +# bf16 == top 16 bits of f32, so widening is a free bit-shift. +@dsl_user_op +def _bf16x2_lo_to_f32(bits, *, loc=None, ip=None) -> Float32: + return _bitcast_i32_to_f32( + (Int32(bits) & Int32(0xFFFF)) << Int32(16), loc=loc, ip=ip + ) + + +@dsl_user_op +def _bf16x2_hi_to_f32(bits, *, loc=None, ip=None) -> Float32: + # `(x >> 16) << 16` == `x & 0xFFFF0000` without a signed literal; the + # left shift zeroes the arithmetic shift's smeared sign bits. + return _bitcast_i32_to_f32((Int32(bits) >> Int32(16)) << Int32(16), loc=loc, ip=ip) + + +# The ``.NaN`` max variants match the standalone quantizers' amax reduction, +# which propagates NaN; plain ``max`` would return the non-NaN operand. +@dsl_user_op +def _max_nan_bf16x2(a, b, *, loc=None, ip=None): + """NaN-propagating packed bf16x2 max.""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Int32(a).ir_value(loc=loc, ip=ip), + cutlass.Int32(b).ir_value(loc=loc, ip=ip), + ], + "max.NaN.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _abs_max_nan_bf16x2(a, b, *, loc=None, ip=None): + """NaN-propagating packed bf16x2 |max|; per-lane sign bits are junk.""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Int32(a).ir_value(loc=loc, ip=ip), + cutlass.Int32(b).ir_value(loc=loc, ip=ip), + ], + "max.NaN.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _prmt_even(a, b, *, loc=None, ip=None): + """Select bytes [0,2,4,6] from a pair of b32 words.""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Int32(a).ir_value(loc=loc, ip=ip), + cutlass.Int32(b).ir_value(loc=loc, ip=ip), + ], + "prmt.b32 $0, $1, $2, 0x6420;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _prmt_odd(a, b, *, loc=None, ip=None): + """Select bytes [1,3,5,7] from a pair of b32 words.""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Int32(a).ir_value(loc=loc, ip=ip), + cutlass.Int32(b).ir_value(loc=loc, ip=ip), + ], + "prmt.b32 $0, $1, $2, 0x7531;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _mul_cvt_2x(w0, w1, s, *, loc=None, ip=None): + """Scale two bf16x2 words by bf16x2 ``s`` and pack four E4M3 bytes into + one b32 store word. ``cvt.rn.satfinite.e4m3x2.bf16x2`` is missing on some + Blackwells (GB300's sm_103a), so keep the bf16 multiply for identical + rounding, widen exactly to f32, and use the portable f32-source cvt. + """ + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), + cutlass.Int32(s).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 a, b, t0_lo, t0_hi, t1_lo, t1_hi;\n" + ".reg .b32 t0, t1;\n" + ".reg .f32 f0_lo, f0_hi, f1_lo, f1_hi;\n" + "mul.rn.bf16x2 t0, $1, $3;\n" + "mul.rn.bf16x2 t1, $2, $3;\n" + "mov.b32 {t0_lo, t0_hi}, t0;\n" + "mov.b32 {t1_lo, t1_hi}, t1;\n" + "cvt.f32.bf16 f0_lo, t0_lo;\n" + "cvt.f32.bf16 f0_hi, t0_hi;\n" + "cvt.f32.bf16 f1_lo, t1_lo;\n" + "cvt.f32.bf16 f1_hi, t1_hi;\n" + "cvt.rn.satfinite.e4m3x2.f32 a, f0_hi, f0_lo;\n" + "cvt.rn.satfinite.e4m3x2.f32 b, f1_hi, f1_lo;\n" + "mov.b32 $0, {a, b}; }", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _pack_bf16x2(hi, lo, *, loc=None, ip=None): + """(hi, lo) f32 -> packed bf16x2 word, RNE. lo occupies bits [15:0].""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + Float32(hi).ir_value(loc=loc, ip=ip), + Float32(lo).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _float_to_e8m0(u: Int32) -> Int32: + """Biased E8M0 byte for a non-negative bf16 amax given as f32 bits. + + Finite: the RCEIL mantissa-carry path, matching what the standalone + quantizers' ``cvt.rp.ue8m0x2.f32`` (no ``.satfinite``) emits for + ``amax / 448``. Non-finite: NaN or Inf amax invalidates the block with + scale byte 255; without the branch, Inf would land on 247 and NaN could + carry into the sign bit. + """ + e = cutlass.max(((u + Int32(0x1F0000)) >> 23) - Int32(8), Int32(0)) + if (u & Int32(0x7F800000)) == Int32(0x7F800000): + e = Int32(255) + return e + + +@cute.jit +def _exp2f_rcp_bf16(e: Int32) -> Int32: + """Inverse scale as bf16 bits (the quantization multiply in + :func:`_mul_cvt_2x` is bf16x2, not f32), matching the standalone + quantizers' ``ue8m0(254 - scale_byte)`` reciprocal: 2^(127 - e) for the + normal range (a clamped byte 0 from a zero or tiny amax descales by + 2^127), and NaN for an invalidated block (byte 255), so every element of + a NaN/Inf-amax block quantizes to the E4M3 NaN code. + """ + b = (Int32(254) - e) << 7 + if e == Int32(255): + b = Int32(0x7FC0) + return b + + +@dsl_user_op +def _sigmoidf(x, *, loc=None, ip=None): + """Sigmoid as ``__frcp_rn(1.0f + __expf(-x))``, emitted as raw PTX, + instruction for instruction:: + + mul.f32 t, x, 0fBFB8AA3B // -x * log2(e) + ex2.approx.f32 t, t + add.f32 t, t, 0f3F800000 + rcp.rn.f32 s, t // correctly rounded + + No higher-level formulation reproduces ``ex2.approx``, and the ``rcp.rn`` + vs ``div.full`` choice shows at a few output codes per million. + """ + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(x).ir_value(loc=loc, ip=ip)], + "{ .reg .f32 t;\n" + "mul.f32 t, $1, 0fBFB8AA3B;\n" + "ex2.approx.f32 t, t;\n" + "add.f32 t, t, 0f3F800000;\n" + "rcp.rn.f32 $0, t; }", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _silu_pair(x0, x1, lin0, lin1, g0, g1, IS_BWD: cutlass.Constexpr): + """SwiGLU activation policy for a pair of elements: ``x`` is the + activation input, ``lin`` the linear multiplier, ``g`` the incoming + gradient (ignored unless IS_BWD):: + + s = sigmoid(x); act = x * s + forward: out_act = act * lin + backward: dact = x*s*(1-s) + s (contracted into one FMA) + out_act = (dact * g) * lin, out_gate = act * g + + Returns f32 ``(out_act0, out_act1, out_gate0, out_gate1)`` (gate pair + zero for forward); callers must round to BF16 immediately, before any + amax, caching, or quantization. + + In the module docstring's terms: ``x`` = gate, ``lin`` = up, ``g`` = + grad_h; backward's ``out_act`` half holds dGate and ``out_gate`` holds + dUp — the *_gate-suffixed symbols throughout the kernel carry dUp (stored + at offset K), not a gate gradient. + """ + one = cutlass.Float32(1.0) + s0 = _sigmoidf(x0) + s1 = _sigmoidf(x1) + act0, act1 = cute.arch.mul_packed_f32x2((x0, x1), (s0, s1)) + if cutlass.const_expr(IS_BWD): + om0, om1 = cute.arch.sub_packed_f32x2((one, one), (s0, s1)) + dact0, dact1 = cute.arch.fma_packed_f32x2((act0, act1), (om0, om1), (s0, s1)) + t0, t1 = cute.arch.mul_packed_f32x2((dact0, dact1), (g0, g1)) + oa0, oa1 = cute.arch.mul_packed_f32x2((t0, t1), (lin0, lin1)) + og0, og1 = cute.arch.mul_packed_f32x2((act0, act1), (g0, g1)) + return oa0, oa1, og0, og1 + else: + oa0, oa1 = cute.arch.mul_packed_f32x2((act0, act1), (lin0, lin1)) + return oa0, oa1, cutlass.Float32(0.0), cutlass.Float32(0.0) + + +def _load_direct_inputs(gXv, gLinv, gGradv, half, blk, bx, grow, IS_BWD): + """Issue one stage's evict-first vector input loads into fresh rmem + tensors; trace-time helper for the direct path. ``rg`` is None in + forward mode.""" + rx = cute.make_rmem_tensor(8, cutlass.Int32) + rl = cute.make_rmem_tensor(8, cutlass.Int32) + cute.autovec_copy( + gXv[(None, half, blk, bx, grow)], + rx, + l1c_evict_priority=EVICT_FIRST, + ) + cute.autovec_copy( + gLinv[(None, half, blk, bx, grow)], + rl, + l1c_evict_priority=EVICT_FIRST, + ) + rg = None + if IS_BWD: + rg = cute.make_rmem_tensor(8, cutlass.Int32) + cute.autovec_copy( + gGradv[(None, half, blk, bx, grow)], + rg, + l1c_evict_priority=EVICT_FIRST, + ) + return rx, rl, rg + + +@cute.jit +def _fold_amax(am: cutlass.Int32) -> cutlass.Int32: + """Reduce a packed bf16x2 amax word to bf16 amax bits in [15:0]. + + The input's per-lane sign bits are junk (see ``_abs_max_nan_bf16x2``); mask + them, then fold the two lanes with the NaN-propagating max. + """ + am = am & cutlass.Int32(0x7FFF7FFF) + am = _max_nan_bf16x2(am, am >> 16) + return am & cutlass.Int32(0xFFFF) + + +@cute.kernel +def gated_act_mxfp8_kernel( + atom_x: cute.CopyAtom, + gX: cute.Tensor, + atom_lin: cute.CopyAtom, + gLin: cute.Tensor, + atom_grad: cute.CopyAtom, + gGrad: cute.Tensor, + gXv: cute.Tensor, + gLinv: cute.Tensor, + gGradv: cute.Tensor, + atom_row_act: cute.CopyAtom, + gRowAct: cute.Tensor, + atom_row_gate: cute.CopyAtom, + gRowGate: cute.Tensor, + atom_col_act: cute.CopyAtom, + gColAct: cute.Tensor, + atom_col_gate: cute.CopyAtom, + gColGate: cute.Tensor, + mRS: cute.Tensor, + mCS: cute.Tensor, + rs_ncb: cutlass.Int32, + rs_stride: cutlass.Int32, + rgate_scol_off: cutlass.Int32, + cs_ncb: cutlass.Int32, + cs_stride: cutlass.Int32, + cgate_col_off: cutlass.Int32, + IS_BWD: cutlass.Constexpr, + ROWWISE: cutlass.Constexpr, + COLWISE: cutlass.Constexpr, + SWIZ: cutlass.Constexpr, + ACT_PAIR: cutlass.Constexpr, + DIRECT: cutlass.Constexpr, + CX: cutlass.Constexpr, + CY: cutlass.Constexpr, + THREADS: cutlass.Constexpr, +): + """The gated-activation MXFP8 kernel, specialized entirely at compile + time. Unused parameters are never referenced: disabled directions receive + dummies from the launcher, and the non-selected DIRECT/staged path's + views are simply ignored.""" + tidx, _, _ = cute.arch.thread_idx() + bx, by, _ = cute.arch.block_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + IS_CACHED_ACT_OP = ROWWISE and COLWISE + ONLY_COLWISE = COLWISE and not ROWWISE + OUT_HALVES_C = 2 if IS_BWD else 1 + # DIRECT reads inputs straight into registers and computes the activation + # exactly once (the staged two-pass structure was latency-bound on + # GB200); the colwise side goes through a padded XOR-swizzled transposed + # shared tile instead of recomputation. + COLWISE_WF = THREADS // CX # thread rows stacked along each column + N_STAGES = CY // BUFF_DIM_Y + LOG2_CX = CX.bit_length() - 1 + # 1x32 scale blocks per chunk row (staged rowwise: also threads per row; + # direct: two threads split each block). + TXR = CX // SCALE_DIM_X + LOG2_TXR = TXR.bit_length() - 1 + BUFF_BYTES = BUFF_DIM_Y * CX * 2 + # Direct colwise path: bf16x2 words hold two adjacent columns, so the + # padded transposed tile is indexed by column pair. + PAIRS = CX // 2 + LOG2_PAIRS = PAIRS.bit_length() - 1 + SP_PAD = BUFF_DIM_Y + 4 + if cutlass.const_expr(DIRECT and COLWISE): + # Reader: TPP threads per (column pair, output half); each thread + # owns NCH of the pair's eight 16B row-chunks. + assert THREADS % (PAIRS * OUT_HALVES_C) == 0 + TPP = THREADS // (PAIRS * OUT_HALVES_C) + assert TPP & (TPP - 1) == 0 and 1 <= TPP <= 8 + LOG2_TPP = TPP.bit_length() - 1 + NCH = 8 // TPP + # Prefetch only where the reader phase hides the load latency; measured + # on GB200, prefetch costs ~16 registers and regressed rowwise-only 13%. + DO_PREFETCH = COLWISE + # Direct loads have no mbarrier chain ordering output-buffer reuse against + # outstanding TMA-store reads; only the colwise path's pre-reader barrier + # can cover that, so rowwise-only direct needs one buffer per stage. + assert not DIRECT or COLWISE or N_STAGES <= BUFFS_NUM + # The staged path has no in-loop TMA-store drain: output buffers are safe + # only because the post-loop drain runs before any reuse, which requires + # the stage count to fit the double buffer. The launcher rejects + # violating geometries; this is the trace-time backstop. + assert DIRECT or N_STAGES <= BUFFS_NUM + # Staged rowwise consumes the colwise pass's cached activation; there is + # no standalone staged rowwise compute. + assert DIRECT or COLWISE or not ROWWISE + + # Shared buffers get a value-typed view (TMA partitioning, scalar access) + # plus a b32-word view (vectorized 8B loads, packed 4B stores). + smem = cutlass.utils.SmemAllocator() + + in_elems = BUFF_DIM_Y * CX * BUFFS_NUM + if cutlass.const_expr(not DIRECT): + in_layout = cute.make_layout( + (BUFF_DIM_Y, CX, BUFFS_NUM), + stride=(CX, 1, BUFF_DIM_Y * CX), + ) + # 8B-vector view: (lane pair, 8B group, row, buff). + in_word_layout = cute.make_layout( + (2, CX // 4, BUFF_DIM_Y, BUFFS_NUM), + stride=(1, 2, CX // 2, BUFF_DIM_Y * CX // 2), + ) + + px = smem.allocate(in_elems * 2, byte_alignment=TMA_SHMEM_ALIGNMENT) + sX = cute.make_tensor(cute.recast_ptr(px, dtype=cutlass.BFloat16), in_layout) + sXw = cute.make_tensor(cute.recast_ptr(px, dtype=cutlass.Int32), in_word_layout) + pl = smem.allocate(in_elems * 2, byte_alignment=TMA_SHMEM_ALIGNMENT) + sLin = cute.make_tensor(cute.recast_ptr(pl, dtype=cutlass.BFloat16), in_layout) + sLinw = cute.make_tensor( + cute.recast_ptr(pl, dtype=cutlass.Int32), in_word_layout + ) + if cutlass.const_expr(IS_BWD): + pg = smem.allocate(in_elems * 2, byte_alignment=TMA_SHMEM_ALIGNMENT) + sGrad = cute.make_tensor( + cute.recast_ptr(pg, dtype=cutlass.BFloat16), in_layout + ) + # Bidirectional mode: the columnwise pass caches the post-activation + # BF16 in the input buffers so the rowwise pass never recomputes it + # (which is why staged ROWWISE without COLWISE is invalid — asserted + # above). + cached_act, cached_actw = sX, sXw + cached_gate, cached_gatew = sLin, sLinw + + if cutlass.const_expr(ROWWISE): + row_out_layout = cute.make_layout( + (BUFF_DIM_Y, CX, BUFFS_NUM), + stride=(CX, 1, BUFF_DIM_Y * CX), + ) + row_out_word_layout = cute.make_layout( + (BUFF_DIM_Y, CX // 4, BUFFS_NUM), + stride=(CX // 4, 1, BUFF_DIM_Y * CX // 4), + ) + # Quad view for the direct path's 16B vector stores. + row_out_quad_layout = cute.make_layout( + (4, CX // 16, BUFF_DIM_Y, BUFFS_NUM), + stride=(1, 4, CX // 4, BUFF_DIM_Y * CX // 4), + ) + pra = smem.allocate(in_elems, byte_alignment=TMA_SHMEM_ALIGNMENT) + sRowAct = cute.make_tensor( + cute.recast_ptr(pra, dtype=cutlass.Float8E4M3FN), row_out_layout + ) + sRowActw = cute.make_tensor( + cute.recast_ptr(pra, dtype=cutlass.Int32), row_out_word_layout + ) + sRowQuad = cute.make_tensor( + cute.recast_ptr(pra, dtype=cutlass.Int32), row_out_quad_layout + ) + if cutlass.const_expr(IS_BWD): + prg = smem.allocate(in_elems, byte_alignment=TMA_SHMEM_ALIGNMENT) + sRowGate = cute.make_tensor( + cute.recast_ptr(prg, dtype=cutlass.Float8E4M3FN), row_out_layout + ) + sRowGatew = cute.make_tensor( + cute.recast_ptr(prg, dtype=cutlass.Int32), row_out_word_layout + ) + sRowGateQuad = cute.make_tensor( + cute.recast_ptr(prg, dtype=cutlass.Int32), row_out_quad_layout + ) + if cutlass.const_expr(COLWISE): + # Transposed (output column, row) tiles; staging + TMA store beats + # direct scattered global stores (measured: direct stores pushed + # L1TEX to 97% of peak and cost +12% on bwd_rc). + col_out_layout = cute.make_layout( + (CX, BUFF_DIM_Y, BUFFS_NUM), + stride=(BUFF_DIM_Y, 1, CX * BUFF_DIM_Y), + ) + col_out_word_layout = cute.make_layout( + (CX, BUFF_DIM_Y // 4, BUFFS_NUM), + stride=(BUFF_DIM_Y // 4, 1, CX * BUFF_DIM_Y // 4), + ) + pca = smem.allocate(in_elems, byte_alignment=TMA_SHMEM_ALIGNMENT) + sColAct = cute.make_tensor( + cute.recast_ptr(pca, dtype=cutlass.Float8E4M3FN), col_out_layout + ) + sColActw = cute.make_tensor( + cute.recast_ptr(pca, dtype=cutlass.Int32), col_out_word_layout + ) + if cutlass.const_expr(IS_BWD): + pcg = smem.allocate(in_elems, byte_alignment=TMA_SHMEM_ALIGNMENT) + sColGate = cute.make_tensor( + cute.recast_ptr(pcg, dtype=cutlass.Float8E4M3FN), col_out_layout + ) + sColGatew = cute.make_tensor( + cute.recast_ptr(pcg, dtype=cutlass.Int32), col_out_word_layout + ) + if cutlass.const_expr(DIRECT): + # Reader slice view: each of a column's TPP threads owns NCH + # contiguous words. + col_slice_layout = cute.make_layout( + (NCH, TPP, CX, BUFFS_NUM), + stride=(1, NCH, BUFF_DIM_Y // 4, CX * BUFF_DIM_Y // 4), + ) + sColSliceA = cute.make_tensor( + cute.recast_ptr(pca, dtype=cutlass.Int32), col_slice_layout + ) + if cutlass.const_expr(IS_BWD): + sColSliceG = cute.make_tensor( + cute.recast_ptr(pcg, dtype=cutlass.Int32), col_slice_layout + ) + + if cutlass.const_expr(DIRECT and COLWISE): + # Padded transposed staging between the compute pass and the colwise + # reader; a pair's 32 rows are eight 16B chunks at addr(pair, row) = + # pair*SP_PAD + 4*((row>>2) ^ ((pair>>3 & 3) << 1)) + (row&3). The + # padding plus XOR swizzle make both the writer's word stores and the + # reader's 16B loads shared-memory bank-conflict-free. + ppad = smem.allocate(PAIRS * SP_PAD * OUT_HALVES_C * 4, byte_alignment=16) + sPadW = cute.make_tensor( + cute.recast_ptr(ppad, dtype=cutlass.Int32), + cute.make_layout( + (PAIRS, 8, 4, OUT_HALVES_C), + stride=(SP_PAD, 4, 1, PAIRS * SP_PAD), + ), + ) + sPadR = cute.make_tensor( + cute.recast_ptr(ppad, dtype=cutlass.Int32), + cute.make_layout( + (4, 8, PAIRS, OUT_HALVES_C), + stride=(1, 4, SP_PAD, PAIRS * SP_PAD), + ), + ) + if cutlass.const_expr(not DIRECT): + mbar = smem.allocate_array(cutlass.Int64, N_STAGES, byte_alignment=8) + if cutlass.const_expr(ONLY_COLWISE and not DIRECT): + # Partial-amax exchange between the COLWISE_WF thread rows; the + # single sSubAmax slot per column assumes exactly one non-zero row. + assert COLWISE_WF == 2 + psub = smem.allocate(CX * 4, byte_alignment=4) + sSubAmax = cute.make_tensor( + cute.recast_ptr(psub, dtype=cutlass.Int32), cute.make_layout(CX) + ) + + if cutlass.const_expr(not DIRECT): + # tma_partition takes the no-multicast CTA coord (0) and layout; the + # smem view groups its buffer modes first, the gmem view is the tile + # shape ``cute.zipped_divide`` produced in the launcher. + tXs, tXg = cpasync.tma_partition( + atom_x, 0, cute.make_layout(1), cute.group_modes(sX, 0, 2), gX + ) + tLs, tLg = cpasync.tma_partition( + atom_lin, 0, cute.make_layout(1), cute.group_modes(sLin, 0, 2), gLin + ) + if cutlass.const_expr(IS_BWD): + tGs, tGg = cpasync.tma_partition( + atom_grad, 0, cute.make_layout(1), cute.group_modes(sGrad, 0, 2), gGrad + ) + if cutlass.const_expr(ROWWISE): + tRAs, tRAg = cpasync.tma_partition( + atom_row_act, + 0, + cute.make_layout(1), + cute.group_modes(sRowAct, 0, 2), + gRowAct, + ) + if cutlass.const_expr(IS_BWD): + tRGs, tRGg = cpasync.tma_partition( + atom_row_gate, + 0, + cute.make_layout(1), + cute.group_modes(sRowGate, 0, 2), + gRowGate, + ) + if cutlass.const_expr(COLWISE): + tCAs, tCAg = cpasync.tma_partition( + atom_col_act, + 0, + cute.make_layout(1), + cute.group_modes(sColAct, 0, 2), + gColAct, + ) + if cutlass.const_expr(IS_BWD): + tCGs, tCGg = cpasync.tma_partition( + atom_col_gate, + 0, + cute.make_layout(1), + cute.group_modes(sColGate, 0, 2), + gColGate, + ) + + if cutlass.const_expr(not DIRECT): + # Barrier arrive count = CTA thread count (every thread arrives once + # per stage; TMA bytes are tracked on top via ``expect_tx``); the + # init fence makes the initialization visible to the TMA async proxy. + if tidx == 0: + for i in cutlass.range_constexpr(N_STAGES): + cute.arch.mbarrier_init(mbar + i, THREADS) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + if cutlass.const_expr(DIRECT): + # Stage-invariant mapping: two threads split each 1x32 block. + half = tidx & 1 + blk = (tidx >> 1) & cutlass.Int32(TXR - 1) + row = tidx >> (1 + LOG2_TXR) + if cutlass.const_expr(COLWISE): + # sPad swizzle key; the thread's pairs are 16*blk + 8*half + j, + # so pair bits 3-4 (hence the key) are stage-invariant. + wchk = (row >> 2) ^ (((blk * 2 + half) & cutlass.Int32(3)) * 2) + wrow = row & cutlass.Int32(3) + rxs = [None] * N_STAGES + rls = [None] * N_STAGES + rgs = [None] * N_STAGES + if cutlass.const_expr(DO_PREFETCH): + grow0 = by * cutlass.Int32(CY) + row + rxs[0], rls[0], rgs[0] = _load_direct_inputs( + gXv, gLinv, gGradv, half, blk, bx, grow0, IS_BWD + ) + + row_tile0 = by * N_STAGES + if cutlass.const_expr(not DIRECT): + copies = ( + (atom_x, tXg[(None, (row_tile0, bx))], tXs[(None, 0)]), + (atom_lin, tLg[(None, (row_tile0, bx))], tLs[(None, 0)]), + ) + if cutlass.const_expr(IS_BWD): + copies += ((atom_grad, tGg[(None, (row_tile0, bx))], tGs[(None, 0)]),) + # The TMA copies must issue under a warp-uniform predicate (a + # single-thread predicate deadlocks the DSL's issuing-lane election); + # thread 0 expects the combined byte count, every thread arrives once. + if warp_idx == 0: + for atom, g, s in copies: + cute.copy(atom, g, s, tma_bar_ptr=mbar) + if tidx == 0: + cute.arch.mbarrier_arrive_and_expect_tx(mbar, len(copies) * BUFF_BYTES) + else: + cute.arch.mbarrier_arrive(mbar) + + for stage in cutlass.range_constexpr(N_STAGES): + buff = stage % BUFFS_NUM + row_tile = by * N_STAGES + stage + + if cutlass.const_expr((not DIRECT) and stage + 1 < N_STAGES): + # Prefetch the next stage's inputs. Output smem is never reused + # inside this loop (N_STAGES <= BUFFS_NUM, asserted above), so + # TMA-store groups need no drain before the post-loop + # wait_group(0); the input refill is ordered by the + # end-of-stage sync. + nbuff = (stage + 1) % BUFFS_NUM + nmbar = mbar + (stage + 1) + row_tile_n = by * N_STAGES + (stage + 1) + copies = ( + (atom_x, tXg[(None, (row_tile_n, bx))], tXs[(None, nbuff)]), + (atom_lin, tLg[(None, (row_tile_n, bx))], tLs[(None, nbuff)]), + ) + if cutlass.const_expr(IS_BWD): + copies += ( + (atom_grad, tGg[(None, (row_tile_n, bx))], tGs[(None, nbuff)]), + ) + if warp_idx == 0: + for atom, g, s in copies: + cute.copy(atom, g, s, tma_bar_ptr=nmbar) + if tidx == 0: + cute.arch.mbarrier_arrive_and_expect_tx(nmbar, len(copies) * BUFF_BYTES) + else: + cute.arch.mbarrier_arrive(nmbar) + + if cutlass.const_expr(not DIRECT): + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.mbarrier_wait(mbar + stage, 0) + + # -- Columnwise producer pass (staged path) ------------------------- + if cutlass.const_expr(COLWISE and not DIRECT): + col = tidx & cutlass.Int32(CX - 1) + ROWS_PER_THREAD = SCALE_DIM_Y // COLWISE_WF + NWORDS = ROWS_PER_THREAD // 2 + + # 1. Compute post-activation values, round to BF16, find amax. + w_act = [None] * NWORDS + w_gate = [None] * NWORDS + am_act = cutlass.Int32(0) + am_gate = cutlass.Int32(0) + for j in cutlass.range_constexpr(NWORDS): + # Contiguous per-thread row split: the amax reduction is + # order-independent, and contiguous rows pack into 4B words. + if cutlass.const_expr(COLWISE_WF == 1): + rlo = 2 * j + rhi = 2 * j + 1 + else: + ty = tidx >> LOG2_CX + rlo = ty * ROWS_PER_THREAD + 2 * j + rhi = rlo + 1 + x0 = sX[(rlo, col, buff)].to(cutlass.Float32) + x1 = sX[(rhi, col, buff)].to(cutlass.Float32) + l0 = sLin[(rlo, col, buff)].to(cutlass.Float32) + l1 = sLin[(rhi, col, buff)].to(cutlass.Float32) + if cutlass.const_expr(IS_BWD): + g0 = sGrad[(rlo, col, buff)].to(cutlass.Float32) + g1 = sGrad[(rhi, col, buff)].to(cutlass.Float32) + else: + g0 = cutlass.Float32(0.0) + g1 = cutlass.Float32(0.0) + oa0, oa1, og0, og1 = ACT_PAIR(x0, x1, l0, l1, g0, g1, IS_BWD) + # Numerical truncation to the input type before anything else. + wa = _pack_bf16x2(oa1, oa0) + w_act[j] = wa + am_act = _abs_max_nan_bf16x2(am_act, wa) + if cutlass.const_expr(IS_BWD): + wg = _pack_bf16x2(og1, og0) + w_gate[j] = wg + am_gate = _abs_max_nan_bf16x2(am_gate, wg) + if cutlass.const_expr(IS_CACHED_ACT_OP): + cached_act[(rlo, col, buff)] = oa0.to(cutlass.BFloat16) + cached_act[(rhi, col, buff)] = oa1.to(cutlass.BFloat16) + if cutlass.const_expr(IS_BWD): + cached_gate[(rlo, col, buff)] = og0.to(cutlass.BFloat16) + cached_gate[(rhi, col, buff)] = og1.to(cutlass.BFloat16) + + am_act = _fold_amax(am_act) + if cutlass.const_expr(IS_BWD): + am_gate = _fold_amax(am_gate) + + # Reduce partial amaxes across the two thread rows (colwise-only + # staged always launches THREADS == 2*CX, so ty is 0 or 1; the + # exchange does not generalize to more rows). + if cutlass.const_expr(ONLY_COLWISE): + ty = tidx >> LOG2_CX + if ty > 0: + sSubAmax[col] = am_act + cute.arch.sync_threads() + if ty == 0: + am_act = _max_nan_bf16x2(am_act, sSubAmax[col]) + sSubAmax[col] = am_act + cute.arch.sync_threads() + am_act = sSubAmax[col] + if cutlass.const_expr(IS_BWD): + # The previous reads must complete before the rewrite. + cute.arch.sync_threads() + if ty > 0: + sSubAmax[col] = am_gate + cute.arch.sync_threads() + if ty == 0: + am_gate = _max_nan_bf16x2(am_gate, sSubAmax[col]) + sSubAmax[col] = am_gate + cute.arch.sync_threads() + am_gate = sSubAmax[col] + + # 2. Compute and store the E8M0 scales (one per 32x1 block). + out_col = bx * cutlass.Int32(CX) + col + mcol = row_tile + u_act = am_act << 16 + e_act = _float_to_e8m0(u_act) + sidx = _scale_idx(out_col, mcol, cs_ncb, cs_stride, SWIZ) + if cutlass.const_expr(ONLY_COLWISE): + if tidx < cutlass.Int32(CX): + mCS[sidx] = e_act.to(cutlass.Uint8) + else: + mCS[sidx] = e_act.to(cutlass.Uint8) + r_act = _exp2f_rcp_bf16(e_act) * cutlass.Int32(0x10001) + + if cutlass.const_expr(IS_BWD): + u_gate = am_gate << 16 + e_gate = _float_to_e8m0(u_gate) + gidx = _scale_idx( + out_col + cgate_col_off, mcol, cs_ncb, cs_stride, SWIZ + ) + if cutlass.const_expr(ONLY_COLWISE): + if tidx < cutlass.Int32(CX): + mCS[gidx] = e_gate.to(cutlass.Uint8) + else: + mCS[gidx] = e_gate.to(cutlass.Uint8) + r_gate = _exp2f_rcp_bf16(e_gate) * cutlass.Int32(0x10001) + + # 3. Scale and pack into the transposed shared output tile. + if cutlass.const_expr(COLWISE_WF == 1): + wbase = 0 + else: + wbase = (tidx >> LOG2_CX) * (NWORDS // 2) + for w in cutlass.range_constexpr(NWORDS // 2): + sColActw[(col, wbase + w, buff)] = _mul_cvt_2x( + w_act[2 * w], w_act[2 * w + 1], r_act + ) + if cutlass.const_expr(IS_BWD): + sColGatew[(col, wbase + w, buff)] = _mul_cvt_2x( + w_gate[2 * w], w_gate[2 * w + 1], r_gate + ) + + # -- Direct single-pass compute ------------------------------------- + if cutlass.const_expr(DIRECT): + # One contiguous 32B evict-first load per stream (each input is + # touched exactly once); the activation feeds both orientations. + grow = by * cutlass.Int32(CY) + cutlass.Int32(stage * BUFF_DIM_Y) + row + + # With prefetch, issue the NEXT stage's loads so their latency + # hides behind this stage's compute and reader phases. + ld = stage + 1 if DO_PREFETCH else stage + if cutlass.const_expr(ld < N_STAGES): + grow_ld = by * cutlass.Int32(CY) + cutlass.Int32(ld * BUFF_DIM_Y) + row + rxs[ld], rls[ld], rgs[ld] = _load_direct_inputs( + gXv, gLinv, gGradv, half, blk, bx, grow_ld, IS_BWD + ) + rx = rxs[stage] + rl = rls[stage] + rg = rgs[stage] + + am_act = cutlass.Int32(0) + am_gate = cutlass.Int32(0) + w_act = [None] * 8 + w_gate = [None] * 8 + for j in cutlass.range_constexpr(8): + x0 = _bf16x2_lo_to_f32(rx[j]) + x1 = _bf16x2_hi_to_f32(rx[j]) + l0 = _bf16x2_lo_to_f32(rl[j]) + l1 = _bf16x2_hi_to_f32(rl[j]) + if cutlass.const_expr(IS_BWD): + g0 = _bf16x2_lo_to_f32(rg[j]) + g1 = _bf16x2_hi_to_f32(rg[j]) + else: + g0 = cutlass.Float32(0.0) + g1 = cutlass.Float32(0.0) + oa0, oa1, og0, og1 = ACT_PAIR(x0, x1, l0, l1, g0, g1, IS_BWD) + w_act[j] = _pack_bf16x2(oa1, oa0) + am_act = _abs_max_nan_bf16x2(am_act, w_act[j]) + if cutlass.const_expr(IS_BWD): + w_gate[j] = _pack_bf16x2(og1, og0) + am_gate = _abs_max_nan_bf16x2(am_gate, w_gate[j]) + # Park each column-pair word in the swizzled transposed tile + # for the columnwise reader. + if cutlass.const_expr(COLWISE): + pair = ( + blk * cutlass.Int32(16) + + half * cutlass.Int32(8) + + cutlass.Int32(j) + ) + sPadW[(pair, wchk, wrow, 0)] = w_act[j] + if cutlass.const_expr(IS_BWD): + sPadW[(pair, wchk, wrow, 1)] = w_gate[j] + + # -- Direct columnwise reader, part 1 (issue) ------------------------- + if cutlass.const_expr(DIRECT and COLWISE): + if cutlass.const_expr(stage >= BUFFS_NUM): + # TMA stores committed at stage - BUFFS_NUM may still read + # the tiles this stage overwrites: warp 0 drains them, the + # sync below releases everyone's stores. + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(BUFFS_NUM - 1, read=True) + cute.arch.sync_threads() + # Loads issued here so their latency hides behind the rowwise + # block; thread tq of a (pair, half) owns NCH 16B row-chunks. + tq = tidx & cutlass.Int32(TPP - 1) + cpr = (tidx >> LOG2_TPP) & cutlass.Int32(PAIRS - 1) + arr = tidx >> (LOG2_TPP + LOG2_PAIRS) + rsw = ((cpr >> 3) & cutlass.Int32(3)) * 2 + + # Undo the writer's chunk swizzle with the same XOR key. + vs = [cute.make_rmem_tensor(4, cutlass.Int32) for _ in range(NCH)] + for c in cutlass.range_constexpr(NCH): + li = tq * cutlass.Int32(NCH) + cutlass.Int32(c) + cute.autovec_copy(sPadR[(None, li ^ rsw, cpr, arr)], vs[c]) + + if cutlass.const_expr(DIRECT and ROWWISE): + # Butterfly-combine the half-block amaxes; the even lane of each + # pair owns the scale-byte store. + scol = bx * cutlass.Int32(TXR) + blk + am_act = am_act & cutlass.Int32(0x7FFF7FFF) + am_act = _max_nan_bf16x2(am_act, cute.arch.shuffle_sync_bfly(am_act, 1)) + am_act = _max_nan_bf16x2(am_act, am_act >> 16) + u_act = (am_act & cutlass.Int32(0xFFFF)) << 16 + e_act = _float_to_e8m0(u_act) + sidx = _scale_idx(grow, scol, rs_ncb, rs_stride, SWIZ) + if half == 0: + mRS[sidx] = e_act.to(cutlass.Uint8) + r_act = _exp2f_rcp_bf16(e_act) * cutlass.Int32(0x10001) + qa = cute.make_rmem_tensor(4, cutlass.Int32) + for q in cutlass.range_constexpr(4): + qa[q] = _mul_cvt_2x(w_act[2 * q], w_act[2 * q + 1], r_act) + wq = blk * 2 + half + cute.autovec_copy(qa, sRowQuad[(None, wq, row, buff)]) + + if cutlass.const_expr(IS_BWD): + am_gate = am_gate & cutlass.Int32(0x7FFF7FFF) + am_gate = _max_nan_bf16x2( + am_gate, cute.arch.shuffle_sync_bfly(am_gate, 1) + ) + am_gate = _max_nan_bf16x2(am_gate, am_gate >> 16) + u_gate = (am_gate & cutlass.Int32(0xFFFF)) << 16 + e_gate = _float_to_e8m0(u_gate) + gidx = _scale_idx(grow, scol + rgate_scol_off, rs_ncb, rs_stride, SWIZ) + if half == 0: + mRS[gidx] = e_gate.to(cutlass.Uint8) + r_gate = _exp2f_rcp_bf16(e_gate) * cutlass.Int32(0x10001) + qg = cute.make_rmem_tensor(4, cutlass.Int32) + for q in cutlass.range_constexpr(4): + qg[q] = _mul_cvt_2x(w_gate[2 * q], w_gate[2 * q + 1], r_gate) + cute.autovec_copy(qg, sRowGateQuad[(None, wq, row, buff)]) + + # -- Direct columnwise reader, part 2 (consume) ----------------------- + if cutlass.const_expr(DIRECT and COLWISE): + ac = cutlass.Int32(0) + for c in cutlass.range_constexpr(NCH): + for t in cutlass.range_constexpr(4): + ac = _abs_max_nan_bf16x2(ac, vs[c][t]) + ac = ac & cutlass.Int32(0x7FFF7FFF) + # Butterfly-combine the TPP partial amaxes (bit-identical to a + # single-thread fold). + for d in cutlass.range_constexpr(LOG2_TPP): + ac = _max_nan_bf16x2(ac, cute.arch.shuffle_sync_bfly(ac, 1 << d)) + + # Independent per-column scales for the two packed lanes. + uc0 = (ac & cutlass.Int32(0xFFFF)) << 16 + uc1 = ac & cutlass.Int32(-65536) + ec0 = _float_to_e8m0(uc0) + ec1 = _float_to_e8m0(uc1) + s01 = _exp2f_rcp_bf16(ec0) | (_exp2f_rcp_bf16(ec1) << 16) + + # arr is 0 for every active thread in forward mode, so the + # gate-half offset term vanishes there. + c_out_col = bx * cutlass.Int32(CX) + cpr * 2 + arr * cgate_col_off + if tq == 0: + ci0 = _scale_idx(c_out_col, row_tile, cs_ncb, cs_stride, SWIZ) + ci1 = _scale_idx(c_out_col + 1, row_tile, cs_ncb, cs_stride, SWIZ) + mCS[ci0] = ec0.to(cutlass.Uint8) + mCS[ci1] = ec1.to(cutlass.Uint8) + + # Quantize and de-interleave the two columns' bytes into the TMA + # tile; the ~2-way bank conflict is inherent (the tile's 32B + # column blocks must stay contiguous for TMA). + fq = cute.make_rmem_tensor(NCH, cutlass.Int32) + fqb = cute.make_rmem_tensor(NCH, cutlass.Int32) + for c in cutlass.range_constexpr(NCH): + a01 = _mul_cvt_2x(vs[c][0], vs[c][1], s01) + a23 = _mul_cvt_2x(vs[c][2], vs[c][3], s01) + fq[c] = _prmt_even(a01, a23) + fqb[c] = _prmt_odd(a01, a23) + col_local = cpr * 2 + if cutlass.const_expr(IS_BWD): + if arr == 0: + cute.autovec_copy(fq, sColSliceA[(None, tq, col_local, buff)]) + cute.autovec_copy(fqb, sColSliceA[(None, tq, col_local + 1, buff)]) + else: + cute.autovec_copy(fq, sColSliceG[(None, tq, col_local, buff)]) + cute.autovec_copy(fqb, sColSliceG[(None, tq, col_local + 1, buff)]) + else: + cute.autovec_copy(fq, sColSliceA[(None, tq, col_local, buff)]) + cute.autovec_copy(fqb, sColSliceA[(None, tq, col_local + 1, buff)]) + + if cutlass.const_expr(ROWWISE and not DIRECT): + row = tidx >> LOG2_TXR + tx = tidx & cutlass.Int32(TXR - 1) + bank_group = (tidx & 31) >> 2 + grow = by * cutlass.Int32(CY) + cutlass.Int32(stage * BUFF_DIM_Y) + row + + # Make the columnwise pass's cache writes visible. + cute.arch.sync_threads() + + # 1. Load the cached post-activation values with the bank-group + # swizzle; each thread owns a whole 1x32 block per output half. + am_act = cutlass.Int32(0) + am_gate = cutlass.Int32(0) + iv_act = [cute.make_rmem_tensor(2, cutlass.Int32) for _ in range(WAVES)] + if cutlass.const_expr(IS_BWD): + iv_gate = [ + cute.make_rmem_tensor(2, cutlass.Int32) for _ in range(WAVES) + ] + for w in cutlass.range_constexpr(WAVES): + grp = tx * cutlass.Int32(WAVES) + ((cutlass.Int32(w) + bank_group) & 7) + cute.autovec_copy(cached_actw[(None, grp, row, buff)], iv_act[w]) + if cutlass.const_expr(IS_BWD): + cute.autovec_copy(cached_gatew[(None, grp, row, buff)], iv_gate[w]) + am_act = _abs_max_nan_bf16x2(am_act, iv_act[w][0]) + am_act = _abs_max_nan_bf16x2(am_act, iv_act[w][1]) + if cutlass.const_expr(IS_BWD): + am_gate = _abs_max_nan_bf16x2(am_gate, iv_gate[w][0]) + am_gate = _abs_max_nan_bf16x2(am_gate, iv_gate[w][1]) + + # 2. One independent E8M0 scale per 1x32 block per output half. + scol = bx * cutlass.Int32(TXR) + tx + u_act = _fold_amax(am_act) << 16 + e_act = _float_to_e8m0(u_act) + sidx = _scale_idx(grow, scol, rs_ncb, rs_stride, SWIZ) + mRS[sidx] = e_act.to(cutlass.Uint8) + r_act = _exp2f_rcp_bf16(e_act) * cutlass.Int32(0x10001) + if cutlass.const_expr(IS_BWD): + u_gate = _fold_amax(am_gate) << 16 + e_gate = _float_to_e8m0(u_gate) + gidx = _scale_idx(grow, scol + rgate_scol_off, rs_ncb, rs_stride, SWIZ) + mRS[gidx] = e_gate.to(cutlass.Uint8) + r_gate = _exp2f_rcp_bf16(e_gate) * cutlass.Int32(0x10001) + + # 3. Scale and pack, storing with the same swizzled traversal. + for w in cutlass.range_constexpr(WAVES): + grp = tx * cutlass.Int32(WAVES) + ((cutlass.Int32(w) + bank_group) & 7) + sRowActw[(row, grp, buff)] = _mul_cvt_2x( + iv_act[w][0], iv_act[w][1], r_act + ) + if cutlass.const_expr(IS_BWD): + sRowGatew[(row, grp, buff)] = _mul_cvt_2x( + iv_gate[w][0], iv_gate[w][1], r_gate + ) + + # Make shared-memory writes visible to the TMA engine, then issue the + # TMA stores under warp 0's warp-uniform predicate (the DSL elects the + # issuing lane) and commit them as one bulk group from the same warp. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + if warp_idx == 0: + if cutlass.const_expr(ROWWISE): + cute.copy( + atom_row_act, tRAs[(None, buff)], tRAg[(None, (row_tile, bx))] + ) + if cutlass.const_expr(IS_BWD): + cute.copy( + atom_row_gate, tRGs[(None, buff)], tRGg[(None, (row_tile, bx))] + ) + if cutlass.const_expr(COLWISE): + cute.copy( + atom_col_act, tCAs[(None, buff)], tCAg[(None, (bx, row_tile))] + ) + if cutlass.const_expr(IS_BWD): + cute.copy( + atom_col_gate, tCGs[(None, buff)], tCGg[(None, (bx, row_tile))] + ) + cute.arch.cp_async_bulk_commit_group() + + # Drain every outstanding TMA-store group before the CTA retires: bulk + # async groups are not implicitly awaited at exit, and the smem tiles the + # in-flight stores read are deallocated with the CTA (a successor CTA may + # reuse them). Warp 0 committed every group, so it alone waits. + if warp_idx == 0: + cute.arch.cp_async_bulk_wait_group(0, read=True) + + if cutlass.const_expr(not DIRECT): + cute.arch.sync_threads() + if tidx == 0: + # ``mbarrier.inval`` is not exposed by the DSL; emit it as raw PTX. + for i in cutlass.range_constexpr(N_STAGES): + llvm.inline_asm( + None, + [Int32((mbar + i).toint()).ir_value()], + "mbarrier.inval.shared::cta.b64 [$0];", + "r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def launcher( + ag: cutlass.Int64, + agi: cutlass.Int64, + arq: cutlass.Int64, + ars: cutlass.Int64, + acq: cutlass.Int64, + acs: cutlass.Int64, + m: cutlass.Int32, + k: cutlass.Int32, + stream, + IS_BWD: cutlass.Constexpr, + ROWWISE: cutlass.Constexpr, + COLWISE: cutlass.Constexpr, + SWIZ: cutlass.Constexpr, + ACT_PAIR: cutlass.Constexpr, + DIRECT: cutlass.Constexpr, + CX: cutlass.Constexpr, + CY: cutlass.Constexpr, +): + """Build the TMA views and launch one (CX, CY)-chunk grid. + + ``x``/``lin`` are the two halves of the packed [M, 2K] input; ``grad`` + is the [M, K] incoming gradient. Rowwise outputs land in the row-major + [M, out_k] tensor at column offsets 0 and K; colwise outputs in the + transposed [out_k, M] storage at row offsets 0 and K. + + Pointer arguments: ``ag`` = grad_h (backward only, else 0); ``agi`` = + gated_input; ``arq``/``ars`` = rowwise quantized-output/scale; + ``acq``/``acs`` = the colwise pair. Disabled directions pass 0. + """ + OUT_HALVES = 2 if cutlass.const_expr(IS_BWD) else 1 + # Direct: two threads per 1x32 block; staged: one thread per column, + # except colwise-only which stacks two thread rows per column. + if cutlass.const_expr(DIRECT): + THREADS = 2 * CX + else: + THREADS = (2 * CX) if cutlass.const_expr(COLWISE and not ROWWISE) else CX + out_k = OUT_HALVES * k + + # The TMA smem layouts describe one buffer, not the full multi-buffer + # allocation. + in_smem = cute.make_layout((BUFF_DIM_Y, CX), stride=(CX, 1)) + in_tiler = (BUFF_DIM_Y, CX) + col_smem = cute.make_layout((CX, BUFF_DIM_Y), stride=(BUFF_DIM_Y, 1)) + col_tiler = (CX, BUFF_DIM_Y) + + px = cute.make_ptr(cutlass.BFloat16, agi, AddressSpace.gmem, assumed_align=16) + plin = px + k + mX = cute.make_tensor(px, cute.make_layout((m, k), stride=(2 * k, 1))) + mLin = cute.make_tensor(plin, cute.make_layout((m, k), stride=(2 * k, 1))) + atom_x, tma_x = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mX, in_smem, in_tiler + ) + atom_lin, tma_lin = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mLin, in_smem, in_tiler + ) + gX = cute.zipped_divide(tma_x, in_tiler) + gLin = cute.zipped_divide(tma_lin, in_tiler) + + if cutlass.const_expr(IS_BWD): + pgrad = cute.make_ptr(cutlass.BFloat16, ag, AddressSpace.gmem, assumed_align=16) + mGrad = cute.make_tensor(pgrad, cute.make_layout((m, k), stride=(k, 1))) + atom_grad, tma_grad = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mGrad, in_smem, in_tiler + ) + gGrad = cute.zipped_divide(tma_grad, in_tiler) + else: + atom_grad, gGrad = atom_x, gX + + # b32 word views for the direct path: (word, block half, 1x32 scale + # block, x-tile, row). The runtime strides k and k//2 would collapse the + # sliced pointer's provable alignment to one word and narrow autovec_copy + # to 32-bit accesses; K % 128 == 0 makes both multiples of 8 words, and + # cute.assume encodes that so the 8-word copies stay 256-bit loads. + kw = cute.assume(k, divby=8) # gated row stride: 2k bf16 = k words + khw = cute.assume(k // 2, divby=8) # Lin base offset / grad row stride + pxw = cute.make_ptr(cutlass.Int32, agi, AddressSpace.gmem, assumed_align=32) + word_layout = cute.make_layout( + (8, 2, CX // SCALE_DIM_X, k // CX, m), stride=(1, 8, 16, CX // 2, kw) + ) + gXv = cute.make_tensor(pxw, word_layout) + gLinv = cute.make_tensor(pxw + khw, word_layout) + if cutlass.const_expr(IS_BWD): + pgw = cute.make_ptr(cutlass.Int32, ag, AddressSpace.gmem, assumed_align=32) + gGradv = cute.make_tensor( + pgw, + cute.make_layout( + (8, 2, CX // SCALE_DIM_X, k // CX, m), + stride=(1, 8, 16, CX // 2, khw), + ), + ) + else: + gGradv = gXv + + if cutlass.const_expr(ROWWISE): + pra = cute.make_ptr( + cutlass.Float8E4M3FN, arq, AddressSpace.gmem, assumed_align=16 + ) + mRowAct = cute.make_tensor(pra, cute.make_layout((m, k), stride=(out_k, 1))) + atom_row_act, tma_ra = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mRowAct, in_smem, in_tiler + ) + gRowAct = cute.zipped_divide(tma_ra, in_tiler) + if cutlass.const_expr(IS_BWD): + prg = pra + k + mRowGate = cute.make_tensor( + prg, cute.make_layout((m, k), stride=(out_k, 1)) + ) + atom_row_gate, tma_rg = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mRowGate, in_smem, in_tiler + ) + gRowGate = cute.zipped_divide(tma_rg, in_tiler) + else: + atom_row_gate, gRowGate = atom_row_act, gRowAct + else: + atom_row_act, gRowAct = atom_x, gX + atom_row_gate, gRowGate = atom_x, gX + + if cutlass.const_expr(COLWISE): + pca = cute.make_ptr( + cutlass.Float8E4M3FN, acq, AddressSpace.gmem, assumed_align=16 + ) + mColAct = cute.make_tensor(pca, cute.make_layout((k, m), stride=(m, 1))) + atom_col_act, tma_ca = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mColAct, col_smem, col_tiler + ) + gColAct = cute.zipped_divide(tma_ca, col_tiler) + if cutlass.const_expr(IS_BWD): + pcg = pca + k * m + mColGate = cute.make_tensor(pcg, cute.make_layout((k, m), stride=(m, 1))) + atom_col_gate, tma_cg = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mColGate, col_smem, col_tiler + ) + gColGate = cute.zipped_divide(tma_cg, col_tiler) + else: + atom_col_gate, gColGate = atom_col_act, gColAct + else: + atom_col_act, gColAct = atom_x, gX + atom_col_gate, gColGate = atom_x, gX + + prs = cute.make_ptr(cutlass.Uint8, ars, AddressSpace.gmem, assumed_align=16) + mRS = cute.make_tensor(prs, cute.make_layout(m * (out_k // 32))) + pcs = cute.make_ptr(cutlass.Uint8, acs, AddressSpace.gmem, assumed_align=16) + mCS = cute.make_tensor(pcs, cute.make_layout(out_k * (m // 32))) + + gated_act_mxfp8_kernel( + atom_x, + gX, + atom_lin, + gLin, + atom_grad, + gGrad, + gXv, + gLinv, + gGradv, + atom_row_act, + gRowAct, + atom_row_gate, + gRowGate, + atom_col_act, + gColAct, + atom_col_gate, + gColGate, + mRS, + mCS, + out_k // 128, # rs_ncb: rowwise 128x4 scale-column blocks + out_k // 32, # rs_stride: rowwise compact-scale row stride + k // 32, # rgate_scol_off: dUp-half rowwise scale-column offset + m // 128, # cs_ncb: colwise 128x4 scale-column blocks + m // 32, # cs_stride: colwise compact-scale row stride + k, # cgate_col_off: dUp-half colwise output-row offset + IS_BWD, + ROWWISE, + COLWISE, + SWIZ, + ACT_PAIR, + DIRECT, + CX, + CY, + THREADS, + ).launch( + grid=(k // CX, m // CY, 1), + block=(THREADS, 1, 1), + stream=stream, + ) + + +@functools.cache +def _compile_kernel( + is_bwd, rowwise, colwise, swizzled_scales, act_pair, direct, cx, cy, device_index +): + """Compile and cache one kernel specialization. ``act_pair`` must be a + module-level function so the cache key stays stable; ``device_index`` is + part of the cache key only (compilation targets the active device).""" + cap = torch.cuda.get_device_capability() + if cap[0] != 10: + raise NotImplementedError( + f"gated_act_mxfp8 requires CUDA SM 10.x (Blackwell); " + f"found sm_{cap[0]}{cap[1]}" + ) + del device_index + from cutlass.cute.runtime import make_fake_stream + + null = cutlass.Int64(0) + dim = cutlass.Int32(128) + return cute.compile( + launcher, + null, + null, + null, + null, + null, + null, + dim, + dim, + make_fake_stream(), + is_bwd, + rowwise, + colwise, + swizzled_scales, + act_pair, + direct, + cx, + cy, + ) + + +def _validate_inputs(gated_input, grad_h=None): + if not gated_input.is_cuda: + raise ValueError("gated_input must be a CUDA tensor") + if gated_input.dtype != torch.bfloat16: + raise TypeError("gated_input must have dtype torch.bfloat16") + if gated_input.ndim != 2 or not gated_input.is_contiguous(): + raise ValueError("gated_input must be contiguous with shape [M, 2K]") + M, two_k = gated_input.shape + if two_k % 2: + raise ValueError("gated_input.shape[1] must be even") + K = two_k // 2 + # Keeps CTA chunks whole and the blocked scale layout padding-free, so + # every element of the scale tensors is written by the kernel. Zero-size + # inputs satisfy every modulus but cannot form a launch grid or a TMA + # descriptor, so they are rejected here instead of failing opaquely. + if M == 0 or K == 0 or M % 128 or K % 128: + raise ValueError("M and K must be nonzero multiples of 128") + # Index arithmetic and scale layouts assume 32-bit offsets; the largest + # offset any layout reaches is 2*K*M - K - 1 elements. + if 2 * K * M - K - 1 > _INT32_MAX: + raise ValueError( + f"M={M}, K={K} exceeds the kernel's 32-bit indexing limit " + f"(needs 2*K*M - K - 1 <= {_INT32_MAX})" + ) + # The launcher passes raw device pointers promised as assumed_align=32 + # (b32 word views backing 256-bit loads) and assumed_align=16 (TMA + # descriptors); contiguity does not imply base alignment for + # storage-offset views. Checked after the shape/int32 gates so + # FakeTensor probes (no data_ptr) exercise those first. + if gated_input.data_ptr() % 32: + raise ValueError( + "gated_input must be 32-byte aligned (data_ptr() % 32 == 0); " + "storage-offset views are not -- pass a fresh copy, e.g. .clone()" + ) + if grad_h is not None: + if ( + not grad_h.is_cuda + or grad_h.dtype != torch.bfloat16 + or not grad_h.is_contiguous() + or tuple(grad_h.shape) != (M, K) + ): + raise ValueError("grad_h must be contiguous BF16 CUDA [M, K]") + if grad_h.device != gated_input.device: + raise ValueError( + f"grad_h is on {grad_h.device} but gated_input is on " + f"{gated_input.device}; both must be on the same CUDA device" + ) + if grad_h.data_ptr() % 32: + raise ValueError( + "grad_h must be 32-byte aligned (data_ptr() % 32 == 0); " + "storage-offset views are not -- pass a fresh copy, e.g. " + ".clone()" + ) + return M, K + + +def _ptr(tensor): + return 0 if tensor is None else tensor.data_ptr() + + +@torch.no_grad() +def _launch_gated_act_mxfp8( + gated_input, grad_h, outputs, rowwise, colwise, geometry=None +): + """Validate, compile the matching specialization, and launch into + ``outputs`` = ``(output_rowwise, output_colwise, scales_rowwise, + scales_colwise)``, caller-allocated. Disabled directions are zero-sized + and not written; scales are always in the blocked (GEMM-swizzled + tcgen05) layout. ``geometry`` overrides the per-mode default + ``(CX, CY, direct)`` 3-tuple (tuning/testing only; staged rowwise + requires the colwise producer — see the kernel's trace-time asserts). + """ + if not (rowwise or colwise): + raise ValueError("at least one of rowwise/colwise must be enabled") + M, K = _validate_inputs(gated_input, grad_h) + output_rowwise, output_colwise, scales_rowwise, scales_colwise = outputs + for out, enabled, name in ( + (output_rowwise, rowwise, "output_rowwise"), + (output_colwise, colwise, "output_colwise"), + ): + if enabled and out.dtype != torch.float8_e4m3fn: + raise TypeError(f"{name} must have dtype torch.float8_e4m3fn") + + # Compile and launch under the input's device: a caller holding cuda:0 + # current while passing a cuda:1 tensor must not launch foreign pointers. + with torch.cuda.device(gated_input.device): + # Wrap per call; caching CUstream handles could alias recycled streams. + stream = CUstream(torch.cuda.current_stream(gated_input.device).cuda_stream) + geom = geometry or _DEFAULT_GEOMETRY[(grad_h is not None, rowwise, colwise)] + cx, cy, direct = geom + # CX feeds bit-mask/shift thread mapping (tidx & (CX-1), >> LOG2_CX) + # and exact-divide layouts; CY feeds the floor grid (k//CX, m//CY) and + # N_STAGES = CY // BUFF_DIM_Y — a violating override silently skips + # columns/rows instead of failing. + if cx & (cx - 1) or cx % SCALE_DIM_X or K % cx: + raise ValueError( + f"geometry CX={cx} must be a power of two, a multiple of " + f"{SCALE_DIM_X}, and divide K={K}" + ) + if cy % BUFF_DIM_Y or M % cy: + raise ValueError( + f"geometry CY={cy} must be a multiple of {BUFF_DIM_Y} and divide M={M}" + ) + # Row chunks ride CUDA grid dim y, which caps at 65535 on every + # compute capability; the int32 element bound alone admits small-K + # shapes past it. + if M // cy > 65535: + raise ValueError( + f"M={M} with geometry CY={cy} needs {M // cy} row chunks, " + "over CUDA's 65535 grid y-dimension limit" + ) + # The staged path drains TMA-store groups only after its stage loop, + # so output smem buffers must never be reused inside it: the stage + # count (CY / 32) is capped at the double buffer's depth. + if not direct and cy // BUFF_DIM_Y > BUFFS_NUM: + raise ValueError( + f"geometry CY={cy} needs {cy // BUFF_DIM_Y} pipeline stages; " + f"the staged path supports at most {BUFFS_NUM} (output smem " + "is double-buffered with no in-loop store drain)" + ) + fn = _compile_kernel( + grad_h is not None, + rowwise, + colwise, + True, # blocked/GEMM-swizzled scales; compact is compile-time only + _silu_pair, + direct, + cx, + cy, + gated_input.device.index, + ) + fn( + _ptr(grad_h), + gated_input.data_ptr(), + _ptr(output_rowwise) if rowwise else 0, + _ptr(scales_rowwise) if rowwise else 0, + _ptr(output_colwise) if colwise else 0, + _ptr(scales_colwise) if colwise else 0, + M, + K, + stream, + ) + + +def _gated_act_mxfp8_outputs( + gated_input: torch.Tensor, + out_k: int, + rowwise: bool, + colwise: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Allocate the fixed four outputs. + + Torch-only, so it serves both the real op and its fake: on a meta input it + returns meta tensors with the shapes and strides the kernel writes. + """ + if not (rowwise or colwise): + raise ValueError("at least one of rowwise/colwise must be enabled") + m = gated_input.shape[0] + empty_qdata = gated_input.new_empty(0, dtype=torch.float8_e4m3fn) + empty_scales = gated_input.new_empty(0, dtype=torch.float8_e8m0fnu) + + if rowwise: + output_rowwise = torch.empty_strided( + (m, out_k), + (out_k, 1), + device=gated_input.device, + dtype=torch.float8_e4m3fn, + ) + scales_rowwise = gated_input.new_empty( + (ceil_div(m, 128) * 128, ceil_div(out_k // 32, 4) * 4), + dtype=torch.float8_e8m0fnu, + ) + else: + output_rowwise, scales_rowwise = empty_qdata, empty_scales + + if colwise: + output_colwise = torch.empty_strided( + (m, out_k), + (1, m), + device=gated_input.device, + dtype=torch.float8_e4m3fn, + ) + # Flat 1D, matching mxfp8_quantize_2d_32x1_cutedsl. + scales_colwise = gated_input.new_empty( + ((ceil_div(out_k, 128) * 128) * (ceil_div(m // 32, 4) * 4),), + dtype=torch.float8_e8m0fnu, + ) + else: + output_colwise, scales_colwise = empty_qdata, empty_scales + + return output_rowwise, output_colwise, scales_rowwise, scales_colwise + + +@torch.library.custom_op("torchao::gated_act_mxfp8_cutedsl_forward", mutates_args=()) +def _gated_act_mxfp8_cutedsl_forward( + gated_input: torch.Tensor, + rowwise: bool = True, + colwise: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + outputs = _gated_act_mxfp8_outputs( + gated_input, gated_input.shape[1] // 2, rowwise, colwise + ) + _launch_gated_act_mxfp8(gated_input, None, outputs, rowwise, colwise) + return outputs + + +@torch.library.custom_op("torchao::gated_act_mxfp8_cutedsl_backward", mutates_args=()) +def _gated_act_mxfp8_cutedsl_backward( + grad_h: torch.Tensor, + gated_input: torch.Tensor, + rowwise: bool = True, + colwise: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + outputs = _gated_act_mxfp8_outputs( + gated_input, gated_input.shape[1], rowwise, colwise + ) + _launch_gated_act_mxfp8(gated_input, grad_h, outputs, rowwise, colwise) + return outputs + + +@_gated_act_mxfp8_cutedsl_forward.register_fake +def _fake_gated_act_mxfp8_cutedsl_forward( + gated_input: torch.Tensor, + rowwise: bool = True, + colwise: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + assert gated_input.ndim == 2, "gated_input must be 2D" + return _gated_act_mxfp8_outputs( + gated_input, gated_input.shape[1] // 2, rowwise, colwise + ) + + +@_gated_act_mxfp8_cutedsl_backward.register_fake +def _fake_gated_act_mxfp8_cutedsl_backward( + grad_h: torch.Tensor, + gated_input: torch.Tensor, + rowwise: bool = True, + colwise: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + assert grad_h.ndim == 2, "grad_h must be 2D" + assert gated_input.ndim == 2, "gated_input must be 2D" + return _gated_act_mxfp8_outputs(gated_input, gated_input.shape[1], rowwise, colwise) + + +def gated_act_mxfp8_cutedsl_forward( + gated_input: torch.Tensor, + *, + rowwise: bool = True, + colwise: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Fuse the gated activation (SwiGLU) forward and its RCEIL MXFP8 cast into + one pass on SM100: ``h = silu(gate) * up`` is quantized without ever + being written to global memory. + + Args: + gated_input: BF16 tensor of shape (M, 2K), ``gate`` in the first K + columns and ``up`` in the last K; M and K multiples of 128. + rowwise: emit 1x32-scaled, row-major output. + colwise: emit 32x1-scaled, column-major (stride ``(1, M)``) output. + + Returns: + ``(output_rowwise, output_colwise, scales_rowwise, scales_colwise)``, + width K, E8M0 scales in the same blocked tcgen05 layouts as the + standalone quantizers. Disabled directions return zero-sized tensors, + so the output arity never varies. Special values follow the standalone + quantizers' contract: a NaN/Inf block amax yields scale byte 0xFF and + all-NaN output codes; zero/tiny amaxes clamp to scale byte 0x00. + """ + # Read the shared availability flag at call time (tests monkeypatch it). + from torchao.prototype.moe_training.kernels.mxfp8 import quant as _quant + + if not _quant._mxfp8_cutedsl_kernels_available: + missing_packages = _missing_cutedsl_runtime_packages() + if missing_packages: + missing = ", ".join(missing_packages) + raise NotImplementedError( + "gated_act_mxfp8_cutedsl_forward requires additional Python " + f"runtime package(s): {missing}. Please install " + "`nvidia-cutlass-dsl` and `apache-tvm-ffi`." + ) + raise NotImplementedError( + "gated_act_mxfp8_cutedsl_forward requires CUDA, SM 10.x, and CUDA 12.8+." + ) + return _gated_act_mxfp8_cutedsl_forward(gated_input, rowwise, colwise) + + +def gated_act_mxfp8_cutedsl_backward( + grad_h: torch.Tensor, + gated_input: torch.Tensor, + *, + rowwise: bool = True, + colwise: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Fuse the gated activation (SwiGLU) backward and its RCEIL MXFP8 cast into + one pass on SM100: quantizes the concatenated ``[dGate | dUp]`` tensor a + fused w13 weight expects for the wgrad GEMM. + + Args: + grad_h: BF16 gradient of shape (M, K). + gated_input: the forward input, shape (M, 2K), as in + :func:`gated_act_mxfp8_cutedsl_forward`. + rowwise: emit 1x32-scaled, row-major output. + colwise: emit 32x1-scaled, column-major (stride ``(1, M)``) output. + + Returns: + Four tensors of width 2K, same order and layouts as + :func:`gated_act_mxfp8_cutedsl_forward`. + """ + from torchao.prototype.moe_training.kernels.mxfp8 import quant as _quant + + if not _quant._mxfp8_cutedsl_kernels_available: + missing_packages = _missing_cutedsl_runtime_packages() + if missing_packages: + missing = ", ".join(missing_packages) + raise NotImplementedError( + "gated_act_mxfp8_cutedsl_backward requires additional Python " + f"runtime package(s): {missing}. Please install " + "`nvidia-cutlass-dsl` and `apache-tvm-ffi`." + ) + raise NotImplementedError( + "gated_act_mxfp8_cutedsl_backward requires CUDA, SM 10.x, and CUDA 12.8+." + ) + return _gated_act_mxfp8_cutedsl_backward(grad_h, gated_input, rowwise, colwise)