blockwise fp8 gemm integration for gfx942 and gfx950 by asdfvg123 · Pull Request #658 · ROCm/TransformerEngine

asdfvg123 · 2026-07-06T22:01:12Z

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

…dant HIP guards, revert unnecessary common.h change

… fix

# Conflicts: # tests/cpp/operator/CMakeLists.txt

…hipkittens_optim # Conflicts: # transformer_engine/common/CMakeLists.txt # transformer_engine/common/gemm/kittens/CMakeLists.txt # transformer_engine/common/gemm/rocm_gemm.cu # transformer_engine/pytorch/quantization.py

…gemm_hipkittens_optim # Conflicts: # transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.cpp # transformer_engine/pytorch/quantization.py

github-actions · 2026-07-07T19:34:12Z

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.

…ubmodule

asdfvg123 · 2026-07-09T14:21:18Z

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

ipanfilo · 2026-07-20T14:33:46Z

Is CDNA3 support expected to be ported to HK?

alextmagro · 2026-07-20T14:41:39Z

Is CDNA3 support expected to be ported to HK?

Currently HK supports CDNA3, CDNA4, and UDNA1. There is currently effort with the external HK team to move away from having separate git branches for each arch to having directory level separation.

ipanfilo · 2026-07-20T15:55:33Z

Is CDNA3 support expected to be ported to HK?

Currently HK supports CDNA3, CDNA4, and UDNA1. There is currently effort with the external HK team to move away from having separate git branches for each arch to having directory level separation.

In this case can we avoid having 2 copies of HK submodule?

alextmagro · 2026-07-20T15:58:45Z

Is CDNA3 support expected to be ported to HK?

Currently HK supports CDNA3, CDNA4, and UDNA1. There is currently effort with the external HK team to move away from having separate git branches for each arch to having directory level separation.

In this case can we avoid having 2 copies of HK submodule?

Yes, looks like they merged it in late last week. I think that will require some integration/cmake changes so that the right headers are imported for each kernel. We'll probably want to move our kittens kernels into cdna3/cdna4/udna1 folders as well, so might be better as a follow up ticket?

alextmagro · 2026-07-20T14:53:40Z

    emulated = get_device_compute_capability() >= (10, 0)
    return supported and not emulated

+def rocm_blockwise_unsupported_reason(


Here generally I think we should return a tuple -- True/False, followed by an optional reason. Maybe rename to rocm_blockwise_is_supported too.

alextmagro · 2026-07-20T14:55:32Z

 ):
    if not fp8_blockwise_gemm_supported():
        pytest.skip("CUDA version does not support blockwise FP8 gemm.")
+    if IS_HIP_EXTENSION:


We could also potentially use the rocm function above here too?

yes, now we are using the function

alextmagro · 2026-07-20T16:03:42Z

    is_w_1d_scaled,
 ) -> None:
    # NOTE: BGRAD epilogue is not supported for fp8.
+    if IS_HIP_EXTENSION:


One of the great things about hipKittens is how easy it is to fuse epilogues into kernels. If you have time it might be worth looking into bgrad/accumulator support.

I will leave this as a future support. The current implementation supports the minimal cases that the test is expecting

alextmagro · 2026-07-20T16:08:16Z

 ) -> None:
    # e5m2 by e5m2 not supported.
+    if IS_HIP_EXTENSION:
+        expected_err_msg = "does not support e5m2 by e5m2 inputs"


e5m2 should be supported by hipkittens, have a look at the mxfp8 kernels to see how we template for that

Same here. The test expects e4m3 x e4m3 / e4m3 x e5m2 / e5m2 x e4m3 but not e5m2 x e5m2 and the current implementation also supports the former 3 but not the e5m2 x e5m2. While this is an easy support, I will leave it as a future PR

alextmagro · 2026-07-20T16:11:02Z

            use_split_accumulator,
            is_x_1d_scaled,
            is_w_1d_scaled,
-            expected_err_msg="dimension requirement",


This change should be hip guarded

fixed. To keep the upstream lines untouched, I guarded early-return block, which makes the code a bit longer. Let me know if you have a cleaner approach.

alextmagro · 2026-07-20T16:30:30Z

