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

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,45 @@
 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
 
+def rocm_blockwise_unsupported_reason(
+    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 "Only 1D by 1D and 1D by 2D block scaling GEMM is supported"
+
+    if x_columnwise and not w_columnwise:
+        return "does not support TT layout"
+
+    return None
+
 
 def cublas_gemm_fp8_blockwise_case(
     x_dtype,
@@ -45,12 +71,19 @@ 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 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:
+        unsupported_reason = rocm_blockwise_unsupported_reason(
+            is_x_1d_scaled, is_w_1d_scaled,
+            x_columnwise=x_columnwise, w_columnwise=w_columnwise,
+        )
+        if unsupported_reason is not None:
+            pytest.skip(unsupported_reason)
     # Setup device and random seed
     device = "cuda"
     seed = 0
@@ -228,6 +261,11 @@ 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:
+        is_1d1d = is_x_1d_scaled and is_w_1d_scaled
+        is_1d2d = is_x_1d_scaled and not is_w_1d_scaled
+        if not (is_1d1d or is_1d2d):
+            expected_err_msg = "Only 1D by 1D and 1D by 2D block scaling GEMM is supported"
     # Setup device and random seed
     device = "cuda"
     seed = 0
@@ -331,8 +369,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 +425,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 +485,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 +547,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 +620,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 +692,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 +718,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 +733,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 +744,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 +771,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 +788,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 +799,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 +835,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 +866,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 +893,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 +908,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 +918,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 +943,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 +973,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"])
@@ -940,7 +1002,13 @@ def test_unaligned_shapes(
     is_w_1d_scaled,
 ) -> None:
     legal = legalX1d if is_x_1d_scaled else legalX2d
+    if IS_HIP_EXTENSION:
+        legal = (K % 16 == 0) and (N % 16 == 0) # M is unconstrained for rocm
     if not legal:
+        if IS_HIP_EXTENSION:
+            expected_err_msg = "must be multiple of 16"
+        else:
+            expected_err_msg = "dimension requirement"
         cublas_gemm_test_constraint_enforced(
             x_dtype,
             w_dtype,
@@ -952,7 +1020,7 @@ def test_unaligned_shapes(
             use_split_accumulator,
             is_x_1d_scaled,
             is_w_1d_scaled,
-            expected_err_msg="dimension requirement",
+            expected_err_msg=expected_err_msg,
         )
     else:
         cublas_gemm_fp8_blockwise_case(

tests/pytorch/test_float8_blockwise_scaling_exact.py

@@ -38,10 +38,7 @@
 tensor_dump_dir_env = os.getenv("NVTE_TEST_BLOCK_CURRENT_SCALING_EXACT_TENSOR_DUMP_DIR")
 if tensor_dump_dir_env is not None:
     TENSOR_DUMP_DIR = pathlib.Path(tensor_dump_dir_env)
-recipe_available, reason_for_no_recipe = te.is_fp8_block_scaling_quantization_available(
-    return_reason=True
-)
-gemm_available, reason_for_no_gemm = te.is_fp8_block_scaling_available(return_reason=True)
+recipe_available, reason_for_no_recipe = te.is_fp8_block_scaling_available(return_reason=True)
 recipe_emulated = get_device_compute_capability() >= (10, 0)
 
 
@@ -390,7 +387,7 @@ def test_quantization_block_tiling_extrema_versus_reference(
 
 
 # FP8 per tesnor current scaling
-@pytest.mark.skipif(not gemm_available, reason=reason_for_no_gemm)
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 class TestFP8BlockScalingRecipeLinear(TestFP8RecipeLinearBase):
 
     @staticmethod
@@ -450,7 +447,7 @@ def test_fp8_current_scaling_with_linear_module(
         )
 
 
-@pytest.mark.skipif(not gemm_available, reason=reason_for_no_gemm)
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 class TestFP8BlockScalingRecipeLayerNormLinear(TestFP8RecipeLayerNormLinearBase):
 
     @staticmethod

tests/pytorch/test_numerics.py

@@ -2135,6 +2135,8 @@ def test_grouped_linear_accuracy(
     if IS_HIP_EXTENSION:
         if dtype not in (torch.float32,) and fuse_wgrad_accumulation and not fp8:
             pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.")
+        if recipe is not None and recipe.float8_block_scaling():
+            pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.")
     if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED:
         pytest.skip("FP8 parameters are not supported in debug mode.")
     if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute:

transformer_engine/common/CMakeLists.txt

@@ -9,7 +9,7 @@ cmake_minimum_required(VERSION 3.21)
 option(USE_ROCM "Use ROCm" ON)
 option(USE_FUSED_ATTN_AOTRITON "Use aotriton backend" ON)
 option(USE_FUSED_ATTN_CK "Use ck backend" ON)
-option(USE_HIPKITTENS_GEMM "Use HipKittens MXFP8 GEMM kernels" ON)
+option(USE_HIPKITTENS_GEMM "Use HipKittens blockwise FP8 and MXFP8 GEMM kernels" ON)
 set(USE_CUDA OFF)
 
 if (USE_ROCM)

transformer_engine/common/common.h

@@ -112,6 +112,10 @@ inline bool is_mxfp_scaling(const NVTEScalingMode &mode) {
 #endif //#ifdef __HIP_PLATFORM_AMD__
 inline bool is_nvfp_scaling(const NVTEScalingMode &mode) { return mode == NVTE_NVFP4_1D_SCALING; }
 
+inline bool is_blockwise_fp8_scaling(const NVTEScalingMode &mode) {
+  return mode == NVTE_BLOCK_SCALING_1D || mode == NVTE_BLOCK_SCALING_2D;
+}
+
 inline size_t product(const std::vector<size_t> &shape, const size_t begin, const size_t end) {
   NVTE_CHECK(begin <= end && end <= shape.size(), "Attempted to access entries ", begin, " to ",
              end, " in a vector with ", shape.size(), " entries");