From 27a6103601be4ac0448d0ba7ec5b5cf21dc40f50 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 17 Aug 2026 15:45:18 -0700 Subject: [PATCH 1/3] Add block-scaled grouped GEMM + SwiGLU + RHT + NVFP4 quantization fusion for Rubin Fused MoE grouped GEMM kernel that computes GEMM + SwiGLU, applies a random Hadamard transform (RHT), and quantizes the result to NVFP4 (with E4M3 or E5M3 block scale factors) in a single kernel, targeting Rubin (SM107). Ported from internal MR 2334. The sf_fp8_dtype_override plumbing it depended on landed separately in #545. Co-authored-by: Ali Hassani Co-authored-by: Kaining Zhong --- .../grouped_gemm_glu_hadamard_quant.md | 330 ++ docs/fe-oss-apis/overview.md | 1 + python/cudnn/__init__.py | 2 + python/cudnn/gemm/cutedsl/grouped/__init__.py | 7 + .../grouped/glu_hadamard_quant/__init__.py | 12 + .../cutedsl/grouped/glu_hadamard_quant/api.py | 946 ++++++ ...kscaled_grouped_gemm_glu_hadamard_quant.py | 2684 +++++++++++++++++ .../glu_hadamard_quant/moe_kernel_helpers.py | 1168 +++++++ .../grouped/glu_hadamard_quant/quant_utils.py | 133 + .../grouped/glu_hadamard_quant/rht_utils.py | 224 ++ .../test_grouped_gemm_glu_hadamard_quant.py | 676 +++++ 11 files changed, 6183 insertions(+) create mode 100644 docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard_quant.md create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/__init__.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard_quant.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard_quant.md new file mode 100644 index 000000000..1ccf0ff63 --- /dev/null +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard_quant.md @@ -0,0 +1,330 @@ +# Grouped GEMM + GLU + Hadamard + Quant (SM100) + +**This is an experimental API and subject to change.** + +## JAX support + +JAX arrays are **not supported**: this fusion is block-scaled-only and its mandatory scale-factor inputs use an MMA-interleaved layout with no row-major (JAX) equivalent. JAX inputs raise a clear `ValueError` at the entry points. The API is otherwise type-erased and torch-lazy. + +## Overview + +**Grouped GEMM + GLU + Hadamard + Quant fusion**: A contiguous grouped block-scaled GEMM fused with a GLU/SReLU epilogue, optional RHT (Hadamard transform) output, and optional NVFP4 output quantization on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. Groups are contiguous in the `M` dimension and described by `padded_offsets`. + +This frontend integration is currently wired for the FP4 input path and exposes the quantized Hadamard fusion under the operation name: + +- `GroupedGemmGluHadamardQuantSm100` +- `grouped_gemm_glu_hadamard_quant_wrapper_sm100` + +This kernel performs: + +1. **Block-scaled grouped GEMM** over contiguous expert ranges +2. **GLU or SReLU epilogue** using per-row `prob` +3. **Optional NVFP4 quantization** of the post-activation `D` output +4. **Optional RHT output** in bf16 or NVFP4 form + +### Shapes + +Let `N_out = N / 2` for `act_func="swiglu"` or `"geglu"` and `N_out = N` for `act_func="srelu"`. Let `SF(rows, cols)` denote the swizzled scale-factor layout `(32, 4, ceil_div(rows, 128), 4, ceil_div(ceil_div(cols, sf_vec_size), 4), 1)`. + +- **Inputs** + - `A`: contiguous activation tensor across all groups, shape `(valid_m, K, 1)` + - `B`: weight tensor across all groups, shape `(N, K, L)` in dense mode + - `b_ptrs` / `sfb_ptrs`: int64 CUDA pointer arrays, shape `(L,)`, in discrete mode + - `SFA`: shape `(32, 4, ceil_div(valid_m, 128), 4, ceil_div(ceil_div(K, sf_vec_size), 4), 1)` + - `SFB`: shape `(32, 4, ceil_div(N, 128), 4, ceil_div(ceil_div(K, sf_vec_size), 4), L)` in dense mode + - `padded_offsets`: cumulative padded group ends, shape `(L,)` + - `alpha`: per-group scaling factors, shape `(L,)` + - `prob`: per-row gating probabilities, shape `(valid_m, 1, 1)` + - `bias` (optional): per-expert bias tensor, shape `(N, L)` with stride `(1, N)` + +- **Outputs** + - `C`: intermediate GEMM result before activation/clamping, shape `(valid_m, N, 1)` + - `D`: post-activation output, logical shape `(valid_m, N_out, 1)` + - `SFD`: swizzled e4m3 scale factors for NVFP4 `D`, shape `SF(valid_m, N_out)`, present only when `D` is NVFP4 + - `RHT`: optional Hadamard-transform output, logical shape `(valid_m, N_out, 1)` + - `SFRHT`: e4m3 scale factors for NVFP4 `RHT`, present only when `RHT` is NVFP4 + - Rowwise RHT: swizzled shape `SF(valid_m, N_out)` + - Colwise RHT: swizzled shape `SF(N_out, valid_m)` + +For packed NVFP4 output tensors (`torch.float4_e2m1fn_x2`), the physical tensor stores two logical values per byte along the innermost dimension. The wrapper therefore allocates packed `D` and `RHT` tensors with physical second dimension `N_out / 2`. Raw `torch.uint8` tensors are not accepted as a packed FP4 container by this fusion. RHT data is always stored in the same logical `(m, feature)` orientation as `D`; `rht_rowwise` only changes the transform axis and, for quantized RHT, the `SFRHT` scale domain. + +`L` is the expert count. `valid_m` is the `M` extent of `a_tensor`; by contract it must match the final cumulative padded offset in `padded_offsets`. + +### Equations + +For rows belonging to expert `g`: + +$$ +C[m, n] = \alpha_g \sum_k \mathrm{dequantize}(A[m, k], SFA) \cdot \mathrm{dequantize}(B[n, k, g], SFB) + \mathrm{bias}_g[n] +$$ + +The `bias` term is omitted when `bias_tensor=None`. `C` stores this unclamped GEMM-plus-bias value. + +Split the `N` dimension into consecutive 32-column gate/up blocks: + +$$ +G_b = C[:, 2bG:(2b+1)G], \quad U_b = C[:, (2b+1)G:(2b+2)G], \quad G = 32 +$$ + +When `glu_limit` is set for `swiglu` or `geglu`, both `G_b` and `U_b` are clamped to `[-glu_limit, glu_limit]` before activation. Let `gamma` be `glu_alpha` when it is set and not `1.0`; otherwise `gamma = 1`. + +For **SwiGLU** (`act_func="swiglu"`): + +$$ +D[:, bG:(b+1)G] = \gamma \cdot \mathrm{prob} \cdot U_b \cdot \left(G_b \cdot \sigma(G_b)\right) +$$ + +For **GeGLU** (`act_func="geglu"`): + +$$ +D[:, bG:(b+1)G] = \gamma \cdot \mathrm{prob} \cdot (U_b + 1) \cdot G_b \cdot \sigma(1.702 \cdot G_b) +$$ + +For **SReLU** (`act_func="srelu"`), the kernel does not split `N` into gate/up halves: + +$$ +D = \mathrm{prob} \cdot \mathrm{ReLU}(C)^2 +$$ + +When requested, the RHT output applies a fixed 16-wide orthonormal Hadamard transform to bf16-rounded `D`, either across feature blocks (`rht_rowwise=True`) or across token blocks (`rht_rowwise=False`). + +When `D` or `RHT` is NVFP4, the kernel emits packed e2m1 data plus e4m3 scale factors. `norm_const` and `rht_norm_const` are the corresponding global encode scales. + +### Diagram + +```text +A (valid_m x K x 1), SFA B (N x K x L), SFB padded_offsets + | | | + | dequantize | | + +-------------+-----------+ | + v v + Grouped GEMM over expert ranges ------> group idx + | + | * alpha[group_idx] + v + C (valid_m x N x 1) + | + | GLU over paired 32-col blocks, or SReLU + | with per-row prob + v + D (valid_m x N_out x 1) + | + +--------+---------+ + | | + v v + optional NVFP4 optional Hadamard/RHT + D + SFD RHT, optional SFRHT +``` + +--- + +## API Usage + +### High-level wrapper + +```python +from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 + +result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=a, + b_tensor=b, + sfa_tensor=sfa, + sfb_tensor=sfb, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + prob_tensor=prob, + bias_tensor=bias, + acc_dtype=torch.float32, + c_dtype=torch.bfloat16, + d_dtype=torch.float4_e2m1fn_x2, + cd_major="n", + rht_output=True, + rht_dtype=torch.float4_e2m1fn_x2, + rht_rowwise=False, + norm_const=norm_const, + rht_norm_const=rht_norm_const, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=16, + sf_fp8_dtype_override=None, + vector_f32=False, + m_aligned=256, + act_func="swiglu", + current_stream=None, +) + +c_tensor = result["c_tensor"] +d_tensor = result["d_tensor"] +sfd_tensor = result["sfd_tensor"] +rht_tensor = result["rht_tensor"] +sfrht_tensor = result["sfrht_tensor"] +``` + +Set `rht_output=False` to skip the Hadamard/RHT output. Set `d_dtype=torch.bfloat16` or `rht_dtype=torch.bfloat16` to request unquantized bf16 outputs for the corresponding path. + +### Class API + +```python +from cudnn import GroupedGemmGluHadamardQuantSm100 + +op = GroupedGemmGluHadamardQuantSm100( + sample_a=a, + sample_b=b, + sample_c=c, + sample_d=d, + sample_sfa=sfa, + sample_sfb=sfb, + sample_padded_offsets=padded_offsets, + sample_alpha=alpha, + sample_prob=prob, + sample_sfd=sfd, + sample_rht=rht, + sample_sfrht=sfrht, + sample_bias=bias, + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=16, + sf_fp8_dtype_override=None, + vector_f32=False, + m_aligned=256, + act_func="swiglu", + rht_rowwise=False, +) +assert op.check_support() +op.compile() +op.execute( + a_tensor=a, + b_tensor=b, + c_tensor=c, + d_tensor=d, + sfa_tensor=sfa, + sfb_tensor=sfb, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + prob_tensor=prob, + sfd_tensor=sfd, + rht_tensor=rht, + sfrht_tensor=sfrht, + bias_tensor=bias, + norm_const=norm_const, + rht_norm_const=rht_norm_const, + current_stream=None, +) +``` + +### Discrete weight mode + +The wrapper also accepts per-expert discrete weight allocations: + +```python +result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=a, + b_ptrs=b_ptrs, + sfa_tensor=sfa, + sfb_ptrs=sfb_ptrs, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + prob_tensor=prob, + n=N, + b_dtype=torch.float4_e2m1fn_x2, + b_major="k", +) +``` + +`b_ptrs` and `sfb_ptrs` must be contiguous int64 CUDA tensors containing device pointers for each expert. Each `b_ptrs` entry points to a logical `(N, K)` FP4 expert weight allocation; each `sfb_ptrs` entry points to that expert's scale-factor allocation with logical shape `(32, 4, ceil_div(N, 128), 4, ceil_div(ceil_div(K, sf_vec_size), 4), 1)`. + +--- + +## Parameters + +### Input/output tensors + +- Input tensor **A**: `a_tensor` (wrapper) or `sample_a` / `a_tensor` (class) + - Shape: `(valid_m, K, 1)` + - Layout: must be `k`-major + - Dtype: `float4_e2m1fn_x2` +- Input tensor **B**: `b_tensor` (wrapper) or `sample_b` / `b_tensor` (class) + - Shape: `(N, K, L)` + - Layout: must be `k`-major + - Dtype: must match `A` +- Input tensor **B pointers**: `b_ptrs` (wrapper execute) or `num_experts` / `b_shape` / `b_dtype` (class construction) + - Shape: `(L,)` + - Dtype: `int64` + - Device: CUDA + - `b_dtype` must match `A`; FP4 discrete mode requires `b_major="k"` +- Input tensor **SFA**: `sfa_tensor` (wrapper) or `sample_sfa` / `sfa_tensor` (class) + - Shape: `(32, 4, ceil_div(valid_m, 128), 4, ceil_div(ceil_div(K, sf_vec_size), 4), 1)` + - Dtype: `{float8_e8m0fnu, float8_e4m3fn}` + - Set `sf_fp8_dtype_override="e5m3"` to reinterpret `float8_e4m3fn` storage as UE5M3 on Rubin. +- Input tensor **SFB**: `sfb_tensor` (wrapper) or `sample_sfb` / `sfb_tensor` (class) + - Shape: `(32, 4, ceil_div(N, 128), 4, ceil_div(ceil_div(K, sf_vec_size), 4), L)` + - Dtype: must match `SFA` +- Input tensor **SFB pointers**: `sfb_ptrs` in discrete mode + - Shape: `(L,)` + - Dtype: `int64` + - Device: CUDA +- Input tensor **padded_offsets** + - Shape: `(L,)` + - Dtype: `int32` +- Input tensor **alpha** + - Shape: `(L,)` + - Dtype: `float32` +- Input tensor **prob** + - Shape: `(valid_m, 1, 1)` + - Dtype: `float32` +- Input tensor **bias** (optional) + - Shape: `(N, L)` + - Stride: `(1, N)` + - Dtype: `{float16, bfloat16, float32}` +- Output tensor **C** + - Shape: `(valid_m, N, 1)` + - Layout: must be `n`-major + - Dtype: `{float16, bfloat16}` +- Output tensor **D** + - Logical shape: `(valid_m, N_out, 1)` + - Layout: must be `n`-major + - Dtype: `{bfloat16, float4_e2m1fn_x2}` + - NVFP4 `D` requires `SFD` +- Output tensor **SFD** (present only with NVFP4 `D`) + - Shape: `SF(valid_m, N_out)` = `(32, 4, ceil_div(valid_m, 128), 4, ceil_div(ceil_div(N_out, sf_vec_size), 4), 1)` + - Layout: swizzled scale-factor layout matching `SFA` + - Dtype: `float8_e4m3fn` +- Output tensor **RHT** (optional) + - Logical shape: `(valid_m, N_out, 1)` + - Layout: must be `n`-major + - Dtype: `{bfloat16, float4_e2m1fn_x2}` + - NVFP4 `RHT` requires `SFRHT` +- Output tensor **SFRHT** (present only with NVFP4 `RHT`) + - Shape: `SF(valid_m, N_out)` when `rht_rowwise=True`; `SF(N_out, valid_m)` when `rht_rowwise=False` + - Layout: swizzled scale-factor layout + - Dtype: `float8_e4m3fn` + +### Configuration + +- `act_func`: `"swiglu"`, `"geglu"`, or `"srelu"` +- `cd_major`: must be `"n"` +- `mma_tiler_mn`: must be `(256, 256)` +- `cluster_shape_mn`: cluster dimensions; for this fixed 2-CTA tiler, `cluster_shape_mn[0]` must be `2` and `cluster_shape_mn[1]` must be a positive power of two no larger than `4` +- `sf_vec_size`: must be `16` +- `sf_fp8_dtype_override`: `None` uses the scale format implied by `SFA`/`SFB` dtype. `"e5m3"` reinterprets `torch.float8_e4m3fn` SFA/SFB storage as UE5M3 input scale factors; this is Rubin-only and does not convert tensor contents. +- `m_aligned`: must be `256` +- `rht_rowwise`: selects feature-blocked Hadamard/RHT (`True`) or token-blocked Hadamard/RHT (`False`); for quantized RHT, the scale grid follows the selected axis +- `glu_alpha`: optional final output scale for `swiglu`/`geglu` +- `glu_limit`: optional clamp limit applied to both gate and up blocks for `swiglu`/`geglu` +- `norm_const`: global encode scale for NVFP4 `D` +- `rht_norm_const`: global encode scale for NVFP4 `RHT` + +### Constraints + +- Requires SM100 or newer. +- `N` must be divisible by `64`. +- `K` must satisfy the FP4 K-major 16-byte alignment requirement, so `K` must be divisible by `32`. +- `valid_m` / `a_tensor.shape[0]` must be divisible by `256`; `padded_offsets` must contain cumulative 256-aligned expert ends and end at `valid_m`. +- `N_out` must be divisible by `32`. +- NVFP4 quantization requires `N_out` divisible by `128`. +- NVFP4 quantization is not supported with `act_func="srelu"`. +- `sf_fp8_dtype_override="e5m3"` requires Rubin (SM107) and `SFA`/`SFB` tensors stored as `torch.float8_e4m3fn`. +- `SFRHT` in colwise quantized mode uses the transposed scale domain `SF(N_out, valid_m)`. +- `expert_cnt` must be `<= 1024`. +- Dense and discrete weight modes are mutually exclusive. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index f11274951..1ef692d5b 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -17,6 +17,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [Grouped GEMM (BF16)](gemm_fusions/grouped_gemm.md) - [Grouped GEMM + GLU (Unified)](gemm_fusions/grouped_gemm_glu.md) - [Grouped GEMM + GLU + Hadamard](gemm_fusions/grouped_gemm_glu_hadamard.md) +- [Grouped GEMM + GLU + Hadamard + Quant](gemm_fusions/grouped_gemm_glu_hadamard_quant.md) - [Grouped GEMM + dGLU (Unified)](gemm_fusions/grouped_gemm_dglu.md) - [Grouped GEMM + SwiGLU (Legacy, Contiguous-only)](gemm_fusions/grouped_gemm_swiglu.md) - [Grouped GEMM + dSwiGLU (Legacy, Contiguous-only)](gemm_fusions/grouped_gemm_dswiglu.md) diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 82d469333..d61884332 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -353,6 +353,8 @@ def _dlopen_cudnn(): "grouped_gemm_glu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_wrapper_sm100"), "GroupedGemmGluHadamardSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluHadamardSm100"), "grouped_gemm_glu_hadamard_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_hadamard_wrapper_sm100"), + "GroupedGemmGluHadamardQuantSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluHadamardQuantSm100"), + "grouped_gemm_glu_hadamard_quant_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_hadamard_quant_wrapper_sm100"), "GroupedGemmDgluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDgluSm100"), "grouped_gemm_dglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_wrapper_sm100"), "GroupedGemmWgradSm100": (".gemm.cutedsl.grouped", "GroupedGemmWgradSm100"), diff --git a/python/cudnn/gemm/cutedsl/grouped/__init__.py b/python/cudnn/gemm/cutedsl/grouped/__init__.py index 8cfba7f10..b9b4ca937 100644 --- a/python/cudnn/gemm/cutedsl/grouped/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/__init__.py @@ -36,6 +36,11 @@ grouped_gemm_glu_hadamard_wrapper_sm100, ) +from .glu_hadamard_quant.api import ( + GroupedGemmGluHadamardQuantSm100, + grouped_gemm_glu_hadamard_quant_wrapper_sm100, +) + from .dglu.api import ( GroupedGemmDgluSm100, grouped_gemm_dglu_wrapper_sm100, @@ -66,6 +71,8 @@ "grouped_gemm_glu_wrapper_sm100", "GroupedGemmGluHadamardSm100", "grouped_gemm_glu_hadamard_wrapper_sm100", + "GroupedGemmGluHadamardQuantSm100", + "grouped_gemm_glu_hadamard_quant_wrapper_sm100", "GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100", "GroupedGemmWgradSm100", diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/__init__.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/__init__.py new file mode 100644 index 000000000..c6f249f6b --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from .api import ( + GroupedGemmGluHadamardQuantSm100, + grouped_gemm_glu_hadamard_quant_wrapper_sm100, +) + +__all__ = [ + "GroupedGemmGluHadamardQuantSm100", + "grouped_gemm_glu_hadamard_quant_wrapper_sm100", +] diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py new file mode 100644 index 000000000..e4c67e16c --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py @@ -0,0 +1,946 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""FE API for grouped GEMM GLU forward fusion with fused RHT output and NVFP4 quantization. + +Output modes are dtype driven, mirroring the kernel: + - D is bf16, or NVFP4 (packed e2m1 data + e4m3 block scales in ``sfd``) + - The optional RHT output is bf16, or NVFP4 (packed e2m1 data + e4m3 block scales + in ``sfrht``) + +The RHT data is always stored at D's own (m, f) orientation — only the SCALE grid +follows the transform orientation: swizzled scale factors for logical (m, f) +when ``rht_rowwise``, and swizzled scale factors for logical (f, m) otherwise. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Optional, Tuple + +from cuda.bindings import driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.runtime import from_dlpack, make_fake_stream + +from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.datatypes import _convert_to_cutlass_data_type + +from ..moe_utils import MoEWeightMode +from .rht_utils import HADAMARD_SIZE +from .moe_blockscaled_grouped_gemm_glu_hadamard_quant import BlockScaledMoEGroupedGemmGluHadamardQuantKernel + +# The GLU + Hadamard + quant fusion is block-scaled only: its mandatory +# scale-factor inputs use an MMA-interleaved 6-D layout with no row-major +# equivalent, so they are not expressible as JAX arrays and the API stays +# torch-only. +_JAX_ERROR = ( + "grouped GEMM GLU hadamard quant is not supported for JAX arrays: the block-scaled " + "scale-factor tensors (sfa/sfb) use an MMA-interleaved layout that is not expressible as JAX arrays; " + "pass torch tensors" +) + + +def _require_torch_inputs(sample: Any, api_name: str) -> None: + from cudnn.tensor_adapter import detect_framework + + framework = detect_framework(sample) + if framework == "jax": + raise ValueError(_JAX_ERROR) + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for {api_name}; pass torch tensors") + + +def _sf_layout_shape(rows: int, cols: int, sf_vec_size: int) -> Tuple[int, int, int, int, int, int]: + return (32, 4, ceil_div(rows, 128), 4, ceil_div(ceil_div(cols, sf_vec_size), 4), 1) + + +class GroupedGemmGluHadamardQuantSm100(APIBase): + """Grouped GEMM GLU forward kernel with fused RHT output and NVFP4 quantization.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_sfa: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_prob: torch.Tensor, + sample_b: Optional[torch.Tensor] = None, + sample_sfb: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + b_major: str = "k", + sample_sfd: Optional[torch.Tensor] = None, + sample_rht: Optional[torch.Tensor] = None, + sample_sfrht: Optional[torch.Tensor] = None, + sample_bias: Optional[torch.Tensor] = None, + acc_dtype: Optional[torch.dtype] = None, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[str] = None, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + rht_rowwise: bool = False, + glu_alpha: Optional[float] = None, + glu_limit: Optional[float] = None, + ): + if sample_a is not None: + _require_torch_inputs(sample_a, "GroupedGemmGluHadamardQuantSm100") + import torch + + if acc_dtype is None: + acc_dtype = torch.float32 + super().__init__() + + self._warn_experimental_api() + self._interpret_uint8_as_fp4x2 = True + self._sample_a_tensor = sample_a + self._sample_b_tensor = sample_b + self._sample_d_tensor = sample_d + self._sample_rht_tensor = sample_rht + + if sample_b is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + if sample_sfb is None: + raise ValueError("sample_sfb is required when sample_b is provided") + elif num_experts is not None and sample_b is None: + self.weight_mode = MoEWeightMode.DISCRETE + if b_shape is None or b_dtype is None: + raise ValueError("b_shape and b_dtype are required in discrete mode") + else: + raise ValueError("Provide either (sample_b, sample_sfb) or (num_experts, b_shape, b_dtype)") + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", interpret_uint8_as_fp4x2=False) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", interpret_uint8_as_fp4x2=False) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.sfd_desc = self._make_tensor_desc(sample_sfd, name="sample_sfd") + self.rht_desc = self._make_tensor_desc(sample_rht, name="sample_rht", interpret_uint8_as_fp4x2=False) + self.sfrht_desc = self._make_tensor_desc(sample_sfrht, name="sample_sfrht") + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") + if self.weight_mode == MoEWeightMode.DENSE: + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", interpret_uint8_as_fp4x2=False) + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.expert_cnt = self.padded_offsets_desc.shape[0] + self.b_shape = None + self.b_dtype = None + self.b_major = b_major + else: + self.b_desc = None + self.sfb_desc = None + self.expert_cnt = num_experts + self.b_shape = b_shape + self.b_dtype = b_dtype + self.b_major = b_major + self._value_error_if( + self.padded_offsets_desc.shape[0] != self.expert_cnt, + f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) must equal num_experts ({self.expert_cnt})", + ) + + self.acc_dtype = acc_dtype + self.mma_tiler_mn = mma_tiler_mn + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn if cluster_shape_mn is not None else ((2, 1) if self.use_2cta_instrs else (1, 1)) + self.sf_vec_size = sf_vec_size + self.sf_fp8_dtype_override = sf_fp8_dtype_override + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.act_func = act_func + self.use_dynamic_sched = use_dynamic_sched + self.rht_rowwise = rht_rowwise + self.glu_alpha = glu_alpha + self.glu_limit = glu_limit + self._kernel = BlockScaledMoEGroupedGemmGluHadamardQuantKernel + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + self._workspace = None + + def check_support(self) -> bool: + import torch + + tensor_m, k, _ = self._tensor_shape(self.a_desc, name="sample_a") + if self.weight_mode == MoEWeightMode.DENSE: + n, _, l = self._tensor_shape(self.b_desc, name="sample_b") + else: + if len(self.b_shape) == 2: + n, b_k = self.b_shape + else: + n, b_k, _ = self.b_shape + self._value_error_if(b_k != k, f"B K dimension ({b_k}) must match A K dimension ({k})") + l = self.expert_cnt + _, n_c, _ = self._tensor_shape(self.c_desc, name="sample_c") + _, n_d, _ = self._tensor_shape(self.d_desc, name="sample_d") + n_out = n if self.act_func == "srelu" else n // 2 + + self._value_error_if(l != self.expert_cnt, f"B L dimension ({l}) must match expert_cnt ({self.expert_cnt})") + self._value_error_if(n % 64 != 0, f"N must be divisible by 64, got {n}") + self._value_error_if(n_out % (2 * HADAMARD_SIZE) != 0, f"D N dimension must be divisible by {2 * HADAMARD_SIZE}, got {n_out}") + + # ---- Output / dump modes (dtype driven, mirroring the kernel) ---- + self.d_quant = self._is_fp4x2(self.d_desc) + self.generate_rht = self.rht_desc is not None + self.rht_quant = self.generate_rht and self._is_fp4x2(self.rht_desc) + self._value_error_if( + self.d_quant != (self.sfd_desc is not None), + "NVFP4 sample_d and sample_sfd must be passed together", + ) + self._value_error_if( + self.rht_quant != (self.sfrht_desc is not None), + "NVFP4 sample_rht and sample_sfrht must be passed together", + ) + self._value_error_if( + (self.d_quant or self.rht_quant) and n_out % (8 * HADAMARD_SIZE) != 0, + f"NVFP4 quantization requires the D N dimension to be divisible by {8 * HADAMARD_SIZE}, got {n_out}", + ) + self._value_error_if( + (self.d_quant or self.rht_quant) and self.act_func == "srelu", + "NVFP4 quantization is not supported with act_func 'srelu'", + ) + + self._check_tensor_shape(self.a_desc, (tensor_m, k, 1), "A") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.b_desc, (n, k, l), "B") + self._check_tensor_shape(self.c_desc, (tensor_m, n, 1), "C") + self._check_tensor_shape(self.d_desc, (tensor_m, n_out, 1), "D") + self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, ceil_div(ceil_div(k, self.sf_vec_size), 4), 1), "SFA") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, ceil_div(ceil_div(k, self.sf_vec_size), 4), l), "SFB") + self._check_tensor_shape(self.padded_offsets_desc, (l,), "padded_offsets") + self._check_tensor_shape(self.alpha_desc, (l,), "alpha") + self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") + self._check_tensor_shape(self.bias_desc, (n, l), "bias") + if self.d_quant: + self._check_tensor_shape(self.sfd_desc, _sf_layout_shape(tensor_m, n_out, self.sf_vec_size), "SFD") + if self.generate_rht: + self._check_tensor_shape(self.rht_desc, (tensor_m, n_out, 1), "RHT") + if self.rht_quant: + if self.rht_rowwise: + self._check_tensor_shape(self.sfrht_desc, _sf_layout_shape(tensor_m, n_out, self.sf_vec_size), "SFRHT") + else: + self._check_tensor_shape(self.sfrht_desc, _sf_layout_shape(n_out, tensor_m, self.sf_vec_size), "SFRHT") + + self._check_tensor_stride(self.a_desc, stride=[(k, 1, tensor_m * k)], name="A", extra_error_msg="A must have k-major layout") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_stride(self.b_desc, stride=[(k, 1, n * k)], name="B", extra_error_msg="B must have k-major layout") + self._check_tensor_stride(self.c_desc, stride=[(n_c, 1, tensor_m * n_c)], name="C", extra_error_msg="C must have n-major layout") + self._check_tensor_stride(self.d_desc, stride=[(n_d, 1, tensor_m * n_d)], name="D", extra_error_msg="D must have n-major layout") + self._check_tensor_stride(self.bias_desc, stride=[(1, n)], name="bias") + if self.generate_rht: + self._check_tensor_stride(self.rht_desc, stride=[(n_d, 1, tensor_m * n_d)], name="RHT", extra_error_msg="RHT must have n-major layout") + + self.ab_dtype = self._check_dtype( + self.a_desc, + dtype=[torch.float4_e2m1fn_x2], + name="A", + ) + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype(self.b_desc, dtype=self.ab_dtype, name="B", extra_error_msg="B must match A dtype") + else: + self._value_error_if(self.b_dtype != self.ab_dtype, f"b_dtype ({self.b_dtype}) must match A dtype ({self.ab_dtype})") + self._value_error_if(self.b_major not in ["k", "n"], f"b_major must be 'k' or 'n', got {self.b_major}") + self._value_error_if(self._is_fp4x2(self.ab_dtype) and self.b_major != "k", "b_major must be 'k' when ab_dtype is fp4") + self.sf_dtype = self._check_dtype(self.sfa_desc, dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], name="SFA") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype(self.sfb_desc, dtype=self.sf_dtype, name="SFB", extra_error_msg="SFB must match SFA dtype") + self.c_dtype = self._check_dtype(self.c_desc, dtype=[torch.float16, torch.bfloat16], name="C") + self.d_dtype = self._check_dtype(self.d_desc, dtype=[torch.bfloat16, torch.float4_e2m1fn_x2], name="D") + self._check_dtype(self.alpha_desc, dtype=torch.float32, name="alpha") + self._check_dtype(self.prob_desc, dtype=torch.float32, name="prob") + self._check_dtype(self.bias_desc, dtype=[torch.float16, torch.bfloat16, torch.float32], name="bias") + if self.d_quant: + self._check_dtype(self.sfd_desc, dtype=torch.float8_e4m3fn, name="SFD") + if self.generate_rht: + self._check_dtype(self.rht_desc, dtype=[torch.bfloat16, torch.float4_e2m1fn_x2], name="RHT") + if self.rht_quant: + self._check_dtype(self.sfrht_desc, dtype=torch.float8_e4m3fn, name="SFRHT") + self._check_dtype(self.acc_dtype, dtype=torch.float32, name="acc_dtype") + + self._value_error_if(self.sf_vec_size != 16, f"sf_vec_size must be 16, got {self.sf_vec_size}") + self._value_error_if( + self.sf_fp8_dtype_override not in (None, "e5m3"), + f"sf_fp8_dtype_override must be None or 'e5m3', got {self.sf_fp8_dtype_override!r}", + ) + if self.sf_fp8_dtype_override == "e5m3": + self._value_error_if( + self.sf_dtype != torch.float8_e4m3fn, + f"sf_fp8_dtype_override='e5m3' requires torch.float8_e4m3fn scale-factor storage at sf_vec_size 16, " + f"got sf_dtype={self.sf_dtype}, sf_vec_size={self.sf_vec_size}", + ) + self._value_error_if( + not self._is_rubin_kernel, + f"sf_fp8_dtype_override='e5m3' requires Rubin (SM107), got device type {self._device_type!r}", + ) + self._value_error_if( + self.act_func not in ["swiglu", "geglu", "srelu"], + f"act_func must be 'swiglu', 'geglu', or 'srelu', got {self.act_func}", + ) + self._value_error_if( + not self.use_2cta_instrs or self.mma_tiler_mn != (256, 256), f"RHT fusion requires mma_tiler_mn=(256, 256), got {self.mma_tiler_mn}" + ) + self._value_error_if(self.cluster_shape_mn[0] % 2 != 0, f"cluster_shape_mn[0] must be divisible by 2, got {self.cluster_shape_mn[0]}") + self._value_error_if( + not ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 + and self.cluster_shape_mn[0] > 0 + and self.cluster_shape_mn[1] > 0 + and self.cluster_shape_mn[0] <= 4 + and self.cluster_shape_mn[1] <= 4 + and is_power_of_2(self.cluster_shape_mn[0]) + and is_power_of_2(self.cluster_shape_mn[1]) + ), + f"Invalid cluster shape: {self.cluster_shape_mn}", + ) + self._value_error_if( + self.m_aligned != BlockScaledMoEGroupedGemmGluHadamardQuantKernel.FIX_PAD_SIZE, + f"m_aligned must be {BlockScaledMoEGroupedGemmGluHadamardQuantKernel.FIX_PAD_SIZE}, got {self.m_aligned}", + ) + self._value_error_if(self.expert_cnt > 1024, f"expert_cnt must be <= 1024, got {self.expert_cnt}") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmGluHadamardQuantSm100 requires SM100+, found SM{compute_capability}") + + if not self._kernel.can_implement( + _convert_to_cutlass_data_type(self.ab_dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2), + _convert_to_cutlass_data_type(self.sf_dtype), + self.sf_vec_size, + _convert_to_cutlass_data_type(self.acc_dtype), + _convert_to_cutlass_data_type(self.d_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2), + self.use_2cta_instrs, + self.mma_tiler_mn, + self.cluster_shape_mn, + self.m_aligned, + n, + k, + l, + "k", + self.b_major, + "n", + self.m_aligned, + rht_quant=self.rht_quant, + d_quant=self.d_quant, + ): + raise RuntimeError("Unsupported grouped GEMM GLU Hadamard Quant configuration") + + self._is_supported = True + return True + + def compile(self) -> None: + import torch + + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + if self.a_desc.shape[0] == 0: + return + + kernel = self._kernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + enable_bias=self.bias_desc is not None, + rht_rowwise=self.rht_rowwise if self.generate_rht else False, + sf_fp8_dtype_override=self.sf_fp8_dtype_override, + glu_alpha=self.glu_alpha, + glu_limit=self.glu_limit, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) + max_active_clusters -= self.num_cluster_overlap_margin + self._value_error_if(max_active_clusters <= 0, "max_active_clusters must be > 0 after overlap margin") + self._workspace = torch.empty(max(kernel.get_workspace_bytes(), 1), dtype=torch.uint8, device="cuda") + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + fake_workspace_ptr = cute.runtime.nullptr(dtype=cutlass.Uint8, assumed_align=128) + cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + + valid_m = cute.sym_int(divisibility=self.m_aligned) + tensor_m_128 = cute.sym_int() + stride_sfa_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + + a_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, self.a_desc.shape[1], 1), + stride_order=self.a_desc.stride_order, + dynamic_mode=self.a_desc.stride_order[0], + divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, + ) + c_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, self.c_desc.shape[1], 1), + stride_order=self.c_desc.stride_order, + dynamic_mode=self.c_desc.stride_order[0], + divisibility=8 if self._is_f16(self.c_desc) else 16, + ) + d_cute_arg = self._make_fake_cute_compact_tensor( + dtype=self.d_desc.dtype, + shape=(valid_m, self.d_desc.shape[1], 1), + stride_order=self.d_desc.stride_order, + dynamic_mode=self.d_desc.stride_order[0], + divisibility=8 if self._is_f16(self.d_desc) else 32, + ) + sfd_cute_arg = None + if self.d_quant: + stride_sfd_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfd_cute_arg = self._make_fake_cute_tensor( + dtype=self.sfd_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_desc.shape[4], 1), + stride=(16, 4, self.sfd_desc.stride[2], 1, 512, stride_sfd_tensor_m_128), + ) + rht_cute_arg = None + if self.generate_rht: + rht_cute_arg = self._make_fake_cute_compact_tensor( + dtype=self.rht_desc.dtype, + shape=(valid_m, self.rht_desc.shape[1], 1), + stride_order=self.rht_desc.stride_order, + dynamic_mode=self.rht_desc.stride_order[0], + divisibility=8 if self._is_f16(self.rht_desc) else 32, + ) + sfrht_cute_arg = None + if self.rht_quant: + if self.rht_rowwise: + stride_sfrht_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfrht_cute_arg = self._make_fake_cute_tensor( + dtype=self.sfrht_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfrht_desc.shape[4], 1), + stride=(16, 4, self.sfrht_desc.stride[2], 1, 512, stride_sfrht_tensor_m_128), + ) + else: + sfrht_rest_m = cute.sym_int() + stride_sfrht_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfrht_l = cute.sym_int(divisibility=32 * 4 * 4) + sfrht_cute_arg = self._make_fake_cute_tensor( + dtype=self.sfrht_desc.dtype, + shape=(32, 4, self.sfrht_desc.shape[2], 4, sfrht_rest_m, 1), + stride=(16, 4, stride_sfrht_rest_m, 1, 512, stride_sfrht_l), + ) + prob_cute_fake = self._make_fake_cute_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, 1, 1), + stride=self.prob_desc.stride, + ) + sfa_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfa_desc.shape[4], 1), + stride=(16, 4, self.sfa_desc.stride[2], 1, 512, stride_sfa_tensor_m_128), + ) + if self.weight_mode == MoEWeightMode.DENSE: + b_cute_arg = self._make_fake_cute_tensor_from_desc(self.b_desc, assumed_align=16) + sfb_cute_arg = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16) + n_arg = cutlass.Int32(0) + k_arg = cutlass.Int32(0) + b_stride_arg = cutlass.Int64(0) + b_major_arg = OperandMajorMode.K + workspace_arg = fake_workspace_ptr + else: + if len(self.b_shape) == 2: + n_compile, k_compile = self.b_shape + else: + n_compile, k_compile, _ = self.b_shape + b_major_arg = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN + b_stride_size = k_compile if self.b_major == "k" else n_compile + b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") + sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") + b_cute_arg = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator + sfb_cute_arg = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + n_arg = cutlass.Int32(n_compile) + k_arg = cutlass.Int32(k_compile) + b_stride_arg = cutlass.Int64(b_stride_size) + workspace_arg = cached_workspace_ptr + self._n = n_compile + self._k = k_compile + self._b_stride_size = b_stride_size + alpha_cute_fake = self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16) + padded_offsets_cute_fake = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) + bias_cute_fake = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) + cached_linear_offset = cutlass.Float32(1.0 if self.act_func == "geglu" else 0.0) + + compiled_kernel = cute.compile( + kernel, + a_cute_fake, + b_cute_arg, + sfa_cute_fake, + sfb_cute_arg, + n_arg, + k_arg, + b_stride_arg, + b_major_arg, + workspace_arg, + c_cute_fake, + d_cute_arg, + sfd_cute_arg, + rht_cute_arg, + sfrht_cute_arg, + padded_offsets_cute_fake, + alpha_cute_fake, + prob_cute_fake, + bias_cute_fake, + max_active_clusters, + fake_stream, + linear_offset=cached_linear_offset, + norm_const=cutlass.Float32(1.0), + rht_norm_const=cutlass.Float32(1.0), + options="--enable-tvm-ffi", + ) + + if self.weight_mode == MoEWeightMode.DENSE: + + def tensor_api( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + sfd_tensor: Optional[torch.Tensor], + rht_tensor: Optional[torch.Tensor], + sfrht_tensor: Optional[torch.Tensor], + bias_tensor: Optional[torch.Tensor], + norm_const: float, + rht_norm_const: float, + stream: cuda.CUstream, + ) -> None: + compiled_kernel( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + cached_workspace_ptr, + c_tensor, + d_tensor, + sfd_tensor, + rht_tensor, + sfrht_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + stream, + cached_linear_offset, + cutlass.Float32(norm_const), + cutlass.Float32(rht_norm_const), + ) + + self._compiled_kernel = tensor_api + else: + cached_n = cutlass.Int32(self._n) + cached_k = cutlass.Int32(self._k) + cached_b_stride = cutlass.Int64(self._b_stride_size) + + def tensor_api( + a_tensor: torch.Tensor, + b_ptrs_device: torch.Tensor, + sfb_ptrs_device: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + sfd_tensor: Optional[torch.Tensor], + rht_tensor: Optional[torch.Tensor], + sfrht_tensor: Optional[torch.Tensor], + bias_tensor: Optional[torch.Tensor], + norm_const: float, + rht_norm_const: float, + stream: cuda.CUstream, + ) -> None: + compiled_kernel( + a_tensor, + int(b_ptrs_device.data_ptr()), + sfa_tensor, + int(sfb_ptrs_device.data_ptr()), + cached_n, + cached_k, + cached_b_stride, + cached_workspace_ptr, + c_tensor, + d_tensor, + sfd_tensor, + rht_tensor, + sfrht_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + stream, + cached_linear_offset, + cutlass.Float32(norm_const), + cutlass.Float32(rht_norm_const), + ) + + self._compiled_kernel = tensor_api + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + sfb_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + sfd_tensor: Optional[torch.Tensor] = None, + rht_tensor: Optional[torch.Tensor] = None, + sfrht_tensor: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + norm_const: float = 1.0, + rht_norm_const: float = 1.0, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + import torch + + self._ensure_support_checked() + if self._compiled_kernel is None: + raise RuntimeError("Kernel has not been compiled") + if a_tensor.shape[0] == 0: + return + if current_stream is None: + current_stream = cuda.CUstream(torch.cuda.current_stream(a_tensor.device).cuda_stream) + if self.d_quant and sfd_tensor is None: + raise ValueError("sfd_tensor must be provided when D is NVFP4") + if self.generate_rht and rht_tensor is None: + raise ValueError("rht_tensor must be provided when the RHT output is enabled") + if self.rht_quant and sfrht_tensor is None: + raise ValueError("sfrht_tensor must be provided when the RHT output is NVFP4") + + if self.weight_mode == MoEWeightMode.DENSE: + if b_tensor is None or sfb_tensor is None: + raise ValueError("b_tensor and sfb_tensor must be provided in dense mode") + self._compiled_kernel( + a_tensor, + b_tensor, + c_tensor, + d_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + sfd_tensor, + rht_tensor, + sfrht_tensor, + bias_tensor, + norm_const, + rht_norm_const, + current_stream, + ) + else: + if b_ptrs is None or sfb_ptrs is None: + raise ValueError("b_ptrs and sfb_ptrs must be provided in discrete mode") + self._compiled_kernel( + a_tensor, + b_ptrs, + sfb_ptrs, + c_tensor, + d_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + sfd_tensor, + rht_tensor, + sfrht_tensor, + bias_tensor, + norm_const, + rht_norm_const, + current_stream, + ) + + +_logger = logging.getLogger(__name__) +_cache_of_GroupedGemmGluHadamardQuantSm100Objects = {} + + +def grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + sfb_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + n: Optional[int] = None, + b_dtype: Optional[torch.dtype] = None, + b_major: str = "k", + bias_tensor: Optional[torch.Tensor] = None, + acc_dtype: Optional[torch.dtype] = None, + c_dtype: Optional[torch.dtype] = None, + d_dtype: Optional[torch.dtype] = None, + cd_major: str = "n", + rht_output: bool = True, + rht_dtype: Optional[torch.dtype] = None, + rht_rowwise: bool = False, + glu_alpha: Optional[float] = None, + glu_limit: Optional[float] = None, + norm_const: float = 1.0, + rht_norm_const: float = 1.0, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[str] = None, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + current_stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """High-level wrapper for grouped GEMM GLU forward fusion with fused RHT output. + + Output modes are dtype driven: ``d_dtype``/``rht_dtype`` of + ``torch.float4_e2m1fn_x2`` emit packed NVFP4 data plus e4m3 block scales + (``sfd_tensor``/``sfrht_tensor``); ``torch.bfloat16`` emits plain bf16. + ``sf_fp8_dtype_override="e5m3"`` reinterprets ``torch.float8_e4m3fn`` + SFA/SFB storage as UE5M3 input scale factors on Rubin. + ``norm_const``/``rht_norm_const`` are the NVFP4 global encode scales + (2688/global_amax, or 1.0). + """ + from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor + + if a_tensor is not None: + _require_torch_inputs(a_tensor, "grouped_gemm_glu_hadamard_quant_wrapper_sm100") + import torch + + if acc_dtype is None: + acc_dtype = torch.float32 + if c_dtype is None: + c_dtype = torch.bfloat16 + if d_dtype is None: + d_dtype = torch.bfloat16 + if rht_dtype is None: + rht_dtype = torch.bfloat16 + if a_tensor.dtype == torch.uint8: + raise ValueError("a_tensor dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + if b_tensor is not None and b_tensor.dtype == torch.uint8: + raise ValueError("b_tensor dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + if b_dtype == torch.uint8: + raise ValueError("b_dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + if d_dtype == torch.uint8: + raise ValueError("d_dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + if rht_dtype == torch.uint8: + raise ValueError("rht_dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + + valid_m = a_tensor.shape[0] + is_dense = b_tensor is not None + is_discrete = b_ptrs is not None + if is_dense and is_discrete: + raise ValueError("Provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs), not both") + if not is_dense and not is_discrete: + raise ValueError("Must provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs)") + + _, k_physical, _ = a_tensor.shape + if is_dense: + weight_mode = MoEWeightMode.DENSE + n_full, _, l = b_tensor.shape + if sfb_tensor is None: + raise ValueError("sfb_tensor is required in dense mode") + else: + weight_mode = MoEWeightMode.DISCRETE + _require_pointer_tensor(b_ptrs, "b_ptrs") + l = b_ptrs.shape[0] + _require_pointer_tensor(sfb_ptrs, "sfb_ptrs", l) + if n is None or b_dtype is None: + raise ValueError("n and b_dtype are required for discrete mode") + n_full = n + k_logical = k_physical * 2 if b_dtype == torch.float4_e2m1fn_x2 else k_physical + b_shape = (n_full, k_logical) + n_out = n_full if act_func == "srelu" else n_full // 2 + + if cd_major != "n": + raise ValueError(f"cd_major must be 'n', got {cd_major}") + + d_quant = d_dtype == torch.float4_e2m1fn_x2 + rht_quant = rht_output and rht_dtype == torch.float4_e2m1fn_x2 + device = a_tensor.device + + def alloc_n_major(rows: int, cols: int, dtype: torch.dtype) -> torch.Tensor: + return torch.empty_strided((rows, cols, 1), (cols, 1, rows * cols), dtype=dtype, device=device) + + def alloc_swizzled_sf(rows: int, cols: int) -> torch.Tensor: + shape = (1, ceil_div(rows, 128), ceil_div(ceil_div(cols, sf_vec_size), 4), 32, 4, 4) + return torch.empty(shape, dtype=torch.float8_e4m3fn, device=device).permute(3, 4, 1, 5, 2, 0) + + c_tensor = alloc_n_major(valid_m, n_full, c_dtype) + if d_quant: + d_tensor = alloc_n_major(valid_m, n_out // 2, d_dtype) + sfd_tensor = alloc_swizzled_sf(valid_m, n_out) + else: + d_tensor = alloc_n_major(valid_m, n_out, d_dtype) + sfd_tensor = None + rht_tensor = None + sfrht_tensor = None + if rht_output: + rht_tensor = alloc_n_major(valid_m, n_out // 2 if rht_quant else n_out, rht_dtype) + if rht_quant: + if rht_rowwise: + sfrht_tensor = alloc_swizzled_sf(valid_m, n_out) + else: + sfrht_tensor = alloc_swizzled_sf(n_out, valid_m) + + if valid_m == 0: + return TupleDict(c_tensor=c_tensor, d_tensor=d_tensor, sfd_tensor=sfd_tensor, rht_tensor=rht_tensor, sfrht_tensor=sfrht_tensor) + + def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple(i for i, _ in sorted(enumerate(tensor.stride()), key=lambda item: item[1])) + + def tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: + if tensor is None: + return None, None, None + return tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype + + def dynamic_m_tensor_signature( + tensor: Optional[torch.Tensor], static_shape_suffix: Optional[Tuple[int, ...]], dynamic_stride_dims: Tuple[int, ...] = () + ) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: + if tensor is None: + return None, None, None + stride_signature = tuple(None if idx in dynamic_stride_dims else value for idx, value in enumerate(tensor.stride())) + return static_shape_suffix, stride_signature, tensor.dtype + + cache_key = ( + weight_mode, + act_func, + a_tensor.shape[1:], + tuple(b_tensor.shape) if is_dense else b_shape, + c_tensor.shape[1:], + a_tensor.dtype, + b_tensor.dtype if is_dense else b_dtype, + c_tensor.dtype, + d_tensor.dtype, + rht_output, + rht_dtype if rht_output else None, + rht_rowwise if rht_output else None, + glu_alpha, + glu_limit, + stride_order(a_tensor), + stride_order(b_tensor) if is_dense else b_major, + stride_order(c_tensor), + *dynamic_m_tensor_signature(sfa_tensor, (sfa_tensor.shape[4], 1), dynamic_stride_dims=(5,)), + *(tensor_signature(sfb_tensor) if is_dense else (tuple(sfb_ptrs.shape), tuple(sfb_ptrs.stride()), sfb_ptrs.dtype)), + *tensor_signature(alpha_tensor), + *dynamic_m_tensor_signature(prob_tensor, (1, 1)), + *tensor_signature(bias_tensor), + *tensor_signature(padded_offsets), + acc_dtype, + mma_tiler_mn, + cluster_shape_mn, + sf_vec_size, + sf_fp8_dtype_override, + vector_f32, + m_aligned, + use_dynamic_sched, + *((tuple(b_ptrs.shape), tuple(b_ptrs.stride()), b_ptrs.dtype, l) if is_discrete else ()), + ) + + if cache_key in _cache_of_GroupedGemmGluHadamardQuantSm100Objects: + api = _cache_of_GroupedGemmGluHadamardQuantSm100Objects[cache_key] + else: + common_kwargs = dict( + sample_a=a_tensor, + sample_c=c_tensor, + sample_d=d_tensor, + sample_sfa=sfa_tensor, + sample_padded_offsets=padded_offsets, + sample_alpha=alpha_tensor, + sample_prob=prob_tensor, + sample_sfd=sfd_tensor, + sample_rht=rht_tensor, + sample_sfrht=sfrht_tensor, + sample_bias=bias_tensor, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=vector_f32, + m_aligned=m_aligned, + act_func=act_func, + use_dynamic_sched=use_dynamic_sched, + rht_rowwise=rht_rowwise, + glu_alpha=glu_alpha, + glu_limit=glu_limit, + ) + if is_dense: + api = GroupedGemmGluHadamardQuantSm100(sample_b=b_tensor, sample_sfb=sfb_tensor, **common_kwargs) + else: + api = GroupedGemmGluHadamardQuantSm100( + num_experts=l, + b_shape=b_shape, + b_dtype=b_dtype, + b_major=b_major, + **common_kwargs, + ) + api.check_support() + api.compile() + _cache_of_GroupedGemmGluHadamardQuantSm100Objects[cache_key] = api + + if is_dense: + api.execute( + a_tensor=a_tensor, + b_tensor=b_tensor, + c_tensor=c_tensor, + d_tensor=d_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + prob_tensor=prob_tensor, + sfd_tensor=sfd_tensor, + rht_tensor=rht_tensor, + sfrht_tensor=sfrht_tensor, + bias_tensor=bias_tensor, + norm_const=norm_const, + rht_norm_const=rht_norm_const, + current_stream=current_stream, + ) + else: + api.execute( + a_tensor=a_tensor, + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + c_tensor=c_tensor, + d_tensor=d_tensor, + sfa_tensor=sfa_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + prob_tensor=prob_tensor, + sfd_tensor=sfd_tensor, + rht_tensor=rht_tensor, + sfrht_tensor=sfrht_tensor, + bias_tensor=bias_tensor, + norm_const=norm_const, + rht_norm_const=rht_norm_const, + current_stream=current_stream, + ) + return TupleDict(c_tensor=c_tensor, d_tensor=d_tensor, sfd_tensor=sfd_tensor, rht_tensor=rht_tensor, sfrht_tensor=sfrht_tensor) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py new file mode 100644 index 000000000..5e656fd78 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py @@ -0,0 +1,2684 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +MoE Block-Scaled Grouped GEMM Kernel with GLU (SwiGLU/GeGLU) + Hadamard Transform Fusion. + +Supports: + - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler) + - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout + - BF16 or NVFP4 (packed e2m1 + e4m3 block scales) D output with GLU activation + - Optional UE5M3 reinterpretation for torch.float8_e4m3fn input scale storage + - Optional C output (pre-activation GLU output) + - Optional fused RHT output (bf16 or NVFP4 + e4m3 block scales) + - GLU activation fusion (SwiGLU / GeGLU) + +Warp assignment (8 epilogue warps, pingpong): + warps 0-3 : ACT warps — TMEM→reg, alpha scale, GLU activation, C/D store + warps 4-7 : RHT store warps — RHT and/or NVFP4 quantization from D SMEM + warp 8 : MMA warp + warp 9 : TMA load warp + warp 10 : Scheduler warp (MoEPersistentTileScheduler) + warp 11 : Bias load warp (optional) + +sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + Validity: tile_info[0] >= 0 (expert_idx == -1 signals end) +""" + +from typing import Type, Tuple, Union, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu import OperandMajorMode +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.typing import Float32, Int32, AddressSpace +from ..moe_persistent_scheduler import ( + MoEPersistentTileScheduler, + MoESchedulerParams, + MoEWorkTileInfo, +) +from ..moe_utils import ( + compute_expert_token_range, + MoEWeightMode, + TensormapWorkspace, + store_tma_desc, +) +from .rht_utils import ( + hadamard_rmem_colwise_fwht, + hadamard_rmem_colwise_fwht_quant, + hadamard_rmem_rowwise_fwht, + load_colwise_pairs_bf16, + HADAMARD_SIZE, +) +from .quant_utils import load_row_bf16, nvfp4_quant_rmem_row +from ..moe_sched_extension import ( + DiscreteWeightScaledGemmSchedExtension, + ContiguousAndConsistentGroupedGemmSchedExtension, +) +from .moe_kernel_helpers import ( + fmin, + fmax, + silu_f32, + silu_f32_geglu_scaled, + compute_grid, + can_implement, +) + + +class BlockScaledMoEGroupedGemmGluHadamardQuantKernel: + """Block-scaled MoE grouped GEMM with GLU activation and Hadamard transform fusion. + + Always uses pingpong epilogue (8 epilogue warps: 4 ACT + 4 RHT-store). + D output is BF16 or NVFP4 (packed e2m1 plus swizzled e4m3 SFD). + + :param sf_vec_size: Scalefactor vector size. + :param mma_tiler_mn: Shape of MMA tile (M, N). + :param cluster_shape_mn: Cluster dimensions (M, N). + :param expert_cnt: Number of experts (compile-time constant). + :param weight_mode: Dense or Discrete weight layout. + :param use_dynamic_sched: Use dynamic tile scheduling. + :param act_func: Activation function ('swiglu', 'geglu', or 'srelu'). + :param enable_bias: Enable bias addition. + """ + + FIX_PAD_SIZE = 256 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + acc_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + rht_quant: bool = False, + d_quant: bool = False, + ) -> bool: + # speical requirements for hadamard fusion + if not use_2cta_instrs or mma_tiler_mn[0] != 256 or mma_tiler_mn[1] != 256: + return False + # NVFP4 quantization: scale-byte rows are stored as one contiguous 8-byte + # store per thread, so f = n/2 must be divisible by 128 (f/16 % 8 == 0; + # same shape gate as the standalone group_rht_cast kernel). + if (rht_quant or d_quant) and (n // 2) % 128 != 0: + return False + return can_implement( + ab_dtype, + sf_dtype, + sf_vec_size, + acc_dtype, + d_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + cd_major, + m_aligned, + fix_pad_size=BlockScaledMoEGroupedGemmGluHadamardQuantKernel.FIX_PAD_SIZE, + ) + + def __init__( + self, + sf_vec_size: int, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + vectorized_f32: bool, + expert_cnt: int, + weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, + use_dynamic_sched: bool = False, + act_func: str = "swiglu", + enable_bias: bool = False, + rht_rowwise: bool = False, + sf_fp8_dtype_override: Optional[str] = None, + glu_alpha: Optional[float] = None, + glu_limit: Optional[float] = None, + ): + mma_tile_m = mma_tiler_mn[0] + if self.FIX_PAD_SIZE % mma_tile_m != 0: + raise ValueError(f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) must be divisible by " f"mma_tiler_mn[0] ({mma_tile_m}).") + if expert_cnt > 1024: + raise ValueError("Expert count > 1024 is not supported.") + if not isinstance(weight_mode, MoEWeightMode): + raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") + + self.sf_vec_size = sf_vec_size + self.expert_cnt = expert_cnt + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.weight_mode = weight_mode + self.use_dynamic_sched = use_dynamic_sched + self.enable_bias = enable_bias + self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None + # RHT dump orientation: False = columnwise (16-token blocks per feature), + # True = rowwise (16-feature blocks per token). Same dump tensor/path either way. + self.rht_rowwise = rht_rowwise + + # Always use pingpong epilogue for Hadamard + self.epilogue_pingpong = True + # Always delay TMA store acquire sync for Hadamard + self.delay_tma_store_acquire_sync = True + + self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.occupancy = 1 + self.threads_per_warp = 32 + + # Warp assignments: 8 epilogue warps (4 ACT + 4 RHT-store) + self.epilog_warp_id = (0, 1, 2, 3, 4, 5, 6, 7) + self.epilog_act_warp_id = (0, 1, 2, 3) + self.epilog_rht_store_warp_id = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.tma_warp_id = 9 + self.sched_warp_id = 10 + self.bias_load_warp_id = 11 if enable_bias else None + + self.epilogue_warp_group_size = len(self.epilog_act_warp_id) # = 4 + + all_warps = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id, self.sched_warp_id] + warps_wo_sched = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id] + if enable_bias: + all_warps.append(self.bias_load_warp_id) + warps_wo_sched.append(self.bias_load_warp_id) + self.threads_per_cta = self.threads_per_warp * len(all_warps) + self.threads_wo_sched = self.threads_per_warp * len(warps_wo_sched) + + # Named barriers + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + tmem_alloc_warp_ids = self.epilog_act_warp_id + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *tmem_alloc_warp_ids)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + # Pingpong barriers (group 0 = ACT warps, group 1 = RHT store warps) + self.epilog_sync_barrier_group0 = pipeline.NamedBarrier( + barrier_id=5, + num_threads=32 * self.epilogue_warp_group_size, + ) + self.epilog_sync_barrier_group1 = pipeline.NamedBarrier( + barrier_id=6, + num_threads=32 * self.epilogue_warp_group_size, + ) + + self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + SM100_TMEM_CAPACITY_COLUMNS = 512 + self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + + self.vectorized_f32 = vectorized_f32 + + self.act_func = act_func + if act_func not in ["swiglu", "geglu", "srelu"]: + raise ValueError(f"Invalid activation function: {act_func}") + + self.glu_alpha = glu_alpha + self.glu_limit = glu_limit + + def _setup_attributes(self): + """Set up configurations dependent on GEMM inputs (called inside __call__).""" + + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + d_tile_n = self.mma_inst_shape_mn[1] if self.act_func == "srelu" else self.mma_inst_shape_mn[1] // 2 + self.mma_tiler_d = ( + self.mma_inst_shape_mn[0], + d_tile_n, + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk_d = ( + self.mma_tiler_d[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_d[1], + self.mma_tiler_d[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = (128, 32) + self.epi_tile_cnt = ( + self.cta_tile_shape_mnk_d[0] // self.epi_tile[0], + self.cta_tile_shape_mnk_d[1] // self.epi_tile[1], + ) + self.epi_tile_c = self.epi_tile if self.act_func == "srelu" else (128, 64) + + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_d_stage, + self.num_tile_stage, + self.num_bias_stage, + self.num_pingpong_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.epi_tile_c, + self.c_dtype, + self.c_layout, + self.d_dtype, + self.d_layout, + self.sf_dtype, + self.sf_vec_size, + self.num_smem_capacity, + self.occupancy, + self.bias_dtype if self.enable_bias else None, + self.rht_dtype if self.generate_rht else None, + self.rht_quant, + self.d_quant, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile_c, + self.num_c_stage, + ) + # sD staging keeps bf16 when D itself is quantized to NVFP4 — the RHT warps + # consume bf16 D rows; the fp4 TMA source is the separate sDq buffer whose + # layout follows the gmem dtype. Without quantization the two layouts coincide. + self.d_smem_dtype = cutlass.BFloat16 if self.d_quant else self.d_dtype + self.d_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_smem_dtype, + self.d_layout, + self.epi_tile, + self.num_d_stage, + ) + self.d_tma_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_dtype, + self.d_layout, + self.epi_tile, + self.num_d_stage, + ) + + if self.enable_bias: + self.bias_smem_layout_staged = cute.make_layout( + (self.mma_tiler[1], self.num_bias_stage), + stride=(1, self.mma_tiler[1]), + ) + else: + self.bias_smem_layout_staged = cute.make_layout((1, 1)) + + self.overlapping_accum = self.num_acc_stage == 1 and self.mma_tiler[1] == 256 + + sf_atom_mn = 32 + self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k + self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + self.num_reserved_tmem_cols = self.num_sf_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_reserved_tmem_cols + ) + + self.epi_tile_n_required = cute.size(self.epi_tile[1]) if self.act_func == "srelu" else 2 * cute.size(self.epi_tile[1]) + self.iter_acc_early_release_in_epilogue = (self.num_reserved_tmem_cols + self.epi_tile_n_required - 1) // self.epi_tile_n_required - 1 + if self.act_func != "srelu": + self.iter_acc_early_release_in_epilogue = self.iter_acc_early_release_in_epilogue * 2 + + def get_desc_workspace_bytes(self) -> int: + """Return descriptor workspace size in bytes.""" + if self.weight_mode == MoEWeightMode.DISCRETE: + from ..moe_utils import DiscreteWeightTensormapConstructor + + return DiscreteWeightTensormapConstructor.get_workspace_size(self.expert_cnt) + return 0 + + def get_workspace_bytes(self) -> int: + """Return total workspace size in bytes.""" + desc_workspace_bytes = self.get_desc_workspace_bytes() + dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + return desc_workspace_bytes + dynamic_sched_bytes + + @cute.jit + def _get_sched_counter_ptr(self, workspace_ptr): + counter_addr = workspace_ptr.toint() + self.get_desc_workspace_bytes() + return cute.make_ptr( + cutlass.Int32, + counter_addr, + AddressSpace.gmem, + assumed_align=4, + ) + + @cute.kernel + def helper_kernel( + self, + ptrs_b: cute.Pointer, + ptrs_sfb: cute.Pointer, + n: Int32, + k: Int32, + b_stride_size: cutlass.Int64, + b_major_mode: cutlass.Constexpr, + workspace_ptr, + tiled_mma_arg: cute.TiledMma, + tiled_mma_sfb_arg: cute.TiledMma, + b_smem_layout_arg, + sfb_smem_layout_arg, + cluster_layout_vmnk_shape_arg: cutlass.Constexpr, + cluster_layout_sfb_vmnk_shape_arg: cutlass.Constexpr, + ): + """Pre-main-kernel: build per-expert TMA descriptors (discrete mode) and/or reset sched counter.""" + expert_idx = cute.arch.block_idx()[0] + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + b_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma_arg.thr_id) + sfb_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma_arg.thr_id) + + # Read per-expert base addresses from the pointer arrays + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + sfb_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + stride_n = b_stride_size + stride_k = c1_64 + else: + stride_n = c1_64 + stride_k = b_stride_size + + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + b_expert = cute.make_tensor( + b_ptr, + cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), + ) + tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B( + b_tma_op_arg, + b_expert, + b_smem_layout_arg, + self.mma_tiler, + tiled_mma_arg, + cluster_layout_vmnk_shape_arg, + ) + + workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"]) + store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) + + sfb_ptr_val = sfb_ptr_tensor[expert_idx] + sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + sfb_expert = cute.make_tensor(sfb_ptr, sfb_layout) + tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B( + sfb_tma_op_arg, + sfb_expert, + sfb_smem_layout_arg, + self.mma_tiler_sfb, + tiled_mma_sfb_arg, + cluster_layout_sfb_vmnk_shape_arg, + internal_type=cutlass.Uint64, + ) + store_tma_desc(tma_atom_sfb, workspace.get_ptr("sfb", expert_idx)) + + if cutlass.const_expr(self.use_dynamic_sched): + if expert_idx == cutlass.Int32(0): + sched_counter = cute.make_tensor( + self._get_sched_counter_ptr(workspace_ptr), + cute.make_layout(1), + ) + sched_counter[0] = cutlass.Int32(0) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b, # Dense: cute.Tensor (N,K,L) | Discrete: cute.Pointer to int64[] + sfa: cute.Tensor, + sfb, # Dense: cute.Tensor | Discrete: cute.Pointer to int64[] + n: Int32, # Ignored for dense mode + k: Int32, # Ignored for dense mode + b_stride_size: cutlass.Int64, # Ignored for dense mode + b_major_mode: cutlass.Constexpr, # Ignored for dense mode + workspace_ptr, + c: cute.Tensor, + d: cute.Tensor, # post-GLU output (bf16, or NVFP4 packed e2m1 with sfd) + sfd: Optional[cute.Tensor], # NVFP4 D block scales (e4m3, swizzled SF layout); required iff d is NVFP4 + rht: Optional[cute.Tensor], # RHT output (bf16 or NVFP4, D layout); None => off + sfrht: Optional[cute.Tensor], # NVFP4 RHT block scales (e4m3, swizzled SF layout); required iff rht is NVFP4 + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + prob: cute.Tensor, + bias: Optional[cute.Tensor], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + linear_offset: cutlass.Float32 = 0.0, + norm_const: cutlass.Float32 = 1.0, # D NVFP4 global encode scale: 2688/global_amax, or 1.0 + rht_norm_const: cutlass.Float32 = 1.0, # RHT NVFP4 global encode scale: 2688/global_amax, or 1.0 + ): + """Execute the MoE GEMM + GLU + Hadamard kernel. + + Dense mode: ``b`` and ``sfb`` are 3-D cute.Tensor (N, K, L). + Discrete mode: ``b`` and ``sfb`` are cute.Pointer to device int64[] + arrays of per-expert base addresses. + """ + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = a.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.d_dtype: Type[cutlass.Numeric] = d.element_type + # UE5M3 has no torch dtype, so callers provide e4m3 storage and ask the + # kernel to reinterpret the scale bytes at compile time. + self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override if cutlass.const_expr(self.sf_dtype_override is not None) else sfa.element_type + self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + self.d_layout = utils.LayoutEnum.from_tensor(d) + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + else: + self.b_major_mode = b_major_mode + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # ---- Output / dump modes (derived from dtypes; before _setup_attributes so + # the smem/stage accounting sees them) ---- + self.generate_rht = rht is not None + self.generate_sfd = sfd is not None + self.generate_sfrht = sfrht is not None + self.d_quant = self.d_dtype == cutlass.Float4E2M1FN + self.rht_dtype = rht.element_type if cutlass.const_expr(self.generate_rht) else self.d_dtype + self.rht_quant = self.generate_rht and self.rht_dtype == cutlass.Float4E2M1FN + # The rht STORE has no orientation logic anywhere in this kernel; only the + # SCALE tensor is orientation-aware below: rowwise uses swizzled SF for + # logical (m, f), while colwise uses swizzled SF for logical (f, m). + # RHT warps run for the RHT output and/or the D quantization (both read sD). + # When neither is on (plain bf16 mode) the warpgroup is COMPILED OUT: no + # pingpong, no ACT<->RHT barriers, and the tile-info pipeline's consumer + # count shrinks by the 4 RHT warps (forget that and the scheduler deadlocks + # after num_tile_stage tiles). + self.run_rht = self.generate_rht or self.d_quant + if cutlass.const_expr(not self.run_rht): + self.threads_wo_sched = self.threads_per_cta - self.threads_per_warp - self.threads_per_warp * len(self.epilog_rht_store_warp_id) + if cutlass.const_expr(self.d_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): + raise ValueError(f"d dtype must be BFloat16 or Float4E2M1FN, got {self.d_dtype}") + if cutlass.const_expr(self.d_quant != self.generate_sfd): + raise ValueError("NVFP4 d and sfd must be passed together") + if cutlass.const_expr(self.generate_rht and self.rht_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): + raise ValueError(f"rht dtype must be BFloat16 or Float4E2M1FN, got {self.rht_dtype}") + if cutlass.const_expr(self.rht_quant != self.generate_sfrht): + raise ValueError("NVFP4 rht and sfrht must be passed together") + if cutlass.const_expr(self.generate_sfd and sfd.element_type != cutlass.Float8E4M3FN): + raise ValueError("sfd element type must be Float8E4M3FN") + if cutlass.const_expr(self.generate_sfrht and sfrht.element_type != cutlass.Float8E4M3FN): + raise ValueError("sfrht element type must be Float8E4M3FN") + if cutlass.const_expr((self.d_quant or self.rht_quant) and self.act_func == "srelu"): + raise ValueError("NVFP4 quantization assumes the GLU subtile pair-step (act_func != srelu)") + + self._setup_attributes() + + # ---- B / SFB setup (mode-dependent) ---- + b_from_call_arg = b + sfb_from_call_arg = sfb + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b.shape, self.sf_vec_size) + sfb = cute.make_tensor(sfb.iterator, sfb_layout) + else: + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + b_template_stride = (b_stride_size, c1_64, c0) + else: + b_template_stride = (c1_64, b_stride_size, c0) + b_template_layout = cute.make_layout((n, k, c1), stride=b_template_stride) + b_ptr_typed = cute.make_ptr(self.b_dtype, b.toint(), AddressSpace.gmem, assumed_align=16) + b = cute.make_tensor(b_ptr_typed, b_template_layout) + + sfb_ptr_typed = cute.make_ptr(self.sf_dtype, sfb.toint(), AddressSpace.gmem, assumed_align=16) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + sfb = cute.make_tensor(sfb_ptr_typed, sfb_layout) + + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, self.sf_vec_size) + sfa = cute.make_tensor(sfa.iterator, sfa_layout) + + # Dump staging dtype follows the dump tensor's element type (fp4: 2KB/stage vs + # 8KB bf16); the layout is CONSTRUCTED f-major like every other epilogue + # output (never derived from the rht gmem tensor) — all FWHT store paths + # pack along features. + self.rht_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.rht_dtype, + utils.LayoutEnum.ROW_MAJOR, + self.epi_tile, + self.num_d_stage, + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA load A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # TMA load SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb = cute.make_tensor( + tma_tensor_sfb.iterator, + cute.make_layout(new_shape, stride=new_stride), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size + + # TMA store C + c_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + c_smem_layout, + self.epi_tile_c, + ) + + # TMA store D (gmem-dtype staging: sDq when D is quantized, sD otherwise) + d_smem_layout = cute.slice_(self.d_tma_smem_layout_staged, (None, None, 0)) + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + d, + d_smem_layout, + self.epi_tile, + ) + + # TMA store RHT — identical tile to D; smem layout follows the RHT output + # element type (== d_smem_layout for bf16, packed fp4 layout in quant mode). + if cutlass.const_expr(self.generate_rht): + rht_smem_layout = cute.slice_(self.rht_smem_layout_staged, (None, None, 0)) + tma_atom_rht, tma_tensor_rht = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + rht, + rht_smem_layout, + self.epi_tile, + ) + else: + tma_atom_rht, tma_tensor_rht = None, None + + # ---- Helper kernel (discrete TMA desc init + dynamic sched counter reset) ---- + _need_helper = cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE or self.use_dynamic_sched) + if cutlass.const_expr(_need_helper): + _helper_grid_x = self.expert_cnt if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else 1 + _helper_args = ( + b_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + sfb_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + n if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + k if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + b_stride_size if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int64(0), + b_major_mode if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else self.b_major_mode, + workspace_ptr, + tiled_mma, + tiled_mma_sfb, + b_smem_layout, + sfb_smem_layout, + self.cluster_layout_vmnk.shape, + self.cluster_layout_sfb_vmnk.shape, + ) + self.helper_kernel(*_helper_args).launch( + grid=(_helper_grid_x, 1, 1), + block=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # ---- Grid computation via MoE scheduler ---- + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + b_n, b_k, b_l = cute.shape(b) + sched_expert_shape = (self.expert_cnt, b_n, b_k) + else: + sched_expert_shape = (self.expert_cnt, n, k) + + sched_params = MoESchedulerParams( + scenario="2Dx3D", + expert_shape=sched_expert_shape, + cta_tile_shape_mnk=self.cta_tile_shape_mnk, + cluster_shape_mn=self.cluster_shape_mn, + use_dynamic_sched=self.use_dynamic_sched, + ) + self.sched_params, grid = compute_grid( + sched_params, + max_active_clusters, + self.use_2cta_instrs, + ) + + self.buffer_align_bytes = 1024 + + # ---- Shared storage ---- + SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_tile_stage, self.use_dynamic_sched) + + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + pingpong_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_pingpong_stage * 2] + if cutlass.const_expr(self.enable_bias): + bias_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_bias_stage * 2] + scheduler: SchedulerStorage + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sD: cute.struct.Align[ + cute.struct.MemRange[self.d_smem_dtype, cute.cosize(self.d_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.d_quant): + # NVFP4 D: packed-fp4 TMA staging + per-thread scale-byte staging. + sDq: cute.struct.Align[ + cute.struct.MemRange[self.d_dtype, cute.cosize(self.d_tma_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSfd: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float8E4M3FN, + self.threads_per_warp * len(self.epilog_rht_store_warp_id) * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), + ], + 16, + ] + if cutlass.const_expr(self.generate_rht): + sRht: cute.struct.Align[ + cute.struct.MemRange[self.rht_dtype, cute.cosize(self.rht_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.rht_quant): + # NVFP4 RHT scale-byte staging, one contiguous row store per thread + # per tile: rowwise rows are (thread=token, cta_tile_n/16 scales); + # colwise rows are (feature-in-tile, 128-token-tile/16). + # Same (128, 8) buffer either way. + sSfRht: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float8E4M3FN, + self.threads_per_warp * len(self.epilog_rht_store_warp_id) * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), + ], + 16, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.enable_bias): + sBias: cute.struct.Align[ + cute.struct.MemRange[self.bias_dtype, cute.cosize(self.bias_smem_layout_staged)], + 16, + ] + + self.shared_storage = SharedStorage + + # Launch main kernel + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + tma_atom_d, + tma_tensor_d, + sfd, + tma_atom_rht, + tma_tensor_rht, + sfrht, + padded_offsets, + alpha, + bias, + prob, + workspace_ptr, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.c_smem_layout_staged, + self.d_smem_layout_staged, + self.d_tma_smem_layout_staged, + self.rht_smem_layout_staged, + self.bias_smem_layout_staged, + self.epi_tile, + self.sched_params, + epilogue_op, + linear_offset, + norm_const, + rht_norm_const, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + max_number_threads=[self.threads_per_cta, 1, 1], + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @cute.jit + def _make_extension(self, workspace_ptr): + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + desc_workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"]) + return DiscreteWeightScaledGemmSchedExtension( + tensormap_ctor=desc_workspace, + sf_vec_size=self.sf_vec_size, + ) + else: + return ContiguousAndConsistentGroupedGemmSchedExtension( + sf_vec_size=self.sf_vec_size, + ) + + @cute.jit + def store_swizzled_sf_row(self, sf_tensor: cute.Tensor, logical_row, sf_col_base, sSf: cute.Tensor, tidx): + """Store one logical scale row into M32x4xrm_K4xrk_L SF layout.""" + row_m0 = logical_row % 32 + row_m1 = (logical_row // 32) % 4 + row_m2 = logical_row // 128 + num_sf = self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE + for vi in cutlass.range_constexpr(num_sf): + sf_col = sf_col_base + vi + sf_tensor[(row_m0, row_m1, row_m2, sf_col % 4, sf_col // 4, 0)] = sSf[(tidx, vi)] + + def mainloop_s2t_copy_and_partition(self, sSF, tSF): + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + @cute.jit + def store_c( + self, + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tTR_rAcc_up, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + prev_subtile_idx, + real_subtile_idx, + ): + c_buffer = prev_subtile_idx % self.num_c_stage + tRS_rC.store(tTR_rAcc.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)]) + tRS_rC.store(tTR_rAcc_up.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 1, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group0.arrive_and_wait() + if warp_idx == self.epilog_act_warp_id[0]: + cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)]) + c_pipeline.producer_commit() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + c_pipeline.producer_acquire() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + self.epilog_sync_barrier_group0.arrive_and_wait() + + @cute.jit + def store_c_unary( + self, + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + prev_subtile_idx, + real_subtile_idx, + ): + c_buffer = prev_subtile_idx % self.num_c_stage + tRS_rC.store(tTR_rAcc.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group0.arrive_and_wait() + if warp_idx == self.epilog_act_warp_id[0]: + cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)]) + c_pipeline.producer_commit() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + c_pipeline.producer_acquire() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + self.epilog_sync_barrier_group0.arrive_and_wait() + + @cute.jit + def geglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset=1.0): + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + scaled_gate_0, scaled_gate_1 = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (1.702, 1.702), + rnd="rn", + ftz=False, + ) + tCompute_log2e = cute.arch.mul_packed_f32x2( + (scaled_gate_0, scaled_gate_1), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.add_packed_f32x2( + (cute.math.exp2(tCompute_log2e[0], fastmath=True), cute.math.exp2(tCompute_log2e[1], fastmath=True)), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate[i], acc_vec_gate[i + 1]), + rnd="rn", + ftz=False, + ) + up0, up1 = cute.arch.add_packed_f32x2( + (linear_offset, linear_offset), + (acc_vec_up[i], acc_vec_up[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (up0, up1), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (self.glu_alpha, self.glu_alpha), + rnd="rn", + ftz=False, + ) + else: + # GeGlu Unpacked Version + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = (acc_vec_up[i] + linear_offset) * silu_f32_geglu_scaled(acc_vec_gate[i], fastmath=True) + tCompute[i] = tCompute[i] * mProb + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + tCompute[i] = tCompute[i] * self.glu_alpha + # + 0.0 canonicalizes -0 -> +0 (negative up x saturated-to-zero sigmoid). SCALAR + # on purpose: the f32 immediate makes it a free FADD with RZ, while f32x2 has no + # immediate form and a live (0, 0) register pair costs regs/spills. + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = tCompute[i] + cutlass.Float32(0.0) + + @cute.jit + def swiglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb): + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + tCompute_log2e = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.add_packed_f32x2( + (cute.math.exp2(tCompute_log2e[0], fastmath=True), cute.math.exp2(tCompute_log2e[1], fastmath=True)), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate[i], acc_vec_gate[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_up[i], acc_vec_up[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (self.glu_alpha, self.glu_alpha), + rnd="rn", + ftz=False, + ) + else: + # SwiGlu Unpacked Version + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = acc_vec_up[i] * silu_f32(acc_vec_gate[i], fastmath=True) + tCompute[i] = tCompute[i] * mProb + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + tCompute[i] = tCompute[i] * self.glu_alpha + # + 0.0 canonicalizes -0 -> +0 (negative up x saturated-to-zero sigmoid). SCALAR + # on purpose: the f32 immediate makes it a free FADD with RZ, while f32x2 has no + # immediate form and a live (0, 0) register pair costs regs/spills. + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = tCompute[i] + cutlass.Float32(0.0) + + @cute.jit + def srelu_act(self, tCompute, acc_vec, mProb): + acc_relu = cute.where(acc_vec > 0, acc_vec, cute.full_like(acc_vec, 0)) + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (acc_relu[i], acc_relu[i + 1]), + (acc_relu[i], acc_relu[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = acc_relu[i] * acc_relu[i] * mProb + + def epilog_tmem_copy_and_partition(self, tidx, tAcc, gD_mnl, epi_tile, use_2cta_instrs): + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, + self.d_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + gD_mnl_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tTR_gC = thr_copy_t2r.partition_D(gD_mnl_epi) + tTR_rAcc_gate = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + tTR_rAcc_up = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc_gate, tTR_rAcc_up + + def epilog_smem_copy_and_partition(self, tiled_copy_t2r, tTR_rC, tidx, sD): + copy_atom_r2s = sm100_utils.get_smem_store_op(self.d_layout, sD.element_type, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = thr_copy_r2s.partition_D(sD) + tRS_rD = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rD, tRS_sD + + def epilog_gmem_copy_and_partition(self, tidx, atom, gD_mnl, epi_tile, sD): + gD_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tma_atom_d = atom + sD_for_tma_partition = cute.group_modes(sD, 0, 2) + gD_for_tma_partition = cute.group_modes(gD_epi, 0, 2) + bSG_sD, bSG_gD = cpasync.tma_partition( + tma_atom_d, + 0, + cute.make_layout(1), + sD_for_tma_partition, + gD_for_tma_partition, + ) + return tma_atom_d, bSG_sD, bSG_gD + + @staticmethod + def _compute_stages( + tiled_mma, + mma_tiler_mnk, + a_dtype, + b_dtype, + epi_tile, + epi_tile_c, + c_dtype, + c_layout, + d_dtype, + d_layout, + sf_dtype, + sf_vec_size, + num_smem_capacity, + occupancy, + bias_dtype, + rht_dtype, # RHT output dtype (None => no RHT output) + rht_quant, # RHT output is NVFP4 (adds sfrht scale-byte staging) + d_quant, # D output is NVFP4 (sD staging stays bf16; adds fp4 + sfd staging) + ): + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + num_c_stage = 1 + num_d_stage = 1 + num_tile_stage = 2 + num_pingpong_stage = mma_tiler_mnk[1] // epi_tile_c[1] + + a_smem_layout_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1) + b_smem_layout_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1) + sfa_smem_layout_one = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, 1) + sfb_smem_layout_one = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, 1) + c_smem_layout_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile_c, 1) + d_smem_layout_one = sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_one) + ) + mbar_helpers_bytes = 1024 + + # One e4m3 scale byte per (1,16) feature block, one row per thread (128 threads). + quant_sf_bytes = 128 * ((mma_tiler_mnk[1] // 2) // HADAMARD_SIZE) + + # sInfo is in SchedulerStorage, not here, so use 4-int sInfo + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes = cute.size_in_bytes(c_dtype, c_smem_layout_one) * num_c_stage + d_bytes = cute.size_in_bytes(d_dtype, d_smem_layout_one) * num_d_stage + + if d_quant: + # sD staging stays bf16 (the RHT/quant warps' source); the gmem-dtype (fp4) + # staging above becomes the TMA source, and sfd rows are staged per thread. + bf16_smem_layout_one = sm100_utils.make_smem_layout_epi(cutlass.BFloat16, d_layout, epi_tile, 1) + d_bytes += cute.size_in_bytes(cutlass.BFloat16, bf16_smem_layout_one) * num_d_stage + d_bytes += quant_sf_bytes + + rht_bytes = 0 + if rht_dtype is not None: + rht_smem_layout_one = sm100_utils.make_smem_layout_epi(rht_dtype, d_layout, epi_tile, 1) + rht_bytes = cute.size_in_bytes(rht_dtype, rht_smem_layout_one) * num_d_stage + if rht_quant: + rht_bytes += quant_sf_bytes + + if bias_dtype is not None: + num_bias_stage = 2 + bias_bytes = mma_tiler_mnk[1] * num_bias_stage * (bias_dtype.width // 8) + else: + num_bias_stage = 0 + bias_bytes = 0 + + epi_bytes = c_bytes + d_bytes + rht_bytes + bias_bytes + + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage + + return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage, num_pingpong_stage + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tma_atom_d: cute.CopyAtom, + mD_mnl: cute.Tensor, + mSfd_mnl: Optional[cute.Tensor], + tma_atom_rht: Optional[cute.CopyAtom], + mRht_mnl: Optional[cute.Tensor], + mSfRht_mnl: Optional[cute.Tensor], + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + mBias_nl: Optional[cute.Tensor], + prob: cute.Tensor, + workspace_ptr, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_tma_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + rht_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + bias_smem_layout_staged: cute.Layout, + epi_tile: cute.Tile, + sched_params: MoESchedulerParams, + epilogue_op: cutlass.Constexpr, + linear_offset: cutlass.Float32 = 0.0, + norm_const: cutlass.Float32 = 1.0, + rht_norm_const: cutlass.Float32 = 1.0, + ): + """GPU device kernel: MoE persistent GEMM + GLU + Hadamard (pingpong epilogue).""" + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + total_token = padded_offsets[self.expert_cnt - 1] + + # Prefetch TMA descriptors + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_sfa) + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + cpasync.prefetch_descriptor(tma_atom_d) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # CTA coordinates + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + # Shared memory allocation + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sched_storage = storage.scheduler + + # AB pipeline + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # ACC pipeline + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_act_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Pingpong pipeline + pingpong_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_act_warp_id) * self.threads_per_warp) + pingpong_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_rht_store_warp_id) * self.threads_per_warp) + pingpong_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.pingpong_mbar_ptr.data_ptr(), + num_stages=self.num_pingpong_stage, + producer_group=pingpong_producer_group, + consumer_group=pingpong_consumer_group, + ) + + # Tile info pipeline (uses SchedulerStorage's barrier) + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=sched_storage.tile_info_mbar.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + # MoE persistent tile scheduler + scheduler = MoEPersistentTileScheduler.create( + sched_params, + padded_offsets, + cute.arch.block_idx(), + cute.arch.grid_dim(), + counter_ptr=self._get_sched_counter_ptr(workspace_ptr), + sched_storage=sched_storage, + ) + scheduler.internal_init() + + # Bias pipeline + if cutlass.const_expr(self.enable_bias): + bias_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp, + ) + bias_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_act_warp_id), + ) + bias_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.bias_mbar_ptr.data_ptr(), + num_stages=self.num_bias_stage, + producer_group=bias_pipeline_producer_group, + consumer_group=bias_pipeline_consumer_group, + ) + sBias = storage.sBias.get_tensor(bias_smem_layout_staged) + gBias_nl = cute.local_tile(mBias_nl, cute.slice_(self.mma_tiler[:2], (0, None)), (None, None)) + + # TMEM allocator + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_act_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # SMEM tensors + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + sD = storage.sD.get_tensor(d_smem_layout_staged.outer, swizzle=d_smem_layout_staged.inner) + _num_sf_per_tile = self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE + _sf_stage_layout = cute.make_layout( + (self.threads_per_warp * len(self.epilog_rht_store_warp_id), _num_sf_per_tile), + stride=(_num_sf_per_tile, 1), + ) + if cutlass.const_expr(self.d_quant): + sDq = storage.sDq.get_tensor(d_tma_smem_layout_staged.outer, swizzle=d_tma_smem_layout_staged.inner) + sSfd = storage.sSfd.get_tensor(_sf_stage_layout) + if cutlass.const_expr(self.generate_rht): + sRht = storage.sRht.get_tensor(rht_smem_layout_staged.outer, swizzle=rht_smem_layout_staged.inner) + if cutlass.const_expr(self.rht_quant): + sSfRht = storage.sSfRht.get_tensor(_sf_stage_layout) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # sInfo from SchedulerStorage + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = sched_storage.sInfo.get_tensor(info_layout) + + # Multicast masks + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1) + + # MMA fragments + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage_overlapped)) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (256 - self.num_reserved_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # Cluster wait / CTA sync + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + self.cta_sync_barrier.arrive_and_wait() + + if total_token <= 0: + cute.arch.nvvm.exit() + + # --------------------------------------------------------------- + # Specialized Scheduler warp (MoEPersistentTileScheduler) + # --------------------------------------------------------------- + if warp_idx == self.sched_warp_id: + work_tile_info = scheduler.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tile_stage) + + while work_tile_info.is_valid_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = work_tile_info.expert_idx + sInfo[(1, tile_info_producer_state.index)] = work_tile_info.tile_m_idx + sInfo[(2, tile_info_producer_state.index)] = work_tile_info.tile_n_idx + sInfo[(3, tile_info_producer_state.index)] = work_tile_info.k_tile_cnt + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + work_tile_info = scheduler.advance_to_next_work() + + # Send invalid signal: expert_idx = -1 + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cutlass.Int32(-1) + sInfo[(1, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(2, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # --------------------------------------------------------------- + # Specialized TMA load warp + # --------------------------------------------------------------- + if warp_idx == self.tma_warp_id: + ext = self._make_extension(workspace_ptr) + + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + k_tile_cnt = work_tile_info.k_tile_cnt + ext.update_expert_info(padded_offsets, work_tile_info.expert_idx) + + real_a, _ = ext.get_gmem_tensor("a", mA_mkl, padded_offsets, work_tile_info) + real_b, desc_ptr_b = ext.get_gmem_tensor("b", mB_nkl, padded_offsets, work_tile_info) + real_sfa, _ = ext.get_gmem_tensor("sfa", mSFA_mkl, padded_offsets, work_tile_info) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor("sfb", mSFB_nkl, padded_offsets, work_tile_info) + + gA_mkl = cute.local_tile(real_a, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gB_nkl = cute.local_tile(real_b, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + gSFA_mkl = cute.local_tile(real_sfa, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gSFB_nkl = cute.local_tile(real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None)) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + sfa_cta_layout = a_cta_layout + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + mma_tile_coord_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape) + mma_tile_coord_n = work_tile_info.tile_n_idx + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, mma_tile_coord_n, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)] + slice_n = mma_tile_coord_n + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_n // 2 + tBgSFB_slice = tBgSFB[(None, slice_n, None, 0)] + + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + tAgA_k = tAgA_slice[(None, ab_producer_state.count)] + tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + tAgSFA_k = tAgSFA_slice[(None, ab_producer_state.count)] + tBgSFB_k = tBgSFB_slice[(None, ab_producer_state.count)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + tAsSFA_pipe = tAsSFA[(None, ab_producer_state.index)] + tBsSFB_pipe = tBsSFB[(None, ab_producer_state.index)] + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + ab_producer_state_next = ab_producer_state.clone() + ab_producer_state_next.advance() + if ab_producer_state_next.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state_next) + else: + peek_ab_empty_status = cutlass.Boolean(1) + + cute.copy(tma_atom_a, tAgA_k, tAsA_pipe, tma_bar_ptr=tma_bar, mcast_mask=a_full_mcast_mask) + cute.copy(tma_atom_b, tBgB_k, tBsB_pipe, tma_bar_ptr=tma_bar, mcast_mask=b_full_mcast_mask, tma_desc_ptr=desc_ptr_b) + cute.copy(tma_atom_sfa, tAgSFA_k, tAsSFA_pipe, tma_bar_ptr=tma_bar, mcast_mask=sfa_full_mcast_mask) + cute.copy(tma_atom_sfb, tBgSFB_k, tBsSFB_pipe, tma_bar_ptr=tma_bar, mcast_mask=sfb_full_mcast_mask, tma_desc_ptr=desc_ptr_sfb) + + ab_producer_state.advance() + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + ab_pipeline.producer_tail(ab_producer_state) + + # --------------------------------------------------------------- + # Specialized MMA warp + # --------------------------------------------------------------- + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acd_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + k_tile_cnt = tile_info[3] + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + acd_producer_state.reset_count() + peek_acc_empty_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state) + + mma_tile_coord_mnl = ( + tile_info[1] // cute.size(tiled_mma.thr_id.shape), + tile_info[2], + tile_info[0], + ) + + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acd_producer_state.phase ^ 1 + else: + acc_stage_index = acd_producer_state.index + + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + + if is_leader_cta: + acc_pipeline.producer_acquire(acd_producer_state, peek_acc_empty_status) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + s2t_stage_coord = (None, None, None, None, ab_consumer_state.index) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_staged, tCtSFA_compact_s2t) + cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_staged, tCtSFB_compact_s2t) + + num_kblocks = cute.size(tCrA, mode=[2]) + ab_consumer_state_next = ab_consumer_state.clone() + ab_consumer_state_next.advance() + if ab_consumer_state_next.count < k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state_next) + + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = (None, None, kblock_idx, ab_consumer_state.index) + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator) + cute.gemm(tiled_mma, tCtAcc, tCrA[kblock_coord], tCrB[kblock_coord], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state = ab_consumer_state_next + + if is_leader_cta: + acc_pipeline.producer_commit(acd_producer_state) + + acd_producer_state.advance() + if acd_producer_state.count < k_tile_cnt: + if is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + acc_pipeline.producer_tail(acd_producer_state) + + # --------------------------------------------------------------- + # Specialized bias load warp + # --------------------------------------------------------------- + if cutlass.const_expr(self.enable_bias): + if warp_idx == self.bias_load_warp_id and total_token > 0: + bias_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_bias_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + bias_elems_per_thread = 128 // self.bias_dtype.width + bias_g2s_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + self.bias_dtype, + num_bits_per_copy=128, + ) + bias_g2s_tiled = cute.make_tiled_copy_tv( + bias_g2s_atom, + cute.make_layout((self.threads_per_warp,)), + cute.make_layout((bias_elems_per_thread,)), + ) + thr_bias_g2s = bias_g2s_tiled.get_slice(cute.arch.lane_idx()) + tBs_sBias = thr_bias_g2s.partition_D(sBias) + + bias_n_total = mBias_nl.shape[0] + tBpBias = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Boolean) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + bias_producer_state.reset_count() + mma_n_coord = tile_info[2] + expert_idx = tile_info[0] + gBias_tile = gBias_nl[(None, mma_n_coord, expert_idx)] + tBs_gBias = thr_bias_g2s.partition_S(gBias_tile) + tBpBias[0] = mma_n_coord * self.mma_tiler[1] + cute.arch.lane_idx() * bias_elems_per_thread < bias_n_total + bias_pipeline.producer_acquire(bias_producer_state) + cute.copy( + bias_g2s_tiled, + tBs_gBias[(None, 0)], + tBs_sBias[(None, 0, bias_producer_state.index)], + pred=tBpBias, + ) + bias_pipeline.producer_commit(bias_producer_state) + bias_producer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + bias_pipeline.producer_tail(bias_producer_state) + + # --------------------------------------------------------------- + # Specialized ACT epilogue warps (0-3): TMEM→regs, alpha, GLU activation, + # C store, hadamard_in + # --------------------------------------------------------------- + if warp_idx < self.epilog_rht_store_warp_id[0] and total_token > 0: + epi_tidx = tidx + + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue (shape-only via mD_mnl for invariant setup) + # + thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v) + gD_mnl_shape = cute.local_tile(mD_mnl, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD_shape = thr_mma_epi.partition_C(gD_mnl_shape) + + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc_gate, + tTR_rAcc_up, + ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgD_shape, epi_tile, use_2cta_instrs) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rC, epi_tidx, sC) + + tTR_rD = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.d_smem_dtype) + tiled_copy_r2s_d, tRS_rD, tRS_sD = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rD, epi_tidx, sD) + + # + # Create per-expert extension (for C/prob tensors inside tile loop) + # + epi_ext = self._make_extension(workspace_ptr) + + # + # Persistent tile scheduling state + # + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + + # + # Pingpong producer state + # + pingpong_act_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_pingpong_stage) + + # Threads/warps participating in TMA store pipeline for C + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_act_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + # NVFP4 D: the fp4 TMA store is issued by the RHT warps (which quantize + # sD); the ACT warps only stage bf16 into sD. + if cutlass.const_expr(not self.d_quant): + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_act_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=d_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_bias_stage) + bias_s2r_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.bias_dtype, num_bits_per_copy=128) + tTR_rBias_gate = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + tTR_rBias_up = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + num_prev_subtiles = cutlass.Int32(0) + num_prev_d_subtiles = cutlass.Int32(0) + while is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + mma_tile_coord_mnl = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx, + cutlass.Int32(0), + ) + + expert_idx = epi_work_tile_info.expert_idx + alpha_val = alpha[expert_idx] + epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx) + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state.reset_count() + bias_pipeline.consumer_wait(bias_consumer_state) + sBias_stage = sBias[(None, bias_consumer_state.index)] + if cutlass.const_expr(self.act_func == "srelu"): + sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(self.epi_tile[1])) + else: + sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(2 * self.epi_tile[1])) + + # + # Get per-expert C tensor inside tile loop + # + real_c, _ = epi_ext.get_gmem_tensor("c", mC_mnl, padded_offsets, epi_work_tile_info) + gC_mnl = cute.local_tile(real_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + thr_mma_epi_loop = tiled_mma.get_slice(mma_tile_coord_v) + tCgC = thr_mma_epi_loop.partition_C(gC_mnl) + _, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, self.epi_tile_c, sC) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Get per-expert D tensor inside tile loop (NVFP4 D: the RHT warps own + # the fp4 D TMA store — no D partition on the ACT side). + # + if cutlass.const_expr(not self.d_quant): + real_d, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + gD_mnl_loop = cute.local_tile(real_d, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD = thr_mma_epi_loop.partition_C(gD_mnl_loop) + _, bSG_sD, bSG_gD_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgD, epi_tile, sD) + bSG_gD = bSG_gD_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gD = cute.group_modes(bSG_gD, 1, cute.rank(bSG_gD)) + + # + # Get per-expert prob tensor inside tile loop + # + real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info) + mPosition = ( + (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0] + + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape)) + + tidx + ) + mProb = real_prob[mPosition, 0, 0] + + # + # Get accumulator stage index + # + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False) + else: + acc_stage_index = acc_consumer_state.index + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + subtile_step = 1 if cutlass.const_expr(self.act_func == "srelu") else 2 + for subtile_idx in cutlass.range(0, subtile_cnt, subtile_step, unroll=1): + real_subtile_idx = subtile_idx if cutlass.const_expr(self.act_func == "srelu") else subtile_idx // 2 + if cutlass.const_expr(self.overlapping_accum): + if reverse_subtile: + real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n_required - 1 - real_subtile_idx + + # + # Load accumulator from tensor memory buffer to register + # + if cutlass.const_expr(self.act_func == "srelu"): + tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx)] + else: + tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx * 2)] + tTR_tAcc_mn_up = tTR_tAcc[(None, None, None, real_subtile_idx * 2 + 1)] + + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_gate, tTR_rAcc_gate) + if cutlass.const_expr(self.act_func != "srelu"): + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_up, tTR_rAcc_up) + + # + # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled + # + if cutlass.const_expr(self.overlapping_accum): + if subtile_idx == self.iter_acc_early_release_in_epilogue: + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Apply alpha (+ bias when enabled) + # + if cutlass.const_expr(self.enable_bias): + if cutlass.const_expr(self.act_func == "srelu"): + sBias_sub = sBias_subtiles[(None, real_subtile_idx)] + cute.copy(bias_s2r_atom, sBias_sub, tTR_rBias_gate) + else: + sBias_pair = sBias_subtiles[(None, real_subtile_idx)] + sBias_sub = cute.flat_divide(sBias_pair, cute.make_layout(self.epi_tile[1])) + cute.copy(bias_s2r_atom, sBias_sub[(None, 0)], tTR_rBias_gate) + bias_vec_gate = tTR_rBias_gate.load() + if cutlass.const_expr(self.act_func != "srelu"): + cute.copy(bias_s2r_atom, sBias_sub[(None, 1)], tTR_rBias_up) + bias_vec_up = tTR_rBias_up.load() + + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2): + bias_gate_f32_0 = bias_vec_gate[i].to(cutlass.Float32) + bias_gate_f32_1 = bias_vec_gate[i + 1].to(cutlass.Float32) + tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + (bias_gate_f32_0, bias_gate_f32_1), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.act_func != "srelu"): + bias_up_f32_0 = bias_vec_up[i].to(cutlass.Float32) + bias_up_f32_1 = bias_vec_up[i + 1].to(cutlass.Float32) + tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + (bias_up_f32_0, bias_up_f32_1), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)): + tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + bias_vec_gate[i].to(cutlass.Float32) + if cutlass.const_expr(self.act_func != "srelu"): + tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + bias_vec_up[i].to(cutlass.Float32) + + last_bias_subtile = subtile_cnt - 1 if cutlass.const_expr(self.act_func == "srelu") else subtile_cnt - 2 + if subtile_idx == last_bias_subtile: + bias_pipeline.consumer_release(bias_consumer_state) + bias_consumer_state.advance() + else: + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2): + tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.act_func != "srelu"): + tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)): + tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + if cutlass.const_expr(self.act_func != "srelu"): + tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + + # + # Store pre-activation output to C tensor for residual/backward. + # + if cutlass.const_expr(self.act_func == "srelu"): + self.store_c_unary( + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc_gate, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + num_prev_subtiles, + real_subtile_idx, + ) + else: + self.store_c( + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc_gate, + tTR_rAcc_up, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + num_prev_subtiles, + real_subtile_idx, + ) + num_prev_subtiles = num_prev_subtiles + 1 + + # + # GeGLU clamp before C store + # + if cutlass.const_expr((self.act_func == "geglu" or self.act_func == "swiglu") and self.glu_limit is not None): + geglu_max_val = cutlass.Float32(self.glu_limit) + geglu_min_val = cutlass.Float32(-self.glu_limit) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_up)): + tTR_rAcc_gate[i] = fmin(tTR_rAcc_gate[i], geglu_max_val) + tTR_rAcc_gate[i] = fmax(tTR_rAcc_gate[i], geglu_min_val) + tTR_rAcc_up[i] = fmin(tTR_rAcc_up[i], geglu_max_val) + tTR_rAcc_up[i] = fmax(tTR_rAcc_up[i], geglu_min_val) + + acc_vec_gate = tTR_rAcc_gate.load() + + # + # Compute activation. + # + tCompute = cute.make_rmem_tensor(acc_vec_gate.shape, self.acc_dtype) + if cutlass.const_expr(self.act_func == "srelu"): + self.srelu_act(tCompute, acc_vec_gate, mProb) + elif cutlass.const_expr(self.act_func == "geglu"): + acc_vec_up = tTR_rAcc_up.load() + self.geglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset) + elif cutlass.const_expr(self.act_func == "swiglu"): + acc_vec_up = tTR_rAcc_up.load() + self.swiglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb) + + # + # Store post-activation output to D staging (bf16 under NVFP4 D — + # the RHT warps quantize it and issue the fp4 TMA store). + # + acc_vec = tiled_copy_r2s_d.retile(tCompute).load() + tRS_rD.store(acc_vec.to(self.d_smem_dtype)) + d_buffer = num_prev_d_subtiles % self.num_d_stage + cute.copy( + tiled_copy_r2s_d, + tRS_rD, + tRS_sD[(None, None, None, d_buffer)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group0.arrive_and_wait() + if cutlass.const_expr(not self.d_quant): + if warp_idx == self.epilog_act_warp_id[0]: + cute.copy( + tma_atom_d, + bSG_sD[(None, d_buffer)], + bSG_gD[(None, real_subtile_idx)], + ) + d_pipeline.producer_commit() + + # + # Signal the RHT epilogue warps that the post-activation D tile is in SMEM. + # + if cutlass.const_expr(self.run_rht): + pingpong_pipeline.producer_acquire(pingpong_act_producer_state) + pingpong_pipeline.producer_commit(pingpong_act_producer_state) + pingpong_act_producer_state.advance() + + num_prev_d_subtiles = num_prev_d_subtiles + 1 + + # + # Delayed TMA store acquire + group sync (always enabled) + # + if cutlass.const_expr(self.delay_tma_store_acquire_sync): + if warp_idx == self.epilog_act_warp_id[0]: + if cutlass.const_expr(not self.d_quant): + d_pipeline.producer_acquire() + c_pipeline.producer_acquire() + self.epilog_sync_barrier_group0.arrive_and_wait() + + if cutlass.const_expr(self.run_rht): + self.epilog_sync_barrier.arrive_and_wait() + + # + # Full epilogue barrier (ACT + RHT must both arrive) + # + if cutlass.const_expr(self.run_rht): + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + if cutlass.const_expr(not self.overlapping_accum): + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier_group0.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store / pingpong complete + # + c_pipeline.producer_tail() + if cutlass.const_expr(not self.d_quant): + d_pipeline.producer_tail() + if cutlass.const_expr(self.run_rht): + pingpong_pipeline.producer_tail(pingpong_act_producer_state) + + # --------------------------------------------------------------- + # Specialized RHT store warps (4-7): RHT and/or NVFP4 quantization from D SMEM + # --------------------------------------------------------------- + if self.run_rht and warp_idx < self.mma_warp_id and warp_idx >= self.epilog_rht_store_warp_id[0] and total_token > 0: + epi_tidx = tidx % 128 + + # + # Pingpong consumer state + # + pingpong_rht_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_pingpong_stage) + + # + # Create per-expert extension (for D tensor inside tile loop) + # + epi_ext = self._make_extension(workspace_ptr) + + # + # RHT output: TMA-store pipeline for the RHT warps (mirrors ACT's D store). + # + if cutlass.const_expr(self.generate_rht): + rht_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_rht_store_warp_id), + ) + rht_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=rht_producer_group, + ) + + # + # NVFP4 D: the RHT warps quantize sD and own the fp4 D TMA store. + # + if cutlass.const_expr(self.d_quant): + dq_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_rht_store_warp_id), + ) + dq_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=dq_producer_group, + ) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + num_prev_d_subtiles = cutlass.Int32(0) + # Mirror ACT's per-tile subtile-column reversal (overlapping_accum phase). + # ACT's first tile uses acc phase 0 => reverse=True, then toggles each tile. + if cutlass.const_expr(self.run_rht and self.overlapping_accum): + rht_reverse = cutlass.Boolean(True) + while is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + mma_tile_coord_mnl = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx, + cutlass.Int32(0), + ) + expert_idx = epi_work_tile_info.expert_idx + epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx) + + # + # RHT output: per-expert RHT gmem tensor + TMA partition (mirrors ACT's D setup). + # + if cutlass.const_expr(self.generate_rht): + thr_mma_epi_rht = tiled_mma.get_slice(mma_tile_coord_v) + real_rht, _ = epi_ext.get_gmem_tensor("d", mRht_mnl, padded_offsets, epi_work_tile_info) + gRht_mnl_loop = cute.local_tile(real_rht, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgRht = thr_mma_epi_rht.partition_C(gRht_mnl_loop) + _, bSG_sRht, bSG_gRht_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_rht, tCgRht, epi_tile, sRht) + bSG_gRht = bSG_gRht_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gRht = cute.group_modes(bSG_gRht, 1, cute.rank(bSG_gRht)) + if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): + # Expert token offset for the colwise (f, m) scale grid's + # tile index (offsets are 256-aligned, divisions exact). + rht_t_off, _rht_t_cnt = compute_expert_token_range(padded_offsets, epi_work_tile_info.expert_idx) + + # + # NVFP4 D: per-expert fp4 D gmem tensor + TMA partition (mirrors ACT's D setup). + # + if cutlass.const_expr(self.d_quant): + thr_mma_epi_dq = tiled_mma.get_slice(mma_tile_coord_v) + real_dq, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + gDq_mnl_loop = cute.local_tile(real_dq, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgDq = thr_mma_epi_dq.partition_C(gDq_mnl_loop) + _, bSG_sDq, bSG_gDq_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgDq, epi_tile, sDq) + bSG_gDq = bSG_gDq_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gDq = cute.group_modes(bSG_gDq, 1, cute.rank(bSG_gDq)) + + # + # NVFP4 scale tensors: each thread owns one logical scale-domain row. + # Per subtile iteration it stages one scale byte per 16-value block + # into smem, then writes those bytes into the swizzled SF atom layout. + # + if cutlass.const_expr(self.rht_quant or self.d_quant): + sf_row = ( + (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0] + + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape)) + + epi_tidx + ) + # + # Consume D subtiles from SMEM + # + subtile_cnt = self.cta_tile_shape_mnk[1] // cute.size(self.epi_tile[1]) + subtile_step = 1 if cutlass.const_expr(self.act_func == "srelu") else 2 + for subtile_idx in cutlass.range(0, subtile_cnt, subtile_step, unroll=1): + # + # Wait for ACT warps to finish writing the post-activation D tile to SMEM. + # + pingpong_pipeline.consumer_wait(pingpong_rht_consumer_state) + + if cutlass.const_expr(self.run_rht): + d_buffer = num_prev_d_subtiles % self.num_d_stage + # Subtile-column reversal (overlapping_accum): computed BEFORE + # the FWHT/quant calls — scale bytes and TMA'd data must land + # in the same reversed gmem feature columns as ACT's D store. + real_subtile_idx = subtile_idx // 2 + if cutlass.const_expr(self.overlapping_accum): + if rht_reverse: + real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n_required - 1 - real_subtile_idx + # + # Load this subtile's sD values to registers, then ARRIVE the + # lockstep barrier immediately: the ACT warps' arrive_and_wait + # releases, and their next subtile overlaps ALL the FWHT/quant + # compute and stores below. The RHT side never waits on this + # barrier (the pingpong gates it). + # + if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.rht_rowwise): + rht_ld = load_row_bf16(sD, d_buffer, epi_tidx) + else: + rht_ld = load_colwise_pairs_bf16(sD, d_buffer, epi_tidx) + if cutlass.const_expr(self.d_quant): + dq_ld = load_row_bf16(sD, d_buffer, epi_tidx) + self.epilog_sync_barrier.arrive() + # RHT output: write the transform (x0.25) to sRht. rht_rowwise + # picks the transform axis (16-feature blocks per token vs + # 16-token blocks per feature); NVFP4 quantization is inferred + # from the sRht dtype inside the FWHT device functions. + if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.rht_quant and self.rht_rowwise): + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, rht_norm_const, sSfRht, real_subtile_idx) + elif cutlass.const_expr(self.rht_quant): + hadamard_rmem_colwise_fwht_quant(rht_ld, d_buffer, epi_tidx, rht_norm_const, sRht, sSfRht, real_subtile_idx * 2 * HADAMARD_SIZE) + elif cutlass.const_expr(self.rht_rowwise): + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht) + else: + hadamard_rmem_colwise_fwht(rht_ld, d_buffer, epi_tidx, sRht) + # NVFP4 D: quantize the bf16 register rows into sDq + sSfd. + if cutlass.const_expr(self.d_quant): + nvfp4_quant_rmem_row(dq_ld, d_buffer, epi_tidx, sDq, norm_const, sSfd, real_subtile_idx) + # + # TMA-store the produced epi-tiles to gmem (mirrors ACT's D + # store, including the overlapping_accum subtile-column reversal). + # + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group1.arrive_and_wait() + if warp_idx == self.epilog_rht_store_warp_id[0]: + if cutlass.const_expr(self.generate_rht): + cute.copy( + tma_atom_rht, + bSG_sRht[(None, d_buffer)], + bSG_gRht[(None, real_subtile_idx)], + ) + rht_pipeline.producer_commit() + rht_pipeline.producer_acquire() + if cutlass.const_expr(self.d_quant): + cute.copy( + tma_atom_d, + bSG_sDq[(None, d_buffer)], + bSG_gDq[(None, real_subtile_idx)], + ) + dq_pipeline.producer_commit() + dq_pipeline.producer_acquire() + self.epilog_sync_barrier_group1.arrive_and_wait() + + # + # Release pingpong consumer slot + # + pingpong_pipeline.consumer_release(pingpong_rht_consumer_state) + pingpong_rht_consumer_state.advance() + num_prev_d_subtiles = num_prev_d_subtiles + 1 + # Lockstep barrier: under run_rht the 128 RHT arrivals already + # happened right after the sD reads (arrive() above); only the + # idle-RHT plain mode still pairs arrive_and_wait here. + if cutlass.const_expr(not self.run_rht): + self.epilog_sync_barrier.arrive_and_wait() + + # + # NVFP4: one contiguous 8-byte scale-row store per thread per tile. + # Rowwise: the thread wrote all its own smem slots (no barrier needed). + # Colwise: rows were filled ACROSS warps — the loop-end epilog_sync + # barrier above already ordered those writes. + # + _num_sf = self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE + if cutlass.const_expr(self.rht_quant and self.rht_rowwise): + self.store_swizzled_sf_row( + mSfRht_mnl, + epi_ext.token_offset + sf_row, + mma_tile_coord_mnl[1] * _num_sf, + sSfRht, + epi_tidx, + ) + if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): + # (f, m) scale domain: thread <-> feature-in-tile; columns are + # 16-token scale blocks, stored in the same swizzled SF atom layout. + sf_feat_row = mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk_d[1] + epi_tidx + self.store_swizzled_sf_row( + mSfRht_mnl, + sf_feat_row, + (rht_t_off + sf_row - epi_tidx) // HADAMARD_SIZE, + sSfRht, + epi_tidx, + ) + if cutlass.const_expr(self.d_quant): + self.store_swizzled_sf_row( + mSfd_mnl, + epi_ext.token_offset + sf_row, + mma_tile_coord_mnl[1] * _num_sf, + sSfd, + epi_tidx, + ) + + # + # Full epilogue barrier (ACT + RHT must both arrive) + # + self.epilog_sync_barrier.arrive_and_wait() + + # + # Toggle the subtile-column reversal for the next tile. + # + if cutlass.const_expr(self.run_rht and self.overlapping_accum): + rht_reverse = not rht_reverse + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Drain the RHT-warp TMA-store pipelines before exit. + # + if cutlass.const_expr(self.generate_rht): + rht_pipeline.producer_tail() + if cutlass.const_expr(self.d_quant): + dq_pipeline.producer_tail() + # END OF KERNEL diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py new file mode 100644 index 000000000..05d1ab870 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py @@ -0,0 +1,1168 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Shared utilities for discrete-weight block-scaled grouped GEMM kernels. + +This module contains: +- Constants shared across kernel variants +- PTX/DSL utility functions (fmin, fmax, warp reduction, atomics) +- Activation functions (sigmoid, silu, geglu-scaled silu) +- CPU-side reference and validation utilities +- Kernel configuration and validation functions +- Kernel helper functions that don't depend on kernel instance state +""" + +from __future__ import annotations + +from typing import Type, Tuple, Union + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass._mlir import ir +from cutlass._mlir.dialects.nvvm import AtomicOpKind + +try: + # DSL <= 4.5 name; Rubin-era DSLs renamed it (same values, MAX included). + from cutlass._mlir.dialects.nvvm import ReduxKind +except ImportError: + from cutlass._mlir.dialects.nvvm import ReductionKind as ReduxKind +from cutlass.cute.typing import Float32, Int32, BFloat16, AddressSpace +from cutlass._mlir.dialects import math, nvvm, llvm, vector, arith +from ..moe_persistent_scheduler import MoESchedulerParams + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +FIX_PAD_SIZE = 256 +"""Fixed pad size for user-side padding, decoupled from the kernel tile size.""" + + +# --------------------------------------------------------------------------- +# PTX / DSL utility functions +# --------------------------------------------------------------------------- + + +@dsl_user_op +def make_dual_ptr_tensors_from_param( + ptx_first_param_name: str, + first_cnt: int, + second_cnt: int, + *, + loc=None, + ip=None, +): + """Create two Int64 tensors from consecutive kernel param-space pointers. + + Kernel parameters are laid out contiguously in param space (8 bytes each + for .u64). Given the PTX name of the first parameter, this function: + 1. Gets the param-space address via ``mov`` + ``cvta.param.u64`` + 2. Creates tensor_a of shape (first_cnt,) starting at offset 0 + 3. Creates tensor_b of shape (second_cnt,) starting at offset first_cnt * 8 + + :param ptx_first_param_name: PTX name of the very first pointer param + :param first_cnt: Number of pointers in the first group (e.g., expert_cnt for b_ptrs) + :param second_cnt: Number of pointers in the second group (e.g., expert_cnt for sfb_ptrs) + :return: (first_tensor, second_tensor) — two (cnt,) Int64 tensors + """ + generic_addr = cutlass.Int64( + llvm.inline_asm( + T.i64(), + [], + f"""{{ + .reg .u64 %paddr, %gaddr; + mov.u64 %paddr, {ptx_first_param_name}; + cvta.param.u64 %gaddr, %paddr; + mov.u64 $0, %gaddr; + }}""", + "=l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + base_ptr = cute.make_ptr( + cutlass.Int64, + generic_addr, + AddressSpace.generic, + assumed_align=8, + loc=loc, + ip=ip, + ) + + first_layout = cute.make_layout((first_cnt,), loc=loc, ip=ip) + first_tensor = cute.make_tensor(base_ptr, first_layout, loc=loc, ip=ip) + + second_ptr = base_ptr + first_cnt + second_layout = cute.make_layout((second_cnt,), loc=loc, ip=ip) + second_tensor = cute.make_tensor(second_ptr, second_layout, loc=loc, ip=ip) + + return first_tensor, second_tensor + + +def fmin(a: Union[float, Float32], b: Union[float, Float32], *, nan=True, loc=None, ip=None) -> Float32: + if nan: + ptx_instr = f"min.NaN.f32 $0, $1, $2;" + else: + ptx_instr = f"min.f32 $0, $1, $2;" + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + f"{ptx_instr}", + f"=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def fmax(a: Union[float, Float32], b: Union[float, Float32], *, nan=True, loc=None, ip=None) -> Float32: + if nan: + ptx_instr = f"max.NaN.f32 $0, $1, $2;" + else: + ptx_instr = f"max.f32 $0, $1, $2;" + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + f"{ptx_instr}", + f"=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def fmin_bf16x2( + a0: BFloat16, + a1: BFloat16, + b0: BFloat16, + b1: BFloat16, + *, + nan: bool = True, + loc=None, + ip=None, +) -> Tuple[BFloat16, BFloat16]: + vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc) + a_vec = vector.from_elements( + vec_bf16x2_type, + [BFloat16(a0).ir_value(loc=loc, ip=ip), BFloat16(a1).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + b_vec = vector.from_elements( + vec_bf16x2_type, + [BFloat16(b0).ir_value(loc=loc, ip=ip), BFloat16(b1).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + a_packed = llvm.bitcast(Int32.mlir_type, a_vec, loc=loc, ip=ip) + b_packed = llvm.bitcast(Int32.mlir_type, b_vec, loc=loc, ip=ip) + + if nan: + ptx_instr = f"min.NaN.bf16x2 $0, $1, $2;" + else: + ptx_instr = f"min.bf16x2 $0, $1, $2;" + + res_packed = llvm.inline_asm( + Int32.mlir_type, + [a_packed, b_packed], + ptx_instr, + "=r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + res_vec = llvm.bitcast(vec_bf16x2_type, res_packed, loc=loc, ip=ip) + + res0 = BFloat16(vector.extract(res_vec, [], [0], loc=loc, ip=ip)) + res1 = BFloat16(vector.extract(res_vec, [], [1], loc=loc, ip=ip)) + return (res0, res1) + + +def fmax_bf16x2( + a0: BFloat16, + a1: BFloat16, + b0: BFloat16, + b1: BFloat16, + *, + nan: bool = True, + loc=None, + ip=None, +) -> Tuple[BFloat16, BFloat16]: + vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc) + a_vec = vector.from_elements( + vec_bf16x2_type, + [BFloat16(a0).ir_value(loc=loc, ip=ip), BFloat16(a1).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + b_vec = vector.from_elements( + vec_bf16x2_type, + [BFloat16(b0).ir_value(loc=loc, ip=ip), BFloat16(b1).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + a_packed = llvm.bitcast(Int32.mlir_type, a_vec, loc=loc, ip=ip) + b_packed = llvm.bitcast(Int32.mlir_type, b_vec, loc=loc, ip=ip) + + if nan: + ptx_instr = f"max.NaN.bf16x2 $0, $1, $2;" + else: + ptx_instr = f"max.bf16x2 $0, $1, $2;" + + res_packed = llvm.inline_asm( + Int32.mlir_type, + [a_packed, b_packed], + ptx_instr, + "=r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + res_vec = llvm.bitcast(vec_bf16x2_type, res_packed, loc=loc, ip=ip) + + res0 = BFloat16(vector.extract(res_vec, [], [0], loc=loc, ip=ip)) + res1 = BFloat16(vector.extract(res_vec, [], [1], loc=loc, ip=ip)) + return (res0, res1) + + +def atomic_add_bf16x2(ptr, val_fp32_lo, val_fp32_hi, *, loc=None, ip=None): + """Packed BF16x2 atomic reduction to global memory.""" + lo_ir = val_fp32_lo.ir_value(loc=loc, ip=ip) + hi_ir = val_fp32_hi.ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ptr, hi_ir, lo_ir], + "{ .reg .b32 packed; cvt.rn.bf16x2.f32 packed, $1, $2; red.global.add.noftz.bf16x2 [$0], packed; }", + "l,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +def warp_redux_sync( + value, + kind, + mask_and_clamp=0xFFFFFFFF, + abs: bool = False, + nan: bool = None, + *, + loc=None, + ip=None, +): + value_type = type(value) + value_ir = value.ir_value(loc=loc, ip=ip) + mask_ir = Int32(mask_and_clamp).ir_value(loc=loc, ip=ip) + ptx_instr = f"redux.sync.max.abs.NaN.f32 $0, $1, $2;" + + return value_type( + llvm.inline_asm( + T.f32(), + [value_ir, mask_ir], + f"{ptx_instr}", + f"=f,f,i", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def atomic_max_float32( + ptr, + value: Float32, + *, + positive_only: bool = True, + loc=None, + ip=None, +) -> Float32: + value_int = llvm.bitcast(T.i32(), value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + + old_value_int = nvvm.atomicrmw( + op=cutlass._mlir.dialects.nvvm.AtomicOpKind.MAX, + ptr=ptr, + a=value_int, + loc=loc, + ip=ip, + ) + + return Float32(llvm.bitcast(T.f32(), old_value_int, loc=loc, ip=ip)) + + +def atomic_add_float32( + ptr, + value: Float32, + *, + loc=None, + ip=None, +) -> Float32: + """Atomic FP32 addition in global memory (used for dprob gradient accumulation).""" + old_value = nvvm.atomicrmw( + op=AtomicOpKind.FADD, + ptr=ptr, + a=value.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + return Float32(llvm.bitcast(T.f32(), old_value, loc=loc, ip=ip)) + + +def cvt_f32x4_to_f8x4_pack_i32(fp32x4, fp8_type, loc=None, ip=None): + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + + src0 = Float32(vector.extract(src_vec4, [], [0])).ir_value(loc=loc, ip=ip) + src1 = Float32(vector.extract(src_vec4, [], [1])).ir_value(loc=loc, ip=ip) + src2 = Float32(vector.extract(src_vec4, [], [2])).ir_value(loc=loc, ip=ip) + src3 = Float32(vector.extract(src_vec4, [], [3])).ir_value(loc=loc, ip=ip) + + cvt_instruction = "" + if cutlass.const_expr(fp8_type == cutlass.Float8E8M0FNU): + cvt_instruction = "cvt.rp.satfinite.ue8m0x2.f32" + elif cutlass.const_expr(fp8_type == cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.e4m3x2.f32" + else: + with cute.arch.elect_one(): + cute.printf("error: unsupported fp8 element type") + return + + asm_tmpl = ( + "{\n" " .reg .b16 lo;\n" " .reg .b16 hi;\n" f" {cvt_instruction} lo, $2, $1;\n" f" {cvt_instruction} hi, $4, $3;\n" " mov.b32 $0, {lo, hi};\n" "}" + ) + packed_i32 = llvm.inline_asm( + T.i32(), + [src0, src1, src2, src3], + asm_tmpl, + "=r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + return packed_i32 + + +def cvt_f32x4_to_f8x4(fp32x4, fp8x4, loc=None, ip=None): + packed_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, fp8x4.element_type) + fp8x4_i32 = cute.recast_tensor(fp8x4, cutlass.Int32) + fp8x4_i32[0] = cutlass.Int32(packed_i32) + + +def cvt_f32_to_f8_to_f32(fp32x1, fp8_type, loc=None, ip=None): + src_fp32 = Float32(fp32x1).ir_value(loc=loc, ip=ip) + + cvt_instruction_downcast = "" + cvt_instruction_upcast = "" + if cutlass.const_expr(fp8_type == cutlass.Float8E8M0FNU): + cvt_instruction_downcast = "cvt.rp.satfinite.ue8m0x2.f32" + cvt_instruction_upcast = "cvt.rn.bf16x2.ue8m0x2" + elif cutlass.const_expr(fp8_type == cutlass.Float8E4M3FN): + cvt_instruction_downcast = "cvt.rn.satfinite.e4m3x2.f32" + cvt_instruction_upcast = "cvt.rn.bf16x2.e4m3x2" + else: + with cute.arch.elect_one(): + cute.printf("error: unsupported fp8 element type") + return + + asm_tmpl = "{\n" " .reg .b16 bf_lo;\n" f" {cvt_instruction_downcast} bf_lo, 0f00000000, $1;\n" f" {cvt_instruction_upcast} $0, bf_lo;\n" "}" + packed_i32 = llvm.inline_asm( + T.i32(), + [src_fp32], + asm_tmpl, + "=r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + vec_bf16_ty = ir.Type.parse("vector<2xbf16>") + bf2_lo = llvm.bitcast(vec_bf16_ty, packed_i32, loc=loc, ip=ip) + h0 = vector.extract(bf2_lo, [], [0], loc=loc, ip=ip) + return arith.extf(Float32.mlir_type, h0, loc=loc, ip=ip) + + +def ceil_div(a, b): + return (a + b - 1) // b + + +# --------------------------------------------------------------------------- +# Activation functions (device-side, used inside kernels) +# --------------------------------------------------------------------------- + + +def sigmoid_f32(a: Union[float, Float32], fastmath: bool = False) -> Union[float, Float32]: + """Compute the sigmoid of the input value.""" + return cute.arch.rcp_approx(1.0 + cute.math.exp(-a, fastmath=fastmath)) + + +def silu_f32(a: Union[float, Float32], fastmath: bool = False) -> Union[float, Float32]: + """Compute the SiLU (Swish) of the input value.""" + return a * sigmoid_f32(a, fastmath=fastmath) + + +def silu_f32_scaled( + a: Union[float, Float32], + alpha: Union[float, Float32] = 1.702, + fastmath: bool = False, +) -> Union[float, Float32]: + """Compute the scaled SiLU ``a * sigmoid(alpha * a)`` of the input.""" + return a * sigmoid_f32(a * alpha, fastmath=fastmath) + + +def silu_f32_geglu_scaled(a: Union[float, Float32], fastmath: bool = False) -> Union[float, Float32]: + """Backwards-compatible wrapper for :func:`silu_f32_scaled` with ``alpha=1.702``.""" + return silu_f32_scaled(a, alpha=1.702, fastmath=fastmath) + + +# --------------------------------------------------------------------------- +# CPU-side reference and validation utilities +# --------------------------------------------------------------------------- + + +def sigmoid(x): + """PyTorch reference sigmoid using exp2 for numerical consistency.""" + import torch + + LOG2_E = 1.4426950408889634 + exp_x = torch.exp2(x * (-LOG2_E)) + ret = 1.0 / (exp_x + 1.0) + return ret + + +def compute_reference_amax(output_tensor: torch.Tensor) -> float: + """ + Compute reference amax value on CPU. + + Args: + output_tensor: torch.Tensor, GEMM output result (CPU tensor) + + Returns: + float: reference amax value + """ + import torch + + if output_tensor.dtype != torch.float32: + output_fp32 = output_tensor.float() + else: + output_fp32 = output_tensor + + reference_amax = torch.amax(torch.abs(output_fp32)) + + return reference_amax.item() + + +def compare_and_report_mismatches( + gpu_tensor, + ref_tensor, + name="Tensor", + atol=1e-05, + rtol=1e-05, + max_mismatches=8, +): + """ + Compare two tensors and report the first N mismatched elements. + + Args: + gpu_tensor: Results computed on GPU + ref_tensor: Reference results (CPU) + name: Name of the tensor + atol: Absolute tolerance + rtol: Relative tolerance + max_mismatches: Maximum number of mismatches to report + """ + import torch + + if gpu_tensor.is_cuda: + gpu_data = gpu_tensor.cpu() + else: + gpu_data = gpu_tensor + + if ref_tensor.is_cuda: + ref_data = ref_tensor.cpu() + else: + ref_data = ref_tensor + + assert gpu_data.shape == ref_data.shape, f"Shape mismatch: {gpu_data.shape} vs {ref_data.shape}" + + if True: + print(f"\n{name} - First 8 elements:") + print(f"{'Index':<6} {'Coordinate':<30} {'GPU Data':<20} {'CPU Data':<20} {'Abs Error':<20}") + print("-" * 100) + print(f"\n") + + flat_gpu = gpu_data.flatten() + flat_ref = ref_data.flatten() + num_elements = min(8, flat_gpu.numel()) + + for i in range(num_elements): + idx_tuple = torch.unravel_index(torch.tensor(i), gpu_data.shape) + coord = tuple(idx.item() for idx in idx_tuple) + gpu_val = gpu_data[coord].item() + ref_val = ref_data[coord].item() + abs_error = abs(gpu_val - ref_val) + print(f"{i + 1:<6} {str(coord):<30} {gpu_val:<20.6f} {ref_val:<20.6f} {abs_error:<20.6f}") + + diff = torch.abs(gpu_data - ref_data) + threshold = atol + rtol * torch.abs(ref_data) + mismatch_mask = diff > threshold + + mismatch_indices = torch.nonzero(mismatch_mask, as_tuple=False) + num_mismatches = mismatch_indices.shape[0] + + if num_mismatches == 0: + print(f"✓ {name} passed validation! All elements are within tolerance.") + return True + else: + print(f"✗ {name} failed validation!") + print( + f" Total {num_mismatches} mismatched elements (total elements: {gpu_data.numel()}, mismatch rate: {100.0 * num_mismatches / gpu_data.numel():.4f}%)" + ) + print(f" Tolerance settings: atol={atol}, rtol={rtol}") + print(f"\nFirst {min(max_mismatches, num_mismatches)} mismatched elements:") + print(f"{'Index':<6} {'Coordinate':<30} {'GPU Data':<20} {'CPU Data':<20} {'Abs Error':<20}") + print("-" * 100) + + for i in range(min(max_mismatches, num_mismatches)): + idx = mismatch_indices[i] + coord = tuple(idx.tolist()) + gpu_val = gpu_data[coord].item() + ref_val = ref_data[coord].item() + abs_error = diff[coord].item() + + print(f"{i + 1:<6} {str(coord):<30} {gpu_val:<20.6f} {ref_val:<20.6f} {abs_error:<20.6f}") + + raise AssertionError(f"{name} validation failed with {num_mismatches} mismatches") + + +# --------------------------------------------------------------------------- +# Kernel configuration and validation functions +# (extracted from BlockScaledDiscreteWeightGroupedGemmBiasKernel @staticmethod) +# --------------------------------------------------------------------------- + + +def get_amax_smem_size(num_epilog_warps=4): + """Shared memory size (bytes) needed for per-warp amax scratch space.""" + return num_epilog_warps * cute.size_in_bytes(cutlass.Float32, cute.make_layout((1,))) + + +def get_dtype_rcp_limits(dtype: Type[cutlass.Numeric]) -> float: + """ + Reciprocal of the maximum representable absolute value for a given data type. + + :param dtype: Data type + :return: 1 / max_abs_value + """ + if dtype == cutlass.Float4E2M1FN: + return 1 / 6.0 + if dtype == cutlass.Float8E4M3FN: + return 1 / 448.0 + if dtype == cutlass.Float8E5M2: + return 1 / 128.0 + return 1.0 + + +def is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + acc_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], +) -> bool: + """ + Check if the data type / scale-factor vector-size combination is valid. + + :return: True if valid, False otherwise + """ + is_valid = True + if ab_dtype not in { + cutlass.Float4E2M1FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + if sf_vec_size not in {16, 32}: + is_valid = False + + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + is_valid = False + + if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32: + is_valid = False + if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16: + is_valid = False + + if acc_dtype not in {cutlass.Float32}: + is_valid = False + + if d_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + cutlass.Float4E2M1FN, + }: + is_valid = False + + return is_valid + + +def is_valid_layouts( + ab_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, +) -> bool: + """Check if layouts and dtypes are valid combinations.""" + is_valid = True + + if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): + is_valid = False + # TODO: Currently we don't support m major output for Float4E2M1FN, + # Need to support it in the future. + if d_dtype is cutlass.Float4E2M1FN and cd_major == "m": + is_valid = False + return is_valid + + +def is_valid_mma_tiler_and_cluster_shape( + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m_aligned: int, + fix_pad_size: int = FIX_PAD_SIZE, +) -> bool: + """ + Check if the MMA tiler and cluster shape are valid. + + :param fix_pad_size: The fixed pad size used by the kernel (default: FIX_PAD_SIZE). + :return: True if valid, False otherwise + """ + is_valid = True + + if not ((not use_2cta_instrs and mma_tiler_mn[0] in [128]) or (use_2cta_instrs and mma_tiler_mn[0] in [256])): + is_valid = False + # Needs to have even iterations with Epi Tile N 64 for swiGeLU fusion + if mma_tiler_mn[1] not in [256]: + is_valid = False + if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: + is_valid = False + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + cluster_tiler_m = (cluster_shape_mn[0] // (2 if use_2cta_instrs else 1)) * mma_tiler_mn[0] + + if cluster_tiler_m not in [128, 256]: + is_valid = False + + if m_aligned % mma_tiler_mn[0] != 0: + is_valid = False + + if m_aligned != fix_pad_size: + is_valid = False + + return is_valid + + +def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, +) -> bool: + """Check if the tensor alignment requirements are met for TMA loads/stores.""" + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(d_dtype, cd_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + +def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + acc_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + fix_pad_size: int = FIX_PAD_SIZE, +) -> bool: + """ + Check if the grouped GEMM can be implemented with the given parameters. + + :param fix_pad_size: The fixed pad size used by the kernel (default: FIX_PAD_SIZE). + :return: True if implementable, False otherwise + """ + result = True + + if m_aligned != fix_pad_size: + result = False + + if not is_valid_dtypes_and_scale_factor_vec_size(ab_dtype, sf_dtype, sf_vec_size, acc_dtype, d_dtype): + result = False + + if not is_valid_layouts(ab_dtype, d_dtype, a_major, b_major, cd_major): + result = False + + if not is_valid_mma_tiler_and_cluster_shape(use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, m_aligned, fix_pad_size): + result = False + + if not is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, cd_major): + result = False + + if not (a_major == "k" and b_major == "k"): + result = False + + if n % 64 != 0 or m % 256 != 0: + result = False + + return result + + +def compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + epi_tile_c: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + d_dtype: Type[cutlass.Numeric], + d_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + num_smem_capacity: int, + occupancy: int, + generate_sfd: bool, + num_epilog_warps: int = 4, + bias_dtype=None, +) -> Tuple[int, int, int, int, int, int]: + """Compute the number of pipeline stages for A/B/D operands based on heuristics. + + :param num_epilog_warps: Number of epilogue warps (default 4, may differ for DGLU). + :param bias_dtype: Bias element type (e.g. cutlass.BFloat16) when bias enabled, None otherwise. + :return: (num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage) + """ + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + num_c_stage = 2 if generate_sfd else 1 + num_d_stage = 2 if generate_sfd else 1 + + num_tile_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, + ) + + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, + ) + + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile_c, + 1, + ) + + d_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + d_dtype, + d_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + d_bytes_per_stage = cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) + d_bytes = d_bytes_per_stage * num_d_stage * (2 if generate_sfd else 1) + amax_bytes = get_amax_smem_size(num_epilog_warps) if d_dtype == cutlass.BFloat16 else 0 + # Bias SMEM stages + if bias_dtype is not None: + num_bias_stage = 2 # double buffer, each stage = full tile_N + bias_epi_tile_n = mma_tiler_mnk[1] + bias_bytes = bias_epi_tile_n * num_bias_stage * (bias_dtype.width // 8) + else: + num_bias_stage = 0 + bias_bytes = 0 + + epi_bytes = c_bytes + d_bytes + amax_bytes + bias_bytes + + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage + + total_bytes = occupancy * (ab_bytes_per_stage * num_ab_stage + epi_bytes + sinfo_bytes + mbar_helpers_bytes) + + return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage + + +def compute_grid( + sched_params: MoESchedulerParams, + max_active_clusters: cutlass.Constexpr, + use_2cta_instrs: bool, +) -> Tuple[MoESchedulerParams, Tuple[int, int, int]]: + """Compute grid shape for MoE persistent tile scheduling. + + The grid Z dimension indexes persistent clusters. Grid X/Y cover + the cluster shape (including 2CTA factor in X). + """ + grid = ( + sched_params.cluster_shape_mn[0], + sched_params.cluster_shape_mn[1], + max_active_clusters, + ) + return sched_params, grid + + +def compute_stages_wgrad( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + num_smem_capacity: int, + occupancy: int, +) -> Tuple[int, int, int]: + """Compute pipeline stages for the grouped GEMM wgrad kernel.""" + num_acc_stage = 2 + num_c_stage = 2 + num_tile_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, 1) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, 1) + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes + sinfo_bytes)) // ab_bytes_per_stage + + return num_acc_stage, num_ab_stage, num_c_stage + + +def get_tma_atom_kind(atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]: + """ + Select the appropriate TMA copy atom based on SM count and multicast flag. + + :raises ValueError: If the atom_sm_cnt / mcast combination is invalid + """ + if atom_sm_cnt == 2 and mcast: + return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO) + elif atom_sm_cnt == 2 and not mcast: + return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO) + elif atom_sm_cnt == 1 and mcast: + return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE) + elif atom_sm_cnt == 1 and not mcast: + return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE) + + raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}") + + +# --------------------------------------------------------------------------- +# Kernel helper functions (no kernel instance state needed) +# --------------------------------------------------------------------------- + + +@cute.jit +def amax_reduction_per_thread(vec_fp32, amax_fp32): + """Per-thread amax reduction over an FP32 register fragment.""" + vec_fp32_ssa = vec_fp32.load() + abs_acc_values_ir = cutlass._mlir.dialects.math.absf(vec_fp32_ssa.ir_value()) + abs_acc_values = type(vec_fp32_ssa)(abs_acc_values_ir, vec_fp32_ssa.shape, vec_fp32_ssa.dtype) + subtile_amax = abs_acc_values.reduce(cute.ReductionOp.MAX, cutlass.Float32(0.0), 0) + return cute.arch.fmax(amax_fp32, subtile_amax) + + +@cute.jit +def quant_sfd_row( + tile_idx, + tiled_copy_r2s, + src, + pvscale, + norm_const, + rcp_limit, + tRSrD, + sf_vec_size, + vectorized_f32, + sf_dtype, + d_dtype, + use_fp8_ptx_cvt, +): + tTR_rAcc_frg = cute.logical_divide(src, cute.make_layout(sf_vec_size)) + acc_frg = tTR_rAcc_frg.load() + abs_acc_frg_ir = cutlass._mlir.dialects.math.absf(acc_frg.ir_value()) + abs_acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + avg_fp32 = abs_acc_frg[None, 0].reduce(cute.ReductionOp.MAX, cutlass.Float32(0.0), 0) * rcp_limit * norm_const + if tile_idx == 0: + pvscale[0] = avg_fp32 + elif tile_idx == 1: + pvscale[1] = avg_fp32 + elif tile_idx == 2: + pvscale[2] = avg_fp32 + elif tile_idx == 3: + pvscale[3] = avg_fp32 + qpvscale_up = cvt_f32_to_f8_to_f32(avg_fp32, sf_dtype) + fp32_max = cutlass.Float32(3.40282346638528859812e38) + acc_scale = norm_const * cute.arch.rcp_approx(qpvscale_up) + acc_scale = fmin(acc_scale, fp32_max, nan=True) + if cutlass.const_expr(vectorized_f32): + vec = tTR_rAcc_frg[None, 0] + for ei in cutlass.range_constexpr(0, sf_vec_size, 2): + vec[ei], vec[ei + 1] = cute.arch.mul_packed_f32x2((vec[ei], vec[ei + 1]), (acc_scale, acc_scale), rnd="rn", ftz=False) + else: + vec = tTR_rAcc_frg[None, 0] + for ei in cutlass.range_constexpr(sf_vec_size): + vec[ei] = vec[ei] * acc_scale + acc_vec = tiled_copy_r2s.retile(src).load() + if cutlass.const_expr(not use_fp8_ptx_cvt): + tRSrD.store(acc_vec.to(d_dtype)) + else: + tRSrD_i32 = cute.recast_tensor(tRSrD, cutlass.Int32) + for ei in cutlass.range_constexpr(0, sf_vec_size, 4): + fp32x4 = cute.make_rmem_tensor(4, cutlass.Float32) + fp32x4[0] = acc_vec[ei + 0] + fp32x4[1] = acc_vec[ei + 1] + fp32x4[2] = acc_vec[ei + 2] + fp32x4[3] = acc_vec[ei + 3] + fp8x4_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, d_dtype) + tRSrD_i32[ei // 4] = cutlass.Int32(fp8x4_i32) + + +@cute.jit +def quant_sfd_col( + tile_idx, + tiled_copy_r2s, + src, + pvscale, + norm_const, + rcp_limit, + tRSrD, + sf_vec_size, + sf_dtype, + d_dtype, + use_fp8_ptx_cvt, +): + tTR_rAcc_frg = cute.logical_divide(src, cute.make_layout(sf_vec_size)) + acc_frg = tTR_rAcc_frg.load() + abs_acc_frg_ir = cutlass._mlir.dialects.math.absf(acc_frg.ir_value()) + acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + avg_fp32 = cutlass.Float32(0.0) + fp32_max = cutlass.Float32(3.40282346638528859812e38) + tidx, _, _ = cute.arch.thread_idx() + for vi in cutlass.range_constexpr(0, acc_frg.shape[0], 4): + max_value0 = cutlass.Float32(warp_redux_sync(value=acc_frg[vi, 0], kind=ReduxKind.MAX, mask_and_clamp=0xFFFFFFFF, nan=True)) + max_value1 = cutlass.Float32(warp_redux_sync(value=acc_frg[vi + 1, 0], kind=ReduxKind.MAX, mask_and_clamp=0xFFFFFFFF, nan=True)) + max_value2 = cutlass.Float32(warp_redux_sync(value=acc_frg[vi + 2, 0], kind=ReduxKind.MAX, mask_and_clamp=0xFFFFFFFF, nan=True)) + max_value3 = cutlass.Float32(warp_redux_sync(value=acc_frg[vi + 3, 0], kind=ReduxKind.MAX, mask_and_clamp=0xFFFFFFFF, nan=True)) + + scale = rcp_limit * norm_const + max_value0, max_value1 = cute.arch.mul_packed_f32x2((max_value0, max_value1), (scale, scale), rnd="rn", ftz=False) + max_value2, max_value3 = cute.arch.mul_packed_f32x2((max_value2, max_value3), (scale, scale), rnd="rn", ftz=False) + + if tidx % 32 == vi: + avg_fp32 = max_value0 + if tidx % 32 == vi + 1: + avg_fp32 = max_value1 + if tidx % 32 == vi + 2: + avg_fp32 = max_value2 + if tidx % 32 == vi + 3: + avg_fp32 = max_value3 + + max_value_tensor = cute.make_rmem_tensor(4, cutlass.Float32) + max_value_tensor[0] = max_value0 + max_value_tensor[1] = max_value1 + max_value_tensor[2] = max_value2 + max_value_tensor[3] = max_value3 + + if cutlass.const_expr(not use_fp8_ptx_cvt): + max_value_vec_f8 = max_value_tensor.load().to(sf_dtype) + else: + max_value_vec_f8 = cute.make_rmem_tensor(4, sf_dtype) + cvt_f32x4_to_f8x4(max_value_tensor, max_value_vec_f8) + max_value_vec_f8 = max_value_vec_f8.load() + max_value_vec_f32_chunked = max_value_vec_f8.to(cutlass.Float32) + max_value0 = max_value_vec_f32_chunked[0] + max_value1 = max_value_vec_f32_chunked[1] + max_value2 = max_value_vec_f32_chunked[2] + max_value3 = max_value_vec_f32_chunked[3] + + max_value_rcp0 = cute.arch.rcp_approx(max_value0) + max_value_rcp1 = cute.arch.rcp_approx(max_value1) + max_value_rcp2 = cute.arch.rcp_approx(max_value2) + max_value_rcp3 = cute.arch.rcp_approx(max_value3) + + max_value_rcp0 = fmin(max_value_rcp0, fp32_max, nan=True) + max_value_rcp1 = fmin(max_value_rcp1, fp32_max, nan=True) + max_value_rcp2 = fmin(max_value_rcp2, fp32_max, nan=True) + max_value_rcp3 = fmin(max_value_rcp3, fp32_max, nan=True) + + acc_scale_col0, acc_scale_col1 = cute.arch.mul_packed_f32x2((norm_const, norm_const), (max_value_rcp0, max_value_rcp1), rnd="rn", ftz=False) + acc_scale_col2, acc_scale_col3 = cute.arch.mul_packed_f32x2((norm_const, norm_const), (max_value_rcp2, max_value_rcp3), rnd="rn", ftz=False) + + tTR_rAcc_frg[vi], tTR_rAcc_frg[vi + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_frg[vi], tTR_rAcc_frg[vi + 1]), (acc_scale_col0, acc_scale_col1), rnd="rn", ftz=False + ) + tTR_rAcc_frg[vi + 2], tTR_rAcc_frg[vi + 3] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_frg[vi + 2], tTR_rAcc_frg[vi + 3]), (acc_scale_col2, acc_scale_col3), rnd="rn", ftz=False + ) + + pvscale[None, None, tile_idx][0] = avg_fp32 + acc_vec = tiled_copy_r2s.retile(src).load() + if cutlass.const_expr(not use_fp8_ptx_cvt): + tRSrD.store(acc_vec.to(d_dtype)) + else: + tRSrD_i32 = cute.recast_tensor(tRSrD, cutlass.Int32) + for ei in cutlass.range_constexpr(0, sf_vec_size, 4): + fp32x4 = cute.make_rmem_tensor(4, cutlass.Float32) + fp32x4[0] = acc_vec[ei + 0] + fp32x4[1] = acc_vec[ei + 1] + fp32x4[2] = acc_vec[ei + 2] + fp32x4[3] = acc_vec[ei + 3] + fp8x4_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, d_dtype) + tRSrD_i32[ei // 4] = cutlass.Int32(fp8x4_i32) + + +def epilog_gmem_copy_and_partition( + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gD_mnl: cute.Tensor, + epi_tile: cute.Tile, + sD: cute.Tensor, +) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Partition shared memory and global memory for TMA epilogue store. + + :param tidx: Thread index in epilogue warp groups (unused, kept for interface compat) + :param atom: TMA copy atom + :param gD_mnl: Global tensor D + :param epi_tile: Epilogue tiler + :param sD: Shared memory tensor + :return: (tma_atom_d, bSG_sD, bSG_gD) + """ + gD_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tma_atom_d = atom + sD_for_tma_partition = cute.group_modes(sD, 0, 2) + gD_for_tma_partition = cute.group_modes(gD_epi, 0, 2) + bSG_sD, bSG_gD = cpasync.tma_partition( + tma_atom_d, + 0, + cute.make_layout(1), + sD_for_tma_partition, + gD_for_tma_partition, + ) + return tma_atom_d, bSG_sD, bSG_gD diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py new file mode 100644 index 000000000..50d55725f --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +"""NVFP4 quantization device functions (flashinfer blockscaled epilogue replicas).""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import math + +from .moe_kernel_helpers import fmin + +HADAMARD_SIZE = 16 + + +# blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:3354-3370 +def get_dtype_rcp_limits(dtype): + if dtype == cutlass.Float4E2M1FN: + return 1 / 6.0 + if dtype == cutlass.Float8E4M3FN: + return 1 / 448.0 + raise ValueError(f"unsupported quantized dtype {dtype}") + + +@cute.jit +def load_row_bf16(sD, d_buffer, tidx): + """One vectorized copy of the thread's full 32-feature bf16 sD row into rmem. + Split out of the quant/FWHT functions so the kernel can signal "sD consumed" + (barrier arrive) right after the loads, before any compute.""" + row_feats = 2 * HADAMARD_SIZE + sD_tiles = cute.zipped_divide(cute.slice_(sD, (None, None, d_buffer)), (1, row_feats)) + src = cute.slice_(sD_tiles, ((None, None), (tidx, 0))) + rmem_bf16 = cute.make_rmem_tensor(src.shape, cutlass.BFloat16) + cute.autovec_copy(src, rmem_bf16) + return rmem_bf16 + + +@cute.jit +def nvfp4_quant_rmem_row(rmem_bf16, d_buffer, tidx, sOut, norm_const, sSf, sf_pair): + """NVFP4-quantize the thread's full 32-feature bf16 row (from load_row_bf16, no + transform) into sOut (packed e2m1) + sSf (e4m3 block scales) — the fused-D + analog of group_rht_cast's rowwise cast path.""" + row_feats = 2 * HADAMARD_SIZE + tCompute = cute.make_rmem_tensor(rmem_bf16.shape, cutlass.Float32) + for i in cutlass.range_constexpr(row_feats): + tCompute[i] = rmem_bf16[i].to(cutlass.Float32) + + _nvfp4_quant_row(tCompute, d_buffer, tidx, sOut, norm_const, sSf, sf_pair) + + +@cute.jit +def _nvfp4_quant_row(tCompute, d_buffer, token, sOut, norm_const, sSf, sf_pair): + """Flashinfer blockscaled-epilogue NVFP4 quantization of one thread's 32-feature + f32 row: packed e2m1 data into sOut, one e4m3 block scale per (1,16) block into + sSf slots (token, num_vecs*sf_pair + vi).""" + row_feats = cute.size(tCompute.shape) + tCompute_flat = cute.make_tensor(tCompute.iterator, cute.make_layout((row_feats,))) + + # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2837-2847 + num_vecs = row_feats // HADAMARD_SIZE + assert num_vecs % 2 == 0, "num_vecs must be even (packed f32x2 pair loops)" + tTR_rAcc_frg = cute.logical_divide(tCompute_flat, cute.make_layout(HADAMARD_SIZE)) + acc_frg = tTR_rAcc_frg.load() + abs_acc_frg_ir = math.absf(acc_frg.ir_value()) + abs_acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + tCrSFC_pvscale = cute.make_rmem_tensor((num_vecs,), cutlass.Float32) + # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2850-2856 + for vi in cutlass.range_constexpr(num_vecs): + tCrSFC_pvscale[vi] = abs_acc_frg[None, vi].reduce( + cute.ReductionOp.MAX, + cutlass.Float32(0.0), + 0, # Use 0.0 as init for abs values + ) + # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2856-2873 + for vi in cutlass.range_constexpr(0, num_vecs, 2): + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = cute.arch.mul_packed_f32x2( + (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), + ( + get_dtype_rcp_limits(cutlass.Float4E2M1FN), + get_dtype_rcp_limits(cutlass.Float4E2M1FN), + ), + ) + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = cute.arch.mul_packed_f32x2( + (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), + (norm_const, norm_const), + ) + + # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2887 + # f32 -> e4m3 as a padded 4-wide vector: nvgpu.cvt_fptrunc requires a + # 32-bit-aligned 1-d vector (4 x f8), never a scalar. Same pattern as + # grouped_gemm_swiglu_quant's pvscale_f32x4 -> sfd_f8x4 round-trip. + assert num_vecs <= 4, "scale convert padding assumes at most 4 blocks per row" + pvscale_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) + for vi in cutlass.range_constexpr(4): + pvscale_f32x4[vi] = tCrSFC_pvscale[min(vi, num_vecs - 1)] + tCrSFC_f8x4 = cute.make_rmem_tensor((4,), cutlass.Float8E4M3FN) + tCrSFC_f8x4.store(pvscale_f32x4.load().to(cutlass.Float8E4M3FN)) + tCrSFC = cute.make_rmem_tensor((num_vecs,), cutlass.Float8E4M3FN) + for vi in cutlass.range_constexpr(num_vecs): + tCrSFC[vi] = tCrSFC_f8x4[vi] + # e4m3 -> f32 widening round-trip, vectorized for the same reason. + tCrSFC_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) + tCrSFC_f32x4.store(tCrSFC_f8x4.load().to(cutlass.Float32)) + + # SFC -> smem + for vi in cutlass.range_constexpr(num_vecs): + sSf[(token, num_vecs * sf_pair + vi)] = tCrSFC[vi] + + # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2900-2921 + fp32_max = cutlass.Float32(3.40282346638528859812e38) + for vi in cutlass.range_constexpr(0, num_vecs, 2): + acc_scale = cute.arch.mul_packed_f32x2( + ( + cute.arch.rcp_approx(tCrSFC_f32x4[vi]), + cute.arch.rcp_approx(tCrSFC_f32x4[vi + 1]), + ), + (norm_const, norm_const), + ) + acc_scale_min0 = fmin(acc_scale[0], fp32_max, nan=True) + acc_scale_min1 = fmin(acc_scale[1], fp32_max, nan=True) + + vec0 = tTR_rAcc_frg[None, vi] + vec1 = tTR_rAcc_frg[None, vi + 1] + for ei in cutlass.range_constexpr(HADAMARD_SIZE): + vec0[ei], vec1[ei] = cute.arch.mul_packed_f32x2( + (vec0[ei], vec1[ei]), + (acc_scale_min0, acc_scale_min1), + ) + + # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2936-2937 + tRS_rC = cute.make_rmem_tensor(tCompute.shape, cutlass.Float4E2M1FN) + tRS_rC.store(tCompute.load().to(cutlass.Float4E2M1FN)) + sOut_tiles = cute.zipped_divide(cute.slice_(sOut, (None, None, d_buffer)), (1, row_feats)) + dst = cute.slice_(sOut_tiles, ((None, None), (token, 0))) + cute.autovec_copy(tRS_rC, dst) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py new file mode 100644 index 000000000..13ddccf86 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +"""SMEM FWHT device functions for the fused RHT output (orthonormal Sylvester H16). + +Split into LOAD helpers (sD -> registers) and rmem COMPUTE functions so the kernel +can signal "sD consumed" (lockstep-barrier arrive) right after the loads — the +FWHT/quant compute and all stores then overlap the ACT warps' next subtile.""" + +import cutlass +import cutlass.cute as cute + +from .moe_kernel_helpers import fmax, fmin +from .quant_utils import ( + _nvfp4_quant_row, + get_dtype_rcp_limits, +) + +HADAMARD_SIZE = 16 + + +@cute.jit +def load_colwise_pairs_bf16(sD, d_buffer, tidx): + """16 vectorized 4-byte pair loads of the thread's TWO adjacent feature columns + of one 16-token block, interleaved token-major: rmem[2*i + c] = + sD(token_base + i, feature + c).""" + token_block = tidx // HADAMARD_SIZE + feat_pair = tidx % HADAMARD_SIZE + token_base = token_block * HADAMARD_SIZE + + rmem_bf16 = cute.make_rmem_tensor((2 * HADAMARD_SIZE,), cutlass.BFloat16) + rmem_pairs = cute.zipped_divide(rmem_bf16, (2,)) + sD_pairs = cute.zipped_divide(cute.slice_(sD, (None, None, d_buffer)), (1, 2)) + for i in cutlass.range_constexpr(HADAMARD_SIZE): + cute.autovec_copy( + cute.slice_(sD_pairs, ((None, None), (token_base + i, feat_pair))), + cute.slice_(rmem_pairs, ((None,), i)), + ) + return rmem_bf16 + + +@cute.jit +def _colwise_fwht_inplace(tCompute): + """In-place natural-order 16-point FWHT over both interleaved token-major + columns (tCompute[2*i + c] = column c, token i).""" + for c in cutlass.range_constexpr(2): + for base in cutlass.range_constexpr(0, HADAMARD_SIZE, 2): + x = tCompute[2 * base + c] + y = tCompute[2 * (base + 1) + c] + tCompute[2 * base + c] = x + y + tCompute[2 * (base + 1) + c] = x - y + for base in cutlass.range_constexpr(0, HADAMARD_SIZE, 4): + for i in cutlass.range_constexpr(2): + x = tCompute[2 * (base + i) + c] + y = tCompute[2 * (base + i + 2) + c] + tCompute[2 * (base + i) + c] = x + y + tCompute[2 * (base + i + 2) + c] = x - y + for base in cutlass.range_constexpr(0, HADAMARD_SIZE, 8): + for i in cutlass.range_constexpr(4): + x = tCompute[2 * (base + i) + c] + y = tCompute[2 * (base + i + 4) + c] + tCompute[2 * (base + i) + c] = x + y + tCompute[2 * (base + i + 4) + c] = x - y + for i in cutlass.range_constexpr(8): + x = tCompute[2 * i + c] + y = tCompute[2 * (i + 8) + c] + tCompute[2 * i + c] = x + y + tCompute[2 * (i + 8) + c] = x - y + + +@cute.jit +def hadamard_rmem_colwise_fwht(rmem_bf16, d_buffer, tidx, sRht): + """Per-feature bf16 FWHT over 16-token blocks from the load_colwise_pairs_bf16 + registers, stored to sRht at the SAME (token, feature) coords the input was + read from as one 4-byte pair store per token row. + (The colwise QUANT path lives in hadamard_rmem_colwise_fwht_quant.) + + In-place natural-order FWHT => output i is the transformed value for token + token_base+i, matching torch_rht's block layout.""" + token_block = tidx // HADAMARD_SIZE + feat_pair = tidx % HADAMARD_SIZE + token_base = token_block * HADAMARD_SIZE + + n_vals = 2 * HADAMARD_SIZE + tCompute = cute.make_rmem_tensor((n_vals,), cutlass.Float32) + for i in cutlass.range_constexpr(n_vals): + tCompute[i] = rmem_bf16[i].to(cutlass.Float32) + + _colwise_fwht_inplace(tCompute) + + # Orthonormal scale, rounded through bf16 (== the bf16 output values). + for i in cutlass.range_constexpr(n_vals): + tCompute[i] = (tCompute[i] * cutlass.Float32(0.25)).to(cutlass.BFloat16).to(cutlass.Float32) + + rmem_st = cute.make_rmem_tensor((2,), cutlass.BFloat16) + sRht_pairs = cute.zipped_divide(cute.slice_(sRht, (None, None, d_buffer)), (1, 2)) + for i in cutlass.range_constexpr(HADAMARD_SIZE): + for c in cutlass.range_constexpr(2): + rmem_st[c] = tCompute[2 * i + c].to(cutlass.BFloat16) + cute.autovec_copy(rmem_st, cute.slice_(sRht_pairs, ((None, None), (token_base + i, feat_pair)))) + + +@cute.jit +def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht, sSf, sf_row_base): + """Colwise FWHT + NVFP4 quantization from the load_colwise_pairs_bf16 registers, + stored at the SAME (token, feature) coords the input was read from (the staging + is f-major like every other output; packed nibbles pair ADJACENT FEATURES of one + token, so each token row is one 1-byte store). Quantization blocks follow the + transform: one (16, 1) token-block scale per feature, staged in the sSf smem + rows (sf_row_base + feature, token_block); the kernel stores each thread's whole + contiguous scale row once per tile.""" + token_block = tidx // HADAMARD_SIZE + feat_pair = tidx % HADAMARD_SIZE + token_base = token_block * HADAMARD_SIZE + feature = 2 * feat_pair + + n_vals = 2 * HADAMARD_SIZE + tCompute = cute.make_rmem_tensor((n_vals,), cutlass.Float32) + for i in cutlass.range_constexpr(n_vals): + tCompute[i] = rmem_bf16[i].to(cutlass.Float32) + + _colwise_fwht_inplace(tCompute) + + # Orthonormal scale, rounded through bf16 (== the bf16 output values). + for i in cutlass.range_constexpr(n_vals): + tCompute[i] = (tCompute[i] * cutlass.Float32(0.25)).to(cutlass.BFloat16).to(cutlass.Float32) + + # group_rht_cast's exact (fast_math=0) op sequence — NOT the flashinfer one: + # gem = ge * (1/6) is pre-folded, and the encode scale is computed with EXACT + # f32 divisions (enc = 1/(dec * gd), gd = 1/ge). The flashinfer + # rcp_approx(dec) * ge form agrees at ge = 1 but drifts 1 ulp at non-dyadic + # ge, flipping e2m1 codes on rounding boundaries. + gem = norm_const * cutlass.Float32(get_dtype_rcp_limits(cutlass.Float4E2M1FN)) + gd = cutlass.Float32(1.0) / norm_const + pv = cute.make_rmem_tensor((2,), cutlass.Float32) + for c in cutlass.range_constexpr(2): + pv[c] = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(HADAMARD_SIZE): + v = tCompute[2 * i + c] + pv[c] = fmax(pv[c], fmax(v, -v)) + pv[0], pv[1] = cute.arch.mul_packed_f32x2((pv[0], pv[1]), (gem, gem)) + + # f32 -> e4m3 as a padded 4-wide vector (and widen back the same way): + # nvgpu.cvt_fptrunc requires a 32-bit-aligned 1-d vector (4 x f8), never a + # scalar. Same pattern as grouped_gemm_swiglu_quant's pvscale round-trip. + pv_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) + for c in cutlass.range_constexpr(4): + pv_f32x4[c] = pv[min(c, 1)] + tCrSFC_f8x4 = cute.make_rmem_tensor((4,), cutlass.Float8E4M3FN) + tCrSFC_f8x4.store(pv_f32x4.load().to(cutlass.Float8E4M3FN)) + tCrSFC_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) + tCrSFC_f32x4.store(tCrSFC_f8x4.load().to(cutlass.Float32)) + for c in cutlass.range_constexpr(2): + sSf[(sf_row_base + feature + c, token_block)] = tCrSFC_f8x4[c] + + fp32_max = cutlass.Float32(3.40282346638528859812e38) + acc_scale_min0 = fmin(cutlass.Float32(1.0) / (tCrSFC_f32x4[0] * gd), fp32_max, nan=True) + acc_scale_min1 = fmin(cutlass.Float32(1.0) / (tCrSFC_f32x4[1] * gd), fp32_max, nan=True) + for i in cutlass.range_constexpr(HADAMARD_SIZE): + tCompute[2 * i], tCompute[2 * i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[2 * i], tCompute[2 * i + 1]), + (acc_scale_min0, acc_scale_min1), + ) + + tRS_rC = cute.make_rmem_tensor(tCompute.shape, cutlass.Float4E2M1FN) + tRS_rC.store(tCompute.load().to(cutlass.Float4E2M1FN)) + src_pairs = cute.zipped_divide(tRS_rC, (2,)) + sRht_pairs = cute.zipped_divide(cute.slice_(sRht, (None, None, d_buffer)), (1, 2)) + for i in cutlass.range_constexpr(HADAMARD_SIZE): + cute.autovec_copy( + cute.slice_(src_pairs, ((None,), i)), + cute.slice_(sRht_pairs, ((None, None), (token_base + i, feat_pair))), + ) + + +@cute.jit +def hadamard_rmem_rowwise_fwht(rmem_bf16, d_buffer, tidx, sRht, norm_const=1.0, sSf=None, sf_pair=0): + """Per-token FWHT over the thread's full 32-feature row (two independent 16-feature + Hadamard blocks) from the load_row_bf16 registers, stored to sRht at the same + (token, feature) coords: bf16, or NVFP4 (packed e2m1 into sRht + e4m3 block scales + into sSf slots (tidx, 2*sf_pair..)) when sRht is fp4-typed.""" + token = tidx + row_feats = 2 * HADAMARD_SIZE + + tCompute = cute.make_rmem_tensor(rmem_bf16.shape, cutlass.Float32) + for i in cutlass.range_constexpr(row_feats): + tCompute[i] = rmem_bf16[i].to(cutlass.Float32) + + for off in cutlass.range_constexpr(0, row_feats, HADAMARD_SIZE): + for base in cutlass.range_constexpr(0, HADAMARD_SIZE, 2): + x = tCompute[off + base] + y = tCompute[off + base + 1] + tCompute[off + base] = x + y + tCompute[off + base + 1] = x - y + for base in cutlass.range_constexpr(0, HADAMARD_SIZE, 4): + for i in cutlass.range_constexpr(2): + x = tCompute[off + base + i] + y = tCompute[off + base + i + 2] + tCompute[off + base + i] = x + y + tCompute[off + base + i + 2] = x - y + for base in cutlass.range_constexpr(0, HADAMARD_SIZE, 8): + for i in cutlass.range_constexpr(4): + x = tCompute[off + base + i] + y = tCompute[off + base + i + 4] + tCompute[off + base + i] = x + y + tCompute[off + base + i + 4] = x - y + for i in cutlass.range_constexpr(8): + x = tCompute[off + i] + y = tCompute[off + i + 8] + tCompute[off + i] = x + y + tCompute[off + i + 8] = x - y + + # Orthonormal scale, rounded through bf16 (== the bf16 output values) either way. + for i in cutlass.range_constexpr(row_feats): + tCompute[i] = (tCompute[i] * cutlass.Float32(0.25)).to(cutlass.BFloat16).to(cutlass.Float32) + + if cutlass.const_expr(sRht.element_type == cutlass.Float4E2M1FN): + _nvfp4_quant_row(tCompute, d_buffer, token, sRht, norm_const, sSf, sf_pair) + else: + rmem_st = cute.make_rmem_tensor(rmem_bf16.shape, cutlass.BFloat16) + for i in cutlass.range_constexpr(row_feats): + rmem_st[i] = tCompute[i].to(cutlass.BFloat16) + sRht_tiles = cute.zipped_divide(cute.slice_(sRht, (None, None, d_buffer)), (1, row_feats)) + dst = cute.slice_(sRht_tiles, ((None, None), (token, 0))) + cute.autovec_copy(rmem_st, dst) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py new file mode 100644 index 000000000..b6c787f67 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py @@ -0,0 +1,676 @@ +"""Tests for grouped GEMM GLU + Hadamard + Quant forward fusion (SM100+).""" + +from typing import Dict, Optional + +import pytest +import torch + +from cudnn.gemm.cutedsl.grouped.glu_hadamard_quant.rht_utils import HADAMARD_SIZE +from test_low_precision_matmul import float4_e2m1fn_x2_to_float32 +from test_utils import torch_fork_set_rng +from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import allocate_discrete_input_tensors +from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import allocate_grouped_gemm_input_tensors, grouped_gemm_swiglu_init +from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES + +FP4_EXECUTION_CASES = [ + (torch.float4_e2m1fn_x2, torch.float8_e4m3fn, 16), + (torch.float4_e2m1fn_x2, torch.float8_e8m0fnu, 16), +] + + +def _make_cfg(request, *, ab_dtype, sf_dtype, sf_vec_size, enable_bias=False) -> Dict: + return grouped_gemm_swiglu_init( + request, + ab_dtype=ab_dtype, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, + vector_f32=False, + discrete_col_sfd=False, + enable_bias=enable_bias, + ) + + +# ============================================================================= +# Reference implementations +# ============================================================================= + + +def _run_grouped_gemm_glu_ref( + inputs: Dict, + act_func: str, + glu_alpha: Optional[float] = None, + glu_limit: Optional[float] = None, +) -> Dict: + n, _, l = inputs["b_ref"].shape + n_out = n // 2 + valid_m = inputs["valid_m"] + aligned_group_m_list = inputs["aligned_group_m_list"] + + ref = torch.empty((1, valid_m, n), dtype=torch.float32, device=inputs["a_ref"].device) + start = 0 + for i, group_m in enumerate(aligned_group_m_list): + end = start + group_m + res_a = torch.einsum("mk,mk->mk", inputs["a_ref"][start:end, :, 0].to(torch.float32), inputs["sfa_ref"][start:end, :, 0].to(torch.float32)) + res_b = torch.einsum("nk,nk->nk", inputs["b_ref"][:, :, i].to(torch.float32), inputs["sfb_ref"][:, :, i].to(torch.float32)) + ref[0, start:end, :] = torch.einsum("mk,nk->mn", res_a, res_b) + start = end + ref = ref.permute((1, 2, 0)) + + start = 0 + for i, group_m in enumerate(aligned_group_m_list): + end = start + group_m + ref[start:end, :, 0] = ref[start:end, :, 0] * inputs["alpha_tensor"][i].item() + start = end + + if inputs.get("bias_tensor") is not None: + start = 0 + for i, group_m in enumerate(aligned_group_m_list): + end = start + group_m + ref[start:end, :, 0] = ref[start:end, :, 0] + inputs["bias_tensor"][:, i].unsqueeze(0).to(torch.float32) + start = end + + group = 32 + assert n % group == 0, "N must be divisible by 32 for GLU block grouping" + num_blocks = n // group + assert num_blocks % 2 == 0, "Number of 32-col blocks must be even" + + cols = torch.arange(n, device=ref.device, dtype=torch.long) + block_cols = cols.view(num_blocks, group) + gate_idx = block_cols[0::2].reshape(-1) + up_idx = block_cols[1::2].reshape(-1) + ref_gate = ref.index_select(1, gate_idx) + ref_up = ref.index_select(1, up_idx) + + # The kernel clamps BOTH gate and up to [-glu_limit, glu_limit] when set + # (after the C store, before the activation). + if glu_limit is not None: + ref_gate = torch.clamp(ref_gate, min=-glu_limit, max=glu_limit) + ref_up = torch.clamp(ref_up, min=-glu_limit, max=glu_limit) + + if act_func == "swiglu": + ref_after_glu = ref_up * (ref_gate * torch.sigmoid(ref_gate)) + elif act_func == "geglu": + ref_after_glu = (ref_up + 1.0) * ref_gate * torch.sigmoid(1.702 * ref_gate) + else: + raise ValueError(f"Unsupported act_func {act_func}") + + ref_after_glu = ref_after_glu * inputs["prob_tensor"].expand(-1, n_out, -1) + # glu_alpha is a final output scale in this kernel. + if glu_alpha is not None and glu_alpha != 1.0: + ref_after_glu = ref_after_glu * glu_alpha + return {"c_ref": ref.clone(), "d_ref": ref_after_glu} + + +def _hadamard16(device: torch.device) -> torch.Tensor: + """Sylvester (natural FWHT order) 16x16 Hadamard matrix.""" + h = torch.ones((1, 1), dtype=torch.float32, device=device) + for _ in range(4): + h = torch.cat([torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], dim=0) + return h + + +def _rht_ref(d_ref: torch.Tensor, rowwise: bool) -> torch.Tensor: + """Orthonormal H16 transform of the bf16-rounded D reference, at D's (m, f) + orientation. rowwise transforms 16-feature blocks per token; colwise + transforms 16-token blocks per feature. Returns (m, f) bf16.""" + x = d_ref[:, :, 0].to(torch.bfloat16).to(torch.float32) + h = _hadamard16(x.device) + m, f = x.shape + if rowwise: + blocks = x.view(m, f // HADAMARD_SIZE, HADAMARD_SIZE) + out = torch.einsum("ij,mbj->mbi", h, blocks).reshape(m, f) + else: + blocks = x.view(m // HADAMARD_SIZE, HADAMARD_SIZE, f) + out = torch.einsum("ij,bjf->bif", h, blocks).reshape(m, f) + return (out * 0.25).to(torch.bfloat16) + + +def _nvfp4_sf_ref(x: torch.Tensor, norm_const: float) -> torch.Tensor: + """e4m3 block scales of x (rows, cols) with (1, 16) blocks along the last dim, + replicating the kernel's op order: amax * (1/6), then * norm_const.""" + xf = x.to(torch.bfloat16).to(torch.float32) + rows, cols = xf.shape + amax = xf.view(rows, cols // HADAMARD_SIZE, HADAMARD_SIZE).abs().amax(dim=-1) + return ((amax * (1.0 / 6.0)) * norm_const).to(torch.float8_e4m3fn) + + +def _swizzled_sf_to_flat(sf_tensor: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Gather M32x4xrm_K4xrk_L scale storage into logical (rows, cols/16).""" + sf_cols = (cols + HADAMARD_SIZE - 1) // HADAMARD_SIZE + row_idx = torch.arange(rows, device=sf_tensor.device, dtype=torch.long).view(rows, 1) + col_idx = torch.arange(sf_cols, device=sf_tensor.device, dtype=torch.long).view(1, sf_cols) + return sf_tensor[ + row_idx % 32, + (row_idx // 32) % 4, + row_idx // 128, + col_idx % 4, + col_idx // 4, + 0, + ] + + +def _check_nvfp4_output( + values: torch.Tensor, + sf_tensor: torch.Tensor, + ref: torch.Tensor, + norm_const: float, + name: str, +) -> None: + """Check unpacked e2m1 values (rows, cols) + e4m3 scales (rows, cols/16), + with (1, 16) quantization blocks along the last dim, against the f32 + reference (rows, cols): scales against the reference scale computation (one + e4m3 step of slack), data via a per-block dequantization error bound (half + the widest e2m1 grid gap, plus saturation headroom).""" + ref_bf16 = ref.to(torch.bfloat16).to(torch.float32) + sf_ref = _nvfp4_sf_ref(ref, norm_const) + sf_flat = _swizzled_sf_to_flat(sf_tensor, ref.shape[0], ref.shape[1]) if sf_tensor.dim() == 6 else sf_tensor + torch.testing.assert_close( + sf_flat.float().cpu(), + sf_ref.float().cpu(), + atol=1e-2, + rtol=0.14, + msg=lambda m: f"{name} block scales mismatch\n{m}", + ) + + decode_scale = sf_flat.float().repeat_interleave(HADAMARD_SIZE, dim=1) / norm_const + dequant = values * decode_scale + err = (dequant - ref_bf16).abs() + bound = 1.5 * decode_scale + 5e-2 + bad = err > bound + assert not bad.any(), f"{name}: {int(bad.sum())} dequantized elements exceed the quantization error bound (max err {err[bad].max().item():.4f})" + + +# ============================================================================= +# Output checking +# ============================================================================= + + +def _check_outputs( + inputs: Dict, + outputs: Dict, + cfg: Dict, + *, + act_func: str, + rht_output: bool, + rht_rowwise: bool, + glu_alpha: Optional[float] = None, + glu_limit: Optional[float] = None, + norm_const: float = 1.0, + rht_norm_const: float = 1.0, + ref_tensors: Optional[Dict] = None, +) -> None: + if ref_tensors is None: + ref_tensors = _run_grouped_gemm_glu_ref(inputs, act_func, glu_alpha=glu_alpha, glu_limit=glu_limit) + valid_m = inputs["valid_m"] + c_ref = ref_tensors["c_ref"] + d_ref = ref_tensors["d_ref"] + + torch.testing.assert_close( + outputs["c_tensor"][:valid_m].cpu().float(), + c_ref.cpu().to(cfg["c_dtype"]).to(torch.float32), + atol=1e-1, + rtol=1e-2, + ) + + if outputs["sfd_tensor"] is None: + torch.testing.assert_close( + outputs["d_tensor"][:valid_m].cpu().float(), + d_ref.cpu().to(torch.bfloat16).to(torch.float32), + atol=1e-1, + rtol=1e-2, + ) + else: + _check_nvfp4_output( + float4_e2m1fn_x2_to_float32(outputs["d_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()), + outputs["sfd_tensor"].cpu(), + d_ref[:, :, 0].cpu(), + norm_const, + "D", + ) + + if not rht_output: + assert outputs["rht_tensor"] is None + assert outputs["sfrht_tensor"] is None + return + + rht_ref = _rht_ref(d_ref.cpu(), rht_rowwise) + if outputs["sfrht_tensor"] is None: + torch.testing.assert_close( + outputs["rht_tensor"][:valid_m, :, 0].cpu().float(), + rht_ref.float(), + atol=1e-1, + rtol=1e-2, + ) + elif rht_rowwise: + _check_nvfp4_output( + float4_e2m1fn_x2_to_float32(outputs["rht_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()), + outputs["sfrht_tensor"].cpu(), + rht_ref.float(), + rht_norm_const, + "RHT", + ) + else: + # Colwise: the packed data stays at D's (m, f) orientation (nibbles pair + # adjacent features), but quantization blocks are (16, 1) token blocks, + # so check through the transposed unpacked values and the swizzled + # SF(N_out, valid_m) scale domain. + _check_nvfp4_output( + float4_e2m1fn_x2_to_float32(outputs["rht_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()).t(), + outputs["sfrht_tensor"].cpu(), + rht_ref.float().t(), + rht_norm_const, + "RHT", + ) + + +# ============================================================================= +# Runners +# ============================================================================= + + +def _run_wrapper( + request, + *, + ab_dtype, + sf_dtype, + sf_vec_size, + act_func="swiglu", + enable_bias=False, + d_dtype=torch.bfloat16, + rht_output=True, + rht_dtype=torch.bfloat16, + rht_rowwise=False, + glu_alpha=None, + glu_limit=None, + sf_fp8_dtype_override=None, +): + cfg = _make_cfg(request, ab_dtype=ab_dtype, sf_dtype=sf_dtype, sf_vec_size=sf_vec_size, enable_bias=enable_bias) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + enable_bias=enable_bias, + ) + + # NVFP4 global encode scales, derived from the reference like production + # does from calibration (2688 = 448 * 6): without them the e4m3 block + # scales saturate on this test data's magnitude. + ref_tensors = _run_grouped_gemm_glu_ref(inputs, act_func, glu_alpha=glu_alpha, glu_limit=glu_limit) + norm_const = 1.0 + rht_norm_const = 1.0 + if d_dtype == torch.float4_e2m1fn_x2: + norm_const = 2688.0 / ref_tensors["d_ref"].to(torch.bfloat16).float().abs().max().item() + if rht_output and rht_dtype == torch.float4_e2m1fn_x2: + rht_ref = _rht_ref(ref_tensors["d_ref"].cpu(), rht_rowwise) + rht_norm_const = 2688.0 / rht_ref.float().abs().max().item() + + from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 + + outputs = grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + bias_tensor=inputs["bias_tensor"], + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=d_dtype, + cd_major=cfg["cd_major"], + rht_output=rht_output, + rht_dtype=rht_dtype, + rht_rowwise=rht_rowwise, + glu_alpha=glu_alpha, + glu_limit=glu_limit, + norm_const=norm_const, + rht_norm_const=rht_norm_const, + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func=act_func, + ) + + _check_outputs( + inputs, + outputs, + cfg, + act_func=act_func, + rht_output=rht_output, + rht_rowwise=rht_rowwise, + glu_alpha=glu_alpha, + glu_limit=glu_limit, + norm_const=norm_const, + rht_norm_const=rht_norm_const, + ref_tensors=ref_tensors, + ) + + +def _run_compile_execute(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func="swiglu", sf_fp8_dtype_override=None): + cfg = _make_cfg(request, ab_dtype=ab_dtype, sf_dtype=sf_dtype, sf_vec_size=sf_vec_size) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + enable_bias=False, + ) + + valid_m = inputs["valid_m"] + n = cfg["n"] + n_out = n // 2 + device = inputs["a_tensor"].device + + def alloc_n_major(rows, cols, dtype): + return torch.empty_strided((rows, cols, 1), (cols, 1, rows * cols), dtype=dtype, device=device) + + outputs = { + "c_tensor": alloc_n_major(valid_m, n, cfg["c_dtype"]), + "d_tensor": alloc_n_major(valid_m, n_out, torch.bfloat16), + "sfd_tensor": None, + "rht_tensor": alloc_n_major(valid_m, n_out, torch.bfloat16), + "sfrht_tensor": None, + } + + from cudnn import GroupedGemmGluHadamardQuantSm100 + + api = GroupedGemmGluHadamardQuantSm100( + sample_a=inputs["a_tensor"], + sample_b=inputs["b_tensor"], + sample_c=outputs["c_tensor"], + sample_d=outputs["d_tensor"], + sample_sfa=inputs["sfa_tensor"], + sample_sfb=inputs["sfb_tensor"], + sample_padded_offsets=inputs["padded_offsets_tensor"], + sample_alpha=inputs["alpha_tensor"], + sample_prob=inputs["prob_tensor"], + sample_rht=outputs["rht_tensor"], + acc_dtype=cfg["acc_dtype"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func=act_func, + ) + api.check_support() + api.compile() + api.execute( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + c_tensor=outputs["c_tensor"], + d_tensor=outputs["d_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + rht_tensor=outputs["rht_tensor"], + ) + + _check_outputs(inputs, outputs, cfg, act_func=act_func, rht_output=True, rht_rowwise=False) + + +def _run_discrete_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func="swiglu", sf_fp8_dtype_override=None): + cfg = _make_cfg(request, ab_dtype=ab_dtype, sf_dtype=sf_dtype, sf_vec_size=sf_vec_size) + inputs = allocate_discrete_input_tensors( + n=cfg["n"], + k=cfg["k"], + num_experts=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + ) + inputs["b_ref"] = torch.cat(inputs["b_ref_list"], dim=2) + inputs["sfb_ref"] = torch.cat(inputs["sfb_ref_list"], dim=2) + + from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 + + outputs = grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_ptrs=inputs["b_ptrs_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_ptrs=inputs["sfb_ptrs_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + n=cfg["n"], + b_dtype=inputs["b_list"][0].dtype, + b_major=cfg["b_major"], + bias_tensor=inputs["bias_tensor"], + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=torch.bfloat16, + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func=act_func, + ) + + _check_outputs(inputs, outputs, cfg, act_func=act_func, rht_output=True, rht_rowwise=False) + + +# ============================================================================= +# Tests +# ============================================================================= + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("ab_dtype,sf_dtype,sf_vec_size", FP4_EXECUTION_CASES) +def test_grouped_gemm_glu_hadamard_quant_compile_execute_fp4(request, ab_dtype, sf_dtype, sf_vec_size): + _run_compile_execute( + request, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("act_func", ["swiglu", "geglu"]) +def test_grouped_gemm_glu_hadamard_quant_wrapper_fp4(request, act_func): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + act_func=act_func, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_rowwise(request): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + rht_rowwise=True, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_no_rht(request): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + rht_output=False, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_with_bias(request): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=16, + enable_bias=True, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_glu_alpha_limit(request): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + glu_alpha=1.702, + glu_limit=7.0, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_d(request): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + d_dtype=torch.float4_e2m1fn_x2, + rht_output=False, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("rht_rowwise", [False, True]) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_rht(request, rht_rowwise): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + rht_dtype=torch.float4_e2m1fn_x2, + rht_rowwise=rht_rowwise, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_full(request): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + d_dtype=torch.float4_e2m1fn_x2, + rht_dtype=torch.float4_e2m1fn_x2, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_discrete_fp4(request): + _run_discrete_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=16, + act_func="swiglu", + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("group_m_list", [[256, 256, 256, 256], DYNAMIC_SHAPES_M_VALUES]) +def test_grouped_gemm_glu_hadamard_quant_wrapper_cache_dynamic_m_smoke(request, monkeypatch, group_m_list): + from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 + import cudnn.gemm.cutedsl.grouped.glu_hadamard_quant.api as grouped_gemm_glu_hadamard_quant_api + + grouped_gemm_glu_hadamard_quant_api._cache_of_GroupedGemmGluHadamardQuantSm100Objects.clear() + + compile_count = {"value": 0} + + def counted_compile(self): + compile_count["value"] += 1 + return None + + monkeypatch.setattr(grouped_gemm_glu_hadamard_quant_api.GroupedGemmGluHadamardQuantSm100, "compile", counted_compile) + monkeypatch.setattr(grouped_gemm_glu_hadamard_quant_api.GroupedGemmGluHadamardQuantSm100, "check_support", lambda self: True) + monkeypatch.setattr(grouped_gemm_glu_hadamard_quant_api.GroupedGemmGluHadamardQuantSm100, "execute", lambda self, **kwargs: None) + + cfg = _make_cfg( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=16, + ) + cfg["group_m_list"] = list(group_m_list) + cfg["l"] = len(group_m_list) + + for _ in range(2): + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + enable_bias=False, + ) + grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func="swiglu", + ) + + assert compile_count["value"] == 1 + assert len(grouped_gemm_glu_hadamard_quant_api._cache_of_GroupedGemmGluHadamardQuantSm100Objects) == 1 + grouped_gemm_glu_hadamard_quant_api._cache_of_GroupedGemmGluHadamardQuantSm100Objects.clear() From 480050c55a6d8625daf3e9ea3db741cc7f043f2a Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 17 Aug 2026 18:23:45 -0700 Subject: [PATCH 2/3] Add e5m3 scale-factor coverage to glu_hadamard_quant tests The kernel plumbs sf_fp8_dtype_override through every entry point but no test ever passed "e5m3". Mirror the #545 test pattern: reencode the e4m3-storage input scales as UE5M3 bytes in place (values exact in both formats, so the fp32 reference stays valid), plus compile-cache separation and unsupported-override rejection tests. Co-Authored-By: Claude Fable 5 --- .../test_grouped_gemm_glu_hadamard_quant.py | 155 +++++++++++++++++- 1 file changed, 154 insertions(+), 1 deletion(-) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py index b6c787f67..96ae2b630 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py @@ -10,7 +10,8 @@ from test_utils import torch_fork_set_rng from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import allocate_discrete_input_tensors from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import allocate_grouped_gemm_input_tensors, grouped_gemm_swiglu_init -from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES +from fe_api.grouped_gemm.test_grouped_gemm_wgrad_utils import _skip_unless_e5m3_supported +from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES, reencode_sf_tensor_as_ue5m3 FP4_EXECUTION_CASES = [ (torch.float4_e2m1fn_x2, torch.float8_e4m3fn, 16), @@ -316,6 +317,12 @@ def _run_wrapper( rht_ref = _rht_ref(ref_tensors["d_ref"].cpu(), rht_rowwise) rht_norm_const = 2688.0 / rht_ref.float().abs().max().item() + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so the fp32 reference stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 outputs = grouped_gemm_glu_hadamard_quant_wrapper_sm100( @@ -377,6 +384,12 @@ def _run_compile_execute(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func=" enable_bias=False, ) + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so the fp32 reference stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + valid_m = inputs["valid_m"] n = cfg["n"] n_out = n // 2 @@ -598,6 +611,146 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_full(request): ) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("rht_rowwise", [False, True], ids=["colwise", "rowwise"]) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_rht_e5m3(request, rht_rowwise): + """quant_rht with the input block scales carried as UE5M3 bytes in e4m3 storage.""" + _skip_unless_e5m3_supported() + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + rht_dtype=torch.float4_e2m1fn_x2, + rht_rowwise=rht_rowwise, + sf_fp8_dtype_override="e5m3", + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_full_e5m3(request): + """quant_full with the input block scales carried as UE5M3 bytes in e4m3 storage.""" + _skip_unless_e5m3_supported() + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + d_dtype=torch.float4_e2m1fn_x2, + rht_dtype=torch.float4_e2m1fn_x2, + sf_fp8_dtype_override="e5m3", + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_compile_execute_e5m3(request): + """Class-API compile/execute path with e5m3-reinterpreted input scales.""" + _skip_unless_e5m3_supported() + _run_compile_execute( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e4m3fn, + sf_vec_size=16, + sf_fp8_dtype_override="e5m3", + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_quant_e5m3_is_not_cached_as_e4m3(request): + """sf_fp8_dtype_override must take part in the compile cache key. + + Identical scale-factor bytes decode to different values under E4M3 and + UE5M3, so if sf_fp8_dtype_override were omitted from the key the second + call would reuse the first kernel and silently return E4M3 results. + """ + _skip_unless_e5m3_supported() + + # One problem, one set of scale-factor bytes, two interpretations. Any + # difference in the output can only come from sf_fp8_dtype_override. + cfg = _make_cfg(request, ab_dtype=torch.float4_e2m1fn_x2, sf_dtype=torch.float8_e4m3fn, sf_vec_size=16) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + enable_bias=False, + ) + + from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 + + def run(sf_fp8_dtype_override): + outputs = grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + bias_tensor=inputs["bias_tensor"], + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=torch.bfloat16, + cd_major=cfg["cd_major"], + rht_output=True, + rht_dtype=torch.bfloat16, + rht_rowwise=False, + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func="swiglu", + ) + return outputs["d_tensor"].float().clone() + + d_e4m3 = run(None) + d_e5m3 = run("e5m3") + + assert not torch.equal( + d_e5m3, d_e4m3 + ), "e5m3 and e4m3 produced identical output from identical scale-factor bytes; sf_fp8_dtype_override is likely missing from the compile cache key" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + "sf_fp8_dtype_override,overrides,expected", + [ + pytest.param( + "e5m3", + dict(sf_dtype=torch.float8_e8m0fnu), + "requires torch.float8_e4m3fn scale-factor storage", + id="e8m0_carrier", + ), + pytest.param("e4m3", {}, "sf_fp8_dtype_override must be", id="e4m3_is_not_an_override"), + pytest.param("e5m2", {}, "sf_fp8_dtype_override must be", id="unknown_format"), + ], +) +def test_grouped_gemm_glu_hadamard_quant_rejects_unsupported_sf_fp8_dtype(request, sf_fp8_dtype_override, overrides, expected): + """e5m3 is only reachable through the FP4xFP4 atom with e4m3-carried scales.""" + if sf_fp8_dtype_override == "e5m3": + _skip_unless_e5m3_supported() + with pytest.raises(ValueError, match=expected): + _run_compile_execute( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=overrides.get("sf_dtype", torch.float8_e4m3fn), + sf_vec_size=overrides.get("sf_vec_size", 16), + sf_fp8_dtype_override=sf_fp8_dtype_override, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_grouped_gemm_glu_hadamard_quant_wrapper_discrete_fp4(request): From a95520a4dda4a0a91e6e7235fd1cd5abdba98318 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 18 Aug 2026 09:25:56 +0000 Subject: [PATCH 3/3] Add Rubin kernel Co-authored-by: Codex Signed-off-by: Tim Moon --- .../cutedsl/grouped/glu_hadamard_quant/api.py | 43 +- ...kscaled_grouped_gemm_glu_hadamard_quant.py | 58 +- ...d_grouped_gemm_glu_hadamard_quant_rubin.py | 3088 +++++++++++++++++ .../glu_hadamard_quant/moe_kernel_helpers.py | 2 + .../grouped/glu_hadamard_quant/quant_utils.py | 44 +- .../grouped/glu_hadamard_quant/rht_utils.py | 23 +- .../test_grouped_gemm_glu_hadamard_quant.py | 103 +- test/python/fe_api/test_fe_api_utils.py | 6 + 8 files changed, 3290 insertions(+), 77 deletions(-) create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant_rubin.py diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py index e4c67e16c..133dda95b 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py @@ -4,9 +4,9 @@ """FE API for grouped GEMM GLU forward fusion with fused RHT output and NVFP4 quantization. Output modes are dtype driven, mirroring the kernel: - - D is bf16, or NVFP4 (packed e2m1 data + e4m3 block scales in ``sfd``) - - The optional RHT output is bf16, or NVFP4 (packed e2m1 data + e4m3 block scales - in ``sfrht``) + - D is bf16, or NVFP4 (packed e2m1 data + e4m3/ue5m3 block scales in ``sfd``) + - The optional RHT output is bf16, or NVFP4 (packed e2m1 data + e4m3/ue5m3 + block scales in ``sfrht``) The RHT data is always stored at D's own (m, f) orientation — only the SCALE grid follows the transform orientation: swizzled scale factors for logical (m, f) @@ -17,7 +17,7 @@ import logging import os -from typing import Any, Optional, Tuple +from typing import Any, Literal, Optional, Tuple from cuda.bindings import driver as cuda import cutlass @@ -25,13 +25,21 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.runtime import from_dlpack, make_fake_stream -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import APIBase, TupleDict, ceil_div, get_device_type, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type from ..moe_utils import MoEWeightMode from .rht_utils import HADAMARD_SIZE from .moe_blockscaled_grouped_gemm_glu_hadamard_quant import BlockScaledMoEGroupedGemmGluHadamardQuantKernel + +def _get_rubin_kernel(): + from .moe_blockscaled_grouped_gemm_glu_hadamard_quant_rubin import ( + BlockScaledMoEGroupedGemmGluHadamardQuantKernel as RubinBlockScaledMoEGroupedGemmGluHadamardQuantKernel, + ) + + return RubinBlockScaledMoEGroupedGemmGluHadamardQuantKernel + # The GLU + Hadamard + quant fusion is block-scaled only: its mandatory # scale-factor inputs use an MMA-interleaved 6-D layout with no row-major # equivalent, so they are not expressible as JAX arrays and the API stays @@ -83,7 +91,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, - sf_fp8_dtype_override: Optional[str] = None, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, act_func: str = "swiglu", @@ -161,7 +169,7 @@ def __init__( self.rht_rowwise = rht_rowwise self.glu_alpha = glu_alpha self.glu_limit = glu_limit - self._kernel = BlockScaledMoEGroupedGemmGluHadamardQuantKernel + self._kernel = _get_rubin_kernel() if self._is_rubin_kernel else BlockScaledMoEGroupedGemmGluHadamardQuantKernel self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) self._workspace = None @@ -258,11 +266,11 @@ def check_support(self) -> bool: self._check_dtype(self.prob_desc, dtype=torch.float32, name="prob") self._check_dtype(self.bias_desc, dtype=[torch.float16, torch.bfloat16, torch.float32], name="bias") if self.d_quant: - self._check_dtype(self.sfd_desc, dtype=torch.float8_e4m3fn, name="SFD") + self._check_dtype(self.sfd_desc, dtype=self.sf_dtype, name="SFD", extra_error_msg="SFD must match SFA dtype") if self.generate_rht: self._check_dtype(self.rht_desc, dtype=[torch.bfloat16, torch.float4_e2m1fn_x2], name="RHT") if self.rht_quant: - self._check_dtype(self.sfrht_desc, dtype=torch.float8_e4m3fn, name="SFRHT") + self._check_dtype(self.sfrht_desc, dtype=self.sf_dtype, name="SFRHT", extra_error_msg="SFRHT must match SFA dtype") self._check_dtype(self.acc_dtype, dtype=torch.float32, name="acc_dtype") self._value_error_if(self.sf_vec_size != 16, f"sf_vec_size must be 16, got {self.sf_vec_size}") @@ -301,8 +309,8 @@ def check_support(self) -> bool: f"Invalid cluster shape: {self.cluster_shape_mn}", ) self._value_error_if( - self.m_aligned != BlockScaledMoEGroupedGemmGluHadamardQuantKernel.FIX_PAD_SIZE, - f"m_aligned must be {BlockScaledMoEGroupedGemmGluHadamardQuantKernel.FIX_PAD_SIZE}, got {self.m_aligned}", + self.m_aligned != self._kernel.FIX_PAD_SIZE, + f"m_aligned must be {self._kernel.FIX_PAD_SIZE}, got {self.m_aligned}", ) self._value_error_if(self.expert_cnt > 1024, f"expert_cnt must be <= 1024, got {self.expert_cnt}") @@ -347,7 +355,7 @@ def compile(self) -> None: if self.a_desc.shape[0] == 0: return - kernel = self._kernel( + kernel_kwargs = dict( sf_vec_size=self.sf_vec_size, acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), use_2cta_instrs=self.use_2cta_instrs, @@ -360,10 +368,12 @@ def compile(self) -> None: act_func=self.act_func, enable_bias=self.bias_desc is not None, rht_rowwise=self.rht_rowwise if self.generate_rht else False, - sf_fp8_dtype_override=self.sf_fp8_dtype_override, glu_alpha=self.glu_alpha, glu_limit=self.glu_limit, ) + if self.sf_fp8_dtype_override == "e5m3": + kernel_kwargs["sf_fp8_dtype_override"] = self.sf_fp8_dtype_override + kernel = self._kernel(**kernel_kwargs) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) @@ -720,7 +730,7 @@ def grouped_gemm_glu_hadamard_quant_wrapper_sm100( """High-level wrapper for grouped GEMM GLU forward fusion with fused RHT output. Output modes are dtype driven: ``d_dtype``/``rht_dtype`` of - ``torch.float4_e2m1fn_x2`` emit packed NVFP4 data plus e4m3 block scales + ``torch.float4_e2m1fn_x2`` emit packed NVFP4 data plus e4m3/ue5m3 block scales (``sfd_tensor``/``sfrht_tensor``); ``torch.bfloat16`` emits plain bf16. ``sf_fp8_dtype_override="e5m3"`` reinterprets ``torch.float8_e4m3fn`` SFA/SFB storage as UE5M3 input scale factors on Rubin. @@ -790,7 +800,7 @@ def alloc_n_major(rows: int, cols: int, dtype: torch.dtype) -> torch.Tensor: def alloc_swizzled_sf(rows: int, cols: int) -> torch.Tensor: shape = (1, ceil_div(rows, 128), ceil_div(ceil_div(cols, sf_vec_size), 4), 32, 4, 4) - return torch.empty(shape, dtype=torch.float8_e4m3fn, device=device).permute(3, 4, 1, 5, 2, 0) + return torch.empty(shape, dtype=sfa_tensor.dtype, device=device).permute(3, 4, 1, 5, 2, 0) c_tensor = alloc_n_major(valid_m, n_full, c_dtype) if d_quant: @@ -828,7 +838,10 @@ def dynamic_m_tensor_signature( stride_signature = tuple(None if idx in dynamic_stride_dims else value for idx, value in enumerate(tensor.stride())) return static_shape_suffix, stride_signature, tensor.dtype + device_type = get_device_type() + cache_key = ( + device_type, weight_mode, act_func, a_tensor.shape[1:], diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py index 5e656fd78..d8afe6883 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant.py @@ -33,7 +33,6 @@ - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler) - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout - BF16 or NVFP4 (packed e2m1 + e4m3 block scales) D output with GLU activation - - Optional UE5M3 reinterpretation for torch.float8_e4m3fn input scale storage - Optional C output (pre-activation GLU output) - Optional fused RHT output (bf16 or NVFP4 + e4m3 block scales) - GLU activation fusion (SwiGLU / GeGLU) @@ -177,7 +176,6 @@ def __init__( act_func: str = "swiglu", enable_bias: bool = False, rht_rowwise: bool = False, - sf_fp8_dtype_override: Optional[str] = None, glu_alpha: Optional[float] = None, glu_limit: Optional[float] = None, ): @@ -198,7 +196,6 @@ def __init__( self.weight_mode = weight_mode self.use_dynamic_sched = use_dynamic_sched self.enable_bias = enable_bias - self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None # RHT dump orientation: False = columnwise (16-token blocks per feature), # True = rowwise (16-feature blocks per token). Same dump tensor/path either way. self.rht_rowwise = rht_rowwise @@ -612,9 +609,7 @@ def __call__( self.b_dtype: Type[cutlass.Numeric] = a.element_type self.c_dtype: Type[cutlass.Numeric] = c.element_type self.d_dtype: Type[cutlass.Numeric] = d.element_type - # UE5M3 has no torch dtype, so callers provide e4m3 storage and ask the - # kernel to reinterpret the scale bytes at compile time. - self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override if cutlass.const_expr(self.sf_dtype_override is not None) else sfa.element_type + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(c) @@ -646,19 +641,27 @@ def __call__( # after num_tile_stage tiles). self.run_rht = self.generate_rht or self.d_quant if cutlass.const_expr(not self.run_rht): - self.threads_wo_sched = self.threads_per_cta - self.threads_per_warp - self.threads_per_warp * len(self.epilog_rht_store_warp_id) + self.threads_wo_sched = ( + self.threads_per_cta + - self.threads_per_warp + - self.threads_per_warp * len(self.epilog_rht_store_warp_id) + ) if cutlass.const_expr(self.d_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): raise ValueError(f"d dtype must be BFloat16 or Float4E2M1FN, got {self.d_dtype}") if cutlass.const_expr(self.d_quant != self.generate_sfd): raise ValueError("NVFP4 d and sfd must be passed together") - if cutlass.const_expr(self.generate_rht and self.rht_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): + if cutlass.const_expr(self.generate_rht + and self.rht_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): raise ValueError(f"rht dtype must be BFloat16 or Float4E2M1FN, got {self.rht_dtype}") if cutlass.const_expr(self.rht_quant != self.generate_sfrht): raise ValueError("NVFP4 rht and sfrht must be passed together") - if cutlass.const_expr(self.generate_sfd and sfd.element_type != cutlass.Float8E4M3FN): - raise ValueError("sfd element type must be Float8E4M3FN") - if cutlass.const_expr(self.generate_sfrht and sfrht.element_type != cutlass.Float8E4M3FN): - raise ValueError("sfrht element type must be Float8E4M3FN") + sf_storage_dtype = ( + cutlass.Float8E4M3FN if self.sf_dtype == cutlass.FloatNV8E5M3FNU else self.sf_dtype + ) + if cutlass.const_expr(self.generate_sfd and sfd.element_type != sf_storage_dtype): + raise ValueError("sfd element type must match scale-factor storage dtype") + if cutlass.const_expr(self.generate_sfrht and sfrht.element_type != sf_storage_dtype): + raise ValueError("sfrht element type must match scale-factor storage dtype") if cutlass.const_expr((self.d_quant or self.rht_quant) and self.act_func == "srelu"): raise ValueError("NVFP4 quantization assumes the GLU subtile pair-step (act_func != srelu)") @@ -905,8 +908,9 @@ class SharedStorage: ] sSfd: cute.struct.Align[ cute.struct.MemRange[ - cutlass.Float8E4M3FN, - self.threads_per_warp * len(self.epilog_rht_store_warp_id) * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), + self.sf_dtype, + self.threads_per_warp * len(self.epilog_rht_store_warp_id) + * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), ], 16, ] @@ -922,8 +926,9 @@ class SharedStorage: # Same (128, 8) buffer either way. sSfRht: cute.struct.Align[ cute.struct.MemRange[ - cutlass.Float8E4M3FN, - self.threads_per_warp * len(self.epilog_rht_store_warp_id) * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), + self.sf_dtype, + self.threads_per_warp * len(self.epilog_rht_store_warp_id) + * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), ], 16, ] @@ -1024,6 +1029,7 @@ def _make_extension(self, workspace_ptr): @cute.jit def store_swizzled_sf_row(self, sf_tensor: cute.Tensor, logical_row, sf_col_base, sSf: cute.Tensor, tidx): """Store one logical scale row into M32x4xrm_K4xrk_L SF layout.""" + sf_tensor = cute.recast_tensor(sf_tensor, self.sf_dtype) row_m0 = logical_row % 32 row_m1 = (logical_row // 32) % 4 row_m2 = logical_row // 128 @@ -2504,7 +2510,8 @@ def kernel( if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): # Expert token offset for the colwise (f, m) scale grid's # tile index (offsets are 256-aligned, divisions exact). - rht_t_off, _rht_t_cnt = compute_expert_token_range(padded_offsets, epi_work_tile_info.expert_idx) + rht_t_off, _rht_t_cnt = compute_expert_token_range( + padded_offsets, epi_work_tile_info.expert_idx) # # NVFP4 D: per-expert fp4 D gmem tensor + TMA partition (mirrors ACT's D setup). @@ -2548,7 +2555,10 @@ def kernel( real_subtile_idx = subtile_idx // 2 if cutlass.const_expr(self.overlapping_accum): if rht_reverse: - real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n_required - 1 - real_subtile_idx + real_subtile_idx = ( + self.cta_tile_shape_mnk[1] // self.epi_tile_n_required + - 1 - real_subtile_idx + ) # # Load this subtile's sD values to registers, then ARRIVE the # lockstep barrier immediately: the ACT warps' arrive_and_wait @@ -2570,16 +2580,20 @@ def kernel( # from the sRht dtype inside the FWHT device functions. if cutlass.const_expr(self.generate_rht): if cutlass.const_expr(self.rht_quant and self.rht_rowwise): - hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, rht_norm_const, sSfRht, real_subtile_idx) + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, + rht_norm_const, sSfRht, real_subtile_idx, self.sf_dtype) elif cutlass.const_expr(self.rht_quant): - hadamard_rmem_colwise_fwht_quant(rht_ld, d_buffer, epi_tidx, rht_norm_const, sRht, sSfRht, real_subtile_idx * 2 * HADAMARD_SIZE) + hadamard_rmem_colwise_fwht_quant( + rht_ld, d_buffer, epi_tidx, rht_norm_const, + sRht, sSfRht, real_subtile_idx * 2 * HADAMARD_SIZE, self.sf_dtype) elif cutlass.const_expr(self.rht_rowwise): - hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht) + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, 1.0, None, 0, self.sf_dtype) else: hadamard_rmem_colwise_fwht(rht_ld, d_buffer, epi_tidx, sRht) # NVFP4 D: quantize the bf16 register rows into sDq + sSfd. if cutlass.const_expr(self.d_quant): - nvfp4_quant_rmem_row(dq_ld, d_buffer, epi_tidx, sDq, norm_const, sSfd, real_subtile_idx) + nvfp4_quant_rmem_row(dq_ld, d_buffer, epi_tidx, sDq, + norm_const, sSfd, real_subtile_idx, self.sf_dtype) # # TMA-store the produced epi-tiles to gmem (mirrors ACT's D # store, including the overlapping_accum subtile-column reversal). diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant_rubin.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant_rubin.py new file mode 100644 index 000000000..ebc371733 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant_rubin.py @@ -0,0 +1,3088 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +MoE Block-Scaled Grouped GEMM Kernel with GLU (SwiGLU/GeGLU) + Hadamard Transform +Fusion — RUBIN (sm_107) port of kernel.py. + +Rubin deltas vs the Blackwell body (the proven recipe, see also +glu_hadamard/kernel_sm107.py): sm107_utils tiled MMAs with explicit (M, N, K_inst) +(K_inst = 128 for FP4xFP4, else 64), a/b_collector_op=DISCARD, identity +atom_layout/permutation; mma_inst_tile_k = 2 for FP4/sf16; SF TMEM columns x +sf_pack_factor (32 // sf_vec_size); SF s2t partition through an appended MN +broadcast-4 mode; sm_107 smem/tmem capacities; TmemAllocator(arch="sm_107"). +No B-reuse (measured slower at the primary shape). + +Supports: + - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler) + - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout + - BF16 or NVFP4 (packed e2m1 + e4m3/ue5m3 block scales) D output with GLU activation + - Optional C output (pre-activation GLU output) + - Optional fused RHT output (bf16 or NVFP4 + e4m3/ue5m3 block scales) + - GLU activation fusion (SwiGLU / GeGLU) + +Warp assignment (8 epilogue warps, pingpong): + warps 0-3 : ACT warps — TMEM→reg, alpha scale, GLU activation, C/D store + warps 4-7 : RHT store warps — RHT and/or NVFP4 quantization from D SMEM + warp 8 : MMA warp + warp 9 : TMA load warp + warp 10 : Scheduler warp (MoEPersistentTileScheduler) + warp 11 : Bias load warp (optional) + +sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + Validity: tile_info[0] >= 0 (expert_idx == -1 signals end) +""" + +from typing import Literal, Type, Tuple, Union, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.nvgpu.tcgen05 import CollectorOp +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.rubin_helpers as sm107_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.typing import Float32, Int32, AddressSpace +from ..moe_persistent_scheduler import ( + MoEPersistentTileScheduler, + MoESchedulerParams, + MoEWorkTileInfo, +) +from ..moe_utils import ( + compute_expert_token_range, + MoEWeightMode, + TensormapWorkspace, + store_tma_desc, +) +from .rht_utils import ( + hadamard_rmem_colwise_fwht, + hadamard_rmem_colwise_fwht_quant, + hadamard_rmem_rowwise_fwht, + load_colwise_pairs_bf16, + HADAMARD_SIZE, +) +from .quant_utils import load_row_bf16, nvfp4_quant_rmem_row +from ..moe_sched_extension import ( + DiscreteWeightScaledGemmSchedExtension, + ContiguousAndConsistentGroupedGemmSchedExtension, +) +from .moe_kernel_helpers import ( + fmin, + fmax, + silu_f32, + silu_f32_geglu_scaled, + compute_grid, + can_implement, +) + + +# Valid launch-config space for THIS kernel (the wrapper's autotune candidates; +# can_implement still prunes per problem). cta shape == mma_tiler_mn; Rubin +# extends the set with the (512, 256) B-reuse tile. +VALID_CTA_SHAPES = ((256, 256), (512, 256)) +VALID_CLUSTER_SHAPES = ((1, 1), (2, 1), (2, 2), (4, 1), (4, 2)) +DEFAULT_CTA_SHAPE = (256, 256) +DEFAULT_CLUSTER_SHAPE = (2, 1) + + +class BlockScaledMoEGroupedGemmGluHadamardQuantKernel: + """Block-scaled MoE grouped GEMM with GLU activation and Hadamard transform fusion. + + Always uses pingpong epilogue (8 epilogue warps: 4 ACT + 4 RHT-store). + D output is BF16 or F16 only (no FP8/FP4, no SFD). + + :param sf_vec_size: Scalefactor vector size. + :param mma_tiler_mn: Shape of MMA tile (M, N). + :param cluster_shape_mn: Cluster dimensions (M, N). + :param expert_cnt: Number of experts (compile-time constant). + :param weight_mode: Dense or Discrete weight layout. + :param use_dynamic_sched: Use dynamic tile scheduling. + :param act_func: Activation function ('swiglu', 'geglu', or 'srelu'). + :param enable_bias: Enable bias addition. + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. + """ + + FIX_PAD_SIZE = 256 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + acc_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + rht_quant: bool = False, + d_quant: bool = False, + ) -> bool: + # NVFP4 quantization: scale-byte rows are stored as one contiguous 8-byte + # store per thread, so f = n/2 must be divisible by 128 (f/16 % 8 == 0; + # same shape gate as the standalone group_rht_cast kernel). + if (rht_quant or d_quant) and (n // 2) % 128 != 0: + return False + # B-reuse: 2CTA + mma tiler M=512 (M-split with B held in the collector). + # RHT fusion keeps the N=256-only constraint. + if use_2cta_instrs and mma_tiler_mn[0] == 512: + from .moe_kernel_helpers import ( + is_valid_dtypes_and_scale_factor_vec_size, + is_valid_layouts, is_valid_tensor_alignment, + ) + if not is_valid_dtypes_and_scale_factor_vec_size(ab_dtype, sf_dtype, sf_vec_size, acc_dtype, d_dtype): + return False + if not is_valid_layouts(ab_dtype, d_dtype, a_major, b_major, cd_major): + return False + if not is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, cd_major): + return False + if mma_tiler_mn[1] != 256: + return False + # Cluster M must be a multiple of 2 (2CTA M-split). + if cluster_shape_mn[0] % 2 != 0: + return False + return True + # special requirements for hadamard fusion + if not use_2cta_instrs or mma_tiler_mn[0] != 256 or mma_tiler_mn[1] != 256: + return False + return can_implement( + ab_dtype, + sf_dtype, + sf_vec_size, + acc_dtype, + d_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + cd_major, + m_aligned, + fix_pad_size=BlockScaledMoEGroupedGemmGluHadamardQuantKernel.FIX_PAD_SIZE, + ) + + def __init__( + self, + sf_vec_size: int, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + vectorized_f32: bool, + expert_cnt: int, + weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, + use_dynamic_sched: bool = False, + act_func: str = "swiglu", + enable_bias: bool = False, + rht_rowwise: bool = False, + sf_fp8_dtype_override: Optional[str] = None, + glu_alpha: Optional[float] = None, + glu_limit: Optional[float] = None, + ): + # B-reuse: mma tiler M = 2x the MMA instruction M (2CTA M-split, B held + # in the tcgen05 collector buffer across the two A halves). + mma_inst_m = 256 if use_2cta_instrs else 128 + enable_breuse = (mma_tiler_mn[0] // mma_inst_m == 2) + if enable_breuse: + if mma_tiler_mn[0] % self.FIX_PAD_SIZE != 0: + raise ValueError( + f"mma_tiler_mn[0] ({mma_tiler_mn[0]}) must be a multiple of " + f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) for breuse. " + f"Also use m_aligned=mma_tiler_mn[0] (={mma_tiler_mn[0]})." + ) + else: + mma_tile_m = mma_tiler_mn[0] + if self.FIX_PAD_SIZE % mma_tile_m != 0: + raise ValueError(f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) must be divisible by " f"mma_tiler_mn[0] ({mma_tile_m}).") + if expert_cnt > 1024: + raise ValueError("Expert count > 1024 is not supported.") + if not isinstance(weight_mode, MoEWeightMode): + raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") + + self.sf_vec_size = sf_vec_size + self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None + self.expert_cnt = expert_cnt + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.enable_breuse = enable_breuse + self.weight_mode = weight_mode + self.use_dynamic_sched = use_dynamic_sched + self.enable_bias = enable_bias + # RHT dump orientation: False = columnwise (16-token blocks per feature), + # True = rowwise (16-feature blocks per token). Same dump tensor/path either way. + self.rht_rowwise = rht_rowwise + + # Always use pingpong epilogue for Hadamard + self.epilogue_pingpong = True + # Always delay TMA store acquire sync for Hadamard + self.delay_tma_store_acquire_sync = True + + self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.occupancy = 1 + self.threads_per_warp = 32 + + # Warp assignments: 8 epilogue warps (4 ACT + 4 RHT-store) + self.epilog_warp_id = (0, 1, 2, 3, 4, 5, 6, 7) + self.epilog_act_warp_id = (0, 1, 2, 3) + self.epilog_rht_store_warp_id = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.tma_warp_id = 9 + self.sched_warp_id = 10 + self.bias_load_warp_id = 11 if enable_bias else None + + self.epilogue_warp_group_size = len(self.epilog_act_warp_id) # = 4 + + all_warps = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id, self.sched_warp_id] + warps_wo_sched = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id] + if enable_bias: + all_warps.append(self.bias_load_warp_id) + warps_wo_sched.append(self.bias_load_warp_id) + self.threads_per_cta = self.threads_per_warp * len(all_warps) + self.threads_wo_sched = self.threads_per_warp * len(warps_wo_sched) + + # Named barriers + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + tmem_alloc_warp_ids = self.epilog_act_warp_id + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *tmem_alloc_warp_ids)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + # Pingpong barriers (group 0 = ACT warps, group 1 = RHT store warps) + self.epilog_sync_barrier_group0 = pipeline.NamedBarrier( + barrier_id=5, + num_threads=32 * self.epilogue_warp_group_size, + ) + self.epilog_sync_barrier_group1 = pipeline.NamedBarrier( + barrier_id=6, + num_threads=32 * self.epilogue_warp_group_size, + ) + + self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_107") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_107") + + self.vectorized_f32 = vectorized_f32 + + self.act_func = act_func + if act_func not in ["swiglu", "geglu", "srelu"]: + raise ValueError(f"Invalid activation function: {act_func}") + + self.glu_alpha = glu_alpha + self.glu_limit = glu_limit + + def _get_mma_permutation_mnk(self, mma_inst_shape_mnk): + """Return MMA permutation for the Bkeep-Breuse pattern (2CTA only). + Only active when enable_breuse=True to avoid breaking the TMA atom + setup for the non-breuse 2CTA case. + """ + if cutlass.const_expr(self.use_2cta_instrs and self.enable_breuse): + m_layout = cute.make_layout( + shape=(mma_inst_shape_mnk[0] // 2, 2, 2), + stride=(1, mma_inst_shape_mnk[0], mma_inst_shape_mnk[0] // 2), + ) + return (m_layout, mma_inst_shape_mnk[1], mma_inst_shape_mnk[2]) + return (1, 1, 1) + + def _setup_attributes(self): + """Set up configurations dependent on GEMM inputs (called inside __call__).""" + + # Hardware MMA instruction M: 2CTA -> 256, 1CTA -> 128 (decoupled from the + # tiler M so the B-reuse (512, 256) tiler maps to 2 M-split instructions). + mma_inst_m = 256 if self.use_2cta_instrs else 128 + self.mma_inst_shape_mn = ( + mma_inst_m, + self.mma_tiler[1], + ) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + # K dim of the MMA instruction shape on Rubin sm107: + # - SM107MmaMXF4NVF4Op (FP4 x FP4) requires K=128 + # - SM107BlockScaledMmaMXF8F6F4Op (FP8 or mixed) requires K=64 + mma_inst_k = ( + 128 + if (self.a_dtype.width == 4 and self.b_dtype.width == 4) + else 64 + ) + mma_inst_shape_mnk = (*self.mma_inst_shape_mn, mma_inst_k) + mma_inst_shape_mnk_sfb = (*self.mma_inst_shape_mn_sfb, mma_inst_k) + + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + mma_inst_shape_mnk, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=self._get_mma_permutation_mnk(mma_inst_shape_mnk), + ) + + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + mma_inst_shape_mnk_sfb, + ) + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 2 if (self.a_dtype.width == 4 and self.sf_vec_size == 16) else 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + d_tile_n = self.mma_inst_shape_mn[1] if self.act_func == "srelu" else self.mma_inst_shape_mn[1] // 2 + # D tiler M = the full mma tiler M (== inst M non-breuse; 512 for breuse, + # so each CTA's epilogue covers both M-halves of its rows). + self.mma_tiler_d = ( + self.mma_tiler[0], + d_tile_n, + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk_d = ( + self.mma_tiler_d[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_d[1], + self.mma_tiler_d[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = (128, 32) + self.epi_tile_cnt = ( + self.cta_tile_shape_mnk_d[0] // self.epi_tile[0], + self.cta_tile_shape_mnk_d[1] // self.epi_tile[1], + ) + self.epi_tile_c = self.epi_tile if self.act_func == "srelu" else (128, 64) + + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_d_stage, + self.num_tile_stage, + self.num_bias_stage, + self.num_pingpong_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.epi_tile_c, + self.c_dtype, + self.c_layout, + self.d_dtype, + self.d_layout, + self.sf_dtype, + self.sf_vec_size, + self.num_smem_capacity, + self.occupancy, + self.bias_dtype if self.enable_bias else None, + self.rht_dtype if self.generate_rht else None, + self.rht_quant, + self.d_quant, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile_c, + self.num_c_stage, + ) + # sD staging keeps bf16 when D itself is quantized to NVFP4 — the RHT warps + # consume bf16 D rows; the fp4 TMA source is the separate sDq buffer whose + # layout follows the gmem dtype. Without quantization the two layouts coincide. + self.d_smem_dtype = cutlass.BFloat16 if self.d_quant else self.d_dtype + self.d_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_smem_dtype, + self.d_layout, + self.epi_tile, + self.num_d_stage, + ) + self.d_tma_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_dtype, + self.d_layout, + self.epi_tile, + self.num_d_stage, + ) + + if self.enable_bias: + self.bias_smem_layout_staged = cute.make_layout( + (self.mma_tiler[1], self.num_bias_stage), + stride=(1, self.mma_tiler[1]), + ) + else: + self.bias_smem_layout_staged = cute.make_layout((1, 1)) + + # B-reuse: single acc stage; overlapping_accum shares SF TMEM with a + # second acc stage and is incompatible with the split accumulator. + if self.enable_breuse: + self.num_acc_stage = 1 + self.overlapping_accum = ( + self.num_acc_stage == 1 and self.mma_tiler[1] == 256 and not self.enable_breuse + ) + + sf_atom_mn = 32 + sf_pack_factor = 32 // self.sf_vec_size + self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k * sf_pack_factor + self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k * sf_pack_factor + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + self.num_reserved_tmem_cols = self.num_sf_tmem_cols + if self.enable_breuse: + # Split accumulator: bkeep + breuse halves (2x cols; 512 + 64 SF = 576 + # = the full sm107 TMEM). + self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage * 2 + else: + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_reserved_tmem_cols + ) + + self.epi_tile_n_required = cute.size(self.epi_tile[1]) if self.act_func == "srelu" else 2 * cute.size(self.epi_tile[1]) + self.iter_acc_early_release_in_epilogue = (self.num_reserved_tmem_cols + self.epi_tile_n_required - 1) // self.epi_tile_n_required - 1 + if self.act_func != "srelu": + self.iter_acc_early_release_in_epilogue = self.iter_acc_early_release_in_epilogue * 2 + + @cute.jit + def store_swizzled_sf_row(self, sf_tensor: cute.Tensor, logical_row, sf_col_base, sSf: cute.Tensor, tidx): + """Store one logical scale row into M32x4xrm_K4xrk_L SF layout.""" + sf_tensor = cute.recast_tensor(sf_tensor, self.sf_dtype) + sf_tensor = cute.recast_tensor(sf_tensor, self.sf_dtype) + row_m0 = logical_row % 32 + row_m1 = (logical_row // 32) % 4 + row_m2 = logical_row // 128 + num_sf = self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE + for vi in cutlass.range_constexpr(num_sf): + sf_col = sf_col_base + vi + sf_tensor[(row_m0, row_m1, row_m2, sf_col % 4, sf_col // 4, 0)] = sSf[(tidx, vi)] + + def mainloop_s2t_copy_and_partition(self, sSF, tSF): + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def get_desc_workspace_bytes(self) -> int: + """Return descriptor workspace size in bytes.""" + if self.weight_mode == MoEWeightMode.DISCRETE: + from ..moe_utils import DiscreteWeightTensormapConstructor + + return DiscreteWeightTensormapConstructor.get_workspace_size(self.expert_cnt) + return 0 + + def get_workspace_bytes(self) -> int: + """Return total workspace size in bytes.""" + desc_workspace_bytes = self.get_desc_workspace_bytes() + dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + return desc_workspace_bytes + dynamic_sched_bytes + + @cute.jit + def _get_sched_counter_ptr(self, workspace_ptr): + counter_addr = workspace_ptr.toint() + self.get_desc_workspace_bytes() + return cute.make_ptr( + cutlass.Int32, + counter_addr, + AddressSpace.gmem, + assumed_align=4, + ) + + @cute.kernel + def helper_kernel( + self, + ptrs_b: cute.Pointer, + ptrs_sfb: cute.Pointer, + n: Int32, + k: Int32, + b_stride_size: cutlass.Int64, + b_major_mode: cutlass.Constexpr, + workspace_ptr, + tiled_mma_arg: cute.TiledMma, + tiled_mma_sfb_arg: cute.TiledMma, + b_smem_layout_arg, + sfb_smem_layout_arg, + cluster_layout_vmnk_shape_arg: cutlass.Constexpr, + cluster_layout_sfb_vmnk_shape_arg: cutlass.Constexpr, + ): + """Pre-main-kernel: build per-expert TMA descriptors (discrete mode) and/or reset sched counter.""" + expert_idx = cute.arch.block_idx()[0] + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + b_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma_arg.thr_id) + sfb_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma_arg.thr_id) + + # Read per-expert base addresses from the pointer arrays + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + sfb_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + stride_n = b_stride_size + stride_k = c1_64 + else: + stride_n = c1_64 + stride_k = b_stride_size + + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + b_expert = cute.make_tensor( + b_ptr, + cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), + ) + tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B( + b_tma_op_arg, + b_expert, + b_smem_layout_arg, + self.mma_tiler, + tiled_mma_arg, + cluster_layout_vmnk_shape_arg, + ) + + workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"]) + store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) + + sfb_ptr_val = sfb_ptr_tensor[expert_idx] + sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + sfb_expert = cute.make_tensor(sfb_ptr, sfb_layout) + tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B( + sfb_tma_op_arg, + sfb_expert, + sfb_smem_layout_arg, + self.mma_tiler_sfb, + tiled_mma_sfb_arg, + cluster_layout_sfb_vmnk_shape_arg, + internal_type=cutlass.Uint64, + ) + store_tma_desc(tma_atom_sfb, workspace.get_ptr("sfb", expert_idx)) + + if cutlass.const_expr(self.use_dynamic_sched): + if expert_idx == cutlass.Int32(0): + sched_counter = cute.make_tensor( + self._get_sched_counter_ptr(workspace_ptr), + cute.make_layout(1), + ) + sched_counter[0] = cutlass.Int32(0) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b, # Dense: cute.Tensor (N,K,L) | Discrete: cute.Pointer to int64[] + sfa: cute.Tensor, + sfb, # Dense: cute.Tensor | Discrete: cute.Pointer to int64[] + n: Int32, # Ignored for dense mode + k: Int32, # Ignored for dense mode + b_stride_size: cutlass.Int64, # Ignored for dense mode + b_major_mode: cutlass.Constexpr, # Ignored for dense mode + workspace_ptr, + c: cute.Tensor, + d: cute.Tensor, # post-GLU output (bf16, or NVFP4 packed e2m1 with sfd) + sfd: Optional[cute.Tensor], # NVFP4 D block scales (e4m3/ue5m3, (m, f/16, l)); required iff d is NVFP4 + rht: Optional[cute.Tensor], # RHT output (bf16 or NVFP4, D layout); None => off + sfrht: Optional[cute.Tensor], # NVFP4 RHT block scales (e4m3/ue5m3, (m, f/16, l)); required iff rht is NVFP4 + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + prob: cute.Tensor, + bias: Optional[cute.Tensor], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + linear_offset: cutlass.Float32 = 0.0, + norm_const: cutlass.Float32 = 1.0, # D NVFP4 global encode scale: 2688/global_amax, or 1.0 + rht_norm_const: cutlass.Float32 = 1.0, # RHT NVFP4 global encode scale: 2688/global_amax, or 1.0 + ): + """Execute the MoE GEMM + GLU + Hadamard kernel. + + Dense mode: ``b`` and ``sfb`` are 3-D cute.Tensor (N, K, L). + Discrete mode: ``b`` and ``sfb`` are cute.Pointer to device int64[] + arrays of per-expert base addresses. + """ + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = a.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.d_dtype: Type[cutlass.Numeric] = d.element_type + if cutlass.const_expr(self.sf_dtype_override is not None): + self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override + else: + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + self.d_layout = utils.LayoutEnum.from_tensor(d) + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + else: + self.b_major_mode = b_major_mode + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # ---- Output / dump modes (derived from dtypes; before _setup_attributes so + # the smem/stage accounting sees them) ---- + self.generate_rht = rht is not None + self.generate_sfd = sfd is not None + self.generate_sfrht = sfrht is not None + self.d_quant = self.d_dtype == cutlass.Float4E2M1FN + self.rht_dtype = rht.element_type if cutlass.const_expr(self.generate_rht) else self.d_dtype + self.rht_quant = self.generate_rht and self.rht_dtype == cutlass.Float4E2M1FN + # The rht STORE has no orientation logic anywhere in this kernel; only the + # SCALE tensor is orientation-aware below (its block grid follows the + # transform: (m, f/16) rowwise vs (f, m/16) colwise — a different logical + # shape, not expressible as strides). + # RHT warps run for the RHT output and/or the D quantization (both read sD). + # When neither is on (plain bf16 mode) the warpgroup is COMPILED OUT: no + # pingpong, no ACT<->RHT barriers, and the tile-info pipeline's consumer + # count shrinks by the 4 RHT warps (forget that and the scheduler deadlocks + # after num_tile_stage tiles). + self.run_rht = self.generate_rht or self.d_quant + if cutlass.const_expr(not self.run_rht): + self.threads_wo_sched = ( + self.threads_per_cta + - self.threads_per_warp + - self.threads_per_warp * len(self.epilog_rht_store_warp_id) + ) + if cutlass.const_expr(self.d_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): + raise ValueError(f"d dtype must be BFloat16 or Float4E2M1FN, got {self.d_dtype}") + if cutlass.const_expr(self.d_quant != self.generate_sfd): + raise ValueError("NVFP4 d and sfd must be passed together") + if cutlass.const_expr(self.generate_rht + and self.rht_dtype not in (cutlass.BFloat16, cutlass.Float4E2M1FN)): + raise ValueError(f"rht dtype must be BFloat16 or Float4E2M1FN, got {self.rht_dtype}") + if cutlass.const_expr(self.rht_quant != self.generate_sfrht): + raise ValueError("NVFP4 rht and sfrht must be passed together") + sf_storage_dtype = cutlass.Float8E4M3FN if self.sf_dtype == cutlass.FloatNV8E5M3FNU else self.sf_dtype + if cutlass.const_expr(self.generate_sfd and sfd.element_type != sf_storage_dtype): + raise ValueError("sfd element type must match scale-factor storage dtype") + if cutlass.const_expr(self.generate_sfrht and sfrht.element_type != sf_storage_dtype): + raise ValueError("sfrht element type must match scale-factor storage dtype") + if cutlass.const_expr((self.d_quant or self.rht_quant) and self.act_func == "srelu"): + raise ValueError("NVFP4 quantization assumes the GLU subtile pair-step (act_func != srelu)") + + self._setup_attributes() + + # ---- B / SFB setup (mode-dependent) ---- + b_from_call_arg = b + sfb_from_call_arg = sfb + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b.shape, self.sf_vec_size) + sfb = cute.make_tensor(sfb.iterator, sfb_layout) + else: + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + b_template_stride = (b_stride_size, c1_64, c0) + else: + b_template_stride = (c1_64, b_stride_size, c0) + b_template_layout = cute.make_layout((n, k, c1), stride=b_template_stride) + b_ptr_typed = cute.make_ptr(self.b_dtype, b.toint(), AddressSpace.gmem, assumed_align=16) + b = cute.make_tensor(b_ptr_typed, b_template_layout) + + sfb_ptr_typed = cute.make_ptr(self.sf_dtype, sfb.toint(), AddressSpace.gmem, assumed_align=16) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + sfb = cute.make_tensor(sfb_ptr_typed, sfb_layout) + + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, self.sf_vec_size) + sfa = cute.make_tensor(sfa.iterator, sfa_layout) + + # Dump staging dtype follows the dump tensor's element type (fp4: 2KB/stage vs + # 8KB bf16); the layout is CONSTRUCTED f-major like every other epilogue + # output (never derived from the rht gmem tensor) — all FWHT store paths + # pack along features. + self.rht_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.rht_dtype, + utils.LayoutEnum.ROW_MAJOR, + self.epi_tile, + self.num_d_stage, + ) + + mma_inst_k = ( + 128 + if (self.a_dtype.width == 4 and self.b_dtype.width == 4) + else 64 + ) + mma_inst_shape_mnk = (*self.mma_inst_shape_mn, mma_inst_k) + permutation_mnk = self._get_mma_permutation_mnk(mma_inst_shape_mnk) + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + mma_inst_shape_mnk, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=permutation_mnk, + ) + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + (*self.mma_inst_shape_mn_sfb, mma_inst_k), + ) + + # B-reuse MMA pair: bkeep FILLs the collector with B, breuse LASTUSEs it. + tiled_mma_bkeep = None + tiled_mma_breuse = None + if cutlass.const_expr(self.enable_breuse): + tiled_mma_bkeep = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.b_dtype, + self.a_major_mode, self.b_major_mode, + self.sf_dtype, self.sf_vec_size, + self.cta_group, mma_inst_shape_mnk, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.FILL, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=permutation_mnk, + ) + tiled_mma_bkeep.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_bkeep.set(tcgen05.Field.NEGATE_B, False) + tiled_mma_breuse = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.b_dtype, + self.a_major_mode, self.b_major_mode, + self.sf_dtype, self.sf_vec_size, + self.cta_group, mma_inst_shape_mnk, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.LASTUSE, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=permutation_mnk, + ) + tiled_mma_breuse.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_breuse.set(tcgen05.Field.NEGATE_B, False) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA load A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # TMA load SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb = cute.make_tensor( + tma_tensor_sfb.iterator, + cute.make_layout(new_shape, stride=new_stride), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size + + # TMA store C + c_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + c_smem_layout, + self.epi_tile_c, + ) + + # TMA store D (gmem-dtype staging: sDq when D is quantized, sD otherwise) + d_smem_layout = cute.slice_(self.d_tma_smem_layout_staged, (None, None, 0)) + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + d, + d_smem_layout, + self.epi_tile, + ) + + # TMA store RHT — identical tile to D; smem layout follows the RHT output + # element type (== d_smem_layout for bf16, packed fp4 layout in quant mode). + if cutlass.const_expr(self.generate_rht): + rht_smem_layout = cute.slice_(self.rht_smem_layout_staged, (None, None, 0)) + tma_atom_rht, tma_tensor_rht = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + rht, + rht_smem_layout, + self.epi_tile, + ) + else: + tma_atom_rht, tma_tensor_rht = None, None + + # ---- Helper kernel (discrete TMA desc init + dynamic sched counter reset) ---- + _need_helper = cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE or self.use_dynamic_sched) + if cutlass.const_expr(_need_helper): + _helper_grid_x = self.expert_cnt if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else 1 + _helper_args = ( + b_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + sfb_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + n if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + k if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + b_stride_size if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int64(0), + b_major_mode if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else self.b_major_mode, + workspace_ptr, + tiled_mma, + tiled_mma_sfb, + b_smem_layout, + sfb_smem_layout, + self.cluster_layout_vmnk.shape, + self.cluster_layout_sfb_vmnk.shape, + ) + self.helper_kernel(*_helper_args).launch( + grid=(_helper_grid_x, 1, 1), + block=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # ---- Grid computation via MoE scheduler ---- + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + b_n, b_k, b_l = cute.shape(b) + sched_expert_shape = (self.expert_cnt, b_n, b_k) + else: + sched_expert_shape = (self.expert_cnt, n, k) + + sched_params = MoESchedulerParams( + scenario="2Dx3D", + expert_shape=sched_expert_shape, + cta_tile_shape_mnk=self.cta_tile_shape_mnk, + cluster_shape_mn=self.cluster_shape_mn, + use_dynamic_sched=self.use_dynamic_sched, + ) + self.sched_params, grid = compute_grid( + sched_params, + max_active_clusters, + self.use_2cta_instrs, + ) + + self.buffer_align_bytes = 1024 + + # ---- Shared storage ---- + SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_tile_stage, self.use_dynamic_sched) + + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + pingpong_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_pingpong_stage * 2] + if cutlass.const_expr(self.enable_bias): + bias_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_bias_stage * 2] + scheduler: SchedulerStorage + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sD: cute.struct.Align[ + cute.struct.MemRange[self.d_smem_dtype, cute.cosize(self.d_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.d_quant): + # NVFP4 D: packed-fp4 TMA staging + per-thread scale-byte staging. + sDq: cute.struct.Align[ + cute.struct.MemRange[self.d_dtype, cute.cosize(self.d_tma_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSfd: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, + self.threads_per_warp * len(self.epilog_rht_store_warp_id) + * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), + ], + 16, + ] + if cutlass.const_expr(self.generate_rht): + sRht: cute.struct.Align[ + cute.struct.MemRange[self.rht_dtype, cute.cosize(self.rht_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.rht_quant): + # NVFP4 RHT scale-byte staging, one contiguous row store per thread + # per tile: rowwise rows are (thread=token, cta_tile_n/16 scales); + # colwise rows are (feature-in-tile, 128-token-tile/16). + # Same (128, 8) buffer either way. + sSfRht: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, + self.threads_per_warp * len(self.epilog_rht_store_warp_id) + * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), + ], + 16, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.enable_bias): + sBias: cute.struct.Align[ + cute.struct.MemRange[self.bias_dtype, cute.cosize(self.bias_smem_layout_staged)], + 16, + ] + + self.shared_storage = SharedStorage + + # Launch main kernel + self.kernel( + tiled_mma, + tiled_mma_sfb, + tiled_mma_bkeep, + tiled_mma_breuse, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + tma_atom_d, + tma_tensor_d, + sfd, + tma_atom_rht, + tma_tensor_rht, + sfrht, + padded_offsets, + alpha, + bias, + prob, + workspace_ptr, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.c_smem_layout_staged, + self.d_smem_layout_staged, + self.d_tma_smem_layout_staged, + self.rht_smem_layout_staged, + self.bias_smem_layout_staged, + self.epi_tile, + self.sched_params, + epilogue_op, + linear_offset, + norm_const, + rht_norm_const, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + max_number_threads=[self.threads_per_cta, 1, 1], + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @cute.jit + def _make_extension(self, workspace_ptr): + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + desc_workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"]) + return DiscreteWeightScaledGemmSchedExtension( + tensormap_ctor=desc_workspace, + sf_vec_size=self.sf_vec_size, + ) + else: + return ContiguousAndConsistentGroupedGemmSchedExtension( + sf_vec_size=self.sf_vec_size, + ) + + def mainloop_s2t_copy_and_partition(self, sSF, tSF): + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # Rubin SF s2t partitions through an appended MN broadcast-4 mode (the + # packed SF TMEM lanes each read the same smem bytes). + def _append_mn_broadcast_mode(smem_layout: cute.Layout): + mn_dim = cute.get(smem_layout, mode=[0, 0]) + mn_dim = cute.append(mn_dim, cute.make_layout((4), stride=(0))) + layout = cute.append( + cute.group_modes(mn_dim, 0), cute.get(smem_layout, mode=[0, 1]) + ) + layout = cute.append( + cute.group_modes(layout, 0), cute.get(smem_layout, mode=[1]) + ) + layout = cute.append(layout, cute.get(smem_layout, mode=[2])) + layout = cute.append(layout, cute.get(smem_layout, mode=[3])) + return layout + + tCsSF_compact_bcast = cute.make_tensor( + tCsSF_compact.iterator, _append_mn_broadcast_mode(tCsSF_compact.layout) + ) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact_bcast) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + @cute.jit + def store_c( + self, + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tTR_rAcc_up, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + prev_subtile_idx, + real_subtile_idx, + ): + c_buffer = prev_subtile_idx % self.num_c_stage + tRS_rC.store(tTR_rAcc.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)]) + tRS_rC.store(tTR_rAcc_up.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 1, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group0.arrive_and_wait() + if warp_idx == self.epilog_act_warp_id[0]: + cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)]) + c_pipeline.producer_commit() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + c_pipeline.producer_acquire() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + self.epilog_sync_barrier_group0.arrive_and_wait() + + @cute.jit + def store_c_unary( + self, + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + prev_subtile_idx, + real_subtile_idx, + ): + c_buffer = prev_subtile_idx % self.num_c_stage + tRS_rC.store(tTR_rAcc.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group0.arrive_and_wait() + if warp_idx == self.epilog_act_warp_id[0]: + cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)]) + c_pipeline.producer_commit() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + c_pipeline.producer_acquire() + if not cutlass.const_expr(self.delay_tma_store_acquire_sync): + self.epilog_sync_barrier_group0.arrive_and_wait() + + @cute.jit + def geglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset=1.0): + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + scaled_gate_0, scaled_gate_1 = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (1.702, 1.702), + rnd="rn", + ftz=False, + ) + tCompute_log2e = cute.arch.mul_packed_f32x2( + (scaled_gate_0, scaled_gate_1), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.add_packed_f32x2( + (cute.math.exp2(tCompute_log2e[0], fastmath=True), cute.math.exp2(tCompute_log2e[1], fastmath=True)), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate[i], acc_vec_gate[i + 1]), + rnd="rn", + ftz=False, + ) + up0, up1 = cute.arch.add_packed_f32x2( + (linear_offset, linear_offset), + (acc_vec_up[i], acc_vec_up[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (up0, up1), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (self.glu_alpha, self.glu_alpha), + rnd="rn", + ftz=False, + ) + else: + # GeGlu Unpacked Version + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = (acc_vec_up[i] + linear_offset) * silu_f32_geglu_scaled(acc_vec_gate[i], fastmath=True) + tCompute[i] = tCompute[i] * mProb + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + tCompute[i] = tCompute[i] * self.glu_alpha + # + 0.0 canonicalizes -0 -> +0 (negative up x saturated-to-zero sigmoid). SCALAR + # on purpose: the f32 immediate makes it a free FADD with RZ, while f32x2 has no + # immediate form and a live (0, 0) register pair costs regs/spills. + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = tCompute[i] + cutlass.Float32(0.0) + + @cute.jit + def swiglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb): + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + tCompute_log2e = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.add_packed_f32x2( + (cute.math.exp2(tCompute_log2e[0], fastmath=True), cute.math.exp2(tCompute_log2e[1], fastmath=True)), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate[i], acc_vec_gate[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_up[i], acc_vec_up[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (self.glu_alpha, self.glu_alpha), + rnd="rn", + ftz=False, + ) + else: + # SwiGlu Unpacked Version + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = acc_vec_up[i] * silu_f32(acc_vec_gate[i], fastmath=True) + tCompute[i] = tCompute[i] * mProb + if cutlass.const_expr(self.glu_alpha is not None and self.glu_alpha != 1.0): + tCompute[i] = tCompute[i] * self.glu_alpha + # + 0.0 canonicalizes -0 -> +0 (negative up x saturated-to-zero sigmoid). SCALAR + # on purpose: the f32 immediate makes it a free FADD with RZ, while f32x2 has no + # immediate form and a live (0, 0) register pair costs regs/spills. + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = tCompute[i] + cutlass.Float32(0.0) + + @cute.jit + def srelu_act(self, tCompute, acc_vec, mProb): + acc_relu = cute.where(acc_vec > 0, acc_vec, cute.full_like(acc_vec, 0)) + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (acc_relu[i], acc_relu[i + 1]), + (acc_relu[i], acc_relu[i + 1]), + rnd="rn", + ftz=False, + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = acc_relu[i] * acc_relu[i] * mProb + + def epilog_tmem_copy_and_partition(self, tidx, tAcc, gD_mnl, epi_tile, use_2cta_instrs, m_half=0): + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, + self.d_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # For breuse, select the requested m_half (0=bkeep, 1=breuse) from the + # split accumulator. m_half is a Python int (resolves at trace time). + if cutlass.const_expr(self.enable_breuse): + tAcc_epi = cute.flat_divide(tAcc[((None, None), m_half, 0, None)], epi_tile) + gD_mnl_epi = cute.flat_divide(gD_mnl[((None, None), m_half, 0, None, None, None)], epi_tile) + else: + tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) + gD_mnl_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + tTR_gC = thr_copy_t2r.partition_D(gD_mnl_epi) + tTR_rAcc_gate = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + tTR_rAcc_up = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc_gate, tTR_rAcc_up + + def epilog_smem_copy_and_partition(self, tiled_copy_t2r, tTR_rC, tidx, sD): + copy_atom_r2s = sm100_utils.get_smem_store_op(self.d_layout, sD.element_type, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = thr_copy_r2s.partition_D(sD) + tRS_rD = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rD, tRS_sD + + def epilog_gmem_copy_and_partition(self, tidx, atom, gD_mnl, epi_tile, sD): + gD_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tma_atom_d = atom + sD_for_tma_partition = cute.group_modes(sD, 0, 2) + gD_for_tma_partition = cute.group_modes(gD_epi, 0, 2) + bSG_sD, bSG_gD = cpasync.tma_partition( + tma_atom_d, + 0, + cute.make_layout(1), + sD_for_tma_partition, + gD_for_tma_partition, + ) + return tma_atom_d, bSG_sD, bSG_gD + + @staticmethod + def _compute_stages( + tiled_mma, + mma_tiler_mnk, + a_dtype, + b_dtype, + epi_tile, + epi_tile_c, + c_dtype, + c_layout, + d_dtype, + d_layout, + sf_dtype, + sf_vec_size, + num_smem_capacity, + occupancy, + bias_dtype, + rht_dtype, # RHT output dtype (None => no RHT output) + rht_quant, # RHT output is NVFP4 (adds sfrht scale-byte staging) + d_quant, # D output is NVFP4 (sD staging stays bf16; adds fp4 + sfd staging) + ): + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + num_c_stage = 1 + # Double-buffer the D staging when the RHT warps consume it (run_rht): + # the pingpong pipeline then lets the ACT warps run up to num_d_stage + # subtiles ahead, taking the whole FWHT/quant/store chain off their + # critical path. Plain D (TMA'd by ACT) keeps the single stage. + num_d_stage = 2 if (rht_dtype is not None or d_quant) else 1 + num_tile_stage = 2 + # The pingpong pipeline protects the sD stages, so its depth must match + # num_d_stage (stage index and d_buffer advance in lockstep). + num_pingpong_stage = num_d_stage + + a_smem_layout_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1) + b_smem_layout_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1) + sfa_smem_layout_one = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, 1) + sfb_smem_layout_one = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, 1) + c_smem_layout_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile_c, 1) + d_smem_layout_one = sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_one) + ) + mbar_helpers_bytes = 1024 + + # One fp8 scale byte per (1,16) feature block, one row per thread (128 threads). + quant_sf_bytes = 128 * ((mma_tiler_mnk[1] // 2) // HADAMARD_SIZE) + + # sInfo is in SchedulerStorage, not here, so use 4-int sInfo + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes = cute.size_in_bytes(c_dtype, c_smem_layout_one) * num_c_stage + d_bytes = cute.size_in_bytes(d_dtype, d_smem_layout_one) * num_d_stage + + if d_quant: + # sD staging stays bf16 (the RHT/quant warps' source); the gmem-dtype (fp4) + # staging above becomes the TMA source, and sfd rows are staged per thread. + bf16_smem_layout_one = sm100_utils.make_smem_layout_epi(cutlass.BFloat16, d_layout, epi_tile, 1) + d_bytes += cute.size_in_bytes(cutlass.BFloat16, bf16_smem_layout_one) * num_d_stage + d_bytes += quant_sf_bytes + + rht_bytes = 0 + if rht_dtype is not None: + rht_smem_layout_one = sm100_utils.make_smem_layout_epi(rht_dtype, d_layout, epi_tile, 1) + rht_bytes = cute.size_in_bytes(rht_dtype, rht_smem_layout_one) * num_d_stage + if rht_quant: + rht_bytes += quant_sf_bytes + + if bias_dtype is not None: + num_bias_stage = 2 + bias_bytes = mma_tiler_mnk[1] * num_bias_stage * (bias_dtype.width // 8) + else: + num_bias_stage = 0 + bias_bytes = 0 + + epi_bytes = c_bytes + d_bytes + rht_bytes + bias_bytes + + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage + + return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage, num_pingpong_stage + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tiled_mma_bkeep: Optional[cute.TiledMma], + tiled_mma_breuse: Optional[cute.TiledMma], + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tma_atom_d: cute.CopyAtom, + mD_mnl: cute.Tensor, + mSfd_mnl: Optional[cute.Tensor], + tma_atom_rht: Optional[cute.CopyAtom], + mRht_mnl: Optional[cute.Tensor], + mSfRht_mnl: Optional[cute.Tensor], + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + mBias_nl: Optional[cute.Tensor], + prob: cute.Tensor, + workspace_ptr, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_tma_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + rht_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + bias_smem_layout_staged: cute.Layout, + epi_tile: cute.Tile, + sched_params: MoESchedulerParams, + epilogue_op: cutlass.Constexpr, + linear_offset: cutlass.Float32 = 0.0, + norm_const: cutlass.Float32 = 1.0, + rht_norm_const: cutlass.Float32 = 1.0, + ): + """GPU device kernel: MoE persistent GEMM + GLU + Hadamard (pingpong epilogue).""" + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + total_token = padded_offsets[self.expert_cnt - 1] + + # Prefetch TMA descriptors + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_sfa) + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + cpasync.prefetch_descriptor(tma_atom_d) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # CTA coordinates + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + # Shared memory allocation + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sched_storage = storage.scheduler + + # AB pipeline + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # ACC pipeline + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_act_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Pingpong pipeline + pingpong_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_act_warp_id) * self.threads_per_warp) + pingpong_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_rht_store_warp_id) * self.threads_per_warp) + pingpong_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.pingpong_mbar_ptr.data_ptr(), + num_stages=self.num_pingpong_stage, + producer_group=pingpong_producer_group, + consumer_group=pingpong_consumer_group, + ) + + # Tile info pipeline (uses SchedulerStorage's barrier) + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=sched_storage.tile_info_mbar.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + # MoE persistent tile scheduler + scheduler = MoEPersistentTileScheduler.create( + sched_params, + padded_offsets, + cute.arch.block_idx(), + cute.arch.grid_dim(), + counter_ptr=self._get_sched_counter_ptr(workspace_ptr), + sched_storage=sched_storage, + ) + scheduler.internal_init() + + # Bias pipeline + if cutlass.const_expr(self.enable_bias): + bias_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp, + ) + bias_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_act_warp_id), + ) + bias_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.bias_mbar_ptr.data_ptr(), + num_stages=self.num_bias_stage, + producer_group=bias_pipeline_producer_group, + consumer_group=bias_pipeline_consumer_group, + ) + sBias = storage.sBias.get_tensor(bias_smem_layout_staged) + gBias_nl = cute.local_tile(mBias_nl, cute.slice_(self.mma_tiler[:2], (0, None)), (None, None)) + + # TMEM allocator + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_act_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch="sm_107", + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # SMEM tensors + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + sD = storage.sD.get_tensor(d_smem_layout_staged.outer, swizzle=d_smem_layout_staged.inner) + _num_sf_per_tile = self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE + _sf_stage_layout = cute.make_layout( + (self.threads_per_warp * len(self.epilog_rht_store_warp_id), _num_sf_per_tile), + stride=(_num_sf_per_tile, 1), + ) + if cutlass.const_expr(self.d_quant): + sDq = storage.sDq.get_tensor(d_tma_smem_layout_staged.outer, swizzle=d_tma_smem_layout_staged.inner) + sSfd = storage.sSfd.get_tensor(_sf_stage_layout) + if cutlass.const_expr(self.generate_rht): + sRht = storage.sRht.get_tensor(rht_smem_layout_staged.outer, swizzle=rht_smem_layout_staged.inner) + if cutlass.const_expr(self.rht_quant): + sSfRht = storage.sSfRht.get_tensor(_sf_stage_layout) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # sInfo from SchedulerStorage + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = sched_storage.sInfo.get_tensor(info_layout) + + # Multicast masks + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1) + + # MMA fragments + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage_overlapped)) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (256 - self.num_reserved_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # Cluster wait / CTA sync + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + self.cta_sync_barrier.arrive_and_wait() + + if total_token <= 0: + cute.arch.nvvm.exit() + + # --------------------------------------------------------------- + # Specialized Scheduler warp (MoEPersistentTileScheduler) + # --------------------------------------------------------------- + if warp_idx == self.sched_warp_id: + work_tile_info = scheduler.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tile_stage) + + while work_tile_info.is_valid_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = work_tile_info.expert_idx + sInfo[(1, tile_info_producer_state.index)] = work_tile_info.tile_m_idx + sInfo[(2, tile_info_producer_state.index)] = work_tile_info.tile_n_idx + sInfo[(3, tile_info_producer_state.index)] = work_tile_info.k_tile_cnt + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + work_tile_info = scheduler.advance_to_next_work() + + # Send invalid signal: expert_idx = -1 + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cutlass.Int32(-1) + sInfo[(1, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(2, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # --------------------------------------------------------------- + # Specialized TMA load warp + # --------------------------------------------------------------- + if warp_idx == self.tma_warp_id: + ext = self._make_extension(workspace_ptr) + + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + k_tile_cnt = work_tile_info.k_tile_cnt + ext.update_expert_info(padded_offsets, work_tile_info.expert_idx) + + real_a, _ = ext.get_gmem_tensor("a", mA_mkl, padded_offsets, work_tile_info) + real_b, desc_ptr_b = ext.get_gmem_tensor("b", mB_nkl, padded_offsets, work_tile_info) + real_sfa, _ = ext.get_gmem_tensor("sfa", mSFA_mkl, padded_offsets, work_tile_info) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor("sfb", mSFB_nkl, padded_offsets, work_tile_info) + + gA_mkl = cute.local_tile(real_a, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gB_nkl = cute.local_tile(real_b, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + gSFA_mkl = cute.local_tile(real_sfa, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gSFB_nkl = cute.local_tile(real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None)) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + sfa_cta_layout = a_cta_layout + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + mma_tile_coord_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape) + mma_tile_coord_n = work_tile_info.tile_n_idx + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, mma_tile_coord_n, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)] + slice_n = mma_tile_coord_n + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_n // 2 + tBgSFB_slice = tBgSFB[(None, slice_n, None, 0)] + + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + tAgA_k = tAgA_slice[(None, ab_producer_state.count)] + tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + tAgSFA_k = tAgSFA_slice[(None, ab_producer_state.count)] + tBgSFB_k = tBgSFB_slice[(None, ab_producer_state.count)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + tAsSFA_pipe = tAsSFA[(None, ab_producer_state.index)] + tBsSFB_pipe = tBsSFB[(None, ab_producer_state.index)] + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + ab_producer_state_next = ab_producer_state.clone() + ab_producer_state_next.advance() + if ab_producer_state_next.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state_next) + else: + peek_ab_empty_status = cutlass.Boolean(1) + + cute.copy(tma_atom_a, tAgA_k, tAsA_pipe, tma_bar_ptr=tma_bar, mcast_mask=a_full_mcast_mask) + cute.copy(tma_atom_b, tBgB_k, tBsB_pipe, tma_bar_ptr=tma_bar, mcast_mask=b_full_mcast_mask, tma_desc_ptr=desc_ptr_b) + cute.copy(tma_atom_sfa, tAgSFA_k, tAsSFA_pipe, tma_bar_ptr=tma_bar, mcast_mask=sfa_full_mcast_mask) + cute.copy(tma_atom_sfb, tBgSFB_k, tBsSFB_pipe, tma_bar_ptr=tma_bar, mcast_mask=sfb_full_mcast_mask, tma_desc_ptr=desc_ptr_sfb) + + ab_producer_state.advance() + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + ab_pipeline.producer_tail(ab_producer_state) + + # --------------------------------------------------------------- + # Specialized MMA warp + # --------------------------------------------------------------- + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acd_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + k_tile_cnt = tile_info[3] + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + acd_producer_state.reset_count() + peek_acc_empty_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state) + + mma_tile_coord_mnl = ( + tile_info[1] // cute.size(tiled_mma.thr_id.shape), + tile_info[2], + tile_info[0], + ) + + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acd_producer_state.phase ^ 1 + else: + acc_stage_index = acd_producer_state.index + + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + + if is_leader_cta: + acc_pipeline.producer_acquire(acd_producer_state, peek_acc_empty_status) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + s2t_stage_coord = (None, None, None, None, ab_consumer_state.index) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_staged, tCtSFA_compact_s2t) + cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_staged, tCtSFB_compact_s2t) + + num_kblocks = cute.size(tCrA, mode=[2]) + ab_consumer_state_next = ab_consumer_state.clone() + ab_consumer_state_next.advance() + if ab_consumer_state_next.count < k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state_next) + + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + if cutlass.const_expr( + self.enable_breuse + and cute.size(tCtAcc.layout, mode=[1]) == 2 + and cute.size(tCtAcc.layout, mode=[2]) == 1 + ): + tCtAcc_bkeep = tCtAcc[(None, 0, 0)] + tCtAcc_breuse = tCtAcc[(None, 1, 0)] + a_kblk_crd_keep = (None, 0, kblock_idx, ab_consumer_state.index) + a_kblk_crd_reuse = (None, 1, kblock_idx, ab_consumer_state.index) + b_kblk_crd = (None, 0, kblock_idx, ab_consumer_state.index) + sfa_kblk_crd_keep = (None, 0, kblock_idx) + sfa_kblk_crd_reuse = (None, 1, kblock_idx) + sfb_kblk_crd = (None, 0, kblock_idx) + # Bkeep: accumulate A_upper x B (keeps B in reuse buffer) + tiled_mma_bkeep.set(tcgen05.Field.ACCUMULATE, k_tile != 0 or kblock_idx != 0) + cute.gemm(tiled_mma_bkeep, tCtAcc_bkeep, + [tCrA[a_kblk_crd_keep], tCtSFA[sfa_kblk_crd_keep]], + [tCrB[b_kblk_crd], tCtSFB_mma[sfb_kblk_crd]], tCtAcc_bkeep) + # Breuse: accumulate A_lower x B (reuses B) + tiled_mma_breuse.set(tcgen05.Field.ACCUMULATE, k_tile != 0 or kblock_idx != 0) + cute.gemm(tiled_mma_breuse, tCtAcc_breuse, + [tCrA[a_kblk_crd_reuse], tCtSFA[sfa_kblk_crd_reuse]], + [tCrB[b_kblk_crd], tCtSFB_mma[sfb_kblk_crd]], tCtAcc_breuse) + else: + kblock_coord = (None, None, kblock_idx, ab_consumer_state.index) + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator) + cute.gemm(tiled_mma, tCtAcc, tCrA[kblock_coord], tCrB[kblock_coord], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state = ab_consumer_state_next + + if is_leader_cta: + acc_pipeline.producer_commit(acd_producer_state) + + acd_producer_state.advance() + if acd_producer_state.count < k_tile_cnt: + if is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + acc_pipeline.producer_tail(acd_producer_state) + + # --------------------------------------------------------------- + # Specialized bias load warp + # --------------------------------------------------------------- + if cutlass.const_expr(self.enable_bias): + if warp_idx == self.bias_load_warp_id and total_token > 0: + bias_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_bias_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + bias_elems_per_thread = 128 // self.bias_dtype.width + bias_g2s_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + self.bias_dtype, + num_bits_per_copy=128, + ) + bias_g2s_tiled = cute.make_tiled_copy_tv( + bias_g2s_atom, + cute.make_layout((self.threads_per_warp,)), + cute.make_layout((bias_elems_per_thread,)), + ) + thr_bias_g2s = bias_g2s_tiled.get_slice(cute.arch.lane_idx()) + tBs_sBias = thr_bias_g2s.partition_D(sBias) + + bias_n_total = mBias_nl.shape[0] + tBpBias = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Boolean) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + bias_producer_state.reset_count() + mma_n_coord = tile_info[2] + expert_idx = tile_info[0] + gBias_tile = gBias_nl[(None, mma_n_coord, expert_idx)] + tBs_gBias = thr_bias_g2s.partition_S(gBias_tile) + tBpBias[0] = mma_n_coord * self.mma_tiler[1] + cute.arch.lane_idx() * bias_elems_per_thread < bias_n_total + bias_pipeline.producer_acquire(bias_producer_state) + cute.copy( + bias_g2s_tiled, + tBs_gBias[(None, 0)], + tBs_sBias[(None, 0, bias_producer_state.index)], + pred=tBpBias, + ) + bias_pipeline.producer_commit(bias_producer_state) + bias_producer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + bias_pipeline.producer_tail(bias_producer_state) + + # --------------------------------------------------------------- + # Specialized ACT epilogue warps (0-3): TMEM→regs, alpha, GLU activation, + # C store, hadamard_in + # --------------------------------------------------------------- + if warp_idx < self.epilog_rht_store_warp_id[0] and total_token > 0: + epi_tidx = tidx + + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue (shape-only via mD_mnl for invariant setup) + # + thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v) + gD_mnl_shape = cute.local_tile(mD_mnl, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD_shape = thr_mma_epi.partition_C(gD_mnl_shape) + + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc_gate, + tTR_rAcc_up, + ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgD_shape, epi_tile, use_2cta_instrs) + # For breuse: second partition for the breuse M-half (defaults to bk + # for JIT scoping; overridden when enable_breuse=True). + tTR_tAcc_base_br = tTR_tAcc_base + if cutlass.const_expr(self.enable_breuse): + ( + _, + tTR_tAcc_base_br, + _tTR_rAcc_gate_br, + _tTR_rAcc_up_br, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgD_shape, epi_tile, use_2cta_instrs, m_half=1 + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rC, epi_tidx, sC) + + tTR_rD = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.d_smem_dtype) + tiled_copy_r2s_d, tRS_rD, tRS_sD = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rD, epi_tidx, sD) + + # + # Create per-expert extension (for C/prob tensors inside tile loop) + # + epi_ext = self._make_extension(workspace_ptr) + + # + # Persistent tile scheduling state + # + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + + # + # Pingpong producer state + # + pingpong_act_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_pingpong_stage) + + # Threads/warps participating in TMA store pipeline for C + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_act_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + # NVFP4 D: the fp4 TMA store is issued by the RHT warps (which quantize + # sD); the ACT warps only stage bf16 into sD. + if cutlass.const_expr(not self.d_quant): + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_act_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=d_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_bias_stage) + bias_s2r_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.bias_dtype, num_bits_per_copy=128) + tTR_rBias_gate = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + tTR_rBias_up = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + num_prev_subtiles = cutlass.Int32(0) + num_prev_d_subtiles = cutlass.Int32(0) + while is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + mma_tile_coord_mnl = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx, + cutlass.Int32(0), + ) + + expert_idx = epi_work_tile_info.expert_idx + alpha_val = alpha[expert_idx] + epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx) + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state.reset_count() + bias_pipeline.consumer_wait(bias_consumer_state) + sBias_stage = sBias[(None, bias_consumer_state.index)] + if cutlass.const_expr(self.act_func == "srelu"): + sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(self.epi_tile[1])) + else: + sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(2 * self.epi_tile[1])) + + # + # Get per-expert C tensor inside tile loop + # + real_c, _ = epi_ext.get_gmem_tensor("c", mC_mnl, padded_offsets, epi_work_tile_info) + gC_mnl = cute.local_tile(real_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + thr_mma_epi_loop = tiled_mma.get_slice(mma_tile_coord_v) + tCgC = thr_mma_epi_loop.partition_C(gC_mnl) + if cutlass.const_expr(self.enable_breuse): + # Select bkeep/breuse halves by indexing mode1 with 0 or 1 + gC_epi_bk = cute.flat_divide(tCgC[((None, None), 0, 0, None, None, None)], self.epi_tile_c) + gC_epi_br = cute.flat_divide(tCgC[((None, None), 1, 0, None, None, None)], self.epi_tile_c) + sC_for_tma = cute.group_modes(sC, 0, 2) + bSG_sC, bSG_gC_partitioned_bk = cpasync.tma_partition( + tma_atom_c, 0, cute.make_layout(1), sC_for_tma, cute.group_modes(gC_epi_bk, 0, 2)) + _, bSG_gC_partitioned_br = cpasync.tma_partition( + tma_atom_c, 0, cute.make_layout(1), sC_for_tma, cute.group_modes(gC_epi_br, 0, 2)) + else: + _, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, self.epi_tile_c, sC) + + # + # Get per-expert D tensor inside tile loop (NVFP4 D: the RHT warps own + # the fp4 D TMA store — no D partition on the ACT side). + # + if cutlass.const_expr(not self.d_quant): + real_d, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + gD_mnl_loop = cute.local_tile(real_d, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD = thr_mma_epi_loop.partition_C(gD_mnl_loop) + if cutlass.const_expr(self.enable_breuse): + gD_epi_bk = cute.flat_divide(tCgD[((None, None), 0, 0, None, None, None)], epi_tile) + gD_epi_br = cute.flat_divide(tCgD[((None, None), 1, 0, None, None, None)], epi_tile) + sD_for_tma = cute.group_modes(sD, 0, 2) + bSG_sD, bSG_gD_partitioned_bk = cpasync.tma_partition( + tma_atom_d, 0, cute.make_layout(1), sD_for_tma, cute.group_modes(gD_epi_bk, 0, 2)) + _, bSG_gD_partitioned_br = cpasync.tma_partition( + tma_atom_d, 0, cute.make_layout(1), sD_for_tma, cute.group_modes(gD_epi_br, 0, 2)) + else: + _, bSG_sD, bSG_gD_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgD, epi_tile, sD) + + # Slice to per-expert tile coords (L=0, domain already offset'd) + def _grp(p): + q = p[(None, None, None, *mma_tile_coord_mnl)] + return cute.group_modes(q, 1, cute.rank(q)) + + if cutlass.const_expr(self.enable_breuse): + bSG_gC_bk = _grp(bSG_gC_partitioned_bk) + bSG_gC_br = _grp(bSG_gC_partitioned_br) + # Placeholders so the JIT sees all names in non-breuse paths too + bSG_gC = bSG_gC_bk + else: + bSG_gC = _grp(bSG_gC_partitioned) + bSG_gC_bk = bSG_gC + bSG_gC_br = bSG_gC + if cutlass.const_expr(not self.d_quant): + if cutlass.const_expr(self.enable_breuse): + bSG_gD_bk = _grp(bSG_gD_partitioned_bk) + bSG_gD_br = _grp(bSG_gD_partitioned_br) + bSG_gD = bSG_gD_bk + else: + bSG_gD = _grp(bSG_gD_partitioned) + bSG_gD_bk = bSG_gD + bSG_gD_br = bSG_gD + + # + # Get per-expert prob tensor inside tile loop + # + real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info) + mPosition = ( + (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0] + + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape)) + + tidx + ) + mProb = real_prob[mPosition, 0, 0] + # For breuse: also fetch the breuse-half prob (second 128-row half + # of this CTA's rows). + mProb_bk = mProb + mProb_br = mProb + if cutlass.const_expr(self.enable_breuse): + mProb_bk = mProb + mProb_br = real_prob[mPosition + (self.cta_tile_shape_mnk[0] // 2), 0, 0] + + # + # Get accumulator stage index + # + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False) + else: + acc_stage_index = acc_consumer_state.index + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # For breuse: process bkeep (m_half=0) then breuse (m_half=1) + # M-halves — each is a normal subtile pass. Non-breuse: one pass. + # (2-space indent below keeps the subtile-loop body untouched.) + _breuse_halves = [0, 1] if self.enable_breuse else [0] + + _tTR_tAcc_base_h = tTR_tAcc_base + _mProb_h = mProb + _bSG_gC_h = bSG_gC + if cutlass.const_expr(not self.d_quant): + _bSG_gD_h = bSG_gD + + for _m_half in _breuse_halves: + if self.enable_breuse: + _tTR_tAcc_base_h = tTR_tAcc_base if _m_half == 0 else tTR_tAcc_base_br + _mProb_h = mProb_bk if _m_half == 0 else mProb_br + _bSG_gC_h = bSG_gC_bk if _m_half == 0 else bSG_gC_br + if cutlass.const_expr(not self.d_quant): + _bSG_gD_h = bSG_gD_bk if _m_half == 0 else bSG_gD_br + else: + _tTR_tAcc_base_h = tTR_tAcc_base + _mProb_h = mProb + _bSG_gC_h = bSG_gC + if cutlass.const_expr(not self.d_quant): + _bSG_gD_h = bSG_gD + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE) + tTR_tAcc = _tTR_tAcc_base_h[(None, None, None, None, None, acc_stage_index)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + subtile_step = 1 if cutlass.const_expr(self.act_func == "srelu") else 2 + for subtile_idx in cutlass.range(0, subtile_cnt, subtile_step, unroll=1): + real_subtile_idx = subtile_idx if cutlass.const_expr(self.act_func == "srelu") else subtile_idx // 2 + if cutlass.const_expr(self.overlapping_accum): + if reverse_subtile: + real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n_required - 1 - real_subtile_idx + + # + # Load accumulator from tensor memory buffer to register + # + if cutlass.const_expr(self.act_func == "srelu"): + tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx)] + else: + tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx * 2)] + tTR_tAcc_mn_up = tTR_tAcc[(None, None, None, real_subtile_idx * 2 + 1)] + + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_gate, tTR_rAcc_gate) + if cutlass.const_expr(self.act_func != "srelu"): + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_up, tTR_rAcc_up) + + # + # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled + # + if cutlass.const_expr(self.overlapping_accum): + if subtile_idx == self.iter_acc_early_release_in_epilogue: + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Apply alpha (+ bias when enabled) + # + if cutlass.const_expr(self.enable_bias): + if cutlass.const_expr(self.act_func == "srelu"): + sBias_sub = sBias_subtiles[(None, real_subtile_idx)] + cute.copy(bias_s2r_atom, sBias_sub, tTR_rBias_gate) + else: + sBias_pair = sBias_subtiles[(None, real_subtile_idx)] + sBias_sub = cute.flat_divide(sBias_pair, cute.make_layout(self.epi_tile[1])) + cute.copy(bias_s2r_atom, sBias_sub[(None, 0)], tTR_rBias_gate) + bias_vec_gate = tTR_rBias_gate.load() + if cutlass.const_expr(self.act_func != "srelu"): + cute.copy(bias_s2r_atom, sBias_sub[(None, 1)], tTR_rBias_up) + bias_vec_up = tTR_rBias_up.load() + + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2): + bias_gate_f32_0 = bias_vec_gate[i].to(cutlass.Float32) + bias_gate_f32_1 = bias_vec_gate[i + 1].to(cutlass.Float32) + tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + (bias_gate_f32_0, bias_gate_f32_1), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.act_func != "srelu"): + bias_up_f32_0 = bias_vec_up[i].to(cutlass.Float32) + bias_up_f32_1 = bias_vec_up[i + 1].to(cutlass.Float32) + tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + (bias_up_f32_0, bias_up_f32_1), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)): + tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + bias_vec_gate[i].to(cutlass.Float32) + if cutlass.const_expr(self.act_func != "srelu"): + tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + bias_vec_up[i].to(cutlass.Float32) + + last_bias_subtile = subtile_cnt - 1 if cutlass.const_expr(self.act_func == "srelu") else subtile_cnt - 2 + # Release once per WORK TILE (the wait precedes the m-half + # loop): only the LAST half may release under B-reuse. + if subtile_idx == last_bias_subtile and _m_half == _breuse_halves[-1]: + bias_pipeline.consumer_release(bias_consumer_state) + bias_consumer_state.advance() + else: + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2): + tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + rnd="rn", + ftz=False, + ) + if cutlass.const_expr(self.act_func != "srelu"): + tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]), + ( + cutlass.Float32(alpha_val), + cutlass.Float32(alpha_val), + ), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)): + tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + if cutlass.const_expr(self.act_func != "srelu"): + tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + + # + # Store pre-activation output to C tensor for residual/backward. + # + if cutlass.const_expr(self.act_func == "srelu"): + self.store_c_unary( + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc_gate, + tRS_rC, + tRS_sC, + _bSG_gC_h, + bSG_sC, + c_pipeline, + num_prev_subtiles, + real_subtile_idx, + ) + else: + self.store_c( + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc_gate, + tTR_rAcc_up, + tRS_rC, + tRS_sC, + _bSG_gC_h, + bSG_sC, + c_pipeline, + num_prev_subtiles, + real_subtile_idx, + ) + num_prev_subtiles = num_prev_subtiles + 1 + + # + # GeGLU clamp before C store + # + if cutlass.const_expr((self.act_func == "geglu" or self.act_func == "swiglu") and self.glu_limit is not None): + geglu_max_val = cutlass.Float32(self.glu_limit) + geglu_min_val = cutlass.Float32(-self.glu_limit) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_up)): + tTR_rAcc_gate[i] = fmin(tTR_rAcc_gate[i], geglu_max_val) + tTR_rAcc_gate[i] = fmax(tTR_rAcc_gate[i], geglu_min_val) + tTR_rAcc_up[i] = fmin(tTR_rAcc_up[i], geglu_max_val) + tTR_rAcc_up[i] = fmax(tTR_rAcc_up[i], geglu_min_val) + + acc_vec_gate = tTR_rAcc_gate.load() + + # + # Compute activation. + # + tCompute = cute.make_rmem_tensor(acc_vec_gate.shape, self.acc_dtype) + if cutlass.const_expr(self.act_func == "srelu"): + self.srelu_act(tCompute, acc_vec_gate, _mProb_h) + elif cutlass.const_expr(self.act_func == "geglu"): + acc_vec_up = tTR_rAcc_up.load() + self.geglu_act(tCompute, acc_vec_up, acc_vec_gate, _mProb_h, linear_offset) + elif cutlass.const_expr(self.act_func == "swiglu"): + acc_vec_up = tTR_rAcc_up.load() + self.swiglu_act(tCompute, acc_vec_up, acc_vec_gate, _mProb_h) + + # + # Store post-activation output to D staging (bf16 under NVFP4 D — + # the RHT warps quantize it and issue the fp4 TMA store). + # + acc_vec = tiled_copy_r2s_d.retile(tCompute).load() + tRS_rD.store(acc_vec.to(self.d_smem_dtype)) + d_buffer = num_prev_d_subtiles % self.num_d_stage + # The sD stage is guarded by the pingpong pipeline: acquire + # BEFORE writing (the RHT warps consumer_release the slot right + # after their sD register loads, up to num_d_stage subtiles + # behind). Replaces the old per-subtile ACT<->RHT lockstep + # barrier — the mbarrier phase bit makes the early release safe + # (the named-barrier arrive() variant raced its release drain). + if cutlass.const_expr(self.run_rht): + pingpong_pipeline.producer_acquire(pingpong_act_producer_state) + cute.copy( + tiled_copy_r2s_d, + tRS_rD, + tRS_sD[(None, None, None, d_buffer)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group0.arrive_and_wait() + if cutlass.const_expr(not self.d_quant): + if warp_idx == self.epilog_act_warp_id[0]: + cute.copy( + tma_atom_d, + bSG_sD[(None, d_buffer)], + _bSG_gD_h[(None, real_subtile_idx)], + ) + d_pipeline.producer_commit() + + # + # Signal the RHT epilogue warps that the post-activation D tile is in SMEM. + # (The matching producer_acquire moved BEFORE the sD write above.) + # + if cutlass.const_expr(self.run_rht): + pingpong_pipeline.producer_commit(pingpong_act_producer_state) + pingpong_act_producer_state.advance() + + num_prev_d_subtiles = num_prev_d_subtiles + 1 + + # + # Delayed TMA store acquire + group sync (always enabled) + # + if cutlass.const_expr(self.delay_tma_store_acquire_sync): + if warp_idx == self.epilog_act_warp_id[0]: + if cutlass.const_expr(not self.d_quant): + d_pipeline.producer_acquire() + c_pipeline.producer_acquire() + self.epilog_sync_barrier_group0.arrive_and_wait() + + # (No per-subtile ACT<->RHT barrier: sD reuse is fully ordered + # by the pingpong producer_acquire above.) + + # + # (No tile-end ACT<->RHT barrier either: the RHT warps touch no + # ACT-owned state across the tile boundary — sD stays pingpong- + # guarded, TMEM/acc is ACT-only, and the sSf staging is RHT-only.) + # + + # + # Async arrive accumulator buffer empty + # + if cutlass.const_expr(not self.overlapping_accum): + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier_group0.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store / pingpong complete + # + c_pipeline.producer_tail() + if cutlass.const_expr(not self.d_quant): + d_pipeline.producer_tail() + if cutlass.const_expr(self.run_rht): + pingpong_pipeline.producer_tail(pingpong_act_producer_state) + + # --------------------------------------------------------------- + # Specialized RHT store warps (4-7): RHT and/or NVFP4 quantization from D SMEM + # --------------------------------------------------------------- + if self.run_rht and warp_idx < self.mma_warp_id and warp_idx >= self.epilog_rht_store_warp_id[0] and total_token > 0: + epi_tidx = tidx % 128 + + # + # Pingpong consumer state + # + pingpong_rht_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_pingpong_stage) + + # + # Create per-expert extension (for D tensor inside tile loop) + # + epi_ext = self._make_extension(workspace_ptr) + + # + # RHT output: TMA-store pipeline for the RHT warps (mirrors ACT's D store). + # + if cutlass.const_expr(self.generate_rht): + rht_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_rht_store_warp_id), + ) + rht_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=rht_producer_group, + ) + + # + # NVFP4 D: the RHT warps quantize sD and own the fp4 D TMA store. + # + if cutlass.const_expr(self.d_quant): + dq_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_rht_store_warp_id), + ) + dq_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=dq_producer_group, + ) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + num_prev_d_subtiles = cutlass.Int32(0) + # Mirror ACT's per-tile subtile-column reversal (overlapping_accum phase). + # ACT's first tile uses acc phase 0 => reverse=True, then toggles each tile. + if cutlass.const_expr(self.run_rht and self.overlapping_accum): + rht_reverse = cutlass.Boolean(True) + while is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + mma_tile_coord_mnl = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx, + cutlass.Int32(0), + ) + expert_idx = epi_work_tile_info.expert_idx + epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx) + + # + # RHT output: per-expert RHT gmem tensor + TMA partition (mirrors ACT's D setup). + # + if cutlass.const_expr(self.generate_rht): + thr_mma_epi_rht = tiled_mma.get_slice(mma_tile_coord_v) + real_rht, _ = epi_ext.get_gmem_tensor("d", mRht_mnl, padded_offsets, epi_work_tile_info) + gRht_mnl_loop = cute.local_tile(real_rht, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgRht = thr_mma_epi_rht.partition_C(gRht_mnl_loop) + if cutlass.const_expr(self.enable_breuse): + gRht_epi_bk = cute.flat_divide(tCgRht[((None, None), 0, 0, None, None, None)], epi_tile) + gRht_epi_br = cute.flat_divide(tCgRht[((None, None), 1, 0, None, None, None)], epi_tile) + sRht_for_tma = cute.group_modes(sRht, 0, 2) + bSG_sRht, bSG_gRht_partitioned_bk = cpasync.tma_partition( + tma_atom_rht, 0, cute.make_layout(1), sRht_for_tma, cute.group_modes(gRht_epi_bk, 0, 2)) + _, bSG_gRht_partitioned_br = cpasync.tma_partition( + tma_atom_rht, 0, cute.make_layout(1), sRht_for_tma, cute.group_modes(gRht_epi_br, 0, 2)) + bSG_gRht_bk = bSG_gRht_partitioned_bk[(None, None, None, *mma_tile_coord_mnl)] + bSG_gRht_bk = cute.group_modes(bSG_gRht_bk, 1, cute.rank(bSG_gRht_bk)) + bSG_gRht_br = bSG_gRht_partitioned_br[(None, None, None, *mma_tile_coord_mnl)] + bSG_gRht_br = cute.group_modes(bSG_gRht_br, 1, cute.rank(bSG_gRht_br)) + bSG_gRht = bSG_gRht_bk + else: + _, bSG_sRht, bSG_gRht_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_rht, tCgRht, epi_tile, sRht) + bSG_gRht = bSG_gRht_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gRht = cute.group_modes(bSG_gRht, 1, cute.rank(bSG_gRht)) + bSG_gRht_bk = bSG_gRht + bSG_gRht_br = bSG_gRht + if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): + # Expert token offset for the colwise (f, m/16) scale grid's + # tile index (offsets are 256-aligned, divisions exact). + rht_t_off, _rht_t_cnt = compute_expert_token_range( + padded_offsets, epi_work_tile_info.expert_idx) + + # + # NVFP4 D: per-expert fp4 D gmem tensor + TMA partition (mirrors ACT's D setup). + # + if cutlass.const_expr(self.d_quant): + thr_mma_epi_dq = tiled_mma.get_slice(mma_tile_coord_v) + real_dq, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + gDq_mnl_loop = cute.local_tile(real_dq, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgDq = thr_mma_epi_dq.partition_C(gDq_mnl_loop) + if cutlass.const_expr(self.enable_breuse): + gDq_epi_bk = cute.flat_divide(tCgDq[((None, None), 0, 0, None, None, None)], epi_tile) + gDq_epi_br = cute.flat_divide(tCgDq[((None, None), 1, 0, None, None, None)], epi_tile) + sDq_for_tma = cute.group_modes(sDq, 0, 2) + bSG_sDq, bSG_gDq_partitioned_bk = cpasync.tma_partition( + tma_atom_d, 0, cute.make_layout(1), sDq_for_tma, cute.group_modes(gDq_epi_bk, 0, 2)) + _, bSG_gDq_partitioned_br = cpasync.tma_partition( + tma_atom_d, 0, cute.make_layout(1), sDq_for_tma, cute.group_modes(gDq_epi_br, 0, 2)) + bSG_gDq_bk = bSG_gDq_partitioned_bk[(None, None, None, *mma_tile_coord_mnl)] + bSG_gDq_bk = cute.group_modes(bSG_gDq_bk, 1, cute.rank(bSG_gDq_bk)) + bSG_gDq_br = bSG_gDq_partitioned_br[(None, None, None, *mma_tile_coord_mnl)] + bSG_gDq_br = cute.group_modes(bSG_gDq_br, 1, cute.rank(bSG_gDq_br)) + bSG_gDq = bSG_gDq_bk + else: + _, bSG_sDq, bSG_gDq_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgDq, epi_tile, sDq) + bSG_gDq = bSG_gDq_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gDq = cute.group_modes(bSG_gDq, 1, cute.rank(bSG_gDq)) + bSG_gDq_bk = bSG_gDq + bSG_gDq_br = bSG_gDq + + # + # NVFP4 scale tensors: per-expert (m, f/16, 1) gmem tensors + this + # thread's row. The "d" extension branch only applies the expert row + # offset + shape rewrite, so it works for the scale tensors as-is. + # Each thread owns one token row (mPosition analog); per subtile + # iteration it stages one scale byte per 16-feature block into smem, + # then stores its whole (contiguous) 8-byte scale row once per tile — + # per-block scalar gmem stores would be ~32x sector-amplified. + # + if cutlass.const_expr(self.rht_quant or self.d_quant): + sf_row = ( + (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0] + + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape)) + + epi_tidx + ) + # For breuse: the ACT warps emit the subtile sequence once per + # M-half (bkeep then breuse) — consume both, with per-half gmem + # partitions and scale rows. The per-half scale-row stores keep + # the sSfRht/sSfd smem staging correct (half 1 reuses it). + # (2-space indent below keeps the subtile-loop body untouched.) + _breuse_halves = [0, 1] if self.enable_breuse else [0] + + if cutlass.const_expr(self.generate_rht): + _bSG_gRht_h = bSG_gRht + if cutlass.const_expr(self.d_quant): + _bSG_gDq_h = bSG_gDq + if cutlass.const_expr(self.rht_quant or self.d_quant): + _sf_row_h = sf_row + + for _m_half in _breuse_halves: + if self.enable_breuse: + if cutlass.const_expr(self.generate_rht): + _bSG_gRht_h = bSG_gRht_bk if _m_half == 0 else bSG_gRht_br + if cutlass.const_expr(self.d_quant): + _bSG_gDq_h = bSG_gDq_bk if _m_half == 0 else bSG_gDq_br + if cutlass.const_expr(self.rht_quant or self.d_quant): + _sf_row_h = sf_row + _m_half * (self.cta_tile_shape_mnk[0] // 2) + else: + if cutlass.const_expr(self.generate_rht): + _bSG_gRht_h = bSG_gRht + if cutlass.const_expr(self.d_quant): + _bSG_gDq_h = bSG_gDq + if cutlass.const_expr(self.rht_quant or self.d_quant): + _sf_row_h = sf_row + + # + # Consume D subtiles from SMEM + # + subtile_cnt = self.cta_tile_shape_mnk[1] // cute.size(self.epi_tile[1]) + subtile_step = 1 if cutlass.const_expr(self.act_func == "srelu") else 2 + for subtile_idx in cutlass.range(0, subtile_cnt, subtile_step, unroll=1): + # + # Wait for ACT warps to finish writing the post-activation D tile to SMEM. + # + pingpong_pipeline.consumer_wait(pingpong_rht_consumer_state) + + if cutlass.const_expr(self.run_rht): + d_buffer = num_prev_d_subtiles % self.num_d_stage + # Subtile-column reversal (overlapping_accum): computed BEFORE + # the FWHT/quant calls — scale bytes and TMA'd data must land + # in the same reversed gmem feature columns as ACT's D store. + real_subtile_idx = subtile_idx // 2 + if cutlass.const_expr(self.overlapping_accum): + if rht_reverse: + real_subtile_idx = ( + self.cta_tile_shape_mnk[1] // self.epi_tile_n_required + - 1 - real_subtile_idx + ) + # + # Load this subtile's sD values to registers, then release + # the pingpong slot IMMEDIATELY: the ACT producers' acquire + # (before their sD write) is the only reuse gate, so the + # whole FWHT/quant/store chain below overlaps the ACT warps' + # next subtiles (up to num_d_stage ahead). The mbarrier + # phase bit makes this early release sound — the old + # named-barrier arrive() variant raced its release drain + # (~1/300 hangs on sm100 silicon). + # + if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.rht_rowwise): + rht_ld = load_row_bf16(sD, d_buffer, epi_tidx) + else: + rht_ld = load_colwise_pairs_bf16(sD, d_buffer, epi_tidx) + if cutlass.const_expr(self.d_quant): + dq_ld = load_row_bf16(sD, d_buffer, epi_tidx) + pingpong_pipeline.consumer_release(pingpong_rht_consumer_state) + pingpong_rht_consumer_state.advance() + # RHT output: write the transform (x0.25) to sRht. rht_rowwise + # picks the transform axis (16-feature blocks per token vs + # 16-token blocks per feature); NVFP4 quantization is inferred + # from the sRht dtype inside the FWHT device functions. + if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.rht_quant and self.rht_rowwise): + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, + rht_norm_const, sSfRht, real_subtile_idx, self.sf_dtype) + elif cutlass.const_expr(self.rht_quant): + hadamard_rmem_colwise_fwht_quant( + rht_ld, d_buffer, epi_tidx, rht_norm_const, + sRht, sSfRht, real_subtile_idx * 2 * HADAMARD_SIZE, self.sf_dtype) + elif cutlass.const_expr(self.rht_rowwise): + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, 1.0, None, 0, self.sf_dtype) + else: + hadamard_rmem_colwise_fwht(rht_ld, d_buffer, epi_tidx, sRht) + # NVFP4 D: quantize the bf16 register rows into sDq + sSfd. + if cutlass.const_expr(self.d_quant): + nvfp4_quant_rmem_row(dq_ld, d_buffer, epi_tidx, sDq, + norm_const, sSfd, real_subtile_idx, self.sf_dtype) + # + # TMA-store the produced epi-tiles to gmem (mirrors ACT's D + # store, including the overlapping_accum subtile-column reversal). + # + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier_group1.arrive_and_wait() + if warp_idx == self.epilog_rht_store_warp_id[0]: + if cutlass.const_expr(self.generate_rht): + cute.copy( + tma_atom_rht, + bSG_sRht[(None, d_buffer)], + _bSG_gRht_h[(None, real_subtile_idx)], + ) + rht_pipeline.producer_commit() + rht_pipeline.producer_acquire() + if cutlass.const_expr(self.d_quant): + cute.copy( + tma_atom_d, + bSG_sDq[(None, d_buffer)], + _bSG_gDq_h[(None, real_subtile_idx)], + ) + dq_pipeline.producer_commit() + dq_pipeline.producer_acquire() + self.epilog_sync_barrier_group1.arrive_and_wait() + + # (Pingpong slot already released right after the sD register + # loads above.) + num_prev_d_subtiles = num_prev_d_subtiles + 1 + + # + # NVFP4: one contiguous 8-byte scale-row store per thread PER + # M-HALF pass (still inside the m-half loop: half 1 reuses the + # sSfRht/sSfd smem staging, so each half must drain it first). + # Rowwise: the thread wrote all its own smem slots (no barrier + # needed). Colwise: rows were filled ACROSS warps — the group1 + # barrier at the end of the subtile loop already ordered those + # writes. + # + _num_sf = self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE + if cutlass.const_expr(self.rht_quant and self.rht_rowwise): + self.store_swizzled_sf_row( + mSfRht_mnl, + epi_ext.token_offset + _sf_row_h, + mma_tile_coord_mnl[1] * _num_sf, + sSfRht, + epi_tidx, + ) + if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): + # (f, m) scale domain: thread <-> feature-in-tile; columns are + # 16-token scale blocks, stored in the same swizzled SF atom layout. + sf_feat_row = mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk_d[1] + epi_tidx + self.store_swizzled_sf_row( + mSfRht_mnl, + sf_feat_row, + (rht_t_off + _sf_row_h - epi_tidx) // HADAMARD_SIZE, + sSfRht, + epi_tidx, + ) + if cutlass.const_expr(self.d_quant): + self.store_swizzled_sf_row( + mSfd_mnl, + epi_ext.token_offset + _sf_row_h, + mma_tile_coord_mnl[1] * _num_sf, + sSfd, + epi_tidx, + ) + + # Under breuse the next M-half rewrites the sSfRht/sSfd staging + # that the (cross-warp for colwise) reads above consume — order + # the halves. Compiled out for non-breuse. + if cutlass.const_expr(self.enable_breuse and (self.rht_quant or self.d_quant)): + self.epilog_sync_barrier_group1.arrive_and_wait() + + # + # (No tile-end ACT<->RHT barrier: sD reuse is pingpong-guarded.) + # The NEXT tile's FWHT/quant rewrites the sSfRht/sSfd staging the + # (cross-warp for colwise) scale-row reads above consume — order + # within the RHT group. Breuse already barriers per half above. + # + if cutlass.const_expr(not self.enable_breuse and (self.rht_quant or self.d_quant)): + self.epilog_sync_barrier_group1.arrive_and_wait() + + # + # Toggle the subtile-column reversal for the next tile. + # + if cutlass.const_expr(self.run_rht and self.overlapping_accum): + rht_reverse = not rht_reverse + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Drain the RHT-warp TMA-store pipelines before exit. + # + if cutlass.const_expr(self.generate_rht): + rht_pipeline.producer_tail() + if cutlass.const_expr(self.d_quant): + dq_pipeline.producer_tail() + # END OF KERNEL diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py index 05d1ab870..158323dd5 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/moe_kernel_helpers.py @@ -598,6 +598,8 @@ def get_dtype_rcp_limits(dtype: Type[cutlass.Numeric]) -> float: return 1 / 448.0 if dtype == cutlass.Float8E5M2: return 1 / 128.0 + if dtype == cutlass.FloatNV8E5M3FNU: + return 1 / 114688.0 return 1.0 diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py index 50d55725f..369d5521a 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/quant_utils.py @@ -17,6 +17,8 @@ def get_dtype_rcp_limits(dtype): return 1 / 6.0 if dtype == cutlass.Float8E4M3FN: return 1 / 448.0 + if dtype == cutlass.FloatNV8E5M3FNU: + return 1 / 114688.0 raise ValueError(f"unsupported quantized dtype {dtype}") @@ -34,22 +36,22 @@ def load_row_bf16(sD, d_buffer, tidx): @cute.jit -def nvfp4_quant_rmem_row(rmem_bf16, d_buffer, tidx, sOut, norm_const, sSf, sf_pair): +def nvfp4_quant_rmem_row(rmem_bf16, d_buffer, tidx, sOut, norm_const, sSf, sf_pair, sf_dtype): """NVFP4-quantize the thread's full 32-feature bf16 row (from load_row_bf16, no - transform) into sOut (packed e2m1) + sSf (e4m3 block scales) — the fused-D + transform) into sOut (packed e2m1) + sSf (fp8 block scales) — the fused-D analog of group_rht_cast's rowwise cast path.""" row_feats = 2 * HADAMARD_SIZE tCompute = cute.make_rmem_tensor(rmem_bf16.shape, cutlass.Float32) for i in cutlass.range_constexpr(row_feats): tCompute[i] = rmem_bf16[i].to(cutlass.Float32) - _nvfp4_quant_row(tCompute, d_buffer, tidx, sOut, norm_const, sSf, sf_pair) + _nvfp4_quant_row(tCompute, d_buffer, tidx, sOut, norm_const, sSf, sf_pair, sf_dtype) @cute.jit -def _nvfp4_quant_row(tCompute, d_buffer, token, sOut, norm_const, sSf, sf_pair): +def _nvfp4_quant_row(tCompute, d_buffer, token, sOut, norm_const, sSf, sf_pair, sf_dtype): """Flashinfer blockscaled-epilogue NVFP4 quantization of one thread's 32-feature - f32 row: packed e2m1 data into sOut, one e4m3 block scale per (1,16) block into + f32 row: packed e2m1 data into sOut, one fp8 block scale per (1,16) block into sSf slots (token, num_vecs*sf_pair + vi).""" row_feats = cute.size(tCompute.shape) tCompute_flat = cute.make_tensor(tCompute.iterator, cute.make_layout((row_feats,))) @@ -71,32 +73,36 @@ def _nvfp4_quant_row(tCompute, d_buffer, token, sOut, norm_const, sSf, sf_pair): ) # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2856-2873 for vi in cutlass.range_constexpr(0, num_vecs, 2): - tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = cute.arch.mul_packed_f32x2( - (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), - ( - get_dtype_rcp_limits(cutlass.Float4E2M1FN), - get_dtype_rcp_limits(cutlass.Float4E2M1FN), - ), + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = ( + cute.arch.mul_packed_f32x2( + (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), + ( + get_dtype_rcp_limits(cutlass.Float4E2M1FN), + get_dtype_rcp_limits(cutlass.Float4E2M1FN), + ), + ) ) - tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = cute.arch.mul_packed_f32x2( - (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), - (norm_const, norm_const), + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = ( + cute.arch.mul_packed_f32x2( + (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), + (norm_const, norm_const), + ) ) # blockscaled_contiguous_gather_grouped_gemm_act_fusion.py:2887 - # f32 -> e4m3 as a padded 4-wide vector: nvgpu.cvt_fptrunc requires a + # f32 -> fp8 scale format as a padded 4-wide vector: nvgpu.cvt_fptrunc requires a # 32-bit-aligned 1-d vector (4 x f8), never a scalar. Same pattern as # grouped_gemm_swiglu_quant's pvscale_f32x4 -> sfd_f8x4 round-trip. assert num_vecs <= 4, "scale convert padding assumes at most 4 blocks per row" pvscale_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) for vi in cutlass.range_constexpr(4): pvscale_f32x4[vi] = tCrSFC_pvscale[min(vi, num_vecs - 1)] - tCrSFC_f8x4 = cute.make_rmem_tensor((4,), cutlass.Float8E4M3FN) - tCrSFC_f8x4.store(pvscale_f32x4.load().to(cutlass.Float8E4M3FN)) - tCrSFC = cute.make_rmem_tensor((num_vecs,), cutlass.Float8E4M3FN) + tCrSFC_f8x4 = cute.make_rmem_tensor((4,), sf_dtype) + tCrSFC_f8x4.store(pvscale_f32x4.load().to(sf_dtype)) + tCrSFC = cute.make_rmem_tensor((num_vecs,), sf_dtype) for vi in cutlass.range_constexpr(num_vecs): tCrSFC[vi] = tCrSFC_f8x4[vi] - # e4m3 -> f32 widening round-trip, vectorized for the same reason. + # fp8 scale -> f32 widening round-trip, vectorized for the same reason. tCrSFC_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) tCrSFC_f32x4.store(tCrSFC_f8x4.load().to(cutlass.Float32)) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py index 13ddccf86..293d0b02c 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/rht_utils.py @@ -100,7 +100,8 @@ def hadamard_rmem_colwise_fwht(rmem_bf16, d_buffer, tidx, sRht): @cute.jit -def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht, sSf, sf_row_base): +def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht, + sSf, sf_row_base, sf_dtype): """Colwise FWHT + NVFP4 quantization from the load_colwise_pairs_bf16 registers, stored at the SAME (token, feature) coords the input was read from (the staging is f-major like every other output; packed nibbles pair ADJACENT FEATURES of one @@ -139,22 +140,24 @@ def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht pv[c] = fmax(pv[c], fmax(v, -v)) pv[0], pv[1] = cute.arch.mul_packed_f32x2((pv[0], pv[1]), (gem, gem)) - # f32 -> e4m3 as a padded 4-wide vector (and widen back the same way): + # f32 -> fp8 scale format as a padded 4-wide vector (and widen back the same way): # nvgpu.cvt_fptrunc requires a 32-bit-aligned 1-d vector (4 x f8), never a # scalar. Same pattern as grouped_gemm_swiglu_quant's pvscale round-trip. pv_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) for c in cutlass.range_constexpr(4): pv_f32x4[c] = pv[min(c, 1)] - tCrSFC_f8x4 = cute.make_rmem_tensor((4,), cutlass.Float8E4M3FN) - tCrSFC_f8x4.store(pv_f32x4.load().to(cutlass.Float8E4M3FN)) + tCrSFC_f8x4 = cute.make_rmem_tensor((4,), sf_dtype) + tCrSFC_f8x4.store(pv_f32x4.load().to(sf_dtype)) tCrSFC_f32x4 = cute.make_rmem_tensor((4,), cutlass.Float32) tCrSFC_f32x4.store(tCrSFC_f8x4.load().to(cutlass.Float32)) for c in cutlass.range_constexpr(2): sSf[(sf_row_base + feature + c, token_block)] = tCrSFC_f8x4[c] fp32_max = cutlass.Float32(3.40282346638528859812e38) - acc_scale_min0 = fmin(cutlass.Float32(1.0) / (tCrSFC_f32x4[0] * gd), fp32_max, nan=True) - acc_scale_min1 = fmin(cutlass.Float32(1.0) / (tCrSFC_f32x4[1] * gd), fp32_max, nan=True) + acc_scale_min0 = fmin( + cutlass.Float32(1.0) / (tCrSFC_f32x4[0] * gd), fp32_max, nan=True) + acc_scale_min1 = fmin( + cutlass.Float32(1.0) / (tCrSFC_f32x4[1] * gd), fp32_max, nan=True) for i in cutlass.range_constexpr(HADAMARD_SIZE): tCompute[2 * i], tCompute[2 * i + 1] = cute.arch.mul_packed_f32x2( (tCompute[2 * i], tCompute[2 * i + 1]), @@ -173,11 +176,11 @@ def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht @cute.jit -def hadamard_rmem_rowwise_fwht(rmem_bf16, d_buffer, tidx, sRht, norm_const=1.0, sSf=None, sf_pair=0): +def hadamard_rmem_rowwise_fwht(rmem_bf16, d_buffer, tidx, sRht, norm_const, sSf, sf_pair, sf_dtype): """Per-token FWHT over the thread's full 32-feature row (two independent 16-feature Hadamard blocks) from the load_row_bf16 registers, stored to sRht at the same - (token, feature) coords: bf16, or NVFP4 (packed e2m1 into sRht + e4m3 block scales - into sSf slots (tidx, 2*sf_pair..)) when sRht is fp4-typed.""" + (token, feature) coords: bf16, or NVFP4 (packed e2m1 into sRht + e4m3/ue5m3 block + scales into sSf slots (tidx, 2*sf_pair..)) when sRht is fp4-typed.""" token = tidx row_feats = 2 * HADAMARD_SIZE @@ -214,7 +217,7 @@ def hadamard_rmem_rowwise_fwht(rmem_bf16, d_buffer, tidx, sRht, norm_const=1.0, tCompute[i] = (tCompute[i] * cutlass.Float32(0.25)).to(cutlass.BFloat16).to(cutlass.Float32) if cutlass.const_expr(sRht.element_type == cutlass.Float4E2M1FN): - _nvfp4_quant_row(tCompute, d_buffer, token, sRht, norm_const, sSf, sf_pair) + _nvfp4_quant_row(tCompute, d_buffer, token, sRht, norm_const, sSf, sf_pair, sf_dtype) else: rmem_st = cute.make_rmem_tensor(rmem_bf16.shape, cutlass.BFloat16) for i in cutlass.range_constexpr(row_feats): diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py index 96ae2b630..757bb3b9d 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py @@ -11,7 +11,11 @@ from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import allocate_discrete_input_tensors from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import allocate_grouped_gemm_input_tensors, grouped_gemm_swiglu_init from fe_api.grouped_gemm.test_grouped_gemm_wgrad_utils import _skip_unless_e5m3_supported -from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES, reencode_sf_tensor_as_ue5m3 +from fe_api.test_fe_api_utils import ( + DYNAMIC_SHAPES_M_VALUES, + reencode_sf_tensor_as_ue5m3, + ue5m3_bytes_to_fp32, +) FP4_EXECUTION_CASES = [ (torch.float4_e2m1fn_x2, torch.float8_e4m3fn, 16), @@ -205,6 +209,7 @@ def _check_outputs( norm_const: float = 1.0, rht_norm_const: float = 1.0, ref_tensors: Optional[Dict] = None, + sf_fp8_dtype_override: Optional[str] = None, ) -> None: if ref_tensors is None: ref_tensors = _run_grouped_gemm_glu_ref(inputs, act_func, glu_alpha=glu_alpha, glu_limit=glu_limit) @@ -227,9 +232,12 @@ def _check_outputs( rtol=1e-2, ) else: + sfd_tensor = outputs["sfd_tensor"] + if sf_fp8_dtype_override == "e5m3": + sfd_tensor = ue5m3_bytes_to_fp32(sfd_tensor) _check_nvfp4_output( float4_e2m1fn_x2_to_float32(outputs["d_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()), - outputs["sfd_tensor"].cpu(), + sfd_tensor.cpu(), d_ref[:, :, 0].cpu(), norm_const, "D", @@ -249,9 +257,12 @@ def _check_outputs( rtol=1e-2, ) elif rht_rowwise: + sfrht_tensor = outputs["sfrht_tensor"] + if sf_fp8_dtype_override == "e5m3": + sfrht_tensor = ue5m3_bytes_to_fp32(sfrht_tensor) _check_nvfp4_output( float4_e2m1fn_x2_to_float32(outputs["rht_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()), - outputs["sfrht_tensor"].cpu(), + sfrht_tensor.cpu(), rht_ref.float(), rht_norm_const, "RHT", @@ -261,9 +272,12 @@ def _check_outputs( # adjacent features), but quantization blocks are (16, 1) token blocks, # so check through the transposed unpacked values and the swizzled # SF(N_out, valid_m) scale domain. + sfrht_tensor = outputs["sfrht_tensor"] + if sf_fp8_dtype_override == "e5m3": + sfrht_tensor = ue5m3_bytes_to_fp32(sfrht_tensor) _check_nvfp4_output( float4_e2m1fn_x2_to_float32(outputs["rht_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()).t(), - outputs["sfrht_tensor"].cpu(), + sfrht_tensor.cpu(), rht_ref.float().t(), rht_norm_const, "RHT", @@ -366,6 +380,7 @@ def _run_wrapper( norm_const=norm_const, rht_norm_const=rht_norm_const, ref_tensors=ref_tensors, + sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -443,7 +458,15 @@ def alloc_n_major(rows, cols, dtype): rht_tensor=outputs["rht_tensor"], ) - _check_outputs(inputs, outputs, cfg, act_func=act_func, rht_output=True, rht_rowwise=False) + _check_outputs( + inputs, + outputs, + cfg, + act_func=act_func, + rht_output=True, + rht_rowwise=False, + sf_fp8_dtype_override=sf_fp8_dtype_override, + ) def _run_discrete_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func="swiglu", sf_fp8_dtype_override=None): @@ -489,7 +512,15 @@ def _run_discrete_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func= act_func=act_func, ) - _check_outputs(inputs, outputs, cfg, act_func=act_func, rht_output=True, rht_rowwise=False) + _check_outputs( + inputs, + outputs, + cfg, + act_func=act_func, + rht_output=True, + rht_rowwise=False, + sf_fp8_dtype_override=sf_fp8_dtype_override, + ) # ============================================================================= @@ -741,13 +772,63 @@ def test_grouped_gemm_glu_hadamard_quant_rejects_unsupported_sf_fp8_dtype(reques """e5m3 is only reachable through the FP4xFP4 atom with e4m3-carried scales.""" if sf_fp8_dtype_override == "e5m3": _skip_unless_e5m3_supported() + + # Construct kernel config + cfg_kwargs = dict( + ab_dtype=torch.float4_e2m1fn_x2, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=16, + sf_dtype=torch.float8_e4m3fn, + vector_f32=False, + discrete_col_sfd=False, + enable_bias=False, + ) + cfg_kwargs.update(overrides) + cfg = grouped_gemm_swiglu_init(request, **cfg_kwargs) + + # Allocate input tensors + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + enable_bias=False, + ) + + from cudnn import grouped_gemm_glu_hadamard_quant_wrapper_sm100 + + # Check that calling kernel API triggers exception with pytest.raises(ValueError, match=expected): - _run_compile_execute( - request, - ab_dtype=torch.float4_e2m1fn_x2, - sf_dtype=overrides.get("sf_dtype", torch.float8_e4m3fn), - sf_vec_size=overrides.get("sf_vec_size", 16), + grouped_gemm_glu_hadamard_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + bias_tensor=inputs["bias_tensor"], + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func="swiglu", ) diff --git a/test/python/fe_api/test_fe_api_utils.py b/test/python/fe_api/test_fe_api_utils.py index 3c6f4ac6a..5abfca8a7 100644 --- a/test/python/fe_api/test_fe_api_utils.py +++ b/test/python/fe_api/test_fe_api_utils.py @@ -177,6 +177,12 @@ def f32_to_ue5m3_bytes(values: torch.Tensor) -> torch.Tensor: return idx.to(torch.uint8).masked_fill(is_nan, _UE5M3_NAN_BYTE).reshape(values.shape) +def ue5m3_bytes_to_fp32(encoded: torch.Tensor) -> torch.Tensor: + """Decode a UE5M3 byte tensor to float.""" + lut = _ue5m3_lut(encoded.device) + return lut[encoded.view(torch.uint8).to(torch.int)] + + def reencode_sf_tensor_as_ue5m3(sf_tensor: torch.Tensor) -> torch.Tensor: """Rewrite an e4m3-valued scale-factor tensor's bytes as UE5M3, in place.""" assert sf_tensor.dtype == torch.float8_e4m3fn, f"expected e4m3 storage, got {sf_tensor.dtype}"