gfx942-only build bugfix and kittens refactor

transformer_engine/common/CMakeLists.txt

@@ -717,12 +717,7 @@ else() # USE_ROCM
   endif()
   if(USE_HIPKITTENS_GEMM)
     target_compile_definitions(transformer_engine PUBLIC USE_HIPKITTENS_GEMM)
-    if(KITTENS_HAVE_CDNA3)
-      target_compile_definitions(transformer_engine PUBLIC KITTENS_HAVE_CDNA3)
-    endif()
-    if(KITTENS_HAVE_CDNA4)
-      target_compile_definitions(transformer_engine PUBLIC KITTENS_HAVE_CDNA4)
-    endif()
+    target_compile_definitions(transformer_engine PUBLIC ${KITTENS_HAVE_DEFS})
     list(APPEND transformer_engine_LINKER_LIBS kittens_gemm)
   endif()
   target_link_libraries(transformer_engine PUBLIC ${transformer_engine_LINKER_LIBS})

transformer_engine/common/gemm/kittens/CMakeLists.txt

@@ -3,23 +3,26 @@
 
 cmake_minimum_required(VERSION 3.21)
 
-list(FIND CMAKE_HIP_ARCHITECTURES "gfx942" _gfx942_index)
-list(FIND CMAKE_HIP_ARCHITECTURES "gfx950" _gfx950_index)
+set(KITTENS_SUPPORTED_ARCHS gfx942 gfx950)
+
+set(_kittens_enabled_archs "")
+foreach(_arch IN LISTS KITTENS_SUPPORTED_ARCHS)
+    if(_arch IN_LIST CMAKE_HIP_ARCHITECTURES)
+        list(APPEND _kittens_enabled_archs ${_arch})
+    endif()
+endforeach()
 
 include(CheckCXXCompilerFlag)
 check_cxx_compiler_flag("-std=c++20" HAS_CXX20)
 
-if(_gfx942_index EQUAL -1 AND _gfx950_index EQUAL -1)
-    message(STATUS "HipKittens GEMM disabled (neither gfx942 nor gfx950 in CMAKE_HIP_ARCHITECTURES)")
+if(NOT _kittens_enabled_archs)
+    message(STATUS "HipKittens GEMM disabled (none of ${KITTENS_SUPPORTED_ARCHS} in CMAKE_HIP_ARCHITECTURES)")
     set(USE_HIPKITTENS_GEMM OFF PARENT_SCOPE)
 elseif(NOT HAS_CXX20)
     message(WARNING "HipKittens GEMMs require C++20")
     set(USE_HIPKITTENS_GEMM OFF PARENT_SCOPE)
 else()
