From cc52e94178597f7917900b817b1e26680c67fc93 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 04:40:43 +0000 Subject: [PATCH 1/7] [Common/PyTorch] Support power-of-2 scales in grouped FP8 block-scaling quantize The default Float8BlockScaling recipe constrains scales to powers of 2, so the fused grouped path must honor the flag to stay numerically consistent with the unfused path. Thread a runtime pow_2_scales argument through the grouped quantize kernels (the shared scale helper already implements the rounding) and drop the force_pow_2_scales rejections. Also add a quantization-config parameter to nvte_group_quantize_dbias, which previously had no way to receive force_pow_2_scales or amax_epsilon on the bgrad path. Signed-off-by: Alp Dener --- .../test_cast_float8blockwise_grouped.cu | 16 ++++++---- tests/cpp/operator/test_cast_mxfp8_grouped.cu | 4 +-- tests/pytorch/test_grouped_tensor.py | 14 +++++++-- .../common/cast/cast_grouped_dbias.cu | 5 ++-- .../common/cast/dispatch/quantize.cuh | 30 +++++++------------ .../group_quantize_fp8_blockwise.cuh | 28 +++++++++-------- .../common/include/transformer_engine/cast.h | 4 ++- .../pytorch/csrc/extensions/cast.cpp | 15 ++++++++-- transformer_engine/pytorch/csrc/quantizer.cpp | 8 ----- 9 files changed, 68 insertions(+), 56 deletions(-) diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu index c7d7a73475..88cd14c805 100644 --- a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu +++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu @@ -367,6 +367,7 @@ struct TestConfig { ScalingDir dir; std::vector first_dims; size_t K; + bool force_pow_2_scales; }; class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam {}; @@ -374,7 +375,7 @@ class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam TEST_P(GroupedFP8BlockwiseTestSuite, Test) { const TestConfig& cfg = GetParam(); perform_test(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 make_configs() { @@ -386,11 +387,13 @@ std::vector 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}); + } } } } @@ -407,6 +410,7 @@ std::string make_name(const ::testing::TestParamInfo& 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; } diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index de72299be1..80b80da0d8 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/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: { diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index eeb6e7a394..4dd52ee2bd 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/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, ) diff --git a/transformer_engine/common/cast/cast_grouped_dbias.cu b/transformer_engine/common/cast/cast_grouped_dbias.cu index 5290255a00..b7ced30b11 100644 --- a/transformer_engine/common/cast/cast_grouped_dbias.cu +++ b/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( - input, activation_input, output, dbias, workspace, nullptr, stream); + input, activation_input, output, dbias, workspace, quant_config, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 031122d966..12de23e0bc 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/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; } diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 1fd1738f93..e43c796d79 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -284,7 +284,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_stride_y, - const float epsilon, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { + const float epsilon, const bool pow_2_scales, const float* __restrict__ noop_ptr, + float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -379,7 +380,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma block_amax = fmaxf(block_amax, warp_amaxes[w]); } const CType scale = - compute_scale_from_types(block_amax, epsilon, /*pow_2_scaling=*/false); + compute_scale_from_types(block_amax, epsilon, pow_2_scales); // The 2D colwise per-expert scale offset requires a CTA-cooperative prefix // sum in the VARYING_FIRST_DIM case, so compute it across all threads before @@ -475,6 +476,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t R_total, const float epsilon, + const bool pow_2_scales, const float* __restrict__ noop_ptr) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -531,7 +533,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) amax = subwarp_reduce_max_broadcast(amax); const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; if (thr_col == 0 && r_global < R_total) { // Per-expert layout: (blocks_X, roundup(M_t, 4)). Compute expert base @@ -572,8 +574,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_t_stride_aligned_K, - const size_t R_total, const float epsilon, const float* __restrict__ noop_ptr, - float* __restrict__ dbias_workspace) { + const size_t R_total, const float epsilon, const bool pow_2_scales, + const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -657,7 +659,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke amax = subwarp_reduce_max_broadcast(amax); const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t r_global = global_row_base + row_local; @@ -734,7 +736,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke amax = subwarp_reduce_max_broadcast(amax); const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t c_global = global_col_base + col_local; @@ -814,7 +816,8 @@ inline GroupedBlockwiseLaunchInfo prepare_grouped_blockwise_launch(const Grouped // reports the [total_row_blocks, K] fp32 shape and returns without launching. inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -883,7 +886,7 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_stride_y, epsilon, - noop_ptr, dbias_workspace); + pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS @@ -907,7 +910,8 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso // per-tile column partial can be computed. inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -967,7 +971,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso reinterpret_cast(output->scale_inv.dptr), info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, - info.R_total, epsilon, noop_ptr); + info.R_total, epsilon, pow_2_scales, noop_ptr); } } else if constexpr (kRowwise || kColwise) { // CW-only, BOTH, or RW-only WITH dbias: smem-cached TMA kernel. @@ -998,7 +1002,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_t_stride_aligned_K, - info.R_total, epsilon, noop_ptr, dbias_workspace); + info.R_total, epsilon, pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 554d8c1ac9..4d6d24ba65 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -161,10 +161,12 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d * \param[in,out] output Output grouped FP8/MXFP8 tensor. * \param[out] dbias Result of the reduction of the input along columns. * \param[out] workspace Workspace tensor. + * \param[in] quant_config Quantization configuration. * \param[in] stream CUDA stream used for the operation. */ 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); /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8d77a9e349..8b1cd384aa 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -500,6 +500,15 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto quantizer_cpp = convert_quantizer(quantizer); + // MXFP8 carries no quantization knobs; FP8 block-scaling reads scale constraints + // off the quantizer, matching the forward group_quantize dispatch. + QuantizationConfigWrapper quant_config_cpp; + if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) { + auto *fp8_block_quantizer_cpp = static_cast(quantizer_cpp.get()); + quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); + quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + } + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), @@ -530,7 +539,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto stream = at::cuda::getCurrentCUDAStream(); NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); if (workspace_nvte.ndim() > 0 && workspace_nvte.numel() > 0) { at::Tensor workspace_torch = allocateSpace(workspace_nvte.shape(), workspace_nvte.dtype()); @@ -539,7 +549,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); return py::make_tuple(py::reinterpret_borrow(grouped_output_py), py::cast(std::move(dbias_torch))); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index d2a3be888f..f168e66ada 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1154,14 +1154,6 @@ std::pair Float8BlockQuantizer::create_grouped const size_t logical_last_dim) const { using namespace pybind11::literals; - // The fused grouped FP8 block-scaling path uses unconstrained FP32 scales and does not - // implement power-of-2 scaling. Reject force_pow_2_scales rather than silently ignoring it; - // the unfused per-tensor split-quantize path still honors it. - NVTE_CHECK(!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) which supports power-of-2 scales."); - const auto tensor_offsets = resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, logical_first_dim, logical_last_dim); From d8f6a08bb501132d82e0272abb7855a29754ba19 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 07:23:37 +0000 Subject: [PATCH 2/7] [PyTorch] Enable fused grouped FP8 block-scaling path in GroupedLinear module Admit Float8BlockQuantizer in the fused GroupedTensor path on Hopper. The existing usage flags already match the Hopper TN-only mapping and the grouped GEMM selects transposed columnwise storage for NN/NT layouts, so only the path predicate changes. The fused path is an explicit opt-in via NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM, so raise on Blackwell (SM100/SM110) instead of silently falling back; the fused path has no MXFP8-broadcast emulation. Extend the fused dbias path (tex.bgrad_group_quantize) to FP8 block scaling when dgrad is required (dbias is computed in the rowwise pass). Add fp8_block_scaling to the fused-path tests with a Hopper-only gate, assert the fused path engages via a group_quantize spy, and add a Blackwell error-path test. Signed-off-by: Alp Dener --- tests/pytorch/test_grouped_linear.py | 73 +++++++++++++++++-- .../pytorch/module/grouped_linear.py | 40 ++++++++-- 2 files changed, 100 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 01a7cf2415..730211ba95 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/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+)." ) @@ -1808,8 +1838,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): @@ -1820,10 +1856,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+)." ) @@ -1928,6 +1967,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, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index c79adbf7e4..91712a277e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -55,7 +55,13 @@ from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias -from ..tensor import Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -103,11 +109,15 @@ def _is_grouped_tensor_path_supported( and be incompatible with CUDA Graphs. Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16 and FP8 per-tensor current scaling. + * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 + block scaling (1D/2D, including power-of-2 scales). * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 per-tensor current scaling. - FP8 delayed scaling and FP8 block scaling are not supported because the - corresponding grouped quantization kernels are missing. + FP8 delayed scaling is not supported because the corresponding grouped + quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and + SM110) raises instead of falling back: the fused path is Hopper-only and has + no MXFP8-broadcast emulation. Architectures outside the fused-path window + (e.g. SM120) fall back to the legacy path like every other recipe. Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization currently requires RHT. @@ -136,6 +146,21 @@ def _is_grouped_tensor_path_supported( if fp8: if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): return True + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM + # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast + # emulation. On Blackwell (SM100/SM110, the only other arch that reaches + # this branch) fail loudly rather than silently falling back to the + # unfused path the user explicitly opted out of. + if get_device_compute_capability() >= (10, 0): + raise RuntimeError( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" + " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" + " FP8 block-scaling path is Hopper-only. Unset" + " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" + " path (emulated via MXFP8 GEMM on Blackwell)." + ) + return True # MXFP8 and NVFP4 require Blackwell+. if not (10, 0) <= get_device_compute_capability() <= (11, 0): return False @@ -762,7 +787,12 @@ def _backward_grouped_tensor( columnwise=ctx.weights_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - if ctx.use_bias and isinstance(grad_output_quantizer, MXFP8Quantizer): + # The grouped FP8 block-scaling bgrad kernel computes dbias in the rowwise + # pass, so the fusion needs rowwise output (i.e. dgrad required). + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.requires_dgrad + ) + if ctx.use_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, From 01056d2ba9f77a230aa312d37d881f22328fe95e Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 07:23:57 +0000 Subject: [PATCH 3/7] [PyTorch] Enable FP8 block-scaling in GroupedLinear fusible op Replace the blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state with a per-op supports_float8_block_scaling flag and opt in the GroupedLinear op. Mirror the module-path predicate and fused-bgrad changes; since the graph-safe flow is default-on here (no env-var opt-in), other architectures fall back to the split-quantize flow instead of raising. Force use_split_accumulator=True for FP8 block-scaling operands in general_grouped_gemm_for_grouped_tensor, matching non-grouped general_gemm: cuBLAS has no fast-accum FP8 block-scaling algorithm, so the ops-layer forward failed algo selection without it. Add fp8_block_scaling coverage to the ops GroupedLinear tests. The CUDA-graph-safe test skips it for now: the replayed wgrad for the last expert diverges between replays depending on process allocation history; under investigation. Graph capture remains covered by the module-path test. Signed-off-by: Alp Dener --- tests/pytorch/test_fusible_ops.py | 23 +++++++++- tests/pytorch/test_grouped_mlp.py | 45 +++++++++++++++++-- tests/pytorch/utils.py | 1 + .../pytorch/cpp_extensions/gemm.py | 12 +++++ .../pytorch/ops/basic/grouped_linear.py | 30 +++++++++---- transformer_engine/pytorch/ops/op.py | 7 ++- 6 files changed, 104 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3c6560e7c3..f349ead414 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -32,6 +32,7 @@ ) from transformer_engine.pytorch import ( QuantizedTensor, + Float8BlockQuantizer, Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, @@ -54,6 +55,9 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -73,6 +77,11 @@ _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") +# GroupedLinear is currently the only basic op that supports FP8 block scaling. +_grouped_linear_quantization_list: list[Optional[str]] = list(_quantization_list) +if fp8_block_scaling_available: + _grouped_linear_quantization_list.append("fp8_block_scaling") + @pytest.fixture(autouse=True, scope="function") def _reset_rng_states_per_test(): @@ -182,6 +191,18 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) + elif quantization == "fp8_block_scaling": + tensor_type = "input" + if quantizer_role is not None: + tensor_type = quantizer_role.tensor_type + # Match the Float8BlockScaling recipe defaults: 2D (128x128) blocks for + # weights, 1D (1x128) otherwise, with power-of-2 scales. + test = Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2 if tensor_type == "weight" else 1, + )(test) elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): tensor_type = "input" if quantizer_role is not None: @@ -2063,7 +2084,7 @@ def test_dropout( @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) - @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantization", _grouped_linear_quantization_list) @pytest.mark.parametrize("quantized_compute", (False, True)) @pytest.mark.parametrize("quantized_weight", (False, True)) @pytest.mark.parametrize("input_requires_grad", (False, True)) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index e24fff9049..9c965fbe01 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -21,6 +21,7 @@ ) from transformer_engine.pytorch import ( QuantizedTensor, + Float8BlockQuantizer, Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, @@ -45,6 +46,9 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -61,6 +65,11 @@ _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") +# GroupedLinear is currently the only basic op that supports FP8 block scaling. +_grouped_linear_quantization_list: list[Optional[str]] = list(_quantization_list) +if fp8_block_scaling_available: + _grouped_linear_quantization_list.append("fp8_block_scaling") + # Quantization recipes supported by grouped MLP fused op _grouped_mlp_quantization_list: list[Optional[str]] = [None] if mxfp8_available: @@ -177,6 +186,18 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) + elif quantization == "fp8_block_scaling": + tensor_type = "input" + if quantizer_role is not None: + tensor_type = quantizer_role.tensor_type + # Match the Float8BlockScaling recipe defaults: 2D (128x128) blocks for + # weights, 1D (1x128) otherwise, with power-of-2 scales. + test = Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2 if tensor_type == "weight" else 1, + )(test) elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): tensor_type = "input" if quantizer_role is not None: @@ -230,7 +251,7 @@ class TestGroupedLinearOp: @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) - @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantization", _grouped_linear_quantization_list) @pytest.mark.parametrize("quantized_compute", (False, True)) @pytest.mark.parametrize("quantized_weight", (False, True)) @pytest.mark.parametrize("input_requires_grad", (False, True)) @@ -437,6 +458,7 @@ def test_grouped_linear( "quantization", [None] + (["fp8_current_scaling"] if fp8_available else []) + + (["fp8_block_scaling"] if fp8_block_scaling_available else []) + (["mxfp8"] if mxfp8_available else []) + (["nvfp4_rht"] if nvfp4_available else []), ) @@ -480,8 +502,25 @@ def test_grouped_linear_cuda_graph_safe( "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)" ) # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path, - # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+). - requires_blackwell = quantization is not None and quantization != "fp8_current_scaling" + # FP8 block scaling is Hopper-only, and MXFP8/NVFP4 grouped quantization kernels + # require Blackwell (SM100+). + if quantization == "fp8_block_scaling": + # Under investigation: the graph-replayed FP8 block-scaling wgrad for the last + # expert diverges between replays in this test (the first replay is correct + # against a float64 reference; the reference step's replay is not). The effect + # depends on the process allocation history and has not been reproduced in a + # standalone equivalent of this flow. Note make_graphed_callables replaces + # module.forward in place, so the "eager" reference step here replays the same + # graphs. The eager ops path is covered by test_grouped_linear and graph capture + # by the module-path test in test_grouped_linear.py. + pytest.skip( + "FP8 block-scaling grouped wgrad diverges between graph replays;" + " under investigation." + ) + requires_blackwell = quantization is not None and quantization not in ( + "fp8_current_scaling", + "fp8_block_scaling", + ) if requires_blackwell and device_capability < (10, 0): pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+. diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index f068f3581e..21601d8cdd 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -114,6 +114,7 @@ def quantization_tols(name: str) -> dict[str, float]: "fp8_delayed_scaling", "fp8_current_scaling", "fp8_blockwise", + "fp8_block_scaling", "mxfp8", "mxfp8_block_scaling", ): diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3b066d50b..f0f43515dc 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -13,6 +13,7 @@ from ..utils import get_sm_count, _empty_tensor from ..quantized_tensor import Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -474,6 +475,17 @@ def general_grouped_gemm_for_grouped_tensor( if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + def _is_fp8_blockwise(operand) -> bool: + if isinstance(operand, (list, tuple)): + return any(isinstance(t, Float8BlockwiseQTensorStorage) for t in operand) + if isinstance(operand, GroupedTensorStorage): + return isinstance(operand.quantizer, Float8BlockQuantizer) + return False + + if _is_fp8_blockwise(A) or _is_fp8_blockwise(B): + # FP8 block-scaling requires split accumulator + use_split_accumulator = True + if is_discrete_out: # wgrad case. grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_out diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 057c9da0e6..918d2bedef 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -27,6 +27,7 @@ from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...quantized_tensor import QuantizedTensorStorage from ...tensor import ( + Float8BlockQuantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer, MXFP8Tensor, @@ -116,6 +117,9 @@ class GroupedLinear(BasicOperation): # Operation expects input split sizes (and optionally scales tensor) num_extra_inputs: int = 1 + # Grouped FP8 block-scaling quantize kernels and grouped GEMM support exist + # for this op (Hopper-only fused path; legacy flow elsewhere). + supports_float8_block_scaling: bool = True def __init__( self, @@ -778,9 +782,11 @@ def _is_graph_safe_path_supported( * FP8 per-tensor current scaling is backed by grouped current-scaling quantization (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling, which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0). - Every other quantization recipe (fp8 delayed scaling, fp8 block scaling, ...) - falls back to the legacy flow because the corresponding grouped quantization kernels are - missing. + * FP8 block scaling runs the graph-safe flow on Hopper (CC 9.0) only, including + power-of-2 scales. On other architectures it falls back to the split-quantize + flow (emulated via MXFP8 GEMM on Blackwell). + Every other quantization recipe (fp8 delayed scaling, ...) falls back to the + legacy flow because the corresponding grouped quantization kernels are missing. * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0) -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it. * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it @@ -793,6 +799,13 @@ def _is_graph_safe_path_supported( # path; the compute-capability range was already checked above. if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): return True + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM + # scale modes are Hopper-only. The graph-safe flow is default-on here + # (no env-var opt-in like the module layer), so on other architectures + # fall back to the split-quantize flow like the other unsupported + # recipe/arch combinations (on Blackwell it emulates via MXFP8 GEMM). + return get_device_compute_capability() < (10, 0) # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. if not (10, 0) <= get_device_compute_capability() <= (11, 0): return False @@ -1612,11 +1625,12 @@ def _fuser_backward_grouped_tensor( ) grad_output_quantizer.optimize_for_gemm = True - if ( - has_bias - and not self._scale_bias - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ): + # The grouped FP8 block-scaling bgrad kernel computes dbias in the rowwise + # pass, so the fusion needs rowwise output (i.e. dgrad required). + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad + ) + if has_bias and not self._scale_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, num_groups, split_sizes ) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 849c900f95..2eee9e3497 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -183,6 +183,9 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): num_extra_inputs: int = 0 # Number of extra tensor outputs num_extra_outputs: int = 0 + # Whether this operation supports the FP8 block-scaling recipe. Most basic + # ops have not been audited for it, so they opt in explicitly. + supports_float8_block_scaling: bool = False def __init__(self) -> None: super().__init__() @@ -275,9 +278,9 @@ def reset_recipe_state( if num_quantizers == 0: continue - if recipe.float8_block_scaling(): + if recipe.float8_block_scaling() and not self.supports_float8_block_scaling: raise NotImplementedError( - "Fusible operations do not support FP8 block scaling recipe" + f"{self.__class__.__name__} does not support FP8 block scaling recipe" ) # Construct quantization recipe state From 98000059868eaf06f1880af6bcb1fe5448e7f79e Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 07:24:17 +0000 Subject: [PATCH 4/7] [PyTorch] Use persistent workspaces in grouped-tensor GEMM general_grouped_gemm_for_grouped_tensor allocated its setup workspace (the cuBLAS per-group pointer/dimension arrays) and its cuBLAS workspace with per-call torch.empty. Under make_graphed_callables the forward and backward graphs share one capture memory pool, and a per-call allocation's block returns to that pool as soon as the Python reference dies, so blocks alias across the two graphs and captured kernels from one graph overwrite the GEMM metadata the other graph reads at replay. Observed as allocation-history-dependent failures in the ops-layer GroupedLinear cuda-graph test: capture-time cublasLtMatmulAlgoGetHeuristic NOT_SUPPORTED errors and corrupted wgrad outputs. This is also the likely mechanism behind the FP8 block-scaling wgrad corruption under CUDA graphs previously observed on Hopper and attributed to cuBLAS. Cache the setup workspace per (device, group size) and reuse the cached per-device cuBLAS workspace from the non-grouped path; consecutive GEMMs reusing one workspace are ordered by the stream. Signed-off-by: Alp Dener --- .../pytorch/cpp_extensions/gemm.py | 32 +++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f0f43515dc..6d713796f4 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -436,6 +436,24 @@ def _get_fp32_zeros_tensor(num_tensors: int, device: torch.device) -> torch.Tens return torch.zeros(num_tensors, dtype=torch.float32, device=device) +@functools.lru_cache(maxsize=None) +def _get_grouped_gemm_setup_workspace(device: int, num_tensors: int) -> torch.Tensor: + """Persistent setup workspace (per-group pointer/dim arrays) for grouped-tensor GEMM. + + This must not be allocated per call: under CUDA-graph capture a per-call + allocation's block returns to the shared capture pool as soon as the Python + reference dies, so the forward and backward graphs can alias the same block + and the captured GEMM's pointer/dimension arrays get overwritten at replay. + Consecutive GEMMs reusing one workspace are ordered by the stream, matching + how the non-grouped path shares its cached cuBLAS workspace. + """ + return torch.empty( + get_grouped_gemm_setup_workspace_size(num_tensors), + dtype=torch.uint8, + device=device, + ) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -525,16 +543,10 @@ def _is_fp8_blockwise(operand) -> bool: if not alpha.is_cuda or not beta.is_cuda: raise ValueError("alpha and beta must be CUDA tensors.") - workspace_setup = torch.empty( - get_grouped_gemm_setup_workspace_size(num_tensors), - dtype=torch.uint8, - device=device, - ) - workspace_cublas = torch.empty( - get_cublas_workspace_size_bytes(), - dtype=torch.uint8, - device=device, - ) + workspace_setup = _get_grouped_gemm_setup_workspace(device.index, num_tensors) + # Shares the cached per-device workspace with the non-grouped GEMM path + # (persistent for the same CUDA-graph reason as the setup workspace). + workspace_cublas = get_cublas_workspace(device.index, False, False) sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) From b35924ccc3280d8789d7ba2b678f0731f9e6a20d Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 19:45:55 +0000 Subject: [PATCH 5/7] [PyTorch] Fix grouped FP8 block-scaling CUDA-graph deadlock via per-role cuBLAS workspaces The grouped-tensor GEMM path shared one persistent cuBLAS workspace across all grouped matmuls. cuBLAS's grouped GEMM keeps a grid-synchronization flag in the first bytes of that workspace and zeros it (via a captured memset) before each matmul. When the dgrad and wgrad grouped matmuls of a GroupedLinear backward share one workspace inside a replayed CUDA graph, that flag is aliased between the two matmuls; on the second graph replay the second matmul's cooperative kernel deadlocks with cuBLAS 13.6 (and corrupts the last expert's wgrad on cuBLAS < 13.6). The two matmuls are strictly stream-ordered (single stream, all-DEFAULT graph edges, no programmatic dependent launch), so this is shared-workspace reuse, not concurrent co-scheduling. Give dgrad/forward (slot 0) and wgrad (slot 1) distinct persistent cuBLAS workspaces, dedicated to the grouped path. Each slot remains a single persistent allocation, so CUDA-graph capture safety is preserved. Also drop the cuBLAS-version gate that skipped the FP8 block-scaling GroupedLinear CUDA-graph test, so it now exercises the fix on all supported cuBLAS versions. Signed-off-by: Alp Dener --- tests/pytorch/test_grouped_mlp.py | 19 ++++--------- .../pytorch/cpp_extensions/gemm.py | 27 ++++++++++++++++--- 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 9c965fbe01..a79d8eeb9a 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -503,20 +503,11 @@ def test_grouped_linear_cuda_graph_safe( ) # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path, # FP8 block scaling is Hopper-only, and MXFP8/NVFP4 grouped quantization kernels - # require Blackwell (SM100+). - if quantization == "fp8_block_scaling": - # Under investigation: the graph-replayed FP8 block-scaling wgrad for the last - # expert diverges between replays in this test (the first replay is correct - # against a float64 reference; the reference step's replay is not). The effect - # depends on the process allocation history and has not been reproduced in a - # standalone equivalent of this flow. Note make_graphed_callables replaces - # module.forward in place, so the "eager" reference step here replays the same - # graphs. The eager ops path is covered by test_grouped_linear and graph capture - # by the module-path test in test_grouped_linear.py. - pytest.skip( - "FP8 block-scaling grouped wgrad diverges between graph replays;" - " under investigation." - ) + # require Blackwell (SM100+). FP8 block-scaling grouped dgrad/wgrad previously + # shared one persistent cuBLAS workspace, which deadlocked (cuBLAS 13.6) or + # corrupted the last expert's wgrad (cuBLAS < 13.6) under CUDA-graph replay; the + # fix in _get_grouped_cublas_workspace (cpp_extensions/gemm.py) gives them distinct + # workspaces, so this test now runs on all supported cuBLAS versions. requires_blackwell = quantization is not None and quantization not in ( "fp8_current_scaling", "fp8_block_scaling", diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 6d713796f4..828875410c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -454,6 +454,27 @@ def _get_grouped_gemm_setup_workspace(device: int, num_tensors: int) -> torch.Te ) +@functools.lru_cache(maxsize=None) +def _get_grouped_cublas_workspace(device: int, slot: int) -> torch.Tensor: + """Persistent cuBLAS workspace for the grouped-tensor GEMM path, one per ``slot``. + + dgrad (``slot`` 0) and wgrad (``slot`` 1) must NOT share a single cuBLAS + workspace. The grouped GEMM keeps a grid-synchronization flag in the first + bytes of the workspace (cuBLAS emits a memset that zeros it before each + matmul). When two grouped matmuls in a replayed CUDA graph share one + workspace, that flag is aliased, and the second matmul's cooperative kernel + deadlocks on its second replay with cuBLAS 13.6 (and corrupts the last + expert's wgrad on cuBLAS < 13.6) -- the two matmuls are strictly stream- + ordered, so this is shared-workspace reuse, not concurrent co-scheduling. + Giving each its own persistent buffer removes the hazard; each slot is a + single persistent allocation, so CUDA-graph capture safety is preserved. + These are dedicated to the grouped path and are intentionally not shared with + the non-grouped GEMM workspace. + """ + assert slot in (0, 1), "grouped cuBLAS workspace slot must be 0 (dgrad/fwd) or 1 (wgrad)" + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -544,9 +565,9 @@ def _is_fp8_blockwise(operand) -> bool: raise ValueError("alpha and beta must be CUDA tensors.") workspace_setup = _get_grouped_gemm_setup_workspace(device.index, num_tensors) - # Shares the cached per-device workspace with the non-grouped GEMM path - # (persistent for the same CUDA-graph reason as the setup workspace). - workspace_cublas = get_cublas_workspace(device.index, False, False) + # dgrad and wgrad must use distinct persistent cuBLAS workspaces: sharing one + # deadlocks under CUDA-graph replay on cuBLAS 13.6 (see _get_grouped_cublas_workspace). + workspace_cublas = _get_grouped_cublas_workspace(device.index, 1 if is_discrete_out else 0) sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) From a65fbc0f1476a367ef659574c695ab78bd65b484 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 20:36:17 +0000 Subject: [PATCH 6/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../group_quantize_fp8_blockwise.cuh | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index e43c796d79..33aadb5bfa 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -379,8 +379,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma for (int w = 1; w < kNumWarps; ++w) { block_amax = fmaxf(block_amax, warp_amaxes[w]); } - const CType scale = - compute_scale_from_types(block_amax, epsilon, pow_2_scales); + const CType scale = compute_scale_from_types(block_amax, epsilon, pow_2_scales); // The 2D colwise per-expert scale offset requires a CTA-cooperative prefix // sum in the VARYING_FIRST_DIM case, so compute it across all threads before @@ -476,8 +475,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t R_total, const float epsilon, - const bool pow_2_scales, - const float* __restrict__ noop_ptr) { + const bool pow_2_scales, const float* __restrict__ noop_ptr) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -532,8 +530,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) CType amax = compute_row_amax(in_vec[it]); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, pow_2_scales); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; if (thr_col == 0 && r_global < R_total) { // Per-expert layout: (blocks_X, roundup(M_t, 4)). Compute expert base @@ -658,8 +655,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType amax = compute_row_amax(in_vec); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, pow_2_scales); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t r_global = global_row_base + row_local; @@ -735,8 +731,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke } amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, pow_2_scales); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t c_global = global_col_base + col_local; From 54e0f2e6443f128079967480813bf318794c3f3f Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 2 Jul 2026 20:55:30 +0000 Subject: [PATCH 7/7] [PyTorch] Address review: document split-accumulator override, fix stale dbias comment - general_grouped_gemm_for_grouped_tensor: expand the comment to state that the fused grouped FP8 block-scaling GEMM forces use_split_accumulator=True and intentionally overrides the caller-supplied value, consistent with the Float8BlockScaling recipe (which fixes it True for fprop/dgrad/wgrad). - Float8BlockScaling recipe docstring: document that FP8 block scaling always uses split accumulation and that the fused grouped GEMM path ignores any caller- or recipe-supplied use_split_accumulator value. - GroupedLinear ops backward: correct the stale "BF16/FP16 path" comment; that branch also handles quantized paths where bgrad fusion did not apply (e.g. FP8 block scaling without a dgrad pass). Signed-off-by: Alp Dener --- transformer_engine/common/recipe/__init__.py | 6 ++++++ transformer_engine/pytorch/cpp_extensions/gemm.py | 5 ++++- transformer_engine/pytorch/ops/basic/grouped_linear.py | 4 +++- 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8a03f2f51a..8c209ace15 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -404,6 +404,12 @@ class Float8BlockScaling(Recipe): NOTE: To relax the default constraint that scales be powers of 2, set env variable NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 to override it for the recipe defaults. + NOTE: FP8 block scaling requires split accumulation for numerical accuracy, so + ``fp8_gemm_fprop``/``fp8_gemm_dgrad``/``fp8_gemm_wgrad`` all fix + ``use_split_accumulator=True`` (enforced in ``__post_init__``). The fused grouped + GEMM path (GroupedLinear) always uses split accumulation for FP8 block scaling and + ignores any caller- or recipe-supplied ``use_split_accumulator`` value. + Parameters ---------- fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.E4M3 diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 828875410c..0a78aa894c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -522,7 +522,10 @@ def _is_fp8_blockwise(operand) -> bool: return False if _is_fp8_blockwise(A) or _is_fp8_blockwise(B): - # FP8 block-scaling requires split accumulator + # The fused grouped FP8 block-scaling GEMM only supports split accumulation, + # so force it on and intentionally override any caller-supplied value. This + # matches the Float8BlockScaling recipe, which fixes use_split_accumulator=True + # for all of fprop/dgrad/wgrad, so no user-configurable setting is discarded. use_split_accumulator = True if is_discrete_out: diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 918d2bedef..a0a90e76c9 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1665,7 +1665,9 @@ def _fuser_backward_grouped_tensor( offsets=base_split_offsets, ) elif dbias_packed is None: - # BF16/FP16 path + # dbias was not fused into the grad-output quantization: either + # unquantized BF16/FP16 compute, or a quantized path where bgrad + # fusion did not apply (e.g. FP8 block scaling without a dgrad pass). dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups) if self.single_grouped_bias: final_bias_grads = [dbias_packed.to(dtype=dtype)]