Add row-scale support to grouped GEMM quant

python/cudnn/grouped_gemm/grouped_gemm_quant/api.py

@@ -90,6 +90,7 @@ def __init__(
         sample_amax: Optional[torch.Tensor] = None,
         sample_norm_const: Optional[torch.Tensor] = None,
         sample_prob: Optional[torch.Tensor] = None,
+        sample_row_scale: Optional[torch.Tensor] = None,
         # Configuration
         acc_dtype: torch.dtype = torch.float32,
         mma_tiler_mn: Tuple[int, int] = (256, 256),
@@ -121,6 +122,9 @@ def __init__(
         :param sample_amax: Optional amax tensor for quantization
         :param sample_norm_const: Optional normalization constant
         :param sample_prob: Optional probability tensor for gating
+        :param sample_row_scale: Optional 1-D FP32 row-scale tensor. When
+            provided, the epilogue scales GEMM accumulators by
+            ``alpha[expert] * row_scale[m]`` before output conversion.
         :param acc_dtype: Accumulator data type
         :param mma_tiler_mn: MMA tiler shape (M, N)
         :param cluster_shape_mn: Cluster shape (M, N)
@@ -174,6 +178,11 @@ def __init__(
             "norm_const",
         )
         self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob")
+        self.row_scale_desc = self._unpad_tensor_to_ndim(
+            self._make_tensor_desc(sample_row_scale, name="sample_row_scale"),
+            1,
+            "row_scale",
+        )
         self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias")
 
         if self.weight_mode == MoEWeightMode.DENSE:
@@ -275,6 +284,7 @@ def check_support(self) -> bool:
             "Pass a tensor of ones with shape (valid_m, 1, 1) if no gating is needed.",
         )
         self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob")
+        self._check_tensor_shape(self.row_scale_desc, (tensor_m,), "row_scale")
         self._check_tensor_shape(self.bias_desc, (n, l), "bias")
         self._check_tensor_shape(self.amax_desc, (self.expert_cnt, 1), "amax")
         self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const")
@@ -312,6 +322,11 @@ def check_support(self) -> bool:
             self.bias_desc,
             stride=[(1, n)],
         )
+        _ = self._check_tensor_stride(
+            self.row_scale_desc,
+            stride=[(1,)],
+            extra_error_msg="row_scale must be a contiguous 1-D tensor",
+        )
 
         self._logger.debug("Checking data types")
         self.ab_dtype = self._check_dtype(
@@ -418,6 +433,12 @@ def check_support(self) -> bool:
             name="D_col",
             extra_error_msg="D_col must have the same dtype as D",
         )
+        self._check_dtype(
+            self.row_scale_desc,
+            dtype=torch.float32,
+            name="row_scale",
+            extra_error_msg="row_scale must be float32",
+        )
 
         self._not_implemented_error_if(
             self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32,
@@ -624,6 +645,13 @@ def _compile_dense(self, gemm_quant, max_active_clusters, fake_stream) -> None:
                     shape=(valid_m, *self.prob_desc.shape[1:]),
                     stride=self.prob_desc.stride,
                 )
+            row_scale_cute_fake = None
+            if self.row_scale_desc is not None:
+                row_scale_cute_fake = self._make_fake_cute_tensor(
+                    dtype=self.row_scale_desc.dtype,
+                    shape=(valid_m,),
+                    stride=self.row_scale_desc.stride,
+                )
 
             sfd_row_fake = None
             sfd_col_fake = None
@@ -712,6 +740,13 @@ def _compile_dense(self, gemm_quant, max_active_clusters, fake_stream) -> None:
                     shape=(valid_m, *self.prob_desc.shape[1:]),
                     stride=self.prob_desc.stride,
                 )
+            row_scale_cute_fake = None
+            if self.row_scale_desc is not None:
+                row_scale_cute_fake = self._make_fake_cute_tensor(
+                    dtype=self.row_scale_desc.dtype,
+                    shape=(valid_m,),
+                    stride=self.row_scale_desc.stride,
+                )
 
             sfd_row_fake = None
             sfd_col_fake = None
@@ -762,6 +797,7 @@ def _compile_dense(self, gemm_quant, max_active_clusters, fake_stream) -> None:
             norm_const_tensor=self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16),
             padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16),
             alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16),
+            row_scale=row_scale_cute_fake,
             bias=bias_cute_fake,
             prob=prob_cute_fake,
             max_active_clusters=max_active_clusters,
