blockwise fp8 gemm integration for gfx942 and gfx950

[asdfvg123]

Description

Integrating blockwise FP8 gemm for gfx942 and gfx950 using hipkittens. Using 2 submodules for now. The hipkittens submodule uses separate branch for CDNA3 and CDNA4.

Fixes # (issue)

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:

• Change A
  separate the kernel for gfx942 and gfx950. The host selects the kernel in runtime.
  The kernel supports 1d2d, 1d1d GEMM with output bf16, fp32, fp16 / input fp8e5m2, fp8e4m3 / and different epilogues.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[github-actions]

Claude Walkthrough

Intent. Adds blockwise FP8 GEMM support on ROCm for both gfx942 (CDNA3) and gfx950 (CDNA4) using HipKittens kernels. Extends the prior gfx950-only HipKittens MXFP8 path with new blockwise-scaled FP8 kernels (1Dx1D and 1Dx2D scaling), and pulls in a separate HipKittens submodule for CDNA3 since the two archs use different kernel branches.

Key changes.

• New git submodule 3rdparty/hipkittens_cdna3 (branch yeonsoo/cdna3_fp8) alongside the existing 3rdparty/hipkittens (CDNA4) — see .gitmodules and .github/workflows/rocm-wheels-build.yml:92.
• New runtime arch-dispatch shim transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.cpp:24 (uses transformer_engine::cuda::sm_arch() to pick 94 vs 95).
• Per-arch kernel dispatch + device code under cdna3/ and cdna4/ (each providing blockwise_fp8_gemm.cpp, blockwise_fp8_gemm.h, and blockwise_fp8_gemm_device.cuh).
• Existing MXFP8 sources moved to transformer_engine/common/gemm/kittens/cdna4/ (arch-specific now).
• transformer_engine/common/gemm/kittens/CMakeLists.txt refactored into a try_enable_hipkittens_gemm() function that builds separate OBJECT libs per arch and links them into a single kittens_gemm shared lib; enables when either gfx942 or gfx950 is in CMAKE_HIP_ARCHITECTURES.
• transformer_engine/common/gemm/rocm_gemm.cu:2010 — cublas_gemm short-circuits to kittens_blockwise_fp8_gemm when both operands are blockwise FP8, with a block of NVTE_CHECK preconditions.
• transformer_engine/common/gemm/rocm_gemm.cu:400 — CanonicalizeGemmInput now allows A and B to have different blockwise scaling modes (1D vs 2D), and threads per-operand A_scaling_mode/B_scaling_mode through GemmParam so the dispatcher can distinguish 1Dx1D from 1Dx2D.
• transformer_engine/common/common.h:115 — new is_blockwise_fp8_scaling() helper.
• transformer_engine/pytorch/quantization.py — removes the separate is_fp8_block_scaling_quantization_available API and folds it into is_fp8_block_scaling_available, which now returns True on ROCm for gfx94x/gfx95x (previously always False on ROCm for GEMM).

Walkthrough.

