Integrate MXFP4 hipblaslt GEMM support by VeeraRajasekhar · Pull Request #697 · ROCm/TransformerEngine

VeeraRajasekhar · 2026-08-07T15:04:06Z

Description

Adds a native MXFP4 GEMM path through hipBLASLt (F4F4 kernels) on gfx950 / ROCm ≥ 7.13 /
hipBLASLt ≥ 1.3, alongside the existing AITER a4w4 backend. Until now an MXFP4 GEMM always
routed to AITER and never reached rocm_gemm.cu; hipBLASLt 1.3 now ships FP4×FP4 + UE8M0
block-32 kernels, so this wires MXFP4 into the hipBLASLt path (mirroring the MXFP8 native path)
behind an opt-in toggle, enabling A/B benchmarking against AITER.

The new path is opt-in and regression-safe: with NVTE_ROCM_USE_HIPBLASLT_MXFP4 unset,
MXFP4 still routes to AITER with the existing shuffled quantization, so current behavior and
tests are unchanged.

Fixes # (N/A — internal ROCm enablement)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

C++ (transformer_engine/common/gemm/rocm_gemm.cu):
- Map DType::kFloat4E2M1 → HIP_R_4F_E2M1 in get_hipblaslt_dtype() and the algo-cache
  type_name_map.
- Route MXFP4 through the shared is_mxfp_scaling canonicalization and the existing
  VEC32_UE8M0 block-scale arm (wires A/B scale pointers/modes); no separate branch needed.
- Gate the FP4 dequant→BF16 fallback on the scaling mode (use_nvfp4) so NVFP4 keeps the
  fallback while MXFP4 stays native.
- Add an MXFP4 capability gate in cublas_gemm: compile-time hipBLASLt ≥ 1.3, runtime
  gfx950, K % 256, M/N % 32, BF16/FP32 output only, no bias/GELU, beta == 0.
Python (pytorch/cpp_extensions/gemm.py): route MXFP4 to hipBLASLt when
NVTE_ROCM_USE_HIPBLASLT_MXFP4=1, otherwise AITER a4w4 (default).
Python (pytorch/quantization.py): MXFP4BlockScalingRecipeState.make_quantizers emits
plain (un-shuffled) FP4 data + plain UE8M0 scales when the toggle is on, so operand layout
matches the GEMM backend by construction (AITER-shuffled otherwise).
Tests (C++, tests/cpp/operator/test_cublaslt_gemm.cu): OperatorTestMXFP4 compares the
native MXFP4 GEMM against a BF16 reference built by CPU-dequantizing the same operands (TN,
BF16/FP32 output, K%256; no MXFP4 nvte_dequantize exists on ROCm).
Tests (pytest, tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py): parametrized over both
backends (routed automatically via monkeypatch, no env var required) vs MXFP4QuantizerRef,
plus a direct hipBLASLt-vs-AITER cross-check.

MXFP4 GEMM Performance: hipBLASLt vs AITER

Device: AMD Instinct MI355X
Backends: AITER a4w4 (shuffled weights + swizzled scales) vs hipBLASLt F4F4 in both plain (VEC32_UE8M0) and pre-swizzled (BLK32_UE8M0_32_8_EXT) scale modes. MXFP4, BF16 output.
Passes: forward (TN), dgrad (NN), wgrad (NT) -- true per-pass shapes (fwd contracts hidden, dgrad contracts out, wgrad contracts tokens); same FLOPs, different M/N/K.
Timing: CUDA events + leading kernel, 15 warmup / 50 iters; TFLOPS = 2·m·n·k / t
seqlen: 2048, mbs: [1, 2, 4]

Speedup = aiter_ms / hipblaslt_ms (> 1.0 means the hipBLASLt variant is faster than AITER).

Summary (hipBLASLt vs AITER)

hipblaslt_plain

Pass	Configs	Median speedup	Geomean speedup	hipBLASLt wins
forward (TN)	30	0.68x	0.60x	0/30
dgrad (NN)	30	0.68x	0.62x	0/30
wgrad (NT)	30	0.66x	0.62x	0/30
Overall	90	0.67x	0.61x	0/90

hipblaslt_swizzled

Pass	Configs	Median speedup	Geomean speedup	hipBLASLt wins
forward (TN)	30	0.95x	0.96x	7/30
dgrad (NN)	30	0.95x	0.95x	5/30
wgrad (NT)	30	0.91x	0.92x	2/30
Overall	90	0.94x	0.94x	14/90

Takeaways