-    set(HIPKITTENS_CDNA3_INCLUDE_DIR
-        "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens/include")
-    set(HIPKITTENS_CDNA4_INCLUDE_DIR
-        "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens/include")
+    set(HIPKITTENS_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens/include")
 
     set(CMAKE_CXX_STANDARD 20)
     project(kittens_gemm LANGUAGES HIP CXX)
@@ -28,53 +31,69 @@ else()
     include_directories("${ROCM_PATH}/include/hip")
 
     set(_kittens_arch_objs "")
+    set(_kittens_have_defs "")
 
-    if(NOT _gfx942_index EQUAL -1)
-        if(NOT EXISTS "${HIPKITTENS_CDNA3_INCLUDE_DIR}/kittens.cuh")
+    function(kittens_add_arch)
+        cmake_parse_arguments(A "" "TAG;GFX" "SOURCES;FLAGS;SOURCE_FLAGS" ${ARGN})
+        if(NOT "${A_GFX}" IN_LIST _kittens_enabled_archs)
+            return()
+        endif()
+        if(NOT EXISTS "${HIPKITTENS_INCLUDE_DIR}/kittens.cuh")
             message(FATAL_ERROR
-                "Could not find HipKittens (CDNA3) headers at ${HIPKITTENS_CDNA3_INCLUDE_DIR}. "
+                "Could not find HipKittens headers at ${HIPKITTENS_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 $<TARGET_OBJECTS:kittens_gemm_cdna3>)
-        set(KITTENS_HAVE_CDNA3 ON PARENT_SCOPE)
-    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'.")
+        string(TOUPPER "${A_TAG}" _tag_upper)
+        set(_target "kittens_gemm_${A_TAG}")
+
+        add_library(${_target} OBJECT ${A_SOURCES})
+        set_source_files_properties(${A_SOURCES} PROPERTIES LANGUAGE HIP)
+
+        # SOURCE_FLAGS is a flat <source> <flags> list
+        list(LENGTH A_SOURCE_FLAGS _sf_len)
+        if(_sf_len GREATER 0)
+            math(EXPR _sf_last "${_sf_len} - 2")
+            foreach(_i RANGE 0 ${_sf_last} 2)
+                math(EXPR _j "${_i} + 1")
+                list(GET A_SOURCE_FLAGS ${_i} _sf_src)
+                list(GET A_SOURCE_FLAGS ${_j} _sf_opts)
+                if(NOT "${_sf_src}" IN_LIST A_SOURCES)
+                    message(FATAL_ERROR
+                        "kittens_add_arch(TAG ${A_TAG}): SOURCE_FLAGS names ${_sf_src}, "
+                        "which is not in SOURCES.")
+                endif()
+                set_source_files_properties(${_sf_src} PROPERTIES COMPILE_FLAGS "${_sf_opts}")
+            endforeach()
         endif()
-        add_library(kittens_gemm_cdna4 OBJECT cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp)
-        set_source_files_properties(cdna4/mxfp8_gemm.cpp
-            PROPERTIES LANGUAGE HIP COMPILE_FLAGS "-gline-tables-only")
-        set_source_files_properties(cdna4/blockwise_fp8_gemm.cpp
-            PROPERTIES LANGUAGE HIP COMPILE_FLAGS "-gline-tables-only -ffast-math")
-        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 $<TARGET_OBJECTS:kittens_gemm_cdna4>)
-        set(KITTENS_HAVE_CDNA4 ON PARENT_SCOPE)
-    endif()
+
+        set_target_properties(${_target} PROPERTIES
+            HIP_ARCHITECTURES "${A_GFX}" POSITION_INDEPENDENT_CODE ON)
+        target_include_directories(${_target} PRIVATE
+            ${HIP_INCLUDE_DIRS} "${HIPKITTENS_INCLUDE_DIR}")
+        target_compile_options(${_target} PRIVATE
+            -DKITTENS_${_tag_upper} -fno-gpu-rdc -O3 ${A_FLAGS})
+        target_link_libraries(${_target} PRIVATE hip::host hip::device)
+
+        set(_kittens_arch_objs ${_kittens_arch_objs} $<TARGET_OBJECTS:${_target}> PARENT_SCOPE)
+        set(_kittens_have_defs ${_kittens_have_defs} KITTENS_HAVE_${_tag_upper} PARENT_SCOPE)
+    endfunction()
+
+    kittens_add_arch(TAG cdna3 GFX gfx942
+                     SOURCES cdna3/blockwise_fp8_gemm.cpp)
+
+    kittens_add_arch(TAG cdna4 GFX gfx950
+                     SOURCES cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp
+                     FLAGS -gline-tables-only
+                     SOURCE_FLAGS cdna4/blockwise_fp8_gemm.cpp "-ffast-math")
 
     add_library(kittens_gemm SHARED ${_kittens_arch_objs})
     set_target_properties(kittens_gemm PROPERTIES LINKER_LANGUAGE HIP)
     target_include_directories(kittens_gemm PRIVATE ${HIP_INCLUDE_DIRS})
     target_link_libraries(kittens_gemm PUBLIC hip::host hip::device)
 
+    set(KITTENS_HAVE_DEFS "${_kittens_have_defs}" PARENT_SCOPE)
+
     install(TARGETS kittens_gemm
         DESTINATION ${CMAKE_INSTALL_PREFIX}/transformer_engine/lib)
 endif()

transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.cpp

@@ -6,8 +6,9 @@
 #include <type_traits>
 #include "kittens.cuh"
 #include "../kittens_common.h"
+#include "../kittens_kernel_common.cuh"
 
-namespace {
+namespace te_kittens::cdna3::blockwise_fp8 {
 
 #include "blockwise_fp8_gemm_helper.cuh"
 
@@ -354,16 +355,6 @@ void micro_tk(const micro_globals<AType, BType, OType> g) {
     store_output(g.c.raw_ptr, C_accum[1], row * 4 + warp_row + WARPS_ROW, col * 4 + warp_col, M, N);
 }
 
-#define BOOL_SWITCH(val, NAME, ...) \
-    if (val) { constexpr bool NAME = true; __VA_ARGS__ } \
-    else { constexpr bool NAME = false; __VA_ARGS__ }
-
-static GemmEpilogue select_epilogue(bool has_bias, bool has_gelu, bool has_beta) {
-    if (has_gelu) return has_beta ? GemmEpilogue::GELU_AUX_BETA : GemmEpilogue::GELU_AUX;
-    if (has_bias) return has_beta ? GemmEpilogue::BIAS_BETA     : GemmEpilogue::BIAS;
-    return has_beta ? GemmEpilogue::BETA : GemmEpilogue::DEFAULT;
-}
-
 template <bool IS_1D2D, typename AType, typename BType, typename OType,
           GemmEpilogue EPILOGUE, bool IS_PARTIAL_K, bool A_E4M3, bool B_E4M3>
 static void dispatch_micro_epilogue(micro_globals<AType, BType, OType> g) {
@@ -374,8 +365,8 @@ static void dispatch_micro_epilogue(micro_globals<AType, BType, OType> g) {
         hipFuncSetAttribute((void*)kern, hipFuncAttributeMaxDynamicSharedMemorySize, mem_size);
         kern<<<g.grid(), g.block(), mem_size, g.stream>>>(g);
     };
-    BOOL_SWITCH(is_partial_m, IS_PARTIAL_M,
-        BOOL_SWITCH(is_partial_n, IS_PARTIAL_N,
+    KITTENS_BOOL_SWITCH(is_partial_m, IS_PARTIAL_M,
+        KITTENS_BOOL_SWITCH(is_partial_n, IS_PARTIAL_N,
             launch(micro_tk<AType, BType, OType, IS_PARTIAL_M, IS_PARTIAL_N, IS_1D2D, EPILOGUE, IS_PARTIAL_K, A_E4M3, B_E4M3>);
         )
     )
@@ -405,7 +396,7 @@ template <bool IS_1D2D, typename AType, typename BType, typename OType,
           bool A_E4M3, bool B_E4M3>
 static void dispatch_micro(micro_globals<AType, BType, OType> g,
                            bool has_bias, bool has_gelu, bool has_beta, bool has_partial_k) {
-    BOOL_SWITCH(has_partial_k, IS_PARTIAL_K,
+    KITTENS_BOOL_SWITCH(has_partial_k, IS_PARTIAL_K,
         dispatch_micro_k<IS_1D2D, AType, BType, OType, IS_PARTIAL_K, A_E4M3, B_E4M3>(g, has_bias, has_gelu, has_beta);
     )
 }
@@ -465,11 +456,9 @@ class BlockwiseGemmCdna3 final : public BlockwiseGemmBackend {
     }
 };
 
-#undef BOOL_SWITCH
-
-}
+}  // namespace te_kittens::cdna3::blockwise_fp8
 
 BlockwiseGemmBackend *BlockwiseGemmBackend::get_cdna3() {
-    static BlockwiseGemmCdna3 impl;
+    static te_kittens::cdna3::blockwise_fp8::BlockwiseGemmCdna3 impl;
     return &impl;
 }

transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh

@@ -8,6 +8,7 @@
 #include <type_traits>
 #include "kittens.cuh"
 #include "../../../util/math.h"
+using namespace te_kittens::blockwise;  // NOLINT(build/namespaces)
 
 typedef int int32x4_lds_t __attribute__((ext_vector_type(4)));
 struct __attribute__((packed)) buf_res { const void *ptr; uint32_t range; uint32_t config; };
@@ -87,12 +88,6 @@ __device__ inline float rtne_bias(float v) {
     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<const __hip_bfloat16 *>(p)[idx]);
-    if (dtype == 5) return __half2float(reinterpret_cast<const __half *>(p)[idx]);
-    return reinterpret_cast<const float *>(p)[idx];
-}
-
 template <typename OType>
 __device__ inline float rtne_cast_roundtrip(float v) {
     if constexpr (std::is_same_v<OType, float>) {
@@ -132,26 +127,6 @@ __device__ inline void store_output(OType *c_ptr, const AccType &Cacc,
     }
 }
 
-enum struct GemmEpilogue {
-    DEFAULT,
-    BIAS,
-    GELU_AUX,
-    BETA,
-    BIAS_BETA,
-    GELU_AUX_BETA,
-};
-
-__host__ __device__ inline constexpr bool epilogue_has_bias(GemmEpilogue e) {
-    return e == GemmEpilogue::BIAS || e == GemmEpilogue::BIAS_BETA;
-}
-__host__ __device__ inline constexpr bool epilogue_has_gelu(GemmEpilogue e) {
-    return e == GemmEpilogue::GELU_AUX || e == GemmEpilogue::GELU_AUX_BETA;
-}
-__host__ __device__ inline constexpr bool epilogue_has_beta(GemmEpilogue e) {
-    return e == GemmEpilogue::BETA || e == GemmEpilogue::BIAS_BETA
-        || e == GemmEpilogue::GELU_AUX_BETA;
-}
-
 template <typename OType, GemmEpilogue EPILOGUE, typename AccType>
 __device__ inline void apply_epilogue(
     AccType &Cacc, int Rtile, int Ctile, int M, int N,

transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.cpp

@@ -9,9 +9,10 @@
 #include <utility>
 #include "kittens.cuh"
 #include "../kittens_common.h"
+#include "../kittens_kernel_common.cuh"
 
 
-namespace {
+namespace te_kittens::cdna4::blockwise_fp8 {
 
 #include "blockwise_fp8_gemm_helper.cuh"
 
@@ -941,16 +942,6 @@ void micro_tk_partial_k(micro_globals<kittens::fp8e4m3, kittens::fp8e4m3, OType>
 }
 
 
-#define BOOL_SWITCH(val, NAME, ...) \
-    if (val) { constexpr bool NAME = true; __VA_ARGS__ } \
-    else { constexpr bool NAME = false; __VA_ARGS__ }
-
-static GemmEpilogue select_epilogue(bool has_bias, bool has_gelu, bool has_beta) {
-    if (has_gelu) return has_beta ? GemmEpilogue::GELU_AUX_BETA : GemmEpilogue::GELU_AUX;
-    if (has_bias) return has_beta ? GemmEpilogue::BIAS_BETA     : GemmEpilogue::BIAS;
-    return has_beta ? GemmEpilogue::BETA : GemmEpilogue::DEFAULT;
-}
-
 template <typename OType, int CBSZ, int BLGP, bool IS_1D2D,
           GemmEpilogue EPILOGUE, bool IS_PARTIAL_K>
 static void dispatch_micro_kernel(micro_globals_fp8<OType> g) {
@@ -991,13 +982,12 @@ static void dispatch_micro_epilogue(int cbsz, int blgp, bool has_bias, bool has_
 template <typename OType>
 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<OType> g) {
-    BOOL_SWITCH(is_1d2d, IS_1D2D,
-        BOOL_SWITCH(has_partial_k, IS_PARTIAL_K,
+    KITTENS_BOOL_SWITCH(is_1d2d, IS_1D2D,
+        KITTENS_BOOL_SWITCH(has_partial_k, IS_PARTIAL_K,
             dispatch_micro_epilogue<OType, IS_1D2D, IS_PARTIAL_K>(cbsz, blgp, has_bias, has_gelu, has_beta, g);
         )
     )
 }
-#undef BOOL_SWITCH
 
 template <typename OType, int CBSZ, int BLGP, GemmEpilogue EPILOGUE, bool B_BROADCAST>
 static void launch_pow2_kernel(const pow2_kernel_args<OType> &a) {
@@ -1049,7 +1039,7 @@ static void launch_pow2(int cbsz, int blgp, bool has_bias, bool has_gelu, bool h
     const int padM = tiles_M * BLOCK_M;
     const int padN = tiles_N * BLOCK_N;
 
-    const size_t sa_bytes = align_up_pow2ws((size_t)k_iters * padM * sizeof(uint32_t));
+    const size_t sa_bytes = kittens_align_up((size_t)k_iters * padM * sizeof(uint32_t), 256);
     uint32_t *packed_sa = reinterpret_cast<uint32_t *>(workspace);
     uint32_t *packed_sb = reinterpret_cast<uint32_t *>((uint8_t *)workspace + sa_bytes);
 
@@ -1117,7 +1107,7 @@ class BlockwiseGemmCdna4 final : public BlockwiseGemmBackend {
         const int k_iters = K / BLOCK_K;
         const int padM = ((kM + BLOCK_M - 1) / BLOCK_M) * BLOCK_M;
         const int padN = ((kN + BLOCK_N - 1) / BLOCK_N) * BLOCK_N;
-        const size_t pow2_ws_bytes = align_up_pow2ws((size_t)k_iters * padM * sizeof(uint32_t)) +
+        const size_t pow2_ws_bytes = kittens_align_up((size_t)k_iters * padM * sizeof(uint32_t), 256) +
                                      (size_t)k_iters * padN * sizeof(uint32_t);
         void *owned_ws = nullptr;
         if (use_pow2 && !has_partial_k &&
@@ -1163,9 +1153,9 @@ class BlockwiseGemmCdna4 final : public BlockwiseGemmBackend {
     }
 };
 
-}
+}  // namespace te_kittens::cdna4::blockwise_fp8
 
 BlockwiseGemmBackend *BlockwiseGemmBackend::get_cdna4() {
-    static BlockwiseGemmCdna4 impl;
+    static te_kittens::cdna4::blockwise_fp8::BlockwiseGemmCdna4 impl;
     return &impl;
 }

transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh

@@ -8,6 +8,7 @@
 #include <type_traits>
 #include "kittens.cuh"
 #include "../../../util/math.h"
+using namespace te_kittens::blockwise;  // NOLINT(build/namespaces)
 
 template <int HEIGHT>
 struct RowScale { float2 v[HEIGHT][2]; };
@@ -77,12 +78,6 @@ __device__ inline void store_output(OType *c_ptr, const AccType &acc,
     }
 }
 
-__device__ inline float read_elem(const void *p, int dtype, int idx) {
-    if (dtype == 6) return __bfloat162float(reinterpret_cast<const __hip_bfloat16 *>(p)[idx]);
-    if (dtype == 5) return __half2float(reinterpret_cast<const __half *>(p)[idx]);
-    return reinterpret_cast<const float *>(p)[idx];
-}
-
 template <typename OType>
 __device__ inline float round_to_out_dtype(float v) {
     if constexpr (std::is_same_v<OType, float>) {
@@ -94,26 +89,6 @@ __device__ inline float round_to_out_dtype(float v) {
     }
 }
 
-enum struct GemmEpilogue {
-    DEFAULT,
-    BIAS,
-    GELU_AUX,
-    BETA,
-    BIAS_BETA,
-    GELU_AUX_BETA,
-};
-
-__host__ __device__ inline constexpr bool epilogue_has_bias(GemmEpilogue e) {
-    return e == GemmEpilogue::BIAS || e == GemmEpilogue::BIAS_BETA;
-}
-__host__ __device__ inline constexpr bool epilogue_has_gelu(GemmEpilogue e) {
-    return e == GemmEpilogue::GELU_AUX || e == GemmEpilogue::GELU_AUX_BETA;
-}
-__host__ __device__ inline constexpr bool epilogue_has_beta(GemmEpilogue e) {
-    return e == GemmEpilogue::BETA || e == GemmEpilogue::BIAS_BETA
-        || e == GemmEpilogue::GELU_AUX_BETA;
-}
-
 template <typename OType, bool HAS_BIAS, bool HAS_GELU, bool HAS_BETA, typename AccType>
 __device__ inline void apply_epilogue(
     AccType &acc, int m_off, int n_off, int M, int N,
@@ -455,4 +430,3 @@ static void launch_pack_scales_pow2(const float *scales, uint32_t *packed, int p
     pack_scales_pow2_kernel<WEIGHT, TRANSPOSE><<<blocks, 256, 0, stream>>>(scales, packed, padded_dim, real_dim, scale_K, k_iters, scale_block);
 }
 
-static inline size_t align_up_pow2ws(size_t x) { return (x + 255) & ~size_t(255); }