Kernel layer. Two independent HipKittens builds are kept side-by-side because CDNA3 and CDNA4 need different HipKittens branches. Each arch has its own blockwise_fp8_gemm.cpp implementing a blockwise_gfx942::kittens_blockwise_fp8_gemm_impl_cdna3 / blockwise_gfx950::kittens_blockwise_fp8_gemm_impl_cdna4 symbol in its own namespace, and a shared blockwise_fp8_gemm_device.cuh is inlined into that namespace (the #include sits inside the namespace blockwise_gfx942 / blockwise_gfx950 block, keeping symbols per-arch). The kernels use different MFMA_K sizes (32 for CDNA3, 128 for CDNA4) reflecting the different MFMA instruction shapes. The top-level blockwise_fp8_gemm.cpp picks between them at runtime — the arch is not known at library build time when both are compiled in.

CMake. Previously the kittens lib was gfx950-only. It is now a function that conditionally adds per-arch OBJECT libraries (kittens_gemm_cdna3 compiled with -DKITTENS_CDNA3 for gfx942, kittens_gemm_cdna4 with -DKITTENS_CDNA4 for gfx950) and combines them into a single kittens_gemm shared library along with the arch-neutral dispatcher. If neither arch is targeted the whole thing turns itself off.

rocm_gemm.cu integration. The blockwise path is checked before mxfp8/hipblasLt. It enforces: output must be bf16/fp32/fp16; B must be 1D-scaled (so only 1Dx1D or 1Dx2D permitted); TT layout is rejected; e5m2 x e5m2 is rejected; gelu epilogue requires grad and bf16 output; bias with grad is rejected; use_split_accumulator is mandatory; k % 16 == 0 and m % 8 == 0 (n % 16 for 1D scaling). The change to CanonicalizeGemmInput is important — it was previously asserting A.scaling_mode == B.scaling_mode, which would have blocked the legitimate mixed-mode 1Dx2D case, so the check is loosened only when both are blockwise. Data/scale pointer selection for blockwise deliberately mirrors the previous manual dispatch (inputA_col = !is_transa, inputB_col = is_transb), noted in comments.

Python side. The prior split between "quantization available" (cast) and "GEMM available" (kernels) is collapsed since ROCm now has both. is_fp8_block_scaling_available returns true on gfx94x/gfx95x; the previously separate is_fp8_block_scaling_quantization_available symbol is removed from pytorch/__init__.py and quantization.py. constants.py swaps is tex.DType.kFloat8E4M3 for == (identity to equality on the enum values).

Testing. No new test files. Existing tests/pytorch/test_float8_blockwise_gemm_exact.py is adapted to run on ROCm: a helper rocm_blockwise_unsupported_reason skips the currently unsupported combos (non-1Dx1D/1Dx2D, TT layout) on gfx94x/gfx95x; hardcoded torch.float8_e4m3fn/torch.float8_e5m2 are replaced with get_torch_float8_e4m3_type()/get_torch_float8_e5m2_type() for portability; constraint tests loosen their expected error class/message on ROCm. tests/pytorch/test_float8_blockwise_scaling_exact.py replaces the two-availability check with the unified one. tests/pytorch/test_numerics.py:2138 skips grouped-linear tests with the blockwise recipe on ROCm (grouped GEMM not yet supported).

Notes for reviewers.

• The new 3rdparty/hipkittens_cdna3 submodule points at a personal fork (https://github.com/asdfvg123/HipKittens.git, branch yeonsoo/cdna3_fp8, pinned to f9805185) — likely needs to be pointed at an upstream/ROCm-owned branch before merge.
• Public Python API surface change: is_fp8_block_scaling_quantization_available is removed from transformer_engine.pytorch. Callers must migrate to is_fp8_block_scaling_available. Since ROCm now supports the GEMM, they collapse to the same answer on gfx94x+, but this is technically a breaking export change.
• Unsupported combos are hard errors via NVTE_CHECK at the GEMM entry — no fallback to hipblasLt for e5m2 x e5m2, TT layout, or gelu-without-grad. Frameworks currently relying on such fallbacks would need explicit guards.
• The two KittensDType / KittensScalingMode enum blocks are duplicated across the top-level blockwise_fp8_gemm.h and the per-arch headers under cdna3/ and cdna4/ — guarded by matching #ifndef sentinels to avoid ODR issues when multiple headers are included together.
• Grouped GEMM blockwise support is explicitly deferred (test skip in test_numerics.py).

---

Generated by Claude. To request a code review, comment /claude review.

[github-actions]

Submodule points to a personal user fork on an unreleased branch — must be relocated before merge.

• url = https://github.com/asdfvg123/HipKittens.git is a personal repo, not a canonical one. The sibling 3rdparty/hipkittens correctly points at HazyResearch/HipKittens.git. Anyone cloning ROCm/TransformerEngine will fail git submodule update --init the moment this repo goes private or gets renamed. It also blocks reproducible ROCm CI wheel builds (this submodule is now in the init list in .github/workflows/rocm-wheels-build.yml).
• The pinned branch yeonsoo/cdna3_fp8 is likewise a personal in-flight branch; the CDNA3 kittens sources need to land in a persistent upstream branch (HazyResearch or a ROCm-owned mirror) and this URL/branch should be updated before merge.

[github-actions]

Error messages have M/N swapped relative to the values being checked.

The check is on m but the message says "N divisible by 8", and the check on n says "M divisible by 16". A user hitting either of these gets a message that names the wrong dimension, which will be very confusing to debug (especially since the arch impls internally swap kM=N/kN=M — a reader seeing the message will not know which side of the swap is being enforced).

Fix: name the dimension after the value actually being checked, e.g.

Suggested change

	NVTE_CHECK((m % 8) == 0,
	"Blockwise FP8 GEMM requires N divisible by 8");
	NVTE_CHECK(inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || (n % 16) == 0,
	"Blockwise FP8 GEMM requires M divisible by 16 for 1D scaling");
	NVTE_CHECK((k % 16) == 0,
	"Blockwise FP8 GEMM requires K divisible by 16");
	NVTE_CHECK((m % 8) == 0,
	"Blockwise FP8 GEMM requires M divisible by 8");
	NVTE_CHECK(inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || (n % 16) == 0,
	"Blockwise FP8 GEMM requires N divisible by 16 for 1D scaling");

(If the intent was to speak in the kernel's post-swap convention, then the two messages need to be swapped with each other — either way the current pairing is wrong.)

Also note: since the preceding NVTE_CHECK(inputB->scaling_mode == NVTE_BLOCK_SCALING_1D, ...) already forces B to be 1D, the inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || guard on the last check is always false-branch — you can drop it and just check (n % 16) == 0 unconditionally.

[github-actions]

This weakens the "constraint enforced" tests to a near-tautology on ROCm.

Setting expected_err_msg = None here means the test now accepts any RuntimeError from the ROCm code path, regardless of whether it comes from the constraint we're supposedly enforcing or from an unrelated failure (setup error, shape mismatch, OOM, etc.). Every call site of cublas_gemm_test_constraint_enforced is paired with a specific constraint (illegal dtype, illegal 2Dx2D, split-accumulator, bgrad, gelu-unsupported…), and none of that specificity is preserved on ROCm.

Prefer either:

• Mapping the cuBLAS-flavored expected_err_msg to the corresponding ROCm one (the NVTE_CHECK strings you added in rocm_gemm.cu above, e.g. "Blockwise FP8 GEMM does not support ..."), so the test still asserts the right thing, or
• Passing a per-callsite ROCm expected message via the caller instead of blanket-nulling it here.

[github-actions]

This change also affects the CUDA path — please classify and document.

Custom_DType_Dict is shared code that both CUDA and ROCm execute. Changing is to == on the pybind11 enum is likely a real bug fix (enum identity via is is brittle across module reloads / repeated pybind imports), not a ROCm-specific tweak. Per the ROCm-fork upstream-compat rule, generic bug fixes to shared code are fine but should be called out explicitly (PR description or a short comment) so it can be upstreamed and future IFU merges know why the line diverges from NVIDIA's TE.

If instead this change is only motivated by a ROCm-observed failure, please add a one-line comment explaining that, or guard it with IS_HIP_EXTENSION.

[github-actions]

This drops is_fp8_block_scaling_quantization_available — that is a public API change.

is_fp8_block_scaling_quantization_available and the underlying check_fp8_block_scaling_quantization_support were being removed together with the corresponding re-export from transformer_engine/pytorch/__init__.py and the __all__ entry above. That symbol was exposed on transformer_engine.pytorch, so any downstream code (or ROCm-side tests / benchmarks / megablocks / megatron patches) that imported is_fp8_block_scaling_quantization_available(...) will now break with ImportError.

Two options:

• If block scaling GEMM being available implies block scaling quantization is available (the new _compute_fp8_block_scaling_support gating on (9,4)/(9,5) suggests yes), keep the old API as a thin alias to is_fp8_block_scaling_available and mark it deprecated, so we don't break importers.
• If we're intentionally taking the API away, please call this out in the PR description as a breaking change.

[github-actions]

Claude review — PR 658

Reviewed the ROCm-only 3-dot diff (HEAD^1...HEAD^2, 23 files, ~2.3k added): new HipKittens-based blockwise FP8 GEMM kernels for gfx942 (CDNA3) and gfx950 (CDNA4), corresponding dispatch in rocm_gemm.cu, ROCm gating in pytorch/quantization.py, test enablement, and a new hipkittens_cdna3 submodule.

Verdict: solid direction, but there are a few things worth fixing before merge. Findings posted inline:

• Blocking: 3rdparty/hipkittens_cdna3 submodule URL points to a personal fork (asdfvg123/HipKittens.git) on an in-flight branch — needs to move to a canonical repo/branch before merge, otherwise clones and ROCm CI wheel builds are one repo-rename away from breaking.
• Correctness: In rocm_gemm.cu, two NVTE_CHECK messages have M/N labels swapped relative to the values being checked.
• Test quality: test_float8_blockwise_gemm_exact.py::cublas_gemm_test_constraint_enforced sets expected_err_msg = None on ROCm, which turns every "constraint enforced" test into "any RuntimeError passes".
• Upstream compat: constants.py is → == on a shared code path — likely a real generic bug fix that should be documented (or guarded) per the ROCm fork rules.
• API: is_fp8_block_scaling_quantization_available is removed from the public transformer_engine.pytorch surface. Either restore as an alias or call out as a breaking change in the PR description.

Copyright headers: OK — every touched file carries an up-to-date AMD 2026 header in the correct format; NVIDIA lines preserved on modified upstream files.

[asdfvg123]

1. check_fp8_block_scaling_quantization_support
   These were temporary APIs introduced in PR enable blockwise FP8 quantization on rocm #609 to enable blockwise FP8 quantization only. Now that GEMM is also supported, they are folded into is_fp8_block_scaling_available and removed. This is an intentional public symbol removal with no known downstream users.

2. constants.py FP8 dtype lookup fix
   Custom_DType_Dict.__missing__ now compares the enum key with == instead of is. This is a bug fix on the rocm side. tex.DType values returned from the C++ bindings (i.e. tensor._fp8_dtype) are fresh pybind11 instances that compare equal to but are not identical to the module-level singletons, so the identity check could cache the e4m3 key as e5m2. The dict is a process-global singleton, so this was a latent bug that surfaced when a C++-Enum lookup was the first e4m3 lookup in the process (seen in test_float8block_linear_fused_adam_master_weights). The Custom_DType_Dict class is created in this fork (upstream does not have the class), so no upstream compatibility concern.