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..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 @@ -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,31 @@ 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, + 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 @@ -185,7 +205,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(), @@ -197,9 +222,62 @@ 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( + 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}", + # 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 # ============================================================================= @@ -213,8 +291,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 +332,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 - rht_ref = _rht_ref(d_ref.cpu(), rht_rowwise) - if outputs["sfrht_tensor"] is None: + 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") + + 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 +386,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 +414,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 +442,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 +462,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 +506,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 +524,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 +546,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 +554,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 +594,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 +610,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 +654,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 +667,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 +705,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 +736,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 +755,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 +772,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 +834,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"],