diff --git a/.github/scripts/aiter_prebuild_upload.sh b/.github/scripts/aiter_prebuild_upload.sh index 473ef1c75..9c1dbb50c 100755 --- a/.github/scripts/aiter_prebuild_upload.sh +++ b/.github/scripts/aiter_prebuild_upload.sh @@ -24,7 +24,11 @@ ROCM_VER=`head -n1 "${ROCM_PATH}/.info/version" | cut -d. -f1` QOLA_DIR="${ROOT_DIR}/3rdparty/QoLA" AITER_DIR="${QOLA_DIR}/3rdparty/aiter" -QOLA_MANIFEST="${ROOT_DIR}/transformer_engine/common/ck_fused_attn/qola_manifest.toml" +QOLA_MANIFEST="${ROOT_DIR}/transformer_engine/common/qola_manifest.toml" +# This cache feeds aiter_prebuilt.cmake, which only consumes the fused-attn +# kernels. The shared manifest also declares the gfx950-only a4w4 GEMM group, +# so the build must be restricted or the gfx942 leg would fail. +QOLA_GROUP="ck_fused_attn" GIT_CONFIG_GLOBAL="$(mktemp /tmp/gitconfig.XXXXXX)" trap 'rm -f "${GIT_CONFIG_GLOBAL}"' EXIT git config --file "${GIT_CONFIG_GLOBAL}" --add safe.directory "${AITER_DIR}" @@ -64,6 +68,7 @@ if [[ "${1:-}" == "--build" ]]; then --manifest "${QOLA_MANIFEST}" \ --aiter-root "${AITER_DIR}" \ --output-dir "${QOLA_BUILD_DIR}" \ + --group "${QOLA_GROUP}" \ "${arch_args[@]}" # Stage QoLA outputs into the cache layout expected by aiter_prebuilt.cmake. diff --git a/3rdparty/QoLA b/3rdparty/QoLA index bbcf79e61..a09731bce 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit bbcf79e610c69b4d84615c0bb51abfee7bad1a8d +Subproject commit a09731bcedb27747b8b8cae625249aa179aa59eb diff --git a/CLAUDE.md b/CLAUDE.md index b040f0a5e..89e2fa2bd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -13,7 +13,7 @@ - **Python import flow**: framework selection happens in `transformer_engine/__init__.py` (driven by `NVTE_FRAMEWORK`); native `.so` resolution happens in `transformer_engine/common/__init__.py`. Start tracing from these two files for any import or framework-selection issue — function names inside may change, but the entry points are stable. - **Build orchestration**: `setup.py` + helpers in `build_tools/` + CMake. ROCm vs CUDA backend detection lives in `build_tools/utils.py` — grep there (and for `NVTE_USE_ROCM`) for current behavior. - **3rdparty submodules**: read `.gitmodules` for the current set and commit pins to determine current submodules. -- **AITER source**: the tree QoLA/TE actually builds is cloned on-demand to `build/cmake/ck_fused_attn/qola/third_party/aiter/` at the commit pinned in `transformer_engine/common/ck_fused_attn/qola_manifest.toml`. The in-tree `3rdparty/*/aiter` checkouts are stale — read the built tree (and check its HEAD) when reasoning about the linked ABI. +- **AITER source**: the tree QoLA/TE actually builds is cloned on-demand to `build/cmake/ck_fused_attn/qola/third_party/aiter/` at the commit pinned in `transformer_engine/common/qola_manifest.toml`. That single manifest serves every QoLA consumer; each one builds its own `[[modules]]` subset via `qola build --group ` (`ck_fused_attn`, `aiter_gemm`). The in-tree `3rdparty/*/aiter` checkouts are stale — read the built tree (and check its HEAD) when reasoning about the linked ABI. ## Hipify convention The build auto-generates HIP files from CUDA sources via `hipify_torch`. Generated files are marked with `// !!! This is a file automatically generated by hipify!!!` at line 1. **Never edit generated files directly** — edit the CUDA source instead. diff --git a/setup.py b/setup.py index 62482d12b..75d13c1c4 100644 --- a/setup.py +++ b/setup.py @@ -97,6 +97,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 b4631c6d3..b9c3e296e 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 blockwise FP8 and MXFP8 GEMM kernels" ON) +option(USE_AITER_GEMM "Use AITER a4w4 (FP4) GEMM backend" ON) set(USE_CUDA OFF) if (USE_ROCM) @@ -384,6 +385,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 @@ -708,6 +710,34 @@ 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. + # + # There is no TE-side wrapper library any more: QoLA exports a C ABI that + # libtransformer_engine.so links against directly, so all this needs is the + # include dir and the built .so files. + 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) + include(${CMAKE_CURRENT_LIST_DIR}/../../3rdparty/QoLA/cmake/QoLA.cmake) + qola_add_modules( + GROUP aiter_gemm + MANIFEST ${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml + BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/aiter_gemm/qola + ARCHS gfx950 + LIBS te_libgemm_a4w4_blockscale.so te_libgemm_a4w4_asm.so + OUT_INCLUDE_DIR __AITER_GEMM_INCLUDE_DIR + OUT_LIB_DIR __AITER_GEMM_LIB_DIR + OUT_LIBS __AITER_GEMM_LIBS + OUT_CONFIG_DIR __AITER_GEMM_CONFIG_DIR) + 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) @@ -726,6 +756,26 @@ 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) + target_include_directories(transformer_engine PRIVATE ${__AITER_GEMM_INCLUDE_DIR}) + target_link_directories(transformer_engine PUBLIC ${__AITER_GEMM_LIB_DIR}) + foreach(__lib ${__AITER_GEMM_LIBS}) + list(APPEND transformer_engine_LINKER_LIBS -l:${__lib}) + endforeach() + set(__AITER_GEMM_INSTALL_DIR "${CMAKE_INSTALL_PREFIX}/transformer_engine/lib") + if(NOT "${__AITER_GEMM_LIB_DIR}" STREQUAL "${__AITER_GEMM_INSTALL_DIR}") + foreach(__lib ${__AITER_GEMM_LIBS}) + install(FILES ${__AITER_GEMM_LIB_DIR}/${__lib} DESTINATION ${__AITER_GEMM_INSTALL_DIR}) + endforeach() + endif() + # Ship the a4w4 tuning table next to the kernels it describes; the Python + # layer reads it instead of importing the `aiter` package. + if(EXISTS "${__AITER_GEMM_CONFIG_DIR}/a4w4_blockscale_tuned_gemm.csv") + install(FILES ${__AITER_GEMM_CONFIG_DIR}/a4w4_blockscale_tuned_gemm.csv + DESTINATION ${__AITER_GEMM_INSTALL_DIR}/configs) + endif() + endif() target_link_libraries(transformer_engine PUBLIC ${transformer_engine_LINKER_LIBS}) endif() @@ -850,7 +900,10 @@ endif() # Install library install(TARGETS transformer_engine DESTINATION .) if (USE_ROCM) - set_target_properties(transformer_engine PROPERTIES INSTALL_RPATH "$ORIGIN/lib;$ORIGIN/transformer_engine/lib") + # "$ORIGIN" covers QoLA-built kernel libraries (e.g. te_libgemm_a4w4_*.so), + # which are installed alongside libtransformer_engine.so rather than under a + # nested lib/ directory. + set_target_properties(transformer_engine PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/lib;$ORIGIN/transformer_engine/lib") file(READ "${ROCM_PATH}/.info/version" ROCM_VER) string(STRIP "${ROCM_VER}" ROCM_VER) string(REGEX MATCH "^[0-9]+\\.[0-9]+" ROCM_VER "${ROCM_VER}") diff --git a/transformer_engine/common/ck_fused_attn/CMakeLists.txt b/transformer_engine/common/ck_fused_attn/CMakeLists.txt index 96b3e7b2f..32b0809a2 100644 --- a/transformer_engine/common/ck_fused_attn/CMakeLists.txt +++ b/transformer_engine/common/ck_fused_attn/CMakeLists.txt @@ -48,7 +48,11 @@ endif() # QoLA-managed AITER source tree to that commit before any consumer reads it # (header validation below, header includes for the .cpp build later, and # QoLA's own kernel build if the prebuilt cache misses). -set(__QOLA_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml") +# TE keeps a single QoLA manifest so the whole build shares one AITER commit, +# one patch set and one checkout. Each consumer builds only its own module +# group (see `qola build --group`). +set(__QOLA_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/../qola_manifest.toml") +set(__QOLA_GROUP "ck_fused_attn") set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${__QOLA_MANIFEST}") file(STRINGS "${__QOLA_MANIFEST}" __AITER_COMMIT_LINES REGEX "^[ \t]*aiter_commit[ \t]*=[ \t]*\"[^\"]+\"") @@ -155,8 +159,11 @@ if(__AITER_CACHE_DIR STREQUAL "") else() set(__CK_JIT_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/ck_jit") endif() + # ck_jit_build.py shells out to `qola build` with a fixed argument list and + # cannot forward --group, so the group is passed through the environment. execute_process( - COMMAND ${Python_EXECUTABLE} "${__CK_JIT_SOURCE_DIR}/ck_jit_build.py" full + COMMAND ${CMAKE_COMMAND} -E env "QOLA_BUILD_GROUPS=${__QOLA_GROUP}" + ${Python_EXECUTABLE} "${__CK_JIT_SOURCE_DIR}/ck_jit_build.py" full --with-qola --qola-dir ${__QOLA_DIR} --qola-manifest ${__QOLA_MANIFEST} @@ -182,6 +189,7 @@ if(__AITER_CACHE_DIR STREQUAL "") --manifest ${__QOLA_MANIFEST} --aiter-root ${__AITER_SOURCE_DIR} --output-dir ${__QOLA_BUILD_DIR} + --group ${__QOLA_GROUP} --arch "${GPU_ARCHS_STR}" --skip-checkout RESULT_VARIABLE QOLA_BUILD_RESULT diff --git a/transformer_engine/common/ck_fused_attn/qola_manifest.toml b/transformer_engine/common/ck_fused_attn/qola_manifest.toml deleted file mode 100644 index 237f03ad4..000000000 --- a/transformer_engine/common/ck_fused_attn/qola_manifest.toml +++ /dev/null @@ -1,19 +0,0 @@ -[qola] -aiter_commit = "e95f7d4e3924e848c1daf4df449e27822b332a09" # pinned AITER submodule commit -namespace = "te" -rocm_versions = ["7.2"] - -[build] -architectures = ["gfx950", "gfx942"] - -[[modules]] -name = "libmha_fwd" -mode = "cpp_itfs" -receipt = 700 -drop_srcs = ["mha_fwd_split.cu", "mha_fwd_batch_prefill.cu"] -drop_directions = ["fwd_splitkv", "batch_prefill"] - -[[modules]] -name = "libmha_bwd" -mode = "cpp_itfs" -receipt = 700 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..e88f3456e --- /dev/null +++ b/transformer_engine/common/gemm/aiter_a4w4_gemm.cpp @@ -0,0 +1,121 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVTE C API for AITER's a4w4 (FP4) GEMM. QoLA exports a C ABI over a +// descriptor that is layout-identical to NVTEAiterGemmTensor, so these are +// pure forwards -- the static_asserts below are what keep that true. +// +// 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 + +#include "qola_gemm_a4w4.h" + +namespace { + +static_assert(sizeof(NVTEAiterGemmTensor) == sizeof(qola_tensor_t), + "NVTEAiterGemmTensor must mirror qola_tensor_t"); +static_assert(offsetof(NVTEAiterGemmTensor, ptr) == offsetof(qola_tensor_t, ptr), + "NVTEAiterGemmTensor::ptr must mirror qola_tensor_t::ptr"); +static_assert(offsetof(NVTEAiterGemmTensor, ndim) == offsetof(qola_tensor_t, ndim), + "NVTEAiterGemmTensor::ndim must mirror qola_tensor_t::ndim"); +static_assert(offsetof(NVTEAiterGemmTensor, dtype) == offsetof(qola_tensor_t, dtype), + "NVTEAiterGemmTensor::dtype must mirror qola_tensor_t::dtype"); +static_assert(offsetof(NVTEAiterGemmTensor, device_id) == offsetof(qola_tensor_t, device_id), + "NVTEAiterGemmTensor::device_id must mirror qola_tensor_t::device_id"); +static_assert(offsetof(NVTEAiterGemmTensor, shape) == offsetof(qola_tensor_t, shape), + "NVTEAiterGemmTensor::shape must mirror qola_tensor_t::shape"); +static_assert(offsetof(NVTEAiterGemmTensor, strides) == offsetof(qola_tensor_t, strides), + "NVTEAiterGemmTensor::strides must mirror qola_tensor_t::strides"); + +// The dtype enumerators are part of the shared ABI too. Compared as ints +// because the two enums are deliberately distinct types. +static_assert(static_cast(kNVTEAiterGemmFP4x2) == static_cast(QOLA_DTYPE_FP4X2), + "a4w4 dtype enum drift"); +static_assert(static_cast(kNVTEAiterGemmE8M0) == static_cast(QOLA_DTYPE_E8M0), + "a4w4 dtype enum drift"); +static_assert(static_cast(kNVTEAiterGemmBF16) == static_cast(QOLA_DTYPE_BF16), + "a4w4 dtype enum drift"); +static_assert(static_cast(kNVTEAiterGemmFP16) == static_cast(QOLA_DTYPE_FP16), + "a4w4 dtype enum drift"); +static_assert(static_cast(kNVTEAiterGemmFP32) == static_cast(QOLA_DTYPE_FP32), + "a4w4 dtype enum drift"); + +inline const qola_tensor_t *as_qola(const NVTEAiterGemmTensor *t) { + return reinterpret_cast(t); +} + +} // 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, char *err_buf, + size_t err_buf_size) { + return QOLA_C(gemm_a4w4_blockscale)(as_qola(XQ), as_qola(WQ), as_qola(x_scale), as_qola(w_scale), + as_qola(Y), split_k, kernel_name, + static_cast(stream), err_buf, err_buf_size); +} + +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, char *err_buf, size_t err_buf_size) { + return QOLA_C(gemm_a4w4_asm)(as_qola(A), as_qola(B), as_qola(a_scale), as_qola(b_scale), + as_qola(out), as_qola(bias), kernel_name, alpha, beta, bpreshuffle, + log2_k_split, static_cast(stream), err_buf, + err_buf_size); +} + +#else // !USE_AITER_GEMM + +#include + +namespace { + +void report_disabled(char *err_buf, size_t err_buf_size) { + const char *msg = "TransformerEngine was built without the AITER a4w4 GEMM backend"; + if (err_buf != nullptr && err_buf_size > 0) { + std::strncpy(err_buf, msg, err_buf_size - 1); + err_buf[err_buf_size - 1] = '\0'; + } +} + +} // namespace + +extern "C" int nvte_aiter_gemm_a4w4_blockscale(const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, + const NVTEAiterGemmTensor *, int, const char *, + void *, char *err_buf, size_t err_buf_size) { + report_disabled(err_buf, err_buf_size); + 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 *, char *err_buf, + size_t err_buf_size) { + report_disabled(err_buf, err_buf_size); + 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..074f7cc20 --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/aiter_gemm.h @@ -0,0 +1,83 @@ +/************************************************************************* + * 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. + * + * The kernels themselves live in QoLA-built shared objects. This header + * mirrors QoLA's public C ABI so the two descriptors are layout-identical + * and the call needs no translation; the mirroring is asserted at compile + * time in the implementation. It is declared independently (rather than + * including QoLA's header) so that this API stays available on builds where + * the AITER a4w4 backend is disabled. + */ + +#ifndef TRANSFORMER_ENGINE_AITER_GEMM_H_ +#define TRANSFORMER_ENGINE_AITER_GEMM_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! \brief Element type of an a4w4 GEMM operand. Values match QoLA's + * qola_dtype_t. */ +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). + * Layout-identical to QoLA's qola_tensor_t. */ +typedef struct { + void *ptr; + int32_t ndim; + int32_t dtype; /*!< one of NVTEAiterGemmDType */ + int32_t device_id; + int32_t reserved; + int64_t shape[8]; + int64_t strides[8]; +} NVTEAiterGemmTensor; + +/*! \brief CK blockscale a4w4 GEMM: Y = XQ @ WQ^T with per-1x32 microscaling. + * + * \param[out] err_buf Optional buffer receiving a failure message. + * \param[in] err_buf_size Size of \p err_buf in bytes; 0 to discard. + * \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, + char *err_buf, size_t err_buf_size); + +/*! \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, char *err_buf, size_t err_buf_size); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TRANSFORMER_ENGINE_AITER_GEMM_H_ diff --git a/transformer_engine/common/qola_manifest.toml b/transformer_engine/common/qola_manifest.toml new file mode 100644 index 000000000..59b618fa9 --- /dev/null +++ b/transformer_engine/common/qola_manifest.toml @@ -0,0 +1,47 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Single QoLA consumer manifest for TransformerEngine. +# +# One manifest means one AITER commit, one patch set and one source checkout +# for the whole build. Consumers select the kernels they need with +# `qola build --group `; the arch list is passed per group from CMake, +# which is why [build].architectures below is the union of what any group +# supports rather than a per-group value. +# +# Groups: +# ck_fused_attn - torch-free MHA fwd/bwd kernels behind the CK fused-attn +# backend (transformer_engine/common/ck_fused_attn). +# aiter_gemm - torch-free a4w4 (FP4) GEMM kernels, gfx950 only. + +[qola] +aiter_commit = "ab55045555071eb951b1bf50e29cdd3d194e1e51" # pinned AITER submodule commit -- zain/torch-free/a4w4-poc +namespace = "te" +rocm_versions = ["7.2"] + +[build] +architectures = ["gfx950", "gfx942"] + +[[modules]] +name = "libmha_fwd" +group = "ck_fused_attn" +mode = "cpp_itfs" +receipt = 700 +drop_srcs = ["mha_fwd_split.cu", "mha_fwd_batch_prefill.cu"] +drop_directions = ["fwd_splitkv", "batch_prefill"] + +[[modules]] +name = "libmha_bwd" +group = "ck_fused_attn" +mode = "cpp_itfs" +receipt = 700 + +[[modules]] +name = "libgemm_a4w4_blockscale" +group = "aiter_gemm" +mode = "cpp_itfs" + +[[modules]] +name = "libgemm_a4w4_asm" +group = "aiter_gemm" +mode = "cpp_itfs" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 6a1e84dfe..e97102c1c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -146,18 +146,82 @@ def _should_use_bf16_output_for_nvfp4_tn( ) +def _shuffle_weight_16x16(x: torch.Tensor) -> torch.Tensor: + """Pre-shuffle an FP4 weight tensor into AITER's (16, 16) MFMA layout. + + Mirrors ``aiter.ops.shuffle.shuffle_weight(x, layout=(16, 16))`` for the + non-interleaved case that the a4w4 GEMM uses. It is a pure view/permute of + the packed bytes, so it is reproduced here rather than pulled from the + ``aiter`` package: QoLA's kernel libraries are torch-free by construction + and cannot host a torch-level helper. + """ + x_type = x.dtype + if x_type == torch.float4_e2m1fn_x2: + x = x.view(torch.uint8) + + IN, IK = 16, 16 + BN = IN + BK = IK * 2 + K = 16 // x.element_size() + assert x.shape[-2] % BN == 0, f"{x.shape[-2]} % {BN} != 0" + assert x.shape[-1] % BK == 0, f"{x.shape[-1]} % {BK} != 0" + + x_ = x.view(-1, x.shape[-2] // BN, BN, x.shape[-1] // BK, BK // K, K) + x_ = x_.permute(0, 1, 3, 4, 2, 5) + x_ = x_.contiguous() + x_ = x_.view(*x.shape) + return x_.view(x_type) + + +@functools.lru_cache(maxsize=1) +def _fp4_tuned_gemm_table(): + """Load the a4w4 tuned-GEMM table staged next to the QoLA kernel libs. + + QoLA copies AITER's tuning CSVs into its build output and TE installs them + alongside the a4w4 shared objects, so the table is versioned with the + kernels it describes. Reading it here keeps the ``aiter`` Python package + out of the runtime dependency set. + + Returns a ``{(M, N, K): (kernel_name, split_k)}`` mapping, empty when the + table is absent (callers then fall back to the AITER heuristic). + """ + import csv + + csv_path = os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname(__file__))), + "lib", + "configs", + "a4w4_blockscale_tuned_gemm.csv", + ) + table = {} + if not os.path.exists(csv_path): + return table + with open(csv_path, newline="", encoding="utf-8") as f: + for row in csv.DictReader(f): + try: + key = (int(row["M"]), int(row["N"]), int(row["K"])) + except (KeyError, ValueError): + continue + kernel_name = (row.get("kernelName") or "").strip() + if not kernel_name: + continue + try: + split_k = int(row.get("splitK") or 0) + except ValueError: + split_k = 0 + table[key] = (kernel_name, split_k) + return table + + def _select_kernel_fp4(layout: str, grad: bool, M: int, N: int, K: int): """Select kernel via tuned CSV lookup, falling back to AITER heuristic.""" - from aiter.ops.gemm_op_a4w4 import get_GEMM_config - kernel_name = "" split_k = 0 if _FP4_USE_TUNED_GEMM: - cfg = get_GEMM_config(M, N, K) + cfg = _fp4_tuned_gemm_table().get((M, N, K)) if cfg is not None: - kernel_name = cfg["kernelName"] - split_k = int(cfg.get("splitK", 0)) + kernel_name, split_k = cfg if _FP4_LOG_SHAPES: print(f"[FP4-GEMM] {layout} grad={grad} M={M} N={N} K={K} " @@ -174,18 +238,17 @@ 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. - """ - import aiter - from aiter.ops.shuffle import shuffle_weight - from aiter.ops.gemm_op_a4w4 import gemm_a4w4_blockscale + Everything runs through TE's torch-free AITER libs (``tex``); the ``aiter`` + Python package is not needed at runtime. + """ _fp4_dtype = torch.float4_e2m1fn_x2 A_fp4 = A_fp4.view(_fp4_dtype) if A_fp4.dtype != _fp4_dtype else A_fp4 B_fp4 = B_fp4.view(_fp4_dtype) if B_fp4.dtype != _fp4_dtype else B_fp4 A_scales_uint8 = A_scales.view(torch.uint8) B_scales_uint8 = B_scales.view(torch.uint8) - B_shuffled = B_fp4 if b_pre_shuffled else shuffle_weight(B_fp4, layout=(16, 16)) + B_shuffled = B_fp4 if b_pre_shuffled else _shuffle_weight_16x16(B_fp4) M = A_fp4.shape[0] N = B_fp4.shape[0] @@ -198,17 +261,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 1ee3b087f..46b1f1cdb 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -188,6 +188,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 6773d19e5..b8a41be59 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 @@ -872,4 +875,87 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p return py::reinterpret_borrow(D); } #endif // !USE_ROCM + +#ifdef USE_ROCM + +namespace { + +// Size of the buffer handed to the a4w4 entry points for failure messages. +constexpr size_t kAiterGemmErrLen = 512; + +NVTEAiterGemmTensor make_aiter_gemm_tensor(const at::Tensor &t, int dtype) { + NVTEAiterGemmTensor d{}; + d.ptr = t.data_ptr(); + d.ndim = static_cast(t.dim()); + NVTE_CHECK(d.ndim <= 8, "AITER a4w4 GEMM operands support at most 8 dimensions"); + for (int i = 0; i < d.ndim; ++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; + char err[kAiterGemmErrLen] = {}; + 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(), err, + sizeof(err)); + }); + NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_blockscale failed (rc=", rc, "): ", err); +} + +// 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; + char err[kAiterGemmErrLen] = {}; + 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(), err, sizeof(err)); + }); + NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_asm failed (rc=", rc, "): ", err); +} + +#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 dd77a2ef3..24723ef82 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -247,6 +247,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"));