From 13ff529113fbcb77c13407dd43c4ea9a71192e2b Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 09:54:19 +0000 Subject: [PATCH 1/2] GGEMM+GLU+RHT+quant outputs column-wise RHT in ragged tensor layout Co-authored-by: Codex Signed-off-by: Tim Moon --- .../grouped_gemm_glu_hadamard_quant.md | 58 ++-- .../cutedsl/grouped/glu_hadamard_quant/api.py | 255 ++++++++++++------ ...kscaled_grouped_gemm_glu_hadamard_quant.py | 70 +++-- ...d_grouped_gemm_glu_hadamard_quant_rubin.py | 79 ++++-- .../grouped/glu_hadamard_quant/rht_utils.py | 124 +++++++-- .../test_grouped_gemm_glu_hadamard_quant.py | 194 ++++++++----- 6 files changed, 544 insertions(+), 236 deletions(-) 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 index 1ccf0ff63..e82c2e251 100644 --- 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 @@ -41,12 +41,14 @@ Let `N_out = N / 2` for `act_func="swiglu"` or `"geglu"` and `N_out = N` for `ac - `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)` + - `RHT rowwise`: optional Hadamard-transform output across feature blocks, logical shape `(valid_m, N_out, 1)` + - `SFRHT rowwise`: scale factors for NVFP4 rowwise `RHT`, shape `SF(valid_m, N_out)`, present only when rowwise `RHT` is NVFP4 + - `RHT colwise`: optional NVFP4 Hadamard-transform output across token blocks, flattened in per-expert transposed order + - `SFRHT colwise`: scale factors for NVFP4 colwise `RHT`, flattened in per-expert ragged order -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. +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 rowwise `RHT` tensors with physical second dimension `N_out / 2`. Raw `torch.uint8` tensors are not accepted as a packed FP4 container by this fusion. + +Colwise `RHT` is NVFP4-only. Its data tensor is a flat physical tensor of length `valid_m * N_out / 2`; after FP4 unpacking, each expert segment is logical `(N_out, expert_m)` with adjacent token values packed together. Expert segments are concatenated in expert order. `SFRHT colwise` is flat length `valid_m * N_out / sf_vec_size`; each expert segment is `SF(N_out, expert_m)` flattened in expert order, so the scale segment for an expert starts at `N_out * expert_m_prefix / sf_vec_size`. `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`. @@ -86,7 +88,7 @@ $$ 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 requested, the RHT output applies a fixed 16-wide orthonormal Hadamard transform to bf16-rounded `D`. Rowwise RHT transforms feature blocks. Colwise RHT transforms token blocks and writes NVFP4 data in per-expert transposed order. 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. @@ -138,9 +140,7 @@ result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( 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, + rht_colwise_dtype=torch.float4_e2m1fn_x2, norm_const=norm_const, rht_norm_const=rht_norm_const, mma_tiler_mn=(256, 256), @@ -156,11 +156,11 @@ result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( 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"] +rht_colwise_tensor = result["rht_colwise_tensor"] +sfrht_colwise_tensor = result["sfrht_colwise_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. +Leave both `rht_rowwise_dtype` and `rht_colwise_dtype` unset to skip the Hadamard/RHT output. Set `rht_rowwise_dtype=torch.bfloat16` for unquantized rowwise RHT, `rht_rowwise_dtype=torch.float4_e2m1fn_x2` for quantized rowwise RHT, or `rht_colwise_dtype=torch.float4_e2m1fn_x2` for wgrad-compatible colwise RHT. ### Class API @@ -178,8 +178,8 @@ op = GroupedGemmGluHadamardQuantSm100( sample_alpha=alpha, sample_prob=prob, sample_sfd=sfd, - sample_rht=rht, - sample_sfrht=sfrht, + sample_rht_colwise=rht_colwise, + sample_sfrht_colwise=sfrht_colwise, sample_bias=bias, acc_dtype=torch.float32, mma_tiler_mn=(256, 256), @@ -189,7 +189,6 @@ op = GroupedGemmGluHadamardQuantSm100( vector_f32=False, m_aligned=256, act_func="swiglu", - rht_rowwise=False, ) assert op.check_support() op.compile() @@ -204,8 +203,8 @@ op.execute( alpha_tensor=alpha, prob_tensor=prob, sfd_tensor=sfd, - rht_tensor=rht, - sfrht_tensor=sfrht, + rht_colwise_tensor=rht_colwise, + sfrht_colwise_tensor=sfrht_colwise, bias_tensor=bias, norm_const=norm_const, rht_norm_const=rht_norm_const, @@ -290,15 +289,24 @@ result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( - 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) +- Output tensor **RHT rowwise** (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` + - NVFP4 rowwise `RHT` requires `SFRHT rowwise` +- Output tensor **SFRHT rowwise** (present only with NVFP4 rowwise `RHT`) + - Shape: `SF(valid_m, N_out)` - Layout: swizzled scale-factor layout - - Dtype: `float8_e4m3fn` + - Dtype: `float8_e4m3fn`, independent of `SFA`/`SFB` dtype +- Output tensor **RHT colwise** (optional) + - Logical unpacked shape: per-expert `(N_out, expert_m)` segments concatenated in expert order + - Physical shape: 1-D packed FP4 tensor with `valid_m * N_out / 2` elements + - Dtype: `float4_e2m1fn_x2` + - NVFP4 colwise `RHT` always requires `SFRHT colwise` +- Output tensor **SFRHT colwise** (present with colwise `RHT`) + - Shape: `(valid_m * N_out / sf_vec_size,)` + - Layout: flattened per-expert `SF(N_out, expert_m)` segments in expert order + - Dtype: `float8_e4m3fn`, independent of `SFA`/`SFB` dtype ### Configuration @@ -309,7 +317,8 @@ result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( - `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 +- `rht_rowwise_dtype`: unset to skip rowwise RHT, `torch.bfloat16` for unquantized rowwise RHT, or `torch.float4_e2m1fn_x2` for NVFP4 rowwise RHT +- `rht_colwise_dtype`: unset to skip colwise RHT, or `torch.float4_e2m1fn_x2` for NVFP4 colwise RHT in per-expert transposed order - `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` @@ -325,6 +334,7 @@ result = grouped_gemm_glu_hadamard_quant_wrapper_sm100( - 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)`. +- At most one of rowwise and colwise RHT may be requested. +- Colwise RHT is NVFP4-only and uses flat per-expert ragged scale storage. - `expert_cnt` must be `<= 1024`. - Dense and discrete weight modes are mutually exclusive. 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 dd41b2c53..f6dbb686e 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard_quant/api.py @@ -5,12 +5,8 @@ Output modes are dtype driven, mirroring the kernel: - 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) -when ``rht_rowwise``, and swizzled scale factors for logical (f, m) otherwise. + - Optional rowwise RHT is bf16, or NVFP4 with swizzled rowwise scales + - Optional colwise RHT is NVFP4 in wgrad-compatible per-expert ragged order """ from __future__ import annotations @@ -85,8 +81,10 @@ def __init__( 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_rht_rowwise: Optional[torch.Tensor] = None, + sample_sfrht_rowwise: Optional[torch.Tensor] = None, + sample_rht_colwise: Optional[torch.Tensor] = None, + sample_sfrht_colwise: Optional[torch.Tensor] = None, sample_bias: Optional[torch.Tensor] = None, acc_dtype: Optional[torch.dtype] = None, mma_tiler_mn: Tuple[int, int] = (256, 256), @@ -97,7 +95,6 @@ def __init__( 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, ): @@ -114,7 +111,16 @@ def __init__( self._sample_a_tensor = sample_a self._sample_b_tensor = sample_b self._sample_d_tensor = sample_d - self._sample_rht_tensor = sample_rht + self.generate_rht_rowwise = sample_rht_rowwise is not None + self.generate_rht_colwise = sample_rht_colwise is not None + if self.generate_rht_rowwise and self.generate_rht_colwise: + raise NotImplementedError("Request at most one of rowwise or colwise RHT output") + if sample_sfrht_rowwise is not None and sample_rht_rowwise is None: + raise ValueError("sample_sfrht_rowwise requires sample_rht_rowwise") + if sample_sfrht_colwise is not None and sample_rht_colwise is None: + raise ValueError("sample_sfrht_colwise requires sample_rht_colwise") + self.rht_rowwise = self.generate_rht_rowwise + self.rht_per_expert = self.generate_rht_colwise if sample_b is not None and num_experts is None: self.weight_mode = MoEWeightMode.DENSE @@ -135,8 +141,12 @@ def __init__( 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.rht_rowwise_desc = self._make_tensor_desc(sample_rht_rowwise, name="sample_rht_rowwise", interpret_uint8_as_fp4x2=False) + self.sfrht_rowwise_desc = self._make_tensor_desc(sample_sfrht_rowwise, name="sample_sfrht_rowwise") + self.rht_colwise_desc = self._make_tensor_desc(sample_rht_colwise, name="sample_rht_colwise", interpret_uint8_as_fp4x2=False) + self.sfrht_colwise_desc = self._make_tensor_desc(sample_sfrht_colwise, name="sample_sfrht_colwise") + self.rht_desc = self.rht_rowwise_desc if self.generate_rht_rowwise else self.rht_colwise_desc + self.sfrht_desc = self.sfrht_rowwise_desc if self.generate_rht_rowwise else self.sfrht_colwise_desc 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) @@ -167,7 +177,6 @@ def __init__( 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 = _get_rubin_kernel() if self._is_rubin_kernel else BlockScaledMoEGroupedGemmGluHadamardQuantKernel @@ -190,6 +199,7 @@ def check_support(self) -> bool: _, 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.n_out = n_out 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}") @@ -197,7 +207,7 @@ def check_support(self) -> bool: # ---- 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.generate_rht = self.generate_rht_rowwise or self.generate_rht_colwise 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), @@ -207,6 +217,10 @@ def check_support(self) -> bool: self.rht_quant != (self.sfrht_desc is not None), "NVFP4 sample_rht and sample_sfrht must be passed together", ) + self._value_error_if( + self.generate_rht_colwise and not self.rht_quant, + "Colwise RHT output is supported only as NVFP4; pass sample_rht_colwise with dtype torch.float4_e2m1fn_x2", + ) 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}", @@ -230,13 +244,13 @@ def check_support(self) -> bool: 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") + if self.generate_rht_rowwise: + self._check_tensor_shape(self.rht_rowwise_desc, (tensor_m, n_out, 1), "RHT rowwise") + if self.rht_quant: + self._check_tensor_shape(self.sfrht_rowwise_desc, _sf_layout_shape(tensor_m, n_out, self.sf_vec_size), "SFRHT rowwise") + if self.generate_rht_colwise: + self._check_tensor_shape(self.rht_colwise_desc, (tensor_m * n_out,), "RHT colwise") + self._check_tensor_shape(self.sfrht_colwise_desc, (tensor_m * n_out // self.sf_vec_size,), "SFRHT colwise") 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: @@ -244,8 +258,21 @@ def check_support(self) -> bool: 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") + if self.generate_rht_rowwise: + self._check_tensor_stride( + self.rht_rowwise_desc, + stride=[(n_d, 1, tensor_m * n_d)], + name="RHT rowwise", + extra_error_msg="RHT rowwise must have n-major layout", + ) + if self.generate_rht_colwise: + self._check_tensor_stride(self.rht_colwise_desc, stride=[(1,)], name="RHT colwise", extra_error_msg="RHT colwise must be flat") + self._check_tensor_stride( + self.sfrht_colwise_desc, + stride=[(1,)], + name="SFRHT colwise", + extra_error_msg="SFRHT colwise must be flat per-expert swizzled scale storage", + ) self.ab_dtype = self._check_dtype( self.a_desc, @@ -268,10 +295,23 @@ def check_support(self) -> bool: 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=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=self.sf_dtype, name="SFRHT", extra_error_msg="SFRHT must match SFA dtype") + if self.generate_rht_rowwise: + self._check_dtype(self.rht_rowwise_desc, dtype=[torch.bfloat16, torch.float4_e2m1fn_x2], name="RHT rowwise") + if self.rht_quant: + self._check_dtype( + self.sfrht_rowwise_desc, + dtype=torch.float8_e4m3fn, + name="SFRHT rowwise", + extra_error_msg="SFRHT rowwise must use NVFP4 scale storage dtype torch.float8_e4m3fn", + ) + if self.generate_rht_colwise: + self._check_dtype(self.rht_colwise_desc, dtype=torch.float4_e2m1fn_x2, name="RHT colwise") + self._check_dtype( + self.sfrht_colwise_desc, + dtype=torch.float8_e4m3fn, + name="SFRHT colwise", + extra_error_msg="SFRHT colwise must use NVFP4 scale storage dtype torch.float8_e4m3fn", + ) 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}") @@ -369,6 +409,7 @@ 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, + rht_per_expert=self.rht_per_expert if self.generate_rht else False, glu_alpha=self.glu_alpha, glu_limit=self.glu_limit, ) @@ -420,30 +461,34 @@ def compile(self) -> None: ) 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, - ) + if self.generate_rht_colwise: + rht_cute_arg = self._make_fake_cute_tensor( + dtype=self.rht_colwise_desc.dtype, + shape=(valid_m * self.n_out,), + stride=(1,), + ) + else: + rht_cute_arg = self._make_fake_cute_compact_tensor( + dtype=self.rht_rowwise_desc.dtype, + shape=(valid_m, self.rht_rowwise_desc.shape[1], 1), + stride_order=self.rht_rowwise_desc.stride_order, + dynamic_mode=self.rht_rowwise_desc.stride_order[0], + divisibility=8 if self._is_f16(self.rht_rowwise_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), + dtype=self.sfrht_rowwise_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfrht_rowwise_desc.shape[4], 1), + stride=(16, 4, self.sfrht_rowwise_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), + dtype=self.sfrht_colwise_desc.dtype, + shape=(valid_m * (self.n_out // self.sf_vec_size),), + stride=(1,), ) prob_cute_fake = self._make_fake_cute_tensor( dtype=self.prob_desc.dtype, @@ -622,8 +667,10 @@ def execute( 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, + rht_rowwise_tensor: Optional[torch.Tensor] = None, + sfrht_rowwise_tensor: Optional[torch.Tensor] = None, + rht_colwise_tensor: Optional[torch.Tensor] = None, + sfrht_colwise_tensor: Optional[torch.Tensor] = None, bias_tensor: Optional[torch.Tensor] = None, norm_const: float = 1.0, rht_norm_const: float = 1.0, @@ -640,10 +687,22 @@ def execute( 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") + rht_tensor = None + sfrht_tensor = None + if self.generate_rht_rowwise: + if rht_rowwise_tensor is None: + raise ValueError("rht_rowwise_tensor must be provided when rowwise RHT output is enabled") + if self.rht_quant and sfrht_rowwise_tensor is None: + raise ValueError("sfrht_rowwise_tensor must be provided when rowwise RHT output is NVFP4") + rht_tensor = rht_rowwise_tensor + sfrht_tensor = sfrht_rowwise_tensor + elif self.generate_rht_colwise: + if rht_colwise_tensor is None: + raise ValueError("rht_colwise_tensor must be provided when colwise RHT output is enabled") + if sfrht_colwise_tensor is None: + raise ValueError("sfrht_colwise_tensor must be provided when colwise RHT output is enabled") + rht_tensor = rht_colwise_tensor + sfrht_tensor = sfrht_colwise_tensor if self.weight_mode == MoEWeightMode.DENSE: if b_tensor is None or sfb_tensor is None: @@ -711,9 +770,8 @@ def grouped_gemm_glu_hadamard_quant_wrapper_sm100( 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, + rht_rowwise_dtype: Optional[torch.dtype] = None, + rht_colwise_dtype: Optional[torch.dtype] = None, glu_alpha: Optional[float] = None, glu_limit: Optional[float] = None, norm_const: float = 1.0, @@ -730,9 +788,10 @@ def grouped_gemm_glu_hadamard_quant_wrapper_sm100( ) -> 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/ue5m3 block scales - (``sfd_tensor``/``sfrht_tensor``); ``torch.bfloat16`` emits plain bf16. + Output modes are dtype driven. ``d_dtype`` or ``rht_rowwise_dtype`` of + ``torch.float4_e2m1fn_x2`` emit packed NVFP4 data plus block scales. + ``rht_colwise_dtype`` is NVFP4-only and returns flat per-expert ragged data + plus flat per-expert ragged scale storage. ``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 @@ -750,8 +809,8 @@ def grouped_gemm_glu_hadamard_quant_wrapper_sm100( c_dtype = torch.bfloat16 if d_dtype is None: d_dtype = torch.bfloat16 - if rht_dtype is None: - rht_dtype = torch.bfloat16 + if rht_rowwise_dtype is not None and rht_colwise_dtype is not None: + raise NotImplementedError("Request at most one of rht_rowwise_dtype or rht_colwise_dtype") 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: @@ -760,8 +819,14 @@ def grouped_gemm_glu_hadamard_quant_wrapper_sm100( 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") + if rht_rowwise_dtype == torch.uint8: + raise ValueError("rht_rowwise_dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + if rht_colwise_dtype == torch.uint8: + raise ValueError("rht_colwise_dtype torch.uint8 is not supported as packed FP4 for this fusion; use torch.float4_e2m1fn_x2") + if rht_rowwise_dtype is not None and rht_rowwise_dtype not in (torch.bfloat16, torch.float4_e2m1fn_x2): + raise ValueError(f"rht_rowwise_dtype must be torch.bfloat16, torch.float4_e2m1fn_x2, or None; got {rht_rowwise_dtype}") + if rht_colwise_dtype is not None and rht_colwise_dtype != torch.float4_e2m1fn_x2: + raise NotImplementedError("Only NVFP4 colwise RHT output is supported; use rht_colwise_dtype=torch.float4_e2m1fn_x2") valid_m = a_tensor.shape[0] is_dense = b_tensor is not None @@ -793,35 +858,45 @@ def grouped_gemm_glu_hadamard_quant_wrapper_sm100( 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 + rht_rowwise_quant = rht_rowwise_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: + def alloc_swizzled_sf(rows: int, cols: int, dtype: torch.dtype) -> 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=sfa_tensor.dtype, device=device).permute(3, 4, 1, 5, 2, 0) + return torch.empty(shape, dtype=dtype, 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) + sfd_tensor = alloc_swizzled_sf(valid_m, n_out, sfa_tensor.dtype) 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) + rht_rowwise_tensor = None + sfrht_rowwise_tensor = None + rht_colwise_tensor = None + sfrht_colwise_tensor = None + if rht_rowwise_dtype is not None: + rht_rowwise_tensor = alloc_n_major(valid_m, n_out // 2 if rht_rowwise_quant else n_out, rht_rowwise_dtype) + if rht_rowwise_quant: + sfrht_rowwise_tensor = alloc_swizzled_sf(valid_m, n_out, torch.float8_e4m3fn) + if rht_colwise_dtype is not None: + rht_colwise_tensor = torch.empty_strided((valid_m * n_out // 2,), (1,), dtype=rht_colwise_dtype, device=device) + sfrht_colwise_tensor = torch.empty_strided((valid_m * n_out // sf_vec_size,), (1,), dtype=torch.float8_e4m3fn, device=device) 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) + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + sfd_tensor=sfd_tensor, + rht_rowwise_tensor=rht_rowwise_tensor, + sfrht_rowwise_tensor=sfrht_rowwise_tensor, + rht_colwise_tensor=rht_colwise_tensor, + sfrht_colwise_tensor=sfrht_colwise_tensor, + ) def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: return tuple(i for i, _ in sorted(enumerate(tensor.stride()), key=lambda item: item[1])) @@ -852,9 +927,8 @@ def dynamic_m_tensor_signature( 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, + rht_rowwise_dtype, + rht_colwise_dtype, glu_alpha, glu_limit, stride_order(a_tensor), @@ -889,8 +963,10 @@ def dynamic_m_tensor_signature( sample_alpha=alpha_tensor, sample_prob=prob_tensor, sample_sfd=sfd_tensor, - sample_rht=rht_tensor, - sample_sfrht=sfrht_tensor, + sample_rht_rowwise=rht_rowwise_tensor, + sample_sfrht_rowwise=sfrht_rowwise_tensor, + sample_rht_colwise=rht_colwise_tensor, + sample_sfrht_colwise=sfrht_colwise_tensor, sample_bias=bias_tensor, acc_dtype=acc_dtype, mma_tiler_mn=mma_tiler_mn, @@ -901,7 +977,6 @@ def dynamic_m_tensor_signature( 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, ) @@ -931,8 +1006,10 @@ def dynamic_m_tensor_signature( alpha_tensor=alpha_tensor, prob_tensor=prob_tensor, sfd_tensor=sfd_tensor, - rht_tensor=rht_tensor, - sfrht_tensor=sfrht_tensor, + rht_rowwise_tensor=rht_rowwise_tensor, + sfrht_rowwise_tensor=sfrht_rowwise_tensor, + rht_colwise_tensor=rht_colwise_tensor, + sfrht_colwise_tensor=sfrht_colwise_tensor, bias_tensor=bias_tensor, norm_const=norm_const, rht_norm_const=rht_norm_const, @@ -950,11 +1027,21 @@ def dynamic_m_tensor_signature( alpha_tensor=alpha_tensor, prob_tensor=prob_tensor, sfd_tensor=sfd_tensor, - rht_tensor=rht_tensor, - sfrht_tensor=sfrht_tensor, + rht_rowwise_tensor=rht_rowwise_tensor, + sfrht_rowwise_tensor=sfrht_rowwise_tensor, + rht_colwise_tensor=rht_colwise_tensor, + sfrht_colwise_tensor=sfrht_colwise_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) + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + sfd_tensor=sfd_tensor, + rht_rowwise_tensor=rht_rowwise_tensor, + sfrht_rowwise_tensor=sfrht_rowwise_tensor, + rht_colwise_tensor=rht_colwise_tensor, + sfrht_colwise_tensor=sfrht_colwise_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 index eece2471a..43315ea9a 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 @@ -176,6 +176,7 @@ def __init__( act_func: str = "swiglu", enable_bias: bool = False, rht_rowwise: bool = False, + rht_per_expert: bool = False, glu_alpha: Optional[float] = None, glu_limit: Optional[float] = None, ): @@ -199,6 +200,9 @@ def __init__( # 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 + self.rht_per_expert = rht_per_expert + if rht_per_expert and rht_rowwise: + raise ValueError("rht_per_expert is a colwise RHT layout") # Always use pingpong epilogue for Hadamard self.epilogue_pingpong = True @@ -610,6 +614,7 @@ def __call__( self.c_dtype: Type[cutlass.Numeric] = c.element_type self.d_dtype: Type[cutlass.Numeric] = d.element_type self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + self.rht_sf_dtype: Type[cutlass.Numeric] = cutlass.Float8E4M3FN 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) @@ -653,8 +658,8 @@ def __call__( 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.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)") @@ -686,12 +691,12 @@ def __call__( 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. + # 8KB bf16). Colwise quantized staging is m-minor because packed bytes pair + # adjacent tokens; per-expert mode writes the same packed vectors directly + # to the flat output and does not TMA-copy this stage. self.rht_smem_layout_staged = sm100_utils.make_smem_layout_epi( self.rht_dtype, - utils.LayoutEnum.ROW_MAJOR, + utils.LayoutEnum.COL_MAJOR if cutlass.const_expr(self.rht_quant and not self.rht_rowwise) else utils.LayoutEnum.ROW_MAJOR, self.epi_tile, self.num_d_stage, ) @@ -811,9 +816,11 @@ def __call__( 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): + # TMA store RHT — identical tile to D unless colwise quantized RHT uses + # the flat per-expert direct-store layout. + if cutlass.const_expr(self.generate_rht and self.rht_per_expert): + tma_atom_rht, tma_tensor_rht = None, rht + elif 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(), @@ -918,7 +925,7 @@ class SharedStorage: # Same (128, 8) buffer either way. sSfRht: cute.struct.Align[ cute.struct.MemRange[ - self.sf_dtype, + self.rht_sf_dtype, self.threads_per_warp * len(self.epilog_rht_store_warp_id) * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), ], 16, @@ -1018,9 +1025,12 @@ 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): + def store_swizzled_sf_row(self, sf_tensor: cute.Tensor, logical_row, sf_col_base, sSf: cute.Tensor, tidx, sf_dtype=None): """Store one logical scale row into M32x4xrm_K4xrk_L SF layout.""" - sf_tensor = cute.recast_tensor(sf_tensor, self.sf_dtype) + if cutlass.const_expr(sf_dtype is None): + sf_tensor = cute.recast_tensor(sf_tensor, self.sf_dtype) + else: + sf_tensor = cute.recast_tensor(sf_tensor, sf_dtype) row_m0 = logical_row % 32 row_m1 = (logical_row // 32) % 4 row_m2 = logical_row // 128 @@ -2490,7 +2500,7 @@ def kernel( # # RHT output: per-expert RHT gmem tensor + TMA partition (mirrors ACT's D setup). # - if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): 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)) @@ -2499,9 +2509,9 @@ def kernel( 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) + # Expert token offset/count for the colwise (f, m) data and + # scale grids. Offsets are 256-aligned, so divisions by 16 are 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). @@ -2567,10 +2577,28 @@ 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, self.sf_dtype) + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, rht_norm_const, sSfRht, real_subtile_idx, self.rht_sf_dtype) + elif cutlass.const_expr(self.rht_quant and self.rht_per_expert): + hadamard_rmem_colwise_fwht_quant( + rht_ld, + d_buffer, + epi_tidx, + rht_norm_const, + None, + sSfRht, + real_subtile_idx * 2 * HADAMARD_SIZE, + self.rht_sf_dtype, + gRht=mRht_mnl, + seg16=(mD_mnl.shape[1] // HADAMARD_SIZE) * rht_t_off, + pitch16=rht_t_cnt // HADAMARD_SIZE, + feat0=mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk_d[1], + tok16=(sf_row - epi_tidx) // HADAMARD_SIZE, + row_blocks=cute.ceil_div(mD_mnl.shape[1], 128), + gSf=mSfRht_mnl, + ) 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 + rht_ld, d_buffer, epi_tidx, rht_norm_const, sRht, sSfRht, real_subtile_idx * 2 * HADAMARD_SIZE, self.rht_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) @@ -2586,7 +2614,7 @@ def kernel( 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): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): cute.copy( tma_atom_rht, bSG_sRht[(None, d_buffer)], @@ -2630,8 +2658,9 @@ def kernel( mma_tile_coord_mnl[1] * _num_sf, sSfRht, epi_tidx, + self.rht_sf_dtype, ) - if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): + if cutlass.const_expr(self.rht_quant and not self.rht_rowwise and not self.rht_per_expert): # (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 @@ -2641,6 +2670,7 @@ def kernel( (rht_t_off + sf_row - epi_tidx) // HADAMARD_SIZE, sSfRht, epi_tidx, + self.rht_sf_dtype, ) if cutlass.const_expr(self.d_quant): self.store_swizzled_sf_row( 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 index fcf66bf37..2cceeab9e 100644 --- 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 @@ -222,6 +222,7 @@ def __init__( act_func: str = "swiglu", enable_bias: bool = False, rht_rowwise: bool = False, + rht_per_expert: bool = False, sf_fp8_dtype_override: Optional[str] = None, glu_alpha: Optional[float] = None, glu_limit: Optional[float] = None, @@ -260,6 +261,9 @@ def __init__( # 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 + self.rht_per_expert = rht_per_expert + if rht_per_expert and rht_rowwise: + raise ValueError("rht_per_expert is a colwise RHT layout") # Always use pingpong epilogue for Hadamard self.epilogue_pingpong = True @@ -562,10 +566,12 @@ def _setup_attributes(self): 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): + def store_swizzled_sf_row(self, sf_tensor: cute.Tensor, logical_row, sf_col_base, sSf: cute.Tensor, tidx, sf_dtype=None): """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) + if cutlass.const_expr(sf_dtype is None): + sf_tensor = cute.recast_tensor(sf_tensor, self.sf_dtype) + else: + sf_tensor = cute.recast_tensor(sf_tensor, sf_dtype) row_m0 = logical_row % 32 row_m1 = (logical_row // 32) % 4 row_m2 = logical_row // 128 @@ -739,6 +745,9 @@ def __call__( self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override else: self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + self.rht_sf_dtype: Type[cutlass.Numeric] = ( + cutlass.FloatNV8E5M3FNU if cutlass.const_expr(self.sf_dtype == cutlass.FloatNV8E5M3FNU) else cutlass.Float8E4M3FN + ) 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) @@ -783,8 +792,8 @@ def __call__( 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.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)") @@ -816,12 +825,12 @@ def __call__( 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. + # 8KB bf16). Colwise quantized staging is m-minor because packed bytes pair + # adjacent tokens; per-expert mode writes the same packed vectors directly + # to the flat output and does not TMA-copy this stage. self.rht_smem_layout_staged = sm100_utils.make_smem_layout_epi( self.rht_dtype, - utils.LayoutEnum.ROW_MAJOR, + utils.LayoutEnum.COL_MAJOR if cutlass.const_expr(self.rht_quant and not self.rht_rowwise) else utils.LayoutEnum.ROW_MAJOR, self.epi_tile, self.num_d_stage, ) @@ -986,9 +995,11 @@ def __call__( 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): + # TMA store RHT — identical tile to D unless colwise quantized RHT uses + # the flat per-expert direct-store layout. + if cutlass.const_expr(self.generate_rht and self.rht_per_expert): + tma_atom_rht, tma_tensor_rht = None, rht + elif 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(), @@ -1093,7 +1104,7 @@ class SharedStorage: # Same (128, 8) buffer either way. sSfRht: cute.struct.Align[ cute.struct.MemRange[ - self.sf_dtype, + self.rht_sf_dtype, self.threads_per_warp * len(self.epilog_rht_store_warp_id) * (self.cta_tile_shape_mnk_d[1] // HADAMARD_SIZE), ], 16, @@ -2798,7 +2809,7 @@ def _grp(p): # # RHT output: per-expert RHT gmem tensor + TMA partition (mirrors ACT's D setup). # - if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): 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)) @@ -2825,9 +2836,9 @@ def _grp(p): 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) + # Expert token offset/count for the colwise (f, m) data and + # scale grids. Offsets are 256-aligned, so divisions by 16 are 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). @@ -2879,7 +2890,7 @@ def _grp(p): # (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): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): _bSG_gRht_h = bSG_gRht if cutlass.const_expr(self.d_quant): _bSG_gDq_h = bSG_gDq @@ -2888,14 +2899,14 @@ def _grp(p): for _m_half in _breuse_halves: if self.enable_breuse: - if cutlass.const_expr(self.generate_rht): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): _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): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): _bSG_gRht_h = bSG_gRht if cutlass.const_expr(self.d_quant): _bSG_gDq_h = bSG_gDq @@ -2947,10 +2958,28 @@ def _grp(p): # 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) + hadamard_rmem_rowwise_fwht(rht_ld, d_buffer, epi_tidx, sRht, rht_norm_const, sSfRht, real_subtile_idx, self.rht_sf_dtype) + elif cutlass.const_expr(self.rht_quant and self.rht_per_expert): + hadamard_rmem_colwise_fwht_quant( + rht_ld, + d_buffer, + epi_tidx, + rht_norm_const, + None, + sSfRht, + real_subtile_idx * 2 * HADAMARD_SIZE, + self.rht_sf_dtype, + gRht=mRht_mnl, + seg16=(mD_mnl.shape[1] // HADAMARD_SIZE) * rht_t_off, + pitch16=rht_t_cnt // HADAMARD_SIZE, + feat0=mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk_d[1], + tok16=(_sf_row_h - epi_tidx) // HADAMARD_SIZE, + row_blocks=cute.ceil_div(mD_mnl.shape[1], 128), + gSf=mSfRht_mnl, + ) 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 + rht_ld, d_buffer, epi_tidx, rht_norm_const, sRht, sSfRht, real_subtile_idx * 2 * HADAMARD_SIZE, self.rht_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) @@ -2966,7 +2995,7 @@ def _grp(p): 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): + if cutlass.const_expr(self.generate_rht and not self.rht_per_expert): cute.copy( tma_atom_rht, bSG_sRht[(None, d_buffer)], @@ -3005,8 +3034,9 @@ def _grp(p): mma_tile_coord_mnl[1] * _num_sf, sSfRht, epi_tidx, + self.rht_sf_dtype, ) - if cutlass.const_expr(self.rht_quant and not self.rht_rowwise): + if cutlass.const_expr(self.rht_quant and not self.rht_rowwise and not self.rht_per_expert): # (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 @@ -3016,6 +3046,7 @@ def _grp(p): (rht_t_off + _sf_row_h - epi_tidx) // HADAMARD_SIZE, sSfRht, epi_tidx, + self.rht_sf_dtype, ) if cutlass.const_expr(self.d_quant): self.store_swizzled_sf_row( 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 3a72b4dae..25d1f42a7 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,19 +100,27 @@ 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, sf_dtype): +def hadamard_rmem_colwise_fwht_quant( + rmem_bf16, + d_buffer, + tidx, + norm_const, + sRht, + sSf, + sf_row_base, + sf_dtype, + gRht=None, + seg16=None, + pitch16=None, + feat0=None, + tok16=None, + row_blocks=None, + gSf=None, +): """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 - + stored transposed: packed nibbles pair adjacent tokens of one feature. The + default target stages into sRht for a transposed TMA store; with gRht set, + stores go directly into a flat per-expert ragged output buffer.""" n_vals = 2 * HADAMARD_SIZE tCompute = cute.make_rmem_tensor((n_vals,), cutlass.Float32) for i in cutlass.range_constexpr(n_vals): @@ -124,6 +132,53 @@ def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht for i in cutlass.range_constexpr(n_vals): tCompute[i] = (tCompute[i] * cutlass.Float32(0.25)).to(cutlass.BFloat16).to(cutlass.Float32) + _nvfp4_quant_colwise_transposed( + tCompute, + d_buffer, + tidx, + norm_const, + sRht, + sSf, + sf_row_base, + sf_dtype, + gRht, + seg16, + pitch16, + feat0, + tok16, + row_blocks, + gSf, + ) + + +@cute.jit +def _nvfp4_quant_colwise_transposed( + tCompute, + d_buffer, + tidx, + norm_const, + sRht, + sSf, + sf_row_base, + sf_dtype, + gRht=None, + seg16=None, + pitch16=None, + feat0=None, + tok16=None, + row_blocks=None, + gSf=None, +): + """Colwise NVFP4 quantization of two 16-token feature columns. + + Data bytes are stored in transposed order, so each feature's 16 transformed + token values are one contiguous 8-byte vector. ``gRht``/``gSf`` enable direct + flat per-expert stores for data and per-expert swizzled scale segments. + """ + token_block = tidx // HADAMARD_SIZE + feat_pair = tidx % HADAMARD_SIZE + feature = 2 * feat_pair + # 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 @@ -149,8 +204,24 @@ def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht 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)) + if cutlass.const_expr(gSf is not None): + gSf = cute.recast_tensor(gSf, sf_dtype) for c in cutlass.range_constexpr(2): - sSf[(sf_row_base + feature + c, token_block)] = tCrSFC_f8x4[c] + if cutlass.const_expr(gSf is not None): + row = feat0 + sf_row_base + feature + c + col = tok16 + token_block + col_blocks = pitch16 // 4 + sf_idx = ( + seg16 + + (row % 32) * 4 * row_blocks * 4 * col_blocks + + ((row // 32) % 4) * row_blocks * 4 * col_blocks + + (row // 128) * 4 * col_blocks + + (col % 4) * col_blocks + + (col // 4) + ) + gSf[sf_idx] = tCrSFC_f8x4[c] + else: + 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) @@ -161,15 +232,24 @@ def hadamard_rmem_colwise_fwht_quant(rmem_bf16, d_buffer, tidx, norm_const, sRht (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))), - ) + tCol = cute.make_rmem_tensor((HADAMARD_SIZE,), cutlass.Float32) + tCol_fp4 = cute.make_rmem_tensor((HADAMARD_SIZE,), cutlass.Float4E2M1FN) + if cutlass.const_expr(gRht is not None): + gRht_cols = cute.zipped_divide(gRht, (HADAMARD_SIZE,)) + else: + sRht_cols = cute.zipped_divide(cute.slice_(sRht, (None, None, d_buffer)), (HADAMARD_SIZE, 1)) + for c in cutlass.range_constexpr(2): + for i in cutlass.range_constexpr(HADAMARD_SIZE): + tCol[i] = tCompute[2 * i + c] + tCol_fp4.store(tCol.load().to(cutlass.Float4E2M1FN)) + if cutlass.const_expr(gRht is not None): + col = seg16 + (feat0 + sf_row_base + feature + c) * pitch16 + tok16 + token_block + cute.autovec_copy(tCol_fp4, cute.slice_(gRht_cols, ((None,), col))) + else: + cute.autovec_copy( + tCol_fp4, + cute.slice_(sRht_cols, ((None, None), (token_block, feature + c))), + ) @cute.jit 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 26bf47a70..d64bc238c 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 @@ -13,6 +13,7 @@ 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, + ceil_div, reencode_sf_tensor_as_ue5m3, ue5m3_bytes_to_fp32, ) @@ -171,12 +172,29 @@ def _swizzled_sf_to_flat(sf_tensor: torch.Tensor, rows: int, cols: int) -> torch ] +def _blocked_sf_to_flat(sf_tensor: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Gather wgrad-style blocked scale backing 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) + col_blocks = sf_tensor.shape[1] // 4 + linear = ( + (row_idx // 128) * col_blocks * 512 + + (col_idx // 4) * 512 + + (row_idx % 32) * 16 + + ((row_idx // 32) % 4) * 4 + + (col_idx % 4) + ) + return sf_tensor.flatten()[linear] + + def _check_nvfp4_output( values: torch.Tensor, sf_tensor: torch.Tensor, ref: torch.Tensor, norm_const: float, name: str, + sf_layout: Optional[str] = None, ) -> None: """Check unpacked e2m1 values (rows, cols) + e4m3 scales (rows, cols/16), with (1, 16) quantization blocks along the last dim, against the f32 @@ -185,7 +203,12 @@ def _check_nvfp4_output( 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 + if sf_tensor.dim() == 6: + sf_flat = _swizzled_sf_to_flat(sf_tensor, ref.shape[0], ref.shape[1]) + elif sf_layout == "blocked": + sf_flat = _blocked_sf_to_flat(sf_tensor, ref.shape[0], ref.shape[1]) + else: + sf_flat = sf_tensor torch.testing.assert_close( sf_flat.float().cpu(), sf_ref.float().cpu(), @@ -202,6 +225,43 @@ def _check_nvfp4_output( assert not bad.any(), f"{name}: {int(bad.sum())} dequantized elements exceed the quantization error bound (max err {err[bad].max().item():.4f})" +def _check_colwise_rht( + inputs: Dict, + outputs: Dict, + rht_ref: torch.Tensor, + rht_norm_const: float, + sf_fp8_dtype_override: Optional[str], +) -> None: + assert outputs["rht_rowwise_tensor"] is None + assert outputs["sfrht_rowwise_tensor"] is None + assert outputs["rht_colwise_tensor"] is not None + assert outputs["sfrht_colwise_tensor"] is not None + + valid_m, n_out = rht_ref.shape + values = float4_e2m1fn_x2_to_float32(outputs["rht_colwise_tensor"].view(torch.uint8).cpu()).reshape(-1) + sf = outputs["sfrht_colwise_tensor"].cpu().reshape(-1) + if sf_fp8_dtype_override == "e5m3": + sf = ue5m3_bytes_to_fp32(sf) + + start_m = 0 + sf_col_offset = 0 + for expert_idx, group_m in enumerate(inputs["aligned_group_m_list"]): + end_m = start_m + group_m + ref_e = rht_ref[start_m:end_m].float().t().contiguous() + value_offset = n_out * start_m + value_count = n_out * group_m + values_e = values[value_offset : value_offset + value_count].reshape(n_out, group_m) + sf_offset = n_out * start_m // HADAMARD_SIZE + sf_count = n_out * group_m // HADAMARD_SIZE + sf_e = sf[sf_offset : sf_offset + sf_count] + sf_e = sf_e.reshape(32, 4, ceil_div(n_out, 128), 4, ceil_div(group_m, 64), 1) + sf_flat_e = _swizzled_sf_to_flat(sf_e, n_out, group_m) + _check_nvfp4_output(values_e, sf_e, ref_e, rht_norm_const, f"RHT colwise expert {expert_idx}") + start_m = end_m + + assert start_m == valid_m + + # ============================================================================= # Output checking # ============================================================================= @@ -213,8 +273,8 @@ def _check_outputs( cfg: Dict, *, act_func: str, - rht_output: bool, - rht_rowwise: bool, + rht_rowwise_dtype: Optional[torch.dtype], + rht_colwise_dtype: Optional[torch.dtype], glu_alpha: Optional[float] = None, glu_limit: Optional[float] = None, norm_const: float = 1.0, @@ -254,44 +314,43 @@ def _check_outputs( "D", ) - if not rht_output: - assert outputs["rht_tensor"] is None - assert outputs["sfrht_tensor"] is None + if rht_rowwise_dtype is None and rht_colwise_dtype is None: + assert outputs["rht_rowwise_tensor"] is None + assert outputs["sfrht_rowwise_tensor"] is None + assert outputs["rht_colwise_tensor"] is None + assert outputs["sfrht_colwise_tensor"] is None + return + + if rht_rowwise_dtype is not None: + assert outputs["rht_colwise_tensor"] is None + assert outputs["sfrht_colwise_tensor"] is None + rht_ref = _rht_ref(d_ref.cpu(), rowwise=True) + + elif rht_colwise_dtype is not None: + assert rht_colwise_dtype == torch.float4_e2m1fn_x2 + rht_ref = _rht_ref(d_ref.cpu(), rowwise=False) + _check_colwise_rht(inputs, outputs, rht_ref, rht_norm_const, sf_fp8_dtype_override) return + else: + raise AssertionError("unreachable RHT mode") - rht_ref = _rht_ref(d_ref.cpu(), rht_rowwise) - if outputs["sfrht_tensor"] is None: + if outputs["sfrht_rowwise_tensor"] is None: torch.testing.assert_close( - outputs["rht_tensor"][:valid_m, :, 0].cpu().float(), + outputs["rht_rowwise_tensor"][:valid_m, :, 0].cpu().float(), rht_ref.float(), atol=1e-1, 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()), - 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. - sfrht_tensor = outputs["sfrht_tensor"] + sfrht_tensor = outputs["sfrht_rowwise_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(), + float4_e2m1fn_x2_to_float32(outputs["rht_rowwise_tensor"][:valid_m, :, 0].view(torch.uint8).cpu()), sfrht_tensor.cpu(), - rht_ref.float().t(), + rht_ref.float(), rht_norm_const, - "RHT", + "RHT rowwise", ) @@ -309,9 +368,8 @@ def _run_wrapper( act_func="swiglu", enable_bias=False, d_dtype=torch.bfloat16, - rht_output=True, - rht_dtype=torch.bfloat16, - rht_rowwise=False, + rht_rowwise_dtype=None, + rht_colwise_dtype=torch.float4_e2m1fn_x2, glu_alpha=None, glu_limit=None, sf_fp8_dtype_override=None, @@ -338,8 +396,11 @@ def _run_wrapper( 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) + if rht_rowwise_dtype == torch.float4_e2m1fn_x2: + rht_ref = _rht_ref(ref_tensors["d_ref"].cpu(), rowwise=True) + rht_norm_const = 2688.0 / rht_ref.float().abs().max().item() + elif rht_colwise_dtype == torch.float4_e2m1fn_x2: + rht_ref = _rht_ref(ref_tensors["d_ref"].cpu(), rowwise=False) rht_norm_const = 2688.0 / rht_ref.float().abs().max().item() if sf_fp8_dtype_override == "e5m3": @@ -363,9 +424,8 @@ def _run_wrapper( 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, + rht_rowwise_dtype=rht_rowwise_dtype, + rht_colwise_dtype=rht_colwise_dtype, glu_alpha=glu_alpha, glu_limit=glu_limit, norm_const=norm_const, @@ -384,8 +444,8 @@ def _run_wrapper( outputs, cfg, act_func=act_func, - rht_output=rht_output, - rht_rowwise=rht_rowwise, + rht_rowwise_dtype=rht_rowwise_dtype, + rht_colwise_dtype=rht_colwise_dtype, glu_alpha=glu_alpha, glu_limit=glu_limit, norm_const=norm_const, @@ -428,8 +488,10 @@ def alloc_n_major(rows, cols, dtype): "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, + "rht_rowwise_tensor": alloc_n_major(valid_m, n_out, torch.bfloat16), + "sfrht_rowwise_tensor": None, + "rht_colwise_tensor": None, + "sfrht_colwise_tensor": None, } from cudnn import GroupedGemmGluHadamardQuantSm100 @@ -444,7 +506,7 @@ def alloc_n_major(rows, cols, dtype): sample_padded_offsets=inputs["padded_offsets_tensor"], sample_alpha=inputs["alpha_tensor"], sample_prob=inputs["prob_tensor"], - sample_rht=outputs["rht_tensor"], + sample_rht_rowwise=outputs["rht_rowwise_tensor"], acc_dtype=cfg["acc_dtype"], mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], @@ -466,7 +528,7 @@ def alloc_n_major(rows, cols, dtype): padded_offsets=inputs["padded_offsets_tensor"], alpha_tensor=inputs["alpha_tensor"], prob_tensor=inputs["prob_tensor"], - rht_tensor=outputs["rht_tensor"], + rht_rowwise_tensor=outputs["rht_rowwise_tensor"], ) _check_outputs( @@ -474,8 +536,8 @@ def alloc_n_major(rows, cols, dtype): outputs, cfg, act_func=act_func, - rht_output=True, - rht_rowwise=False, + rht_rowwise_dtype=torch.bfloat16, + rht_colwise_dtype=None, sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -514,6 +576,8 @@ def _run_discrete_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func= c_dtype=cfg["c_dtype"], d_dtype=torch.bfloat16, cd_major=cfg["cd_major"], + rht_rowwise_dtype=torch.bfloat16, + rht_colwise_dtype=None, mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], @@ -528,8 +592,8 @@ def _run_discrete_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func= outputs, cfg, act_func=act_func, - rht_output=True, - rht_rowwise=False, + rht_rowwise_dtype=torch.bfloat16, + rht_colwise_dtype=None, sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -572,7 +636,8 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_rowwise(request): ab_dtype=torch.float4_e2m1fn_x2, sf_dtype=torch.float8_e4m3fn, sf_vec_size=16, - rht_rowwise=True, + rht_rowwise_dtype=torch.bfloat16, + rht_colwise_dtype=None, ) @@ -584,7 +649,7 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_no_rht(request): ab_dtype=torch.float4_e2m1fn_x2, sf_dtype=torch.float8_e4m3fn, sf_vec_size=16, - rht_output=False, + rht_colwise_dtype=None, ) @@ -622,21 +687,25 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_d(request): sf_dtype=torch.float8_e4m3fn, sf_vec_size=16, d_dtype=torch.float4_e2m1fn_x2, - rht_output=False, + rht_colwise_dtype=None, ) @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): +@pytest.mark.parametrize( + "rht_rowwise_dtype,rht_colwise_dtype", + [(None, torch.float4_e2m1fn_x2), (torch.float4_e2m1fn_x2, None)], + ids=["colwise", "rowwise"], +) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_rht(request, rht_rowwise_dtype, rht_colwise_dtype): _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, + rht_rowwise_dtype=rht_rowwise_dtype, + rht_colwise_dtype=rht_colwise_dtype, ) @@ -649,14 +718,18 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_full(request): sf_dtype=torch.float8_e4m3fn, sf_vec_size=16, d_dtype=torch.float4_e2m1fn_x2, - rht_dtype=torch.float4_e2m1fn_x2, + rht_colwise_dtype=torch.float4_e2m1fn_x2, ) @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): +@pytest.mark.parametrize( + "rht_rowwise_dtype,rht_colwise_dtype", + [(None, torch.float4_e2m1fn_x2), (torch.float4_e2m1fn_x2, None)], + ids=["colwise", "rowwise"], +) +def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_rht_e5m3(request, rht_rowwise_dtype, rht_colwise_dtype): """quant_rht with the input block scales carried as UE5M3 bytes in e4m3 storage.""" _skip_unless_e5m3_supported() _run_wrapper( @@ -664,8 +737,8 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_rht_e5m3(request, rht_row 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, + rht_rowwise_dtype=rht_rowwise_dtype, + rht_colwise_dtype=rht_colwise_dtype, sf_fp8_dtype_override="e5m3", ) @@ -681,7 +754,7 @@ def test_grouped_gemm_glu_hadamard_quant_wrapper_quant_full_e5m3(request): sf_dtype=torch.float8_e4m3fn, sf_vec_size=16, d_dtype=torch.float4_e2m1fn_x2, - rht_dtype=torch.float4_e2m1fn_x2, + rht_colwise_dtype=torch.float4_e2m1fn_x2, sf_fp8_dtype_override="e5m3", ) @@ -743,9 +816,6 @@ def run(sf_fp8_dtype_override): 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"], From c25942a93e2a4fdb3257a2a6f926b3685460cf8d Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 25 Aug 2026 06:50:53 +0000 Subject: [PATCH 2/2] Debug tests Loosen RHT tols since it can accumulate error from multiple BF16 casts. Skip relative error check when scales have been driven to zero (scales are verified separately from FP4 values). Signed-off-by: Tim Moon --- .../test_grouped_gemm_glu_hadamard_quant.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) 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 d64bc238c..98a998ae6 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 @@ -195,6 +195,8 @@ def _check_nvfp4_output( norm_const: float, name: str, sf_layout: Optional[str] = None, + rtol: float = 0.25, + atol: float = 5e-2, ) -> None: """Check unpacked e2m1 values (rows, cols) + e4m3 scales (rows, cols/16), with (1, 16) quantization blocks along the last dim, against the f32 @@ -220,9 +222,15 @@ def _check_nvfp4_output( 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})" + bound = rtol * (6 * decode_scale) + atol + bad = torch.logical_and(err > bound, decode_scale != 0) + if bad.any(): + max_abs_err = err[bad].max().item() + max_rel_err = torch.nanquantile(err[bad] / ref_bf16[bad].abs(), 1.0).item() + raise RuntimeError( + f"{name}: {int(bad.sum())} dequantized elements exceed the quantization error bound " + f"(max abs err {max_abs_err:.4f}, max rel err {max_rel_err:.4f})" + ) def _check_colwise_rht( @@ -256,7 +264,17 @@ def _check_colwise_rht( sf_e = sf[sf_offset : sf_offset + sf_count] sf_e = sf_e.reshape(32, 4, ceil_div(n_out, 128), 4, ceil_div(group_m, 64), 1) sf_flat_e = _swizzled_sf_to_flat(sf_e, n_out, group_m) - _check_nvfp4_output(values_e, sf_e, ref_e, rht_norm_const, f"RHT colwise expert {expert_idx}") + _check_nvfp4_output( + values_e, + sf_e, + ref_e, + rht_norm_const, + f"RHT colwise expert {expert_idx}", + # RHT is performed on BF16-quantized values, so loosen + # tols to handle extra accumulated error. + rtol=0.5, + atol=0.1, + ) start_m = end_m assert start_m == valid_m