Support FP32 output and dynamic M in row-scaled FP4 grouped GEMM

docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md

@@ -340,7 +340,6 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 - `acc_dtype` must be `float32`
 - `sf_dtype=float8_e4m3fn` is incompatible with `sf_vec_size=32`
 - FP8 `ab_dtype` is incompatible with `sf_vec_size=16`
-- FP4 `ab_dtype` with `sf_vec_size=16` and `d_dtype=float32` is not supported
 
 ### Scale Factor Output Requirements
 

python/cudnn/gemm/cutedsl/grouped/quant/api.py

@@ -431,11 +431,6 @@ def check_support(self) -> bool:
                 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,
-            "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. Please use sf_vec_size 32 or d_dtype bf16 instead",
-        )
-
         if self.weight_mode == MoEWeightMode.DISCRETE:
             self._value_error_if(
                 self.b_major not in ["k", "n"],
@@ -951,7 +946,14 @@ 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)
+        row_scale_tensor = None
+        if self.row_scale_desc is not None:
+            row_scale_tensor = self._make_fake_cute_tensor(
+                dtype=self.row_scale_desc.dtype,
+                shape=(valid_m,),
+                stride=self.row_scale_desc.stride,
+                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")

test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py

@@ -300,6 +300,33 @@ def test_grouped_gemm_quant_discrete_wrapper_cache_dynamic_m_smoke(request, monk
     assert cache_entries == 1
 
 
+@pytest.mark.L0
+@torch_fork_set_rng(seed=8)
+def test_grouped_gemm_quant_discrete_wrapper_cache_dynamic_m_row_scale(request, monkeypatch):
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("Requires SM100+ for grouped GEMM quant kernel.")
+    if torch.cuda.get_device_capability() == (10, 7):
+        pytest.skip("Row-scale fusion is not supported on SM107.")
+
+    compile_count, cache_entries = _test_grouped_gemm_quant_discrete_wrapper_dynamic_m_cache_behavior(
+        request=request,
+        monkeypatch=monkeypatch,
+        d_dtype=torch.float32,
+        sf_dtype=torch.float8_e4m3fn,
+        row_scale=True,
+        use_dynamic_sched=True,
+        cfg_overrides={"n": 256, "k": 64, "l": 4},
+        group_m_lists=[
+            [0, 512, 512, 512],
+            [0, 1024, 1024, 1024],
+        ],
+        skip_unsupported=False,
+    )
+
+    assert compile_count == 1
+    assert cache_entries == 1
+
+
 @pytest.mark.L0
 @torch_fork_set_rng(seed=0)
 @with_scheduler_modes
@@ -518,6 +545,42 @@ def test_grouped_gemm_quant_discrete_wrapper_fp4_row_scale(request):
     )
 
 
+@pytest.mark.L0
+@torch_fork_set_rng(seed=7)
+@pytest.mark.parametrize("d_dtype", [torch.bfloat16, torch.float32], ids=["bf16", "fp32"])
+@pytest.mark.parametrize("enable_bias", [False, True], ids=["no_bias", "bias"])
+def test_grouped_gemm_quant_discrete_wrapper_fp4_row_scale_small_k(d_dtype, enable_bias, request):
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("Requires SM100+ for grouped GEMM quant kernel.")
+    if torch.cuda.get_device_capability() == (10, 7):
+        pytest.skip("Row-scale fusion is not supported on SM107.")
+
+    _test_grouped_gemm_quant_discrete_wrapper(
+        ab_dtype=torch.float4_e2m1fn_x2,
+        c_dtype=torch.bfloat16,
+        d_dtype=d_dtype,
+        b_major="k",
+        cd_major="n",
+        acc_dtype=torch.float32,
+        mma_tiler_mn=(256, 256),
+        cluster_shape_mn=(2, 1),
+        sf_vec_size=16,
+        sf_dtype=torch.float8_e4m3fn,
+        vector_f32=False,
+        discrete_col_sfd=False,
+        request=request,
+        use_dynamic_sched=True,
+        row_scale=True,
+        cfg_overrides={
+            "n": 256,
+            "k": 64,
+            "group_m_list": [0, 512, 512, 512],
+        },
+        enable_bias=enable_bias,
+        skip_unsupported=False,
+    )
+
+
 @pytest.mark.L0
 @torch_fork_set_rng(seed=0)
 def test_grouped_gemm_quant_wrapper_without_prob_tensor(request):
@@ -899,6 +962,13 @@ def counted_compile(self):
 def _test_grouped_gemm_quant_discrete_wrapper_dynamic_m_cache_behavior(
     request,
     monkeypatch,
+    d_dtype=torch.bfloat16,
+    sf_dtype=torch.float8_e8m0fnu,
+    row_scale=False,
+    use_dynamic_sched=False,
+    cfg_overrides=None,
+    group_m_lists=None,
+    skip_unsupported=True,
 ):
     try:
         from cudnn import grouped_gemm_quant_wrapper_sm100
@@ -923,33 +993,38 @@ def counted_compile(self):
         request=request,
         ab_dtype=torch.float4_e2m1fn_x2,
         c_dtype=torch.bfloat16,
-        d_dtype=torch.bfloat16,
+        d_dtype=d_dtype,
         cd_major="n",
         acc_dtype=torch.float32,
         mma_tiler_mn=(256, 256),
         cluster_shape_mn=(2, 1),
         sf_vec_size=16,
-        sf_dtype=torch.float8_e8m0fnu,
+        sf_dtype=sf_dtype,
         vector_f32=False,
         discrete_col_sfd=False,
         b_major="k",
     )
