blockwise fp8 gemm integration for gfx942 and gfx950

.github/workflows/rocm-wheels-build.yml

@@ -89,6 +89,7 @@ jobs:
             3rdparty/ck_jit \
             3rdparty/hipify_torch \
             3rdparty/hipkittens \
+            3rdparty/hipkittens_cdna3 \
             3rdparty/composable_kernel
 
       - name: Derive Docker image tag

.gitmodules

@@ -33,3 +33,7 @@
 [submodule "3rdparty/ck_jit"]
 	path = 3rdparty/ck_jit
 	url = https://github.com/ROCm/ck-jit.git
+[submodule "3rdparty/hipkittens_cdna3"]
+	path = 3rdparty/hipkittens_cdna3
+	url = https://github.com/asdfvg123/HipKittens.git
+	branch = yeonsoo/cdna3_fp8

README.rst

@@ -281,6 +281,21 @@ To enable MXFP8 support, use NVTE_ROCM_ENABLE_MXFP8 environment variable which c
 * 1 - enable MXFP8 support in fp8;
 * 2 - make MXFP8 a default fp8 recipe.
 
+Blockwise FP8 GEMM support on ROCm (gfx942 and gfx950)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Blockwise FP8 GEMM is supported on gfx942 and gfx950 GPUs through the HipKittens backend and is enabled by default.
+It supports 1Dx2D and 1Dx1D scaling, e4m3/e5m2 inputs in any combination except e5m2 x e5m2, bf16/fp32/fp16 output, and bias/beta/DGELU epilogues.
+The K and N dimensions must be multiples of 16.
+The backend shares the existing USE_HIPKITTENS_GEMM build option with MXFP8:
+
+* OFF - disable the HipKittens blockwise FP8 and MXFP8 GEMM backend;
+* ON - enable the HipKittens GEMM backend (default).
+
+On gfx950 there are two blockwise FP8 GEMM kernels, selected at runtime by the NVTE_BLOCKWISE_FP8_POWER_OF_2_SCALE environment variable:
+
+* 0 - compute the GEMM with the FP32 scales directly;
+* 1 - cast the incoming FP32 scales to E8M0 (power-of-2) and compute the GEMM with those scales (default).
+
 Two-stage amax Kernel
 ^^^^^^^^^^^^^^^^^^^^^
 

tests/pytorch/test_float8_blockwise_gemm_exact.py

@@ -1,3 +1,5 @@
+# This file was modified for portability to AMDGPU
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
@@ -6,21 +8,46 @@
 import torch
 import transformer_engine.pytorch as te
 import transformer_engine_torch as tex
+from torch.utils.cpp_extension import IS_HIP_EXTENSION
 
 from transformer_engine.pytorch.constants import TE_DType
 from transformer_engine.pytorch import (
     Float8BlockQuantizer,
     get_device_compute_capability,
 )
+from transformer_engine.pytorch.utils import (
+    get_torch_float8_e4m3_type,
+    get_torch_float8_e5m2_type,
+)
 from references.blockwise_quantizer_reference import CuBLASScaleMunger
 from references.blockwise_fp8_gemm_reference import CuBLASRefBlockwiseGemm
 
+fp8_e4m3_type = get_torch_float8_e4m3_type()
+fp8_e5m2_type = get_torch_float8_e5m2_type()
 
 def fp8_blockwise_gemm_supported() -> bool:
     supported = te.is_fp8_block_scaling_available()
     emulated = get_device_compute_capability() >= (10, 0)
     return supported and not emulated
 
