Integrate MXFP4 hipblaslt GEMM support - #697
Conversation
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
|
|
||
| # 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 |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
Will keep this, and make this change in a separate PR.
I am currently doing performance runs and see which to enable by default. |
There was a problem hiding this comment.
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
MXFP4Quantizerdirectly with explicit shuffle flags, so they would still pass if this environment/recipe mapping regressed. Add a recipe-state or module-level test that togglesNVTE_ROCM_USE_HIPBLASLT_MXFP4and verifies both plain anduse_swizzled_scales=Truequantizers (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
|
Planning to rebase and merge after the approvals. Thanks |
…veergopu/maxfp4-hipblaslt-integration
Description
Adds a native MXFP4 GEMM path through hipBLASLt (F4F4 kernels) on gfx950 / ROCm ≥ 7.13 /
hipBLASLt ≥ 1.3, alongside the existing AITER
a4w4backend. Until now an MXFP4 GEMM alwaysrouted to AITER and never reached
rocm_gemm.cu; hipBLASLt 1.3 now ships FP4×FP4 + UE8M0block-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_MXFP4unset,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
Changes
Please list the changes introduced in this PR:
transformer_engine/common/gemm/rocm_gemm.cu):DType::kFloat4E2M1 → HIP_R_4F_E2M1inget_hipblaslt_dtype()and the algo-cachetype_name_map.is_mxfp_scalingcanonicalization and the existingVEC32_UE8M0block-scale arm (wires A/B scale pointers/modes); no separate branch needed.use_nvfp4) so NVFP4 keeps thefallback while MXFP4 stays native.
cublas_gemm: compile-time hipBLASLt ≥ 1.3, runtimegfx950,
K % 256,M/N % 32, BF16/FP32 output only, no bias/GELU,beta == 0.pytorch/cpp_extensions/gemm.py): route MXFP4 to hipBLASLt whenNVTE_ROCM_USE_HIPBLASLT_MXFP4=1, otherwise AITERa4w4(default).pytorch/quantization.py):MXFP4BlockScalingRecipeState.make_quantizersemitsplain (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/cpp/operator/test_cublaslt_gemm.cu):OperatorTestMXFP4compares thenative MXFP4 GEMM against a BF16 reference built by CPU-dequantizing the same operands (TN,
BF16/FP32 output, K%256; no MXFP4
nvte_dequantizeexists on ROCm).tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py): parametrized over bothbackends (routed automatically via
monkeypatch, no env var required) vsMXFP4QuantizerRef,plus a direct hipBLASLt-vs-AITER cross-check.
MXFP4 GEMM Performance: hipBLASLt vs AITER
Speedup = aiter_ms / hipblaslt_ms (> 1.0 means the hipBLASLt variant is faster than AITER).
Summary (hipBLASLt vs AITER)
hipblaslt_plain
hipblaslt_swizzled
Takeaways
NVTE_ROCM_USE_HIPBLASLT_MXFP4unset) 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