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: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.

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.

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

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

transformer_engine/pytorch/quantization.py:1765

The new backend-dependent quantizer selection is not exercised by the added GEMM tests: those tests instantiate MXFP4Quantizer directly with explicit shuffle flags, so they would still pass if this environment/recipe mapping regressed. Add a recipe-state or module-level test that toggles NVTE_ROCM_USE_HIPBLASLT_MXFP4 and verifies both plain and use_swizzled_scales=True quantizers (including forward weight/activation and backward slots).

 use_hipblaslt = bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0")))
use_swizzled = use_hipblaslt and self.recipe.use_swizzled_scales
# AITER path swizzles scales; hipBLASLt path swizzles only when the recipe opts in.
# FP4 data shuffle stays off on the hipBLASLt path regardless
swizzled_scales = use_swizzled if use_hipblaslt else True

VeeraRajasekhar · 2026-08-26T19:02:27Z

Planning to rebase and merge after the approvals. Thanks

matthiasdiener · 2026-08-27T20:38:12Z

 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);
+ const bool use_mxfp4 = is_mxfp4_scaling(inputA->scaling_mode);


nit: should this also check for the B type (like the fp8 above)?

they are asserted to be equal in line 471

…veergopu/maxfp4-hipblaslt-integration

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

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

CopilotAI reviewed Aug 11, 2026
View reviewed changes

Comment threadtests/cpp/operator/test_cublaslt_gemm.cu Outdated
Comment threadtests/pytorch/mxfp4/test_mxfp4_gemm_exact.py Outdated
Comment threadtransformer_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 threadtransformer_engine/common/gemm/rocm_gemm.cu Outdated

aris134 reviewed Aug 12, 2026
View reviewed changes

Comment threadtransformer_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

Comment threadtests/cpp/operator/test_cublaslt_gemm.cu Outdated
Comment threadtests/cpp/operator/test_cublaslt_gemm.cu Outdated
Comment threadtests/cpp/operator/test_cublaslt_gemm.cu Outdated

Addressed reviews
b6fc33b

VeeraRajasekhar requested review from aris134 and ipanfilo and a balanced review from Copilot August 24, 2026 13:29

Copilotstarted reviewing on behalf of VeeraRajasekhar August 24, 2026 13:30 View session

CopilotAI reviewed Aug 24, 2026
View reviewed changes

Comment threadtransformer_engine/common/gemm/rocm_gemm.cu
Comment threadtransformer_engine/common/gemm/rocm_fp4_e2m1_table.h

Addressed reviews
5f5a224

matthiasdiener reviewed Aug 27, 2026
View reviewed changes

Merge branch 'dev' ofhttps://github.com/ROCm/TransformerEngineinto …
35322ca
…veergopu/maxfp4-hipblaslt-integration

VeeraRajasekhar requested a review from matthiasdiener August 28, 2026 02:23

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!

CopilotAI 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

Uh oh!

Uh oh!

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

VeeraRajasekhar commented Aug 26, 2026

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