From 10a17bdc30a8ffbcfc422cb47ce1374c3f9c45d8 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Tue, 7 Jul 2026 19:51:25 +0000 Subject: [PATCH 1/5] 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")); From e2f1bea463ef2b0166020620550b62323234619c Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Tue, 28 Jul 2026 17:50:36 +0000 Subject: [PATCH 2/5] Updated to qola integration branch w/ main merge --- 3rdparty/QoLA | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/QoLA b/3rdparty/QoLA index bbcf79e61..99026d4b8 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit bbcf79e610c69b4d84615c0bb51abfee7bad1a8d +Subproject commit 99026d4b8b682235f4d442a986bd60a6a6be479c From 9bc7ad5e5c144ecdff284007ddb4f650fea48952 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Wed, 29 Jul 2026 16:16:58 +0000 Subject: [PATCH 3/5] Refactor w/ QoLA update --- .github/scripts/aiter_prebuild_upload.sh | 7 +- 3rdparty/QoLA | 2 +- CLAUDE.md | 2 +- transformer_engine/common/CMakeLists.txt | 39 +++++- .../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/CMakeLists.txt | 12 +- .../common/ck_fused_attn/qola_manifest.toml | 19 --- .../common/gemm/aiter_a4w4_gemm.cpp | 120 ++++++++++------- .../include/transformer_engine/aiter_gemm.h | 30 +++-- transformer_engine/common/qola_manifest.toml | 47 +++++++ .../pytorch/cpp_extensions/gemm.py | 82 +++++++++-- .../pytorch/csrc/extensions/gemm.cpp | 21 ++- 15 files changed, 283 insertions(+), 451 deletions(-) delete mode 100644 transformer_engine/common/aiter_gemm/CMakeLists.txt delete mode 100644 transformer_engine/common/aiter_gemm/include/aiter_gemm/aiter_gemm.hpp delete mode 100644 transformer_engine/common/aiter_gemm/qola_manifest.toml delete mode 100644 transformer_engine/common/aiter_gemm/src/aiter_gemm_a4w4.cpp delete mode 100644 transformer_engine/common/ck_fused_attn/qola_manifest.toml create mode 100644 transformer_engine/common/qola_manifest.toml 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 99026d4b8..7502c6e65 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit 99026d4b8b682235f4d442a986bd60a6a6be479c +Subproject commit 7502c6e653db66f1da5ee7014605bb6616074ba9 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/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index daa44f5dc..b9c3e296e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -713,10 +713,24 @@ else() # USE_ROCM # 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) - add_subdirectory(aiter_gemm ${CMAKE_CURRENT_BINARY_DIR}/aiter_gemm) + 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) @@ -744,7 +758,23 @@ else() # USE_ROCM endif() if(__BUILD_AITER_GEMM) target_compile_definitions(transformer_engine PUBLIC USE_AITER_GEMM) - list(APPEND transformer_engine_LINKER_LIBS 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() @@ -870,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/aiter_gemm/CMakeLists.txt b/transformer_engine/common/aiter_gemm/CMakeLists.txt deleted file mode 100644 index 8ad5eb8f2..000000000 --- a/transformer_engine/common/aiter_gemm/CMakeLists.txt +++ /dev/null @@ -1,127 +0,0 @@ -# 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 deleted file mode 100644 index 598b11477..000000000 --- a/transformer_engine/common/aiter_gemm/include/aiter_gemm/aiter_gemm.hpp +++ /dev/null @@ -1,82 +0,0 @@ -/************************************************************************* - * 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 deleted file mode 100644 index 7a6d5e4b6..000000000 --- a/transformer_engine/common/aiter_gemm/qola_manifest.toml +++ /dev/null @@ -1,25 +0,0 @@ -# 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 deleted file mode 100644 index 049d9fb8b..000000000 --- a/transformer_engine/common/aiter_gemm/src/aiter_gemm_a4w4.cpp +++ /dev/null @@ -1,119 +0,0 @@ -/************************************************************************* - * 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/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 746bdbbf0..000000000 --- a/transformer_engine/common/ck_fused_attn/qola_manifest.toml +++ /dev/null @@ -1,19 +0,0 @@ -[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" -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 index a08530d07..e88f3456e 100644 --- a/transformer_engine/common/gemm/aiter_a4w4_gemm.cpp +++ b/transformer_engine/common/gemm/aiter_a4w4_gemm.cpp @@ -4,11 +4,13 @@ * 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. +// 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" @@ -16,21 +18,42 @@ #include -#include "aiter_gemm/aiter_gemm.hpp" +#include + +#include "qola_gemm_a4w4.h" 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; +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 @@ -40,15 +63,11 @@ extern "C" int nvte_aiter_gemm_a4w4_blockscale(const NVTEAiterGemmTensor *XQ, 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; + 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, @@ -57,36 +76,45 @@ extern "C" int nvte_aiter_gemm_a4w4_asm(const NVTEAiterGemmTensor *A, const NVTE 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; + 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 -extern "C" int nvte_aiter_gemm_a4w4_blockscale(const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, - const NVTEAiterGemmTensor *, const NVTEAiterGemmTensor *, - const NVTEAiterGemmTensor *, int, const char *, void *) { +#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 *) { + const char *, float, float, int, int, void *, char *err_buf, + size_t err_buf_size) { + report_disabled(err_buf, err_buf_size); return -1; } diff --git a/transformer_engine/common/include/transformer_engine/aiter_gemm.h b/transformer_engine/common/include/transformer_engine/aiter_gemm.h index ca302de47..074f7cc20 100644 --- a/transformer_engine/common/include/transformer_engine/aiter_gemm.h +++ b/transformer_engine/common/include/transformer_engine/aiter_gemm.h @@ -10,19 +10,27 @@ * 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 the internal - * aiter_gemm::DType mapping. */ +/*! \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) */ @@ -33,24 +41,30 @@ typedef enum { kNVTEAiterGemmI8 = 6, } NVTEAiterGemmDType; -/*! \brief Lightweight device-tensor descriptor (raw pointer + layout). */ +/*! \brief Lightweight device-tensor descriptor (raw pointer + layout). + * Layout-identical to QoLA's qola_tensor_t. */ typedef struct { void *ptr; - int ndim; + int32_t ndim; + int32_t dtype; /*!< one of NVTEAiterGemmDType */ + int32_t device_id; + int32_t reserved; 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. + * + * \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); + 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 @@ -60,7 +74,7 @@ int nvte_aiter_gemm_a4w4_asm(const NVTEAiterGemmTensor *A, const NVTEAiterGemmTe 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); + int log2_k_split, void *stream, char *err_buf, size_t err_buf_size); #ifdef __cplusplus } // extern "C" 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 0fd4b1d93..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} " @@ -175,18 +239,16 @@ def _fp4_gemm_core(A_fp4, A_scales, B_fp4, B_scales, out_dtype=torch.bfloat16, (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. + Everything runs through TE's torch-free AITER libs (``tex``); the ``aiter`` + Python package is not needed at runtime. """ - from aiter.ops.shuffle import shuffle_weight - _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] diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 81cda7ef4..b8a41be59 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -880,16 +880,20 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p 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()); - for (int i = 0; i < d.ndim && i < 8; ++i) { + 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()); + d.device_id = static_cast(t.device().index()); return d; } @@ -917,11 +921,13 @@ void gemm_a4w4_blockscale(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at:: 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()); + kernel_name.c_str(), at::cuda::getCurrentCUDAStream(), err, + sizeof(err)); }); - NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_blockscale failed (rc=", rc, ")"); + NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_blockscale failed (rc=", rc, "): ", err); } // ASM (f4gemm) a4w4 GEMM. `bias` optional; kernel_name empty -> ASM heuristic. @@ -940,13 +946,14 @@ void gemm_a4w4_asm(at::Tensor A, at::Tensor B, at::Tensor a_scale, at::Tensor b_ 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()); + at::cuda::getCurrentCUDAStream(), err, sizeof(err)); }); - NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_asm failed (rc=", rc, ")"); + NVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_asm failed (rc=", rc, "): ", err); } #endif // USE_ROCM From c6aab5224f4bf54eb7a3fe974f45e115433840b3 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Wed, 29 Jul 2026 20:36:20 +0000 Subject: [PATCH 4/5] Updated qola --- 3rdparty/QoLA | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/QoLA b/3rdparty/QoLA index 7502c6e65..820cec9e1 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit 7502c6e653db66f1da5ee7014605bb6616074ba9 +Subproject commit 820cec9e14a1527350516a12b858bcdd73f9c364 From 93f938b84036b7d5c629ba3e181ed537abe00e01 Mon Sep 17 00:00:00 2001 From: Meekail Zain Date: Wed, 29 Jul 2026 21:01:03 +0000 Subject: [PATCH 5/5] Added error guard --- 3rdparty/QoLA | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/QoLA b/3rdparty/QoLA index 820cec9e1..a09731bce 160000 --- a/3rdparty/QoLA +++ b/3rdparty/QoLA @@ -1 +1 @@ -Subproject commit 820cec9e14a1527350516a12b858bcdd73f9c364 +Subproject commit a09731bcedb27747b8b8cae625249aa179aa59eb