+
+#ifndef KITTENS_DTYPE_ENUM_DEFINED
+#define KITTENS_DTYPE_ENUM_DEFINED
+enum KittensDType {


I'd rather not have to redefine these. I think we can have a parent folder header file that contains these enums and any other shared values.

fixed. Added kittens_gemm_enums.h

alextmagro · 2026-07-20T16:31:58Z

+#endif  // KITTENS_SCALING_MODE_DEFINED
+
+namespace blockwise_gfx942 {
+void kittens_blockwise_fp8_gemm_impl_cdna3(


Rather than a namespace for gfx942 and gfx950, etc, I think we would be better off with a purely virtual class at the parent folder level, with child class definitions for functions within each architectural folder. That way we can track more easily what is available and what is not for each architecture.

alextmagro · 2026-07-20T16:40:25Z

@@ -0,0 +1,1459 @@
+/*************************************************************************


Most of my comments regarding the cdna3 kernel applies here as well -- otherwise the kernel looks good.

alextmagro · 2026-07-20T16:43:07Z

+option(NVTE_KITTENS_USE_POWER_OF_2_SCALE
+    "Use HW E8M0 MFMA scaling for power-of-2 blockwise FP8 scales (gfx950)" ON)
+
+function(try_enable_hipkittens_gemm)


Not sure this needs to bea function if we are only calling it once?

alextmagro · 2026-07-20T16:45:06Z

+      return;
+    }
+  }
+#endif


if hipkittens is not being used, we should add an nvte_error here for blockwise usage

asdfvg123 added 27 commits June 3, 2026 17:15

enable blockwise FP8 quantization on rocm

8335488

enable blockwise FP8 C++ tests on ROCm, fix wave64 bugs, remove redun…

6226301

…dant HIP guards, revert unnecessary common.h change

fix constexpr chain, arch guard

bdf905e

guard kThreadsPerWarp for wave32, derive fp8 max on AMD host

676d1f0

revert unnecessary guard, add license

f8a0fc5

use non-sync __shfl on ROCm, drop kFullWarpMask

231e381

fix fp32 LDS overflow in vector blockwise on gfx942

e158d3e

drop WarpSyncMask, use __shfl_xor reduce, guard gfx1250 wave32, minor…

70c35df

… fix

use explicit gfx942/host guard

7ede21d

fix guard

8895166

Merge remote-tracking branch 'origin/dev' into yeonsoo/blockwise_fp8

da49b3c

# Conflicts: # tests/cpp/operator/CMakeLists.txt

use shared rocm_subwarp_allreduce for warp amax reduction

64f92a6

guard blockwise fp8 GEMM

21c4f10

add hipkittens blockwise fp8 GEMM for gfx942

96ab7d0

revert the blockwise fp8 gemm guard

b542efc

optimize kernel

1a5c4a5

Merge remote-tracking branch 'origin/dev' into yeonsoo/blockwise_fp8

039ff90

rename submodule

0c4e2dd

integrate gfx950 blockwise fp8 gemm

0803cc4

refactor gemm dispatch, optimize kernel

812fb91

refactor dispatch for gfx942

29b2605

refactor code

a07f0b3

refactor code

37f1e2f

support e5m2, fp32 output, add 1d1d kernel

34720d5

refactor code

9c8a379

refactor code

1f14306

asdfvg123 requested a review from alextmagro July 6, 2026 22:01

Merge remote-tracking branch 'origin/dev' into yeonsoo/blockwise_fp8_…

8224e6b

…gemm_hipkittens_optim # Conflicts: # transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.cpp # transformer_engine/pytorch/quantization.py

asdfvg123 added the ci-level 1 CI test level 1 label Jul 7, 2026

github-actions Bot removed the claude-review label Jul 7, 2026

asdfvg123 added 2 commits July 8, 2026 19:55

drop orphan composable_kernel submodule gitlink in hipkittens_cdna3 s…

2199903

…ubmodule

guard fix

e7241e8

asdfvg123 requested a review from matthiasdiener July 9, 2026 14:12

asdfvg123 marked this pull request as ready for review July 9, 2026 14:12

asdfvg123 requested review from ipanfilo, wangye805 and wenchenvincent as code owners July 9, 2026 14:12

asdfvg123 added 8 commits July 13, 2026 19:28

optimize kernel

0ddeefb

add pow2 scale blockwise gemm kernel

71c02a8

fix spill issue

c502029

clean code

f6a9ae0

clean code

cb96f47

clean code

449d620

add guard for blockwise group gemm

1c172ca

clean code remove debug guard

562c6fe

alextmagro requested changes Jul 20, 2026

View reviewed changes

asdfvg123 added 6 commits July 20, 2026 18:38

fix OOM

6fb0712

refactor code

1f3ee3c

switch gemm dispatch

bd4e737

merge 1d2d and 1d1d kernel for mi300

d6b2005

replace arch namespaces with virtual backend class

7687bde

fix guard

d3ef59c

asdfvg123 requested a review from alextmagro July 22, 2026 02:19

		@@ -0,0 +1,1459 @@
		/*************************************************************************

Uh oh!

Conversation

asdfvg123 commented Jul 6, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

github-actions Bot commented Jul 7, 2026

Claude review — PR 658

Uh oh!

asdfvg123 commented Jul 9, 2026

Uh oh!

ipanfilo commented Jul 20, 2026

Uh oh!

alextmagro commented Jul 20, 2026

Uh oh!

ipanfilo commented Jul 20, 2026

Uh oh!

alextmagro commented Jul 20, 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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants