[PyTorch] Enable fused FP8 block-scaling path in GroupedLinear module

tests/cpp/operator/test_cast_float8blockwise_grouped.cu

@@ -368,14 +368,15 @@ struct TestConfig {
   ScalingDir dir;
   std::vector<size_t> first_dims;
   size_t K;
+  bool force_pow_2_scales;
 };
 
 class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam<TestConfig> {};
 
 TEST_P(GroupedFP8BlockwiseTestSuite, Test) {
   const TestConfig& cfg = GetParam();
   perform_test<bf16, fp8e4m3>(cfg.shape_rep, cfg.block_dim, cfg.dir, cfg.first_dims, cfg.K,
-                              /*force_pow_2_scales=*/false, /*epsilon=*/0.0f);
+                              cfg.force_pow_2_scales, /*epsilon=*/0.0f);
 }
 
 std::vector<TestConfig> make_configs() {
@@ -387,11 +388,13 @@ std::vector<TestConfig> make_configs() {
   for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) {
     for (auto dir : {ScalingDir::ROWWISE, ScalingDir::COLWISE, ScalingDir::BOTH}) {
       for (size_t K : Ks) {
-        for (const auto& v : uniform) {
-          configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K});
-        }
-        for (const auto& v : jagged) {
-          configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K});
+        for (bool pow2 : {false, true}) {
+          for (const auto& v : uniform) {
+            configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K, pow2});
+          }
+          for (const auto& v : jagged) {
+            configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K, pow2});
+          }
         }
       }
     }
@@ -408,6 +411,7 @@ std::string make_name(const ::testing::TestParamInfo<TestConfig>& info) {
   s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size());
   s += "_M";
   for (size_t m : c.first_dims) s += "_" + std::to_string(m);
+  s += (c.force_pow_2_scales ? "_POW2" : "_FP32SC");
   return s;
 }
 

tests/cpp/operator/test_cast_mxfp8_grouped.cu

@@ -509,9 +509,9 @@ void performTest(const ProcessingMethod processing_method,
             break;
         }
         case ProcessingMethod::CAST_DBIAS: {
-            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0);
+            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0);
             workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype());
-            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0);
+            nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0);
             break;
         }
         case ProcessingMethod::CAST_DBIAS_DACT: {

tests/pytorch/test_grouped_linear.py

@@ -1499,6 +1499,9 @@ def test_fp8_grouped_gemm(shape, accumulate):
 _fp8_available, _reason_for_no_fp8 = fp8_available, reason_for_no_fp8
 _mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8
 _nvfp4_available, _reason_for_no_nvfp4 = nvfp4_available, reason_for_no_nvfp4
+_fp8_block_scaling_available, _reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available(
+    return_reason=True
+)
 
 
 @pytest.fixture(autouse=True)
@@ -1536,6 +1539,20 @@ def _run_grouped_linear_path(
     out_features = weights[0].size(0)
     use_fp8 = fp8_recipe is not None
 
+    # Spy on the grouped quantize entry point so a predicate that silently
+    # declines the fused path fails the test instead of vacuously matching the
+    # legacy reference.
+    group_quantize_calls = 0
+    if enable_grouped_tensor_path and use_fp8:
+        real_group_quantize = tex.group_quantize
+
+        def _counting_group_quantize(*args, **kwargs):
+            nonlocal group_quantize_calls
+            group_quantize_calls += 1
+            return real_group_quantize(*args, **kwargs)
+
+        monkeypatch.setattr(tex, "group_quantize", _counting_group_quantize)
+
     x = x_base.detach().clone().requires_grad_(True)
     with quantized_model_init(enabled=fp8_model_params, recipe=fp8_recipe):
         grouped_linear = GroupedLinear(
@@ -1566,6 +1583,10 @@ def _run_grouped_linear_path(
     if delay_wgrad_compute:
         grouped_linear.backward_dw()
 
+    if enable_grouped_tensor_path and use_fp8:
+        monkeypatch.setattr(tex, "group_quantize", real_group_quantize)
+        assert group_quantize_calls > 0, "fused GroupedTensor path was not taken"
+
     outputs = [y, x.grad]
     for i in range(num_gemms):
         outputs.append(getattr(grouped_linear, f"weight{i}").grad)
@@ -1590,8 +1611,14 @@ def _run_grouped_linear_path(
             recipe.NVFP4BlockScaling(disable_stochastic_rounding=True),
             marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4),
         ),
+        pytest.param(
+            recipe.Float8BlockScaling(),
+            marks=pytest.mark.skipif(
+                not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling
+            ),
+        ),
     ],
-    ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"],
+    ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"],
 )
 @pytest.mark.parametrize("bias", _ALL_BOOLEAN)
 @pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN)
@@ -1605,10 +1632,13 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy(
         pytest.skip(
             "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)."
         )
-    # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor
-    # current scaling also runs on the Hopper grouped GEMM path.
+    # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor
+    # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only.
     is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling()
-    if use_fp8 and not is_current_scaling and device_capability < (10, 0):
+    is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling()
+    if is_block_scaling and not (9, 0) <= device_capability < (10, 0):
+        pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).")
+    if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0):
         pytest.skip(
             "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)."
         )
@@ -1810,8 +1840,14 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
             recipe.NVFP4BlockScaling(disable_stochastic_rounding=True),
             marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4),
         ),
+        pytest.param(
+            recipe.Float8BlockScaling(),
+            marks=pytest.mark.skipif(
+                not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling
+            ),
+        ),
     ],
-    ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"],
+    ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"],
 )
 @pytest.mark.parametrize("bias", _ALL_BOOLEAN)
 def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch):
@@ -1822,10 +1858,13 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch
         pytest.skip(
             "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)."
         )
-    # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor
-    # current scaling also runs on the Hopper grouped GEMM path.
+    # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor
+    # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only.
     is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling()
-    if use_fp8 and not is_current_scaling and device_capability < (10, 0):
+    is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling()
+    if is_block_scaling and not (9, 0) <= device_capability < (10, 0):
+        pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).")
+    if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0):
         pytest.skip(
             "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)."
         )
@@ -1932,6 +1971,24 @@ def _train_step(x, dy, out_buf, *, use_graphed):
         torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols)
 
 
+@pytest.mark.skipif(not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling)
+@pytest.mark.skipif(
+    not (10, 0) <= torch.cuda.get_device_capability() <= (11, 0),
+    reason="Error path only triggers on Blackwell (SM100/SM110).",
+)
+def test_grouped_linear_fused_path_fp8_block_scaling_blackwell_error(monkeypatch):
+    """FP8BS + fused env var on Blackwell must raise, not silently fall back."""
+    monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1")
+    FP8GlobalStateManager.reset()
+    dtype = torch.bfloat16
+    grouped_linear = GroupedLinear(2, 128, 128, bias=False, params_dtype=dtype, device="cuda")
+    x = torch.randn(256, 128, device="cuda", dtype=dtype, requires_grad=True)
+    m_splits = torch.tensor([128, 128], dtype=torch.int64, device="cuda")
+    with pytest.raises(RuntimeError, match="Hopper-only"):
+        with autocast(enabled=True, recipe=recipe.Float8BlockScaling()):
+            grouped_linear(x, m_splits)
+
+
 @pytest.mark.parametrize("swizzle_type", ["mxfp8_rowwise", "mxfp8_columnwise", "nvfp4"])
 def test_swizzle_scales_and_pack_ptrs_for_discrete_weights(
     swizzle_type: str,

tests/pytorch/test_grouped_tensor.py

@@ -1008,11 +1008,19 @@ def _assert_fp8_cs_group_quantize_matches_reference(
     @pytest.mark.parametrize("shape_case", ["uniform", "varying_first"])
     @pytest.mark.parametrize("direction", ["rowwise", "columnwise", "both"])
     @pytest.mark.parametrize("output_dbias", [False, True])
+    @pytest.mark.parametrize(
+        "force_pow_2_scales", [False, True], ids=["fp32_scales", "pow2_scales"]
+    )
     @pytest.mark.skipif(
         not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped
     )
     def test_quantize_grouped_fp8_blockwise(
-        self, block_scaling_dim: int, shape_case: str, direction: str, output_dbias: bool
+        self,
+        block_scaling_dim: int,
+        shape_case: str,
+        direction: str,
+        output_dbias: bool,
+        force_pow_2_scales: bool,
     ) -> None:
         """Test grouped FP8 block-scaling quantization against per-tensor quantization.
 
@@ -1061,7 +1069,7 @@ def test_quantize_grouped_fp8_blockwise(
             fp8_dtype=tex.DType.kFloat8E4M3,
             rowwise=rowwise,
             columnwise=columnwise,
-            force_pow_2_scales=False,
+            force_pow_2_scales=force_pow_2_scales,
             amax_epsilon=0.0,
             block_scaling_dim=block_scaling_dim,
         )
@@ -1081,7 +1089,7 @@ def test_quantize_grouped_fp8_blockwise(
             fp8_dtype=tex.DType.kFloat8E4M3,
             rowwise=True,
             columnwise=True,
-            force_pow_2_scales=False,
+            force_pow_2_scales=force_pow_2_scales,
             amax_epsilon=0.0,
             block_scaling_dim=block_scaling_dim,
         )

transformer_engine/common/cast/cast_grouped_dbias.cu

@@ -11,7 +11,8 @@
 #include "dispatch/quantize.cuh"
 
 void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output,
-                               NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) {
+                               NVTEGroupedTensor dbias, NVTETensor workspace,
+                               const NVTEQuantizationConfig quant_config, cudaStream_t stream) {
   NVTE_API_CALL(nvte_group_quantize_dbias);
   using namespace transformer_engine;
 
@@ -20,5 +21,5 @@ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor
   constexpr const NVTEGroupedTensor activation_input = nullptr;
 
   dispatch::group_quantize_bwd_helper<IS_DBIAS, IS_DACT, Empty, nullptr>(
-      input, activation_input, output, dbias, workspace, nullptr, stream);
+      input, activation_input, output, dbias, workspace, quant_config, stream);
 }

transformer_engine/common/cast/dispatch/quantize.cuh

@@ -475,22 +475,16 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor
     }
     case NVTE_BLOCK_SCALING_1D: {
       NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_1D.");
-      NVTE_CHECK(!quant_config_cpp.force_pow_2_scales,
-                 "Fused grouped FP8 block-scaling quantize does not support "
-                 "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused "
-                 "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0).");
       fp8_blockwise::group_quantize_blockwise_1d(input_tensor, output_tensor, noop_tensor,
-                                                 quant_config_cpp.amax_epsilon, stream);
+                                                 quant_config_cpp.amax_epsilon,
+                                                 quant_config_cpp.force_pow_2_scales, stream);
       break;
     }
     case NVTE_BLOCK_SCALING_2D: {
       NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_2D.");
-      NVTE_CHECK(!quant_config_cpp.force_pow_2_scales,
-                 "Fused grouped FP8 block-scaling quantize does not support "
-                 "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused "
-                 "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0).");
       fp8_blockwise::group_quantize_blockwise_2d(input_tensor, output_tensor, noop_tensor,
-                                                 quant_config_cpp.amax_epsilon, stream);
+                                                 quant_config_cpp.amax_epsilon,
+                                                 quant_config_cpp.force_pow_2_scales, stream);
       break;
     }
     default:
@@ -537,22 +531,18 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe
     case NVTE_BLOCK_SCALING_1D:
     case NVTE_BLOCK_SCALING_2D: {
       NVTE_CHECK(!IS_DACT, "IS_DACT is not implemented for grouped FP8 block scaling.");
-      NVTE_CHECK(!quant_config_cpp.force_pow_2_scales,
-                 "Fused grouped FP8 block-scaling quantize does not support "
-                 "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused "
-                 "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0).");
       // dbias is computed in-kernel and reduced per-expert inside group_quantize_blockwise_{1d,2d}
       // (mirrors MXFP8); those also handle the two-call workspace sizing protocol.
       GroupedTensor *dbias_arg = IS_DBIAS ? dbias_tensor : nullptr;
       Tensor *workspace_arg = IS_DBIAS ? workspace_tensor : nullptr;
       if (scaling_mode == NVTE_BLOCK_SCALING_1D) {
-        fp8_blockwise::group_quantize_blockwise_1d(grad_tensor, output_tensor, noop_tensor,
-                                                   quant_config_cpp.amax_epsilon, stream, dbias_arg,
-                                                   workspace_arg);
+        fp8_blockwise::group_quantize_blockwise_1d(
+            grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon,
+            quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg);
       } else {
-        fp8_blockwise::group_quantize_blockwise_2d(grad_tensor, output_tensor, noop_tensor,
-                                                   quant_config_cpp.amax_epsilon, stream, dbias_arg,
-                                                   workspace_arg);
+        fp8_blockwise::group_quantize_blockwise_2d(
+            grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon,
+            quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg);
       }
       break;
     }