From f8a1ac98439b1783f0e1ab4964d9247901db4c10 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 16 Aug 2026 21:12:47 -0700 Subject: [PATCH 01/11] Add MXFP8 grouped-MLP kernel family foundation and the wgrad kernel Adds the shared infrastructure for three fused MXFP8 kernels on the routed expert path (SM100a), plus the first complete kernel. Shared: grouped_mlp_validation.py host preconditions at the custom-op boundary grouped_mlp_epilogue.py device primitives: RCEIL E8M0, NaN-propagating packed amax, tcgen05 blocked-scale indexing, SwiGLU/dSwiGLU policy grouped_mlp_ops.py the three custom ops and their fake/meta impls grouped_gemm_config.py frozen configs, support predicate, the published T2R and epilogue protocols grouped_gemm_core.py ragged blockscaled GEMM: TMA descriptors, tcgen05 mainloop, TMEM, accumulator-to-register handoff epilogue_quant.py rowwise 1x32 and columnwise 32x1 quantization Kernel C (mxfp8_grouped_gemm_wgrad): kernel_wgrad.py BF16 store epilogue and cached launcher cutedsl_grouped_mlp.py launcher facade (A and B land here next) Per-expert row counts are multiples of 128, so no tile straddles an expert boundary. That removes the per-group tensormap updates and the persistent tile scheduler the general grouped-GEMM path needs: an expert is selected by an integer tile-coordinate base. The blockscaled scale layout's K stride was measured to be a constant 512 bytes for every row block, which is what makes that indexing correct. Columnwise blocked scales use whole-matrix to_blocked rather than the per-group K-groups form. The two encodings have identical byte counts and coincide when N <= 128, so the difference is invisible to any length check and to small-shape tests; mixing them corrupts the weight gradient (cosine 0.82). Validation: the quantize chain is bitwise-identical to torchao's RCEIL quantizers over 23.7M qdata bytes and all special values; the GEMM core is bitwise exact on both ragged orientations including zero-token experts, strict inactive tails and the pipeline stage wrap; Kernel C is bitwise exact against a float64 oracle on a production shape (23,068,672 elements) and clean at the A/B -> C seam. Note: cute.testing.assert_ is compiled out unless CUTE_DSL_ENABLE_ASSERTIONS=1, so device-side offset checks are a debugging aid only and the host validation is the sole enforcement of the alignment precondition. Co-Authored-By: Claude Opus 5 (1M context) --- .../kernels/mxfp8/cutedsl_grouped_mlp.py | 23 + .../kernels/mxfp8/epilogue_quant.py | 414 +++++++++ .../kernels/mxfp8/grouped_gemm_config.py | 463 ++++++++++ .../kernels/mxfp8/grouped_gemm_core.py | 825 ++++++++++++++++++ .../kernels/mxfp8/grouped_mlp_epilogue.py | 448 ++++++++++ .../kernels/mxfp8/grouped_mlp_ops.py | 451 ++++++++++ .../kernels/mxfp8/grouped_mlp_validation.py | 245 ++++++ .../kernels/mxfp8/kernel_wgrad.py | 347 ++++++++ 8 files changed, 3216 insertions(+) create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py new file mode 100644 index 0000000000..6036860cf9 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py @@ -0,0 +1,23 @@ +# 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. + +"""Launcher facade for the MXFP8 routed-expert grouped-MLP CuTe DSL kernels. + +``grouped_mlp_ops`` imports its three launchers from this module and from +nowhere else, so this name and these three signatures are the seam between the +custom-op layer and the kernels. Keep it a pure re-export: anything defined here +would be code the ops layer depends on but that no kernel test covers. + +The import is deliberately per-launcher rather than a package-level ``*``, so a +kernel that has not landed yet costs an ImportError naming that kernel instead of +breaking the two that have. +""" + +from torchao.prototype.moe_training.kernels.mxfp8.kernel_wgrad import ( + launch_grouped_gemm_wgrad, +) + +__all__ = ["launch_grouped_gemm_wgrad"] diff --git a/torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py b/torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py new file mode 100644 index 0000000000..f750ec6b7e --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py @@ -0,0 +1,414 @@ +# 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. + +"""Rowwise 1x32 and columnwise 32x1 quantizing epilogue for the MXFP8 grouped-MLP kernels. + +Direction-agnostic and GEMM-agnostic: everything here consumes a run of packed +``bf16x2`` words held **one output row per thread**, which is exactly what the +tcgen05 TMEM->register copy delivers (``Ld32x32bOp`` has ``ThrID 32:1`` and a +destination TV layout ``(32, EPI_N_ACC):(EPI_N_ACC, 1)``, so thread ``t`` owns +row ``t`` of the epilogue subtile and all of its columns, contiguously). Two +consequences are load-bearing and are asserted by the geometry below: + +* the rowwise 1x32 amax runs along N inside a single thread -- no cross-lane op; +* the columnwise 32x1 amax runs along M inside a single warp, and a 32-row MX + block is never split across warps or CTAs, because the CTA tile is 128 rows, + ``32 | 128``, and every expert boundary is 128-aligned. + +WORD CONTRACT (every entry point below assumes it): word ``j`` of ``words`` +holds output column ``col_base + 2*j`` in its **low** bf16 half and column +``col_base + 2*j + 1`` in its **high** half, and 16 words are exactly one +32-value scale block. Callers must have rounded to BF16 already -- both because +the kernel contract defines correctness at that boundary and because +:func:`float_to_e8m0` is exact only for a BF16-valued amax. + +The columnwise destination has stride ``(1, R)``, i.e. it is physically a +row-major ``[cols, R]`` buffer, so the epilogue must transpose. It does that +through ``sPad``, a shared-memory staging tile of **BF16** values (not quantized +bytes): the reader's 16-byte load has to come out as four consecutive rows of +one column, which ``mul_cvt_2x`` + ``prmt_even``/``prmt_odd`` produce from two +words x two columns; staging bytes instead would force a stride-2 byte gather. + +Both columnwise scale orientations use whole-matrix ``to_blocked`` coordinates +(feature index as the blocked row, row-block index as the blocked column). That +is this family's choice and it deliberately differs from torchao's per-group +``triton_mx_block_rearrange_2d_K_groups``. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.utils import SmemAllocator + +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_epilogue import ( + SCALE_BLOCK, + abs_max_nan_bf16x2, + blocked_scale_idx, + e8m0_reciprocal_bf16, + float_to_e8m0, + fold_amax, + max_nan_bf16x2, + mul_cvt_2x, + prmt_even, + prmt_odd, +) + +__all__ = [ + "WORDS_PER_SCALE_BLOCK", + "COLWISE_PAIRS", + "COLWISE_ROWS_PER_WARP", + "COLWISE_PAIR_PITCH", + "COLWISE_WORDS_PER_WARP", + "EPILOGUE_ROWS_PER_WARP", + "spad_words", + "spad_bytes", + "alloc_spad", + "rowwise_quant_block", + "rowwise_quant_store", + "rowwise_scale_flush", + "colwise_quant_store", +] + +# 16 bf16x2 words == 32 values == one E8M0 scale block. +WORDS_PER_SCALE_BLOCK = SCALE_BLOCK // 2 +# One columnwise staging chunk is 32 output columns == 16 column pairs. +COLWISE_PAIRS = WORDS_PER_SCALE_BLOCK +# Rows one epilogue warp owns; equal to the 32x1 block height, which is what +# keeps the whole columnwise round trip inside a warp. +COLWISE_ROWS_PER_WARP = SCALE_BLOCK +EPILOGUE_ROWS_PER_WARP = COLWISE_ROWS_PER_WARP +# Word pitch between two column-pair slabs. The +4 pad is the entire swizzle: +# the reader's 128-bit shared loads are serviced in four phases of eight lanes, +# and phase 0 covers pairs 0..7 at word `36*pair + 4*chunk`, i.e. banks +# `4*pair + 4*chunk (mod 32)` -- eight distinct 4-word groups, all 32 banks, no +# conflict. At pitch 32 every one of those eight lanes lands on the same bank +# group and the load degenerates to an 8-way conflict. +COLWISE_PAIR_PITCH = COLWISE_ROWS_PER_WARP + 4 +COLWISE_WORDS_PER_WARP = COLWISE_PAIRS * COLWISE_PAIR_PITCH +# Every sPad access is a 16-byte vector; 144 * pair + 64 * tq + 16 * chunk and +# the 2304-byte per-warp stride are all multiples of 16, so this holds for the +# allocation too. +_VEC_BYTES = 16 +_VEC_WORDS = _VEC_BYTES // 4 + + +def spad_words(num_epilogue_warps: int = 4) -> int: + """Int32 word count of the columnwise staging tile. + + Sized by the warp count rather than by a config object so that this module + stays importable, and testable, without the GEMM core. Callers holding a + ``GroupedGemmConfig`` pass ``len(config.epilogue_warp_ids)``. + """ + return num_epilogue_warps * COLWISE_WORDS_PER_WARP + + +def spad_bytes(num_epilogue_warps: int = 4) -> int: + """Shared-memory bytes the columnwise transpose costs (9216 B at 4 warps). + + That is 0.27 of one 128x128x128 E4M3 AB pipeline stage, i.e. it does not + change the achievable stage count. Feed it to + ``config.smem_bytes(epilogue_smem_bytes=...)``. + """ + return 4 * spad_words(num_epilogue_warps) # 4 bytes per Int32 word + + +@cute.jit +def alloc_spad( + allocator: SmemAllocator, NUM_EPILOGUE_WARPS: cutlass.Constexpr = 4 +) -> cute.Tensor: + """Allocate the columnwise staging tile as a flat Int32 shared tensor.""" + return allocator.allocate_tensor( + Int32, + cute.make_layout(spad_words(NUM_EPILOGUE_WARPS)), + byte_alignment=_VEC_BYTES, + ) + + +@cute.jit +def _store_vec_i32(dst: cute.Tensor, word_offset: Int32, src: cute.Tensor): + """Vectorized store of a register word run into a flat byte destination. + + `dst` is any 1-byte-element flat view; the run is reinterpreted as Int32, so + `word_offset` counts 4-byte words from the start of the buffer. + """ + cute.autovec_copy( + src, + cute.make_tensor( + (cute.recast_ptr(dst.iterator, dtype=Int32) + word_offset).align( + _VEC_BYTES + ), + cute.make_layout(cute.size(src)), + ), + ) + + +@cute.jit +def _load_spad_quad(spad: cute.Tensor, word_offset: Int32) -> cute.Tensor: + """One 16-byte shared load: four consecutive rows of one column pair.""" + quad = cute.make_rmem_tensor((_VEC_WORDS,), Int32) + cute.autovec_copy( + cute.make_tensor( + (spad.iterator + word_offset).align(_VEC_BYTES), + cute.make_layout(_VEC_WORDS), + ), + quad, + ) + return quad + + +@cute.jit +def rowwise_quant_block(words: cute.Tensor, WORD_BASE: cutlass.Constexpr = 0): + """Quantize one 1x32 rowwise block held entirely in one thread's registers. + + Returns ``(qwords, scale_byte)``: eight Int32 words holding the 32 E4M3 + bytes in column order, and the E8M0 exponent byte for the block. + + The amax reduction is intra-thread by construction (see the module + docstring), so there is no shuffle and no partial amax carried between + epilogue subtiles. + """ + amax_packed = words[WORD_BASE] + for w in cutlass.range_constexpr(1, WORDS_PER_SCALE_BLOCK): + amax_packed = abs_max_nan_bf16x2(amax_packed, words[WORD_BASE + w]) + # fold_amax masks the junk sign bits xorsign leaves behind, then folds the + # two bf16 lanes with the NaN-propagating max. + scale_byte = float_to_e8m0(fold_amax(amax_packed) << Int32(16)) + inv = e8m0_reciprocal_bf16(scale_byte) + inv_packed = inv | (inv << Int32(16)) + + qwords = cute.make_rmem_tensor((WORDS_PER_SCALE_BLOCK // 2,), Int32) + for q in cutlass.range_constexpr(WORDS_PER_SCALE_BLOCK // 2): + # mul_cvt_2x emits bytes [w0.lo, w0.hi, w1.lo, w1.hi], i.e. four + # consecutive columns in increasing address order. + qwords[q] = mul_cvt_2x( + words[WORD_BASE + 2 * q], words[WORD_BASE + 2 * q + 1], inv_packed + ) + return qwords, scale_byte + + +@cute.jit +def rowwise_quant_store( + words: cute.Tensor, + qdata: cute.Tensor, + row: Int32, + col: Int32, + row_stride: Int32, + WORD_BASE: cutlass.Constexpr = 0, +) -> Int32: + """Quantize one 1x32 block and store its 32 E4M3 bytes; return the scale byte. + + `qdata` is a flat 1-byte-element view of a row-major destination whose row + pitch is `row_stride`. `row_stride` is a multiple of 128 and `col` a + multiple of 32 under the kernel contract, so the destination byte offset is + 32-byte aligned and the 32 bytes leave as two STG.128. + + The scale byte is returned rather than stored: consecutive subtiles produce + consecutive blocked-scale columns, so the caller buffers them and flushes + once per CTA tile through :func:`rowwise_scale_flush`. + """ + qwords, scale_byte = rowwise_quant_block(words, WORD_BASE=WORD_BASE) + _store_vec_i32(qdata, (row * row_stride + col) >> Int32(2), qwords) + return scale_byte + + +@cute.jit +def rowwise_scale_flush( + scales: cute.Tensor, + row: Int32, + scale_col_base: Int32, + scale_bytes: cute.Tensor, + num_scale_col_blocks: Int32, + NUM_BYTES: cutlass.Constexpr, +): + """Write `NUM_BYTES` consecutive rowwise blocked-scale bytes for one row. + + Scale columns that share ``scale_col >> 2`` are contiguous bytes in the + tcgen05 blocked layout (they differ only in the low two bits of the flat + index), so a run of 4 starting at a 4-aligned `scale_col_base` is one 4-byte + store. `scale_col_base` is `tile_n * NUM_BYTES` in both kernels, hence + aligned to whichever width is selected here. + + `scales` must be a flat uint8 view of the blocked buffer. + """ + if cutlass.const_expr(NUM_BYTES % 4 == 0): + _flush_scale_run( + scales, row, scale_col_base, scale_bytes, num_scale_col_blocks, 4, NUM_BYTES + ) + elif cutlass.const_expr(NUM_BYTES % 2 == 0): + _flush_scale_run( + scales, row, scale_col_base, scale_bytes, num_scale_col_blocks, 2, NUM_BYTES + ) + else: + for i in cutlass.range_constexpr(NUM_BYTES): + idx = blocked_scale_idx(row, scale_col_base + i, num_scale_col_blocks) + scales[idx] = cutlass.Uint8(scale_bytes[i]) + + +@cute.jit +def _flush_scale_run( + scales: cute.Tensor, + row: Int32, + scale_col_base: Int32, + scale_bytes: cute.Tensor, + num_scale_col_blocks: Int32, + WIDTH: cutlass.Constexpr, + NUM_BYTES: cutlass.Constexpr, +): + """Emit `NUM_BYTES` scale bytes as `NUM_BYTES // WIDTH` packed stores. + + A packed store addresses `scales` in units of WIDTH bytes, so the byte index + it lands on is `(idx // WIDTH) * WIDTH`. That equals `idx` only when the + index is WIDTH-aligned; otherwise the store both corrupts a neighbouring + block's scale byte and leaves the intended one unwritten. + + In the blocked layout every term of `blocked_scale_idx` except `scale_col & 3` + is a multiple of 4, so `idx % WIDTH == scale_col % WIDTH`. The requirement is + therefore exactly `scale_col_base % WIDTH == 0`, a property of the caller's + argument rather than of the data. Today's callers pass `2 * tile_n` (WIDTH 2) + and `8 * tile_n` (WIDTH 4), whose multipliers are multiples of WIDTH, so + alignment holds structurally for any tile index. + + The `assert_` below only fires in an assertions-enabled build + (`CUTE_DSL_ENABLE_ASSERTIONS=1`); it is a debugging aid, not the guarantee. + The guarantee is the caller contract above, which matters because the + design's generalization ("buffer min(4, CTA_N/64) bytes") would put a + computed expression here. + """ + packed_ty = cutlass.Uint32 if cutlass.const_expr(WIDTH == 4) else cutlass.Uint16 + packed_ptr = cute.recast_ptr(scales.iterator, dtype=packed_ty) + for run in cutlass.range_constexpr(NUM_BYTES // WIDTH): + acc = Int32(0) + for i in cutlass.range_constexpr(WIDTH): + acc = acc | ((scale_bytes[run * WIDTH + i] & Int32(0xFF)) << Int32(8 * i)) + idx = blocked_scale_idx(row, scale_col_base + run * WIDTH, num_scale_col_blocks) + cute.testing.assert_( + idx % WIDTH == 0, + "packed scale store is not WIDTH-aligned: scale_col_base must be a " + "multiple of the store width", + ) + dst = cute.make_tensor(packed_ptr + (idx // WIDTH), cute.make_layout(1)) + dst[0] = packed_ty(acc) + + +@cute.jit +def colwise_quant_store( + words: cute.Tensor, + spad: cute.Tensor, + qdata: cute.Tensor, + scales: cute.Tensor, + tidx: Int32, + row_base: Int32, + col_base: Int32, + num_rows: Int32, + num_scale_col_blocks: Int32, + WORD_BASE: cutlass.Constexpr = 0, +): + """Transpose-quantize one 32-column chunk into the ``(1, R)`` destination. + + Every epilogue thread contributes its own row's 32 columns through `words` + and then reads back a *different* slice -- 16 rows of one column pair -- + which is the transpose. Writer and reader sets are the same 32 threads, so + both hazards are covered by ``sync_warp`` rather than a CTA barrier. + + `qdata` is a flat 1-byte-element view of the ``(1, num_rows)``-strided + destination, i.e. physically row-major ``[cols, num_rows]``. `scales` is a + flat uint8 view of the whole-matrix blocked columnwise scale buffer, indexed + with transposed coordinates (feature, row-block). + + `tidx` is the EPILOGUE-LOCAL thread index in + ``[0, 32 * num_epilogue_warps)``: a kernel whose epilogue runs on warps 4-7 + of 256 threads passes ``thread_idx - 128``. The epilogue's first thread must + be 32-aligned so that epilogue warp `w` is one physical warp -- otherwise + ``sync_warp`` and the butterfly shuffle no longer cover the writer/reader + set and the transpose silently mixes rows from two warps. + + `row_base` must be a multiple of 128 and `col_base` a multiple of 32. + """ + lane = tidx % Int32(COLWISE_ROWS_PER_WARP) + warp = tidx // Int32(COLWISE_ROWS_PER_WARP) + warp_base = warp * Int32(COLWISE_WORDS_PER_WARP) + + # Writer: for a fixed pair the 32 lanes hit 32 consecutive words, so every + # store is conflict-free without a swizzle. + for p in cutlass.range_constexpr(COLWISE_PAIRS): + spad[warp_base + Int32(p * COLWISE_PAIR_PITCH) + lane] = words[WORD_BASE + p] + cute.arch.sync_warp() + + # Reader: thread (pair, tq) owns column pair `pair` and the 16 rows + # [16*tq, 16*tq+16) of this warp's 32-row slab. `half` numbers those 16-row + # halves across the whole CTA tile, so `half` and `warp` are the same index + # at two granularities and the row-block below needs no extra arithmetic. + pair = tidx % Int32(COLWISE_PAIRS) + half = tidx // Int32(COLWISE_PAIRS) + tq = half % Int32(2) + feature = col_base + Int32(2) * pair + pair_base = ( + warp_base + + pair * Int32(COLWISE_PAIR_PITCH) + + tq * Int32(COLWISE_ROWS_PER_WARP // 2) + ) + + quads = [] + amax_packed = Int32(0) + for c in cutlass.range_constexpr(COLWISE_ROWS_PER_WARP // 2 // _VEC_WORDS): + quad = _load_spad_quad(spad, pair_base + Int32(c * _VEC_WORDS)) + quads.append(quad) + for t in cutlass.range_constexpr(_VEC_WORDS): + amax_packed = abs_max_nan_bf16x2(amax_packed, quad[t]) + cute.arch.sync_warp() + + # The two threads holding the halves of one 32-row block differ exactly in + # bit 4 of tidx, which is a lane bit, so one butterfly at distance 16 + # completes the amax. Masking first drops the junk signs xorsign leaves. + amax_packed = amax_packed & Int32(0x7FFF7FFF) + amax_packed = max_nan_bf16x2( + amax_packed, cute.arch.shuffle_sync_bfly(amax_packed, COLWISE_PAIRS) + ) + + # Low half is column `2*pair`'s amax, high half is `2*pair+1`'s: two + # independent scales out of one reduction. Both are BF16-valued widened to + # f32, which is what makes float_to_e8m0 exact. + lo_byte = float_to_e8m0((amax_packed & Int32(0xFFFF)) << Int32(16)) + hi_byte = float_to_e8m0((amax_packed >> Int32(16)) << Int32(16)) + inv_packed = e8m0_reciprocal_bf16(lo_byte) | ( + e8m0_reciprocal_bf16(hi_byte) << Int32(16) + ) + + # One of the two threads per (pair, row-block) owns the scale bytes. They + # are not adjacent in the blocked layout (they differ in the feature index, + # hence by 16 bytes), so this stays two 1-byte stores. + if tq == Int32(0): + row_block = (row_base // Int32(SCALE_BLOCK)) + warp + scales[blocked_scale_idx(feature, row_block, num_scale_col_blocks)] = ( + cutlass.Uint8(lo_byte) + ) + scales[ + blocked_scale_idx(feature + Int32(1), row_block, num_scale_col_blocks) + ] = cutlass.Uint8(hi_byte) + + # Quantize and transpose: mul_cvt_2x turns two rows x two columns into + # bytes [r0c0, r0c1, r1c0, r1c1], and the byte permutes split that into one + # word per column holding four consecutive rows. This is where prmt_even / + # prmt_odd belong -- not in any gate/up de-interleave, where the pair is + # already in two separate FP32 registers. + num_quads = cutlass.const_expr(COLWISE_ROWS_PER_WARP // 2 // _VEC_WORDS) + col_lo = cute.make_rmem_tensor((num_quads,), Int32) + col_hi = cute.make_rmem_tensor((num_quads,), Int32) + for c in cutlass.range_constexpr(num_quads): + quad = quads[c] + pack01 = mul_cvt_2x(quad[0], quad[1], inv_packed) + pack23 = mul_cvt_2x(quad[2], quad[3], inv_packed) + col_lo[c] = prmt_even(pack01, pack23) + col_hi[c] = prmt_odd(pack01, pack23) + + # col_lo/col_hi are 16 consecutive rows of one column: 16 contiguous bytes + # of the [cols, num_rows] buffer. num_rows and row_base are multiples of 128 + # and the row offset is a multiple of 16, so both stores are STG.128. + row_offset = row_base + half * Int32(COLWISE_ROWS_PER_WARP // 2) + _store_vec_i32(qdata, (feature * num_rows + row_offset) >> Int32(2), col_lo) + _store_vec_i32( + qdata, ((feature + Int32(1)) * num_rows + row_offset) >> Int32(2), col_hi + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py new file mode 100644 index 0000000000..11c448fba5 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py @@ -0,0 +1,463 @@ +# 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. + +"""Frozen configuration for the shared MXFP8 grouped blockscaled GEMM core. + +One tiling, one pipeline shape, one warp assignment for all three kernels in the +family (FC1 GEMM+SwiGLU, FC2 dgrad+dSwiGLU, grouped wgrad). Kernels differ only +in their epilogue and in ``epi_n_acc``; nothing else here is a per-kernel knob. + +Why each value is what it is: + +``cta_group = ONE`` / ``cluster_shape_mn = (1, 1)`` + Every per-expert row count is a multiple of 128 and ``cta_tile_m`` is 128, so + an M tile never straddles an expert boundary. 2CTA forces a 256-row cluster + tile, which reintroduces partial tiles for ``m[g] == 128 (mod 256)`` and + would need hand-predication in every quantized store and every scale byte. + With one CTA per MMA there is also nothing to multicast, so the cluster is + trivial and no multicast masks or cluster launch barriers exist. + +``mma_tiler = (128, 128, 128)`` + M=128 is mandatory for ``CtaGroup.ONE``. K=128 is the smallest multiple of + ``sf_vec_size * 4`` and makes one K tile exactly one scale-factor atom. N=128 + keeps TMEM at 160 of 512 columns and avoids both the N=64 SFB column-shift + and the N=256 overlapping-accumulator machinery. + +``num_acc_stage = 1`` + One CTA produces exactly one output tile (the grid is data-independent and + there is no persistent tile scheduler), so there is no second accumulator to + overlap with. + +``threads = 256`` + Warp 0 loads (TMA), warp 1 issues the MMA, warps 4-7 run the epilogue, warps + 2-3 idle. 128 epilogue threads is what makes ``tmem_warp_shape_mn = (4, 1)`` + and the thread-owns-one-row identity hold; see :data:`T2R_PARTITION_DOC`. + +``epi_tile = (128, epi_n_acc)`` + ``epi_tile_M == cta_tile_M == 128`` is mandatory: every gate/up pairing and + every scale index below assumes a single epilogue tile along M. + +CTA_N = 256 is a legal alternative (TMEM 256+16+32 = 304 <= 512 with one +accumulator stage) and changes only the rowwise scale store width, but it halves +``num_ab_stage``; nothing structural depends on the choice. + +WARNING -- ``cta_tile_k`` is pinned at 128 and raising it is not a free tuning +knob. The grouped wgrad kernel selects an expert's K range with an integer K-tile +index base rather than a per-expert TMA descriptor, which is exact only because +``token_offset`` (a multiple of 128) is a multiple of ``cta_tile_k``. At +``cta_tile_k = 256`` a group boundary can land mid-tile, and the kernel would +need per-expert descriptors or explicit K-tail predication brought back in. The +constant-512-byte ``Rest_K`` stride that makes the index base correct is a +property of the 128-element scale-factor granule, not of the tile size, so it +does not rescue a larger tile either. +""" + +import enum +from dataclasses import dataclass +from typing import Tuple + +__all__ = [ + "RaggedAxis", + "GroupedGemmConfig", + "SWIGLU_FWD_CONFIG", + "DSWIGLU_BWD_CONFIG", + "WGRAD_CONFIG", + "SMEM_CAPACITY_BYTES", + "TMEM_TOTAL_COLS", + "is_supported", + "check_supported", + "T2R_PARTITION_DOC", + "EPILOGUE_PROTOCOL_DOC", +] + +# sm_100 usable dynamic shared memory per CTA, (228 - 1) KiB. +SMEM_CAPACITY_BYTES = 232448 +# The TMEM allocator requires a power-of-two multiple of 32 columns; shared +# memory already pins us to one CTA per SM, so allocating the whole array costs +# nothing. +TMEM_TOTAL_COLS = 512 +# Row-count granularity every expert group, and the allocation itself, respect. +GROUP_ALIGNMENT = 128 +# MXFP8 scaling block: 32 values share one E8M0 scale. +SF_VEC_SIZE = 32 + + +class RaggedAxis(enum.Enum): + """Which GEMM axis the per-expert offsets partition. + + ``M`` is kernels A and B: the ragged axis is the token axis, the grid is + ``(R // 128, N // CTA_N, 1)``, and the expert is looked up per CTA from the + absolute row base. ``K`` is kernel C: the ragged axis is the contraction, the + grid is ``(N // 128, K // 128, G)``, and only the K-loop trip count is + data-dependent. + """ + + M = 0 + K = 1 + + +@dataclass(frozen=True) +class GroupedGemmConfig: + """Constexpr configuration of the shared grouped blockscaled GEMM core. + + Every field is a trace-time constant. No kernel body may hard-code any of + these numbers; read them from the config so a retune cannot silently + desynchronize the mainloop from an epilogue. + """ + + # Per-kernel: accumulator columns handed to the epilogue per subtile. + # 64 for the FC1 SwiGLU forward (one output column consumes an adjacent + # gate/up accumulator pair, so 64 accumulator columns are 32 output columns + # == exactly one rowwise 1x32 block per thread), 32 for the dgrad backward + # and for wgrad. + epi_n_acc: int + # Which axis the offsets partition. + ragged_axis: RaggedAxis + + # --- frozen for every kernel in the family ------------------------------- + cta_tile_m: int = 128 + cta_tile_n: int = 128 + cta_tile_k: int = 128 + num_ab_stage: int = 6 + num_acc_stage: int = 1 + cluster_shape_mn: Tuple[int, int] = (1, 1) + sf_vec_size: int = SF_VEC_SIZE + threads: int = 256 + tma_warp_id: int = 0 + mma_warp_id: int = 1 + epilogue_warp_ids: Tuple[int, ...] = (4, 5, 6, 7) + # Named barrier ids. 0 is left free for the DSL's own use. + epilogue_sync_barrier_id: int = 1 + tmem_alloc_barrier_id: int = 2 + + def __post_init__(self): + if self.cta_tile_m != 128: + raise ValueError( + "cta_tile_m is pinned at 128: CtaGroup.ONE requires it and the " + "no-partial-M-tile argument depends on it" + ) + if self.cta_tile_k != 128: + raise ValueError( + "cta_tile_k is pinned at 128; see this module's docstring for why " + "raising it reintroduces per-expert descriptors" + ) + if self.cta_tile_n % 128 != 0: + # SFB's MN extent is round_up(N, 128). Only when that equals N is the + # SFB tiled MMA identical to the data one, which is what lets the core + # build a single tiled MMA and skip the N=64 TMEM column-shift path. + raise ValueError( + f"cta_tile_n must be a multiple of 128, got {self.cta_tile_n}; " + "a smaller N needs a separate SFB tiled MMA and a TMEM column shift" + ) + if self.cta_tile_n % self.epi_n_acc != 0: + raise ValueError( + f"cta_tile_n ({self.cta_tile_n}) must be a multiple of epi_n_acc " + f"({self.epi_n_acc})" + ) + if self.cta_tile_m % (4 * SF_VEC_SIZE) != 0: + # 4 epilogue warps x 32 rows: a columnwise 32x1 block must never be + # split across warps. + raise ValueError( + "cta_tile_m must be a multiple of 128 for the 4-warp epilogue" + ) + if len(self.epilogue_warp_ids) * 32 != self.cta_tile_m: + raise ValueError( + "the epilogue must have exactly one thread per row of the CTA tile" + ) + if self.epilogue_warp_ids[0] % 4 != 0: + # tcgen05.ld selects its TMEM datapath sub-partition from the + # PHYSICAL warp id, so the epilogue must start on an aligned warp + # quad. A misaligned block still launches and still returns + # plausible numbers -- every 128-row tile comes back with its four + # 32-row datapath groups rotated -- and no shape, byte or NaN check + # detects a pure row permutation. Reject it here. + raise ValueError( + f"epilogue_warp_ids must start on an aligned warp quad, got " + f"{self.epilogue_warp_ids}: warp {self.epilogue_warp_ids[0]} is " + f"{self.epilogue_warp_ids[0] % 4} past a multiple of 4; tcgen05.ld " + "would silently permute the 32-row groups of every tile" + ) + # The kernel selects the epilogue with `warp_idx >= epilogue_warp_ids[0]`, + # so they must be the last contiguous block of warps in the CTA. + if tuple(self.epilogue_warp_ids) != tuple( + range(self.threads // 32 - len(self.epilogue_warp_ids), self.threads // 32) + ): + raise ValueError( + f"epilogue_warp_ids {self.epilogue_warp_ids} must be the last " + f"{len(self.epilogue_warp_ids)} warps of the {self.threads // 32} " + "in the CTA" + ) + if self.tma_warp_id in self.epilogue_warp_ids or ( + self.mma_warp_id in self.epilogue_warp_ids + ): + raise ValueError("the TMA and MMA warps must not be epilogue warps") + + # --- derived, all trace-time --------------------------------------------- + + @property + def mma_tiler_mnk(self) -> Tuple[int, int, int]: + return (self.cta_tile_m, self.cta_tile_n, self.cta_tile_k) + + @property + def cta_tile_shape_mnk(self) -> Tuple[int, int, int]: + return (self.cta_tile_m, self.cta_tile_n, self.cta_tile_k) + + @property + def epi_tile(self) -> Tuple[int, int]: + return (self.cta_tile_m, self.epi_n_acc) + + @property + def num_epi_subtiles(self) -> int: + return self.cta_tile_n // self.epi_n_acc + + @property + def num_epilogue_threads(self) -> int: + return 32 * len(self.epilogue_warp_ids) + + @property + def first_epilogue_thread(self) -> int: + return 32 * self.epilogue_warp_ids[0] + + # --- shared memory budget ------------------------------------------------ + + @property + def ab_stage_bytes(self) -> int: + """Bytes of one mainloop stage: A and B tiles plus both scale atoms. + + At (128, 128, 128) E4M3 that is 16384 + 16384 + 512 + 512 = 33792, which + is also the exact ``tx_count`` the TMA pipeline barrier must expect. + """ + a = self.cta_tile_m * self.cta_tile_k # E4M3, 1 byte per element + b = self.cta_tile_n * self.cta_tile_k + # One E8M0 byte per 32 contracted elements per MN row, i.e. exactly one + # 128x4 blocked tile (512 B) per operand at the frozen shape. SFB's MN + # extent is round_up(N, 128). + sfa = self.cta_tile_m * (self.cta_tile_k // SF_VEC_SIZE) + sfb = max(self.cta_tile_n, 128) * (self.cta_tile_k // SF_VEC_SIZE) + return a + b + sfa + sfb + + @property + def mbarrier_bytes(self) -> int: + """AB pipeline (full+empty) and accumulator handoff (full+empty) mbarriers.""" + return 8 * (2 * self.num_ab_stage + 2 * self.num_acc_stage) + + def smem_bytes(self, epilogue_smem_bytes: int = 0) -> int: + return ( + self.num_ab_stage * self.ab_stage_bytes + + self.mbarrier_bytes + + epilogue_smem_bytes + # tmem holding buffer plus struct alignment slack + + 256 + ) + + def max_ab_stages(self, epilogue_smem_bytes: int = 0) -> int: + """Stages that fit alongside the epilogue's shared-memory request.""" + fixed = self.mbarrier_bytes + epilogue_smem_bytes + 256 + return (SMEM_CAPACITY_BYTES - fixed) // self.ab_stage_bytes + + # --- tensor memory budget ------------------------------------------------ + + @property + def acc_tmem_cols(self) -> int: + return self.cta_tile_n * self.num_acc_stage + + @property + def sfa_tmem_cols(self) -> int: + return (self.cta_tile_m // SF_VEC_SIZE) * 4 + + @property + def sfb_tmem_cols(self) -> int: + # SFB's MN extent is round_up(N, 128); at N=128 that is N itself. + return (max(self.cta_tile_n, 128) // SF_VEC_SIZE) * 4 + + @property + def used_tmem_cols(self) -> int: + """160 of 512 at the frozen shape.""" + return self.acc_tmem_cols + self.sfa_tmem_cols + self.sfb_tmem_cols + + +# Kernel A. EPI_N_ACC=64 is not negotiable: accumulator column 2f is gate_f and +# 2f+1 is up_f, so 64 accumulator columns are 32 output columns, exactly one +# rowwise 1x32 block per thread. 32 would give half a block and force a partial +# amax carried across subtiles. +SWIGLU_FWD_CONFIG = GroupedGemmConfig(epi_n_acc=64, ragged_axis=RaggedAxis.M) +# Kernel B. The accumulator N axis is F (one column per feature) and each column +# produces two interleaved output columns, so 32 accumulator columns are 64 +# output columns = two rowwise 1x32 blocks per thread per subtile. +DSWIGLU_BWD_CONFIG = GroupedGemmConfig(epi_n_acc=32, ragged_axis=RaggedAxis.M) +# Kernel C. 32 FP32 accumulators are 32 contiguous BF16 outputs = 64 contiguous +# bytes per thread, 64-byte aligned. +WGRAD_CONFIG = GroupedGemmConfig(epi_n_acc=32, ragged_axis=RaggedAxis.K) + + +def is_supported( + model_dim: int, hidden_dim: int, allocated_rows: int, num_groups: int +) -> bool: + """The initial optimized support predicate, as a pure boolean. + + Mirrors the kernel contract's shape predicate and the host validators. The + caller is expected to fall back to the unfused torchao path when this is + false rather than launching -- feeding non-128-aligned groups to the blocked + scale path produces a wrong-sized buffer and an unusable CUDA context, not an + error. + + The per-expert row counts live in device memory and are not checkable here; + they are asserted on device on every launch. + """ + return ( + num_groups >= 1 + and model_dim > 0 + and hidden_dim > 0 + and allocated_rows > 0 + and model_dim % GROUP_ALIGNMENT == 0 + and hidden_dim % GROUP_ALIGNMENT == 0 + and allocated_rows % GROUP_ALIGNMENT == 0 + ) + + +def check_supported( + model_dim: int, hidden_dim: int, allocated_rows: int, num_groups: int +) -> None: + """:func:`is_supported` with a message naming the offending value. + + Raises ``ValueError``, never ``assert``: ``python -O`` strips assertions and + would reintroduce the silent-corruption path. + """ + if num_groups < 1: + raise ValueError(f"G must be at least 1, got {num_groups}") + for name, value in ( + ("D", model_dim), + ("F", hidden_dim), + ("R", allocated_rows), + ): + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + if value % GROUP_ALIGNMENT != 0: + raise ValueError( + f"{name} must be a multiple of {GROUP_ALIGNMENT}, got {value}" + ) + + +# --------------------------------------------------------------------------- +# The two frozen interfaces lanes 2 and 3 compile against. +# --------------------------------------------------------------------------- + +T2R_PARTITION_DOC = """\ +grouped_gemm_core.t2r_partition(tidx, tAcc, config) -> (tiled_copy_t2r, +tTR_tAcc, tTR_rAcc, tTR_cAcc) + + tidx Int32 epilogue-local thread index in [0, 128). NOT the raw threadIdx.x; + subtract config.first_epilogue_thread first. + tAcc the (MMA, MMA_M, MMA_N, ACC_STAGE) TMEM accumulator tensor. + config a GroupedGemmConfig, which fixes epi_tile = (128, epi_n_acc). + +Returns, with EPI_M == 1 always: + + tiled_copy_t2r the tcgen05 TMEM->register tiled copy. + tTR_tAcc (T2R, T2R_M, T2R_N, EPI_M, EPI_N) TMEM source, sliced per + epilogue subtile s as tTR_tAcc[(None, None, None, 0, s)]. + tTR_rAcc (T2R, T2R_M, T2R_N) FP32 register destination for one subtile. + tTR_cAcc (T2R, T2R_M, T2R_N, EPI_M, EPI_N) of (row, col) coordinates in + the 128 x cta_tile_n CTA tile, partitioned identically. + +The copy atom is built with elem_ty_d = Float32 even though the real outputs are +E4M3. Passing an 8-bit d type steers get_tmem_load_op into the tmem_dp=16 +layouts, which are shaped for a direct FP8 TMA store and are not what a +dual-quantization epilogue wants. + +Structural consequence, measured on GB200 with a real TMEM accumulator (not just +from the copy atom's thread-value layout): thread t owns row t of the 128-row CTA +tile and all epi_n_acc accumulator columns of the subtile, contiguously, and the +raw register order is linear -- register v of thread t holds accumulator element +(t, v) of the subtile. Warp w therefore owns rows [32w, 32w+32) = exactly one +32-row MX block. So a rowwise 1x32 amax is intra-thread and a columnwise 32x1 +amax is intra-warp, and a columnwise scale block is never split across warps or +CTAs. + +Even so, derive every index -- the gate/up de-interleave and every scale +coordinate -- from tTR_cAcc rather than from a raw register number. The linear +order is a measured property of one copy atom at one epi_tile, not a guarantee, +and tTR_cAcc costs nothing because it folds at trace time. +""" + +EPILOGUE_PROTOCOL_DOC = """\ +An epilogue is a module-level function passed to the core as a Constexpr. Two +requirements, both load-bearing: + +1. It must be decorated ``@cute.jit``. The DSL preprocessor only rewrites + decorated functions, so an undecorated epilogue cannot use + ``cutlass.range_constexpr`` (it raises "range_constexpr should be preprocessed + by preprocessor") and, worse, a dynamic ``if`` in one would be evaluated as a + Python truth test instead of becoming a predicated region. +2. It must be a module-level function object. The DSL keys its compile cache on + function identity, so a lambda or a closure built per call recompiles every + launch. + + @cute.jit + def my_epilogue( + tTR_rAcc, # (T2R, T2R_M, T2R_N) FP32 register fragment, one subtile + tTR_cAcc_s, # (T2R, T2R_M, T2R_N) matching (row, col) coordinates in + # the 128 x cta_tile_n CTA tile + tiled_copy_t2r,# for cute.make_tiled_copy_D / retile, if needed + epi_tidx, # Int32 in [0, 128); equals the CTA-tile row this thread owns + subtile_idx, # Constexpr int in [0, config.num_epi_subtiles) + tile, # TileCoords: see below + epi_smem, # cute.Pointer(Int32) to the requested scratch, or None + out, # the tuple of destination tensors the launcher passed + cfg, # Constexpr GroupedGemmConfig + ) -> None + +TileCoords fields, all Int32 and all CTA-uniform: + + tile_m absolute M-tile index (kernels A/B: over [0, R/128)) + tile_n absolute N-tile index + expert selected expert index. Meaningless on an inactive-tail tile, where + the scan saturates at G-1; nothing may depend on it there, since + that tile's output is defined to be zero. + row_base tile_m * 128 + col_base tile_n * cta_tile_n + k_cnt mainloop trip count; 0 on a tail tile or a zero-token expert + +The epilogue is called once per subtile, num_epi_subtiles times per CTA, always +by all 128 epilogue threads and always with k_cnt CTA-uniform. + +Tail rule, stated precisely because the two halves pull in opposite directions: + + * NEVER predicate a STORE on the inactive tail. A tail tile arrives with a + zeroed accumulator and the unmodified store path then emits exactly the zeros + the contract requires (zero qdata bytes, zero E8M0 scale bytes). Skipping a + store is how a destination element stops being written. + + * ALWAYS predicate an extra GMEM INPUT LOAD on `tile.k_cnt == 0`, substituting + zeros. This is not optional. Kernel B reads the saved `z_bf16` in its + epilogue, and rows [A, R) of that tensor are read-forbidden precisely because + they may hold anything. Measured, feeding a tail `z` into the dSwiGLU with + dh == 0 as the zeroed accumulator delivers it: + z = 0 -> dz 0x00000000 correct + z = NaN or +Inf -> dz 0x7fff7fff (dgate/dup = NaN) WRONG + z = uninit 0xDEADBEEF -> dz 0x80008000 (dgate/dup = -0.0) WRONG + A NaN tail makes that block's scale byte 0xFF and every qdata byte 0x7F; + even benign garbage yields qdata 0x80 rather than 0x00. Both violate the + read-forbidden and write-zero halves of the ragged-tail contract. + +In short: the accumulator is already zeroed for you, so trust it and store +unconditionally; anything you load yourself must be gated on k_cnt. + +No cross-subtile state: EPILOGUE is called once per subtile and every call gets +fresh registers. There is deliberately no way to carry a value from subtile s to +s+1, so rowwise scale bytes cannot be buffered across subtiles and emitted as one +wide store per CTA tile. DECIDED 2026-08-16: emit ONE scale store per subtile +(`rowwise_scale_flush(..., NUM_BYTES=1)`). That costs 2 stores instead of 1 for +Kernel A and 8 instead of 2 for Kernel B, and is verified bitwise clean. Store-count +reduction is a tuning-stage concern; buying it here would mean either +trace-time module-level state across the unrolled loop or staging through +epi_smem, both of which trade a correctness-critical interface for a few stores. + +Shared-memory scratch: the epilogue declares its byte count to the launcher, +which passes back a 128-byte-aligned pointer of that size. The mainloop's stage +count is computed against that request, so it is not free -- but 9216 bytes (the +columnwise transpose staging) still leaves 6 stages. +""" diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py new file mode 100644 index 0000000000..174620a213 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py @@ -0,0 +1,825 @@ +# 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. + +"""Shared MXFP8 grouped blockscaled GEMM core: descriptors, mainloop, TMEM, T2R. + +Everything between "torch tensors" and "an FP32 accumulator tile in registers". +The three kernels in this family plug different epilogues into this one mainloop +by passing a module-level function as the ``EPILOGUE`` Constexpr; see +``grouped_gemm_config.EPILOGUE_PROTOCOL_DOC``. + +Three structural decisions, all descending from the per-expert row counts being +multiples of 128: + +*No per-group tensormaps, anywhere.* Every operand is one host-built static TMA +descriptor over the whole tensor. Per-expert selection is an integer coordinate: +an L coordinate for the 3-D weight operands, and a K-tile index base for the +wgrad kernel's ragged contraction. The latter is exact because the blockscaled +scale-factor layout's ``Rest_K`` stride is a constant 512 bytes for *every* MN +row-block, so advancing the K tile index advances every row-block's byte address +by the same amount. That was verified on this wheel (probe V1) and it is what +deletes the tensormap workspace, the descriptor-init kernel, the descriptor +fences, and the padded-offset prefix sum that the reference kernels carry. + +*No tile scheduler.* With the ragged axis tile-aligned, kernels A/B enumerate all +of ``[0, R/128)`` M tiles and kernel C's ``(N/128, K/128, G)`` grid is fully +static; only C's K-loop trip count is data-dependent. So there is no +``max_active_clusters`` query, no persistent loop, no offsets prefix scan, and no +work-tile shared-memory pipeline. + +*The inactive tail needs no special code path.* A tile whose row base is at or +past the active row count runs with ``k_cnt == 0``: no TMA loads are issued (so +no inactive row is ever read), the accumulator fragment is zeroed in registers, +and the *unmodified* epilogue emits the zeros the contract requires. ``k_cnt`` is +CTA-uniform, so no barrier arrival count can desynchronize. + +Expert lookup for the ragged-M kernels is an unrolled G-way scan over the +device-side offsets -- no host synchronization and no ``.item()``. +""" + +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Float32, Int32 +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.utils import LayoutEnum + +from torchao.prototype.moe_training.kernels.mxfp8.grouped_gemm_config import ( + SMEM_CAPACITY_BYTES, + TMEM_TOTAL_COLS, + GroupedGemmConfig, + RaggedAxis, +) +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_epilogue import ( + validate_group_offsets_device, +) + +__all__ = [ + "TileCoords", + "activation_gemm_view", + "weight_gemm_view", + "make_operand_views", + "make_tiled_mma", + "make_sf_gemm_tensor", + "t2r_partition", + "dump_accumulator_epilogue", + "grouped_gemm_kernel", + "launch_grouped_gemm", +] + + +# --------------------------------------------------------------------------- +# Host-side operand views. Pure torch, no copies: every one of these is a +# restride of the caller's storage into the (MN, K, L) GEMM domain the core +# wants, with K contiguous. +# --------------------------------------------------------------------------- + + +def activation_gemm_view(t: torch.Tensor) -> torch.Tensor: + """``[MN, K]`` K-contiguous -> ``(MN, K, 1)`` with a defined batch stride. + + Covers ``x_q``/``do_q`` directly, and both wgrad operands via their free + transpose: a logical ``[R, N]`` with stride ``(1, R)`` *is* a K-contiguous + ``[N, R]``, so pass ``t.t()``. + """ + mn, k = t.shape + if t.stride() != (k, 1): + raise ValueError( + f"GEMM operand must be K-contiguous with stride {(k, 1)}, got {t.stride()}" + ) + return torch.as_strided(t, (mn, k, 1), (k, 1, mn * k)) + + +def weight_gemm_view(w: torch.Tensor) -> torch.Tensor: + """``[G, K, N]`` stride ``(K*N, 1, K)`` -> ``(N, K, G)`` stride ``(K, 1, K*N)``. + + That is the prequantized weight layout both kernels A and B receive, and the + permute is exactly the K-major B operand the MMA wants, so the expert becomes + an L coordinate and no descriptor is ever rebuilt. + """ + g, k, n = w.shape + if w.stride() != (k * n, 1, k): + raise ValueError( + f"grouped weight must have stride {(k * n, 1, k)}, got {w.stride()}" + ) + return w.permute(2, 1, 0) + + +def make_operand_views(a: torch.Tensor, b: torch.Tensor): + """``(mA, mB)`` in the GEMM domain, choosing the view from ``b``'s rank. + + A 3-D ``b`` is a grouped weight (expert becomes the L coordinate); a 2-D one + is the wgrad case, where both operands are ungrouped and the expert is a + K-tile index base instead. + """ + return activation_gemm_view(a), ( + weight_gemm_view(b) if b.ndim == 3 else activation_gemm_view(b) + ) + + +@dataclass +class TileCoords: + """The per-CTA tile description handed to the epilogue. All fields CTA-uniform. + + ``k_cnt == 0`` marks both cases where the mainloop is skipped: an inactive + tail tile (ragged M) and a zero-token expert (ragged K). The epilogue must + not branch on it -- the accumulator is already zero. + """ + + tile_m: Int32 + tile_n: Int32 + expert: Int32 + row_base: Int32 + col_base: Int32 + k_cnt: Int32 + + +def make_tiled_mma(cfg: GroupedGemmConfig, a_dtype, b_dtype, sf_dtype): + """The one blockscaled tiled MMA, K-major on both operands. + + No operand in this family is MN-major, so the 8-bit MN-major N-step and + transpose-swizzle caveats never apply. ``MmaMXF8F6F4Op`` hard-wires FP32 + accumulation and instruction K=32, so a 128-element K tile is four MMA-K + instructions and exactly one scale-factor atom. + """ + return sm100_utils.make_blockscaled_trivial_tiled_mma( + a_dtype, + b_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + sf_dtype, + cfg.sf_vec_size, + tcgen05.CtaGroup.ONE, + (cfg.cta_tile_m, cfg.cta_tile_n), + ) + + +def make_sf_gemm_tensor( + flat_sf: cute.Tensor, mn: int, k: int, l: int, sf_vec_size: int +): + """Retile a flat blocked E8M0 buffer into the GEMM-domain scale-factor layout. + + The buffer is carried flat by ABI and may arrive as either uint8 or + float8_e8m0fnu, so the pointer is recast unconditionally -- the MMA rejects a + scale operand whose element type is not E8M0. + + ``tile_atom_to_shape_SF`` builds kernel IR, so this must be called inside a + trace even though the shapes are static and every evaluation folds. + """ + return cute.make_tensor( + cute.recast_ptr(flat_sf.iterator, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.tile_atom_to_shape_SF((mn, k, l), sf_vec_size), + ) + + +def t2r_partition(tidx, tAcc_base: cute.Tensor, cfg: GroupedGemmConfig): + """Accumulator -> register handoff. See ``config.T2R_PARTITION_DOC``. + + ``elem_ty_d`` is Float32 even though the real outputs are E4M3: passing an + 8-bit d type steers ``get_tmem_load_op`` into the tmem_dp=16 layouts, which + are shaped for a direct FP8 TMA store, not for a dual-quantization epilogue. + + ``tTR_cAcc`` carries each register's ``(row, col)`` in the CTA tile. Deriving + every index from it, rather than from a raw register number, is what makes + the gate/up de-interleave and the scale addressing correct by construction + whatever the copy atom's internal value order is. + """ + copy_atom_t2r = sm100_utils.get_tmem_load_op( + cfg.cta_tile_shape_mnk, + LayoutEnum.ROW_MAJOR, + Float32, + Float32, + cfg.epi_tile, + False, + ) + # (MMA, MMA_M, MMA_N, ACC_STAGE) -> (CTA_M, CTA_N); one accumulator stage. + tAcc_mn = tAcc_base[((None, None), 0, 0, 0)] + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N) + tAcc_epi = cute.flat_divide(tAcc_mn, cfg.epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + cAcc_epi = cute.flat_divide( + cute.make_identity_tensor((cfg.cta_tile_m, cfg.cta_tile_n)), cfg.epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N), values are (row, col) in the CTA tile + tTR_cAcc = thr_copy_t2r.partition_D(cAcc_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor(tTR_cAcc[(None, None, None, 0, 0)].shape, Float32) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc + + +def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor): + """SMEM -> TMEM scale-factor copy, issued once per K tile from the MMA warp. + + ``Cp4x32x128bOp`` carries the warpx4 broadcast qualifier: issue it as a plain + ``cute.copy`` and never wrap it in ``elect_one()``, which deadlocks because + the compiler already inserts the election. + """ + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), sSF.element_type + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, thr_copy_s2t.partition_S(tCsSF_compact) + ) + tCtSF_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_s2t, tCtSF_s2t + + +@cute.jit +def dump_accumulator_epilogue( + tTR_rAcc, + tTR_cAcc_s, + tiled_copy_t2r, + epi_tidx, + subtile_idx: cutlass.Constexpr, + tile: TileCoords, + epi_smem, + out, + cfg: cutlass.Constexpr, +): + """The M2a gate epilogue: write the raw FP32 accumulator to ``out[0]``. + + ``out[0]`` is ``[M, N, L]`` FP32, where L is 1 for a ragged-M kernel and the + expert count for a ragged-K one. It exists to prove the blockscaled MMA and + the scale-factor addressing before any real epilogue does; it is deliberately + a scalar store loop, not a vectorized one. + """ + gD = out[0] + if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.K): + out_l = tile.expert + else: + out_l = Int32(0) + for v in cutlass.range_constexpr(cute.size(tTR_rAcc)): + crd = tTR_cAcc_s[v] + gD[(tile.row_base + crd[0], tile.col_base + crd[1], out_l)] = tTR_rAcc[v] + + +@cute.kernel +def grouped_gemm_kernel( + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB: cute.Tensor, + offs: cute.Tensor, + out, + a_smem_layout: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + sfa_smem_layout: cute.Layout, + sfb_smem_layout: cute.Layout, + cfg: cutlass.Constexpr, + storage_type: cutlass.Constexpr, + EPILOGUE: cutlass.Constexpr, + EPI_SMEM_BYTES: cutlass.Constexpr, + VALIDATE_OFFSETS: cutlass.Constexpr, +): + """One CTA computes one 128 x cta_tile_n output tile. Warps: 0 TMA, 1 MMA, + 4-7 epilogue, 2-3 idle.""" + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, bidz = cute.arch.block_idx() + + num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) + num_k_tiles_full = cutlass.const_expr(cute.size(mA, mode=[1]) // cfg.cta_tile_k) + + # ------------------------------------------------------------------ + # Device-side precondition on the offset VALUES, which the host cannot + # check without synchronizing. One block, one warp, one lane. + # ------------------------------------------------------------------ + if cutlass.const_expr(VALIDATE_OFFSETS): + if bidx == 0 and bidy == 0 and bidz == 0: + if warp_idx == 0: + with cute.arch.elect_one(): + validate_group_offsets_device( + offs, + Int32(cute.size(mA, mode=[0])) + if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.M) + else Int32(cute.size(mA, mode=[1])), + ) + + # ------------------------------------------------------------------ + # Tile coordinates. No scheduler: the grid IS the tile enumeration. + # ------------------------------------------------------------------ + tile_m = Int32(bidx) + tile_n = Int32(bidy) + if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.M): + row_base = tile_m * cfg.cta_tile_m + # Unrolled G-way scan: the owning expert is the number of groups that end + # at or before this tile's row base. A zero-token expert is never + # selected, since its end equals its start. + expert = Int32(0) + for g in cutlass.range_constexpr(num_groups - 1): + expert += Int32(offs[g] <= row_base) + # Branch-free tail predicate: rows at or past offs[-1] belong to no + # expert, so the whole mainloop is skipped for them. + is_active = Int32(offs[num_groups - 1] > row_base) + k_base = Int32(0) + k_cnt = is_active * Int32(num_k_tiles_full) + l_a = Int32(0) + l_b = expert + else: + expert = Int32(bidz) + row_base = tile_m * cfg.cta_tile_m + # offs[expert - 1], clamped for expert 0; offsets are nonnegative so the + # multiply is a legal select. + prev = offs[cutlass.max(expert - Int32(1), Int32(0))] * Int32(expert > Int32(0)) + # Exact, not a ceil: every group boundary is a multiple of cta_tile_k. + k_base = prev // cfg.cta_tile_k + k_cnt = (offs[expert] - prev) // cfg.cta_tile_k + l_a = Int32(0) + l_b = Int32(0) + tile = TileCoords( + tile_m=tile_m, + tile_n=tile_n, + expert=expert, + row_base=row_base, + col_base=tile_n * cfg.cta_tile_n, + k_cnt=k_cnt, + ) + + # ------------------------------------------------------------------ + # Shared memory and pipelines + # ------------------------------------------------------------------ + smem = utils.SmemAllocator() + storage = smem.allocate(storage_type) + + sA = storage.sA.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) + sB = storage.sB.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout) + epi_smem = None + if cutlass.const_expr(EPI_SMEM_BYTES > 0): + epi_smem = storage.sEpi.data_ptr() + + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*cfg.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) + ) + + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar.data_ptr(), + num_stages=cfg.num_ab_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + # Must be the EXACT byte count of all four TMA copies in one stage: too + # small and the MMA consumes a partially arrived stage. + tx_count=cutlass.const_expr(cfg.ab_stage_bytes), + cta_layout_vmnk=cluster_layout_vmnk, + ) + # One accumulator stage, so this is a single mbarrier pair; there is no + # inter-tile pipelining to overlap with. + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar.data_ptr(), + num_stages=cfg.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cfg.num_epilogue_threads + ), + cta_layout_vmnk=cluster_layout_vmnk, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=cfg.tmem_alloc_barrier_id, + num_threads=32 * (1 + len(cfg.epilogue_warp_ids)), + ) + epilogue_barrier = pipeline.NamedBarrier( + barrier_id=cfg.epilogue_sync_barrier_id, + num_threads=cfg.num_epilogue_threads, + ) + tmem_alloc = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=cfg.epilogue_warp_ids[0], + is_two_cta=False, + ) + + # ------------------------------------------------------------------ + # Tile the global tensors. One static descriptor per operand; the expert is + # an L coordinate (ragged M) or a K-tile index base (ragged K). + # ------------------------------------------------------------------ + mma_tiler = cfg.mma_tiler_mnk + gA = cute.local_tile( + mA, cute.slice_(mma_tiler, (None, 0, None)), (None, None, None) + ) + gB = cute.local_tile( + mB, cute.slice_(mma_tiler, (0, None, None)), (None, None, None) + ) + gSFA = cute.local_tile( + mSFA, cute.slice_(mma_tiler, (None, 0, None)), (None, None, None) + ) + gSFB = cute.local_tile( + mSFB, cute.slice_(mma_tiler, (0, None, None)), (None, None, None) + ) + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA) + tCgB = thr_mma.partition_B(gB) + tCgSFA = thr_mma.partition_A(gSFA) + tCgSFB = thr_mma.partition_B(gSFB) + + trivial_cta_layout = cute.make_layout(1) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + trivial_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + trivial_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + 0, + trivial_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + # Strip the stride-0 sf_vec_size sub-mode: the 512-byte scale atom is + # contiguous and TMA moves it as 8-byte elements. + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + 0, + trivial_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + tAgA_slice = tAgA[(None, tile_m, None, l_a)] + tBgB_slice = tBgB[(None, tile_n, None, l_b)] + tAgSFA_slice = tAgSFA[(None, tile_m, None, l_a)] + tBgSFB_slice = tBgSFB[(None, tile_n, None, l_b)] + + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, cfg.num_acc_stage)) + + # ------------------------------------------------------------------ + # Warp 0: TMA producer + # ------------------------------------------------------------------ + if warp_idx == cfg.tma_warp_id: + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, cfg.num_ab_stage + ) + for _ in cutlass.range(0, k_cnt, 1, unroll=1): + ab_pipeline.producer_acquire(ab_producer_state) + k_idx = k_base + ab_producer_state.count + bar = ab_pipeline.producer_get_barrier(ab_producer_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_idx)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, k_idx)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, k_idx)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, k_idx)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + ab_producer_state.advance() + ab_pipeline.producer_tail(ab_producer_state) + + # ------------------------------------------------------------------ + # Warp 1: MMA + # ------------------------------------------------------------------ + if warp_idx == cfg.mma_warp_id: + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + # The MMA warp joins the allocation barrier but must never allocate. + tmem_alloc.wait_for_alloc() + acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + dtype=sSFA.element_type, + ) + tCtSFA = cute.make_tensor( + sfa_tmem_ptr, + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + mma_tiler, + cfg.sf_vec_size, + cute.slice_(sfa_smem_layout, (None, None, None, 0)), + ), + ) + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=sSFB.element_type, + ) + tCtSFB = cute.make_tensor( + sfb_tmem_ptr, + blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + mma_tiler, + cfg.sf_vec_size, + cute.slice_(sfb_smem_layout, (None, None, None, 0)), + ), + ) + s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_copy_and_partition(sSFA, tCtSFA) + s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, cfg.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, cfg.num_acc_stage + ) + tCtAcc = tCtAcc_base[(None, None, None, 0)] + + # Acquire and commit unconditionally, including when k_cnt == 0, so the + # accumulator handoff barrier stays balanced on tail tiles. + acc_pipeline.producer_acquire(acc_producer_state) + for k_tile in cutlass.range(0, k_cnt, 1, unroll=1): + ab_pipeline.consumer_wait(ab_consumer_state) + stage_crd = (None, None, None, None, ab_consumer_state.index) + cute.copy(s2t_sfa, tCsSFA_s2t[stage_crd], tCtSFA_s2t) + cute.copy(s2t_sfb, tCsSFB_s2t[stage_crd], tCtSFB_s2t) + # ACCUMULATE=False on the first K tile is what zeroes the + # accumulator; there is no separate TMEM clear. + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + mma_crd = (None, None, None, ab_consumer_state.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[mma_crd], tCtSFA], + [tCrB[mma_crd], tCtSFB], + tCtAcc, + ) + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state.advance() + acc_pipeline.producer_commit(acc_producer_state) + + # ------------------------------------------------------------------ + # Warps 4-7: epilogue + # ------------------------------------------------------------------ + if warp_idx >= cfg.epilogue_warp_ids[0]: + # A power-of-two multiple of 32 columns is required; shared memory + # already pins us to one CTA per SM, so taking the whole array is free. + tmem_alloc.allocate(TMEM_TOTAL_COLS) + tmem_alloc.wait_for_alloc() + acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + epi_tidx = tidx - cfg.first_epilogue_thread + tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc = t2r_partition( + epi_tidx, tCtAcc_base, cfg + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, cfg.num_acc_stage + ) + acc_pipeline.consumer_wait(acc_consumer_state) + + for s in cutlass.range_constexpr(cfg.num_epi_subtiles): + if k_cnt == Int32(0): + # Tail tile or zero-token expert: nothing was accumulated, so the + # fragment is zeroed here and the epilogue runs unchanged. + for v in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[v] = Float32(0.0) + else: + cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, 0, s)], tTR_rAcc) + EPILOGUE( + tTR_rAcc, + tTR_cAcc[(None, None, None, 0, s)], + tiled_copy_t2r, + epi_tidx, + s, + tile, + epi_smem, + out, + cfg, + ) + + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + tmem_alloc.relinquish_alloc_permit() + epilogue_barrier.arrive_and_wait() + tmem_alloc.free(acc_tmem_ptr) + + +@cute.jit +def launch_grouped_gemm( + mA: cute.Tensor, + mB: cute.Tensor, + sfa_flat: cute.Tensor, + sfb_flat: cute.Tensor, + offs: cute.Tensor, + out, + stream, + cfg: cutlass.Constexpr, + EPILOGUE: cutlass.Constexpr, + EPI_SMEM_BYTES: cutlass.Constexpr = 0, + VALIDATE_OFFSETS: cutlass.Constexpr = True, +): + """Build the four static TMA descriptors and launch. Grid is data-independent. + + ``mA`` is the GEMM-domain ``(M, K, L)`` operand and ``mB`` the ``(N, K, L)`` + one, both K-major. ``sfa_flat`` / ``sfb_flat`` are the flat blocked E8M0 + buffers; they are retiled here, not by the caller. ``out`` is whatever tuple + of destinations the epilogue expects. + + This is a trace body, not a launcher. Calling it directly retraces the whole + kernel on every invocation -- 130 ms measured at these shapes. A public + launcher must wrap it in its own ``@cute.jit`` entry point taking only the + dynamic tensors, ``cute.compile`` that once behind a ``functools.cache``, and + call the compiled executor (35 us). The Constexpr arguments must not be + passed again to that executor; hand it the dynamic arguments only, or it + raises "cannot be converted to pointer". + + ``VALIDATE_OFFSETS`` emits the device-side precondition check on the offset + *values*, which the host cannot see without synchronizing. Note what it does + and does not buy: ``cute.testing.assert_`` is compiled out entirely unless + ``CUTE_DSL_ENABLE_ASSERTIONS=1`` is set in the environment, and when it does + fire it traps the kernel and leaves the CUDA context unusable + (``unspecified launch failure``) rather than raising cleanly. It is a + debugging aid, not a guardrail; the host validators are the guardrail. + """ + a_dtype = mA.element_type + b_dtype = mB.element_type + sf_dtype = cutlass.Float8E8M0FNU + + gemm_m = cutlass.const_expr(cute.size(mA, mode=[0])) + gemm_k = cutlass.const_expr(cute.size(mA, mode=[1])) + gemm_n = cutlass.const_expr(cute.size(mB, mode=[0])) + l_a = cutlass.const_expr(cute.size(mA, mode=[2])) + l_b = cutlass.const_expr(cute.size(mB, mode=[2])) + num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) + + if cutlass.const_expr(gemm_m % cfg.cta_tile_m != 0): + raise ValueError(f"GEMM M {gemm_m} must be a multiple of {cfg.cta_tile_m}") + if cutlass.const_expr(gemm_n % cfg.cta_tile_n != 0): + raise ValueError(f"GEMM N {gemm_n} must be a multiple of {cfg.cta_tile_n}") + if cutlass.const_expr(gemm_k % cfg.cta_tile_k != 0): + raise ValueError(f"GEMM K {gemm_k} must be a multiple of {cfg.cta_tile_k}") + + mSFA = make_sf_gemm_tensor(sfa_flat, gemm_m, gemm_k, l_a, cfg.sf_vec_size) + mSFB = make_sf_gemm_tensor(sfb_flat, gemm_n, gemm_k, l_b, cfg.sf_vec_size) + + tiled_mma = make_tiled_mma(cfg, a_dtype, b_dtype, sf_dtype) + mma_tiler = cfg.mma_tiler_mnk + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*cfg.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) + ) + + a_smem_layout = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler, a_dtype, cfg.num_ab_stage + ) + b_smem_layout = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler, b_dtype, cfg.num_ab_stage + ) + sfa_smem_layout = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, mma_tiler, cfg.sf_vec_size, cfg.num_ab_stage + ) + sfb_smem_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, mma_tiler, cfg.sf_vec_size, cfg.num_ab_stage + ) + + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + sm100_utils.cluster_shape_to_tma_atom_A(cfg.cluster_shape_mn, tiled_mma.thr_id), + mA, + cute.slice_(a_smem_layout, (None, None, None, 0)), + mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + sm100_utils.cluster_shape_to_tma_atom_B(cfg.cluster_shape_mn, tiled_mma.thr_id), + mB, + cute.slice_(b_smem_layout, (None, None, None, 0)), + mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + # The 512-byte scale atom is contiguous; TMA must move it as 8-byte elements. + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sm100_utils.cluster_shape_to_tma_atom_A(cfg.cluster_shape_mn, tiled_mma.thr_id), + mSFA, + cute.slice_(sfa_smem_layout, (None, None, None, 0)), + mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sm100_utils.cluster_shape_to_tma_atom_SFB( + cfg.cluster_shape_mn, tiled_mma.thr_id + ), + mSFB, + cute.slice_(sfb_smem_layout, (None, None, None, 0)), + mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # Round UP: flooring would hand back fewer bytes than the epilogue asked for + # and its last store would land in sA. + epi_words = cutlass.const_expr(max((EPI_SMEM_BYTES + 3) // 4, 1)) + + @cute.struct + class SharedStorage: + ab_full_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_ab_stage] + ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_ab_stage] + acc_full_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_acc_stage] + acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_acc_stage] + tmem_holding_buf: cutlass.Int32 + sEpi: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, epi_words], 128] + sA: cute.struct.Align[ + cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024 + ] + sB: cute.struct.Align[ + cute.struct.MemRange[b_dtype, cute.cosize(b_smem_layout.outer)], 1024 + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[sf_dtype, cute.cosize(sfa_smem_layout)], 1024 + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[sf_dtype, cute.cosize(sfb_smem_layout)], 1024 + ] + + smem_bytes = cutlass.const_expr(SharedStorage.size_in_bytes()) + if cutlass.const_expr(smem_bytes > SMEM_CAPACITY_BYTES): + # The struct's 1024-byte operand alignment costs a little more than the + # config's arithmetic, so check the real number rather than the estimate. + raise ValueError( + f"shared memory request {smem_bytes} B exceeds the sm_100 capacity " + f"{SMEM_CAPACITY_BYTES} B: lower num_ab_stage (currently " + f"{cfg.num_ab_stage}) or the epilogue's {EPI_SMEM_BYTES} B request" + ) + + if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.M): + grid = (gemm_m // cfg.cta_tile_m, gemm_n // cfg.cta_tile_n, 1) + else: + grid = (gemm_m // cfg.cta_tile_m, gemm_n // cfg.cta_tile_n, num_groups) + + grouped_gemm_kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + offs, + out, + a_smem_layout, + b_smem_layout, + sfa_smem_layout, + sfb_smem_layout, + cfg, + SharedStorage, + EPILOGUE, + EPI_SMEM_BYTES, + VALIDATE_OFFSETS, + ).launch( + grid=grid, + block=(cfg.threads, 1, 1), + cluster=(*cfg.cluster_shape_mn, 1), + smem=smem_bytes, + stream=stream, + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py new file mode 100644 index 0000000000..1e2d7332cc --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py @@ -0,0 +1,448 @@ +# 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. + +"""Device-side epilogue primitives shared by the MXFP8 grouped-MLP kernels. + +These are the pieces the three grouped kernels (FC1 GEMM + SwiGLU + dual quant, +FC2 dgrad + dSwiGLU + dual quant, and grouped wgrad) have in common: blocked +scale addressing, the RCEIL E8M0 conversion, NaN-propagating packed amax, and +the gated-activation policy. They are lifted from the activation-only gated +kernel, whose numerics are already validated bitwise against the standalone +quantizers. + +NUMERICAL PRECONDITION (load-bearing, read before reusing anything here): + + :func:`float_to_e8m0` is integer bit math whose rounding constant assumes + its input is a BF16 value widened to FP32. It is exact for a BF16 amax and + is NOT exact for a general FP32 amax -- for example FP32 0x40600001 gives + 120 where the canonical conversion gives 121. + +A fused GEMM epilogue holds FP32 accumulators, so it must round to BF16 *before* +taking the amax, which is what the kernel contract already requires at every +activation boundary. Do not "optimize" that rounding away while keeping this +conversion; if an FP32-amax path is ever wanted, use the canonical +``cute_utils.compute_scale_rceil`` (a real ``cvt.rp`` instruction) instead. + +Scale semantics follow torchao #4725: RCEIL with saturation disabled, a +non-finite amax yielding scale byte 255 and an all-NaN block, a zero block +yielding scale byte 0 (which dequantizes to 2^-127, not 1.0), and +``inv_scale = e8m0(254 - byte)``. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +__all__ = [ + "SCALE_BLOCK", + "SCALE_TILE_ROWS", + "SCALE_TILE_COLS", + "SCALE_TILE_BYTES", + "blocked_scale_idx", + "float_to_e8m0", + "e8m0_reciprocal_bf16", + "max_nan_bf16x2", + "abs_max_nan_bf16x2", + "fold_amax", + "pack_bf16x2", + "bf16x2_lo_to_f32", + "bf16x2_hi_to_f32", + "mul_cvt_2x", + "prmt_even", + "prmt_odd", + "sigmoidf", + "silu_pair", + "validate_group_offsets_device", +] + +# MXFP8 scaling block: 32 values share one E8M0 scale. +SCALE_BLOCK = 32 +# tcgen05 blocked scale tile geometry: 128 scale rows x 4 scale columns = 512 B. +SCALE_TILE_ROWS = 128 +SCALE_TILE_COLS = 4 +SCALE_TILE_BYTES = SCALE_TILE_ROWS * SCALE_TILE_COLS +# Every TMA-accessed shared-memory buffer must be 128-byte aligned. +TMA_SHMEM_ALIGNMENT = 128 + + +def blocked_scale_idx(row, scale_col, num_scale_col_blocks): + """Flat index of one scale byte in the tcgen05 blocked (128x4) layout. + + The logical ``[rows, cols/32]`` scale matrix is stored as 512-byte tiles of + 128 rows x 4 scale columns (cuBLAS "128x4 block scaling factors layout"), + tiles ordered ``row_block * num_scale_col_blocks + col_block``. + ``num_scale_col_blocks`` is ``ceil_div(num_scale_cols, 4)``. + + Coordinates are ABSOLUTE, never per-group. That is legal for both + orientations this family emits: + + * rowwise, where the ragged axis is the scale row -- because every expert + row count is a multiple of 128, no 128-row tile straddles a group + boundary, so per-group blocking and whole-matrix blocking are the same + bytes; + * columnwise, where the ragged axis is the scale column -- because this + family defines those buffers as whole-matrix ``to_blocked`` (see the + kernel contract 4.2.1). Note this deliberately differs from torchao's + ``triton_mx_block_rearrange_2d_K_groups``, which pads per group. + + For columnwise scales pass transposed coordinates (feature index as ``row``, + row-block index as ``scale_col``). + """ + return ( + ((row >> 7) * num_scale_col_blocks + (scale_col >> 2)) * SCALE_TILE_BYTES + + (row & 31) * 16 + + ((row >> 5) & 3) * 4 + + (scale_col & 3) + ) + + +@dsl_user_op +def _bitcast_i32_to_f32(val: Int32, *, loc=None, ip=None) -> Float32: + """Bitcast int32 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) + + +# Each packed-bf16x2 op below is written out explicitly rather than produced by a +# factory: the DSL keys its compile cache on function identity and name, so ops +# sharing a `__name__` are a cache hazard. +# +# The `.NaN` max variants match the standalone quantizers' amax reduction, which +# propagates NaN; a plain max would return the non-NaN operand and silently +# rescue a block that must be invalidated. +@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, + ) + ) + + +@cute.jit +def fold_amax(am: Int32) -> Int32: + """Reduce a packed bf16x2 amax word to bf16 amax bits in [15:0]. + + The input's per-lane sign bits are junk (see :func:`abs_max_nan_bf16x2`); + mask them, then fold the two lanes with the NaN-propagating max. + """ + am = am & Int32(0x7FFF7FFF) + am = max_nan_bf16x2(am, am >> 16) + return am & Int32(0xFFFF) + + +@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. Saturation to +/-448 comes from + this conversion; do not add an explicit clamp. + """ + 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, round-to-nearest-even. + + This is the mandatory "truncate to BF16 before any amax" step for a GEMM + epilogue holding FP32 accumulators; see the module docstring. + """ + 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 RCEIL scale byte for a non-negative BF16 amax, as f32 bits. + + Pass the RAW amax. The division by 448 is folded into the constants -- the + ``- 8`` shifts the exponent by 256 and the ``+ 0x1F0000`` mantissa offset + supplies the remaining 1.75 factor (448 = 256 * 1.75) together with the + RCEIL round-up carry. Pre-dividing by 448 before calling this double-counts + the division and yields a scale 8-9 codes too low. + + Finite: the RCEIL mantissa-carry path, matching what ``cvt.rp.ue8m0x2.f32`` + (no ``.satfinite``) emits for ``amax / 448``. Non-finite: a 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. + + Exact only for a BF16-valued amax -- see the module docstring. + """ + 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 e8m0_reciprocal_bf16(e: Int32) -> Int32: + """Inverse scale as bf16 bits, matching ``ue8m0(254 - scale_byte)``. + + The quantization multiply in :func:`mul_cvt_2x` is bf16x2, not f32, so the + reciprocal is synthesized directly in bf16: 2^(127 - e) over the normal + range (a byte-0 block from a zero or tiny amax descales by 2^127), and NaN + for an invalidated block (byte 255) so every element quantizes to the E4M3 + NaN code. + + This is NOT a general ``ue8m0`` helper: byte 254 yields +0.0 rather than + 2^-127, and bytes above 254 go negative. That is safe here only because a + scale byte produced by this family can never exceed 247 (amax is divided by + 448 first). Do not reuse it to dequantize externally produced scale bytes. + """ + 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 / dSwiGLU policy for a pair of elements. + + ``x`` is the gate (activation input), ``lin`` the up (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 -> dGate + out_gate = act * g -> dUp + + Returns f32 ``(out_act0, out_act1, out_gate0, out_gate1)``, the gate pair + zero in forward. Callers MUST round to BF16 immediately, before any amax, + caching, or quantization -- both because the kernel contract defines + correctness at that boundary and because :func:`float_to_e8m0` requires it. + + Pass this as a Constexpr kernel parameter. It must stay a module-level + function: the DSL keys its compile cache on the function object, so a lambda + or closure built per call misses the cache and recompiles every launch. + """ + one = 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, Float32(0.0), Float32(0.0) + + +@cute.jit +def validate_group_offsets_device(offs: cute.Tensor, allocated_rows: Int32): + """Device-side precondition check on the exclusive-end group offsets. + + Checks what the host cannot see without synchronizing: every per-expert row + count is a nonnegative multiple of 128, the offsets are nondecreasing, and + the active row count does not exceed the allocation. + + DO NOT RELY ON THIS AS A GUARDRAIL. ``cute.testing.assert_`` is compiled out + unless ``CUTE_DSL_ENABLE_ASSERTIONS=1``, so in a default build this function + is a no-op -- measured directly: an always-false assertion in four + placements (plain, warp-0, elected, block-0-elected) lets the kernel run to + completion and write its output. torchao's own ``validate_group_sizes`` has + the same property. Even with assertions enabled the failure mode is the + message followed by ``unspecified launch failure``, i.e. a dead CUDA + context, not a catchable error. + + The real enforcement of the 128-multiple precondition is the host-side + metadata validation in ``grouped_mlp_validation`` at the custom-op boundary, + which raises ``ValueError`` before any launch. This function is a debugging + aid for assertion-enabled builds. It matters that the distinction is + explicit: with malformed offsets the ragged-K path (Kernel C) silently + returns a WRONG weight gradient rather than crashing, so "it did not fault" + is not evidence that the offsets were valid. + """ + num_groups = offs.shape[0] + prev = Int32(0) + for i in range(num_groups): + end = offs[i] + size = end - prev + cute.testing.assert_(size >= 0, "Group offsets must be nondecreasing") + cute.testing.assert_(size % 128 == 0, "Group sizes must be multiples of 128") + prev = end + cute.testing.assert_( + prev <= allocated_rows, + "Active row count offsets[-1] must not exceed the allocated row count", + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py new file mode 100644 index 0000000000..c608d78c68 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py @@ -0,0 +1,451 @@ +# 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. + +"""Public custom-op surface for the MXFP8 routed-expert grouped-MLP kernels. + +Three ops, one per fused kernel: + +* ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` -- FC1 grouped GEMM + SwiGLU + + rowwise/columnwise MXFP8 quantization +* ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` -- FC2 dgrad grouped GEMM + + dSwiGLU + rowwise/columnwise MXFP8 quantization +* ``torchao::mxfp8_grouped_gemm_wgrad`` -- grouped MXFP8 weight-gradient + GEMM, invoked once for FC1 and once for FC2 + +Each op allocates its destinations through the ``_allocate_*_outputs`` helpers +below, which are pure torch and are shared by the real implementation and by +``register_fake``. Meta shapes and strides therefore cannot drift from eager -- +a class of bug that matters here because several outputs are column-major and a +row-major fake would silently change what ``torch.compile`` traces. + +Layout, alignment and numerical semantics are defined by the kernel contract; +the load-bearing preconditions (every per-expert row count a multiple of 128, +exact strides, blocked E8M0 scales in the tcgen05 128x4 layout) are enforced by +``grouped_mlp_validation`` before any launch, because violating them otherwise +corrupts the CUDA context rather than raising. +""" + +from typing import Tuple + +import torch + +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( + blocked_scale_numel, + validate_allocated_rows, + validate_blocked_scales, + validate_feature_dims, + validate_group_offsets, + validate_grouped_operand, +) + +__all__ = [ + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_wgrad", +] + +_E4M3 = torch.float8_e4m3fn +_E8M0 = torch.float8_e8m0fnu +_SCALE_BLOCK = 32 + + +def _empty_blocked_scales( + logical_rows: int, logical_cols: int, *, device, groups: int = 1 +) -> torch.Tensor: + """Allocate a flat blocked E8M0 scale buffer. + + The buffer is flat by ABI: its logical shape is metadata. Kernels write every + byte, including the inactive-tail rows, so an uninitialized allocation is + safe here -- but only because that write obligation is part of the contract. + """ + numel = groups * blocked_scale_numel(logical_rows, logical_cols) + shape = (groups, numel // groups) if groups > 1 else (numel,) + return torch.empty(shape, dtype=_E8M0, device=device) + + +# -------------------------------------------------------------------------- +# Kernel A: FC1 grouped GEMM + SwiGLU + dual quantization +# -------------------------------------------------------------------------- + + +def _allocate_swiglu_fwd_outputs( + rows: int, hidden: int, device +) -> Tuple[torch.Tensor, ...]: + z = torch.empty_strided( + (rows, hidden, 2), (2 * hidden, 2, 1), dtype=torch.bfloat16, device=device + ) + h_row_q = torch.empty_strided( + (rows, hidden), (hidden, 1), dtype=_E4M3, device=device + ) + h_row_sf = _empty_blocked_scales(rows, hidden // _SCALE_BLOCK, device=device) + # Column-major: the 32x1 quantized operand is consumed as its own transpose. + h_col_q = torch.empty_strided((rows, hidden), (1, rows), dtype=_E4M3, device=device) + h_col_sf = _empty_blocked_scales(hidden, rows // _SCALE_BLOCK, device=device) + return z, h_row_q, h_row_sf, h_col_q, h_col_sf + + +def _validate_swiglu_fwd_inputs( + x_q, x_sf, w13_t_q, w13_t_sf, offsets +) -> Tuple[int, int, int, int]: + if x_q.ndim != 2: + raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") + if w13_t_q.ndim != 3: + raise ValueError( + f"w13_t_q must be 3D [G, D, 2F], got shape {tuple(w13_t_q.shape)}" + ) + rows, model_dim = x_q.shape + groups, w_k, two_hidden = w13_t_q.shape + if w_k != model_dim: + raise ValueError( + f"w13_t_q contraction dim {w_k} must match x_q's D {model_dim}" + ) + if two_hidden % 2 != 0: + raise ValueError( + f"w13_t_q's N dim must be 2F with interleaved gate/up channels, got {two_hidden}" + ) + hidden = two_hidden // 2 + device = x_q.device + + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + validate_grouped_operand( + x_q, + name="x_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + w13_t_q, + name="w13_t_q", + shape=(groups, model_dim, two_hidden), + stride=(model_dim * two_hidden, 1, model_dim), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + x_sf, + name="x_sf", + logical_rows=rows, + logical_cols=model_dim // _SCALE_BLOCK, + device=device, + ) + validate_blocked_scales( + w13_t_sf, + name="w13_t_sf", + logical_rows=two_hidden, + logical_cols=model_dim // _SCALE_BLOCK, + device=device, + groups=groups, + ) + return rows, model_dim, hidden, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_swiglu_fwd", mutates_args=()) +def _mxfp8_grouped_gemm_swiglu_fwd( + x_q: torch.Tensor, + x_sf: torch.Tensor, + w13_t_q: torch.Tensor, + w13_t_sf: torch.Tensor, + offsets: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rows, _model_dim, hidden, _groups = _validate_swiglu_fwd_inputs( + x_q, x_sf, w13_t_q, w13_t_sf, offsets + ) + outputs = _allocate_swiglu_fwd_outputs(rows, hidden, x_q.device) + + if rows == 0: + # R == 0 is a required correctness case. Every destination is empty in + # its row dimension and both scale buffers are zero-length, so there is + # nothing to write; launching would build a degenerate (0, D, 1) layout + # and fail inside the SF layout builder with an opaque MLIR error. + return outputs + + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + launch_grouped_gemm_swiglu_fwd, + ) + + launch_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_t_q, w13_t_sf, offsets, *outputs) + return outputs + + +@_mxfp8_grouped_gemm_swiglu_fwd.register_fake +def _(x_q, x_sf, w13_t_q, w13_t_sf, offsets): + rows, _model_dim, hidden, _groups = _validate_swiglu_fwd_inputs( + x_q, x_sf, w13_t_q, w13_t_sf, offsets + ) + return _allocate_swiglu_fwd_outputs(rows, hidden, x_q.device) + + +# -------------------------------------------------------------------------- +# Kernel B: FC2 dgrad grouped GEMM + dSwiGLU + dual quantization +# -------------------------------------------------------------------------- + + +def _allocate_dswiglu_bwd_outputs( + rows: int, hidden: int, device +) -> Tuple[torch.Tensor, ...]: + two_hidden = 2 * hidden + dz_row_q = torch.empty_strided( + (rows, two_hidden), (two_hidden, 1), dtype=_E4M3, device=device + ) + dz_row_sf = _empty_blocked_scales(rows, two_hidden // _SCALE_BLOCK, device=device) + dz_col_q = torch.empty_strided( + (rows, two_hidden), (1, rows), dtype=_E4M3, device=device + ) + dz_col_sf = _empty_blocked_scales(two_hidden, rows // _SCALE_BLOCK, device=device) + return dz_row_q, dz_row_sf, dz_col_q, dz_col_sf + + +def _validate_dswiglu_bwd_inputs(do_q, do_sf, w2_q, w2_sf, z_bf16, offsets): + if do_q.ndim != 2: + raise ValueError(f"do_q must be 2D [R, D], got shape {tuple(do_q.shape)}") + if w2_q.ndim != 3: + raise ValueError( + f"w2_dgrad_q must be 3D [G, D, F], got shape {tuple(w2_q.shape)}" + ) + if z_bf16.ndim != 3 or z_bf16.shape[-1] != 2: + raise ValueError( + f"z_bf16 must be [R, F, 2] with gate at index 0 and up at index 1, " + f"got shape {tuple(z_bf16.shape)}" + ) + rows, model_dim = do_q.shape + groups, w_k, hidden = w2_q.shape + if w_k != model_dim: + raise ValueError( + f"w2_dgrad_q contraction dim {w_k} must match do_q's D {model_dim}" + ) + if tuple(z_bf16.shape) != (rows, hidden, 2): + raise ValueError( + f"z_bf16 must be [{rows}, {hidden}, 2] to match do_q and w2_dgrad_q, " + f"got {tuple(z_bf16.shape)}" + ) + device = do_q.device + + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + validate_grouped_operand( + do_q, + name="do_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + w2_q, + name="w2_dgrad_q", + shape=(groups, model_dim, hidden), + stride=(model_dim * hidden, 1, model_dim), + dtype=_E4M3, + device=device, + ) + # z_bf16 is the exact destination Kernel A wrote, so its stride is pinned too. + validate_grouped_operand( + z_bf16, + name="z_bf16", + shape=(rows, hidden, 2), + stride=(2 * hidden, 2, 1), + dtype=torch.bfloat16, + device=device, + ) + validate_blocked_scales( + do_sf, + name="do_sf", + logical_rows=rows, + logical_cols=model_dim // _SCALE_BLOCK, + device=device, + ) + validate_blocked_scales( + w2_sf, + name="w2_dgrad_sf", + logical_rows=hidden, + logical_cols=model_dim // _SCALE_BLOCK, + device=device, + groups=groups, + ) + return rows, model_dim, hidden, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_dswiglu_bwd", mutates_args=()) +def _mxfp8_grouped_gemm_dswiglu_bwd( + do_q: torch.Tensor, + do_sf: torch.Tensor, + w2_dgrad_q: torch.Tensor, + w2_dgrad_sf: torch.Tensor, + z_bf16: torch.Tensor, + offsets: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rows, _model_dim, hidden, _groups = _validate_dswiglu_bwd_inputs( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets + ) + outputs = _allocate_dswiglu_bwd_outputs(rows, hidden, do_q.device) + + if rows == 0: + # See the R == 0 note in the forward op: nothing to write, and launching + # would build a degenerate SF layout. + return outputs + + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + launch_grouped_gemm_dswiglu_bwd, + ) + + launch_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets, *outputs + ) + return outputs + + +@_mxfp8_grouped_gemm_dswiglu_bwd.register_fake +def _(do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets): + rows, _model_dim, hidden, _groups = _validate_dswiglu_bwd_inputs( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets + ) + return _allocate_dswiglu_bwd_outputs(rows, hidden, do_q.device) + + +# -------------------------------------------------------------------------- +# Kernel C: grouped MXFP8 wgrad +# -------------------------------------------------------------------------- + + +def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + if dy_col_q.ndim != 2 or x_col_q.ndim != 2: + raise ValueError( + "dy_col_q and x_col_q must both be 2D logical [R, N] / [R, K], got " + f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" + ) + rows, out_features = dy_col_q.shape + x_rows, in_features = x_col_q.shape + if x_rows != rows: + raise ValueError( + f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" + ) + groups = offsets.numel() + device = dy_col_q.device + + validate_allocated_rows(rows) + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + # Both operands are column-major so their transposes are free. + validate_grouped_operand( + dy_col_q, + name="dy_col_q", + shape=(rows, out_features), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + x_col_q, + name="x_col_q", + shape=(rows, in_features), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + dy_col_sf, + name="dy_col_sf", + logical_rows=out_features, + logical_cols=rows // _SCALE_BLOCK, + device=device, + ) + validate_blocked_scales( + x_col_sf, + name="x_col_sf", + logical_rows=in_features, + logical_cols=rows // _SCALE_BLOCK, + device=device, + ) + return rows, out_features, in_features, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_wgrad", mutates_args=()) +def _mxfp8_grouped_gemm_wgrad( + dy_col_q: torch.Tensor, + dy_col_sf: torch.Tensor, + x_col_q: torch.Tensor, + x_col_sf: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + rows, out_features, in_features, groups = _validate_wgrad_inputs( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) + dw = torch.empty( + (groups, out_features, in_features), + dtype=torch.bfloat16, + device=dy_col_q.device, + ) + + if rows == 0: + # Unlike A and B, this destination is NOT empty at R == 0: every expert + # has zero rows, and an expert with zero rows is defined to produce an + # all-zero slice. Zero it here rather than launching over an empty + # contraction. + return dw.zero_() + + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + launch_grouped_gemm_wgrad, + ) + + launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw) + return dw + + +@_mxfp8_grouped_gemm_wgrad.register_fake +def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + _rows, out_features, in_features, groups = _validate_wgrad_inputs( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) + return torch.empty( + (groups, out_features, in_features), + dtype=torch.bfloat16, + device=dy_col_q.device, + ) + + +# -------------------------------------------------------------------------- +# Public wrappers +# -------------------------------------------------------------------------- + + +def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_t_q, w13_t_sf, offsets): + """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. + + Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``. ``z_bf16`` is the + BF16 pre-activation saved for backward; pass it unchanged to + :func:`mxfp8_grouped_gemm_dswiglu_bwd`. + """ + return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_t_q, w13_t_sf, offsets + ) + + +def mxfp8_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets +): + """FC2 dgrad grouped GEMM + dSwiGLU + rowwise/columnwise MXFP8 quantization. + + Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)`` with gate/up gradients + element-interleaved along the 2F axis. + """ + return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets + ) + + +def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + """Grouped MXFP8 weight-gradient GEMM, returning BF16 ``[G, N, K]``. + + Used once for FC1 (``N=2F, K=D``) and once for FC2 (``N=D, K=F``). An expert + with zero rows yields an all-zero slice. + """ + return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py new file mode 100644 index 0000000000..c9671d8ed6 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py @@ -0,0 +1,245 @@ +"""Host-side precondition validation for the MXFP8 routed-expert grouped-MLP kernels. + +The grouped MXFP8 kernels require every per-expert row count to be a multiple of +128. That is not a convenience: the tcgen05 blocked scale layout permutes in +128-row tiles, so a group boundary off a 128 multiple splits a tile and the +blocked buffer for the group no longer matches what the GEMM reads. Feeding such +offsets to the existing CuTe DSL quantizer produces a device-side assertion, a +wrong-sized scale buffer, and an unusable CUDA context rather than a clean error. + +Everything here is metadata-only so it costs no host/device synchronization and +stays traceable under torch.compile. The per-expert counts live in device memory +and are validated on device by the kernels themselves; set +TORCHAO_MXFP8_VALIDATE_OFFSETS=1 to additionally check them on the host while +debugging, at the cost of a D2H copy. + +Checks raise ValueError rather than asserting, so `python -O` cannot strip them. +""" + +import os + +import torch + +__all__ = [ + "SCALE_BLOCK_SIZE", + "SCALE_TILE_ROWS", + "SCALE_TILE_COLS", + "blocked_scale_numel", + "host_offsets_validation_enabled", + "validate_group_offsets", + "validate_grouped_operand", + "validate_blocked_scales", + "validate_destination", +] + +# MXFP8 scaling block: 32 values share one E8M0 scale. +SCALE_BLOCK_SIZE = 32 +# tcgen05 blocked scale tile: 128 rows x 4 columns, 512 bytes. +SCALE_TILE_ROWS = 128 +SCALE_TILE_COLS = 4 +# Row-count granularity every expert group must respect. +GROUP_ALIGNMENT = 128 +# Byte alignment the launchers promise for TMA/vectorized accesses. +_PTR_ALIGNMENT = 32 + + +def _round_up(x: int, to: int) -> int: + return ((x + to - 1) // to) * to + + +def blocked_scale_numel(rows: int, cols: int) -> int: + """Element count of the blocked E8M0 buffer for a logical [rows, cols] scale matrix. + + `cols` is a count of scale values, i.e. the reduced dimension divided by 32. + """ + return _round_up(rows, SCALE_TILE_ROWS) * _round_up(cols, SCALE_TILE_COLS) + + +def host_offsets_validation_enabled() -> bool: + """Opt-in host-side offset validation. Off by default: it forces a D2H sync.""" + return os.environ.get("TORCHAO_MXFP8_VALIDATE_OFFSETS", "0") == "1" + + +def validate_group_offsets( + offsets: torch.Tensor, + *, + num_groups: int, + allocated_rows: int, + name: str = "offsets", +) -> None: + """Validate the exclusive-end group offsets tensor. + + Metadata is always checked. The offset *values* are checked only when + host_offsets_validation_enabled(), because reading them synchronizes; the + kernels assert the same invariants on device on every launch. + """ + if not isinstance(offsets, torch.Tensor): + raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") + if offsets.dtype != torch.int32: + raise ValueError(f"{name} must be int32, got {offsets.dtype}") + if not offsets.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") + if offsets.ndim != 1: + raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") + if offsets.numel() != num_groups: + raise ValueError( + f"{name} must have one entry per local expert: expected {num_groups}, " + f"got {offsets.numel()}" + ) + if not offsets.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {offsets.stride()}") + + if not host_offsets_validation_enabled(): + return + + values = offsets.tolist() # d2h sync; opt-in debugging path only + previous = 0 + for group, end in enumerate(values): + if end < previous: + raise ValueError( + f"{name} must be nondecreasing, but entry {group} is {end} after {previous}" + ) + size = end - previous + if size % GROUP_ALIGNMENT != 0: + raise ValueError( + f"per-expert row counts must be multiples of {GROUP_ALIGNMENT}: " + f"expert {group} has {size} rows (offsets {previous} -> {end})" + ) + previous = end + if previous > allocated_rows: + raise ValueError( + f"{name}[-1] ({previous}) exceeds the allocated row count ({allocated_rows})" + ) + + +def validate_grouped_operand( + tensor: torch.Tensor, + *, + name: str, + shape: tuple, + stride: tuple, + dtype: torch.dtype, + device: torch.device, + check_pointer_alignment: bool = True, +) -> None: + """Validate one quantized operand's dtype, shape, exact stride, device, alignment. + + Order matters: every metadata gate runs before the data_ptr() gate so that + FakeTensor tracing exercises the same checks (a fake tensor has no pointer). + """ + if tensor.dtype != dtype: + raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != tuple(shape): + raise ValueError( + f"{name} must have shape {tuple(shape)}, got {tuple(tensor.shape)}" + ) + if tuple(tensor.stride()) != tuple(stride): + raise ValueError( + f"{name} must have stride {tuple(stride)}, got {tuple(tensor.stride())}. " + "This layout is part of the ABI; a values-equal tensor with a different " + "stride is not interchangeable." + ) + if tensor.device != device: + raise ValueError( + f"{name} must be on {device}, got {tensor.device}; all operands and " + "destinations must share one CUDA device" + ) + if check_pointer_alignment and not _is_fake(tensor): + if tensor.data_ptr() % _PTR_ALIGNMENT != 0: + raise ValueError( + f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data pointer is " + f"{tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an aligned address. A " + "contiguous view with a nonzero storage offset can violate this." + ) + + +def validate_blocked_scales( + scales: torch.Tensor, + *, + name: str, + logical_rows: int, + logical_cols: int, + device: torch.device, + groups: int = 1, +) -> None: + """Validate a blocked E8M0 scale buffer's dtype, element count, and device. + + The buffer is carried flat: its logical shape is metadata, not its physical + shape, so only the element count is constrained. `groups` > 1 describes the + per-expert weight buffers, which are [G, per_group_numel]. + """ + if scales.dtype not in (torch.uint8, torch.float8_e8m0fnu): + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), got {scales.dtype}" + ) + expected = groups * blocked_scale_numel(logical_rows, logical_cols) + if scales.numel() != expected: + raise ValueError( + f"{name} must hold {expected} blocked scale bytes for a logical " + f"[{logical_rows}, {logical_cols}] scale matrix" + + (f" across {groups} experts" if groups > 1 else "") + + f", got {scales.numel()}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + + +def validate_destination( + tensor: torch.Tensor, + *, + name: str, + shape: tuple, + stride: tuple, + dtype: torch.dtype, + device: torch.device, +) -> None: + """Validate a caller-allocated destination in the destination-passing entry points. + + Destinations are validated exactly like inputs. Skipping this is how a private + entry point turns a caller's shape mistake into an out-of-bounds write. + """ + validate_grouped_operand( + tensor, + name=name, + shape=shape, + stride=stride, + dtype=dtype, + device=device, + ) + + +def validate_feature_dims(*, model_dim: int, hidden_dim: int) -> None: + """D and F must both be multiples of 128 for the initial supported predicate.""" + if model_dim % GROUP_ALIGNMENT != 0: + raise ValueError( + f"model dimension D must be a multiple of {GROUP_ALIGNMENT}, got {model_dim}" + ) + if hidden_dim % GROUP_ALIGNMENT != 0: + raise ValueError( + f"routed-expert hidden dimension F must be a multiple of {GROUP_ALIGNMENT}, " + f"got {hidden_dim}" + ) + + +def validate_allocated_rows(rows: int, *, name: str = "R") -> None: + """The allocated row count must itself be 128-aligned. + + Group sizes are multiples of 128 and the active row count is their sum, so a + non-128 allocation can only describe an inactive tail that no legal offsets + vector can reach; rejecting it early keeps the tail contract simple. + """ + if rows % GROUP_ALIGNMENT != 0: + raise ValueError(f"{name} must be a multiple of {GROUP_ALIGNMENT}, got {rows}") + + +def _is_fake(tensor: torch.Tensor) -> bool: + """True for meta/fake tensors, which have no usable data pointer.""" + if tensor.device.type == "meta": + return True + try: + from torch._subclasses.fake_tensor import FakeTensor + except ImportError: + return False + return isinstance(tensor, FakeTensor) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py b/torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py new file mode 100644 index 0000000000..0050c4033d --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py @@ -0,0 +1,347 @@ +# 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. + +"""Kernel C: grouped MXFP8 weight-gradient GEMM (``mxfp8_grouped_gemm_wgrad``). + +One Definition covers both call sites -- FC1 (``N = 2F``, ``K = D``) and FC2 +(``N = D``, ``K = F``) -- with no mode flag; the shapes come from the tensors. + +Per expert ``g`` over rows ``[offsets[g-1], offsets[g])``:: + + dw[g] = dequant(dy_col[rows]).T @ dequant(x_col[rows]) + +FP32 accumulation, BF16 output. Neither transpose is materialized: both operands +arrive logically ``[R, N]`` / ``[R, K]`` with stride ``(1, R)``, which *is* a +K-contiguous row-major ``[N, R]`` / ``[K, R]``, so the free transpose is a +restride on the host and the ragged axis lands on the GEMM's contraction. The +expert is then an integer K-tile index base rather than a per-expert TMA +descriptor, exact because every per-expert row count is a multiple of +``cta_tile_k``. + +Two things about this kernel that are easy to get wrong: + +*Both scale buffers are WHOLE-MATRIX* ``to_blocked``, not torchao's per-group +K-groups form (``triton_mx_block_rearrange_2d_K_groups``). The two orderings +differ whenever ``N > 128`` -- whole-matrix orders blocked tiles by +``row_block * ncb_total + col_block``, per-group by ``row_block * ncb_g + +col_block`` within each group -- so feeding a per-group buffer here produces a +*block-permuted* ``dw``, which is large and structured and reads like a GEMM bug +rather than like a layout bug. The producers of these buffers are kernels A and +B in this same family, which emit the whole-matrix form. + +*The epilogue never predicates its store.* A zero-token expert arrives with +``k_cnt == 0``, the core hands the epilogue a zeroed register fragment, and the +unmodified store path writes the all-zero ``dw[g]`` the contract requires. The +grid enumerates every ``(tile_m, tile_n, expert)``, so every element of ``dw`` is +written on every call with no memset. There is also no gmem input load in this +epilogue, so the ``k_cnt``-gated-load half of the tail rule has nothing to cover +here. + +No output TMA and no epilogue shared memory: one row per thread means the 32 +FP32 accumulators of a subtile are 32 contiguous BF16 values == 64 naturally +aligned contiguous bytes of ``dw``, so a direct vectorized ``STG`` is both +correct and fully sector-efficient. +""" + +import functools +from typing import Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Int32 +from cutlass.cute.runtime import from_dlpack + +from torchao.prototype.moe_training.kernels.mxfp8.grouped_gemm_config import ( + SF_VEC_SIZE, + WGRAD_CONFIG, +) +from torchao.prototype.moe_training.kernels.mxfp8.grouped_gemm_core import ( + activation_gemm_view, + launch_grouped_gemm, +) +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_epilogue import ( + pack_bf16x2, +) +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( + validate_allocated_rows, + validate_blocked_scales, + validate_destination, + validate_group_offsets, + validate_grouped_operand, +) + +__all__ = ["bf16_store_epilogue", "launch_grouped_gemm_wgrad"] + +_E4M3 = torch.float8_e4m3fn +_BF16 = torch.bfloat16 +# Widest vector the epilogue's store is allowed to assume. The run is 64 B and +# 64-B aligned, so this only has to be a divisor of that. +_VEC_BYTES = 16 + + +@cute.jit +def _store_bf16_run(dst: cute.Tensor, elem_offset: Int32, words: cute.Tensor): + """Store a run of packed ``bf16x2`` words at a BF16 element offset. + + ``dst`` is any BF16 destination; the run is reinterpreted as Int32, so the + element offset must be even. Every caller's offset is a multiple of 32 + elements (see :func:`bf16_store_epilogue`), which also makes the address + 64-byte aligned and the copy four ``STG.128``. + """ + cute.autovec_copy( + words, + cute.make_tensor( + ( + cute.recast_ptr(dst.iterator, dtype=Int32) + (elem_offset >> Int32(1)) + ).align(_VEC_BYTES), + cute.make_layout(cute.size(words)), + ), + ) + + +@cute.jit +def bf16_store_epilogue( + tTR_rAcc, + tTR_cAcc_s, + tiled_copy_t2r, + epi_tidx, + subtile_idx: cutlass.Constexpr, + tile, + epi_smem, + out, + cfg: cutlass.Constexpr, +): + """Round the FP32 accumulator subtile to BF16 and store it. No quantization. + + ``out`` is ``(mDw,)`` with ``mDw`` the ``(N, K, G)`` view of the contiguous + ``[G, N, K]`` destination. + + The register-to-column map is taken from ``tTR_cAcc``, whose column + coordinates fold to Python ints at trace time, and the run is required to be + contiguous and increasing. That check is what licenses both the ``pack`` + pairing and the single vectorized store: if a thread owned more than one row + of the epilogue tile its columns would repeat instead of forming a run, so + this also establishes the one-row-per-thread property the store address + assumes, rather than trusting it. + """ + num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) + cols = [] + for v in cutlass.range_constexpr(num_acc): + cols.append(tTR_cAcc_s[v][1]) + frag_col = cols[0] + if cutlass.const_expr( + num_acc % 2 != 0 + or not all(isinstance(c, int) for c in cols) + or tuple(cols) != tuple(range(frag_col, frag_col + num_acc)) + ): + raise ValueError( + f"the wgrad epilogue needs an even, contiguous, increasing column run " + f"per thread to pack and store BF16 pairs, but tTR_cAcc gave {cols}" + ) + + # cvt.rn.bf16x2.f32 packs the second source into the low half, so column + # frag_col + 2j lands at the lower address -- row-major order for dw. + words = cute.make_rmem_tensor((num_acc // 2,), Int32) + for j in cutlass.range_constexpr(num_acc // 2): + words[j] = pack_bf16x2(tTR_rAcc[2 * j + 1], tTR_rAcc[2 * j]) + + # Address from the destination's real strides, which are static ints here, + # rather than from its extents: the only layout property the store actually + # needs is that the K axis is contiguous. + gDw = out[0] + strides = gDw.stride + if cutlass.const_expr(strides[1] != 1): + raise ValueError( + f"the wgrad epilogue stores a contiguous run along K, so dw's K " + f"stride must be 1, got layout {gDw.layout}" + ) + # Row from tTR_cAcc, not from epi_tidx: the contiguity check above proves + # the fragment is one row, and this is the coordinate that names it. + row = tile.row_base + tTR_cAcc_s[0][0] + elem = ( + row * Int32(strides[0]) + + (tile.col_base + Int32(frag_col)) + + tile.expert * Int32(strides[2]) + ) + _store_bf16_run(gDw, elem, words) + + +@cute.jit +def _wgrad_entry(mA, mB, sfa, sfb, offs, mDw, stream): + """Trace entry point: dynamic tensors only, config and epilogue closed over. + + ``launch_grouped_gemm`` is a trace body, so calling it directly retraces the + whole kernel on every launch. Everything Constexpr is bound here so that + :func:`cute.compile` can hand back an executor that takes only these + arguments; passing a Constexpr to that executor raises "cannot be converted + to pointer". + """ + launch_grouped_gemm( + mA, + mB, + sfa, + sfb, + offs, + (mDw,), + stream, + WGRAD_CONFIG, + bf16_store_epilogue, + ) + + +@functools.cache +def _wgrad_executor_slot(key: Tuple) -> list: + """One memo slot per compiled shape. + + The executor cannot be built from ``key`` alone: the shared core needs static + shapes (the blocked scale-factor layout and the grid are both built from + them), so a symbolic ``cute.sym_int`` compile is not available and the first + real call's tensors are what gets compiled. ``functools.cache`` therefore + keys the slot and the caller fills it once. + """ + return [] + + +def _validate_wgrad_operands(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw): + """Host-only precondition check. Metadata and pointers, never offset values. + + The custom op validates its own inputs, but this launcher is also reachable + directly from the DSL and it owns the destination, which the op does not + check. Every gate here is metadata-derivable, so it costs no synchronization + and stays traceable. + """ + if dy_col_q.ndim != 2 or x_col_q.ndim != 2: + raise ValueError( + "dy_col_q and x_col_q must be 2D logical [R, N] and [R, K], got " + f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" + ) + rows, out_features = dy_col_q.shape + x_rows, in_features = x_col_q.shape + if x_rows != rows: + raise ValueError( + f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" + ) + groups = offsets.numel() + device = dy_col_q.device + + validate_allocated_rows(rows) + # N and K are the GEMM's two free axes; both are tiled with no tail path, so + # reject a non-multiple here rather than from inside the trace. + for name, value, tile in ( + ("dy_col_q's N", out_features, WGRAD_CONFIG.cta_tile_m), + ("x_col_q's K", in_features, WGRAD_CONFIG.cta_tile_n), + ): + if value <= 0 or value % tile != 0: + raise ValueError( + f"{name} must be a positive multiple of {tile}, got {value}" + ) + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + # Column-major, so the transposes below are free restrides. + validate_grouped_operand( + dy_col_q, + name="dy_col_q", + shape=(rows, out_features), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + x_col_q, + name="x_col_q", + shape=(rows, in_features), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + dy_col_sf, + name="dy_col_sf", + logical_rows=out_features, + logical_cols=rows // SF_VEC_SIZE, + device=device, + ) + validate_blocked_scales( + x_col_sf, + name="x_col_sf", + logical_rows=in_features, + logical_cols=rows // SF_VEC_SIZE, + device=device, + ) + # validate_blocked_scales checks dtype, length and device but not the + # pointer, and both scale buffers are TMA operands: a contiguous view with a + # storage offset can be 2-byte aligned. Only the launcher promises the TMA + # alignment, so only the launcher can require it. + for name, buf in (("dy_col_sf", dy_col_sf), ("x_col_sf", x_col_sf)): + if buf.data_ptr() % 32 != 0: + raise ValueError( + f"{name} must be 32-byte aligned for its TMA descriptor, but its " + f"data pointer is {buf.data_ptr() % 32} bytes past an aligned " + "address" + ) + validate_destination( + dw, + name="dw_bf16", + shape=(groups, out_features, in_features), + stride=(out_features * in_features, in_features, 1), + dtype=_BF16, + device=device, + ) + # The epilogue computes its destination element index in Int32. + if groups * out_features * in_features >= 2**31: + raise ValueError( + f"dw_bf16 has {groups * out_features * in_features} elements, which " + "does not fit the epilogue's int32 element index; the store address " + "arithmetic would wrap" + ) + return rows, out_features, in_features, groups + + +def launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw): + """Grouped MXFP8 wgrad into a caller-allocated BF16 ``[G, N, K]`` destination. + + Inputs are the columnwise-quantized outputs of kernels A and B (or of the + standalone 32x1 quantizer): ``dy_col_q`` E4M3 logical ``[R, N]`` stride + ``(1, R)`` with ``dy_col_sf`` blocked for logical ``[N, R/32]``, and + ``x_col_q`` / ``x_col_sf`` likewise for ``[R, K]``. ``offsets`` is the int32 + CUDA ``[G]`` vector of exclusive group ends and is never read on the host. + + Every element of ``dw`` is written, including the all-zero slice of a + zero-token expert. + """ + import cuda.bindings.driver as cuda + + rows, out_features, in_features, groups = _validate_wgrad_operands( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw + ) + if rows == 0: + # Every expert has zero rows, so every slice is the zero matrix. Unlike + # the other two kernels the destination is NOT empty here, and the + # contraction is, so this cannot be expressed as a launch. + dw.zero_() + return + + stream = cuda.CUstream(int(torch.cuda.current_stream().cuda_stream)) + args = ( + # The free transpose: logical [R, N] stride (1, R) IS a K-contiguous + # [N, R], so both operands become ordinary K-major GEMM operands and the + # ragged axis becomes the contraction. + from_dlpack(activation_gemm_view(dy_col_q.t()), assumed_align=16), + from_dlpack(activation_gemm_view(x_col_q.t()), assumed_align=16), + # Carried flat and recast to E8M0 inside the trace; E8M0 has no DLPack + # dtype, so hand over the raw bytes. + from_dlpack(dy_col_sf.view(torch.uint8), assumed_align=16), + from_dlpack(x_col_sf.view(torch.uint8), assumed_align=16), + from_dlpack(offsets, assumed_align=4), + # (G, N, K) contiguous -> (N, K, G): the expert is the L coordinate. + from_dlpack(dw.permute(1, 2, 0), assumed_align=16), + stream, + ) + + slot = _wgrad_executor_slot((rows, out_features, in_features, groups)) + if not slot: + slot.append(cute.compile(_wgrad_entry, *args)) + slot[0](*args) From b822dbe6feaa925eb3c9b586edfb6d7fee6a54a9 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 16 Aug 2026 23:28:15 -0700 Subject: [PATCH 02/11] Rewrite the MXFP8 grouped-MLP family: public-DSL kernels A/B/C, ops, tests, bench Replaces the foundation commit's kernel layer with a correctness-first implementation of all three physically fused kernels on public CuTe DSL 4.7.0 API only, and completes the operator surface: mxfp8_grouped_gemm_swiglu_fwd FC1 ragged grouped GEMM + SwiGLU + dual (1x32 rowwise, 32x1 columnwise) MXFP8 RCEIL quantization + BF16 preactivation save, one kernel launch mxfp8_grouped_gemm_dswiglu_bwd FC2 dgrad grouped GEMM + dSwiGLU + dual quantization, one launch mxfp8_grouped_gemm_wgrad generic ragged-K grouped wgrad (FC1 and FC2), BF16 out, one launch per call No inline PTX, no llvm.inline_asm, no private cutlass._mlir/nvvm interfaces, no dsl_user_op asm wrappers remain. The replacements are public API and were verified against the prior implementations: Float8E8M0FNU conversion is natively round-upward (RCEIL), the public f32->E4M3 conversion is byte-identical to torch's cast, fmax(abs=True, nan=True) provides the NaN-propagating amax, and sigmoid composed as 1/(1+exp(-x)) matches torch.sigmoid bitwise. Epilogues stage through shared memory and derive every index from the T2R partitioner's coordinate tensor, removing the previous thread-to-row-ownership and physical-layout assumptions. The config/protocol framework, the published epilogue plug-in docs, and the decorative device-side offset assertions are gone; offset values are documented caller invariants with an opt-in synchronized host validator. Structure now follows moe_training conventions: public functional wrappers with availability detection in torchao/prototype/moe_training/mxfp8_grouped_mlp.py (exported from the package __init__, so a normal import registers the ops), custom ops + output allocation + fakes in kernels/mxfp8/grouped_mlp_ops.py, host validation in kernels/mxfp8/grouped_mlp_validation.py, and all three kernels + launchers in kernels/mxfp8/cutedsl_grouped_mlp.py. Launcher compile caches key on device index, compute capability, dtypes, shapes, and DSL version; streams come from the input tensor's device; G == 0 is rejected and R == 0 short-circuits without a launch. Checked-in tests (test/prototype/moe_training/test_mxfp8_grouped_mlp.py) cover numerics against pure-torch to_mx/to_blocked references (bitwise quantization via saturated-gate exact-product constructions and the shared MXFP8 semantic cases; SQNR gates for the BF16 GEMM stages), inactive-tail and zero-token semantics, FakeTensor/compile contracts, validation negatives, and a warmed torch.profiler launch-count test proving one kernel per op. The benchmark (benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py) compares the decomposed torchao path, TransformerEngine's modular and fused grouped-MLP lanes, and these kernels on identical inputs. The kernels, tests, and benchmark do not import cute_utils and run on the public nvidia-cutlass-dsl 4.7.0 wheel with no local patches. Validation (GB200 SM100, driver 580.173.02, CUDA 13.4 fwd-compat, torch 2.14.0a0, nvidia-cutlass-dsl 4.7.0, pristine cute_utils.py): test suite 46 passed + 1 skipped on GPU (the cross-device negative passes with two visible devices) and 7 passed CPU-only. The rewritten wgrad is bitwise-identical to the previous kernel on zero-token, strict-tail, stage-wrap and 16B FC1-shape configs; kernels A/B are bitwise against the to_mx(RCEIL)+to_blocked references at a production shape (0 mismatches over 2.16M and 4.33M qdata bytes), with kernel A bitwise even on random inputs. torch.profiler confirms exactly one CUDA kernel per warmed call. Capped-clock (1200 MHz) relative timings vs the decomposed torchao lane at R=2048/D=2048/F=1408/G=8: A 194 vs 339 us, B 192 vs 434 us, wgrad ~118 vs ~271 us per call. TE's tuned fused kernels remain substantially faster; tuning is the follow-up phase and the operator contracts are frozen for it. Co-Authored-By: Claude Fable 5 --- .../moe_training/mxfp8/bench_grouped_mlp.py | 759 ++++++++ .../moe_training/test_mxfp8_grouped_mlp.py | 1029 ++++++++++ torchao/prototype/moe_training/__init__.py | 12 + .../moe_training/kernels/mxfp8/__init__.py | 7 + .../kernels/mxfp8/cutedsl_grouped_mlp.py | 1722 ++++++++++++++++- .../kernels/mxfp8/epilogue_quant.py | 414 ---- .../kernels/mxfp8/grouped_gemm_config.py | 463 ----- .../kernels/mxfp8/grouped_gemm_core.py | 825 -------- .../kernels/mxfp8/grouped_mlp_epilogue.py | 448 ----- .../kernels/mxfp8/grouped_mlp_ops.py | 211 +- .../kernels/mxfp8/grouped_mlp_validation.py | 44 +- .../kernels/mxfp8/kernel_wgrad.py | 347 ---- .../moe_training/mxfp8_grouped_mlp.py | 220 +++ 13 files changed, 3927 insertions(+), 2574 deletions(-) create mode 100644 benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py create mode 100644 test/prototype/moe_training/test_mxfp8_grouped_mlp.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py create mode 100644 torchao/prototype/moe_training/mxfp8_grouped_mlp.py diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py b/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py new file mode 100644 index 0000000000..0ae723aea7 --- /dev/null +++ b/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py @@ -0,0 +1,759 @@ +# 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. +"""Benchmark the fused MXFP8 grouped-MLP kernel family against the existing +decomposed torchao path and TransformerEngine, on identical inputs, offsets, +shapes and RCEIL scale mode. + +Lanes (``--lane``): + +* ``torchao`` -- the existing decomposed SM100 path: triton/CUDA quantizers + + ``torch._scaled_grouped_mm`` + eager SwiGLU/dSwiGLU, staged to match each + fused kernel's covered work. Uses no CuTe DSL code. +* ``ours`` -- the three fused ops ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` + / ``_dswiglu_bwd`` / ``_wgrad`` (one kernel launch each). +* ``te`` -- TransformerEngine: the fused CuTe-DSL lane + (``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``; per-kernel times recovered from a + profiler pass by kernel-name fragment) plus the modular lane's single-kernel + ``tex.swiglu`` / ``tex.dswiglu`` gated-activation+dual-quantize points. +* ``all`` -- ``torchao`` + ``ours``. + +The TE lane must run in a separate process from ``ours``: our kernels need the +public ``nvidia-cutlass-dsl`` 4.7.0 wheel on the user site, while TE's fused +lane uses the container-native cuDNN/cutlass stack, which that wheel shadows. +Invocation on the GB200 dev host:: + + # ours / torchao lanes + bash ./run_te.sh env PYTHONUSERBASE=/.local PYTHONPATH=/ao \\ + python /ao/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py --lane all + # TE lane (no PYTHONUSERBASE) + bash ./run_te.sh env PYTHONPATH=/ao \\ + python /ao/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py --lane te + +Caveats stated up front so the numbers are read honestly: + +* The decomposed wgrad stage includes its own dim1 (columnwise) quantization of + both operands, because that is what the existing path pays; the fused wgrad + consumes columnwise operands produced by kernels A/B. The A/B stages of both + lanes start from identically prequantized GEMM inputs. +* The TE fused forward also applies per-token router probs in-kernel; we pass + probs = 1 so the work matches. +* Eager launch timing (``do_bench`` median) only. No CUDA-graph replay column. +* Absolute microseconds from a clock-capped host (this dev box pins app clocks + at 1200 MHz) are not publishable; the startup banner prints the clocks. + +``MXFP8_BENCH_VALIDATE=1`` cross-checks the fused outputs against pure-torch +``to_mx``/``to_blocked`` references (SQNR gates; the checked-in test suite owns +the bitwise contracts). +""" + +import argparse +import os +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional + +import torch +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 ( + grouped_mlp_ops, # noqa: F401 (registers the three fused ops) + mx_block_rearrange_2d_M_groups_cuda, + triton_mx_block_rearrange_2d_K_groups, +) +from torchao.prototype.moe_training.utils import generate_jagged_offs +from torchao.prototype.mx_formats.config import ( + MXFP8Dim1CastKernelChoice, + ScaleCalculationMode, +) +from torchao.prototype.mx_formats.kernels import ( + mxfp8_quantize_cuda, + triton_to_mxfp8_dim0, +) +from torchao.prototype.mx_formats.mx_tensor import to_mx +from torchao.prototype.mx_formats.utils import ( + _to_mxfp8_dim1_kernel_wrapper, + from_blocked, + to_blocked, +) +from torchao.quantization.quantize_.common import KernelPreference + +device = torch.device("cuda") +VALIDATE = os.environ.get("MXFP8_BENCH_VALIDATE", "0") == "1" +BLOCK = 32 +RCEIL = ScaleCalculationMode.RCEIL + + +# -------------------------------------------------------------------------- +# Configs +# -------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ExperimentConfig: + rows: int # R: padded token rows (sum of per-expert rows) + model_dim: int # D + hidden_dim: int # F + num_groups: int # G (local experts) + distribution: str # "balanced" | "skewed" + + +@dataclass(frozen=True) +class ExperimentResult: + # Per fused-kernel-equivalent stage, microseconds (median). + a_us: float + b_us: float + c_fc1_us: float + c_fc2_us: float + seq_us: float + # Derived TFLOP/s for the GEMM in each stage. + a_tflops: float + b_tflops: float + c_fc1_tflops: float + c_fc2_tflops: float + + +@dataclass(frozen=True) +class Experiment: + lane: str + config: ExperimentConfig + result: ExperimentResult + + +def get_configs(args) -> List[ExperimentConfig]: + if args.shape is not None: + r, d, f, g = (int(v) for v in args.shape.split(",")) + return [ExperimentConfig(r, d, f, g, dist) for dist in args.dists] + shapes = [ + # smoke + (512, 256, 256, 2), + # DeepSeekV3 16B class (D=2048, F=1408, G=8 local experts) + (2048, 2048, 1408, 8), + (8192, 2048, 1408, 8), + (16384, 2048, 1408, 8), + # DeepSeekV3 671B class (D=7168, F=2048, G=4 local experts) + (2048, 7168, 2048, 4), + (8192, 7168, 2048, 4), + (16384, 7168, 2048, 4), + ] + return [ + ExperimentConfig(r, d, f, g, dist) + for (r, d, f, g) in shapes + for dist in args.dists + ] + + +def make_offsets(cfg: ExperimentConfig) -> torch.Tensor: + """Exclusive per-expert end offsets, every group a multiple of 128.""" + r, g = cfg.rows, cfg.num_groups + if cfg.distribution == "balanced": + per = r // g + if per % 128 != 0 or per * g != r: + raise ValueError( + f"balanced distribution needs R/G to be a 128 multiple, got {r}/{g}" + ) + return torch.arange(1, g + 1, device=device, dtype=torch.int32) * per + return generate_jagged_offs(g, r, multiple_of=128, device=device) + + +# -------------------------------------------------------------------------- +# Pure-torch quantization recipes (input prep + validation only, never timed) +# -------------------------------------------------------------------------- + + +def ref_quantize_rowwise_1x32(x: torch.Tensor): + """[M, K] -> (E4M3 [M, K] row-major, flat blocked E8M0 for [M, K/32]).""" + scale, q = to_mx(x, torch.float8_e4m3fn, BLOCK, scaling_mode=RCEIL) + return q, to_blocked(scale) + + +def ref_quantize_colwise_32x1(x: torch.Tensor): + """[R, N] -> (E4M3 [R, N] stride (1, R), flat blocked E8M0 for [N, R/32]).""" + scale_t, q_t = to_mx( + x.t().contiguous(), torch.float8_e4m3fn, BLOCK, scaling_mode=RCEIL + ) + return q_t.t(), to_blocked(scale_t) + + +def ref_dequant_colwise(q_col: torch.Tensor, sf_blocked: torch.Tensor): + """FP32 dequant of a columnwise operand (for validation oracles).""" + rows, cols = q_col.shape + logical = from_blocked(sf_blocked, cols, rows // BLOCK) # [N, R/32] + scales = logical.t().to(torch.float32).repeat_interleave(BLOCK, dim=0) + return q_col.to(torch.float32) * scales + + +def ref_dequant_rowwise(q_row: torch.Tensor, sf_blocked: torch.Tensor): + """FP32 dequant of a rowwise 1x32-quantized operand (for validation oracles).""" + rows, cols = q_row.shape + scales = from_blocked(sf_blocked, rows, cols // BLOCK).to(torch.float32) + return q_row.to(torch.float32) * scales.repeat_interleave(BLOCK, dim=1) + + +def sqnr(ref: torch.Tensor, actual: torch.Tensor) -> float: + err = (ref - actual).float().pow(2).mean() + if err == 0: + return float("inf") + return (10 * torch.log10(ref.float().pow(2).mean() / err)).item() + + +# -------------------------------------------------------------------------- +# Shared input bundle +# -------------------------------------------------------------------------- + + +class Inputs: + """All operands both non-TE lanes consume, prepared once per config. + + GEMM operands are prequantized identically for both lanes (kernels A and B + take prequantized inputs by contract, and the decomposed path accepts + prequantized MX operands at the same seam). + """ + + def __init__(self, cfg: ExperimentConfig): + r, d, f, g = cfg.rows, cfg.model_dim, cfg.hidden_dim, cfg.num_groups + torch.manual_seed(0) + self.offsets = make_offsets(cfg) + self.offsets_host = self.offsets.tolist() + + self.x = torch.randn(r, d, device=device, dtype=torch.bfloat16) / d**0.5 + self.do = torch.randn(r, d, device=device, dtype=torch.bfloat16) / d**0.5 + # Element-interleaved gate/up FC1 weight [G, 2F, D] and FC2-dgrad + # weight-view source [G, F, D]; both quantized along D (the GEMM + # contraction), then freely transposed into the K-major ABI layouts. + self.w13i = ( + torch.randn(g, 2 * f, d, device=device, dtype=torch.bfloat16) / d**0.5 + ) + self.w2d = torch.randn(g, f, d, device=device, dtype=torch.bfloat16) / d**0.5 + + self.x_q, self.x_sf = ref_quantize_rowwise_1x32(self.x) + self.do_q, self.do_sf = ref_quantize_rowwise_1x32(self.do) + + w13_q, w13_sf = zip( + *(ref_quantize_rowwise_1x32(self.w13i[i]) for i in range(g)) + ) + self.w13_t_q = torch.stack(list(w13_q)).transpose(-2, -1) # [G, D, 2F] + self.w13_t_sf = torch.stack(list(w13_sf)) + w2_q, w2_sf = zip(*(ref_quantize_rowwise_1x32(self.w2d[i]) for i in range(g))) + self.w2_t_q = torch.stack(list(w2_q)).transpose(-2, -1) # [G, D, F] + self.w2_t_sf = torch.stack(list(w2_sf)) + + # Reference forward intermediates (bf16), computed per expert from the + # DEQUANTIZED operands — the values the fused kernels consume by + # contract — so the validation SQNR isolates each kernel's own work + # instead of stacking input-quantization error on top of it. The same + # z then feeds kernel B and the decomposed dSwiGLU identically. + x_f32 = ref_dequant_rowwise(self.x_q, self.x_sf) + do_f32 = ref_dequant_rowwise(self.do_q, self.do_sf) + z = torch.zeros(r, 2 * f, device=device, dtype=torch.bfloat16) + prev = 0 + for i in range(g): + end = self.offsets_host[i] + if end > prev: + w13_f32 = ref_dequant_rowwise(w13_q[i], w13_sf[i]) + z[prev:end] = (x_f32[prev:end] @ w13_f32.t()).to(torch.bfloat16) + prev = end + self.z_flat = z + self.z_bf16 = z.view(r, f, 2) + gate = self.z_bf16[..., 0].float() + up = self.z_bf16[..., 1].float() + self.h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) + sig = torch.sigmoid(gate) + dh = torch.zeros(r, f, device=device, dtype=torch.bfloat16) + prev = 0 + for i in range(g): + end = self.offsets_host[i] + if end > prev: + # do [m, D] contracts with the [D, F] dgrad weight view; the + # previous `do @ w2d[i]` only type-checked when D == F and + # computed the transpose of the intended dgrad. + w2_f32 = ref_dequant_rowwise(w2_q[i], w2_sf[i]) + dh[prev:end] = (do_f32[prev:end] @ w2_f32.t()).to(torch.bfloat16) + prev = end + dhf = dh.float() + dgate = (dhf * up * (sig * (1.0 + gate * (1.0 - sig)))).to(torch.bfloat16) + dup = (dhf * (gate * sig)).to(torch.bfloat16) + self.dz_flat = torch.stack((dgate, dup), dim=-1).view(r, 2 * f) + + # Columnwise operands for the two wgrad calls (produced by A/B in the + # fused regime, by standalone quantizers in the decomposed one). + self.dz_col_q, self.dz_col_sf = ref_quantize_colwise_32x1(self.dz_flat) + self.x_col_q, self.x_col_sf = ref_quantize_colwise_32x1(self.x) + self.do_col_q, self.do_col_sf = ref_quantize_colwise_32x1(self.do) + self.h_col_q, self.h_col_sf = ref_quantize_colwise_32x1(self.h) + + +def stage_flops(cfg: ExperimentConfig) -> Dict[str, float]: + r, d, f = cfg.rows, cfg.model_dim, cfg.hidden_dim + return { + "a": 2.0 * r * d * 2 * f, + "b": 2.0 * r * d * f, + "c_fc1": 2.0 * r * 2 * f * d, + "c_fc2": 2.0 * r * d * f, + } + + +# -------------------------------------------------------------------------- +# Lane: ours (the three fused ops) +# -------------------------------------------------------------------------- + + +def lane_ours(cfg: ExperimentConfig, inp: Inputs) -> Dict[str, Callable]: + ops = torch.ops.torchao + + def a(): + return ops.mxfp8_grouped_gemm_swiglu_fwd( + inp.x_q, inp.x_sf, inp.w13_t_q, inp.w13_t_sf, inp.offsets + ) + + def b(): + return ops.mxfp8_grouped_gemm_dswiglu_bwd( + inp.do_q, inp.do_sf, inp.w2_t_q, inp.w2_t_sf, inp.z_bf16, inp.offsets + ) + + def c_fc1(): + return ops.mxfp8_grouped_gemm_wgrad( + inp.dz_col_q, inp.dz_col_sf, inp.x_col_q, inp.x_col_sf, inp.offsets + ) + + def c_fc2(): + return ops.mxfp8_grouped_gemm_wgrad( + inp.do_col_q, inp.do_col_sf, inp.h_col_q, inp.h_col_sf, inp.offsets + ) + + def seq(): + _, _, _, h_col_q, h_col_sf = a() + _, _, dz_col_q, dz_col_sf = b() + ops.mxfp8_grouped_gemm_wgrad( + dz_col_q, dz_col_sf, inp.x_col_q, inp.x_col_sf, inp.offsets + ) + ops.mxfp8_grouped_gemm_wgrad( + inp.do_col_q, inp.do_col_sf, h_col_q, h_col_sf, inp.offsets + ) + + return {"a": a, "b": b, "c_fc1": c_fc1, "c_fc2": c_fc2, "seq": seq} + + +def validate_ours(cfg: ExperimentConfig, inp: Inputs, stages) -> None: + """SQNR cross-checks against pure-torch references. Bitwise contracts are + owned by test_mxfp8_grouped_mlp.py; this is a sanity gate for benching.""" + z_k, h_row_q, h_row_sf, h_col_q, h_col_sf = stages["a"]() + assert z_k.shape == inp.z_bf16.shape and z_k.stride() == inp.z_bf16.stride() + s = sqnr(inp.z_flat.float(), z_k.reshape(cfg.rows, -1).float()) + assert s >= 27.0, f"A z SQNR {s:.1f} < 27" + h_deq = ref_dequant_colwise(h_col_q, h_col_sf) + s = sqnr(inp.h.float(), h_deq) + assert s >= 27.0, f"A h (dequant colwise) SQNR {s:.1f} < 27" + row_deq = h_row_q.float() * ( + from_blocked(h_row_sf, cfg.rows, cfg.hidden_dim // BLOCK) + .to(torch.float32) + .repeat_interleave(BLOCK, dim=1) + ) + s = sqnr(inp.h.float(), row_deq) + assert s >= 27.0, f"A h (dequant rowwise) SQNR {s:.1f} < 27" + + dz_row_q, dz_row_sf, dz_col_q, dz_col_sf = stages["b"]() + dz_deq = ref_dequant_colwise(dz_col_q, dz_col_sf) + s = sqnr(inp.dz_flat.float(), dz_deq) + assert s >= 25.0, f"B dz SQNR {s:.1f} < 25" + + dw = stages["c_fc1"]() + dy_f32 = ref_dequant_colwise(inp.dz_col_q, inp.dz_col_sf) + x_f32 = ref_dequant_colwise(inp.x_col_q, inp.x_col_sf) + prev = 0 + ref = torch.zeros_like(dw, dtype=torch.float32) + for i in range(cfg.num_groups): + end = inp.offsets_host[i] + if end > prev: + ref[i] = dy_f32[prev:end].t() @ x_f32[prev:end] + prev = end + s = sqnr(ref, dw.float()) + assert s >= 24.0, f"C dw SQNR {s:.1f} < 24" + print(" validate(ours): OK (A z/h, B dz, C dw)") + + +# -------------------------------------------------------------------------- +# Lane: torchao decomposed (existing path; no CuTe DSL) +# -------------------------------------------------------------------------- + + +def lane_torchao(cfg: ExperimentConfig, inp: Inputs) -> Dict[str, Callable]: + r, f = cfg.rows, cfg.hidden_dim + offs = inp.offsets + + def dual_quantize(t: torch.Tensor): + # Rowwise via the triton dim0 quantizer (the CUDA kernel is + # colwise-only today), colwise via the CUDA quantizer, plus the two + # scale rearranges the existing SM100 path performs. + out_row, s_row = triton_to_mxfp8_dim0(t, BLOCK, "rceil") + s_row_blocked = mx_block_rearrange_2d_M_groups_cuda( + s_row.view(torch.uint8), offs + ) + _, out_col, _, s_col = mxfp8_quantize_cuda( + t, rowwise=False, colwise=True, scaling_mode="rceil" + ) + s_col_blocked = triton_mx_block_rearrange_2d_K_groups( + s_col.view(torch.uint8), offs // BLOCK + ) + return out_row, s_row_blocked, out_col, s_col_blocked + + def a(): + # FC1 grouped GEMM (prequantized inputs) -> eager SwiGLU -> dual quant. + z = torch._scaled_grouped_mm( + inp.x_q, + inp.w13_t_q, + inp.x_sf.view(r, -1), + inp.w13_t_sf.view(cfg.num_groups, -1), + offs=offs, + out_dtype=torch.bfloat16, + ) + zv = z.view(r, f, 2) + h = (torch.nn.functional.silu(zv[..., 0].float()) * zv[..., 1].float()).to( + torch.bfloat16 + ) + return dual_quantize(h) + + def b(): + dh = torch._scaled_grouped_mm( + inp.do_q, + inp.w2_t_q, + inp.do_sf.view(r, -1), + inp.w2_t_sf.view(cfg.num_groups, -1), + offs=offs, + out_dtype=torch.bfloat16, + ) + gate = inp.z_bf16[..., 0].float() + up = inp.z_bf16[..., 1].float() + sig = torch.sigmoid(gate) + dhf = dh.float() + dgate = (dhf * up * (sig * (1.0 + gate * (1.0 - sig)))).to(torch.bfloat16) + dup = (dhf * (gate * sig)).to(torch.bfloat16) + dz = torch.stack((dgate, dup), dim=-1).view(r, 2 * f) + return dual_quantize(dz) + + def make_wgrad(dy: torch.Tensor, x: torch.Tensor): + # Verbatim shape of the existing wgrad stage: CUDA dim1 quantization of + # both operands + K-groups scale rearranges + scaled grouped GEMM. + # (The fused kernel C instead consumes columnwise operands produced by + # kernels A/B, so this stage's quantization cost is the decomposed + # path's own.) + def run(): + dy_t_mx = _to_mxfp8_dim1_kernel_wrapper( + dy, + BLOCK, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=dy.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=RCEIL, + ) + x_t_mx = _to_mxfp8_dim1_kernel_wrapper( + x, + BLOCK, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=x.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=RCEIL, + ) + scale_offs = offs // BLOCK + dy_scales_blocked = triton_mx_block_rearrange_2d_K_groups( + dy_t_mx.scale, scale_offs + ) + x_scales_blocked = triton_mx_block_rearrange_2d_K_groups( + x_t_mx.scale, scale_offs + ) + return torch._scaled_grouped_mm( + dy_t_mx.qdata, + x_t_mx.qdata.transpose(-2, -1), + dy_scales_blocked, + x_scales_blocked, + offs=offs, + out_dtype=torch.bfloat16, + ) + + return run + + c_fc1 = make_wgrad(inp.dz_flat, inp.x) + c_fc2 = make_wgrad(inp.do, inp.h) + + def seq(): + a() + b() + c_fc1() + c_fc2() + + return {"a": a, "b": b, "c_fc1": c_fc1, "c_fc2": c_fc2, "seq": seq} + + +# -------------------------------------------------------------------------- +# Lane: TransformerEngine +# -------------------------------------------------------------------------- + + +def run_te_lane(cfg: ExperimentConfig) -> None: + """TE fused lane (one CuTe kernel per stage) + modular-lane single-kernel + gated-activation points. Prints its own tables; per-kernel CUDA times come + from a profiler pass filtered by kernel-name fragment.""" + os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1" + import transformer_engine.pytorch as te + import transformer_engine_torch as tex + from transformer_engine.common.recipe import MXFP8BlockScaling + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + r, d, f, g = cfg.rows, cfg.model_dim, cfg.hidden_dim, cfg.num_groups + torch.manual_seed(0) + offsets = make_offsets(cfg) + sizes = torch.diff(offsets, prepend=torch.zeros(1, device=device).int()) + sizes = sizes.to(torch.int32) + if (sizes % 128 != 0).any(): + raise ValueError( + "TE fused lane hard-crashes the CUDA context on per-expert sizes " + f"that are not multiples of 128, got {sizes.tolist()}" + ) + recipe = MXFP8BlockScaling() + + # Fused lane: one Sequential MLP, per-kernel attribution by name fragment. + fc1 = te.ops.GroupedLinear( + g, d, 2 * f, bias=False, device="cuda", dtype=torch.bfloat16 + ) + act = te.ops.ScaledSwiGLU(glu_interleave_size=32) + fc2 = te.ops.GroupedLinear(g, f, d, bias=False, device="cuda", dtype=torch.bfloat16) + mlp = te.ops.Sequential(fc1, act, fc2) + x = torch.randn(r, d, device=device, dtype=torch.bfloat16, requires_grad=True) + probs = torch.ones(r, device=device, dtype=torch.bfloat16) + + def fwd(): + with te.autocast(enabled=True, recipe=recipe): + return mlp(x, sizes, probs, sizes) + + y = fwd() + dy = torch.randn_like(y) + + def fwd_bwd(): + out = fwd() + out.backward(dy) + + fwd_bwd() # warmup / lazy init + torch.cuda.synchronize() + fwd_us = benchmark_cuda_function_in_microseconds(fwd) + fwd_bwd_us = benchmark_cuda_function_in_microseconds(fwd_bwd) + + fragments = { + "A analog (GroupedGemmGlu)": "GroupedGemmGlu", + "B analog (GroupedGemmDglu)": "GroupedGemmDglu", + "C analog (GroupedGemmWgrad)": "GroupedGemmWgrad", + "plain GEMM (GroupedGemmQuant)": "GroupedGemmQuant", + } + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA] + ) as prof: + for _ in range(5): + fwd_bwd() + torch.cuda.synchronize() + sums = dict.fromkeys(fragments, 0.0) + counts = dict.fromkeys(fragments, 0) + for evt in prof.key_averages(): + for label, frag in fragments.items(): + if frag in evt.key and "helper" not in evt.key: + sums[label] += evt.self_device_time_total + counts[label] += evt.count + rows = [ + [label, counts[label] / 5.0, sums[label] / 5.0] + for label in fragments + if counts[label] + ] + print(f"\nTE fused lane (NVTE_CUTEDSL_FUSED_GROUPED_MLP=1) {cfg}") + print(f" fwd wall: {fwd_us:.1f} us fwd+bwd wall: {fwd_bwd_us:.1f} us") + print( + tabulate( + rows, + headers=["main kernel", "launches/iter", "device us/iter"], + floatfmt=".1f", + ) + ) + print( + " (fwd+bwd wall also covers FC2-fwd/FC1-dgrad GEMMs, input/dy " + "quantize and offsets prep, matching a full MLP step)" + ) + + # Modular lane: the single fused gated-act+dual-quant kernels. + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + z = torch.randn(r, 2 * f, device=device, dtype=torch.bfloat16) + dh = torch.randn(r, f, device=device, dtype=torch.bfloat16) + tex.swiglu(z, q) + tex.dswiglu(dh, z, q) + swiglu_us = benchmark_cuda_function_in_microseconds(tex.swiglu, z, q) + dswiglu_us = benchmark_cuda_function_in_microseconds(tex.dswiglu, dh, z, q) + print( + f"TE modular lane single kernels: tex.swiglu {swiglu_us:.1f} us, " + f"tex.dswiglu {dswiglu_us:.1f} us (gated activation + dual MXFP8 " + "quantize only; GEMMs are per-expert cuBLASLt in this lane)" + ) + + +# -------------------------------------------------------------------------- +# Driver +# -------------------------------------------------------------------------- + + +def run_experiment( + lane: str, cfg: ExperimentConfig, inp: Inputs +) -> Optional[ExperimentResult]: + stages = (lane_ours if lane == "ours" else lane_torchao)(cfg, inp) + if VALIDATE and lane == "ours": + validate_ours(cfg, inp, stages) + times: Dict[str, float] = {} + for name, fn in stages.items(): + fn() # warmup + lazy compile + torch.cuda.synchronize() + times[name] = benchmark_cuda_function_in_microseconds(fn) + flops = stage_flops(cfg) + return ExperimentResult( + a_us=times["a"], + b_us=times["b"], + c_fc1_us=times["c_fc1"], + c_fc2_us=times["c_fc2"], + seq_us=times["seq"], + a_tflops=flops["a"] / times["a"] / 1e6, + b_tflops=flops["b"] / times["b"] / 1e6, + c_fc1_tflops=flops["c_fc1"] / times["c_fc1"] / 1e6, + c_fc2_tflops=flops["c_fc2"] / times["c_fc2"] / 1e6, + ) + + +def print_banner() -> None: + props = torch.cuda.get_device_properties(device) + clocks = os.popen( + "nvidia-smi --query-gpu=clocks.applications.graphics,clocks.max.graphics " + "--format=csv,noheader 2>/dev/null" + ).read() + try: + import cutlass + + dsl = cutlass.__version__ + except Exception: + dsl = "n/a" + print( + f"device: {props.name} (cc {props.major}.{props.minor}, index " + f"{torch.cuda.current_device()}), torch {torch.__version__}, CUDA " + f"{torch.version.cuda}, nvidia-cutlass-dsl {dsl}" + ) + print(f"app clocks / max (per GPU):\n{clocks.strip()}") + print( + "NOTE: if app clocks are capped below max (e.g. 1200 MHz on the GB200 " + "dev hosts), absolute microseconds are NOT publishable; use ratios." + ) + + +def print_results(experiments: List[Experiment]) -> None: + headers = [ + "lane", + "R", + "D", + "F", + "G", + "dist", + "A us", + "B us", + "C_fc1 us", + "C_fc2 us", + "seq us", + "A TF/s", + "B TF/s", + "C1 TF/s", + "C2 TF/s", + ] + rows = [] + for e in experiments: + c, r = e.config, e.result + rows.append( + [ + e.lane, + c.rows, + c.model_dim, + c.hidden_dim, + c.num_groups, + c.distribution, + f"{r.a_us:.1f}", + f"{r.b_us:.1f}", + f"{r.c_fc1_us:.1f}", + f"{r.c_fc2_us:.1f}", + f"{r.seq_us:.1f}", + f"{r.a_tflops:.1f}", + f"{r.b_tflops:.1f}", + f"{r.c_fc1_tflops:.1f}", + f"{r.c_fc2_tflops:.1f}", + ] + ) + print(tabulate(rows, headers=headers)) + print( + "stage coverage: A = FC1 grouped GEMM + SwiGLU + dual MXFP8 quantize; " + "B = FC2 dgrad + dSwiGLU + dual quantize; C_* = grouped wgrad " + "(decomposed lane's C includes its own dim1 operand quantization); " + "seq = A;B;C_fc1;C_fc2." + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--lane", + choices=["torchao", "ours", "te", "all"], + default="all", + help="'all' = torchao + ours; 'te' must run in its own process " + "(container-native stack, no PYTHONUSERBASE)", + ) + parser.add_argument( + "--shape", + default=None, + help="single shape as 'R,D,F,G' instead of the built-in sweep", + ) + parser.add_argument( + "--dist", + choices=["balanced", "skewed", "both"], + default="balanced", + dest="dist", + ) + parser.add_argument( + "--profile", + action="store_true", + help="export a chrome trace of the fused-op sequence per config", + ) + args = parser.parse_args() + args.dists = ["balanced", "skewed"] if args.dist == "both" else [args.dist] + + print_banner() + configs = get_configs(args) + + if args.lane == "te": + for cfg in configs: + run_te_lane(cfg) + return + + lanes = ["torchao", "ours"] if args.lane == "all" else [args.lane] + experiments: List[Experiment] = [] + for cfg in tqdm(configs): + inp = Inputs(cfg) + for lane in lanes: + result = run_experiment(lane, cfg, inp) + experiments.append(Experiment(lane, cfg, result)) + if args.profile and "ours" in lanes: + from benchmarks.utils import profile_fn + + stages = lane_ours(cfg, inp) + profile_fn( + stages["seq"], + profile_name=f"grouped_mlp_seq_R{cfg.rows}_D{cfg.model_dim}" + f"_F{cfg.hidden_dim}_G{cfg.num_groups}", + ) + del inp + torch.cuda.empty_cache() + print_results(experiments) + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py new file mode 100644 index 0000000000..06f3996f52 --- /dev/null +++ b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py @@ -0,0 +1,1029 @@ +# 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. + +"""Tests for the fused MXFP8 grouped-MLP kernel family (SM100). + +Three custom ops, one physical kernel launch each: + +* ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` (A) FC1 grouped GEMM + SwiGLU + + rowwise 1x32 and columnwise 32x1 MXFP8 RCEIL quantization +* ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` (B) FC2 dgrad grouped GEMM + + dSwiGLU + dual quantization +* ``torchao::mxfp8_grouped_gemm_wgrad`` (C) grouped MXFP8 wgrad + +References are deliberately bridge-free: eager per-expert BF16/FP64 matmuls, +``F.silu`` / the closed-form dSwiGLU, and the pure-torch ``to_mx`` (RCEIL) + +``to_blocked`` quantizers. The CuTe DSL standalone quantizer ops and +``cute_utils`` are never imported. + +Numerical strategy: + +* GEMM outputs (``z``, ``dh``, ``dw``): SQNR / tolerance vs an FP64 oracle + (reduction order is free), plus BITWISE equality on exact-integer operand + configs where every partial sum is exactly representable in FP32. +* Quantized outputs: bitwise vs ``to_mx`` wherever the activation is exact. + ``silu(g) == g`` exactly for ``g >= 128`` (sigmoid saturates to 1.0f), so a + saturated gate turns the fused activation into exact products and the + quantization stage must match ``to_mx`` byte-for-byte, including special + values. Random-input forward comparisons are ALSO bitwise (measured zero + mismatches: the kernel's sigmoid composition matches torch's float32 silu + exactly); backward random inputs are SQNR-gated because the reference dh + comes from an FP64 oracle whose BF16 rounding can differ at reduction-order + boundaries, with bitwise coverage provided by the exact-dh configuration. + +FakeTensor and validation tests run without a GPU; kernel tests require SM100. +""" + +import random + +import pytest + +torch = pytest.importorskip("torch") + +import torch.nn.functional as F # noqa: E402 +from torch._subclasses.fake_tensor import FakeTensorMode # noqa: E402 + +from torchao.float8.float8_utils import compute_error # noqa: E402 +from torchao.prototype.moe_training.utils import generate_jagged_offs # noqa: E402 +from torchao.prototype.mx_formats.config import ScaleCalculationMode # noqa: E402 +from torchao.prototype.mx_formats.mx_tensor import to_mx # noqa: E402 +from torchao.prototype.mx_formats.utils import from_blocked, to_blocked # noqa: E402 +from torchao.testing._mxfp8_test_utils import make_mxfp8_semantic_cases # noqa: E402 + +# Importing the ops module registers the three custom ops. The public wrapper +# module is preferred once it exists; both expose the same wrapper names. +try: + from torchao.prototype.moe_training import mxfp8_grouped_mlp as _api +except ImportError: + from torchao.prototype.moe_training.kernels.mxfp8 import grouped_mlp_ops as _api + +_E4M3 = torch.float8_e4m3fn +_E8M0 = torch.float8_e8m0fnu +_BLOCK = 32 +_RCEIL = ScaleCalculationMode.RCEIL + +_OP_NAMES = ( + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_wgrad", +) + + +def _is_sm_10x() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +_gpu = pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") + +torch._dynamo.config.cache_size_limit = 1000 + + +# --------------------------------------------------------------------------- +# Reference helpers (pure torch) +# --------------------------------------------------------------------------- + + +def _round_up(x: int, to: int) -> int: + return ((x + to - 1) // to) * to + + +def _blocked_numel(rows: int, cols: int) -> int: + return _round_up(rows, 128) * _round_up(cols, 4) + + +def _quantize_rowwise_ref(x: torch.Tensor): + """[M, K] high precision -> (qdata [M, K] row-major, flat blocked scales).""" + scale, q = to_mx(x, _E4M3, _BLOCK, scaling_mode=_RCEIL) + return q, to_blocked(scale) + + +def _quantize_colwise_ref(x: torch.Tensor): + """[R, N] high precision -> (qdata [R, N] stride (1, R), flat blocked scales + for the logical [N, R/32] scale matrix). Recipe from the repo bench.""" + scale_t, q_t = to_mx(x.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + return q_t.t(), to_blocked(scale_t) + + +def _dequant_rowwise(q: torch.Tensor, sf_flat: torch.Tensor, dtype=torch.float64): + """Dequantize a row-major [M, K] E4M3 tensor with flat blocked scales.""" + m, k = q.shape + scale = from_blocked(sf_flat.view(_E8M0).reshape(-1), m, k // _BLOCK) + return q.to(dtype) * scale.to(dtype).repeat_interleave(_BLOCK, dim=1) + + +def _dequant_colwise(q_col: torch.Tensor, sf_flat: torch.Tensor, dtype=torch.float64): + """Dequantize a [R, N] stride-(1, R) E4M3 tensor (scales logical [N, R/32]).""" + r, n = q_col.shape + scale = from_blocked(sf_flat.view(_E8M0).reshape(-1), n, r // _BLOCK) + return (q_col.t().to(dtype) * scale.to(dtype).repeat_interleave(_BLOCK, dim=1)).t() + + +def _dswiglu_closed_form(dh: torch.Tensor, gate: torch.Tensor, up: torch.Tensor): + """CONTRACT section 6.3 normative math, all fp32 in, (dgate, dup) fp32 out.""" + sig = torch.sigmoid(gate) + silu = gate * sig + dsilu = sig * (1.0 + gate * (1.0 - sig)) + return dh * up * dsilu, dh * silu + + +def _bytes(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint8) + + +def _mismatch_rate(a: torch.Tensor, b: torch.Tensor) -> float: + assert a.shape == b.shape + return (a != b).float().mean().item() + + +# --------------------------------------------------------------------------- +# Input builders. Tests own the offsets, so references never need a D2H sync. +# --------------------------------------------------------------------------- + + +def _mk_offsets(sizes, device): + ends = torch.tensor(sizes, dtype=torch.int64).cumsum(0) + return ends.to(torch.int32).to(device) + + +def _pack_grouped_weight(q_list, sf_list, device): + """Per-expert row-major [N, K] qdata -> [G, K, N] stride (K*N, 1, K) + [G, sf].""" + g = len(q_list) + n, k = q_list[0].shape + w = torch.empty_strided((g, k, n), (k * n, 1, k), dtype=_E4M3, device=device) + for i, q in enumerate(q_list): + w[i].copy_(q.t()) + sf = torch.stack([s.reshape(-1) for s in sf_list]) + return w, sf + + +def _random_grouped_weight(g, n, k, device, exact_int=False): + """Random per-expert [N, K] weights; returns (packed q, packed sf, dequants).""" + qs, sfs, deqs = [], [], [] + for _ in range(g): + if exact_int: + q = torch.randint(-6, 7, (n, k), device=device).to(_E4M3) + logical = torch.full( + (n, k // _BLOCK), 127, dtype=torch.uint8, device=device + ) + sf = to_blocked(logical.view(_E8M0)) + else: + w = torch.randn(n, k, device=device, dtype=torch.bfloat16) + q, sf = _quantize_rowwise_ref(w) + qs.append(q) + sfs.append(sf) + deqs.append(_dequant_rowwise(q, sf)) + packed_q, packed_sf = _pack_grouped_weight(qs, sfs, device) + return packed_q, packed_sf, deqs + + +def _random_activation(r, k, device, exact_int=False): + if exact_int: + q = torch.randint(-6, 7, (r, k), device=device).to(_E4M3) + logical = torch.full((r, k // _BLOCK), 127, dtype=torch.uint8, device=device) + sf = to_blocked(logical.view(_E8M0)) + else: + x = torch.randn(r, k, device=device, dtype=torch.bfloat16) + q, sf = _quantize_rowwise_ref(x) + return q, sf, _dequant_rowwise(q, sf) + + +def _ref_grouped_gemm(x_deq, w_deqs, offsets_sizes): + """Per-expert fp64 x @ w.T over test-owned split sizes; inactive tail = 0.""" + r = x_deq.shape[0] + out = torch.zeros(r, w_deqs[0].shape[0], dtype=torch.float64, device=x_deq.device) + start = 0 + for g, size in enumerate(offsets_sizes): + if size: + out[start : start + size] = x_deq[start : start + size] @ w_deqs[g].t() + start += size + return out + + +def _make_a_inputs(r, d, f, sizes, device, exact_int=False): + x_q, x_sf, x_deq = _random_activation(r, d, device, exact_int) + w_q, w_sf, w_deqs = _random_grouped_weight(len(sizes), 2 * f, d, device, exact_int) + offsets = _mk_offsets(sizes, device) + z_ref = _ref_grouped_gemm(x_deq, w_deqs, sizes) + return (x_q, x_sf, w_q, w_sf, offsets), z_ref + + +def _make_b_inputs(r, d, f, sizes, device, exact_int=False, z=None): + do_q, do_sf, do_deq = _random_activation(r, d, device, exact_int) + w_q, w_sf, w_deqs = _random_grouped_weight(len(sizes), f, d, device, exact_int) + offsets = _mk_offsets(sizes, device) + if z is None: + z = torch.randn(r, f, 2, device=device, dtype=torch.bfloat16) + active = sum(sizes) + if active < r: + # The inactive tail of z is read-forbidden: poison it so any read + # shows up as NaN contamination in the outputs. + z[active:] = float("nan") + dh_ref = _ref_grouped_gemm(do_deq, w_deqs, sizes) + return (do_q, do_sf, w_q, w_sf, z, offsets), dh_ref + + +def _make_c_inputs(r, n, k, sizes, device, exact_int=False): + def colwise(rows, cols): + if exact_int: + q_rm = torch.randint(-6, 7, (cols, rows), device=device).to(_E4M3) + logical = torch.full( + (cols, rows // _BLOCK), 127, dtype=torch.uint8, device=device + ) + sf = to_blocked(logical.view(_E8M0)) + return q_rm.t(), sf + x = torch.randn(rows, cols, device=device, dtype=torch.bfloat16) + return _quantize_colwise_ref(x) + + dy_q, dy_sf = colwise(r, n) + x_q, x_sf = colwise(r, k) + offsets = _mk_offsets(sizes, device) + return dy_q, dy_sf, x_q, x_sf, offsets + + +def _ref_wgrad(dy_q, dy_sf, x_q, x_sf, sizes): + dy = _dequant_colwise(dy_q, dy_sf) + x = _dequant_colwise(x_q, x_sf) + g = len(sizes) + n, k = dy.shape[1], x.shape[1] + dw = torch.zeros(g, n, k, dtype=torch.float64, device=dy.device) + start = 0 + for i, size in enumerate(sizes): + if size: + dw[i] = dy[start : start + size].t() @ x[start : start + size] + start += size + return dw.to(torch.bfloat16) + + +def _b_reference_dz(dh_bf16, z): + """CONTRACT section 6.3: bf16 dh + saved z -> interleaved dz (bf16 [R, 2F]).""" + dgate, dup = _dswiglu_closed_form( + dh_bf16.float(), z[..., 0].float(), z[..., 1].float() + ) + dz = torch.stack((dgate.bfloat16(), dup.bfloat16()), dim=-1) + return dz.reshape(dh_bf16.shape[0], -1) + + +def _assert_quantized_pair( + q_row, sf_row, q_col, sf_col, ref_bf16, max_qdata_rate=0.0, max_scale_rate=0.0 +): + """Compare both fused quantized orientations against to_mx of ``ref_bf16``.""" + ref_row_q, ref_row_sf = _quantize_rowwise_ref(ref_bf16) + ref_col_q, ref_col_sf = _quantize_colwise_ref(ref_bf16) + checks = ( + ("h_row_q", _bytes(q_row), _bytes(ref_row_q), max_qdata_rate), + ("h_row_sf", _bytes(sf_row.reshape(-1)), _bytes(ref_row_sf), max_scale_rate), + # Column-major output: compare bytes of the same logical view without + # forcing contiguity (the stride IS the ABI). + ("h_col_q", _bytes(q_col.t()), _bytes(ref_col_q.t()), max_qdata_rate), + ("h_col_sf", _bytes(sf_col.reshape(-1)), _bytes(ref_col_sf), max_scale_rate), + ) + for name, got, want, budget in checks: + rate = _mismatch_rate(got, want) + assert rate <= budget, f"{name}: byte mismatch rate {rate} > {budget}" + + +# --------------------------------------------------------------------------- +# 1. Registration and public surface +# --------------------------------------------------------------------------- + + +def test_ops_registered(): + for name in _OP_NAMES: + assert hasattr(torch.ops.torchao, name), name + assert hasattr(_api, name), f"{_api.__name__} must export {name}" + + +# --------------------------------------------------------------------------- +# 2. Fake / meta output contracts (no GPU required) +# --------------------------------------------------------------------------- + + +def _fake_a_inputs(r, d, f, g, device="cuda"): + x_q = torch.empty(r, d, dtype=_E4M3, device=device) + x_sf = torch.empty(_blocked_numel(r, d // _BLOCK), dtype=_E8M0, device=device) + w_q = torch.empty_strided( + (g, d, 2 * f), (d * 2 * f, 1, d), dtype=_E4M3, device=device + ) + w_sf = torch.empty( + g, _blocked_numel(2 * f, d // _BLOCK), dtype=_E8M0, device=device + ) + offs = torch.empty(g, dtype=torch.int32, device=device) + return x_q, x_sf, w_q, w_sf, offs + + +def _fake_b_inputs(r, d, f, g, device="cuda"): + do_q = torch.empty(r, d, dtype=_E4M3, device=device) + do_sf = torch.empty(_blocked_numel(r, d // _BLOCK), dtype=_E8M0, device=device) + w_q = torch.empty_strided((g, d, f), (d * f, 1, d), dtype=_E4M3, device=device) + w_sf = torch.empty(g, _blocked_numel(f, d // _BLOCK), dtype=_E8M0, device=device) + z = torch.empty_strided( + (r, f, 2), (2 * f, 2, 1), dtype=torch.bfloat16, device=device + ) + offs = torch.empty(g, dtype=torch.int32, device=device) + return do_q, do_sf, w_q, w_sf, z, offs + + +def _fake_c_inputs(r, n, k, g, device="cuda"): + dy = torch.empty_strided((r, n), (1, r), dtype=_E4M3, device=device) + dy_sf = torch.empty(_blocked_numel(n, r // _BLOCK), dtype=_E8M0, device=device) + x = torch.empty_strided((r, k), (1, r), dtype=_E4M3, device=device) + x_sf = torch.empty(_blocked_numel(k, r // _BLOCK), dtype=_E8M0, device=device) + offs = torch.empty(g, dtype=torch.int32, device=device) + return dy, dy_sf, x, x_sf, offs + + +def test_fake_swiglu_fwd_contract(): + r, d, f, g = 256, 256, 128, 2 + with FakeTensorMode(): + z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + *_fake_a_inputs(r, d, f, g) + ) + assert z.shape == (r, f, 2) and z.stride() == (2 * f, 2, 1) + assert z.dtype == torch.bfloat16 + assert hq.shape == (r, f) and hq.stride() == (f, 1) and hq.dtype == _E4M3 + assert hsf.numel() == _blocked_numel(r, f // _BLOCK) + assert hcq.shape == (r, f) and hcq.stride() == (1, r) and hcq.dtype == _E4M3 + assert hcsf.numel() == _blocked_numel(f, r // _BLOCK) + + +def test_fake_dswiglu_bwd_contract(): + r, d, f, g = 256, 256, 128, 2 + with FakeTensorMode(): + dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + *_fake_b_inputs(r, d, f, g) + ) + assert dzq.shape == (r, 2 * f) and dzq.stride() == (2 * f, 1) + assert dzsf.numel() == _blocked_numel(r, 2 * f // _BLOCK) + assert dzcq.shape == (r, 2 * f) and dzcq.stride() == (1, r) + assert dzcsf.numel() == _blocked_numel(2 * f, r // _BLOCK) + + +def test_fake_wgrad_contract(): + r, n, k, g = 256, 256, 128, 2 + with FakeTensorMode(): + dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*_fake_c_inputs(r, n, k, g)) + assert dw.shape == (g, n, k) + assert dw.stride() == (n * k, k, 1) + assert dw.dtype == torch.bfloat16 + + +def test_fake_validation_rejects_bad_metadata(): + with FakeTensorMode(): + args = list(_fake_c_inputs(256, 192, 128, 2)) # N = 192, not 128-multiple + with pytest.raises(ValueError, match="multiple of 128"): + torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + a_args = list(_fake_a_inputs(192, 256, 128, 2)) # R = 192 + with pytest.raises(ValueError, match="multiple of 128"): + torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*a_args) + + +# --------------------------------------------------------------------------- +# 3. Validation negatives (real tensors, small shapes) +# --------------------------------------------------------------------------- + + +def _valid_c_args(device): + return list(_make_c_inputs(256, 256, 128, [128, 128], device)) + + +_NEGATIVE_CASES = [ + "dtype", + "row_major_colwise_operand", + "scale_numel", + "offsets_int64", + "offsets_cpu", + "offsets_2d", + "offsets_numel", + "offsets_noncontig", + "n_not_128", + "misaligned_view", + "z_stride", +] + + +@_gpu +@pytest.mark.parametrize("case", _NEGATIVE_CASES) +def test_validation_negatives(case): + device = "cuda" + args = _valid_c_args(device) + if case == "dtype": + args[0] = torch.empty_strided( + (256, 256), (1, 256), dtype=torch.bfloat16, device=device + ) + err = "float8_e4m3fn" + elif case == "row_major_colwise_operand": + args[0] = torch.empty(256, 256, dtype=_E4M3, device=device) + err = "stride" + elif case == "scale_numel": + args[1] = args[1].reshape(-1)[:-1] + err = "blocked scale bytes" + elif case == "offsets_int64": + args[4] = args[4].to(torch.int64) + err = "int32" + elif case == "offsets_cpu": + args[4] = args[4].cpu() + err = "CUDA" + elif case == "offsets_2d": + args[4] = args[4].reshape(1, -1) + err = "1D" + elif case == "offsets_numel": + args[4] = torch.tensor([128, 128, 256], dtype=torch.int32, device=device) + # wgrad takes G from offsets, so a numel change alone is legal there; + # use kernel A where G comes from the weight tensor instead. + a_args = list(_fake_a_inputs(256, 256, 128, 2, device=device)) + a_args[4] = args[4] + with pytest.raises(ValueError, match="one entry per"): + torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*a_args) + return + elif case == "offsets_noncontig": + base = torch.zeros(4, dtype=torch.int32, device=device) + args[4] = base.as_strided((2,), (2,)) + err = "contiguous" + elif case == "n_not_128": + args = _valid_c_args(device) + dy = torch.empty_strided((256, 192), (1, 256), dtype=_E4M3, device=device) + dy_sf = torch.empty( + _blocked_numel(192, 256 // _BLOCK), dtype=_E8M0, device=device + ) + args[0], args[1] = dy, dy_sf + err = "multiple of 128" + elif case == "misaligned_view": + base = torch.zeros(256 * 256 + 32, dtype=_E4M3, device=device) + args[0] = base.as_strided((256, 256), (1, 256), 2) + err = "aligned" + elif case == "z_stride": + b_args = list(_fake_b_inputs(256, 256, 128, 2, device=device)) + # materialize real tensors with a wrong z layout + b_args = [torch.empty_like(t) if t.is_cuda else t for t in b_args] + b_args[4] = torch.empty( + 256, 2, 128, dtype=torch.bfloat16, device=device + ).permute(0, 2, 1) + with pytest.raises(ValueError, match="stride"): + torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd(*b_args) + return + with pytest.raises(ValueError, match=err): + torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + + +@_gpu +def test_validation_rejects_g0(): + device = "cuda" + args = _valid_c_args(device) + args[4] = torch.empty(0, dtype=torch.int32, device=device) + with pytest.raises(ValueError): + torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + + +@pytest.mark.skipif( + not (_is_sm_10x() and torch.cuda.device_count() >= 2), + reason="needs two CUDA devices", +) +def test_validation_rejects_cross_device(): + args = _valid_c_args("cuda:0") + args[2] = args[2].to("cuda:1") + with pytest.raises(ValueError, match="device"): + torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + + +# --------------------------------------------------------------------------- +# 4. R == 0 and zero-token experts +# --------------------------------------------------------------------------- + + +@_gpu +def test_r0_all_ops(): + device = "cuda" + d, f, g = 256, 128, 2 + a_args = list(_fake_a_inputs(0, d, f, g, device=device)) + a_args = [torch.empty_like(t) for t in a_args] + a_args[4] = torch.zeros(g, dtype=torch.int32, device=device) + outs = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*a_args) + assert all(o.shape[0] == 0 or o.numel() == 0 for o in outs) + + b_args = list(_fake_b_inputs(0, d, f, g, device=device)) + b_args = [torch.empty_like(t) for t in b_args] + b_args[5] = torch.zeros(g, dtype=torch.int32, device=device) + outs = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd(*b_args) + assert all(o.numel() == 0 for o in outs) + + c_args = list(_fake_c_inputs(0, 256, 128, g, device=device)) + c_args = [torch.empty_like(t) for t in c_args] + c_args[4] = torch.zeros(g, dtype=torch.int32, device=device) + dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*c_args) + assert dw.shape == (g, 256, 128) + assert (dw == 0).all() + + +@_gpu +def test_wgrad_zero_token_expert(): + torch.manual_seed(0) + device = "cuda" + sizes = [128, 0, 256] + args = _make_c_inputs(384, 256, 128, sizes, device) + dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + assert (dw[1] == 0).all(), "zero-token expert must produce an all-zero slice" + ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) + assert compute_error(ref[0].float(), dw[0].float()) >= 24.0 + assert compute_error(ref[2].float(), dw[2].float()) >= 24.0 + + +# --------------------------------------------------------------------------- +# 5. Kernel C numerics +# --------------------------------------------------------------------------- + + +@_gpu +@pytest.mark.parametrize( + "r,n,k,sizes", + [ + (256, 256, 128, [128, 128]), # FC1-like: N = 2F, K = D (small) + (1536, 2816, 2048, [512, 0, 640, 384]), # 16B FC1 wgrad class + (1536, 2048, 1408, [512, 0, 640, 384]), # 16B FC2 wgrad class + ], + ids=["small", "fc1_16b", "fc2_16b"], +) +def test_wgrad_numerics_random(r, n, k, sizes): + torch.manual_seed(1) + args = _make_c_inputs(r, n, k, sizes, "cuda") + dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) + assert dw.dtype == torch.bfloat16 and dw.shape == (len(sizes), n, k) + torch.testing.assert_close(dw.float(), ref.float(), atol=2e-3, rtol=0.01) + assert compute_error(ref.float(), dw.float()) >= 24.0 + + +@_gpu +def test_wgrad_bitwise_exact_integers(): + torch.manual_seed(2) + sizes = [256, 128, 384] + args = _make_c_inputs(768, 256, 256, sizes, "cuda", exact_int=True) + dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) + assert torch.equal(_bytes(dw), _bytes(ref)), "exact-integer wgrad must be bitwise" + + +@_gpu +def test_wgrad_rejects_kgroups_scale_ordering_by_result(): + """Whole-matrix vs per-group K-groups blocked scales differ whenever N > 128; + no length check can catch it, so prove the kernel is sensitive to it.""" + torch.manual_seed(3) + device = "cuda" + sizes = [128, 128] + r, n, k = 256, 256, 128 + dy = torch.randn(r, n, device=device, dtype=torch.bfloat16) + x = torch.randn(r, k, device=device, dtype=torch.bfloat16) + dy_q, dy_sf = _quantize_colwise_ref(dy) + x_q, x_sf = _quantize_colwise_ref(x) + offsets = _mk_offsets(sizes, device) + + good = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf, x_q, x_sf, offsets) + ref = _ref_wgrad(dy_q, dy_sf, x_q, x_sf, sizes) + torch.testing.assert_close(good.float(), ref.float(), atol=2e-3, rtol=0.01) + + # Re-encode dy's scales per group (torchao K-groups form) and rerun. + scale_t, _ = to_mx(dy.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + per_group = torch.cat( + [ + to_blocked(scale_t[:, s // _BLOCK : e // _BLOCK]).reshape(-1) + for s, e in ((0, 128), (128, 256)) + ] + ) + assert per_group.numel() == dy_sf.numel() + assert not torch.equal(_bytes(per_group), _bytes(dy_sf.reshape(-1))) + bad = torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_q, per_group, x_q, x_sf, offsets + ) + assert not torch.equal(bad, good), ( + "kernel must consume whole-matrix blocked scales; a K-groups buffer of " + "identical length must change the result" + ) + + +# --------------------------------------------------------------------------- +# 6. Kernel A numerics +# --------------------------------------------------------------------------- + +_A_SHAPES = [ + (256, 256, 128, [128, 128]), + (1024, 512, 256, [256, 0, 512, 256]), # zero-token expert + ragged + (1536, 2048, 1408, [512, 128, 640, 256]), # 16B class +] + + +@_gpu +@pytest.mark.parametrize("r,d,f,sizes", _A_SHAPES, ids=["small", "ragged", "16b"]) +def test_swiglu_fwd_random(r, d, f, sizes): + torch.manual_seed(4) + args, z_ref = _make_a_inputs(r, d, f, sizes, "cuda") + z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) + active = sum(sizes) + assert ( + compute_error(z_ref[:active].float(), z[:active].reshape(active, -1).float()) + >= 27.0 + ) + # Quantized outputs compare against the normative chain applied to the + # kernel's own z (removes GEMM reduction-order noise from the comparison). + # Measured bitwise-identical on all of these shapes (the kernel's sigmoid + # composition matches torch's float32 silu exactly), so no mismatch budget. + gate = z[..., 0].float() + up = z[..., 1].float() + h_ref = (F.silu(gate) * up).bfloat16() + _assert_quantized_pair(hq, hsf, hcq, hcsf, h_ref) + + +@_gpu +def test_swiglu_fwd_bitwise_exact_integers(): + torch.manual_seed(5) + r, d, f, sizes = 512, 256, 128, [256, 256] + args, z_ref = _make_a_inputs(r, d, f, sizes, "cuda", exact_int=True) + z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) + assert torch.equal( + _bytes(z.reshape(r, -1)), _bytes(z_ref.bfloat16().reshape(r, -1)) + ), "integer-exact z must be bitwise" + + +@_gpu +def test_swiglu_fwd_saturated_gate_bitwise(): + """gate == 128 makes silu exact in any implementation, so the fused dual + quantization must match to_mx byte-for-byte, including special values.""" + torch.manual_seed(6) + device = "cuda" + r, d, f = 256, 128, 128 + sizes = [256] + # x = identity blocks: row r selects weight column r % d. + x_q = torch.zeros(r, d, dtype=torch.uint8, device=device) + x_q[torch.arange(r), torch.arange(r) % d] = 0x38 # 1.0 + x_q = x_q.view(_E4M3) + x_logical = torch.full((r, d // _BLOCK), 127, dtype=torch.uint8, device=device) + x_sf = to_blocked(x_logical.view(_E8M0)) + + # Up rows: random E4M3 bytes (finite lanes), then crafted special features. + # Feature f's specials land in rowwise scale block f // 32 of every row. + up_bytes = torch.randint(0, 0x7E, (f, d), dtype=torch.uint8, device=device) + up_scales = torch.randint( + 120, 134, (f, d // _BLOCK), dtype=torch.uint8, device=device + ) + up_bytes[0, :] = 0x7F # NaN up -> h NaN in block 0 + up_bytes[1, :] = 0x00 # zero up + up_bytes[40, :] = 0x7E # 448 * 2^119: z_up ~ 2.98e38 (finite bf16); + up_scales[40, :] = 127 + 119 # h = 128 * z_up overflows f32 -> +Inf, block 1 + w_q2f = torch.zeros(2 * f, d, dtype=torch.uint8, device=device) + w_scale2f = torch.zeros(2 * f, d // _BLOCK, dtype=torch.uint8, device=device) + w_q2f[0::2] = 0x38 + w_scale2f[0::2] = 127 + 7 + w_q2f[1::2] = up_bytes + w_scale2f[1::2] = up_scales + w_q = w_q2f.view(_E4M3) + w_sf = to_blocked(w_scale2f.view(_E8M0)) + w_packed, w_sf_packed = _pack_grouped_weight([w_q], [w_sf], device) + offsets = _mk_offsets(sizes, device) + + z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w_packed, w_sf_packed, offsets + ) + gate = z[..., 0].float() + up = z[..., 1].float() + assert torch.equal(gate, torch.full_like(gate, 128.0)), "gate must be exactly 128" + h_ref = (128.0 * up).bfloat16() # silu(128) == 128 exactly + _assert_quantized_pair(hq, hsf, hcq, hcsf, h_ref) # zero mismatch budget + # Special-value spot checks straight from the RCEIL table: + hq_bytes = _bytes(hq).reshape(r, f) + hsf_logical = _bytes( + from_blocked(hsf.view(_E8M0).reshape(-1), r, f // _BLOCK) + ).reshape(r, f // _BLOCK) + nan_blocks = torch.isnan(h_ref).view(r, f // _BLOCK, _BLOCK).any(-1) + inf_blocks = torch.isinf(h_ref).view(r, f // _BLOCK, _BLOCK).any(-1) + assert nan_blocks.any() and inf_blocks.any(), "crafted specials must appear" + assert (hsf_logical[nan_blocks | inf_blocks] == 0xFF).all() + nonfinite_cols = (nan_blocks | inf_blocks).repeat_interleave(_BLOCK, dim=1) + assert (hq_bytes[nonfinite_cols] == 0x7F).all() + + +@_gpu +def test_swiglu_fwd_tail_and_poison(): + torch.manual_seed(7) + device = "cuda" + r, d, f, sizes = 512, 256, 128, [128, 256] # active 384, tail 128 + args, _ = _make_a_inputs(r, d, f, sizes, device) + z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) + active = sum(sizes) + assert (_bytes(z.reshape(r, -1))[active:] == 0).all(), "z tail must be zero bytes" + assert (_bytes(hq)[active:] == 0).all() + assert (_bytes(hcq.t())[:, active:] == 0).all() + hsf_logical = _bytes(from_blocked(hsf.view(_E8M0).reshape(-1), r, f // _BLOCK)) + assert (hsf_logical.reshape(r, -1)[active:] == 0).all() + hcsf_logical = _bytes(from_blocked(hcsf.view(_E8M0).reshape(-1), f, r // _BLOCK)) + assert (hcsf_logical.reshape(f, -1)[:, active // _BLOCK :] == 0).all() + + +# --------------------------------------------------------------------------- +# 7. Kernel B numerics +# --------------------------------------------------------------------------- + + +@_gpu +@pytest.mark.parametrize( + "r,d,f,sizes", + [ + (256, 256, 128, [128, 128]), + (1024, 512, 256, [256, 0, 512, 128]), # zero-token + strict tail + (1536, 2048, 1408, [512, 128, 640, 256]), # 16B class + ], + ids=["small", "ragged_tail", "16b"], +) +def test_dswiglu_bwd_random(r, d, f, sizes): + torch.manual_seed(8) + args, dh_ref = _make_b_inputs(r, d, f, sizes, "cuda") + dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd(*args) + active = sum(sizes) + z = args[4] + dz_ref = _b_reference_dz(dh_ref.bfloat16(), z) + got = _dequant_rowwise(dzq, dzsf, dtype=torch.float32) + assert compute_error(dz_ref[:active].float(), got[:active]) >= 25.0 + got_col = _dequant_colwise(dzcq, dzcsf, dtype=torch.float32) + assert compute_error(dz_ref[:active].float(), got_col[:active]) >= 25.0 + if active < r: + assert (_bytes(dzq)[active:] == 0).all(), "dz tail must be zero bytes" + assert (_bytes(dzcq.t())[:, active:] == 0).all() + + +@_gpu +def test_dswiglu_bwd_bitwise_exact(): + """Exact dh (integer GEMM) + saturated/zero gates: dz is exact products, so + all four outputs must match to_mx byte-for-byte; interleave order checked.""" + torch.manual_seed(9) + device = "cuda" + r, d, f, sizes = 256, 128, 128, [256] + # dh == 1.0 exactly: do = identity rows, w2 = all ones. + do_q = torch.zeros(r, d, dtype=torch.uint8, device=device) + do_q[torch.arange(r), torch.arange(r) % d] = 0x38 + do_q = do_q.view(_E4M3) + do_sf = to_blocked( + torch.full((r, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) + ) + w_q = torch.full((f, d), 0x38, dtype=torch.uint8, device=device).view(_E4M3) + w_sf = to_blocked( + torch.full((f, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) + ) + w_packed, w_sf_packed = _pack_grouped_weight([w_q], [w_sf], device) + + # Saturated gates make dsilu == 1 and silu == gate exactly; gate rows of z + # also cover 0 (silu(0) == 0, dsilu(0) == 0.5 -- both exact). + z = torch.zeros(r, f, 2, device=device, dtype=torch.bfloat16) + z[..., 0] = 128.0 + z[: r // 2, :, 1] = torch.randn(r // 2, f, device=device).bfloat16() + z[r // 2 :, :, 0] = 0.0 # gate 0 rows: dgate = 0.5 * up, dup = 0 + z[r // 2 :, :, 1] = ( + torch.randn(r - r // 2, f, device=device).bfloat16().float() * 2.0 + ).bfloat16() + z[0, 0:4, 0] = float("nan") # NaN gate+up block + z[0, 0:4, 1] = float("nan") + z[1, 0:16, 1] = 448.0 # uniform 448 dgate lanes + offsets = _mk_offsets(sizes, device) + + dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w_packed, w_sf_packed, z, offsets + ) + dh = torch.ones(r, f, device=device, dtype=torch.bfloat16) + dz_ref = _b_reference_dz(dh, z) + # Interleave check on exact lanes (bitwise: NaN blocks are present). + dgate_ref, dup_ref = _dswiglu_closed_form( + dh.float(), z[..., 0].float(), z[..., 1].float() + ) + assert torch.equal(_bytes(dz_ref[:, 0::2]), _bytes(dgate_ref.bfloat16())) + assert torch.equal(_bytes(dz_ref[:, 1::2]), _bytes(dup_ref.bfloat16())) + _assert_quantized_pair(dzq, dzsf, dzcq, dzcsf, dz_ref) # zero budget + + +@_gpu +def test_dswiglu_bwd_semantic_blocks(): + """Drive the shared MXFP8 semantic contract through the fused backward. + + Uniform cases with |value| >= 128 are constructed exactly (gate = up = + value under a saturated gate gives a uniform dz block); the rest are + covered transitively: kernel bytes must equal to_mx bytes on blocks + containing the case values, and to_mx itself is asserted against the + shared table in test_to_mx_matches_semantic_table.""" + device = "cuda" + cases = make_mxfp8_semantic_cases(torch.bfloat16, _RCEIL, device=device) + n_cases = len(cases.names) + r, d, f = 128, 128, max(128, _round_up(n_cases * 16, 128)) + do_q = torch.zeros(r, d, dtype=torch.uint8, device=device) + do_q[torch.arange(r), torch.arange(r) % d] = 0x38 + do_q = do_q.view(_E4M3) + do_sf = to_blocked( + torch.full((r, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) + ) + w_q = torch.full((f, d), 0x38, dtype=torch.uint8, device=device).view(_E4M3) + w_sf = to_blocked( + torch.full((f, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) + ) + w_packed, w_sf_packed = _pack_grouped_weight([w_q], [w_sf], device) + + # Row 0: each case occupies 16 features -> one 32-wide dz block. + z = torch.zeros(r, f, 2, device=device, dtype=torch.bfloat16) + z[..., 0] = 128.0 + direct_table = [] + for idx in range(n_cases): + vals = cases.inputs[idx].to(device) + f0 = idx * 16 + uniform = bool((vals == vals[0]).all()) and not torch.isnan(vals).any() + if uniform and abs(float(vals[0])) >= 128.0: + # gate = up = v: dgate = v, dup = gate = |v|-signed... both = v + # requires gate == v which needs v >= 128; negatives via up lane. + v = float(vals[0]) + if v >= 128.0: + z[0, f0 : f0 + 16, 0] = v + z[0, f0 : f0 + 16, 1] = v + direct_table.append((idx, None)) + else: + z[0, f0 : f0 + 16, 0] = -v + z[0, f0 : f0 + 16, 1] = v + direct_table.append((idx, "even_only")) + else: + # Transitive: dgate lanes carry the case's even values, dup lanes + # its odd values scaled through the saturated gate where possible; + # fall back to plain interleave of the case into up lanes. + z[0, f0 : f0 + 16, 1] = vals[0::2] + z[0, f0 : f0 + 16, 0] = torch.where( + torch.isfinite(vals[1::2]) & (vals[1::2].abs() >= 128), + vals[1::2], + torch.full_like(vals[1::2], 128.0), + ) + offsets = _mk_offsets([r], device) + dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w_packed, w_sf_packed, z, offsets + ) + dh = torch.ones(r, f, device=device, dtype=torch.bfloat16) + dz_ref = _b_reference_dz(dh, z) + _assert_quantized_pair(dzq, dzsf, dzcq, dzcsf, dz_ref) # bitwise transitivity + # Direct table assertions where the block is exactly the case input. + row_sf_logical = _bytes( + from_blocked(dzsf.view(_E8M0).reshape(-1), r, 2 * f // _BLOCK) + ).reshape(r, -1) + dz_bytes = _bytes(dzq).reshape(r, -1) + for idx, mode in direct_table: + blk = slice(idx * _BLOCK, (idx + 1) * _BLOCK) + want_scale = int(cases.expected_scales[idx]) + want_data = cases.expected_data[idx].to(device) + assert row_sf_logical[0, idx] == want_scale, cases.names[idx] + got = dz_bytes[0, blk] + if mode is None: + assert torch.equal(got, want_data.to(got.device)), cases.names[idx] + else: + assert torch.equal(got[0::2], want_data.to(got.device)[0::2]), cases.names[ + idx + ] + + +def _semantic_table_dtype_check(device): + cases = make_mxfp8_semantic_cases(torch.bfloat16, _RCEIL, device=device) + scale, q = to_mx(cases.inputs, _E4M3, _BLOCK, scaling_mode=_RCEIL) + assert torch.equal( + _bytes(q).cpu().reshape(len(cases.names), 32), cases.expected_data + ) + assert torch.equal( + _bytes(scale.reshape(-1)).cpu(), cases.expected_scales.reshape(-1) + ) + + +@_gpu +def test_to_mx_matches_semantic_table(): + """Anchors the transitive comparisons above: to_mx == the shared table.""" + _semantic_table_dtype_check("cuda") + + +# --------------------------------------------------------------------------- +# 8. torch.compile +# --------------------------------------------------------------------------- + + +@_gpu +@pytest.mark.parametrize("op_name", _OP_NAMES) +def test_compile_matches_eager(op_name): + torch.manual_seed(10) + device = "cuda" + if op_name == "mxfp8_grouped_gemm_wgrad": + args = _make_c_inputs(256, 256, 128, [128, 128], device) + elif op_name == "mxfp8_grouped_gemm_swiglu_fwd": + args, _ = _make_a_inputs(256, 256, 128, [128, 128], device) + else: + args, _ = _make_b_inputs(256, 256, 128, [128, 128], device) + op = getattr(torch.ops.torchao, op_name) + eager = op(*args) + compiled_fn = torch.compile(lambda *a: op(*a), fullgraph=True) + compiled = compiled_fn(*args) + eager = eager if isinstance(eager, tuple) else (eager,) + compiled = compiled if isinstance(compiled, tuple) else (compiled,) + for e, c in zip(eager, compiled): + assert e.stride() == c.stride() + assert torch.equal(_bytes(e.reshape(-1)), _bytes(c.reshape(-1))) + + +# --------------------------------------------------------------------------- +# 9. One physical launch per op (profiler evidence) +# --------------------------------------------------------------------------- + + +def _device_kernel_names(fn): + fn() + torch.cuda.synchronize() + fn() + torch.cuda.synchronize() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA] + ) as prof: + fn() + torch.cuda.synchronize() + names = [] + for evt in prof.key_averages(): + if evt.device_type != torch.autograd.DeviceType.CUDA: + continue + if "Memset" in evt.key or "Memcpy" in evt.key: + continue + names.extend([evt.key] * evt.count) + return names + + +@_gpu +@pytest.mark.parametrize("op_name", _OP_NAMES) +def test_one_physical_launch(op_name): + torch.manual_seed(11) + device = "cuda" + if op_name == "mxfp8_grouped_gemm_wgrad": + args = _make_c_inputs(256, 256, 128, [128, 128], device) + elif op_name == "mxfp8_grouped_gemm_swiglu_fwd": + args, _ = _make_a_inputs(256, 256, 128, [128, 128], device) + else: + args, _ = _make_b_inputs(256, 256, 128, [128, 128], device) + op = getattr(torch.ops.torchao, op_name) + names = _device_kernel_names(lambda: op(*args)) + assert len(names) == 1, f"{op_name} must be ONE kernel launch, saw: {names}" + banned = ("quantize", "silu", "elementwise", "vectorized") + assert not any(b in names[0].lower() for b in banned), names + + +# --------------------------------------------------------------------------- +# 10. Large DSv3-class shapes (memory gated) +# --------------------------------------------------------------------------- + + +def _enough_memory(bytes_needed: int) -> bool: + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_properties(0).total_memory >= bytes_needed + + +@_gpu +@pytest.mark.skipif(not _enough_memory(16 << 30), reason="needs >= 16 GiB") +def test_wgrad_dsv3_671b_class(): + torch.manual_seed(12) + sizes = [256, 0, 512, 256] + args = _make_c_inputs(1024, 4096, 7168, sizes, "cuda") + dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) + ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) + assert compute_error(ref.float(), dw.float()) >= 24.0 + + +@_gpu +@pytest.mark.skipif(not _enough_memory(16 << 30), reason="needs >= 16 GiB") +def test_swiglu_fwd_dsv3_671b_class(): + torch.manual_seed(13) + args, z_ref = _make_a_inputs(1024, 7168, 2048, [256, 0, 512, 256], "cuda") + z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) + active = 1024 + assert ( + compute_error(z_ref[:active].float(), z[:active].reshape(active, -1).float()) + >= 27.0 + ) + + +# --------------------------------------------------------------------------- +# 11. Availability / unsupported environments +# --------------------------------------------------------------------------- + + +def test_wrapper_unavailable_raises_cleanly(monkeypatch): + mod = pytest.importorskip("torchao.prototype.moe_training.mxfp8_grouped_mlp") + flag = "_mxfp8_grouped_mlp_kernels_available" + if not hasattr(mod, flag): + pytest.skip("availability flag not exposed") + monkeypatch.setattr(mod, flag, False) + with pytest.raises(NotImplementedError): + with FakeTensorMode(): + mod.mxfp8_grouped_gemm_wgrad(*_fake_c_inputs(256, 256, 128, 2)) + + +def test_jagged_offs_generator_contract(): + """The test/bench offsets helper must produce 128-multiples when asked.""" + random.seed(0) + if torch.cuda.is_available(): + offs = generate_jagged_offs(4, 1024, multiple_of=128) + else: + offs = generate_jagged_offs(4, 1024, multiple_of=128, device="cpu") + sizes = torch.diff(offs.cpu(), prepend=torch.tensor([0], dtype=offs.dtype)) + assert (sizes % 128 == 0).all() + assert offs[-1].item() == 1024 diff --git a/torchao/prototype/moe_training/__init__.py b/torchao/prototype/moe_training/__init__.py index e58461a80c..cd3bf7abc3 100644 --- a/torchao/prototype/moe_training/__init__.py +++ b/torchao/prototype/moe_training/__init__.py @@ -1,6 +1,14 @@ from torchao.prototype.moe_training.fp8_grouped_mm import ( _to_fp8_rowwise_then_scaled_grouped_mm, ) +from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( + is_supported as mxfp8_grouped_mlp_is_supported, +) +from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( + mxfp8_grouped_gemm_dswiglu_bwd, + mxfp8_grouped_gemm_swiglu_fwd, + mxfp8_grouped_gemm_wgrad, +) from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _to_mxfp8_then_scaled_grouped_mm, ) @@ -8,4 +16,8 @@ __all__ = [ "_to_mxfp8_then_scaled_grouped_mm", "_to_fp8_rowwise_then_scaled_grouped_mm", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_wgrad", + "mxfp8_grouped_mlp_is_supported", ] diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index 0ed217b43a..da39d890d3 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,3 +1,10 @@ +# Importing grouped_mlp_ops registers the fused grouped-MLP custom ops +# (torchao::mxfp8_grouped_gemm_{swiglu_fwd,dswiglu_bwd,wgrad}). The module is +# importable with no CuTe DSL installed; kernel imports are deferred into the +# op bodies. +from torchao.prototype.moe_training.kernels.mxfp8 import ( + grouped_mlp_ops, # noqa: F401 +) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, # noqa: F401 _mxfp8_flydsl_kernels_available, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py index 6036860cf9..3947cc6970 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py @@ -4,20 +4,1720 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. -"""Launcher facade for the MXFP8 routed-expert grouped-MLP CuTe DSL kernels. +"""MXFP8 routed-expert grouped-MLP kernels for SM100 (CuTe DSL, public API only). -``grouped_mlp_ops`` imports its three launchers from this module and from -nowhere else, so this name and these three signatures are the seam between the -custom-op layer and the kernels. Keep it a pure re-export: anything defined here -would be code the ops layer depends on but that no kernel test covers. +Three physically fused kernels, one launch each: -The import is deliberately per-launcher rather than a package-level ``*``, so a -kernel that has not landed yet costs an ImportError naming that kernel instead of -breaking the two that have. +* ``launch_grouped_gemm_swiglu_fwd`` -- FC1 ragged grouped GEMM + SwiGLU + + rowwise (1x32) and columnwise (32x1) MXFP8 RCEIL quantization, plus the BF16 + pre-activation save. +* ``launch_grouped_gemm_dswiglu_bwd`` -- FC2 dgrad ragged grouped GEMM + + dSwiGLU + the same dual quantization of the FC1 input gradient. +* ``launch_grouped_gemm_wgrad`` -- generic ragged-K grouped weight + gradient, BF16 output; called once for FC1 and once for FC2. + +They share one blockscaled tcgen05 mainloop. Three structural decisions, all +descending from every per-expert row count being a multiple of 128: + +*No per-group tensormaps.* Every operand is one host-built static TMA +descriptor over the whole tensor. Per-expert selection is an integer +coordinate: an L coordinate for the 3-D weight operands, and a K-tile index +base for the wgrad kernel's ragged contraction. The latter is plain layout +algebra: the scale-factor tensor is retiled with +``blockscaled_utils.tile_atom_to_shape_SF``, whose K-tile mode is uniform by +construction, so slicing the partitioned tensor at ``k_base + i`` addresses +expert data exactly, with no per-expert descriptor rebuilds. + +*No tile scheduler.* With the ragged axis tile-aligned, the forward/backward +kernels enumerate all of ``[0, R/128)`` M tiles and the wgrad grid +``(N/128, K/128, G)`` is fully static; only wgrad's K-loop trip count is +data-dependent. + +*The inactive tail needs no special code path.* A tile whose row base is at or +past the active row count runs with ``k_cnt == 0``: no TMA loads are issued, +the accumulator fragment is zeroed in registers, and the unmodified epilogue +emits the zero bytes the contract requires. Epilogue *stores* are never +predicated; the one epilogue-side gmem *input* (the backward kernel's saved +``z_bf16``) is loaded only when ``k_cnt > 0`` because tail rows of that tensor +are read-forbidden. + +Offset contract (documented caller invariants -- the offset VALUES live on +device and cannot be checked on the host without a synchronization): offsets +are exclusive per-expert end indices, int32, CUDA, 1-D, contiguous, +nondecreasing, every per-expert row count a multiple of 128, and +``offsets[-1] <= R``. The launchers validate all metadata and reject the rest +of the malformed-input space with ``ValueError``; set +``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` to additionally validate the values on +the host while debugging, at the cost of a D2H copy. There is deliberately no +device-side assertion in the default build: assertions are compiled out of +CuTe DSL kernels unless ``CUTE_DSL_ENABLE_ASSERTIONS=1``, so they cannot serve +as a production guard, and with malformed offsets the wgrad kernel returns a +wrong result rather than faulting -- "it did not crash" is not evidence the +offsets were valid. """ -from torchao.prototype.moe_training.kernels.mxfp8.kernel_wgrad import ( - launch_grouped_gemm_wgrad, +import functools +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +from cutlass import Float32, Int32 +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.utils import LayoutEnum + +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( + _is_fake, + validate_allocated_rows, + validate_blocked_scales, + validate_destination, + validate_feature_dims, + validate_group_offsets, + validate_grouped_operand, ) -__all__ = ["launch_grouped_gemm_wgrad"] +__all__ = [ + "launch_grouped_gemm_swiglu_fwd", + "launch_grouped_gemm_dswiglu_bwd", + "launch_grouped_gemm_wgrad", +] + +# --------------------------------------------------------------------------- +# Frozen configuration. One tiling, one pipeline shape, one warp assignment. +# --------------------------------------------------------------------------- + +# MXFP8 scaling block: 32 values share one E8M0 scale. +_SF_VEC_SIZE = 32 +# cta_tile_m = 128 is required by CtaGroup.ONE and by the no-partial-M-tile +# argument. cta_tile_k = 128 is pinned: the wgrad kernel selects an expert's K +# range with an integer K-tile index base, exact only because every group +# boundary (a multiple of 128 rows) is a multiple of the K tile. +_CTA_M = 128 +_CTA_N = 128 +_CTA_K = 128 +_MMA_TILER = (_CTA_M, _CTA_N, _CTA_K) +# One stage: A tile + B tile (E4M3) + one 512-byte scale atom per operand. +# This is also the exact tx_count the TMA pipeline barrier must expect. +_AB_STAGE_BYTES = _CTA_M * _CTA_K + _CTA_N * _CTA_K + 2 * (128 * (_CTA_K // 32)) +_NUM_AB_STAGE = 6 +_NUM_ACC_STAGE = 1 +# Warp 0 loads (TMA), warp 1 issues the MMA, warps 4-7 run the epilogue, warps +# 2-3 idle. The epilogue must start on a warp quad (warp id multiple of 4): +# tcgen05.ld selects its TMEM datapath sub-partition from the physical warp +# id, and a misaligned epilogue block would read every 128-row tile with its +# 32-row groups rotated. +_THREADS = 256 +_TMA_WARP_ID = 0 +_MMA_WARP_ID = 1 +_FIRST_EPI_WARP = 4 +_NUM_EPI_THREADS = 128 +_FIRST_EPI_THREAD = 32 * _FIRST_EPI_WARP +# Named barrier ids (0 is left free for the DSL's own use). +_EPI_FINAL_BARRIER_ID = 1 +_TMEM_ALLOC_BARRIER_ID = 2 +_EPI_STAGE_BARRIER_ID = 3 +# sm_100 usable dynamic shared memory per CTA, (228 - 1) KiB. +_SMEM_CAPACITY_BYTES = 232448 +# The TMEM allocator requires a power-of-two multiple of 32 columns; shared +# memory already pins us to one CTA per SM, so taking the whole array is free. +_TMEM_TOTAL_COLS = 512 + + +@dataclass(frozen=True) +class _KernelConfig: + """Per-kernel trace-time constants (everything else is module-frozen).""" + + # Accumulator columns handed to the epilogue per subtile. 64 for the + # forward (an adjacent gate/up accumulator pair per output column, so 64 + # accumulator columns are 32 output columns = one 1x32 block per row), 32 + # for the backward and for wgrad. + epi_n_acc: int + # False: offsets partition the GEMM M axis (forward/backward). True: they + # partition the contraction (wgrad). + ragged_k: bool + # Columns of the [128, cols] BF16 epilogue staging tile (0 = no staging). + # Padded to an odd count so columnwise reads don't serialize on banks. + epi_smem_cols: int + + @property + def epi_tile(self): + return (_CTA_M, self.epi_n_acc) + + @property + def num_epi_subtiles(self) -> int: + return _CTA_N // self.epi_n_acc + + @property + def epi_smem_elems(self) -> int: + # A zero-size struct field would be degenerate; keep a tiny slab. + return max(_CTA_M * self.epi_smem_cols, 8) + + +_SWIGLU_FWD_CONFIG = _KernelConfig(epi_n_acc=64, ragged_k=False, epi_smem_cols=33) +_DSWIGLU_BWD_CONFIG = _KernelConfig(epi_n_acc=32, ragged_k=False, epi_smem_cols=65) +_WGRAD_CONFIG = _KernelConfig(epi_n_acc=32, ragged_k=True, epi_smem_cols=0) + + +# --------------------------------------------------------------------------- +# Host-side operand views. Pure torch restrides into the (MN, K, L) GEMM +# domain, K contiguous. +# --------------------------------------------------------------------------- + + +def activation_gemm_view(t: torch.Tensor) -> torch.Tensor: + """``[MN, K]`` K-contiguous -> ``(MN, K, 1)`` with a defined batch stride.""" + mn, k = t.shape + if t.stride() != (k, 1): + raise ValueError( + f"GEMM operand must be K-contiguous with stride {(k, 1)}, got {t.stride()}" + ) + return torch.as_strided(t, (mn, k, 1), (k, 1, mn * k)) + + +def weight_gemm_view(w: torch.Tensor) -> torch.Tensor: + """``[G, K, N]`` stride ``(K*N, 1, K)`` -> ``(N, K, G)``: expert becomes L.""" + g, k, n = w.shape + if w.stride() != (k * n, 1, k): + raise ValueError( + f"grouped weight must have stride {(k * n, 1, k)}, got {w.stride()}" + ) + return w.permute(2, 1, 0) + + +@dataclass +class TileCoords: + """Per-CTA tile description handed to the epilogue. All fields CTA-uniform. + + ``k_cnt == 0`` marks both an inactive tail tile (ragged M) and a + zero-token expert (ragged K); the accumulator arrives zeroed and the + epilogue must store unconditionally. + """ + + tile_m: Int32 + tile_n: Int32 + expert: Int32 + row_base: Int32 + col_base: Int32 + k_cnt: Int32 + + +# --------------------------------------------------------------------------- +# Mainloop building blocks (all public CuTe DSL API). +# --------------------------------------------------------------------------- + + +def _make_tiled_mma(a_dtype, b_dtype, sf_dtype): + """The one blockscaled tiled MMA, K-major on both operands, FP32 acc.""" + return sm100_utils.make_blockscaled_trivial_tiled_mma( + a_dtype, + b_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + sf_dtype, + _SF_VEC_SIZE, + tcgen05.CtaGroup.ONE, + (_CTA_M, _CTA_N), + ) + + +def _make_sf_gemm_tensor(flat_sf: cute.Tensor, mn: int, k: int, l: int): + """Retile a flat blocked E8M0 buffer into the GEMM-domain SF layout. + + The buffer travels flat by ABI and may arrive as raw uint8, so the pointer + is recast: the MMA rejects a scale operand that is not E8M0. + ``tile_atom_to_shape_SF`` builds kernel IR, so this runs inside the trace. + """ + return cute.make_tensor( + cute.recast_ptr(flat_sf.iterator, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.tile_atom_to_shape_SF((mn, k, l), _SF_VEC_SIZE), + ) + + +def _t2r_partition(tidx, tAcc_base: cute.Tensor, cfg): + """TMEM accumulator -> register handoff. + + ``tTR_cAcc`` carries each register's ``(row, col)`` coordinate in the CTA + tile; every epilogue index below is derived from it rather than from a raw + register number or an assumed thread-to-row mapping, so the addressing is + correct by construction for whatever value order the copy atom uses. + ``elem_ty_d`` stays Float32: an 8-bit d type would steer + ``get_tmem_load_op`` into layouts shaped for a direct FP8 TMA store. + """ + copy_atom_t2r = sm100_utils.get_tmem_load_op( + _MMA_TILER, + LayoutEnum.ROW_MAJOR, + Float32, + Float32, + cfg.epi_tile, + False, + ) + tAcc_mn = tAcc_base[((None, None), 0, 0, 0)] + tAcc_epi = cute.flat_divide(tAcc_mn, cfg.epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + cAcc_epi = cute.flat_divide( + cute.make_identity_tensor((_CTA_M, _CTA_N)), cfg.epi_tile + ) + tTR_cAcc = thr_copy_t2r.partition_D(cAcc_epi) + tTR_rAcc = cute.make_rmem_tensor(tTR_cAcc[(None, None, None, 0, 0)].shape, Float32) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc + + +def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor): + """SMEM -> TMEM scale-factor copy, issued once per K tile from the MMA warp. + + ``Cp4x32x128bOp`` must be issued as a plain ``cute.copy`` -- the DSL + inserts the single-thread election itself, and wrapping it in + ``elect_one()`` deadlocks. + """ + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), sSF.element_type + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, thr_copy_s2t.partition_S(tCsSF_compact) + ) + tCtSF_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_s2t, tCtSF_s2t + + +# --------------------------------------------------------------------------- +# Quantization math. Public conversions only; every step mirrors the torchao +# reference (`to_mx(..., RCEIL)` + `to_blocked`) so the fused outputs are +# byte-identical to the standalone quantizers on the same BF16 input. +# --------------------------------------------------------------------------- + + +@cute.jit +def _sigmoid_f32(x: Float32) -> Float32: + """sigmoid(x) = 1 / (1 + exp(-x)), accurate mode. + + Composed exactly like torch's float32 sigmoid (default-mode ``exp`` plus a + true divide); measured bit-identical to ``torch.sigmoid`` over 1e6 values. + """ + return Float32(1.0) / (Float32(1.0) + cute.math.exp(Float32(0.0) - x)) + + +@cute.jit +def _blocked_scale_idx(row: Int32, scale_col: Int32, ncb: Int32) -> Int32: + """Flat byte index in the tcgen05 blocked (128x4) scale layout. + + The logical ``[rows, cols]`` scale matrix is stored as 512-byte tiles of + 128 rows x 4 scale columns, tiles ordered ``row_block * ncb + col_block`` + with ``ncb = ceil_div(cols, 4)`` -- torchao ``to_blocked``, whole-matrix. + Coordinates are ABSOLUTE. For columnwise scales pass transposed + coordinates (feature index as ``row``, 32-row block index as ``col``). + """ + return ( + ((row >> 7) * ncb + (scale_col >> 2)) * Int32(512) + + (row & Int32(31)) * Int32(16) + + ((row >> 5) & Int32(3)) * Int32(4) + + (scale_col & Int32(3)) + ) + + +@cute.jit +def _store_frag(dst: cute.Tensor, elem_offset: Int32, frag: cute.Tensor): + """Store a register fragment as one contiguous run of ``dst`` elements. + + Callers guarantee the element offset is at least 16-byte aligned relative + to the (validated, 32-byte-aligned) base, so the copy vectorizes. + """ + cute.autovec_copy( + frag, + cute.make_tensor( + (dst.iterator + elem_offset).align(16), + cute.make_layout(cute.size(frag)), + ), + ) + + +@cute.jit +def _quant_block_from_smem( + sEpi: cute.Tensor, + base_row: Int32, + base_col: Int32, + q_dst: cute.Tensor, + q_offset: Int32, + sf_dst: cute.Tensor, + sf_idx: Int32, + COLWISE: cutlass.Constexpr, +): + """Quantize one 32-value MX block from the BF16 staging tile. + + Reads ``sEpi[base_row, base_col + i]`` (rowwise) or + ``sEpi[base_row + i, base_col]`` (columnwise), then: + + * amax: NaN-propagating |max| chain, so a NaN element invalidates the + block exactly like the torchao reference. + * scale: ``descale = amax / 448``; the public Float32 -> Float8E8M0FNU + conversion rounds toward +inf (RCEIL) by construction. A non-finite + amax (Inf would otherwise clamp to byte 254) is overridden to byte 255. + * reciprocal, in the E8M0 byte domain: ``254 - byte`` reinterpreted as + E8M0 and widened exactly to f32. Byte 0 (zero/tiny block) descales by + 2^127; byte 255 gives a NaN reciprocal so every element of an + invalidated block quantizes to the E4M3 NaN code. + * qdata: one f32 multiply per element, then the public saturating-RNE + Float32 -> Float8E4M3FN conversion (byte-identical to torch's cast; no + explicit clamp -- saturation is part of the conversion's contract). + + 32 qdata bytes are one contiguous run of ``q_dst`` in both orientations + (row-major rowwise output; column-major columnwise output), stored + vectorized; the scale byte is stored individually at ``sf_idx``. + """ + vals = [] + for i in cutlass.range_constexpr(_SF_VEC_SIZE): + if cutlass.const_expr(COLWISE): + v = sEpi[base_row + Int32(i), base_col] + else: + v = sEpi[base_row, base_col + Int32(i)] + vals.append(Float32(v)) + + amax = cute.arch.fmax(vals[0], vals[1], abs=True, nan=True) + for i in cutlass.range_constexpr(2, _SF_VEC_SIZE): + amax = cute.arch.fmax(amax, vals[i], abs=True, nan=True) + + scale_byte = Int32( + cutlass.Float8E8M0FNU(amax / Float32(448.0)).bitcast(cutlass.Int8) + ) & Int32(0xFF) + amax_bits = Float32(amax).bitcast(Int32) + if (amax_bits & Int32(0x7F800000)) == Int32(0x7F800000): + scale_byte = Int32(255) + + recip_byte = (Int32(254) - scale_byte) & Int32(0xFF) + recip = Float32(cutlass.Uint8(recip_byte).bitcast(cutlass.Float8E8M0FNU)) + + qfrag = cute.make_rmem_tensor(cute.make_layout(_SF_VEC_SIZE), cutlass.Float8E4M3FN) + for i in cutlass.range_constexpr(_SF_VEC_SIZE): + qfrag[i] = cutlass.Float8E4M3FN(vals[i] * recip) + _store_frag(q_dst, q_offset, qfrag) + sf_dst[sf_idx] = cutlass.Uint8(scale_byte) + + +@cute.jit +def _epilogue_column_run( + tTR_cAcc_s, num_acc: cutlass.Constexpr, even: cutlass.Constexpr +): + """Trace-time proof that a thread's fragment is one contiguous column run. + + Returns the (static) first column. The column coordinates fold to Python + ints at trace time; requiring one even-based, contiguous, increasing run + also proves the fragment covers a single row (a two-row fragment would + repeat columns), so this CHECKS the layout the epilogues need instead of + assuming a physical thread-to-row mapping. + """ + cols = [] + for v in cutlass.range_constexpr(num_acc): + cols.append(tTR_cAcc_s[v][1]) + first = cols[0] + if cutlass.const_expr( + not all(isinstance(c, int) for c in cols) + or tuple(cols) != tuple(range(first, first + num_acc)) + or (even and first % 2 != 0) + ): + raise ValueError( + "the epilogue needs one contiguous, increasing" + + (", even-based" if even else "") + + f" column run per thread, but tTR_cAcc gave {cols}" + ) + return first + + +# --------------------------------------------------------------------------- +# Epilogues. Called once per subtile by all 128 epilogue threads with a +# CTA-uniform k_cnt; stores are never predicated (a tail tile's zeroed +# accumulator produces exactly the zero bytes the contract requires). +# --------------------------------------------------------------------------- + + +@cute.jit +def _wgrad_epilogue( + tTR_rAcc, + tTR_cAcc_s, + epi_tidx, + subtile_idx: cutlass.Constexpr, + tile: TileCoords, + sEpi, + out, + R: cutlass.Constexpr, + N: cutlass.Constexpr, +): + """Round the FP32 accumulator subtile to BF16 and store it (no quant). + + ``out`` is ``(mDw,)`` with ``mDw`` the ``(N, K, G)`` view of the + contiguous ``[G, N, K]`` destination. + """ + num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) + frag_col = _epilogue_column_run(tTR_cAcc_s, num_acc, even=False) + + frag = cute.make_rmem_tensor(cute.make_layout(num_acc), cutlass.BFloat16) + for v in cutlass.range_constexpr(num_acc): + frag[v] = cutlass.BFloat16(tTR_rAcc[v]) + + gDw = out[0] + strides = gDw.stride + if cutlass.const_expr(strides[1] != 1): + raise ValueError( + f"the wgrad epilogue stores a contiguous run along K, so dw's K " + f"stride must be 1, got layout {gDw.layout}" + ) + row = tile.row_base + tTR_cAcc_s[0][0] + elem = ( + row * Int32(strides[0]) + + (tile.col_base + Int32(frag_col)) + + tile.expert * Int32(strides[2]) + ) + _store_frag(gDw, elem, frag) + + +@cute.jit +def _swiglu_fwd_epilogue( + tTR_rAcc, + tTR_cAcc_s, + epi_tidx, + subtile_idx: cutlass.Constexpr, + tile: TileCoords, + sEpi, + out, + R: cutlass.Constexpr, + N: cutlass.Constexpr, +): + """SwiGLU + BF16 pre-activation save + dual MXFP8 quantization. + + ``out`` = flat views ``(z [R*2F] bf16, h_row_q [R*F] e4m3, + h_row_sf uint8, h_col_q [F*R] e4m3 in column-major storage order, + h_col_sf uint8)``. ``N == 2F`` is the GEMM (and z) column count. + + Per the kernel contract: the accumulator is rounded to BF16 first (that IS + z), SwiGLU is evaluated once from the rounded values, h is rounded to BF16 + once, and both quantizers consume the same staged BF16 h. + """ + num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) # 64 + half = cutlass.const_expr(num_acc // 2) # 32 h columns per subtile + F = cutlass.const_expr(N // 2) + frag_col = _epilogue_column_run(tTR_cAcc_s, num_acc, even=True) + + mZ = out[0] + mHrowQ = out[1] + mHrowSF = out[2] + mHcolQ = out[3] + mHcolSF = out[4] + + # ---- stage 1: z store + h compute into the staging tile -------------- + lrow = tTR_cAcc_s[0][0] # CTA-tile-local row of this thread's fragment + row_g = tile.row_base + lrow + zfrag = cute.make_rmem_tensor(cute.make_layout(num_acc), cutlass.BFloat16) + for v in cutlass.range_constexpr(num_acc): + zfrag[v] = cutlass.BFloat16(tTR_rAcc[v]) + _store_frag(mZ, row_g * Int32(N) + tile.col_base + Int32(frag_col), zfrag) + + for j in cutlass.range_constexpr(half): + gate = Float32(zfrag[2 * j]) + up = Float32(zfrag[2 * j + 1]) + sig = _sigmoid_f32(gate) + sEpi[lrow, Int32(j)] = cutlass.BFloat16((gate * sig) * up) + + cute.arch.barrier( + barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS + ) + + # ---- stage 2: dual quantization off the staging tile ----------------- + # Global h column base of this subtile; frag_col is static (64 * + # subtile_idx), so hbase is divisible by 32. + hbase = (tile.col_base + Int32(frag_col)) >> 1 + ncb_row = cutlass.const_expr((F // _SF_VEC_SIZE + 3) // 4) + ncb_col = cutlass.const_expr((R // _SF_VEC_SIZE + 3) // 4) + + # Rowwise 1x32: one block per thread (row = epi_tidx of the staging tile). + q_row = tile.row_base + epi_tidx + _quant_block_from_smem( + sEpi, + epi_tidx, + Int32(0), + mHrowQ, + q_row * Int32(F) + hbase, + mHrowSF, + _blocked_scale_idx( + q_row, (tile.col_base + Int32(frag_col)) >> 6, Int32(ncb_row) + ), + COLWISE=False, + ) + + # Columnwise 32x1: 32 columns x 4 row-blocks = one block per thread. + col_l = epi_tidx & Int32(31) + blk = epi_tidx >> 5 + col_g = hbase + col_l + _quant_block_from_smem( + sEpi, + blk * Int32(32), + col_l, + mHcolQ, + col_g * Int32(R) + tile.row_base + blk * Int32(32), + mHcolSF, + _blocked_scale_idx(col_g, (tile.row_base >> 5) + blk, Int32(ncb_col)), + COLWISE=True, + ) + + # The next subtile reuses the staging tile. + cute.arch.barrier( + barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS + ) + + +@cute.jit +def _dswiglu_bwd_epilogue( + tTR_rAcc, + tTR_cAcc_s, + epi_tidx, + subtile_idx: cutlass.Constexpr, + tile: TileCoords, + sEpi, + out, + R: cutlass.Constexpr, + N: cutlass.Constexpr, +): + """dSwiGLU from the saved z + dual MXFP8 quantization of dz. + + ``out`` = ``(z [R*2F] bf16 INPUT, dz_row_q [R*2F] e4m3, dz_row_sf uint8, + dz_col_q [2F*R] e4m3 column-major storage, dz_col_sf uint8)``. ``N == F`` + is the dgrad GEMM column count; dz has 2F element-interleaved columns. + + Per the kernel contract: dh is rounded to BF16 first; gate/up come from + the saved BF16 z; dgate/dup are each rounded to BF16 before interleaving; + both quantizers consume the same staged BF16 dz. The z load is the one + epilogue-side gmem input in the family and is predicated on + ``k_cnt == 0`` -- tail rows of z are read-forbidden and contribute zeros. + """ + num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) # 32 + two_f = cutlass.const_expr(2 * N) + frag_col = _epilogue_column_run(tTR_cAcc_s, num_acc, even=False) + + mZ = out[0] + mDzRowQ = out[1] + mDzRowSF = out[2] + mDzColQ = out[3] + mDzColSF = out[4] + + # ---- stage 1: z load (predicated) + dSwiGLU into the staging tile ---- + lrow = tTR_cAcc_s[0][0] + row_g = tile.row_base + lrow + # dz columns covered by this subtile: [dzbase, dzbase + 64). + dzbase = (tile.col_base + Int32(frag_col)) * Int32(2) + + zfrag = cute.make_rmem_tensor(cute.make_layout(2 * num_acc), cutlass.BFloat16) + if tile.k_cnt > Int32(0): + cute.autovec_copy( + cute.make_tensor( + (mZ.iterator + (row_g * Int32(two_f) + dzbase)).align(16), + cute.make_layout(2 * num_acc), + ), + zfrag, + ) + else: + for i in cutlass.range_constexpr(2 * num_acc): + zfrag[i] = cutlass.BFloat16(0.0) + + for j in cutlass.range_constexpr(num_acc): + gate = Float32(zfrag[2 * j]) + up = Float32(zfrag[2 * j + 1]) + dh = Float32(cutlass.BFloat16(tTR_rAcc[j])) + sig = _sigmoid_f32(gate) + silu = gate * sig + dsilu = sig * (Float32(1.0) + gate * (Float32(1.0) - sig)) + sEpi[lrow, Int32(2 * j)] = cutlass.BFloat16((dh * up) * dsilu) + sEpi[lrow, Int32(2 * j + 1)] = cutlass.BFloat16(dh * silu) + + cute.arch.barrier( + barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS + ) + + # ---- stage 2: dual quantization off the staging tile ----------------- + ncb_row = cutlass.const_expr((two_f // _SF_VEC_SIZE + 3) // 4) + ncb_col = cutlass.const_expr((R // _SF_VEC_SIZE + 3) // 4) + + # Rowwise 1x32: 128 rows x 2 blocks = two tasks per thread. + q_row = tile.row_base + epi_tidx + for blk in cutlass.range_constexpr(2): + _quant_block_from_smem( + sEpi, + epi_tidx, + Int32(blk * _SF_VEC_SIZE), + mDzRowQ, + q_row * Int32(two_f) + dzbase + Int32(blk * _SF_VEC_SIZE), + mDzRowSF, + _blocked_scale_idx(q_row, (dzbase >> 5) + Int32(blk), Int32(ncb_row)), + COLWISE=False, + ) + + # Columnwise 32x1: 64 columns x 4 row-blocks = two tasks per thread. + for k in cutlass.range_constexpr(2): + task = epi_tidx + Int32(128 * k) + col_l = task & Int32(63) + blk = task >> 6 + col_g = dzbase + col_l + _quant_block_from_smem( + sEpi, + blk * Int32(32), + col_l, + mDzColQ, + col_g * Int32(R) + tile.row_base + blk * Int32(32), + mDzColSF, + _blocked_scale_idx(col_g, (tile.row_base >> 5) + blk, Int32(ncb_col)), + COLWISE=True, + ) + + cute.arch.barrier( + barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS + ) + + +# --------------------------------------------------------------------------- +# The shared kernel and launch builder. +# --------------------------------------------------------------------------- + + +@cute.kernel +def _grouped_gemm_kernel( + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB: cute.Tensor, + offs: cute.Tensor, + out, + a_smem_layout: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + sfa_smem_layout: cute.Layout, + sfb_smem_layout: cute.Layout, + cfg: cutlass.Constexpr, + storage_type: cutlass.Constexpr, + EPILOGUE: cutlass.Constexpr, +): + """One CTA computes one 128 x 128 output tile. Warps: 0 TMA, 1 MMA, + 4-7 epilogue, 2-3 idle.""" + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, bidz = cute.arch.block_idx() + + num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) + num_k_tiles_full = cutlass.const_expr(cute.size(mA, mode=[1]) // _CTA_K) + gemm_m = cutlass.const_expr(cute.size(mA, mode=[0])) + gemm_n = cutlass.const_expr(cute.size(mB, mode=[0])) + + # ------------------------------------------------------------------ + # Tile coordinates. No scheduler: the grid IS the tile enumeration. + # ------------------------------------------------------------------ + tile_m = Int32(bidx) + tile_n = Int32(bidy) + if cutlass.const_expr(not cfg.ragged_k): + row_base = tile_m * _CTA_M + # Unrolled G-way scan: the owning expert is the number of groups that + # end at or before this tile's row base; a zero-token expert is never + # selected since its end equals its start. + expert = Int32(0) + for g in cutlass.range_constexpr(num_groups - 1): + expert += Int32(offs[g] <= row_base) + is_active = Int32(offs[num_groups - 1] > row_base) + k_base = Int32(0) + k_cnt = is_active * Int32(num_k_tiles_full) + l_a = Int32(0) + l_b = expert + else: + expert = Int32(bidz) + row_base = tile_m * _CTA_M + prev = offs[cutlass.max(expert - Int32(1), Int32(0))] * Int32(expert > Int32(0)) + # Exact, not a ceil: every group boundary is a multiple of cta_tile_k. + k_base = prev // _CTA_K + # Clamp: nonmonotone offsets (undefined behavior per the contract, and + # only device-checkable) would otherwise make k_cnt negative -- the + # loops still run zero trips, but the epilogue would read TMEM the MMA + # never wrote. Clamped, a malformed expert degrades to an all-zero + # slice instead. + k_cnt = cutlass.max((offs[expert] - prev) // _CTA_K, Int32(0)) + l_a = Int32(0) + l_b = Int32(0) + tile = TileCoords( + tile_m=tile_m, + tile_n=tile_n, + expert=expert, + row_base=row_base, + col_base=tile_n * _CTA_N, + k_cnt=k_cnt, + ) + + # ------------------------------------------------------------------ + # Shared memory and pipelines + # ------------------------------------------------------------------ + smem = utils.SmemAllocator() + storage = smem.allocate(storage_type) + + sA = storage.sA.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) + sB = storage.sB.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout) + sEpi = None + if cutlass.const_expr(cfg.epi_smem_cols > 0): + sEpi = storage.sEpi.get_tensor( + cute.make_layout((_CTA_M, cfg.epi_smem_cols), stride=(cfg.epi_smem_cols, 1)) + ) + + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar.data_ptr(), + num_stages=_NUM_AB_STAGE, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + # The EXACT byte count of the four TMA copies of one stage: too small + # and the MMA consumes a partially arrived stage. + tx_count=_AB_STAGE_BYTES, + cta_layout_vmnk=cluster_layout_vmnk, + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar.data_ptr(), + num_stages=_NUM_ACC_STAGE, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, _NUM_EPI_THREADS + ), + cta_layout_vmnk=cluster_layout_vmnk, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=_TMEM_ALLOC_BARRIER_ID, + num_threads=32 * 5, # the MMA warp joins the four epilogue warps + ) + epilogue_barrier = pipeline.NamedBarrier( + barrier_id=_EPI_FINAL_BARRIER_ID, + num_threads=_NUM_EPI_THREADS, + ) + tmem_alloc = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=_FIRST_EPI_WARP, + is_two_cta=False, + ) + + # ------------------------------------------------------------------ + # Tile the global tensors. One static descriptor per operand; the expert + # is an L coordinate (ragged M) or a K-tile index base (ragged K). + # ------------------------------------------------------------------ + gA = cute.local_tile( + mA, cute.slice_(_MMA_TILER, (None, 0, None)), (None, None, None) + ) + gB = cute.local_tile( + mB, cute.slice_(_MMA_TILER, (0, None, None)), (None, None, None) + ) + gSFA = cute.local_tile( + mSFA, cute.slice_(_MMA_TILER, (None, 0, None)), (None, None, None) + ) + gSFB = cute.local_tile( + mSFB, cute.slice_(_MMA_TILER, (0, None, None)), (None, None, None) + ) + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA) + tCgB = thr_mma.partition_B(gB) + tCgSFA = thr_mma.partition_A(gSFA) + tCgSFB = thr_mma.partition_B(gSFB) + + trivial_cta_layout = cute.make_layout(1) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + trivial_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + trivial_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + 0, + trivial_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + 0, + trivial_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + tAgA_slice = tAgA[(None, tile_m, None, l_a)] + tBgB_slice = tBgB[(None, tile_n, None, l_b)] + tAgSFA_slice = tAgSFA[(None, tile_m, None, l_a)] + tBgSFB_slice = tBgSFB[(None, tile_n, None, l_b)] + + acc_shape = tiled_mma.partition_shape_C(_MMA_TILER[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, _NUM_ACC_STAGE)) + + # ------------------------------------------------------------------ + # Warp 0: TMA producer + # ------------------------------------------------------------------ + if warp_idx == _TMA_WARP_ID: + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, _NUM_AB_STAGE + ) + for _ in cutlass.range(0, k_cnt, 1, unroll=1): + ab_pipeline.producer_acquire(ab_producer_state) + k_idx = k_base + ab_producer_state.count + bar = ab_pipeline.producer_get_barrier(ab_producer_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_idx)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, k_idx)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, k_idx)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, k_idx)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=bar, + ) + ab_producer_state.advance() + ab_pipeline.producer_tail(ab_producer_state) + + # ------------------------------------------------------------------ + # Warp 1: MMA + # ------------------------------------------------------------------ + if warp_idx == _MMA_WARP_ID: + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + # The MMA warp joins the allocation barrier but must never allocate. + tmem_alloc.wait_for_alloc() + acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + dtype=sSFA.element_type, + ) + tCtSFA = cute.make_tensor( + sfa_tmem_ptr, + blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + _MMA_TILER, + _SF_VEC_SIZE, + cute.slice_(sfa_smem_layout, (None, None, None, 0)), + ), + ) + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=sSFB.element_type, + ) + tCtSFB = cute.make_tensor( + sfb_tmem_ptr, + blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + _MMA_TILER, + _SF_VEC_SIZE, + cute.slice_(sfb_smem_layout, (None, None, None, 0)), + ), + ) + s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_copy_and_partition(sSFA, tCtSFA) + s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, _NUM_AB_STAGE + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, _NUM_ACC_STAGE + ) + tCtAcc = tCtAcc_base[(None, None, None, 0)] + + # Acquire and commit unconditionally, including when k_cnt == 0, so + # the accumulator handoff barrier stays balanced on tail tiles. + acc_pipeline.producer_acquire(acc_producer_state) + for k_tile in cutlass.range(0, k_cnt, 1, unroll=1): + ab_pipeline.consumer_wait(ab_consumer_state) + stage_crd = (None, None, None, None, ab_consumer_state.index) + cute.copy(s2t_sfa, tCsSFA_s2t[stage_crd], tCtSFA_s2t) + cute.copy(s2t_sfb, tCsSFB_s2t[stage_crd], tCtSFB_s2t) + # ACCUMULATE=False on the first K tile is what zeroes the + # accumulator; there is no separate TMEM clear. + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + mma_crd = (None, None, None, ab_consumer_state.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[mma_crd], tCtSFA], + [tCrB[mma_crd], tCtSFB], + tCtAcc, + ) + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state.advance() + acc_pipeline.producer_commit(acc_producer_state) + + # ------------------------------------------------------------------ + # Warps 4-7: epilogue + # ------------------------------------------------------------------ + if warp_idx >= _FIRST_EPI_WARP: + tmem_alloc.allocate(_TMEM_TOTAL_COLS) + tmem_alloc.wait_for_alloc() + acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + epi_tidx = tidx - _FIRST_EPI_THREAD + tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc = _t2r_partition( + epi_tidx, tCtAcc_base, cfg + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, _NUM_ACC_STAGE + ) + acc_pipeline.consumer_wait(acc_consumer_state) + + for s in cutlass.range_constexpr(cfg.num_epi_subtiles): + if k_cnt == Int32(0): + # Tail tile or zero-token expert: nothing was accumulated, so + # the fragment is zeroed here and the epilogue runs unchanged. + for v in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[v] = Float32(0.0) + else: + cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, 0, s)], tTR_rAcc) + EPILOGUE( + tTR_rAcc, + tTR_cAcc[(None, None, None, 0, s)], + epi_tidx, + s, + tile, + sEpi, + out, + gemm_m, + gemm_n, + ) + + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + tmem_alloc.relinquish_alloc_permit() + epilogue_barrier.arrive_and_wait() + tmem_alloc.free(acc_tmem_ptr) + + +@cute.jit +def _launch_grouped_gemm( + mA: cute.Tensor, + mB: cute.Tensor, + sfa_flat: cute.Tensor, + sfb_flat: cute.Tensor, + offs: cute.Tensor, + out, + stream, + cfg: cutlass.Constexpr, + EPILOGUE: cutlass.Constexpr, +): + """Build the four static TMA descriptors and launch. Grid is data-independent. + + This is a trace body: calling it directly retraces the kernel on every + invocation. The public launchers below wrap it in per-kernel ``@cute.jit`` + entry points taking only dynamic tensors, ``cute.compile`` those once per + shape key, and call the compiled executor. + """ + a_dtype = mA.element_type + b_dtype = mB.element_type + sf_dtype = cutlass.Float8E8M0FNU + + gemm_m = cutlass.const_expr(cute.size(mA, mode=[0])) + gemm_k = cutlass.const_expr(cute.size(mA, mode=[1])) + gemm_n = cutlass.const_expr(cute.size(mB, mode=[0])) + l_a = cutlass.const_expr(cute.size(mA, mode=[2])) + l_b = cutlass.const_expr(cute.size(mB, mode=[2])) + num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) + + if cutlass.const_expr( + gemm_m % _CTA_M != 0 or gemm_n % _CTA_N != 0 or gemm_k % _CTA_K != 0 + ): + raise ValueError( + f"GEMM extents ({gemm_m}, {gemm_n}, {gemm_k}) must be multiples of " + f"({_CTA_M}, {_CTA_N}, {_CTA_K})" + ) + + mSFA = _make_sf_gemm_tensor(sfa_flat, gemm_m, gemm_k, l_a) + mSFB = _make_sf_gemm_tensor(sfb_flat, gemm_n, gemm_k, l_b) + + tiled_mma = _make_tiled_mma(a_dtype, b_dtype, sf_dtype) + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) + ) + + a_smem_layout = sm100_utils.make_smem_layout_a( + tiled_mma, _MMA_TILER, a_dtype, _NUM_AB_STAGE + ) + b_smem_layout = sm100_utils.make_smem_layout_b( + tiled_mma, _MMA_TILER, b_dtype, _NUM_AB_STAGE + ) + sfa_smem_layout = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, _MMA_TILER, _SF_VEC_SIZE, _NUM_AB_STAGE + ) + sfb_smem_layout = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, _MMA_TILER, _SF_VEC_SIZE, _NUM_AB_STAGE + ) + + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id), + mA, + cute.slice_(a_smem_layout, (None, None, None, 0)), + _MMA_TILER, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + sm100_utils.cluster_shape_to_tma_atom_B((1, 1), tiled_mma.thr_id), + mB, + cute.slice_(b_smem_layout, (None, None, None, 0)), + _MMA_TILER, + tiled_mma, + cluster_layout_vmnk.shape, + ) + # The 512-byte scale atom is contiguous; TMA moves it as 8-byte elements. + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id), + mSFA, + cute.slice_(sfa_smem_layout, (None, None, None, 0)), + _MMA_TILER, + tiled_mma, + cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sm100_utils.cluster_shape_to_tma_atom_SFB((1, 1), tiled_mma.thr_id), + mSFB, + cute.slice_(sfb_smem_layout, (None, None, None, 0)), + _MMA_TILER, + tiled_mma, + cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + @cute.struct + class SharedStorage: + ab_full_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_AB_STAGE] + ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_AB_STAGE] + acc_full_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_ACC_STAGE] + acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_ACC_STAGE] + tmem_holding_buf: cutlass.Int32 + sEpi: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, cfg.epi_smem_elems], 128 + ] + sA: cute.struct.Align[ + cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024 + ] + sB: cute.struct.Align[ + cute.struct.MemRange[b_dtype, cute.cosize(b_smem_layout.outer)], 1024 + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[sf_dtype, cute.cosize(sfa_smem_layout)], 1024 + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[sf_dtype, cute.cosize(sfb_smem_layout)], 1024 + ] + + smem_bytes = cutlass.const_expr(SharedStorage.size_in_bytes()) + if cutlass.const_expr(smem_bytes > _SMEM_CAPACITY_BYTES): + raise ValueError( + f"shared memory request {smem_bytes} B exceeds the sm_100 capacity " + f"{_SMEM_CAPACITY_BYTES} B" + ) + + if cutlass.const_expr(not cfg.ragged_k): + grid = (gemm_m // _CTA_M, gemm_n // _CTA_N, 1) + else: + grid = (gemm_m // _CTA_M, gemm_n // _CTA_N, num_groups) + + _grouped_gemm_kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + offs, + out, + a_smem_layout, + b_smem_layout, + sfa_smem_layout, + sfb_smem_layout, + cfg, + SharedStorage, + EPILOGUE, + ).launch( + grid=grid, + block=(_THREADS, 1, 1), + cluster=(1, 1, 1), + smem=smem_bytes, + stream=stream, + ) + + +# --------------------------------------------------------------------------- +# Compile-once entry points (dynamic tensors only; Constexprs closed over). +# --------------------------------------------------------------------------- + + +@cute.jit +def _swiglu_fwd_entry(mA, mB, sfa, sfb, offs, mZ, mHrq, mHrs, mHcq, mHcs, stream): + _launch_grouped_gemm( + mA, + mB, + sfa, + sfb, + offs, + (mZ, mHrq, mHrs, mHcq, mHcs), + stream, + _SWIGLU_FWD_CONFIG, + _swiglu_fwd_epilogue, + ) + + +@cute.jit +def _dswiglu_bwd_entry(mA, mB, sfa, sfb, offs, mZ, mDrq, mDrs, mDcq, mDcs, stream): + _launch_grouped_gemm( + mA, + mB, + sfa, + sfb, + offs, + (mZ, mDrq, mDrs, mDcq, mDcs), + stream, + _DSWIGLU_BWD_CONFIG, + _dswiglu_bwd_epilogue, + ) + + +@cute.jit +def _wgrad_entry(mA, mB, sfa, sfb, offs, mDw, stream): + _launch_grouped_gemm( + mA, + mB, + sfa, + sfb, + offs, + (mDw,), + stream, + _WGRAD_CONFIG, + _wgrad_epilogue, + ) + + +@functools.cache +def _executor_slot(key: tuple) -> list: + """One memo slot per (kernel, shape, device, dtype, DSL version) key. + + The executor cannot be built from the key alone -- the shared trace needs + static shapes (the blocked scale layout and the grid are built from them) + -- so the first real call's tensors are what gets compiled and the caller + fills the slot once. + """ + return [] + + +def _cache_key(kind: str, dims: tuple, tensors: tuple, device) -> tuple: + return ( + kind, + dims, + tuple(str(t.dtype) for t in tensors), + device.index, + torch.cuda.get_device_capability(device), + cutlass.__version__, + ) + + +def _common_launch_checks(name: str, device, tensors, groups: int): + """Support and safety gates shared by the three launchers.""" + if any(_is_fake(t) for t in tensors): + raise ValueError( + f"{name} cannot run on fake/meta tensors; call the corresponding " + "torchao::* op instead, whose register_fake handles tracing" + ) + if groups < 1: + raise ValueError(f"G must be at least 1, got {groups}") + if device.type != "cuda": + raise ValueError(f"{name} requires CUDA tensors, got device {device}") + major, _minor = torch.cuda.get_device_capability(device) + if major != 10: + raise NotImplementedError( + f"{name} requires an SM100-class GPU (compute capability 10.x), " + f"got {torch.cuda.get_device_capability(device)}" + ) + + +def _stream_for(device): + import cuda.bindings.driver as cuda + + return cuda.CUstream(int(torch.cuda.current_stream(device).cuda_stream)) + + +def _check_sf_pointer_alignment(name: str, buf: torch.Tensor): + # Scale buffers feed TMA descriptors (inputs) or vectorized stores; a + # contiguous view with a storage offset can be 2-byte aligned, and only + # the launcher promises the alignment. + if buf.data_ptr() % 32 != 0: + raise ValueError( + f"{name} must be 32-byte aligned, but its data pointer is " + f"{buf.data_ptr() % 32} bytes past an aligned address" + ) + + +_E4M3 = torch.float8_e4m3fn +_BF16 = torch.bfloat16 + + +def launch_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_t_q, w13_t_sf, offsets, z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf +): + """FC1 grouped GEMM + SwiGLU + dual MXFP8 quantization, one kernel launch. + + Inputs are prequantized: ``x_q`` E4M3 ``[R, D]`` row-major with blocked + ``x_sf``; ``w13_t_q`` E4M3 ``[G, D, 2F]`` stride ``(2F*D, 1, D)`` with + per-expert blocked ``w13_t_sf`` (the 2F axis is element-interleaved + gate/up). Destinations are caller-allocated: ``z_bf16 [R, F, 2]``, + ``h_row_q [R, F]`` row-major + ``h_row_sf``, ``h_col_q [R, F]`` + column-major + ``h_col_sf`` (whole-matrix blocked for logical + ``[F, R/32]``). Every destination byte is written, including the + inactive-tail zeros. + """ + if x_q.ndim != 2: + raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") + rows, model_dim = x_q.shape + if w13_t_q.ndim != 3 or w13_t_q.shape[1] != model_dim or w13_t_q.shape[2] % 2: + raise ValueError( + f"w13_t_q must be [G, D, 2F] with D == {model_dim} and even 2F, " + f"got shape {tuple(w13_t_q.shape)}" + ) + groups, _, two_hidden = w13_t_q.shape + hidden = two_hidden // 2 + device = x_q.device + tensors = ( + x_q, + x_sf, + w13_t_q, + w13_t_sf, + offsets, + z_bf16, + h_row_q, + h_row_sf, + h_col_q, + h_col_sf, + ) + _common_launch_checks("launch_grouped_gemm_swiglu_fwd", device, tensors, groups) + if rows == 0: + raise ValueError( + "R == 0 is handled by the op layer (empty destinations, no launch)" + ) + + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + validate_grouped_operand( + x_q, + name="x_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + w13_t_q, + name="w13_t_q", + shape=(groups, model_dim, two_hidden), + stride=(model_dim * two_hidden, 1, model_dim), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + x_sf, + name="x_sf", + logical_rows=rows, + logical_cols=model_dim // _SF_VEC_SIZE, + device=device, + ) + validate_blocked_scales( + w13_t_sf, + name="w13_t_sf", + logical_rows=two_hidden, + logical_cols=model_dim // _SF_VEC_SIZE, + device=device, + groups=groups, + ) + _check_sf_pointer_alignment("x_sf", x_sf) + _check_sf_pointer_alignment("w13_t_sf", w13_t_sf) + validate_destination( + z_bf16, + name="z_bf16", + shape=(rows, hidden, 2), + stride=(two_hidden, 2, 1), + dtype=_BF16, + device=device, + ) + validate_destination( + h_row_q, + name="h_row_q", + shape=(rows, hidden), + stride=(hidden, 1), + dtype=_E4M3, + device=device, + ) + validate_destination( + h_col_q, + name="h_col_q", + shape=(rows, hidden), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + h_row_sf, + name="h_row_sf", + logical_rows=rows, + logical_cols=hidden // _SF_VEC_SIZE, + device=device, + ) + validate_blocked_scales( + h_col_sf, + name="h_col_sf", + logical_rows=hidden, + logical_cols=rows // _SF_VEC_SIZE, + device=device, + ) + _check_sf_pointer_alignment("h_row_sf", h_row_sf) + _check_sf_pointer_alignment("h_col_sf", h_col_sf) + # The epilogue computes flat element offsets in Int32. + if rows * two_hidden >= 2**31: + raise ValueError( + f"R * 2F = {rows * two_hidden} does not fit the epilogue's int32 " + "element indexing" + ) + if two_hidden // _CTA_N > 65535: + raise ValueError(f"2F = {two_hidden} exceeds the launch grid's Y limit") + + stream = _stream_for(device) + args = ( + from_dlpack(activation_gemm_view(x_q), assumed_align=16), + from_dlpack(weight_gemm_view(w13_t_q), assumed_align=16), + from_dlpack(x_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(w13_t_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(offsets, assumed_align=4), + from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), + from_dlpack(h_row_q.view(-1), assumed_align=16), + from_dlpack(h_row_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(h_col_q.t().reshape(-1), assumed_align=16), + from_dlpack(h_col_sf.view(torch.uint8).view(-1), assumed_align=16), + stream, + ) + key = _cache_key( + "swiglu_fwd", (rows, model_dim, hidden, groups), tensors[:5], device + ) + slot = _executor_slot(key) + if not slot: + slot.append(cute.compile(_swiglu_fwd_entry, *args)) + slot[0](*args) + + +def launch_grouped_gemm_dswiglu_bwd( + do_q, + do_sf, + w2_dgrad_q, + w2_dgrad_sf, + z_bf16, + offsets, + dz_row_q, + dz_row_sf, + dz_col_q, + dz_col_sf, +): + """FC2 dgrad grouped GEMM + dSwiGLU + dual MXFP8 quantization, one launch. + + ``do_q`` E4M3 ``[R, D]`` row-major + blocked ``do_sf``; ``w2_dgrad_q`` + E4M3 ``[G, D, F]`` stride ``(D*F, 1, D)`` + per-expert blocked + ``w2_dgrad_sf``; ``z_bf16`` is the exact ``[R, F, 2]`` tensor the forward + kernel wrote (tail rows are never read). Destinations: ``dz_row_q + [R, 2F]`` row-major + ``dz_row_sf``, ``dz_col_q [R, 2F]`` column-major + + ``dz_col_sf`` (whole-matrix blocked for logical ``[2F, R/32]``), gate/up + gradients element-interleaved. + """ + if do_q.ndim != 2: + raise ValueError(f"do_q must be 2D [R, D], got shape {tuple(do_q.shape)}") + rows, model_dim = do_q.shape + if w2_dgrad_q.ndim != 3 or w2_dgrad_q.shape[1] != model_dim: + raise ValueError( + f"w2_dgrad_q must be [G, D, F] with D == {model_dim}, got shape " + f"{tuple(w2_dgrad_q.shape)}" + ) + groups, _, hidden = w2_dgrad_q.shape + two_hidden = 2 * hidden + device = do_q.device + tensors = ( + do_q, + do_sf, + w2_dgrad_q, + w2_dgrad_sf, + z_bf16, + offsets, + dz_row_q, + dz_row_sf, + dz_col_q, + dz_col_sf, + ) + _common_launch_checks("launch_grouped_gemm_dswiglu_bwd", device, tensors, groups) + if rows == 0: + raise ValueError( + "R == 0 is handled by the op layer (empty destinations, no launch)" + ) + + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + validate_grouped_operand( + do_q, + name="do_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + w2_dgrad_q, + name="w2_dgrad_q", + shape=(groups, model_dim, hidden), + stride=(model_dim * hidden, 1, model_dim), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + z_bf16, + name="z_bf16", + shape=(rows, hidden, 2), + stride=(two_hidden, 2, 1), + dtype=_BF16, + device=device, + ) + validate_blocked_scales( + do_sf, + name="do_sf", + logical_rows=rows, + logical_cols=model_dim // _SF_VEC_SIZE, + device=device, + ) + validate_blocked_scales( + w2_dgrad_sf, + name="w2_dgrad_sf", + logical_rows=hidden, + logical_cols=model_dim // _SF_VEC_SIZE, + device=device, + groups=groups, + ) + _check_sf_pointer_alignment("do_sf", do_sf) + _check_sf_pointer_alignment("w2_dgrad_sf", w2_dgrad_sf) + validate_destination( + dz_row_q, + name="dz_row_q", + shape=(rows, two_hidden), + stride=(two_hidden, 1), + dtype=_E4M3, + device=device, + ) + validate_destination( + dz_col_q, + name="dz_col_q", + shape=(rows, two_hidden), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + dz_row_sf, + name="dz_row_sf", + logical_rows=rows, + logical_cols=two_hidden // _SF_VEC_SIZE, + device=device, + ) + validate_blocked_scales( + dz_col_sf, + name="dz_col_sf", + logical_rows=two_hidden, + logical_cols=rows // _SF_VEC_SIZE, + device=device, + ) + _check_sf_pointer_alignment("dz_row_sf", dz_row_sf) + _check_sf_pointer_alignment("dz_col_sf", dz_col_sf) + if rows * two_hidden >= 2**31: + raise ValueError( + f"R * 2F = {rows * two_hidden} does not fit the epilogue's int32 " + "element indexing" + ) + if hidden // _CTA_N > 65535: + raise ValueError(f"F = {hidden} exceeds the launch grid's Y limit") + + stream = _stream_for(device) + args = ( + from_dlpack(activation_gemm_view(do_q), assumed_align=16), + from_dlpack(weight_gemm_view(w2_dgrad_q), assumed_align=16), + from_dlpack(do_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(w2_dgrad_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(offsets, assumed_align=4), + from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), + from_dlpack(dz_row_q.view(-1), assumed_align=16), + from_dlpack(dz_row_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(dz_col_q.t().reshape(-1), assumed_align=16), + from_dlpack(dz_col_sf.view(torch.uint8).view(-1), assumed_align=16), + stream, + ) + key = _cache_key( + "dswiglu_bwd", (rows, model_dim, hidden, groups), tensors[:6], device + ) + slot = _executor_slot(key) + if not slot: + slot.append(cute.compile(_dswiglu_bwd_entry, *args)) + slot[0](*args) + + +def launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw): + """Grouped MXFP8 wgrad into a caller-allocated BF16 ``[G, N, K]``, one launch. + + Inputs are the columnwise-quantized outputs of the forward/backward + kernels: ``dy_col_q`` E4M3 logical ``[R, N]`` stride ``(1, R)`` with + ``dy_col_sf`` whole-matrix blocked for logical ``[N, R/32]``, and + ``x_col_q`` / ``x_col_sf`` likewise for ``[R, K]``. Every element of + ``dw`` is written, including the all-zero slice of a zero-token expert. + """ + if dy_col_q.ndim != 2 or x_col_q.ndim != 2: + raise ValueError( + "dy_col_q and x_col_q must be 2D logical [R, N] and [R, K], got " + f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" + ) + rows, out_features = dy_col_q.shape + x_rows, in_features = x_col_q.shape + if x_rows != rows: + raise ValueError( + f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" + ) + groups = offsets.numel() + device = dy_col_q.device + tensors = (dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw) + _common_launch_checks("launch_grouped_gemm_wgrad", device, tensors, groups) + + validate_allocated_rows(rows) + for name, value in (("dy_col_q's N", out_features), ("x_col_q's K", in_features)): + if value <= 0 or value % 128 != 0: + raise ValueError(f"{name} must be a positive multiple of 128, got {value}") + validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + validate_grouped_operand( + dy_col_q, + name="dy_col_q", + shape=(rows, out_features), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_grouped_operand( + x_col_q, + name="x_col_q", + shape=(rows, in_features), + stride=(1, rows), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + dy_col_sf, + name="dy_col_sf", + logical_rows=out_features, + logical_cols=rows // _SF_VEC_SIZE, + device=device, + ) + validate_blocked_scales( + x_col_sf, + name="x_col_sf", + logical_rows=in_features, + logical_cols=rows // _SF_VEC_SIZE, + device=device, + ) + _check_sf_pointer_alignment("dy_col_sf", dy_col_sf) + _check_sf_pointer_alignment("x_col_sf", x_col_sf) + validate_destination( + dw, + name="dw_bf16", + shape=(groups, out_features, in_features), + stride=(out_features * in_features, in_features, 1), + dtype=_BF16, + device=device, + ) + if groups * out_features * in_features >= 2**31: + raise ValueError( + f"dw_bf16 has {groups * out_features * in_features} elements, which " + "does not fit the epilogue's int32 element index" + ) + if in_features // _CTA_N > 65535: + raise ValueError(f"K = {in_features} exceeds the launch grid's Y limit") + if groups > 65535: + raise ValueError(f"G = {groups} exceeds the launch grid's Z limit") + + if rows == 0: + # Every expert has zero rows: every slice is the zero matrix. The + # destination is NOT empty here, and the contraction is. + dw.zero_() + return + + stream = _stream_for(device) + args = ( + # The free transpose: logical [R, N] stride (1, R) IS a K-contiguous + # [N, R], so the ragged axis becomes the contraction. + from_dlpack(activation_gemm_view(dy_col_q.t()), assumed_align=16), + from_dlpack(activation_gemm_view(x_col_q.t()), assumed_align=16), + from_dlpack(dy_col_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(x_col_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(offsets, assumed_align=4), + from_dlpack(dw.permute(1, 2, 0), assumed_align=16), + stream, + ) + key = _cache_key( + "wgrad", (rows, out_features, in_features, groups), tensors[:5], device + ) + slot = _executor_slot(key) + if not slot: + slot.append(cute.compile(_wgrad_entry, *args)) + slot[0](*args) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py b/torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py deleted file mode 100644 index f750ec6b7e..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/epilogue_quant.py +++ /dev/null @@ -1,414 +0,0 @@ -# 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. - -"""Rowwise 1x32 and columnwise 32x1 quantizing epilogue for the MXFP8 grouped-MLP kernels. - -Direction-agnostic and GEMM-agnostic: everything here consumes a run of packed -``bf16x2`` words held **one output row per thread**, which is exactly what the -tcgen05 TMEM->register copy delivers (``Ld32x32bOp`` has ``ThrID 32:1`` and a -destination TV layout ``(32, EPI_N_ACC):(EPI_N_ACC, 1)``, so thread ``t`` owns -row ``t`` of the epilogue subtile and all of its columns, contiguously). Two -consequences are load-bearing and are asserted by the geometry below: - -* the rowwise 1x32 amax runs along N inside a single thread -- no cross-lane op; -* the columnwise 32x1 amax runs along M inside a single warp, and a 32-row MX - block is never split across warps or CTAs, because the CTA tile is 128 rows, - ``32 | 128``, and every expert boundary is 128-aligned. - -WORD CONTRACT (every entry point below assumes it): word ``j`` of ``words`` -holds output column ``col_base + 2*j`` in its **low** bf16 half and column -``col_base + 2*j + 1`` in its **high** half, and 16 words are exactly one -32-value scale block. Callers must have rounded to BF16 already -- both because -the kernel contract defines correctness at that boundary and because -:func:`float_to_e8m0` is exact only for a BF16-valued amax. - -The columnwise destination has stride ``(1, R)``, i.e. it is physically a -row-major ``[cols, R]`` buffer, so the epilogue must transpose. It does that -through ``sPad``, a shared-memory staging tile of **BF16** values (not quantized -bytes): the reader's 16-byte load has to come out as four consecutive rows of -one column, which ``mul_cvt_2x`` + ``prmt_even``/``prmt_odd`` produce from two -words x two columns; staging bytes instead would force a stride-2 byte gather. - -Both columnwise scale orientations use whole-matrix ``to_blocked`` coordinates -(feature index as the blocked row, row-block index as the blocked column). That -is this family's choice and it deliberately differs from torchao's per-group -``triton_mx_block_rearrange_2d_K_groups``. -""" - -import cutlass -import cutlass.cute as cute -from cutlass import Int32 -from cutlass.utils import SmemAllocator - -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_epilogue import ( - SCALE_BLOCK, - abs_max_nan_bf16x2, - blocked_scale_idx, - e8m0_reciprocal_bf16, - float_to_e8m0, - fold_amax, - max_nan_bf16x2, - mul_cvt_2x, - prmt_even, - prmt_odd, -) - -__all__ = [ - "WORDS_PER_SCALE_BLOCK", - "COLWISE_PAIRS", - "COLWISE_ROWS_PER_WARP", - "COLWISE_PAIR_PITCH", - "COLWISE_WORDS_PER_WARP", - "EPILOGUE_ROWS_PER_WARP", - "spad_words", - "spad_bytes", - "alloc_spad", - "rowwise_quant_block", - "rowwise_quant_store", - "rowwise_scale_flush", - "colwise_quant_store", -] - -# 16 bf16x2 words == 32 values == one E8M0 scale block. -WORDS_PER_SCALE_BLOCK = SCALE_BLOCK // 2 -# One columnwise staging chunk is 32 output columns == 16 column pairs. -COLWISE_PAIRS = WORDS_PER_SCALE_BLOCK -# Rows one epilogue warp owns; equal to the 32x1 block height, which is what -# keeps the whole columnwise round trip inside a warp. -COLWISE_ROWS_PER_WARP = SCALE_BLOCK -EPILOGUE_ROWS_PER_WARP = COLWISE_ROWS_PER_WARP -# Word pitch between two column-pair slabs. The +4 pad is the entire swizzle: -# the reader's 128-bit shared loads are serviced in four phases of eight lanes, -# and phase 0 covers pairs 0..7 at word `36*pair + 4*chunk`, i.e. banks -# `4*pair + 4*chunk (mod 32)` -- eight distinct 4-word groups, all 32 banks, no -# conflict. At pitch 32 every one of those eight lanes lands on the same bank -# group and the load degenerates to an 8-way conflict. -COLWISE_PAIR_PITCH = COLWISE_ROWS_PER_WARP + 4 -COLWISE_WORDS_PER_WARP = COLWISE_PAIRS * COLWISE_PAIR_PITCH -# Every sPad access is a 16-byte vector; 144 * pair + 64 * tq + 16 * chunk and -# the 2304-byte per-warp stride are all multiples of 16, so this holds for the -# allocation too. -_VEC_BYTES = 16 -_VEC_WORDS = _VEC_BYTES // 4 - - -def spad_words(num_epilogue_warps: int = 4) -> int: - """Int32 word count of the columnwise staging tile. - - Sized by the warp count rather than by a config object so that this module - stays importable, and testable, without the GEMM core. Callers holding a - ``GroupedGemmConfig`` pass ``len(config.epilogue_warp_ids)``. - """ - return num_epilogue_warps * COLWISE_WORDS_PER_WARP - - -def spad_bytes(num_epilogue_warps: int = 4) -> int: - """Shared-memory bytes the columnwise transpose costs (9216 B at 4 warps). - - That is 0.27 of one 128x128x128 E4M3 AB pipeline stage, i.e. it does not - change the achievable stage count. Feed it to - ``config.smem_bytes(epilogue_smem_bytes=...)``. - """ - return 4 * spad_words(num_epilogue_warps) # 4 bytes per Int32 word - - -@cute.jit -def alloc_spad( - allocator: SmemAllocator, NUM_EPILOGUE_WARPS: cutlass.Constexpr = 4 -) -> cute.Tensor: - """Allocate the columnwise staging tile as a flat Int32 shared tensor.""" - return allocator.allocate_tensor( - Int32, - cute.make_layout(spad_words(NUM_EPILOGUE_WARPS)), - byte_alignment=_VEC_BYTES, - ) - - -@cute.jit -def _store_vec_i32(dst: cute.Tensor, word_offset: Int32, src: cute.Tensor): - """Vectorized store of a register word run into a flat byte destination. - - `dst` is any 1-byte-element flat view; the run is reinterpreted as Int32, so - `word_offset` counts 4-byte words from the start of the buffer. - """ - cute.autovec_copy( - src, - cute.make_tensor( - (cute.recast_ptr(dst.iterator, dtype=Int32) + word_offset).align( - _VEC_BYTES - ), - cute.make_layout(cute.size(src)), - ), - ) - - -@cute.jit -def _load_spad_quad(spad: cute.Tensor, word_offset: Int32) -> cute.Tensor: - """One 16-byte shared load: four consecutive rows of one column pair.""" - quad = cute.make_rmem_tensor((_VEC_WORDS,), Int32) - cute.autovec_copy( - cute.make_tensor( - (spad.iterator + word_offset).align(_VEC_BYTES), - cute.make_layout(_VEC_WORDS), - ), - quad, - ) - return quad - - -@cute.jit -def rowwise_quant_block(words: cute.Tensor, WORD_BASE: cutlass.Constexpr = 0): - """Quantize one 1x32 rowwise block held entirely in one thread's registers. - - Returns ``(qwords, scale_byte)``: eight Int32 words holding the 32 E4M3 - bytes in column order, and the E8M0 exponent byte for the block. - - The amax reduction is intra-thread by construction (see the module - docstring), so there is no shuffle and no partial amax carried between - epilogue subtiles. - """ - amax_packed = words[WORD_BASE] - for w in cutlass.range_constexpr(1, WORDS_PER_SCALE_BLOCK): - amax_packed = abs_max_nan_bf16x2(amax_packed, words[WORD_BASE + w]) - # fold_amax masks the junk sign bits xorsign leaves behind, then folds the - # two bf16 lanes with the NaN-propagating max. - scale_byte = float_to_e8m0(fold_amax(amax_packed) << Int32(16)) - inv = e8m0_reciprocal_bf16(scale_byte) - inv_packed = inv | (inv << Int32(16)) - - qwords = cute.make_rmem_tensor((WORDS_PER_SCALE_BLOCK // 2,), Int32) - for q in cutlass.range_constexpr(WORDS_PER_SCALE_BLOCK // 2): - # mul_cvt_2x emits bytes [w0.lo, w0.hi, w1.lo, w1.hi], i.e. four - # consecutive columns in increasing address order. - qwords[q] = mul_cvt_2x( - words[WORD_BASE + 2 * q], words[WORD_BASE + 2 * q + 1], inv_packed - ) - return qwords, scale_byte - - -@cute.jit -def rowwise_quant_store( - words: cute.Tensor, - qdata: cute.Tensor, - row: Int32, - col: Int32, - row_stride: Int32, - WORD_BASE: cutlass.Constexpr = 0, -) -> Int32: - """Quantize one 1x32 block and store its 32 E4M3 bytes; return the scale byte. - - `qdata` is a flat 1-byte-element view of a row-major destination whose row - pitch is `row_stride`. `row_stride` is a multiple of 128 and `col` a - multiple of 32 under the kernel contract, so the destination byte offset is - 32-byte aligned and the 32 bytes leave as two STG.128. - - The scale byte is returned rather than stored: consecutive subtiles produce - consecutive blocked-scale columns, so the caller buffers them and flushes - once per CTA tile through :func:`rowwise_scale_flush`. - """ - qwords, scale_byte = rowwise_quant_block(words, WORD_BASE=WORD_BASE) - _store_vec_i32(qdata, (row * row_stride + col) >> Int32(2), qwords) - return scale_byte - - -@cute.jit -def rowwise_scale_flush( - scales: cute.Tensor, - row: Int32, - scale_col_base: Int32, - scale_bytes: cute.Tensor, - num_scale_col_blocks: Int32, - NUM_BYTES: cutlass.Constexpr, -): - """Write `NUM_BYTES` consecutive rowwise blocked-scale bytes for one row. - - Scale columns that share ``scale_col >> 2`` are contiguous bytes in the - tcgen05 blocked layout (they differ only in the low two bits of the flat - index), so a run of 4 starting at a 4-aligned `scale_col_base` is one 4-byte - store. `scale_col_base` is `tile_n * NUM_BYTES` in both kernels, hence - aligned to whichever width is selected here. - - `scales` must be a flat uint8 view of the blocked buffer. - """ - if cutlass.const_expr(NUM_BYTES % 4 == 0): - _flush_scale_run( - scales, row, scale_col_base, scale_bytes, num_scale_col_blocks, 4, NUM_BYTES - ) - elif cutlass.const_expr(NUM_BYTES % 2 == 0): - _flush_scale_run( - scales, row, scale_col_base, scale_bytes, num_scale_col_blocks, 2, NUM_BYTES - ) - else: - for i in cutlass.range_constexpr(NUM_BYTES): - idx = blocked_scale_idx(row, scale_col_base + i, num_scale_col_blocks) - scales[idx] = cutlass.Uint8(scale_bytes[i]) - - -@cute.jit -def _flush_scale_run( - scales: cute.Tensor, - row: Int32, - scale_col_base: Int32, - scale_bytes: cute.Tensor, - num_scale_col_blocks: Int32, - WIDTH: cutlass.Constexpr, - NUM_BYTES: cutlass.Constexpr, -): - """Emit `NUM_BYTES` scale bytes as `NUM_BYTES // WIDTH` packed stores. - - A packed store addresses `scales` in units of WIDTH bytes, so the byte index - it lands on is `(idx // WIDTH) * WIDTH`. That equals `idx` only when the - index is WIDTH-aligned; otherwise the store both corrupts a neighbouring - block's scale byte and leaves the intended one unwritten. - - In the blocked layout every term of `blocked_scale_idx` except `scale_col & 3` - is a multiple of 4, so `idx % WIDTH == scale_col % WIDTH`. The requirement is - therefore exactly `scale_col_base % WIDTH == 0`, a property of the caller's - argument rather than of the data. Today's callers pass `2 * tile_n` (WIDTH 2) - and `8 * tile_n` (WIDTH 4), whose multipliers are multiples of WIDTH, so - alignment holds structurally for any tile index. - - The `assert_` below only fires in an assertions-enabled build - (`CUTE_DSL_ENABLE_ASSERTIONS=1`); it is a debugging aid, not the guarantee. - The guarantee is the caller contract above, which matters because the - design's generalization ("buffer min(4, CTA_N/64) bytes") would put a - computed expression here. - """ - packed_ty = cutlass.Uint32 if cutlass.const_expr(WIDTH == 4) else cutlass.Uint16 - packed_ptr = cute.recast_ptr(scales.iterator, dtype=packed_ty) - for run in cutlass.range_constexpr(NUM_BYTES // WIDTH): - acc = Int32(0) - for i in cutlass.range_constexpr(WIDTH): - acc = acc | ((scale_bytes[run * WIDTH + i] & Int32(0xFF)) << Int32(8 * i)) - idx = blocked_scale_idx(row, scale_col_base + run * WIDTH, num_scale_col_blocks) - cute.testing.assert_( - idx % WIDTH == 0, - "packed scale store is not WIDTH-aligned: scale_col_base must be a " - "multiple of the store width", - ) - dst = cute.make_tensor(packed_ptr + (idx // WIDTH), cute.make_layout(1)) - dst[0] = packed_ty(acc) - - -@cute.jit -def colwise_quant_store( - words: cute.Tensor, - spad: cute.Tensor, - qdata: cute.Tensor, - scales: cute.Tensor, - tidx: Int32, - row_base: Int32, - col_base: Int32, - num_rows: Int32, - num_scale_col_blocks: Int32, - WORD_BASE: cutlass.Constexpr = 0, -): - """Transpose-quantize one 32-column chunk into the ``(1, R)`` destination. - - Every epilogue thread contributes its own row's 32 columns through `words` - and then reads back a *different* slice -- 16 rows of one column pair -- - which is the transpose. Writer and reader sets are the same 32 threads, so - both hazards are covered by ``sync_warp`` rather than a CTA barrier. - - `qdata` is a flat 1-byte-element view of the ``(1, num_rows)``-strided - destination, i.e. physically row-major ``[cols, num_rows]``. `scales` is a - flat uint8 view of the whole-matrix blocked columnwise scale buffer, indexed - with transposed coordinates (feature, row-block). - - `tidx` is the EPILOGUE-LOCAL thread index in - ``[0, 32 * num_epilogue_warps)``: a kernel whose epilogue runs on warps 4-7 - of 256 threads passes ``thread_idx - 128``. The epilogue's first thread must - be 32-aligned so that epilogue warp `w` is one physical warp -- otherwise - ``sync_warp`` and the butterfly shuffle no longer cover the writer/reader - set and the transpose silently mixes rows from two warps. - - `row_base` must be a multiple of 128 and `col_base` a multiple of 32. - """ - lane = tidx % Int32(COLWISE_ROWS_PER_WARP) - warp = tidx // Int32(COLWISE_ROWS_PER_WARP) - warp_base = warp * Int32(COLWISE_WORDS_PER_WARP) - - # Writer: for a fixed pair the 32 lanes hit 32 consecutive words, so every - # store is conflict-free without a swizzle. - for p in cutlass.range_constexpr(COLWISE_PAIRS): - spad[warp_base + Int32(p * COLWISE_PAIR_PITCH) + lane] = words[WORD_BASE + p] - cute.arch.sync_warp() - - # Reader: thread (pair, tq) owns column pair `pair` and the 16 rows - # [16*tq, 16*tq+16) of this warp's 32-row slab. `half` numbers those 16-row - # halves across the whole CTA tile, so `half` and `warp` are the same index - # at two granularities and the row-block below needs no extra arithmetic. - pair = tidx % Int32(COLWISE_PAIRS) - half = tidx // Int32(COLWISE_PAIRS) - tq = half % Int32(2) - feature = col_base + Int32(2) * pair - pair_base = ( - warp_base - + pair * Int32(COLWISE_PAIR_PITCH) - + tq * Int32(COLWISE_ROWS_PER_WARP // 2) - ) - - quads = [] - amax_packed = Int32(0) - for c in cutlass.range_constexpr(COLWISE_ROWS_PER_WARP // 2 // _VEC_WORDS): - quad = _load_spad_quad(spad, pair_base + Int32(c * _VEC_WORDS)) - quads.append(quad) - for t in cutlass.range_constexpr(_VEC_WORDS): - amax_packed = abs_max_nan_bf16x2(amax_packed, quad[t]) - cute.arch.sync_warp() - - # The two threads holding the halves of one 32-row block differ exactly in - # bit 4 of tidx, which is a lane bit, so one butterfly at distance 16 - # completes the amax. Masking first drops the junk signs xorsign leaves. - amax_packed = amax_packed & Int32(0x7FFF7FFF) - amax_packed = max_nan_bf16x2( - amax_packed, cute.arch.shuffle_sync_bfly(amax_packed, COLWISE_PAIRS) - ) - - # Low half is column `2*pair`'s amax, high half is `2*pair+1`'s: two - # independent scales out of one reduction. Both are BF16-valued widened to - # f32, which is what makes float_to_e8m0 exact. - lo_byte = float_to_e8m0((amax_packed & Int32(0xFFFF)) << Int32(16)) - hi_byte = float_to_e8m0((amax_packed >> Int32(16)) << Int32(16)) - inv_packed = e8m0_reciprocal_bf16(lo_byte) | ( - e8m0_reciprocal_bf16(hi_byte) << Int32(16) - ) - - # One of the two threads per (pair, row-block) owns the scale bytes. They - # are not adjacent in the blocked layout (they differ in the feature index, - # hence by 16 bytes), so this stays two 1-byte stores. - if tq == Int32(0): - row_block = (row_base // Int32(SCALE_BLOCK)) + warp - scales[blocked_scale_idx(feature, row_block, num_scale_col_blocks)] = ( - cutlass.Uint8(lo_byte) - ) - scales[ - blocked_scale_idx(feature + Int32(1), row_block, num_scale_col_blocks) - ] = cutlass.Uint8(hi_byte) - - # Quantize and transpose: mul_cvt_2x turns two rows x two columns into - # bytes [r0c0, r0c1, r1c0, r1c1], and the byte permutes split that into one - # word per column holding four consecutive rows. This is where prmt_even / - # prmt_odd belong -- not in any gate/up de-interleave, where the pair is - # already in two separate FP32 registers. - num_quads = cutlass.const_expr(COLWISE_ROWS_PER_WARP // 2 // _VEC_WORDS) - col_lo = cute.make_rmem_tensor((num_quads,), Int32) - col_hi = cute.make_rmem_tensor((num_quads,), Int32) - for c in cutlass.range_constexpr(num_quads): - quad = quads[c] - pack01 = mul_cvt_2x(quad[0], quad[1], inv_packed) - pack23 = mul_cvt_2x(quad[2], quad[3], inv_packed) - col_lo[c] = prmt_even(pack01, pack23) - col_hi[c] = prmt_odd(pack01, pack23) - - # col_lo/col_hi are 16 consecutive rows of one column: 16 contiguous bytes - # of the [cols, num_rows] buffer. num_rows and row_base are multiples of 128 - # and the row offset is a multiple of 16, so both stores are STG.128. - row_offset = row_base + half * Int32(COLWISE_ROWS_PER_WARP // 2) - _store_vec_i32(qdata, (feature * num_rows + row_offset) >> Int32(2), col_lo) - _store_vec_i32( - qdata, ((feature + Int32(1)) * num_rows + row_offset) >> Int32(2), col_hi - ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py deleted file mode 100644 index 11c448fba5..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_config.py +++ /dev/null @@ -1,463 +0,0 @@ -# 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. - -"""Frozen configuration for the shared MXFP8 grouped blockscaled GEMM core. - -One tiling, one pipeline shape, one warp assignment for all three kernels in the -family (FC1 GEMM+SwiGLU, FC2 dgrad+dSwiGLU, grouped wgrad). Kernels differ only -in their epilogue and in ``epi_n_acc``; nothing else here is a per-kernel knob. - -Why each value is what it is: - -``cta_group = ONE`` / ``cluster_shape_mn = (1, 1)`` - Every per-expert row count is a multiple of 128 and ``cta_tile_m`` is 128, so - an M tile never straddles an expert boundary. 2CTA forces a 256-row cluster - tile, which reintroduces partial tiles for ``m[g] == 128 (mod 256)`` and - would need hand-predication in every quantized store and every scale byte. - With one CTA per MMA there is also nothing to multicast, so the cluster is - trivial and no multicast masks or cluster launch barriers exist. - -``mma_tiler = (128, 128, 128)`` - M=128 is mandatory for ``CtaGroup.ONE``. K=128 is the smallest multiple of - ``sf_vec_size * 4`` and makes one K tile exactly one scale-factor atom. N=128 - keeps TMEM at 160 of 512 columns and avoids both the N=64 SFB column-shift - and the N=256 overlapping-accumulator machinery. - -``num_acc_stage = 1`` - One CTA produces exactly one output tile (the grid is data-independent and - there is no persistent tile scheduler), so there is no second accumulator to - overlap with. - -``threads = 256`` - Warp 0 loads (TMA), warp 1 issues the MMA, warps 4-7 run the epilogue, warps - 2-3 idle. 128 epilogue threads is what makes ``tmem_warp_shape_mn = (4, 1)`` - and the thread-owns-one-row identity hold; see :data:`T2R_PARTITION_DOC`. - -``epi_tile = (128, epi_n_acc)`` - ``epi_tile_M == cta_tile_M == 128`` is mandatory: every gate/up pairing and - every scale index below assumes a single epilogue tile along M. - -CTA_N = 256 is a legal alternative (TMEM 256+16+32 = 304 <= 512 with one -accumulator stage) and changes only the rowwise scale store width, but it halves -``num_ab_stage``; nothing structural depends on the choice. - -WARNING -- ``cta_tile_k`` is pinned at 128 and raising it is not a free tuning -knob. The grouped wgrad kernel selects an expert's K range with an integer K-tile -index base rather than a per-expert TMA descriptor, which is exact only because -``token_offset`` (a multiple of 128) is a multiple of ``cta_tile_k``. At -``cta_tile_k = 256`` a group boundary can land mid-tile, and the kernel would -need per-expert descriptors or explicit K-tail predication brought back in. The -constant-512-byte ``Rest_K`` stride that makes the index base correct is a -property of the 128-element scale-factor granule, not of the tile size, so it -does not rescue a larger tile either. -""" - -import enum -from dataclasses import dataclass -from typing import Tuple - -__all__ = [ - "RaggedAxis", - "GroupedGemmConfig", - "SWIGLU_FWD_CONFIG", - "DSWIGLU_BWD_CONFIG", - "WGRAD_CONFIG", - "SMEM_CAPACITY_BYTES", - "TMEM_TOTAL_COLS", - "is_supported", - "check_supported", - "T2R_PARTITION_DOC", - "EPILOGUE_PROTOCOL_DOC", -] - -# sm_100 usable dynamic shared memory per CTA, (228 - 1) KiB. -SMEM_CAPACITY_BYTES = 232448 -# The TMEM allocator requires a power-of-two multiple of 32 columns; shared -# memory already pins us to one CTA per SM, so allocating the whole array costs -# nothing. -TMEM_TOTAL_COLS = 512 -# Row-count granularity every expert group, and the allocation itself, respect. -GROUP_ALIGNMENT = 128 -# MXFP8 scaling block: 32 values share one E8M0 scale. -SF_VEC_SIZE = 32 - - -class RaggedAxis(enum.Enum): - """Which GEMM axis the per-expert offsets partition. - - ``M`` is kernels A and B: the ragged axis is the token axis, the grid is - ``(R // 128, N // CTA_N, 1)``, and the expert is looked up per CTA from the - absolute row base. ``K`` is kernel C: the ragged axis is the contraction, the - grid is ``(N // 128, K // 128, G)``, and only the K-loop trip count is - data-dependent. - """ - - M = 0 - K = 1 - - -@dataclass(frozen=True) -class GroupedGemmConfig: - """Constexpr configuration of the shared grouped blockscaled GEMM core. - - Every field is a trace-time constant. No kernel body may hard-code any of - these numbers; read them from the config so a retune cannot silently - desynchronize the mainloop from an epilogue. - """ - - # Per-kernel: accumulator columns handed to the epilogue per subtile. - # 64 for the FC1 SwiGLU forward (one output column consumes an adjacent - # gate/up accumulator pair, so 64 accumulator columns are 32 output columns - # == exactly one rowwise 1x32 block per thread), 32 for the dgrad backward - # and for wgrad. - epi_n_acc: int - # Which axis the offsets partition. - ragged_axis: RaggedAxis - - # --- frozen for every kernel in the family ------------------------------- - cta_tile_m: int = 128 - cta_tile_n: int = 128 - cta_tile_k: int = 128 - num_ab_stage: int = 6 - num_acc_stage: int = 1 - cluster_shape_mn: Tuple[int, int] = (1, 1) - sf_vec_size: int = SF_VEC_SIZE - threads: int = 256 - tma_warp_id: int = 0 - mma_warp_id: int = 1 - epilogue_warp_ids: Tuple[int, ...] = (4, 5, 6, 7) - # Named barrier ids. 0 is left free for the DSL's own use. - epilogue_sync_barrier_id: int = 1 - tmem_alloc_barrier_id: int = 2 - - def __post_init__(self): - if self.cta_tile_m != 128: - raise ValueError( - "cta_tile_m is pinned at 128: CtaGroup.ONE requires it and the " - "no-partial-M-tile argument depends on it" - ) - if self.cta_tile_k != 128: - raise ValueError( - "cta_tile_k is pinned at 128; see this module's docstring for why " - "raising it reintroduces per-expert descriptors" - ) - if self.cta_tile_n % 128 != 0: - # SFB's MN extent is round_up(N, 128). Only when that equals N is the - # SFB tiled MMA identical to the data one, which is what lets the core - # build a single tiled MMA and skip the N=64 TMEM column-shift path. - raise ValueError( - f"cta_tile_n must be a multiple of 128, got {self.cta_tile_n}; " - "a smaller N needs a separate SFB tiled MMA and a TMEM column shift" - ) - if self.cta_tile_n % self.epi_n_acc != 0: - raise ValueError( - f"cta_tile_n ({self.cta_tile_n}) must be a multiple of epi_n_acc " - f"({self.epi_n_acc})" - ) - if self.cta_tile_m % (4 * SF_VEC_SIZE) != 0: - # 4 epilogue warps x 32 rows: a columnwise 32x1 block must never be - # split across warps. - raise ValueError( - "cta_tile_m must be a multiple of 128 for the 4-warp epilogue" - ) - if len(self.epilogue_warp_ids) * 32 != self.cta_tile_m: - raise ValueError( - "the epilogue must have exactly one thread per row of the CTA tile" - ) - if self.epilogue_warp_ids[0] % 4 != 0: - # tcgen05.ld selects its TMEM datapath sub-partition from the - # PHYSICAL warp id, so the epilogue must start on an aligned warp - # quad. A misaligned block still launches and still returns - # plausible numbers -- every 128-row tile comes back with its four - # 32-row datapath groups rotated -- and no shape, byte or NaN check - # detects a pure row permutation. Reject it here. - raise ValueError( - f"epilogue_warp_ids must start on an aligned warp quad, got " - f"{self.epilogue_warp_ids}: warp {self.epilogue_warp_ids[0]} is " - f"{self.epilogue_warp_ids[0] % 4} past a multiple of 4; tcgen05.ld " - "would silently permute the 32-row groups of every tile" - ) - # The kernel selects the epilogue with `warp_idx >= epilogue_warp_ids[0]`, - # so they must be the last contiguous block of warps in the CTA. - if tuple(self.epilogue_warp_ids) != tuple( - range(self.threads // 32 - len(self.epilogue_warp_ids), self.threads // 32) - ): - raise ValueError( - f"epilogue_warp_ids {self.epilogue_warp_ids} must be the last " - f"{len(self.epilogue_warp_ids)} warps of the {self.threads // 32} " - "in the CTA" - ) - if self.tma_warp_id in self.epilogue_warp_ids or ( - self.mma_warp_id in self.epilogue_warp_ids - ): - raise ValueError("the TMA and MMA warps must not be epilogue warps") - - # --- derived, all trace-time --------------------------------------------- - - @property - def mma_tiler_mnk(self) -> Tuple[int, int, int]: - return (self.cta_tile_m, self.cta_tile_n, self.cta_tile_k) - - @property - def cta_tile_shape_mnk(self) -> Tuple[int, int, int]: - return (self.cta_tile_m, self.cta_tile_n, self.cta_tile_k) - - @property - def epi_tile(self) -> Tuple[int, int]: - return (self.cta_tile_m, self.epi_n_acc) - - @property - def num_epi_subtiles(self) -> int: - return self.cta_tile_n // self.epi_n_acc - - @property - def num_epilogue_threads(self) -> int: - return 32 * len(self.epilogue_warp_ids) - - @property - def first_epilogue_thread(self) -> int: - return 32 * self.epilogue_warp_ids[0] - - # --- shared memory budget ------------------------------------------------ - - @property - def ab_stage_bytes(self) -> int: - """Bytes of one mainloop stage: A and B tiles plus both scale atoms. - - At (128, 128, 128) E4M3 that is 16384 + 16384 + 512 + 512 = 33792, which - is also the exact ``tx_count`` the TMA pipeline barrier must expect. - """ - a = self.cta_tile_m * self.cta_tile_k # E4M3, 1 byte per element - b = self.cta_tile_n * self.cta_tile_k - # One E8M0 byte per 32 contracted elements per MN row, i.e. exactly one - # 128x4 blocked tile (512 B) per operand at the frozen shape. SFB's MN - # extent is round_up(N, 128). - sfa = self.cta_tile_m * (self.cta_tile_k // SF_VEC_SIZE) - sfb = max(self.cta_tile_n, 128) * (self.cta_tile_k // SF_VEC_SIZE) - return a + b + sfa + sfb - - @property - def mbarrier_bytes(self) -> int: - """AB pipeline (full+empty) and accumulator handoff (full+empty) mbarriers.""" - return 8 * (2 * self.num_ab_stage + 2 * self.num_acc_stage) - - def smem_bytes(self, epilogue_smem_bytes: int = 0) -> int: - return ( - self.num_ab_stage * self.ab_stage_bytes - + self.mbarrier_bytes - + epilogue_smem_bytes - # tmem holding buffer plus struct alignment slack - + 256 - ) - - def max_ab_stages(self, epilogue_smem_bytes: int = 0) -> int: - """Stages that fit alongside the epilogue's shared-memory request.""" - fixed = self.mbarrier_bytes + epilogue_smem_bytes + 256 - return (SMEM_CAPACITY_BYTES - fixed) // self.ab_stage_bytes - - # --- tensor memory budget ------------------------------------------------ - - @property - def acc_tmem_cols(self) -> int: - return self.cta_tile_n * self.num_acc_stage - - @property - def sfa_tmem_cols(self) -> int: - return (self.cta_tile_m // SF_VEC_SIZE) * 4 - - @property - def sfb_tmem_cols(self) -> int: - # SFB's MN extent is round_up(N, 128); at N=128 that is N itself. - return (max(self.cta_tile_n, 128) // SF_VEC_SIZE) * 4 - - @property - def used_tmem_cols(self) -> int: - """160 of 512 at the frozen shape.""" - return self.acc_tmem_cols + self.sfa_tmem_cols + self.sfb_tmem_cols - - -# Kernel A. EPI_N_ACC=64 is not negotiable: accumulator column 2f is gate_f and -# 2f+1 is up_f, so 64 accumulator columns are 32 output columns, exactly one -# rowwise 1x32 block per thread. 32 would give half a block and force a partial -# amax carried across subtiles. -SWIGLU_FWD_CONFIG = GroupedGemmConfig(epi_n_acc=64, ragged_axis=RaggedAxis.M) -# Kernel B. The accumulator N axis is F (one column per feature) and each column -# produces two interleaved output columns, so 32 accumulator columns are 64 -# output columns = two rowwise 1x32 blocks per thread per subtile. -DSWIGLU_BWD_CONFIG = GroupedGemmConfig(epi_n_acc=32, ragged_axis=RaggedAxis.M) -# Kernel C. 32 FP32 accumulators are 32 contiguous BF16 outputs = 64 contiguous -# bytes per thread, 64-byte aligned. -WGRAD_CONFIG = GroupedGemmConfig(epi_n_acc=32, ragged_axis=RaggedAxis.K) - - -def is_supported( - model_dim: int, hidden_dim: int, allocated_rows: int, num_groups: int -) -> bool: - """The initial optimized support predicate, as a pure boolean. - - Mirrors the kernel contract's shape predicate and the host validators. The - caller is expected to fall back to the unfused torchao path when this is - false rather than launching -- feeding non-128-aligned groups to the blocked - scale path produces a wrong-sized buffer and an unusable CUDA context, not an - error. - - The per-expert row counts live in device memory and are not checkable here; - they are asserted on device on every launch. - """ - return ( - num_groups >= 1 - and model_dim > 0 - and hidden_dim > 0 - and allocated_rows > 0 - and model_dim % GROUP_ALIGNMENT == 0 - and hidden_dim % GROUP_ALIGNMENT == 0 - and allocated_rows % GROUP_ALIGNMENT == 0 - ) - - -def check_supported( - model_dim: int, hidden_dim: int, allocated_rows: int, num_groups: int -) -> None: - """:func:`is_supported` with a message naming the offending value. - - Raises ``ValueError``, never ``assert``: ``python -O`` strips assertions and - would reintroduce the silent-corruption path. - """ - if num_groups < 1: - raise ValueError(f"G must be at least 1, got {num_groups}") - for name, value in ( - ("D", model_dim), - ("F", hidden_dim), - ("R", allocated_rows), - ): - if value <= 0: - raise ValueError(f"{name} must be positive, got {value}") - if value % GROUP_ALIGNMENT != 0: - raise ValueError( - f"{name} must be a multiple of {GROUP_ALIGNMENT}, got {value}" - ) - - -# --------------------------------------------------------------------------- -# The two frozen interfaces lanes 2 and 3 compile against. -# --------------------------------------------------------------------------- - -T2R_PARTITION_DOC = """\ -grouped_gemm_core.t2r_partition(tidx, tAcc, config) -> (tiled_copy_t2r, -tTR_tAcc, tTR_rAcc, tTR_cAcc) - - tidx Int32 epilogue-local thread index in [0, 128). NOT the raw threadIdx.x; - subtract config.first_epilogue_thread first. - tAcc the (MMA, MMA_M, MMA_N, ACC_STAGE) TMEM accumulator tensor. - config a GroupedGemmConfig, which fixes epi_tile = (128, epi_n_acc). - -Returns, with EPI_M == 1 always: - - tiled_copy_t2r the tcgen05 TMEM->register tiled copy. - tTR_tAcc (T2R, T2R_M, T2R_N, EPI_M, EPI_N) TMEM source, sliced per - epilogue subtile s as tTR_tAcc[(None, None, None, 0, s)]. - tTR_rAcc (T2R, T2R_M, T2R_N) FP32 register destination for one subtile. - tTR_cAcc (T2R, T2R_M, T2R_N, EPI_M, EPI_N) of (row, col) coordinates in - the 128 x cta_tile_n CTA tile, partitioned identically. - -The copy atom is built with elem_ty_d = Float32 even though the real outputs are -E4M3. Passing an 8-bit d type steers get_tmem_load_op into the tmem_dp=16 -layouts, which are shaped for a direct FP8 TMA store and are not what a -dual-quantization epilogue wants. - -Structural consequence, measured on GB200 with a real TMEM accumulator (not just -from the copy atom's thread-value layout): thread t owns row t of the 128-row CTA -tile and all epi_n_acc accumulator columns of the subtile, contiguously, and the -raw register order is linear -- register v of thread t holds accumulator element -(t, v) of the subtile. Warp w therefore owns rows [32w, 32w+32) = exactly one -32-row MX block. So a rowwise 1x32 amax is intra-thread and a columnwise 32x1 -amax is intra-warp, and a columnwise scale block is never split across warps or -CTAs. - -Even so, derive every index -- the gate/up de-interleave and every scale -coordinate -- from tTR_cAcc rather than from a raw register number. The linear -order is a measured property of one copy atom at one epi_tile, not a guarantee, -and tTR_cAcc costs nothing because it folds at trace time. -""" - -EPILOGUE_PROTOCOL_DOC = """\ -An epilogue is a module-level function passed to the core as a Constexpr. Two -requirements, both load-bearing: - -1. It must be decorated ``@cute.jit``. The DSL preprocessor only rewrites - decorated functions, so an undecorated epilogue cannot use - ``cutlass.range_constexpr`` (it raises "range_constexpr should be preprocessed - by preprocessor") and, worse, a dynamic ``if`` in one would be evaluated as a - Python truth test instead of becoming a predicated region. -2. It must be a module-level function object. The DSL keys its compile cache on - function identity, so a lambda or a closure built per call recompiles every - launch. - - @cute.jit - def my_epilogue( - tTR_rAcc, # (T2R, T2R_M, T2R_N) FP32 register fragment, one subtile - tTR_cAcc_s, # (T2R, T2R_M, T2R_N) matching (row, col) coordinates in - # the 128 x cta_tile_n CTA tile - tiled_copy_t2r,# for cute.make_tiled_copy_D / retile, if needed - epi_tidx, # Int32 in [0, 128); equals the CTA-tile row this thread owns - subtile_idx, # Constexpr int in [0, config.num_epi_subtiles) - tile, # TileCoords: see below - epi_smem, # cute.Pointer(Int32) to the requested scratch, or None - out, # the tuple of destination tensors the launcher passed - cfg, # Constexpr GroupedGemmConfig - ) -> None - -TileCoords fields, all Int32 and all CTA-uniform: - - tile_m absolute M-tile index (kernels A/B: over [0, R/128)) - tile_n absolute N-tile index - expert selected expert index. Meaningless on an inactive-tail tile, where - the scan saturates at G-1; nothing may depend on it there, since - that tile's output is defined to be zero. - row_base tile_m * 128 - col_base tile_n * cta_tile_n - k_cnt mainloop trip count; 0 on a tail tile or a zero-token expert - -The epilogue is called once per subtile, num_epi_subtiles times per CTA, always -by all 128 epilogue threads and always with k_cnt CTA-uniform. - -Tail rule, stated precisely because the two halves pull in opposite directions: - - * NEVER predicate a STORE on the inactive tail. A tail tile arrives with a - zeroed accumulator and the unmodified store path then emits exactly the zeros - the contract requires (zero qdata bytes, zero E8M0 scale bytes). Skipping a - store is how a destination element stops being written. - - * ALWAYS predicate an extra GMEM INPUT LOAD on `tile.k_cnt == 0`, substituting - zeros. This is not optional. Kernel B reads the saved `z_bf16` in its - epilogue, and rows [A, R) of that tensor are read-forbidden precisely because - they may hold anything. Measured, feeding a tail `z` into the dSwiGLU with - dh == 0 as the zeroed accumulator delivers it: - z = 0 -> dz 0x00000000 correct - z = NaN or +Inf -> dz 0x7fff7fff (dgate/dup = NaN) WRONG - z = uninit 0xDEADBEEF -> dz 0x80008000 (dgate/dup = -0.0) WRONG - A NaN tail makes that block's scale byte 0xFF and every qdata byte 0x7F; - even benign garbage yields qdata 0x80 rather than 0x00. Both violate the - read-forbidden and write-zero halves of the ragged-tail contract. - -In short: the accumulator is already zeroed for you, so trust it and store -unconditionally; anything you load yourself must be gated on k_cnt. - -No cross-subtile state: EPILOGUE is called once per subtile and every call gets -fresh registers. There is deliberately no way to carry a value from subtile s to -s+1, so rowwise scale bytes cannot be buffered across subtiles and emitted as one -wide store per CTA tile. DECIDED 2026-08-16: emit ONE scale store per subtile -(`rowwise_scale_flush(..., NUM_BYTES=1)`). That costs 2 stores instead of 1 for -Kernel A and 8 instead of 2 for Kernel B, and is verified bitwise clean. Store-count -reduction is a tuning-stage concern; buying it here would mean either -trace-time module-level state across the unrolled loop or staging through -epi_smem, both of which trade a correctness-critical interface for a few stores. - -Shared-memory scratch: the epilogue declares its byte count to the launcher, -which passes back a 128-byte-aligned pointer of that size. The mainloop's stage -count is computed against that request, so it is not free -- but 9216 bytes (the -columnwise transpose staging) still leaves 6 stages. -""" diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py deleted file mode 100644 index 174620a213..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_gemm_core.py +++ /dev/null @@ -1,825 +0,0 @@ -# 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. - -"""Shared MXFP8 grouped blockscaled GEMM core: descriptors, mainloop, TMEM, T2R. - -Everything between "torch tensors" and "an FP32 accumulator tile in registers". -The three kernels in this family plug different epilogues into this one mainloop -by passing a module-level function as the ``EPILOGUE`` Constexpr; see -``grouped_gemm_config.EPILOGUE_PROTOCOL_DOC``. - -Three structural decisions, all descending from the per-expert row counts being -multiples of 128: - -*No per-group tensormaps, anywhere.* Every operand is one host-built static TMA -descriptor over the whole tensor. Per-expert selection is an integer coordinate: -an L coordinate for the 3-D weight operands, and a K-tile index base for the -wgrad kernel's ragged contraction. The latter is exact because the blockscaled -scale-factor layout's ``Rest_K`` stride is a constant 512 bytes for *every* MN -row-block, so advancing the K tile index advances every row-block's byte address -by the same amount. That was verified on this wheel (probe V1) and it is what -deletes the tensormap workspace, the descriptor-init kernel, the descriptor -fences, and the padded-offset prefix sum that the reference kernels carry. - -*No tile scheduler.* With the ragged axis tile-aligned, kernels A/B enumerate all -of ``[0, R/128)`` M tiles and kernel C's ``(N/128, K/128, G)`` grid is fully -static; only C's K-loop trip count is data-dependent. So there is no -``max_active_clusters`` query, no persistent loop, no offsets prefix scan, and no -work-tile shared-memory pipeline. - -*The inactive tail needs no special code path.* A tile whose row base is at or -past the active row count runs with ``k_cnt == 0``: no TMA loads are issued (so -no inactive row is ever read), the accumulator fragment is zeroed in registers, -and the *unmodified* epilogue emits the zeros the contract requires. ``k_cnt`` is -CTA-uniform, so no barrier arrival count can desynchronize. - -Expert lookup for the ragged-M kernels is an unrolled G-way scan over the -device-side offsets -- no host synchronization and no ``.item()``. -""" - -from dataclasses import dataclass - -import cutlass -import cutlass.cute as cute -import cutlass.pipeline as pipeline -import cutlass.utils as utils -import cutlass.utils.blackwell_helpers as sm100_utils -import cutlass.utils.blockscaled_layout as blockscaled_utils -import torch -from cutlass import Float32, Int32 -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.utils import LayoutEnum - -from torchao.prototype.moe_training.kernels.mxfp8.grouped_gemm_config import ( - SMEM_CAPACITY_BYTES, - TMEM_TOTAL_COLS, - GroupedGemmConfig, - RaggedAxis, -) -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_epilogue import ( - validate_group_offsets_device, -) - -__all__ = [ - "TileCoords", - "activation_gemm_view", - "weight_gemm_view", - "make_operand_views", - "make_tiled_mma", - "make_sf_gemm_tensor", - "t2r_partition", - "dump_accumulator_epilogue", - "grouped_gemm_kernel", - "launch_grouped_gemm", -] - - -# --------------------------------------------------------------------------- -# Host-side operand views. Pure torch, no copies: every one of these is a -# restride of the caller's storage into the (MN, K, L) GEMM domain the core -# wants, with K contiguous. -# --------------------------------------------------------------------------- - - -def activation_gemm_view(t: torch.Tensor) -> torch.Tensor: - """``[MN, K]`` K-contiguous -> ``(MN, K, 1)`` with a defined batch stride. - - Covers ``x_q``/``do_q`` directly, and both wgrad operands via their free - transpose: a logical ``[R, N]`` with stride ``(1, R)`` *is* a K-contiguous - ``[N, R]``, so pass ``t.t()``. - """ - mn, k = t.shape - if t.stride() != (k, 1): - raise ValueError( - f"GEMM operand must be K-contiguous with stride {(k, 1)}, got {t.stride()}" - ) - return torch.as_strided(t, (mn, k, 1), (k, 1, mn * k)) - - -def weight_gemm_view(w: torch.Tensor) -> torch.Tensor: - """``[G, K, N]`` stride ``(K*N, 1, K)`` -> ``(N, K, G)`` stride ``(K, 1, K*N)``. - - That is the prequantized weight layout both kernels A and B receive, and the - permute is exactly the K-major B operand the MMA wants, so the expert becomes - an L coordinate and no descriptor is ever rebuilt. - """ - g, k, n = w.shape - if w.stride() != (k * n, 1, k): - raise ValueError( - f"grouped weight must have stride {(k * n, 1, k)}, got {w.stride()}" - ) - return w.permute(2, 1, 0) - - -def make_operand_views(a: torch.Tensor, b: torch.Tensor): - """``(mA, mB)`` in the GEMM domain, choosing the view from ``b``'s rank. - - A 3-D ``b`` is a grouped weight (expert becomes the L coordinate); a 2-D one - is the wgrad case, where both operands are ungrouped and the expert is a - K-tile index base instead. - """ - return activation_gemm_view(a), ( - weight_gemm_view(b) if b.ndim == 3 else activation_gemm_view(b) - ) - - -@dataclass -class TileCoords: - """The per-CTA tile description handed to the epilogue. All fields CTA-uniform. - - ``k_cnt == 0`` marks both cases where the mainloop is skipped: an inactive - tail tile (ragged M) and a zero-token expert (ragged K). The epilogue must - not branch on it -- the accumulator is already zero. - """ - - tile_m: Int32 - tile_n: Int32 - expert: Int32 - row_base: Int32 - col_base: Int32 - k_cnt: Int32 - - -def make_tiled_mma(cfg: GroupedGemmConfig, a_dtype, b_dtype, sf_dtype): - """The one blockscaled tiled MMA, K-major on both operands. - - No operand in this family is MN-major, so the 8-bit MN-major N-step and - transpose-swizzle caveats never apply. ``MmaMXF8F6F4Op`` hard-wires FP32 - accumulation and instruction K=32, so a 128-element K tile is four MMA-K - instructions and exactly one scale-factor atom. - """ - return sm100_utils.make_blockscaled_trivial_tiled_mma( - a_dtype, - b_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.K, - sf_dtype, - cfg.sf_vec_size, - tcgen05.CtaGroup.ONE, - (cfg.cta_tile_m, cfg.cta_tile_n), - ) - - -def make_sf_gemm_tensor( - flat_sf: cute.Tensor, mn: int, k: int, l: int, sf_vec_size: int -): - """Retile a flat blocked E8M0 buffer into the GEMM-domain scale-factor layout. - - The buffer is carried flat by ABI and may arrive as either uint8 or - float8_e8m0fnu, so the pointer is recast unconditionally -- the MMA rejects a - scale operand whose element type is not E8M0. - - ``tile_atom_to_shape_SF`` builds kernel IR, so this must be called inside a - trace even though the shapes are static and every evaluation folds. - """ - return cute.make_tensor( - cute.recast_ptr(flat_sf.iterator, dtype=cutlass.Float8E8M0FNU), - blockscaled_utils.tile_atom_to_shape_SF((mn, k, l), sf_vec_size), - ) - - -def t2r_partition(tidx, tAcc_base: cute.Tensor, cfg: GroupedGemmConfig): - """Accumulator -> register handoff. See ``config.T2R_PARTITION_DOC``. - - ``elem_ty_d`` is Float32 even though the real outputs are E4M3: passing an - 8-bit d type steers ``get_tmem_load_op`` into the tmem_dp=16 layouts, which - are shaped for a direct FP8 TMA store, not for a dual-quantization epilogue. - - ``tTR_cAcc`` carries each register's ``(row, col)`` in the CTA tile. Deriving - every index from it, rather than from a raw register number, is what makes - the gate/up de-interleave and the scale addressing correct by construction - whatever the copy atom's internal value order is. - """ - copy_atom_t2r = sm100_utils.get_tmem_load_op( - cfg.cta_tile_shape_mnk, - LayoutEnum.ROW_MAJOR, - Float32, - Float32, - cfg.epi_tile, - False, - ) - # (MMA, MMA_M, MMA_N, ACC_STAGE) -> (CTA_M, CTA_N); one accumulator stage. - tAcc_mn = tAcc_base[((None, None), 0, 0, 0)] - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N) - tAcc_epi = cute.flat_divide(tAcc_mn, cfg.epi_tile) - tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) - thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - cAcc_epi = cute.flat_divide( - cute.make_identity_tensor((cfg.cta_tile_m, cfg.cta_tile_n)), cfg.epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N), values are (row, col) in the CTA tile - tTR_cAcc = thr_copy_t2r.partition_D(cAcc_epi) - # (T2R, T2R_M, T2R_N) - tTR_rAcc = cute.make_rmem_tensor(tTR_cAcc[(None, None, None, 0, 0)].shape, Float32) - return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc - - -def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor): - """SMEM -> TMEM scale-factor copy, issued once per K tile from the MMA warp. - - ``Cp4x32x128bOp`` carries the warpx4 broadcast qualifier: issue it as a plain - ``cute.copy`` and never wrap it in ``elect_one()``, which deadlocks because - the compiler already inserts the election. - """ - tCsSF_compact = cute.filter_zeros(sSF) - tCtSF_compact = cute.filter_zeros(tSF) - copy_atom_s2t = cute.make_copy_atom( - tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), sSF.element_type - ) - tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) - thr_copy_s2t = tiled_copy_s2t.get_slice(0) - tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor( - tiled_copy_s2t, thr_copy_s2t.partition_S(tCsSF_compact) - ) - tCtSF_s2t = thr_copy_s2t.partition_D(tCtSF_compact) - return tiled_copy_s2t, tCsSF_s2t, tCtSF_s2t - - -@cute.jit -def dump_accumulator_epilogue( - tTR_rAcc, - tTR_cAcc_s, - tiled_copy_t2r, - epi_tidx, - subtile_idx: cutlass.Constexpr, - tile: TileCoords, - epi_smem, - out, - cfg: cutlass.Constexpr, -): - """The M2a gate epilogue: write the raw FP32 accumulator to ``out[0]``. - - ``out[0]`` is ``[M, N, L]`` FP32, where L is 1 for a ragged-M kernel and the - expert count for a ragged-K one. It exists to prove the blockscaled MMA and - the scale-factor addressing before any real epilogue does; it is deliberately - a scalar store loop, not a vectorized one. - """ - gD = out[0] - if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.K): - out_l = tile.expert - else: - out_l = Int32(0) - for v in cutlass.range_constexpr(cute.size(tTR_rAcc)): - crd = tTR_cAcc_s[v] - gD[(tile.row_base + crd[0], tile.col_base + crd[1], out_l)] = tTR_rAcc[v] - - -@cute.kernel -def grouped_gemm_kernel( - tiled_mma: cute.TiledMma, - tma_atom_a: cute.CopyAtom, - mA: cute.Tensor, - tma_atom_b: cute.CopyAtom, - mB: cute.Tensor, - tma_atom_sfa: cute.CopyAtom, - mSFA: cute.Tensor, - tma_atom_sfb: cute.CopyAtom, - mSFB: cute.Tensor, - offs: cute.Tensor, - out, - a_smem_layout: cute.ComposedLayout, - b_smem_layout: cute.ComposedLayout, - sfa_smem_layout: cute.Layout, - sfb_smem_layout: cute.Layout, - cfg: cutlass.Constexpr, - storage_type: cutlass.Constexpr, - EPILOGUE: cutlass.Constexpr, - EPI_SMEM_BYTES: cutlass.Constexpr, - VALIDATE_OFFSETS: cutlass.Constexpr, -): - """One CTA computes one 128 x cta_tile_n output tile. Warps: 0 TMA, 1 MMA, - 4-7 epilogue, 2-3 idle.""" - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - bidx, bidy, bidz = cute.arch.block_idx() - - num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) - num_k_tiles_full = cutlass.const_expr(cute.size(mA, mode=[1]) // cfg.cta_tile_k) - - # ------------------------------------------------------------------ - # Device-side precondition on the offset VALUES, which the host cannot - # check without synchronizing. One block, one warp, one lane. - # ------------------------------------------------------------------ - if cutlass.const_expr(VALIDATE_OFFSETS): - if bidx == 0 and bidy == 0 and bidz == 0: - if warp_idx == 0: - with cute.arch.elect_one(): - validate_group_offsets_device( - offs, - Int32(cute.size(mA, mode=[0])) - if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.M) - else Int32(cute.size(mA, mode=[1])), - ) - - # ------------------------------------------------------------------ - # Tile coordinates. No scheduler: the grid IS the tile enumeration. - # ------------------------------------------------------------------ - tile_m = Int32(bidx) - tile_n = Int32(bidy) - if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.M): - row_base = tile_m * cfg.cta_tile_m - # Unrolled G-way scan: the owning expert is the number of groups that end - # at or before this tile's row base. A zero-token expert is never - # selected, since its end equals its start. - expert = Int32(0) - for g in cutlass.range_constexpr(num_groups - 1): - expert += Int32(offs[g] <= row_base) - # Branch-free tail predicate: rows at or past offs[-1] belong to no - # expert, so the whole mainloop is skipped for them. - is_active = Int32(offs[num_groups - 1] > row_base) - k_base = Int32(0) - k_cnt = is_active * Int32(num_k_tiles_full) - l_a = Int32(0) - l_b = expert - else: - expert = Int32(bidz) - row_base = tile_m * cfg.cta_tile_m - # offs[expert - 1], clamped for expert 0; offsets are nonnegative so the - # multiply is a legal select. - prev = offs[cutlass.max(expert - Int32(1), Int32(0))] * Int32(expert > Int32(0)) - # Exact, not a ceil: every group boundary is a multiple of cta_tile_k. - k_base = prev // cfg.cta_tile_k - k_cnt = (offs[expert] - prev) // cfg.cta_tile_k - l_a = Int32(0) - l_b = Int32(0) - tile = TileCoords( - tile_m=tile_m, - tile_n=tile_n, - expert=expert, - row_base=row_base, - col_base=tile_n * cfg.cta_tile_n, - k_cnt=k_cnt, - ) - - # ------------------------------------------------------------------ - # Shared memory and pipelines - # ------------------------------------------------------------------ - smem = utils.SmemAllocator() - storage = smem.allocate(storage_type) - - sA = storage.sA.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) - sB = storage.sB.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) - sSFA = storage.sSFA.get_tensor(sfa_smem_layout) - sSFB = storage.sSFB.get_tensor(sfb_smem_layout) - epi_smem = None - if cutlass.const_expr(EPI_SMEM_BYTES > 0): - epi_smem = storage.sEpi.data_ptr() - - cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((*cfg.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) - ) - - ab_pipeline = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar.data_ptr(), - num_stages=cfg.num_ab_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - # Must be the EXACT byte count of all four TMA copies in one stage: too - # small and the MMA consumes a partially arrived stage. - tx_count=cutlass.const_expr(cfg.ab_stage_bytes), - cta_layout_vmnk=cluster_layout_vmnk, - ) - # One accumulator stage, so this is a single mbarrier pair; there is no - # inter-tile pipelining to overlap with. - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar.data_ptr(), - num_stages=cfg.num_acc_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, cfg.num_epilogue_threads - ), - cta_layout_vmnk=cluster_layout_vmnk, - ) - - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=cfg.tmem_alloc_barrier_id, - num_threads=32 * (1 + len(cfg.epilogue_warp_ids)), - ) - epilogue_barrier = pipeline.NamedBarrier( - barrier_id=cfg.epilogue_sync_barrier_id, - num_threads=cfg.num_epilogue_threads, - ) - tmem_alloc = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=cfg.epilogue_warp_ids[0], - is_two_cta=False, - ) - - # ------------------------------------------------------------------ - # Tile the global tensors. One static descriptor per operand; the expert is - # an L coordinate (ragged M) or a K-tile index base (ragged K). - # ------------------------------------------------------------------ - mma_tiler = cfg.mma_tiler_mnk - gA = cute.local_tile( - mA, cute.slice_(mma_tiler, (None, 0, None)), (None, None, None) - ) - gB = cute.local_tile( - mB, cute.slice_(mma_tiler, (0, None, None)), (None, None, None) - ) - gSFA = cute.local_tile( - mSFA, cute.slice_(mma_tiler, (None, 0, None)), (None, None, None) - ) - gSFB = cute.local_tile( - mSFB, cute.slice_(mma_tiler, (0, None, None)), (None, None, None) - ) - - thr_mma = tiled_mma.get_slice(0) - tCgA = thr_mma.partition_A(gA) - tCgB = thr_mma.partition_B(gB) - tCgSFA = thr_mma.partition_A(gSFA) - tCgSFB = thr_mma.partition_B(gSFB) - - trivial_cta_layout = cute.make_layout(1) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, - 0, - trivial_cta_layout, - cute.group_modes(sA, 0, 3), - cute.group_modes(tCgA, 0, 3), - ) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, - 0, - trivial_cta_layout, - cute.group_modes(sB, 0, 3), - cute.group_modes(tCgB, 0, 3), - ) - tAsSFA, tAgSFA = cpasync.tma_partition( - tma_atom_sfa, - 0, - trivial_cta_layout, - cute.group_modes(sSFA, 0, 3), - cute.group_modes(tCgSFA, 0, 3), - ) - # Strip the stride-0 sf_vec_size sub-mode: the 512-byte scale atom is - # contiguous and TMA moves it as 8-byte elements. - tAsSFA = cute.filter_zeros(tAsSFA) - tAgSFA = cute.filter_zeros(tAgSFA) - tBsSFB, tBgSFB = cpasync.tma_partition( - tma_atom_sfb, - 0, - trivial_cta_layout, - cute.group_modes(sSFB, 0, 3), - cute.group_modes(tCgSFB, 0, 3), - ) - tBsSFB = cute.filter_zeros(tBsSFB) - tBgSFB = cute.filter_zeros(tBgSFB) - - tAgA_slice = tAgA[(None, tile_m, None, l_a)] - tBgB_slice = tBgB[(None, tile_n, None, l_b)] - tAgSFA_slice = tAgSFA[(None, tile_m, None, l_a)] - tBgSFB_slice = tBgSFB[(None, tile_n, None, l_b)] - - acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, cfg.num_acc_stage)) - - # ------------------------------------------------------------------ - # Warp 0: TMA producer - # ------------------------------------------------------------------ - if warp_idx == cfg.tma_warp_id: - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, cfg.num_ab_stage - ) - for _ in cutlass.range(0, k_cnt, 1, unroll=1): - ab_pipeline.producer_acquire(ab_producer_state) - k_idx = k_base + ab_producer_state.count - bar = ab_pipeline.producer_get_barrier(ab_producer_state) - cute.copy( - tma_atom_a, - tAgA_slice[(None, k_idx)], - tAsA[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - cute.copy( - tma_atom_b, - tBgB_slice[(None, k_idx)], - tBsB[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - cute.copy( - tma_atom_sfa, - tAgSFA_slice[(None, k_idx)], - tAsSFA[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - cute.copy( - tma_atom_sfb, - tBgSFB_slice[(None, k_idx)], - tBsSFB[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - ab_producer_state.advance() - ab_pipeline.producer_tail(ab_producer_state) - - # ------------------------------------------------------------------ - # Warp 1: MMA - # ------------------------------------------------------------------ - if warp_idx == cfg.mma_warp_id: - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - - # The MMA warp joins the allocation barrier but must never allocate. - tmem_alloc.wait_for_alloc() - acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - - sfa_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), - dtype=sSFA.element_type, - ) - tCtSFA = cute.make_tensor( - sfa_tmem_ptr, - blockscaled_utils.make_tmem_layout_sfa( - tiled_mma, - mma_tiler, - cfg.sf_vec_size, - cute.slice_(sfa_smem_layout, (None, None, None, 0)), - ), - ) - sfb_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr - + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) - + tcgen05.find_tmem_tensor_col_offset(tCtSFA), - dtype=sSFB.element_type, - ) - tCtSFB = cute.make_tensor( - sfb_tmem_ptr, - blockscaled_utils.make_tmem_layout_sfb( - tiled_mma, - mma_tiler, - cfg.sf_vec_size, - cute.slice_(sfb_smem_layout, (None, None, None, 0)), - ), - ) - s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_copy_and_partition(sSFA, tCtSFA) - s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_copy_and_partition(sSFB, tCtSFB) - - ab_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, cfg.num_ab_stage - ) - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, cfg.num_acc_stage - ) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - # Acquire and commit unconditionally, including when k_cnt == 0, so the - # accumulator handoff barrier stays balanced on tail tiles. - acc_pipeline.producer_acquire(acc_producer_state) - for k_tile in cutlass.range(0, k_cnt, 1, unroll=1): - ab_pipeline.consumer_wait(ab_consumer_state) - stage_crd = (None, None, None, None, ab_consumer_state.index) - cute.copy(s2t_sfa, tCsSFA_s2t[stage_crd], tCtSFA_s2t) - cute.copy(s2t_sfb, tCsSFB_s2t[stage_crd], tCtSFB_s2t) - # ACCUMULATE=False on the first K tile is what zeroes the - # accumulator; there is no separate TMEM clear. - tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) - mma_crd = (None, None, None, ab_consumer_state.index) - cute.gemm( - tiled_mma, - tCtAcc, - [tCrA[mma_crd], tCtSFA], - [tCrB[mma_crd], tCtSFB], - tCtAcc, - ) - ab_pipeline.consumer_release(ab_consumer_state) - ab_consumer_state.advance() - acc_pipeline.producer_commit(acc_producer_state) - - # ------------------------------------------------------------------ - # Warps 4-7: epilogue - # ------------------------------------------------------------------ - if warp_idx >= cfg.epilogue_warp_ids[0]: - # A power-of-two multiple of 32 columns is required; shared memory - # already pins us to one CTA per SM, so taking the whole array is free. - tmem_alloc.allocate(TMEM_TOTAL_COLS) - tmem_alloc.wait_for_alloc() - acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - - epi_tidx = tidx - cfg.first_epilogue_thread - tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc = t2r_partition( - epi_tidx, tCtAcc_base, cfg - ) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, cfg.num_acc_stage - ) - acc_pipeline.consumer_wait(acc_consumer_state) - - for s in cutlass.range_constexpr(cfg.num_epi_subtiles): - if k_cnt == Int32(0): - # Tail tile or zero-token expert: nothing was accumulated, so the - # fragment is zeroed here and the epilogue runs unchanged. - for v in cutlass.range_constexpr(cute.size(tTR_rAcc)): - tTR_rAcc[v] = Float32(0.0) - else: - cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, 0, s)], tTR_rAcc) - EPILOGUE( - tTR_rAcc, - tTR_cAcc[(None, None, None, 0, s)], - tiled_copy_t2r, - epi_tidx, - s, - tile, - epi_smem, - out, - cfg, - ) - - cute.arch.fence_view_async_tmem_load() - acc_pipeline.consumer_release(acc_consumer_state) - tmem_alloc.relinquish_alloc_permit() - epilogue_barrier.arrive_and_wait() - tmem_alloc.free(acc_tmem_ptr) - - -@cute.jit -def launch_grouped_gemm( - mA: cute.Tensor, - mB: cute.Tensor, - sfa_flat: cute.Tensor, - sfb_flat: cute.Tensor, - offs: cute.Tensor, - out, - stream, - cfg: cutlass.Constexpr, - EPILOGUE: cutlass.Constexpr, - EPI_SMEM_BYTES: cutlass.Constexpr = 0, - VALIDATE_OFFSETS: cutlass.Constexpr = True, -): - """Build the four static TMA descriptors and launch. Grid is data-independent. - - ``mA`` is the GEMM-domain ``(M, K, L)`` operand and ``mB`` the ``(N, K, L)`` - one, both K-major. ``sfa_flat`` / ``sfb_flat`` are the flat blocked E8M0 - buffers; they are retiled here, not by the caller. ``out`` is whatever tuple - of destinations the epilogue expects. - - This is a trace body, not a launcher. Calling it directly retraces the whole - kernel on every invocation -- 130 ms measured at these shapes. A public - launcher must wrap it in its own ``@cute.jit`` entry point taking only the - dynamic tensors, ``cute.compile`` that once behind a ``functools.cache``, and - call the compiled executor (35 us). The Constexpr arguments must not be - passed again to that executor; hand it the dynamic arguments only, or it - raises "cannot be converted to pointer". - - ``VALIDATE_OFFSETS`` emits the device-side precondition check on the offset - *values*, which the host cannot see without synchronizing. Note what it does - and does not buy: ``cute.testing.assert_`` is compiled out entirely unless - ``CUTE_DSL_ENABLE_ASSERTIONS=1`` is set in the environment, and when it does - fire it traps the kernel and leaves the CUDA context unusable - (``unspecified launch failure``) rather than raising cleanly. It is a - debugging aid, not a guardrail; the host validators are the guardrail. - """ - a_dtype = mA.element_type - b_dtype = mB.element_type - sf_dtype = cutlass.Float8E8M0FNU - - gemm_m = cutlass.const_expr(cute.size(mA, mode=[0])) - gemm_k = cutlass.const_expr(cute.size(mA, mode=[1])) - gemm_n = cutlass.const_expr(cute.size(mB, mode=[0])) - l_a = cutlass.const_expr(cute.size(mA, mode=[2])) - l_b = cutlass.const_expr(cute.size(mB, mode=[2])) - num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) - - if cutlass.const_expr(gemm_m % cfg.cta_tile_m != 0): - raise ValueError(f"GEMM M {gemm_m} must be a multiple of {cfg.cta_tile_m}") - if cutlass.const_expr(gemm_n % cfg.cta_tile_n != 0): - raise ValueError(f"GEMM N {gemm_n} must be a multiple of {cfg.cta_tile_n}") - if cutlass.const_expr(gemm_k % cfg.cta_tile_k != 0): - raise ValueError(f"GEMM K {gemm_k} must be a multiple of {cfg.cta_tile_k}") - - mSFA = make_sf_gemm_tensor(sfa_flat, gemm_m, gemm_k, l_a, cfg.sf_vec_size) - mSFB = make_sf_gemm_tensor(sfb_flat, gemm_n, gemm_k, l_b, cfg.sf_vec_size) - - tiled_mma = make_tiled_mma(cfg, a_dtype, b_dtype, sf_dtype) - mma_tiler = cfg.mma_tiler_mnk - cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((*cfg.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) - ) - - a_smem_layout = sm100_utils.make_smem_layout_a( - tiled_mma, mma_tiler, a_dtype, cfg.num_ab_stage - ) - b_smem_layout = sm100_utils.make_smem_layout_b( - tiled_mma, mma_tiler, b_dtype, cfg.num_ab_stage - ) - sfa_smem_layout = blockscaled_utils.make_smem_layout_sfa( - tiled_mma, mma_tiler, cfg.sf_vec_size, cfg.num_ab_stage - ) - sfb_smem_layout = blockscaled_utils.make_smem_layout_sfb( - tiled_mma, mma_tiler, cfg.sf_vec_size, cfg.num_ab_stage - ) - - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - sm100_utils.cluster_shape_to_tma_atom_A(cfg.cluster_shape_mn, tiled_mma.thr_id), - mA, - cute.slice_(a_smem_layout, (None, None, None, 0)), - mma_tiler, - tiled_mma, - cluster_layout_vmnk.shape, - ) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - sm100_utils.cluster_shape_to_tma_atom_B(cfg.cluster_shape_mn, tiled_mma.thr_id), - mB, - cute.slice_(b_smem_layout, (None, None, None, 0)), - mma_tiler, - tiled_mma, - cluster_layout_vmnk.shape, - ) - # The 512-byte scale atom is contiguous; TMA must move it as 8-byte elements. - tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( - sm100_utils.cluster_shape_to_tma_atom_A(cfg.cluster_shape_mn, tiled_mma.thr_id), - mSFA, - cute.slice_(sfa_smem_layout, (None, None, None, 0)), - mma_tiler, - tiled_mma, - cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64, - ) - tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( - sm100_utils.cluster_shape_to_tma_atom_SFB( - cfg.cluster_shape_mn, tiled_mma.thr_id - ), - mSFB, - cute.slice_(sfb_smem_layout, (None, None, None, 0)), - mma_tiler, - tiled_mma, - cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64, - ) - - # Round UP: flooring would hand back fewer bytes than the epilogue asked for - # and its last store would land in sA. - epi_words = cutlass.const_expr(max((EPI_SMEM_BYTES + 3) // 4, 1)) - - @cute.struct - class SharedStorage: - ab_full_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_ab_stage] - ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_ab_stage] - acc_full_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_acc_stage] - acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, cfg.num_acc_stage] - tmem_holding_buf: cutlass.Int32 - sEpi: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, epi_words], 128] - sA: cute.struct.Align[ - cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024 - ] - sB: cute.struct.Align[ - cute.struct.MemRange[b_dtype, cute.cosize(b_smem_layout.outer)], 1024 - ] - sSFA: cute.struct.Align[ - cute.struct.MemRange[sf_dtype, cute.cosize(sfa_smem_layout)], 1024 - ] - sSFB: cute.struct.Align[ - cute.struct.MemRange[sf_dtype, cute.cosize(sfb_smem_layout)], 1024 - ] - - smem_bytes = cutlass.const_expr(SharedStorage.size_in_bytes()) - if cutlass.const_expr(smem_bytes > SMEM_CAPACITY_BYTES): - # The struct's 1024-byte operand alignment costs a little more than the - # config's arithmetic, so check the real number rather than the estimate. - raise ValueError( - f"shared memory request {smem_bytes} B exceeds the sm_100 capacity " - f"{SMEM_CAPACITY_BYTES} B: lower num_ab_stage (currently " - f"{cfg.num_ab_stage}) or the epilogue's {EPI_SMEM_BYTES} B request" - ) - - if cutlass.const_expr(cfg.ragged_axis is RaggedAxis.M): - grid = (gemm_m // cfg.cta_tile_m, gemm_n // cfg.cta_tile_n, 1) - else: - grid = (gemm_m // cfg.cta_tile_m, gemm_n // cfg.cta_tile_n, num_groups) - - grouped_gemm_kernel( - tiled_mma, - tma_atom_a, - tma_tensor_a, - tma_atom_b, - tma_tensor_b, - tma_atom_sfa, - tma_tensor_sfa, - tma_atom_sfb, - tma_tensor_sfb, - offs, - out, - a_smem_layout, - b_smem_layout, - sfa_smem_layout, - sfb_smem_layout, - cfg, - SharedStorage, - EPILOGUE, - EPI_SMEM_BYTES, - VALIDATE_OFFSETS, - ).launch( - grid=grid, - block=(cfg.threads, 1, 1), - cluster=(*cfg.cluster_shape_mn, 1), - smem=smem_bytes, - stream=stream, - ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py deleted file mode 100644 index 1e2d7332cc..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_epilogue.py +++ /dev/null @@ -1,448 +0,0 @@ -# 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. - -"""Device-side epilogue primitives shared by the MXFP8 grouped-MLP kernels. - -These are the pieces the three grouped kernels (FC1 GEMM + SwiGLU + dual quant, -FC2 dgrad + dSwiGLU + dual quant, and grouped wgrad) have in common: blocked -scale addressing, the RCEIL E8M0 conversion, NaN-propagating packed amax, and -the gated-activation policy. They are lifted from the activation-only gated -kernel, whose numerics are already validated bitwise against the standalone -quantizers. - -NUMERICAL PRECONDITION (load-bearing, read before reusing anything here): - - :func:`float_to_e8m0` is integer bit math whose rounding constant assumes - its input is a BF16 value widened to FP32. It is exact for a BF16 amax and - is NOT exact for a general FP32 amax -- for example FP32 0x40600001 gives - 120 where the canonical conversion gives 121. - -A fused GEMM epilogue holds FP32 accumulators, so it must round to BF16 *before* -taking the amax, which is what the kernel contract already requires at every -activation boundary. Do not "optimize" that rounding away while keeping this -conversion; if an FP32-amax path is ever wanted, use the canonical -``cute_utils.compute_scale_rceil`` (a real ``cvt.rp`` instruction) instead. - -Scale semantics follow torchao #4725: RCEIL with saturation disabled, a -non-finite amax yielding scale byte 255 and an all-NaN block, a zero block -yielding scale byte 0 (which dequantizes to 2^-127, not 1.0), and -``inv_scale = e8m0(254 - byte)``. -""" - -import cutlass -import cutlass.cute as cute -from cutlass import Float32, Int32 -from cutlass._mlir.dialects import arith as mlir_arith -from cutlass._mlir.dialects import llvm -from cutlass.cutlass_dsl import T, dsl_user_op - -__all__ = [ - "SCALE_BLOCK", - "SCALE_TILE_ROWS", - "SCALE_TILE_COLS", - "SCALE_TILE_BYTES", - "blocked_scale_idx", - "float_to_e8m0", - "e8m0_reciprocal_bf16", - "max_nan_bf16x2", - "abs_max_nan_bf16x2", - "fold_amax", - "pack_bf16x2", - "bf16x2_lo_to_f32", - "bf16x2_hi_to_f32", - "mul_cvt_2x", - "prmt_even", - "prmt_odd", - "sigmoidf", - "silu_pair", - "validate_group_offsets_device", -] - -# MXFP8 scaling block: 32 values share one E8M0 scale. -SCALE_BLOCK = 32 -# tcgen05 blocked scale tile geometry: 128 scale rows x 4 scale columns = 512 B. -SCALE_TILE_ROWS = 128 -SCALE_TILE_COLS = 4 -SCALE_TILE_BYTES = SCALE_TILE_ROWS * SCALE_TILE_COLS -# Every TMA-accessed shared-memory buffer must be 128-byte aligned. -TMA_SHMEM_ALIGNMENT = 128 - - -def blocked_scale_idx(row, scale_col, num_scale_col_blocks): - """Flat index of one scale byte in the tcgen05 blocked (128x4) layout. - - The logical ``[rows, cols/32]`` scale matrix is stored as 512-byte tiles of - 128 rows x 4 scale columns (cuBLAS "128x4 block scaling factors layout"), - tiles ordered ``row_block * num_scale_col_blocks + col_block``. - ``num_scale_col_blocks`` is ``ceil_div(num_scale_cols, 4)``. - - Coordinates are ABSOLUTE, never per-group. That is legal for both - orientations this family emits: - - * rowwise, where the ragged axis is the scale row -- because every expert - row count is a multiple of 128, no 128-row tile straddles a group - boundary, so per-group blocking and whole-matrix blocking are the same - bytes; - * columnwise, where the ragged axis is the scale column -- because this - family defines those buffers as whole-matrix ``to_blocked`` (see the - kernel contract 4.2.1). Note this deliberately differs from torchao's - ``triton_mx_block_rearrange_2d_K_groups``, which pads per group. - - For columnwise scales pass transposed coordinates (feature index as ``row``, - row-block index as ``scale_col``). - """ - return ( - ((row >> 7) * num_scale_col_blocks + (scale_col >> 2)) * SCALE_TILE_BYTES - + (row & 31) * 16 - + ((row >> 5) & 3) * 4 - + (scale_col & 3) - ) - - -@dsl_user_op -def _bitcast_i32_to_f32(val: Int32, *, loc=None, ip=None) -> Float32: - """Bitcast int32 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) - - -# Each packed-bf16x2 op below is written out explicitly rather than produced by a -# factory: the DSL keys its compile cache on function identity and name, so ops -# sharing a `__name__` are a cache hazard. -# -# The `.NaN` max variants match the standalone quantizers' amax reduction, which -# propagates NaN; a plain max would return the non-NaN operand and silently -# rescue a block that must be invalidated. -@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, - ) - ) - - -@cute.jit -def fold_amax(am: Int32) -> Int32: - """Reduce a packed bf16x2 amax word to bf16 amax bits in [15:0]. - - The input's per-lane sign bits are junk (see :func:`abs_max_nan_bf16x2`); - mask them, then fold the two lanes with the NaN-propagating max. - """ - am = am & Int32(0x7FFF7FFF) - am = max_nan_bf16x2(am, am >> 16) - return am & Int32(0xFFFF) - - -@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. Saturation to +/-448 comes from - this conversion; do not add an explicit clamp. - """ - 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, round-to-nearest-even. - - This is the mandatory "truncate to BF16 before any amax" step for a GEMM - epilogue holding FP32 accumulators; see the module docstring. - """ - 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 RCEIL scale byte for a non-negative BF16 amax, as f32 bits. - - Pass the RAW amax. The division by 448 is folded into the constants -- the - ``- 8`` shifts the exponent by 256 and the ``+ 0x1F0000`` mantissa offset - supplies the remaining 1.75 factor (448 = 256 * 1.75) together with the - RCEIL round-up carry. Pre-dividing by 448 before calling this double-counts - the division and yields a scale 8-9 codes too low. - - Finite: the RCEIL mantissa-carry path, matching what ``cvt.rp.ue8m0x2.f32`` - (no ``.satfinite``) emits for ``amax / 448``. Non-finite: a 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. - - Exact only for a BF16-valued amax -- see the module docstring. - """ - 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 e8m0_reciprocal_bf16(e: Int32) -> Int32: - """Inverse scale as bf16 bits, matching ``ue8m0(254 - scale_byte)``. - - The quantization multiply in :func:`mul_cvt_2x` is bf16x2, not f32, so the - reciprocal is synthesized directly in bf16: 2^(127 - e) over the normal - range (a byte-0 block from a zero or tiny amax descales by 2^127), and NaN - for an invalidated block (byte 255) so every element quantizes to the E4M3 - NaN code. - - This is NOT a general ``ue8m0`` helper: byte 254 yields +0.0 rather than - 2^-127, and bytes above 254 go negative. That is safe here only because a - scale byte produced by this family can never exceed 247 (amax is divided by - 448 first). Do not reuse it to dequantize externally produced scale bytes. - """ - 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 / dSwiGLU policy for a pair of elements. - - ``x`` is the gate (activation input), ``lin`` the up (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 -> dGate - out_gate = act * g -> dUp - - Returns f32 ``(out_act0, out_act1, out_gate0, out_gate1)``, the gate pair - zero in forward. Callers MUST round to BF16 immediately, before any amax, - caching, or quantization -- both because the kernel contract defines - correctness at that boundary and because :func:`float_to_e8m0` requires it. - - Pass this as a Constexpr kernel parameter. It must stay a module-level - function: the DSL keys its compile cache on the function object, so a lambda - or closure built per call misses the cache and recompiles every launch. - """ - one = 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, Float32(0.0), Float32(0.0) - - -@cute.jit -def validate_group_offsets_device(offs: cute.Tensor, allocated_rows: Int32): - """Device-side precondition check on the exclusive-end group offsets. - - Checks what the host cannot see without synchronizing: every per-expert row - count is a nonnegative multiple of 128, the offsets are nondecreasing, and - the active row count does not exceed the allocation. - - DO NOT RELY ON THIS AS A GUARDRAIL. ``cute.testing.assert_`` is compiled out - unless ``CUTE_DSL_ENABLE_ASSERTIONS=1``, so in a default build this function - is a no-op -- measured directly: an always-false assertion in four - placements (plain, warp-0, elected, block-0-elected) lets the kernel run to - completion and write its output. torchao's own ``validate_group_sizes`` has - the same property. Even with assertions enabled the failure mode is the - message followed by ``unspecified launch failure``, i.e. a dead CUDA - context, not a catchable error. - - The real enforcement of the 128-multiple precondition is the host-side - metadata validation in ``grouped_mlp_validation`` at the custom-op boundary, - which raises ``ValueError`` before any launch. This function is a debugging - aid for assertion-enabled builds. It matters that the distinction is - explicit: with malformed offsets the ragged-K path (Kernel C) silently - returns a WRONG weight gradient rather than crashing, so "it did not fault" - is not evidence that the offsets were valid. - """ - num_groups = offs.shape[0] - prev = Int32(0) - for i in range(num_groups): - end = offs[i] - size = end - prev - cute.testing.assert_(size >= 0, "Group offsets must be nondecreasing") - cute.testing.assert_(size % 128 == 0, "Group sizes must be multiples of 128") - prev = end - cute.testing.assert_( - prev <= allocated_rows, - "Active row count offsets[-1] must not exceed the allocated row count", - ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py index c608d78c68..744e32f16e 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. -"""Public custom-op surface for the MXFP8 routed-expert grouped-MLP kernels. +"""Custom-op surface for the MXFP8 routed-expert grouped-MLP kernels. Three ops, one per fused kernel: @@ -21,11 +21,16 @@ a class of bug that matters here because several outputs are column-major and a row-major fake would silently change what ``torch.compile`` traces. -Layout, alignment and numerical semantics are defined by the kernel contract; -the load-bearing preconditions (every per-expert row count a multiple of 128, -exact strides, blocked E8M0 scales in the tcgen05 128x4 layout) are enforced by -``grouped_mlp_validation`` before any launch, because violating them otherwise -corrupts the CUDA context rather than raising. +Every op validates its inputs through a shared ``_validate_*_inputs`` helper +that also backs ``register_fake``, so ``torch.compile`` rejects an unsupported +call at graph capture rather than mid-training from a compiled region. The +checks are metadata-only (no host/device sync). The per-expert offset VALUES +are a documented caller invariant -- see ``grouped_mlp_validation`` for what +is and is not enforced. + +The user-facing functional wrappers live in +``torchao.prototype.moe_training.mxfp8_grouped_mlp``; importing that module (or +this one) registers the ops. """ from typing import Tuple @@ -33,23 +38,46 @@ import torch from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( + GROUP_ALIGNMENT, + SCALE_BLOCK_SIZE, + _is_fake, blocked_scale_numel, validate_allocated_rows, validate_blocked_scales, + validate_destination, validate_feature_dims, validate_group_offsets, validate_grouped_operand, ) __all__ = [ - "mxfp8_grouped_gemm_swiglu_fwd", - "mxfp8_grouped_gemm_dswiglu_bwd", - "mxfp8_grouped_gemm_wgrad", + # Validation surface re-exported for the kernel launchers, which check the + # caller-allocated destinations that the ops (allocating their own) do not. + "GROUP_ALIGNMENT", + "SCALE_BLOCK_SIZE", + "_is_fake", + "blocked_scale_numel", + "validate_allocated_rows", + "validate_blocked_scales", + "validate_destination", + "validate_feature_dims", + "validate_group_offsets", + "validate_grouped_operand", ] _E4M3 = torch.float8_e4m3fn _E8M0 = torch.float8_e8m0fnu -_SCALE_BLOCK = 32 +_SCALE_BLOCK = SCALE_BLOCK_SIZE + + +def _require_cuda_device(device: torch.device, name: str) -> None: + """All operands must live on one CUDA device; CPU tensors get a clean error + here instead of a launcher failure.""" + if device.type != "cuda": + raise ValueError( + f"{name} must be a CUDA tensor, got device {device}; these kernels " + "run only on CUDA SM100 devices" + ) def _empty_blocked_scales( @@ -109,9 +137,20 @@ def _validate_swiglu_fwd_inputs( hidden = two_hidden // 2 device = x_q.device + _require_cuda_device(device, "x_q") validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) validate_allocated_rows(rows) - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + # The epilogue computes z/h element offsets as row * 2F + col in int32. + # Checked here as well as in the launcher so torch.compile tracing rejects + # the shape at graph capture instead of mid-training. + if rows * two_hidden >= 2**31: + raise ValueError( + f"R * 2F = {rows * two_hidden} does not fit the epilogue's int32 " + "element index; the store address arithmetic would wrap" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) validate_grouped_operand( x_q, name="x_q", @@ -154,6 +193,35 @@ def _mxfp8_grouped_gemm_swiglu_fwd( w13_t_sf: torch.Tensor, offsets: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FC1 grouped GEMM + SwiGLU + dual MXFP8 RCEIL quantization, one launch. + + Inputs (all CUDA, same device; prequantized outside, never requantized here): + x_q E4M3 ``[R, D]`` stride ``(D, 1)``, rowwise 1x32 quantized. + x_sf flat blocked E8M0 scales for logical ``[R, D/32]`` + (``round_up(R,128) * round_up(D/32,4)`` bytes). + w13_t_q E4M3 ``[G, D, 2F]`` stride ``(2*D*F, 1, D)`` -- the quantized view + of ``w13_bf16.reshape(G, 2F, D).transpose(-2, -1)``; the 2F axis + is ELEMENT-interleaved gate/up (gate even, up odd). + w13_t_sf blocked E8M0 ``[G, round_up(2F,128) * round_up(D/32,4)]``. + offsets int32 CUDA ``[G]`` exclusive group end rows; every per-expert row + count must be a nonnegative multiple of 128 and + ``offsets[-1] <= R`` (caller invariant, see grouped_mlp_validation). + + Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``: + z_bf16 BF16 ``[R, F, 2]`` stride ``(2F, 2, 1)`` -- the pre-activation + rounded to BF16 BEFORE SwiGLU; gate at index 0, up at index 1; + saved for backward and consumed unchanged by the dswiglu op. + h_row_q E4M3 ``[R, F]`` stride ``(F, 1)``; h_row_sf blocked scales for + logical ``[R, F/32]``. + h_col_q E4M3 ``[R, F]`` COLUMN-MAJOR stride ``(1, R)``; h_col_sf + whole-matrix blocked scales for logical ``[F, R/32]``. + + ``h = silu(gate) * up`` is evaluated once from the BF16-rounded z and + rounded to BF16 before BOTH quantizers. Inactive tail rows + ``[offsets[-1], R)`` of every output are written as zero bytes. ``R == 0`` + returns empty outputs without launching. A zero-token expert contributes no + rows. G == 0 is rejected. + """ rows, _model_dim, hidden, _groups = _validate_swiglu_fwd_inputs( x_q, x_sf, w13_t_q, w13_t_sf, offsets ) @@ -227,9 +295,18 @@ def _validate_dswiglu_bwd_inputs(do_q, do_sf, w2_q, w2_sf, z_bf16, offsets): ) device = do_q.device + _require_cuda_device(device, "do_q") validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) validate_allocated_rows(rows) - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + # Same int32 element-index bound as the forward: dz is [R, 2F]. + if rows * 2 * hidden >= 2**31: + raise ValueError( + f"R * 2F = {rows * 2 * hidden} does not fit the epilogue's int32 " + "element index; the store address arithmetic would wrap" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) validate_grouped_operand( do_q, name="do_q", @@ -282,6 +359,36 @@ def _mxfp8_grouped_gemm_dswiglu_bwd( z_bf16: torch.Tensor, offsets: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FC2 dgrad grouped GEMM + dSwiGLU + dual MXFP8 RCEIL quantization, one launch. + + Inputs (all CUDA, same device; GEMM operands prequantized outside): + do_q E4M3 ``[R, D]`` stride ``(D, 1)``, rowwise 1x32 quantized + FC2 output-gradient. + do_sf flat blocked E8M0 scales for logical ``[R, D/32]``. + w2_dgrad_q E4M3 ``[G, D, F]`` stride ``(D*F, 1, D)`` -- the dgrad + orientation of w2. + w2_dgrad_sf blocked E8M0 ``[G, round_up(F,128) * round_up(D/32,4)]`` for + per-expert logical ``[F, D/32]``. + z_bf16 BF16 ``[R, F, 2]`` stride ``(2F, 2, 1)`` -- the EXACT saved + output of the swiglu_fwd op (not recomputed). Rows past + ``offsets[-1]`` are never read. + offsets int32 CUDA ``[G]`` exclusive group ends (same contract as the + forward op). + + Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)``: + dz_row_q E4M3 ``[R, 2F]`` stride ``(2F, 1)`` with ELEMENT-interleaved + ``[dgate_0, dup_0, ...]`` channels; dz_row_sf blocked scales for + logical ``[R, 2F/32]``. + dz_col_q E4M3 ``[R, 2F]`` COLUMN-MAJOR stride ``(1, R)``, same logical + order; dz_col_sf whole-matrix blocked scales for logical + ``[2F, R/32]``. + + ``dh`` is rounded to BF16 before dSwiGLU; ``dgate = dh * up * dsilu`` and + ``dup = dh * silu`` are each rounded to BF16 before interleaving, and both + quantizers consume the same BF16 ``dz``. Tail rows of every output are + written as zero bytes; ``R == 0`` returns empty outputs without launching; + G == 0 is rejected. + """ rows, _model_dim, hidden, _groups = _validate_dswiglu_bwd_inputs( do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets ) @@ -327,11 +434,23 @@ def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): raise ValueError( f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" ) - groups = offsets.numel() + groups = offsets.numel() if isinstance(offsets, torch.Tensor) else 0 device = dy_col_q.device + _require_cuda_device(device, "dy_col_q") validate_allocated_rows(rows) - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) + # N and K are the GEMM's free axes and are tiled with no tail path. The + # launcher checks this too, but it must ALSO live here because this + # function backs register_fake: without it, torch.compile traces a shape + # the real op rejects and the ValueError fires from inside a compiled + # region mid-training instead of at graph capture, defeating the caller's + # fallback predicate. + for name, value in (("dy_col_q's N", out_features), ("x_col_q's K", in_features)): + if value <= 0 or value % 128 != 0: + raise ValueError(f"{name} must be a positive multiple of 128, got {value}") + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) # Both operands are column-major so their transposes are free. validate_grouped_operand( dy_col_q, @@ -374,6 +493,29 @@ def _mxfp8_grouped_gemm_wgrad( x_col_sf: torch.Tensor, offsets: torch.Tensor, ) -> torch.Tensor: + """Grouped MXFP8 weight-gradient GEMM: ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + Inputs (all CUDA, same device): + dy_col_q E4M3 logical ``[R, N]`` COLUMN-MAJOR stride ``(1, R)`` -- the + columnwise (32x1) quantized output of the swiglu_fwd or + dswiglu_bwd op. + dy_col_sf whole-matrix blocked E8M0 scales for logical ``[N, R/32]``. + x_col_q E4M3 logical ``[R, K]`` stride ``(1, R)``; x_col_sf likewise + for logical ``[K, R/32]``. + offsets int32 CUDA ``[G]`` exclusive group ends over the shared row + (contraction) axis. + + Both scale buffers must be WHOLE-MATRIX ``to_blocked``, not torchao's + per-group K-groups rearrangement: the two orderings have identical byte + counts and differ whenever N > 128, so no length check can tell them apart, + and mixing them silently produces a block-permuted (wrong) ``dw``. + + Returns contiguous BF16 ``dw [G, N, K]`` with FP32 accumulation. Reusable + with no mode flag: FC1 wgrad passes ``N=2F, K=D``; FC2 wgrad passes ``N=D, + K=F``. Every element of ``dw`` is written on every call; a zero-token + expert (and the ``R == 0`` / all-empty cases) yields an all-zero slice. + G == 0 is rejected. + """ rows, out_features, in_features, groups = _validate_wgrad_inputs( dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets ) @@ -408,44 +550,3 @@ def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): dtype=torch.bfloat16, device=dy_col_q.device, ) - - -# -------------------------------------------------------------------------- -# Public wrappers -# -------------------------------------------------------------------------- - - -def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_t_q, w13_t_sf, offsets): - """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. - - Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``. ``z_bf16`` is the - BF16 pre-activation saved for backward; pass it unchanged to - :func:`mxfp8_grouped_gemm_dswiglu_bwd`. - """ - return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( - x_q, x_sf, w13_t_q, w13_t_sf, offsets - ) - - -def mxfp8_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets -): - """FC2 dgrad grouped GEMM + dSwiGLU + rowwise/columnwise MXFP8 quantization. - - Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)`` with gate/up gradients - element-interleaved along the 2F axis. - """ - return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets - ) - - -def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): - """Grouped MXFP8 weight-gradient GEMM, returning BF16 ``[G, N, K]``. - - Used once for FC1 (``N=2F, K=D``) and once for FC2 (``N=D, K=F``). An expert - with zero rows yields an all-zero slice. - """ - return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( - dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets - ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py index c9671d8ed6..40269c7275 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py @@ -1,22 +1,32 @@ +# 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. + """Host-side precondition validation for the MXFP8 routed-expert grouped-MLP kernels. The grouped MXFP8 kernels require every per-expert row count to be a multiple of 128. That is not a convenience: the tcgen05 blocked scale layout permutes in 128-row tiles, so a group boundary off a 128 multiple splits a tile and the -blocked buffer for the group no longer matches what the GEMM reads. Feeding such -offsets to the existing CuTe DSL quantizer produces a device-side assertion, a -wrong-sized scale buffer, and an unusable CUDA context rather than a clean error. +blocked buffer for the group no longer matches what the GEMM reads. Everything here is metadata-only so it costs no host/device synchronization and -stays traceable under torch.compile. The per-expert counts live in device memory -and are validated on device by the kernels themselves; set -TORCHAO_MXFP8_VALIDATE_OFFSETS=1 to additionally check them on the host while -debugging, at the cost of a D2H copy. +stays traceable under torch.compile. The per-expert offset VALUES live in device +memory and are a documented caller invariant, not something these checks can +enforce: reading them requires a D2H sync, so they are validated only by the +opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS=1 debugging path below, and there is NO +device-side enforcement in a default build. Malformed offsets are not guaranteed +to fault, either -- the ragged-K weight-gradient kernel in particular can return +a wrong result from a clean-looking launch. Callers (and any future selection +predicate) must guarantee the offsets contract, e.g. by padding every expert +group to a multiple of 128 rows at dispatch time. Checks raise ValueError rather than asserting, so `python -O` cannot strip them. """ import os +from typing import Optional import torch @@ -65,20 +75,32 @@ def validate_group_offsets( *, num_groups: int, allocated_rows: int, + device: Optional[torch.device] = None, name: str = "offsets", ) -> None: - """Validate the exclusive-end group offsets tensor. + """Validate the exclusive-end group offsets tensor's metadata. - Metadata is always checked. The offset *values* are checked only when - host_offsets_validation_enabled(), because reading them synchronizes; the - kernels assert the same invariants on device on every launch. + Metadata is always checked, including that at least one expert group + exists. The offset *values* are checked only when + host_offsets_validation_enabled(), because reading them forces a D2H sync; + otherwise they are a documented caller invariant with no default-build + enforcement anywhere (see the module docstring). """ if not isinstance(offsets, torch.Tensor): raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") + if num_groups < 1: + raise ValueError( + f"{name} must describe at least one expert group, got G={num_groups}" + ) if offsets.dtype != torch.int32: raise ValueError(f"{name} must be int32, got {offsets.dtype}") if not offsets.is_cuda: raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") + if device is not None and offsets.device != device: + raise ValueError( + f"{name} must be on {device}, got {offsets.device}; all operands and " + "destinations must share one CUDA device" + ) if offsets.ndim != 1: raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") if offsets.numel() != num_groups: diff --git a/torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py b/torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py deleted file mode 100644 index 0050c4033d..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/kernel_wgrad.py +++ /dev/null @@ -1,347 +0,0 @@ -# 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. - -"""Kernel C: grouped MXFP8 weight-gradient GEMM (``mxfp8_grouped_gemm_wgrad``). - -One Definition covers both call sites -- FC1 (``N = 2F``, ``K = D``) and FC2 -(``N = D``, ``K = F``) -- with no mode flag; the shapes come from the tensors. - -Per expert ``g`` over rows ``[offsets[g-1], offsets[g])``:: - - dw[g] = dequant(dy_col[rows]).T @ dequant(x_col[rows]) - -FP32 accumulation, BF16 output. Neither transpose is materialized: both operands -arrive logically ``[R, N]`` / ``[R, K]`` with stride ``(1, R)``, which *is* a -K-contiguous row-major ``[N, R]`` / ``[K, R]``, so the free transpose is a -restride on the host and the ragged axis lands on the GEMM's contraction. The -expert is then an integer K-tile index base rather than a per-expert TMA -descriptor, exact because every per-expert row count is a multiple of -``cta_tile_k``. - -Two things about this kernel that are easy to get wrong: - -*Both scale buffers are WHOLE-MATRIX* ``to_blocked``, not torchao's per-group -K-groups form (``triton_mx_block_rearrange_2d_K_groups``). The two orderings -differ whenever ``N > 128`` -- whole-matrix orders blocked tiles by -``row_block * ncb_total + col_block``, per-group by ``row_block * ncb_g + -col_block`` within each group -- so feeding a per-group buffer here produces a -*block-permuted* ``dw``, which is large and structured and reads like a GEMM bug -rather than like a layout bug. The producers of these buffers are kernels A and -B in this same family, which emit the whole-matrix form. - -*The epilogue never predicates its store.* A zero-token expert arrives with -``k_cnt == 0``, the core hands the epilogue a zeroed register fragment, and the -unmodified store path writes the all-zero ``dw[g]`` the contract requires. The -grid enumerates every ``(tile_m, tile_n, expert)``, so every element of ``dw`` is -written on every call with no memset. There is also no gmem input load in this -epilogue, so the ``k_cnt``-gated-load half of the tail rule has nothing to cover -here. - -No output TMA and no epilogue shared memory: one row per thread means the 32 -FP32 accumulators of a subtile are 32 contiguous BF16 values == 64 naturally -aligned contiguous bytes of ``dw``, so a direct vectorized ``STG`` is both -correct and fully sector-efficient. -""" - -import functools -from typing import Tuple - -import cutlass -import cutlass.cute as cute -import torch -from cutlass import Int32 -from cutlass.cute.runtime import from_dlpack - -from torchao.prototype.moe_training.kernels.mxfp8.grouped_gemm_config import ( - SF_VEC_SIZE, - WGRAD_CONFIG, -) -from torchao.prototype.moe_training.kernels.mxfp8.grouped_gemm_core import ( - activation_gemm_view, - launch_grouped_gemm, -) -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_epilogue import ( - pack_bf16x2, -) -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( - validate_allocated_rows, - validate_blocked_scales, - validate_destination, - validate_group_offsets, - validate_grouped_operand, -) - -__all__ = ["bf16_store_epilogue", "launch_grouped_gemm_wgrad"] - -_E4M3 = torch.float8_e4m3fn -_BF16 = torch.bfloat16 -# Widest vector the epilogue's store is allowed to assume. The run is 64 B and -# 64-B aligned, so this only has to be a divisor of that. -_VEC_BYTES = 16 - - -@cute.jit -def _store_bf16_run(dst: cute.Tensor, elem_offset: Int32, words: cute.Tensor): - """Store a run of packed ``bf16x2`` words at a BF16 element offset. - - ``dst`` is any BF16 destination; the run is reinterpreted as Int32, so the - element offset must be even. Every caller's offset is a multiple of 32 - elements (see :func:`bf16_store_epilogue`), which also makes the address - 64-byte aligned and the copy four ``STG.128``. - """ - cute.autovec_copy( - words, - cute.make_tensor( - ( - cute.recast_ptr(dst.iterator, dtype=Int32) + (elem_offset >> Int32(1)) - ).align(_VEC_BYTES), - cute.make_layout(cute.size(words)), - ), - ) - - -@cute.jit -def bf16_store_epilogue( - tTR_rAcc, - tTR_cAcc_s, - tiled_copy_t2r, - epi_tidx, - subtile_idx: cutlass.Constexpr, - tile, - epi_smem, - out, - cfg: cutlass.Constexpr, -): - """Round the FP32 accumulator subtile to BF16 and store it. No quantization. - - ``out`` is ``(mDw,)`` with ``mDw`` the ``(N, K, G)`` view of the contiguous - ``[G, N, K]`` destination. - - The register-to-column map is taken from ``tTR_cAcc``, whose column - coordinates fold to Python ints at trace time, and the run is required to be - contiguous and increasing. That check is what licenses both the ``pack`` - pairing and the single vectorized store: if a thread owned more than one row - of the epilogue tile its columns would repeat instead of forming a run, so - this also establishes the one-row-per-thread property the store address - assumes, rather than trusting it. - """ - num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) - cols = [] - for v in cutlass.range_constexpr(num_acc): - cols.append(tTR_cAcc_s[v][1]) - frag_col = cols[0] - if cutlass.const_expr( - num_acc % 2 != 0 - or not all(isinstance(c, int) for c in cols) - or tuple(cols) != tuple(range(frag_col, frag_col + num_acc)) - ): - raise ValueError( - f"the wgrad epilogue needs an even, contiguous, increasing column run " - f"per thread to pack and store BF16 pairs, but tTR_cAcc gave {cols}" - ) - - # cvt.rn.bf16x2.f32 packs the second source into the low half, so column - # frag_col + 2j lands at the lower address -- row-major order for dw. - words = cute.make_rmem_tensor((num_acc // 2,), Int32) - for j in cutlass.range_constexpr(num_acc // 2): - words[j] = pack_bf16x2(tTR_rAcc[2 * j + 1], tTR_rAcc[2 * j]) - - # Address from the destination's real strides, which are static ints here, - # rather than from its extents: the only layout property the store actually - # needs is that the K axis is contiguous. - gDw = out[0] - strides = gDw.stride - if cutlass.const_expr(strides[1] != 1): - raise ValueError( - f"the wgrad epilogue stores a contiguous run along K, so dw's K " - f"stride must be 1, got layout {gDw.layout}" - ) - # Row from tTR_cAcc, not from epi_tidx: the contiguity check above proves - # the fragment is one row, and this is the coordinate that names it. - row = tile.row_base + tTR_cAcc_s[0][0] - elem = ( - row * Int32(strides[0]) - + (tile.col_base + Int32(frag_col)) - + tile.expert * Int32(strides[2]) - ) - _store_bf16_run(gDw, elem, words) - - -@cute.jit -def _wgrad_entry(mA, mB, sfa, sfb, offs, mDw, stream): - """Trace entry point: dynamic tensors only, config and epilogue closed over. - - ``launch_grouped_gemm`` is a trace body, so calling it directly retraces the - whole kernel on every launch. Everything Constexpr is bound here so that - :func:`cute.compile` can hand back an executor that takes only these - arguments; passing a Constexpr to that executor raises "cannot be converted - to pointer". - """ - launch_grouped_gemm( - mA, - mB, - sfa, - sfb, - offs, - (mDw,), - stream, - WGRAD_CONFIG, - bf16_store_epilogue, - ) - - -@functools.cache -def _wgrad_executor_slot(key: Tuple) -> list: - """One memo slot per compiled shape. - - The executor cannot be built from ``key`` alone: the shared core needs static - shapes (the blocked scale-factor layout and the grid are both built from - them), so a symbolic ``cute.sym_int`` compile is not available and the first - real call's tensors are what gets compiled. ``functools.cache`` therefore - keys the slot and the caller fills it once. - """ - return [] - - -def _validate_wgrad_operands(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw): - """Host-only precondition check. Metadata and pointers, never offset values. - - The custom op validates its own inputs, but this launcher is also reachable - directly from the DSL and it owns the destination, which the op does not - check. Every gate here is metadata-derivable, so it costs no synchronization - and stays traceable. - """ - if dy_col_q.ndim != 2 or x_col_q.ndim != 2: - raise ValueError( - "dy_col_q and x_col_q must be 2D logical [R, N] and [R, K], got " - f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" - ) - rows, out_features = dy_col_q.shape - x_rows, in_features = x_col_q.shape - if x_rows != rows: - raise ValueError( - f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" - ) - groups = offsets.numel() - device = dy_col_q.device - - validate_allocated_rows(rows) - # N and K are the GEMM's two free axes; both are tiled with no tail path, so - # reject a non-multiple here rather than from inside the trace. - for name, value, tile in ( - ("dy_col_q's N", out_features, WGRAD_CONFIG.cta_tile_m), - ("x_col_q's K", in_features, WGRAD_CONFIG.cta_tile_n), - ): - if value <= 0 or value % tile != 0: - raise ValueError( - f"{name} must be a positive multiple of {tile}, got {value}" - ) - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) - # Column-major, so the transposes below are free restrides. - validate_grouped_operand( - dy_col_q, - name="dy_col_q", - shape=(rows, out_features), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - x_col_q, - name="x_col_q", - shape=(rows, in_features), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - dy_col_sf, - name="dy_col_sf", - logical_rows=out_features, - logical_cols=rows // SF_VEC_SIZE, - device=device, - ) - validate_blocked_scales( - x_col_sf, - name="x_col_sf", - logical_rows=in_features, - logical_cols=rows // SF_VEC_SIZE, - device=device, - ) - # validate_blocked_scales checks dtype, length and device but not the - # pointer, and both scale buffers are TMA operands: a contiguous view with a - # storage offset can be 2-byte aligned. Only the launcher promises the TMA - # alignment, so only the launcher can require it. - for name, buf in (("dy_col_sf", dy_col_sf), ("x_col_sf", x_col_sf)): - if buf.data_ptr() % 32 != 0: - raise ValueError( - f"{name} must be 32-byte aligned for its TMA descriptor, but its " - f"data pointer is {buf.data_ptr() % 32} bytes past an aligned " - "address" - ) - validate_destination( - dw, - name="dw_bf16", - shape=(groups, out_features, in_features), - stride=(out_features * in_features, in_features, 1), - dtype=_BF16, - device=device, - ) - # The epilogue computes its destination element index in Int32. - if groups * out_features * in_features >= 2**31: - raise ValueError( - f"dw_bf16 has {groups * out_features * in_features} elements, which " - "does not fit the epilogue's int32 element index; the store address " - "arithmetic would wrap" - ) - return rows, out_features, in_features, groups - - -def launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw): - """Grouped MXFP8 wgrad into a caller-allocated BF16 ``[G, N, K]`` destination. - - Inputs are the columnwise-quantized outputs of kernels A and B (or of the - standalone 32x1 quantizer): ``dy_col_q`` E4M3 logical ``[R, N]`` stride - ``(1, R)`` with ``dy_col_sf`` blocked for logical ``[N, R/32]``, and - ``x_col_q`` / ``x_col_sf`` likewise for ``[R, K]``. ``offsets`` is the int32 - CUDA ``[G]`` vector of exclusive group ends and is never read on the host. - - Every element of ``dw`` is written, including the all-zero slice of a - zero-token expert. - """ - import cuda.bindings.driver as cuda - - rows, out_features, in_features, groups = _validate_wgrad_operands( - dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw - ) - if rows == 0: - # Every expert has zero rows, so every slice is the zero matrix. Unlike - # the other two kernels the destination is NOT empty here, and the - # contraction is, so this cannot be expressed as a launch. - dw.zero_() - return - - stream = cuda.CUstream(int(torch.cuda.current_stream().cuda_stream)) - args = ( - # The free transpose: logical [R, N] stride (1, R) IS a K-contiguous - # [N, R], so both operands become ordinary K-major GEMM operands and the - # ragged axis becomes the contraction. - from_dlpack(activation_gemm_view(dy_col_q.t()), assumed_align=16), - from_dlpack(activation_gemm_view(x_col_q.t()), assumed_align=16), - # Carried flat and recast to E8M0 inside the trace; E8M0 has no DLPack - # dtype, so hand over the raw bytes. - from_dlpack(dy_col_sf.view(torch.uint8), assumed_align=16), - from_dlpack(x_col_sf.view(torch.uint8), assumed_align=16), - from_dlpack(offsets, assumed_align=4), - # (G, N, K) contiguous -> (N, K, G): the expert is the L coordinate. - from_dlpack(dw.permute(1, 2, 0), assumed_align=16), - stream, - ) - - slot = _wgrad_executor_slot((rows, out_features, in_features, groups)) - if not slot: - slot.append(cute.compile(_wgrad_entry, *args)) - slot[0](*args) diff --git a/torchao/prototype/moe_training/mxfp8_grouped_mlp.py b/torchao/prototype/moe_training/mxfp8_grouped_mlp.py new file mode 100644 index 0000000000..b33cd7a4aa --- /dev/null +++ b/torchao/prototype/moe_training/mxfp8_grouped_mlp.py @@ -0,0 +1,220 @@ +# 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 MXFP8 grouped-MLP operations for the routed-expert training path. + +Three physically fused CuTe DSL kernels for Blackwell (SM 10.x), each exactly +one GPU kernel launch for a nonempty supported input: + +* :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM (MXFP8 + operands, FP32 accumulation) + BF16 pre-activation save + SwiGLU + rowwise + 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization of the activation. +* :func:`mxfp8_grouped_gemm_dswiglu_bwd` -- FC2 dgrad ragged grouped GEMM + + dSwiGLU (from the saved pre-activation) + dual MXFP8 quantization of the + FC1 gradient. +* :func:`mxfp8_grouped_gemm_wgrad` -- generic ragged-reduction grouped + weight gradient, called once for FC1 and once for FC2. + +This module owns the kernels and their operator wrappers only. Trainer +integration (converter selection, the autograd composite, saved-activation +ownership, expert padding configuration) is follow-up work in the consumer; +:func:`is_supported` is the shape predicate that integration should call +before selecting this operator family. + +Importing this module registers the three ``torchao::`` custom ops. +""" + +import importlib.util + +import torch + +# Importing the ops module registers the custom ops as a side effect. It is +# importable with no CuTe DSL installed; the DSL is imported lazily inside the +# op bodies at first real launch. +from torchao.prototype.moe_training.kernels.mxfp8 import ( + grouped_mlp_ops as _grouped_mlp_ops, # noqa: F401 +) +from torchao.utils import is_cuda_version_at_least + +__all__ = [ + "is_supported", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_wgrad", +] + +# Every per-expert row count, the row allocation, and both feature dims must be +# multiples of this: the tcgen05 blocked scale layout permutes in 128-row tiles +# and the kernels tile all three GEMM axes at 128 with no tail path. +GROUP_ALIGNMENT = 128 + +# Runtime package detection, deliberately independent of the quantizer +# modules' availability flags: probing specs never imports the DSL. +_CUTEDSL_RUNTIME_PACKAGES = { + "cuda.bindings.driver": "cuda-python", + "cutlass": "nvidia-cutlass-dsl", + "cutlass.cute": "nvidia-cutlass-dsl", + "tvm_ffi": "apache-tvm-ffi", +} + + +def _missing_cutedsl_runtime_packages() -> list: + """Names of the pip packages required by the CuTe DSL runtime but absent.""" + missing = [] + for module_name, package_name in _CUTEDSL_RUNTIME_PACKAGES.items(): + try: + spec = importlib.util.find_spec(module_name) + except (ModuleNotFoundError, ValueError): + spec = None + if spec is None and package_name not in missing: + missing.append(package_name) + return missing + + +def _is_sm_10x() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +_mxfp8_grouped_mlp_kernels_available = ( + _is_sm_10x() + and is_cuda_version_at_least(12, 8) + and not _missing_cutedsl_runtime_packages() +) + + +def _require_available() -> None: + """Raise a clean NotImplementedError when the kernels cannot run here.""" + if _mxfp8_grouped_mlp_kernels_available: + return + reasons = [] + if not torch.cuda.is_available(): + reasons.append("CUDA is not available") + elif not _is_sm_10x(): + reasons.append( + "requires an SM 10.x (Blackwell) GPU, found compute capability " + f"{torch.cuda.get_device_capability()}" + ) + if torch.cuda.is_available() and not is_cuda_version_at_least(12, 8): + reasons.append(f"requires CUDA >= 12.8, found {torch.version.cuda}") + missing = _missing_cutedsl_runtime_packages() + if missing: + reasons.append("missing required packages: " + ", ".join(missing)) + if not reasons: + reasons.append("kernels are disabled on this system") + raise NotImplementedError( + "MXFP8 grouped-MLP kernels are unavailable: " + "; ".join(reasons) + ) + + +def is_supported( + model_dim: int, hidden_dim: int, allocated_rows: int, num_groups: int +) -> bool: + """Pure shape predicate for selecting this operator family. + + True when the static shapes satisfy the kernel contract: at least one + expert group and D, F, R all positive multiples of 128. Integration code + (e.g. a quantization converter choosing between this fused path and the + unfused grouped-mm path) should call this BEFORE selecting the ops and + fall back when it is False -- and must also guarantee the runtime offsets + invariant, i.e. configure expert padding to 128 rows, because per-expert + row counts live in device memory and are not host-checkable here. + + Environment availability (CUDA, SM 10.x, the CuTe DSL runtime) is a + separate concern: combine with ``_mxfp8_grouped_mlp_kernels_available``. + """ + return ( + num_groups >= 1 + and model_dim > 0 + and hidden_dim > 0 + and allocated_rows > 0 + and model_dim % GROUP_ALIGNMENT == 0 + and hidden_dim % GROUP_ALIGNMENT == 0 + and allocated_rows % GROUP_ALIGNMENT == 0 + ) + + +def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_t_q, w13_t_sf, offsets): + """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. + + One kernel launch for any nonempty supported input. Arguments (all CUDA + tensors on one device, prequantized by the caller): + + * ``x_q``: E4M3 ``[R, D]``, stride ``(D, 1)`` -- rowwise 1x32-quantized + activations, expert-major packed rows. + * ``x_sf``: flat blocked E8M0 scales for logical ``[R, D/32]``. + * ``w13_t_q``: E4M3 ``[G, D, 2F]``, stride ``(2*D*F, 1, D)`` -- quantized + ``w13.reshape(G, 2F, D).transpose(-2, -1)``; the ``2F`` axis is + element-interleaved gate/up (gate at even indices). + * ``w13_t_sf``: blocked E8M0 ``[G, round_up(2F,128)*round_up(D/32,4)]``. + * ``offsets``: int32 ``[G]`` exclusive per-expert end rows; every expert's + row count must be a nonnegative multiple of 128 and + ``offsets[-1] <= R`` (documented caller invariant; see + ``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` for the opt-in synchronized check). + + Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``; ``z_bf16`` + ``[R, F, 2]`` (gate index 0, up index 1) is the BF16 pre-activation saved + for backward -- pass it unchanged to + :func:`mxfp8_grouped_gemm_dswiglu_bwd`. ``h_col_q`` is column-major with + whole-matrix blocked scales, ready for + :func:`mxfp8_grouped_gemm_wgrad`. Inactive tail rows of every output are + written as zeros. ``R == 0`` returns empty outputs; ``G == 0`` raises + ``ValueError``. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_t_q, w13_t_sf, offsets + ) + + +def mxfp8_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets +): + """FC2 dgrad grouped GEMM + dSwiGLU + rowwise/columnwise MXFP8 quantization. + + One kernel launch for any nonempty supported input. Arguments: + + * ``do_q`` / ``do_sf``: rowwise 1x32-quantized FC2 output-gradient, E4M3 + ``[R, D]`` stride ``(D, 1)`` with blocked scales for ``[R, D/32]``. + * ``w2_dgrad_q`` / ``w2_dgrad_sf``: E4M3 ``[G, D, F]`` stride + ``(D*F, 1, D)`` (dgrad orientation of w2) with per-expert blocked scales + for logical ``[F, D/32]``. + * ``z_bf16``: the exact ``[R, F, 2]`` pre-activation saved by + :func:`mxfp8_grouped_gemm_swiglu_fwd`; rows past ``offsets[-1]`` are + never read. + * ``offsets``: as in the forward op. + + Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)`` -- the FC1 gradient + ``[R, 2F]`` with gate/up gradients element-interleaved to match ``z_bf16``, + quantized both rowwise (row-major qdata) and columnwise (column-major + qdata, whole-matrix blocked scales, ready for the FC1 wgrad call). + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets + ) + + +def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + """Grouped MXFP8 weight gradient: ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + One kernel launch per nonempty invocation, with the per-expert row ranges + of the shared ``R`` axis forming the ragged reduction. Both operands are + columnwise (32x1) quantized: E4M3 logical ``[R, N]`` / ``[R, K]`` with + column-major stride ``(1, R)`` and WHOLE-MATRIX blocked E8M0 scales for + logical ``[N, R/32]`` / ``[K, R/32]`` -- exactly what the forward and + backward ops emit. Do not feed torchao's per-group K-groups scale + rearrangement here: it has the same byte count but a different block + order, and produces a silently wrong ``dw``. + + Generic over both call sites with no mode flag: FC1 wgrad is ``N=2F, + K=D``; FC2 wgrad is ``N=D, K=F``. Returns contiguous BF16 ``[G, N, K]`` + (FP32 accumulation); zero-token experts yield all-zero slices, and + ``R == 0`` returns an all-zero result without launching. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) From 758dd049a7aaa6e75f42ab4d2ba65dc3ceb8e1fa Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Mon, 17 Aug 2026 13:23:59 -0700 Subject: [PATCH 03/11] Simplify cutedsl_grouped_mlp internals: fewer helpers, flatter bodies - Delete the three _*_entry wrappers; the launchers cute.compile the shared _launch_grouped_gemm directly, with the output tensors passed as a tuple argument (probe-verified: no compile-cache aliasing across the three kernels, outputs bitwise-identical to the entry path, still one launch). - Inline the single/dual-use trace-time builders at their call sites: _t2r_partition, _make_tiled_mma, _make_sf_gemm_tensor, _sigmoid_f32. - Drop dead code: TileCoords.tile_m/tile_n, the unused subtile_idx epilogue parameter, the constant-folded kernel-local l_a, and _KernelConfig's three single-consumer properties (inlined as expressions). - Compress narration comments to contracts; the offset contract and the quantization numeric bullets are kept verbatim. 1723 -> 1606 lines; module-level defs 29 -> 22. Public API (__all__, the three launch_* signatures) and the ops/validation/wrapper modules are byte-identical. Generated cubins and SASS are sha256-identical pre/post at both tested shapes, so no machine code changed. Full suite 46 passed + 1 skipped (GPU) / 7 passed (CPU), green after every fold stage; one-launch property preserved; MXFP8_BENCH_VALIDATE=1 bench clean and within run-to-run noise at the 16B-class shape (1200 MHz-capped host, ratios only). Co-Authored-By: Claude Fable 5 --- .../kernels/mxfp8/cutedsl_grouped_mlp.py | 365 ++++++------------ 1 file changed, 124 insertions(+), 241 deletions(-) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py index 3947cc6970..78f1fcb028 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py @@ -16,30 +16,21 @@ * ``launch_grouped_gemm_wgrad`` -- generic ragged-K grouped weight gradient, BF16 output; called once for FC1 and once for FC2. -They share one blockscaled tcgen05 mainloop. Three structural decisions, all +They share one blockscaled tcgen05 mainloop. Three structural invariants, all descending from every per-expert row count being a multiple of 128: -*No per-group tensormaps.* Every operand is one host-built static TMA -descriptor over the whole tensor. Per-expert selection is an integer -coordinate: an L coordinate for the 3-D weight operands, and a K-tile index -base for the wgrad kernel's ragged contraction. The latter is plain layout -algebra: the scale-factor tensor is retiled with -``blockscaled_utils.tile_atom_to_shape_SF``, whose K-tile mode is uniform by -construction, so slicing the partitioned tensor at ``k_base + i`` addresses -expert data exactly, with no per-expert descriptor rebuilds. - -*No tile scheduler.* With the ragged axis tile-aligned, the forward/backward -kernels enumerate all of ``[0, R/128)`` M tiles and the wgrad grid -``(N/128, K/128, G)`` is fully static; only wgrad's K-loop trip count is -data-dependent. - -*The inactive tail needs no special code path.* A tile whose row base is at or -past the active row count runs with ``k_cnt == 0``: no TMA loads are issued, -the accumulator fragment is zeroed in registers, and the unmodified epilogue -emits the zero bytes the contract requires. Epilogue *stores* are never -predicated; the one epilogue-side gmem *input* (the backward kernel's saved -``z_bf16``) is loaded only when ``k_cnt > 0`` because tail rows of that tensor -are read-forbidden. +* No per-group tensormaps: one host-built static TMA descriptor per operand; + per-expert selection is an integer coordinate (an L coordinate for the 3-D + weights, a K-tile index base for wgrad's ragged contraction -- exact + because ``tile_atom_to_shape_SF``'s K-tile mode is uniform). +* No tile scheduler: forward/backward enumerate all ``[0, R/128)`` M tiles + and the wgrad grid ``(N/128, K/128, G)`` is fully static; only wgrad's + K-loop trip count is data-dependent. +* No special tail path: an inactive tile runs with ``k_cnt == 0`` -- no TMA + loads, a register-zeroed accumulator, and the unmodified epilogue emits the + zero bytes the contract requires. Stores are never predicated; the one + epilogue-side gmem input (the backward kernel's saved ``z_bf16``) is loaded + only when ``k_cnt > 0`` because its tail rows are read-forbidden. Offset contract (documented caller invariants -- the offset VALUES live on device and cannot be checked on the host without a synchronization): offsets @@ -87,9 +78,7 @@ "launch_grouped_gemm_wgrad", ] -# --------------------------------------------------------------------------- -# Frozen configuration. One tiling, one pipeline shape, one warp assignment. -# --------------------------------------------------------------------------- +# --- Frozen configuration: one tiling, one pipeline shape, one warp assignment. # MXFP8 scaling block: 32 values share one E8M0 scale. _SF_VEC_SIZE = 32 @@ -144,29 +133,14 @@ class _KernelConfig: # Padded to an odd count so columnwise reads don't serialize on banks. epi_smem_cols: int - @property - def epi_tile(self): - return (_CTA_M, self.epi_n_acc) - - @property - def num_epi_subtiles(self) -> int: - return _CTA_N // self.epi_n_acc - - @property - def epi_smem_elems(self) -> int: - # A zero-size struct field would be degenerate; keep a tiny slab. - return max(_CTA_M * self.epi_smem_cols, 8) - _SWIGLU_FWD_CONFIG = _KernelConfig(epi_n_acc=64, ragged_k=False, epi_smem_cols=33) _DSWIGLU_BWD_CONFIG = _KernelConfig(epi_n_acc=32, ragged_k=False, epi_smem_cols=65) _WGRAD_CONFIG = _KernelConfig(epi_n_acc=32, ragged_k=True, epi_smem_cols=0) -# --------------------------------------------------------------------------- -# Host-side operand views. Pure torch restrides into the (MN, K, L) GEMM +# --- Host-side operand views: pure torch restrides into the (MN, K, L) GEMM # domain, K contiguous. -# --------------------------------------------------------------------------- def activation_gemm_view(t: torch.Tensor) -> torch.Tensor: @@ -198,77 +172,12 @@ class TileCoords: epilogue must store unconditionally. """ - tile_m: Int32 - tile_n: Int32 expert: Int32 row_base: Int32 col_base: Int32 k_cnt: Int32 -# --------------------------------------------------------------------------- -# Mainloop building blocks (all public CuTe DSL API). -# --------------------------------------------------------------------------- - - -def _make_tiled_mma(a_dtype, b_dtype, sf_dtype): - """The one blockscaled tiled MMA, K-major on both operands, FP32 acc.""" - return sm100_utils.make_blockscaled_trivial_tiled_mma( - a_dtype, - b_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.K, - sf_dtype, - _SF_VEC_SIZE, - tcgen05.CtaGroup.ONE, - (_CTA_M, _CTA_N), - ) - - -def _make_sf_gemm_tensor(flat_sf: cute.Tensor, mn: int, k: int, l: int): - """Retile a flat blocked E8M0 buffer into the GEMM-domain SF layout. - - The buffer travels flat by ABI and may arrive as raw uint8, so the pointer - is recast: the MMA rejects a scale operand that is not E8M0. - ``tile_atom_to_shape_SF`` builds kernel IR, so this runs inside the trace. - """ - return cute.make_tensor( - cute.recast_ptr(flat_sf.iterator, dtype=cutlass.Float8E8M0FNU), - blockscaled_utils.tile_atom_to_shape_SF((mn, k, l), _SF_VEC_SIZE), - ) - - -def _t2r_partition(tidx, tAcc_base: cute.Tensor, cfg): - """TMEM accumulator -> register handoff. - - ``tTR_cAcc`` carries each register's ``(row, col)`` coordinate in the CTA - tile; every epilogue index below is derived from it rather than from a raw - register number or an assumed thread-to-row mapping, so the addressing is - correct by construction for whatever value order the copy atom uses. - ``elem_ty_d`` stays Float32: an 8-bit d type would steer - ``get_tmem_load_op`` into layouts shaped for a direct FP8 TMA store. - """ - copy_atom_t2r = sm100_utils.get_tmem_load_op( - _MMA_TILER, - LayoutEnum.ROW_MAJOR, - Float32, - Float32, - cfg.epi_tile, - False, - ) - tAcc_mn = tAcc_base[((None, None), 0, 0, 0)] - tAcc_epi = cute.flat_divide(tAcc_mn, cfg.epi_tile) - tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) - thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - cAcc_epi = cute.flat_divide( - cute.make_identity_tensor((_CTA_M, _CTA_N)), cfg.epi_tile - ) - tTR_cAcc = thr_copy_t2r.partition_D(cAcc_epi) - tTR_rAcc = cute.make_rmem_tensor(tTR_cAcc[(None, None, None, 0, 0)].shape, Float32) - return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc - - def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor): """SMEM -> TMEM scale-factor copy, issued once per K tile from the MMA warp. @@ -290,21 +199,9 @@ def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor): return tiled_copy_s2t, tCsSF_s2t, tCtSF_s2t -# --------------------------------------------------------------------------- -# Quantization math. Public conversions only; every step mirrors the torchao +# --- Quantization math: public conversions only; every step mirrors the torchao # reference (`to_mx(..., RCEIL)` + `to_blocked`) so the fused outputs are # byte-identical to the standalone quantizers on the same BF16 input. -# --------------------------------------------------------------------------- - - -@cute.jit -def _sigmoid_f32(x: Float32) -> Float32: - """sigmoid(x) = 1 / (1 + exp(-x)), accurate mode. - - Composed exactly like torch's float32 sigmoid (default-mode ``exp`` plus a - true divide); measured bit-identical to ``torch.sigmoid`` over 1e6 values. - """ - return Float32(1.0) / (Float32(1.0) + cute.math.exp(Float32(0.0) - x)) @cute.jit @@ -352,10 +249,9 @@ def _quant_block_from_smem( sf_idx: Int32, COLWISE: cutlass.Constexpr, ): - """Quantize one 32-value MX block from the BF16 staging tile. - - Reads ``sEpi[base_row, base_col + i]`` (rowwise) or - ``sEpi[base_row + i, base_col]`` (columnwise), then: + """Quantize one 32-value MX block read from the BF16 staging tile + (``sEpi[base_row, base_col + i]`` rowwise, ``sEpi[base_row + i, base_col]`` + columnwise): * amax: NaN-propagating |max| chain, so a NaN element invalidates the block exactly like the torchao reference. @@ -367,12 +263,10 @@ def _quant_block_from_smem( 2^127; byte 255 gives a NaN reciprocal so every element of an invalidated block quantizes to the E4M3 NaN code. * qdata: one f32 multiply per element, then the public saturating-RNE - Float32 -> Float8E4M3FN conversion (byte-identical to torch's cast; no - explicit clamp -- saturation is part of the conversion's contract). + Float32 -> Float8E4M3FN conversion (byte-identical to torch's cast). - 32 qdata bytes are one contiguous run of ``q_dst`` in both orientations - (row-major rowwise output; column-major columnwise output), stored - vectorized; the scale byte is stored individually at ``sf_idx``. + The 32 qdata bytes are one contiguous ``q_dst`` run in both orientations, + stored vectorized; the scale byte is stored individually at ``sf_idx``. """ vals = [] for i in cutlass.range_constexpr(_SF_VEC_SIZE): @@ -432,11 +326,9 @@ def _epilogue_column_run( return first -# --------------------------------------------------------------------------- -# Epilogues. Called once per subtile by all 128 epilogue threads with a +# --- Epilogues: called once per subtile by all 128 epilogue threads with a # CTA-uniform k_cnt; stores are never predicated (a tail tile's zeroed # accumulator produces exactly the zero bytes the contract requires). -# --------------------------------------------------------------------------- @cute.jit @@ -444,7 +336,6 @@ def _wgrad_epilogue( tTR_rAcc, tTR_cAcc_s, epi_tidx, - subtile_idx: cutlass.Constexpr, tile: TileCoords, sEpi, out, @@ -484,7 +375,6 @@ def _swiglu_fwd_epilogue( tTR_rAcc, tTR_cAcc_s, epi_tidx, - subtile_idx: cutlass.Constexpr, tile: TileCoords, sEpi, out, @@ -495,11 +385,10 @@ def _swiglu_fwd_epilogue( ``out`` = flat views ``(z [R*2F] bf16, h_row_q [R*F] e4m3, h_row_sf uint8, h_col_q [F*R] e4m3 in column-major storage order, - h_col_sf uint8)``. ``N == 2F`` is the GEMM (and z) column count. - - Per the kernel contract: the accumulator is rounded to BF16 first (that IS - z), SwiGLU is evaluated once from the rounded values, h is rounded to BF16 - once, and both quantizers consume the same staged BF16 h. + h_col_sf uint8)``; ``N == 2F`` is the GEMM (and z) column count. Contract: + the accumulator is rounded to BF16 first (that IS z), SwiGLU is evaluated + once from the rounded values, h is rounded to BF16 once, and both + quantizers consume the same staged BF16 h. """ num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) # 64 half = cutlass.const_expr(num_acc // 2) # 32 h columns per subtile @@ -523,7 +412,9 @@ def _swiglu_fwd_epilogue( for j in cutlass.range_constexpr(half): gate = Float32(zfrag[2 * j]) up = Float32(zfrag[2 * j + 1]) - sig = _sigmoid_f32(gate) + # sigmoid composed exactly like torch's float32 sigmoid (default-mode + # exp plus a true divide); measured bit-identical to torch.sigmoid. + sig = Float32(1.0) / (Float32(1.0) + cute.math.exp(Float32(0.0) - gate)) sEpi[lrow, Int32(j)] = cutlass.BFloat16((gate * sig) * up) cute.arch.barrier( @@ -531,8 +422,8 @@ def _swiglu_fwd_epilogue( ) # ---- stage 2: dual quantization off the staging tile ----------------- - # Global h column base of this subtile; frag_col is static (64 * - # subtile_idx), so hbase is divisible by 32. + # Global h column base of this subtile; frag_col is static, 64 per + # subtile, so hbase is divisible by 32. hbase = (tile.col_base + Int32(frag_col)) >> 1 ncb_row = cutlass.const_expr((F // _SF_VEC_SIZE + 3) // 4) ncb_col = cutlass.const_expr((R // _SF_VEC_SIZE + 3) // 4) @@ -578,7 +469,6 @@ def _dswiglu_bwd_epilogue( tTR_rAcc, tTR_cAcc_s, epi_tidx, - subtile_idx: cutlass.Constexpr, tile: TileCoords, sEpi, out, @@ -588,14 +478,12 @@ def _dswiglu_bwd_epilogue( """dSwiGLU from the saved z + dual MXFP8 quantization of dz. ``out`` = ``(z [R*2F] bf16 INPUT, dz_row_q [R*2F] e4m3, dz_row_sf uint8, - dz_col_q [2F*R] e4m3 column-major storage, dz_col_sf uint8)``. ``N == F`` + dz_col_q [2F*R] e4m3 column-major storage, dz_col_sf uint8)``; ``N == F`` is the dgrad GEMM column count; dz has 2F element-interleaved columns. - - Per the kernel contract: dh is rounded to BF16 first; gate/up come from - the saved BF16 z; dgate/dup are each rounded to BF16 before interleaving; - both quantizers consume the same staged BF16 dz. The z load is the one - epilogue-side gmem input in the family and is predicated on - ``k_cnt == 0`` -- tail rows of z are read-forbidden and contribute zeros. + Contract: dh is rounded to BF16 first; gate/up come from the saved BF16 z; + dgate/dup are each rounded to BF16 before interleaving; both quantizers + consume the same staged BF16 dz. The z load is predicated on ``k_cnt`` -- + tail rows of z are read-forbidden and contribute zeros. """ num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) # 32 two_f = cutlass.const_expr(2 * N) @@ -630,7 +518,7 @@ def _dswiglu_bwd_epilogue( gate = Float32(zfrag[2 * j]) up = Float32(zfrag[2 * j + 1]) dh = Float32(cutlass.BFloat16(tTR_rAcc[j])) - sig = _sigmoid_f32(gate) + sig = Float32(1.0) / (Float32(1.0) + cute.math.exp(Float32(0.0) - gate)) silu = gate * sig dsilu = sig * (Float32(1.0) + gate * (Float32(1.0) - sig)) sEpi[lrow, Int32(2 * j)] = cutlass.BFloat16((dh * up) * dsilu) @@ -680,9 +568,7 @@ def _dswiglu_bwd_epilogue( ) -# --------------------------------------------------------------------------- -# The shared kernel and launch builder. -# --------------------------------------------------------------------------- +# --- The shared kernel and launch builder. @cute.kernel @@ -733,7 +619,6 @@ def _grouped_gemm_kernel( is_active = Int32(offs[num_groups - 1] > row_base) k_base = Int32(0) k_cnt = is_active * Int32(num_k_tiles_full) - l_a = Int32(0) l_b = expert else: expert = Int32(bidz) @@ -747,11 +632,8 @@ def _grouped_gemm_kernel( # never wrote. Clamped, a malformed expert degrades to an all-zero # slice instead. k_cnt = cutlass.max((offs[expert] - prev) // _CTA_K, Int32(0)) - l_a = Int32(0) l_b = Int32(0) tile = TileCoords( - tile_m=tile_m, - tile_n=tile_n, expert=expert, row_base=row_base, col_base=tile_n * _CTA_N, @@ -870,9 +752,9 @@ def _grouped_gemm_kernel( tBsSFB = cute.filter_zeros(tBsSFB) tBgSFB = cute.filter_zeros(tBgSFB) - tAgA_slice = tAgA[(None, tile_m, None, l_a)] + tAgA_slice = tAgA[(None, tile_m, None, 0)] tBgB_slice = tBgB[(None, tile_n, None, l_b)] - tAgSFA_slice = tAgSFA[(None, tile_m, None, l_a)] + tAgSFA_slice = tAgSFA[(None, tile_m, None, 0)] tBgSFB_slice = tBgSFB[(None, tile_n, None, l_b)] acc_shape = tiled_mma.partition_shape_C(_MMA_TILER[:2]) @@ -1000,8 +882,28 @@ def _grouped_gemm_kernel( tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) epi_tidx = tidx - _FIRST_EPI_THREAD - tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_cAcc = _t2r_partition( - epi_tidx, tCtAcc_base, cfg + # TMEM -> register handoff. tTR_cAcc carries each register's (row, col) + # coordinate in the CTA tile; every epilogue index is derived from it, + # not from an assumed thread-to-row mapping. The d element type stays + # Float32: an 8-bit d would steer get_tmem_load_op into layouts shaped + # for a direct FP8 TMA store. + epi_tile = (_CTA_M, cfg.epi_n_acc) + copy_atom_t2r = sm100_utils.get_tmem_load_op( + _MMA_TILER, LayoutEnum.ROW_MAJOR, Float32, Float32, epi_tile, False + ) + tAcc_mn = tCtAcc_base[((None, None), 0, 0, 0)] + tAcc_epi = cute.flat_divide(tAcc_mn, epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0)] + ) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + cAcc_epi = cute.flat_divide( + cute.make_identity_tensor((_CTA_M, _CTA_N)), epi_tile + ) + tTR_cAcc = thr_copy_t2r.partition_D(cAcc_epi) + tTR_rAcc = cute.make_rmem_tensor( + tTR_cAcc[(None, None, None, 0, 0)].shape, Float32 ) acc_consumer_state = pipeline.make_pipeline_state( @@ -1009,7 +911,7 @@ def _grouped_gemm_kernel( ) acc_pipeline.consumer_wait(acc_consumer_state) - for s in cutlass.range_constexpr(cfg.num_epi_subtiles): + for s in cutlass.range_constexpr(_CTA_N // cfg.epi_n_acc): if k_cnt == Int32(0): # Tail tile or zero-token expert: nothing was accumulated, so # the fragment is zeroed here and the epilogue runs unchanged. @@ -1021,7 +923,6 @@ def _grouped_gemm_kernel( tTR_rAcc, tTR_cAcc[(None, None, None, 0, s)], epi_tidx, - s, tile, sEpi, out, @@ -1051,9 +952,9 @@ def _launch_grouped_gemm( """Build the four static TMA descriptors and launch. Grid is data-independent. This is a trace body: calling it directly retraces the kernel on every - invocation. The public launchers below wrap it in per-kernel ``@cute.jit`` - entry points taking only dynamic tensors, ``cute.compile`` those once per - shape key, and call the compiled executor. + invocation. The public launchers ``cute.compile`` it once per shape key, + passing ``cfg`` and ``EPILOGUE`` as trailing Constexpr args, then call the + compiled executor with the runtime args only. """ a_dtype = mA.element_type b_dtype = mB.element_type @@ -1074,10 +975,29 @@ def _launch_grouped_gemm( f"({_CTA_M}, {_CTA_N}, {_CTA_K})" ) - mSFA = _make_sf_gemm_tensor(sfa_flat, gemm_m, gemm_k, l_a) - mSFB = _make_sf_gemm_tensor(sfb_flat, gemm_n, gemm_k, l_b) + # Retile the flat blocked E8M0 buffers into the GEMM-domain SF layout. The + # buffers travel flat by ABI and may arrive as raw uint8, so the pointer is + # recast: the MMA rejects a scale operand that is not E8M0. + mSFA = cute.make_tensor( + cute.recast_ptr(sfa_flat.iterator, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.tile_atom_to_shape_SF((gemm_m, gemm_k, l_a), _SF_VEC_SIZE), + ) + mSFB = cute.make_tensor( + cute.recast_ptr(sfb_flat.iterator, dtype=cutlass.Float8E8M0FNU), + blockscaled_utils.tile_atom_to_shape_SF((gemm_n, gemm_k, l_b), _SF_VEC_SIZE), + ) - tiled_mma = _make_tiled_mma(a_dtype, b_dtype, sf_dtype) + # The one blockscaled tiled MMA, K-major on both operands, FP32 acc. + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + a_dtype, + b_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + sf_dtype, + _SF_VEC_SIZE, + tcgen05.CtaGroup.ONE, + (_CTA_M, _CTA_N), + ) cluster_layout_vmnk = cute.tiled_divide( cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) ) @@ -1133,13 +1053,17 @@ def _launch_grouped_gemm( @cute.struct class SharedStorage: + # The *_empty ranges are live storage: each Pipeline*.create consumes + # 2 x num_stages barriers starting at the *_full pointer. ab_full_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_AB_STAGE] ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_AB_STAGE] acc_full_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_ACC_STAGE] acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_ACC_STAGE] tmem_holding_buf: cutlass.Int32 + # A zero-size struct field would be degenerate; keep a tiny slab. sEpi: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, cfg.epi_smem_elems], 128 + cute.struct.MemRange[cutlass.BFloat16, max(_CTA_M * cfg.epi_smem_cols, 8)], + 128, ] sA: cute.struct.Align[ cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024 @@ -1194,65 +1118,10 @@ class SharedStorage: ) -# --------------------------------------------------------------------------- -# Compile-once entry points (dynamic tensors only; Constexprs closed over). -# --------------------------------------------------------------------------- - - -@cute.jit -def _swiglu_fwd_entry(mA, mB, sfa, sfb, offs, mZ, mHrq, mHrs, mHcq, mHcs, stream): - _launch_grouped_gemm( - mA, - mB, - sfa, - sfb, - offs, - (mZ, mHrq, mHrs, mHcq, mHcs), - stream, - _SWIGLU_FWD_CONFIG, - _swiglu_fwd_epilogue, - ) - - -@cute.jit -def _dswiglu_bwd_entry(mA, mB, sfa, sfb, offs, mZ, mDrq, mDrs, mDcq, mDcs, stream): - _launch_grouped_gemm( - mA, - mB, - sfa, - sfb, - offs, - (mZ, mDrq, mDrs, mDcq, mDcs), - stream, - _DSWIGLU_BWD_CONFIG, - _dswiglu_bwd_epilogue, - ) - - -@cute.jit -def _wgrad_entry(mA, mB, sfa, sfb, offs, mDw, stream): - _launch_grouped_gemm( - mA, - mB, - sfa, - sfb, - offs, - (mDw,), - stream, - _WGRAD_CONFIG, - _wgrad_epilogue, - ) - - @functools.cache def _executor_slot(key: tuple) -> list: - """One memo slot per (kernel, shape, device, dtype, DSL version) key. - - The executor cannot be built from the key alone -- the shared trace needs - static shapes (the blocked scale layout and the grid are built from them) - -- so the first real call's tensors are what gets compiled and the caller - fills the slot once. - """ + """One memo slot per (kernel, shape, device, dtype, DSL version) key; the + trace needs real tensors, so the first caller compiles and fills it once.""" return [] @@ -1442,11 +1311,13 @@ def launch_grouped_gemm_swiglu_fwd( from_dlpack(x_sf.view(torch.uint8).view(-1), assumed_align=16), from_dlpack(w13_t_sf.view(torch.uint8).view(-1), assumed_align=16), from_dlpack(offsets, assumed_align=4), - from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), - from_dlpack(h_row_q.view(-1), assumed_align=16), - from_dlpack(h_row_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(h_col_q.t().reshape(-1), assumed_align=16), - from_dlpack(h_col_sf.view(torch.uint8).view(-1), assumed_align=16), + ( + from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), + from_dlpack(h_row_q.view(-1), assumed_align=16), + from_dlpack(h_row_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(h_col_q.t().reshape(-1), assumed_align=16), + from_dlpack(h_col_sf.view(torch.uint8).view(-1), assumed_align=16), + ), stream, ) key = _cache_key( @@ -1454,7 +1325,11 @@ def launch_grouped_gemm_swiglu_fwd( ) slot = _executor_slot(key) if not slot: - slot.append(cute.compile(_swiglu_fwd_entry, *args)) + slot.append( + cute.compile( + _launch_grouped_gemm, *args, _SWIGLU_FWD_CONFIG, _swiglu_fwd_epilogue + ) + ) slot[0](*args) @@ -1600,11 +1475,13 @@ def launch_grouped_gemm_dswiglu_bwd( from_dlpack(do_sf.view(torch.uint8).view(-1), assumed_align=16), from_dlpack(w2_dgrad_sf.view(torch.uint8).view(-1), assumed_align=16), from_dlpack(offsets, assumed_align=4), - from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), - from_dlpack(dz_row_q.view(-1), assumed_align=16), - from_dlpack(dz_row_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(dz_col_q.t().reshape(-1), assumed_align=16), - from_dlpack(dz_col_sf.view(torch.uint8).view(-1), assumed_align=16), + ( + from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), + from_dlpack(dz_row_q.view(-1), assumed_align=16), + from_dlpack(dz_row_sf.view(torch.uint8).view(-1), assumed_align=16), + from_dlpack(dz_col_q.t().reshape(-1), assumed_align=16), + from_dlpack(dz_col_sf.view(torch.uint8).view(-1), assumed_align=16), + ), stream, ) key = _cache_key( @@ -1612,7 +1489,11 @@ def launch_grouped_gemm_dswiglu_bwd( ) slot = _executor_slot(key) if not slot: - slot.append(cute.compile(_dswiglu_bwd_entry, *args)) + slot.append( + cute.compile( + _launch_grouped_gemm, *args, _DSWIGLU_BWD_CONFIG, _dswiglu_bwd_epilogue + ) + ) slot[0](*args) @@ -1711,7 +1592,7 @@ def launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, d from_dlpack(dy_col_sf.view(torch.uint8).view(-1), assumed_align=16), from_dlpack(x_col_sf.view(torch.uint8).view(-1), assumed_align=16), from_dlpack(offsets, assumed_align=4), - from_dlpack(dw.permute(1, 2, 0), assumed_align=16), + (from_dlpack(dw.permute(1, 2, 0), assumed_align=16),), stream, ) key = _cache_key( @@ -1719,5 +1600,7 @@ def launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, d ) slot = _executor_slot(key) if not slot: - slot.append(cute.compile(_wgrad_entry, *args)) + slot.append( + cute.compile(_launch_grouped_gemm, *args, _WGRAD_CONFIG, _wgrad_epilogue) + ) slot[0](*args) From c6fa589c40f65665e07bcc04a2ec0bc46eb4622d Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 16:36:31 -0700 Subject: [PATCH 04/11] Add cuDNN-frontend MXFP8 grouped-MLP custom ops (glu/mm/dglu/wgrad) Four torchao:: custom ops for the routed-expert grouped MLP, each one launch of a cudnn.grouped_gemm_*_wrapper_sm100 CuTe DSL kernel from the standalone cudnn-frontend python package (>= 1.27, SM 10.x; no TransformerEngine dependency): mxfp8_cudnn_grouped_mlp_fwd FC1 GEMM + SwiGLU + dual MXFP8 RCEIL quant + BF16 pre-GLU (32-block GLU row order) mxfp8_cudnn_grouped_mm grouped GEMM on prequantized operands -> BF16 (FC2 forward and FC1 dgrad; b [G,N,K] along K) mxfp8_cudnn_grouped_mlp_bwd FC2 dgrad + dSwiGLU + dual quant of dz mxfp8_cudnn_grouped_mlp_wgrad ragged-reduction weight gradient, dense mode Contract highlights (probe-derived, see agent_scratch/cudnn_fe_torchao): - per-expert row counts and R must be multiples of 256 (cuDNN FE FIX_PAD_SIZE; 128-only groups corrupt results silently and NONDETERMINISTICALLY). Two-tier validation: always-on metadata checks + opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS=1 value checks. - flat blocked E8M0 scale ABI; activation colwise scales are PER-GROUP blocked (K-groups layout) and sized by offsets[-1], which may be < R. - colwise qdata accepted in ANY major (dim1-native transposed memory, kernel un-transposed bytes, and mixes -- every combination probe-proven). - rows past offsets[-1]: caller-allocated outputs untouched; kernel-allocated outputs garbage and read-forbidden (verified with NaN-poisoned tails). Self-gate (GB200, FE 1.27.0/backend 92500, TE image build 401656373): 46/46 gates green -- full fwd+bwd chains at dbg/D!=F/G=1/16B shapes (GEMM outputs 87-159 dB vs dequantized-operand references, quant outputs 31.5 dB), A= 1.27, +Blackwell SM 10.x; no TransformerEngine dependency): + +* :func:`mxfp8_cudnn_grouped_mlp_fwd` -- FC1 ragged grouped GEMM + SwiGLU + + rowwise 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization + BF16 pre-GLU. +* :func:`mxfp8_cudnn_grouped_mm` -- ragged grouped GEMM on prequantized + MXFP8 operands to BF16 (FC2 forward and FC1 dgrad). +* :func:`mxfp8_cudnn_grouped_mlp_bwd` -- FC2 dgrad + dSwiGLU + dual MXFP8 + quantization of the FC1 gradient. +* :func:`mxfp8_cudnn_grouped_mlp_wgrad` -- ragged-reduction grouped weight + gradient (FC1 and FC2). + +CONTRACT (stricter than the archived custom-kernel family): every per-expert +row count and the allocated row count must be multiples of **256** — the cuDNN +FE kernels hard-code ``FIX_PAD_SIZE = 256``, and groups that are only 128-row +aligned corrupt results SILENTLY and NONDETERMINISTICALLY (the corruption +locus migrates between identical-input reruns; no smoke test can prove a +misaligned config safe). Use a token dispatcher with ``pad_multiple=256`` and +see ``cudnn_grouped_mlp_validation`` for the two-tier enforcement +(``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` for the opt-in synchronized check). + +The FC1 weight must be provided in the cuDNN 32-block GLU row order +``[gate0(32) | up0(32) | gate1(32) | ...]`` along the 2F axis (gate = the +SiLU'd operand). Trainer integration (converter selection, the autograd +composite, weight-layout remaps, expert padding) lives in the consumer; +:func:`is_supported` is the static shape predicate to call before selecting +this family. + +Importing this module registers the four ``torchao::`` custom ops. The +``cudnn`` package itself is imported lazily inside the op bodies. +""" + +import importlib.util + +import torch + +# Importing the ops module registers the custom ops as a side effect. It is +# importable with no cudnn-frontend installed; `import cudnn` is deferred into +# the op bodies. +from torchao.prototype.moe_training.kernels.mxfp8 import ( + cudnn_grouped_mlp_ops as _cudnn_grouped_mlp_ops, # noqa: F401 +) +from torchao.prototype.moe_training.kernels.mxfp8.cudnn_grouped_mlp_validation import ( + DIM_ALIGNMENT, + ROW_GROUP_ALIGNMENT, +) + +__all__ = [ + "DIM_ALIGNMENT", + "ROW_GROUP_ALIGNMENT", + "is_supported", + "mxfp8_cudnn_grouped_mlp_bwd", + "mxfp8_cudnn_grouped_mlp_fwd", + "mxfp8_cudnn_grouped_mlp_wgrad", + "mxfp8_cudnn_grouped_mm", +] + +_REQUIRED_WRAPPERS = ( + "grouped_gemm_glu_wrapper_sm100", + "grouped_gemm_quant_wrapper_sm100", + "grouped_gemm_dglu_wrapper_sm100", + "grouped_gemm_wgrad_wrapper_sm100", +) +# 1.27 is required: earlier frontends reject prob_tensor=None. +_MIN_FE_VERSION = (1, 27) + + +def _fe_version_tuple(version: str) -> tuple: + """Numeric prefix of a version string as a tuple ('1.27.0' -> (1, 27, 0)). + + Never compare version STRINGS: '1.100' < '1.27' lexicographically. + """ + parts = [] + for piece in version.split("."): + digits = "" + for ch in piece: + if not ch.isdigit(): + break + digits += ch + if not digits: + break + parts.append(int(digits)) + return tuple(parts) + + +def _is_sm_10x() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +def _probe_cudnn_frontend() -> str: + """Empty string when usable; else the reason it is not.""" + if importlib.util.find_spec("cudnn") is None: + return "the cudnn-frontend python package ('cudnn') is not installed" + try: + import cudnn + except Exception as exc: # pragma: no cover - environment-specific + return f"'import cudnn' failed: {exc!r}" + version = getattr(cudnn, "__version__", "0") + if _fe_version_tuple(version) < _MIN_FE_VERSION: + return ( + f"cudnn-frontend {version} is too old; >= " + f"{'.'.join(map(str, _MIN_FE_VERSION))} is required " + "(prob_tensor=None support)" + ) + missing = [name for name in _REQUIRED_WRAPPERS if not hasattr(cudnn, name)] + if missing: + return "cudnn-frontend lacks required wrappers: " + ", ".join(missing) + return "" + + +_cudnn_unavailable_reason = ( + _probe_cudnn_frontend() + if _is_sm_10x() + else ( + "requires an SM 10.x (Blackwell) GPU" + if torch.cuda.is_available() + else "CUDA is not available" + ) +) +_cudnn_grouped_mlp_available = _cudnn_unavailable_reason == "" + + +def _require_available() -> None: + """Raise a clean NotImplementedError when the kernels cannot run here.""" + if not _cudnn_grouped_mlp_available: + raise NotImplementedError( + "cuDNN-frontend MXFP8 grouped-MLP kernels are unavailable: " + + _cudnn_unavailable_reason + ) + + +def is_supported(model_dim: int, hidden_dim: int) -> bool: + """Static shape predicate for selecting this operator family. + + True when D and F are positive multiples of 128. Integration code must + ALSO guarantee the runtime row contract (per-expert groups and the row + allocation padded to multiples of 256, e.g. dispatcher pad_multiple=256): + row counts live in device memory and are not checkable here. + + Environment availability (cudnn-frontend >= 1.27, SM 10.x) is a separate + concern: combine with ``_cudnn_grouped_mlp_available``. + """ + return ( + model_dim > 0 + and hidden_dim > 0 + and model_dim % DIM_ALIGNMENT == 0 + and hidden_dim % DIM_ALIGNMENT == 0 + ) + + +def mxfp8_cudnn_grouped_mlp_fwd(x_q, x_sf, w13_q, w13_sf, offsets): + """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. + + See ``torchao::mxfp8_cudnn_grouped_mlp_fwd`` for the full ABI. ``w13_q`` + is E4M3 ``[G, 2F, D]`` contiguous with rows in 32-block GLU order; returns + ``(z_bf16 [R, 2F], h_row_q [R, F], h_row_sf, h_col_q [R, F], h_col_sf)`` + where the columnwise scales are PER-GROUP blocked. Rows past + ``offsets[-1]`` of every output are garbage and read-forbidden. + """ + _require_available() + return torch.ops.torchao.mxfp8_cudnn_grouped_mlp_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + + +def mxfp8_cudnn_grouped_mm(a_q, a_sf, b_q, b_sf, offsets): + """Ragged grouped GEMM on prequantized MXFP8 operands, BF16 output. + + ``b_q`` is ``[G, N, K]``-logical quantized along K with free strides + (rowwise casts as-is; dim1-colwise casts transposed into this + orientation); ``b_sf`` is always the per-group blocked ``[N, K/32]`` + orientation. Returns BF16 ``[R, N]`` with rows past ``offsets[-1]`` + uninitialized. + """ + _require_available() + return torch.ops.torchao.mxfp8_cudnn_grouped_mm(a_q, a_sf, b_q, b_sf, offsets) + + +def mxfp8_cudnn_grouped_mlp_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + """FC2 dgrad + dSwiGLU + dual MXFP8 quantization of the FC1 gradient. + + ``z_bf16`` must be the exact fwd-op output. Returns + ``(dz_row_q [R, 2F], dz_row_sf, dz_col_q [R, 2F], dz_col_sf)`` in the same + 32-block order. + """ + _require_available() + return torch.ops.torchao.mxfp8_cudnn_grouped_mlp_bwd( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + + +def mxfp8_cudnn_grouped_mlp_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + Both operands columnwise (32x1) quantized with PER-GROUP blocked scales + (never whole-matrix ``to_blocked`` — same byte count, silently wrong + block order). Returns contiguous BF16 ``[G, N, K]``. + """ + _require_available() + return torch.ops.torchao.mxfp8_cudnn_grouped_mlp_wgrad( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index da39d890d3..938a6e7dc2 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,8 +1,10 @@ -# Importing grouped_mlp_ops registers the fused grouped-MLP custom ops -# (torchao::mxfp8_grouped_gemm_{swiglu_fwd,dswiglu_bwd,wgrad}). The module is -# importable with no CuTe DSL installed; kernel imports are deferred into the -# op bodies. +# Importing grouped_mlp_ops / cudnn_grouped_mlp_ops registers the fused +# grouped-MLP custom ops (torchao::mxfp8_grouped_gemm_{swiglu_fwd,dswiglu_bwd, +# wgrad} and torchao::mxfp8_cudnn_grouped_{mlp_fwd,mm,mlp_bwd,mlp_wgrad}). +# Both modules are importable with no CuTe DSL / cudnn-frontend installed; +# kernel imports are deferred into the op bodies. from torchao.prototype.moe_training.kernels.mxfp8 import ( + cudnn_grouped_mlp_ops, # noqa: F401 grouped_mlp_ops, # noqa: F401 ) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_ops.py new file mode 100644 index 0000000000..ceb8bc49d5 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_ops.py @@ -0,0 +1,732 @@ +# 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. + +"""Custom-op surface for the cuDNN-frontend MXFP8 routed-expert grouped MLP. + +Four ops, each wrapping one ``cudnn.grouped_gemm_*_wrapper_sm100`` CuTe DSL +kernel from the standalone cudnn-frontend python package (>= 1.27; no +TransformerEngine involvement): + +* ``torchao::mxfp8_cudnn_grouped_mlp_fwd`` -- FC1 ragged grouped GEMM + + SwiGLU + rowwise AND columnwise MXFP8 RCEIL quantization + BF16 pre-GLU save + (``grouped_gemm_glu_wrapper_sm100``). +* ``torchao::mxfp8_cudnn_grouped_mm`` -- ragged grouped GEMM on + prequantized operands to BF16 (``grouped_gemm_quant_wrapper_sm100``); used + for both FC2 forward and FC1 dgrad. +* ``torchao::mxfp8_cudnn_grouped_mlp_bwd`` -- FC2 dgrad + dSwiGLU + dual + MXFP8 quantization of dz (``grouped_gemm_dglu_wrapper_sm100``). +* ``torchao::mxfp8_cudnn_grouped_mlp_wgrad`` -- ragged-reduction grouped + weight gradient (``grouped_gemm_wgrad_wrapper_sm100``, dense output mode); + called once for FC1 and once for FC2. + +All scale arguments are FLAT blocked E8M0 buffers (uint8 or float8_e8m0fnu); +the ops build the kernel-native 6-D / 2-D views internally with probe-proven +recipes. ``offsets`` is int32 CUDA ``[G]`` exclusive-end rows; per-expert row +counts must be multiples of 256 (cuDNN FE ``FIX_PAD_SIZE``; see the validation +module for the two-tier enforcement and the misalignment hazard). Rows in +``[offsets[-1], R)``: caller-allocated outputs (the grouped-mm result and the +weight gradients) keep their tails untouched, while kernel-allocated outputs +(z, h, dz and their scales) carry garbage tails that are read-forbidden -- +both behaviors probe-verified with NaN-poisoned tails. + +Shared ``_validate_*`` helpers back ``register_fake`` so torch.compile rejects +unsupported calls at capture time, and shared output-spec helpers keep fake +and eager metadata identical (eager normalizes the wrapper's returned tensors +and checks them against the same spec the fake allocates from). + +The user-facing wrappers live in +``torchao.prototype.moe_training.cudnn_grouped_mlp``; importing that module +(or this one) registers the ops. ``import cudnn`` happens lazily inside op +bodies at first real launch. +""" + +from typing import Tuple + +import torch + +from torchao.prototype.moe_training.kernels.mxfp8.cudnn_grouped_mlp_validation import ( + ROW_GROUP_ALIGNMENT, + SCALE_BLOCK_SIZE, + validate_allocated_rows, + validate_blocked_scales, + validate_feature_dims, + validate_group_offsets, + validate_operand, + validate_ragged_colwise_scales, +) + +__all__ = ["ROW_GROUP_ALIGNMENT", "SCALE_BLOCK_SIZE"] + +_E4M3 = torch.float8_e4m3fn +_E8M0 = torch.float8_e8m0fnu +_BLOCK = SCALE_BLOCK_SIZE + +# Small per-(groups, dtype, device) caches for the kernels' alpha/beta and +# norm-const tensors. Never cached: the CUDA stream (looked up per call). +_ones_cache: dict = {} + + +def _cached_ones(numel: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + key = (numel, dtype, device) + out = _ones_cache.get(key) + if out is None: + out = torch.ones(numel, dtype=dtype, device=device) + _ones_cache[key] = out + return out + + +def _require_cuda_device(device: torch.device, name: str) -> None: + if device.type != "cuda": + raise ValueError( + f"{name} must be a CUDA tensor, got device {device}; these kernels " + "run only on CUDA SM100 devices" + ) + + +def _as_e8m0(scales: torch.Tensor) -> torch.Tensor: + return scales if scales.dtype == _E8M0 else scales.view(_E8M0) + + +def _act_scale_view(sf_flat: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Flat blocked scales of a logical [rows, cols/32] matrix -> the wrapper's + 6-D activation view (32, 4, rows/128, 4, cols/128, 1).""" + return ( + _as_e8m0(sf_flat) + .view(1, rows // 128, cols // 128, 32, 4, 4) + .permute(3, 4, 1, 5, 2, 0) + ) + + +def _weight_scale_view( + sf_flat: torch.Tensor, groups: int, n: int, k: int +) -> torch.Tensor: + """Per-group-concat flat blocked scales of logical [n, k/32] per expert -> + the wrapper's 6-D weight view (32, 4, n/128, 4, k/128, G).""" + return ( + _as_e8m0(sf_flat) + .view(groups, n // 128, k // 128, 32, 4, 4) + .permute(3, 4, 1, 5, 2, 0) + ) + + +def _flat_scales(sf_6d: torch.Tensor) -> torch.Tensor: + """Kernel-returned 6-D scale view -> the flat blocked buffer (a free view: + the inverse permute restores the allocation's contiguous order).""" + return sf_6d.permute(5, 2, 4, 0, 1, 3).reshape(-1) + + +def _check_normalized( + tensor: torch.Tensor, *, name: str, shape: tuple, dtype: torch.dtype +) -> torch.Tensor: + """Guard against wrapper-output metadata drifting from the fake spec.""" + if tuple(tensor.shape) != tuple(shape) or tensor.dtype != dtype: + raise RuntimeError( + f"cudnn wrapper output {name} has shape {tuple(tensor.shape)} dtype " + f"{tensor.dtype}; expected {tuple(shape)} {dtype}. The installed " + "cudnn-frontend's output contract changed; the registered fake no " + "longer matches eager." + ) + if not tensor.is_contiguous(): + return tensor.contiguous() + return tensor + + +def _stream() -> int: + return torch.cuda.current_stream().cuda_stream + + +# -------------------------------------------------------------------------- +# Op 1: FC1 grouped GEMM + SwiGLU + dual quantization (glu wrapper) +# -------------------------------------------------------------------------- + + +def _fwd_output_specs(rows: int, hidden: int): + two_hidden = 2 * hidden + return ( + ("z_bf16", (rows, two_hidden), torch.bfloat16), + ("h_row_q", (rows, hidden), _E4M3), + ("h_row_sf", (rows * hidden // _BLOCK,), _E8M0), + ("h_col_q", (rows, hidden), _E4M3), + ("h_col_sf", (hidden * rows // _BLOCK,), _E8M0), + ) + + +def _allocate_from_specs(specs, device) -> Tuple[torch.Tensor, ...]: + return tuple( + torch.empty(shape, dtype=dtype, device=device) for _, shape, dtype in specs + ) + + +def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): + if x_q.ndim != 2: + raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") + if w13_q.ndim != 3: + raise ValueError(f"w13_q must be 3D [G, 2F, D], got shape {tuple(w13_q.shape)}") + rows, model_dim = x_q.shape + groups, two_hidden, w_k = w13_q.shape + if w_k != model_dim: + raise ValueError(f"w13_q contraction dim {w_k} must match x_q's D {model_dim}") + if two_hidden % 2 != 0: + raise ValueError( + f"w13_q's row dim must be 2F (32-block interleaved gate/up), " + f"got {two_hidden}" + ) + hidden = two_hidden // 2 + device = x_q.device + + _require_cuda_device(device, "x_q") + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + if rows * two_hidden >= 2**31: + raise ValueError( + f"R * 2F = {rows * two_hidden} does not fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + validate_operand( + x_q, + name="x_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + # The rowwise weight cast delivers a contiguous [G, 2F, D] stack; the + # kernel-facing (2F, D, G) view is built from exactly that layout. + validate_operand( + w13_q, + name="w13_q", + shape=(groups, two_hidden, model_dim), + stride=(two_hidden * model_dim, model_dim, 1), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + x_sf, + name="x_sf", + logical_rows=rows, + logical_cols=model_dim // _BLOCK, + device=device, + ) + validate_blocked_scales( + w13_sf, + name="w13_sf", + logical_rows=two_hidden, + logical_cols=model_dim // _BLOCK, + device=device, + groups=groups, + ) + return rows, model_dim, hidden, groups + + +@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mlp_fwd", mutates_args=()) +def _mxfp8_cudnn_grouped_mlp_fwd( + x_q: torch.Tensor, + x_sf: torch.Tensor, + w13_q: torch.Tensor, + w13_sf: torch.Tensor, + offsets: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FC1 grouped GEMM + SwiGLU + dual MXFP8 RCEIL quantization (one cuDNN launch). + + Inputs (all CUDA, one device; prequantized outside): + x_q E4M3 ``[R, D]`` stride ``(D, 1)``, rowwise 1x32 quantized. + x_sf flat blocked E8M0 scales for logical ``[R, D/32]`` (whole-matrix + blocked == per-group concat because R and every group are %256). + w13_q E4M3 ``[G, 2F, D]`` contiguous, rowwise quantized, rows in the + cuDNN 32-BLOCK GLU order ``[gate0(32) | up0(32) | gate1 | ...]``. + w13_sf per-group flat blocked E8M0, logical ``[2F, D/32]`` per expert. + offsets int32 CUDA ``[G]`` exclusive end rows; per-expert counts %256 + (caller invariant; see the validation module). + + Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``: + z_bf16 BF16 ``[R, 2F]`` contiguous pre-activation in the same 32-block + order; consumed unchanged by the bwd op. + h_row_q E4M3 ``[R, F]`` contiguous; h_row_sf flat blocked for ``[R, F/32]``. + h_col_q E4M3 ``[R, F]`` contiguous columnwise-quantized bytes + (un-transposed kernel layout); h_col_sf PER-GROUP flat blocked + for ``[F, rows_g/32]`` per expert. + + Rows past ``offsets[-1]`` of every output are GARBAGE (kernel-computed from + the quantized input tail) and read-forbidden. ``R == 0`` returns empty + outputs without touching cudnn; ``G == 0`` raises ValueError. + """ + rows, model_dim, hidden, groups = _validate_fwd_inputs( + x_q, x_sf, w13_q, w13_sf, offsets + ) + specs = _fwd_output_specs(rows, hidden) + if rows == 0: + return _allocate_from_specs(specs, x_q.device) + + import cudnn + + out = cudnn.grouped_gemm_glu_wrapper_sm100( + a_tensor=x_q.unsqueeze(0).permute(1, 2, 0), + sfa_tensor=_act_scale_view(x_sf, rows, model_dim), + padded_offsets=offsets, + alpha_tensor=_cached_ones(groups, torch.bfloat16, x_q.device), + b_tensor=w13_q.permute(1, 2, 0), + sfb_tensor=_weight_scale_view(w13_sf, groups, 2 * hidden, model_dim), + norm_const_tensor=_cached_ones(1, torch.float32, x_q.device), + prob_tensor=None, + acc_dtype=torch.float32, + c_dtype=torch.bfloat16, + d_dtype=_E4M3, + cd_major="n", + sf_vec_size=_BLOCK, + act_func="swiglu", + discrete_col_sfd=True, + use_dynamic_sched=True, + current_stream=_stream(), + ) + results = ( + out["c_tensor"].view(rows, 2 * hidden), + out["d_tensor"].view(rows, hidden), + _flat_scales(out["sfd_row_tensor"]), + out["d_col_tensor"].view(rows, hidden), + _flat_scales(out["sfd_col_tensor"]), + ) + return tuple( + _check_normalized(t, name=spec[0], shape=spec[1], dtype=spec[2]) + for t, spec in zip(results, specs) + ) + + +@_mxfp8_cudnn_grouped_mlp_fwd.register_fake +def _(x_q, x_sf, w13_q, w13_sf, offsets): + rows, _model_dim, hidden, _groups = _validate_fwd_inputs( + x_q, x_sf, w13_q, w13_sf, offsets + ) + return _allocate_from_specs(_fwd_output_specs(rows, hidden), x_q.device) + + +# -------------------------------------------------------------------------- +# Op 2: grouped GEMM on prequantized operands -> BF16 (quant wrapper) +# -------------------------------------------------------------------------- + + +def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): + if a_q.ndim != 2: + raise ValueError(f"a_q must be 2D [R, K], got shape {tuple(a_q.shape)}") + if b_q.ndim != 3: + raise ValueError(f"b_q must be 3D [G, N, K], got shape {tuple(b_q.shape)}") + rows, contraction = a_q.shape + groups, out_features, b_k = b_q.shape + if b_k != contraction: + raise ValueError(f"b_q contraction dim {b_k} must match a_q's K {contraction}") + device = a_q.device + + _require_cuda_device(device, "a_q") + # N and K are both feature dims here (D/F/2F at the two call sites). + validate_feature_dims(model_dim=out_features, hidden_dim=contraction) + validate_allocated_rows(rows) + if rows * max(out_features, contraction) >= 2**31: + raise ValueError( + f"R * max(N, K) = {rows * max(out_features, contraction)} does not " + "fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + validate_operand( + a_q, + name="a_q", + shape=(rows, contraction), + stride=(contraction, 1), + dtype=_E4M3, + device=device, + ) + # b_q strides are free: rowwise weight casts arrive [G, N, K] contiguous + # and dim1-colwise casts arrive transposed to [G, N, K] (also row-major in + # this orientation); the wrapper reads the strides (both probe-proven). + validate_operand( + b_q, + name="b_q", + shape=(groups, out_features, contraction), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + a_sf, + name="a_sf", + logical_rows=rows, + logical_cols=contraction // _BLOCK, + device=device, + ) + validate_blocked_scales( + b_sf, + name="b_sf", + logical_rows=out_features, + logical_cols=contraction // _BLOCK, + device=device, + groups=groups, + ) + return rows, out_features, contraction, groups + + +@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mm", mutates_args=()) +def _mxfp8_cudnn_grouped_mm( + a_q: torch.Tensor, + a_sf: torch.Tensor, + b_q: torch.Tensor, + b_sf: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + """Ragged grouped GEMM ``out[r] = dequant(a[r]) @ dequant(b[g(r)]).T`` -> BF16. + + Inputs: + a_q E4M3 ``[R, K]`` stride ``(K, 1)``, rowwise 1x32 quantized; a_sf flat + blocked for logical ``[R, K/32]``. + b_q E4M3 ``[G, N, K]``-logical, quantized ALONG K, any strides (rowwise + weight casts pass as-is; dim1-colwise casts pass transposed into + this orientation). + b_sf per-group flat blocked for the ``[N, K/32]``-oriented scale matrix + (uniform for both quantization axes). + offsets int32 CUDA ``[G]`` exclusive end rows. + + Covers FC2 forward (b = w2 rowwise: N=D, K=F) and FC1 dgrad (b = w13 + colwise: N=D, K=2F). Returns contiguous BF16 ``[R, N]``; rows past + ``offsets[-1]`` are left uninitialized (probe-verified untouched). + ``R == 0`` returns an empty output without touching cudnn. + """ + rows, out_features, contraction, groups = _validate_mm_inputs( + a_q, a_sf, b_q, b_sf, offsets + ) + out = torch.empty(rows, out_features, dtype=torch.bfloat16, device=a_q.device) + if rows == 0: + return out + + import cudnn + + cudnn.grouped_gemm_quant_wrapper_sm100( + a_tensor=a_q.unsqueeze(0).permute(1, 2, 0), + sfa_tensor=_act_scale_view(a_sf, rows, contraction), + padded_offsets=offsets, + alpha_tensor=_cached_ones(groups, torch.bfloat16, a_q.device), + b_tensor=b_q.permute(1, 2, 0), + sfb_tensor=_weight_scale_view(b_sf, groups, out_features, contraction), + norm_const_tensor=None, + prob_tensor=None, + acc_dtype=torch.float32, + d_dtype=torch.bfloat16, + d_tensor=out.as_strided( + (rows, out_features, 1), (out_features, 1, rows * out_features) + ), + cd_major="n", + sf_vec_size=_BLOCK, + use_dynamic_sched=True, + current_stream=_stream(), + ) + return out + + +@_mxfp8_cudnn_grouped_mm.register_fake +def _(a_q, a_sf, b_q, b_sf, offsets): + rows, out_features, _contraction, _groups = _validate_mm_inputs( + a_q, a_sf, b_q, b_sf, offsets + ) + return torch.empty(rows, out_features, dtype=torch.bfloat16, device=a_q.device) + + +# -------------------------------------------------------------------------- +# Op 3: FC2 dgrad + dSwiGLU + dual quantization (dglu wrapper) +# -------------------------------------------------------------------------- + + +def _bwd_output_specs(rows: int, hidden: int): + two_hidden = 2 * hidden + return ( + ("dz_row_q", (rows, two_hidden), _E4M3), + ("dz_row_sf", (rows * two_hidden // _BLOCK,), _E8M0), + ("dz_col_q", (rows, two_hidden), _E4M3), + ("dz_col_sf", (two_hidden * rows // _BLOCK,), _E8M0), + ) + + +def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + if dy_q.ndim != 2: + raise ValueError(f"dy_q must be 2D [R, D], got shape {tuple(dy_q.shape)}") + if w2_col_q.ndim != 3: + raise ValueError( + f"w2_col_q must be 3D [G, D, F], got shape {tuple(w2_col_q.shape)}" + ) + rows, model_dim = dy_q.shape + groups, w_d, hidden = w2_col_q.shape + if w_d != model_dim: + raise ValueError(f"w2_col_q's D dim {w_d} must match dy_q's D {model_dim}") + if z_bf16.ndim != 2 or tuple(z_bf16.shape) != (rows, 2 * hidden): + raise ValueError( + f"z_bf16 must be [{rows}, {2 * hidden}] (32-block interleaved, the " + f"exact fwd-op output), got shape {tuple(z_bf16.shape)}" + ) + device = dy_q.device + + _require_cuda_device(device, "dy_q") + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + if rows * 2 * hidden >= 2**31: + raise ValueError( + f"R * 2F = {rows * 2 * hidden} does not fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + validate_operand( + dy_q, + name="dy_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + # Colwise-quantized w2; strides free (dim1-native layout probe-proven). + validate_operand( + w2_col_q, + name="w2_col_q", + shape=(groups, model_dim, hidden), + dtype=_E4M3, + device=device, + ) + validate_operand( + z_bf16, + name="z_bf16", + shape=(rows, 2 * hidden), + stride=(2 * hidden, 1), + dtype=torch.bfloat16, + device=device, + ) + validate_blocked_scales( + dy_sf, + name="dy_sf", + logical_rows=rows, + logical_cols=model_dim // _BLOCK, + device=device, + ) + # Colwise weight scales: logical [F, D/32] per expert. + validate_blocked_scales( + w2_col_sf, + name="w2_col_sf", + logical_rows=hidden, + logical_cols=model_dim // _BLOCK, + device=device, + groups=groups, + ) + return rows, model_dim, hidden, groups + + +@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mlp_bwd", mutates_args=()) +def _mxfp8_cudnn_grouped_mlp_bwd( + dy_q: torch.Tensor, + dy_sf: torch.Tensor, + w2_col_q: torch.Tensor, + w2_col_sf: torch.Tensor, + z_bf16: torch.Tensor, + offsets: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FC2 dgrad grouped GEMM + dSwiGLU + dual MXFP8 quantization (one launch). + + Inputs: + dy_q / dy_sf rowwise-quantized FC2 output gradient ``[R, D]``. + w2_col_q E4M3 ``[G, D, F]``-logical, quantized along D, any + strides (dim1-native accepted). + w2_col_sf per-group flat blocked for logical ``[F, D/32]``. + z_bf16 the EXACT ``[R, 2F]`` output of the fwd op (32-block + interleaved). Rows past ``offsets[-1]`` never read. + offsets int32 CUDA ``[G]``. + + Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)``: the FC1 gradient + ``[R, 2F]`` in the same 32-block order, rowwise + columnwise quantized + (columnwise: un-transposed kernel bytes, PER-GROUP flat blocked scales). + Tails garbage/read-forbidden as in the fwd op. + """ + rows, model_dim, hidden, groups = _validate_bwd_inputs( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + specs = _bwd_output_specs(rows, hidden) + if rows == 0: + return _allocate_from_specs(specs, dy_q.device) + + import cudnn + + out = cudnn.grouped_gemm_dglu_wrapper_sm100( + a_tensor=dy_q.unsqueeze(0).permute(1, 2, 0), + c_tensor=z_bf16.unsqueeze(0).permute(1, 2, 0), + sfa_tensor=_act_scale_view(dy_sf, rows, model_dim), + padded_offsets=offsets, + alpha_tensor=_cached_ones(groups, torch.bfloat16, dy_q.device), + beta_tensor=_cached_ones(groups, torch.bfloat16, dy_q.device), + prob_tensor=None, + dprob_tensor=None, + b_tensor=w2_col_q.permute(2, 1, 0), + sfb_tensor=_weight_scale_view(w2_col_sf, groups, hidden, model_dim), + norm_const_tensor=_cached_ones(1, torch.float32, dy_q.device), + acc_dtype=torch.float32, + d_dtype=_E4M3, + cd_major="n", + sf_vec_size=_BLOCK, + act_func="dswiglu", + discrete_col_sfd=True, + use_dynamic_sched=True, + current_stream=_stream(), + ) + results = ( + out["d_row_tensor"].view(rows, 2 * hidden), + _flat_scales(out["sfd_row_tensor"]), + out["d_col_tensor"].view(rows, 2 * hidden), + _flat_scales(out["sfd_col_tensor"]), + ) + return tuple( + _check_normalized(t, name=spec[0], shape=spec[1], dtype=spec[2]) + for t, spec in zip(results, specs) + ) + + +@_mxfp8_cudnn_grouped_mlp_bwd.register_fake +def _(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + rows, _model_dim, hidden, _groups = _validate_bwd_inputs( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + return _allocate_from_specs(_bwd_output_specs(rows, hidden), dy_q.device) + + +# -------------------------------------------------------------------------- +# Op 4: grouped weight gradient (wgrad wrapper, dense output mode) +# -------------------------------------------------------------------------- + + +def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + if dy_col_q.ndim != 2 or x_col_q.ndim != 2: + raise ValueError( + "dy_col_q and x_col_q must both be 2D logical [R, N] / [R, K], got " + f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" + ) + rows, out_features = dy_col_q.shape + x_rows, in_features = x_col_q.shape + if x_rows != rows: + raise ValueError( + f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" + ) + groups = offsets.numel() if isinstance(offsets, torch.Tensor) else 0 + device = dy_col_q.device + + _require_cuda_device(device, "dy_col_q") + validate_allocated_rows(rows) + validate_feature_dims(model_dim=out_features, hidden_dim=in_features) + if rows * max(out_features, in_features) >= 2**31: + raise ValueError( + f"R * max(N, K) = {rows * max(out_features, in_features)} does not " + "fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + # Both operands accept ANY major: dim1-native transposed memory, the fwd/ + # bwd ops' un-transposed kernel bytes, and mixes -- all four combinations + # probe-proven. + validate_operand( + dy_col_q, + name="dy_col_q", + shape=(rows, out_features), + dtype=_E4M3, + device=device, + ) + validate_operand( + x_col_q, + name="x_col_q", + shape=(rows, in_features), + dtype=_E4M3, + device=device, + ) + # Columnwise scale buffers are sized by offsets[-1] (a device value), not + # by R: only dtype/device/divisibility are host-checkable. + validate_ragged_colwise_scales( + dy_col_sf, + name="dy_col_sf", + features=out_features, + allocated_rows=rows, + device=device, + ) + validate_ragged_colwise_scales( + x_col_sf, + name="x_col_sf", + features=in_features, + allocated_rows=rows, + device=device, + ) + # No cross-buffer size check: a kernel-produced operand's scales are sized + # by the ALLOCATED rows while a composite-produced operand's are sized by + # the ROUTED total offsets[-1] -- mixing the two is legitimate and + # probe-proven (tail case); the kernel reads only within offsets. + return rows, out_features, in_features, groups + + +@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mlp_wgrad", mutates_args=()) +def _mxfp8_cudnn_grouped_mlp_wgrad( + dy_col_q: torch.Tensor, + dy_col_sf: torch.Tensor, + x_col_q: torch.Tensor, + x_col_sf: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + Inputs: + dy_col_q E4M3 logical ``[R, N]``, columnwise (32x1) quantized, ANY major. + dy_col_sf PER-GROUP flat blocked scales (each expert's ``[N, rows_g/32]`` + block concatenated; the K-groups layout). Sized by the routed + total ``offsets[-1]``, which may be < R. + x_col_q / x_col_sf likewise for logical ``[R, K]``. + offsets int32 CUDA ``[G]`` exclusive ends over the shared row axis. + + Do NOT feed whole-matrix ``to_blocked`` scales here: the per-group and + whole-matrix orders coincide in byte count but not content whenever G > 1, + and the mismatch is silent (probe: 2-5 dB instead of 100+). + + Returns contiguous BF16 ``dw [G, N, K]`` (FP32 accumulation). Zero-token + experts ARE written (all-zero slices, probe-verified); ``R == 0`` returns + zeros without launching. FC1 wgrad: N=2F, K=D. FC2 wgrad: N=D, K=F. + """ + rows, out_features, in_features, groups = _validate_wgrad_inputs( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) + dw = torch.empty( + (groups, out_features, in_features), + dtype=torch.bfloat16, + device=dy_col_q.device, + ) + if rows == 0: + return dw.zero_() + + import cudnn + + cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=dy_col_q.t(), + b_tensor=x_col_q, + sfa_tensor=_as_e8m0(dy_col_sf).view(out_features, -1), + sfb_tensor=_as_e8m0(x_col_sf).view(in_features, -1), + offsets_tensor=offsets, + acc_dtype=torch.float32, + sf_vec_size=_BLOCK, + accumulate_on_output=False, + output_mode="dense", + wgrad_tensor=dw, + wgrad_dtype=torch.bfloat16, + current_stream=_stream(), + ) + return dw + + +@_mxfp8_cudnn_grouped_mlp_wgrad.register_fake +def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + _rows, out_features, in_features, groups = _validate_wgrad_inputs( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) + return torch.empty( + (groups, out_features, in_features), + dtype=torch.bfloat16, + device=dy_col_q.device, + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_validation.py b/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_validation.py new file mode 100644 index 0000000000..11bba82e4a --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_validation.py @@ -0,0 +1,310 @@ +# 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. + +"""Host-side precondition validation for the cuDNN-frontend MXFP8 grouped-MLP ops. + +The cuDNN FE grouped kernels hard-code a 256-row group granularity +(``FIX_PAD_SIZE = 256``): per-expert row counts that are only multiples of 128 +SILENTLY and NONDETERMINISTICALLY corrupt results (the corruption locus +migrates between identical-input reruns, consistent with reads of stale memory +at sub-256 group boundaries). No smoke test can prove a misaligned +configuration safe, so the alignment contract is enforced in two tiers: + +* ALWAYS-ON metadata-only checks (no host/device sync, FakeTensor-safe): + dims, dtypes, devices, strides, the allocated row count ``R % 256 == 0``. +* OPT-IN offset-VALUE checks behind ``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` + (one D2H sync per call; validation runs only, skipped for fake tensors): + offsets nondecreasing, every per-expert row count ``% 256 == 0``, and + ``offsets[-1] <= R``. In a default build the offset values are a documented + caller invariant provided by a pad_multiple=256 token dispatcher; there is + NO device-side enforcement. + +Checks raise ValueError rather than asserting, so ``python -O`` cannot strip +them. Metadata gates run before any ``data_ptr()`` gate so FakeTensor tracing +exercises the same checks. +""" + +import os +from typing import Optional + +import torch + +__all__ = [ + "DIM_ALIGNMENT", + "ROW_GROUP_ALIGNMENT", + "SCALE_BLOCK_SIZE", + "blocked_scale_numel", + "host_offsets_validation_enabled", + "validate_allocated_rows", + "validate_blocked_scales", + "validate_feature_dims", + "validate_group_offsets", + "validate_operand", + "validate_ragged_colwise_scales", + "_is_fake", +] + +# MXFP8 scaling block: 32 values share one E8M0 scale. +SCALE_BLOCK_SIZE = 32 +# tcgen05 blocked scale tile: 128 rows x 4 columns, 512 bytes. +SCALE_TILE_ROWS = 128 +SCALE_TILE_COLS = 4 +# Feature-dimension granularity (D and F). +DIM_ALIGNMENT = 128 +# Row-count granularity: per-expert groups AND the allocated row count. This is +# the cuDNN FE kernels' FIX_PAD_SIZE, stricter than the archived family's 128. +ROW_GROUP_ALIGNMENT = 256 +# Byte alignment for TMA/vectorized accesses. +_PTR_ALIGNMENT = 16 + +_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) + + +def _round_up(x: int, to: int) -> int: + return ((x + to - 1) // to) * to + + +def blocked_scale_numel(rows: int, cols: int) -> int: + """Element count of the blocked E8M0 buffer for a logical [rows, cols] scale matrix. + + ``cols`` counts scale values, i.e. the reduced dimension divided by 32. + """ + return _round_up(rows, SCALE_TILE_ROWS) * _round_up(cols, SCALE_TILE_COLS) + + +def host_offsets_validation_enabled() -> bool: + """Opt-in host-side offset validation. Off by default: it forces a D2H sync.""" + return os.environ.get("TORCHAO_MXFP8_VALIDATE_OFFSETS", "0") == "1" + + +def _is_fake(tensor: torch.Tensor) -> bool: + """True for meta/fake tensors, which have no usable data pointer or values.""" + if tensor.device.type == "meta": + return True + try: + from torch._subclasses.fake_tensor import FakeTensor + except ImportError: + return False + return isinstance(tensor, FakeTensor) + + +def validate_group_offsets( + offsets: torch.Tensor, + *, + num_groups: int, + allocated_rows: int, + device: Optional[torch.device] = None, + name: str = "offsets", +) -> None: + """Validate the exclusive-end group offsets tensor. + + Metadata is always checked. The offset VALUES (nondecreasing, per-expert + row counts % 256, offsets[-1] <= R) are checked only when + ``host_offsets_validation_enabled()`` and the tensor is not fake, because + reading them forces a D2H sync. + """ + if not isinstance(offsets, torch.Tensor): + raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") + if num_groups < 1: + raise ValueError( + f"{name} must describe at least one expert group, got G={num_groups}" + ) + if offsets.dtype != torch.int32: + raise ValueError(f"{name} must be int32, got {offsets.dtype}") + if not offsets.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") + if device is not None and offsets.device != device: + raise ValueError( + f"{name} must be on {device}, got {offsets.device}; all operands and " + "destinations must share one CUDA device" + ) + if offsets.ndim != 1: + raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") + if offsets.numel() != num_groups: + raise ValueError( + f"{name} must have one entry per local expert: expected {num_groups}, " + f"got {offsets.numel()}" + ) + if not offsets.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {offsets.stride()}") + + if not host_offsets_validation_enabled() or _is_fake(offsets): + return + + values = offsets.tolist() # d2h sync; opt-in debugging path only + previous = 0 + for group, end in enumerate(values): + if end < previous: + raise ValueError( + f"{name} must be nondecreasing, but entry {group} is {end} " + f"after {previous}" + ) + size = end - previous + if size % ROW_GROUP_ALIGNMENT != 0: + raise ValueError( + f"per-expert row counts must be multiples of {ROW_GROUP_ALIGNMENT} " + f"(cuDNN FE FIX_PAD_SIZE; sub-256 groups corrupt results " + f"nondeterministically): expert {group} has {size} rows " + f"(offsets {previous} -> {end})" + ) + previous = end + if previous > allocated_rows: + raise ValueError( + f"{name}[-1] ({previous}) exceeds the allocated row count " + f"({allocated_rows})" + ) + + +def validate_operand( + tensor: torch.Tensor, + *, + name: str, + shape: tuple, + dtype: torch.dtype, + device: torch.device, + stride: Optional[tuple] = None, + check_pointer_alignment: bool = True, +) -> None: + """Validate one operand's dtype, shape, device, optional exact stride, alignment. + + ``stride=None`` accepts any strides (the cuDNN FE wrappers consume both + row-major and transposed-memory colwise operands; every combination the + composite produces is probe-proven). Metadata gates run before the + ``data_ptr()`` gate so FakeTensor tracing exercises the same checks. + """ + if tensor.dtype != dtype: + raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != tuple(shape): + raise ValueError( + f"{name} must have shape {tuple(shape)}, got {tuple(tensor.shape)}" + ) + if stride is not None and tuple(tensor.stride()) != tuple(stride): + raise ValueError( + f"{name} must have stride {tuple(stride)}, got {tuple(tensor.stride())}. " + "This layout is part of the ABI; a values-equal tensor with a " + "different stride is not interchangeable." + ) + if tensor.device != device: + raise ValueError( + f"{name} must be on {device}, got {tensor.device}; all operands and " + "destinations must share one CUDA device" + ) + if check_pointer_alignment and not _is_fake(tensor): + if tensor.data_ptr() % _PTR_ALIGNMENT != 0: + raise ValueError( + f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data " + f"pointer is {tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an " + "aligned address. A contiguous view with a nonzero storage " + "offset can violate this." + ) + + +def validate_blocked_scales( + scales: torch.Tensor, + *, + name: str, + logical_rows: int, + logical_cols: int, + device: torch.device, + groups: int = 1, +) -> None: + """Validate a flat blocked E8M0 scale buffer with a statically known size. + + The buffer is carried flat (uint8 or float8_e8m0fnu); its logical shape is + metadata. ``groups > 1`` describes per-expert weight buffers whose per-group + blocks are concatenated. + """ + if scales.dtype not in _SCALE_DTYPES: + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " + f"got {scales.dtype}" + ) + expected = groups * blocked_scale_numel(logical_rows, logical_cols) + if scales.numel() != expected: + raise ValueError( + f"{name} must hold {expected} blocked scale bytes for a logical " + f"[{logical_rows}, {logical_cols}] scale matrix" + + (f" across {groups} experts" if groups > 1 else "") + + f", got {scales.numel()}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + + +def validate_ragged_colwise_scales( + scales: torch.Tensor, + *, + name: str, + features: int, + allocated_rows: int, + device: torch.device, +) -> None: + """Validate a per-group columnwise scale buffer whose size depends on offsets. + + Columnwise activation scale buffers are sized by the ROUTED row total + ``offsets[-1]`` — a device value — not by the allocated ``R``: at + ``offsets[-1] < R`` they legitimately cover only ``offsets[-1]/32`` scale + columns (probe-verified). So only dtype/device/contiguity and divisibility + are checked: the numel must be a nonnegative multiple of + ``features * SCALE_TILE_COLS`` rows-block granularity and must not exceed + the buffer implied by the allocated rows. + """ + if scales.dtype not in _SCALE_DTYPES: + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " + f"got {scales.dtype}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + rows_pad = _round_up(features, SCALE_TILE_ROWS) + # Each 256-row group contributes features_pad * (group_rows/32) bytes and + # group_rows/32 is a multiple of 8, so the buffer is a multiple of + # rows_pad * 8 bytes. + granule = rows_pad * (ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE) + if scales.numel() % granule != 0: + raise ValueError( + f"{name} numel {scales.numel()} is not a multiple of {granule} " + f"(= round_up({features},128) x {ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE} " + "scale columns per 256-row group)" + ) + max_numel = rows_pad * (allocated_rows // SCALE_BLOCK_SIZE) + if scales.numel() > max_numel: + raise ValueError( + f"{name} numel {scales.numel()} exceeds the maximum {max_numel} implied " + f"by the allocated row count {allocated_rows}" + ) + + +def validate_feature_dims(*, model_dim: int, hidden_dim: int) -> None: + """D and F must both be positive multiples of 128.""" + if model_dim <= 0 or model_dim % DIM_ALIGNMENT != 0: + raise ValueError( + f"model dimension D must be a positive multiple of {DIM_ALIGNMENT}, " + f"got {model_dim}" + ) + if hidden_dim <= 0 or hidden_dim % DIM_ALIGNMENT != 0: + raise ValueError( + f"routed-expert hidden dimension F must be a positive multiple of " + f"{DIM_ALIGNMENT}, got {hidden_dim}" + ) + + +def validate_allocated_rows(rows: int, *, name: str = "R") -> None: + """The allocated row count must be a multiple of 256 (may be zero). + + Per-expert groups are multiples of 256 (cuDNN FE FIX_PAD_SIZE) and the + allocation must be reachable by a legal offsets vector plus an inactive + tail; a non-256 allocation additionally breaks the whole-matrix == + per-group-concat identity of the rowwise blocked scales. + """ + if rows % ROW_GROUP_ALIGNMENT != 0: + raise ValueError( + f"{name} must be a multiple of {ROW_GROUP_ALIGNMENT}, got {rows}" + ) From 1c5bceb1efa2cd521d5057a8ba4545edc1849821 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 16:49:08 -0700 Subject: [PATCH 05/11] Add unit tests for the cuDNN-frontend MXFP8 grouped-MLP ops 29 tests, ~35 s on one GB200: full-chain numerics at four shape classes (debug with a zero-token expert, D!=F, G=1, 16B-class D=2048/F=1408/G=8) against two references per stage -- dequantized-operand refs with gates derived from the measured FP32 reduction-order band (-12 dB, capped 60 dB; ops measure 93-188 dB) and an independent no-quantization chain from the original BF16 tensors with gates at the measured MXFP8 band -6 dB (z 28.5, y 23.7 dB; op outputs land on the band). Also: the full 2x2 wgrad operand major-mode matrix, A torch.Tensor: + return t.contiguous().view(torch.uint8) + + +def _e8m0_to_f64(s: torch.Tensor) -> torch.Tensor: + u = s.view(torch.uint8).to(torch.int32) + out = torch.exp2((u - 127).to(torch.float64)) + return torch.where(u == 255, torch.full_like(out, float("nan")), out) + + +def _quant_rowwise(x: torch.Tensor): + """[M, K] -> (qdata [M, K] e4m3 row-major, flat blocked scales).""" + s, q = to_mx(x, _E4M3, _BLOCK, scaling_mode=_RCEIL) + return q, to_blocked(s.view(_E8M0)).view(_E8M0) + + +def _quant_colwise(x: torch.Tensor, native: bool): + """[M, K] quantized along M in 32-blocks. + + native=False: un-transposed row-major [M, K] bytes ("rowmajor"). + native=True: the dim1-quantizer layout, [M, K] logical with (1, M) strides. + Scales: flat blocked of the transposed [K, M/32] scale matrix (one group). + """ + M, K = x.shape + if M == 0: + return ( + torch.empty(0, K, dtype=_E4M3, device=x.device), + torch.empty(0, dtype=_E8M0, device=x.device), + ) + s_t, q_t = to_mx(x.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + q = q_t.t() if native else q_t.t().contiguous() + return q, to_blocked(s_t.view(_E8M0)).view(_E8M0) + + +def _cat8(ts, dim=0): + dt = ts[0].dtype + return torch.cat([t.view(torch.uint8) for t in ts], dim).view(dt) + + +def _quant_colwise_grouped(x: torch.Tensor, sizes, native: bool): + """Ragged [R, K]: per-group colwise quantization, per-group blocked scales.""" + qs, sfs = [], [] + off = 0 + for m in sizes: + q, sf = _quant_colwise(x[off : off + m], native=False) + qs.append(q) + sfs.append(sf.reshape(-1)) + off += m + q = _cat8(qs, 0) + if native: + q = q.t().contiguous().t() # values identical; (1, R) strides + return q, _cat8(sfs) + + +def _quant_weight_rowwise(w: torch.Tensor): + """[G, N, K] quantized along K -> (contiguous stack, per-group blocked).""" + qs, sfs = [], [] + for g in range(w.shape[0]): + q, sf = _quant_rowwise(w[g]) + qs.append(q.view(torch.uint8)) + sfs.append(sf.reshape(-1)) + return torch.stack(qs).view(_E4M3), _cat8(sfs) + + +def _quant_weight_colwise(w: torch.Tensor): + """[G, N, K] quantized along N (dim1-native strides per group).""" + qs, sfs = [], [] + for g in range(w.shape[0]): + q, sf = _quant_colwise(w[g], native=False) + qs.append(q.view(torch.uint8)) + sfs.append(sf.reshape(-1)) + return torch.stack(qs).view(_E4M3), _cat8(sfs) + + +def _dequant_rowwise(q: torch.Tensor, sf_flat: torch.Tensor): + M, K = q.shape + s = _e8m0_to_f64(from_blocked(sf_flat.view(_E8M0), M, K // _BLOCK)) + return (q.to(torch.float64) * s.repeat_interleave(_BLOCK, dim=1)).to(torch.float32) + + +def _dequant_colwise_grouped(q: torch.Tensor, sf_flat: torch.Tensor, sizes, K: int): + """q [R, K] logical (any major), per-group blocked scales -> f32 [R, K].""" + R = q.shape[0] + out = torch.empty(R, K, dtype=torch.float32, device=q.device) + off, soff = 0, 0 + for m in sizes: + if m == 0: + continue + n = K * (m // _BLOCK) + s_t = _e8m0_to_f64( + from_blocked(sf_flat[soff : soff + n].view(_E8M0), K, m // _BLOCK) + ) + block = q[off : off + m].to(torch.float64) + out[off : off + m] = (block * s_t.t().repeat_interleave(_BLOCK, dim=0)).to( + torch.float32 + ) + off += m + soff += n + return out + + +def _mk_offsets(sizes, device): + return torch.cumsum( + torch.tensor(sizes, dtype=torch.int64, device=device), dim=0 + ).to(torch.int32) + + +def _zsplit(z: torch.Tensor, hidden: int): + """32-block interleaved [R, 2F] -> (gate [R, F], up [R, F]).""" + v = z.view(z.shape[0], hidden // _BLOCK, 2, _BLOCK) + return ( + v[:, :, 0, :].reshape(z.shape[0], hidden), + v[:, :, 1, :].reshape(z.shape[0], hidden), + ) + + +def _to_32block(w13_elem: torch.Tensor) -> torch.Tensor: + """Element-interleaved [G, F, 2, D] -> cuDNN 32-block GLU order [G, 2F, D].""" + G, hidden, _, D = w13_elem.shape + return ( + w13_elem.view(G, hidden // _BLOCK, _BLOCK, 2, D) + .permute(0, 1, 3, 2, 4) + .reshape(G, 2 * hidden, D) + .contiguous() + ) + + +def _grouped_matmul(a_f32, b_f32_per_group, sizes, transpose_b: bool, chunks: int = 1): + """Per-group f32 matmul with an optional chunked-K reduction order.""" + R = a_f32.shape[0] + N = b_f32_per_group[0].shape[0 if transpose_b else 1] + out = torch.zeros(R, N, dtype=torch.float32, device=a_f32.device) + off = 0 + for g, m in enumerate(sizes): + b = b_f32_per_group[g] + bt = b.t() if transpose_b else b + if chunks == 1: + out[off : off + m] = a_f32[off : off + m] @ bt + else: + K = a_f32.shape[1] + step = K // chunks + acc = torch.zeros(m, N, dtype=torch.float32, device=a_f32.device) + for c in range(chunks): + lo, hi = c * step, K if c == chunks - 1 else (c + 1) * step + acc += a_f32[off : off + m, lo:hi] @ bt[lo:hi] + out[off : off + m] = acc + off += m + return out + + +def _refA_gate(ref_whole: torch.Tensor, ref_chunked: torch.Tensor) -> float: + """GEMM-exactness gate from the reduction-order variability band - 12 dB.""" + band = compute_error(ref_whole.bfloat16(), ref_chunked.bfloat16()).item() + return min(band - 12.0, 60.0) + + +def _dswiglu(dh, gate, up): + s = torch.sigmoid(gate) + return dh * up * (s * (1 + gate * (1 - s))), dh * F.silu(gate) + + +# --------------------------------------------------------------------------- +# Case construction: quantize everything once per case, with references. +# --------------------------------------------------------------------------- + +_CASES = { + # name: (D, F, sizes) + "dbg_zero_token": (256, 256, [256, 0, 512, 256]), + "dnef": (256, 384, [512, 256]), + "g1": (256, 256, [512]), + "16b": (2048, 1408, [256] * 8), +} + + +def _build_case(D, hidden, sizes, device="cuda", seed=0): + torch.manual_seed(seed) + G = len(sizes) + R = sum(sizes) + c = {} + c["sizes"], c["G"], c["R"], c["D"], c["F"] = sizes, G, R, D, hidden + c["offsets"] = _mk_offsets(sizes, device) + c["x"] = torch.randn(R, D, dtype=torch.bfloat16, device=device) * 0.5 + w13_elem = torch.randn(G, hidden, 2, D, dtype=torch.bfloat16, device=device) * 0.02 + c["w13"] = _to_32block(w13_elem) + c["w2"] = torch.randn(G, D, hidden, dtype=torch.bfloat16, device=device) * 0.02 + c["dy"] = torch.randn(R, D, dtype=torch.bfloat16, device=device) * 0.5 + + c["x_q"], c["x_sf"] = _quant_rowwise(c["x"]) + c["w13_q"], c["w13_sf"] = _quant_weight_rowwise(c["w13"]) + c["w2_q"], c["w2_sf"] = _quant_weight_rowwise(c["w2"]) + c["w13c_q"], c["w13c_sf"] = _quant_weight_colwise(c["w13"]) + c["w2c_q"], c["w2c_sf"] = _quant_weight_colwise(c["w2"]) + c["dy_q"], c["dy_sf"] = _quant_rowwise(c["dy"]) + c["x_colq"], c["x_col_sf"] = _quant_colwise_grouped(c["x"], sizes, native=True) + c["dy_colq"], c["dy_col_sf"] = _quant_colwise_grouped(c["dy"], sizes, native=True) + + c["x_deq"] = _dequant_rowwise(c["x_q"], c["x_sf"]) + c["w13_deq"] = [ + _dequant_rowwise(c["w13_q"][g], c["w13_sf"].view(G, -1)[g]) for g in range(G) + ] + c["w2_deq"] = [ + _dequant_rowwise(c["w2_q"][g], c["w2_sf"].view(G, -1)[g]) for g in range(G) + ] + return c + + +@pytest.fixture(scope="module") +def dbg(): + return _build_case(*_CASES["dbg_zero_token"]) + + +# --------------------------------------------------------------------------- +# Registration / availability / fakes (no GPU launch). +# --------------------------------------------------------------------------- + + +def test_ops_registered(): + for name in ( + "mxfp8_cudnn_grouped_mlp_fwd", + "mxfp8_cudnn_grouped_mm", + "mxfp8_cudnn_grouped_mlp_bwd", + "mxfp8_cudnn_grouped_mlp_wgrad", + ): + assert hasattr(_OPS, name), f"torchao::{name} is not registered" + + +def test_is_supported(): + assert is_supported(2048, 1408) + assert is_supported(256, 256) + assert not is_supported(192, 256) + assert not is_supported(256, 64) + assert not is_supported(0, 256) + + +def _fake_chain_shapes(R=512, D=256, hidden=256, G=2): + N1 = 2 * hidden + dev = "cuda" + x_q = torch.empty(R, D, dtype=_E4M3, device=dev) + x_sf = torch.empty(R * D // _BLOCK, dtype=_E8M0, device=dev) + w13_q = torch.empty(G, N1, D, dtype=_E4M3, device=dev) + w13_sf = torch.empty(G * N1 * D // _BLOCK, dtype=_E8M0, device=dev) + offsets = torch.empty(G, dtype=torch.int32, device=dev) + outs = {} + outs["fwd"] = _OPS.mxfp8_cudnn_grouped_mlp_fwd(x_q, x_sf, w13_q, w13_sf, offsets) + w2_q = torch.empty(G, D, hidden, dtype=_E4M3, device=dev) + w2_sf = torch.empty(G * D * hidden // _BLOCK, dtype=_E8M0, device=dev) + outs["mm"] = _OPS.mxfp8_cudnn_grouped_mm( + outs["fwd"][1], outs["fwd"][2], w2_q, w2_sf, offsets + ) + dy_q = torch.empty(R, D, dtype=_E4M3, device=dev) + dy_sf = torch.empty(R * D // _BLOCK, dtype=_E8M0, device=dev) + w2c_q = torch.empty(G, D, hidden, dtype=_E4M3, device=dev) + w2c_sf = torch.empty(G * D * hidden // _BLOCK, dtype=_E8M0, device=dev) + outs["bwd"] = _OPS.mxfp8_cudnn_grouped_mlp_bwd( + dy_q, dy_sf, w2c_q, w2c_sf, outs["fwd"][0], offsets + ) + outs["wgrad"] = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + torch.empty(R, D, dtype=_E4M3, device=dev), + torch.empty(D * R // _BLOCK, dtype=_E8M0, device=dev), + torch.empty(R, hidden, dtype=_E4M3, device=dev), + torch.empty( + ((hidden + 127) // 128 * 128) * R // _BLOCK, dtype=_E8M0, device=dev + ), + offsets, + ) + return outs + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device object") +def test_fake_contracts_match_specs(): + """All four fakes produce the documented shapes/dtypes/contiguity.""" + with FakeTensorMode(): + outs = _fake_chain_shapes() + R, D, hidden, N1 = 512, 256, 256, 512 + z, hq, hsf, hcq, hcsf = outs["fwd"] + assert tuple(z.shape) == (R, N1) and z.dtype == torch.bfloat16 + assert tuple(hq.shape) == (R, hidden) and hq.dtype == _E4M3 + assert hsf.numel() == R * hidden // _BLOCK and hsf.dtype == _E8M0 + assert tuple(hcq.shape) == (R, hidden) and hcq.dtype == _E4M3 + assert hcsf.numel() == hidden * R // _BLOCK + assert all(t.is_contiguous() for t in outs["fwd"]) + y = outs["mm"] + assert tuple(y.shape) == (R, D) and y.dtype == torch.bfloat16 and y.is_contiguous() + dz_q, dz_sf, dzc_q, dzc_sf = outs["bwd"] + assert tuple(dz_q.shape) == (R, N1) and dz_sf.numel() == R * N1 // _BLOCK + assert tuple(dzc_q.shape) == (R, N1) and dzc_sf.numel() == N1 * R // _BLOCK + dw = outs["wgrad"] + assert tuple(dw.shape) == (2, D, hidden) and dw.dtype == torch.bfloat16 + + +# --------------------------------------------------------------------------- +# Full-chain numerics: two references per stage, derived gates. +# --------------------------------------------------------------------------- + + +def _run_chain(c): + """fwd -> FC2 mm -> bwd -> FC1-dgrad mm -> wgrad x2, production layouts.""" + r = {} + r["z"], r["h_q"], r["h_sf"], r["h_colq"], r["h_col_sf"] = ( + _OPS.mxfp8_cudnn_grouped_mlp_fwd( + c["x_q"], c["x_sf"], c["w13_q"], c["w13_sf"], c["offsets"] + ) + ) + r["y"] = _OPS.mxfp8_cudnn_grouped_mm( + r["h_q"], r["h_sf"], c["w2_q"], c["w2_sf"], c["offsets"] + ) + r["dz_q"], r["dz_sf"], r["dz_colq"], r["dz_col_sf"] = ( + _OPS.mxfp8_cudnn_grouped_mlp_bwd( + c["dy_q"], c["dy_sf"], c["w2c_q"], c["w2c_sf"], r["z"], c["offsets"] + ) + ) + # FC1 dgrad: colwise weight cast enters op 2 TRANSPOSED into [G, N=D, K=2F]. + r["dx"] = _OPS.mxfp8_cudnn_grouped_mm( + r["dz_q"], + r["dz_sf"], + c["w13c_q"].transpose(-2, -1), + c["w13c_sf"], + c["offsets"], + ) + # Production wgrad layout mixes: native dy x kernel h; kernel dz x native x. + r["dw2"] = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + c["dy_colq"], c["dy_col_sf"], r["h_colq"], r["h_col_sf"], c["offsets"] + ) + r["dw13"] = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + r["dz_colq"], r["dz_col_sf"], c["x_colq"], c["x_col_sf"], c["offsets"] + ) + return r + + +@_gpu +@pytest.mark.parametrize("case", list(_CASES)) +def test_chain_numerics(case): + D, hidden, sizes = _CASES[case] + c = _build_case(D, hidden, sizes) + G, R = c["G"], c["R"] + r = _run_chain(c) + + # --- z: refA (dequantized operands, two reduction orders) + z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], sizes, transpose_b=True) + z_ref2 = _grouped_matmul( + c["x_deq"], c["w13_deq"], sizes, transpose_b=True, chunks=4 + ) + gate_a = _refA_gate(z_ref, z_ref2) + z_db = compute_error(z_ref.bfloat16(), r["z"]).item() + assert z_db >= gate_a, f"z {z_db:.1f} dB < derived refA gate {gate_a:.1f}" + + # --- z: refB (exact chain from ORIGINAL bf16 tensors; no quant helpers) + z_exact = _grouped_matmul( + c["x"].float(), [w.float() for w in c["w13"]], sizes, transpose_b=True + ) + band_b = compute_error(z_exact, z_ref).item() # quantization band + z_db_b = compute_error(z_exact.bfloat16(), r["z"]).item() + assert z_db_b >= band_b - 6.0, ( + f"z vs independent exact chain {z_db_b:.1f} dB < band {band_b:.1f} - 6" + ) + + # --- h (both quantized orientations) vs silu ref from the KERNEL's z + gate_f, up_f = _zsplit(r["z"].float(), hidden) + h_ref = F.silu(gate_f) * up_f + band_h = compute_error( + h_ref, _dequant_rowwise(*_quant_rowwise(h_ref.bfloat16())) + ).item() + h_deq = _dequant_rowwise(r["h_q"], r["h_sf"]) + h_db = compute_error(h_ref, h_deq).item() + assert h_db >= band_h - 6.0, f"h {h_db:.1f} dB < requant band {band_h:.1f} - 6" + h_col_deq = _dequant_colwise_grouped(r["h_colq"], r["h_col_sf"], sizes, hidden) + h_col_db = compute_error(h_ref, h_col_deq).item() + assert h_col_db >= band_h - 6.0, ( + f"h_col {h_col_db:.1f} dB < requant band {band_h:.1f} - 6 " + "(a whole-matrix-vs-per-group scale layout bug lands at 2-5 dB)" + ) + + # --- y: refA from the op's own quantized h + refB independent chain + w2_deq = c["w2_deq"] + y_ref = _grouped_matmul(h_deq, w2_deq, sizes, transpose_b=True) + y_ref2 = _grouped_matmul(h_deq, w2_deq, sizes, transpose_b=True, chunks=4) + y_gate = _refA_gate(y_ref, y_ref2) + y_db = compute_error(y_ref.bfloat16(), r["y"]).item() + assert y_db >= y_gate, f"y {y_db:.1f} dB < derived refA gate {y_gate:.1f}" + # refB for y: the whole forward computed from ORIGINAL bf16 tensors only. + gate_x, up_x = _zsplit(z_exact, hidden) + y_exact = _grouped_matmul( + F.silu(gate_x) * up_x, [w.float() for w in c["w2"]], sizes, transpose_b=True + ) + y_band_b = compute_error(y_exact, y_ref).item() + y_db_b = compute_error(y_exact.bfloat16(), r["y"]).item() + assert y_db_b >= y_band_b - 6.0, ( + f"y vs independent chain {y_db_b:.1f} dB < band {y_band_b:.1f} - 6" + ) + + # --- dz vs closed-form dSwiGLU from the kernel's z + dy_deq = _dequant_rowwise(c["dy_q"], c["dy_sf"]) + w2c_deq = [ + _dequant_colwise_grouped(c["w2c_q"][g], c["w2c_sf"].view(G, -1)[g], [D], hidden) + for g in range(G) + ] + dh_ref = _grouped_matmul(dy_deq, w2c_deq, sizes, transpose_b=False) + dgate, dup = _dswiglu(dh_ref, gate_f, up_f) + dz_ref = torch.empty(R, 2 * hidden, dtype=torch.float32, device="cuda") + v = dz_ref.view(R, hidden // _BLOCK, 2, _BLOCK) + v[:, :, 0, :] = dgate.view(R, hidden // _BLOCK, _BLOCK) + v[:, :, 1, :] = dup.view(R, hidden // _BLOCK, _BLOCK) + band_dz = compute_error( + dz_ref, _dequant_rowwise(*_quant_rowwise(dz_ref.bfloat16())) + ).item() + dz_deq = _dequant_rowwise(r["dz_q"], r["dz_sf"]) + dz_db = compute_error(dz_ref, dz_deq).item() + assert dz_db >= band_dz - 6.0, f"dz {dz_db:.1f} dB < band {band_dz:.1f} - 6" + + # --- dx refA + w13c_deq = [ + _dequant_colwise_grouped( + c["w13c_q"][g], c["w13c_sf"].view(G, -1)[g], [2 * hidden], D + ) + for g in range(G) + ] + dx_ref = _grouped_matmul(dz_deq, w13c_deq, sizes, transpose_b=False) + dx_ref2 = _grouped_matmul(dz_deq, w13c_deq, sizes, transpose_b=False, chunks=4) + dx_gate = _refA_gate(dx_ref, dx_ref2) + dx_db = compute_error(dx_ref.bfloat16(), r["dx"]).item() + assert dx_db >= dx_gate, f"dx {dx_db:.1f} dB < derived refA gate {dx_gate:.1f}" + + # --- wgrads refA (production layout mixes) + dy_col_deq = _dequant_colwise_grouped(c["dy_colq"], c["dy_col_sf"], sizes, D) + dz_col_deq = _dequant_colwise_grouped( + r["dz_colq"], r["dz_col_sf"], sizes, 2 * hidden + ) + x_col_deq = _dequant_colwise_grouped(c["x_colq"], c["x_col_sf"], sizes, D) + off = 0 + dw2_ref = torch.zeros(G, D, hidden, dtype=torch.float32, device="cuda") + dw13_ref = torch.zeros(G, 2 * hidden, D, dtype=torch.float32, device="cuda") + for g, m in enumerate(sizes): + dw2_ref[g] = dy_col_deq[off : off + m].t() @ h_col_deq[off : off + m] + dw13_ref[g] = dz_col_deq[off : off + m].t() @ x_col_deq[off : off + m] + off += m + dw2_db = compute_error(dw2_ref.bfloat16(), r["dw2"]).item() + dw13_db = compute_error(dw13_ref.bfloat16(), r["dw13"]).item() + assert dw2_db >= 50.0, f"dw2 {dw2_db:.1f} dB < 50 (probe level: 98-155)" + assert dw13_db >= 50.0, f"dw13 {dw13_db:.1f} dB < 50 (probe level: 91-160)" + + # zero-token experts must come back written as exact zeros + for g, m in enumerate(sizes): + if m == 0: + assert (r["dw2"][g] == 0).all() and (r["dw13"][g] == 0).all(), ( + f"zero-token expert {g} weight gradients must be exactly zero" + ) + + print( + f"\n[{case}] derived gates/bands (dB): " + f"z refA_gate={gate_a:.1f} (got {z_db:.1f}), " + f"z refB band={band_b:.1f} (got {z_db_b:.1f}), " + f"h band={band_h:.1f} (row {h_db:.1f}, col {h_col_db:.1f}), " + f"y refA_gate={y_gate:.1f} (got {y_db:.1f}), " + f"y refB band={y_band_b:.1f} (got {y_db_b:.1f}), " + f"dz band={band_dz:.1f} (got {dz_db:.1f}), " + f"dx refA_gate={dx_gate:.1f} (got {dx_db:.1f}), " + f"dw2 {dw2_db:.1f}, dw13 {dw13_db:.1f}" + ) + + +# --------------------------------------------------------------------------- +# Wgrad stride matrix: both operands in each major, all four combinations. +# --------------------------------------------------------------------------- + + +@_gpu +@pytest.mark.parametrize("a_native", [False, True], ids=["aRM", "aNat"]) +@pytest.mark.parametrize("b_native", [False, True], ids=["bRM", "bNat"]) +def test_wgrad_stride_matrix(dbg, a_native, b_native): + c = dbg + sizes, D = c["sizes"], c["D"] + dy_q, dy_sf = _quant_colwise_grouped(c["dy"], sizes, native=a_native) + x_q, x_sf = _quant_colwise_grouped(c["x"], sizes, native=b_native) + dw = _OPS.mxfp8_cudnn_grouped_mlp_wgrad(dy_q, dy_sf, x_q, x_sf, c["offsets"]) + dy_deq = _dequant_colwise_grouped(dy_q, dy_sf, sizes, D) + x_deq = _dequant_colwise_grouped(x_q, x_sf, sizes, D) + ref = torch.zeros(c["G"], D, D, dtype=torch.float32, device="cuda") + off = 0 + for g, m in enumerate(sizes): + ref[g] = dy_deq[off : off + m].t() @ x_deq[off : off + m] + off += m + db = compute_error(ref.bfloat16(), dw).item() + assert db >= 50.0, f"wgrad[{a_native=} {b_native=}] {db:.1f} dB < 50" + + +# --------------------------------------------------------------------------- +# A < R strict tail with planted garbage. +# --------------------------------------------------------------------------- + + +@_gpu +def test_tail_a_lt_r_poisoned(): + D = hidden = 256 + sizes = [256, 0, 512, 256] + A, R = sum(sizes), 1280 + torch.manual_seed(3) + dev = "cuda" + offsets = _mk_offsets(sizes, dev) + x = torch.randn(R, D, dtype=torch.bfloat16, device=dev) * 0.5 + dy = torch.randn(R, D, dtype=torch.bfloat16, device=dev) * 0.5 + x[A:] = float("nan") + dy[A::2] = float("inf") + dy[A + 1 :: 2] = float("nan") + w13 = _to_32block( + torch.randn(4, hidden, 2, D, dtype=torch.bfloat16, device=dev) * 0.02 + ) + w2 = torch.randn(4, D, hidden, dtype=torch.bfloat16, device=dev) * 0.02 + + x_q, x_sf = _quant_rowwise(x) + dy_q, dy_sf = _quant_rowwise(dy) + w13_q, w13_sf = _quant_weight_rowwise(w13) + w2_q, w2_sf = _quant_weight_rowwise(w2) + w2c_q, w2c_sf = _quant_weight_colwise(w2) + + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + assert not z[:A].isnan().any(), "active z rows contaminated by the poisoned tail" + y = _OPS.mxfp8_cudnn_grouped_mm(h_q, h_sf, w2_q, w2_sf, offsets) + assert not y[:A].isnan().any(), "active y rows contaminated" + dz_q, dz_sf, dz_colq, dz_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_bwd( + dy_q, dy_sf, w2c_q, w2c_sf, z, offsets + ) + w13c_q, w13c_sf = _quant_weight_colwise(w13) + dx = _OPS.mxfp8_cudnn_grouped_mm( + dz_q, dz_sf, w13c_q.transpose(-2, -1), w13c_sf, offsets + ) + assert not dx[:A].isnan().any(), "active dx rows contaminated" + + # wgrad: colwise scales cover only the routed A rows; the qdata tail is + # additionally poisoned with NaN bytes and must never be read. + dy_colq, dy_col_sf = _quant_colwise_grouped(dy[:A], sizes, native=True) + dy_colq_full = _cat8( + [ + dy_colq.contiguous(), + torch.full((R - A, D), 0x7F, dtype=torch.uint8, device=dev).view(_E4M3), + ], + 0, + ) + dw2 = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + dy_colq_full, dy_col_sf, h_colq, h_col_sf, offsets + ) + assert not dw2.isnan().any(), "wgrad read the NaN-poisoned inactive tail" + dy_col_deq = _dequant_colwise_grouped(dy_colq, dy_col_sf, sizes, D) + h_col_deq = _dequant_colwise_grouped(h_colq[:A], h_col_sf, sizes, hidden) + ref = torch.zeros(4, D, hidden, dtype=torch.float32, device=dev) + off = 0 + for g, m in enumerate(sizes): + ref[g] = dy_col_deq[off : off + m].t() @ h_col_deq[off : off + m] + off += m + db = compute_error(ref.bfloat16(), dw2).item() + assert db >= 50.0, f"tail-poisoned dw2 {db:.1f} dB < 50" + + +# --------------------------------------------------------------------------- +# Determinism, compile, R == 0. +# --------------------------------------------------------------------------- + + +@_gpu +def test_determinism_all_ops_bitwise(dbg): + c = dbg + r1 = _run_chain(c) + r2 = _run_chain(c) + for key in r1: + assert torch.equal(_bytes(r1[key]), _bytes(r2[key])), ( + f"{key} is not bitwise deterministic across identical launches" + ) + + +@_gpu +def test_compile_fullgraph_bitwise(dbg): + c = dbg + + def fwd_then_mm(x_q, x_sf, w13_q, w13_sf, w2_q, w2_sf, offsets): + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + y = _OPS.mxfp8_cudnn_grouped_mm(h_q, h_sf, w2_q, w2_sf, offsets) + return z, h_q, y + + eager = fwd_then_mm( + c["x_q"], + c["x_sf"], + c["w13_q"], + c["w13_sf"], + c["w2_q"], + c["w2_sf"], + c["offsets"], + ) + compiled = torch.compile(fwd_then_mm, fullgraph=True)( + c["x_q"], + c["x_sf"], + c["w13_q"], + c["w13_sf"], + c["w2_q"], + c["w2_sf"], + c["offsets"], + ) + for e, co, name in zip(eager, compiled, ("z", "h_q", "y")): + assert torch.equal(_bytes(e), _bytes(co)), f"compiled {name} != eager" + + +@_gpu +def test_r0_all_ops(): + dev = "cuda" + D = hidden = 256 + offsets = torch.zeros(2, dtype=torch.int32, device=dev) + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + torch.empty(0, D, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.zeros(2, 2 * hidden, D, dtype=torch.uint8, device=dev).view(_E4M3), + torch.empty(2 * 2 * hidden * D // _BLOCK, dtype=_E8M0, device=dev), + offsets, + ) + assert z.shape == (0, 2 * hidden) and h_q.shape == (0, hidden) + assert h_sf.numel() == 0 and h_col_sf.numel() == 0 + dw = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + torch.empty(0, D, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.empty(0, hidden, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + offsets, + ) + assert dw.shape == (2, D, hidden) and (dw == 0).all() + + +# --------------------------------------------------------------------------- +# Negative controls: each sabotage must fail the numerics gates decisively. +# --------------------------------------------------------------------------- + + +@_gpu +def test_negative_control_whole_matrix_colwise_scales(dbg): + """Whole-matrix to_blocked colwise scales: same bytes, silently wrong order.""" + c = dbg + sizes, D, G = c["sizes"], c["D"], c["G"] + dy_q, dy_sf_pg = _quant_colwise_grouped(c["dy"], sizes, native=False) + x_q, x_sf_pg = _quant_colwise_grouped(c["x"], sizes, native=False) + # Rebuild the SAME logical scales in whole-matrix blocked order. + s_t, _ = to_mx(c["dy"].t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + dy_sf_wm = to_blocked(s_t.view(_E8M0)).view(_E8M0) + assert dy_sf_wm.numel() == dy_sf_pg.numel() + good = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + dy_q, dy_sf_pg, x_q, x_sf_pg, c["offsets"] + ) + bad = _OPS.mxfp8_cudnn_grouped_mlp_wgrad(dy_q, dy_sf_wm, x_q, x_sf_pg, c["offsets"]) + dy_deq = _dequant_colwise_grouped(dy_q, dy_sf_pg, sizes, D) + x_deq = _dequant_colwise_grouped(x_q, x_sf_pg, sizes, D) + ref = torch.zeros(G, D, D, dtype=torch.float32, device="cuda") + off = 0 + for g, m in enumerate(sizes): + ref[g] = dy_deq[off : off + m].t() @ x_deq[off : off + m] + off += m + good_db = compute_error(ref.bfloat16(), good).item() + bad_db = compute_error(ref.bfloat16(), bad).item() + assert good_db >= 50.0 + assert bad_db < 25.0, ( + f"whole-matrix colwise scales scored {bad_db:.1f} dB -- the negative " + f"control lost its teeth (good arm: {good_db:.1f})" + ) + + +@_gpu +def test_negative_control_gate_up_swap(dbg): + """Swapping the gate/up 32-blocks must collapse h against the correct ref.""" + c = dbg + hidden, G, D = c["F"], c["G"], c["D"] + w13_sw = ( + c["w13"] + .view(G, hidden // _BLOCK, 2, _BLOCK, D) + .flip(2) + .reshape(G, 2 * hidden, D) + .contiguous() + ) + w13_sw_q, w13_sw_sf = _quant_weight_rowwise(w13_sw) + _, h_q, h_sf, _, _ = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + c["x_q"], c["x_sf"], w13_sw_q, w13_sw_sf, c["offsets"] + ) + z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True) + gate_f, up_f = _zsplit(z_ref, hidden) + h_ref = F.silu(gate_f) * up_f + good_db = compute_error( + h_ref, + _dequant_rowwise( + *( + _OPS.mxfp8_cudnn_grouped_mlp_fwd( + c["x_q"], c["x_sf"], c["w13_q"], c["w13_sf"], c["offsets"] + )[1:3] + ) + ), + ).item() + bad_db = compute_error(h_ref, _dequant_rowwise(h_q, h_sf)).item() + assert bad_db < good_db - 10.0, ( + f"gate/up swap only moved h from {good_db:.1f} to {bad_db:.1f} dB -- " + "the 32-block order convention is not actually being exercised" + ) + + +@_gpu +def test_negative_control_scale_byte_flip(dbg): + """One +2-code E8M0 flip (x4) in the weight scales must break refA.""" + c = dbg + sf_bad = c["w13_sf"].view(torch.uint8).clone() + sf_bad[sf_bad.numel() // 2] += 2 + z_bad = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + c["x_q"], c["x_sf"], c["w13_q"], sf_bad, c["offsets"] + )[0] + z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True) + z_ref2 = _grouped_matmul( + c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True, chunks=4 + ) + gate = _refA_gate(z_ref, z_ref2) + bad_db = compute_error(z_ref.bfloat16(), z_bad).item() + assert bad_db < gate, ( + f"single scale-byte flip still passes refA ({bad_db:.1f} >= {gate:.1f} dB)" + ) + + +@_gpu +def test_kernel_scale_mode_is_rceil(dbg): + """The fwd op's h scale bytes must match RCEIL, and not FLOOR, quantization.""" + c = dbg + r = _run_chain(c) + gate_f, up_f = _zsplit(r["z"].float(), c["F"]) + h_ref = (F.silu(gate_f) * up_f).bfloat16() + s_rceil, _ = to_mx(h_ref, _E4M3, _BLOCK, scaling_mode=_RCEIL) + s_floor, _ = to_mx(h_ref, _E4M3, _BLOCK, scaling_mode=ScaleCalculationMode.FLOOR) + got = from_blocked(r["h_sf"].view(_E8M0), c["R"], c["F"] // _BLOCK).view( + torch.uint8 + ) + rceil_frac = (got == s_rceil.view(torch.uint8)).float().mean().item() + floor_frac = (got == s_floor.view(torch.uint8)).float().mean().item() + assert rceil_frac > 0.98, f"h scales match RCEIL on only {rceil_frac:.3f}" + assert rceil_frac > floor_frac + 0.1, ( + f"RCEIL ({rceil_frac:.3f}) does not dominate FLOOR ({floor_frac:.3f})" + ) + + +# --------------------------------------------------------------------------- +# Validation: rejection matrix and the opt-in offsets path. +# --------------------------------------------------------------------------- + + +def _valid_fwd_args(device="cuda", R=512, D=256, hidden=256, G=2): + N1 = 2 * hidden + return dict( + x_q=torch.zeros(R, D, dtype=_E4M3, device=device), + x_sf=torch.zeros(R * D // _BLOCK, dtype=_E8M0, device=device), + w13_q=torch.zeros(G, N1, D, dtype=_E4M3, device=device), + w13_sf=torch.zeros(G * N1 * D // _BLOCK, dtype=_E8M0, device=device), + offsets=torch.tensor([256, 512], dtype=torch.int32, device=device), + ) + + +_NEGATIVES = [ + ("x_q_dtype", lambda a: a.update(x_q=a["x_q"].view(torch.int8)), "float8_e4m3fn"), + ("x_q_cpu", lambda a: a.update(x_q=a["x_q"].cpu()), "CUDA"), + ( + "r_not_256", + lambda a: a.update( + x_q=torch.zeros(384, 256, dtype=_E4M3, device="cuda"), + x_sf=torch.zeros(384 * 8, dtype=_E8M0, device="cuda"), + ), + "multiple of 256", + ), + ( + "d_192", + lambda a: a.update( + x_q=torch.zeros(512, 192, dtype=_E4M3, device="cuda"), + x_sf=torch.zeros(512 * 6, dtype=_E8M0, device="cuda"), + w13_q=torch.zeros(2, 512, 192, dtype=_E4M3, device="cuda"), + w13_sf=torch.zeros(2 * 512 * 6, dtype=_E8M0, device="cuda"), + ), + "multiple of 128", + ), + ( + "offsets_i64", + lambda a: a.update(offsets=a["offsets"].to(torch.int64)), + "int32", + ), + ( + "offsets_wrong_len", + lambda a: a.update(offsets=a["offsets"][:1]), + "one entry per local expert", + ), + ( + "g0", + lambda a: a.update( + w13_q=torch.zeros(0, 512, 256, dtype=_E4M3, device="cuda"), + w13_sf=torch.zeros(0, dtype=_E8M0, device="cuda"), + offsets=torch.zeros(0, dtype=torch.int32, device="cuda"), + ), + "at least one expert group", + ), + ( + "x_sf_short", + lambda a: a.update(x_sf=a["x_sf"][:-8].clone()), + "blocked scale bytes", + ), + ( + "w13_stride", + lambda a: a.update( + w13_q=a["w13_q"].transpose(-2, -1).contiguous().transpose(-2, -1) + ), + "stride", + ), +] + + +@_gpu +@pytest.mark.parametrize("case", _NEGATIVES, ids=[c[0] for c in _NEGATIVES]) +def test_validation_negatives(case): + _name, mutate, needle = case + args = _valid_fwd_args() + mutate(args) + with pytest.raises(ValueError) as exc_info: + _OPS.mxfp8_cudnn_grouped_mlp_fwd(**args) + assert needle.lower() in str(exc_info.value).lower(), ( + f"rejection message {str(exc_info.value)!r} does not name the defect " + f"({needle!r})" + ) + + +@_gpu +def test_optin_offsets_validation(monkeypatch): + args = _valid_fwd_args() + bad = dict(args) + bad["offsets"] = torch.tensor([128, 512], dtype=torch.int32, device="cuda") + + # Default build: metadata-only, misaligned VALUES are not (and cannot be) + # caught without a D2H sync. + monkeypatch.delenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", raising=False) + _OPS.mxfp8_cudnn_grouped_mlp_fwd(**bad) + + monkeypatch.setenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", "1") + with pytest.raises(ValueError, match="FIX_PAD_SIZE"): + _OPS.mxfp8_cudnn_grouped_mlp_fwd(**bad) + + over = dict(args) + over["offsets"] = torch.tensor([256, 768], dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="exceeds the allocated row count"): + _OPS.mxfp8_cudnn_grouped_mlp_fwd(**over) + + # The opt-in check must not break fake tracing (no values to read). + with FakeTensorMode(): + _fake_chain_shapes() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From e03489ada4e8724797e603bffbfe741fb5372718 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 20:30:36 -0700 Subject: [PATCH 06/11] Remove the archived pure-CuTe-DSL MXFP8 grouped-MLP kernel family (A/B/C) The cuDNN-frontend grouped-MLP ops (cudnn_grouped_mlp_{ops,validation} + the public wrapper) supersede the in-repo CuTe DSL kernels: same fused FC1+SwiGLU+dual-quant / FC2-dgrad+dSwiGLU / ragged-wgrad surface, with no dependency on this repo's CuTe DSL runtime shims. Delete the kernels (cutedsl_grouped_mlp.py), their ops/validation modules, the public mxfp8_grouped_mlp wrapper, tests, and bench; unwire both __init__ files. The quantizer/swizzle kernels in this package (quant.py, cutedsl/flydsl quantizers, cute_utils) are untouched -- the unfused baseline path and the cudnn composite's casts still use them. torchao.prototype.moe_training import smoke and the 29 cudnn grouped-MLP tests stay green. Co-Authored-By: Claude Fable 5 --- .../moe_training/mxfp8/bench_grouped_mlp.py | 759 -------- .../moe_training/test_mxfp8_grouped_mlp.py | 1029 ----------- torchao/prototype/moe_training/__init__.py | 12 - .../moe_training/kernels/mxfp8/__init__.py | 10 +- .../kernels/mxfp8/cutedsl_grouped_mlp.py | 1606 ----------------- .../kernels/mxfp8/grouped_mlp_ops.py | 552 ------ .../kernels/mxfp8/grouped_mlp_validation.py | 267 --- .../moe_training/mxfp8_grouped_mlp.py | 220 --- 8 files changed, 4 insertions(+), 4451 deletions(-) delete mode 100644 benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py delete mode 100644 test/prototype/moe_training/test_mxfp8_grouped_mlp.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py delete mode 100644 torchao/prototype/moe_training/mxfp8_grouped_mlp.py diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py b/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py deleted file mode 100644 index 0ae723aea7..0000000000 --- a/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py +++ /dev/null @@ -1,759 +0,0 @@ -# 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. -"""Benchmark the fused MXFP8 grouped-MLP kernel family against the existing -decomposed torchao path and TransformerEngine, on identical inputs, offsets, -shapes and RCEIL scale mode. - -Lanes (``--lane``): - -* ``torchao`` -- the existing decomposed SM100 path: triton/CUDA quantizers + - ``torch._scaled_grouped_mm`` + eager SwiGLU/dSwiGLU, staged to match each - fused kernel's covered work. Uses no CuTe DSL code. -* ``ours`` -- the three fused ops ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` - / ``_dswiglu_bwd`` / ``_wgrad`` (one kernel launch each). -* ``te`` -- TransformerEngine: the fused CuTe-DSL lane - (``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``; per-kernel times recovered from a - profiler pass by kernel-name fragment) plus the modular lane's single-kernel - ``tex.swiglu`` / ``tex.dswiglu`` gated-activation+dual-quantize points. -* ``all`` -- ``torchao`` + ``ours``. - -The TE lane must run in a separate process from ``ours``: our kernels need the -public ``nvidia-cutlass-dsl`` 4.7.0 wheel on the user site, while TE's fused -lane uses the container-native cuDNN/cutlass stack, which that wheel shadows. -Invocation on the GB200 dev host:: - - # ours / torchao lanes - bash ./run_te.sh env PYTHONUSERBASE=/.local PYTHONPATH=/ao \\ - python /ao/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py --lane all - # TE lane (no PYTHONUSERBASE) - bash ./run_te.sh env PYTHONPATH=/ao \\ - python /ao/benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py --lane te - -Caveats stated up front so the numbers are read honestly: - -* The decomposed wgrad stage includes its own dim1 (columnwise) quantization of - both operands, because that is what the existing path pays; the fused wgrad - consumes columnwise operands produced by kernels A/B. The A/B stages of both - lanes start from identically prequantized GEMM inputs. -* The TE fused forward also applies per-token router probs in-kernel; we pass - probs = 1 so the work matches. -* Eager launch timing (``do_bench`` median) only. No CUDA-graph replay column. -* Absolute microseconds from a clock-capped host (this dev box pins app clocks - at 1200 MHz) are not publishable; the startup banner prints the clocks. - -``MXFP8_BENCH_VALIDATE=1`` cross-checks the fused outputs against pure-torch -``to_mx``/``to_blocked`` references (SQNR gates; the checked-in test suite owns -the bitwise contracts). -""" - -import argparse -import os -from dataclasses import dataclass -from typing import Callable, Dict, List, Optional - -import torch -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 ( - grouped_mlp_ops, # noqa: F401 (registers the three fused ops) - mx_block_rearrange_2d_M_groups_cuda, - triton_mx_block_rearrange_2d_K_groups, -) -from torchao.prototype.moe_training.utils import generate_jagged_offs -from torchao.prototype.mx_formats.config import ( - MXFP8Dim1CastKernelChoice, - ScaleCalculationMode, -) -from torchao.prototype.mx_formats.kernels import ( - mxfp8_quantize_cuda, - triton_to_mxfp8_dim0, -) -from torchao.prototype.mx_formats.mx_tensor import to_mx -from torchao.prototype.mx_formats.utils import ( - _to_mxfp8_dim1_kernel_wrapper, - from_blocked, - to_blocked, -) -from torchao.quantization.quantize_.common import KernelPreference - -device = torch.device("cuda") -VALIDATE = os.environ.get("MXFP8_BENCH_VALIDATE", "0") == "1" -BLOCK = 32 -RCEIL = ScaleCalculationMode.RCEIL - - -# -------------------------------------------------------------------------- -# Configs -# -------------------------------------------------------------------------- - - -@dataclass(frozen=True) -class ExperimentConfig: - rows: int # R: padded token rows (sum of per-expert rows) - model_dim: int # D - hidden_dim: int # F - num_groups: int # G (local experts) - distribution: str # "balanced" | "skewed" - - -@dataclass(frozen=True) -class ExperimentResult: - # Per fused-kernel-equivalent stage, microseconds (median). - a_us: float - b_us: float - c_fc1_us: float - c_fc2_us: float - seq_us: float - # Derived TFLOP/s for the GEMM in each stage. - a_tflops: float - b_tflops: float - c_fc1_tflops: float - c_fc2_tflops: float - - -@dataclass(frozen=True) -class Experiment: - lane: str - config: ExperimentConfig - result: ExperimentResult - - -def get_configs(args) -> List[ExperimentConfig]: - if args.shape is not None: - r, d, f, g = (int(v) for v in args.shape.split(",")) - return [ExperimentConfig(r, d, f, g, dist) for dist in args.dists] - shapes = [ - # smoke - (512, 256, 256, 2), - # DeepSeekV3 16B class (D=2048, F=1408, G=8 local experts) - (2048, 2048, 1408, 8), - (8192, 2048, 1408, 8), - (16384, 2048, 1408, 8), - # DeepSeekV3 671B class (D=7168, F=2048, G=4 local experts) - (2048, 7168, 2048, 4), - (8192, 7168, 2048, 4), - (16384, 7168, 2048, 4), - ] - return [ - ExperimentConfig(r, d, f, g, dist) - for (r, d, f, g) in shapes - for dist in args.dists - ] - - -def make_offsets(cfg: ExperimentConfig) -> torch.Tensor: - """Exclusive per-expert end offsets, every group a multiple of 128.""" - r, g = cfg.rows, cfg.num_groups - if cfg.distribution == "balanced": - per = r // g - if per % 128 != 0 or per * g != r: - raise ValueError( - f"balanced distribution needs R/G to be a 128 multiple, got {r}/{g}" - ) - return torch.arange(1, g + 1, device=device, dtype=torch.int32) * per - return generate_jagged_offs(g, r, multiple_of=128, device=device) - - -# -------------------------------------------------------------------------- -# Pure-torch quantization recipes (input prep + validation only, never timed) -# -------------------------------------------------------------------------- - - -def ref_quantize_rowwise_1x32(x: torch.Tensor): - """[M, K] -> (E4M3 [M, K] row-major, flat blocked E8M0 for [M, K/32]).""" - scale, q = to_mx(x, torch.float8_e4m3fn, BLOCK, scaling_mode=RCEIL) - return q, to_blocked(scale) - - -def ref_quantize_colwise_32x1(x: torch.Tensor): - """[R, N] -> (E4M3 [R, N] stride (1, R), flat blocked E8M0 for [N, R/32]).""" - scale_t, q_t = to_mx( - x.t().contiguous(), torch.float8_e4m3fn, BLOCK, scaling_mode=RCEIL - ) - return q_t.t(), to_blocked(scale_t) - - -def ref_dequant_colwise(q_col: torch.Tensor, sf_blocked: torch.Tensor): - """FP32 dequant of a columnwise operand (for validation oracles).""" - rows, cols = q_col.shape - logical = from_blocked(sf_blocked, cols, rows // BLOCK) # [N, R/32] - scales = logical.t().to(torch.float32).repeat_interleave(BLOCK, dim=0) - return q_col.to(torch.float32) * scales - - -def ref_dequant_rowwise(q_row: torch.Tensor, sf_blocked: torch.Tensor): - """FP32 dequant of a rowwise 1x32-quantized operand (for validation oracles).""" - rows, cols = q_row.shape - scales = from_blocked(sf_blocked, rows, cols // BLOCK).to(torch.float32) - return q_row.to(torch.float32) * scales.repeat_interleave(BLOCK, dim=1) - - -def sqnr(ref: torch.Tensor, actual: torch.Tensor) -> float: - err = (ref - actual).float().pow(2).mean() - if err == 0: - return float("inf") - return (10 * torch.log10(ref.float().pow(2).mean() / err)).item() - - -# -------------------------------------------------------------------------- -# Shared input bundle -# -------------------------------------------------------------------------- - - -class Inputs: - """All operands both non-TE lanes consume, prepared once per config. - - GEMM operands are prequantized identically for both lanes (kernels A and B - take prequantized inputs by contract, and the decomposed path accepts - prequantized MX operands at the same seam). - """ - - def __init__(self, cfg: ExperimentConfig): - r, d, f, g = cfg.rows, cfg.model_dim, cfg.hidden_dim, cfg.num_groups - torch.manual_seed(0) - self.offsets = make_offsets(cfg) - self.offsets_host = self.offsets.tolist() - - self.x = torch.randn(r, d, device=device, dtype=torch.bfloat16) / d**0.5 - self.do = torch.randn(r, d, device=device, dtype=torch.bfloat16) / d**0.5 - # Element-interleaved gate/up FC1 weight [G, 2F, D] and FC2-dgrad - # weight-view source [G, F, D]; both quantized along D (the GEMM - # contraction), then freely transposed into the K-major ABI layouts. - self.w13i = ( - torch.randn(g, 2 * f, d, device=device, dtype=torch.bfloat16) / d**0.5 - ) - self.w2d = torch.randn(g, f, d, device=device, dtype=torch.bfloat16) / d**0.5 - - self.x_q, self.x_sf = ref_quantize_rowwise_1x32(self.x) - self.do_q, self.do_sf = ref_quantize_rowwise_1x32(self.do) - - w13_q, w13_sf = zip( - *(ref_quantize_rowwise_1x32(self.w13i[i]) for i in range(g)) - ) - self.w13_t_q = torch.stack(list(w13_q)).transpose(-2, -1) # [G, D, 2F] - self.w13_t_sf = torch.stack(list(w13_sf)) - w2_q, w2_sf = zip(*(ref_quantize_rowwise_1x32(self.w2d[i]) for i in range(g))) - self.w2_t_q = torch.stack(list(w2_q)).transpose(-2, -1) # [G, D, F] - self.w2_t_sf = torch.stack(list(w2_sf)) - - # Reference forward intermediates (bf16), computed per expert from the - # DEQUANTIZED operands — the values the fused kernels consume by - # contract — so the validation SQNR isolates each kernel's own work - # instead of stacking input-quantization error on top of it. The same - # z then feeds kernel B and the decomposed dSwiGLU identically. - x_f32 = ref_dequant_rowwise(self.x_q, self.x_sf) - do_f32 = ref_dequant_rowwise(self.do_q, self.do_sf) - z = torch.zeros(r, 2 * f, device=device, dtype=torch.bfloat16) - prev = 0 - for i in range(g): - end = self.offsets_host[i] - if end > prev: - w13_f32 = ref_dequant_rowwise(w13_q[i], w13_sf[i]) - z[prev:end] = (x_f32[prev:end] @ w13_f32.t()).to(torch.bfloat16) - prev = end - self.z_flat = z - self.z_bf16 = z.view(r, f, 2) - gate = self.z_bf16[..., 0].float() - up = self.z_bf16[..., 1].float() - self.h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) - sig = torch.sigmoid(gate) - dh = torch.zeros(r, f, device=device, dtype=torch.bfloat16) - prev = 0 - for i in range(g): - end = self.offsets_host[i] - if end > prev: - # do [m, D] contracts with the [D, F] dgrad weight view; the - # previous `do @ w2d[i]` only type-checked when D == F and - # computed the transpose of the intended dgrad. - w2_f32 = ref_dequant_rowwise(w2_q[i], w2_sf[i]) - dh[prev:end] = (do_f32[prev:end] @ w2_f32.t()).to(torch.bfloat16) - prev = end - dhf = dh.float() - dgate = (dhf * up * (sig * (1.0 + gate * (1.0 - sig)))).to(torch.bfloat16) - dup = (dhf * (gate * sig)).to(torch.bfloat16) - self.dz_flat = torch.stack((dgate, dup), dim=-1).view(r, 2 * f) - - # Columnwise operands for the two wgrad calls (produced by A/B in the - # fused regime, by standalone quantizers in the decomposed one). - self.dz_col_q, self.dz_col_sf = ref_quantize_colwise_32x1(self.dz_flat) - self.x_col_q, self.x_col_sf = ref_quantize_colwise_32x1(self.x) - self.do_col_q, self.do_col_sf = ref_quantize_colwise_32x1(self.do) - self.h_col_q, self.h_col_sf = ref_quantize_colwise_32x1(self.h) - - -def stage_flops(cfg: ExperimentConfig) -> Dict[str, float]: - r, d, f = cfg.rows, cfg.model_dim, cfg.hidden_dim - return { - "a": 2.0 * r * d * 2 * f, - "b": 2.0 * r * d * f, - "c_fc1": 2.0 * r * 2 * f * d, - "c_fc2": 2.0 * r * d * f, - } - - -# -------------------------------------------------------------------------- -# Lane: ours (the three fused ops) -# -------------------------------------------------------------------------- - - -def lane_ours(cfg: ExperimentConfig, inp: Inputs) -> Dict[str, Callable]: - ops = torch.ops.torchao - - def a(): - return ops.mxfp8_grouped_gemm_swiglu_fwd( - inp.x_q, inp.x_sf, inp.w13_t_q, inp.w13_t_sf, inp.offsets - ) - - def b(): - return ops.mxfp8_grouped_gemm_dswiglu_bwd( - inp.do_q, inp.do_sf, inp.w2_t_q, inp.w2_t_sf, inp.z_bf16, inp.offsets - ) - - def c_fc1(): - return ops.mxfp8_grouped_gemm_wgrad( - inp.dz_col_q, inp.dz_col_sf, inp.x_col_q, inp.x_col_sf, inp.offsets - ) - - def c_fc2(): - return ops.mxfp8_grouped_gemm_wgrad( - inp.do_col_q, inp.do_col_sf, inp.h_col_q, inp.h_col_sf, inp.offsets - ) - - def seq(): - _, _, _, h_col_q, h_col_sf = a() - _, _, dz_col_q, dz_col_sf = b() - ops.mxfp8_grouped_gemm_wgrad( - dz_col_q, dz_col_sf, inp.x_col_q, inp.x_col_sf, inp.offsets - ) - ops.mxfp8_grouped_gemm_wgrad( - inp.do_col_q, inp.do_col_sf, h_col_q, h_col_sf, inp.offsets - ) - - return {"a": a, "b": b, "c_fc1": c_fc1, "c_fc2": c_fc2, "seq": seq} - - -def validate_ours(cfg: ExperimentConfig, inp: Inputs, stages) -> None: - """SQNR cross-checks against pure-torch references. Bitwise contracts are - owned by test_mxfp8_grouped_mlp.py; this is a sanity gate for benching.""" - z_k, h_row_q, h_row_sf, h_col_q, h_col_sf = stages["a"]() - assert z_k.shape == inp.z_bf16.shape and z_k.stride() == inp.z_bf16.stride() - s = sqnr(inp.z_flat.float(), z_k.reshape(cfg.rows, -1).float()) - assert s >= 27.0, f"A z SQNR {s:.1f} < 27" - h_deq = ref_dequant_colwise(h_col_q, h_col_sf) - s = sqnr(inp.h.float(), h_deq) - assert s >= 27.0, f"A h (dequant colwise) SQNR {s:.1f} < 27" - row_deq = h_row_q.float() * ( - from_blocked(h_row_sf, cfg.rows, cfg.hidden_dim // BLOCK) - .to(torch.float32) - .repeat_interleave(BLOCK, dim=1) - ) - s = sqnr(inp.h.float(), row_deq) - assert s >= 27.0, f"A h (dequant rowwise) SQNR {s:.1f} < 27" - - dz_row_q, dz_row_sf, dz_col_q, dz_col_sf = stages["b"]() - dz_deq = ref_dequant_colwise(dz_col_q, dz_col_sf) - s = sqnr(inp.dz_flat.float(), dz_deq) - assert s >= 25.0, f"B dz SQNR {s:.1f} < 25" - - dw = stages["c_fc1"]() - dy_f32 = ref_dequant_colwise(inp.dz_col_q, inp.dz_col_sf) - x_f32 = ref_dequant_colwise(inp.x_col_q, inp.x_col_sf) - prev = 0 - ref = torch.zeros_like(dw, dtype=torch.float32) - for i in range(cfg.num_groups): - end = inp.offsets_host[i] - if end > prev: - ref[i] = dy_f32[prev:end].t() @ x_f32[prev:end] - prev = end - s = sqnr(ref, dw.float()) - assert s >= 24.0, f"C dw SQNR {s:.1f} < 24" - print(" validate(ours): OK (A z/h, B dz, C dw)") - - -# -------------------------------------------------------------------------- -# Lane: torchao decomposed (existing path; no CuTe DSL) -# -------------------------------------------------------------------------- - - -def lane_torchao(cfg: ExperimentConfig, inp: Inputs) -> Dict[str, Callable]: - r, f = cfg.rows, cfg.hidden_dim - offs = inp.offsets - - def dual_quantize(t: torch.Tensor): - # Rowwise via the triton dim0 quantizer (the CUDA kernel is - # colwise-only today), colwise via the CUDA quantizer, plus the two - # scale rearranges the existing SM100 path performs. - out_row, s_row = triton_to_mxfp8_dim0(t, BLOCK, "rceil") - s_row_blocked = mx_block_rearrange_2d_M_groups_cuda( - s_row.view(torch.uint8), offs - ) - _, out_col, _, s_col = mxfp8_quantize_cuda( - t, rowwise=False, colwise=True, scaling_mode="rceil" - ) - s_col_blocked = triton_mx_block_rearrange_2d_K_groups( - s_col.view(torch.uint8), offs // BLOCK - ) - return out_row, s_row_blocked, out_col, s_col_blocked - - def a(): - # FC1 grouped GEMM (prequantized inputs) -> eager SwiGLU -> dual quant. - z = torch._scaled_grouped_mm( - inp.x_q, - inp.w13_t_q, - inp.x_sf.view(r, -1), - inp.w13_t_sf.view(cfg.num_groups, -1), - offs=offs, - out_dtype=torch.bfloat16, - ) - zv = z.view(r, f, 2) - h = (torch.nn.functional.silu(zv[..., 0].float()) * zv[..., 1].float()).to( - torch.bfloat16 - ) - return dual_quantize(h) - - def b(): - dh = torch._scaled_grouped_mm( - inp.do_q, - inp.w2_t_q, - inp.do_sf.view(r, -1), - inp.w2_t_sf.view(cfg.num_groups, -1), - offs=offs, - out_dtype=torch.bfloat16, - ) - gate = inp.z_bf16[..., 0].float() - up = inp.z_bf16[..., 1].float() - sig = torch.sigmoid(gate) - dhf = dh.float() - dgate = (dhf * up * (sig * (1.0 + gate * (1.0 - sig)))).to(torch.bfloat16) - dup = (dhf * (gate * sig)).to(torch.bfloat16) - dz = torch.stack((dgate, dup), dim=-1).view(r, 2 * f) - return dual_quantize(dz) - - def make_wgrad(dy: torch.Tensor, x: torch.Tensor): - # Verbatim shape of the existing wgrad stage: CUDA dim1 quantization of - # both operands + K-groups scale rearranges + scaled grouped GEMM. - # (The fused kernel C instead consumes columnwise operands produced by - # kernels A/B, so this stage's quantization cost is the decomposed - # path's own.) - def run(): - dy_t_mx = _to_mxfp8_dim1_kernel_wrapper( - dy, - BLOCK, - elem_dtype=torch.float8_e4m3fn, - hp_dtype=dy.dtype, - kernel_preference=KernelPreference.AUTO, - cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, - scale_calculation_mode=RCEIL, - ) - x_t_mx = _to_mxfp8_dim1_kernel_wrapper( - x, - BLOCK, - elem_dtype=torch.float8_e4m3fn, - hp_dtype=x.dtype, - kernel_preference=KernelPreference.AUTO, - cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, - scale_calculation_mode=RCEIL, - ) - scale_offs = offs // BLOCK - dy_scales_blocked = triton_mx_block_rearrange_2d_K_groups( - dy_t_mx.scale, scale_offs - ) - x_scales_blocked = triton_mx_block_rearrange_2d_K_groups( - x_t_mx.scale, scale_offs - ) - return torch._scaled_grouped_mm( - dy_t_mx.qdata, - x_t_mx.qdata.transpose(-2, -1), - dy_scales_blocked, - x_scales_blocked, - offs=offs, - out_dtype=torch.bfloat16, - ) - - return run - - c_fc1 = make_wgrad(inp.dz_flat, inp.x) - c_fc2 = make_wgrad(inp.do, inp.h) - - def seq(): - a() - b() - c_fc1() - c_fc2() - - return {"a": a, "b": b, "c_fc1": c_fc1, "c_fc2": c_fc2, "seq": seq} - - -# -------------------------------------------------------------------------- -# Lane: TransformerEngine -# -------------------------------------------------------------------------- - - -def run_te_lane(cfg: ExperimentConfig) -> None: - """TE fused lane (one CuTe kernel per stage) + modular-lane single-kernel - gated-activation points. Prints its own tables; per-kernel CUDA times come - from a profiler pass filtered by kernel-name fragment.""" - os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1" - import transformer_engine.pytorch as te - import transformer_engine_torch as tex - from transformer_engine.common.recipe import MXFP8BlockScaling - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - r, d, f, g = cfg.rows, cfg.model_dim, cfg.hidden_dim, cfg.num_groups - torch.manual_seed(0) - offsets = make_offsets(cfg) - sizes = torch.diff(offsets, prepend=torch.zeros(1, device=device).int()) - sizes = sizes.to(torch.int32) - if (sizes % 128 != 0).any(): - raise ValueError( - "TE fused lane hard-crashes the CUDA context on per-expert sizes " - f"that are not multiples of 128, got {sizes.tolist()}" - ) - recipe = MXFP8BlockScaling() - - # Fused lane: one Sequential MLP, per-kernel attribution by name fragment. - fc1 = te.ops.GroupedLinear( - g, d, 2 * f, bias=False, device="cuda", dtype=torch.bfloat16 - ) - act = te.ops.ScaledSwiGLU(glu_interleave_size=32) - fc2 = te.ops.GroupedLinear(g, f, d, bias=False, device="cuda", dtype=torch.bfloat16) - mlp = te.ops.Sequential(fc1, act, fc2) - x = torch.randn(r, d, device=device, dtype=torch.bfloat16, requires_grad=True) - probs = torch.ones(r, device=device, dtype=torch.bfloat16) - - def fwd(): - with te.autocast(enabled=True, recipe=recipe): - return mlp(x, sizes, probs, sizes) - - y = fwd() - dy = torch.randn_like(y) - - def fwd_bwd(): - out = fwd() - out.backward(dy) - - fwd_bwd() # warmup / lazy init - torch.cuda.synchronize() - fwd_us = benchmark_cuda_function_in_microseconds(fwd) - fwd_bwd_us = benchmark_cuda_function_in_microseconds(fwd_bwd) - - fragments = { - "A analog (GroupedGemmGlu)": "GroupedGemmGlu", - "B analog (GroupedGemmDglu)": "GroupedGemmDglu", - "C analog (GroupedGemmWgrad)": "GroupedGemmWgrad", - "plain GEMM (GroupedGemmQuant)": "GroupedGemmQuant", - } - with torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CUDA] - ) as prof: - for _ in range(5): - fwd_bwd() - torch.cuda.synchronize() - sums = dict.fromkeys(fragments, 0.0) - counts = dict.fromkeys(fragments, 0) - for evt in prof.key_averages(): - for label, frag in fragments.items(): - if frag in evt.key and "helper" not in evt.key: - sums[label] += evt.self_device_time_total - counts[label] += evt.count - rows = [ - [label, counts[label] / 5.0, sums[label] / 5.0] - for label in fragments - if counts[label] - ] - print(f"\nTE fused lane (NVTE_CUTEDSL_FUSED_GROUPED_MLP=1) {cfg}") - print(f" fwd wall: {fwd_us:.1f} us fwd+bwd wall: {fwd_bwd_us:.1f} us") - print( - tabulate( - rows, - headers=["main kernel", "launches/iter", "device us/iter"], - floatfmt=".1f", - ) - ) - print( - " (fwd+bwd wall also covers FC2-fwd/FC1-dgrad GEMMs, input/dy " - "quantize and offsets prep, matching a full MLP step)" - ) - - # Modular lane: the single fused gated-act+dual-quant kernels. - q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) - z = torch.randn(r, 2 * f, device=device, dtype=torch.bfloat16) - dh = torch.randn(r, f, device=device, dtype=torch.bfloat16) - tex.swiglu(z, q) - tex.dswiglu(dh, z, q) - swiglu_us = benchmark_cuda_function_in_microseconds(tex.swiglu, z, q) - dswiglu_us = benchmark_cuda_function_in_microseconds(tex.dswiglu, dh, z, q) - print( - f"TE modular lane single kernels: tex.swiglu {swiglu_us:.1f} us, " - f"tex.dswiglu {dswiglu_us:.1f} us (gated activation + dual MXFP8 " - "quantize only; GEMMs are per-expert cuBLASLt in this lane)" - ) - - -# -------------------------------------------------------------------------- -# Driver -# -------------------------------------------------------------------------- - - -def run_experiment( - lane: str, cfg: ExperimentConfig, inp: Inputs -) -> Optional[ExperimentResult]: - stages = (lane_ours if lane == "ours" else lane_torchao)(cfg, inp) - if VALIDATE and lane == "ours": - validate_ours(cfg, inp, stages) - times: Dict[str, float] = {} - for name, fn in stages.items(): - fn() # warmup + lazy compile - torch.cuda.synchronize() - times[name] = benchmark_cuda_function_in_microseconds(fn) - flops = stage_flops(cfg) - return ExperimentResult( - a_us=times["a"], - b_us=times["b"], - c_fc1_us=times["c_fc1"], - c_fc2_us=times["c_fc2"], - seq_us=times["seq"], - a_tflops=flops["a"] / times["a"] / 1e6, - b_tflops=flops["b"] / times["b"] / 1e6, - c_fc1_tflops=flops["c_fc1"] / times["c_fc1"] / 1e6, - c_fc2_tflops=flops["c_fc2"] / times["c_fc2"] / 1e6, - ) - - -def print_banner() -> None: - props = torch.cuda.get_device_properties(device) - clocks = os.popen( - "nvidia-smi --query-gpu=clocks.applications.graphics,clocks.max.graphics " - "--format=csv,noheader 2>/dev/null" - ).read() - try: - import cutlass - - dsl = cutlass.__version__ - except Exception: - dsl = "n/a" - print( - f"device: {props.name} (cc {props.major}.{props.minor}, index " - f"{torch.cuda.current_device()}), torch {torch.__version__}, CUDA " - f"{torch.version.cuda}, nvidia-cutlass-dsl {dsl}" - ) - print(f"app clocks / max (per GPU):\n{clocks.strip()}") - print( - "NOTE: if app clocks are capped below max (e.g. 1200 MHz on the GB200 " - "dev hosts), absolute microseconds are NOT publishable; use ratios." - ) - - -def print_results(experiments: List[Experiment]) -> None: - headers = [ - "lane", - "R", - "D", - "F", - "G", - "dist", - "A us", - "B us", - "C_fc1 us", - "C_fc2 us", - "seq us", - "A TF/s", - "B TF/s", - "C1 TF/s", - "C2 TF/s", - ] - rows = [] - for e in experiments: - c, r = e.config, e.result - rows.append( - [ - e.lane, - c.rows, - c.model_dim, - c.hidden_dim, - c.num_groups, - c.distribution, - f"{r.a_us:.1f}", - f"{r.b_us:.1f}", - f"{r.c_fc1_us:.1f}", - f"{r.c_fc2_us:.1f}", - f"{r.seq_us:.1f}", - f"{r.a_tflops:.1f}", - f"{r.b_tflops:.1f}", - f"{r.c_fc1_tflops:.1f}", - f"{r.c_fc2_tflops:.1f}", - ] - ) - print(tabulate(rows, headers=headers)) - print( - "stage coverage: A = FC1 grouped GEMM + SwiGLU + dual MXFP8 quantize; " - "B = FC2 dgrad + dSwiGLU + dual quantize; C_* = grouped wgrad " - "(decomposed lane's C includes its own dim1 operand quantization); " - "seq = A;B;C_fc1;C_fc2." - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--lane", - choices=["torchao", "ours", "te", "all"], - default="all", - help="'all' = torchao + ours; 'te' must run in its own process " - "(container-native stack, no PYTHONUSERBASE)", - ) - parser.add_argument( - "--shape", - default=None, - help="single shape as 'R,D,F,G' instead of the built-in sweep", - ) - parser.add_argument( - "--dist", - choices=["balanced", "skewed", "both"], - default="balanced", - dest="dist", - ) - parser.add_argument( - "--profile", - action="store_true", - help="export a chrome trace of the fused-op sequence per config", - ) - args = parser.parse_args() - args.dists = ["balanced", "skewed"] if args.dist == "both" else [args.dist] - - print_banner() - configs = get_configs(args) - - if args.lane == "te": - for cfg in configs: - run_te_lane(cfg) - return - - lanes = ["torchao", "ours"] if args.lane == "all" else [args.lane] - experiments: List[Experiment] = [] - for cfg in tqdm(configs): - inp = Inputs(cfg) - for lane in lanes: - result = run_experiment(lane, cfg, inp) - experiments.append(Experiment(lane, cfg, result)) - if args.profile and "ours" in lanes: - from benchmarks.utils import profile_fn - - stages = lane_ours(cfg, inp) - profile_fn( - stages["seq"], - profile_name=f"grouped_mlp_seq_R{cfg.rows}_D{cfg.model_dim}" - f"_F{cfg.hidden_dim}_G{cfg.num_groups}", - ) - del inp - torch.cuda.empty_cache() - print_results(experiments) - - -if __name__ == "__main__": - main() diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py deleted file mode 100644 index 06f3996f52..0000000000 --- a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py +++ /dev/null @@ -1,1029 +0,0 @@ -# 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. - -"""Tests for the fused MXFP8 grouped-MLP kernel family (SM100). - -Three custom ops, one physical kernel launch each: - -* ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` (A) FC1 grouped GEMM + SwiGLU + - rowwise 1x32 and columnwise 32x1 MXFP8 RCEIL quantization -* ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` (B) FC2 dgrad grouped GEMM + - dSwiGLU + dual quantization -* ``torchao::mxfp8_grouped_gemm_wgrad`` (C) grouped MXFP8 wgrad - -References are deliberately bridge-free: eager per-expert BF16/FP64 matmuls, -``F.silu`` / the closed-form dSwiGLU, and the pure-torch ``to_mx`` (RCEIL) + -``to_blocked`` quantizers. The CuTe DSL standalone quantizer ops and -``cute_utils`` are never imported. - -Numerical strategy: - -* GEMM outputs (``z``, ``dh``, ``dw``): SQNR / tolerance vs an FP64 oracle - (reduction order is free), plus BITWISE equality on exact-integer operand - configs where every partial sum is exactly representable in FP32. -* Quantized outputs: bitwise vs ``to_mx`` wherever the activation is exact. - ``silu(g) == g`` exactly for ``g >= 128`` (sigmoid saturates to 1.0f), so a - saturated gate turns the fused activation into exact products and the - quantization stage must match ``to_mx`` byte-for-byte, including special - values. Random-input forward comparisons are ALSO bitwise (measured zero - mismatches: the kernel's sigmoid composition matches torch's float32 silu - exactly); backward random inputs are SQNR-gated because the reference dh - comes from an FP64 oracle whose BF16 rounding can differ at reduction-order - boundaries, with bitwise coverage provided by the exact-dh configuration. - -FakeTensor and validation tests run without a GPU; kernel tests require SM100. -""" - -import random - -import pytest - -torch = pytest.importorskip("torch") - -import torch.nn.functional as F # noqa: E402 -from torch._subclasses.fake_tensor import FakeTensorMode # noqa: E402 - -from torchao.float8.float8_utils import compute_error # noqa: E402 -from torchao.prototype.moe_training.utils import generate_jagged_offs # noqa: E402 -from torchao.prototype.mx_formats.config import ScaleCalculationMode # noqa: E402 -from torchao.prototype.mx_formats.mx_tensor import to_mx # noqa: E402 -from torchao.prototype.mx_formats.utils import from_blocked, to_blocked # noqa: E402 -from torchao.testing._mxfp8_test_utils import make_mxfp8_semantic_cases # noqa: E402 - -# Importing the ops module registers the three custom ops. The public wrapper -# module is preferred once it exists; both expose the same wrapper names. -try: - from torchao.prototype.moe_training import mxfp8_grouped_mlp as _api -except ImportError: - from torchao.prototype.moe_training.kernels.mxfp8 import grouped_mlp_ops as _api - -_E4M3 = torch.float8_e4m3fn -_E8M0 = torch.float8_e8m0fnu -_BLOCK = 32 -_RCEIL = ScaleCalculationMode.RCEIL - -_OP_NAMES = ( - "mxfp8_grouped_gemm_swiglu_fwd", - "mxfp8_grouped_gemm_dswiglu_bwd", - "mxfp8_grouped_gemm_wgrad", -) - - -def _is_sm_10x() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 - - -_gpu = pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") - -torch._dynamo.config.cache_size_limit = 1000 - - -# --------------------------------------------------------------------------- -# Reference helpers (pure torch) -# --------------------------------------------------------------------------- - - -def _round_up(x: int, to: int) -> int: - return ((x + to - 1) // to) * to - - -def _blocked_numel(rows: int, cols: int) -> int: - return _round_up(rows, 128) * _round_up(cols, 4) - - -def _quantize_rowwise_ref(x: torch.Tensor): - """[M, K] high precision -> (qdata [M, K] row-major, flat blocked scales).""" - scale, q = to_mx(x, _E4M3, _BLOCK, scaling_mode=_RCEIL) - return q, to_blocked(scale) - - -def _quantize_colwise_ref(x: torch.Tensor): - """[R, N] high precision -> (qdata [R, N] stride (1, R), flat blocked scales - for the logical [N, R/32] scale matrix). Recipe from the repo bench.""" - scale_t, q_t = to_mx(x.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) - return q_t.t(), to_blocked(scale_t) - - -def _dequant_rowwise(q: torch.Tensor, sf_flat: torch.Tensor, dtype=torch.float64): - """Dequantize a row-major [M, K] E4M3 tensor with flat blocked scales.""" - m, k = q.shape - scale = from_blocked(sf_flat.view(_E8M0).reshape(-1), m, k // _BLOCK) - return q.to(dtype) * scale.to(dtype).repeat_interleave(_BLOCK, dim=1) - - -def _dequant_colwise(q_col: torch.Tensor, sf_flat: torch.Tensor, dtype=torch.float64): - """Dequantize a [R, N] stride-(1, R) E4M3 tensor (scales logical [N, R/32]).""" - r, n = q_col.shape - scale = from_blocked(sf_flat.view(_E8M0).reshape(-1), n, r // _BLOCK) - return (q_col.t().to(dtype) * scale.to(dtype).repeat_interleave(_BLOCK, dim=1)).t() - - -def _dswiglu_closed_form(dh: torch.Tensor, gate: torch.Tensor, up: torch.Tensor): - """CONTRACT section 6.3 normative math, all fp32 in, (dgate, dup) fp32 out.""" - sig = torch.sigmoid(gate) - silu = gate * sig - dsilu = sig * (1.0 + gate * (1.0 - sig)) - return dh * up * dsilu, dh * silu - - -def _bytes(t: torch.Tensor) -> torch.Tensor: - return t.contiguous().view(torch.uint8) - - -def _mismatch_rate(a: torch.Tensor, b: torch.Tensor) -> float: - assert a.shape == b.shape - return (a != b).float().mean().item() - - -# --------------------------------------------------------------------------- -# Input builders. Tests own the offsets, so references never need a D2H sync. -# --------------------------------------------------------------------------- - - -def _mk_offsets(sizes, device): - ends = torch.tensor(sizes, dtype=torch.int64).cumsum(0) - return ends.to(torch.int32).to(device) - - -def _pack_grouped_weight(q_list, sf_list, device): - """Per-expert row-major [N, K] qdata -> [G, K, N] stride (K*N, 1, K) + [G, sf].""" - g = len(q_list) - n, k = q_list[0].shape - w = torch.empty_strided((g, k, n), (k * n, 1, k), dtype=_E4M3, device=device) - for i, q in enumerate(q_list): - w[i].copy_(q.t()) - sf = torch.stack([s.reshape(-1) for s in sf_list]) - return w, sf - - -def _random_grouped_weight(g, n, k, device, exact_int=False): - """Random per-expert [N, K] weights; returns (packed q, packed sf, dequants).""" - qs, sfs, deqs = [], [], [] - for _ in range(g): - if exact_int: - q = torch.randint(-6, 7, (n, k), device=device).to(_E4M3) - logical = torch.full( - (n, k // _BLOCK), 127, dtype=torch.uint8, device=device - ) - sf = to_blocked(logical.view(_E8M0)) - else: - w = torch.randn(n, k, device=device, dtype=torch.bfloat16) - q, sf = _quantize_rowwise_ref(w) - qs.append(q) - sfs.append(sf) - deqs.append(_dequant_rowwise(q, sf)) - packed_q, packed_sf = _pack_grouped_weight(qs, sfs, device) - return packed_q, packed_sf, deqs - - -def _random_activation(r, k, device, exact_int=False): - if exact_int: - q = torch.randint(-6, 7, (r, k), device=device).to(_E4M3) - logical = torch.full((r, k // _BLOCK), 127, dtype=torch.uint8, device=device) - sf = to_blocked(logical.view(_E8M0)) - else: - x = torch.randn(r, k, device=device, dtype=torch.bfloat16) - q, sf = _quantize_rowwise_ref(x) - return q, sf, _dequant_rowwise(q, sf) - - -def _ref_grouped_gemm(x_deq, w_deqs, offsets_sizes): - """Per-expert fp64 x @ w.T over test-owned split sizes; inactive tail = 0.""" - r = x_deq.shape[0] - out = torch.zeros(r, w_deqs[0].shape[0], dtype=torch.float64, device=x_deq.device) - start = 0 - for g, size in enumerate(offsets_sizes): - if size: - out[start : start + size] = x_deq[start : start + size] @ w_deqs[g].t() - start += size - return out - - -def _make_a_inputs(r, d, f, sizes, device, exact_int=False): - x_q, x_sf, x_deq = _random_activation(r, d, device, exact_int) - w_q, w_sf, w_deqs = _random_grouped_weight(len(sizes), 2 * f, d, device, exact_int) - offsets = _mk_offsets(sizes, device) - z_ref = _ref_grouped_gemm(x_deq, w_deqs, sizes) - return (x_q, x_sf, w_q, w_sf, offsets), z_ref - - -def _make_b_inputs(r, d, f, sizes, device, exact_int=False, z=None): - do_q, do_sf, do_deq = _random_activation(r, d, device, exact_int) - w_q, w_sf, w_deqs = _random_grouped_weight(len(sizes), f, d, device, exact_int) - offsets = _mk_offsets(sizes, device) - if z is None: - z = torch.randn(r, f, 2, device=device, dtype=torch.bfloat16) - active = sum(sizes) - if active < r: - # The inactive tail of z is read-forbidden: poison it so any read - # shows up as NaN contamination in the outputs. - z[active:] = float("nan") - dh_ref = _ref_grouped_gemm(do_deq, w_deqs, sizes) - return (do_q, do_sf, w_q, w_sf, z, offsets), dh_ref - - -def _make_c_inputs(r, n, k, sizes, device, exact_int=False): - def colwise(rows, cols): - if exact_int: - q_rm = torch.randint(-6, 7, (cols, rows), device=device).to(_E4M3) - logical = torch.full( - (cols, rows // _BLOCK), 127, dtype=torch.uint8, device=device - ) - sf = to_blocked(logical.view(_E8M0)) - return q_rm.t(), sf - x = torch.randn(rows, cols, device=device, dtype=torch.bfloat16) - return _quantize_colwise_ref(x) - - dy_q, dy_sf = colwise(r, n) - x_q, x_sf = colwise(r, k) - offsets = _mk_offsets(sizes, device) - return dy_q, dy_sf, x_q, x_sf, offsets - - -def _ref_wgrad(dy_q, dy_sf, x_q, x_sf, sizes): - dy = _dequant_colwise(dy_q, dy_sf) - x = _dequant_colwise(x_q, x_sf) - g = len(sizes) - n, k = dy.shape[1], x.shape[1] - dw = torch.zeros(g, n, k, dtype=torch.float64, device=dy.device) - start = 0 - for i, size in enumerate(sizes): - if size: - dw[i] = dy[start : start + size].t() @ x[start : start + size] - start += size - return dw.to(torch.bfloat16) - - -def _b_reference_dz(dh_bf16, z): - """CONTRACT section 6.3: bf16 dh + saved z -> interleaved dz (bf16 [R, 2F]).""" - dgate, dup = _dswiglu_closed_form( - dh_bf16.float(), z[..., 0].float(), z[..., 1].float() - ) - dz = torch.stack((dgate.bfloat16(), dup.bfloat16()), dim=-1) - return dz.reshape(dh_bf16.shape[0], -1) - - -def _assert_quantized_pair( - q_row, sf_row, q_col, sf_col, ref_bf16, max_qdata_rate=0.0, max_scale_rate=0.0 -): - """Compare both fused quantized orientations against to_mx of ``ref_bf16``.""" - ref_row_q, ref_row_sf = _quantize_rowwise_ref(ref_bf16) - ref_col_q, ref_col_sf = _quantize_colwise_ref(ref_bf16) - checks = ( - ("h_row_q", _bytes(q_row), _bytes(ref_row_q), max_qdata_rate), - ("h_row_sf", _bytes(sf_row.reshape(-1)), _bytes(ref_row_sf), max_scale_rate), - # Column-major output: compare bytes of the same logical view without - # forcing contiguity (the stride IS the ABI). - ("h_col_q", _bytes(q_col.t()), _bytes(ref_col_q.t()), max_qdata_rate), - ("h_col_sf", _bytes(sf_col.reshape(-1)), _bytes(ref_col_sf), max_scale_rate), - ) - for name, got, want, budget in checks: - rate = _mismatch_rate(got, want) - assert rate <= budget, f"{name}: byte mismatch rate {rate} > {budget}" - - -# --------------------------------------------------------------------------- -# 1. Registration and public surface -# --------------------------------------------------------------------------- - - -def test_ops_registered(): - for name in _OP_NAMES: - assert hasattr(torch.ops.torchao, name), name - assert hasattr(_api, name), f"{_api.__name__} must export {name}" - - -# --------------------------------------------------------------------------- -# 2. Fake / meta output contracts (no GPU required) -# --------------------------------------------------------------------------- - - -def _fake_a_inputs(r, d, f, g, device="cuda"): - x_q = torch.empty(r, d, dtype=_E4M3, device=device) - x_sf = torch.empty(_blocked_numel(r, d // _BLOCK), dtype=_E8M0, device=device) - w_q = torch.empty_strided( - (g, d, 2 * f), (d * 2 * f, 1, d), dtype=_E4M3, device=device - ) - w_sf = torch.empty( - g, _blocked_numel(2 * f, d // _BLOCK), dtype=_E8M0, device=device - ) - offs = torch.empty(g, dtype=torch.int32, device=device) - return x_q, x_sf, w_q, w_sf, offs - - -def _fake_b_inputs(r, d, f, g, device="cuda"): - do_q = torch.empty(r, d, dtype=_E4M3, device=device) - do_sf = torch.empty(_blocked_numel(r, d // _BLOCK), dtype=_E8M0, device=device) - w_q = torch.empty_strided((g, d, f), (d * f, 1, d), dtype=_E4M3, device=device) - w_sf = torch.empty(g, _blocked_numel(f, d // _BLOCK), dtype=_E8M0, device=device) - z = torch.empty_strided( - (r, f, 2), (2 * f, 2, 1), dtype=torch.bfloat16, device=device - ) - offs = torch.empty(g, dtype=torch.int32, device=device) - return do_q, do_sf, w_q, w_sf, z, offs - - -def _fake_c_inputs(r, n, k, g, device="cuda"): - dy = torch.empty_strided((r, n), (1, r), dtype=_E4M3, device=device) - dy_sf = torch.empty(_blocked_numel(n, r // _BLOCK), dtype=_E8M0, device=device) - x = torch.empty_strided((r, k), (1, r), dtype=_E4M3, device=device) - x_sf = torch.empty(_blocked_numel(k, r // _BLOCK), dtype=_E8M0, device=device) - offs = torch.empty(g, dtype=torch.int32, device=device) - return dy, dy_sf, x, x_sf, offs - - -def test_fake_swiglu_fwd_contract(): - r, d, f, g = 256, 256, 128, 2 - with FakeTensorMode(): - z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( - *_fake_a_inputs(r, d, f, g) - ) - assert z.shape == (r, f, 2) and z.stride() == (2 * f, 2, 1) - assert z.dtype == torch.bfloat16 - assert hq.shape == (r, f) and hq.stride() == (f, 1) and hq.dtype == _E4M3 - assert hsf.numel() == _blocked_numel(r, f // _BLOCK) - assert hcq.shape == (r, f) and hcq.stride() == (1, r) and hcq.dtype == _E4M3 - assert hcsf.numel() == _blocked_numel(f, r // _BLOCK) - - -def test_fake_dswiglu_bwd_contract(): - r, d, f, g = 256, 256, 128, 2 - with FakeTensorMode(): - dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( - *_fake_b_inputs(r, d, f, g) - ) - assert dzq.shape == (r, 2 * f) and dzq.stride() == (2 * f, 1) - assert dzsf.numel() == _blocked_numel(r, 2 * f // _BLOCK) - assert dzcq.shape == (r, 2 * f) and dzcq.stride() == (1, r) - assert dzcsf.numel() == _blocked_numel(2 * f, r // _BLOCK) - - -def test_fake_wgrad_contract(): - r, n, k, g = 256, 256, 128, 2 - with FakeTensorMode(): - dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*_fake_c_inputs(r, n, k, g)) - assert dw.shape == (g, n, k) - assert dw.stride() == (n * k, k, 1) - assert dw.dtype == torch.bfloat16 - - -def test_fake_validation_rejects_bad_metadata(): - with FakeTensorMode(): - args = list(_fake_c_inputs(256, 192, 128, 2)) # N = 192, not 128-multiple - with pytest.raises(ValueError, match="multiple of 128"): - torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - a_args = list(_fake_a_inputs(192, 256, 128, 2)) # R = 192 - with pytest.raises(ValueError, match="multiple of 128"): - torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*a_args) - - -# --------------------------------------------------------------------------- -# 3. Validation negatives (real tensors, small shapes) -# --------------------------------------------------------------------------- - - -def _valid_c_args(device): - return list(_make_c_inputs(256, 256, 128, [128, 128], device)) - - -_NEGATIVE_CASES = [ - "dtype", - "row_major_colwise_operand", - "scale_numel", - "offsets_int64", - "offsets_cpu", - "offsets_2d", - "offsets_numel", - "offsets_noncontig", - "n_not_128", - "misaligned_view", - "z_stride", -] - - -@_gpu -@pytest.mark.parametrize("case", _NEGATIVE_CASES) -def test_validation_negatives(case): - device = "cuda" - args = _valid_c_args(device) - if case == "dtype": - args[0] = torch.empty_strided( - (256, 256), (1, 256), dtype=torch.bfloat16, device=device - ) - err = "float8_e4m3fn" - elif case == "row_major_colwise_operand": - args[0] = torch.empty(256, 256, dtype=_E4M3, device=device) - err = "stride" - elif case == "scale_numel": - args[1] = args[1].reshape(-1)[:-1] - err = "blocked scale bytes" - elif case == "offsets_int64": - args[4] = args[4].to(torch.int64) - err = "int32" - elif case == "offsets_cpu": - args[4] = args[4].cpu() - err = "CUDA" - elif case == "offsets_2d": - args[4] = args[4].reshape(1, -1) - err = "1D" - elif case == "offsets_numel": - args[4] = torch.tensor([128, 128, 256], dtype=torch.int32, device=device) - # wgrad takes G from offsets, so a numel change alone is legal there; - # use kernel A where G comes from the weight tensor instead. - a_args = list(_fake_a_inputs(256, 256, 128, 2, device=device)) - a_args[4] = args[4] - with pytest.raises(ValueError, match="one entry per"): - torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*a_args) - return - elif case == "offsets_noncontig": - base = torch.zeros(4, dtype=torch.int32, device=device) - args[4] = base.as_strided((2,), (2,)) - err = "contiguous" - elif case == "n_not_128": - args = _valid_c_args(device) - dy = torch.empty_strided((256, 192), (1, 256), dtype=_E4M3, device=device) - dy_sf = torch.empty( - _blocked_numel(192, 256 // _BLOCK), dtype=_E8M0, device=device - ) - args[0], args[1] = dy, dy_sf - err = "multiple of 128" - elif case == "misaligned_view": - base = torch.zeros(256 * 256 + 32, dtype=_E4M3, device=device) - args[0] = base.as_strided((256, 256), (1, 256), 2) - err = "aligned" - elif case == "z_stride": - b_args = list(_fake_b_inputs(256, 256, 128, 2, device=device)) - # materialize real tensors with a wrong z layout - b_args = [torch.empty_like(t) if t.is_cuda else t for t in b_args] - b_args[4] = torch.empty( - 256, 2, 128, dtype=torch.bfloat16, device=device - ).permute(0, 2, 1) - with pytest.raises(ValueError, match="stride"): - torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd(*b_args) - return - with pytest.raises(ValueError, match=err): - torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - - -@_gpu -def test_validation_rejects_g0(): - device = "cuda" - args = _valid_c_args(device) - args[4] = torch.empty(0, dtype=torch.int32, device=device) - with pytest.raises(ValueError): - torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - - -@pytest.mark.skipif( - not (_is_sm_10x() and torch.cuda.device_count() >= 2), - reason="needs two CUDA devices", -) -def test_validation_rejects_cross_device(): - args = _valid_c_args("cuda:0") - args[2] = args[2].to("cuda:1") - with pytest.raises(ValueError, match="device"): - torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - - -# --------------------------------------------------------------------------- -# 4. R == 0 and zero-token experts -# --------------------------------------------------------------------------- - - -@_gpu -def test_r0_all_ops(): - device = "cuda" - d, f, g = 256, 128, 2 - a_args = list(_fake_a_inputs(0, d, f, g, device=device)) - a_args = [torch.empty_like(t) for t in a_args] - a_args[4] = torch.zeros(g, dtype=torch.int32, device=device) - outs = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*a_args) - assert all(o.shape[0] == 0 or o.numel() == 0 for o in outs) - - b_args = list(_fake_b_inputs(0, d, f, g, device=device)) - b_args = [torch.empty_like(t) for t in b_args] - b_args[5] = torch.zeros(g, dtype=torch.int32, device=device) - outs = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd(*b_args) - assert all(o.numel() == 0 for o in outs) - - c_args = list(_fake_c_inputs(0, 256, 128, g, device=device)) - c_args = [torch.empty_like(t) for t in c_args] - c_args[4] = torch.zeros(g, dtype=torch.int32, device=device) - dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*c_args) - assert dw.shape == (g, 256, 128) - assert (dw == 0).all() - - -@_gpu -def test_wgrad_zero_token_expert(): - torch.manual_seed(0) - device = "cuda" - sizes = [128, 0, 256] - args = _make_c_inputs(384, 256, 128, sizes, device) - dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - assert (dw[1] == 0).all(), "zero-token expert must produce an all-zero slice" - ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) - assert compute_error(ref[0].float(), dw[0].float()) >= 24.0 - assert compute_error(ref[2].float(), dw[2].float()) >= 24.0 - - -# --------------------------------------------------------------------------- -# 5. Kernel C numerics -# --------------------------------------------------------------------------- - - -@_gpu -@pytest.mark.parametrize( - "r,n,k,sizes", - [ - (256, 256, 128, [128, 128]), # FC1-like: N = 2F, K = D (small) - (1536, 2816, 2048, [512, 0, 640, 384]), # 16B FC1 wgrad class - (1536, 2048, 1408, [512, 0, 640, 384]), # 16B FC2 wgrad class - ], - ids=["small", "fc1_16b", "fc2_16b"], -) -def test_wgrad_numerics_random(r, n, k, sizes): - torch.manual_seed(1) - args = _make_c_inputs(r, n, k, sizes, "cuda") - dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) - assert dw.dtype == torch.bfloat16 and dw.shape == (len(sizes), n, k) - torch.testing.assert_close(dw.float(), ref.float(), atol=2e-3, rtol=0.01) - assert compute_error(ref.float(), dw.float()) >= 24.0 - - -@_gpu -def test_wgrad_bitwise_exact_integers(): - torch.manual_seed(2) - sizes = [256, 128, 384] - args = _make_c_inputs(768, 256, 256, sizes, "cuda", exact_int=True) - dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) - assert torch.equal(_bytes(dw), _bytes(ref)), "exact-integer wgrad must be bitwise" - - -@_gpu -def test_wgrad_rejects_kgroups_scale_ordering_by_result(): - """Whole-matrix vs per-group K-groups blocked scales differ whenever N > 128; - no length check can catch it, so prove the kernel is sensitive to it.""" - torch.manual_seed(3) - device = "cuda" - sizes = [128, 128] - r, n, k = 256, 256, 128 - dy = torch.randn(r, n, device=device, dtype=torch.bfloat16) - x = torch.randn(r, k, device=device, dtype=torch.bfloat16) - dy_q, dy_sf = _quantize_colwise_ref(dy) - x_q, x_sf = _quantize_colwise_ref(x) - offsets = _mk_offsets(sizes, device) - - good = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf, x_q, x_sf, offsets) - ref = _ref_wgrad(dy_q, dy_sf, x_q, x_sf, sizes) - torch.testing.assert_close(good.float(), ref.float(), atol=2e-3, rtol=0.01) - - # Re-encode dy's scales per group (torchao K-groups form) and rerun. - scale_t, _ = to_mx(dy.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) - per_group = torch.cat( - [ - to_blocked(scale_t[:, s // _BLOCK : e // _BLOCK]).reshape(-1) - for s, e in ((0, 128), (128, 256)) - ] - ) - assert per_group.numel() == dy_sf.numel() - assert not torch.equal(_bytes(per_group), _bytes(dy_sf.reshape(-1))) - bad = torch.ops.torchao.mxfp8_grouped_gemm_wgrad( - dy_q, per_group, x_q, x_sf, offsets - ) - assert not torch.equal(bad, good), ( - "kernel must consume whole-matrix blocked scales; a K-groups buffer of " - "identical length must change the result" - ) - - -# --------------------------------------------------------------------------- -# 6. Kernel A numerics -# --------------------------------------------------------------------------- - -_A_SHAPES = [ - (256, 256, 128, [128, 128]), - (1024, 512, 256, [256, 0, 512, 256]), # zero-token expert + ragged - (1536, 2048, 1408, [512, 128, 640, 256]), # 16B class -] - - -@_gpu -@pytest.mark.parametrize("r,d,f,sizes", _A_SHAPES, ids=["small", "ragged", "16b"]) -def test_swiglu_fwd_random(r, d, f, sizes): - torch.manual_seed(4) - args, z_ref = _make_a_inputs(r, d, f, sizes, "cuda") - z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) - active = sum(sizes) - assert ( - compute_error(z_ref[:active].float(), z[:active].reshape(active, -1).float()) - >= 27.0 - ) - # Quantized outputs compare against the normative chain applied to the - # kernel's own z (removes GEMM reduction-order noise from the comparison). - # Measured bitwise-identical on all of these shapes (the kernel's sigmoid - # composition matches torch's float32 silu exactly), so no mismatch budget. - gate = z[..., 0].float() - up = z[..., 1].float() - h_ref = (F.silu(gate) * up).bfloat16() - _assert_quantized_pair(hq, hsf, hcq, hcsf, h_ref) - - -@_gpu -def test_swiglu_fwd_bitwise_exact_integers(): - torch.manual_seed(5) - r, d, f, sizes = 512, 256, 128, [256, 256] - args, z_ref = _make_a_inputs(r, d, f, sizes, "cuda", exact_int=True) - z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) - assert torch.equal( - _bytes(z.reshape(r, -1)), _bytes(z_ref.bfloat16().reshape(r, -1)) - ), "integer-exact z must be bitwise" - - -@_gpu -def test_swiglu_fwd_saturated_gate_bitwise(): - """gate == 128 makes silu exact in any implementation, so the fused dual - quantization must match to_mx byte-for-byte, including special values.""" - torch.manual_seed(6) - device = "cuda" - r, d, f = 256, 128, 128 - sizes = [256] - # x = identity blocks: row r selects weight column r % d. - x_q = torch.zeros(r, d, dtype=torch.uint8, device=device) - x_q[torch.arange(r), torch.arange(r) % d] = 0x38 # 1.0 - x_q = x_q.view(_E4M3) - x_logical = torch.full((r, d // _BLOCK), 127, dtype=torch.uint8, device=device) - x_sf = to_blocked(x_logical.view(_E8M0)) - - # Up rows: random E4M3 bytes (finite lanes), then crafted special features. - # Feature f's specials land in rowwise scale block f // 32 of every row. - up_bytes = torch.randint(0, 0x7E, (f, d), dtype=torch.uint8, device=device) - up_scales = torch.randint( - 120, 134, (f, d // _BLOCK), dtype=torch.uint8, device=device - ) - up_bytes[0, :] = 0x7F # NaN up -> h NaN in block 0 - up_bytes[1, :] = 0x00 # zero up - up_bytes[40, :] = 0x7E # 448 * 2^119: z_up ~ 2.98e38 (finite bf16); - up_scales[40, :] = 127 + 119 # h = 128 * z_up overflows f32 -> +Inf, block 1 - w_q2f = torch.zeros(2 * f, d, dtype=torch.uint8, device=device) - w_scale2f = torch.zeros(2 * f, d // _BLOCK, dtype=torch.uint8, device=device) - w_q2f[0::2] = 0x38 - w_scale2f[0::2] = 127 + 7 - w_q2f[1::2] = up_bytes - w_scale2f[1::2] = up_scales - w_q = w_q2f.view(_E4M3) - w_sf = to_blocked(w_scale2f.view(_E8M0)) - w_packed, w_sf_packed = _pack_grouped_weight([w_q], [w_sf], device) - offsets = _mk_offsets(sizes, device) - - z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( - x_q, x_sf, w_packed, w_sf_packed, offsets - ) - gate = z[..., 0].float() - up = z[..., 1].float() - assert torch.equal(gate, torch.full_like(gate, 128.0)), "gate must be exactly 128" - h_ref = (128.0 * up).bfloat16() # silu(128) == 128 exactly - _assert_quantized_pair(hq, hsf, hcq, hcsf, h_ref) # zero mismatch budget - # Special-value spot checks straight from the RCEIL table: - hq_bytes = _bytes(hq).reshape(r, f) - hsf_logical = _bytes( - from_blocked(hsf.view(_E8M0).reshape(-1), r, f // _BLOCK) - ).reshape(r, f // _BLOCK) - nan_blocks = torch.isnan(h_ref).view(r, f // _BLOCK, _BLOCK).any(-1) - inf_blocks = torch.isinf(h_ref).view(r, f // _BLOCK, _BLOCK).any(-1) - assert nan_blocks.any() and inf_blocks.any(), "crafted specials must appear" - assert (hsf_logical[nan_blocks | inf_blocks] == 0xFF).all() - nonfinite_cols = (nan_blocks | inf_blocks).repeat_interleave(_BLOCK, dim=1) - assert (hq_bytes[nonfinite_cols] == 0x7F).all() - - -@_gpu -def test_swiglu_fwd_tail_and_poison(): - torch.manual_seed(7) - device = "cuda" - r, d, f, sizes = 512, 256, 128, [128, 256] # active 384, tail 128 - args, _ = _make_a_inputs(r, d, f, sizes, device) - z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) - active = sum(sizes) - assert (_bytes(z.reshape(r, -1))[active:] == 0).all(), "z tail must be zero bytes" - assert (_bytes(hq)[active:] == 0).all() - assert (_bytes(hcq.t())[:, active:] == 0).all() - hsf_logical = _bytes(from_blocked(hsf.view(_E8M0).reshape(-1), r, f // _BLOCK)) - assert (hsf_logical.reshape(r, -1)[active:] == 0).all() - hcsf_logical = _bytes(from_blocked(hcsf.view(_E8M0).reshape(-1), f, r // _BLOCK)) - assert (hcsf_logical.reshape(f, -1)[:, active // _BLOCK :] == 0).all() - - -# --------------------------------------------------------------------------- -# 7. Kernel B numerics -# --------------------------------------------------------------------------- - - -@_gpu -@pytest.mark.parametrize( - "r,d,f,sizes", - [ - (256, 256, 128, [128, 128]), - (1024, 512, 256, [256, 0, 512, 128]), # zero-token + strict tail - (1536, 2048, 1408, [512, 128, 640, 256]), # 16B class - ], - ids=["small", "ragged_tail", "16b"], -) -def test_dswiglu_bwd_random(r, d, f, sizes): - torch.manual_seed(8) - args, dh_ref = _make_b_inputs(r, d, f, sizes, "cuda") - dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd(*args) - active = sum(sizes) - z = args[4] - dz_ref = _b_reference_dz(dh_ref.bfloat16(), z) - got = _dequant_rowwise(dzq, dzsf, dtype=torch.float32) - assert compute_error(dz_ref[:active].float(), got[:active]) >= 25.0 - got_col = _dequant_colwise(dzcq, dzcsf, dtype=torch.float32) - assert compute_error(dz_ref[:active].float(), got_col[:active]) >= 25.0 - if active < r: - assert (_bytes(dzq)[active:] == 0).all(), "dz tail must be zero bytes" - assert (_bytes(dzcq.t())[:, active:] == 0).all() - - -@_gpu -def test_dswiglu_bwd_bitwise_exact(): - """Exact dh (integer GEMM) + saturated/zero gates: dz is exact products, so - all four outputs must match to_mx byte-for-byte; interleave order checked.""" - torch.manual_seed(9) - device = "cuda" - r, d, f, sizes = 256, 128, 128, [256] - # dh == 1.0 exactly: do = identity rows, w2 = all ones. - do_q = torch.zeros(r, d, dtype=torch.uint8, device=device) - do_q[torch.arange(r), torch.arange(r) % d] = 0x38 - do_q = do_q.view(_E4M3) - do_sf = to_blocked( - torch.full((r, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) - ) - w_q = torch.full((f, d), 0x38, dtype=torch.uint8, device=device).view(_E4M3) - w_sf = to_blocked( - torch.full((f, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) - ) - w_packed, w_sf_packed = _pack_grouped_weight([w_q], [w_sf], device) - - # Saturated gates make dsilu == 1 and silu == gate exactly; gate rows of z - # also cover 0 (silu(0) == 0, dsilu(0) == 0.5 -- both exact). - z = torch.zeros(r, f, 2, device=device, dtype=torch.bfloat16) - z[..., 0] = 128.0 - z[: r // 2, :, 1] = torch.randn(r // 2, f, device=device).bfloat16() - z[r // 2 :, :, 0] = 0.0 # gate 0 rows: dgate = 0.5 * up, dup = 0 - z[r // 2 :, :, 1] = ( - torch.randn(r - r // 2, f, device=device).bfloat16().float() * 2.0 - ).bfloat16() - z[0, 0:4, 0] = float("nan") # NaN gate+up block - z[0, 0:4, 1] = float("nan") - z[1, 0:16, 1] = 448.0 # uniform 448 dgate lanes - offsets = _mk_offsets(sizes, device) - - dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w_packed, w_sf_packed, z, offsets - ) - dh = torch.ones(r, f, device=device, dtype=torch.bfloat16) - dz_ref = _b_reference_dz(dh, z) - # Interleave check on exact lanes (bitwise: NaN blocks are present). - dgate_ref, dup_ref = _dswiglu_closed_form( - dh.float(), z[..., 0].float(), z[..., 1].float() - ) - assert torch.equal(_bytes(dz_ref[:, 0::2]), _bytes(dgate_ref.bfloat16())) - assert torch.equal(_bytes(dz_ref[:, 1::2]), _bytes(dup_ref.bfloat16())) - _assert_quantized_pair(dzq, dzsf, dzcq, dzcsf, dz_ref) # zero budget - - -@_gpu -def test_dswiglu_bwd_semantic_blocks(): - """Drive the shared MXFP8 semantic contract through the fused backward. - - Uniform cases with |value| >= 128 are constructed exactly (gate = up = - value under a saturated gate gives a uniform dz block); the rest are - covered transitively: kernel bytes must equal to_mx bytes on blocks - containing the case values, and to_mx itself is asserted against the - shared table in test_to_mx_matches_semantic_table.""" - device = "cuda" - cases = make_mxfp8_semantic_cases(torch.bfloat16, _RCEIL, device=device) - n_cases = len(cases.names) - r, d, f = 128, 128, max(128, _round_up(n_cases * 16, 128)) - do_q = torch.zeros(r, d, dtype=torch.uint8, device=device) - do_q[torch.arange(r), torch.arange(r) % d] = 0x38 - do_q = do_q.view(_E4M3) - do_sf = to_blocked( - torch.full((r, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) - ) - w_q = torch.full((f, d), 0x38, dtype=torch.uint8, device=device).view(_E4M3) - w_sf = to_blocked( - torch.full((f, d // _BLOCK), 127, dtype=torch.uint8, device=device).view(_E8M0) - ) - w_packed, w_sf_packed = _pack_grouped_weight([w_q], [w_sf], device) - - # Row 0: each case occupies 16 features -> one 32-wide dz block. - z = torch.zeros(r, f, 2, device=device, dtype=torch.bfloat16) - z[..., 0] = 128.0 - direct_table = [] - for idx in range(n_cases): - vals = cases.inputs[idx].to(device) - f0 = idx * 16 - uniform = bool((vals == vals[0]).all()) and not torch.isnan(vals).any() - if uniform and abs(float(vals[0])) >= 128.0: - # gate = up = v: dgate = v, dup = gate = |v|-signed... both = v - # requires gate == v which needs v >= 128; negatives via up lane. - v = float(vals[0]) - if v >= 128.0: - z[0, f0 : f0 + 16, 0] = v - z[0, f0 : f0 + 16, 1] = v - direct_table.append((idx, None)) - else: - z[0, f0 : f0 + 16, 0] = -v - z[0, f0 : f0 + 16, 1] = v - direct_table.append((idx, "even_only")) - else: - # Transitive: dgate lanes carry the case's even values, dup lanes - # its odd values scaled through the saturated gate where possible; - # fall back to plain interleave of the case into up lanes. - z[0, f0 : f0 + 16, 1] = vals[0::2] - z[0, f0 : f0 + 16, 0] = torch.where( - torch.isfinite(vals[1::2]) & (vals[1::2].abs() >= 128), - vals[1::2], - torch.full_like(vals[1::2], 128.0), - ) - offsets = _mk_offsets([r], device) - dzq, dzsf, dzcq, dzcsf = torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w_packed, w_sf_packed, z, offsets - ) - dh = torch.ones(r, f, device=device, dtype=torch.bfloat16) - dz_ref = _b_reference_dz(dh, z) - _assert_quantized_pair(dzq, dzsf, dzcq, dzcsf, dz_ref) # bitwise transitivity - # Direct table assertions where the block is exactly the case input. - row_sf_logical = _bytes( - from_blocked(dzsf.view(_E8M0).reshape(-1), r, 2 * f // _BLOCK) - ).reshape(r, -1) - dz_bytes = _bytes(dzq).reshape(r, -1) - for idx, mode in direct_table: - blk = slice(idx * _BLOCK, (idx + 1) * _BLOCK) - want_scale = int(cases.expected_scales[idx]) - want_data = cases.expected_data[idx].to(device) - assert row_sf_logical[0, idx] == want_scale, cases.names[idx] - got = dz_bytes[0, blk] - if mode is None: - assert torch.equal(got, want_data.to(got.device)), cases.names[idx] - else: - assert torch.equal(got[0::2], want_data.to(got.device)[0::2]), cases.names[ - idx - ] - - -def _semantic_table_dtype_check(device): - cases = make_mxfp8_semantic_cases(torch.bfloat16, _RCEIL, device=device) - scale, q = to_mx(cases.inputs, _E4M3, _BLOCK, scaling_mode=_RCEIL) - assert torch.equal( - _bytes(q).cpu().reshape(len(cases.names), 32), cases.expected_data - ) - assert torch.equal( - _bytes(scale.reshape(-1)).cpu(), cases.expected_scales.reshape(-1) - ) - - -@_gpu -def test_to_mx_matches_semantic_table(): - """Anchors the transitive comparisons above: to_mx == the shared table.""" - _semantic_table_dtype_check("cuda") - - -# --------------------------------------------------------------------------- -# 8. torch.compile -# --------------------------------------------------------------------------- - - -@_gpu -@pytest.mark.parametrize("op_name", _OP_NAMES) -def test_compile_matches_eager(op_name): - torch.manual_seed(10) - device = "cuda" - if op_name == "mxfp8_grouped_gemm_wgrad": - args = _make_c_inputs(256, 256, 128, [128, 128], device) - elif op_name == "mxfp8_grouped_gemm_swiglu_fwd": - args, _ = _make_a_inputs(256, 256, 128, [128, 128], device) - else: - args, _ = _make_b_inputs(256, 256, 128, [128, 128], device) - op = getattr(torch.ops.torchao, op_name) - eager = op(*args) - compiled_fn = torch.compile(lambda *a: op(*a), fullgraph=True) - compiled = compiled_fn(*args) - eager = eager if isinstance(eager, tuple) else (eager,) - compiled = compiled if isinstance(compiled, tuple) else (compiled,) - for e, c in zip(eager, compiled): - assert e.stride() == c.stride() - assert torch.equal(_bytes(e.reshape(-1)), _bytes(c.reshape(-1))) - - -# --------------------------------------------------------------------------- -# 9. One physical launch per op (profiler evidence) -# --------------------------------------------------------------------------- - - -def _device_kernel_names(fn): - fn() - torch.cuda.synchronize() - fn() - torch.cuda.synchronize() - with torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CUDA] - ) as prof: - fn() - torch.cuda.synchronize() - names = [] - for evt in prof.key_averages(): - if evt.device_type != torch.autograd.DeviceType.CUDA: - continue - if "Memset" in evt.key or "Memcpy" in evt.key: - continue - names.extend([evt.key] * evt.count) - return names - - -@_gpu -@pytest.mark.parametrize("op_name", _OP_NAMES) -def test_one_physical_launch(op_name): - torch.manual_seed(11) - device = "cuda" - if op_name == "mxfp8_grouped_gemm_wgrad": - args = _make_c_inputs(256, 256, 128, [128, 128], device) - elif op_name == "mxfp8_grouped_gemm_swiglu_fwd": - args, _ = _make_a_inputs(256, 256, 128, [128, 128], device) - else: - args, _ = _make_b_inputs(256, 256, 128, [128, 128], device) - op = getattr(torch.ops.torchao, op_name) - names = _device_kernel_names(lambda: op(*args)) - assert len(names) == 1, f"{op_name} must be ONE kernel launch, saw: {names}" - banned = ("quantize", "silu", "elementwise", "vectorized") - assert not any(b in names[0].lower() for b in banned), names - - -# --------------------------------------------------------------------------- -# 10. Large DSv3-class shapes (memory gated) -# --------------------------------------------------------------------------- - - -def _enough_memory(bytes_needed: int) -> bool: - if not torch.cuda.is_available(): - return False - return torch.cuda.get_device_properties(0).total_memory >= bytes_needed - - -@_gpu -@pytest.mark.skipif(not _enough_memory(16 << 30), reason="needs >= 16 GiB") -def test_wgrad_dsv3_671b_class(): - torch.manual_seed(12) - sizes = [256, 0, 512, 256] - args = _make_c_inputs(1024, 4096, 7168, sizes, "cuda") - dw = torch.ops.torchao.mxfp8_grouped_gemm_wgrad(*args) - ref = _ref_wgrad(args[0], args[1], args[2], args[3], sizes) - assert compute_error(ref.float(), dw.float()) >= 24.0 - - -@_gpu -@pytest.mark.skipif(not _enough_memory(16 << 30), reason="needs >= 16 GiB") -def test_swiglu_fwd_dsv3_671b_class(): - torch.manual_seed(13) - args, z_ref = _make_a_inputs(1024, 7168, 2048, [256, 0, 512, 256], "cuda") - z, hq, hsf, hcq, hcsf = torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd(*args) - active = 1024 - assert ( - compute_error(z_ref[:active].float(), z[:active].reshape(active, -1).float()) - >= 27.0 - ) - - -# --------------------------------------------------------------------------- -# 11. Availability / unsupported environments -# --------------------------------------------------------------------------- - - -def test_wrapper_unavailable_raises_cleanly(monkeypatch): - mod = pytest.importorskip("torchao.prototype.moe_training.mxfp8_grouped_mlp") - flag = "_mxfp8_grouped_mlp_kernels_available" - if not hasattr(mod, flag): - pytest.skip("availability flag not exposed") - monkeypatch.setattr(mod, flag, False) - with pytest.raises(NotImplementedError): - with FakeTensorMode(): - mod.mxfp8_grouped_gemm_wgrad(*_fake_c_inputs(256, 256, 128, 2)) - - -def test_jagged_offs_generator_contract(): - """The test/bench offsets helper must produce 128-multiples when asked.""" - random.seed(0) - if torch.cuda.is_available(): - offs = generate_jagged_offs(4, 1024, multiple_of=128) - else: - offs = generate_jagged_offs(4, 1024, multiple_of=128, device="cpu") - sizes = torch.diff(offs.cpu(), prepend=torch.tensor([0], dtype=offs.dtype)) - assert (sizes % 128 == 0).all() - assert offs[-1].item() == 1024 diff --git a/torchao/prototype/moe_training/__init__.py b/torchao/prototype/moe_training/__init__.py index cd3bf7abc3..e58461a80c 100644 --- a/torchao/prototype/moe_training/__init__.py +++ b/torchao/prototype/moe_training/__init__.py @@ -1,14 +1,6 @@ from torchao.prototype.moe_training.fp8_grouped_mm import ( _to_fp8_rowwise_then_scaled_grouped_mm, ) -from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( - is_supported as mxfp8_grouped_mlp_is_supported, -) -from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( - mxfp8_grouped_gemm_dswiglu_bwd, - mxfp8_grouped_gemm_swiglu_fwd, - mxfp8_grouped_gemm_wgrad, -) from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _to_mxfp8_then_scaled_grouped_mm, ) @@ -16,8 +8,4 @@ __all__ = [ "_to_mxfp8_then_scaled_grouped_mm", "_to_fp8_rowwise_then_scaled_grouped_mm", - "mxfp8_grouped_gemm_swiglu_fwd", - "mxfp8_grouped_gemm_dswiglu_bwd", - "mxfp8_grouped_gemm_wgrad", - "mxfp8_grouped_mlp_is_supported", ] diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index 938a6e7dc2..350ed187f7 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,11 +1,9 @@ -# Importing grouped_mlp_ops / cudnn_grouped_mlp_ops registers the fused -# grouped-MLP custom ops (torchao::mxfp8_grouped_gemm_{swiglu_fwd,dswiglu_bwd, -# wgrad} and torchao::mxfp8_cudnn_grouped_{mlp_fwd,mm,mlp_bwd,mlp_wgrad}). -# Both modules are importable with no CuTe DSL / cudnn-frontend installed; -# kernel imports are deferred into the op bodies. +# Importing cudnn_grouped_mlp_ops registers the fused grouped-MLP custom ops +# (torchao::mxfp8_cudnn_grouped_{mlp_fwd,mm,mlp_bwd,mlp_wgrad}). The module is +# importable with no cudnn-frontend installed; `import cudnn` is deferred into +# the op bodies. from torchao.prototype.moe_training.kernels.mxfp8 import ( cudnn_grouped_mlp_ops, # noqa: F401 - grouped_mlp_ops, # noqa: F401 ) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py deleted file mode 100644 index 78f1fcb028..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py +++ /dev/null @@ -1,1606 +0,0 @@ -# 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. - -"""MXFP8 routed-expert grouped-MLP kernels for SM100 (CuTe DSL, public API only). - -Three physically fused kernels, one launch each: - -* ``launch_grouped_gemm_swiglu_fwd`` -- FC1 ragged grouped GEMM + SwiGLU + - rowwise (1x32) and columnwise (32x1) MXFP8 RCEIL quantization, plus the BF16 - pre-activation save. -* ``launch_grouped_gemm_dswiglu_bwd`` -- FC2 dgrad ragged grouped GEMM + - dSwiGLU + the same dual quantization of the FC1 input gradient. -* ``launch_grouped_gemm_wgrad`` -- generic ragged-K grouped weight - gradient, BF16 output; called once for FC1 and once for FC2. - -They share one blockscaled tcgen05 mainloop. Three structural invariants, all -descending from every per-expert row count being a multiple of 128: - -* No per-group tensormaps: one host-built static TMA descriptor per operand; - per-expert selection is an integer coordinate (an L coordinate for the 3-D - weights, a K-tile index base for wgrad's ragged contraction -- exact - because ``tile_atom_to_shape_SF``'s K-tile mode is uniform). -* No tile scheduler: forward/backward enumerate all ``[0, R/128)`` M tiles - and the wgrad grid ``(N/128, K/128, G)`` is fully static; only wgrad's - K-loop trip count is data-dependent. -* No special tail path: an inactive tile runs with ``k_cnt == 0`` -- no TMA - loads, a register-zeroed accumulator, and the unmodified epilogue emits the - zero bytes the contract requires. Stores are never predicated; the one - epilogue-side gmem input (the backward kernel's saved ``z_bf16``) is loaded - only when ``k_cnt > 0`` because its tail rows are read-forbidden. - -Offset contract (documented caller invariants -- the offset VALUES live on -device and cannot be checked on the host without a synchronization): offsets -are exclusive per-expert end indices, int32, CUDA, 1-D, contiguous, -nondecreasing, every per-expert row count a multiple of 128, and -``offsets[-1] <= R``. The launchers validate all metadata and reject the rest -of the malformed-input space with ``ValueError``; set -``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` to additionally validate the values on -the host while debugging, at the cost of a D2H copy. There is deliberately no -device-side assertion in the default build: assertions are compiled out of -CuTe DSL kernels unless ``CUTE_DSL_ENABLE_ASSERTIONS=1``, so they cannot serve -as a production guard, and with malformed offsets the wgrad kernel returns a -wrong result rather than faulting -- "it did not crash" is not evidence the -offsets were valid. -""" - -import functools -from dataclasses import dataclass - -import cutlass -import cutlass.cute as cute -import cutlass.pipeline as pipeline -import cutlass.utils as utils -import cutlass.utils.blackwell_helpers as sm100_utils -import cutlass.utils.blockscaled_layout as blockscaled_utils -import torch -from cutlass import Float32, Int32 -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.runtime import from_dlpack -from cutlass.utils import LayoutEnum - -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( - _is_fake, - validate_allocated_rows, - validate_blocked_scales, - validate_destination, - validate_feature_dims, - validate_group_offsets, - validate_grouped_operand, -) - -__all__ = [ - "launch_grouped_gemm_swiglu_fwd", - "launch_grouped_gemm_dswiglu_bwd", - "launch_grouped_gemm_wgrad", -] - -# --- Frozen configuration: one tiling, one pipeline shape, one warp assignment. - -# MXFP8 scaling block: 32 values share one E8M0 scale. -_SF_VEC_SIZE = 32 -# cta_tile_m = 128 is required by CtaGroup.ONE and by the no-partial-M-tile -# argument. cta_tile_k = 128 is pinned: the wgrad kernel selects an expert's K -# range with an integer K-tile index base, exact only because every group -# boundary (a multiple of 128 rows) is a multiple of the K tile. -_CTA_M = 128 -_CTA_N = 128 -_CTA_K = 128 -_MMA_TILER = (_CTA_M, _CTA_N, _CTA_K) -# One stage: A tile + B tile (E4M3) + one 512-byte scale atom per operand. -# This is also the exact tx_count the TMA pipeline barrier must expect. -_AB_STAGE_BYTES = _CTA_M * _CTA_K + _CTA_N * _CTA_K + 2 * (128 * (_CTA_K // 32)) -_NUM_AB_STAGE = 6 -_NUM_ACC_STAGE = 1 -# Warp 0 loads (TMA), warp 1 issues the MMA, warps 4-7 run the epilogue, warps -# 2-3 idle. The epilogue must start on a warp quad (warp id multiple of 4): -# tcgen05.ld selects its TMEM datapath sub-partition from the physical warp -# id, and a misaligned epilogue block would read every 128-row tile with its -# 32-row groups rotated. -_THREADS = 256 -_TMA_WARP_ID = 0 -_MMA_WARP_ID = 1 -_FIRST_EPI_WARP = 4 -_NUM_EPI_THREADS = 128 -_FIRST_EPI_THREAD = 32 * _FIRST_EPI_WARP -# Named barrier ids (0 is left free for the DSL's own use). -_EPI_FINAL_BARRIER_ID = 1 -_TMEM_ALLOC_BARRIER_ID = 2 -_EPI_STAGE_BARRIER_ID = 3 -# sm_100 usable dynamic shared memory per CTA, (228 - 1) KiB. -_SMEM_CAPACITY_BYTES = 232448 -# The TMEM allocator requires a power-of-two multiple of 32 columns; shared -# memory already pins us to one CTA per SM, so taking the whole array is free. -_TMEM_TOTAL_COLS = 512 - - -@dataclass(frozen=True) -class _KernelConfig: - """Per-kernel trace-time constants (everything else is module-frozen).""" - - # Accumulator columns handed to the epilogue per subtile. 64 for the - # forward (an adjacent gate/up accumulator pair per output column, so 64 - # accumulator columns are 32 output columns = one 1x32 block per row), 32 - # for the backward and for wgrad. - epi_n_acc: int - # False: offsets partition the GEMM M axis (forward/backward). True: they - # partition the contraction (wgrad). - ragged_k: bool - # Columns of the [128, cols] BF16 epilogue staging tile (0 = no staging). - # Padded to an odd count so columnwise reads don't serialize on banks. - epi_smem_cols: int - - -_SWIGLU_FWD_CONFIG = _KernelConfig(epi_n_acc=64, ragged_k=False, epi_smem_cols=33) -_DSWIGLU_BWD_CONFIG = _KernelConfig(epi_n_acc=32, ragged_k=False, epi_smem_cols=65) -_WGRAD_CONFIG = _KernelConfig(epi_n_acc=32, ragged_k=True, epi_smem_cols=0) - - -# --- Host-side operand views: pure torch restrides into the (MN, K, L) GEMM -# domain, K contiguous. - - -def activation_gemm_view(t: torch.Tensor) -> torch.Tensor: - """``[MN, K]`` K-contiguous -> ``(MN, K, 1)`` with a defined batch stride.""" - mn, k = t.shape - if t.stride() != (k, 1): - raise ValueError( - f"GEMM operand must be K-contiguous with stride {(k, 1)}, got {t.stride()}" - ) - return torch.as_strided(t, (mn, k, 1), (k, 1, mn * k)) - - -def weight_gemm_view(w: torch.Tensor) -> torch.Tensor: - """``[G, K, N]`` stride ``(K*N, 1, K)`` -> ``(N, K, G)``: expert becomes L.""" - g, k, n = w.shape - if w.stride() != (k * n, 1, k): - raise ValueError( - f"grouped weight must have stride {(k * n, 1, k)}, got {w.stride()}" - ) - return w.permute(2, 1, 0) - - -@dataclass -class TileCoords: - """Per-CTA tile description handed to the epilogue. All fields CTA-uniform. - - ``k_cnt == 0`` marks both an inactive tail tile (ragged M) and a - zero-token expert (ragged K); the accumulator arrives zeroed and the - epilogue must store unconditionally. - """ - - expert: Int32 - row_base: Int32 - col_base: Int32 - k_cnt: Int32 - - -def _s2t_copy_and_partition(sSF: cute.Tensor, tSF: cute.Tensor): - """SMEM -> TMEM scale-factor copy, issued once per K tile from the MMA warp. - - ``Cp4x32x128bOp`` must be issued as a plain ``cute.copy`` -- the DSL - inserts the single-thread election itself, and wrapping it in - ``elect_one()`` deadlocks. - """ - tCsSF_compact = cute.filter_zeros(sSF) - tCtSF_compact = cute.filter_zeros(tSF) - copy_atom_s2t = cute.make_copy_atom( - tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), sSF.element_type - ) - tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) - thr_copy_s2t = tiled_copy_s2t.get_slice(0) - tCsSF_s2t = tcgen05.get_s2t_smem_desc_tensor( - tiled_copy_s2t, thr_copy_s2t.partition_S(tCsSF_compact) - ) - tCtSF_s2t = thr_copy_s2t.partition_D(tCtSF_compact) - return tiled_copy_s2t, tCsSF_s2t, tCtSF_s2t - - -# --- Quantization math: public conversions only; every step mirrors the torchao -# reference (`to_mx(..., RCEIL)` + `to_blocked`) so the fused outputs are -# byte-identical to the standalone quantizers on the same BF16 input. - - -@cute.jit -def _blocked_scale_idx(row: Int32, scale_col: Int32, ncb: Int32) -> Int32: - """Flat byte index in the tcgen05 blocked (128x4) scale layout. - - The logical ``[rows, cols]`` scale matrix is stored as 512-byte tiles of - 128 rows x 4 scale columns, tiles ordered ``row_block * ncb + col_block`` - with ``ncb = ceil_div(cols, 4)`` -- torchao ``to_blocked``, whole-matrix. - Coordinates are ABSOLUTE. For columnwise scales pass transposed - coordinates (feature index as ``row``, 32-row block index as ``col``). - """ - return ( - ((row >> 7) * ncb + (scale_col >> 2)) * Int32(512) - + (row & Int32(31)) * Int32(16) - + ((row >> 5) & Int32(3)) * Int32(4) - + (scale_col & Int32(3)) - ) - - -@cute.jit -def _store_frag(dst: cute.Tensor, elem_offset: Int32, frag: cute.Tensor): - """Store a register fragment as one contiguous run of ``dst`` elements. - - Callers guarantee the element offset is at least 16-byte aligned relative - to the (validated, 32-byte-aligned) base, so the copy vectorizes. - """ - cute.autovec_copy( - frag, - cute.make_tensor( - (dst.iterator + elem_offset).align(16), - cute.make_layout(cute.size(frag)), - ), - ) - - -@cute.jit -def _quant_block_from_smem( - sEpi: cute.Tensor, - base_row: Int32, - base_col: Int32, - q_dst: cute.Tensor, - q_offset: Int32, - sf_dst: cute.Tensor, - sf_idx: Int32, - COLWISE: cutlass.Constexpr, -): - """Quantize one 32-value MX block read from the BF16 staging tile - (``sEpi[base_row, base_col + i]`` rowwise, ``sEpi[base_row + i, base_col]`` - columnwise): - - * amax: NaN-propagating |max| chain, so a NaN element invalidates the - block exactly like the torchao reference. - * scale: ``descale = amax / 448``; the public Float32 -> Float8E8M0FNU - conversion rounds toward +inf (RCEIL) by construction. A non-finite - amax (Inf would otherwise clamp to byte 254) is overridden to byte 255. - * reciprocal, in the E8M0 byte domain: ``254 - byte`` reinterpreted as - E8M0 and widened exactly to f32. Byte 0 (zero/tiny block) descales by - 2^127; byte 255 gives a NaN reciprocal so every element of an - invalidated block quantizes to the E4M3 NaN code. - * qdata: one f32 multiply per element, then the public saturating-RNE - Float32 -> Float8E4M3FN conversion (byte-identical to torch's cast). - - The 32 qdata bytes are one contiguous ``q_dst`` run in both orientations, - stored vectorized; the scale byte is stored individually at ``sf_idx``. - """ - vals = [] - for i in cutlass.range_constexpr(_SF_VEC_SIZE): - if cutlass.const_expr(COLWISE): - v = sEpi[base_row + Int32(i), base_col] - else: - v = sEpi[base_row, base_col + Int32(i)] - vals.append(Float32(v)) - - amax = cute.arch.fmax(vals[0], vals[1], abs=True, nan=True) - for i in cutlass.range_constexpr(2, _SF_VEC_SIZE): - amax = cute.arch.fmax(amax, vals[i], abs=True, nan=True) - - scale_byte = Int32( - cutlass.Float8E8M0FNU(amax / Float32(448.0)).bitcast(cutlass.Int8) - ) & Int32(0xFF) - amax_bits = Float32(amax).bitcast(Int32) - if (amax_bits & Int32(0x7F800000)) == Int32(0x7F800000): - scale_byte = Int32(255) - - recip_byte = (Int32(254) - scale_byte) & Int32(0xFF) - recip = Float32(cutlass.Uint8(recip_byte).bitcast(cutlass.Float8E8M0FNU)) - - qfrag = cute.make_rmem_tensor(cute.make_layout(_SF_VEC_SIZE), cutlass.Float8E4M3FN) - for i in cutlass.range_constexpr(_SF_VEC_SIZE): - qfrag[i] = cutlass.Float8E4M3FN(vals[i] * recip) - _store_frag(q_dst, q_offset, qfrag) - sf_dst[sf_idx] = cutlass.Uint8(scale_byte) - - -@cute.jit -def _epilogue_column_run( - tTR_cAcc_s, num_acc: cutlass.Constexpr, even: cutlass.Constexpr -): - """Trace-time proof that a thread's fragment is one contiguous column run. - - Returns the (static) first column. The column coordinates fold to Python - ints at trace time; requiring one even-based, contiguous, increasing run - also proves the fragment covers a single row (a two-row fragment would - repeat columns), so this CHECKS the layout the epilogues need instead of - assuming a physical thread-to-row mapping. - """ - cols = [] - for v in cutlass.range_constexpr(num_acc): - cols.append(tTR_cAcc_s[v][1]) - first = cols[0] - if cutlass.const_expr( - not all(isinstance(c, int) for c in cols) - or tuple(cols) != tuple(range(first, first + num_acc)) - or (even and first % 2 != 0) - ): - raise ValueError( - "the epilogue needs one contiguous, increasing" - + (", even-based" if even else "") - + f" column run per thread, but tTR_cAcc gave {cols}" - ) - return first - - -# --- Epilogues: called once per subtile by all 128 epilogue threads with a -# CTA-uniform k_cnt; stores are never predicated (a tail tile's zeroed -# accumulator produces exactly the zero bytes the contract requires). - - -@cute.jit -def _wgrad_epilogue( - tTR_rAcc, - tTR_cAcc_s, - epi_tidx, - tile: TileCoords, - sEpi, - out, - R: cutlass.Constexpr, - N: cutlass.Constexpr, -): - """Round the FP32 accumulator subtile to BF16 and store it (no quant). - - ``out`` is ``(mDw,)`` with ``mDw`` the ``(N, K, G)`` view of the - contiguous ``[G, N, K]`` destination. - """ - num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) - frag_col = _epilogue_column_run(tTR_cAcc_s, num_acc, even=False) - - frag = cute.make_rmem_tensor(cute.make_layout(num_acc), cutlass.BFloat16) - for v in cutlass.range_constexpr(num_acc): - frag[v] = cutlass.BFloat16(tTR_rAcc[v]) - - gDw = out[0] - strides = gDw.stride - if cutlass.const_expr(strides[1] != 1): - raise ValueError( - f"the wgrad epilogue stores a contiguous run along K, so dw's K " - f"stride must be 1, got layout {gDw.layout}" - ) - row = tile.row_base + tTR_cAcc_s[0][0] - elem = ( - row * Int32(strides[0]) - + (tile.col_base + Int32(frag_col)) - + tile.expert * Int32(strides[2]) - ) - _store_frag(gDw, elem, frag) - - -@cute.jit -def _swiglu_fwd_epilogue( - tTR_rAcc, - tTR_cAcc_s, - epi_tidx, - tile: TileCoords, - sEpi, - out, - R: cutlass.Constexpr, - N: cutlass.Constexpr, -): - """SwiGLU + BF16 pre-activation save + dual MXFP8 quantization. - - ``out`` = flat views ``(z [R*2F] bf16, h_row_q [R*F] e4m3, - h_row_sf uint8, h_col_q [F*R] e4m3 in column-major storage order, - h_col_sf uint8)``; ``N == 2F`` is the GEMM (and z) column count. Contract: - the accumulator is rounded to BF16 first (that IS z), SwiGLU is evaluated - once from the rounded values, h is rounded to BF16 once, and both - quantizers consume the same staged BF16 h. - """ - num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) # 64 - half = cutlass.const_expr(num_acc // 2) # 32 h columns per subtile - F = cutlass.const_expr(N // 2) - frag_col = _epilogue_column_run(tTR_cAcc_s, num_acc, even=True) - - mZ = out[0] - mHrowQ = out[1] - mHrowSF = out[2] - mHcolQ = out[3] - mHcolSF = out[4] - - # ---- stage 1: z store + h compute into the staging tile -------------- - lrow = tTR_cAcc_s[0][0] # CTA-tile-local row of this thread's fragment - row_g = tile.row_base + lrow - zfrag = cute.make_rmem_tensor(cute.make_layout(num_acc), cutlass.BFloat16) - for v in cutlass.range_constexpr(num_acc): - zfrag[v] = cutlass.BFloat16(tTR_rAcc[v]) - _store_frag(mZ, row_g * Int32(N) + tile.col_base + Int32(frag_col), zfrag) - - for j in cutlass.range_constexpr(half): - gate = Float32(zfrag[2 * j]) - up = Float32(zfrag[2 * j + 1]) - # sigmoid composed exactly like torch's float32 sigmoid (default-mode - # exp plus a true divide); measured bit-identical to torch.sigmoid. - sig = Float32(1.0) / (Float32(1.0) + cute.math.exp(Float32(0.0) - gate)) - sEpi[lrow, Int32(j)] = cutlass.BFloat16((gate * sig) * up) - - cute.arch.barrier( - barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS - ) - - # ---- stage 2: dual quantization off the staging tile ----------------- - # Global h column base of this subtile; frag_col is static, 64 per - # subtile, so hbase is divisible by 32. - hbase = (tile.col_base + Int32(frag_col)) >> 1 - ncb_row = cutlass.const_expr((F // _SF_VEC_SIZE + 3) // 4) - ncb_col = cutlass.const_expr((R // _SF_VEC_SIZE + 3) // 4) - - # Rowwise 1x32: one block per thread (row = epi_tidx of the staging tile). - q_row = tile.row_base + epi_tidx - _quant_block_from_smem( - sEpi, - epi_tidx, - Int32(0), - mHrowQ, - q_row * Int32(F) + hbase, - mHrowSF, - _blocked_scale_idx( - q_row, (tile.col_base + Int32(frag_col)) >> 6, Int32(ncb_row) - ), - COLWISE=False, - ) - - # Columnwise 32x1: 32 columns x 4 row-blocks = one block per thread. - col_l = epi_tidx & Int32(31) - blk = epi_tidx >> 5 - col_g = hbase + col_l - _quant_block_from_smem( - sEpi, - blk * Int32(32), - col_l, - mHcolQ, - col_g * Int32(R) + tile.row_base + blk * Int32(32), - mHcolSF, - _blocked_scale_idx(col_g, (tile.row_base >> 5) + blk, Int32(ncb_col)), - COLWISE=True, - ) - - # The next subtile reuses the staging tile. - cute.arch.barrier( - barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS - ) - - -@cute.jit -def _dswiglu_bwd_epilogue( - tTR_rAcc, - tTR_cAcc_s, - epi_tidx, - tile: TileCoords, - sEpi, - out, - R: cutlass.Constexpr, - N: cutlass.Constexpr, -): - """dSwiGLU from the saved z + dual MXFP8 quantization of dz. - - ``out`` = ``(z [R*2F] bf16 INPUT, dz_row_q [R*2F] e4m3, dz_row_sf uint8, - dz_col_q [2F*R] e4m3 column-major storage, dz_col_sf uint8)``; ``N == F`` - is the dgrad GEMM column count; dz has 2F element-interleaved columns. - Contract: dh is rounded to BF16 first; gate/up come from the saved BF16 z; - dgate/dup are each rounded to BF16 before interleaving; both quantizers - consume the same staged BF16 dz. The z load is predicated on ``k_cnt`` -- - tail rows of z are read-forbidden and contribute zeros. - """ - num_acc = cutlass.const_expr(cute.size(tTR_rAcc)) # 32 - two_f = cutlass.const_expr(2 * N) - frag_col = _epilogue_column_run(tTR_cAcc_s, num_acc, even=False) - - mZ = out[0] - mDzRowQ = out[1] - mDzRowSF = out[2] - mDzColQ = out[3] - mDzColSF = out[4] - - # ---- stage 1: z load (predicated) + dSwiGLU into the staging tile ---- - lrow = tTR_cAcc_s[0][0] - row_g = tile.row_base + lrow - # dz columns covered by this subtile: [dzbase, dzbase + 64). - dzbase = (tile.col_base + Int32(frag_col)) * Int32(2) - - zfrag = cute.make_rmem_tensor(cute.make_layout(2 * num_acc), cutlass.BFloat16) - if tile.k_cnt > Int32(0): - cute.autovec_copy( - cute.make_tensor( - (mZ.iterator + (row_g * Int32(two_f) + dzbase)).align(16), - cute.make_layout(2 * num_acc), - ), - zfrag, - ) - else: - for i in cutlass.range_constexpr(2 * num_acc): - zfrag[i] = cutlass.BFloat16(0.0) - - for j in cutlass.range_constexpr(num_acc): - gate = Float32(zfrag[2 * j]) - up = Float32(zfrag[2 * j + 1]) - dh = Float32(cutlass.BFloat16(tTR_rAcc[j])) - sig = Float32(1.0) / (Float32(1.0) + cute.math.exp(Float32(0.0) - gate)) - silu = gate * sig - dsilu = sig * (Float32(1.0) + gate * (Float32(1.0) - sig)) - sEpi[lrow, Int32(2 * j)] = cutlass.BFloat16((dh * up) * dsilu) - sEpi[lrow, Int32(2 * j + 1)] = cutlass.BFloat16(dh * silu) - - cute.arch.barrier( - barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS - ) - - # ---- stage 2: dual quantization off the staging tile ----------------- - ncb_row = cutlass.const_expr((two_f // _SF_VEC_SIZE + 3) // 4) - ncb_col = cutlass.const_expr((R // _SF_VEC_SIZE + 3) // 4) - - # Rowwise 1x32: 128 rows x 2 blocks = two tasks per thread. - q_row = tile.row_base + epi_tidx - for blk in cutlass.range_constexpr(2): - _quant_block_from_smem( - sEpi, - epi_tidx, - Int32(blk * _SF_VEC_SIZE), - mDzRowQ, - q_row * Int32(two_f) + dzbase + Int32(blk * _SF_VEC_SIZE), - mDzRowSF, - _blocked_scale_idx(q_row, (dzbase >> 5) + Int32(blk), Int32(ncb_row)), - COLWISE=False, - ) - - # Columnwise 32x1: 64 columns x 4 row-blocks = two tasks per thread. - for k in cutlass.range_constexpr(2): - task = epi_tidx + Int32(128 * k) - col_l = task & Int32(63) - blk = task >> 6 - col_g = dzbase + col_l - _quant_block_from_smem( - sEpi, - blk * Int32(32), - col_l, - mDzColQ, - col_g * Int32(R) + tile.row_base + blk * Int32(32), - mDzColSF, - _blocked_scale_idx(col_g, (tile.row_base >> 5) + blk, Int32(ncb_col)), - COLWISE=True, - ) - - cute.arch.barrier( - barrier_id=_EPI_STAGE_BARRIER_ID, number_of_threads=_NUM_EPI_THREADS - ) - - -# --- The shared kernel and launch builder. - - -@cute.kernel -def _grouped_gemm_kernel( - tiled_mma: cute.TiledMma, - tma_atom_a: cute.CopyAtom, - mA: cute.Tensor, - tma_atom_b: cute.CopyAtom, - mB: cute.Tensor, - tma_atom_sfa: cute.CopyAtom, - mSFA: cute.Tensor, - tma_atom_sfb: cute.CopyAtom, - mSFB: cute.Tensor, - offs: cute.Tensor, - out, - a_smem_layout: cute.ComposedLayout, - b_smem_layout: cute.ComposedLayout, - sfa_smem_layout: cute.Layout, - sfb_smem_layout: cute.Layout, - cfg: cutlass.Constexpr, - storage_type: cutlass.Constexpr, - EPILOGUE: cutlass.Constexpr, -): - """One CTA computes one 128 x 128 output tile. Warps: 0 TMA, 1 MMA, - 4-7 epilogue, 2-3 idle.""" - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - bidx, bidy, bidz = cute.arch.block_idx() - - num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) - num_k_tiles_full = cutlass.const_expr(cute.size(mA, mode=[1]) // _CTA_K) - gemm_m = cutlass.const_expr(cute.size(mA, mode=[0])) - gemm_n = cutlass.const_expr(cute.size(mB, mode=[0])) - - # ------------------------------------------------------------------ - # Tile coordinates. No scheduler: the grid IS the tile enumeration. - # ------------------------------------------------------------------ - tile_m = Int32(bidx) - tile_n = Int32(bidy) - if cutlass.const_expr(not cfg.ragged_k): - row_base = tile_m * _CTA_M - # Unrolled G-way scan: the owning expert is the number of groups that - # end at or before this tile's row base; a zero-token expert is never - # selected since its end equals its start. - expert = Int32(0) - for g in cutlass.range_constexpr(num_groups - 1): - expert += Int32(offs[g] <= row_base) - is_active = Int32(offs[num_groups - 1] > row_base) - k_base = Int32(0) - k_cnt = is_active * Int32(num_k_tiles_full) - l_b = expert - else: - expert = Int32(bidz) - row_base = tile_m * _CTA_M - prev = offs[cutlass.max(expert - Int32(1), Int32(0))] * Int32(expert > Int32(0)) - # Exact, not a ceil: every group boundary is a multiple of cta_tile_k. - k_base = prev // _CTA_K - # Clamp: nonmonotone offsets (undefined behavior per the contract, and - # only device-checkable) would otherwise make k_cnt negative -- the - # loops still run zero trips, but the epilogue would read TMEM the MMA - # never wrote. Clamped, a malformed expert degrades to an all-zero - # slice instead. - k_cnt = cutlass.max((offs[expert] - prev) // _CTA_K, Int32(0)) - l_b = Int32(0) - tile = TileCoords( - expert=expert, - row_base=row_base, - col_base=tile_n * _CTA_N, - k_cnt=k_cnt, - ) - - # ------------------------------------------------------------------ - # Shared memory and pipelines - # ------------------------------------------------------------------ - smem = utils.SmemAllocator() - storage = smem.allocate(storage_type) - - sA = storage.sA.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) - sB = storage.sB.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) - sSFA = storage.sSFA.get_tensor(sfa_smem_layout) - sSFB = storage.sSFB.get_tensor(sfb_smem_layout) - sEpi = None - if cutlass.const_expr(cfg.epi_smem_cols > 0): - sEpi = storage.sEpi.get_tensor( - cute.make_layout((_CTA_M, cfg.epi_smem_cols), stride=(cfg.epi_smem_cols, 1)) - ) - - cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) - ) - - ab_pipeline = pipeline.PipelineTmaUmma.create( - barrier_storage=storage.ab_full_mbar.data_ptr(), - num_stages=_NUM_AB_STAGE, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - # The EXACT byte count of the four TMA copies of one stage: too small - # and the MMA consumes a partially arrived stage. - tx_count=_AB_STAGE_BYTES, - cta_layout_vmnk=cluster_layout_vmnk, - ) - acc_pipeline = pipeline.PipelineUmmaAsync.create( - barrier_storage=storage.acc_full_mbar.data_ptr(), - num_stages=_NUM_ACC_STAGE, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, _NUM_EPI_THREADS - ), - cta_layout_vmnk=cluster_layout_vmnk, - ) - - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=_TMEM_ALLOC_BARRIER_ID, - num_threads=32 * 5, # the MMA warp joins the four epilogue warps - ) - epilogue_barrier = pipeline.NamedBarrier( - barrier_id=_EPI_FINAL_BARRIER_ID, - num_threads=_NUM_EPI_THREADS, - ) - tmem_alloc = utils.TmemAllocator( - storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=_FIRST_EPI_WARP, - is_two_cta=False, - ) - - # ------------------------------------------------------------------ - # Tile the global tensors. One static descriptor per operand; the expert - # is an L coordinate (ragged M) or a K-tile index base (ragged K). - # ------------------------------------------------------------------ - gA = cute.local_tile( - mA, cute.slice_(_MMA_TILER, (None, 0, None)), (None, None, None) - ) - gB = cute.local_tile( - mB, cute.slice_(_MMA_TILER, (0, None, None)), (None, None, None) - ) - gSFA = cute.local_tile( - mSFA, cute.slice_(_MMA_TILER, (None, 0, None)), (None, None, None) - ) - gSFB = cute.local_tile( - mSFB, cute.slice_(_MMA_TILER, (0, None, None)), (None, None, None) - ) - - thr_mma = tiled_mma.get_slice(0) - tCgA = thr_mma.partition_A(gA) - tCgB = thr_mma.partition_B(gB) - tCgSFA = thr_mma.partition_A(gSFA) - tCgSFB = thr_mma.partition_B(gSFB) - - trivial_cta_layout = cute.make_layout(1) - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, - 0, - trivial_cta_layout, - cute.group_modes(sA, 0, 3), - cute.group_modes(tCgA, 0, 3), - ) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, - 0, - trivial_cta_layout, - cute.group_modes(sB, 0, 3), - cute.group_modes(tCgB, 0, 3), - ) - tAsSFA, tAgSFA = cpasync.tma_partition( - tma_atom_sfa, - 0, - trivial_cta_layout, - cute.group_modes(sSFA, 0, 3), - cute.group_modes(tCgSFA, 0, 3), - ) - tAsSFA = cute.filter_zeros(tAsSFA) - tAgSFA = cute.filter_zeros(tAgSFA) - tBsSFB, tBgSFB = cpasync.tma_partition( - tma_atom_sfb, - 0, - trivial_cta_layout, - cute.group_modes(sSFB, 0, 3), - cute.group_modes(tCgSFB, 0, 3), - ) - tBsSFB = cute.filter_zeros(tBsSFB) - tBgSFB = cute.filter_zeros(tBgSFB) - - tAgA_slice = tAgA[(None, tile_m, None, 0)] - tBgB_slice = tBgB[(None, tile_n, None, l_b)] - tAgSFA_slice = tAgSFA[(None, tile_m, None, 0)] - tBgSFB_slice = tBgSFB[(None, tile_n, None, l_b)] - - acc_shape = tiled_mma.partition_shape_C(_MMA_TILER[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, _NUM_ACC_STAGE)) - - # ------------------------------------------------------------------ - # Warp 0: TMA producer - # ------------------------------------------------------------------ - if warp_idx == _TMA_WARP_ID: - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, _NUM_AB_STAGE - ) - for _ in cutlass.range(0, k_cnt, 1, unroll=1): - ab_pipeline.producer_acquire(ab_producer_state) - k_idx = k_base + ab_producer_state.count - bar = ab_pipeline.producer_get_barrier(ab_producer_state) - cute.copy( - tma_atom_a, - tAgA_slice[(None, k_idx)], - tAsA[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - cute.copy( - tma_atom_b, - tBgB_slice[(None, k_idx)], - tBsB[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - cute.copy( - tma_atom_sfa, - tAgSFA_slice[(None, k_idx)], - tAsSFA[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - cute.copy( - tma_atom_sfb, - tBgSFB_slice[(None, k_idx)], - tBsSFB[(None, ab_producer_state.index)], - tma_bar_ptr=bar, - ) - ab_producer_state.advance() - ab_pipeline.producer_tail(ab_producer_state) - - # ------------------------------------------------------------------ - # Warp 1: MMA - # ------------------------------------------------------------------ - if warp_idx == _MMA_WARP_ID: - tCrA = tiled_mma.make_fragment_A(sA) - tCrB = tiled_mma.make_fragment_B(sB) - - # The MMA warp joins the allocation barrier but must never allocate. - tmem_alloc.wait_for_alloc() - acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - - sfa_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), - dtype=sSFA.element_type, - ) - tCtSFA = cute.make_tensor( - sfa_tmem_ptr, - blockscaled_utils.make_tmem_layout_sfa( - tiled_mma, - _MMA_TILER, - _SF_VEC_SIZE, - cute.slice_(sfa_smem_layout, (None, None, None, 0)), - ), - ) - sfb_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr - + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) - + tcgen05.find_tmem_tensor_col_offset(tCtSFA), - dtype=sSFB.element_type, - ) - tCtSFB = cute.make_tensor( - sfb_tmem_ptr, - blockscaled_utils.make_tmem_layout_sfb( - tiled_mma, - _MMA_TILER, - _SF_VEC_SIZE, - cute.slice_(sfb_smem_layout, (None, None, None, 0)), - ), - ) - s2t_sfa, tCsSFA_s2t, tCtSFA_s2t = _s2t_copy_and_partition(sSFA, tCtSFA) - s2t_sfb, tCsSFB_s2t, tCtSFB_s2t = _s2t_copy_and_partition(sSFB, tCtSFB) - - ab_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, _NUM_AB_STAGE - ) - acc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, _NUM_ACC_STAGE - ) - tCtAcc = tCtAcc_base[(None, None, None, 0)] - - # Acquire and commit unconditionally, including when k_cnt == 0, so - # the accumulator handoff barrier stays balanced on tail tiles. - acc_pipeline.producer_acquire(acc_producer_state) - for k_tile in cutlass.range(0, k_cnt, 1, unroll=1): - ab_pipeline.consumer_wait(ab_consumer_state) - stage_crd = (None, None, None, None, ab_consumer_state.index) - cute.copy(s2t_sfa, tCsSFA_s2t[stage_crd], tCtSFA_s2t) - cute.copy(s2t_sfb, tCsSFB_s2t[stage_crd], tCtSFB_s2t) - # ACCUMULATE=False on the first K tile is what zeroes the - # accumulator; there is no separate TMEM clear. - tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) - mma_crd = (None, None, None, ab_consumer_state.index) - cute.gemm( - tiled_mma, - tCtAcc, - [tCrA[mma_crd], tCtSFA], - [tCrB[mma_crd], tCtSFB], - tCtAcc, - ) - ab_pipeline.consumer_release(ab_consumer_state) - ab_consumer_state.advance() - acc_pipeline.producer_commit(acc_producer_state) - - # ------------------------------------------------------------------ - # Warps 4-7: epilogue - # ------------------------------------------------------------------ - if warp_idx >= _FIRST_EPI_WARP: - tmem_alloc.allocate(_TMEM_TOTAL_COLS) - tmem_alloc.wait_for_alloc() - acc_tmem_ptr = tmem_alloc.retrieve_ptr(Float32) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - - epi_tidx = tidx - _FIRST_EPI_THREAD - # TMEM -> register handoff. tTR_cAcc carries each register's (row, col) - # coordinate in the CTA tile; every epilogue index is derived from it, - # not from an assumed thread-to-row mapping. The d element type stays - # Float32: an 8-bit d would steer get_tmem_load_op into layouts shaped - # for a direct FP8 TMA store. - epi_tile = (_CTA_M, cfg.epi_n_acc) - copy_atom_t2r = sm100_utils.get_tmem_load_op( - _MMA_TILER, LayoutEnum.ROW_MAJOR, Float32, Float32, epi_tile, False - ) - tAcc_mn = tCtAcc_base[((None, None), 0, 0, 0)] - tAcc_epi = cute.flat_divide(tAcc_mn, epi_tile) - tiled_copy_t2r = tcgen05.make_tmem_copy( - copy_atom_t2r, tAcc_epi[(None, None, 0, 0)] - ) - thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - cAcc_epi = cute.flat_divide( - cute.make_identity_tensor((_CTA_M, _CTA_N)), epi_tile - ) - tTR_cAcc = thr_copy_t2r.partition_D(cAcc_epi) - tTR_rAcc = cute.make_rmem_tensor( - tTR_cAcc[(None, None, None, 0, 0)].shape, Float32 - ) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, _NUM_ACC_STAGE - ) - acc_pipeline.consumer_wait(acc_consumer_state) - - for s in cutlass.range_constexpr(_CTA_N // cfg.epi_n_acc): - if k_cnt == Int32(0): - # Tail tile or zero-token expert: nothing was accumulated, so - # the fragment is zeroed here and the epilogue runs unchanged. - for v in cutlass.range_constexpr(cute.size(tTR_rAcc)): - tTR_rAcc[v] = Float32(0.0) - else: - cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, 0, s)], tTR_rAcc) - EPILOGUE( - tTR_rAcc, - tTR_cAcc[(None, None, None, 0, s)], - epi_tidx, - tile, - sEpi, - out, - gemm_m, - gemm_n, - ) - - cute.arch.fence_view_async_tmem_load() - acc_pipeline.consumer_release(acc_consumer_state) - tmem_alloc.relinquish_alloc_permit() - epilogue_barrier.arrive_and_wait() - tmem_alloc.free(acc_tmem_ptr) - - -@cute.jit -def _launch_grouped_gemm( - mA: cute.Tensor, - mB: cute.Tensor, - sfa_flat: cute.Tensor, - sfb_flat: cute.Tensor, - offs: cute.Tensor, - out, - stream, - cfg: cutlass.Constexpr, - EPILOGUE: cutlass.Constexpr, -): - """Build the four static TMA descriptors and launch. Grid is data-independent. - - This is a trace body: calling it directly retraces the kernel on every - invocation. The public launchers ``cute.compile`` it once per shape key, - passing ``cfg`` and ``EPILOGUE`` as trailing Constexpr args, then call the - compiled executor with the runtime args only. - """ - a_dtype = mA.element_type - b_dtype = mB.element_type - sf_dtype = cutlass.Float8E8M0FNU - - gemm_m = cutlass.const_expr(cute.size(mA, mode=[0])) - gemm_k = cutlass.const_expr(cute.size(mA, mode=[1])) - gemm_n = cutlass.const_expr(cute.size(mB, mode=[0])) - l_a = cutlass.const_expr(cute.size(mA, mode=[2])) - l_b = cutlass.const_expr(cute.size(mB, mode=[2])) - num_groups = cutlass.const_expr(cute.size(offs, mode=[0])) - - if cutlass.const_expr( - gemm_m % _CTA_M != 0 or gemm_n % _CTA_N != 0 or gemm_k % _CTA_K != 0 - ): - raise ValueError( - f"GEMM extents ({gemm_m}, {gemm_n}, {gemm_k}) must be multiples of " - f"({_CTA_M}, {_CTA_N}, {_CTA_K})" - ) - - # Retile the flat blocked E8M0 buffers into the GEMM-domain SF layout. The - # buffers travel flat by ABI and may arrive as raw uint8, so the pointer is - # recast: the MMA rejects a scale operand that is not E8M0. - mSFA = cute.make_tensor( - cute.recast_ptr(sfa_flat.iterator, dtype=cutlass.Float8E8M0FNU), - blockscaled_utils.tile_atom_to_shape_SF((gemm_m, gemm_k, l_a), _SF_VEC_SIZE), - ) - mSFB = cute.make_tensor( - cute.recast_ptr(sfb_flat.iterator, dtype=cutlass.Float8E8M0FNU), - blockscaled_utils.tile_atom_to_shape_SF((gemm_n, gemm_k, l_b), _SF_VEC_SIZE), - ) - - # The one blockscaled tiled MMA, K-major on both operands, FP32 acc. - tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( - a_dtype, - b_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.K, - sf_dtype, - _SF_VEC_SIZE, - tcgen05.CtaGroup.ONE, - (_CTA_M, _CTA_N), - ) - cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,) - ) - - a_smem_layout = sm100_utils.make_smem_layout_a( - tiled_mma, _MMA_TILER, a_dtype, _NUM_AB_STAGE - ) - b_smem_layout = sm100_utils.make_smem_layout_b( - tiled_mma, _MMA_TILER, b_dtype, _NUM_AB_STAGE - ) - sfa_smem_layout = blockscaled_utils.make_smem_layout_sfa( - tiled_mma, _MMA_TILER, _SF_VEC_SIZE, _NUM_AB_STAGE - ) - sfb_smem_layout = blockscaled_utils.make_smem_layout_sfb( - tiled_mma, _MMA_TILER, _SF_VEC_SIZE, _NUM_AB_STAGE - ) - - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id), - mA, - cute.slice_(a_smem_layout, (None, None, None, 0)), - _MMA_TILER, - tiled_mma, - cluster_layout_vmnk.shape, - ) - tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( - sm100_utils.cluster_shape_to_tma_atom_B((1, 1), tiled_mma.thr_id), - mB, - cute.slice_(b_smem_layout, (None, None, None, 0)), - _MMA_TILER, - tiled_mma, - cluster_layout_vmnk.shape, - ) - # The 512-byte scale atom is contiguous; TMA moves it as 8-byte elements. - tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( - sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id), - mSFA, - cute.slice_(sfa_smem_layout, (None, None, None, 0)), - _MMA_TILER, - tiled_mma, - cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64, - ) - tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( - sm100_utils.cluster_shape_to_tma_atom_SFB((1, 1), tiled_mma.thr_id), - mSFB, - cute.slice_(sfb_smem_layout, (None, None, None, 0)), - _MMA_TILER, - tiled_mma, - cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64, - ) - - @cute.struct - class SharedStorage: - # The *_empty ranges are live storage: each Pipeline*.create consumes - # 2 x num_stages barriers starting at the *_full pointer. - ab_full_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_AB_STAGE] - ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_AB_STAGE] - acc_full_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_ACC_STAGE] - acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, _NUM_ACC_STAGE] - tmem_holding_buf: cutlass.Int32 - # A zero-size struct field would be degenerate; keep a tiny slab. - sEpi: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, max(_CTA_M * cfg.epi_smem_cols, 8)], - 128, - ] - sA: cute.struct.Align[ - cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024 - ] - sB: cute.struct.Align[ - cute.struct.MemRange[b_dtype, cute.cosize(b_smem_layout.outer)], 1024 - ] - sSFA: cute.struct.Align[ - cute.struct.MemRange[sf_dtype, cute.cosize(sfa_smem_layout)], 1024 - ] - sSFB: cute.struct.Align[ - cute.struct.MemRange[sf_dtype, cute.cosize(sfb_smem_layout)], 1024 - ] - - smem_bytes = cutlass.const_expr(SharedStorage.size_in_bytes()) - if cutlass.const_expr(smem_bytes > _SMEM_CAPACITY_BYTES): - raise ValueError( - f"shared memory request {smem_bytes} B exceeds the sm_100 capacity " - f"{_SMEM_CAPACITY_BYTES} B" - ) - - if cutlass.const_expr(not cfg.ragged_k): - grid = (gemm_m // _CTA_M, gemm_n // _CTA_N, 1) - else: - grid = (gemm_m // _CTA_M, gemm_n // _CTA_N, num_groups) - - _grouped_gemm_kernel( - tiled_mma, - tma_atom_a, - tma_tensor_a, - tma_atom_b, - tma_tensor_b, - tma_atom_sfa, - tma_tensor_sfa, - tma_atom_sfb, - tma_tensor_sfb, - offs, - out, - a_smem_layout, - b_smem_layout, - sfa_smem_layout, - sfb_smem_layout, - cfg, - SharedStorage, - EPILOGUE, - ).launch( - grid=grid, - block=(_THREADS, 1, 1), - cluster=(1, 1, 1), - smem=smem_bytes, - stream=stream, - ) - - -@functools.cache -def _executor_slot(key: tuple) -> list: - """One memo slot per (kernel, shape, device, dtype, DSL version) key; the - trace needs real tensors, so the first caller compiles and fills it once.""" - return [] - - -def _cache_key(kind: str, dims: tuple, tensors: tuple, device) -> tuple: - return ( - kind, - dims, - tuple(str(t.dtype) for t in tensors), - device.index, - torch.cuda.get_device_capability(device), - cutlass.__version__, - ) - - -def _common_launch_checks(name: str, device, tensors, groups: int): - """Support and safety gates shared by the three launchers.""" - if any(_is_fake(t) for t in tensors): - raise ValueError( - f"{name} cannot run on fake/meta tensors; call the corresponding " - "torchao::* op instead, whose register_fake handles tracing" - ) - if groups < 1: - raise ValueError(f"G must be at least 1, got {groups}") - if device.type != "cuda": - raise ValueError(f"{name} requires CUDA tensors, got device {device}") - major, _minor = torch.cuda.get_device_capability(device) - if major != 10: - raise NotImplementedError( - f"{name} requires an SM100-class GPU (compute capability 10.x), " - f"got {torch.cuda.get_device_capability(device)}" - ) - - -def _stream_for(device): - import cuda.bindings.driver as cuda - - return cuda.CUstream(int(torch.cuda.current_stream(device).cuda_stream)) - - -def _check_sf_pointer_alignment(name: str, buf: torch.Tensor): - # Scale buffers feed TMA descriptors (inputs) or vectorized stores; a - # contiguous view with a storage offset can be 2-byte aligned, and only - # the launcher promises the alignment. - if buf.data_ptr() % 32 != 0: - raise ValueError( - f"{name} must be 32-byte aligned, but its data pointer is " - f"{buf.data_ptr() % 32} bytes past an aligned address" - ) - - -_E4M3 = torch.float8_e4m3fn -_BF16 = torch.bfloat16 - - -def launch_grouped_gemm_swiglu_fwd( - x_q, x_sf, w13_t_q, w13_t_sf, offsets, z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf -): - """FC1 grouped GEMM + SwiGLU + dual MXFP8 quantization, one kernel launch. - - Inputs are prequantized: ``x_q`` E4M3 ``[R, D]`` row-major with blocked - ``x_sf``; ``w13_t_q`` E4M3 ``[G, D, 2F]`` stride ``(2F*D, 1, D)`` with - per-expert blocked ``w13_t_sf`` (the 2F axis is element-interleaved - gate/up). Destinations are caller-allocated: ``z_bf16 [R, F, 2]``, - ``h_row_q [R, F]`` row-major + ``h_row_sf``, ``h_col_q [R, F]`` - column-major + ``h_col_sf`` (whole-matrix blocked for logical - ``[F, R/32]``). Every destination byte is written, including the - inactive-tail zeros. - """ - if x_q.ndim != 2: - raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") - rows, model_dim = x_q.shape - if w13_t_q.ndim != 3 or w13_t_q.shape[1] != model_dim or w13_t_q.shape[2] % 2: - raise ValueError( - f"w13_t_q must be [G, D, 2F] with D == {model_dim} and even 2F, " - f"got shape {tuple(w13_t_q.shape)}" - ) - groups, _, two_hidden = w13_t_q.shape - hidden = two_hidden // 2 - device = x_q.device - tensors = ( - x_q, - x_sf, - w13_t_q, - w13_t_sf, - offsets, - z_bf16, - h_row_q, - h_row_sf, - h_col_q, - h_col_sf, - ) - _common_launch_checks("launch_grouped_gemm_swiglu_fwd", device, tensors, groups) - if rows == 0: - raise ValueError( - "R == 0 is handled by the op layer (empty destinations, no launch)" - ) - - validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) - validate_allocated_rows(rows) - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) - validate_grouped_operand( - x_q, - name="x_q", - shape=(rows, model_dim), - stride=(model_dim, 1), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - w13_t_q, - name="w13_t_q", - shape=(groups, model_dim, two_hidden), - stride=(model_dim * two_hidden, 1, model_dim), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - x_sf, - name="x_sf", - logical_rows=rows, - logical_cols=model_dim // _SF_VEC_SIZE, - device=device, - ) - validate_blocked_scales( - w13_t_sf, - name="w13_t_sf", - logical_rows=two_hidden, - logical_cols=model_dim // _SF_VEC_SIZE, - device=device, - groups=groups, - ) - _check_sf_pointer_alignment("x_sf", x_sf) - _check_sf_pointer_alignment("w13_t_sf", w13_t_sf) - validate_destination( - z_bf16, - name="z_bf16", - shape=(rows, hidden, 2), - stride=(two_hidden, 2, 1), - dtype=_BF16, - device=device, - ) - validate_destination( - h_row_q, - name="h_row_q", - shape=(rows, hidden), - stride=(hidden, 1), - dtype=_E4M3, - device=device, - ) - validate_destination( - h_col_q, - name="h_col_q", - shape=(rows, hidden), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - h_row_sf, - name="h_row_sf", - logical_rows=rows, - logical_cols=hidden // _SF_VEC_SIZE, - device=device, - ) - validate_blocked_scales( - h_col_sf, - name="h_col_sf", - logical_rows=hidden, - logical_cols=rows // _SF_VEC_SIZE, - device=device, - ) - _check_sf_pointer_alignment("h_row_sf", h_row_sf) - _check_sf_pointer_alignment("h_col_sf", h_col_sf) - # The epilogue computes flat element offsets in Int32. - if rows * two_hidden >= 2**31: - raise ValueError( - f"R * 2F = {rows * two_hidden} does not fit the epilogue's int32 " - "element indexing" - ) - if two_hidden // _CTA_N > 65535: - raise ValueError(f"2F = {two_hidden} exceeds the launch grid's Y limit") - - stream = _stream_for(device) - args = ( - from_dlpack(activation_gemm_view(x_q), assumed_align=16), - from_dlpack(weight_gemm_view(w13_t_q), assumed_align=16), - from_dlpack(x_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(w13_t_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(offsets, assumed_align=4), - ( - from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), - from_dlpack(h_row_q.view(-1), assumed_align=16), - from_dlpack(h_row_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(h_col_q.t().reshape(-1), assumed_align=16), - from_dlpack(h_col_sf.view(torch.uint8).view(-1), assumed_align=16), - ), - stream, - ) - key = _cache_key( - "swiglu_fwd", (rows, model_dim, hidden, groups), tensors[:5], device - ) - slot = _executor_slot(key) - if not slot: - slot.append( - cute.compile( - _launch_grouped_gemm, *args, _SWIGLU_FWD_CONFIG, _swiglu_fwd_epilogue - ) - ) - slot[0](*args) - - -def launch_grouped_gemm_dswiglu_bwd( - do_q, - do_sf, - w2_dgrad_q, - w2_dgrad_sf, - z_bf16, - offsets, - dz_row_q, - dz_row_sf, - dz_col_q, - dz_col_sf, -): - """FC2 dgrad grouped GEMM + dSwiGLU + dual MXFP8 quantization, one launch. - - ``do_q`` E4M3 ``[R, D]`` row-major + blocked ``do_sf``; ``w2_dgrad_q`` - E4M3 ``[G, D, F]`` stride ``(D*F, 1, D)`` + per-expert blocked - ``w2_dgrad_sf``; ``z_bf16`` is the exact ``[R, F, 2]`` tensor the forward - kernel wrote (tail rows are never read). Destinations: ``dz_row_q - [R, 2F]`` row-major + ``dz_row_sf``, ``dz_col_q [R, 2F]`` column-major + - ``dz_col_sf`` (whole-matrix blocked for logical ``[2F, R/32]``), gate/up - gradients element-interleaved. - """ - if do_q.ndim != 2: - raise ValueError(f"do_q must be 2D [R, D], got shape {tuple(do_q.shape)}") - rows, model_dim = do_q.shape - if w2_dgrad_q.ndim != 3 or w2_dgrad_q.shape[1] != model_dim: - raise ValueError( - f"w2_dgrad_q must be [G, D, F] with D == {model_dim}, got shape " - f"{tuple(w2_dgrad_q.shape)}" - ) - groups, _, hidden = w2_dgrad_q.shape - two_hidden = 2 * hidden - device = do_q.device - tensors = ( - do_q, - do_sf, - w2_dgrad_q, - w2_dgrad_sf, - z_bf16, - offsets, - dz_row_q, - dz_row_sf, - dz_col_q, - dz_col_sf, - ) - _common_launch_checks("launch_grouped_gemm_dswiglu_bwd", device, tensors, groups) - if rows == 0: - raise ValueError( - "R == 0 is handled by the op layer (empty destinations, no launch)" - ) - - validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) - validate_allocated_rows(rows) - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) - validate_grouped_operand( - do_q, - name="do_q", - shape=(rows, model_dim), - stride=(model_dim, 1), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - w2_dgrad_q, - name="w2_dgrad_q", - shape=(groups, model_dim, hidden), - stride=(model_dim * hidden, 1, model_dim), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - z_bf16, - name="z_bf16", - shape=(rows, hidden, 2), - stride=(two_hidden, 2, 1), - dtype=_BF16, - device=device, - ) - validate_blocked_scales( - do_sf, - name="do_sf", - logical_rows=rows, - logical_cols=model_dim // _SF_VEC_SIZE, - device=device, - ) - validate_blocked_scales( - w2_dgrad_sf, - name="w2_dgrad_sf", - logical_rows=hidden, - logical_cols=model_dim // _SF_VEC_SIZE, - device=device, - groups=groups, - ) - _check_sf_pointer_alignment("do_sf", do_sf) - _check_sf_pointer_alignment("w2_dgrad_sf", w2_dgrad_sf) - validate_destination( - dz_row_q, - name="dz_row_q", - shape=(rows, two_hidden), - stride=(two_hidden, 1), - dtype=_E4M3, - device=device, - ) - validate_destination( - dz_col_q, - name="dz_col_q", - shape=(rows, two_hidden), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - dz_row_sf, - name="dz_row_sf", - logical_rows=rows, - logical_cols=two_hidden // _SF_VEC_SIZE, - device=device, - ) - validate_blocked_scales( - dz_col_sf, - name="dz_col_sf", - logical_rows=two_hidden, - logical_cols=rows // _SF_VEC_SIZE, - device=device, - ) - _check_sf_pointer_alignment("dz_row_sf", dz_row_sf) - _check_sf_pointer_alignment("dz_col_sf", dz_col_sf) - if rows * two_hidden >= 2**31: - raise ValueError( - f"R * 2F = {rows * two_hidden} does not fit the epilogue's int32 " - "element indexing" - ) - if hidden // _CTA_N > 65535: - raise ValueError(f"F = {hidden} exceeds the launch grid's Y limit") - - stream = _stream_for(device) - args = ( - from_dlpack(activation_gemm_view(do_q), assumed_align=16), - from_dlpack(weight_gemm_view(w2_dgrad_q), assumed_align=16), - from_dlpack(do_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(w2_dgrad_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(offsets, assumed_align=4), - ( - from_dlpack(z_bf16.view(rows, two_hidden).view(-1), assumed_align=16), - from_dlpack(dz_row_q.view(-1), assumed_align=16), - from_dlpack(dz_row_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(dz_col_q.t().reshape(-1), assumed_align=16), - from_dlpack(dz_col_sf.view(torch.uint8).view(-1), assumed_align=16), - ), - stream, - ) - key = _cache_key( - "dswiglu_bwd", (rows, model_dim, hidden, groups), tensors[:6], device - ) - slot = _executor_slot(key) - if not slot: - slot.append( - cute.compile( - _launch_grouped_gemm, *args, _DSWIGLU_BWD_CONFIG, _dswiglu_bwd_epilogue - ) - ) - slot[0](*args) - - -def launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw): - """Grouped MXFP8 wgrad into a caller-allocated BF16 ``[G, N, K]``, one launch. - - Inputs are the columnwise-quantized outputs of the forward/backward - kernels: ``dy_col_q`` E4M3 logical ``[R, N]`` stride ``(1, R)`` with - ``dy_col_sf`` whole-matrix blocked for logical ``[N, R/32]``, and - ``x_col_q`` / ``x_col_sf`` likewise for ``[R, K]``. Every element of - ``dw`` is written, including the all-zero slice of a zero-token expert. - """ - if dy_col_q.ndim != 2 or x_col_q.ndim != 2: - raise ValueError( - "dy_col_q and x_col_q must be 2D logical [R, N] and [R, K], got " - f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" - ) - rows, out_features = dy_col_q.shape - x_rows, in_features = x_col_q.shape - if x_rows != rows: - raise ValueError( - f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" - ) - groups = offsets.numel() - device = dy_col_q.device - tensors = (dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw) - _common_launch_checks("launch_grouped_gemm_wgrad", device, tensors, groups) - - validate_allocated_rows(rows) - for name, value in (("dy_col_q's N", out_features), ("x_col_q's K", in_features)): - if value <= 0 or value % 128 != 0: - raise ValueError(f"{name} must be a positive multiple of 128, got {value}") - validate_group_offsets(offsets, num_groups=groups, allocated_rows=rows) - validate_grouped_operand( - dy_col_q, - name="dy_col_q", - shape=(rows, out_features), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - x_col_q, - name="x_col_q", - shape=(rows, in_features), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - dy_col_sf, - name="dy_col_sf", - logical_rows=out_features, - logical_cols=rows // _SF_VEC_SIZE, - device=device, - ) - validate_blocked_scales( - x_col_sf, - name="x_col_sf", - logical_rows=in_features, - logical_cols=rows // _SF_VEC_SIZE, - device=device, - ) - _check_sf_pointer_alignment("dy_col_sf", dy_col_sf) - _check_sf_pointer_alignment("x_col_sf", x_col_sf) - validate_destination( - dw, - name="dw_bf16", - shape=(groups, out_features, in_features), - stride=(out_features * in_features, in_features, 1), - dtype=_BF16, - device=device, - ) - if groups * out_features * in_features >= 2**31: - raise ValueError( - f"dw_bf16 has {groups * out_features * in_features} elements, which " - "does not fit the epilogue's int32 element index" - ) - if in_features // _CTA_N > 65535: - raise ValueError(f"K = {in_features} exceeds the launch grid's Y limit") - if groups > 65535: - raise ValueError(f"G = {groups} exceeds the launch grid's Z limit") - - if rows == 0: - # Every expert has zero rows: every slice is the zero matrix. The - # destination is NOT empty here, and the contraction is. - dw.zero_() - return - - stream = _stream_for(device) - args = ( - # The free transpose: logical [R, N] stride (1, R) IS a K-contiguous - # [N, R], so the ragged axis becomes the contraction. - from_dlpack(activation_gemm_view(dy_col_q.t()), assumed_align=16), - from_dlpack(activation_gemm_view(x_col_q.t()), assumed_align=16), - from_dlpack(dy_col_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(x_col_sf.view(torch.uint8).view(-1), assumed_align=16), - from_dlpack(offsets, assumed_align=4), - (from_dlpack(dw.permute(1, 2, 0), assumed_align=16),), - stream, - ) - key = _cache_key( - "wgrad", (rows, out_features, in_features, groups), tensors[:5], device - ) - slot = _executor_slot(key) - if not slot: - slot.append( - cute.compile(_launch_grouped_gemm, *args, _WGRAD_CONFIG, _wgrad_epilogue) - ) - slot[0](*args) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py deleted file mode 100644 index 744e32f16e..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py +++ /dev/null @@ -1,552 +0,0 @@ -# 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. - -"""Custom-op surface for the MXFP8 routed-expert grouped-MLP kernels. - -Three ops, one per fused kernel: - -* ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` -- FC1 grouped GEMM + SwiGLU + - rowwise/columnwise MXFP8 quantization -* ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` -- FC2 dgrad grouped GEMM + - dSwiGLU + rowwise/columnwise MXFP8 quantization -* ``torchao::mxfp8_grouped_gemm_wgrad`` -- grouped MXFP8 weight-gradient - GEMM, invoked once for FC1 and once for FC2 - -Each op allocates its destinations through the ``_allocate_*_outputs`` helpers -below, which are pure torch and are shared by the real implementation and by -``register_fake``. Meta shapes and strides therefore cannot drift from eager -- -a class of bug that matters here because several outputs are column-major and a -row-major fake would silently change what ``torch.compile`` traces. - -Every op validates its inputs through a shared ``_validate_*_inputs`` helper -that also backs ``register_fake``, so ``torch.compile`` rejects an unsupported -call at graph capture rather than mid-training from a compiled region. The -checks are metadata-only (no host/device sync). The per-expert offset VALUES -are a documented caller invariant -- see ``grouped_mlp_validation`` for what -is and is not enforced. - -The user-facing functional wrappers live in -``torchao.prototype.moe_training.mxfp8_grouped_mlp``; importing that module (or -this one) registers the ops. -""" - -from typing import Tuple - -import torch - -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( - GROUP_ALIGNMENT, - SCALE_BLOCK_SIZE, - _is_fake, - blocked_scale_numel, - validate_allocated_rows, - validate_blocked_scales, - validate_destination, - validate_feature_dims, - validate_group_offsets, - validate_grouped_operand, -) - -__all__ = [ - # Validation surface re-exported for the kernel launchers, which check the - # caller-allocated destinations that the ops (allocating their own) do not. - "GROUP_ALIGNMENT", - "SCALE_BLOCK_SIZE", - "_is_fake", - "blocked_scale_numel", - "validate_allocated_rows", - "validate_blocked_scales", - "validate_destination", - "validate_feature_dims", - "validate_group_offsets", - "validate_grouped_operand", -] - -_E4M3 = torch.float8_e4m3fn -_E8M0 = torch.float8_e8m0fnu -_SCALE_BLOCK = SCALE_BLOCK_SIZE - - -def _require_cuda_device(device: torch.device, name: str) -> None: - """All operands must live on one CUDA device; CPU tensors get a clean error - here instead of a launcher failure.""" - if device.type != "cuda": - raise ValueError( - f"{name} must be a CUDA tensor, got device {device}; these kernels " - "run only on CUDA SM100 devices" - ) - - -def _empty_blocked_scales( - logical_rows: int, logical_cols: int, *, device, groups: int = 1 -) -> torch.Tensor: - """Allocate a flat blocked E8M0 scale buffer. - - The buffer is flat by ABI: its logical shape is metadata. Kernels write every - byte, including the inactive-tail rows, so an uninitialized allocation is - safe here -- but only because that write obligation is part of the contract. - """ - numel = groups * blocked_scale_numel(logical_rows, logical_cols) - shape = (groups, numel // groups) if groups > 1 else (numel,) - return torch.empty(shape, dtype=_E8M0, device=device) - - -# -------------------------------------------------------------------------- -# Kernel A: FC1 grouped GEMM + SwiGLU + dual quantization -# -------------------------------------------------------------------------- - - -def _allocate_swiglu_fwd_outputs( - rows: int, hidden: int, device -) -> Tuple[torch.Tensor, ...]: - z = torch.empty_strided( - (rows, hidden, 2), (2 * hidden, 2, 1), dtype=torch.bfloat16, device=device - ) - h_row_q = torch.empty_strided( - (rows, hidden), (hidden, 1), dtype=_E4M3, device=device - ) - h_row_sf = _empty_blocked_scales(rows, hidden // _SCALE_BLOCK, device=device) - # Column-major: the 32x1 quantized operand is consumed as its own transpose. - h_col_q = torch.empty_strided((rows, hidden), (1, rows), dtype=_E4M3, device=device) - h_col_sf = _empty_blocked_scales(hidden, rows // _SCALE_BLOCK, device=device) - return z, h_row_q, h_row_sf, h_col_q, h_col_sf - - -def _validate_swiglu_fwd_inputs( - x_q, x_sf, w13_t_q, w13_t_sf, offsets -) -> Tuple[int, int, int, int]: - if x_q.ndim != 2: - raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") - if w13_t_q.ndim != 3: - raise ValueError( - f"w13_t_q must be 3D [G, D, 2F], got shape {tuple(w13_t_q.shape)}" - ) - rows, model_dim = x_q.shape - groups, w_k, two_hidden = w13_t_q.shape - if w_k != model_dim: - raise ValueError( - f"w13_t_q contraction dim {w_k} must match x_q's D {model_dim}" - ) - if two_hidden % 2 != 0: - raise ValueError( - f"w13_t_q's N dim must be 2F with interleaved gate/up channels, got {two_hidden}" - ) - hidden = two_hidden // 2 - device = x_q.device - - _require_cuda_device(device, "x_q") - validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) - validate_allocated_rows(rows) - # The epilogue computes z/h element offsets as row * 2F + col in int32. - # Checked here as well as in the launcher so torch.compile tracing rejects - # the shape at graph capture instead of mid-training. - if rows * two_hidden >= 2**31: - raise ValueError( - f"R * 2F = {rows * two_hidden} does not fit the epilogue's int32 " - "element index; the store address arithmetic would wrap" - ) - validate_group_offsets( - offsets, num_groups=groups, allocated_rows=rows, device=device - ) - validate_grouped_operand( - x_q, - name="x_q", - shape=(rows, model_dim), - stride=(model_dim, 1), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - w13_t_q, - name="w13_t_q", - shape=(groups, model_dim, two_hidden), - stride=(model_dim * two_hidden, 1, model_dim), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - x_sf, - name="x_sf", - logical_rows=rows, - logical_cols=model_dim // _SCALE_BLOCK, - device=device, - ) - validate_blocked_scales( - w13_t_sf, - name="w13_t_sf", - logical_rows=two_hidden, - logical_cols=model_dim // _SCALE_BLOCK, - device=device, - groups=groups, - ) - return rows, model_dim, hidden, groups - - -@torch.library.custom_op("torchao::mxfp8_grouped_gemm_swiglu_fwd", mutates_args=()) -def _mxfp8_grouped_gemm_swiglu_fwd( - x_q: torch.Tensor, - x_sf: torch.Tensor, - w13_t_q: torch.Tensor, - w13_t_sf: torch.Tensor, - offsets: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """FC1 grouped GEMM + SwiGLU + dual MXFP8 RCEIL quantization, one launch. - - Inputs (all CUDA, same device; prequantized outside, never requantized here): - x_q E4M3 ``[R, D]`` stride ``(D, 1)``, rowwise 1x32 quantized. - x_sf flat blocked E8M0 scales for logical ``[R, D/32]`` - (``round_up(R,128) * round_up(D/32,4)`` bytes). - w13_t_q E4M3 ``[G, D, 2F]`` stride ``(2*D*F, 1, D)`` -- the quantized view - of ``w13_bf16.reshape(G, 2F, D).transpose(-2, -1)``; the 2F axis - is ELEMENT-interleaved gate/up (gate even, up odd). - w13_t_sf blocked E8M0 ``[G, round_up(2F,128) * round_up(D/32,4)]``. - offsets int32 CUDA ``[G]`` exclusive group end rows; every per-expert row - count must be a nonnegative multiple of 128 and - ``offsets[-1] <= R`` (caller invariant, see grouped_mlp_validation). - - Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``: - z_bf16 BF16 ``[R, F, 2]`` stride ``(2F, 2, 1)`` -- the pre-activation - rounded to BF16 BEFORE SwiGLU; gate at index 0, up at index 1; - saved for backward and consumed unchanged by the dswiglu op. - h_row_q E4M3 ``[R, F]`` stride ``(F, 1)``; h_row_sf blocked scales for - logical ``[R, F/32]``. - h_col_q E4M3 ``[R, F]`` COLUMN-MAJOR stride ``(1, R)``; h_col_sf - whole-matrix blocked scales for logical ``[F, R/32]``. - - ``h = silu(gate) * up`` is evaluated once from the BF16-rounded z and - rounded to BF16 before BOTH quantizers. Inactive tail rows - ``[offsets[-1], R)`` of every output are written as zero bytes. ``R == 0`` - returns empty outputs without launching. A zero-token expert contributes no - rows. G == 0 is rejected. - """ - rows, _model_dim, hidden, _groups = _validate_swiglu_fwd_inputs( - x_q, x_sf, w13_t_q, w13_t_sf, offsets - ) - outputs = _allocate_swiglu_fwd_outputs(rows, hidden, x_q.device) - - if rows == 0: - # R == 0 is a required correctness case. Every destination is empty in - # its row dimension and both scale buffers are zero-length, so there is - # nothing to write; launching would build a degenerate (0, D, 1) layout - # and fail inside the SF layout builder with an opaque MLIR error. - return outputs - - from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( - launch_grouped_gemm_swiglu_fwd, - ) - - launch_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_t_q, w13_t_sf, offsets, *outputs) - return outputs - - -@_mxfp8_grouped_gemm_swiglu_fwd.register_fake -def _(x_q, x_sf, w13_t_q, w13_t_sf, offsets): - rows, _model_dim, hidden, _groups = _validate_swiglu_fwd_inputs( - x_q, x_sf, w13_t_q, w13_t_sf, offsets - ) - return _allocate_swiglu_fwd_outputs(rows, hidden, x_q.device) - - -# -------------------------------------------------------------------------- -# Kernel B: FC2 dgrad grouped GEMM + dSwiGLU + dual quantization -# -------------------------------------------------------------------------- - - -def _allocate_dswiglu_bwd_outputs( - rows: int, hidden: int, device -) -> Tuple[torch.Tensor, ...]: - two_hidden = 2 * hidden - dz_row_q = torch.empty_strided( - (rows, two_hidden), (two_hidden, 1), dtype=_E4M3, device=device - ) - dz_row_sf = _empty_blocked_scales(rows, two_hidden // _SCALE_BLOCK, device=device) - dz_col_q = torch.empty_strided( - (rows, two_hidden), (1, rows), dtype=_E4M3, device=device - ) - dz_col_sf = _empty_blocked_scales(two_hidden, rows // _SCALE_BLOCK, device=device) - return dz_row_q, dz_row_sf, dz_col_q, dz_col_sf - - -def _validate_dswiglu_bwd_inputs(do_q, do_sf, w2_q, w2_sf, z_bf16, offsets): - if do_q.ndim != 2: - raise ValueError(f"do_q must be 2D [R, D], got shape {tuple(do_q.shape)}") - if w2_q.ndim != 3: - raise ValueError( - f"w2_dgrad_q must be 3D [G, D, F], got shape {tuple(w2_q.shape)}" - ) - if z_bf16.ndim != 3 or z_bf16.shape[-1] != 2: - raise ValueError( - f"z_bf16 must be [R, F, 2] with gate at index 0 and up at index 1, " - f"got shape {tuple(z_bf16.shape)}" - ) - rows, model_dim = do_q.shape - groups, w_k, hidden = w2_q.shape - if w_k != model_dim: - raise ValueError( - f"w2_dgrad_q contraction dim {w_k} must match do_q's D {model_dim}" - ) - if tuple(z_bf16.shape) != (rows, hidden, 2): - raise ValueError( - f"z_bf16 must be [{rows}, {hidden}, 2] to match do_q and w2_dgrad_q, " - f"got {tuple(z_bf16.shape)}" - ) - device = do_q.device - - _require_cuda_device(device, "do_q") - validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) - validate_allocated_rows(rows) - # Same int32 element-index bound as the forward: dz is [R, 2F]. - if rows * 2 * hidden >= 2**31: - raise ValueError( - f"R * 2F = {rows * 2 * hidden} does not fit the epilogue's int32 " - "element index; the store address arithmetic would wrap" - ) - validate_group_offsets( - offsets, num_groups=groups, allocated_rows=rows, device=device - ) - validate_grouped_operand( - do_q, - name="do_q", - shape=(rows, model_dim), - stride=(model_dim, 1), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - w2_q, - name="w2_dgrad_q", - shape=(groups, model_dim, hidden), - stride=(model_dim * hidden, 1, model_dim), - dtype=_E4M3, - device=device, - ) - # z_bf16 is the exact destination Kernel A wrote, so its stride is pinned too. - validate_grouped_operand( - z_bf16, - name="z_bf16", - shape=(rows, hidden, 2), - stride=(2 * hidden, 2, 1), - dtype=torch.bfloat16, - device=device, - ) - validate_blocked_scales( - do_sf, - name="do_sf", - logical_rows=rows, - logical_cols=model_dim // _SCALE_BLOCK, - device=device, - ) - validate_blocked_scales( - w2_sf, - name="w2_dgrad_sf", - logical_rows=hidden, - logical_cols=model_dim // _SCALE_BLOCK, - device=device, - groups=groups, - ) - return rows, model_dim, hidden, groups - - -@torch.library.custom_op("torchao::mxfp8_grouped_gemm_dswiglu_bwd", mutates_args=()) -def _mxfp8_grouped_gemm_dswiglu_bwd( - do_q: torch.Tensor, - do_sf: torch.Tensor, - w2_dgrad_q: torch.Tensor, - w2_dgrad_sf: torch.Tensor, - z_bf16: torch.Tensor, - offsets: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """FC2 dgrad grouped GEMM + dSwiGLU + dual MXFP8 RCEIL quantization, one launch. - - Inputs (all CUDA, same device; GEMM operands prequantized outside): - do_q E4M3 ``[R, D]`` stride ``(D, 1)``, rowwise 1x32 quantized - FC2 output-gradient. - do_sf flat blocked E8M0 scales for logical ``[R, D/32]``. - w2_dgrad_q E4M3 ``[G, D, F]`` stride ``(D*F, 1, D)`` -- the dgrad - orientation of w2. - w2_dgrad_sf blocked E8M0 ``[G, round_up(F,128) * round_up(D/32,4)]`` for - per-expert logical ``[F, D/32]``. - z_bf16 BF16 ``[R, F, 2]`` stride ``(2F, 2, 1)`` -- the EXACT saved - output of the swiglu_fwd op (not recomputed). Rows past - ``offsets[-1]`` are never read. - offsets int32 CUDA ``[G]`` exclusive group ends (same contract as the - forward op). - - Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)``: - dz_row_q E4M3 ``[R, 2F]`` stride ``(2F, 1)`` with ELEMENT-interleaved - ``[dgate_0, dup_0, ...]`` channels; dz_row_sf blocked scales for - logical ``[R, 2F/32]``. - dz_col_q E4M3 ``[R, 2F]`` COLUMN-MAJOR stride ``(1, R)``, same logical - order; dz_col_sf whole-matrix blocked scales for logical - ``[2F, R/32]``. - - ``dh`` is rounded to BF16 before dSwiGLU; ``dgate = dh * up * dsilu`` and - ``dup = dh * silu`` are each rounded to BF16 before interleaving, and both - quantizers consume the same BF16 ``dz``. Tail rows of every output are - written as zero bytes; ``R == 0`` returns empty outputs without launching; - G == 0 is rejected. - """ - rows, _model_dim, hidden, _groups = _validate_dswiglu_bwd_inputs( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets - ) - outputs = _allocate_dswiglu_bwd_outputs(rows, hidden, do_q.device) - - if rows == 0: - # See the R == 0 note in the forward op: nothing to write, and launching - # would build a degenerate SF layout. - return outputs - - from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( - launch_grouped_gemm_dswiglu_bwd, - ) - - launch_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets, *outputs - ) - return outputs - - -@_mxfp8_grouped_gemm_dswiglu_bwd.register_fake -def _(do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets): - rows, _model_dim, hidden, _groups = _validate_dswiglu_bwd_inputs( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets - ) - return _allocate_dswiglu_bwd_outputs(rows, hidden, do_q.device) - - -# -------------------------------------------------------------------------- -# Kernel C: grouped MXFP8 wgrad -# -------------------------------------------------------------------------- - - -def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): - if dy_col_q.ndim != 2 or x_col_q.ndim != 2: - raise ValueError( - "dy_col_q and x_col_q must both be 2D logical [R, N] / [R, K], got " - f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" - ) - rows, out_features = dy_col_q.shape - x_rows, in_features = x_col_q.shape - if x_rows != rows: - raise ValueError( - f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" - ) - groups = offsets.numel() if isinstance(offsets, torch.Tensor) else 0 - device = dy_col_q.device - - _require_cuda_device(device, "dy_col_q") - validate_allocated_rows(rows) - # N and K are the GEMM's free axes and are tiled with no tail path. The - # launcher checks this too, but it must ALSO live here because this - # function backs register_fake: without it, torch.compile traces a shape - # the real op rejects and the ValueError fires from inside a compiled - # region mid-training instead of at graph capture, defeating the caller's - # fallback predicate. - for name, value in (("dy_col_q's N", out_features), ("x_col_q's K", in_features)): - if value <= 0 or value % 128 != 0: - raise ValueError(f"{name} must be a positive multiple of 128, got {value}") - validate_group_offsets( - offsets, num_groups=groups, allocated_rows=rows, device=device - ) - # Both operands are column-major so their transposes are free. - validate_grouped_operand( - dy_col_q, - name="dy_col_q", - shape=(rows, out_features), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_grouped_operand( - x_col_q, - name="x_col_q", - shape=(rows, in_features), - stride=(1, rows), - dtype=_E4M3, - device=device, - ) - validate_blocked_scales( - dy_col_sf, - name="dy_col_sf", - logical_rows=out_features, - logical_cols=rows // _SCALE_BLOCK, - device=device, - ) - validate_blocked_scales( - x_col_sf, - name="x_col_sf", - logical_rows=in_features, - logical_cols=rows // _SCALE_BLOCK, - device=device, - ) - return rows, out_features, in_features, groups - - -@torch.library.custom_op("torchao::mxfp8_grouped_gemm_wgrad", mutates_args=()) -def _mxfp8_grouped_gemm_wgrad( - dy_col_q: torch.Tensor, - dy_col_sf: torch.Tensor, - x_col_q: torch.Tensor, - x_col_sf: torch.Tensor, - offsets: torch.Tensor, -) -> torch.Tensor: - """Grouped MXFP8 weight-gradient GEMM: ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. - - Inputs (all CUDA, same device): - dy_col_q E4M3 logical ``[R, N]`` COLUMN-MAJOR stride ``(1, R)`` -- the - columnwise (32x1) quantized output of the swiglu_fwd or - dswiglu_bwd op. - dy_col_sf whole-matrix blocked E8M0 scales for logical ``[N, R/32]``. - x_col_q E4M3 logical ``[R, K]`` stride ``(1, R)``; x_col_sf likewise - for logical ``[K, R/32]``. - offsets int32 CUDA ``[G]`` exclusive group ends over the shared row - (contraction) axis. - - Both scale buffers must be WHOLE-MATRIX ``to_blocked``, not torchao's - per-group K-groups rearrangement: the two orderings have identical byte - counts and differ whenever N > 128, so no length check can tell them apart, - and mixing them silently produces a block-permuted (wrong) ``dw``. - - Returns contiguous BF16 ``dw [G, N, K]`` with FP32 accumulation. Reusable - with no mode flag: FC1 wgrad passes ``N=2F, K=D``; FC2 wgrad passes ``N=D, - K=F``. Every element of ``dw`` is written on every call; a zero-token - expert (and the ``R == 0`` / all-empty cases) yields an all-zero slice. - G == 0 is rejected. - """ - rows, out_features, in_features, groups = _validate_wgrad_inputs( - dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets - ) - dw = torch.empty( - (groups, out_features, in_features), - dtype=torch.bfloat16, - device=dy_col_q.device, - ) - - if rows == 0: - # Unlike A and B, this destination is NOT empty at R == 0: every expert - # has zero rows, and an expert with zero rows is defined to produce an - # all-zero slice. Zero it here rather than launching over an empty - # contraction. - return dw.zero_() - - from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( - launch_grouped_gemm_wgrad, - ) - - launch_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets, dw) - return dw - - -@_mxfp8_grouped_gemm_wgrad.register_fake -def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): - _rows, out_features, in_features, groups = _validate_wgrad_inputs( - dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets - ) - return torch.empty( - (groups, out_features, in_features), - dtype=torch.bfloat16, - device=dy_col_q.device, - ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py deleted file mode 100644 index 40269c7275..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py +++ /dev/null @@ -1,267 +0,0 @@ -# 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. - -"""Host-side precondition validation for the MXFP8 routed-expert grouped-MLP kernels. - -The grouped MXFP8 kernels require every per-expert row count to be a multiple of -128. That is not a convenience: the tcgen05 blocked scale layout permutes in -128-row tiles, so a group boundary off a 128 multiple splits a tile and the -blocked buffer for the group no longer matches what the GEMM reads. - -Everything here is metadata-only so it costs no host/device synchronization and -stays traceable under torch.compile. The per-expert offset VALUES live in device -memory and are a documented caller invariant, not something these checks can -enforce: reading them requires a D2H sync, so they are validated only by the -opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS=1 debugging path below, and there is NO -device-side enforcement in a default build. Malformed offsets are not guaranteed -to fault, either -- the ragged-K weight-gradient kernel in particular can return -a wrong result from a clean-looking launch. Callers (and any future selection -predicate) must guarantee the offsets contract, e.g. by padding every expert -group to a multiple of 128 rows at dispatch time. - -Checks raise ValueError rather than asserting, so `python -O` cannot strip them. -""" - -import os -from typing import Optional - -import torch - -__all__ = [ - "SCALE_BLOCK_SIZE", - "SCALE_TILE_ROWS", - "SCALE_TILE_COLS", - "blocked_scale_numel", - "host_offsets_validation_enabled", - "validate_group_offsets", - "validate_grouped_operand", - "validate_blocked_scales", - "validate_destination", -] - -# MXFP8 scaling block: 32 values share one E8M0 scale. -SCALE_BLOCK_SIZE = 32 -# tcgen05 blocked scale tile: 128 rows x 4 columns, 512 bytes. -SCALE_TILE_ROWS = 128 -SCALE_TILE_COLS = 4 -# Row-count granularity every expert group must respect. -GROUP_ALIGNMENT = 128 -# Byte alignment the launchers promise for TMA/vectorized accesses. -_PTR_ALIGNMENT = 32 - - -def _round_up(x: int, to: int) -> int: - return ((x + to - 1) // to) * to - - -def blocked_scale_numel(rows: int, cols: int) -> int: - """Element count of the blocked E8M0 buffer for a logical [rows, cols] scale matrix. - - `cols` is a count of scale values, i.e. the reduced dimension divided by 32. - """ - return _round_up(rows, SCALE_TILE_ROWS) * _round_up(cols, SCALE_TILE_COLS) - - -def host_offsets_validation_enabled() -> bool: - """Opt-in host-side offset validation. Off by default: it forces a D2H sync.""" - return os.environ.get("TORCHAO_MXFP8_VALIDATE_OFFSETS", "0") == "1" - - -def validate_group_offsets( - offsets: torch.Tensor, - *, - num_groups: int, - allocated_rows: int, - device: Optional[torch.device] = None, - name: str = "offsets", -) -> None: - """Validate the exclusive-end group offsets tensor's metadata. - - Metadata is always checked, including that at least one expert group - exists. The offset *values* are checked only when - host_offsets_validation_enabled(), because reading them forces a D2H sync; - otherwise they are a documented caller invariant with no default-build - enforcement anywhere (see the module docstring). - """ - if not isinstance(offsets, torch.Tensor): - raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") - if num_groups < 1: - raise ValueError( - f"{name} must describe at least one expert group, got G={num_groups}" - ) - if offsets.dtype != torch.int32: - raise ValueError(f"{name} must be int32, got {offsets.dtype}") - if not offsets.is_cuda: - raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") - if device is not None and offsets.device != device: - raise ValueError( - f"{name} must be on {device}, got {offsets.device}; all operands and " - "destinations must share one CUDA device" - ) - if offsets.ndim != 1: - raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") - if offsets.numel() != num_groups: - raise ValueError( - f"{name} must have one entry per local expert: expected {num_groups}, " - f"got {offsets.numel()}" - ) - if not offsets.is_contiguous(): - raise ValueError(f"{name} must be contiguous, got stride {offsets.stride()}") - - if not host_offsets_validation_enabled(): - return - - values = offsets.tolist() # d2h sync; opt-in debugging path only - previous = 0 - for group, end in enumerate(values): - if end < previous: - raise ValueError( - f"{name} must be nondecreasing, but entry {group} is {end} after {previous}" - ) - size = end - previous - if size % GROUP_ALIGNMENT != 0: - raise ValueError( - f"per-expert row counts must be multiples of {GROUP_ALIGNMENT}: " - f"expert {group} has {size} rows (offsets {previous} -> {end})" - ) - previous = end - if previous > allocated_rows: - raise ValueError( - f"{name}[-1] ({previous}) exceeds the allocated row count ({allocated_rows})" - ) - - -def validate_grouped_operand( - tensor: torch.Tensor, - *, - name: str, - shape: tuple, - stride: tuple, - dtype: torch.dtype, - device: torch.device, - check_pointer_alignment: bool = True, -) -> None: - """Validate one quantized operand's dtype, shape, exact stride, device, alignment. - - Order matters: every metadata gate runs before the data_ptr() gate so that - FakeTensor tracing exercises the same checks (a fake tensor has no pointer). - """ - if tensor.dtype != dtype: - raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") - if tuple(tensor.shape) != tuple(shape): - raise ValueError( - f"{name} must have shape {tuple(shape)}, got {tuple(tensor.shape)}" - ) - if tuple(tensor.stride()) != tuple(stride): - raise ValueError( - f"{name} must have stride {tuple(stride)}, got {tuple(tensor.stride())}. " - "This layout is part of the ABI; a values-equal tensor with a different " - "stride is not interchangeable." - ) - if tensor.device != device: - raise ValueError( - f"{name} must be on {device}, got {tensor.device}; all operands and " - "destinations must share one CUDA device" - ) - if check_pointer_alignment and not _is_fake(tensor): - if tensor.data_ptr() % _PTR_ALIGNMENT != 0: - raise ValueError( - f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data pointer is " - f"{tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an aligned address. A " - "contiguous view with a nonzero storage offset can violate this." - ) - - -def validate_blocked_scales( - scales: torch.Tensor, - *, - name: str, - logical_rows: int, - logical_cols: int, - device: torch.device, - groups: int = 1, -) -> None: - """Validate a blocked E8M0 scale buffer's dtype, element count, and device. - - The buffer is carried flat: its logical shape is metadata, not its physical - shape, so only the element count is constrained. `groups` > 1 describes the - per-expert weight buffers, which are [G, per_group_numel]. - """ - if scales.dtype not in (torch.uint8, torch.float8_e8m0fnu): - raise ValueError( - f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), got {scales.dtype}" - ) - expected = groups * blocked_scale_numel(logical_rows, logical_cols) - if scales.numel() != expected: - raise ValueError( - f"{name} must hold {expected} blocked scale bytes for a logical " - f"[{logical_rows}, {logical_cols}] scale matrix" - + (f" across {groups} experts" if groups > 1 else "") - + f", got {scales.numel()}" - ) - if not scales.is_contiguous(): - raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") - if scales.device != device: - raise ValueError(f"{name} must be on {device}, got {scales.device}") - - -def validate_destination( - tensor: torch.Tensor, - *, - name: str, - shape: tuple, - stride: tuple, - dtype: torch.dtype, - device: torch.device, -) -> None: - """Validate a caller-allocated destination in the destination-passing entry points. - - Destinations are validated exactly like inputs. Skipping this is how a private - entry point turns a caller's shape mistake into an out-of-bounds write. - """ - validate_grouped_operand( - tensor, - name=name, - shape=shape, - stride=stride, - dtype=dtype, - device=device, - ) - - -def validate_feature_dims(*, model_dim: int, hidden_dim: int) -> None: - """D and F must both be multiples of 128 for the initial supported predicate.""" - if model_dim % GROUP_ALIGNMENT != 0: - raise ValueError( - f"model dimension D must be a multiple of {GROUP_ALIGNMENT}, got {model_dim}" - ) - if hidden_dim % GROUP_ALIGNMENT != 0: - raise ValueError( - f"routed-expert hidden dimension F must be a multiple of {GROUP_ALIGNMENT}, " - f"got {hidden_dim}" - ) - - -def validate_allocated_rows(rows: int, *, name: str = "R") -> None: - """The allocated row count must itself be 128-aligned. - - Group sizes are multiples of 128 and the active row count is their sum, so a - non-128 allocation can only describe an inactive tail that no legal offsets - vector can reach; rejecting it early keeps the tail contract simple. - """ - if rows % GROUP_ALIGNMENT != 0: - raise ValueError(f"{name} must be a multiple of {GROUP_ALIGNMENT}, got {rows}") - - -def _is_fake(tensor: torch.Tensor) -> bool: - """True for meta/fake tensors, which have no usable data pointer.""" - if tensor.device.type == "meta": - return True - try: - from torch._subclasses.fake_tensor import FakeTensor - except ImportError: - return False - return isinstance(tensor, FakeTensor) diff --git a/torchao/prototype/moe_training/mxfp8_grouped_mlp.py b/torchao/prototype/moe_training/mxfp8_grouped_mlp.py deleted file mode 100644 index b33cd7a4aa..0000000000 --- a/torchao/prototype/moe_training/mxfp8_grouped_mlp.py +++ /dev/null @@ -1,220 +0,0 @@ -# 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 MXFP8 grouped-MLP operations for the routed-expert training path. - -Three physically fused CuTe DSL kernels for Blackwell (SM 10.x), each exactly -one GPU kernel launch for a nonempty supported input: - -* :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM (MXFP8 - operands, FP32 accumulation) + BF16 pre-activation save + SwiGLU + rowwise - 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization of the activation. -* :func:`mxfp8_grouped_gemm_dswiglu_bwd` -- FC2 dgrad ragged grouped GEMM + - dSwiGLU (from the saved pre-activation) + dual MXFP8 quantization of the - FC1 gradient. -* :func:`mxfp8_grouped_gemm_wgrad` -- generic ragged-reduction grouped - weight gradient, called once for FC1 and once for FC2. - -This module owns the kernels and their operator wrappers only. Trainer -integration (converter selection, the autograd composite, saved-activation -ownership, expert padding configuration) is follow-up work in the consumer; -:func:`is_supported` is the shape predicate that integration should call -before selecting this operator family. - -Importing this module registers the three ``torchao::`` custom ops. -""" - -import importlib.util - -import torch - -# Importing the ops module registers the custom ops as a side effect. It is -# importable with no CuTe DSL installed; the DSL is imported lazily inside the -# op bodies at first real launch. -from torchao.prototype.moe_training.kernels.mxfp8 import ( - grouped_mlp_ops as _grouped_mlp_ops, # noqa: F401 -) -from torchao.utils import is_cuda_version_at_least - -__all__ = [ - "is_supported", - "mxfp8_grouped_gemm_dswiglu_bwd", - "mxfp8_grouped_gemm_swiglu_fwd", - "mxfp8_grouped_gemm_wgrad", -] - -# Every per-expert row count, the row allocation, and both feature dims must be -# multiples of this: the tcgen05 blocked scale layout permutes in 128-row tiles -# and the kernels tile all three GEMM axes at 128 with no tail path. -GROUP_ALIGNMENT = 128 - -# Runtime package detection, deliberately independent of the quantizer -# modules' availability flags: probing specs never imports the DSL. -_CUTEDSL_RUNTIME_PACKAGES = { - "cuda.bindings.driver": "cuda-python", - "cutlass": "nvidia-cutlass-dsl", - "cutlass.cute": "nvidia-cutlass-dsl", - "tvm_ffi": "apache-tvm-ffi", -} - - -def _missing_cutedsl_runtime_packages() -> list: - """Names of the pip packages required by the CuTe DSL runtime but absent.""" - missing = [] - for module_name, package_name in _CUTEDSL_RUNTIME_PACKAGES.items(): - try: - spec = importlib.util.find_spec(module_name) - except (ModuleNotFoundError, ValueError): - spec = None - if spec is None and package_name not in missing: - missing.append(package_name) - return missing - - -def _is_sm_10x() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 - - -_mxfp8_grouped_mlp_kernels_available = ( - _is_sm_10x() - and is_cuda_version_at_least(12, 8) - and not _missing_cutedsl_runtime_packages() -) - - -def _require_available() -> None: - """Raise a clean NotImplementedError when the kernels cannot run here.""" - if _mxfp8_grouped_mlp_kernels_available: - return - reasons = [] - if not torch.cuda.is_available(): - reasons.append("CUDA is not available") - elif not _is_sm_10x(): - reasons.append( - "requires an SM 10.x (Blackwell) GPU, found compute capability " - f"{torch.cuda.get_device_capability()}" - ) - if torch.cuda.is_available() and not is_cuda_version_at_least(12, 8): - reasons.append(f"requires CUDA >= 12.8, found {torch.version.cuda}") - missing = _missing_cutedsl_runtime_packages() - if missing: - reasons.append("missing required packages: " + ", ".join(missing)) - if not reasons: - reasons.append("kernels are disabled on this system") - raise NotImplementedError( - "MXFP8 grouped-MLP kernels are unavailable: " + "; ".join(reasons) - ) - - -def is_supported( - model_dim: int, hidden_dim: int, allocated_rows: int, num_groups: int -) -> bool: - """Pure shape predicate for selecting this operator family. - - True when the static shapes satisfy the kernel contract: at least one - expert group and D, F, R all positive multiples of 128. Integration code - (e.g. a quantization converter choosing between this fused path and the - unfused grouped-mm path) should call this BEFORE selecting the ops and - fall back when it is False -- and must also guarantee the runtime offsets - invariant, i.e. configure expert padding to 128 rows, because per-expert - row counts live in device memory and are not host-checkable here. - - Environment availability (CUDA, SM 10.x, the CuTe DSL runtime) is a - separate concern: combine with ``_mxfp8_grouped_mlp_kernels_available``. - """ - return ( - num_groups >= 1 - and model_dim > 0 - and hidden_dim > 0 - and allocated_rows > 0 - and model_dim % GROUP_ALIGNMENT == 0 - and hidden_dim % GROUP_ALIGNMENT == 0 - and allocated_rows % GROUP_ALIGNMENT == 0 - ) - - -def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_t_q, w13_t_sf, offsets): - """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. - - One kernel launch for any nonempty supported input. Arguments (all CUDA - tensors on one device, prequantized by the caller): - - * ``x_q``: E4M3 ``[R, D]``, stride ``(D, 1)`` -- rowwise 1x32-quantized - activations, expert-major packed rows. - * ``x_sf``: flat blocked E8M0 scales for logical ``[R, D/32]``. - * ``w13_t_q``: E4M3 ``[G, D, 2F]``, stride ``(2*D*F, 1, D)`` -- quantized - ``w13.reshape(G, 2F, D).transpose(-2, -1)``; the ``2F`` axis is - element-interleaved gate/up (gate at even indices). - * ``w13_t_sf``: blocked E8M0 ``[G, round_up(2F,128)*round_up(D/32,4)]``. - * ``offsets``: int32 ``[G]`` exclusive per-expert end rows; every expert's - row count must be a nonnegative multiple of 128 and - ``offsets[-1] <= R`` (documented caller invariant; see - ``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` for the opt-in synchronized check). - - Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``; ``z_bf16`` - ``[R, F, 2]`` (gate index 0, up index 1) is the BF16 pre-activation saved - for backward -- pass it unchanged to - :func:`mxfp8_grouped_gemm_dswiglu_bwd`. ``h_col_q`` is column-major with - whole-matrix blocked scales, ready for - :func:`mxfp8_grouped_gemm_wgrad`. Inactive tail rows of every output are - written as zeros. ``R == 0`` returns empty outputs; ``G == 0`` raises - ``ValueError``. - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( - x_q, x_sf, w13_t_q, w13_t_sf, offsets - ) - - -def mxfp8_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets -): - """FC2 dgrad grouped GEMM + dSwiGLU + rowwise/columnwise MXFP8 quantization. - - One kernel launch for any nonempty supported input. Arguments: - - * ``do_q`` / ``do_sf``: rowwise 1x32-quantized FC2 output-gradient, E4M3 - ``[R, D]`` stride ``(D, 1)`` with blocked scales for ``[R, D/32]``. - * ``w2_dgrad_q`` / ``w2_dgrad_sf``: E4M3 ``[G, D, F]`` stride - ``(D*F, 1, D)`` (dgrad orientation of w2) with per-expert blocked scales - for logical ``[F, D/32]``. - * ``z_bf16``: the exact ``[R, F, 2]`` pre-activation saved by - :func:`mxfp8_grouped_gemm_swiglu_fwd`; rows past ``offsets[-1]`` are - never read. - * ``offsets``: as in the forward op. - - Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)`` -- the FC1 gradient - ``[R, 2F]`` with gate/up gradients element-interleaved to match ``z_bf16``, - quantized both rowwise (row-major qdata) and columnwise (column-major - qdata, whole-matrix blocked scales, ready for the FC1 wgrad call). - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( - do_q, do_sf, w2_dgrad_q, w2_dgrad_sf, z_bf16, offsets - ) - - -def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): - """Grouped MXFP8 weight gradient: ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. - - One kernel launch per nonempty invocation, with the per-expert row ranges - of the shared ``R`` axis forming the ragged reduction. Both operands are - columnwise (32x1) quantized: E4M3 logical ``[R, N]`` / ``[R, K]`` with - column-major stride ``(1, R)`` and WHOLE-MATRIX blocked E8M0 scales for - logical ``[N, R/32]`` / ``[K, R/32]`` -- exactly what the forward and - backward ops emit. Do not feed torchao's per-group K-groups scale - rearrangement here: it has the same byte count but a different block - order, and produces a silently wrong ``dw``. - - Generic over both call sites with no mode flag: FC1 wgrad is ``N=2F, - K=D``; FC2 wgrad is ``N=D, K=F``. Returns contiguous BF16 ``[G, N, K]`` - (FP32 accumulation); zero-token experts yield all-zero slices, and - ``R == 0`` returns an all-zero result without launching. - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( - dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets - ) From c478ba551dafe659d6287a6ce53d52d95bd0fcc6 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 21:10:51 -0700 Subject: [PATCH 07/11] Rename the fused grouped-MLP surface off the cudnn prefix The four custom ops wrapping the cudnn-frontend CuTe DSL grouped-GEMM kernels take over the naming the removed in-repo kernel family vacated: - torchao::mxfp8_cudnn_grouped_mlp_fwd -> torchao::mxfp8_grouped_gemm_swiglu_fwd - torchao::mxfp8_cudnn_grouped_mm -> torchao::mxfp8_grouped_gemm - torchao::mxfp8_cudnn_grouped_mlp_bwd -> torchao::mxfp8_grouped_gemm_dswiglu_bwd - torchao::mxfp8_cudnn_grouped_mlp_wgrad -> torchao::mxfp8_grouped_gemm_wgrad Files follow: kernels/mxfp8/cudnn_grouped_mlp_{ops,validation}.py -> grouped_mlp_{ops,validation}.py; the public wrapper cudnn_grouped_mlp.py -> mxfp8_grouped_mlp.py (availability flag now _mxfp8_grouped_mlp_kernels_available); moe_training/__init__ re-exports the four ops plus is_supported. cudnn-frontend remains named in docstrings as the kernel provenance; identifiers no longer carry it. Also restyle the test module to match the other moe_training MXFP8 tests: module-level capability/availability skips (allow_module_level) instead of a custom availability marker, no per-import noqa, no __main__ block. All 29 tests semantically unchanged and green. Co-Authored-By: Claude Fable 5 --- ...ouped_mlp.py => test_mxfp8_grouped_mlp.py} | 161 ++++++++---------- torchao/prototype/moe_training/__init__.py | 14 ++ .../moe_training/kernels/mxfp8/__init__.py | 10 +- ..._grouped_mlp_ops.py => grouped_mlp_ops.py} | 36 ++-- ...alidation.py => grouped_mlp_validation.py} | 0 ...nn_grouped_mlp.py => mxfp8_grouped_mlp.py} | 50 +++--- 6 files changed, 134 insertions(+), 137 deletions(-) rename test/prototype/moe_training/{test_cudnn_grouped_mlp.py => test_mxfp8_grouped_mlp.py} (88%) rename torchao/prototype/moe_training/kernels/mxfp8/{cudnn_grouped_mlp_ops.py => grouped_mlp_ops.py} (96%) rename torchao/prototype/moe_training/kernels/mxfp8/{cudnn_grouped_mlp_validation.py => grouped_mlp_validation.py} (100%) rename torchao/prototype/moe_training/{cudnn_grouped_mlp.py => mxfp8_grouped_mlp.py} (80%) diff --git a/test/prototype/moe_training/test_cudnn_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py similarity index 88% rename from test/prototype/moe_training/test_cudnn_grouped_mlp.py rename to test/prototype/moe_training/test_mxfp8_grouped_mlp.py index 1d8b5ced6e..62108eefdb 100644 --- a/test/prototype/moe_training/test_cudnn_grouped_mlp.py +++ b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py @@ -4,7 +4,10 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. -"""Unit tests for the cuDNN-frontend MXFP8 grouped-MLP custom ops. +"""Unit tests for the MXFP8 fused grouped-MLP custom ops. + +The four ops wrap the cudnn-frontend package's CuTe DSL grouped-GEMM kernels +(``cudnn.grouped_gemm_{glu,quant,dglu,wgrad}_wrapper_sm100``). Every numerics gate is DERIVED at test time, never hard-coded: @@ -20,32 +23,46 @@ (wrong scale blocking built and decoded the same wrong way) cannot pass it. Measured band: ~30-40 dB at these shapes (pure MXFP8 requantization error). -Layout vocabulary (probe-derived, see agent_scratch/cudnn_fe_torchao in the -development environment): columnwise operands are accepted in ANY major -- -"rowmajor" here means un-transposed ``[R, N]`` row-major bytes (also the -layout the fwd/bwd ops emit for their columnwise outputs) and "native" means -the transposed-memory layout torchao's dim1 quantizers produce. Columnwise -scale buffers are PER-GROUP blocked (each expert's block ``to_blocked``-ed -independently, concatenated); whole-matrix blocking has the same byte count -and is silently wrong -- a dedicated negative control asserts the gap. +Layout vocabulary (probe-derived): columnwise operands are accepted in ANY +major -- "rowmajor" here means un-transposed ``[R, N]`` row-major bytes (also +the layout the fwd/bwd ops emit for their columnwise outputs) and "native" +means the transposed-memory layout torchao's dim1 quantizers produce. +Columnwise scale buffers are PER-GROUP blocked (each expert's block +``to_blocked``-ed independently, concatenated); whole-matrix blocking has the +same byte count and is silently wrong -- a dedicated negative control asserts +the gap. """ import pytest +import torch +import torch.nn.functional as F +from torch._subclasses.fake_tensor import FakeTensorMode -torch = pytest.importorskip("torch") +from torchao.utils import is_sm_version -import torch.nn.functional as F # noqa: E402 -from torch._subclasses.fake_tensor import FakeTensorMode # noqa: E402 +if not (torch.cuda.is_available() and is_sm_version(10, 0)): + pytest.skip( + "MXFP8 fused grouped MLP requires CUDA SM100", + allow_module_level=True, + ) -from torchao.float8.float8_utils import compute_error # noqa: E402 -from torchao.prototype.moe_training.cudnn_grouped_mlp import ( # noqa: E402 - _cudnn_grouped_mlp_available, - _cudnn_unavailable_reason, +from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( + _mxfp8_grouped_mlp_kernels_available, + _mxfp8_grouped_mlp_unavailable_reason, is_supported, ) -from torchao.prototype.mx_formats.config import ScaleCalculationMode # noqa: E402 -from torchao.prototype.mx_formats.mx_tensor import to_mx # noqa: E402 -from torchao.prototype.mx_formats.utils import from_blocked, to_blocked # noqa: E402 + +if not _mxfp8_grouped_mlp_kernels_available: + pytest.skip( + f"cudnn-frontend grouped-GEMM wrappers unavailable: " + f"{_mxfp8_grouped_mlp_unavailable_reason}", + allow_module_level=True, + ) + +from torchao.float8.float8_utils import compute_error +from torchao.prototype.mx_formats.config import ScaleCalculationMode +from torchao.prototype.mx_formats.mx_tensor import to_mx +from torchao.prototype.mx_formats.utils import from_blocked, to_blocked _E4M3 = torch.float8_e4m3fn _E8M0 = torch.float8_e8m0fnu @@ -54,10 +71,6 @@ _OPS = torch.ops.torchao -_gpu = pytest.mark.skipif( - not _cudnn_grouped_mlp_available, - reason=f"cuDNN-frontend grouped MLP unavailable: {_cudnn_unavailable_reason}", -) # --------------------------------------------------------------------------- # Pure-torch quantization / dequantization helpers. @@ -181,7 +194,7 @@ def _zsplit(z: torch.Tensor, hidden: int): def _to_32block(w13_elem: torch.Tensor) -> torch.Tensor: - """Element-interleaved [G, F, 2, D] -> cuDNN 32-block GLU order [G, 2F, D].""" + """Element-interleaved [G, F, 2, D] -> 32-block GLU order [G, 2F, D].""" G, hidden, _, D = w13_elem.shape return ( w13_elem.view(G, hidden // _BLOCK, _BLOCK, 2, D) @@ -282,10 +295,10 @@ def dbg(): def test_ops_registered(): for name in ( - "mxfp8_cudnn_grouped_mlp_fwd", - "mxfp8_cudnn_grouped_mm", - "mxfp8_cudnn_grouped_mlp_bwd", - "mxfp8_cudnn_grouped_mlp_wgrad", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_wgrad", ): assert hasattr(_OPS, name), f"torchao::{name} is not registered" @@ -307,20 +320,20 @@ def _fake_chain_shapes(R=512, D=256, hidden=256, G=2): w13_sf = torch.empty(G * N1 * D // _BLOCK, dtype=_E8M0, device=dev) offsets = torch.empty(G, dtype=torch.int32, device=dev) outs = {} - outs["fwd"] = _OPS.mxfp8_cudnn_grouped_mlp_fwd(x_q, x_sf, w13_q, w13_sf, offsets) + outs["fwd"] = _OPS.mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_q, w13_sf, offsets) w2_q = torch.empty(G, D, hidden, dtype=_E4M3, device=dev) w2_sf = torch.empty(G * D * hidden // _BLOCK, dtype=_E8M0, device=dev) - outs["mm"] = _OPS.mxfp8_cudnn_grouped_mm( + outs["mm"] = _OPS.mxfp8_grouped_gemm( outs["fwd"][1], outs["fwd"][2], w2_q, w2_sf, offsets ) dy_q = torch.empty(R, D, dtype=_E4M3, device=dev) dy_sf = torch.empty(R * D // _BLOCK, dtype=_E8M0, device=dev) w2c_q = torch.empty(G, D, hidden, dtype=_E4M3, device=dev) w2c_sf = torch.empty(G * D * hidden // _BLOCK, dtype=_E8M0, device=dev) - outs["bwd"] = _OPS.mxfp8_cudnn_grouped_mlp_bwd( + outs["bwd"] = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( dy_q, dy_sf, w2c_q, w2c_sf, outs["fwd"][0], offsets ) - outs["wgrad"] = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + outs["wgrad"] = _OPS.mxfp8_grouped_gemm_wgrad( torch.empty(R, D, dtype=_E4M3, device=dev), torch.empty(D * R // _BLOCK, dtype=_E8M0, device=dev), torch.empty(R, hidden, dtype=_E4M3, device=dev), @@ -332,7 +345,6 @@ def _fake_chain_shapes(R=512, D=256, hidden=256, G=2): return outs -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device object") def test_fake_contracts_match_specs(): """All four fakes produce the documented shapes/dtypes/contiguity.""" with FakeTensorMode(): @@ -363,20 +375,21 @@ def _run_chain(c): """fwd -> FC2 mm -> bwd -> FC1-dgrad mm -> wgrad x2, production layouts.""" r = {} r["z"], r["h_q"], r["h_sf"], r["h_colq"], r["h_col_sf"] = ( - _OPS.mxfp8_cudnn_grouped_mlp_fwd( + _OPS.mxfp8_grouped_gemm_swiglu_fwd( c["x_q"], c["x_sf"], c["w13_q"], c["w13_sf"], c["offsets"] ) ) - r["y"] = _OPS.mxfp8_cudnn_grouped_mm( + r["y"] = _OPS.mxfp8_grouped_gemm( r["h_q"], r["h_sf"], c["w2_q"], c["w2_sf"], c["offsets"] ) r["dz_q"], r["dz_sf"], r["dz_colq"], r["dz_col_sf"] = ( - _OPS.mxfp8_cudnn_grouped_mlp_bwd( + _OPS.mxfp8_grouped_gemm_dswiglu_bwd( c["dy_q"], c["dy_sf"], c["w2c_q"], c["w2c_sf"], r["z"], c["offsets"] ) ) - # FC1 dgrad: colwise weight cast enters op 2 TRANSPOSED into [G, N=D, K=2F]. - r["dx"] = _OPS.mxfp8_cudnn_grouped_mm( + # FC1 dgrad: colwise weight cast enters the mm op TRANSPOSED into + # [G, N=D, K=2F]. + r["dx"] = _OPS.mxfp8_grouped_gemm( r["dz_q"], r["dz_sf"], c["w13c_q"].transpose(-2, -1), @@ -384,16 +397,15 @@ def _run_chain(c): c["offsets"], ) # Production wgrad layout mixes: native dy x kernel h; kernel dz x native x. - r["dw2"] = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + r["dw2"] = _OPS.mxfp8_grouped_gemm_wgrad( c["dy_colq"], c["dy_col_sf"], r["h_colq"], r["h_col_sf"], c["offsets"] ) - r["dw13"] = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + r["dw13"] = _OPS.mxfp8_grouped_gemm_wgrad( r["dz_colq"], r["dz_col_sf"], c["x_colq"], c["x_col_sf"], c["offsets"] ) return r -@_gpu @pytest.mark.parametrize("case", list(_CASES)) def test_chain_numerics(case): D, hidden, sizes = _CASES[case] @@ -511,25 +523,12 @@ def test_chain_numerics(case): f"zero-token expert {g} weight gradients must be exactly zero" ) - print( - f"\n[{case}] derived gates/bands (dB): " - f"z refA_gate={gate_a:.1f} (got {z_db:.1f}), " - f"z refB band={band_b:.1f} (got {z_db_b:.1f}), " - f"h band={band_h:.1f} (row {h_db:.1f}, col {h_col_db:.1f}), " - f"y refA_gate={y_gate:.1f} (got {y_db:.1f}), " - f"y refB band={y_band_b:.1f} (got {y_db_b:.1f}), " - f"dz band={band_dz:.1f} (got {dz_db:.1f}), " - f"dx refA_gate={dx_gate:.1f} (got {dx_db:.1f}), " - f"dw2 {dw2_db:.1f}, dw13 {dw13_db:.1f}" - ) - # --------------------------------------------------------------------------- # Wgrad stride matrix: both operands in each major, all four combinations. # --------------------------------------------------------------------------- -@_gpu @pytest.mark.parametrize("a_native", [False, True], ids=["aRM", "aNat"]) @pytest.mark.parametrize("b_native", [False, True], ids=["bRM", "bNat"]) def test_wgrad_stride_matrix(dbg, a_native, b_native): @@ -537,7 +536,7 @@ def test_wgrad_stride_matrix(dbg, a_native, b_native): sizes, D = c["sizes"], c["D"] dy_q, dy_sf = _quant_colwise_grouped(c["dy"], sizes, native=a_native) x_q, x_sf = _quant_colwise_grouped(c["x"], sizes, native=b_native) - dw = _OPS.mxfp8_cudnn_grouped_mlp_wgrad(dy_q, dy_sf, x_q, x_sf, c["offsets"]) + dw = _OPS.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf, x_q, x_sf, c["offsets"]) dy_deq = _dequant_colwise_grouped(dy_q, dy_sf, sizes, D) x_deq = _dequant_colwise_grouped(x_q, x_sf, sizes, D) ref = torch.zeros(c["G"], D, D, dtype=torch.float32, device="cuda") @@ -554,7 +553,6 @@ def test_wgrad_stride_matrix(dbg, a_native, b_native): # --------------------------------------------------------------------------- -@_gpu def test_tail_a_lt_r_poisoned(): D = hidden = 256 sizes = [256, 0, 512, 256] @@ -578,17 +576,17 @@ def test_tail_a_lt_r_poisoned(): w2_q, w2_sf = _quant_weight_rowwise(w2) w2c_q, w2c_sf = _quant_weight_colwise(w2) - z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_grouped_gemm_swiglu_fwd( x_q, x_sf, w13_q, w13_sf, offsets ) assert not z[:A].isnan().any(), "active z rows contaminated by the poisoned tail" - y = _OPS.mxfp8_cudnn_grouped_mm(h_q, h_sf, w2_q, w2_sf, offsets) + y = _OPS.mxfp8_grouped_gemm(h_q, h_sf, w2_q, w2_sf, offsets) assert not y[:A].isnan().any(), "active y rows contaminated" - dz_q, dz_sf, dz_colq, dz_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_bwd( + dz_q, dz_sf, dz_colq, dz_col_sf = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( dy_q, dy_sf, w2c_q, w2c_sf, z, offsets ) w13c_q, w13c_sf = _quant_weight_colwise(w13) - dx = _OPS.mxfp8_cudnn_grouped_mm( + dx = _OPS.mxfp8_grouped_gemm( dz_q, dz_sf, w13c_q.transpose(-2, -1), w13c_sf, offsets ) assert not dx[:A].isnan().any(), "active dx rows contaminated" @@ -603,7 +601,7 @@ def test_tail_a_lt_r_poisoned(): ], 0, ) - dw2 = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + dw2 = _OPS.mxfp8_grouped_gemm_wgrad( dy_colq_full, dy_col_sf, h_colq, h_col_sf, offsets ) assert not dw2.isnan().any(), "wgrad read the NaN-poisoned inactive tail" @@ -623,7 +621,6 @@ def test_tail_a_lt_r_poisoned(): # --------------------------------------------------------------------------- -@_gpu def test_determinism_all_ops_bitwise(dbg): c = dbg r1 = _run_chain(c) @@ -634,15 +631,14 @@ def test_determinism_all_ops_bitwise(dbg): ) -@_gpu def test_compile_fullgraph_bitwise(dbg): c = dbg def fwd_then_mm(x_q, x_sf, w13_q, w13_sf, w2_q, w2_sf, offsets): - z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_grouped_gemm_swiglu_fwd( x_q, x_sf, w13_q, w13_sf, offsets ) - y = _OPS.mxfp8_cudnn_grouped_mm(h_q, h_sf, w2_q, w2_sf, offsets) + y = _OPS.mxfp8_grouped_gemm(h_q, h_sf, w2_q, w2_sf, offsets) return z, h_q, y eager = fwd_then_mm( @@ -667,12 +663,11 @@ def fwd_then_mm(x_q, x_sf, w13_q, w13_sf, w2_q, w2_sf, offsets): assert torch.equal(_bytes(e), _bytes(co)), f"compiled {name} != eager" -@_gpu def test_r0_all_ops(): dev = "cuda" D = hidden = 256 offsets = torch.zeros(2, dtype=torch.int32, device=dev) - z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_grouped_gemm_swiglu_fwd( torch.empty(0, D, dtype=_E4M3, device=dev), torch.empty(0, dtype=_E8M0, device=dev), torch.zeros(2, 2 * hidden, D, dtype=torch.uint8, device=dev).view(_E4M3), @@ -681,7 +676,7 @@ def test_r0_all_ops(): ) assert z.shape == (0, 2 * hidden) and h_q.shape == (0, hidden) assert h_sf.numel() == 0 and h_col_sf.numel() == 0 - dw = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( + dw = _OPS.mxfp8_grouped_gemm_wgrad( torch.empty(0, D, dtype=_E4M3, device=dev), torch.empty(0, dtype=_E8M0, device=dev), torch.empty(0, hidden, dtype=_E4M3, device=dev), @@ -696,7 +691,6 @@ def test_r0_all_ops(): # --------------------------------------------------------------------------- -@_gpu def test_negative_control_whole_matrix_colwise_scales(dbg): """Whole-matrix to_blocked colwise scales: same bytes, silently wrong order.""" c = dbg @@ -707,10 +701,8 @@ def test_negative_control_whole_matrix_colwise_scales(dbg): s_t, _ = to_mx(c["dy"].t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) dy_sf_wm = to_blocked(s_t.view(_E8M0)).view(_E8M0) assert dy_sf_wm.numel() == dy_sf_pg.numel() - good = _OPS.mxfp8_cudnn_grouped_mlp_wgrad( - dy_q, dy_sf_pg, x_q, x_sf_pg, c["offsets"] - ) - bad = _OPS.mxfp8_cudnn_grouped_mlp_wgrad(dy_q, dy_sf_wm, x_q, x_sf_pg, c["offsets"]) + good = _OPS.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf_pg, x_q, x_sf_pg, c["offsets"]) + bad = _OPS.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf_wm, x_q, x_sf_pg, c["offsets"]) dy_deq = _dequant_colwise_grouped(dy_q, dy_sf_pg, sizes, D) x_deq = _dequant_colwise_grouped(x_q, x_sf_pg, sizes, D) ref = torch.zeros(G, D, D, dtype=torch.float32, device="cuda") @@ -727,7 +719,6 @@ def test_negative_control_whole_matrix_colwise_scales(dbg): ) -@_gpu def test_negative_control_gate_up_swap(dbg): """Swapping the gate/up 32-blocks must collapse h against the correct ref.""" c = dbg @@ -740,7 +731,7 @@ def test_negative_control_gate_up_swap(dbg): .contiguous() ) w13_sw_q, w13_sw_sf = _quant_weight_rowwise(w13_sw) - _, h_q, h_sf, _, _ = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + _, h_q, h_sf, _, _ = _OPS.mxfp8_grouped_gemm_swiglu_fwd( c["x_q"], c["x_sf"], w13_sw_q, w13_sw_sf, c["offsets"] ) z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True) @@ -750,7 +741,7 @@ def test_negative_control_gate_up_swap(dbg): h_ref, _dequant_rowwise( *( - _OPS.mxfp8_cudnn_grouped_mlp_fwd( + _OPS.mxfp8_grouped_gemm_swiglu_fwd( c["x_q"], c["x_sf"], c["w13_q"], c["w13_sf"], c["offsets"] )[1:3] ) @@ -763,13 +754,12 @@ def test_negative_control_gate_up_swap(dbg): ) -@_gpu def test_negative_control_scale_byte_flip(dbg): """One +2-code E8M0 flip (x4) in the weight scales must break refA.""" c = dbg sf_bad = c["w13_sf"].view(torch.uint8).clone() sf_bad[sf_bad.numel() // 2] += 2 - z_bad = _OPS.mxfp8_cudnn_grouped_mlp_fwd( + z_bad = _OPS.mxfp8_grouped_gemm_swiglu_fwd( c["x_q"], c["x_sf"], c["w13_q"], sf_bad, c["offsets"] )[0] z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True) @@ -783,7 +773,6 @@ def test_negative_control_scale_byte_flip(dbg): ) -@_gpu def test_kernel_scale_mode_is_rceil(dbg): """The fwd op's h scale bytes must match RCEIL, and not FLOOR, quantization.""" c = dbg @@ -874,21 +863,19 @@ def _valid_fwd_args(device="cuda", R=512, D=256, hidden=256, G=2): ] -@_gpu @pytest.mark.parametrize("case", _NEGATIVES, ids=[c[0] for c in _NEGATIVES]) def test_validation_negatives(case): _name, mutate, needle = case args = _valid_fwd_args() mutate(args) with pytest.raises(ValueError) as exc_info: - _OPS.mxfp8_cudnn_grouped_mlp_fwd(**args) + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**args) assert needle.lower() in str(exc_info.value).lower(), ( f"rejection message {str(exc_info.value)!r} does not name the defect " f"({needle!r})" ) -@_gpu def test_optin_offsets_validation(monkeypatch): args = _valid_fwd_args() bad = dict(args) @@ -897,21 +884,17 @@ def test_optin_offsets_validation(monkeypatch): # Default build: metadata-only, misaligned VALUES are not (and cannot be) # caught without a D2H sync. monkeypatch.delenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", raising=False) - _OPS.mxfp8_cudnn_grouped_mlp_fwd(**bad) + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**bad) monkeypatch.setenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", "1") with pytest.raises(ValueError, match="FIX_PAD_SIZE"): - _OPS.mxfp8_cudnn_grouped_mlp_fwd(**bad) + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**bad) over = dict(args) over["offsets"] = torch.tensor([256, 768], dtype=torch.int32, device="cuda") with pytest.raises(ValueError, match="exceeds the allocated row count"): - _OPS.mxfp8_cudnn_grouped_mlp_fwd(**over) + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**over) # The opt-in check must not break fake tracing (no values to read). with FakeTensorMode(): _fake_chain_shapes() - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/torchao/prototype/moe_training/__init__.py b/torchao/prototype/moe_training/__init__.py index e58461a80c..7e99290de2 100644 --- a/torchao/prototype/moe_training/__init__.py +++ b/torchao/prototype/moe_training/__init__.py @@ -1,6 +1,15 @@ from torchao.prototype.moe_training.fp8_grouped_mm import ( _to_fp8_rowwise_then_scaled_grouped_mm, ) +from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( + is_supported as mxfp8_grouped_mlp_is_supported, +) +from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( + mxfp8_grouped_gemm, + mxfp8_grouped_gemm_dswiglu_bwd, + mxfp8_grouped_gemm_swiglu_fwd, + mxfp8_grouped_gemm_wgrad, +) from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _to_mxfp8_then_scaled_grouped_mm, ) @@ -8,4 +17,9 @@ __all__ = [ "_to_mxfp8_then_scaled_grouped_mm", "_to_fp8_rowwise_then_scaled_grouped_mm", + "mxfp8_grouped_gemm", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_wgrad", + "mxfp8_grouped_mlp_is_supported", ] diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index 350ed187f7..a83cb3fd29 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,9 +1,9 @@ -# Importing cudnn_grouped_mlp_ops registers the fused grouped-MLP custom ops -# (torchao::mxfp8_cudnn_grouped_{mlp_fwd,mm,mlp_bwd,mlp_wgrad}). The module is -# importable with no cudnn-frontend installed; `import cudnn` is deferred into -# the op bodies. +# Importing grouped_mlp_ops registers the fused grouped-MLP custom ops +# (torchao::mxfp8_grouped_gemm{_swiglu_fwd,,_dswiglu_bwd,_wgrad}). The module +# is importable with no cudnn-frontend installed; `import cudnn` is deferred +# into the op bodies. from torchao.prototype.moe_training.kernels.mxfp8 import ( - cudnn_grouped_mlp_ops, # noqa: F401 + grouped_mlp_ops, # noqa: F401 ) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py similarity index 96% rename from torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_ops.py rename to torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py index ceb8bc49d5..7f28b3ed1d 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_ops.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py @@ -10,15 +10,15 @@ kernel from the standalone cudnn-frontend python package (>= 1.27; no TransformerEngine involvement): -* ``torchao::mxfp8_cudnn_grouped_mlp_fwd`` -- FC1 ragged grouped GEMM + +* ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` -- FC1 ragged grouped GEMM + SwiGLU + rowwise AND columnwise MXFP8 RCEIL quantization + BF16 pre-GLU save (``grouped_gemm_glu_wrapper_sm100``). -* ``torchao::mxfp8_cudnn_grouped_mm`` -- ragged grouped GEMM on +* ``torchao::mxfp8_grouped_gemm`` -- ragged grouped GEMM on prequantized operands to BF16 (``grouped_gemm_quant_wrapper_sm100``); used for both FC2 forward and FC1 dgrad. -* ``torchao::mxfp8_cudnn_grouped_mlp_bwd`` -- FC2 dgrad + dSwiGLU + dual +* ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` -- FC2 dgrad + dSwiGLU + dual MXFP8 quantization of dz (``grouped_gemm_dglu_wrapper_sm100``). -* ``torchao::mxfp8_cudnn_grouped_mlp_wgrad`` -- ragged-reduction grouped +* ``torchao::mxfp8_grouped_gemm_wgrad`` -- ragged-reduction grouped weight gradient (``grouped_gemm_wgrad_wrapper_sm100``, dense output mode); called once for FC1 and once for FC2. @@ -38,7 +38,7 @@ and checks them against the same spec the fake allocates from). The user-facing wrappers live in -``torchao.prototype.moe_training.cudnn_grouped_mlp``; importing that module +``torchao.prototype.moe_training.mxfp8_grouped_mlp``; importing that module (or this one) registers the ops. ``import cudnn`` happens lazily inside op bodies at first real launch. """ @@ -47,7 +47,7 @@ import torch -from torchao.prototype.moe_training.kernels.mxfp8.cudnn_grouped_mlp_validation import ( +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( ROW_GROUP_ALIGNMENT, SCALE_BLOCK_SIZE, validate_allocated_rows, @@ -223,8 +223,8 @@ def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): return rows, model_dim, hidden, groups -@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mlp_fwd", mutates_args=()) -def _mxfp8_cudnn_grouped_mlp_fwd( +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_swiglu_fwd", mutates_args=()) +def _mxfp8_grouped_gemm_swiglu_fwd( x_q: torch.Tensor, x_sf: torch.Tensor, w13_q: torch.Tensor, @@ -296,7 +296,7 @@ def _mxfp8_cudnn_grouped_mlp_fwd( ) -@_mxfp8_cudnn_grouped_mlp_fwd.register_fake +@_mxfp8_grouped_gemm_swiglu_fwd.register_fake def _(x_q, x_sf, w13_q, w13_sf, offsets): rows, _model_dim, hidden, _groups = _validate_fwd_inputs( x_q, x_sf, w13_q, w13_sf, offsets @@ -368,8 +368,8 @@ def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): return rows, out_features, contraction, groups -@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mm", mutates_args=()) -def _mxfp8_cudnn_grouped_mm( +@torch.library.custom_op("torchao::mxfp8_grouped_gemm", mutates_args=()) +def _mxfp8_grouped_gemm( a_q: torch.Tensor, a_sf: torch.Tensor, b_q: torch.Tensor, @@ -424,7 +424,7 @@ def _mxfp8_cudnn_grouped_mm( return out -@_mxfp8_cudnn_grouped_mm.register_fake +@_mxfp8_grouped_gemm.register_fake def _(a_q, a_sf, b_q, b_sf, offsets): rows, out_features, _contraction, _groups = _validate_mm_inputs( a_q, a_sf, b_q, b_sf, offsets @@ -518,8 +518,8 @@ def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): return rows, model_dim, hidden, groups -@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mlp_bwd", mutates_args=()) -def _mxfp8_cudnn_grouped_mlp_bwd( +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_dswiglu_bwd", mutates_args=()) +def _mxfp8_grouped_gemm_dswiglu_bwd( dy_q: torch.Tensor, dy_sf: torch.Tensor, w2_col_q: torch.Tensor, @@ -585,7 +585,7 @@ def _mxfp8_cudnn_grouped_mlp_bwd( ) -@_mxfp8_cudnn_grouped_mlp_bwd.register_fake +@_mxfp8_grouped_gemm_dswiglu_bwd.register_fake def _(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): rows, _model_dim, hidden, _groups = _validate_bwd_inputs( dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets @@ -664,8 +664,8 @@ def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): return rows, out_features, in_features, groups -@torch.library.custom_op("torchao::mxfp8_cudnn_grouped_mlp_wgrad", mutates_args=()) -def _mxfp8_cudnn_grouped_mlp_wgrad( +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_wgrad", mutates_args=()) +def _mxfp8_grouped_gemm_wgrad( dy_col_q: torch.Tensor, dy_col_sf: torch.Tensor, x_col_q: torch.Tensor, @@ -720,7 +720,7 @@ def _mxfp8_cudnn_grouped_mlp_wgrad( return dw -@_mxfp8_cudnn_grouped_mlp_wgrad.register_fake +@_mxfp8_grouped_gemm_wgrad.register_fake def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): _rows, out_features, in_features, groups = _validate_wgrad_inputs( dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_validation.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py similarity index 100% rename from torchao/prototype/moe_training/kernels/mxfp8/cudnn_grouped_mlp_validation.py rename to torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py diff --git a/torchao/prototype/moe_training/cudnn_grouped_mlp.py b/torchao/prototype/moe_training/mxfp8_grouped_mlp.py similarity index 80% rename from torchao/prototype/moe_training/cudnn_grouped_mlp.py rename to torchao/prototype/moe_training/mxfp8_grouped_mlp.py index d6719a9883..dfe7010c37 100644 --- a/torchao/prototype/moe_training/cudnn_grouped_mlp.py +++ b/torchao/prototype/moe_training/mxfp8_grouped_mlp.py @@ -10,13 +10,13 @@ CuTe DSL kernel from the standalone cudnn-frontend python package (>= 1.27, Blackwell SM 10.x; no TransformerEngine dependency): -* :func:`mxfp8_cudnn_grouped_mlp_fwd` -- FC1 ragged grouped GEMM + SwiGLU + +* :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM + SwiGLU + rowwise 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization + BF16 pre-GLU. -* :func:`mxfp8_cudnn_grouped_mm` -- ragged grouped GEMM on prequantized +* :func:`mxfp8_grouped_gemm` -- ragged grouped GEMM on prequantized MXFP8 operands to BF16 (FC2 forward and FC1 dgrad). -* :func:`mxfp8_cudnn_grouped_mlp_bwd` -- FC2 dgrad + dSwiGLU + dual MXFP8 +* :func:`mxfp8_grouped_gemm_dswiglu_bwd` -- FC2 dgrad + dSwiGLU + dual MXFP8 quantization of the FC1 gradient. -* :func:`mxfp8_cudnn_grouped_mlp_wgrad` -- ragged-reduction grouped weight +* :func:`mxfp8_grouped_gemm_wgrad` -- ragged-reduction grouped weight gradient (FC1 and FC2). CONTRACT (stricter than the archived custom-kernel family): every per-expert @@ -25,7 +25,7 @@ aligned corrupt results SILENTLY and NONDETERMINISTICALLY (the corruption locus migrates between identical-input reruns; no smoke test can prove a misaligned config safe). Use a token dispatcher with ``pad_multiple=256`` and -see ``cudnn_grouped_mlp_validation`` for the two-tier enforcement +see ``grouped_mlp_validation`` for the two-tier enforcement (``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` for the opt-in synchronized check). The FC1 weight must be provided in the cuDNN 32-block GLU row order @@ -47,9 +47,9 @@ # importable with no cudnn-frontend installed; `import cudnn` is deferred into # the op bodies. from torchao.prototype.moe_training.kernels.mxfp8 import ( - cudnn_grouped_mlp_ops as _cudnn_grouped_mlp_ops, # noqa: F401 + grouped_mlp_ops as _grouped_mlp_ops, # noqa: F401 ) -from torchao.prototype.moe_training.kernels.mxfp8.cudnn_grouped_mlp_validation import ( +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( DIM_ALIGNMENT, ROW_GROUP_ALIGNMENT, ) @@ -58,10 +58,10 @@ "DIM_ALIGNMENT", "ROW_GROUP_ALIGNMENT", "is_supported", - "mxfp8_cudnn_grouped_mlp_bwd", - "mxfp8_cudnn_grouped_mlp_fwd", - "mxfp8_cudnn_grouped_mlp_wgrad", - "mxfp8_cudnn_grouped_mm", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_wgrad", + "mxfp8_grouped_gemm", ] _REQUIRED_WRAPPERS = ( @@ -117,7 +117,7 @@ def _probe_cudnn_frontend() -> str: return "" -_cudnn_unavailable_reason = ( +_mxfp8_grouped_mlp_unavailable_reason = ( _probe_cudnn_frontend() if _is_sm_10x() else ( @@ -126,15 +126,15 @@ def _probe_cudnn_frontend() -> str: else "CUDA is not available" ) ) -_cudnn_grouped_mlp_available = _cudnn_unavailable_reason == "" +_mxfp8_grouped_mlp_kernels_available = _mxfp8_grouped_mlp_unavailable_reason == "" def _require_available() -> None: """Raise a clean NotImplementedError when the kernels cannot run here.""" - if not _cudnn_grouped_mlp_available: + if not _mxfp8_grouped_mlp_kernels_available: raise NotImplementedError( "cuDNN-frontend MXFP8 grouped-MLP kernels are unavailable: " - + _cudnn_unavailable_reason + + _mxfp8_grouped_mlp_unavailable_reason ) @@ -147,7 +147,7 @@ def is_supported(model_dim: int, hidden_dim: int) -> bool: row counts live in device memory and are not checkable here. Environment availability (cudnn-frontend >= 1.27, SM 10.x) is a separate - concern: combine with ``_cudnn_grouped_mlp_available``. + concern: combine with ``_mxfp8_grouped_mlp_kernels_available``. """ return ( model_dim > 0 @@ -157,22 +157,22 @@ def is_supported(model_dim: int, hidden_dim: int) -> bool: ) -def mxfp8_cudnn_grouped_mlp_fwd(x_q, x_sf, w13_q, w13_sf, offsets): +def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_q, w13_sf, offsets): """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. - See ``torchao::mxfp8_cudnn_grouped_mlp_fwd`` for the full ABI. ``w13_q`` + See ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` for the full ABI. ``w13_q`` is E4M3 ``[G, 2F, D]`` contiguous with rows in 32-block GLU order; returns ``(z_bf16 [R, 2F], h_row_q [R, F], h_row_sf, h_col_q [R, F], h_col_sf)`` where the columnwise scales are PER-GROUP blocked. Rows past ``offsets[-1]`` of every output are garbage and read-forbidden. """ _require_available() - return torch.ops.torchao.mxfp8_cudnn_grouped_mlp_fwd( + return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( x_q, x_sf, w13_q, w13_sf, offsets ) -def mxfp8_cudnn_grouped_mm(a_q, a_sf, b_q, b_sf, offsets): +def mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets): """Ragged grouped GEMM on prequantized MXFP8 operands, BF16 output. ``b_q`` is ``[G, N, K]``-logical quantized along K with free strides @@ -182,10 +182,10 @@ def mxfp8_cudnn_grouped_mm(a_q, a_sf, b_q, b_sf, offsets): uninitialized. """ _require_available() - return torch.ops.torchao.mxfp8_cudnn_grouped_mm(a_q, a_sf, b_q, b_sf, offsets) + return torch.ops.torchao.mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets) -def mxfp8_cudnn_grouped_mlp_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): +def mxfp8_grouped_gemm_dswiglu_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): """FC2 dgrad + dSwiGLU + dual MXFP8 quantization of the FC1 gradient. ``z_bf16`` must be the exact fwd-op output. Returns @@ -193,12 +193,12 @@ def mxfp8_cudnn_grouped_mlp_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offset 32-block order. """ _require_available() - return torch.ops.torchao.mxfp8_cudnn_grouped_mlp_bwd( + return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets ) -def mxfp8_cudnn_grouped_mlp_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): +def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. Both operands columnwise (32x1) quantized with PER-GROUP blocked scales @@ -206,6 +206,6 @@ def mxfp8_cudnn_grouped_mlp_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offset block order). Returns contiguous BF16 ``[G, N, K]``. """ _require_available() - return torch.ops.torchao.mxfp8_cudnn_grouped_mlp_wgrad( + return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets ) From 683b370896f724e3897ca61dfa8391b655370f84 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 21:29:49 -0700 Subject: [PATCH 08/11] Memoize the grouped-MLP ops' metadata validation by signature The always-on validation tier is metadata-only, so its verdict is a pure function of the operands' metadata -- yet the training loop re-ran the full battery on every op call (hundreds per step with identical metadata, including tensors the previous op in the chain had just produced). Each validator now records a (shapes, strides, dtype, device, storage_offset) signature after PASSING and skips straight to the derived dims on repeats: 8.2 -> 3.4 us/call on the 16B fwd signature. First-call and torch.compile capture-time rejection behavior is unchanged (a rejected call never records its signature), the pointer-alignment gate stays covered because storage_offset is part of the signature, and the opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS values check still runs on every call while enabled (values are not metadata). Signature recording caps at 4096 entries; beyond that new signatures simply validate every time. 29/29 ao tests and 13/13 torchtitan composite tests green. Co-Authored-By: Claude Fable 5 --- .../kernels/mxfp8/grouped_mlp_ops.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py index 7f28b3ed1d..827b056a88 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py @@ -50,6 +50,7 @@ from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( ROW_GROUP_ALIGNMENT, SCALE_BLOCK_SIZE, + host_offsets_validation_enabled, validate_allocated_rows, validate_blocked_scales, validate_feature_dims, @@ -78,6 +79,32 @@ def _cached_ones(numel: int, dtype: torch.dtype, device: torch.device) -> torch. return out +# The always-on validation tier is metadata-only, so its verdict is a pure +# function of the operands' metadata (the pointer-alignment gate is covered +# by storage_offset: torch's CUDA caching allocator hands out aligned storage +# bases). A training step calls each op hundreds of times with identical +# metadata; the full battery runs once per distinct signature and repeats +# skip straight to the derived dims. Signatures are recorded only AFTER +# validation passes, so a rejected call never poisons the cache. The opt-in +# offsets-VALUES check (TORCHAO_MXFP8_VALIDATE_OFFSETS) reads data, not +# metadata, so it runs on every call while enabled. +_validated_sigs: set = set() +_VALIDATED_SIGS_CAP = 4096 + + +def _meta_sig(tag: str, *tensors: torch.Tensor) -> tuple: + # torch.Size and stride() are hashable tuples; device/dtype hash directly. + return (tag,) + tuple( + (t.shape, t.stride(), t.dtype, t.device, t.storage_offset()) + for t in tensors + ) + + +def _remember_sig(sig: tuple) -> None: + if len(_validated_sigs) < _VALIDATED_SIGS_CAP: + _validated_sigs.add(sig) + + def _require_cuda_device(device: torch.device, name: str) -> None: if device.type != "cuda": raise ValueError( @@ -161,6 +188,15 @@ def _allocate_from_specs(specs, device) -> Tuple[torch.Tensor, ...]: def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): + sig = _meta_sig("fwd", x_q, x_sf, w13_q, w13_sf, offsets) + if sig in _validated_sigs: + rows, model_dim = x_q.shape + groups, two_hidden, _ = w13_q.shape + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=x_q.device + ) + return rows, model_dim, two_hidden // 2, groups if x_q.ndim != 2: raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") if w13_q.ndim != 3: @@ -220,6 +256,7 @@ def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): device=device, groups=groups, ) + _remember_sig(sig) return rows, model_dim, hidden, groups @@ -310,6 +347,15 @@ def _(x_q, x_sf, w13_q, w13_sf, offsets): def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): + sig = _meta_sig("mm", a_q, a_sf, b_q, b_sf, offsets) + if sig in _validated_sigs: + rows, contraction = a_q.shape + groups, out_features, _ = b_q.shape + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=a_q.device + ) + return rows, out_features, contraction, groups if a_q.ndim != 2: raise ValueError(f"a_q must be 2D [R, K], got shape {tuple(a_q.shape)}") if b_q.ndim != 3: @@ -365,6 +411,7 @@ def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): device=device, groups=groups, ) + _remember_sig(sig) return rows, out_features, contraction, groups @@ -448,6 +495,15 @@ def _bwd_output_specs(rows: int, hidden: int): def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + sig = _meta_sig("bwd", dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets) + if sig in _validated_sigs: + rows, model_dim = dy_q.shape + groups, _, hidden = w2_col_q.shape + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=dy_q.device + ) + return rows, model_dim, hidden, groups if dy_q.ndim != 2: raise ValueError(f"dy_q must be 2D [R, D], got shape {tuple(dy_q.shape)}") if w2_col_q.ndim != 3: @@ -515,6 +571,7 @@ def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): device=device, groups=groups, ) + _remember_sig(sig) return rows, model_dim, hidden, groups @@ -599,6 +656,19 @@ def _(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + sig = _meta_sig("wgrad", dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets) + if sig in _validated_sigs: + rows, out_features = dy_col_q.shape + in_features = x_col_q.shape[1] + groups = offsets.numel() + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, + num_groups=groups, + allocated_rows=rows, + device=dy_col_q.device, + ) + return rows, out_features, in_features, groups if dy_col_q.ndim != 2 or x_col_q.ndim != 2: raise ValueError( "dy_col_q and x_col_q must both be 2D logical [R, N] / [R, K], got " @@ -661,6 +731,7 @@ def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): # by the ALLOCATED rows while a composite-produced operand's are sized by # the ROUTED total offsets[-1] -- mixing the two is legitimate and # probe-proven (tail case); the kernel reads only within offsets. + _remember_sig(sig) return rows, out_features, in_features, groups From 7bc1755535a2097b119ead10e096256041701d72 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 22:23:50 -0700 Subject: [PATCH 09/11] Fold the grouped-MLP validation and public wrappers into the ops module Mirror the gated-act (SwiGLU) + MXFP8 PR's single-module structure: the metadata-validation helpers (grouped_mlp_validation.py) and the public wrapper module (moe_training/mxfp8_grouped_mlp.py) fold into kernels/mxfp8/grouped_mlp_ops.py, which is now self-contained -- availability probe, is_supported, condensed validation helpers, the four custom ops with their fakes, and the availability-gated public wrappers. Both package __init__s return to their upstream state (importing grouped_mlp_ops registers the ops; consumers import it directly). All validation logic and error messages are unchanged; docstrings condensed to contracts. Net: 3 modules / 1,363 lines -> 1 module / 1,196 lines. 29/29 tests green (rejection-message and opt-in offsets gates included); ruff 0.11.6 lint+format clean. Co-Authored-By: Claude Fable 5 --- .../moe_training/test_mxfp8_grouped_mlp.py | 2 +- torchao/prototype/moe_training/__init__.py | 14 - .../moe_training/kernels/mxfp8/__init__.py | 7 - .../kernels/mxfp8/grouped_mlp_ops.py | 493 ++++++++++++++++-- .../kernels/mxfp8/grouped_mlp_validation.py | 310 ----------- .../moe_training/mxfp8_grouped_mlp.py | 211 -------- 6 files changed, 444 insertions(+), 593 deletions(-) delete mode 100644 torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py delete mode 100644 torchao/prototype/moe_training/mxfp8_grouped_mlp.py diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py index 62108eefdb..bf3e332555 100644 --- a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py +++ b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py @@ -46,7 +46,7 @@ allow_module_level=True, ) -from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( +from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_ops import ( _mxfp8_grouped_mlp_kernels_available, _mxfp8_grouped_mlp_unavailable_reason, is_supported, diff --git a/torchao/prototype/moe_training/__init__.py b/torchao/prototype/moe_training/__init__.py index 7e99290de2..e58461a80c 100644 --- a/torchao/prototype/moe_training/__init__.py +++ b/torchao/prototype/moe_training/__init__.py @@ -1,15 +1,6 @@ from torchao.prototype.moe_training.fp8_grouped_mm import ( _to_fp8_rowwise_then_scaled_grouped_mm, ) -from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( - is_supported as mxfp8_grouped_mlp_is_supported, -) -from torchao.prototype.moe_training.mxfp8_grouped_mlp import ( - mxfp8_grouped_gemm, - mxfp8_grouped_gemm_dswiglu_bwd, - mxfp8_grouped_gemm_swiglu_fwd, - mxfp8_grouped_gemm_wgrad, -) from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _to_mxfp8_then_scaled_grouped_mm, ) @@ -17,9 +8,4 @@ __all__ = [ "_to_mxfp8_then_scaled_grouped_mm", "_to_fp8_rowwise_then_scaled_grouped_mm", - "mxfp8_grouped_gemm", - "mxfp8_grouped_gemm_swiglu_fwd", - "mxfp8_grouped_gemm_dswiglu_bwd", - "mxfp8_grouped_gemm_wgrad", - "mxfp8_grouped_mlp_is_supported", ] diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index a83cb3fd29..0ed217b43a 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,10 +1,3 @@ -# Importing grouped_mlp_ops registers the fused grouped-MLP custom ops -# (torchao::mxfp8_grouped_gemm{_swiglu_fwd,,_dswiglu_bwd,_wgrad}). The module -# is importable with no cudnn-frontend installed; `import cudnn` is deferred -# into the op bodies. -from torchao.prototype.moe_training.kernels.mxfp8 import ( - grouped_mlp_ops, # noqa: F401 -) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, # noqa: F401 _mxfp8_flydsl_kernels_available, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py index 827b056a88..ffb8526a97 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py @@ -4,67 +4,402 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. -"""Custom-op surface for the cuDNN-frontend MXFP8 routed-expert grouped MLP. - -Four ops, each wrapping one ``cudnn.grouped_gemm_*_wrapper_sm100`` CuTe DSL -kernel from the standalone cudnn-frontend python package (>= 1.27; no -TransformerEngine involvement): - -* ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` -- FC1 ragged grouped GEMM + - SwiGLU + rowwise AND columnwise MXFP8 RCEIL quantization + BF16 pre-GLU save - (``grouped_gemm_glu_wrapper_sm100``). -* ``torchao::mxfp8_grouped_gemm`` -- ragged grouped GEMM on - prequantized operands to BF16 (``grouped_gemm_quant_wrapper_sm100``); used - for both FC2 forward and FC1 dgrad. -* ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` -- FC2 dgrad + dSwiGLU + dual - MXFP8 quantization of dz (``grouped_gemm_dglu_wrapper_sm100``). -* ``torchao::mxfp8_grouped_gemm_wgrad`` -- ragged-reduction grouped - weight gradient (``grouped_gemm_wgrad_wrapper_sm100``, dense output mode); - called once for FC1 and once for FC2. +"""MXFP8 routed-expert grouped-MLP ops over the cuDNN-frontend CuTe DSL kernels. + +Four custom ops, each one launch of a ``cudnn.grouped_gemm_*_wrapper_sm100`` +kernel from the standalone cudnn-frontend python package (>= 1.27, Blackwell +SM 10.x; no TransformerEngine dependency); the matching public wrappers live +at the bottom of this module: + +* :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM + SwiGLU + + rowwise 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization + BF16 pre-GLU. +* :func:`mxfp8_grouped_gemm` -- ragged grouped GEMM on + prequantized operands to BF16 (FC2 forward and FC1 dgrad). +* :func:`mxfp8_grouped_gemm_dswiglu_bwd` -- FC2 dgrad + dSwiGLU + dual MXFP8 + quantization of the FC1 gradient. +* :func:`mxfp8_grouped_gemm_wgrad` -- ragged-reduction grouped weight + gradient (dense output mode; called once for FC1 and once for FC2). + +CONTRACT: every per-expert row count and the allocated row count must be +multiples of **256** -- the cuDNN FE kernels hard-code ``FIX_PAD_SIZE = 256``, +and groups that are only 128-row aligned corrupt results SILENTLY and +NONDETERMINISTICALLY (the corruption locus migrates between identical-input +reruns; no smoke test can prove a misaligned config safe). Use a token +dispatcher with ``pad_multiple=256``. Enforcement is two-tier: metadata-only +checks always run (memoized per signature, FakeTensor-safe, back +``register_fake`` so torch.compile rejects at capture time); the offset +VALUES (nondecreasing, per-expert %256, ``offsets[-1] <= R``) are checked +only under ``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` because reading them forces a +D2H sync. Checks raise ValueError, never assert, so ``python -O`` cannot +strip them. All scale arguments are FLAT blocked E8M0 buffers (uint8 or float8_e8m0fnu); the ops build the kernel-native 6-D / 2-D views internally with probe-proven -recipes. ``offsets`` is int32 CUDA ``[G]`` exclusive-end rows; per-expert row -counts must be multiples of 256 (cuDNN FE ``FIX_PAD_SIZE``; see the validation -module for the two-tier enforcement and the misalignment hazard). Rows in -``[offsets[-1], R)``: caller-allocated outputs (the grouped-mm result and the -weight gradients) keep their tails untouched, while kernel-allocated outputs -(z, h, dz and their scales) carry garbage tails that are read-forbidden -- -both behaviors probe-verified with NaN-poisoned tails. - -Shared ``_validate_*`` helpers back ``register_fake`` so torch.compile rejects -unsupported calls at capture time, and shared output-spec helpers keep fake -and eager metadata identical (eager normalizes the wrapper's returned tensors -and checks them against the same spec the fake allocates from). - -The user-facing wrappers live in -``torchao.prototype.moe_training.mxfp8_grouped_mlp``; importing that module -(or this one) registers the ops. ``import cudnn`` happens lazily inside op -bodies at first real launch. +recipes. The FC1 weight is E4M3 ``[G, 2F, D]`` with rows in the cuDNN +32-block GLU order ``[gate0(32) | up0(32) | gate1(32) | ...]`` (gate = the +SiLU'd operand). ``offsets`` is int32 CUDA ``[G]`` exclusive-end rows. Rows +in ``[offsets[-1], R)``: caller-allocated outputs (the grouped-mm result and +the weight gradients) keep their tails untouched, while kernel-allocated +outputs (z, h, dz and their scales) carry garbage tails that are +read-forbidden -- both behaviors probe-verified with NaN-poisoned tails. + +Importing this module registers the four ``torchao::`` custom ops; the +``cudnn`` package itself is imported lazily inside the op bodies at first +real launch. :func:`is_supported` is the static shape predicate to call +before selecting this family. """ -from typing import Tuple +import importlib.util +import os +from typing import Optional, Tuple import torch -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( - ROW_GROUP_ALIGNMENT, - SCALE_BLOCK_SIZE, - host_offsets_validation_enabled, - validate_allocated_rows, - validate_blocked_scales, - validate_feature_dims, - validate_group_offsets, - validate_operand, - validate_ragged_colwise_scales, -) - -__all__ = ["ROW_GROUP_ALIGNMENT", "SCALE_BLOCK_SIZE"] +__all__ = [ + "DIM_ALIGNMENT", + "ROW_GROUP_ALIGNMENT", + "SCALE_BLOCK_SIZE", + "is_supported", + "mxfp8_grouped_gemm", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_wgrad", +] + +# MXFP8 scaling block: 32 values share one E8M0 scale. +SCALE_BLOCK_SIZE = 32 +# tcgen05 blocked scale tile: 128 rows x 4 columns, 512 bytes. +SCALE_TILE_ROWS = 128 +SCALE_TILE_COLS = 4 +# Feature-dimension granularity (D and F). +DIM_ALIGNMENT = 128 +# Row-count granularity: per-expert groups AND the allocated row count (the +# cuDNN FE kernels' FIX_PAD_SIZE). +ROW_GROUP_ALIGNMENT = 256 +# Byte alignment for TMA/vectorized accesses. +_PTR_ALIGNMENT = 16 + +_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) _E4M3 = torch.float8_e4m3fn _E8M0 = torch.float8_e8m0fnu _BLOCK = SCALE_BLOCK_SIZE + +# -------------------------------------------------------------------------- +# Availability probe and the static shape predicate. +# -------------------------------------------------------------------------- + +_REQUIRED_WRAPPERS = ( + "grouped_gemm_glu_wrapper_sm100", + "grouped_gemm_quant_wrapper_sm100", + "grouped_gemm_dglu_wrapper_sm100", + "grouped_gemm_wgrad_wrapper_sm100", +) +# 1.27 is required: earlier frontends reject prob_tensor=None. +_MIN_FE_VERSION = (1, 27) + + +def _fe_version_tuple(version: str) -> tuple: + """Numeric prefix as a tuple ('1.27.0' -> (1, 27, 0)); never compare + version STRINGS ('1.100' < '1.27' lexicographically).""" + parts = [] + for piece in version.split("."): + digits = "" + for ch in piece: + if not ch.isdigit(): + break + digits += ch + if not digits: + break + parts.append(int(digits)) + return tuple(parts) + + +def _is_sm_10x() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +def _probe_cudnn_frontend() -> str: + """Empty string when usable; else the reason it is not.""" + if importlib.util.find_spec("cudnn") is None: + return "the cudnn-frontend python package ('cudnn') is not installed" + try: + import cudnn + except Exception as exc: # pragma: no cover - environment-specific + return f"'import cudnn' failed: {exc!r}" + version = getattr(cudnn, "__version__", "0") + if _fe_version_tuple(version) < _MIN_FE_VERSION: + return ( + f"cudnn-frontend {version} is too old; >= " + f"{'.'.join(map(str, _MIN_FE_VERSION))} is required " + "(prob_tensor=None support)" + ) + missing = [name for name in _REQUIRED_WRAPPERS if not hasattr(cudnn, name)] + if missing: + return "cudnn-frontend lacks required wrappers: " + ", ".join(missing) + return "" + + +_mxfp8_grouped_mlp_unavailable_reason = ( + _probe_cudnn_frontend() + if _is_sm_10x() + else ( + "requires an SM 10.x (Blackwell) GPU" + if torch.cuda.is_available() + else "CUDA is not available" + ) +) +_mxfp8_grouped_mlp_kernels_available = _mxfp8_grouped_mlp_unavailable_reason == "" + + +def _require_available() -> None: + if not _mxfp8_grouped_mlp_kernels_available: + raise NotImplementedError( + "cuDNN-frontend MXFP8 grouped-MLP kernels are unavailable: " + + _mxfp8_grouped_mlp_unavailable_reason + ) + + +def is_supported(model_dim: int, hidden_dim: int) -> bool: + """True when D and F are positive multiples of 128. Integration code must + ALSO guarantee the runtime row contract (per-expert groups and the row + allocation padded to multiples of 256): row counts live in device memory + and are not checkable here. Environment availability is a separate + concern (``_mxfp8_grouped_mlp_kernels_available``).""" + return ( + model_dim > 0 + and hidden_dim > 0 + and model_dim % DIM_ALIGNMENT == 0 + and hidden_dim % DIM_ALIGNMENT == 0 + ) + + +# -------------------------------------------------------------------------- +# Metadata validation helpers (see the module docstring for the two tiers). +# -------------------------------------------------------------------------- + + +def _round_up(x: int, to: int) -> int: + return ((x + to - 1) // to) * to + + +def blocked_scale_numel(rows: int, cols: int) -> int: + """Blocked-buffer element count for a logical [rows, cols] scale matrix + (``cols`` counts scale values: the reduced dimension divided by 32).""" + return _round_up(rows, SCALE_TILE_ROWS) * _round_up(cols, SCALE_TILE_COLS) + + +def host_offsets_validation_enabled() -> bool: + """Opt-in offset-VALUES validation; off by default (forces a D2H sync).""" + return os.environ.get("TORCHAO_MXFP8_VALIDATE_OFFSETS", "0") == "1" + + +def _is_fake(tensor: torch.Tensor) -> bool: + """True for meta/fake tensors (no usable data pointer or values).""" + if tensor.device.type == "meta": + return True + try: + from torch._subclasses.fake_tensor import FakeTensor + except ImportError: + return False + return isinstance(tensor, FakeTensor) + + +def validate_group_offsets( + offsets: torch.Tensor, + *, + num_groups: int, + allocated_rows: int, + device: Optional[torch.device] = None, + name: str = "offsets", +) -> None: + """Metadata always; VALUES only when opted in and the tensor is real.""" + if not isinstance(offsets, torch.Tensor): + raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") + if num_groups < 1: + raise ValueError( + f"{name} must describe at least one expert group, got G={num_groups}" + ) + if offsets.dtype != torch.int32: + raise ValueError(f"{name} must be int32, got {offsets.dtype}") + if not offsets.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") + if device is not None and offsets.device != device: + raise ValueError( + f"{name} must be on {device}, got {offsets.device}; all operands and " + "destinations must share one CUDA device" + ) + if offsets.ndim != 1: + raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") + if offsets.numel() != num_groups: + raise ValueError( + f"{name} must have one entry per local expert: expected {num_groups}, " + f"got {offsets.numel()}" + ) + if not offsets.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {offsets.stride()}") + + if not host_offsets_validation_enabled() or _is_fake(offsets): + return + + values = offsets.tolist() # d2h sync; opt-in debugging path only + previous = 0 + for group, end in enumerate(values): + if end < previous: + raise ValueError( + f"{name} must be nondecreasing, but entry {group} is {end} " + f"after {previous}" + ) + size = end - previous + if size % ROW_GROUP_ALIGNMENT != 0: + raise ValueError( + f"per-expert row counts must be multiples of {ROW_GROUP_ALIGNMENT} " + f"(cuDNN FE FIX_PAD_SIZE; sub-256 groups corrupt results " + f"nondeterministically): expert {group} has {size} rows " + f"(offsets {previous} -> {end})" + ) + previous = end + if previous > allocated_rows: + raise ValueError( + f"{name}[-1] ({previous}) exceeds the allocated row count " + f"({allocated_rows})" + ) + + +def validate_operand( + tensor: torch.Tensor, + *, + name: str, + shape: tuple, + dtype: torch.dtype, + device: torch.device, + stride: Optional[tuple] = None, + check_pointer_alignment: bool = True, +) -> None: + """dtype/shape/device, optional EXACT stride (None = any: the wrappers + consume both majors, every composite combination probe-proven), pointer + alignment. Metadata gates run before the ``data_ptr()`` gate so + FakeTensor tracing exercises the same checks.""" + if tensor.dtype != dtype: + raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != tuple(shape): + raise ValueError( + f"{name} must have shape {tuple(shape)}, got {tuple(tensor.shape)}" + ) + if stride is not None and tuple(tensor.stride()) != tuple(stride): + raise ValueError( + f"{name} must have stride {tuple(stride)}, got {tuple(tensor.stride())}. " + "This layout is part of the ABI; a values-equal tensor with a " + "different stride is not interchangeable." + ) + if tensor.device != device: + raise ValueError( + f"{name} must be on {device}, got {tensor.device}; all operands and " + "destinations must share one CUDA device" + ) + if check_pointer_alignment and not _is_fake(tensor): + if tensor.data_ptr() % _PTR_ALIGNMENT != 0: + raise ValueError( + f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data " + f"pointer is {tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an " + "aligned address. A contiguous view with a nonzero storage " + "offset can violate this." + ) + + +def validate_blocked_scales( + scales: torch.Tensor, + *, + name: str, + logical_rows: int, + logical_cols: int, + device: torch.device, + groups: int = 1, +) -> None: + """Flat blocked E8M0 buffer with a statically known size; ``groups > 1`` + means per-expert blocks concatenated.""" + if scales.dtype not in _SCALE_DTYPES: + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " + f"got {scales.dtype}" + ) + expected = groups * blocked_scale_numel(logical_rows, logical_cols) + if scales.numel() != expected: + raise ValueError( + f"{name} must hold {expected} blocked scale bytes for a logical " + f"[{logical_rows}, {logical_cols}] scale matrix" + + (f" across {groups} experts" if groups > 1 else "") + + f", got {scales.numel()}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + + +def validate_ragged_colwise_scales( + scales: torch.Tensor, + *, + name: str, + features: int, + allocated_rows: int, + device: torch.device, +) -> None: + """Per-group columnwise scale buffer sized by ``offsets[-1]`` -- a device + value -- so only dtype/device/contiguity, granule divisibility, and the + allocated-rows maximum are host-checkable (an ``offsets[-1] < R`` buffer + legitimately covers fewer scale columns, probe-verified).""" + if scales.dtype not in _SCALE_DTYPES: + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " + f"got {scales.dtype}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + rows_pad = _round_up(features, SCALE_TILE_ROWS) + # Each 256-row group contributes features_pad * (group_rows/32) bytes and + # group_rows/32 is a multiple of 8. + granule = rows_pad * (ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE) + if scales.numel() % granule != 0: + raise ValueError( + f"{name} numel {scales.numel()} is not a multiple of {granule} " + f"(= round_up({features},128) x {ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE} " + "scale columns per 256-row group)" + ) + max_numel = rows_pad * (allocated_rows // SCALE_BLOCK_SIZE) + if scales.numel() > max_numel: + raise ValueError( + f"{name} numel {scales.numel()} exceeds the maximum {max_numel} implied " + f"by the allocated row count {allocated_rows}" + ) + + +def validate_feature_dims(*, model_dim: int, hidden_dim: int) -> None: + if model_dim <= 0 or model_dim % DIM_ALIGNMENT != 0: + raise ValueError( + f"model dimension D must be a positive multiple of {DIM_ALIGNMENT}, " + f"got {model_dim}" + ) + if hidden_dim <= 0 or hidden_dim % DIM_ALIGNMENT != 0: + raise ValueError( + f"routed-expert hidden dimension F must be a positive multiple of " + f"{DIM_ALIGNMENT}, got {hidden_dim}" + ) + + +def validate_allocated_rows(rows: int, *, name: str = "R") -> None: + """%256 (may be zero): the allocation must be reachable by a legal offsets + vector plus an inactive tail, and a non-256 allocation also breaks the + whole-matrix == per-group-concat identity of the rowwise blocked scales.""" + if rows % ROW_GROUP_ALIGNMENT != 0: + raise ValueError( + f"{name} must be a multiple of {ROW_GROUP_ALIGNMENT}, got {rows}" + ) + + # Small per-(groups, dtype, device) caches for the kernels' alpha/beta and # norm-const tensors. Never cached: the CUDA stream (looked up per call). _ones_cache: dict = {} @@ -95,8 +430,7 @@ def _cached_ones(numel: int, dtype: torch.dtype, device: torch.device) -> torch. def _meta_sig(tag: str, *tensors: torch.Tensor) -> tuple: # torch.Size and stride() are hashable tuples; device/dtype hash directly. return (tag,) + tuple( - (t.shape, t.stride(), t.dtype, t.device, t.storage_offset()) - for t in tensors + (t.shape, t.stride(), t.dtype, t.device, t.storage_offset()) for t in tensors ) @@ -801,3 +1135,62 @@ def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): dtype=torch.bfloat16, device=dy_col_q.device, ) + + +# -------------------------------------------------------------------------- +# Public wrappers: availability-gated entry points over the four custom ops. +# -------------------------------------------------------------------------- + + +def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_q, w13_sf, offsets): + """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. + + See ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` for the full ABI. ``w13_q`` + is E4M3 ``[G, 2F, D]`` contiguous with rows in 32-block GLU order; returns + ``(z_bf16 [R, 2F], h_row_q [R, F], h_row_sf, h_col_q [R, F], h_col_sf)`` + where the columnwise scales are PER-GROUP blocked. Rows past + ``offsets[-1]`` of every output are garbage and read-forbidden. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + + +def mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets): + """Ragged grouped GEMM on prequantized MXFP8 operands, BF16 output. + + ``b_q`` is ``[G, N, K]``-logical quantized along K with free strides + (rowwise casts as-is; dim1-colwise casts transposed into this + orientation); ``b_sf`` is always the per-group blocked ``[N, K/32]`` + orientation. Returns BF16 ``[R, N]`` with rows past ``offsets[-1]`` + uninitialized. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets) + + +def mxfp8_grouped_gemm_dswiglu_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + """FC2 dgrad + dSwiGLU + dual MXFP8 quantization of the FC1 gradient. + + ``z_bf16`` must be the exact fwd-op output. Returns + ``(dz_row_q [R, 2F], dz_row_sf, dz_col_q [R, 2F], dz_col_sf)`` in the same + 32-block order. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + + +def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + Both operands columnwise (32x1) quantized with PER-GROUP blocked scales + (never whole-matrix ``to_blocked`` -- same byte count, silently wrong + block order). Returns contiguous BF16 ``[G, N, K]``. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py b/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py deleted file mode 100644 index 11bba82e4a..0000000000 --- a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_validation.py +++ /dev/null @@ -1,310 +0,0 @@ -# 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. - -"""Host-side precondition validation for the cuDNN-frontend MXFP8 grouped-MLP ops. - -The cuDNN FE grouped kernels hard-code a 256-row group granularity -(``FIX_PAD_SIZE = 256``): per-expert row counts that are only multiples of 128 -SILENTLY and NONDETERMINISTICALLY corrupt results (the corruption locus -migrates between identical-input reruns, consistent with reads of stale memory -at sub-256 group boundaries). No smoke test can prove a misaligned -configuration safe, so the alignment contract is enforced in two tiers: - -* ALWAYS-ON metadata-only checks (no host/device sync, FakeTensor-safe): - dims, dtypes, devices, strides, the allocated row count ``R % 256 == 0``. -* OPT-IN offset-VALUE checks behind ``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` - (one D2H sync per call; validation runs only, skipped for fake tensors): - offsets nondecreasing, every per-expert row count ``% 256 == 0``, and - ``offsets[-1] <= R``. In a default build the offset values are a documented - caller invariant provided by a pad_multiple=256 token dispatcher; there is - NO device-side enforcement. - -Checks raise ValueError rather than asserting, so ``python -O`` cannot strip -them. Metadata gates run before any ``data_ptr()`` gate so FakeTensor tracing -exercises the same checks. -""" - -import os -from typing import Optional - -import torch - -__all__ = [ - "DIM_ALIGNMENT", - "ROW_GROUP_ALIGNMENT", - "SCALE_BLOCK_SIZE", - "blocked_scale_numel", - "host_offsets_validation_enabled", - "validate_allocated_rows", - "validate_blocked_scales", - "validate_feature_dims", - "validate_group_offsets", - "validate_operand", - "validate_ragged_colwise_scales", - "_is_fake", -] - -# MXFP8 scaling block: 32 values share one E8M0 scale. -SCALE_BLOCK_SIZE = 32 -# tcgen05 blocked scale tile: 128 rows x 4 columns, 512 bytes. -SCALE_TILE_ROWS = 128 -SCALE_TILE_COLS = 4 -# Feature-dimension granularity (D and F). -DIM_ALIGNMENT = 128 -# Row-count granularity: per-expert groups AND the allocated row count. This is -# the cuDNN FE kernels' FIX_PAD_SIZE, stricter than the archived family's 128. -ROW_GROUP_ALIGNMENT = 256 -# Byte alignment for TMA/vectorized accesses. -_PTR_ALIGNMENT = 16 - -_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) - - -def _round_up(x: int, to: int) -> int: - return ((x + to - 1) // to) * to - - -def blocked_scale_numel(rows: int, cols: int) -> int: - """Element count of the blocked E8M0 buffer for a logical [rows, cols] scale matrix. - - ``cols`` counts scale values, i.e. the reduced dimension divided by 32. - """ - return _round_up(rows, SCALE_TILE_ROWS) * _round_up(cols, SCALE_TILE_COLS) - - -def host_offsets_validation_enabled() -> bool: - """Opt-in host-side offset validation. Off by default: it forces a D2H sync.""" - return os.environ.get("TORCHAO_MXFP8_VALIDATE_OFFSETS", "0") == "1" - - -def _is_fake(tensor: torch.Tensor) -> bool: - """True for meta/fake tensors, which have no usable data pointer or values.""" - if tensor.device.type == "meta": - return True - try: - from torch._subclasses.fake_tensor import FakeTensor - except ImportError: - return False - return isinstance(tensor, FakeTensor) - - -def validate_group_offsets( - offsets: torch.Tensor, - *, - num_groups: int, - allocated_rows: int, - device: Optional[torch.device] = None, - name: str = "offsets", -) -> None: - """Validate the exclusive-end group offsets tensor. - - Metadata is always checked. The offset VALUES (nondecreasing, per-expert - row counts % 256, offsets[-1] <= R) are checked only when - ``host_offsets_validation_enabled()`` and the tensor is not fake, because - reading them forces a D2H sync. - """ - if not isinstance(offsets, torch.Tensor): - raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") - if num_groups < 1: - raise ValueError( - f"{name} must describe at least one expert group, got G={num_groups}" - ) - if offsets.dtype != torch.int32: - raise ValueError(f"{name} must be int32, got {offsets.dtype}") - if not offsets.is_cuda: - raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") - if device is not None and offsets.device != device: - raise ValueError( - f"{name} must be on {device}, got {offsets.device}; all operands and " - "destinations must share one CUDA device" - ) - if offsets.ndim != 1: - raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") - if offsets.numel() != num_groups: - raise ValueError( - f"{name} must have one entry per local expert: expected {num_groups}, " - f"got {offsets.numel()}" - ) - if not offsets.is_contiguous(): - raise ValueError(f"{name} must be contiguous, got stride {offsets.stride()}") - - if not host_offsets_validation_enabled() or _is_fake(offsets): - return - - values = offsets.tolist() # d2h sync; opt-in debugging path only - previous = 0 - for group, end in enumerate(values): - if end < previous: - raise ValueError( - f"{name} must be nondecreasing, but entry {group} is {end} " - f"after {previous}" - ) - size = end - previous - if size % ROW_GROUP_ALIGNMENT != 0: - raise ValueError( - f"per-expert row counts must be multiples of {ROW_GROUP_ALIGNMENT} " - f"(cuDNN FE FIX_PAD_SIZE; sub-256 groups corrupt results " - f"nondeterministically): expert {group} has {size} rows " - f"(offsets {previous} -> {end})" - ) - previous = end - if previous > allocated_rows: - raise ValueError( - f"{name}[-1] ({previous}) exceeds the allocated row count " - f"({allocated_rows})" - ) - - -def validate_operand( - tensor: torch.Tensor, - *, - name: str, - shape: tuple, - dtype: torch.dtype, - device: torch.device, - stride: Optional[tuple] = None, - check_pointer_alignment: bool = True, -) -> None: - """Validate one operand's dtype, shape, device, optional exact stride, alignment. - - ``stride=None`` accepts any strides (the cuDNN FE wrappers consume both - row-major and transposed-memory colwise operands; every combination the - composite produces is probe-proven). Metadata gates run before the - ``data_ptr()`` gate so FakeTensor tracing exercises the same checks. - """ - if tensor.dtype != dtype: - raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") - if tuple(tensor.shape) != tuple(shape): - raise ValueError( - f"{name} must have shape {tuple(shape)}, got {tuple(tensor.shape)}" - ) - if stride is not None and tuple(tensor.stride()) != tuple(stride): - raise ValueError( - f"{name} must have stride {tuple(stride)}, got {tuple(tensor.stride())}. " - "This layout is part of the ABI; a values-equal tensor with a " - "different stride is not interchangeable." - ) - if tensor.device != device: - raise ValueError( - f"{name} must be on {device}, got {tensor.device}; all operands and " - "destinations must share one CUDA device" - ) - if check_pointer_alignment and not _is_fake(tensor): - if tensor.data_ptr() % _PTR_ALIGNMENT != 0: - raise ValueError( - f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data " - f"pointer is {tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an " - "aligned address. A contiguous view with a nonzero storage " - "offset can violate this." - ) - - -def validate_blocked_scales( - scales: torch.Tensor, - *, - name: str, - logical_rows: int, - logical_cols: int, - device: torch.device, - groups: int = 1, -) -> None: - """Validate a flat blocked E8M0 scale buffer with a statically known size. - - The buffer is carried flat (uint8 or float8_e8m0fnu); its logical shape is - metadata. ``groups > 1`` describes per-expert weight buffers whose per-group - blocks are concatenated. - """ - if scales.dtype not in _SCALE_DTYPES: - raise ValueError( - f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " - f"got {scales.dtype}" - ) - expected = groups * blocked_scale_numel(logical_rows, logical_cols) - if scales.numel() != expected: - raise ValueError( - f"{name} must hold {expected} blocked scale bytes for a logical " - f"[{logical_rows}, {logical_cols}] scale matrix" - + (f" across {groups} experts" if groups > 1 else "") - + f", got {scales.numel()}" - ) - if not scales.is_contiguous(): - raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") - if scales.device != device: - raise ValueError(f"{name} must be on {device}, got {scales.device}") - - -def validate_ragged_colwise_scales( - scales: torch.Tensor, - *, - name: str, - features: int, - allocated_rows: int, - device: torch.device, -) -> None: - """Validate a per-group columnwise scale buffer whose size depends on offsets. - - Columnwise activation scale buffers are sized by the ROUTED row total - ``offsets[-1]`` — a device value — not by the allocated ``R``: at - ``offsets[-1] < R`` they legitimately cover only ``offsets[-1]/32`` scale - columns (probe-verified). So only dtype/device/contiguity and divisibility - are checked: the numel must be a nonnegative multiple of - ``features * SCALE_TILE_COLS`` rows-block granularity and must not exceed - the buffer implied by the allocated rows. - """ - if scales.dtype not in _SCALE_DTYPES: - raise ValueError( - f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " - f"got {scales.dtype}" - ) - if not scales.is_contiguous(): - raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") - if scales.device != device: - raise ValueError(f"{name} must be on {device}, got {scales.device}") - rows_pad = _round_up(features, SCALE_TILE_ROWS) - # Each 256-row group contributes features_pad * (group_rows/32) bytes and - # group_rows/32 is a multiple of 8, so the buffer is a multiple of - # rows_pad * 8 bytes. - granule = rows_pad * (ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE) - if scales.numel() % granule != 0: - raise ValueError( - f"{name} numel {scales.numel()} is not a multiple of {granule} " - f"(= round_up({features},128) x {ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE} " - "scale columns per 256-row group)" - ) - max_numel = rows_pad * (allocated_rows // SCALE_BLOCK_SIZE) - if scales.numel() > max_numel: - raise ValueError( - f"{name} numel {scales.numel()} exceeds the maximum {max_numel} implied " - f"by the allocated row count {allocated_rows}" - ) - - -def validate_feature_dims(*, model_dim: int, hidden_dim: int) -> None: - """D and F must both be positive multiples of 128.""" - if model_dim <= 0 or model_dim % DIM_ALIGNMENT != 0: - raise ValueError( - f"model dimension D must be a positive multiple of {DIM_ALIGNMENT}, " - f"got {model_dim}" - ) - if hidden_dim <= 0 or hidden_dim % DIM_ALIGNMENT != 0: - raise ValueError( - f"routed-expert hidden dimension F must be a positive multiple of " - f"{DIM_ALIGNMENT}, got {hidden_dim}" - ) - - -def validate_allocated_rows(rows: int, *, name: str = "R") -> None: - """The allocated row count must be a multiple of 256 (may be zero). - - Per-expert groups are multiples of 256 (cuDNN FE FIX_PAD_SIZE) and the - allocation must be reachable by a legal offsets vector plus an inactive - tail; a non-256 allocation additionally breaks the whole-matrix == - per-group-concat identity of the rowwise blocked scales. - """ - if rows % ROW_GROUP_ALIGNMENT != 0: - raise ValueError( - f"{name} must be a multiple of {ROW_GROUP_ALIGNMENT}, got {rows}" - ) diff --git a/torchao/prototype/moe_training/mxfp8_grouped_mlp.py b/torchao/prototype/moe_training/mxfp8_grouped_mlp.py deleted file mode 100644 index dfe7010c37..0000000000 --- a/torchao/prototype/moe_training/mxfp8_grouped_mlp.py +++ /dev/null @@ -1,211 +0,0 @@ -# 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. - -"""cuDNN-frontend MXFP8 grouped-MLP operations for routed-expert training. - -Four custom ops, each one launch of a ``cudnn.grouped_gemm_*_wrapper_sm100`` -CuTe DSL kernel from the standalone cudnn-frontend python package (>= 1.27, -Blackwell SM 10.x; no TransformerEngine dependency): - -* :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM + SwiGLU + - rowwise 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization + BF16 pre-GLU. -* :func:`mxfp8_grouped_gemm` -- ragged grouped GEMM on prequantized - MXFP8 operands to BF16 (FC2 forward and FC1 dgrad). -* :func:`mxfp8_grouped_gemm_dswiglu_bwd` -- FC2 dgrad + dSwiGLU + dual MXFP8 - quantization of the FC1 gradient. -* :func:`mxfp8_grouped_gemm_wgrad` -- ragged-reduction grouped weight - gradient (FC1 and FC2). - -CONTRACT (stricter than the archived custom-kernel family): every per-expert -row count and the allocated row count must be multiples of **256** — the cuDNN -FE kernels hard-code ``FIX_PAD_SIZE = 256``, and groups that are only 128-row -aligned corrupt results SILENTLY and NONDETERMINISTICALLY (the corruption -locus migrates between identical-input reruns; no smoke test can prove a -misaligned config safe). Use a token dispatcher with ``pad_multiple=256`` and -see ``grouped_mlp_validation`` for the two-tier enforcement -(``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` for the opt-in synchronized check). - -The FC1 weight must be provided in the cuDNN 32-block GLU row order -``[gate0(32) | up0(32) | gate1(32) | ...]`` along the 2F axis (gate = the -SiLU'd operand). Trainer integration (converter selection, the autograd -composite, weight-layout remaps, expert padding) lives in the consumer; -:func:`is_supported` is the static shape predicate to call before selecting -this family. - -Importing this module registers the four ``torchao::`` custom ops. The -``cudnn`` package itself is imported lazily inside the op bodies. -""" - -import importlib.util - -import torch - -# Importing the ops module registers the custom ops as a side effect. It is -# importable with no cudnn-frontend installed; `import cudnn` is deferred into -# the op bodies. -from torchao.prototype.moe_training.kernels.mxfp8 import ( - grouped_mlp_ops as _grouped_mlp_ops, # noqa: F401 -) -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_validation import ( - DIM_ALIGNMENT, - ROW_GROUP_ALIGNMENT, -) - -__all__ = [ - "DIM_ALIGNMENT", - "ROW_GROUP_ALIGNMENT", - "is_supported", - "mxfp8_grouped_gemm_dswiglu_bwd", - "mxfp8_grouped_gemm_swiglu_fwd", - "mxfp8_grouped_gemm_wgrad", - "mxfp8_grouped_gemm", -] - -_REQUIRED_WRAPPERS = ( - "grouped_gemm_glu_wrapper_sm100", - "grouped_gemm_quant_wrapper_sm100", - "grouped_gemm_dglu_wrapper_sm100", - "grouped_gemm_wgrad_wrapper_sm100", -) -# 1.27 is required: earlier frontends reject prob_tensor=None. -_MIN_FE_VERSION = (1, 27) - - -def _fe_version_tuple(version: str) -> tuple: - """Numeric prefix of a version string as a tuple ('1.27.0' -> (1, 27, 0)). - - Never compare version STRINGS: '1.100' < '1.27' lexicographically. - """ - parts = [] - for piece in version.split("."): - digits = "" - for ch in piece: - if not ch.isdigit(): - break - digits += ch - if not digits: - break - parts.append(int(digits)) - return tuple(parts) - - -def _is_sm_10x() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 - - -def _probe_cudnn_frontend() -> str: - """Empty string when usable; else the reason it is not.""" - if importlib.util.find_spec("cudnn") is None: - return "the cudnn-frontend python package ('cudnn') is not installed" - try: - import cudnn - except Exception as exc: # pragma: no cover - environment-specific - return f"'import cudnn' failed: {exc!r}" - version = getattr(cudnn, "__version__", "0") - if _fe_version_tuple(version) < _MIN_FE_VERSION: - return ( - f"cudnn-frontend {version} is too old; >= " - f"{'.'.join(map(str, _MIN_FE_VERSION))} is required " - "(prob_tensor=None support)" - ) - missing = [name for name in _REQUIRED_WRAPPERS if not hasattr(cudnn, name)] - if missing: - return "cudnn-frontend lacks required wrappers: " + ", ".join(missing) - return "" - - -_mxfp8_grouped_mlp_unavailable_reason = ( - _probe_cudnn_frontend() - if _is_sm_10x() - else ( - "requires an SM 10.x (Blackwell) GPU" - if torch.cuda.is_available() - else "CUDA is not available" - ) -) -_mxfp8_grouped_mlp_kernels_available = _mxfp8_grouped_mlp_unavailable_reason == "" - - -def _require_available() -> None: - """Raise a clean NotImplementedError when the kernels cannot run here.""" - if not _mxfp8_grouped_mlp_kernels_available: - raise NotImplementedError( - "cuDNN-frontend MXFP8 grouped-MLP kernels are unavailable: " - + _mxfp8_grouped_mlp_unavailable_reason - ) - - -def is_supported(model_dim: int, hidden_dim: int) -> bool: - """Static shape predicate for selecting this operator family. - - True when D and F are positive multiples of 128. Integration code must - ALSO guarantee the runtime row contract (per-expert groups and the row - allocation padded to multiples of 256, e.g. dispatcher pad_multiple=256): - row counts live in device memory and are not checkable here. - - Environment availability (cudnn-frontend >= 1.27, SM 10.x) is a separate - concern: combine with ``_mxfp8_grouped_mlp_kernels_available``. - """ - return ( - model_dim > 0 - and hidden_dim > 0 - and model_dim % DIM_ALIGNMENT == 0 - and hidden_dim % DIM_ALIGNMENT == 0 - ) - - -def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_q, w13_sf, offsets): - """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. - - See ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` for the full ABI. ``w13_q`` - is E4M3 ``[G, 2F, D]`` contiguous with rows in 32-block GLU order; returns - ``(z_bf16 [R, 2F], h_row_q [R, F], h_row_sf, h_col_q [R, F], h_col_sf)`` - where the columnwise scales are PER-GROUP blocked. Rows past - ``offsets[-1]`` of every output are garbage and read-forbidden. - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( - x_q, x_sf, w13_q, w13_sf, offsets - ) - - -def mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets): - """Ragged grouped GEMM on prequantized MXFP8 operands, BF16 output. - - ``b_q`` is ``[G, N, K]``-logical quantized along K with free strides - (rowwise casts as-is; dim1-colwise casts transposed into this - orientation); ``b_sf`` is always the per-group blocked ``[N, K/32]`` - orientation. Returns BF16 ``[R, N]`` with rows past ``offsets[-1]`` - uninitialized. - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets) - - -def mxfp8_grouped_gemm_dswiglu_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): - """FC2 dgrad + dSwiGLU + dual MXFP8 quantization of the FC1 gradient. - - ``z_bf16`` must be the exact fwd-op output. Returns - ``(dz_row_q [R, 2F], dz_row_sf, dz_col_q [R, 2F], dz_col_sf)`` in the same - 32-block order. - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( - dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets - ) - - -def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): - """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. - - Both operands columnwise (32x1) quantized with PER-GROUP blocked scales - (never whole-matrix ``to_blocked`` — same byte count, silently wrong - block order). Returns contiguous BF16 ``[G, N, K]``. - """ - _require_available() - return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( - dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets - ) From b3510d64fc599099fc468343501bc8f894d0c818 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Tue, 18 Aug 2026 23:24:05 -0700 Subject: [PATCH 10/11] Rename grouped_mlp_ops.py to cutedsl_grouped_mlp.py Follow the sibling file convention (cutedsl_gated_act_mxfp8.py, cutedsl_quantize_*.py) for the self-contained grouped-MLP module; the name also matches the slot the removed in-repo kernel family vacated. Content unchanged. 29/29 tests green. Co-Authored-By: Claude Fable 5 --- test/prototype/moe_training/test_mxfp8_grouped_mlp.py | 2 +- .../mxfp8/{grouped_mlp_ops.py => cutedsl_grouped_mlp.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename torchao/prototype/moe_training/kernels/mxfp8/{grouped_mlp_ops.py => cutedsl_grouped_mlp.py} (100%) diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py index bf3e332555..8f94e1eece 100644 --- a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py +++ b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py @@ -46,7 +46,7 @@ allow_module_level=True, ) -from torchao.prototype.moe_training.kernels.mxfp8.grouped_mlp_ops import ( +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( _mxfp8_grouped_mlp_kernels_available, _mxfp8_grouped_mlp_unavailable_reason, is_supported, diff --git a/torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py similarity index 100% rename from torchao/prototype/moe_training/kernels/mxfp8/grouped_mlp_ops.py rename to torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py From 5a50196f8770f42f3b7a389607cdc3afd0f46cf8 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Wed, 19 Aug 2026 15:53:26 -0700 Subject: [PATCH 11/11] Fix review findings in the grouped-MLP ops module and its tests Ops module (cutedsl_grouped_mlp.py): - Skip validation memoization when any dim/stride/storage_offset is a SymInt (unhashable under dynamic-shape compile); validation still runs. - Record signatures only on REAL-tensor passes so the first real call always runs the 16-byte data_ptr gate that fake passes skip. - Widen the fwd/bwd int32 element-index guards to R * max(D, 2F), mirroring the mm/wgrad validators' R * max(N, K). - Apply the 16-byte data_ptr gate (skipped for fakes) to validate_blocked_scales and validate_ragged_colwise_scales via a shared _check_pointer_alignment helper. - Make _cached_ones cross-stream safe: record a CUDA event after the first fill and have every cache hit wait on it from the consuming stream (the buffer is immutable once filled, so one event suffices). - Parametrize validate_feature_dims' dim names so mm/wgrad call sites report their actual N/K dims instead of D/F. - Tighten the availability gate to exactly capability (10, 0) (_is_sm_10x -> _is_sm100): the cudnn wrappers are *_sm100-specific. Tests (test_mxfp8_grouped_mlp.py): - Guard the cutedsl_grouped_mlp import with a module-level skip so a stale installed torchao skips instead of erroring at collection. - Extend test_r0_all_ops to the R==0 early returns of all four ops. - Restructure test_optin_offsets_validation to assert the default-build non-rejection on the validator directly, never launching a kernel with out-of-contract 128-row offsets; add a nondecreasing-offsets rejection. - Fix the _quant_weight_colwise docstring (it builds contiguous row-major bytes, not dim1-native strides), add a native=True path, and cover the production dim1-native weight major for ops 2 and 3. Co-Authored-By: Claude Fable 5 --- .../moe_training/test_mxfp8_grouped_mlp.py | 105 +++++++++++-- .../kernels/mxfp8/cutedsl_grouped_mlp.py | 140 +++++++++++++----- 2 files changed, 191 insertions(+), 54 deletions(-) diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py index 8f94e1eece..f53c71d174 100644 --- a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py +++ b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py @@ -40,17 +40,26 @@ from torchao.utils import is_sm_version +# Exactly SM 10.0, matching the ops module's availability gate: the wrapped +# cudnn kernels are *_sm100-specific. if not (torch.cuda.is_available() and is_sm_version(10, 0)): pytest.skip( "MXFP8 fused grouped MLP requires CUDA SM100", allow_module_level=True, ) -from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( - _mxfp8_grouped_mlp_kernels_available, - _mxfp8_grouped_mlp_unavailable_reason, - is_supported, -) +try: + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + _mxfp8_grouped_mlp_kernels_available, + _mxfp8_grouped_mlp_unavailable_reason, + is_supported, + validate_group_offsets, + ) +except ImportError: + pytest.skip( + "installed torchao does not provide the cutedsl_grouped_mlp module", + allow_module_level=True, + ) if not _mxfp8_grouped_mlp_kernels_available: pytest.skip( @@ -141,14 +150,22 @@ def _quant_weight_rowwise(w: torch.Tensor): return torch.stack(qs).view(_E4M3), _cat8(sfs) -def _quant_weight_colwise(w: torch.Tensor): - """[G, N, K] quantized along N (dim1-native strides per group).""" +def _quant_weight_colwise(w: torch.Tensor, native: bool = False): + """[G, N, K] quantized along N. + + native=False: contiguous row-major [G, N, K] bytes. + native=True: the dim1-quantizer memory-transposed major -- [G, N, K] + logical with per-group (1, N) strides (values identical). + """ qs, sfs = [], [] for g in range(w.shape[0]): q, sf = _quant_colwise(w[g], native=False) qs.append(q.view(torch.uint8)) sfs.append(sf.reshape(-1)) - return torch.stack(qs).view(_E4M3), _cat8(sfs) + q = torch.stack(qs).view(_E4M3) + if native: + q = q.transpose(-2, -1).contiguous().transpose(-2, -1) + return q, _cat8(sfs) def _dequant_rowwise(q: torch.Tensor, sf_flat: torch.Tensor): @@ -548,6 +565,34 @@ def test_wgrad_stride_matrix(dbg, a_native, b_native): assert db >= 50.0, f"wgrad[{a_native=} {b_native=}] {db:.1f} dB < 50" +def test_native_weight_major_mm_bwd(dbg): + """Ops 2 and 3 accept the production dim1-native (memory-transposed) + colwise weight major. Both majors carry identical logical values, so each + native arm must agree with the rowmajor arm far above any + reduction-order band.""" + c = dbg + r = _run_chain(c) + # Op 3: dim1-native w2 colwise major. + w2c_nat, w2c_nat_sf = _quant_weight_colwise(c["w2"], native=True) + assert not w2c_nat.is_contiguous() + assert torch.equal(_bytes(w2c_nat), _bytes(c["w2c_q"])) + dz_q, dz_sf, _, _ = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + c["dy_q"], c["dy_sf"], w2c_nat, w2c_nat_sf, r["z"], c["offsets"] + ) + db = compute_error( + _dequant_rowwise(r["dz_q"], r["dz_sf"]), _dequant_rowwise(dz_q, dz_sf) + ).item() + assert db >= 50.0, f"native-major w2 dz vs rowmajor arm: {db:.1f} dB < 50" + # Op 2 (FC1 dgrad): dim1-native w13 colwise major, transposed into + # [G, N=D, K=2F] exactly like the production call. + w13c_nat, w13c_nat_sf = _quant_weight_colwise(c["w13"], native=True) + dx_nat = _OPS.mxfp8_grouped_gemm( + r["dz_q"], r["dz_sf"], w13c_nat.transpose(-2, -1), w13c_nat_sf, c["offsets"] + ) + db = compute_error(r["dx"].float(), dx_nat.float()).item() + assert db >= 50.0, f"native-major w13 dx vs rowmajor arm: {db:.1f} dB < 50" + + # --------------------------------------------------------------------------- # A < R strict tail with planted garbage. # --------------------------------------------------------------------------- @@ -676,6 +721,24 @@ def test_r0_all_ops(): ) assert z.shape == (0, 2 * hidden) and h_q.shape == (0, hidden) assert h_sf.numel() == 0 and h_col_sf.numel() == 0 + y = _OPS.mxfp8_grouped_gemm( + torch.empty(0, hidden, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.zeros(2, D, hidden, dtype=torch.uint8, device=dev).view(_E4M3), + torch.empty(2 * D * hidden // _BLOCK, dtype=_E8M0, device=dev), + offsets, + ) + assert y.shape == (0, D) and y.dtype == torch.bfloat16 + dz_q, dz_sf, dz_colq, dz_col_sf = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + torch.empty(0, D, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.zeros(2, D, hidden, dtype=torch.uint8, device=dev).view(_E4M3), + torch.empty(2 * hidden * D // _BLOCK, dtype=_E8M0, device=dev), + torch.empty(0, 2 * hidden, dtype=torch.bfloat16, device=dev), + offsets, + ) + assert dz_q.shape == (0, 2 * hidden) and dz_sf.numel() == 0 + assert dz_colq.shape == (0, 2 * hidden) and dz_col_sf.numel() == 0 dw = _OPS.mxfp8_grouped_gemm_wgrad( torch.empty(0, D, dtype=_E4M3, device=dev), torch.empty(0, dtype=_E8M0, device=dev), @@ -877,21 +940,37 @@ def test_validation_negatives(case): def test_optin_offsets_validation(monkeypatch): + """Offset VALUES are checked only under TORCHAO_MXFP8_VALIDATE_OFFSETS=1. + + The default-build non-rejection of a 128-row group is asserted on the + validator directly: launching a kernel with misaligned offsets is the + exact out-of-contract config the module documents as corrupting silently + and nondeterministically, so this test must never perform that launch. + The opt-in rejections DO go through the ops, which raise before any + launch. + """ args = _valid_fwd_args() - bad = dict(args) - bad["offsets"] = torch.tensor([128, 512], dtype=torch.int32, device="cuda") + bad_offsets = torch.tensor([128, 512], dtype=torch.int32, device="cuda") # Default build: metadata-only, misaligned VALUES are not (and cannot be) # caught without a D2H sync. monkeypatch.delenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", raising=False) - _OPS.mxfp8_grouped_gemm_swiglu_fwd(**bad) + validate_group_offsets( + bad_offsets, num_groups=2, allocated_rows=512, device=bad_offsets.device + ) monkeypatch.setenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", "1") + bad = dict(args, offsets=bad_offsets) with pytest.raises(ValueError, match="FIX_PAD_SIZE"): _OPS.mxfp8_grouped_gemm_swiglu_fwd(**bad) - over = dict(args) - over["offsets"] = torch.tensor([256, 768], dtype=torch.int32, device="cuda") + dec = dict(args, offsets=torch.tensor([512, 256], dtype=torch.int32, device="cuda")) + with pytest.raises(ValueError, match="nondecreasing"): + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**dec) + + over = dict( + args, offsets=torch.tensor([256, 768], dtype=torch.int32, device="cuda") + ) with pytest.raises(ValueError, match="exceeds the allocated row count"): _OPS.mxfp8_grouped_gemm_swiglu_fwd(**over) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py index ffb8526a97..1c9abd34e5 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py @@ -8,8 +8,8 @@ Four custom ops, each one launch of a ``cudnn.grouped_gemm_*_wrapper_sm100`` kernel from the standalone cudnn-frontend python package (>= 1.27, Blackwell -SM 10.x; no TransformerEngine dependency); the matching public wrappers live -at the bottom of this module: +SM 10.0 exactly -- the wrappers are sm100-specific; no TransformerEngine +dependency); the matching public wrappers live at the bottom of this module: * :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM + SwiGLU + rowwise 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization + BF16 pre-GLU. @@ -116,8 +116,10 @@ def _fe_version_tuple(version: str) -> tuple: return tuple(parts) -def _is_sm_10x() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 +def _is_sm100() -> bool: + # Exactly capability (10, 0): the cudnn wrappers are *_sm100-specific and + # unproven on other SM 10.x parts. + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) def _probe_cudnn_frontend() -> str: @@ -143,9 +145,9 @@ def _probe_cudnn_frontend() -> str: _mxfp8_grouped_mlp_unavailable_reason = ( _probe_cudnn_frontend() - if _is_sm_10x() + if _is_sm100() else ( - "requires an SM 10.x (Blackwell) GPU" + "requires an SM 10.0 (Blackwell) GPU; the cudnn wrappers are sm100-specific" if torch.cuda.is_available() else "CUDA is not available" ) @@ -267,6 +269,19 @@ def validate_group_offsets( ) +def _check_pointer_alignment(tensor: torch.Tensor, *, name: str) -> None: + """16-byte data_ptr gate (TMA/vectorized accesses); fakes have no pointer.""" + if _is_fake(tensor): + return + if tensor.data_ptr() % _PTR_ALIGNMENT != 0: + raise ValueError( + f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data " + f"pointer is {tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an " + "aligned address. A contiguous view with a nonzero storage " + "offset can violate this." + ) + + def validate_operand( tensor: torch.Tensor, *, @@ -298,14 +313,8 @@ def validate_operand( f"{name} must be on {device}, got {tensor.device}; all operands and " "destinations must share one CUDA device" ) - if check_pointer_alignment and not _is_fake(tensor): - if tensor.data_ptr() % _PTR_ALIGNMENT != 0: - raise ValueError( - f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data " - f"pointer is {tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an " - "aligned address. A contiguous view with a nonzero storage " - "offset can violate this." - ) + if check_pointer_alignment: + _check_pointer_alignment(tensor, name=name) def validate_blocked_scales( @@ -336,6 +345,7 @@ def validate_blocked_scales( raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") if scales.device != device: raise ValueError(f"{name} must be on {device}, got {scales.device}") + _check_pointer_alignment(scales, name=name) def validate_ragged_colwise_scales( @@ -375,17 +385,26 @@ def validate_ragged_colwise_scales( f"{name} numel {scales.numel()} exceeds the maximum {max_numel} implied " f"by the allocated row count {allocated_rows}" ) + _check_pointer_alignment(scales, name=name) -def validate_feature_dims(*, model_dim: int, hidden_dim: int) -> None: +def validate_feature_dims( + *, + model_dim: int, + hidden_dim: int, + model_dim_name: str = "model dimension D", + hidden_dim_name: str = "routed-expert hidden dimension F", +) -> None: + """The name arguments let mm/wgrad call sites report their generic N/K + dims instead of the fwd/bwd ops' D/F.""" if model_dim <= 0 or model_dim % DIM_ALIGNMENT != 0: raise ValueError( - f"model dimension D must be a positive multiple of {DIM_ALIGNMENT}, " + f"{model_dim_name} must be a positive multiple of {DIM_ALIGNMENT}, " f"got {model_dim}" ) if hidden_dim <= 0 or hidden_dim % DIM_ALIGNMENT != 0: raise ValueError( - f"routed-expert hidden dimension F must be a positive multiple of " + f"{hidden_dim_name} must be a positive multiple of " f"{DIM_ALIGNMENT}, got {hidden_dim}" ) @@ -401,16 +420,25 @@ def validate_allocated_rows(rows: int, *, name: str = "R") -> None: # Small per-(groups, dtype, device) caches for the kernels' alpha/beta and -# norm-const tensors. Never cached: the CUDA stream (looked up per call). +# norm-const tensors, each stored with the event recorded after its fill: +# the fill runs on the FIRST caller's stream, so a cache hit on any other +# stream must order after it (the buffer is immutable once filled, so one +# event covers every later consumer). Never cached: the CUDA stream itself +# (looked up per call). _ones_cache: dict = {} def _cached_ones(numel: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: key = (numel, dtype, device) - out = _ones_cache.get(key) - if out is None: + hit = _ones_cache.get(key) + if hit is None: out = torch.ones(numel, dtype=dtype, device=device) - _ones_cache[key] = out + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(device)) + _ones_cache[key] = (out, event) + return out + out, event = hit + event.wait(torch.cuda.current_stream(device)) return out @@ -419,22 +447,40 @@ def _cached_ones(numel: int, dtype: torch.dtype, device: torch.device) -> torch. # by storage_offset: torch's CUDA caching allocator hands out aligned storage # bases). A training step calls each op hundreds of times with identical # metadata; the full battery runs once per distinct signature and repeats -# skip straight to the derived dims. Signatures are recorded only AFTER -# validation passes, so a rejected call never poisons the cache. The opt-in -# offsets-VALUES check (TORCHAO_MXFP8_VALIDATE_OFFSETS) reads data, not -# metadata, so it runs on every call while enabled. +# skip straight to the derived dims. Signatures are recorded only AFTER a +# REAL-tensor pass (a rejected call never poisons the cache; a fake pass has +# no data pointer to prove alignment). The opt-in offsets-VALUES check +# (TORCHAO_MXFP8_VALIDATE_OFFSETS) reads data, not metadata, so it runs on +# every call while enabled. _validated_sigs: set = set() _VALIDATED_SIGS_CAP = 4096 -def _meta_sig(tag: str, *tensors: torch.Tensor) -> tuple: +# SymInt ships with every torch new enough to compile these ops; the empty +# tuple keeps the isinstance gate a no-op elsewhere. +_SYMBOLIC_TYPES = (torch.SymInt,) if hasattr(torch, "SymInt") else () + + +def _meta_sig(tag: str, *tensors: torch.Tensor) -> Optional[tuple]: # torch.Size and stride() are hashable tuples; device/dtype hash directly. + # Symbolic metadata (SymInt dims/strides/offsets under dynamic-shape + # compile) is unhashable, so those calls get no signature and never touch + # the memo; the full battery still runs. + for t in tensors: + for d in (*t.shape, *t.stride(), t.storage_offset()): + if isinstance(d, _SYMBOLIC_TYPES): + return None return (tag,) + tuple( (t.shape, t.stride(), t.dtype, t.device, t.storage_offset()) for t in tensors ) -def _remember_sig(sig: tuple) -> None: +def _remember_sig(sig: Optional[tuple], *tensors: torch.Tensor) -> None: + # Fake passes skip the data_ptr alignment gates, so a fake-recorded + # signature would exempt the first REAL call from them: record only + # real-tensor passes (fakes revalidate every time; metadata is cheap). + if sig is None or any(_is_fake(t) for t in tensors): + return if len(_validated_sigs) < _VALIDATED_SIGS_CAP: _validated_sigs.add(sig) @@ -523,7 +569,7 @@ def _allocate_from_specs(specs, device) -> Tuple[torch.Tensor, ...]: def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): sig = _meta_sig("fwd", x_q, x_sf, w13_q, w13_sf, offsets) - if sig in _validated_sigs: + if sig is not None and sig in _validated_sigs: rows, model_dim = x_q.shape groups, two_hidden, _ = w13_q.shape if host_offsets_validation_enabled(): @@ -550,9 +596,10 @@ def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): _require_cuda_device(device, "x_q") validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) validate_allocated_rows(rows) - if rows * two_hidden >= 2**31: + if rows * max(model_dim, two_hidden) >= 2**31: raise ValueError( - f"R * 2F = {rows * two_hidden} does not fit an int32 element index" + f"R * max(D, 2F) = {rows * max(model_dim, two_hidden)} does not " + "fit an int32 element index" ) validate_group_offsets( offsets, num_groups=groups, allocated_rows=rows, device=device @@ -590,7 +637,7 @@ def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): device=device, groups=groups, ) - _remember_sig(sig) + _remember_sig(sig, x_q, x_sf, w13_q, w13_sf, offsets) return rows, model_dim, hidden, groups @@ -682,7 +729,7 @@ def _(x_q, x_sf, w13_q, w13_sf, offsets): def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): sig = _meta_sig("mm", a_q, a_sf, b_q, b_sf, offsets) - if sig in _validated_sigs: + if sig is not None and sig in _validated_sigs: rows, contraction = a_q.shape groups, out_features, _ = b_q.shape if host_offsets_validation_enabled(): @@ -702,7 +749,12 @@ def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): _require_cuda_device(device, "a_q") # N and K are both feature dims here (D/F/2F at the two call sites). - validate_feature_dims(model_dim=out_features, hidden_dim=contraction) + validate_feature_dims( + model_dim=out_features, + hidden_dim=contraction, + model_dim_name="b_q's output feature dim N", + hidden_dim_name="the contraction dim K", + ) validate_allocated_rows(rows) if rows * max(out_features, contraction) >= 2**31: raise ValueError( @@ -745,7 +797,7 @@ def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): device=device, groups=groups, ) - _remember_sig(sig) + _remember_sig(sig, a_q, a_sf, b_q, b_sf, offsets) return rows, out_features, contraction, groups @@ -830,7 +882,7 @@ def _bwd_output_specs(rows: int, hidden: int): def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): sig = _meta_sig("bwd", dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets) - if sig in _validated_sigs: + if sig is not None and sig in _validated_sigs: rows, model_dim = dy_q.shape groups, _, hidden = w2_col_q.shape if host_offsets_validation_enabled(): @@ -858,9 +910,10 @@ def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): _require_cuda_device(device, "dy_q") validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) validate_allocated_rows(rows) - if rows * 2 * hidden >= 2**31: + if rows * max(model_dim, 2 * hidden) >= 2**31: raise ValueError( - f"R * 2F = {rows * 2 * hidden} does not fit an int32 element index" + f"R * max(D, 2F) = {rows * max(model_dim, 2 * hidden)} does not " + "fit an int32 element index" ) validate_group_offsets( offsets, num_groups=groups, allocated_rows=rows, device=device @@ -905,7 +958,7 @@ def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): device=device, groups=groups, ) - _remember_sig(sig) + _remember_sig(sig, dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets) return rows, model_dim, hidden, groups @@ -991,7 +1044,7 @@ def _(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): sig = _meta_sig("wgrad", dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets) - if sig in _validated_sigs: + if sig is not None and sig in _validated_sigs: rows, out_features = dy_col_q.shape in_features = x_col_q.shape[1] groups = offsets.numel() @@ -1019,7 +1072,12 @@ def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): _require_cuda_device(device, "dy_col_q") validate_allocated_rows(rows) - validate_feature_dims(model_dim=out_features, hidden_dim=in_features) + validate_feature_dims( + model_dim=out_features, + hidden_dim=in_features, + model_dim_name="dy_col_q's feature dim N", + hidden_dim_name="x_col_q's feature dim K", + ) if rows * max(out_features, in_features) >= 2**31: raise ValueError( f"R * max(N, K) = {rows * max(out_features, in_features)} does not " @@ -1065,7 +1123,7 @@ def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): # by the ALLOCATED rows while a composite-produced operand's are sized by # the ROUTED total offsets[-1] -- mixing the two is legitimate and # probe-proven (tail case); the kernel reads only within offsets. - _remember_sig(sig) + _remember_sig(sig, dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets) return rows, out_features, in_features, groups