@@ -784,6 +820,7 @@ def tensor_api(
             norm_const_tensor: Optional[torch.Tensor],
             padded_offsets: torch.Tensor,
             alpha_tensor: torch.Tensor,
+            row_scale_tensor: Optional[torch.Tensor],
             prob_tensor: Optional[torch.Tensor],
             bias_tensor: Optional[torch.Tensor],
             stream: cuda.CUstream,
@@ -806,6 +843,7 @@ def tensor_api(
                 norm_const_tensor,
                 padded_offsets,
                 alpha_tensor,
+                row_scale_tensor,
                 bias_tensor,
                 prob_tensor,
                 stream,
@@ -886,6 +924,7 @@ def _compile_discrete(self, gemm_quant, max_active_clusters, fake_stream) -> Non
             stride=self.prob_desc.stride,
             assumed_align=16,
         )
+        row_scale_tensor = self._make_fake_cute_tensor_from_desc(self.row_scale_desc, assumed_align=16)
         bias_cute_fake = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16)
 
         b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda")
@@ -914,6 +953,7 @@ def _compile_discrete(self, gemm_quant, max_active_clusters, fake_stream) -> Non
             norm_const_tensor=norm_const_tensor_cute,
             padded_offsets=padded_offsets_tensor,
             alpha=alpha_tensor,
+            row_scale=row_scale_tensor,
             bias=bias_cute_fake,
             prob=prob_tensor,
             max_active_clusters=max_active_clusters,