+if IS_HIP_EXTENSION:
+    def rocm_blockwise_is_supported(
+        is_x_1d_scaled,
+        is_w_1d_scaled,
+        *,
+        x_columnwise: bool = False,
+        w_columnwise: bool = False,
+    ):
+        is_1d2d = is_x_1d_scaled and not is_w_1d_scaled
+        is_1d1d = is_x_1d_scaled and is_w_1d_scaled
+        if not (is_1d2d or is_1d1d):
+            return False, "Only 1D by 1D and 1D by 2D block scaling GEMM is supported"
+
+        if x_columnwise and not w_columnwise:
+            return False, "does not support TT layout"
+
+        return True, None
+
 
 def cublas_gemm_fp8_blockwise_case(
     x_dtype,
@@ -45,12 +72,22 @@ def cublas_gemm_fp8_blockwise_case(
     atol: float = 0.0,
     rtol: float = 0.0
 ):
-    if x_dtype == torch.float8_e5m2 and w_dtype == torch.float8_e5m2:
+    if IS_HIP_EXTENSION:
+        atol = 1e-5
+        rtol = 1.3e-6
+    if x_dtype == fp8_e5m2_type and w_dtype == fp8_e5m2_type:
         pytest.skip("FP8 GEMM doesn't support both a and b types being torch.float8_e5m2")
     if not (is_x_1d_scaled or is_w_1d_scaled):
         pytest.skip("FP8 GEMM doesn't support 2dimensional qtile by 2dimensional qtile")
     if not fp8_blockwise_gemm_supported():
         pytest.skip("CUDA version does not support blockwise FP8 gemm.")
+    if IS_HIP_EXTENSION:
+        supported, reason = rocm_blockwise_is_supported(
+            is_x_1d_scaled, is_w_1d_scaled,
+            x_columnwise=x_columnwise, w_columnwise=w_columnwise,
+        )
+        if not supported:
+            pytest.skip(reason)
     # Setup device and random seed
     device = "cuda"
     seed = 0
@@ -228,6 +265,13 @@ def cublas_gemm_test_constraint_enforced(
 ):
     if not fp8_blockwise_gemm_supported():
         pytest.skip("CUDA version does not support blockwise FP8 gemm.")
+    if IS_HIP_EXTENSION:
+        supported, reason = rocm_blockwise_is_supported(
+            is_x_1d_scaled, is_w_1d_scaled,
+            x_columnwise=x_columnwise, w_columnwise=w_columnwise,
+        )
+        if not supported:
+            expected_err_msg = reason
     # Setup device and random seed
     device = "cuda"
     seed = 0
@@ -331,8 +375,8 @@ def cublas_gemm_test_constraint_enforced(
         (1024, 4096, 1024),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("noise_type", ["normal"], ids=str)
 @pytest.mark.parametrize("x_magnitude", [1], ids=str)
@@ -387,8 +431,8 @@ def test_cublas_gemm_fp8_blockwise_shape_varying(
         (320, 256, 336),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("noise_type", ["normal", "uniform"], ids=str)
 @pytest.mark.parametrize("x_magnitude", [1e-28, 1, 1e3], ids=str)
@@ -447,8 +491,8 @@ def test_cublas_gemm_fp8_blockwise_accumulate_magnitude_varying(
         (256, 256, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("noise_type", ["normal"], ids=str)
 @pytest.mark.parametrize("x_magnitude", [1e-3], ids=str)
@@ -509,8 +553,8 @@ def test_cublas_gemm_fp8_blockwise_bias(
         (4096, 128, 4096),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type, fp8_e5m2_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("noise_type", ["normal"], ids=str)
 @pytest.mark.parametrize("x_magnitude", [1], ids=str)
@@ -582,8 +626,8 @@ def test_cublas_gemm_fp8_blockwise_columnwise(
         (256, 256, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str)
 @pytest.mark.parametrize("noise_type", ["normal"], ids=str)
 @pytest.mark.parametrize("x_magnitude", [1], ids=str)
@@ -654,8 +698,8 @@ def test_cublas_gemm_fp8_gelu(
         (256, 128, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("accumulate", [True, False], ids=["accumulate", "no_accumulate"])
 @pytest.mark.parametrize("use_split_accumulator", [False], ids=["split_acc"])
@@ -680,6 +724,10 @@ def test_split_accumulator_enforced(
     is_x_1d_scaled,
     is_w_1d_scaled,
 ) -> None:
+    if IS_HIP_EXTENSION:
+        expected_err_msg = "requires split accumulator"
+    else:
+        expected_err_msg = "CUBLAS_STATUS_NOT_SUPPORTED"
     cublas_gemm_test_constraint_enforced(
         x_dtype,
         w_dtype,
@@ -691,6 +739,7 @@ def test_split_accumulator_enforced(
         use_split_accumulator,
         is_x_1d_scaled,
         is_w_1d_scaled,
+        expected_err_msg=expected_err_msg,
     )
 
 
@@ -701,8 +750,8 @@ def test_split_accumulator_enforced(
         (256, 128, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("accumulate", [True, False], ids=["accumulate", "no_accumulate"])
 @pytest.mark.parametrize("use_split_accumulator", [True], ids=["split_acc"])
@@ -728,6 +777,10 @@ def test_bgrad_not_supported(
     is_w_1d_scaled,
 ) -> None:
     # NOTE: BGRAD epilogue is not supported for fp8.
+    if IS_HIP_EXTENSION:
+        expected_err_msg = "does not support bias with grad"
+    else:
+        expected_err_msg = "Epilogue requested outside of the available"
     cublas_gemm_test_constraint_enforced(
         x_dtype,
         w_dtype,
@@ -741,7 +794,7 @@ def test_bgrad_not_supported(
         is_w_1d_scaled,
         use_grad=True,
         use_bias=True,
-        expected_err_msg="Epilogue requested outside of the available",
+        expected_err_msg=expected_err_msg,
     )
 
 
@@ -752,8 +805,8 @@ def test_bgrad_not_supported(
         (256, 128, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("accumulate", [True, False], ids=["accumulate", "no_accumulate"])
 @pytest.mark.parametrize("use_bias", [True, False], ids=["bias", "no_bias"])
@@ -788,6 +841,13 @@ def test_gelu_unsupported_cases_error(
         expected_err = "an unsupported value or parameter was passed"
     else:
         expected_err = "Epilogue requested outside of the available"
+    if IS_HIP_EXTENSION:
+        if use_grad and not use_bias:
+            expected_err = "DGELU epilogue only supports bfloat16 output"
+        elif not use_grad:
+            expected_err = "only supports DGELU grad epilogue"
+        else:
+            expected_err = "does not support bias with grad"
     cublas_gemm_test_constraint_enforced(
         x_dtype,
         w_dtype,
@@ -812,8 +872,8 @@ def test_gelu_unsupported_cases_error(
         (256, 128, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e5m2], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e5m2], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e5m2_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e5m2_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("accumulate", [True, False], ids=["accumulate", "no_accumulate"])
 @pytest.mark.parametrize("use_split_accumulator", [True], ids=["split_acc"])
@@ -839,6 +899,10 @@ def test_illegal_dtype_enforced(
     is_w_1d_scaled,
 ) -> None:
     # e5m2 by e5m2 not supported.
+    if IS_HIP_EXTENSION:
+        expected_err_msg = "does not support e5m2 by e5m2 inputs"
+    else:
+        expected_err_msg = "CUBLAS_STATUS_NOT_SUPPORTED"
     cublas_gemm_test_constraint_enforced(
         x_dtype,
         w_dtype,
@@ -850,6 +914,7 @@ def test_illegal_dtype_enforced(
         use_split_accumulator,
         is_x_1d_scaled,
         is_w_1d_scaled,
+        expected_err_msg=expected_err_msg,
     )
 
 
@@ -859,8 +924,8 @@ def test_illegal_dtype_enforced(
         (256, 128, 256),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
 @pytest.mark.parametrize("accumulate", [True, False], ids=["accumulate", "no_accumulate"])
 @pytest.mark.parametrize("use_split_accumulator", [True], ids=["split_acc"])
@@ -884,7 +949,10 @@ def test_illegal_2D_by_2D_enforced(
     is_w_1d_scaled,
 ) -> None:
     # 2D block quantization by 2D block quantization is not supported.
-    expected_err_msg = "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling GEMM is supported"
+    if IS_HIP_EXTENSION:
+        expected_err_msg = "Only 1D by 1D and 1D by 2D block scaling GEMM is supported"
+    else:
+        expected_err_msg = "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling GEMM is supported"
     cublas_gemm_test_constraint_enforced(
         x_dtype,
         w_dtype,
@@ -911,8 +979,8 @@ def test_illegal_2D_by_2D_enforced(
         (256, 128, 252, False, False),
     ],
 )
-@pytest.mark.parametrize("x_dtype", [torch.float8_e4m3fn], ids=str)
-@pytest.mark.parametrize("w_dtype", [torch.float8_e4m3fn], ids=str)
+@pytest.mark.parametrize("x_dtype", [fp8_e4m3_type], ids=str)
+@pytest.mark.parametrize("w_dtype", [fp8_e4m3_type], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str)
 @pytest.mark.parametrize("accumulate", [False], ids=["no_accumulate"])
 @pytest.mark.parametrize("use_split_accumulator", [True], ids=["split_acc"])
@@ -939,6 +1007,39 @@ def test_unaligned_shapes(
     is_x_1d_scaled,
     is_w_1d_scaled,
 ) -> None:
+    if IS_HIP_EXTENSION:
+        legal = (K % 16 == 0) and (N % 16 == 0)  # M is unconstrained for rocm
+        if not legal:
+            cublas_gemm_test_constraint_enforced(
+                x_dtype,
+                w_dtype,
+                out_dtype,
+                M,
+                K,
+                N,
+                accumulate,
+                use_split_accumulator,
+                is_x_1d_scaled,
+                is_w_1d_scaled,
+                expected_err_msg="must be multiple of 16",
+            )
+        else:
+            cublas_gemm_fp8_blockwise_case(
+                x_dtype,
+                w_dtype,
+                out_dtype,
+                M,
+                K,
+                N,
+                "uniform",  # noise type
+                1.0,  # x_magnitude
+                1.0,  # w_magnitude
+                accumulate,
+                use_split_accumulator,
+                is_x_1d_scaled,
+                is_w_1d_scaled,
+            )
+        return
     legal = legalX1d if is_x_1d_scaled else legalX2d
     if not legal:
         cublas_gemm_test_constraint_enforced(