hipBLASLt with pre-swizzled scales (mode 1001) is ~0.94x AITER (geomean) -- effectively on par with the hand-tuned a4w4 ASM kernels.
hipBLASLt with plain UE8M0 scales is ~0.61x AITER (geomean); the pre-swizzled scale path is essential for competitive performance.
AITER remains marginally faster on aggregate, so keeping it the default (NVTE_ROCM_USE_HIPBLASLT_MXFP4 unset) is justified; hipBLASLt (swizzled) is a competitive alternative that additionally offers native BF16/FP32 output and the full four-layout set including TT.

Figures

Geometric-mean speedup vs AITER by pass (dashed line = AITER parity).

Per-config hipBLASLt-swizzled speedup vs AITER across contraction sizes (K); most configs land within ~10% of AITER, with several at or above parity.

Script:
benchmark_mxfp4_hipblaslt_vs_aiter.py

Copilot

Pull request overview

Adds opt-in native MXFP4 GEMM support through hipBLASLt while retaining AITER as the default backend.

Changes:

Adds MXFP4 datatype, scaling, capability gating, and dispatch.
Introduces backend-aware quantization layouts.
Adds native and reference GEMM tests.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`transformer_engine/pytorch/quantization.py`	Selects MXFP4 layouts by backend.
`transformer_engine/pytorch/cpp_extensions/gemm.py`	Routes opted-in MXFP4 GEMMs to hipBLASLt.
`transformer_engine/common/recipe/__init__.py`	Adds swizzled-scale configuration.
`transformer_engine/common/gemm/rocm_gemm.cu`	Implements native MXFP4 hipBLASLt support.
`tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py`	Tests Python MXFP4 backend results.
`tests/cpp/operator/test_cublaslt_gemm.cu`	Tests native MXFP4 GEMMs.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

matthiasdiener

Out of curiosity, why is it disabled by default?

Also, for visibility, #675 does a few of the same things as this PR, but I don't see any conflict.

matthiasdiener · 2026-08-11T16:06:17Z

+        # the 32x8 tile order that hipBLASLt's BLK32_UE8M0_32_8_EXT (mode 1001) reads in place
+        # -- opt in via the recipe's use_swizzled_scales flag. The AITER a4w4 backend always
+        # needs the 16x16 weight shuffle and swizzled scales.
+        use_hipblaslt = os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") == "1"


This could perhaps be

Suggested change

use_hipblaslt = os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") == "1"

use_hipblaslt = bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0")))

so that other values besides "1" are also accepted.

matthiasdiener · 2026-08-11T16:09:59Z

      }
    }
  } else if (is_mxfp_scaling(B.scaling_mode)) {
-    // MXFP8


I'd suggest just repeating the adjusted comment you have above, instead of deleting it.

matthiasdiener · 2026-08-11T16:38:03Z

-            accumulate=accumulate,
-        )
-        return result, None, None, None
+        if os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") != "1":


Same recommendation as below regarding the bool change.

matthiasdiener · 2026-08-11T16:44:41Z

+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};


We have a similar table in rocm_gemm.cu, perhaps this can be factored out (e.g., into a header file).

ipanfilo · 2026-08-11T23:18:13Z

+  }
+
+// 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.


What does "§3.5/§10.1" refer to ?

Oo, it referred to my design doc. Let me correct it.

ipanfilo · 2026-08-11T23:19:37Z

