Repository navigation
Conversation
Micky774
changed the title
Claude WalkthroughIntent. Initial integration of the AITER a4w4 (FP4 × FP4) GEMM into TE, wired through a new torch-free C++ backend instead of the Key changes.
Walkthrough.
Testing. No new tests in this PR. The existing FP4 GEMM code paths in Notes for reviewers.
Generated by Claude. To request a code review, comment |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| NVTEAiterGemmTensor bias_t{}; | ||
| NVTEAiterGemmTensor *bias_ptr = nullptr; | ||
| if (bias.has_value() && bias->defined()) { | ||
| bias_t = make_aiter_gemm_tensor(*bias, kNVTEAiterGemmFP32); |
There was a problem hiding this comment.
The bias dtype is hard-coded to kNVTEAiterGemmFP32 regardless of the actual at::Tensor scalar type. Today mxfp4_gemm only ever calls this with bias=None, so this branch is unreachable — but if a future caller passes a bf16/fp16 bias, AITER's ASM kernel will reinterpret the byte pattern as fp32 and produce silently wrong results.
Two low-cost hedges:
NVTE_CHECK(bias->scalar_type() == at::kFloat, "AITER a4w4 GEMM bias must be fp32");- or derive the dtype from
bias->scalar_type()the same wayaiter_gemm_out_dtypedoes for the output tensor.
Either is fine; leaving the current line silently mis-typing is the concerning shape.
| for (int i = 0; i < d.ndim; ++i) { | ||
| t.shape[i] = d.shape[i]; | ||
| t.strides[i] = d.strides[i]; | ||
| numel *= static_cast<size_t>(d.shape[i]); | ||
| } |
There was a problem hiding this comment.
d.shape / d.strides are fixed-size int64_t[8] in TensorDesc, but this loop trusts the caller-supplied d.ndim without clamping. If d.ndim > 8 (unlikely for a4w4 GEMM but possible via the public C API), we'll OOB-read past the descriptor and write past t.shape / t.strides in aiter_tensor_t. The NVTE C API wrapper (aiter_a4w4_gemm.cpp::to_desc) already clamps to 8; matching that here would make the wrapper library equally safe against a raw C API caller:
| for (int i = 0; i < d.ndim; ++i) { | |
| t.shape[i] = d.shape[i]; | |
| t.strides[i] = d.strides[i]; | |
| numel *= static_cast<size_t>(d.shape[i]); | |
| } | |
| size_t numel = (d.ndim > 0) ? 1 : 0; | |
| const int loop_ndim = d.ndim < 8 ? d.ndim : 8; | |
| for (int i = 0; i < loop_ndim; ++i) { | |
| t.shape[i] = d.shape[i]; | |
| t.strides[i] = d.strides[i]; | |
| numel *= static_cast<size_t>(d.shape[i]); | |
| } |
| 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, | ||
| ) |
There was a problem hiding this comment.
Behavior change worth being deliberate about: the pre-PR call site invoked gemm_a4w4_blockscale(..., splitK=log2_k_split) without kernel_name, so on the CK branch AITER's own heuristic picked the kernel. This PR now forwards kernel_name (the tuned-CSV lookup). Per the wrapper header docstring, "non-empty must exist in the compiled registry" — so if the tuned CSV holds a stale/renamed CK kernel name that used to be silently ignored, it will now fail loudly at runtime.
If that's the intended fix (tuned CSV wasn't actually reaching CK before), fine — please confirm. Otherwise consider a fall-through (empty kernel_name on miss).
Also nit: the local is still called log2_k_split but the value flowing in is a linear splitK from _select_kernel_fp4 (cfg.get("splitK", 0)). The CK function takes split_k (linear); the ASM function takes log2_k_split (log2). Passing the same variable to both only works if _select_kernel_fp4's output is already backend-appropriate — worth a comment either way to make future misuse harder.
| 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() |
There was a problem hiding this comment.
The QoLA build runs via execute_process at CMake configure time. CMAKE_CONFIGURE_DEPENDS on qola_manifest.toml (line 44) covers manifest edits, but a modified AITER source tree (e.g. NVTE_AITER_SOURCE_DIR pointing at a local checkout that the developer is iterating on) will not trigger a rebuild — only re-running cmake .. will. The aiter_gemm shared library also has no add_dependencies link back to the QoLA output, so .sos that get replaced under ${__QOLA_BUILD_DIR}/lib do not cause relinking.
For a first integration this is likely acceptable, but a follow-up switching to ExternalProject_Add (or a custom target with DEPENDS/BYPRODUCTS on the te_libgemm_a4w4_*.sos) would make iterative development on gfx950 far less error-prone. Worth at least a TODO comment.
| #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; | ||
| } | ||
|
|
There was a problem hiding this comment.
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.cpp unconditionally calls NVTE_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 cryptic rc=-1 error rather than a "a4w4 GEMM not built for this arch" message.
Two options that would materially improve the experience with almost zero code:
- use a distinctive negative code (e.g.
-2= "backend not compiled in") plus a matching branch ingemm.cpp's check that produces a targeted error, or - expose a build-state query (
extern "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.
Claude reviewReviewed the initial a4w4 (FP4) GEMM integration + QoLA/AITER version bump. Diff scope: 14 files, +652/-10, entirely ROCm-side (all C++ additions guarded by Verdict: looks close to landable. No blockers; five inline notes on latent correctness/UX rough edges:
Copyright headers: OK (all modified files preserve NVIDIA headers unchanged and carry AMD end-year 2026; new ROCm-only files follow the existing AMD-only "License for AMD contributions = MIT" convention). |
wangye805
left a comment
wangye805
left a comment
There was a problem hiding this comment.
I presume you are integrating aiter_gemm in the same way as ck_fused_attn. But this add one extra layer of wrapping.
The separate libck_fused_attn.so I created was a history problem. There were several incompatible apis/arguments so I hope this ck_fused_attn can be taken by CK team eventually. Actually our qola is exactly what I want in the beginning. And I didn't have time to refactor to get rid of this ck_fused_attn layer.
So I would suggest to get rid of extra layer and integrate a4w4 gemm directly from qola. What do you think
Conflict resolutions: - 3rdparty/QoLA: keep ours (920ecf9); it already contains dev's 100b315. - common/CMakeLists.txt: hoist dev's NVTE_QOLA_DIR override to the parent scope as TE_QOLA_DIR, matching how the shared AITER checkout was hoisted. - ck_fused_attn/CMakeLists.txt: consume TE_QOLA_DIR/TE_AITER_SOURCE_DIR from the parent. Dev's local NVTE_AITER_SOURCE_DIR block is redundant -- qola_checkout_aiter() handles that override. - qola_manifest.toml: dropped the ck_fused_attn copy (moved to common/) and folded in dev's rocm_versions 7.14 and libmha_fwd -DFA_WITH_NATIVE_SPLITKV=1, which pairs with QoLA patches 0008/0010. Kept our aiter pin: dev's a50a62a lacks the a4w4 torch-free sources this branch builds; rebasing that lineage is separate work. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
QoLA patch 0009-fmha-sink-fixes stopped applying to the pinned AITER tree (CK f33252ce refactored the sink_value site), hard-aborting CMake configure. Verified independent of the AITER pin -- the full series fails identically at dev's a50a62a and at our ab55045. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
QoLA rekeyed its cpp_itfs a4w4 recipes from libgemm_a4w4_* to AITER's pybind module_gemm_a4w4_* names; the manifest still asked for the old keys and configure aborted with "No cpp_itfs variant available". AITER's optCompilerConfig.json still carries both key sets, so only the recipe lookup was stale. md_name pins the artifact name back to te_libgemm_a4w4_*.so, which common/CMakeLists.txt links by name. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
# Conflicts: # transformer_engine/common/ck_fused_attn/CMakeLists.txt # transformer_engine/common/ck_fused_attn/qola_manifest.toml
| 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. |
There was a problem hiding this comment.
Why aren't them installed to lib?
There was a problem hiding this comment.
The comment is incorrect, they indeed are. I'll remove the comment
| use_ck = bool(kernel_name) and kernel_name.find("_ZN") == -1 | ||
| if use_ck: | ||
| result = gemm_a4w4_blockscale( | ||
| tex.gemm_a4w4_blockscale( |
There was a problem hiding this comment.
With new abstraction layer added, consider wrap this blockscale/asm selection into TE core code and have single API exposed
There was a problem hiding this comment.
That's essentially what the mxfp4_gemm entry point is meant to be -- do you mean an additional neutral tex.gemm_a4w4 that handles the dispatch internally?
There was a problem hiding this comment.
Yes. Current implementation passes AITER API and QoLA (which I think are copies of AITER) datatypes several levels of abstraction up. It made sense when TE pytorch layer called AITER directly. Now, when everything goes through TE core the latter can be dispatch point.
There was a problem hiding this comment.
Done, moved the entry point into a unified tex.gemm_a4w4 and updated the backend dispatch.
| list(APPEND __QOLA_GROUPS ck_fused_attn) | ||
| endif() | ||
| if(__BUILD_AITER_GEMM) | ||
| qola_prebuilt_dir(aiter_gemm __AITER_GEMM_PREBUILT) |
There was a problem hiding this comment.
What about flag to skip AITER checkout/patching?
There was a problem hiding this comment.
You can use QOLA_PREBUILT_DIR_<GROUP>=/path/to/prebuit so e.g. QOLA_PREBUILT_DIR_AITER_GEMM
There was a problem hiding this comment.
Prebuilt dir has binaries (and headers). I'm talking about AITER source dir that has to be only built but not patched. CK-JIT build pipeline is (AITER checkout + QoLA patch -> CK-JIT patch + enable interception -> QoLA build -> CK-JIT patches revert)
There was a problem hiding this comment.
In that case I think QOLA_AITER_SOURCE_DIR=/path/to/aiter is sufficient for that case.
…ation # Conflicts: # 3rdparty/QoLA
QoLA patch 0001 rebound blob_gen_cmd inside build_module's nested MainFunc without declaring it nonlocal, so every non-cached AITER build failed with "cannot access local variable 'blob_gen_cmd'". Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
matthiasdiener
left a comment
matthiasdiener
left a comment
There was a problem hiding this comment.
Just a nit; I'd also suggest updating the PR title to mention Aiter.
Co-authored-by: Matthias Diener <matthias.diener@amd.com>
Micky774
changed the title
…ation # Conflicts: # transformer_engine/common/ck_fused_attn/CMakeLists.txt
…ation # Conflicts: # .github/scripts/aiter_prebuild_upload.sh # 3rdparty/QoLA # 3rdparty/ck_jit # transformer_engine/common/ck_fused_attn/CMakeLists.txt # transformer_engine/common/ck_fused_attn/qola_manifest.toml
| fmha_args.p_undrop = 1.0 - args.dropout_probability; | ||
| fmha_args.drop_seed_offset = std::pair<const void*, const void*>{args.philox_seed_ptr, args.philox_offset_ptr}; | ||
| // gfx1250 deterministic only: number of dq_acc planes; 0 lets aiter pick from the shape. | ||
| fmha_args.det_split = 0; |
There was a problem hiding this comment.
It duplicates the very next assginment
| [qola] | ||
| aiter_commit = "16979d75774f00e5e8905bcf5f2216230093e441" # AITER main; includes the a4w4 torch-free work (#4124) | ||
| namespace = "te" | ||
| rocm_versions = ["7.14"] |
There was a problem hiding this comment.
Is it still needed? The value does not reflect actual ROCm version
There was a problem hiding this comment.
No, this is a vestige of the original design pattern of having pre-built images provided via QoLA CI. I've removed it from the TE manifest, and will remove it from QoLA's example soon.
| rocm_versions = ["7.14"] | ||
|
|
||
| [build] | ||
| architectures = ["gfx950", "gfx942", "gfx1250"] |
There was a problem hiding this comment.
what about gfx1250-strict? TE dev should support building for it that requires AITER patches. Has it been validated with aiter_gemm module?
There was a problem hiding this comment.
The AITER a4w4 is supported only on gfx950 and skips/stubs otherwise.
| add_subdirectory(gemm/kittens ${CMAKE_CURRENT_BINARY_DIR}/kittens) | ||
| endif() | ||
|
|
||
| # There is no TE-side wrapper library any more: QoLA exports a C ABI that |
There was a problem hiding this comment.
This is rather change comment than code comment
There was a problem hiding this comment.
Can QoLA PR be merged so all QoLA main fixes are in place?
There was a problem hiding this comment.
Yes that'll be done last after this PR is finalized, since this PR's CI is used to validate the QoLA PR and any changes that might crop up during review.
- qola_manifest.toml: drop unused rocm_versions; add gfx1250-strict arch - ck_fused_attn_bwd.cpp: remove duplicate det_split assignment - CMakeLists.txt: drop stale comment about TE-side wrapper library Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Description
Companion and integration point to ROCm/aiter#4124 and ROCm/QoLA#10
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: