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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

denera merged 13 commits into NVIDIA:main from denera:pytorch/grouped-linear-fused-fp8bs

Jul 22, 2026

tests/cpp/operator/test_cast_float8blockwise_grouped.cu

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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

-Original file line number
+Diff line change
@@ Expand Up / @@ -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: {
@@ Expand Down @@

tests/pytorch/test_grouped_linear.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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)

    @@ -1590,8 +1593,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 +1614,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+)."

            )

    @@ -1906,8 +1918,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):

    @@ -1918,10 +1936,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+)."

            )

    @@ -2037,6 +2058,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

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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

-Original file line number
+Diff line change
@@ Expand Up / @@ -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;
@@ Expand All @@
       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

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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;

        }

Navigation Menu

Uh oh!

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

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!