+    cfg = _apply_grouped_gemm_cfg_overrides(cfg, cfg_overrides)
 
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    if group_m_lists is None:
+        group_m_lists = [[group_m] * cfg["l"] for group_m in DYNAMIC_SHAPES_M_VALUES]
 
     try:
-        for group_m in DYNAMIC_SHAPES_M_VALUES:
+        for group_m_list in group_m_lists:
             inputs = allocate_discrete_input_tensors(
                 n=cfg["n"],
                 k=cfg["k"],
                 num_experts=cfg["l"],
-                group_m_list=[group_m] * cfg["l"],
+                group_m_list=group_m_list,
                 ab_dtype=cfg["ab_dtype"],
                 sf_dtype=cfg["sf_dtype"],
                 sf_vec_size=cfg["sf_vec_size"],
                 m_aligned=cfg["m_aligned"],
                 b_major=cfg["b_major"],
             )
+            if row_scale:
+                _add_row_scale(inputs)
 
             grouped_gemm_quant_wrapper_sm100(
                 a_tensor=inputs["a_tensor"],
@@ -963,6 +1038,7 @@ def counted_compile(self):
                 b_major=cfg["b_major"],
                 norm_const_tensor=inputs.get("norm_const_tensor"),
                 prob_tensor=inputs["prob_tensor"],
+                row_scale_tensor=inputs.get("row_scale_tensor"),
                 acc_dtype=cfg["acc_dtype"],
                 d_dtype=cfg["d_dtype"],
                 cd_major=cfg["cd_major"],
@@ -972,11 +1048,14 @@ def counted_compile(self):
                 vector_f32=cfg["vector_f32"],
                 m_aligned=cfg["m_aligned"],
                 discrete_col_sfd=cfg["discrete_col_sfd"],
+                use_dynamic_sched=use_dynamic_sched,
                 current_stream=stream,
             )
             torch.cuda.synchronize()
     except (ValueError, NotImplementedError) as e:
-        pytest.skip(f"Unsupported testcase: {e}")
+        if skip_unsupported:
+            pytest.skip(f"Unsupported testcase: {e}")
+        raise
     finally:
         cache_entries = len(grouped_gemm_quant_api._cache_of_GroupedGemmQuantSm100Objects)
         grouped_gemm_quant_api._cache_of_GroupedGemmQuantSm100Objects.clear()
@@ -1392,6 +1471,9 @@ def _test_grouped_gemm_quant_discrete_wrapper(
     request,
     use_dynamic_sched=False,
     row_scale=False,
+    cfg_overrides=None,
+    enable_bias=False,
+    skip_unsupported=True,
 ):
     try:
         from cudnn import grouped_gemm_quant_wrapper_sm100
@@ -1414,6 +1496,7 @@ def _test_grouped_gemm_quant_discrete_wrapper(
         discrete_col_sfd,
         b_major,
     )
+    cfg = _apply_grouped_gemm_cfg_overrides(cfg, cfg_overrides)
 
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
@@ -1427,9 +1510,12 @@ def _test_grouped_gemm_quant_discrete_wrapper(
         sf_vec_size=cfg["sf_vec_size"],
         m_aligned=cfg["m_aligned"],
         b_major=cfg["b_major"],
+        enable_bias=enable_bias,
     )
     if row_scale:
         _add_row_scale(inputs)
+    if enable_bias:
+        inputs["bias_ref"] = inputs["bias_tensor"]
 
     try:
         for _ in range(2):
@@ -1440,6 +1526,7 @@ def _test_grouped_gemm_quant_discrete_wrapper(
                 alpha_tensor=inputs["alpha_tensor"],
                 b_ptrs=inputs["b_ptrs_tensor"],
                 sfb_ptrs=inputs["sfb_ptrs_tensor"],
+                bias_tensor=inputs["bias_tensor"],
                 n=cfg["n"],
                 b_dtype=cfg["ab_dtype"],
                 b_major=cfg["b_major"],
@@ -1459,7 +1546,9 @@ def _test_grouped_gemm_quant_discrete_wrapper(
                 current_stream=stream,
             )
     except (ValueError, NotImplementedError) as e:
-        pytest.skip(f"Unsupported testcase: {e}")
+        if skip_unsupported:
+            pytest.skip(f"Unsupported testcase: {e}")
+        raise
 
     _check_ref_grouped_gemm_quant_discrete(
         inputs,

test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py

@@ -346,7 +346,10 @@ def run_grouped_gemm_quant_ref(
         start = 0
         for i, group_m in enumerate(aligned_group_m_list):
             end = start + group_m
-            amax_ref[i] = compute_reference_amax(ref[start:end, :, 0].clone())
+            if group_m == 0:
+                amax_ref[i] = float("-inf")
+            else:
+                amax_ref[i] = compute_reference_amax(ref[start:end, :, 0].clone())
             start = end
         ref_tensors["amax_ref"] = amax_ref