From 9d42b33a545b08a83f08084a64706bd943b51bb2 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Wed, 22 Jul 2026 03:32:33 -0700 Subject: [PATCH 1/2] [common][PyTorch] NVFP4: enable row-scaled transpose quantization for backward The #2931 row-scaled NVFP4 path only produced the rowwise forward activation; its columnwise/transpose output was rejected. That made the per-token activation unusable in the backward wgrad GEMM, so row-scaled training had to fall back to a dequantized/high-precision backward. This change extends the existing row-scaled path to also emit the columnwise (transpose) direction, so a training Linear with row_scaled_activation=True now quantizes the forward activation row-scaled in both directions and the wgrad GEMM consumes the row-scaled transpose directly. It is a minimal extension of #2931 (no new CUTLASS kernels, no grouped path, no RHT/4over6 transpose fusion). Signed-off-by: Cael Ling --- .../cpp/operator/test_cast_nvfp4_transpose.cu | 154 ++++++++++++++++++ tests/cpp/test_common.cu | 8 +- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 60 +++++++ .../pytorch/nvfp4/test_nvfp4_module_exact.py | 113 ++++++++++++- .../nvfp4/test_nvfp4_quantize_exact.py | 7 +- .../common/cast/dispatch/quantize.cuh | 38 ++++- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 81 +++++++++ .../quantize_transpose_nvfp4_tuned_1D.cuh | 32 ++-- .../pytorch/cpp_extensions/gemm.py | 74 +++++---- transformer_engine/pytorch/csrc/quantizer.cpp | 16 +- .../custom_recipes/quantization_ref_nvfp4.py | 31 +++- 11 files changed, 545 insertions(+), 69 deletions(-) diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index c23a1c23f9..2d9b3073a2 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -1143,6 +1143,138 @@ std::vector Activation_types = { ActivationType::Identity }; +// Element-level FP4 code differences between two compact NVFP4 buffers of the +// same logical (rows, cols) shape (cols must be even; FP4 is packed 2/byte). +inline size_t count_nvfp4_code_mismatches(const fp4e2m1* test_data, const fp4e2m1* ref_data, + int rows, int cols) { + size_t mismatches = 0; + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < cols; j += 2) { + const int idx = i * cols + j; + const double2 t = cvt_fp4x2_to_double2(*reinterpret_cast(&test_data[idx / 2])); + const double2 r = cvt_fp4x2_to_double2(*reinterpret_cast(&ref_data[idx / 2])); + if (t.x != r.x) ++mismatches; + if (t.y != r.y) ++mismatches; + } + } + return mismatches; +} + +// Row-scaled transpose NVFP4 cast: emits rowwise (per-row amax) and columnwise +// (per-col amax) NVFP4 outputs in one pass. Cross-validated against the merged +// row-scaled 1D kernel (PR #2931, NVFP4ScalingMode::RowScaled1D): rowwise vs +// RowScaled1D(input), columnwise vs RowScaled1D(input^T). Amaxes and FP8 block +// scales must match bitwise; FP4 codes may differ only at rounding-midpoint +// ties (< 0.1%). bf16 input, rows/cols multiples of 128. +template +void performTestRowScaledTranspose(const std::vector& shape) { + using namespace test; + + const DType itype = TypeInfo::dtype; + const DType otype = DType::kFloat4E2M1; + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + + Tensor input("input", shape, itype); + fillCase(&input, InputsFillCase::uniform); + + // System under test: row-scaled NVFP4 with both directions requested in one + // call. The transpose kernel is selected by the generic nvte_quantize_v2 + // dispatch when a row-scaled tensor (set_row_scaled_nvfp4) also allocates a + // columnwise output -- no dedicated config flag. + Tensor output("output", shape, otype, /*rowwise=*/true, /*columnwise=*/true, + NVTE_NVFP4_1D_SCALING); + // Marking the tensor row-scaled with a columnwise output allocated selects the + // transpose kernel and sizes the per-row (rowwise) and per-col (columnwise) + // amax vectors (default is a single scalar amax). + output.set_row_scaled_nvfp4(true); + QuantizationConfigWrapper quant_config; + quant_config.set_stochastic_rounding(false); + nvte_quantize_v2(input.data(), output.data(), quant_config, 0); + + // Reference (rowwise direction): trusted row-scaled 1D kernel on the input. + Tensor ref_row("ref_row", shape, otype, /*rowwise=*/true, /*columnwise=*/false, + NVTE_NVFP4_1D_SCALING); + ref_row.set_row_scaled_nvfp4(true); + QuantizationConfigWrapper ref_config; + ref_config.set_stochastic_rounding(false); + nvte_quantize_v2(input.data(), ref_row.data(), ref_config, 0); + + // Reference (columnwise direction): trusted row-scaled 1D kernel on input^T. + input.to_cpu(); + std::vector input_t_host = + create_transpose(input.rowwise_cpu_dptr(), rows, cols); + const std::vector shape_t = {cols, rows}; + Tensor input_t("input_t", shape_t, itype); + std::copy(input_t_host.begin(), input_t_host.end(), input_t.rowwise_cpu_dptr()); + input_t.from_cpu(); + Tensor ref_col("ref_col", shape_t, otype, /*rowwise=*/true, /*columnwise=*/false, + NVTE_NVFP4_1D_SCALING); + ref_col.set_row_scaled_nvfp4(true); + nvte_quantize_v2(input_t.data(), ref_col.data(), ref_config, 0); + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + output.to_cpu(); + ref_row.to_cpu(); + ref_col.to_cpu(); + + // FP4 codes must match the trusted kernel except at rounding-midpoint ties. + // Cap the tolerated disagreement well below what a real bug would produce. + constexpr double kMaxCodeMismatchRate = 5e-3; // 0.5% (observed < 0.1%) + const size_t row_mismatches = count_nvfp4_code_mismatches( + output.rowwise_cpu_dptr(), ref_row.rowwise_cpu_dptr(), + static_cast(rows), static_cast(cols)); + EXPECT_LT(static_cast(row_mismatches) / static_cast(rows * cols), + kMaxCodeMismatchRate) + << "rowwise FP4 disagreement " << row_mismatches << "/" << (rows * cols); + const size_t col_mismatches = count_nvfp4_code_mismatches( + output.columnwise_cpu_dptr(), ref_col.rowwise_cpu_dptr(), + static_cast(cols), static_cast(rows)); + EXPECT_LT(static_cast(col_mismatches) / static_cast(cols * rows), + kMaxCodeMismatchRate) + << "columnwise FP4 disagreement " << col_mismatches << "/" << (cols * rows); + + // FP8 e4m3 block scale factors must match (compact layout on both sides). + const std::array sd = get_scale_tensor_dims(rows, cols, 1, 16); + const std::array sd_t = get_scale_tensor_dims(cols, rows, 1, 16); + size_t scale_mismatches = 0; + compare_scaling_factors("rowwise_scales", + output.rowwise_cpu_scale_inv_ptr(), + ref_row.rowwise_cpu_scale_inv_ptr(), + sd[0], sd[1], sd[3], scale_mismatches); + compare_scaling_factors("columnwise_scales", + output.columnwise_cpu_scale_inv_ptr(), + ref_col.rowwise_cpu_scale_inv_ptr(), + sd_t[0], sd_t[1], sd_t[3], scale_mismatches); + ASSERT_EQ(scale_mismatches, 0u); + + // Per-row / per-col amaxes must match the reference amaxes exactly. + ASSERT_EQ(output.rowwise_amax_size(), rows); + const float* row_amax = output.cpu_rowwise_amax_ptr(); + const float* ref_row_amax = ref_row.cpu_rowwise_amax_ptr(); + for (size_t i = 0; i < rows; ++i) { + ASSERT_EQ(row_amax[i], ref_row_amax[i]) << "rowwise amax mismatch at row " << i; + } + const float* col_amax = output.cpu_columnwise_amax_ptr(); + const float* ref_col_amax = ref_col.cpu_rowwise_amax_ptr(); + for (size_t j = 0; j < cols; ++j) { + ASSERT_EQ(col_amax[j], ref_col_amax[j]) << "columnwise amax mismatch at col " << j; + } +} + +// Row-scaled transpose requires 128-aligned rows and cols (bf16 only). +std::vector> row_scaled_transpose_dims = { + {128, 128}, + {256, 256}, + {128, 256}, + {256, 512}, + {512, 512}, + {384, 1024}, + {2048, 256}, +}; + } // namespace class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam @@ -1323,3 +1455,25 @@ INSTANTIATE_TEST_SUITE_P( } return name; }); + +class CastNVFP4RowScaledTransposeTestSuite : public ::testing::TestWithParam> {}; + +TEST_P(CastNVFP4RowScaledTransposeTestSuite, MatchesRowScaledReference) { + // The row-scaled transpose NVFP4 cast kernel requires Blackwell. + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performTestRowScaledTranspose(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + CastNVFP4RowScaledTransposeTestSuite, + ::testing::ValuesIn(row_scaled_transpose_dims), + [](const testing::TestParamInfo& info) { + std::string name; + for (const auto& s : info.param) { + name += "X" + std::to_string(s); + } + return name; + }); diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index fc41d44720..e1468ef981 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -420,13 +420,17 @@ void Tensor::set_row_scaled_nvfp4(bool row_scaled_nvfp4) { // Update amax tensor if (row_scaled_nvfp4) { - // Row-scaled NVFP4 has amax matching number of rows NVTE_CHECK(rowwise_, "Row-scaled NVFP4 requires row-wise data."); - NVTE_CHECK(!columnwise_, "Row-scaled NVFP4 does not support column-wise data."); auto shape = tensor_.shape(); const size_t rows = product(shape, 0, shape.ndim - 1); + const size_t cols = shape.data[shape.ndim - 1]; amax_rowwise_.emplace(rows, DType::kFloat32); tensor_.set_amax(amax_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{rows}); + if (columnwise_) { + amax_columnwise_.emplace(cols, DType::kFloat32); + tensor_.set_columnwise_amax(amax_columnwise_->gpu_buffer(), DType::kFloat32, + std::vector{cols}); + } } else { // Tensor-scaled NVFP4 has single amax if (rowwise_) { diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 63e69beafc..67dbe1d1d5 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -11,6 +11,7 @@ from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -417,6 +418,65 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( torch.testing.assert_close(y_row_scaled, y_emulated, atol=3.0517578125e-5, rtol=0.0) +def _dequantize_nvfp4_usage( + tensor: NVFP4TensorStorage, + *, + columnwise: bool, +) -> torch.Tensor: + """Dequantize one independently-quantized tensor orientation.""" + if not columnwise: + return tensor.dequantize(dtype=torch.float32) + + metadata = tensor.get_metadata() + metadata["rowwise_data"] = metadata["columnwise_data"] + metadata["rowwise_scale_inv"] = metadata["columnwise_scale_inv"] + metadata["amax_rowwise"] = metadata["amax_columnwise"] + metadata["columnwise_data"] = None + metadata["columnwise_scale_inv"] = None + metadata["amax_columnwise"] = None + return NVFP4TensorStorage(**metadata).dequantize(dtype=torch.float32) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "layout,a_shape,b_shape", + [ + ("TN", (64, 64), (96, 64)), + ("NN", (64, 96), (128, 64)), + ("NT", (128, 64), (128, 96)), + ], +) +def test_nvfp4_bilateral_row_scaled_gemm_matches_dequantized( + layout: str, + a_shape: tuple[int, int], + b_shape: tuple[int, int], +) -> None: + """Check per-tensor GEMM plus bilateral FP32 post-scaling.""" + torch.manual_seed(41) + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + a = torch.randn(a_shape, dtype=torch.bfloat16, device="cuda") + b = torch.randn(b_shape, dtype=torch.bfloat16, device="cuda") + a_nvfp4 = quantizer(a) + b_nvfp4 = quantizer(b) + + actual = general_gemm(a_nvfp4, b_nvfp4, out_dtype=torch.float32, layout=layout)[0] + transa, transb = layout[0] == "T", layout[1] == "T" + a_dequant = _dequantize_nvfp4_usage(a_nvfp4, columnwise=not transa) + b_dequant = _dequantize_nvfp4_usage(b_nvfp4, columnwise=transb) + expected = b_dequant @ a_dequant.T + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, K, N", diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index b57b78eb13..66ac0181e3 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -53,7 +53,23 @@ def nvfp4_rht_and_2d_quantization(): return nvfp4_recipe @staticmethod - def nvfp4_recipe_to_test(with_rht: bool = False, with_2d_quantization: bool = False): + def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling() + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + # Emit row-scaled (per-token) NVFP4 for the forward activation. In a + # training Linear this also drives the row-scaled columnwise (transpose) + # of the activation that the wgrad GEMM consumes in the backward pass. + nvfp4_recipe.row_scaled_activation = True + return nvfp4_recipe + + @staticmethod + def nvfp4_recipe_to_test( + with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False + ): + if row_scaled: + return GetRecipes.nvfp4_row_scaled() if with_rht and with_2d_quantization: return GetRecipes.nvfp4_rht_and_2d_quantization() elif with_rht: @@ -64,7 +80,9 @@ def nvfp4_recipe_to_test(with_rht: bool = False, with_2d_quantization: bool = Fa return GetRecipes.nvfp4_vanilla() -def get_nvfp4_quantizer_factory(with_rht: bool = False, with_2d_quantization: bool = False): +def get_nvfp4_quantizer_factory( + with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False +): """ Create a quantizer factory for NVFP4 reference implementation. @@ -96,7 +114,15 @@ def factory(role): if role is None: return _default_quantizer() if role.tensor_type == "input": - return _default_quantizer() + # Only the forward activation is row-scaled, mirroring the production + # wiring in quantization.py (mode == "forward" and tensor_type != "weight"). + return quantization_ref_nvfp4.NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + quant_tile_shape=(1, 16), + pow_2_scales=False, + with_rht=with_rht, + row_scaled_nvfp4=row_scaled, + ) if role.tensor_type == "weight": return quantization_ref_nvfp4.NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, @@ -117,6 +143,23 @@ def reset_rng_states(): torch.cuda.manual_seed(seed) +def assert_close_relative(name: str, native, ref, step: int, rel_tol: float) -> None: + """Compare two tensors by relative Frobenius-norm error. + + Used for the row-scaled NVFP4 path, where native (cuBLAS + post-scale) and the + dequantized reference (torch matmul) share the same 4-bit quantization error but + differ at the kernel/impl level (accumulation order, post-scaling). The relative + error therefore stays small even though the absolute NVFP4 error is large. + """ + if native is None or ref is None: + return + ref_f = ref.detach().float() + diff = (native.detach().float() - ref_f).norm().item() + denom = ref_f.norm().item() + rel = diff / denom if denom > 0 else diff + assert rel <= rel_tol, f"{name} relative error {rel:.5f} > {rel_tol} at step {step}" + + def check_nvfp4_module_versus_reference( module_class, in_features: int, @@ -126,6 +169,8 @@ def check_nvfp4_module_versus_reference( num_steps: int = 1, with_rht: bool = False, with_2d_quantization: bool = False, + row_scaled: bool = False, + rel_tol: float = None, ): """ Compare native NVFP4 module against reference implementation. @@ -137,6 +182,10 @@ def check_nvfp4_module_versus_reference( bias: Whether to use bias x_dtype: Input tensor dtype num_steps: Number of forward/backward steps to test + row_scaled: Enable row-scaled (per-token) NVFP4 activation quantization, + exercising the row-scaled columnwise (transpose) path consumed by wgrad. + rel_tol: If set, compare tensors by relative Frobenius-norm error instead of + the default bitwise-tight ``assert_close`` (used for the row-scaled path). """ device = "cuda" batch_size = 32 @@ -203,8 +252,8 @@ def check_nvfp4_module_versus_reference( ref_module.layer_norm_bias.copy_(native_module.layer_norm_bias) # Create recipes for native and reference implementations - nvfp4_recipe = GetRecipes.nvfp4_recipe_to_test(with_rht, with_2d_quantization) - nvfp4_ref_factory = get_nvfp4_quantizer_factory(with_rht, with_2d_quantization) + nvfp4_recipe = GetRecipes.nvfp4_recipe_to_test(with_rht, with_2d_quantization, row_scaled) + nvfp4_ref_factory = get_nvfp4_quantizer_factory(with_rht, with_2d_quantization, row_scaled) nvfp4_ref_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_factory) # Training loop comparison @@ -279,6 +328,22 @@ def check_nvfp4_module_versus_reference( native_out = native_outputs[step] ref_out = ref_outputs[step] + if rel_tol is not None: + # Row-scaled path: native cuBLAS + post-scale vs dequantized torch + # reference share the same 4-bit error, so compare by relative norm. + assert_close_relative("Output", native_out["output"], ref_out["output"], step, rel_tol) + assert_close_relative( + "Input gradient", native_out["input_grad"], ref_out["input_grad"], step, rel_tol + ) + assert_close_relative( + "Weight gradient", native_out["weight_grad"], ref_out["weight_grad"], step, rel_tol + ) + if bias: + assert_close_relative( + "Bias gradient", native_out["bias_grad"], ref_out["bias_grad"], step, rel_tol + ) + continue + # Compare outputs torch.testing.assert_close( native_out["output"], @@ -367,6 +432,44 @@ def test_nvfp4_linear_versus_reference( ) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "in_features, out_features", + [ + (128, 256), + (256, 128), + (512, 512), + (768, 3072), + ], +) +@pytest.mark.parametrize("num_steps", [1, 3], ids=["single_step", "multi_step"]) +def test_nvfp4_linear_row_scaled_versus_reference( + in_features: int, + out_features: int, + num_steps: int, +): + """End-to-end row-scaled (per-token) NVFP4 Linear forward + backward. + + Exercises the row-scaled columnwise (transpose) activation quantization that + this PR unblocks: the forward activation is quantized row-scaled with both + rowwise (fprop) and columnwise (wgrad) directions, so ``backward()`` drives the + new transpose path through the wgrad GEMM. Compared against the dequantized + reference quantizer by relative norm (both share the same 4-bit error). + """ + check_nvfp4_module_versus_reference( + module_class=te.Linear, + in_features=in_features, + out_features=out_features, + bias=False, + x_dtype=torch.bfloat16, + num_steps=num_steps, + with_rht=False, + with_2d_quantization=False, + row_scaled=True, + rel_tol=2e-2, + ) + + def check_nvfp4_layernorm_linear_versus_reference( in_features: int, out_features: int, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 85953687a7..4a3d4d3413 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -87,7 +87,12 @@ def maybe_skip_row_scaled_unsupported_quantization( if not row_scaled_nvfp4: return if return_transpose: - pytest.skip("Row-scaled NVFP4 does not support columnwise usage") + if use_4over6: + pytest.skip("Row-scaled NVFP4 transpose does not support 4over6 mode") + if x_dtype != torch.bfloat16 or M is None or N is None or M % 32 != 0 or N % 32 != 0: + pytest.skip( + "Row-scaled NVFP4 transpose requires BF16 input and dimensions divisible by 32" + ) if with_2d_quantization: pytest.skip("Row-scaled NVFP4 does not support 2D quantization") diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 033d464bcf..87f1ba68b2 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -111,9 +111,20 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, if (row_scaled_nvfp4) { NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); - NVTE_CHECK(!output_tensor->has_columnwise_data(), - "Row-scaled NVFP4 quantization does not produce columnwise output."); + NVTE_CHECK( + !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()), + "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 " + "kernel does not consume the per-row/per-column amaxes, so the columnwise output would " + "be incorrect."); + NVTE_CHECK( + !output_tensor->has_columnwise_data() || + (dtype == DType::kBFloat16 && rows % 32 == 0 && cols % 32 == 0), + "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions that are " + "multiples of 32."); nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream); + if (output_tensor->has_columnwise_data()) { + nvfp4::compute_columnwise_amax(*input_tensor, noop_tensor, output_tensor, stream); + } } // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and // both-directions keep their existing routing. Columnwise-only 1D and non-bf16 fall back to @@ -268,13 +279,30 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens // Choose kernel const auto [rows, cols] = grad_tensor->flat_2d_dims(); auto dtype = grad_tensor->dtype(); + const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, "Non-4over6 NVFP4 quantization requires E4M3 max 448."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); - NVTE_CHECK(!output_tensor->row_scaled_nvfp4, - "Backward NVFP4 quantization does not support row-scaled outputs."); + if (row_scaled_nvfp4) { + NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK( + !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()), + "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 " + "kernel does not consume the per-row/per-column amaxes, so the columnwise output would " + "be incorrect."); + NVTE_CHECK( + !output_tensor->has_columnwise_data() || + (dtype == DType::kBFloat16 && rows % 32 == 0 && cols % 32 == 0), + "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions that are " + "multiples of 32."); + nvfp4::compute_rowwise_amax(*grad_tensor, noop_tensor, output_tensor, stream); + if (output_tensor->has_columnwise_data()) { + nvfp4::compute_columnwise_amax(*grad_tensor, noop_tensor, output_tensor, stream); + } + } // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and // both-directions keep their existing routing. Columnwise-only 1D and non-bf16 fall back to // quantize_transpose_vector_blockwise_fp4. @@ -314,7 +342,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, - /*row_scaled_nvfp4=*/false, + /*row_scaled_nvfp4=*/row_scaled_nvfp4, /*noop_tensor=*/noop_tensor->data, /*stream=*/stream); } diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 3a34c76de8..e3971329d2 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -99,6 +99,39 @@ __launch_bounds__(BLOCK_SIZE) #endif } +template +__global__ void +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__launch_bounds__(BLOCK_SIZE) +#endif + compute_columnwise_amax_kernel(const int num_rows, const int num_cols, + const IType *__restrict__ input, + float *__restrict__ output_columnwise_amax, + const float *__restrict__ noop) { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) + NVTE_DEVICE_ERROR("SM 10.0+ is required."); +#else + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int col_idx = blockIdx.x; + if (col_idx >= num_cols) return; + + float thread_max = 0.0f; + for (int row_idx = threadIdx.x; row_idx < num_rows; row_idx += BLOCK_SIZE) { + thread_max = + fmaxf(thread_max, fabsf(static_cast(input[row_idx * num_cols + col_idx]))); + } + const float col_amax = + reduce_max(thread_max, threadIdx.x / THREADS_PER_WARP); + + if (threadIdx.x == 0) { + output_columnwise_amax[col_idx] = col_amax; + } +#endif +} + template void launch_compute_rowwise_amax(const int num_rows, const int num_cols, const IType *input, float *output_rowwise_amax, cudaStream_t stream, @@ -113,6 +146,20 @@ void launch_compute_rowwise_amax(const int num_rows, const int num_cols, const I NVTE_CHECK_CUDA(cudaGetLastError()); } +template +void launch_compute_columnwise_amax(const int num_rows, const int num_cols, const IType *input, + float *output_columnwise_amax, cudaStream_t stream, + const float *noop = nullptr) { + if (num_rows == 0 || num_cols == 0) return; + + dim3 grid(num_cols); + dim3 block(ROWWISE_AMAX_BLOCK_SIZE); + + compute_columnwise_amax_kernel + <<>>(num_rows, num_cols, input, output_columnwise_amax, noop); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + #endif // FP4_TYPE_SUPPORTED } // namespace rowwise_amax_kernel @@ -155,6 +202,40 @@ inline void compute_rowwise_amax(const Tensor &input, const Tensor *noop, Tensor #endif // FP4_TYPE_SUPPORTED } +inline void compute_columnwise_amax(const Tensor &input, const Tensor *noop, Tensor *output, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace rowwise_amax_kernel; + + const auto [rows, cols] = input.flat_2d_dims(); + auto *amax_ptr = reinterpret_cast(output->columnwise_amax.dptr); + NVTE_CHECK(amax_ptr != nullptr, "Row-scaled columnwise amax tensor must be allocated."); + NVTE_CHECK(output->columnwise_amax.numel() == cols, "Row-scaled columnwise amax must have ", cols, + " entries, got ", output->columnwise_amax.shape, "."); + + const auto *noop_ptr = reinterpret_cast(noop->data.dptr); + if (input.dtype() == DType::kBFloat16) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_columnwise_amax<__nv_bfloat16>( + static_cast(rows), static_cast(cols), input_ptr, amax_ptr, stream, noop_ptr); + } else if (input.dtype() == DType::kFloat16) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_columnwise_amax(static_cast(rows), static_cast(cols), input_ptr, + amax_ptr, stream, noop_ptr); + } else if (input.dtype() == DType::kFloat32) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_columnwise_amax(static_cast(rows), static_cast(cols), input_ptr, + amax_ptr, stream, noop_ptr); + } else { + NVTE_ERROR( + "Unsupported input dtype for row-scaled NVFP4 quantization. " + "Expected BFloat16, Float16, or Float32."); + } +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + namespace quantize_transpose_kernel { using namespace quantization_and_transposition_SF; diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 8f37229fd5..637abe7a0f 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -184,14 +184,16 @@ compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const f return static_cast(scale_rcp); } -template +template __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_tr_ptr, nvfp4_scale_t *__restrict__ sSFcolwise_ptr, const float S_enc_colwise, const int stage_Y, const int stage_X, const int buff_in, - const int buff_out_tr, RNG_t &rng, - uint4 &random_uint4, int &rnd_idx) { + const int buff_out_tr, + const float *amax_colwise_ptr, + const size_t col_offset, const size_t cols, + RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn2x = *reinterpret_cast(sIn_ptr); @@ -231,13 +233,23 @@ __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_pt static_cast(__habs(thread_amax_2x.y))}; #pragma unroll for (int w = 0; w < 2; ++w) { - const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax[w], S_enc_colwise); + float S_enc_colwise_block = S_enc_colwise; + if constexpr (ROW_SCALED_NVFP4) { + const size_t col_idx = + col_offset + stage_X * TILE_DIM_X + thread_offset_X_colwise + w; + S_enc_colwise_block = + col_idx < cols ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx]) + : 1.0f; + } + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax[w], S_enc_colwise_block); // Store scaling factors to SMEM buffer (R2S) sSFcolwise[scale_tr_offset_Y + w][scale_tr_offset_X] = S_dec_b_fp8; const scaling_coeff_type SFcoefficient = - compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_colwise); + compute_nvfp4_scaling_coefficient(S_dec_b_fp8, + S_enc_colwise_block); // Scale elements __align__(8) uint32_t rOut[SCALE_DIM / 8]; @@ -432,7 +444,7 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); const float S_enc_colwise = - (amax_colwise_ptr == nullptr) + (amax_colwise_ptr == nullptr || ROW_SCALED_NVFP4) ? S_enc_rowwise : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); @@ -587,9 +599,9 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D amax_rowwise_ptr, block_offset_Y, rows, rng, random_uint4, rnd_idx); if constexpr (RETURN_TRANSPOSE) { - colwise_scaling( + colwise_scaling( sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, S_enc_colwise, stage_Y, stage_X, buff_in, - buff_out_tr, rng, random_uint4, rnd_idx); + buff_out_tr, amax_colwise_ptr, block_offset_X, cols, rng, random_uint4, rnd_idx); } // Wait for shared memory writes to be visible to TMA engine @@ -708,14 +720,14 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr, "Row-scaled NVFP4 quantization requires rowwise amax."); - NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), - "Row-scaled NVFP4 quantization does not produce columnwise output."); if (return_transpose) { NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), "Transposed output must have FP4 type."); NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Transposed scaling tensor must be allocated"); + NVTE_CHECK(!row_scaled_nvfp4 || output->columnwise_amax.dptr != nullptr, + "Row-scaled NVFP4 transpose quantization requires columnwise amax."); } const auto [rows, cols] = input.flat_2d_dims(); diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3b066d50b..e87b402177 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -80,25 +80,29 @@ def _nvfp4_row_scaled_gemm_inputs( B: NVFP4TensorStorage, *, transa: bool, -) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor]: - """Return GEMM aliases and FP32 output scales for row-scaled NVFP4.""" + transb: bool, +) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor, torch.Tensor]: + """Return per-tensor GEMM aliases and row/column FP32 output scales.""" A_metadata = A.get_metadata() - weight_amax = A._amax_rowwise if transa else A._amax_columnwise - assert weight_amax is not None and weight_amax.numel() == 1 - A_metadata["amax_rowwise" if transa else "amax_columnwise"] = weight_amax.new_ones(1) + a_amax_key = "amax_rowwise" if transa else "amax_columnwise" + output_col_scales = A_metadata[a_amax_key] + assert output_col_scales is not None + A_metadata[a_amax_key] = output_col_scales.new_ones(1) A_metadata["row_scaled_nvfp4"] = False B_metadata = B.get_metadata() - rhs_rowwise_amax = B._amax_rowwise - assert rhs_rowwise_amax is not None - B_metadata["amax_rowwise"] = rhs_rowwise_amax.new_ones(1) + b_amax_key = "amax_columnwise" if transb else "amax_rowwise" + output_row_scales = B_metadata[b_amax_key] + assert output_row_scales is not None + B_metadata[b_amax_key] = output_row_scales.new_ones(1) B_metadata["row_scaled_nvfp4"] = False - assert rhs_rowwise_amax.dtype == torch.float32 and weight_amax.dtype == torch.float32 + assert output_row_scales.dtype == torch.float32 and output_col_scales.dtype == torch.float32 return ( NVFP4TensorStorage(**A_metadata), NVFP4TensorStorage(**B_metadata), - (rhs_rowwise_amax * weight_amax).view(-1, 1), + output_row_scales.view(-1, 1), + output_col_scales.view(1, -1), ) @@ -210,16 +214,7 @@ def general_gemm( if not _is_nvfp4_row_scaled_tensor(A) and not _is_nvfp4_row_scaled_tensor(B): out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) else: - if _is_nvfp4_row_scaled_tensor(A): - raise NotImplementedError("Row-scaled NVFP4 GEMM does not support row-scaled A.") - assert layout[1] == "N", "Row-scaled NVFP4 GEMM currently supports N-layout B only." - if grad: - raise RuntimeError( - "Row-scaled NVFP4 GEMM currently supports fprop only. " - "Backward NVFP4 gradient quantizers should use scalar global amax." - ) assert not gelu, "Row-scaled NVFP4 GEMM currently does not support fused GELU." - assert not accumulate, "Row-scaled NVFP4 GEMM currently does not support accumulation." assert ( quantization_params is None ), "Row-scaled NVFP4 GEMM currently does not support output quantization." @@ -234,9 +229,14 @@ def general_gemm( assert isinstance( A, NVFP4TensorStorage ), "Row-scaled NVFP4 GEMM currently requires NVFP4 A." - # cuBLAS folds NVFP4 global amax values into GEMM alpha. Keep the row-scaled - # recipe's global scales out of alpha and apply them in FP32 below. - gemm_A, gemm_B, rowwise_global_scales = _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa) + assert isinstance( + B, NVFP4TensorStorage + ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." + # Reuse the per-tensor GEMM and apply selected row/column global scales + # to the FP32 output. This extends #2931 without a dedicated GEMM kernel. + gemm_A, gemm_B, output_row_scales, output_col_scales = ( + _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa, transb=transb) + ) requested_out, requested_out_dtype = out, out_dtype fp32_out = ( @@ -251,18 +251,36 @@ def general_gemm( gemm_args[5] = None # quantization_params gemm_args[6] = TE_DType[torch.float32] # out_dtype gemm_args[7] = None # bias - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **kwargs) + gemm_args[14] = False # accumulate after applying the outer scales + gemm_kwargs = dict(kwargs) + gemm_kwargs["beta"] = 0.0 + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **gemm_kwargs) out_2d = out.reshape(-1, out.shape[-1]) - assert rowwise_global_scales.dtype == torch.float32 and out.dtype == torch.float32 - assert rowwise_global_scales.numel() == out_2d.shape[0] - - out_2d.mul_(rowwise_global_scales) + assert output_row_scales.numel() in (1, out_2d.shape[0]) + assert output_col_scales.numel() in (1, out_2d.shape[1]) + assert out.dtype == torch.float32 + # When one side is a scalar global amax (e.g. fprop weight), fold both + # scales into a single factor before the multiply. This reproduces + # #2931's fused `out * (row_amax * col_amax)` arithmetic bit-for-bit; + # only the true bilateral case (both per-row and per-col) needs the + # two-step outer-product scaling. + if output_col_scales.numel() == 1: + out_2d.mul_(output_row_scales * output_col_scales) + elif output_row_scales.numel() == 1: + out_2d.mul_(output_col_scales * output_row_scales) + else: + out_2d.mul_(output_row_scales) + out_2d.mul_(output_col_scales) if bias is not None: + assert not grad, "Row-scaled NVFP4 backward does not support fused bias gradient." out_2d.add_(bias.to(dtype=torch.float32)) if requested_out is not None: - requested_out.copy_(out.to(dtype=requested_out.dtype)) + if accumulate: + requested_out.add_(out.to(dtype=requested_out.dtype)) + else: + requested_out.copy_(out.to(dtype=requested_out.dtype)) out = requested_out elif requested_out_dtype is not None and requested_out_dtype != torch.float32: out = out.to(dtype=requested_out_dtype) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4308464c52..28e790fa02 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1972,8 +1972,6 @@ std::pair NVFP4Quantizer::create_tensor( const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 quantization does not support columnwise usage."); } const auto rowwise_scale_inv_shape = get_scale_shape(shape, false); const auto columnwise_scale_inv_shape = get_scale_shape(shape, true); @@ -2008,7 +2006,8 @@ std::pair NVFP4Quantizer::create_tensor( columnwise_scale_inv_tensor = at::empty(scale_inv_shape_int64, bit8_tensor_opts); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_columnwise = at::empty({1}, bit32_tensor_opts); + const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; + amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); } // Convert tensors to Python @@ -2100,7 +2099,7 @@ std::pair NVFP4Quantizer::create_tensor( out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, columnwise_scale_inv_shape); out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(amax_columnwise)); } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); @@ -2299,8 +2298,6 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 quantization does not support columnwise usage."); } tensor.attr("_row_scaled_nvfp4") = row_scaled_nvfp4; tensor.attr("_with_gemm_swizzled_scales") = with_gemm_swizzled_scales; @@ -2366,11 +2363,12 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( columnwise_scale_inv = at::empty(scale_inv_shape_int64, opts); tensor.attr("_columnwise_scale_inv") = *columnwise_scale_inv; } - if (!amax_columnwise) { + const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; + if (!amax_columnwise || amax_columnwise->numel() != amax_cols) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_columnwise = at::empty({1}, opts); + amax_columnwise = at::empty({amax_cols}, opts); tensor.attr("_amax_columnwise") = *amax_columnwise; } } else { // columnwise_usage == false @@ -2406,7 +2404,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*columnwise_scale_inv)); out_cpp.set_columnwise_amax(amax_columnwise->data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(*amax_columnwise)); } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py index d09b95ace3..a2b21465af 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py @@ -363,10 +363,6 @@ def __init__( if row_scaled_nvfp4: if not rowwise: raise ValueError("Row-scaled NVFP4 reference quantization requires rowwise usage.") - if columnwise: - raise ValueError( - "Row-scaled NVFP4 reference quantization does not support columnwise usage." - ) if nvfp4_use_4over6: if nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError(f"Unsupported NVFP4 4over6 error mode: {nvfp4_4over6_err_mode}.") @@ -883,7 +879,14 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ f"got {self.quant_tile_shape}" ) global_amax_row = torch.max(torch.abs(row_input), dim=1).values.to(torch.float32) - global_amax_col = global_amax_row + # Columnwise (transpose) uses per-row-of-transpose amax, i.e. the + # per-column amax of the original input. When columnwise output is + # not requested, keep it aliased to the rowwise amax as before. + global_amax_col = ( + torch.max(torch.abs(col_input), dim=1).values.to(torch.float32) + if self.columnwise_usage + else global_amax_row + ) else: # Compute amax for rowwise and columnwise paths separately global_amax_row = torch.max(torch.abs(row_input)).to(torch.float32).view(1) @@ -936,6 +939,7 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ self.quant_tile_shape[1], self.quant_tile_shape[0], pow_2_scales=self.pow_2_scales, + row_scaled_nvfp4=self.row_scaled_nvfp4, nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, @@ -1167,12 +1171,21 @@ def qgemm( if gemm_type == quantization.GEMMType.WGRAD: partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col + # A row-scaled operand contributes a per-output-column (N) vector + # here, so broadcast along the last axis. Selecting the axis from + # gemm_type (rather than matching numel against M) avoids the + # square-matrix ambiguity where M == N. + if partial_alpha.numel() > 1: + partial_alpha = partial_alpha.reshape(1, -1) + else: + partial_alpha = partial_alpha.squeeze(-1) else: partial_alpha = qresult_x.global_amax_row * qresult_w.global_amax_row - if partial_alpha.numel() > 1 and partial_alpha.numel() == high_precision_x.shape[0]: - partial_alpha = partial_alpha.view(-1, 1) - else: - partial_alpha = partial_alpha.squeeze(-1) + # A row-scaled operand contributes a per-output-row (M) vector. + if partial_alpha.numel() > 1: + partial_alpha = partial_alpha.reshape(-1, 1) + else: + partial_alpha = partial_alpha.squeeze(-1) alpha = torch.div(partial_alpha, factor) M, K = high_precision_x.shape From 86ee9094f649bd71b62cf91803d6979e20c0f1b9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:30:27 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 7 +++---- .../quantize_transpose_nvfp4_tuned_1D.cuh | 20 +++++++------------ .../pytorch/cpp_extensions/gemm.py | 4 ++-- 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index e3971329d2..a38a620ebe 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -120,8 +120,7 @@ __launch_bounds__(BLOCK_SIZE) float thread_max = 0.0f; for (int row_idx = threadIdx.x; row_idx < num_rows; row_idx += BLOCK_SIZE) { - thread_max = - fmaxf(thread_max, fabsf(static_cast(input[row_idx * num_cols + col_idx]))); + thread_max = fmaxf(thread_max, fabsf(static_cast(input[row_idx * num_cols + col_idx]))); } const float col_amax = reduce_max(thread_max, threadIdx.x / THREADS_PER_WARP); @@ -216,8 +215,8 @@ inline void compute_columnwise_amax(const Tensor &input, const Tensor *noop, Ten const auto *noop_ptr = reinterpret_cast(noop->data.dptr); if (input.dtype() == DType::kBFloat16) { const auto *input_ptr = reinterpret_cast(input.data.dptr); - launch_compute_columnwise_amax<__nv_bfloat16>( - static_cast(rows), static_cast(cols), input_ptr, amax_ptr, stream, noop_ptr); + launch_compute_columnwise_amax<__nv_bfloat16>(static_cast(rows), static_cast(cols), + input_ptr, amax_ptr, stream, noop_ptr); } else if (input.dtype() == DType::kFloat16) { const auto *input_ptr = reinterpret_cast(input.data.dptr); launch_compute_columnwise_amax(static_cast(rows), static_cast(cols), input_ptr, diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 637abe7a0f..ad21486368 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -185,15 +185,11 @@ compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const f } template -__device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_ptr, - fp4e2m1x2 *__restrict__ sOut_tr_ptr, - nvfp4_scale_t *__restrict__ sSFcolwise_ptr, - const float S_enc_colwise, const int stage_Y, - const int stage_X, const int buff_in, - const int buff_out_tr, - const float *amax_colwise_ptr, - const size_t col_offset, const size_t cols, - RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { +__device__ __forceinline__ void colwise_scaling( + const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_tr_ptr, + nvfp4_scale_t *__restrict__ sSFcolwise_ptr, const float S_enc_colwise, const int stage_Y, + const int stage_X, const int buff_in, const int buff_out_tr, const float *amax_colwise_ptr, + const size_t col_offset, const size_t cols, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn2x = *reinterpret_cast(sIn_ptr); @@ -235,8 +231,7 @@ __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_pt for (int w = 0; w < 2; ++w) { float S_enc_colwise_block = S_enc_colwise; if constexpr (ROW_SCALED_NVFP4) { - const size_t col_idx = - col_offset + stage_X * TILE_DIM_X + thread_offset_X_colwise + w; + const size_t col_idx = col_offset + stage_X * TILE_DIM_X + thread_offset_X_colwise + w; S_enc_colwise_block = col_idx < cols ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx]) : 1.0f; @@ -248,8 +243,7 @@ __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_pt sSFcolwise[scale_tr_offset_Y + w][scale_tr_offset_X] = S_dec_b_fp8; const scaling_coeff_type SFcoefficient = - compute_nvfp4_scaling_coefficient(S_dec_b_fp8, - S_enc_colwise_block); + compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_colwise_block); // Scale elements __align__(8) uint32_t rOut[SCALE_DIM / 8]; diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index e87b402177..d56727e3a6 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -234,8 +234,8 @@ def general_gemm( ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." # Reuse the per-tensor GEMM and apply selected row/column global scales # to the FP32 output. This extends #2931 without a dedicated GEMM kernel. - gemm_A, gemm_B, output_row_scales, output_col_scales = ( - _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa, transb=transb) + gemm_A, gemm_B, output_row_scales, output_col_scales = _nvfp4_row_scaled_gemm_inputs( + A, B, transa=transa, transb=transb ) requested_out, requested_out_dtype = out, out_dtype