From b3701172d96e92e9cafc69aa371bbbcd5001ec69 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Fri, 7 Aug 2026 05:57:55 +0000 Subject: [PATCH 1/2] Integrate MXFP4 hipblaslt GEMM support --- tests/cpp/operator/test_cublaslt_gemm.cu | 163 ++++++++++++++ tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py | 198 ++++++++++++------ transformer_engine/common/gemm/rocm_gemm.cu | 53 ++++- .../pytorch/cpp_extensions/gemm.py | 26 +-- transformer_engine/pytorch/quantization.py | 15 +- 5 files changed, 363 insertions(+), 92 deletions(-) diff --git a/tests/cpp/operator/test_cublaslt_gemm.cu b/tests/cpp/operator/test_cublaslt_gemm.cu index aeb40d8967..339d27bfbb 100644 --- a/tests/cpp/operator/test_cublaslt_gemm.cu +++ b/tests/cpp/operator/test_cublaslt_gemm.cu @@ -36,6 +36,16 @@ std::vector> test_case_sizes_mxfp8 = { {4096, 16384, 4096}, }; +// MXFP4 (m, k, n): M/N multiples of 32 (block size), K a multiple of 256 (see rocm_gemm.cu +// gate; K%128-not-%256 runs but is numerically wrong through TE's padded-scale layout). +// Square + non-square. +std::vector> test_case_sizes_mxfp4 = { + {256, 256, 256}, + {128, 256, 512}, + {768, 3072, 4096}, + {4096, 512, 3072}, +}; + // ============================================================================ // Production LLM MXFP8 GEMM shapes. // ============================================================================ @@ -853,6 +863,122 @@ void performDqTest(const TestParams ¶ms) { } #endif // __HIP_PLATFORM_AMD__ +#ifdef __HIP_PLATFORM_AMD__ +// FP4 E2M1 value table (matches rocm_gemm.cu / OCP microscaling). +static const float kHostFP4E2M1Table[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f,-0.5f,-1.0f,-1.5f,-2.0f,-3.0f,-4.0f,-6.0f}; + +// CPU-dequantize the row-wise data of an MXFP4 test::Tensor (E2M1 packed two-per-byte + +// UE8M0 block-32 scales, padded scale layout) into the row-wise BF16 buffer of dst, then +// upload. Used to build a high-precision reference for the native MXFP4 GEMM. +static void dequantize_mxfp4_rowwise_to_bf16(test::Tensor &src_fp4, test::Tensor &dst_bf16) { + const NVTEShape data_shape = src_fp4.rowwise_shape(); // logical [R, C] + NVTE_CHECK(data_shape.ndim == 2, "Expected 2D MXFP4 data"); + const size_t R = data_shape.data[0]; + const size_t C = data_shape.data[1]; + NVTE_CHECK((C % 2) == 0, "MXFP4 columns must be even (two-per-byte packing)"); + + const size_t packed_bytes = R * (C / 2); + std::vector h_data(packed_bytes); + NVTE_CHECK_CUDA(cudaMemcpy(h_data.data(), src_fp4.rowwise_dptr(), packed_bytes, + cudaMemcpyDeviceToHost)); + + const NVTEShape scale_shape = src_fp4.rowwise_scale_inv_shape(); // padded [Ypad, Xpad] + NVTE_CHECK(scale_shape.ndim == 2, "Expected 2D MXFP4 scale_inv"); + const size_t x_pad = scale_shape.data[1]; + const size_t num_scales = scale_shape.data[0] * scale_shape.data[1]; + std::vector h_scale(num_scales); + NVTE_CHECK_CUDA(cudaMemcpy(h_scale.data(), src_fp4.rowwise_scale_inv_dptr(), num_scales, + cudaMemcpyDeviceToHost)); + + bf16 *dst = dst_bf16.rowwise_cpu_dptr(); + for (size_t r = 0; r < R; ++r) { + for (size_t c = 0; c < C; ++c) { + const uint8_t byte = h_data[r * (C / 2) + c / 2]; + const uint8_t nib = (c & 1) ? (byte >> 4) : (byte & 0xF); + const uint8_t e8m0 = h_scale[r * x_pad + c / 32]; + const float scale = exp2f(static_cast(e8m0) - 127.0f); + dst[r * C + c] = static_cast(kHostFP4E2M1Table[nib] * scale); + } + } + dst_bf16.from_cpu(); +} + +// Native MXFP4 (hipBLASLt) GEMM vs a BF16 reference GEMM built by dequantizing the same +// MXFP4 operands. Restricted to TN layout (the F4F4 kernels are TN in practice), so both +// operands consume row-wise data. +template +void performMxfp4Test(const TestParams ¶ms) { + DType dtype = TypeInfo::dtype; + + cudaDeviceProp prop; + (void)cudaGetDeviceProperties(&prop, 0); + + // hipBLASLt native MXFP4 GEMM requires gfx950 (ROCm >= 7.13 / hipBLASLt >= 1.3). + if (!(prop.major == 9 && prop.minor == 5)) { + GTEST_SKIP() << "MXFP4 GEMM is only supported on gfx950"; + } + if (!(params.transa && !params.transb)) { + GTEST_SKIP() << "MXFP4 GEMM test only covers TN layout"; + } + if (params.m % 32 || params.n % 32 || params.k % 256) { + GTEST_SKIP() << "MXFP4 requires M, N multiples of 32 and K a multiple of 256"; + } + + // TN: A is [m, k], B is [n, k]; both consumed row-wise. + TShape a_shape = TShape{params.m, params.k}; + TShape b_shape = TShape{params.n, params.k}; + + Tensor A_src("A", a_shape, DType::kBFloat16); + Tensor B_src("B", b_shape, DType::kBFloat16); + fillUniform(&A_src); + fillUniform(&B_src); + + Tensor A_fp4("A_fp4", a_shape, DType::kFloat4E2M1, /*rowwise=*/true, /*columnwise=*/false, + NVTE_MXFP4_1D_SCALING); + Tensor B_fp4("B_fp4", b_shape, DType::kFloat4E2M1, /*rowwise=*/true, /*columnwise=*/false, + NVTE_MXFP4_1D_SCALING); + nvte_quantize(A_src.data(), A_fp4.data(), 0); + nvte_quantize(B_src.data(), B_fp4.data(), 0); + + // High-precision reference operands = exact dequant of the MXFP4 data. + Tensor A_ref("A_ref", a_shape, DType::kBFloat16); + Tensor B_ref("B_ref", b_shape, DType::kBFloat16); + dequantize_mxfp4_rowwise_to_bf16(A_fp4, A_ref); + dequantize_mxfp4_rowwise_to_bf16(B_fp4, B_ref); + + Tensor bias; + if (params.use_bias) { + bias = Tensor("bias", TShape{params.m}, dtype); + fillUniform(&bias); + } + Tensor pre_gelu_out; + Tensor Workspace("Workspace", TShape{67'108'864}, DType::kByte); + + Tensor D("D", TShape{params.n, params.m}, dtype); + nvte_cublas_gemm(A_fp4.data(), B_fp4.data(), D.data(), bias.data(), pre_gelu_out.data(), + params.transa, params.transb, false, Workspace.data(), false, false, + prop.multiProcessorCount, 0); + D.to_cpu(); + + Tensor D_ref("D_ref", TShape{params.n, params.m}, dtype); + nvte_cublas_gemm(A_ref.data(), B_ref.data(), D_ref.data(), bias.data(), pre_gelu_out.data(), + params.transa, params.transb, false, Workspace.data(), false, false, + prop.multiProcessorCount, 0); + D_ref.to_cpu(); + + (void)cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + // FP4 is coarse; the native kernel and the BF16 reference differ mainly in accumulation. + const double atol = 3e-2; + const double rtol = 6e-2; + compareResults("D", D, D_ref.rowwise_cpu_dptr(), true, atol, rtol); +} +#endif // __HIP_PLATFORM_AMD__ + #define MAKE_TEST_PARAMS(P_) \ bool force_hipblaslt_ = std::get<5>(GetParam()); \ if (force_hipblaslt_) { \ @@ -988,6 +1114,43 @@ INSTANTIATE_TEST_SUITE_P(OperatorTestMXFP8, DqGEMMTestSuite, (std::get<5>(info.param) ? "HB" : "HK"); }); +// ============================================================================ +// Native MXFP4 (hipBLASLt) GEMM tests +// ============================================================================ +class Mxfp4GEMMTestSuite + : public ::testing::TestWithParam< + std::tuple, bool, Layout>> {}; + +#define MAKE_MXFP4_GEMM_TEST(NAME_, D_) \ + TEST_P(Mxfp4GEMMTestSuite, NAME_) { \ + const auto shape = std::get<0>(GetParam()); \ + TestParams params = {.m = std::get<0>(shape), \ + .k = std::get<1>(shape), \ + .n = std::get<2>(shape), \ + .use_bias = std::get<1>(GetParam()), \ + .use_gelu = false, \ + .transa = std::get<2>(GetParam()).first, \ + .transb = std::get<2>(GetParam()).second, \ + .scaling_mode = NVTEScalingMode::NVTE_MXFP4_1D_SCALING, \ + .force_hipblaslt = false}; \ + performMxfp4Test(params); \ + } + +// TE MXFP4 path supports only BF16/FP32 output (FP16 F4H kernel exists in hipBLASLt but TE's +// descriptor finds "no suitable algorithms" for it -- see §3.5/§10.1) and no bias/GELU epilogue. +MAKE_MXFP4_GEMM_TEST(Testbf16, bf16) +MAKE_MXFP4_GEMM_TEST(Testfp32, fp32) + +INSTANTIATE_TEST_SUITE_P(OperatorTestMXFP4, Mxfp4GEMMTestSuite, + ::testing::Combine(::testing::ValuesIn(test_case_sizes_mxfp4), + ::testing::Values(false), // bias unsupported + ::testing::Values(kTN)), // TN only + [](const testing::TestParamInfo& info) { + return MKN(std::get<0>(info.param)) + "x" + + (std::get<1>(info.param) ? "bias" : "nobias") + "x" + + TN(std::get<2>(info.param)); + }); + // ============================================================================ // Production GEMM shape instantiations (run with --gtest_filter='ProdGemm*') // ============================================================================ diff --git a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py index df0de3d77a..5031a9939b 100644 --- a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py +++ b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py @@ -2,11 +2,14 @@ # # See LICENSE for license information. -"""MXFP4 GEMM tests: native AITER a4w4 GEMM vs Python reference GEMM. +"""MXFP4 GEMM tests: native GEMM (AITER a4w4 or hipBLASLt) vs Python reference GEMM. -Requires the aiter package (ROCm gfx950 only). +Requires the aiter package (ROCm gfx950 only). The hipBLASLt backend is exercised by +setting NVTE_ROCM_USE_HIPBLASLT_MXFP4=1, which routes MXFP4 GEMMs through the native +hipBLASLt path (rocm_gemm.cu) instead of AITER. """ +import os import pytest import torch import transformer_engine.pytorch as te @@ -27,60 +30,40 @@ BLOCK_SIZE = 32 -def check_mxfp4_gemm_versus_reference( - x_dtype: torch.dtype, - w_dtype: torch.dtype, - out_dtype: torch.dtype, - M: int, - K: int, - N: int, - accumulate: bool, -): +def _quantize_pair(x, w, M, K, N, x_dtype, w_dtype, device, *, shuffle: bool): + """Quantize (x, w) to MXFP4 with either AITER-shuffled or plain (hipBLASLt) layout.""" te_dtype = tex.DType.kFloat4E2M1 - device = "cuda" - seed = 0 - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - - x = torch.randn((M, K), dtype=x_dtype, device=device) - w = torch.randn((N, K), dtype=w_dtype, device=device) - - if accumulate: - out = torch.randn((M, N), dtype=out_dtype, device=device) - else: - out = None - - # Native MXFP4 quantization (shuffled for AITER GEMM) + # hipBLASLt consumes plain (un-shuffled) FP4 data + plain scales; AITER needs the + # 16x16 weight shuffle and swizzled scales. x_quantizer = MXFP4Quantizer( fp4_dtype=te_dtype, rowwise=True, columnwise=True, shuffle_rowwise_data=False, - shuffle_columnwise_data=False, - with_gemm_swizzled_scales=True, + shuffle_columnwise_data=shuffle, + with_gemm_swizzled_scales=shuffle, use_hadamard=False, ) w_quantizer = MXFP4Quantizer( fp4_dtype=te_dtype, rowwise=True, columnwise=True, - shuffle_rowwise_data=True, - shuffle_columnwise_data=True, - with_gemm_swizzled_scales=True, + shuffle_rowwise_data=shuffle, + shuffle_columnwise_data=shuffle, + with_gemm_swizzled_scales=shuffle, use_hadamard=False, ) + x_mxfp4 = x_quantizer.make_empty((M, K), dtype=x_dtype, device=device, requires_grad=False) + x_mxfp4 = x_quantizer.update_quantized(x, x_mxfp4) + w_mxfp4 = w_quantizer.make_empty((N, K), dtype=w_dtype, device=device, requires_grad=False) + w_mxfp4 = w_quantizer.update_quantized(w, w_mxfp4) + return x_mxfp4, w_mxfp4 - # Reference quantization (plain layout, no shuffle) - x_quantizer_ref = MXFP4Quantizer( - fp4_dtype=te_dtype, - rowwise=True, - columnwise=True, - shuffle_rowwise_data=False, - shuffle_columnwise_data=False, - with_gemm_swizzled_scales=False, - use_hadamard=False, - ) - w_quantizer_ref = MXFP4Quantizer( + +def _reference_gemm(x, w, M, K, N, out_dtype, out, accumulate, device): + """MXFP4 reference GEMM: quantize plain, dequant-matmul via MXFP4QuantizerRef.""" + te_dtype = tex.DType.kFloat4E2M1 + ref_quantizer_cfg = dict( fp4_dtype=te_dtype, rowwise=True, columnwise=True, @@ -89,30 +72,21 @@ def check_mxfp4_gemm_versus_reference( with_gemm_swizzled_scales=False, use_hadamard=False, ) - - x_mxfp4 = x_quantizer.make_empty((M, K), dtype=x_dtype, device=device, requires_grad=False) - x_mxfp4 = x_quantizer.update_quantized(x, x_mxfp4) - w_mxfp4 = w_quantizer.make_empty((N, K), dtype=w_dtype, device=device, requires_grad=False) - w_mxfp4 = w_quantizer.update_quantized(w, w_mxfp4) - - x_mxfp4_ref = x_quantizer_ref.make_empty((M, K), dtype=x_dtype, device=device, requires_grad=False) - x_mxfp4_ref = x_quantizer_ref.update_quantized(x, x_mxfp4_ref) - w_mxfp4_ref = w_quantizer_ref.make_empty((N, K), dtype=w_dtype, device=device, requires_grad=False) - w_mxfp4_ref = w_quantizer_ref.update_quantized(w, w_mxfp4_ref) - - # Extract un-shuffled quantized data for the reference GEMM - qx_data = x_mxfp4_ref._rowwise_data.view(dtype=torch.uint8)[:M, :] - qw_data = w_mxfp4_ref._rowwise_data.view(dtype=torch.uint8)[:N, :] - sx_native = x_mxfp4_ref._rowwise_scale_inv - sw_native = w_mxfp4_ref._rowwise_scale_inv - + x_ref_q = MXFP4Quantizer(**ref_quantizer_cfg) + w_ref_q = MXFP4Quantizer(**ref_quantizer_cfg) + x_ref = x_ref_q.make_empty((M, K), dtype=x.dtype, device=device, requires_grad=False) + x_ref = x_ref_q.update_quantized(x, x_ref) + w_ref = w_ref_q.make_empty((N, K), dtype=w.dtype, device=device, requires_grad=False) + w_ref = w_ref_q.update_quantized(w, w_ref) + + qx_data = x_ref._rowwise_data.view(dtype=torch.uint8)[:M, :] + qw_data = w_ref._rowwise_data.view(dtype=torch.uint8)[:N, :] expected_scale_cols = K // BLOCK_SIZE - sx_trimmed = sx_native[:M, :expected_scale_cols] - sw_trimmed = sw_native[:N, :expected_scale_cols] + sx_trimmed = x_ref._rowwise_scale_inv[:M, :expected_scale_cols] + sw_trimmed = w_ref._rowwise_scale_inv[:N, :expected_scale_cols] - # Reference GEMM ref_quantizer = MXFP4QuantizerRef(rowwise=True, columnwise=True) - y_ref = ref_quantizer.qgemm( + return ref_quantizer.qgemm( qx=qx_data, qw=qw_data, m_params=None, # MMParams not used in reference @@ -124,7 +98,8 @@ def check_mxfp4_gemm_versus_reference( accumulate=accumulate, ) - # Native AITER GEMM via general_gemm + +def _native_gemm(w_mxfp4, x_mxfp4, out_dtype, out, accumulate): from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm y_native, *_ = general_gemm( @@ -137,14 +112,48 @@ def check_mxfp4_gemm_versus_reference( out=out.clone() if accumulate else None, accumulate=accumulate, ) + return y_native + + +def _sanitize(t): + return torch.where(t.isnan(), torch.zeros_like(t), t) + + +def check_mxfp4_gemm_versus_reference( + x_dtype: torch.dtype, + w_dtype: torch.dtype, + out_dtype: torch.dtype, + M: int, + K: int, + N: int, + accumulate: bool, + backend: str, +): + device = "cuda" + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + x = torch.randn((M, K), dtype=x_dtype, device=device) + w = torch.randn((N, K), dtype=w_dtype, device=device) + out = torch.randn((M, N), dtype=out_dtype, device=device) if accumulate else None + + shuffle = backend == "aiter" + x_mxfp4, w_mxfp4 = _quantize_pair( + x, w, M, K, N, x_dtype, w_dtype, device, shuffle=shuffle + ) + + y_ref = _reference_gemm(x, w, M, K, N, out_dtype, out, accumulate, device) + y_native = _native_gemm(w_mxfp4, x_mxfp4, out_dtype, out, accumulate) assert y_ref is not y_native assert not torch.isnan(y_ref.float()).all(), "All reference elements are NaN" - y_ref = torch.where(y_ref.isnan(), torch.zeros_like(y_ref), y_ref) - y_native = torch.where(y_native.isnan(), torch.zeros_like(y_native), y_native) + torch.testing.assert_close(_sanitize(y_native), _sanitize(y_ref), atol=8e-3, rtol=8e-3) + - torch.testing.assert_close(y_native, y_ref, atol=8e-3, rtol=8e-3) +_BACKENDS = ["aiter"] +if _use_hipblaslt := os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") == "1": + _BACKENDS = ["hipblaslt"] @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @@ -163,6 +172,7 @@ def check_mxfp4_gemm_versus_reference( @pytest.mark.parametrize("w_dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("accumulate", [True, False], ids=["accumulate", "no_accumulate"]) +@pytest.mark.parametrize("backend", _BACKENDS) def test_mxfp4_gemm_versus_reference( M: int, K: int, @@ -170,8 +180,17 @@ def test_mxfp4_gemm_versus_reference( x_dtype: torch.dtype, w_dtype: torch.dtype, out_dtype: torch.dtype, - accumulate: bool + accumulate: bool, + backend: str, ): + # The hipBLASLt MXFP4 path requires K % 256 == 0 (TE's padded UE8M0 scale layout; K%128 runs + # but is numerically wrong -- design doc §3.5/§10) and does not support accumulate yet + # (§10 #4). AITER covers the full matrix. + if backend == "hipblaslt": + if K % 256 != 0: + pytest.skip("hipBLASLt MXFP4 currently requires K to be a multiple of 256") + if accumulate: + pytest.skip("hipBLASLt MXFP4 does not support accumulate yet") check_mxfp4_gemm_versus_reference( x_dtype=x_dtype, w_dtype=w_dtype, @@ -180,4 +199,49 @@ def test_mxfp4_gemm_versus_reference( K=K, N=N, accumulate=accumulate, + backend=backend, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not _aiter_available, reason="aiter package not available") +@pytest.mark.skipif( + os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") != "1", + reason="hipBLASLt MXFP4 backend not enabled (set NVTE_ROCM_USE_HIPBLASLT_MXFP4=1)", +) +@pytest.mark.parametrize( + "M, K, N", + [ + (256, 256, 256), + (1024, 1024, 1024), + (4096, 512, 3072), + ], +) +def test_mxfp4_gemm_hipblaslt_matches_aiter(M: int, K: int, N: int): + """Cross-check: same inputs, hipBLASLt vs AITER MXFP4 GEMM produce close results.""" + device = "cuda" + dtype = torch.bfloat16 + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + x = torch.randn((M, K), dtype=dtype, device=device) + w = torch.randn((N, K), dtype=dtype, device=device) + + prev = os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4") + try: + # AITER path (shuffled operands) + os.environ["NVTE_ROCM_USE_HIPBLASLT_MXFP4"] = "0" + x_a, w_a = _quantize_pair(x, w, M, K, N, dtype, dtype, device, shuffle=True) + y_aiter = _native_gemm(w_a, x_a, dtype, None, False) + + # hipBLASLt path (plain operands) + os.environ["NVTE_ROCM_USE_HIPBLASLT_MXFP4"] = "1" + x_h, w_h = _quantize_pair(x, w, M, K, N, dtype, dtype, device, shuffle=False) + y_hip = _native_gemm(w_h, x_h, dtype, None, False) + finally: + if prev is None: + os.environ.pop("NVTE_ROCM_USE_HIPBLASLT_MXFP4", None) + else: + os.environ["NVTE_ROCM_USE_HIPBLASLT_MXFP4"] = prev + + torch.testing.assert_close(_sanitize(y_hip), _sanitize(y_aiter), atol=8e-3, rtol=8e-3) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index 082c96ee7a..287a78d09e 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -188,6 +188,8 @@ static hipDataType get_hipblaslt_dtype(const transformer_engine::DType t) { return te_fp8_fnuz() ? HIP_R_8F_E4M3_FNUZ : HIP_R_8F_E4M3; case DType::kFloat8E5M2: return te_fp8_fnuz() ? HIP_R_8F_E5M2_FNUZ: HIP_R_8F_E5M2; + case DType::kFloat4E2M1: + return HIP_R_4F_E2M1; default: NVTE_ERROR("Invalid type"); } @@ -435,7 +437,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla } } } else if (is_mxfp_scaling(A.scaling_mode)) { - // MXFP8 + // MXFP8, MXFP4 // Note: Row-wise and column-wise data are scaled along different // dimensions (with matrix interpreted in row-major order). if (is_A_transposed) { @@ -487,7 +489,9 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla } } } else if (is_mxfp_scaling(B.scaling_mode)) { - // MXFP8 + // MXFP8 and MXFP4 (is_mxfp_scaling() covers both). The FP4 data dtype flows + // through B.data.dtype and leading dims are in elements (hipBLASLt handles the + // two-per-byte FP4 packing), so both formats share this canonicalization. // Note: Row-wise and column-wise data are scaled along different // dimensions (with matrix interpreted in row-major order). if (is_B_transposed) { @@ -801,6 +805,7 @@ static std::unordered_map type_name_map = { {HIP_R_8F_E5M2_FNUZ, "float8e5m2"}, {HIP_R_8F_E4M3, "float8e4m3"}, {HIP_R_8F_E5M2, "float8e5m2"}, + {HIP_R_4F_E2M1, "float4e2m1"}, }; static NameMapper typeNameMapper(type_name_map); @@ -1291,10 +1296,15 @@ void hipblaslt_gemm(const Tensor *inputA, // alpha'[i] = alpha * amax_A * amax_B / (fp4_max^2 * fp8_max^2) // Alpha is passed as a device vector of length m via // HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST. Beta stays on host. - const bool use_fp4 = is_fp4_dtype(param.Atype) || is_fp4_dtype(param.Btype); + // + // Only NVFP4 uses this dequant fallback. MXFP4 (also an FP4 data dtype) has a native + // hipBLASLt path, so gate on the scaling mode rather than the FP4 data dtype (which + // would match both). + const bool use_nvfp4 = + is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode); const void* alpha_ptr = static_cast(&alpha); const void* beta_ptr = static_cast(&beta); - if (use_fp4) { + if (use_nvfp4) { dequant_fp4_gemm_inputs(param, *inputA, transa, *inputB, transb, m, n, k, alpha, workspace, workspaceSize, &alpha_ptr, stream); @@ -1335,6 +1345,8 @@ void hipblaslt_gemm(const Tensor *inputA, void *pre_gelu_out = outputPreGelu->data.dptr; const bool gelu = pre_gelu_out != nullptr; const bool use_fp8 = is_fp8_dtype(param.Atype) || is_fp8_dtype(param.Btype); + // MXFP4 native path: FP4 E2M1 data + UE8M0 block-32 scales through hipBLASLt. + const bool use_mxfp4 = is_mxfp4_scaling(inputA->scaling_mode); const hipDataType A_type = get_hipblaslt_dtype(param.Atype); const hipDataType B_type = get_hipblaslt_dtype(param.Btype); @@ -1417,7 +1429,7 @@ void hipblaslt_gemm(const Tensor *inputA, #else constexpr int scaling_mode = 0; #endif - if (use_fp8) { + if (use_fp8 || use_mxfp4) { // Split accumulator. const int8_t fastAccuMode = (use_split_accumulator) ? 0 : 1; /* @@ -1519,18 +1531,18 @@ void hipblaslt_gemm(const Tensor *inputA, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); - if (use_fp4) { + if (use_nvfp4) { int32_t pointer_mode = HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST; NVTE_CHECK_HIPBLASLT(hipblasLtMatmulDescSetAttribute( operationDesc, HIPBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode))); } - GemmAlgoCache::Key gemm_cfg(algoCache.device_cap(device_id), A_type, B_type, D_type, - use_fp8 ? bias_type : (hipDataType)-1, - (use_fp8 && gelu) ? aux_type : (hipDataType)-1, + GemmAlgoCache::Key gemm_cfg(algoCache.device_cap(device_id), A_type, B_type, D_type, + (use_fp8 || use_mxfp4) ? bias_type : (hipDataType)-1, + ((use_fp8 || use_mxfp4) && gelu) ? aux_type : (hipDataType)-1, m, n, k, param.lda, param.ldb, ldd, param.transA, param.transB, scaling_mode, epilogue, - use_fp4); + use_nvfp4); GemmAlgoCache::Algo cached_algo; if (algoCache.find(gemm_cfg, workspaceSize, cached_algo) == 0 || !cached_algo.algo.has_value()) { @@ -1976,6 +1988,27 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK(inputBias->data.dptr == nullptr, "hipBLASlt MXFP8 GEMM does not support bias."); #endif } + // MXFP4 GEMM capability gate. hipBLASLt shipped the MXFP4 kernels in >=1.3. + if (inputA->scaling_mode == NVTE_MXFP4_1D_SCALING || + inputB->scaling_mode == NVTE_MXFP4_1D_SCALING) { +#if (HIPBLASLT_VERSION_MAJOR > 1) || (HIPBLASLT_VERSION_MAJOR == 1 && HIPBLASLT_VERSION_MINOR >= 3) + NVTE_CHECK(cuda::sm_arch() == 95, "MXFP4 GEMM is only supported on gfx950"); + NVTE_CHECK((k % 256) == 0, + "hipBLASLt MXFP4 GEMM requires K to be a multiple of 256 (got K=", k, ")"); + NVTE_CHECK((m % 32) == 0, "hipBLASLt MXFP4 GEMM requires M to be a multiple of 32 (got M=", m, ")"); + NVTE_CHECK((n % 32) == 0, "hipBLASLt MXFP4 GEMM requires N to be a multiple of 32 (got N=", n, ")"); + NVTE_CHECK(outputD->data.dtype == DType::kBFloat16 || outputD->data.dtype == DType::kFloat32, + "hipBLASLt MXFP4 GEMM supports only BF16 or FP32 output"); + NVTE_CHECK(inputBias->data.dptr == nullptr, + "hipBLASLt MXFP4 GEMM does not support fused bias"); + NVTE_CHECK(outputPreGelu->data.dptr == nullptr, + "hipBLASLt MXFP4 GEMM does not support fused GELU"); + NVTE_CHECK(*reinterpret_cast(beta_ptr) == 0.0f, + "hipBLASLt MXFP4 GEMM does not support accumulate (beta != 0)"); +#else + NVTE_ERROR("MXFP4 GEMM requires hipBLASLt >= 1.3"); +#endif + } const int lda = is_transa ? k : m; const int ldb = is_transb ? n : k; diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 6a1e84dfeb..e943654fa8 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -427,21 +427,23 @@ def general_gemm( # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] - # MXFP4 GEMM: route to AITER a4w4 ASM kernels + # MXFP4 GEMM: route to AITER a4w4 ASM kernels, unless the hipBLASLt backend is + # opted in via NVTE_ROCM_USE_HIPBLASLT_MXFP4 from ..tensor.storage.mxfp4_tensor_storage import MXFP4TensorStorage if isinstance(A, MXFP4TensorStorage) or isinstance(B, MXFP4TensorStorage): - result = mxfp4_gemm( - A, - B, - layout=layout, - out_dtype=out_dtype if out_dtype is not None else torch.bfloat16, - bias=bias, - out=out, - grad=grad, - accumulate=accumulate, - ) - return result, None, None, None + if os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") != "1": + result = mxfp4_gemm( + A, + B, + layout=layout, + out_dtype=out_dtype if out_dtype is not None else torch.bfloat16, + bias=bias, + out=out, + grad=grad, + accumulate=accumulate, + ) + return result, None, None, None if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance(B, Float8BlockwiseQTensorStorage): # FP8 block-scaling requires split accumulator diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 196cd31dcc..431c25dab7 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1753,12 +1753,21 @@ def make_quantizers(self) -> list: use_hadamard = self.recipe.use_hadamard + # The hipBLASLt MXFP4 GEMM path consumes plain (un-shuffled) FP4 data and plain + # (non-swizzled) UE8M0 scales; only the AITER a4w4 backend needs the 16x16 weight + # shuffle and swizzled scales + use_hipblaslt = os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") == "1" + swizzled_scales = not use_hipblaslt + if self.mode == "forward": def _make_quantizer(idx: int): is_activation = idx % 3 == 0 is_weight = idx % 3 == 1 - if is_activation: + if use_hipblaslt: + shuffle_rowwise_data = False + shuffle_columnwise_data = False + elif is_activation: shuffle_rowwise_data = False shuffle_columnwise_data = True elif is_weight: @@ -1773,7 +1782,7 @@ def _make_quantizer(idx: int): columnwise=True, shuffle_rowwise_data=shuffle_rowwise_data, shuffle_columnwise_data=shuffle_columnwise_data, - with_gemm_swizzled_scales=True, + with_gemm_swizzled_scales=swizzled_scales, use_hadamard=use_hadamard, ) @@ -1787,7 +1796,7 @@ def _make_quantizer(idx: int): columnwise=True, shuffle_rowwise_data=False, shuffle_columnwise_data=False, - with_gemm_swizzled_scales=True, + with_gemm_swizzled_scales=swizzled_scales, use_hadamard=use_hadamard, ) for _ in range(self.num_quantizers) From 9312928c5d708e4cf44639a2560dd1d3156ca4aa Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Fri, 7 Aug 2026 19:33:58 +0000 Subject: [PATCH 2/2] Added tests --- tests/cpp/operator/test_cublaslt_gemm.cu | 11 +-- tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py | 70 +++++++++----------- 2 files changed, 33 insertions(+), 48 deletions(-) diff --git a/tests/cpp/operator/test_cublaslt_gemm.cu b/tests/cpp/operator/test_cublaslt_gemm.cu index 339d27bfbb..a00e2fc559 100644 --- a/tests/cpp/operator/test_cublaslt_gemm.cu +++ b/tests/cpp/operator/test_cublaslt_gemm.cu @@ -37,8 +37,7 @@ std::vector> test_case_sizes_mxfp8 = { }; // MXFP4 (m, k, n): M/N multiples of 32 (block size), K a multiple of 256 (see rocm_gemm.cu -// gate; K%128-not-%256 runs but is numerically wrong through TE's padded-scale layout). -// Square + non-square. +// gate). std::vector> test_case_sizes_mxfp4 = { {256, 256, 256}, {128, 256, 512}, @@ -869,9 +868,7 @@ static const float kHostFP4E2M1Table[16] = { 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, -0.0f,-0.5f,-1.0f,-1.5f,-2.0f,-3.0f,-4.0f,-6.0f}; -// CPU-dequantize the row-wise data of an MXFP4 test::Tensor (E2M1 packed two-per-byte + -// UE8M0 block-32 scales, padded scale layout) into the row-wise BF16 buffer of dst, then -// upload. Used to build a high-precision reference for the native MXFP4 GEMM. +// CPU-dequantize the row-wise data of an MXFP4 test::Tensor static void dequantize_mxfp4_rowwise_to_bf16(test::Tensor &src_fp4, test::Tensor &dst_bf16) { const NVTEShape data_shape = src_fp4.rowwise_shape(); // logical [R, C] NVTE_CHECK(data_shape.ndim == 2, "Expected 2D MXFP4 data"); @@ -905,9 +902,7 @@ static void dequantize_mxfp4_rowwise_to_bf16(test::Tensor &src_fp4, test::Tensor dst_bf16.from_cpu(); } -// Native MXFP4 (hipBLASLt) GEMM vs a BF16 reference GEMM built by dequantizing the same -// MXFP4 operands. Restricted to TN layout (the F4F4 kernels are TN in practice), so both -// operands consume row-wise data. +// Native MXFP4 (hipBLASLt) GEMM vs a BF16 reference GEMM built by dequantizing the same MXFP4 operands. template void performMxfp4Test(const TestParams ¶ms) { DType dtype = TypeInfo::dtype; diff --git a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py index 5031a9939b..e913409377 100644 --- a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py +++ b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py @@ -2,14 +2,19 @@ # # See LICENSE for license information. -"""MXFP4 GEMM tests: native GEMM (AITER a4w4 or hipBLASLt) vs Python reference GEMM. +"""MXFP4 GEMM tests. -Requires the aiter package (ROCm gfx950 only). The hipBLASLt backend is exercised by -setting NVTE_ROCM_USE_HIPBLASLT_MXFP4=1, which routes MXFP4 GEMMs through the native -hipBLASLt path (rocm_gemm.cu) instead of AITER. +This module tests native MXFP4 GEMM implementations against Python reference GEMMs: +- AITER a4w4 kernels +- hipBLASLt F4F4 kernels + +Both backends are exercised automatically (parametrized). Each test toggles +NVTE_ROCM_USE_HIPBLASLT_MXFP4 via monkeypatch to route general_gemm to the backend +under test, so the caller does not need to set any environment variable. + +Requires the aiter package (ROCm gfx950 only). """ -import os import pytest import torch import transformer_engine.pytorch as te @@ -115,10 +120,6 @@ def _native_gemm(w_mxfp4, x_mxfp4, out_dtype, out, accumulate): return y_native -def _sanitize(t): - return torch.where(t.isnan(), torch.zeros_like(t), t) - - def check_mxfp4_gemm_versus_reference( x_dtype: torch.dtype, w_dtype: torch.dtype, @@ -148,14 +149,13 @@ def check_mxfp4_gemm_versus_reference( assert y_ref is not y_native assert not torch.isnan(y_ref.float()).all(), "All reference elements are NaN" - torch.testing.assert_close(_sanitize(y_native), _sanitize(y_ref), atol=8e-3, rtol=8e-3) + y_ref = torch.where(y_ref.isnan(), torch.zeros_like(y_ref), y_ref) + y_native = torch.where(y_native.isnan(), torch.zeros_like(y_native), y_native) - -_BACKENDS = ["aiter"] -if _use_hipblaslt := os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") == "1": - _BACKENDS = ["hipblaslt"] + torch.testing.assert_close(y_native, y_ref, atol=8e-3, rtol=8e-3) +_BACKENDS = ["aiter", "hipblaslt"] @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.skipif(not _aiter_available, reason="aiter package not available") @pytest.mark.parametrize( @@ -182,15 +182,14 @@ def test_mxfp4_gemm_versus_reference( out_dtype: torch.dtype, accumulate: bool, backend: str, + monkeypatch, ): - # The hipBLASLt MXFP4 path requires K % 256 == 0 (TE's padded UE8M0 scale layout; K%128 runs - # but is numerically wrong -- design doc §3.5/§10) and does not support accumulate yet - # (§10 #4). AITER covers the full matrix. if backend == "hipblaslt": if K % 256 != 0: pytest.skip("hipBLASLt MXFP4 currently requires K to be a multiple of 256") if accumulate: pytest.skip("hipBLASLt MXFP4 does not support accumulate yet") + monkeypatch.setenv("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "1" if backend == "hipblaslt" else "0") check_mxfp4_gemm_versus_reference( x_dtype=x_dtype, w_dtype=w_dtype, @@ -205,10 +204,6 @@ def test_mxfp4_gemm_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.skipif(not _aiter_available, reason="aiter package not available") -@pytest.mark.skipif( - os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") != "1", - reason="hipBLASLt MXFP4 backend not enabled (set NVTE_ROCM_USE_HIPBLASLT_MXFP4=1)", -) @pytest.mark.parametrize( "M, K, N", [ @@ -217,7 +212,7 @@ def test_mxfp4_gemm_versus_reference( (4096, 512, 3072), ], ) -def test_mxfp4_gemm_hipblaslt_matches_aiter(M: int, K: int, N: int): +def test_mxfp4_gemm_hipblaslt_matches_aiter(M: int, K: int, N: int, monkeypatch): """Cross-check: same inputs, hipBLASLt vs AITER MXFP4 GEMM produce close results.""" device = "cuda" dtype = torch.bfloat16 @@ -227,21 +222,16 @@ def test_mxfp4_gemm_hipblaslt_matches_aiter(M: int, K: int, N: int): x = torch.randn((M, K), dtype=dtype, device=device) w = torch.randn((N, K), dtype=dtype, device=device) - prev = os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4") - try: - # AITER path (shuffled operands) - os.environ["NVTE_ROCM_USE_HIPBLASLT_MXFP4"] = "0" - x_a, w_a = _quantize_pair(x, w, M, K, N, dtype, dtype, device, shuffle=True) - y_aiter = _native_gemm(w_a, x_a, dtype, None, False) - - # hipBLASLt path (plain operands) - os.environ["NVTE_ROCM_USE_HIPBLASLT_MXFP4"] = "1" - x_h, w_h = _quantize_pair(x, w, M, K, N, dtype, dtype, device, shuffle=False) - y_hip = _native_gemm(w_h, x_h, dtype, None, False) - finally: - if prev is None: - os.environ.pop("NVTE_ROCM_USE_HIPBLASLT_MXFP4", None) - else: - os.environ["NVTE_ROCM_USE_HIPBLASLT_MXFP4"] = prev - - torch.testing.assert_close(_sanitize(y_hip), _sanitize(y_aiter), atol=8e-3, rtol=8e-3) + # AITER path (shuffled operands) + monkeypatch.setenv("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") + x_a, w_a = _quantize_pair(x, w, M, K, N, dtype, dtype, device, shuffle=True) + y_aiter = _native_gemm(w_a, x_a, dtype, None, False) + + # hipBLASLt path (plain operands) + monkeypatch.setenv("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "1") + x_h, w_h = _quantize_pair(x, w, M, K, N, dtype, dtype, device, shuffle=False) + y_hip = _native_gemm(w_h, x_h, dtype, None, False) + + y_aiter = torch.where(y_aiter.isnan(), torch.zeros_like(y_aiter), y_aiter) + y_hip = torch.where(y_hip.isnan(), torch.zeros_like(y_hip), y_hip) + torch.testing.assert_close(y_hip, y_aiter, atol=8e-3, rtol=8e-3)