Skip to content

gfx942-only build bugfix and kittens refactor #698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alextmagro wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

gfx942-only build bugfix and kittens refactor #698

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 1 addition & 6 deletions transformer_engine/common/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand Down
90 changes: 46 additions & 44 deletions transformer_engine/common/gemm/kittens/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -28,53 +31,52 @@ 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")
message(FATAL_ERROR
"Could not find HipKittens (CDNA3) headers at ${HIPKITTENS_CDNA3_INCLUDE_DIR}. "
"Try running 'git submodule update --init --recursive'.")
function(kittens_add_arch)
cmake_parse_arguments(A "" "TAG;GFX" "SOURCES;FLAGS" ${ARGN})
if(NOT "${A_GFX}" IN_LIST _kittens_enabled_archs)
return()
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")
if(NOT EXISTS "${HIPKITTENS_INCLUDE_DIR}/kittens.cuh")
message(FATAL_ERROR
"Could not find HipKittens (CDNA4) headers at ${HIPKITTENS_CDNA4_INCLUDE_DIR}. "
"Could not find HipKittens headers at ${HIPKITTENS_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/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()

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)
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)
set_source_files_properties(cdna4/blockwise_fp8_gemm.cpp
PROPERTIES COMPILE_FLAGS "-ffast-math")
Comment on lines +70 to +71

@github-actions github-actions Bot Aug 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the one thing the new kittens_add_arch abstraction doesn't cover, and it sits outside the arch guard: -ffast-math gets applied unconditionally, including on gfx942-only builds where cdna4/blockwise_fp8_gemm.cpp is never compiled. Harmless today (a source property on an unused file is a no-op), but it means the source list now lives in two places and has to be kept in sync by hand — exactly the coupling the refactor is trying to remove.

Verified the flags themselves are equivalent to before: -gline-tables-only moved from a per-source COMPILE_FLAGS to FLAGS (target-wide) and now covers both cdna4 sources, which matches the old behavior since both had it. Only -ffast-math is genuinely per-source.

Consider a SOURCE_FLAGS pair argument on kittens_add_arch so the per-source case is handled inside the guarded region too, or at minimum move this call inside the if(gfx950 IN_LIST ...) scope.


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()
25 changes: 7 additions & 18 deletions transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,9 @@
#include <type_traits>
#include "kittens.cuh"
#include "../kittens_common.h"
#include "../kittens_kernel_common.cuh"

namespace {
namespace te_kittens::cdna3 {

#include "blockwise_fp8_gemm_helper.cuh"

Expand Down Expand Up @@ -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) {
Expand All @@ -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>);
)
)
Expand Down Expand Up @@ -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);
)
}
Expand Down Expand Up @@ -465,11 +456,9 @@ class BlockwiseGemmCdna3 final : public BlockwiseGemmBackend {
}
};

#undef BOOL_SWITCH

}
} // namespace te_kittens::cdna3

BlockwiseGemmBackend *BlockwiseGemmBackend::get_cdna3() {
static BlockwiseGemmCdna3 impl;
static te_kittens::cdna3::BlockwiseGemmCdna3 impl;
return &impl;
}
27 changes: 1 addition & 26 deletions transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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; };
Expand Down Expand Up @@ -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>) {
Expand Down Expand Up @@ -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,
Expand Down
26 changes: 8 additions & 18 deletions transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,10 @@
#include <utility>
#include "kittens.cuh"
#include "../kittens_common.h"
#include "../kittens_kernel_common.cuh"


namespace {
namespace te_kittens::cdna4 {

#include "blockwise_fp8_gemm_helper.cuh"

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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 &&
Expand Down Expand Up @@ -1163,9 +1153,9 @@ class BlockwiseGemmCdna4 final : public BlockwiseGemmBackend {
}
};

}
} // namespace te_kittens::cdna4

BlockwiseGemmBackend *BlockwiseGemmBackend::get_cdna4() {
static BlockwiseGemmCdna4 impl;
static te_kittens::cdna4::BlockwiseGemmCdna4 impl;
return &impl;
}
28 changes: 1 addition & 27 deletions transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]; };
Expand Down Expand Up @@ -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>) {
Expand All @@ -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,
Expand Down Expand Up @@ -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); }
Loading
Loading