@@ -940,6 +980,7 @@ def tensor_api(
             norm_const_tensor: Optional[torch.Tensor],
             padded_offsets: torch.Tensor,
             alpha_tensor: torch.Tensor,
+            row_scale_tensor: Optional[torch.Tensor],
             prob_tensor: Optional[torch.Tensor],
             bias_tensor: Optional[torch.Tensor],
             stream: cuda.CUstream,
@@ -964,6 +1005,7 @@ def tensor_api(
                 norm_const_tensor,
                 padded_offsets,
                 alpha_tensor,
+                row_scale_tensor,
                 bias_tensor,
                 prob_tensor,
                 stream,
@@ -991,6 +1033,7 @@ def execute(
         amax_tensor: Optional[torch.Tensor] = None,
         norm_const_tensor: Optional[torch.Tensor] = None,
         prob_tensor: Optional[torch.Tensor] = None,
+        row_scale_tensor: Optional[torch.Tensor] = None,
         current_stream: Optional[cuda.CUstream] = None,
     ) -> None:
         """Execute the compiled kernel.
@@ -1013,6 +1056,10 @@ def execute(
         :param amax_tensor: Optional amax tensor
         :param norm_const_tensor: Optional normalization constant
         :param prob_tensor: Probability tensor for per-row gating. Required.
+        :param row_scale_tensor: Optional contiguous FP32 tensor of shape ``(valid_m,)``.
+            When provided, the epilogue multiplies accumulators by
+            ``alpha_tensor[expert] * row_scale_tensor[m]`` before output
+            conversion.
         :param current_stream: CUDA stream
         """
         self._logger.debug("Entering execute")
@@ -1047,6 +1094,16 @@ def execute(
                 bias_tensor is not None,
                 "bias_tensor must be omitted at execute() when the API was compiled without sample_bias",
             )
+        if self.row_scale_desc is None:
+            self._value_error_if(
+                row_scale_tensor is not None,
+                "row_scale_tensor must be omitted at execute() when the API was compiled without sample_row_scale",
+            )
+        else:
+            self._value_error_if(
+                row_scale_tensor is None,
+                "row_scale_tensor must be provided at execute() when the API was compiled with sample_row_scale",
+            )
 
         self._logger.debug("Executing grouped_gemm_quant kernel")
         if self.weight_mode == MoEWeightMode.DENSE:
@@ -1063,6 +1120,7 @@ def execute(
                 norm_const_tensor=norm_const_tensor,
                 padded_offsets=padded_offsets,
                 alpha_tensor=alpha_tensor,
+                row_scale_tensor=row_scale_tensor,
                 prob_tensor=prob_tensor,
                 bias_tensor=bias_tensor,
                 stream=current_stream,
@@ -1081,6 +1139,7 @@ def execute(
                 norm_const_tensor=norm_const_tensor,
                 padded_offsets=padded_offsets,
                 alpha_tensor=alpha_tensor,
+                row_scale_tensor=row_scale_tensor,
                 prob_tensor=prob_tensor,
                 bias_tensor=bias_tensor,
                 stream=current_stream,
@@ -1110,6 +1169,7 @@ def grouped_gemm_quant_wrapper_sm100(
     b_major: str = "k",
     norm_const_tensor: Optional[torch.Tensor] = None,
     prob_tensor: Optional[torch.Tensor] = None,
+    row_scale_tensor: Optional[torch.Tensor] = None,
     acc_dtype: torch.dtype = torch.float32,
     d_dtype: torch.dtype = torch.bfloat16,
     cd_major: str = "n",
@@ -1146,6 +1206,10 @@ def grouped_gemm_quant_wrapper_sm100(
             Should be None for FP4/BF16 input configurations.
         prob_tensor: Probability tensor for per-row gating (shape `(valid_m, 1, 1)`).
             This argument is required. Pass a tensor of ones when no gating is needed.
+        row_scale_tensor: Optional FP32 tensor of shape `(valid_m,)`.
+            When provided, the epilogue multiplies accumulators by
+            `alpha_tensor[expert] * row_scale_tensor[m]` before output
+            conversion.
         acc_dtype: Accumulator data type
         d_dtype: Output D tensor data type
         cd_major: CD major dimension (only "n"-major layout is supported)
@@ -1284,6 +1348,13 @@ def grouped_gemm_quant_wrapper_sm100(
             "prob_tensor is required: the kernel unconditionally multiplies output by per-row gating probability. "
             "Pass a tensor of ones with shape (valid_m, 1, 1) if no gating is needed."
         )
+    if row_scale_tensor is not None:
+        if row_scale_tensor.dtype != torch.float32:
+            raise ValueError(f"row_scale_tensor must be float32, got {row_scale_tensor.dtype}")
+        if tuple(row_scale_tensor.shape) != (valid_m,):
+            raise ValueError(f"row_scale_tensor must have shape {(valid_m,)}, got {tuple(row_scale_tensor.shape)}")
+        if tuple(row_scale_tensor.stride()) != (1,):
+            raise ValueError(f"row_scale_tensor must be contiguous with stride (1,), got {tuple(row_scale_tensor.stride())}")
 
     if valid_m == 0:
         _logger.debug("grouped_gemm_quant_wrapper_sm100: valid_m is zero, skipping kernel execution")
@@ -1340,6 +1411,7 @@ def dynamic_m_tensor_signature(
             *tensor_signature(alpha_tensor),
             *tensor_signature(norm_const_tensor),
             *dynamic_m_tensor_signature(prob_tensor, (1, 1)),
+            *dynamic_m_tensor_signature(row_scale_tensor, ()),
             tuple(padded_offsets.shape),
             tuple(padded_offsets.stride()),
             padded_offsets.dtype,
@@ -1369,6 +1441,7 @@ def dynamic_m_tensor_signature(
             *tensor_signature(alpha_tensor),
             *tensor_signature(norm_const_tensor),
             *dynamic_m_tensor_signature(prob_tensor, (1, 1)),
+            *dynamic_m_tensor_signature(row_scale_tensor, ()),
             tuple(b_ptrs.shape),
             tuple(b_ptrs.stride()),
             b_ptrs.dtype,
@@ -1413,6 +1486,7 @@ def dynamic_m_tensor_signature(
                 sample_sfd_col=sfd_col_tensor,
                 sample_norm_const=norm_const_tensor,
                 sample_prob=prob_tensor,
+                sample_row_scale=row_scale_tensor,
                 acc_dtype=acc_dtype,
                 mma_tiler_mn=mma_tiler_mn,
                 cluster_shape_mn=cluster_shape_mn,
@@ -1439,6 +1513,7 @@ def dynamic_m_tensor_signature(
                 sample_sfd_col=sfd_col_tensor,
                 sample_norm_const=norm_const_tensor,
                 sample_prob=prob_tensor,
+                sample_row_scale=row_scale_tensor,
                 acc_dtype=acc_dtype,
                 mma_tiler_mn=mma_tiler_mn,
                 cluster_shape_mn=cluster_shape_mn,
@@ -1469,6 +1544,7 @@ def dynamic_m_tensor_signature(
             amax_tensor=amax_tensor,
             norm_const_tensor=norm_const_tensor,
             prob_tensor=prob_tensor,
+            row_scale_tensor=row_scale_tensor,
             bias_tensor=bias_tensor,
             current_stream=current_stream,
         )
@@ -1487,6 +1563,7 @@ def dynamic_m_tensor_signature(
             amax_tensor=amax_tensor,
             norm_const_tensor=norm_const_tensor,
             prob_tensor=prob_tensor,
+            row_scale_tensor=row_scale_tensor,
             bias_tensor=bias_tensor,
             current_stream=current_stream,
         )

python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py

@@ -565,6 +565,7 @@ def __call__(
         norm_const_tensor: Optional[cute.Tensor],
         padded_offsets: cute.Tensor,
         alpha: cute.Tensor,
+        row_scale: Optional[cute.Tensor],
         bias: Optional[cute.Tensor],
         prob: cute.Tensor,
         max_active_clusters: cutlass.Constexpr,
@@ -863,6 +864,7 @@ class SharedStorage:
             amax_tensor,
             padded_offsets,
             alpha,
+            row_scale,
             bias,
             prob,
             workspace_ptr,
@@ -1097,6 +1099,7 @@ def kernel(
         mAmax_tensor: Optional[cute.Tensor],
         padded_offsets: cute.Tensor,
         alpha: cute.Tensor,
+        row_scale: Optional[cute.Tensor],
         mBias_nl: Optional[cute.Tensor],
         prob: cute.Tensor,
         workspace_ptr,
@@ -1810,6 +1813,17 @@ def kernel(
                 mPosition = epi_work_tile_info.tile_m_idx * self.cta_tile_shape_mnk[0] + tidx
                 real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info)
                 mProb = real_prob[mPosition, 0, 0]
+                acc_scale = cutlass.Float32(alpha_val)
+                # Optional per-row epilogue scale, applied together with the
+                # per-expert alpha before output conversion.
+                if cutlass.const_expr(row_scale is not None):
+                    real_row_scale, _ = epi_ext.get_gmem_tensor(
+                        "row_scale",
+                        row_scale,
+                        padded_offsets,
+                        epi_work_tile_info,
+                    )
+                    acc_scale = acc_scale * real_row_scale[mPosition]
 
                 # C1 fix: phase-based acc stage indexing for overlapping_accum
                 if cutlass.const_expr(self.overlapping_accum):
@@ -1859,26 +1873,26 @@ def kernel(
                                 )
                                 tTR_rAcc[i], tTR_rAcc[i + 1] = cute.arch.fma_packed_f32x2(
                                     (tTR_rAcc[i], tTR_rAcc[i + 1]),
-                                    (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)),
+                                    (acc_scale, acc_scale),
                                     (bias_f32_0, bias_f32_1),
                                     rnd="rn",
                                     ftz=False,
                                 )
                         else:
                             for i in cutlass.range_constexpr(cute.size(tTR_rAcc)):
-                                tTR_rAcc[i] = tTR_rAcc[i] * cutlass.Float32(alpha_val) + bias_vec[i].to(cutlass.Float32) * mProb
+                                tTR_rAcc[i] = tTR_rAcc[i] * acc_scale + bias_vec[i].to(cutlass.Float32) * mProb
                     else:
                         if cutlass.const_expr(self.vectorized_f32):
                             for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2):
                                 tTR_rAcc[i], tTR_rAcc[i + 1] = cute.arch.mul_packed_f32x2(
                                     (tTR_rAcc[i], tTR_rAcc[i + 1]),
-                                    (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)),
+                                    (acc_scale, acc_scale),
                                     rnd="rn",
                                     ftz=False,
                                 )
                         else:
                             for i in cutlass.range_constexpr(cute.size(tTR_rAcc)):
-                                tTR_rAcc[i] = tTR_rAcc[i] * cutlass.Float32(alpha_val)
+                                tTR_rAcc[i] = tTR_rAcc[i] * acc_scale
 
                     acc_vec = tTR_rAcc.load()
                     if cutlass.const_expr(not self.enable_bias):

python/cudnn/grouped_gemm/moe_sched_extension.py

@@ -101,7 +101,8 @@ class DiscreteWeightScaledGemmSchedExtension(MoESchedExtension):
     MoE scheduler extension for discrete-weight block-scaled grouped GEMM
     with GLU and quantization fusion.
 
-    Handles domain conversion for: a, b, c, d, d_col, prob, dprob, sfa, sfd, sfd_col, sfb.
+    Handles domain conversion for: a, b, c, d, d_col, prob, dprob,
+    row_scale, sfa, sfd, sfd_col, sfb.
 
     B and SFB are discrete (per-expert pointer arrays) → use expert-wise
     TMA descriptors from workspace.
@@ -180,6 +181,11 @@ def get_gmem_tensor(
             real = rewrite_tensor_shape(real, (shape[0], c1))
             return (real, None)
 
+        elif cutlass.const_expr(tensor_name == "row_scale"):
+            real = cute.domain_offset((token_offset,), gmem_tensor_in_moe_view)
+            real = rewrite_tensor_shape(real, (tokens_i,))
+            return (real, None)
+
         elif cutlass.const_expr(tensor_name in ("c", "d", "d_col", "d_srelu", "prob", "dprob")):
             # C/D/D_col/prob: contiguous M, offset by token_offset, global desc
             real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view)
@@ -298,6 +304,11 @@ def get_gmem_tensor(
             real = rewrite_tensor_shape(real, (shape[0], c1))
             return (real, None)
 
+        elif cutlass.const_expr(tensor_name == "row_scale"):
+            real = cute.domain_offset((token_offset,), gmem_tensor_in_moe_view)
+            real = rewrite_tensor_shape(real, (tokens_i,))
+            return (real, None)
+
         elif cutlass.const_expr(tensor_name in ("c", "d", "d_col", "d_srelu", "prob", "dprob")):
             real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view)
             real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1))