From 10a17bdc30a8ffbcfc422cb47ce1374c3f9c45d8 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Tue, 7 Jul 2026 19:51:25 +0000 Subject: [PATCH] Initial integration of a4w4 GEMM and QoLA version bump --- 3rdparty/QoLA | 2 +- setup.py | 7 + transformer_engine/common/CMakeLists.txt | 20 +++ .../common/aiter_gemm/CMakeLists.txt | 127 ++++++++++++++++++ .../include/aiter_gemm/aiter_gemm.hpp | 82 +++++++++++ .../common/aiter_gemm/qola_manifest.toml | 25 ++++ .../common/aiter_gemm/src/aiter_gemm_a4w4.cpp | 119 ++++++++++++++++ .../common/ck_fused_attn/qola_manifest.toml | 2 +- .../common/gemm/aiter_a4w4_gemm.cpp | 93 +++++++++++++ .../include/transformer_engine/aiter_gemm.h | 69 ++++++++++ .../pytorch/cpp_extensions/gemm.py | 18 +-- transformer_engine/pytorch/csrc/extensions.h | 11 ++ .../pytorch/csrc/extensions/gemm.cpp | 78 +++++++++++ .../pytorch/csrc/extensions/pybind.cpp | 9 ++ 14 files changed, 652 insertions(+), 10 deletions(-) create mode 100644 transformer_engine/common/aiter_gemm/CMakeLists.txt create mode 100644 transformer_engine/common/aiter_gemm/include/aiter_gemm/aiter_gemm.hpp create mode 100644 transformer_engine/common/aiter_gemm/qola_manifest.toml create mode 100644 transformer_engine/common/aiter_gemm/src/aiter_gemm_a4w4.cpp create mode 100644 transformer_engine/common/gemm/aiter_a4w4_gemm.cpp create mode 100644 transformer_engine/common/include/transformer_engine/aiter_gemm.h diff --git a/3rdparty/QoLA b/3rdparty/QoLA index 47e0f5501..b9ee37496 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit 47e0f5501af51521c8a911fbbefa4e24b92dd7b6 +Subproject commit b9ee374969e365758e27b06dd35cc919d1e098b4 diff --git a/setup.py b/setup.py index 2f4ae06e0..79a1eb044 100644 --- a/setup.py +++ b/setup.py @@ -96,6 +96,13 @@ def setup_common_extension() -> CMakeExtension: elif os.getenv("NVTE_FUSED_ATTN_CK") or os.getenv("NVTE_FUSED_ATTN"): cmake_flags.append("-DUSE_FUSED_ATTN_CK=ON") + # AITER a4w4 (FP4) GEMM backend (gfx950-only; CMake disables it cleanly + # on other arches). + if int(os.getenv("NVTE_AITER_GEMM", "1"))==0: + cmake_flags.append("-DUSE_AITER_GEMM=OFF") + else: + cmake_flags.append("-DUSE_AITER_GEMM=ON") + if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))) and os.getenv("NVTE_ENABLE_ROCSHMEM") is None: os.environ["NVTE_ENABLE_ROCSHMEM"] = '1' os.environ["NVTE_ENABLE_NVSHMEM"] = '0' diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 60b9f444d..097716558 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -10,6 +10,7 @@ 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_AITER_GEMM "Use AITER a4w4 (FP4) GEMM backend" ON) set(USE_CUDA OFF) if (USE_ROCM) @@ -336,6 +337,7 @@ if(USE_ROCM) fused_attn_rocm/fused_attn_aotriton.cpp fused_attn_rocm/fused_attn_ck.cpp fused_attn_rocm/utils.cpp + gemm/aiter_a4w4_gemm.cpp gemm/ck_grouped_gemm/ck_grouped_gemm.cpp gemm/ck_grouped_gemm/ck_grouped_gemm_fp8.cpp gemm/ck_grouped_gemm/ck_grouped_gemm_fp16.cpp @@ -582,6 +584,20 @@ else() # USE_ROCM add_subdirectory(gemm/kittens ${CMAKE_CURRENT_BINARY_DIR}/kittens) endif() + # AITER a4w4 (FP4) GEMM is gfx950-only; skip cleanly on other arches so a + # default-on USE_AITER_GEMM does not break gfx942/gfx1250 builds. The NVTE + # C API in gemm/aiter_a4w4_gemm.cpp keeps stub symbols when disabled. + list(FIND CMAKE_HIP_ARCHITECTURES "gfx950" _aiter_gemm_gfx950_idx) + if(USE_AITER_GEMM AND NOT _aiter_gemm_gfx950_idx EQUAL -1) + set(__BUILD_AITER_GEMM TRUE) + add_subdirectory(aiter_gemm ${CMAKE_CURRENT_BINARY_DIR}/aiter_gemm) + else() + set(__BUILD_AITER_GEMM FALSE) + if(USE_AITER_GEMM) + message(STATUS "USE_AITER_GEMM requested but no gfx950 target present; a4w4 GEMM backend disabled.") + endif() + endif() + find_package(hip) list(APPEND transformer_engine_LINKER_LIBS hip::host hip::device roctx64) find_package(hiprtc) @@ -600,6 +616,10 @@ else() # USE_ROCM target_compile_definitions(transformer_engine PUBLIC USE_HIPKITTENS_GEMM) list(APPEND transformer_engine_LINKER_LIBS kittens_gemm) endif() + if(__BUILD_AITER_GEMM) + target_compile_definitions(transformer_engine PUBLIC USE_AITER_GEMM) + list(APPEND transformer_engine_LINKER_LIBS aiter_gemm) + endif() target_link_libraries(transformer_engine PUBLIC ${transformer_engine_LINKER_LIBS}) endif() diff --git a/transformer_engine/common/aiter_gemm/CMakeLists.txt b/transformer_engine/common/aiter_gemm/CMakeLists.txt new file mode 100644 index 000000000..8ad5eb8f2 --- /dev/null +++ b/transformer_engine/common/aiter_gemm/CMakeLists.txt @@ -0,0 +1,127 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT + +cmake_minimum_required(VERSION 3.21) +set(CMAKE_CXX_STANDARD 17) +project(aiter_gemm LANGUAGES HIP CXX) + +set(AITER_GEMM_INSTALL_DIR "${CMAKE_INSTALL_PREFIX}/transformer_engine/lib") + +# a4w4 kernels are gfx950-only (ASM f4gemm .co blobs ship for gfx950). Drop +# unsupported arches; a runtime guard covers dispatch. +set(__AG_SUPPORTED_ARCHS "gfx950") +set(__AG_ARCHS) +foreach(__arch ${CMAKE_HIP_ARCHITECTURES}) + if(__arch IN_LIST __AG_SUPPORTED_ARCHS) + list(APPEND __AG_ARCHS ${__arch}) + else() + message(WARNING "[aiter_gemm] Skipping unsupported arch ${__arch} for a4w4 GEMM build.") + endif() +endforeach() +if(NOT __AG_ARCHS) + message(FATAL_ERROR + "[aiter_gemm] No supported architectures (need one of: ${__AG_SUPPORTED_ARCHS}). " + "Re-run the build with NVTE_AITER_GEMM=0 to disable the a4w4 GEMM backend.") +endif() + +set(__QOLA_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/QoLA") +set(__AITER_SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/qola/third_party/aiter") +if(DEFINED ENV{NVTE_AITER_SOURCE_DIR} AND NOT $ENV{NVTE_AITER_SOURCE_DIR} STREQUAL "") + set(__AITER_SOURCE_DIR $ENV{NVTE_AITER_SOURCE_DIR}) + message(STATUS "[aiter_gemm] Using AITER source from NVTE_AITER_SOURCE_DIR=${__AITER_SOURCE_DIR}. Disable AITER checkout.") + set(__SKIP_AITER_CHECKOUT TRUE) +else() + set(__SKIP_AITER_CHECKOUT FALSE) +endif() +set(AITER_INCLUDE_DIR "${__AITER_SOURCE_DIR}/csrc/include") + +if(NOT Python_EXECUTABLE) + find_package(Python COMPONENTS Interpreter QUIET) +endif() + +# Resolve the manifest-pinned AITER commit (defines AITER_SHA). +set(__QOLA_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml") +set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${__QOLA_MANIFEST}") +file(STRINGS "${__QOLA_MANIFEST}" __AITER_COMMIT_LINES + REGEX "^[ \t]*aiter_commit[ \t]*=[ \t]*\"[^\"]+\"") +list(LENGTH __AITER_COMMIT_LINES __AITER_COMMIT_COUNT) +if(NOT __AITER_COMMIT_COUNT EQUAL 1) + message(FATAL_ERROR + "Expected exactly one 'aiter_commit = \"...\"' line in ${__QOLA_MANIFEST}.") +endif() +list(GET __AITER_COMMIT_LINES 0 __AITER_COMMIT_LINE) +string(REGEX MATCH "\"([^\"]+)\"" _UNUSED "${__AITER_COMMIT_LINE}") +set(AITER_SHA "${CMAKE_MATCH_1}") + +if(Python_EXECUTABLE AND NOT __SKIP_AITER_CHECKOUT) + execute_process( + COMMAND sh -c + "PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \ + --manifest '${__QOLA_MANIFEST}' \ + --aiter-root '${__AITER_SOURCE_DIR}'" + RESULT_VARIABLE AITER_CHECKOUT_RESULT + OUTPUT_VARIABLE AITER_CHECKOUT_OUTPUT + ERROR_VARIABLE AITER_CHECKOUT_ERROR) + if(NOT AITER_CHECKOUT_RESULT EQUAL 0) + message(FATAL_ERROR + "[aiter_gemm] Failed to sync AITER source at ${__AITER_SOURCE_DIR} to ${AITER_SHA}.\n" + "${AITER_CHECKOUT_OUTPUT}\n${AITER_CHECKOUT_ERROR}") + endif() + message(STATUS "[aiter_gemm] Synced ${__AITER_SOURCE_DIR} to ${AITER_SHA}") +endif() + +if(NOT EXISTS "${AITER_INCLUDE_DIR}") + message(FATAL_ERROR + "[aiter_gemm] Could not find AITER API at ${AITER_INCLUDE_DIR}.") +endif() + +# Obtain the torch-free kernel libs: prebuilt bypass or build from source (QoLA). +set(__AITER_GEMM_LIB_PATH "") +set(__QOLA_INCLUDE_DIR "") +if(DEFINED ENV{AITER_GEMM_PATH}) + message(STATUS "[aiter_gemm] Using prebuilt libs from AITER_GEMM_PATH=$ENV{AITER_GEMM_PATH}") + set(__AITER_GEMM_LIB_PATH "$ENV{AITER_GEMM_PATH}/lib") + set(__QOLA_INCLUDE_DIR "$ENV{AITER_GEMM_PATH}/include") +else() + list(JOIN __AG_ARCHS ";" GPU_ARCHS_STR) + set(__QOLA_BUILD_DIR "${CMAKE_CURRENT_BINARY_DIR}/qola") + message(STATUS "[aiter_gemm] Building a4w4 GEMM kernels for ${GPU_ARCHS_STR} via QoLA.") + execute_process( + COMMAND ${CMAKE_COMMAND} -E env "PYTHONPATH=${__QOLA_DIR}:$ENV{PYTHONPATH}" + ${Python_EXECUTABLE} -m qola.cli build + --manifest ${__QOLA_MANIFEST} + --aiter-root ${__AITER_SOURCE_DIR} + --output-dir ${__QOLA_BUILD_DIR} + --arch "${GPU_ARCHS_STR}" + --skip-checkout + RESULT_VARIABLE QOLA_BUILD_RESULT) + if(NOT QOLA_BUILD_RESULT EQUAL 0) + message(FATAL_ERROR "[aiter_gemm] QoLA build failed.") + endif() + set(__AITER_GEMM_LIB_PATH "${__QOLA_BUILD_DIR}/lib") + set(__QOLA_INCLUDE_DIR "${__QOLA_BUILD_DIR}/include") +endif() + +if(NOT EXISTS "${__QOLA_INCLUDE_DIR}/qola_config.h") + message(FATAL_ERROR "[aiter_gemm] Could not find QoLA public headers at ${__QOLA_INCLUDE_DIR}.") +endif() + +add_library(aiter_gemm SHARED src/aiter_gemm_a4w4.cpp) +target_include_directories(aiter_gemm PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/include") +target_include_directories(aiter_gemm PRIVATE ${AITER_INCLUDE_DIR} ${__QOLA_INCLUDE_DIR}) + +find_package(hip) +target_link_directories(aiter_gemm PUBLIC ${__AITER_GEMM_LIB_PATH}) +target_link_libraries(aiter_gemm PUBLIC + hip::host hip::device + -l:te_libgemm_a4w4_blockscale.so + -l:te_libgemm_a4w4_asm.so) +set_target_properties(aiter_gemm PROPERTIES INSTALL_RPATH "$ORIGIN") + +if(NOT "${__AITER_GEMM_LIB_PATH}" STREQUAL "${AITER_GEMM_INSTALL_DIR}") + install(FILES + ${__AITER_GEMM_LIB_PATH}/te_libgemm_a4w4_blockscale.so + ${__AITER_GEMM_LIB_PATH}/te_libgemm_a4w4_asm.so + DESTINATION ${AITER_GEMM_INSTALL_DIR}) +endif() +install(TARGETS aiter_gemm DESTINATION ${AITER_GEMM_INSTALL_DIR}) diff --git a/transformer_engine/common/aiter_gemm/include/aiter_gemm/aiter_gemm.hpp b/transformer_engine/common/aiter_gemm/include/aiter_gemm/aiter_gemm.hpp new file mode 100644 index 000000000..598b11477 --- /dev/null +++ b/transformer_engine/common/aiter_gemm/include/aiter_gemm/aiter_gemm.hpp @@ -0,0 +1,82 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#ifndef AITER_GEMM_H +#define AITER_GEMM_H + +#include +#include + +// TE-side wrapper around AITER's torch-free a4w4 (FP4 x FP4) GEMM kernels. +// +// This header is intentionally free of any AITER / QoLA headers so that +// libtransformer_engine.so can consume it without pulling in the AITER kernel +// headers. The translation to AITER's aiter_tensor_t POD lives in the .cpp. +// +// Kernel selection (tuned-CSV lookup) and weight/scale pre-shuffling are the +// caller's responsibility -- these entry points are thin executors that take a +// resolved kernel name and already-shuffled inputs. +namespace aiter_gemm { + +// Mirrors AiterDtype in AITER's aiter_enum.h (only the subset a4w4 needs). +enum class DType { + fp4x2 = 0, /*!< two packed FP4 (E2M1) values per byte */ + e8m0 = 1, /*!< 8-bit exponent-only microscaling factor (1 byte) */ + bf16 = 2, + fp16 = 3, + fp32 = 4, + u8 = 5, + i8 = 6, +}; + +// Lightweight tensor descriptor (raw device pointer + layout). The caller owns +// the storage; the descriptor must outlive the call but not the storage. +struct TensorDesc { + void* ptr = nullptr; + int ndim = 0; + int64_t shape[8] = {0}; + int64_t strides[8] = {0}; + DType dtype = DType::fp4x2; + int device_id = 0; +}; + +// CK blockscale a4w4 GEMM: Y = XQ @ WQ^T with per-1x32 microscaling. +// XQ [M, K/2] fp4x2 +// WQ [N, K/2] fp4x2 +// x_scale [M, K/32] e8m0 +// w_scale [N, K/32] e8m0 +// Y [M, N] bf16 / fp16 (output, pre-allocated) +// `kernel_name` empty -> default heuristic; non-empty must exist in the +// compiled registry. Returns hipSuccess on success, hipErrorUnknown on a +// kernel-side failure (message logged). +hipError_t gemm_a4w4_blockscale(const TensorDesc& XQ, + const TensorDesc& WQ, + const TensorDesc& x_scale, + const TensorDesc& w_scale, + const TensorDesc& Y, + int split_k, + const char* kernel_name, + hipStream_t stream); + +// ASM (f4gemm) a4w4 GEMM: D = alpha*A*B + beta*C. +// A/B/scales/out layout as above; `bias` may be null. +// `kernel_name` empty -> ASM heuristic. +hipError_t gemm_a4w4_asm(const TensorDesc& A, + const TensorDesc& B, + const TensorDesc& a_scale, + const TensorDesc& b_scale, + const TensorDesc& out, + const TensorDesc* bias, + const char* kernel_name, + float alpha, + float beta, + int bpreshuffle, + int log2_k_split, + hipStream_t stream); + +} // namespace aiter_gemm + +#endif // AITER_GEMM_H diff --git a/transformer_engine/common/aiter_gemm/qola_manifest.toml b/transformer_engine/common/aiter_gemm/qola_manifest.toml new file mode 100644 index 000000000..7a6d5e4b6 --- /dev/null +++ b/transformer_engine/common/aiter_gemm/qola_manifest.toml @@ -0,0 +1,25 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# +# QoLA consumer manifest for TE's AITER a4w4 (FP4) GEMM integration. +# Builds the torch-free CK-blockscale and ASM kernels as standalone +# te_libgemm_a4w4_*.so shared libraries (cpp_itfs mode). + +[qola] +# NB: with NVTE_AITER_SOURCE_DIR set the build skips `qola checkout` and uses +# that tree directly; this pin documents the intended AITER commit for the +# reproducible (checkout) path. +aiter_commit = "773c8f67b9bfd99be95700a527d887afb4d8ba6a" +namespace = "te" +rocm_versions = ["7.2"] + +[build] +architectures = ["gfx950"] + +[[modules]] +name = "libgemm_a4w4_blockscale" +mode = "cpp_itfs" + +[[modules]] +name = "libgemm_a4w4_asm" +mode = "cpp_itfs" diff --git a/transformer_engine/common/aiter_gemm/src/aiter_gemm_a4w4.cpp b/transformer_engine/common/aiter_gemm/src/aiter_gemm_a4w4.cpp new file mode 100644 index 000000000..049d9fb8b --- /dev/null +++ b/transformer_engine/common/aiter_gemm/src/aiter_gemm_a4w4.cpp @@ -0,0 +1,119 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include "aiter_gemm/aiter_gemm.hpp" + +#include +#include +#include + +#include "aiter_tensor.h" // aiter_tensor_t, AiterDtype +#include "qola_gemm_a4w4_asm.h" +#include "qola_gemm_a4w4_blockscale.h" + +namespace aiter_gemm { + +namespace { + +AiterDtype to_aiter_dtype(DType dt) { + switch (dt) { + case DType::fp4x2: + return AITER_DTYPE_fp4x2; + case DType::e8m0: + return AITER_DTYPE_fp8_e8m0; + case DType::bf16: + return AITER_DTYPE_bf16; + case DType::fp16: + return AITER_DTYPE_fp16; + case DType::fp32: + return AITER_DTYPE_fp32; + case DType::u8: + return AITER_DTYPE_u8; + case DType::i8: + return AITER_DTYPE_i8; + } + return AITER_DTYPE_u8; +} + +// Build an aiter_tensor_t POD from a TE-side descriptor. Shares the caller's +// device pointer; no ownership is transferred. +aiter_tensor_t to_aiter_tensor(const TensorDesc& d) { + aiter_tensor_t t{}; + t.ptr = d.ptr; + t.ndim = d.ndim; + size_t numel = (d.ndim > 0) ? 1 : 0; + for (int i = 0; i < d.ndim; ++i) { + t.shape[i] = d.shape[i]; + t.strides[i] = d.strides[i]; + numel *= static_cast(d.shape[i]); + } + t.numel_ = numel; + t.dtype_ = to_aiter_dtype(d.dtype); + t.device_id = d.device_id; + return t; +} + +} // namespace + +hipError_t gemm_a4w4_blockscale(const TensorDesc& XQ, + const TensorDesc& WQ, + const TensorDesc& x_scale, + const TensorDesc& w_scale, + const TensorDesc& Y, + int split_k, + const char* kernel_name, + hipStream_t stream) { + try { + aiter_tensor_t a_xq = to_aiter_tensor(XQ); + aiter_tensor_t a_wq = to_aiter_tensor(WQ); + aiter_tensor_t a_xs = to_aiter_tensor(x_scale); + aiter_tensor_t a_ws = to_aiter_tensor(w_scale); + aiter_tensor_t a_y = to_aiter_tensor(Y); + QOLA_NS(gemm_a4w4_blockscale) + (a_xq, a_wq, a_xs, a_ws, a_y, split_k, stream, + kernel_name ? std::string(kernel_name) : std::string()); + return hipSuccess; + } catch (const std::exception& e) { + std::fprintf(stderr, "[aiter_gemm] gemm_a4w4_blockscale failed: %s\n", e.what()); + return hipErrorUnknown; + } +} + +hipError_t gemm_a4w4_asm(const TensorDesc& A, + const TensorDesc& B, + const TensorDesc& a_scale, + const TensorDesc& b_scale, + const TensorDesc& out, + const TensorDesc* bias, + const char* kernel_name, + float alpha, + float beta, + int bpreshuffle, + int log2_k_split, + hipStream_t stream) { + try { + aiter_tensor_t a_a = to_aiter_tensor(A); + aiter_tensor_t a_b = to_aiter_tensor(B); + aiter_tensor_t a_as = to_aiter_tensor(a_scale); + aiter_tensor_t a_bs = to_aiter_tensor(b_scale); + aiter_tensor_t a_out = to_aiter_tensor(out); + aiter_tensor_t a_bias; + aiter_tensor_t* a_bias_ptr = nullptr; + if (bias != nullptr && bias->ptr != nullptr) { + a_bias = to_aiter_tensor(*bias); + a_bias_ptr = &a_bias; + } + QOLA_NS(gemm_a4w4_asm) + (&a_a, &a_b, &a_as, &a_bs, &a_out, kernel_name ? kernel_name : "", a_bias_ptr, alpha, beta, + bpreshuffle, log2_k_split, stream); + return hipSuccess; + } catch (const std::exception& e) { + std::fprintf(stderr, "[aiter_gemm] gemm_a4w4_asm failed: %s\n", e.what()); + return hipErrorUnknown; + } +} + +} // namespace aiter_gemm diff --git a/transformer_engine/common/ck_fused_attn/qola_manifest.toml b/transformer_engine/common/ck_fused_attn/qola_manifest.toml index bba5ccc22..14db75aa0 100644 --- a/transformer_engine/common/ck_fused_attn/qola_manifest.toml +++ b/transformer_engine/common/ck_fused_attn/qola_manifest.toml @@ -1,5 +1,5 @@ [qola] -aiter_commit = "d32b0cb62ecc32bb1a858e8437d58eb9b3856af6" # pinned AITER submodule commit +aiter_commit = "773c8f67b9bfd99be95700a527d887afb4d8ba6a" # pinned AITER submodule commit namespace = "te" rocm_versions = ["7.2"] diff --git a/transformer_engine/common/gemm/aiter_a4w4_gemm.cpp b/transformer_engine/common/gemm/aiter_a4w4_gemm.cpp new file mode 100644 index 000000000..a08530d07 --- /dev/null +++ b/transformer_engine/common/gemm/aiter_a4w4_gemm.cpp @@ -0,0 +1,93 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVTE C API for AITER's a4w4 (FP4) GEMM. Translates the C descriptor into +// the aiter_gemm wrapper library's C++ types. When TE is built without the +// AITER a4w4 backend (USE_AITER_GEMM undefined -- e.g. non-gfx950), the +// symbols still exist but report failure so the framework layer degrades +// gracefully instead of failing to link. + +#include "transformer_engine/aiter_gemm.h" + +#ifdef USE_AITER_GEMM + +#include + +#include "aiter_gemm/aiter_gemm.hpp" + +namespace { + +aiter_gemm::TensorDesc to_desc(const NVTEAiterGemmTensor *t) { + aiter_gemm::TensorDesc d; + d.ptr = t->ptr; + d.ndim = t->ndim; + for (int i = 0; i < t->ndim && i < 8; ++i) { + d.shape[i] = t->shape[i]; + d.strides[i] = t->strides[i]; + } + d.dtype = static_cast(t->dtype); + d.device_id = t->device_id; + return d; +} + +} // namespace + +extern "C" int nvte_aiter_gemm_a4w4_blockscale(const NVTEAiterGemmTensor *XQ, + const NVTEAiterGemmTensor *WQ, + const NVTEAiterGemmTensor *x_scale, + const NVTEAiterGemmTensor *w_scale, + const NVTEAiterGemmTensor *Y, int split_k, + const char *kernel_name, void *stream) { + aiter_gemm::TensorDesc xq = to_desc(XQ); + aiter_gemm::TensorDesc wq = to_desc(WQ); + aiter_gemm::TensorDesc xs = to_desc(x_scale); + aiter_gemm::TensorDesc ws = to_desc(w_scale); + aiter_gemm::TensorDesc y = to_desc(Y); + hipError_t err = aiter_gemm::gemm_a4w4_blockscale( + xq, wq, xs, ws, y, split_k, kernel_name, static_cast(stream)); + return err == hipSuccess ? 0 : 1; +} + +extern "C" int nvte_aiter_gemm_a4w4_asm(const NVTEAiterGemmTensor *A, const NVTEAiterGemmTensor *B, + const NVTEAiterGemmTensor *a_scale, + const NVTEAiterGemmTensor *b_scale, + const NVTEAiterGemmTensor *out, + const NVTEAiterGemmTensor *bias, const char *kernel_name, + float alpha, float beta, int bpreshuffle, int log2_k_split, + void *stream) { + aiter_gemm::TensorDesc a = to_desc(A); + aiter_gemm::TensorDesc b = to_desc(B); + aiter_gemm::TensorDesc as = to_desc(a_scale); + aiter_gemm::TensorDesc bs = to_desc(b_scale); + aiter_gemm::TensorDesc o = to_desc(out); + aiter_gemm::TensorDesc bias_desc; + const aiter_gemm::TensorDesc *bias_ptr = nullptr; + if (bias != nullptr && bias->ptr != nullptr) { + bias_desc = to_desc(bias); + bias_ptr = &bias_desc; + } + hipError_t err = aiter_gemm::gemm_a4w4_asm(a, b, as, bs, o, bias_ptr, kernel_name, alpha, beta, + bpreshuffle, log2_k_split, + static_cast(stream)); + return err == hipSuccess ? 0 : 1; +} + +#else // !USE_AITER_GEMM + +extern "C" int nvte_aiter_gemm_a4w4_blockscale(const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, int, const char *, void *) { + return -1; +} + +extern "C" int nvte_aiter_gemm_a4w4_asm(const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, + const char *, float, float, int, int, void *) { + return -1; +} + +#endif // USE_AITER_GEMM diff --git a/transformer_engine/common/include/transformer_engine/aiter_gemm.h b/transformer_engine/common/include/transformer_engine/aiter_gemm.h new file mode 100644 index 000000000..ca302de47 --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/aiter_gemm.h @@ -0,0 +1,69 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file aiter_gemm.h + * \brief C API for AITER a4w4 (FP4 x FP4) GEMM (ROCm only). + * + * Thin executor entry points: kernel selection (tuned-CSV lookup) and + * weight/scale pre-shuffling happen in the framework layer; these functions + * take a resolved kernel name and already-shuffled inputs. + */ + +#ifndef TRANSFORMER_ENGINE_AITER_GEMM_H_ +#define TRANSFORMER_ENGINE_AITER_GEMM_H_ + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! \brief Element type of an a4w4 GEMM operand. Values match the internal + * aiter_gemm::DType mapping. */ +typedef enum { + kNVTEAiterGemmFP4x2 = 0, /*!< two packed FP4 (E2M1) values per byte */ + kNVTEAiterGemmE8M0 = 1, /*!< 8-bit exponent-only microscale (1 byte) */ + kNVTEAiterGemmBF16 = 2, + kNVTEAiterGemmFP16 = 3, + kNVTEAiterGemmFP32 = 4, + kNVTEAiterGemmU8 = 5, + kNVTEAiterGemmI8 = 6, +} NVTEAiterGemmDType; + +/*! \brief Lightweight device-tensor descriptor (raw pointer + layout). */ +typedef struct { + void *ptr; + int ndim; + int64_t shape[8]; + int64_t strides[8]; + int dtype; /*!< one of NVTEAiterGemmDType */ + int device_id; +} NVTEAiterGemmTensor; + +/*! \brief CK blockscale a4w4 GEMM: Y = XQ @ WQ^T with per-1x32 microscaling. + * \return 0 on success, nonzero on failure (or if TE was built without the + * AITER a4w4 backend). + */ +int nvte_aiter_gemm_a4w4_blockscale(const NVTEAiterGemmTensor *XQ, const NVTEAiterGemmTensor *WQ, + const NVTEAiterGemmTensor *x_scale, + const NVTEAiterGemmTensor *w_scale, const NVTEAiterGemmTensor *Y, + int split_k, const char *kernel_name, void *stream); + +/*! \brief ASM (f4gemm) a4w4 GEMM: D = alpha*A*B + beta*C. `bias` may be NULL. + * \return 0 on success, nonzero on failure (or if TE was built without the + * AITER a4w4 backend). + */ +int nvte_aiter_gemm_a4w4_asm(const NVTEAiterGemmTensor *A, const NVTEAiterGemmTensor *B, + const NVTEAiterGemmTensor *a_scale, const NVTEAiterGemmTensor *b_scale, + const NVTEAiterGemmTensor *out, const NVTEAiterGemmTensor *bias, + const char *kernel_name, float alpha, float beta, int bpreshuffle, + int log2_k_split, void *stream); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TRANSFORMER_ENGINE_AITER_GEMM_H_ diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 3e787f820..f5dbb1188 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -174,10 +174,11 @@ def _fp4_gemm_core(A_fp4, A_scales, B_fp4, B_scales, out_dtype=torch.bfloat16, Routes to the ASM backend when ``kernel_name`` is an ASM-mangled symbol (starts with ``_ZN``) or empty (heuristic). Otherwise routes to the CK blockscale backend, matching AITER's own ``gemm_a4w4`` dispatcher. + + The GEMM itself runs through TE's torch-free AITER libs (``tex``); only the + weight/scale shuffle helper is still sourced from the ``aiter`` package. """ - import aiter from aiter.ops.shuffle import shuffle_weight - from aiter.ops.gemm_op_a4w4 import gemm_a4w4_blockscale _fp4_dtype = torch.float4_e2m1fn_x2 A_fp4 = A_fp4.view(_fp4_dtype) if A_fp4.dtype != _fp4_dtype else A_fp4 @@ -198,17 +199,18 @@ def _fp4_gemm_core(A_fp4, A_scales, B_fp4, B_scales, out_dtype=torch.bfloat16, use_ck = bool(kernel_name) and kernel_name.find("_ZN") == -1 if use_ck: - result = gemm_a4w4_blockscale( + tex.gemm_a4w4_blockscale( A_fp4, B_shuffled, A_scales_uint8, B_scales_uint8, out_hp, - splitK=log2_k_split, + split_k=log2_k_split, kernel_name=kernel_name, ) else: - result = aiter.gemm_a4w4_asm( - A_fp4, B_shuffled, A_scales_uint8, B_scales_uint8, - out_hp, kernel_name, None, - bpreshuffle=True, log2_k_split=log2_k_split, + tex.gemm_a4w4_asm( + A_fp4, B_shuffled, A_scales_uint8, B_scales_uint8, out_hp, + bias=None, kernel_name=kernel_name, + alpha=1.0, beta=0.0, bpreshuffle=True, log2_k_split=log2_k_split, ) + result = out_hp return result[:M, :] if result.shape[0] > M else result diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index e1e7866a7..a5bb67d7c 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -170,6 +170,17 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p at::Tensor workspace_cublas, bool use_split_accumulator, int math_sm_count); +#ifdef USE_ROCM +// AITER a4w4 (FP4) GEMM executors. Kernel selection + weight/scale shuffle are +// done in Python; these take a resolved kernel name and already-shuffled inputs. +void gemm_a4w4_blockscale(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, + at::Tensor Y, int64_t split_k, std::string kernel_name); + +void gemm_a4w4_asm(at::Tensor A, at::Tensor B, at::Tensor a_scale, at::Tensor b_scale, + at::Tensor out, std::optional bias, std::string kernel_name, + double alpha, double beta, bool bpreshuffle, int64_t log2_k_split); +#endif // USE_ROCM + /*************************************************************************************************** * Transpose **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 59b88e995..8f1fee19b 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -16,6 +16,9 @@ #include "common/util/system.h" #include "pybind.h" #include "transformer_engine/transformer_engine.h" +#ifdef USE_ROCM +#include "transformer_engine/aiter_gemm.h" +#endif #include "util.h" #include @@ -822,4 +825,79 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p return py::reinterpret_borrow(D); } +#ifdef USE_ROCM + +namespace { + +NVTEAiterGemmTensor make_aiter_gemm_tensor(const at::Tensor &t, int dtype) { + NVTEAiterGemmTensor d{}; + d.ptr = t.data_ptr(); + d.ndim = static_cast(t.dim()); + for (int i = 0; i < d.ndim && i < 8; ++i) { + d.shape[i] = t.size(i); + d.strides[i] = t.stride(i); + } + d.dtype = dtype; + d.device_id = static_cast(t.device().index()); + return d; +} + +int aiter_gemm_out_dtype(const at::Tensor &t) { + switch (t.scalar_type()) { + case at::kBFloat16: + return kNVTEAiterGemmBF16; + case at::kHalf: + return kNVTEAiterGemmFP16; + default: + NVTE_CHECK(false, "AITER a4w4 GEMM output must be bf16 or fp16"); + return kNVTEAiterGemmBF16; + } +} + +} // namespace + +// CK blockscale a4w4 GEMM. Inputs are already FP4-packed / pre-shuffled and the +// kernel is already selected in Python (kernel_name may be empty -> heuristic). +void gemm_a4w4_blockscale(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, + at::Tensor Y, int64_t split_k, std::string kernel_name) { + NVTEAiterGemmTensor xq = make_aiter_gemm_tensor(XQ, kNVTEAiterGemmFP4x2); + NVTEAiterGemmTensor wq = make_aiter_gemm_tensor(WQ, kNVTEAiterGemmFP4x2); + NVTEAiterGemmTensor xs = make_aiter_gemm_tensor(x_scale, kNVTEAiterGemmE8M0); + NVTEAiterGemmTensor ws = make_aiter_gemm_tensor(w_scale, kNVTEAiterGemmE8M0); + NVTEAiterGemmTensor y = make_aiter_gemm_tensor(Y, aiter_gemm_out_dtype(Y)); + int rc; + NVTE_SCOPED_GIL_RELEASE({ + rc = nvte_aiter_gemm_a4w4_blockscale(&xq, &wq, &xs, &ws, &y, static_cast(split_k), + kernel_name.c_str(), at::cuda::getCurrentCUDAStream()); + }); + NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_blockscale failed (rc=", rc, ")"); +} + +// ASM (f4gemm) a4w4 GEMM. `bias` optional; kernel_name empty -> ASM heuristic. +void gemm_a4w4_asm(at::Tensor A, at::Tensor B, at::Tensor a_scale, at::Tensor b_scale, + at::Tensor out, std::optional bias, std::string kernel_name, + double alpha, double beta, bool bpreshuffle, int64_t log2_k_split) { + NVTEAiterGemmTensor a = make_aiter_gemm_tensor(A, kNVTEAiterGemmFP4x2); + NVTEAiterGemmTensor b = make_aiter_gemm_tensor(B, kNVTEAiterGemmFP4x2); + NVTEAiterGemmTensor as = make_aiter_gemm_tensor(a_scale, kNVTEAiterGemmE8M0); + NVTEAiterGemmTensor bs = make_aiter_gemm_tensor(b_scale, kNVTEAiterGemmE8M0); + NVTEAiterGemmTensor o = make_aiter_gemm_tensor(out, aiter_gemm_out_dtype(out)); + NVTEAiterGemmTensor bias_t{}; + NVTEAiterGemmTensor *bias_ptr = nullptr; + if (bias.has_value() && bias->defined()) { + bias_t = make_aiter_gemm_tensor(*bias, kNVTEAiterGemmFP32); + bias_ptr = &bias_t; + } + int rc; + NVTE_SCOPED_GIL_RELEASE({ + rc = nvte_aiter_gemm_a4w4_asm(&a, &b, &as, &bs, &o, bias_ptr, kernel_name.c_str(), + static_cast(alpha), static_cast(beta), + bpreshuffle ? 1 : 0, static_cast(log2_k_split), + at::cuda::getCurrentCUDAStream()); + }); + NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_asm failed (rc=", rc, ")"); +} + +#endif // USE_ROCM + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index cc90ff396..0747eb774 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -182,6 +182,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("comm_overlap") = nullptr, py::arg("comm_type") = std::nullopt, py::arg("extra_output") = std::nullopt, py::arg("bulk_overlap") = false, py::arg("alpha") = 1.0f, py::arg("beta") = std::nullopt); +#ifdef USE_ROCM + m.def("gemm_a4w4_blockscale", &transformer_engine::pytorch::gemm_a4w4_blockscale, + "AITER a4w4 CK blockscale GEMM", py::arg("XQ"), py::arg("WQ"), py::arg("x_scale"), + py::arg("w_scale"), py::arg("Y"), py::arg("split_k") = 0, py::arg("kernel_name") = ""); + m.def("gemm_a4w4_asm", &transformer_engine::pytorch::gemm_a4w4_asm, "AITER a4w4 ASM (f4gemm) GEMM", + py::arg("A"), py::arg("B"), py::arg("a_scale"), py::arg("b_scale"), py::arg("out"), + py::arg("bias") = std::nullopt, py::arg("kernel_name") = "", py::arg("alpha") = 1.0, + py::arg("beta") = 0.0, py::arg("bpreshuffle") = true, py::arg("log2_k_split") = 0); +#endif // USE_ROCM /* GLU (sigmoid gate) */ m.def("glu", transformer_engine::pytorch::glu, "GLU activation", py::arg("input"), py::arg("quantizer"));