Initial integration of a4w4 GEMM#663
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). |
2 participants
Description
Companion and integration point to ROCm/aiter#4124
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: