-
Notifications
You must be signed in to change notification settings - Fork 34
Initial integration of a4w4 GEMM #663
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Draft
Micky774
wants to merge
6
commits into
dev
Choose a base branch
from
zain/aiter/a4w4-integration
base: dev
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
Changes from all commits
Commits
Show all changes
6 commits
Select commit
Hold shift + click to select a range
10a17bd
Initial integration of a4w4 GEMM and QoLA version bump
Micky774 b44f1d1
Merge branch 'dev' into zain/aiter/a4w4-integration
Micky774 e2f1bea
Updated to qola integration branch w/ main merge
Micky774 9bc7ad5
Refactor w/ QoLA update
Micky774 c6aab52
Updated qola
Micky774 93f938b
Added error guard
Micky774 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Submodule QoLA
updated
16 files
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
19 changes: 0 additions & 19 deletions
19
transformer_engine/common/ck_fused_attn/qola_manifest.toml
This file was deleted.
Oops, something went wrong.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 <hip/hip_runtime.h> | ||
|
|
||
| #include <cstddef> | ||
|
|
||
| #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<int>(kNVTEAiterGemmFP4x2) == static_cast<int>(QOLA_DTYPE_FP4X2), | ||
| "a4w4 dtype enum drift"); | ||
| static_assert(static_cast<int>(kNVTEAiterGemmE8M0) == static_cast<int>(QOLA_DTYPE_E8M0), | ||
| "a4w4 dtype enum drift"); | ||
| static_assert(static_cast<int>(kNVTEAiterGemmBF16) == static_cast<int>(QOLA_DTYPE_BF16), | ||
| "a4w4 dtype enum drift"); | ||
| static_assert(static_cast<int>(kNVTEAiterGemmFP16) == static_cast<int>(QOLA_DTYPE_FP16), | ||
| "a4w4 dtype enum drift"); | ||
| static_assert(static_cast<int>(kNVTEAiterGemmFP32) == static_cast<int>(QOLA_DTYPE_FP32), | ||
| "a4w4 dtype enum drift"); | ||
|
|
||
| inline const qola_tensor_t *as_qola(const NVTEAiterGemmTensor *t) { | ||
| return reinterpret_cast<const qola_tensor_t *>(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<hipStream_t>(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<hipStream_t>(stream), err_buf, | ||
| err_buf_size); | ||
| } | ||
|
|
||
| #else // !USE_AITER_GEMM | ||
|
|
||
| #include <cstring> | ||
|
|
||
| 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 | ||
83 changes: 83 additions & 0 deletions
83
transformer_engine/common/include/transformer_engine/aiter_gemm.h
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 <stddef.h> | ||
| #include <stdint.h> | ||
|
|
||
| #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_ |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The "framework layer degrades gracefully" comment at the top of the file (line 11) is a bit optimistic given how these stubs are consumed on the pytorch side:
gemm.cppunconditionally callsNVTE_CHECK(rc == 0, "nvte_aiter_gemm_a4w4_blockscale failed (rc=", rc, ")")and throws. A user building on gfx942 who accidentally exercises the a4w4 path will see a rather crypticrc=-1error rather than a "a4w4 GEMM not built for this arch" message.Two options that would materially improve the experience with almost zero code:
-2= "backend not compiled in") plus a matching branch ingemm.cpp's check that produces a targeted error, orextern "C" int nvte_aiter_gemm_a4w4_available();) so the pytorch layer can guard the entry points and give the framework ahasattr-like probe from Python.Not a blocker; landing this as-is is fine so long as the intent is understood.