+INSTANTIATE_TEST_SUITE_P(OperatorTestMXFP4, Mxfp4GEMMTestSuite,
+                         ::testing::Combine(::testing::ValuesIn(test_case_sizes_mxfp4),
+                                            ::testing::Values(false),         // bias unsupported
+                                            ::testing::Values(kTN),           // TN only


Is support going to extend soon with bias and non TN?

Removed the bias, and added support for all layouts TN,NN,NT,TT and corresponding tests in both cpp and torch

ipanfilo · 2026-08-11T23:43:15Z


-    # 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


We already have NVTE_ROCM_USE_HIPBLASLT_MXFP8. To limit env vars better to combine them to NVTE_ROCM_FORCE_HIPBLASLT. Might be separate PR though

Will keep this, and make this change in a separate PR.

ipanfilo · 2026-08-11T23:45:33Z

+      // MXFP4 can additionally use hipBLASLt's pre-swizzled UE8M0 scale mode
+      // (BLK32_UE8M0_32_8_EXT, "1001").
+      const bool mxfp4_swizzled = use_mxfp4 && inputA->with_gemm_swizzled_scales &&
+                                  inputB->with_gemm_swizzled_scales;


nit: inputB checking is redundant, as they asserted to be equal

VeeraRajasekhar · 2026-08-12T19:27:07Z

Out of curiosity, why is it disabled by default?

Also, for visibility, #675 does a few of the same things as this PR, but I don't see any conflict.

I am currently doing performance runs and see which to enable by default.

aris134 · 2026-08-21T14:53:51Z

+// writing dst[r, c] = CW[c, r] yields the logical [R, C] operand -- i.e. what nvte_dequantize
+// produces for MXFP8. Used for non-TN layouts, where CanonicalizeGemmInput consumes the
+// column-wise buffer for the non-transposed operand.
+static void dequantize_mxfp4_columnwise_to_bf16(test::Tensor &src_fp4, test::Tensor &dst_bf16,


This function is very similar to dequantize_mxfp4_rowwise_to_bf16. The only real difference is which axis is packed/scaled. I think we could have a single parameterized helper

ipanfilo · 2026-08-21T22:55:26Z

 }
 #endif // __HIP_PLATFORM_AMD__

+#ifdef __HIP_PLATFORM_AMD__


nit: back to back #ifdef HIP_PLATFORM_AMD can be merged

ipanfilo · 2026-08-21T23:06:26Z

+  (void)cudaGetDeviceProperties(&prop, 0);
+
+#if !((HIPBLASLT_VERSION_MAJOR > 1) || (HIPBLASLT_VERSION_MAJOR == 1 && HIPBLASLT_VERSION_MINOR >= 3))
+  GTEST_SKIP() << "MXFP4 GEMM requires hipBLASLt >= 1.3";


It is compile time condition so you can guard the test instantiation instead of unconditionally instantiate it and then skip

ipanfilo · 2026-08-21T23:07:52Z

+  if (!(prop.major == 9 && prop.minor == 5)) {
+    GTEST_SKIP() << "MXFP4 GEMM is only supported on gfx950";
+  }
+  if (params.m % 32 || params.n % 32 || params.k % 256) {


it is called with test_case_sizes_mxfp4 so they should already have only supported sizes

Integrate MXFP4 hipblaslt GEMM support

b370117

VeeraRajasekhar self-assigned this Aug 7, 2026

Added tests

9312928

VeeraRajasekhar marked this pull request as ready for review August 9, 2026 17:02

VeeraRajasekhar requested review from ipanfilo, wangye805 and wenchenvincent as code owners August 9, 2026 17:02

Added support for swizzle scales for mxfp4 hipblaslt

d3c01bb

VeeraRajasekhar added the ci-level 3 CI test level 3 label Aug 10, 2026

Updated comments

5caa054

VeeraRajasekhar requested review from alextmagro, aris134 and matthiasdiener and a balanced review from Copilot August 11, 2026 15:48

Copilot started reviewing on behalf of VeeraRajasekhar August 11, 2026 15:50 View session

Copilot AI reviewed Aug 11, 2026

View reviewed changes

Comment thread tests/cpp/operator/test_cublaslt_gemm.cu Outdated

Comment thread tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py Outdated

Comment thread transformer_engine/common/gemm/rocm_gemm.cu Outdated

matthiasdiener reviewed Aug 11, 2026

View reviewed changes

ipanfilo reviewed Aug 11, 2026

View reviewed changes

VeeraRajasekhar added 2 commits August 12, 2026 17:26

Addressed comments

359a250

Enable all transpose layouts (TN/NN/NT/TT) for MXFP4 hipBLASLt GEMM

86a4dc4

aris134 reviewed Aug 12, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/rocm_gemm.cu Outdated

aris134 reviewed Aug 12, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/cpp_extensions/gemm.py Outdated

Addressed reviews

964b2b3

VeeraRajasekhar requested review from aris134, ipanfilo and matthiasdiener August 12, 2026 19:41

aris134 reviewed Aug 21, 2026

View reviewed changes

ipanfilo reviewed Aug 21, 2026

View reviewed changes

	use_hipblaslt = os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0") == "1"
	use_hipblaslt = bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0")))

Uh oh!

Conversation

VeeraRajasekhar commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

MXFP4 GEMM Performance: hipBLASLt vs AITER

Summary (hipBLASLt vs AITER)

hipblaslt_plain

hipblaslt_swizzled

Takeaways

Figures

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

matthiasdiener left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

matthiasdiener Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

VeeraRajasekhar commented Aug 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

VeeraRajasekhar commented Aug 7, 2026 •

edited

Loading

matthiasdiener left a comment •

edited

Loading

matthiasdiener Aug 11, 2026 •

edited

Loading