diff --git a/.github/workflows/rocm-wheels-build.yml b/.github/workflows/rocm-wheels-build.yml index 006123d5e..4ec2e157c 100644 --- a/.github/workflows/rocm-wheels-build.yml +++ b/.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 diff --git a/.gitmodules b/.gitmodules index 87c8f2d0a..3f88ef67c 100644 --- a/.gitmodules +++ b/.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 diff --git a/3rdparty/hipkittens_cdna3 b/3rdparty/hipkittens_cdna3 new file mode 160000 index 000000000..17d808f61 --- /dev/null +++ b/3rdparty/hipkittens_cdna3 @@ -0,0 +1 @@ +Subproject commit 17d808f61b972106d1e113a8591fc70bff4387d5 diff --git a/tests/pytorch/test_float8_blockwise_gemm_exact.py b/tests/pytorch/test_float8_blockwise_gemm_exact.py index eff571b5c..9215955fe 100644 --- a/tests/pytorch/test_float8_blockwise_gemm_exact.py +++ b/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( diff --git a/tests/pytorch/test_float8_blockwise_scaling_exact.py b/tests/pytorch/test_float8_blockwise_scaling_exact.py index 434a3e6a8..078d1cf2b 100644 --- a/tests/pytorch/test_float8_blockwise_scaling_exact.py +++ b/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 diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 5c9686f15..d2afc9493 100644 --- a/tests/pytorch/test_numerics.py +++ b/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: diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 60b9f444d..1254a0749 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/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) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index c13a10dc5..5311a6a5e 100644 --- a/transformer_engine/common/common.h +++ b/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 &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"); diff --git a/transformer_engine/common/gemm/kittens/CMakeLists.txt b/transformer_engine/common/gemm/kittens/CMakeLists.txt index fb6238ae6..56b421de9 100644 --- a/transformer_engine/common/gemm/kittens/CMakeLists.txt +++ b/transformer_engine/common/gemm/kittens/CMakeLists.txt @@ -3,50 +3,78 @@ cmake_minimum_required(VERSION 3.21) -list(FIND CMAKE_HIP_ARCHITECTURES "gfx950" _gfx950_index) -if(_gfx950_index EQUAL -1) - message(STATUS "HipKittens GEMM disabled (gfx950 not in CMAKE_HIP_ARCHITECTURES)") - set(USE_HIPKITTENS_GEMM OFF CACHE BOOL "Use HipKittens MXFP8 GEMM kernels" FORCE) - return() -endif() - -include(CheckCXXCompilerFlag) -check_cxx_compiler_flag("-std=c++20" HAS_CXX20) -if(NOT HAS_CXX20) - message(WARNING "HipKittens GEMMs require C++20") - set(USE_HIPKITTENS_GEMM OFF CACHE BOOL "Use HipKittens MXFP8 GEMM kernels" FORCE) - return() -endif() - -set(HIPKITTENS_INCLUDE_DIR - "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens/include") -if(NOT EXISTS "${HIPKITTENS_INCLUDE_DIR}/kittens.cuh") - message(FATAL_ERROR - "Could not find HipKittens headers at ${HIPKITTENS_INCLUDE_DIR}. " - "Try running 'git submodule update --init --recursive' " - "within the Transformer Engine source.") -endif() - -set(CMAKE_CXX_STANDARD 20) -project(kittens_gemm LANGUAGES HIP CXX) - -set(kittens_gemm_SOURCES - mxfp8_gemm.cpp) - -add_library(kittens_gemm SHARED ${kittens_gemm_SOURCES}) -set_source_files_properties(${kittens_gemm_SOURCES} PROPERTIES LANGUAGE HIP) -set_target_properties(kittens_gemm PROPERTIES HIP_ARCHITECTURES "gfx950") - -set(KITTENS_GEMM_COMPILE_OPTIONS - -DKITTENS_CDNA4 - -fno-gpu-rdc - -O3) - -find_package(hip) -target_include_directories(kittens_gemm PRIVATE ${HIPKITTENS_INCLUDE_DIR}) -target_include_directories(kittens_gemm PRIVATE ${HIP_INCLUDE_DIRS}) -target_link_libraries(kittens_gemm PUBLIC hip::host hip::device) -target_compile_options(kittens_gemm PRIVATE ${KITTENS_GEMM_COMPILE_OPTIONS}) - -install(TARGETS kittens_gemm - DESTINATION ${CMAKE_INSTALL_PREFIX}/transformer_engine/lib) +function(try_enable_hipkittens_gemm) + list(FIND CMAKE_HIP_ARCHITECTURES "gfx942" _gfx942_index) + list(FIND CMAKE_HIP_ARCHITECTURES "gfx950" _gfx950_index) + if(_gfx942_index EQUAL -1 AND _gfx950_index EQUAL -1) + message(STATUS "HipKittens GEMM disabled (neither gfx942 nor gfx950 in CMAKE_HIP_ARCHITECTURES)") + set(USE_HIPKITTENS_GEMM OFF PARENT_SCOPE) + return() + endif() + + include(CheckCXXCompilerFlag) + check_cxx_compiler_flag("-std=c++20" HAS_CXX20) + if(NOT HAS_CXX20) + message(WARNING "HipKittens GEMMs require C++20") + set(USE_HIPKITTENS_GEMM OFF PARENT_SCOPE) + return() + endif() + + set(HIPKITTENS_CDNA3_INCLUDE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens_cdna3/include") + set(HIPKITTENS_CDNA4_INCLUDE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens/include") + + set(CMAKE_CXX_STANDARD 20) + project(kittens_gemm LANGUAGES HIP CXX) + find_package(hip) + + set(_kittens_arch_objs "") + + if(NOT _gfx942_index EQUAL -1) + if(NOT EXISTS "${HIPKITTENS_CDNA3_INCLUDE_DIR}/kittens.cuh") + message(FATAL_ERROR + "Could not find HipKittens (CDNA3) headers at ${HIPKITTENS_CDNA3_INCLUDE_DIR}. " + "Try running 'git submodule update --init --recursive'.") + endif() + add_library(kittens_gemm_cdna3 OBJECT cdna3/blockwise_fp8_gemm.cpp) + set_source_files_properties(cdna3/blockwise_fp8_gemm.cpp PROPERTIES LANGUAGE HIP) + set_target_properties(kittens_gemm_cdna3 PROPERTIES + HIP_ARCHITECTURES "gfx942" POSITION_INDEPENDENT_CODE ON) + target_include_directories(kittens_gemm_cdna3 PRIVATE + ${HIP_INCLUDE_DIRS} "${HIPKITTENS_CDNA3_INCLUDE_DIR}") + target_compile_options(kittens_gemm_cdna3 PRIVATE + -DKITTENS_CDNA3 -fno-gpu-rdc -O3) + target_link_libraries(kittens_gemm_cdna3 PRIVATE hip::host hip::device) + list(APPEND _kittens_arch_objs $) + endif() + + if(NOT _gfx950_index EQUAL -1) + if(NOT EXISTS "${HIPKITTENS_CDNA4_INCLUDE_DIR}/kittens.cuh") + message(FATAL_ERROR + "Could not find HipKittens (CDNA4) headers at ${HIPKITTENS_CDNA4_INCLUDE_DIR}. " + "Try running 'git submodule update --init --recursive'.") + endif() + add_library(kittens_gemm_cdna4 OBJECT cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp) + set_source_files_properties(cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp + PROPERTIES LANGUAGE HIP) + set_target_properties(kittens_gemm_cdna4 PROPERTIES + HIP_ARCHITECTURES "gfx950" POSITION_INDEPENDENT_CODE ON) + target_include_directories(kittens_gemm_cdna4 PRIVATE + ${HIP_INCLUDE_DIRS} "${HIPKITTENS_CDNA4_INCLUDE_DIR}") + target_compile_options(kittens_gemm_cdna4 PRIVATE + -DKITTENS_CDNA4 -fno-gpu-rdc -O3) + target_link_libraries(kittens_gemm_cdna4 PRIVATE hip::host hip::device) + list(APPEND _kittens_arch_objs $) + endif() + + add_library(kittens_gemm SHARED blockwise_fp8_gemm.cpp ${_kittens_arch_objs}) + set_source_files_properties(blockwise_fp8_gemm.cpp PROPERTIES LANGUAGE HIP) + target_include_directories(kittens_gemm PRIVATE ${HIP_INCLUDE_DIRS}) + target_link_libraries(kittens_gemm PUBLIC hip::host hip::device) + + install(TARGETS kittens_gemm + DESTINATION ${CMAKE_INSTALL_PREFIX}/transformer_engine/lib) +endfunction() + +try_enable_hipkittens_gemm() diff --git a/transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.cpp b/transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.cpp new file mode 100644 index 000000000..aa8fc04b7 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.cpp @@ -0,0 +1,44 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#include +#include +#include +#include "blockwise_fp8_gemm.h" +#include "cdna3/blockwise_fp8_gemm.h" +#include "cdna4/blockwise_fp8_gemm.h" +#include "../../util/hip_runtime.h" + +void kittens_blockwise_fp8_gemm( + const void *A, const void *B, void *C, + const void *scale_A, const void *scale_B, + int M, int N, int K, + bool transa, bool transb, + int a_dtype, int b_dtype, + int a_scaling_mode, int b_scaling_mode, + int out_dtype, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const void *c_in, float beta, + hipStream_t stream) { + const int arch = transformer_engine::cuda::sm_arch(); + if (arch != 94 && arch != 95) { + throw std::runtime_error( + "kittens_blockwise_fp8_gemm: not implemented for this GPU arch (sm_arch=" + + std::to_string(arch) + "). Only gfx942 and gfx950 are supported."); + } + + if (arch == 95) { + blockwise_gfx950::kittens_blockwise_fp8_gemm_impl_cdna4( + A, B, C, scale_A, scale_B, M, N, K, + a_dtype, b_dtype, a_scaling_mode, b_scaling_mode, out_dtype, + bias, bias_dtype, gelu_aux, gelu_aux_dtype, c_in, beta, stream); + } else { + blockwise_gfx942::kittens_blockwise_fp8_gemm_impl_cdna3( + A, B, C, scale_A, scale_B, M, N, K, transa, transb, + a_dtype, b_dtype, a_scaling_mode, b_scaling_mode, out_dtype, + bias, bias_dtype, gelu_aux, gelu_aux_dtype, c_in, beta, stream); + } +} diff --git a/transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.h b/transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.h new file mode 100644 index 000000000..2569b9fd6 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/blockwise_fp8_gemm.h @@ -0,0 +1,41 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include +#include + +#ifndef KITTENS_DTYPE_ENUM_DEFINED +#define KITTENS_DTYPE_ENUM_DEFINED +enum KittensDType { + KITTENS_FLOAT32 = 4, + KITTENS_FLOAT16 = 5, + KITTENS_BFLOAT16 = 6, + KITTENS_FP8E4M3 = 7, + KITTENS_FP8E5M2 = 8, +}; +#endif // KITTENS_DTYPE_ENUM_DEFINED + +#ifndef KITTENS_SCALING_MODE_DEFINED +#define KITTENS_SCALING_MODE_DEFINED +enum KittensScalingMode { + KITTENS_BLOCK_SCALING_1D = 2, + KITTENS_BLOCK_SCALING_2D = 3, +}; +#endif // KITTENS_SCALING_MODE_DEFINED + +void kittens_blockwise_fp8_gemm( + const void *A, const void *B, void *C, + const void *scale_A, const void *scale_B, + int M, int N, int K, + bool transa, bool transb, + int a_dtype, int b_dtype, + int a_scaling_mode, int b_scaling_mode, + int out_dtype, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const void *c_in, float beta, + hipStream_t stream); diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.cpp b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.cpp new file mode 100644 index 000000000..2c01fc5bd --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.cpp @@ -0,0 +1,693 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#include +#include "kittens.cuh" +#include "blockwise_fp8_gemm.h" + +namespace blockwise_gfx942 { + +#include "blockwise_fp8_gemm_device.cuh" + +constexpr int NUM_WARPS = 8; +constexpr int WARPS_ROW = 2; +constexpr int WARPS_COL = 4; +constexpr int BLOCK_M = 128; +constexpr int BLOCK_N = 256; +constexpr int BLOCK_K = 128; +constexpr int REG_M = BLOCK_M / 4; +constexpr int REG_N = BLOCK_N / 4; +constexpr int MFMA_K = 32; +constexpr int SCALE_BLOCK = 128; // blockwise scale granularity +constexpr int NUM_THREADS = NUM_WARPS * kittens::WARP_THREADS; + +template using _gl_A_t = kittens::gl; +template using _gl_B_t = kittens::gl; +template using _gl_C_t = kittens::gl; +using _gl_SA = kittens::gl; +using _gl_SB = kittens::gl; + +using G = kittens::group; + +template +struct micro_globals { + _gl_A_t a; + _gl_B_t b; + _gl_C_t c; + _gl_SA scale_a; + _gl_SB scale_b; + hipStream_t stream; + const void *bias; + int bias_dtype; + const void *gelu_aux; + int gelu_aux_dtype; + const OType *c_in; + float beta; + int M() const { return (int)c.rows(); } + int N() const { return (int)c.cols(); } + int K() const { return (int)a.cols(); } + dim3 grid() { return dim3(((N() + BLOCK_N - 1) / BLOCK_N) * ((M() + BLOCK_M - 1) / BLOCK_M)); } + dim3 block() { return dim3(NUM_THREADS); } + size_t dynamic_shared_memory() { return 49152; } +}; + +template +__global__ __launch_bounds__(NUM_THREADS, 2) +void micro_tk_1d2d(const micro_globals g) { + extern __shared__ kittens::alignment_dummy __shm[]; + kittens::shared_allocator al((int*)&__shm[0]); + kittens::st (&As) = al.allocate>(); + kittens::st (&Bs) = al.allocate>(); + __shared__ float smem_sa[2][BLOCK_M]; + + kittens::rt at[5]; + kittens::rt bt[3]; + kittens::rt_fl C_accum[2]; + kittens::rt_fl partial[2]; + for (int i = 0; i < 2; i++) { kittens::zero(C_accum[i]); } + + const int M = (int)g.c.rows(); + const int N = (int)g.c.cols(); + const int K = (int)g.a.cols(); + + int wgid = (blockIdx.y * gridDim.x) + blockIdx.x; + const int NUM_WGS = gridDim.x * gridDim.y; + constexpr int WGM = 4; + wgid = kittens::chiplet_transform_chunked(wgid, NUM_WGS, kittens::NUM_XCDS, WGM*WGM); + + const int num_pid_m = kittens::ceil_div(M, BLOCK_M); + const int num_pid_n = kittens::ceil_div(N, BLOCK_N); + int num_wgid_in_group = WGM * num_pid_n; + int group_id = wgid / num_wgid_in_group; + int first_pid_m = group_id * WGM; + int group_size_m = min(num_pid_m - first_pid_m, WGM); + int pid_m = first_pid_m + ((wgid % num_wgid_in_group) % group_size_m); + int pid_n = (wgid % num_wgid_in_group) / group_size_m; + const int row = pid_m; + const int col = pid_n; + + const bool is_last_m = IS_PARTIAL_M && (row * BLOCK_M + BLOCK_M > M); + const bool is_last_n = IS_PARTIAL_N && (col * BLOCK_N + BLOCK_N > N); + + const int warp_id = kittens::warpid(); + const int warp_row = warp_id / WARPS_COL; + const int warp_col = warp_id % WARPS_COL; + + const int num_k_steps = kittens::ceil_div(K, BLOCK_K); + const bool is_k_partial = IS_PARTIAL_K && (K % BLOCK_K != 0); + + const float *sa_block = g.scale_a.raw_ptr + row * BLOCK_M; + + const int n_scale_blocks = kittens::ceil_div(N, SCALE_BLOCK); + const int sb_block0 = col * (BLOCK_N / SCALE_BLOCK) + warp_col / 2; + const bool sb_valid = (!is_last_n) || (sb_block0 < n_scale_blocks); + const float *sb_base = g.scale_b.raw_ptr + (sb_valid ? sb_block0 : 0) * num_k_steps; + int32x4_lds_t sb_srsrc = make_buf_res((const void*)sb_base, (uint32_t)num_k_steps * 4); + const int local_m0 = warp_row * REG_M; + const int local_m1 = (warp_row + WARPS_ROW) * REG_M; + const int tid = threadIdx.x; + + const uint32_t sa_range = (uint32_t)((M - row * BLOCK_M) * 4); + + const bool is_first_k_partial = is_k_partial && (num_k_steps == 1); + if (is_first_k_partial || is_last_m) load_tile_masked(As, g.a, row, 0, M, K); + else G::load(As, g.a, {0, 0, row, 0}); + if (is_first_k_partial || is_last_n) load_tile_masked(Bs, g.b, col, 0, N, K); + else G::load(Bs, g.b, {0, 0, col, 0}); + + // Prologue + float sb_cur = llvm_amdgcn_s_buffer_load_f32(sb_srsrc, 0, 0); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + if (warp_row == 1) { + __builtin_amdgcn_s_barrier(); + } + + #pragma unroll + for (int k_step = 0; k_step < num_k_steps - 1; ++k_step) { + + constexpr int A_ELEMS_PER_THREAD = (BLOCK_M * BLOCK_K) / NUM_THREADS; + constexpr int B_ELEMS_PER_THREAD = (BLOCK_N * BLOCK_K) / NUM_THREADS; + float4 a_buffer_next[A_ELEMS_PER_THREAD * sizeof(AType) / sizeof(float4)]; + float4 b_buffer_next[B_ELEMS_PER_THREAD * sizeof(BType) / sizeof(float4)]; + + kittens::zero(partial[0]); kittens::zero(partial[1]); + + const bool is_next_k_partial = is_k_partial && (k_step + 1 == num_k_steps - 1); + float sa_reg0[REG_M / 16 * 4]; + float sa_reg1[REG_M / 16 * 4]; + + // Cluster 0 + if (!is_last_n && !is_next_k_partial) + kittens::load_global_to_register_buffer<2, false, NUM_THREADS>(b_buffer_next, B_ELEMS_PER_THREAD, g.b, {0, 0, col, k_step + 1}, Bs); + float sb_next; + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 0})); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 0})); + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 0})); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 1 + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 2 + kittens::load(bt[1], kittens::subtile_inplace(Bs, {warp_col, 1})); + kittens::load(at[2], kittens::subtile_inplace(As, {warp_row, 1})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 1})); + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 2})); + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 2})); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 3 + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[2], bt[1], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[1], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 4 + if (!is_last_m && !is_next_k_partial) + kittens::load_global_to_register_buffer<2, false, NUM_THREADS>(a_buffer_next, A_ELEMS_PER_THREAD, g.a, {0, 0, row, k_step + 1}, As); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 2})); + kittens::load(bt[2], kittens::subtile_inplace(Bs, {warp_col, 3})); + kittens::load(at[4], kittens::subtile_inplace(As, {warp_row, 3})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 3})); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 5 + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 6 + asm volatile("s_waitcnt lgkmcnt(0)"); + if (is_next_k_partial || is_last_m) load_tile_masked(As, g.a, row, k_step + 1, M, K); + else kittens::store_register_buffer_to_shared(As, a_buffer_next); + if (is_next_k_partial || is_last_n) load_tile_masked(Bs, g.b, col, k_step + 1, N, K); + load_scale_global_reg(sa_reg0, sa_block + k_step * M, local_m0, sa_range); + load_scale_global_reg(sa_reg1, sa_block + k_step * M, local_m1, sa_range); + sb_next = llvm_amdgcn_s_buffer_load_f32(sb_srsrc, (k_step + 1) * 4, 0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 7 + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[4], bt[2], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[2], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 8 + if (!(is_next_k_partial || is_last_n)) + kittens::store_register_buffer_to_shared(Bs, b_buffer_next); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 9 + apply_block_scale_1d2d(C_accum[0], partial[0], sa_reg0, sb_cur); + apply_block_scale_1d2d(C_accum[1], partial[1], sa_reg1, sb_cur); + sb_cur = sb_next; + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + } + + // Epilogue + kittens::zero(partial[0]); kittens::zero(partial[1]); + __builtin_amdgcn_sched_barrier(0); + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 0})); + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 0})); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 0})); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + kittens::load(bt[1], kittens::subtile_inplace(Bs, {warp_col, 1})); + kittens::load(at[2], kittens::subtile_inplace(As, {warp_row, 1})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 1})); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[2], bt[1], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[1], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 2})); + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 2})); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 2})); + kittens::load(bt[1], kittens::subtile_inplace(Bs, {warp_col, 3})); + kittens::load(at[2], kittens::subtile_inplace(As, {warp_row, 3})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 3})); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[2], bt[1], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[1], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + { + const int last = num_k_steps - 1; + float sa_reg0[REG_M / 16 * 4]; + float sa_reg1[REG_M / 16 * 4]; + load_scale_global_reg(sa_reg0, sa_block + last * M, local_m0, sa_range); + load_scale_global_reg(sa_reg1, sa_block + last * M, local_m1, sa_range); + apply_block_scale_1d2d(C_accum[0], partial[0], sa_reg0, sb_cur); + apply_block_scale_1d2d(C_accum[1], partial[1], sa_reg1, sb_cur); + } + + if (warp_row == 0) { + __builtin_amdgcn_s_barrier(); + } + + if constexpr (HAS_BIAS || HAS_GELU || HAS_BETA) { + apply_epilogue( + C_accum[0], row * 4 + warp_row, col * 4 + warp_col, M, N, + g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue( + C_accum[1], row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col, M, N, + g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + } + + if constexpr (std::is_same_v) { + apply_rtne_bias(C_accum[0]); + apply_rtne_bias(C_accum[1]); + } + if (is_last_m || is_last_n) { + store_masked(g.c.raw_ptr, C_accum[0], row * 4 + warp_row, col * 4 + warp_col, M, N); + store_masked(g.c.raw_ptr, C_accum[1], row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col, M, N); + } else { + kittens::store(g.c, C_accum[0], {0, 0, row * 4 + warp_row, col * 4 + warp_col}); + kittens::store(g.c, C_accum[1], {0, 0, row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col}); + } +} + +template +__global__ __launch_bounds__(NUM_THREADS, 2) +void micro_tk_1d1d(const micro_globals g) { + extern __shared__ kittens::alignment_dummy __shm[]; + kittens::shared_allocator al((int*)&__shm[0]); + kittens::st (&As) = al.allocate>(); + kittens::st (&Bs) = al.allocate>(); + + kittens::rt at[5]; + kittens::rt bt[3]; + kittens::rt_fl C_accum[2]; + kittens::rt_fl partial[2]; + for (int i = 0; i < 2; i++) { kittens::zero(C_accum[i]); } + + const int M = (int)g.c.rows(); + const int N = (int)g.c.cols(); + const int K = (int)g.a.cols(); + + int wgid = (blockIdx.y * gridDim.x) + blockIdx.x; + const int NUM_WGS = gridDim.x * gridDim.y; + constexpr int WGM = 4; + wgid = kittens::chiplet_transform_chunked(wgid, NUM_WGS, kittens::NUM_XCDS, WGM*WGM); + + const int num_pid_m = kittens::ceil_div(M, BLOCK_M); + const int num_pid_n = kittens::ceil_div(N, BLOCK_N); + int num_wgid_in_group = WGM * num_pid_n; + int group_id = wgid / num_wgid_in_group; + int first_pid_m = group_id * WGM; + int group_size_m = min(num_pid_m - first_pid_m, WGM); + int pid_m = first_pid_m + ((wgid % num_wgid_in_group) % group_size_m); + int pid_n = (wgid % num_wgid_in_group) / group_size_m; + + const int row = pid_m; + const int col = pid_n; + + const bool is_last_m = IS_PARTIAL_M && (row * BLOCK_M + BLOCK_M > M); + const bool is_last_n = IS_PARTIAL_N && (col * BLOCK_N + BLOCK_N > N); + + const int warp_id = kittens::warpid(); + const int warp_row = warp_id / WARPS_COL; + const int warp_col = warp_id % WARPS_COL; + + const int num_k_steps = kittens::ceil_div(K, BLOCK_K); + const bool is_k_partial = IS_PARTIAL_K && (K % BLOCK_K != 0); + + const float *sa_block = g.scale_a.raw_ptr + row * BLOCK_M; + + const int sb_col0 = col * BLOCK_N; + const float *sb_block = g.scale_b.raw_ptr + sb_col0; + const int local_m0 = warp_row * REG_M; + const int local_m1 = (warp_row + WARPS_ROW) * REG_M; + const int local_n = warp_col * REG_N; + const int tid = threadIdx.x; + const uint32_t sa_range = (uint32_t)((M - row * BLOCK_M) * 4); + const uint32_t sb_range = (uint32_t)((N - sb_col0) * 4); + + // Prologue + const bool is_first_k_partial = is_k_partial && (num_k_steps == 1); + if (is_first_k_partial || is_last_m) load_tile_masked(As, g.a, row, 0, M, K); + else G::load(As, g.a, {0, 0, row, 0}); + if (is_first_k_partial || is_last_n) load_tile_masked(Bs, g.b, col, 0, N, K); + else G::load(Bs, g.b, {0, 0, col, 0}); + __builtin_amdgcn_s_barrier(); + + if (warp_row == 1) { + __builtin_amdgcn_s_barrier(); + } + + #pragma unroll + for (int k_step = 0; k_step < num_k_steps - 1; ++k_step) { + + constexpr int A_ELEMS_PER_THREAD = (BLOCK_M * BLOCK_K) / NUM_THREADS; + constexpr int B_ELEMS_PER_THREAD = (BLOCK_N * BLOCK_K) / NUM_THREADS; + float4 a_buffer_next[A_ELEMS_PER_THREAD * sizeof(AType) / sizeof(float4)]; + float4 b_buffer_next[B_ELEMS_PER_THREAD * sizeof(BType) / sizeof(float4)]; + + kittens::zero(partial[0]); kittens::zero(partial[1]); + + const bool is_next_k_partial = is_k_partial && (k_step + 1 == num_k_steps - 1); + + // Cluster 0 + if (!is_last_n && !is_next_k_partial) + kittens::load_global_to_register_buffer<2, false, NUM_THREADS>(b_buffer_next, B_ELEMS_PER_THREAD, g.b, {0, 0, col, k_step + 1}, Bs); + float sa_reg0[REG_M / 16 * 4]; + float sa_reg1[REG_M / 16 * 4]; + float sb_reg[REG_N / 16]; + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 0})); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 0})); + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 0})); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 1 + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 2 + kittens::load(bt[1], kittens::subtile_inplace(Bs, {warp_col, 1})); + kittens::load(at[2], kittens::subtile_inplace(As, {warp_row, 1})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 1})); + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 2})); + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 2})); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 3 + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[2], bt[1], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[1], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 4 + if (!is_last_m && !is_next_k_partial) + kittens::load_global_to_register_buffer<2, false, NUM_THREADS>(a_buffer_next, A_ELEMS_PER_THREAD, g.a, {0, 0, row, k_step + 1}, As); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 2})); + kittens::load(bt[2], kittens::subtile_inplace(Bs, {warp_col, 3})); + kittens::load(at[4], kittens::subtile_inplace(As, {warp_row, 3})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 3})); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 5 + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 6 + asm volatile("s_waitcnt lgkmcnt(0)"); + if (is_next_k_partial || is_last_m) load_tile_masked(As, g.a, row, k_step + 1, M, K); + else kittens::store_register_buffer_to_shared(As, a_buffer_next); + if (is_next_k_partial || is_last_n) load_tile_masked(Bs, g.b, col, k_step + 1, N, K); + load_scale_global_reg(sa_reg0, sa_block + k_step * M, local_m0, sa_range); + load_scale_global_reg(sa_reg1, sa_block + k_step * M, local_m1, sa_range); + load_scaleB_global_reg(sb_reg, sb_block + k_step * N, local_n, sb_range); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 7 + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[4], bt[2], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[2], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 8 + if (!(is_next_k_partial || is_last_n)) + kittens::store_register_buffer_to_shared(Bs, b_buffer_next); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 9 + apply_block_scale_1d1d(C_accum[0], partial[0], sa_reg0, sb_reg); + apply_block_scale_1d1d(C_accum[1], partial[1], sa_reg1, sb_reg); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + } + + // Epilogue + kittens::zero(partial[0]); kittens::zero(partial[1]); + __builtin_amdgcn_sched_barrier(0); + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 0})); + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 0})); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 0})); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + kittens::load(bt[1], kittens::subtile_inplace(Bs, {warp_col, 1})); + kittens::load(at[2], kittens::subtile_inplace(As, {warp_row, 1})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 1})); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[2], bt[1], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[1], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + kittens::load(bt[0], kittens::subtile_inplace(Bs, {warp_col, 2})); + kittens::load(at[0], kittens::subtile_inplace(As, {warp_row, 2})); + kittens::load(at[1], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 2})); + kittens::load(bt[1], kittens::subtile_inplace(Bs, {warp_col, 3})); + kittens::load(at[2], kittens::subtile_inplace(As, {warp_row, 3})); + kittens::load(at[3], kittens::subtile_inplace(As, {warp_row + WARPS_ROW, 3})); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[0], bt[0], partial[0]); + kittens::mma_ABt(partial[1], at[1], bt[0], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + __builtin_amdgcn_s_setprio(1); + kittens::mma_ABt(partial[0], at[2], bt[1], partial[0]); + kittens::mma_ABt(partial[1], at[3], bt[1], partial[1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + { + const int last = num_k_steps - 1; + float sa_reg0[REG_M / 16 * 4]; + float sa_reg1[REG_M / 16 * 4]; + float sb_reg[REG_N / 16]; + load_scale_global_reg(sa_reg0, sa_block + last * M, local_m0, sa_range); + load_scale_global_reg(sa_reg1, sa_block + last * M, local_m1, sa_range); + load_scaleB_global_reg(sb_reg, sb_block + last * N, local_n, sb_range); + apply_block_scale_1d1d(C_accum[0], partial[0], sa_reg0, sb_reg); + apply_block_scale_1d1d(C_accum[1], partial[1], sa_reg1, sb_reg); + } + + if (warp_row == 0) { + __builtin_amdgcn_s_barrier(); + } + + if constexpr (HAS_BIAS || HAS_GELU || HAS_BETA) { + apply_epilogue( + C_accum[0], row * 4 + warp_row, col * 4 + warp_col, M, N, + g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue( + C_accum[1], row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col, M, N, + g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + } + + if constexpr (std::is_same_v) { + apply_rtne_bias(C_accum[0]); + apply_rtne_bias(C_accum[1]); + } + if (is_last_m || is_last_n) { + store_masked(g.c.raw_ptr, C_accum[0], row * 4 + warp_row, col * 4 + warp_col, M, N); + store_masked(g.c.raw_ptr, C_accum[1], row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col, M, N); + } else { + kittens::store(g.c, C_accum[0], {0, 0, row * 4 + warp_row, col * 4 + warp_col}); + kittens::store(g.c, C_accum[1], {0, 0, row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col}); + } +} + +template +static void dispatch_micro_epilogue(micro_globals g) { + unsigned long mem_size = g.dynamic_shared_memory(); + const bool pm = (g.M() % BLOCK_M != 0); + const bool pn = (g.N() % BLOCK_N != 0); + auto launch = [&](auto kern) { + hipFuncSetAttribute((void*)kern, hipFuncAttributeMaxDynamicSharedMemorySize, mem_size); + kern<<>>(g); + }; + auto kernel = [&]() { + if constexpr (IS_1D2D) + return micro_tk_1d2d; + else + return micro_tk_1d1d; + }; + if (!pm && !pn) launch(kernel.template operator()()); + else if ( pm && !pn) launch(kernel.template operator()()); + else if (!pm && pn) launch(kernel.template operator()()); + else launch(kernel.template operator()()); +} + +template +static void dispatch_micro_k(micro_globals g, + bool has_bias, bool has_gelu, bool has_beta) { + if (has_gelu) { + if (has_beta) dispatch_micro_epilogue(g); + else dispatch_micro_epilogue(g); + } else if (has_bias) { + if (has_beta) dispatch_micro_epilogue(g); + else dispatch_micro_epilogue(g); + } else { + if (has_beta) dispatch_micro_epilogue(g); + else dispatch_micro_epilogue(g); + } +} + +template +static void dispatch_micro(micro_globals g, + bool has_bias, bool has_gelu, bool has_beta, bool has_partial_k) { + if (has_partial_k) dispatch_micro_k(g, has_bias, has_gelu, has_beta); + else dispatch_micro_k(g, has_bias, has_gelu, has_beta); +} + +void kittens_blockwise_fp8_gemm_impl_cdna3( + const void *A, const void *B, void *C, + const void *scale_A, const void *scale_B, + int M, int N, int K, + bool transa, bool transb, + int a_dtype, int b_dtype, + int a_scaling_mode, int b_scaling_mode, + int out_dtype, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const void *c_in, float beta, + hipStream_t stream) { + + // Dispatch passes canonical (A=weight/2D, B=activation/1D, M/N=user) + // The kernel uses swapped layout + const void *kA = B, *kB = A; + const void *ksa = scale_B, *ksb = scale_A; + void *kC = C; + const int kM = N, kN = M; + const int ka_mode = b_scaling_mode, kb_mode = a_scaling_mode; + const int ka_dtype = b_dtype, kb_dtype = a_dtype; + (void)transa; (void)transb; + + const bool is_1d2d = (kb_mode == KITTENS_BLOCK_SCALING_2D); + const bool has_bias = (bias != nullptr); + const bool has_gelu = (gelu_aux != nullptr); + const bool has_beta = (c_in != nullptr); + const bool has_partial_k = (K % BLOCK_K != 0); + const int k_blocks = (K + BLOCK_K - 1) / BLOCK_K; + + auto run = [&]() { + micro_globals g = { + _gl_A_t(reinterpret_cast(const_cast(kA)), 1, 1, kM, K), + _gl_B_t(reinterpret_cast(const_cast(kB)), 1, 1, kN, K), + _gl_C_t(reinterpret_cast(kC), 1, 1, kM, kN), + _gl_SA(reinterpret_cast(const_cast(ksa)), 1, 1, k_blocks, kM), + is_1d2d + ? _gl_SB(reinterpret_cast(const_cast(ksb)), 1, 1, kittens::ceil_div(kN, SCALE_BLOCK), k_blocks) + : _gl_SB(reinterpret_cast(const_cast(ksb)), 1, 1, k_blocks, kN), + stream, + bias, bias_dtype, gelu_aux, gelu_aux_dtype, + reinterpret_cast(c_in), beta, + }; + if (is_1d2d) dispatch_micro(g, has_bias, has_gelu, has_beta, has_partial_k); + else dispatch_micro(g, has_bias, has_gelu, has_beta, has_partial_k); + }; + + const bool a_e5m2 = (ka_dtype == KITTENS_FP8E5M2); + const bool b_e5m2 = (kb_dtype == KITTENS_FP8E5M2); + auto run_ab = [&]() { + if (!a_e5m2 && !b_e5m2) run.template operator()(); + else if ( a_e5m2 && !b_e5m2) run.template operator()(); + else run.template operator()(); + }; + if (out_dtype == KITTENS_FLOAT32) run_ab.template operator()(); + else if (out_dtype == KITTENS_FLOAT16) run_ab.template operator()(); + else run_ab.template operator()(); +} + +} // namespace blockwise_gfx942 diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.h b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.h new file mode 100644 index 000000000..9c8c3b79f --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.h @@ -0,0 +1,43 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include +#include + +#ifndef KITTENS_DTYPE_ENUM_DEFINED +#define KITTENS_DTYPE_ENUM_DEFINED +enum KittensDType { + KITTENS_FLOAT32 = 4, + KITTENS_FLOAT16 = 5, + KITTENS_BFLOAT16 = 6, + KITTENS_FP8E4M3 = 7, + KITTENS_FP8E5M2 = 8, +}; +#endif // KITTENS_DTYPE_ENUM_DEFINED + +#ifndef KITTENS_SCALING_MODE_DEFINED +#define KITTENS_SCALING_MODE_DEFINED +enum KittensScalingMode { + KITTENS_BLOCK_SCALING_1D = 2, + KITTENS_BLOCK_SCALING_2D = 3, +}; +#endif // KITTENS_SCALING_MODE_DEFINED + +namespace blockwise_gfx942 { +void kittens_blockwise_fp8_gemm_impl_cdna3( + const void *A, const void *B, void *C, + const void *scale_A, const void *scale_B, + int M, int N, int K, + bool transa, bool transb, + int a_dtype, int b_dtype, + int a_scaling_mode, int b_scaling_mode, + int out_dtype, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const void *c_in, float beta, + hipStream_t stream); +} // namespace blockwise_gfx942 diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_device.cuh b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_device.cuh new file mode 100644 index 000000000..35cf70de4 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_device.cuh @@ -0,0 +1,229 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include +#include "kittens.cuh" +#include "../../../util/math.h" + +typedef int int32x4_lds_t __attribute__((ext_vector_type(4))); +struct __attribute__((packed)) buf_res { const void *ptr; uint32_t range; uint32_t config; }; +__device__ inline int32x4_lds_t make_buf_res(const void *ptr, uint32_t size) { + buf_res r{ptr, size, 0x00020000u}; + return __builtin_bit_cast(int32x4_lds_t, r); +} +extern "C" __device__ __uint128_t +llvm_amdgcn_raw_buffer_load_b128(int32x4_lds_t rsrc, int voffset, int soffset, + int aux) __asm("llvm.amdgcn.raw.buffer.load.v4f32"); +extern "C" __device__ float +llvm_amdgcn_s_buffer_load_f32(int32x4_lds_t rsrc, int offset, + int cachepolicy) __asm("llvm.amdgcn.s.buffer.load.f32"); +extern "C" __device__ float +llvm_amdgcn_raw_buffer_load_f32(int32x4_lds_t rsrc, int voffset, int soffset, + int aux) __asm("llvm.amdgcn.raw.buffer.load.f32"); + + +template +__device__ inline void load_scale_global_reg(float (&sa_reg)[HEIGHT * 4], const float *sa_base, + int local_m_base, uint32_t range_bytes) { + const int lane = kittens::laneid(); + const int row_g = 4 * (lane / 16); + int32x4_lds_t srsrc = make_buf_res((const void*)sa_base, range_bytes); + #pragma unroll + for (int i = 0; i < HEIGHT; i++) { + const int m0 = local_m_base + i * 16 + row_g; + __uint128_t raw = llvm_amdgcn_raw_buffer_load_b128(srsrc, m0 * 4, 0, 0); + *reinterpret_cast(&sa_reg[i * 4]) = *reinterpret_cast(&raw); + } +} + +template +__device__ inline void load_scaleB_global_reg(float (&sb_reg)[WIDTH], const float *sb_base, + int local_n_base, uint32_t range_bytes) { + const int lane = kittens::laneid(); + const int col_l = lane % 16; + int32x4_lds_t srsrc = make_buf_res((const void*)sb_base, range_bytes); + #pragma unroll + for (int j = 0; j < WIDTH; j++) { + const int n0 = local_n_base + j * 16 + col_l; + sb_reg[j] = llvm_amdgcn_raw_buffer_load_f32(srsrc, n0 * 4, 0, 0); + } +} + +__device__ inline float rtne_bias(float v) { + uint32_t bits = __builtin_bit_cast(uint32_t, v); + if ((bits & 0x7f800000u) == 0x7f800000u) return v; + bits += 0x7fffu + ((bits >> 16) & 1u); + return __builtin_bit_cast(float, bits); +} + +__device__ inline float read_elem(const void *p, int dtype, int idx) { + if (dtype == 6) return __bfloat162float(reinterpret_cast(p)[idx]); + if (dtype == 5) return __half2float(reinterpret_cast(p)[idx]); + return reinterpret_cast(p)[idx]; +} + +template +__device__ inline float rtne_cast_roundtrip(float v) { + if constexpr (std::is_same_v) { + return v; + } else { + return static_cast(kittens::base_types::convertor::convert(rtne_bias(v))); + } +} + +template +__device__ inline void apply_rtne_bias(AccType &Cacc) { + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + Cacc.tiles[i][j].data[0].x = rtne_bias(Cacc.tiles[i][j].data[0].x); + Cacc.tiles[i][j].data[0].y = rtne_bias(Cacc.tiles[i][j].data[0].y); + Cacc.tiles[i][j].data[1].x = rtne_bias(Cacc.tiles[i][j].data[1].x); + Cacc.tiles[i][j].data[1].y = rtne_bias(Cacc.tiles[i][j].data[1].y); + } + } +} + +template +__device__ inline void store_masked(OType *c_ptr, const AccType &Cacc, + int Rtile, int Ctile, int M, int N) { + const int lane = kittens::laneid(); + const int m_base = Rtile * AccType::rows + 4 * (lane / 16); + const int n_base = Ctile * AccType::cols + (lane % 16); + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const int m0 = m_base + i * 16; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const int col = n_base + j * 16; + if (col >= N) continue; + const float v0 = Cacc.tiles[i][j].data[0].x; + const float v1 = Cacc.tiles[i][j].data[0].y; + const float v2 = Cacc.tiles[i][j].data[1].x; + const float v3 = Cacc.tiles[i][j].data[1].y; + if (m0 + 0 < M) c_ptr[(m0 + 0) * N + col] = kittens::base_types::convertor::convert(v0); + if (m0 + 1 < M) c_ptr[(m0 + 1) * N + col] = kittens::base_types::convertor::convert(v1); + if (m0 + 2 < M) c_ptr[(m0 + 2) * N + col] = kittens::base_types::convertor::convert(v2); + if (m0 + 3 < M) c_ptr[(m0 + 3) * N + col] = kittens::base_types::convertor::convert(v3); + } + } +} + +template +__device__ inline void apply_epilogue( + AccType &Cacc, int Rtile, int Ctile, int M, int N, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const OType *c_in, float beta) { + const int lane = kittens::laneid(); + const int m_base = Rtile * AccType::rows + 4 * (lane / 16); + const int n_base = Ctile * AccType::cols + (lane % 16); + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const int m0 = m_base + i * 16; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const int col = n_base + j * 16; + if (col >= N) continue; + float v[4] = { + Cacc.tiles[i][j].data[0].x, Cacc.tiles[i][j].data[0].y, + Cacc.tiles[i][j].data[1].x, Cacc.tiles[i][j].data[1].y, + }; + float bias_v = 0.f; + if constexpr (HAS_BIAS) bias_v = read_elem(bias, bias_dtype, col); + #pragma unroll + for (int r = 0; r < 4; r++) { + const int m_g = m0 + r; + if (m_g >= M) continue; + float x = v[r]; + if constexpr (HAS_BIAS) x += bias_v; + if constexpr (HAS_BETA) { + x = rtne_cast_roundtrip(x); + x += beta * static_cast(c_in[m_g * N + col]); + } + if constexpr (HAS_GELU) { + x *= transformer_engine::dgelu( + read_elem(gelu_aux, gelu_aux_dtype, m_g * N + col), {}); + } + v[r] = x; + } + Cacc.tiles[i][j].data[0].x = v[0]; + Cacc.tiles[i][j].data[0].y = v[1]; + Cacc.tiles[i][j].data[1].x = v[2]; + Cacc.tiles[i][j].data[1].y = v[3]; + } + } +} + +template +__device__ inline void load_tile_masked(ST &dst, const GL &src, int row_blk, + int k_blk, int row_dim, int K) { + using T = typename ST::dtype; + constexpr int elem_per_memcpy = sizeof(float4) / sizeof(T); + constexpr int elem_per_half_memcpy = sizeof(float2) / sizeof(T); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total = (ST::rows * ST::cols) / elem_per_memcpy; + const int row_stride = src.template stride<2>(); + const int row_base = row_blk * ST::rows; + const int k_base = k_blk * ST::cols; + kittens::coord<> uc = kittens::coord(0, 0, row_blk, k_blk).template unit_coord<2, 3>(); + T *src_ptr = (T *)&src[uc]; + uint32_t dst_ptr = reinterpret_cast(&dst.data[0]); + const int tid = threadIdx.x; + #pragma unroll + for (int idx = tid; idx < total; idx += NUM_THREADS) { + const int row = idx / memcpy_per_row; + const int col = (idx % memcpy_per_row) * elem_per_memcpy; + float4 v = {0.f, 0.f, 0.f, 0.f}; + if (row_base + row < row_dim && k_base + col < K) + v = kittens::load_global_vec4((float4 *)(src_ptr + (row * row_stride + col))); + kittens::store_shared_vec(dst.idx(dst_ptr, {row, col}), {v.x, v.y}); + kittens::store_shared_vec(dst.idx(dst_ptr, {row, col + elem_per_half_memcpy}), {v.z, v.w}); + } + asm volatile("s_waitcnt lgkmcnt(0)"); +} + +template +__device__ inline void apply_block_scale_1d2d( + AccType &Cacc, const AccType &partial, const float (&sa_reg)[AccType::height * 4], float sb) { + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const float s0 = sa_reg[i * 4 + 0] * sb; + const float s1 = sa_reg[i * 4 + 1] * sb; + const float s2 = sa_reg[i * 4 + 2] * sb; + const float s3 = sa_reg[i * 4 + 3] * sb; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + Cacc.tiles[i][j].data[0].x += partial.tiles[i][j].data[0].x * s0; + Cacc.tiles[i][j].data[0].y += partial.tiles[i][j].data[0].y * s1; + Cacc.tiles[i][j].data[1].x += partial.tiles[i][j].data[1].x * s2; + Cacc.tiles[i][j].data[1].y += partial.tiles[i][j].data[1].y * s3; + } + } +} + +template +__device__ inline void apply_block_scale_1d1d( + AccType &Cacc, const AccType &partial, const float (&sa_reg)[AccType::height * 4], + const float (&sb_reg)[AccType::width]) { + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const float a0 = sa_reg[i * 4 + 0]; + const float a1 = sa_reg[i * 4 + 1]; + const float a2 = sa_reg[i * 4 + 2]; + const float a3 = sa_reg[i * 4 + 3]; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const float sb = sb_reg[j]; + Cacc.tiles[i][j].data[0].x += partial.tiles[i][j].data[0].x * (a0 * sb); + Cacc.tiles[i][j].data[0].y += partial.tiles[i][j].data[0].y * (a1 * sb); + Cacc.tiles[i][j].data[1].x += partial.tiles[i][j].data[1].x * (a2 * sb); + Cacc.tiles[i][j].data[1].y += partial.tiles[i][j].data[1].y * (a3 * sb); + } + } +} diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.cpp b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.cpp new file mode 100644 index 000000000..46e7c8140 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.cpp @@ -0,0 +1,948 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#include +#include "kittens.cuh" +#include "blockwise_fp8_gemm.h" + + +namespace blockwise_gfx950 { + +#include "blockwise_fp8_gemm_device.cuh" + +constexpr int NUM_WARPS = 8; +constexpr int WARPS_ROW = 2; +constexpr int WARPS_COL = 4; +constexpr int BLOCK_M = 128; +constexpr int BLOCK_N = 256; +constexpr int BLOCK_K = 128; +constexpr int HALF_ROW = BLOCK_M / 2; +constexpr int HALF_COL = BLOCK_N / 2; +constexpr int REG_M = BLOCK_M / WARPS_ROW / 2; +constexpr int REG_N = BLOCK_N / WARPS_COL / 2; +constexpr int MFMA_K = 128; +constexpr int SCALE_BLOCK = 128; +constexpr int NUM_THREADS = NUM_WARPS * kittens::WARP_THREADS; + +template using _gl_A_t = kittens::gl; +template using _gl_B_t = kittens::gl; +template using _gl_C_t = kittens::gl; +using _gl_SA = kittens::gl; +using _gl_SB = kittens::gl; + +using G = kittens::group; + +template +struct micro_globals { + _gl_A_t a; + _gl_B_t b; + _gl_C_t c; + _gl_SA scale_a; + _gl_SB scale_b; + const void *bias; + int bias_dtype; + const void *gelu_aux; + int gelu_aux_dtype; + const OType *c_in; + float beta; + hipStream_t stream; + int M() const { return (int)c.rows(); } + int N() const { return (int)c.cols(); } + int K() const { return (int)a.cols(); } + dim3 grid() { return dim3(((M() + BLOCK_M - 1) / BLOCK_M) * ((N() + BLOCK_N - 1) / BLOCK_N)); } + dim3 block() { return dim3(NUM_THREADS); } +}; + +template +__global__ __launch_bounds__(NUM_THREADS, 2) +void micro_tk_1d2d(micro_globals g) { + const auto A = g.a; + const auto B = g.b; + const auto C = g.c; + const float *scale_A = g.scale_a.raw_ptr; + const float *scale_B = g.scale_b.raw_ptr; + const int M = (int)g.c.rows(), N = (int)g.c.cols(), K = (int)g.a.cols(); + const int k_iters = K / BLOCK_K; + const int scale_K = K / SCALE_BLOCK; + const int blocks_per_col = (N + BLOCK_N - 1) / BLOCK_N; + + using ST_A = kittens::st_fp8e4m3; + using ST_B = kittens::st_fp8e4m3; + using RT_A = kittens::rt_fp8e4m3; + using RT_B = kittens::rt_fp8e4m3; + using RT_C = kittens::rt_fl; + + __shared__ float smem_sa[2][BLOCK_M]; + __shared__ ST_A As[2][2]; + __shared__ ST_B Bs[2][2]; + + RT_A a; + RT_B b0, b1; + RT_C cA, cB, cC, cD; + RT_C p; + + const int global_block_id = blockIdx.x; + const int block_row = global_block_id / blocks_per_col; + const int block_col = global_block_id % blocks_per_col; + + const int warp_m = kittens::warpid() / WARPS_COL; + const int warp_n = kittens::warpid() % WARPS_COL; + const int tid = threadIdx.x; + + int tic = 0, toc = 1; + + using T = kittens::fp8e4m3; + constexpr int bpt = ST_A::underlying_subtile_bytes_per_thread; + constexpr int bpm = bpt * NUM_THREADS; + constexpr int memcpy_A = HALF_ROW * BLOCK_K * sizeof(T) / bpm; + constexpr int memcpy_B = HALF_COL * BLOCK_K * sizeof(T) / bpm; + uint32_t sw_A[memcpy_A], sw_B[memcpy_B]; + G::prefill_swizzled_offsets(As[tic][0], A, sw_A); + G::prefill_swizzled_offsets(Bs[tic][0], B, sw_B); + + const T *a_base = (const T *)&A[{0, 0, 0, 0}]; + const T *b_base = (const T *)&B[{0, 0, 0, 0}]; + const int a_row_stride = A.template stride<2>() * sizeof(T); + const int b_row_stride = B.template stride<2>() * sizeof(T); + kittens::i32x4 a_srd = kittens::make_srsrc(a_base, M * a_row_stride, a_row_stride); + kittens::i32x4 b_srd = kittens::make_srsrc(b_base, N * b_row_stride, b_row_stride); + + const int wid = kittens::warpid() % NUM_WARPS; + constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS; + uint32_t a_lds[2][2], b_lds[2][2]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 2; j++) { + a_lds[i][j] = __builtin_amdgcn_readfirstlane(static_cast( + reinterpret_cast(&As[i][j].data[0]) + wid * elem_per_warp * sizeof(T))); + b_lds[i][j] = __builtin_amdgcn_readfirstlane(static_cast( + reinterpret_cast(&Bs[i][j].data[0]) + wid * elem_per_warp * sizeof(T))); + } + + const kittens::fp8e8m0_4 unit = 0x7F7F7F7Fu; + + const int n_scale_blocks = (N + SCALE_BLOCK - 1) / SCALE_BLOCK; + const int nb0 = min(block_col * 2 + 0, n_scale_blocks - 1); + const int nb1 = min(block_col * 2 + 1, n_scale_blocks - 1); + const float *sa_row = scale_A + block_row * BLOCK_M; + const float *sb0 = scale_B + nb0 * scale_K; + const float *sb1 = scale_B + nb1 * scale_K; + const int local_m0 = warp_m * REG_M; + const int local_m1 = HALF_ROW + warp_m * REG_M; + const int m_valid = M - block_row * BLOCK_M; + + const int local_n0 = block_col * BLOCK_N + warp_n * REG_N; + const int local_n1 = block_col * BLOCK_N + HALF_COL + warp_n * REG_N; + + kittens::zero(cA); kittens::zero(cB); kittens::zero(cC); kittens::zero(cD); + + const int sa_tid_p = tid % BLOCK_M; + + G::load(Bs[tic][0], B, {0, 0, block_col * 2, 0}, sw_B, b_srd, b_base, b_lds[tic][0]); + G::load(As[tic][0], A, {0, 0, block_row * 2, 0}, sw_A, a_srd, a_base, a_lds[tic][0]); + G::load(Bs[tic][1], B, {0, 0, block_col * 2 + 1, 0}, sw_B, b_srd, b_base, b_lds[tic][1]); + G::load(As[tic][1], A, {0, 0, block_row * 2 + 1, 0}, sw_A, a_srd, a_base, a_lds[tic][1]); + const float sa0_reg = sa_tid_p < m_valid ? sa_row[0 * M + sa_tid_p] : 0.f; + + if (warp_m == 1) { + __builtin_amdgcn_s_barrier(); + } + + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + G::load(Bs[toc][0], B, {0, 0, block_col * 2, 1}, sw_B, b_srd, b_base, b_lds[toc][0]); + G::load(As[toc][0], A, {0, 0, block_row * 2, 1}, sw_A, a_srd, a_base, a_lds[toc][0]); + G::load(Bs[toc][1], B, {0, 0, block_col * 2 + 1, 1}, sw_B, b_srd, b_base, b_lds[toc][1]); + const float sa1_reg = (k_iters > 1 && sa_tid_p < m_valid) ? sa_row[1 * M + sa_tid_p] : 0.f; + + asm volatile("s_waitcnt vmcnt(6)"); + if (tid < BLOCK_M) { + smem_sa[tic][tid] = sa0_reg; + if (k_iters > 1) smem_sa[toc][tid] = sa1_reg; + } + __builtin_amdgcn_s_barrier(); + + #pragma unroll 2 + for (int k = 0; k < k_iters - 2; k++, tic ^= 1, toc ^= 1) { + const float sb0_k = __builtin_bit_cast(float, __builtin_amdgcn_readfirstlane(__builtin_bit_cast(int, sb0[k]))); + const float sb1_k = __builtin_bit_cast(float, __builtin_amdgcn_readfirstlane(__builtin_bit_cast(int, sb1[k]))); + + const int sa_tid = tid % BLOCK_M; + const float sa_next = sa_tid < m_valid ? sa_row[(k + 1) * M + sa_tid] : 0.f; + + auto bs0 = kittens::subtile_inplace(Bs[tic][0], {warp_n, 0}); + kittens::load(b0, bs0); + auto as0 = kittens::subtile_inplace(As[tic][0], {warp_m, 0}); + kittens::load(a, as0); + G::load(As[toc][1], A, {0, 0, block_row * 2 + 1, k + 1}, sw_A, a_srd, a_base, a_lds[toc][1]); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + + const auto rs0 = load_row_scale_lds(smem_sa[tic], local_m0); + const auto rs1 = load_row_scale_lds(smem_sa[tic], local_m1); + + if (tid < BLOCK_M) smem_sa[toc][tid] = sa_next; + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate(cA, p, rs0, sb0_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + auto bs1 = kittens::subtile_inplace(Bs[tic][1], {warp_n, 0}); + kittens::load(b1, bs1); + G::load(Bs[tic][0], B, {0, 0, block_col * 2, k + 2}, sw_B, b_srd, b_base, b_lds[tic][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate(cB, p, rs0, sb1_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto as1 = kittens::subtile_inplace(As[tic][1], {warp_m, 0}); + kittens::load(a, as1); + G::load(As[tic][0], A, {0, 0, block_row * 2, k + 2}, sw_A, a_srd, a_base, a_lds[tic][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate(cC, p, rs1, sb0_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + G::load(Bs[tic][1], B, {0, 0, block_col * 2 + 1, k + 2}, sw_B, b_srd, b_base, b_lds[tic][1]); + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate(cD, p, rs1, sb1_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + { + const int k = k_iters - 2; + const float sb0_k = __builtin_bit_cast(float, __builtin_amdgcn_readfirstlane(__builtin_bit_cast(int, sb0[k]))); + const float sb1_k = __builtin_bit_cast(float, __builtin_amdgcn_readfirstlane(__builtin_bit_cast(int, sb1[k]))); + const float *sa_k = sa_row + k * M; + const auto rs0 = load_row_scale(sa_k, local_m0, m_valid); + const auto rs1 = load_row_scale(sa_k, local_m1, m_valid); + + auto bs0 = kittens::subtile_inplace(Bs[tic][0], {warp_n, 0}); + kittens::load(b0, bs0); + auto as0 = kittens::subtile_inplace(As[tic][0], {warp_m, 0}); + kittens::load(a, as0); + G::load(As[toc][1], A, {0, 0, block_row * 2 + 1, k + 1}, sw_A, a_srd, a_base, a_lds[toc][1]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate(cA, p, rs0, sb0_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + auto bs1 = kittens::subtile_inplace(Bs[tic][1], {warp_n, 0}); + kittens::load(b1, bs1); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate(cB, p, rs0, sb1_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto as1 = kittens::subtile_inplace(As[tic][1], {warp_m, 0}); + kittens::load(a, as1); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate(cC, p, rs1, sb0_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + bs0 = kittens::subtile_inplace(Bs[toc][0], {warp_n, 0}); + kittens::load(b0, bs0); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate(cD, p, rs1, sb1_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + tic ^= 1; toc ^= 1; + } + + { + const int k = k_iters - 1; + const float sb0_k = __builtin_bit_cast(float, __builtin_amdgcn_readfirstlane(__builtin_bit_cast(int, sb0[k]))); + const float sb1_k = __builtin_bit_cast(float, __builtin_amdgcn_readfirstlane(__builtin_bit_cast(int, sb1[k]))); + const float *sa_k = sa_row + k * M; + const auto rs0 = load_row_scale(sa_k, local_m0, m_valid); + const auto rs1 = load_row_scale(sa_k, local_m1, m_valid); + + auto as0 = kittens::subtile_inplace(As[tic][0], {warp_m, 0}); + kittens::load(a, as0); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate(cA, p, rs0, sb0_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto bs1 = kittens::subtile_inplace(Bs[tic][1], {warp_n, 0}); + kittens::load(b1, bs1); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate(cB, p, rs0, sb1_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto as1 = kittens::subtile_inplace(As[tic][1], {warp_m, 0}); + kittens::load(a, as1); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate(cC, p, rs1, sb0_k); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate(cD, p, rs1, sb1_k); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + if (warp_m == 0) { + __builtin_amdgcn_s_barrier(); + } + + const int m_off0 = block_row * BLOCK_M + warp_m * REG_M; + const int m_off1 = block_row * BLOCK_M + HALF_ROW + warp_m * REG_M; + const int n_off0 = block_col * BLOCK_N + warp_n * REG_N; + const int n_off1 = block_col * BLOCK_N + HALF_COL + warp_n * REG_N; + + if constexpr (HAS_BIAS || HAS_GELU || HAS_BETA) { + apply_epilogue(cA, m_off0, n_off0, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cB, m_off0, n_off1, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cC, m_off1, n_off0, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cD, m_off1, n_off1, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + } + + if constexpr (std::is_same_v) { + apply_rtne_bias(cA); apply_rtne_bias(cB); apply_rtne_bias(cC); apply_rtne_bias(cD); + } + + OType *c_ptr = C.raw_ptr; + const int ca = block_row * WARPS_ROW * 2 + warp_m; + const int cc = block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m; + const int cn0 = block_col * WARPS_COL * 2 + warp_n; + const int cn1 = block_col * WARPS_COL * 2 + WARPS_COL + warp_n; + + const bool full = (block_row + 1) * BLOCK_M <= M && (block_col + 1) * BLOCK_N <= N; + if (full) { + kittens::store(C, cA, {0, 0, ca, cn0}); + kittens::store(C, cB, {0, 0, ca, cn1}); + kittens::store(C, cC, {0, 0, cc, cn0}); + kittens::store(C, cD, {0, 0, cc, cn1}); + } else { + store_masked(c_ptr, cA, m_off0, n_off0, M, N); + store_masked(c_ptr, cB, m_off0, n_off1, M, N); + store_masked(c_ptr, cC, m_off1, n_off0, M, N); + store_masked(c_ptr, cD, m_off1, n_off1, M, N); + } +} + +template +__global__ __launch_bounds__(NUM_THREADS, 2) +void micro_tk_1d1d(micro_globals g) { + const auto A = g.a; + const auto B = g.b; + const auto C = g.c; + const float *scale_A = g.scale_a.raw_ptr; + const float *scale_B = g.scale_b.raw_ptr; + const int M = (int)g.c.rows(), N = (int)g.c.cols(), K = (int)g.a.cols(); + const int k_iters = K / BLOCK_K; + const int scale_K = K / SCALE_BLOCK; + const int blocks_per_col = (N + BLOCK_N - 1) / BLOCK_N; + + using ST_A = kittens::st_fp8e4m3; + using ST_B = kittens::st_fp8e4m3; + using RT_A = kittens::rt_fp8e4m3; + using RT_B = kittens::rt_fp8e4m3; + using RT_C = kittens::rt_fl; + + __shared__ float smem_sa[2][BLOCK_M]; + __shared__ ST_A As[2][2]; + __shared__ ST_B Bs[2][2]; + + RT_A a; + RT_B b0, b1; + RT_C cA, cB, cC, cD; + RT_C p; + + const int global_block_id = blockIdx.x; + const int block_row = global_block_id / blocks_per_col; + const int block_col = global_block_id % blocks_per_col; + + const int warp_m = kittens::warpid() / WARPS_COL; + const int warp_n = kittens::warpid() % WARPS_COL; + const int tid = threadIdx.x; + + int tic = 0, toc = 1; + + using T = kittens::fp8e4m3; + constexpr int bpt = ST_A::underlying_subtile_bytes_per_thread; + constexpr int bpm = bpt * NUM_THREADS; + constexpr int memcpy_A = HALF_ROW * BLOCK_K * sizeof(T) / bpm; + constexpr int memcpy_B = HALF_COL * BLOCK_K * sizeof(T) / bpm; + uint32_t sw_A[memcpy_A], sw_B[memcpy_B]; + G::prefill_swizzled_offsets(As[tic][0], A, sw_A); + G::prefill_swizzled_offsets(Bs[tic][0], B, sw_B); + + const T *a_base = (const T *)&A[{0, 0, 0, 0}]; + const T *b_base = (const T *)&B[{0, 0, 0, 0}]; + const int a_row_stride = A.template stride<2>() * sizeof(T); + const int b_row_stride = B.template stride<2>() * sizeof(T); + kittens::i32x4 a_srd = kittens::make_srsrc(a_base, M * a_row_stride, a_row_stride); + kittens::i32x4 b_srd = kittens::make_srsrc(b_base, N * b_row_stride, b_row_stride); + + const int wid = kittens::warpid() % NUM_WARPS; + constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS; + uint32_t a_lds[2][2], b_lds[2][2]; + #pragma unroll + for (int i = 0; i < 2; i++) + #pragma unroll + for (int j = 0; j < 2; j++) { + a_lds[i][j] = __builtin_amdgcn_readfirstlane(static_cast( + reinterpret_cast(&As[i][j].data[0]) + wid * elem_per_warp * sizeof(T))); + b_lds[i][j] = __builtin_amdgcn_readfirstlane(static_cast( + reinterpret_cast(&Bs[i][j].data[0]) + wid * elem_per_warp * sizeof(T))); + } + + const kittens::fp8e8m0_4 unit = 0x7F7F7F7Fu; + + const float *sa_row = scale_A + block_row * BLOCK_M; + const int local_m0 = warp_m * REG_M; + const int local_m1 = HALF_ROW + warp_m * REG_M; + const int m_valid = M - block_row * BLOCK_M; + + const int local_n0 = block_col * BLOCK_N + warp_n * REG_N; + const int local_n1 = block_col * BLOCK_N + HALF_COL + warp_n * REG_N; + + kittens::zero(cA); kittens::zero(cB); kittens::zero(cC); kittens::zero(cD); + + const int sa_tid_p = tid % BLOCK_M; + + G::load(Bs[tic][0], B, {0, 0, block_col * 2, 0}, sw_B, b_srd, b_base, b_lds[tic][0]); + G::load(As[tic][0], A, {0, 0, block_row * 2, 0}, sw_A, a_srd, a_base, a_lds[tic][0]); + G::load(Bs[tic][1], B, {0, 0, block_col * 2 + 1, 0}, sw_B, b_srd, b_base, b_lds[tic][1]); + G::load(As[tic][1], A, {0, 0, block_row * 2 + 1, 0}, sw_A, a_srd, a_base, a_lds[tic][1]); + const float sa0_reg = sa_tid_p < m_valid ? sa_row[0 * M + sa_tid_p] : 0.f; + + if (warp_m == 1) { + __builtin_amdgcn_s_barrier(); + } + + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + G::load(Bs[toc][0], B, {0, 0, block_col * 2, 1}, sw_B, b_srd, b_base, b_lds[toc][0]); + G::load(As[toc][0], A, {0, 0, block_row * 2, 1}, sw_A, a_srd, a_base, a_lds[toc][0]); + G::load(Bs[toc][1], B, {0, 0, block_col * 2 + 1, 1}, sw_B, b_srd, b_base, b_lds[toc][1]); + const float sa1_reg = (k_iters > 1 && sa_tid_p < m_valid) ? sa_row[1 * M + sa_tid_p] : 0.f; + + asm volatile("s_waitcnt vmcnt(6)"); + if (tid < BLOCK_M) { + smem_sa[tic][tid] = sa0_reg; + if (k_iters > 1) smem_sa[toc][tid] = sa1_reg; + } + __builtin_amdgcn_s_barrier(); + + #pragma unroll 2 + for (int k = 0; k < k_iters - 2; k++, tic ^= 1, toc ^= 1) { + const float *sb_col = scale_B + k * N; + const ColScale cs0 = load_col_scale(sb_col, local_n0, N); + const ColScale cs1 = load_col_scale(sb_col, local_n1, N); + + const int sa_tid = tid % BLOCK_M; + const float sa_next = sa_tid < m_valid ? sa_row[(k + 1) * M + sa_tid] : 0.f; + + auto bs0 = kittens::subtile_inplace(Bs[tic][0], {warp_n, 0}); + kittens::load(b0, bs0); + auto as0 = kittens::subtile_inplace(As[tic][0], {warp_m, 0}); + kittens::load(a, as0); + G::load(As[toc][1], A, {0, 0, block_row * 2 + 1, k + 1}, sw_A, a_srd, a_base, a_lds[toc][1]); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + + const auto rs0 = load_row_scale_lds(smem_sa[tic], local_m0); + const auto rs1 = load_row_scale_lds(smem_sa[tic], local_m1); + + if (tid < BLOCK_M) smem_sa[toc][tid] = sa_next; + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate_1d1d(cA, p, rs0, cs0); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + auto bs1 = kittens::subtile_inplace(Bs[tic][1], {warp_n, 0}); + kittens::load(b1, bs1); + G::load(Bs[tic][0], B, {0, 0, block_col * 2, k + 2}, sw_B, b_srd, b_base, b_lds[tic][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate_1d1d(cB, p, rs0, cs1); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto as1 = kittens::subtile_inplace(As[tic][1], {warp_m, 0}); + kittens::load(a, as1); + G::load(As[tic][0], A, {0, 0, block_row * 2, k + 2}, sw_A, a_srd, a_base, a_lds[tic][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate_1d1d(cC, p, rs1, cs0); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + G::load(Bs[tic][1], B, {0, 0, block_col * 2 + 1, k + 2}, sw_B, b_srd, b_base, b_lds[tic][1]); + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate_1d1d(cD, p, rs1, cs1); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + { + const int k = k_iters - 2; + const float *sb_col = scale_B + k * N; + const ColScale cs0 = load_col_scale(sb_col, local_n0, N); + const ColScale cs1 = load_col_scale(sb_col, local_n1, N); + const float *sa_k = sa_row + k * M; + const auto rs0 = load_row_scale(sa_k, local_m0, m_valid); + const auto rs1 = load_row_scale(sa_k, local_m1, m_valid); + + auto bs0 = kittens::subtile_inplace(Bs[tic][0], {warp_n, 0}); + kittens::load(b0, bs0); + auto as0 = kittens::subtile_inplace(As[tic][0], {warp_m, 0}); + kittens::load(a, as0); + G::load(As[toc][1], A, {0, 0, block_row * 2 + 1, k + 1}, sw_A, a_srd, a_base, a_lds[toc][1]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate_1d1d(cA, p, rs0, cs0); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + auto bs1 = kittens::subtile_inplace(Bs[tic][1], {warp_n, 0}); + kittens::load(b1, bs1); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate_1d1d(cB, p, rs0, cs1); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto as1 = kittens::subtile_inplace(As[tic][1], {warp_m, 0}); + kittens::load(a, as1); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate_1d1d(cC, p, rs1, cs0); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + bs0 = kittens::subtile_inplace(Bs[toc][0], {warp_n, 0}); + kittens::load(b0, bs0); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate_1d1d(cD, p, rs1, cs1); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + tic ^= 1; toc ^= 1; + } + + { + const int k = k_iters - 1; + const float *sb_col = scale_B + k * N; + const ColScale cs0 = load_col_scale(sb_col, local_n0, N); + const ColScale cs1 = load_col_scale(sb_col, local_n1, N); + const float *sa_k = sa_row + k * M; + const auto rs0 = load_row_scale(sa_k, local_m0, m_valid); + const auto rs1 = load_row_scale(sa_k, local_m1, m_valid); + + auto as0 = kittens::subtile_inplace(As[tic][0], {warp_m, 0}); + kittens::load(a, as0); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate_1d1d(cA, p, rs0, cs0); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto bs1 = kittens::subtile_inplace(Bs[tic][1], {warp_n, 0}); + kittens::load(b1, bs1); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate_1d1d(cB, p, rs0, cs1); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + auto as1 = kittens::subtile_inplace(As[tic][1], {warp_m, 0}); + kittens::load(a, as1); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + scale_accumulate_1d1d(cC, p, rs1, cs0); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + scale_accumulate_1d1d(cD, p, rs1, cs1); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + if (warp_m == 0) { + __builtin_amdgcn_s_barrier(); + } + + const int m_off0 = block_row * BLOCK_M + warp_m * REG_M; + const int m_off1 = block_row * BLOCK_M + HALF_ROW + warp_m * REG_M; + const int n_off0 = block_col * BLOCK_N + warp_n * REG_N; + const int n_off1 = block_col * BLOCK_N + HALF_COL + warp_n * REG_N; + + if constexpr (HAS_BIAS || HAS_GELU || HAS_BETA) { + apply_epilogue(cA, m_off0, n_off0, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cB, m_off0, n_off1, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cC, m_off1, n_off0, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cD, m_off1, n_off1, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + } + + if constexpr (std::is_same_v) { + apply_rtne_bias(cA); apply_rtne_bias(cB); apply_rtne_bias(cC); apply_rtne_bias(cD); + } + + OType *c_ptr = C.raw_ptr; + const int ca = block_row * WARPS_ROW * 2 + warp_m; + const int cc = block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m; + const int cn0 = block_col * WARPS_COL * 2 + warp_n; + const int cn1 = block_col * WARPS_COL * 2 + WARPS_COL + warp_n; + + const bool full = (block_row + 1) * BLOCK_M <= M && (block_col + 1) * BLOCK_N <= N; + if (full) { + kittens::store(C, cA, {0, 0, ca, cn0}); + kittens::store(C, cB, {0, 0, ca, cn1}); + kittens::store(C, cC, {0, 0, cc, cn0}); + kittens::store(C, cD, {0, 0, cc, cn1}); + } else { + store_masked(c_ptr, cA, m_off0, n_off0, M, N); + store_masked(c_ptr, cB, m_off0, n_off1, M, N); + store_masked(c_ptr, cC, m_off1, n_off0, M, N); + store_masked(c_ptr, cD, m_off1, n_off1, M, N); + } +} + +template +__global__ __launch_bounds__(NUM_THREADS, 2) +void micro_tk_smallk(micro_globals g) { + const auto A = g.a; + const auto B = g.b; + const auto C = g.c; + const float *scale_A = g.scale_a.raw_ptr; + const float *scale_B = g.scale_b.raw_ptr; + const int M = (int)g.c.rows(), N = (int)g.c.cols(), K = (int)g.a.cols(); + const int k_blocks = (K + BLOCK_K - 1) / BLOCK_K; + const int scale_K = k_blocks; + const int blocks_per_col = (N + BLOCK_N - 1) / BLOCK_N; + + using ST_A = kittens::st_fp8e4m3; + using ST_B = kittens::st_fp8e4m3; + using RT_A = kittens::rt_fp8e4m3; + using RT_B = kittens::rt_fp8e4m3; + using RT_C = kittens::rt_fl; + + __shared__ ST_A As[2]; + __shared__ ST_B Bs[2]; + + RT_A a; + RT_B b0, b1; + RT_C cA, cB, cC, cD; + RT_C p; + + const int global_block_id = blockIdx.x; + const int block_row = global_block_id / blocks_per_col; + const int block_col = global_block_id % blocks_per_col; + + const int warp_m = kittens::warpid() / WARPS_COL; + const int warp_n = kittens::warpid() % WARPS_COL; + const int tid = threadIdx.x; + + using T = kittens::fp8e4m3; + const kittens::fp8e8m0_4 unit = 0x7F7F7F7Fu; + + const T *a_base = (const T *)&A[{0, 0, 0, 0}]; + const T *b_base = (const T *)&B[{0, 0, 0, 0}]; + const int a_row_stride = A.template stride<2>(); + const int b_row_stride = B.template stride<2>(); + + const int n_scale_blocks = (N + SCALE_BLOCK - 1) / SCALE_BLOCK; + const int nb0 = min(block_col * 2 + 0, n_scale_blocks - 1); + const int nb1 = min(block_col * 2 + 1, n_scale_blocks - 1); + const float *sa_row = scale_A + block_row * BLOCK_M; + const float *sb0 = scale_B + nb0 * scale_K; + const float *sb1 = scale_B + nb1 * scale_K; + const int local_m0 = warp_m * REG_M; + const int local_m1 = HALF_ROW + warp_m * REG_M; + const int m_valid = M - block_row * BLOCK_M; + const int local_n0 = block_col * BLOCK_N + warp_n * REG_N; + const int local_n1 = block_col * BLOCK_N + HALF_COL + warp_n * REG_N; + + kittens::zero(cA); kittens::zero(cB); kittens::zero(cC); kittens::zero(cD); + + for (int k = 0; k < k_blocks; k++) { + __builtin_amdgcn_s_barrier(); + load_tile_masked(As[0], a_base, a_row_stride, block_row * 2 + 0, k, M, K); + load_tile_masked(As[1], a_base, a_row_stride, block_row * 2 + 1, k, M, K); + load_tile_masked(Bs[0], b_base, b_row_stride, block_col * 2 + 0, k, N, K); + load_tile_masked(Bs[1], b_base, b_row_stride, block_col * 2 + 1, k, N, K); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + const auto rs0 = load_row_scale(sa_row + k * M, local_m0, m_valid); + const auto rs1 = load_row_scale(sa_row + k * M, local_m1, m_valid); + float sb0_k, sb1_k; + ColScale cs0, cs1; + if constexpr (IS_1D2D) { + sb0_k = sb0[k]; + sb1_k = sb1[k]; + } else { + const float *sb_col = scale_B + k * N; + cs0 = load_col_scale(sb_col, local_n0, N); + cs1 = load_col_scale(sb_col, local_n1, N); + } + + auto as0 = kittens::subtile_inplace(As[0], {warp_m, 0}); + kittens::load(a, as0); + auto bs0 = kittens::subtile_inplace(Bs[0], {warp_n, 0}); + kittens::load(b0, bs0); + auto bs1 = kittens::subtile_inplace(Bs[1], {warp_n, 0}); + kittens::load(b1, bs1); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + if constexpr (IS_1D2D) scale_accumulate(cA, p, rs0, sb0_k); + else scale_accumulate_1d1d(cA, p, rs0, cs0); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + if constexpr (IS_1D2D) scale_accumulate(cB, p, rs0, sb1_k); + else scale_accumulate_1d1d(cB, p, rs0, cs1); + __builtin_amdgcn_sched_barrier(0); + + auto as1 = kittens::subtile_inplace(As[1], {warp_m, 0}); + kittens::load(a, as1); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + kittens::zero(p); mma_ABt_scaled(p, a, b0, p, &unit, &unit); + if constexpr (IS_1D2D) scale_accumulate(cC, p, rs1, sb0_k); + else scale_accumulate_1d1d(cC, p, rs1, cs0); + kittens::zero(p); mma_ABt_scaled(p, a, b1, p, &unit, &unit); + if constexpr (IS_1D2D) scale_accumulate(cD, p, rs1, sb1_k); + else scale_accumulate_1d1d(cD, p, rs1, cs1); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + } + + const int m_off0 = block_row * BLOCK_M + warp_m * REG_M; + const int m_off1 = block_row * BLOCK_M + HALF_ROW + warp_m * REG_M; + const int n_off0 = block_col * BLOCK_N + warp_n * REG_N; + const int n_off1 = block_col * BLOCK_N + HALF_COL + warp_n * REG_N; + + if constexpr (HAS_BIAS || HAS_GELU || HAS_BETA) { + apply_epilogue(cA, m_off0, n_off0, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cB, m_off0, n_off1, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cC, m_off1, n_off0, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + apply_epilogue(cD, m_off1, n_off1, M, N, g.bias, g.bias_dtype, g.gelu_aux, g.gelu_aux_dtype, g.c_in, g.beta); + } + + if constexpr (std::is_same_v) { + apply_rtne_bias(cA); apply_rtne_bias(cB); apply_rtne_bias(cC); apply_rtne_bias(cD); + } + + OType *c_ptr = C.raw_ptr; + const int ca = block_row * WARPS_ROW * 2 + warp_m; + const int cc = block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m; + const int cn0 = block_col * WARPS_COL * 2 + warp_n; + const int cn1 = block_col * WARPS_COL * 2 + WARPS_COL + warp_n; + + const bool full = (block_row + 1) * BLOCK_M <= M && (block_col + 1) * BLOCK_N <= N; + if (full) { + kittens::store(C, cA, {0, 0, ca, cn0}); + kittens::store(C, cB, {0, 0, ca, cn1}); + kittens::store(C, cC, {0, 0, cc, cn0}); + kittens::store(C, cD, {0, 0, cc, cn1}); + } else { + store_masked(c_ptr, cA, m_off0, n_off0, M, N); + store_masked(c_ptr, cB, m_off0, n_off1, M, N); + store_masked(c_ptr, cC, m_off1, n_off0, M, N); + store_masked(c_ptr, cD, m_off1, n_off1, M, N); + } +} + +template +using micro_globals_fp8 = micro_globals; + +template +static void dispatch_micro_kernel(micro_globals_fp8 g) { + if constexpr (IS_PARTIAL_K) { + micro_tk_smallk<<>>(g); + } else if constexpr (IS_1D2D) { + micro_tk_1d2d<<>>(g); + } else { + micro_tk_1d1d<<>>(g); + } +} + +template +static void dispatch_micro_dtype(int cbsz, int blgp, micro_globals_fp8 g) { + if (cbsz == 0 && blgp == 0) dispatch_micro_kernel(g); + else if (cbsz == 0 && blgp == 1) dispatch_micro_kernel(g); + else if (cbsz == 1 && blgp == 0) dispatch_micro_kernel(g); + else dispatch_micro_kernel(g); +} + +template +static void dispatch_micro_epilogue(int cbsz, int blgp, bool has_bias, bool has_gelu, bool has_beta, + micro_globals_fp8 g) { + if (has_gelu) { + if (has_beta) dispatch_micro_dtype(cbsz, blgp, g); + else dispatch_micro_dtype(cbsz, blgp, g); + } else if (has_bias) { + if (has_beta) dispatch_micro_dtype(cbsz, blgp, g); + else dispatch_micro_dtype(cbsz, blgp, g); + } else { + if (has_beta) dispatch_micro_dtype(cbsz, blgp, g); + else dispatch_micro_dtype(cbsz, blgp, g); + } +} + +template +static void dispatch_micro_k(int cbsz, int blgp, bool has_bias, bool has_gelu, bool has_beta, + bool has_partial_k, micro_globals_fp8 g) { + if (has_partial_k) dispatch_micro_epilogue(cbsz, blgp, has_bias, has_gelu, has_beta, g); + else dispatch_micro_epilogue(cbsz, blgp, has_bias, has_gelu, has_beta, g); +} + +template +static void dispatch_micro(bool is_1d2d, int cbsz, int blgp, bool has_bias, bool has_gelu, bool has_beta, + bool has_partial_k, micro_globals_fp8 g) { + if (is_1d2d) dispatch_micro_k(cbsz, blgp, has_bias, has_gelu, has_beta, has_partial_k, g); + else dispatch_micro_k(cbsz, blgp, has_bias, has_gelu, has_beta, has_partial_k, g); +} + +void kittens_blockwise_fp8_gemm_impl_cdna4( + const void *A, const void *B, void *C, + const void *scale_A, const void *scale_B, + int M, int N, int K, + int a_dtype, int b_dtype, + int a_scaling_mode, int b_scaling_mode, + int out_dtype, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const void *c_in, float beta, + hipStream_t stream) { + const bool has_bias = (bias != nullptr); + const bool has_gelu = (gelu_aux != nullptr); + const bool has_beta = (c_in != nullptr); + const bool has_partial_k = (K < 2 * BLOCK_K || K % BLOCK_K != 0); + + const void *kA = B, *kB = A; + const void *ksa = scale_B, *ksb = scale_A; + const int kM = N, kN = M; + const int ka_mode = b_scaling_mode, kb_mode = a_scaling_mode; + const int ka_dtype = b_dtype, kb_dtype = a_dtype; + + const bool is_1d2d = (kb_mode == KITTENS_BLOCK_SCALING_2D); + const int cbsz = (ka_dtype == KITTENS_FP8E5M2) ? 1 : 0; + const int blgp = (kb_dtype == KITTENS_FP8E5M2) ? 1 : 0; + float *sa = reinterpret_cast(const_cast(ksa)); + float *sb = reinterpret_cast(const_cast(ksb)); + + auto run = [&]() { + micro_globals_fp8 g{ + _gl_A_t((kittens::fp8e4m3 *)const_cast(kA), 1, 1, kM, K), + _gl_B_t((kittens::fp8e4m3 *)const_cast(kB), 1, 1, kN, K), + _gl_C_t((OType *)C, 1, 1, kM, kN), + _gl_SA(sa, 1, 1, 1, kM * K), + _gl_SB(sb, 1, 1, 1, kN * K), + bias, bias_dtype, gelu_aux, gelu_aux_dtype, + reinterpret_cast(c_in), beta, stream}; + dispatch_micro(is_1d2d, cbsz, blgp, has_bias, has_gelu, has_beta, has_partial_k, g); + }; + + if (out_dtype == KITTENS_FLOAT32) run.template operator()(); + else if (out_dtype == KITTENS_FLOAT16) run.template operator()(); + else run.template operator()(); +} + +} // namespace blockwise_gfx950 diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.h b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.h new file mode 100644 index 000000000..3932f4db4 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.h @@ -0,0 +1,42 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include +#include + +#ifndef KITTENS_DTYPE_ENUM_DEFINED +#define KITTENS_DTYPE_ENUM_DEFINED +enum KittensDType { + KITTENS_FLOAT32 = 4, + KITTENS_FLOAT16 = 5, + KITTENS_BFLOAT16 = 6, + KITTENS_FP8E4M3 = 7, + KITTENS_FP8E5M2 = 8, +}; +#endif // KITTENS_DTYPE_ENUM_DEFINED + +#ifndef KITTENS_SCALING_MODE_DEFINED +#define KITTENS_SCALING_MODE_DEFINED +enum KittensScalingMode { + KITTENS_BLOCK_SCALING_1D = 2, + KITTENS_BLOCK_SCALING_2D = 3, +}; +#endif // KITTENS_SCALING_MODE_DEFINED + +namespace blockwise_gfx950 { +void kittens_blockwise_fp8_gemm_impl_cdna4( + const void *A, const void *B, void *C, + const void *scale_A, const void *scale_B, + int M, int N, int K, + int a_dtype, int b_dtype, + int a_scaling_mode, int b_scaling_mode, + int out_dtype, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const void *c_in, float beta, + hipStream_t stream); +} // namespace blockwise_gfx950 diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_device.cuh b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_device.cuh new file mode 100644 index 000000000..11411ba52 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_device.cuh @@ -0,0 +1,239 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include +#include "kittens.cuh" +#include "../../../util/math.h" + +__device__ inline float rtne_bias(float v) { + uint32_t bits = __builtin_bit_cast(uint32_t, v); + if ((bits & 0x7f800000u) == 0x7f800000u) return v; + bits += 0x7fffu + ((bits >> 16) & 1u); + return __builtin_bit_cast(float, bits); +} + +template +__device__ inline void apply_rtne_bias(AccType &acc) { + #pragma unroll + for (int i = 0; i < AccType::height; i++) + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + acc.tiles[i][j].data[0].x = rtne_bias(acc.tiles[i][j].data[0].x); + acc.tiles[i][j].data[0].y = rtne_bias(acc.tiles[i][j].data[0].y); + acc.tiles[i][j].data[1].x = rtne_bias(acc.tiles[i][j].data[1].x); + acc.tiles[i][j].data[1].y = rtne_bias(acc.tiles[i][j].data[1].y); + } +} + +template +__device__ inline void store_masked(OType *c_ptr, const AccType &acc, + int m_off, int n_off, int M, int N) { + const int lane = kittens::laneid(); + const int row_g = 4 * (lane / 16); + const int col_g = lane % 16; + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const int m0 = m_off + i * 16 + row_g; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const int col = n_off + j * 16 + col_g; + if (col >= N) continue; + if (m0 + 0 < M) c_ptr[(m0 + 0) * N + col] = kittens::base_types::convertor::convert(acc.tiles[i][j].data[0].x); + if (m0 + 1 < M) c_ptr[(m0 + 1) * N + col] = kittens::base_types::convertor::convert(acc.tiles[i][j].data[0].y); + if (m0 + 2 < M) c_ptr[(m0 + 2) * N + col] = kittens::base_types::convertor::convert(acc.tiles[i][j].data[1].x); + if (m0 + 3 < M) c_ptr[(m0 + 3) * N + col] = kittens::base_types::convertor::convert(acc.tiles[i][j].data[1].y); + } + } +} + +__device__ inline float read_elem(const void *p, int dtype, int idx) { + if (dtype == 6) return __bfloat162float(reinterpret_cast(p)[idx]); + if (dtype == 5) return __half2float(reinterpret_cast(p)[idx]); + return reinterpret_cast(p)[idx]; +} + +template +__device__ inline float rtne_cast_roundtrip(float v) { + if constexpr (std::is_same_v) { + return v; + } else { + return static_cast(kittens::base_types::convertor::convert(rtne_bias(v))); + } +} + +template +__device__ inline void apply_epilogue( + AccType &acc, int m_off, int n_off, int M, int N, + const void *bias, int bias_dtype, + const void *gelu_aux, int gelu_aux_dtype, + const OType *c_in, float beta) { + const int lane = kittens::laneid(); + const int row_g = 4 * (lane / 16); + const int col_g = lane % 16; + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const int m0 = m_off + i * 16 + row_g; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const int col = n_off + j * 16 + col_g; + if (col >= N) continue; + float v[4] = { + acc.tiles[i][j].data[0].x, acc.tiles[i][j].data[0].y, + acc.tiles[i][j].data[1].x, acc.tiles[i][j].data[1].y, + }; + float bias_v = 0.f; + if constexpr (HAS_BIAS) bias_v = read_elem(bias, bias_dtype, col); + #pragma unroll + for (int r = 0; r < 4; r++) { + const int m_g = m0 + r; + if (m_g >= M) continue; + float x = v[r]; + if constexpr (HAS_BIAS) x += bias_v; + if constexpr (HAS_BETA) { + x = rtne_cast_roundtrip(x); + x += beta * static_cast(c_in[m_g * N + col]); + } + if constexpr (HAS_GELU) { + x *= transformer_engine::dgelu( + read_elem(gelu_aux, gelu_aux_dtype, m_g * N + col), {}); + } + v[r] = x; + } + acc.tiles[i][j].data[0].x = v[0]; + acc.tiles[i][j].data[0].y = v[1]; + acc.tiles[i][j].data[1].x = v[2]; + acc.tiles[i][j].data[1].y = v[3]; + } + } +} + +template +struct RowScale { + float v[HEIGHT][4]; +}; + +template +__device__ inline RowScale load_row_scale_lds( + const float *sa_lds, int local_m_base) { + RowScale rs; + const int row_g = 4 * (kittens::laneid() / 16); + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const int m0 = local_m_base + i * 16 + row_g; + rs.v[i][0] = sa_lds[m0 + 0]; + rs.v[i][1] = sa_lds[m0 + 1]; + rs.v[i][2] = sa_lds[m0 + 2]; + rs.v[i][3] = sa_lds[m0 + 3]; + } + return rs; +} + +template +__device__ inline RowScale load_row_scale( + const float *sa_row_k, int local_m_base, int m_valid) { + RowScale rs; + const int row_g = 4 * (kittens::laneid() / 16); + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const int m0 = local_m_base + i * 16 + row_g; + const int c0 = m0 + 0 < m_valid ? m0 + 0 : m_valid - 1; + const int c1 = m0 + 1 < m_valid ? m0 + 1 : m_valid - 1; + const int c2 = m0 + 2 < m_valid ? m0 + 2 : m_valid - 1; + const int c3 = m0 + 3 < m_valid ? m0 + 3 : m_valid - 1; + rs.v[i][0] = sa_row_k[c0]; + rs.v[i][1] = sa_row_k[c1]; + rs.v[i][2] = sa_row_k[c2]; + rs.v[i][3] = sa_row_k[c3]; + } + return rs; +} + +template +__device__ inline void scale_accumulate( + AccType &acc, const AccType &partial, const RowScale &rs, float sb_tile) { + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + const float s0 = rs.v[i][0] * sb_tile; + const float s1 = rs.v[i][1] * sb_tile; + const float s2 = rs.v[i][2] * sb_tile; + const float s3 = rs.v[i][3] * sb_tile; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + acc.tiles[i][j].data[0].x += partial.tiles[i][j].data[0].x * s0; + acc.tiles[i][j].data[0].y += partial.tiles[i][j].data[0].y * s1; + acc.tiles[i][j].data[1].x += partial.tiles[i][j].data[1].x * s2; + acc.tiles[i][j].data[1].y += partial.tiles[i][j].data[1].y * s3; + } + } +} + +template +struct ColScale { float v[WIDTH]; }; + +template +__device__ inline ColScale load_col_scale( + const float *sb_col_k, int local_n_base, int n_valid) { + ColScale cs; + const int col_g = kittens::laneid() % 16; + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const int n0 = local_n_base + j * 16 + col_g; + cs.v[j] = n0 < n_valid ? sb_col_k[n0] : 0.f; + } + return cs; +} + +template +__device__ inline void scale_accumulate_1d1d( + AccType &acc, const AccType &partial, const RowScale &rs, + const ColScale &cs) { + #pragma unroll + for (int i = 0; i < AccType::height; i++) { + #pragma unroll + for (int j = 0; j < AccType::width; j++) { + const float sc = cs.v[j]; + acc.tiles[i][j].data[0].x += partial.tiles[i][j].data[0].x * (rs.v[i][0] * sc); + acc.tiles[i][j].data[0].y += partial.tiles[i][j].data[0].y * (rs.v[i][1] * sc); + acc.tiles[i][j].data[1].x += partial.tiles[i][j].data[1].x * (rs.v[i][2] * sc); + acc.tiles[i][j].data[1].y += partial.tiles[i][j].data[1].y * (rs.v[i][3] * sc); + } + } +} + +template +__device__ inline void load_tile_masked(ST &dst, const T *src_base, int row_stride_elems, + int row_blk, int k_blk, int row_dim, int K) { + constexpr int elem_per_memcpy = sizeof(float4) / sizeof(T); + constexpr int elem_per_half = sizeof(float2) / sizeof(T); + constexpr int memcpy_per_row = ST::cols / elem_per_memcpy; + constexpr int total = (ST::rows * ST::cols) / elem_per_memcpy; + const int row_base = row_blk * ST::rows; + const int k_base = k_blk * ST::cols; + const T *src_ptr = src_base + (size_t)row_base * row_stride_elems + k_base; + const uint32_t dst_ptr = reinterpret_cast(&dst.data[0]); + constexpr int sub_rows = ST::underlying_subtile_rows; + constexpr int sub_cols = ST::underlying_subtile_cols; + constexpr int sub_bytes = ST::underlying_subtile_bytes; + constexpr int subs_per_row = ST::underlying_subtiles_per_row; + const int tid = threadIdx.x; + #pragma unroll + for (int idx = tid; idx < total; idx += NUM_THREADS) { + const int row = idx / memcpy_per_row; + const int col = (idx % memcpy_per_row) * elem_per_memcpy; + float4 v = {0.f, 0.f, 0.f, 0.f}; + if (row_base + row < row_dim && k_base + col < K) + v = kittens::load_global_vec4( + reinterpret_cast(src_ptr + (size_t)row * row_stride_elems + col)); + const int sub_id = (row / sub_rows) * subs_per_row + (col / sub_cols); + const int sub_off = sub_id * sub_bytes; + const int r = row % sub_rows, c = col % sub_cols; + kittens::store_shared_vec(dst_ptr + sub_off + dst.swizzle({r, c}), {v.x, v.y}); + kittens::store_shared_vec(dst_ptr + sub_off + dst.swizzle({r, c + elem_per_half}), {v.z, v.w}); + } + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); +} diff --git a/transformer_engine/common/gemm/kittens/mxfp8_gemm.cpp b/transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.cpp similarity index 99% rename from transformer_engine/common/gemm/kittens/mxfp8_gemm.cpp rename to transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.cpp index fb2dcb7f4..55df16aa0 100644 --- a/transformer_engine/common/gemm/kittens/mxfp8_gemm.cpp +++ b/transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.cpp @@ -125,7 +125,6 @@ void mxfp8_gemm_tn_kernel( const void *__restrict__ bias, int bias_dtype, int M, int N, int K) { - int k_iters = K / BLOCK_K; int tiles_M = M / BLOCK_ROW; int tiles_N = N / BLOCK_COL; diff --git a/transformer_engine/common/gemm/kittens/mxfp8_gemm.h b/transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.h similarity index 89% rename from transformer_engine/common/gemm/kittens/mxfp8_gemm.h rename to transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.h index ca39b21f9..8048882fc 100644 --- a/transformer_engine/common/gemm/kittens/mxfp8_gemm.h +++ b/transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.h @@ -9,6 +9,8 @@ #include // Values match NVTEDType in transformer_engine.h +#ifndef KITTENS_DTYPE_ENUM_DEFINED +#define KITTENS_DTYPE_ENUM_DEFINED enum KittensDType { KITTENS_FLOAT32 = 4, KITTENS_FLOAT16 = 5, @@ -16,6 +18,7 @@ enum KittensDType { KITTENS_FP8E4M3 = 7, KITTENS_FP8E5M2 = 8, }; +#endif // KITTENS_DTYPE_ENUM_DEFINED bool kittens_mxfp8_gemm( const void *A, const void *B, void *C, diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index 7e11c8054..deabe5222 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -33,7 +33,8 @@ #include "../util/logging.h" #ifdef USE_HIPKITTENS_GEMM -#include "kittens/mxfp8_gemm.h" +#include "kittens/blockwise_fp8_gemm.h" +#include "kittens/cdna4/mxfp8_gemm.h" #endif namespace transformer_engine { @@ -202,6 +203,10 @@ struct GemmParam { void *B_scale_inv = nullptr; int lda = 0; // A column strides int ldb = 0; // B column strides + // Blockwise FP8 only: per-operand scaling mode (NVTE_BLOCK_SCALING_1D/2D). + // Other scaling paths leave these at -1. + int A_scaling_mode = -1; + int B_scaling_mode = -1; }; constexpr int kMXFP8BlockSize = 32; @@ -397,7 +402,11 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla const transformer_engine::Tensor &B, const cublasOperation_t transB, const int m, const int n, const int k) { using namespace transformer_engine; - NVTE_CHECK(A.scaling_mode == B.scaling_mode, + // Blockwise FP8 legitimately mixes 1D (activation) and 2D (weight) scaling + // between A and B, so only require matching modes for non-blockwise inputs. + const bool a_blockwise = is_blockwise_fp8_scaling(A.scaling_mode); + const bool b_blockwise = is_blockwise_fp8_scaling(B.scaling_mode); + NVTE_CHECK((a_blockwise && b_blockwise) || A.scaling_mode == B.scaling_mode, "Inputs A and B to GEMM need to have the same scaling mode!"); NVTE_CHECK(A.has_data() || A.has_columnwise_data(), "Input A does not hold any data!"); NVTE_CHECK(B.has_data() || B.has_columnwise_data(), "Input B does not hold any data!"); @@ -448,6 +457,15 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype; ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr; ret.lda = k; + } else if (is_blockwise_fp8_scaling(A.scaling_mode)) { + // Blockwise FP8: transposed uses row-wise data/scale, non-transposed uses + // column-wise (matches the previous manual dispatch: inputA_col = !is_transa). + ret.A = is_A_transposed ? A.data.dptr : A.columnwise_data.dptr; + ret.transA = transA; + ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype; + ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr; + ret.A_scaling_mode = static_cast(A.scaling_mode); + ret.lda = is_A_transposed ? k : m; } else { NVTE_ERROR("A has unsupported scaling mode"); } @@ -493,6 +511,15 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype; ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr; ret.ldb = k; + } else if (is_blockwise_fp8_scaling(B.scaling_mode)) { + // Blockwise FP8: transposed uses column-wise data/scale, non-transposed uses + // row-wise (matches the previous manual dispatch: inputB_col = is_transb). + ret.B = is_B_transposed ? B.columnwise_data.dptr : B.data.dptr; + ret.transB = transB; + ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype; + ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr; + ret.B_scaling_mode = static_cast(B.scaling_mode); + ret.ldb = is_B_transposed ? n : k; } else { NVTE_ERROR("B has unsupported scaling mode"); } @@ -1980,6 +2007,69 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, handle = hipblaslt_handles[compute_stream_offset]; } +#ifdef USE_HIPKITTENS_GEMM + { + const bool inputA_blockwise = is_blockwise_fp8_scaling(inputA->scaling_mode); + const bool inputB_blockwise = is_blockwise_fp8_scaling(inputB->scaling_mode); + if (inputA_blockwise && inputB_blockwise) { + const bool has_bias = (inputBias->data.dptr != nullptr); + const bool has_gelu = (outputPreGelu->data.dptr != nullptr); + + NVTE_CHECK(outputD->data.dtype == DType::kBFloat16 || + outputD->data.dtype == DType::kFloat32 || + outputD->data.dtype == DType::kFloat16, + "Blockwise FP8 GEMM only supports bfloat16/float32/float16 output"); + NVTE_CHECK(inputB->scaling_mode == NVTE_BLOCK_SCALING_1D, + "Only 1D by 1D and 1D by 2D block scaling GEMM is supported"); + NVTE_CHECK(!(is_transa && is_transb), + "Blockwise FP8 GEMM does not support TT layout"); + NVTE_CHECK(!(inputA->dtype() == DType::kFloat8E5M2 && + inputB->dtype() == DType::kFloat8E5M2), + "Blockwise FP8 GEMM does not support e5m2 by e5m2 inputs"); + NVTE_CHECK(!has_gelu || grad, + "Blockwise FP8 GEMM only supports DGELU grad epilogue"); + NVTE_CHECK(!(has_bias && grad), + "Blockwise FP8 GEMM does not support bias with grad"); + NVTE_CHECK(!has_gelu || outputD->data.dtype == DType::kBFloat16, + "Blockwise FP8 GEMM DGELU epilogue only supports bfloat16 output"); + NVTE_CHECK(use_split_accumulator, + "Blockwise FP8 GEMM requires split accumulator"); + NVTE_CHECK((k % 16) == 0, + "GEMM K dimension must be multiple of 16 for blockwise FP8 scaling (got K=", k, ")"); + NVTE_CHECK((m % 16) == 0, + "GEMM M dimension must be multiple of 16 for blockwise FP8 scaling (got M=", m, ")"); + + const bool has_accum = (beta != 0.0f); + const void *bias = has_bias ? inputBias->data.dptr : nullptr; + const int bias_dtype = static_cast(inputBias->data.dtype); + const void *gelu_aux = has_gelu ? outputPreGelu->data.dptr : nullptr; + const int gelu_aux_dtype = static_cast(outputPreGelu->data.dtype); + const void *c_in = has_accum ? outputD->data.dptr : nullptr; + hipStream_t s = use_service_stream ? ss_ctl.stream : stream; + + // Canonical A/B/M/N (no swap); each arch impl absorbs its own swap. + GemmParam p = CanonicalizeGemmInput(*inputA, transa, *inputB, transb, m, n, k); + NVTE_CHECK(p.A != nullptr && p.A_scale_inv != nullptr && + p.B != nullptr && p.B_scale_inv != nullptr, + "Blockwise FP8 GEMM: missing rowwise or columnwise data/scale pointer."); + kittens_blockwise_fp8_gemm( + p.A, p.B, outputD->data.dptr, + p.A_scale_inv, p.B_scale_inv, + m, n, k, is_transa, is_transb, + static_cast(p.Atype), static_cast(p.Btype), + p.A_scaling_mode, p.B_scaling_mode, + static_cast(outputD->data.dtype), + bias, bias_dtype, gelu_aux, gelu_aux_dtype, c_in, beta, + s); + + if (use_service_stream) + { + release_service_stream(stream, ss_ctl); + } + return; + } + } +#endif bool is_mxfp8 = inputA->scaling_mode == NVTE_MXFP8_1D_SCALING || inputB->scaling_mode == NVTE_MXFP8_1D_SCALING; diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index b1648c711..956529e9d 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -48,7 +48,6 @@ from transformer_engine.pytorch.quantization import is_fp8_available from transformer_engine.pytorch.quantization import is_mxfp8_available from transformer_engine.pytorch.quantization import is_fp8_block_scaling_available -from transformer_engine.pytorch.quantization import is_fp8_block_scaling_quantization_available from transformer_engine.pytorch.quantization import is_nvfp4_available from transformer_engine.pytorch.quantization import is_mxfp4_available from transformer_engine.pytorch.quantization import get_default_recipe diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index 5f4a95bb8..8d1b98137 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -37,7 +37,7 @@ class Custom_DType_Dict(dict): def __missing__(self, key): if key in _FP8_KEYS: value = ( - get_torch_float8_e4m3_type() if key is tex.DType.kFloat8E4M3 + get_torch_float8_e4m3_type() if key == tex.DType.kFloat8E4M3 else get_torch_float8_e5m2_type() ) self[key] = value diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 2a2b035f2..2bcbc0c19 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -44,7 +44,6 @@ "is_fp8_available", "is_mxfp8_available", "is_fp8_block_scaling_available", - "is_fp8_block_scaling_quantization_available", "is_nvfp4_available", "get_default_recipe", "get_align_size_for_quantization", @@ -107,7 +106,10 @@ def _compute_nvfp4_support() -> Tuple[bool, str]: def _compute_fp8_block_scaling_support() -> Tuple[bool, str]: """Return if fp8 block scaling support is available""" if IS_HIP_EXTENSION: - return False, "FP8 block scaled gemm not yet supported for ROCm" + gpu_arch = get_device_compute_capability() + if gpu_arch in ((9, 4), (9, 5)): # TODO: enabled for gfx1250 when ready + return True, "" + return False, "Device arch gfx94x or newer is required for FP8 block scaling execution." if get_device_compute_capability() >= (9, 0) and float(torch.version.cuda) >= 12.9: return True, "" return ( @@ -115,15 +117,6 @@ def _compute_fp8_block_scaling_support() -> Tuple[bool, str]: "FP8 block scaled GEMM requires compute capability 9.0 or higher and CUDA >= 12.9.", ) -@functools.lru_cache(maxsize=None) -def check_fp8_block_scaling_quantization_support() -> Tuple[bool, str]: - """Return if fp8 block scaling quantization (cast only, no GEMM) is available""" - if IS_HIP_EXTENSION: - if get_device_compute_capability() >= (9, 4): - return True, "" - return False, "Device arch gfx94x or newer is required for FP8 block scaling quantization." - return check_fp8_block_scaling_support() - @torch.compiler.assume_constant_result def check_fp8_support() -> Tuple[bool, str]: """Return if fp8 support is available.""" @@ -316,23 +309,6 @@ def is_fp8_block_scaling_available(return_reason: bool = False) -> Union[bool, T return check_fp8_block_scaling_support() return check_fp8_block_scaling_support()[0] -def is_fp8_block_scaling_quantization_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: - """ - Determine if support is available for FP8 block scaling quantization (cast only, no GEMM). - - Parameters - ---------- - return_reason : bool, optional - If ``False`` (default), return only a boolean indicating availability. - If ``True``, return a tuple ``(is_available, reason)`` where ``reason`` provides - a human-readable explanation when required support is not available. The reason - will be an empty string if support for FP8 block scaling quantization is available. - - """ - if return_reason: - return check_fp8_block_scaling_quantization_support() - return check_fp8_block_scaling_quantization_support()[0] - def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: """ Determine if support is available for the NVFP4 recipe.