feat(ck-tile): aquant GEMM TE to dispatcher bridge#9980
Conversation
Add a registry-bypass, single-kernel-per-.so TE->Dispatcher bridge for the
gemm_aquant (A-only quantized) block-scale GEMM operator, mirroring the merged
grouped_gemm_bquant bridge.
- codegen/unified_gemm_aquant_codegen.py: per-kernel .hpp generator using
GemmAQuantPipelineProblem + AQuantGemmPipelineAgBgCrMem/CompV3 and
QuantType::AQuantGrouped (QuantGemmHostArgs, aq_ptr set / bq_ptr null).
- bindings/ctypes/gemm_aquant_ctypes_lib.cpp: direct-launch C API
(dispatcher_run_aquant_gemm) with pk_int4 A permute + AQ shuffle for
APreshuffleQuant; runtime arch check throws on unknown arch (no gfx942 default).
- python/gemm_aquant_utils.py: byte-exact config/runner/build helpers; arch
validated, never silently defaulted.
- codegen_common.py: shared make_aquant_kernel_name / aquant_effective_epilogue.
- CMake: guarded dispatcher_gemm_aquant_lib target.
- scripts/aquant_selftest.py: codegen + name-parity + optional hipcc build runner.
dtype x layout matrix matches Old-TE gemm_aquant_quantgrouped*.cpp exactly:
fp8/bf8/fp8i4/bf8i4 x {rcr,rrr,crr,ccr} decode (16) and {rcr,rrr,crr}
preshufflequant (12). All 28 kernels name-byte-exact and hipcc-compile on gfx950.
Reviewer pass (aquant bridge)Reviewed Parity confirmed: 28 kernels = 4 dtypes {fp8,bf8,fp8i4,bf8i4} x 4 decode layouts {rcr,rrr,crr,ccr} (16) + 4 dtypes x 3 preshufflequant layouts {rcr,rrr,crr} (12). ccr-exclusion for preshufflequant is real (Old-TE line 1052 gates the ccr branch on Blocking
Should-fix
Nit
Arch handling (get_arch + throw, no gfx942 default) is correct in both C++ ( |
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
…des) Reviewer feedback on draft PR #9980: - BLOCKING: derive stride_AQ from the AQ layout on both C++ and Python sides. The ccr layout makes the A-scale tensor column-major, so its leading dim is M (row count), not QK_A. Previously stride_AQ was hardcoded to QK_A, silently mis-indexing the ccr A-scale. Codegen now also emits a global `using AQLayout` so the ctypes lib can key off it. - SHOULD-FIX: build dispatcher_gemm_aquant_lib with a main-include-only include set instead of add_ctypes_library, which pulled in dispatcher/include and the generated_tile_backend.hpp run()/launch() conflict the Python build path deliberately avoids. - SHOULD-FIX: pass warp_tile_k through the default_*_config helpers instead of hardcoding 128 (dead named param). - NITs: static_assert ties the shuffle_aq row-major assumption to the ccr preshufflequant exclusion; CPU tests assert ccr column-major AQ stride, the layout scope, the 28-kernel count, and gfx90a validation.
Round-2 fixesThanks for the review. Pushed as BLOCKING — ccr AQ scale stride.
SHOULD-FIX #1 — CMake include set. SHOULD-FIX #2 — dead NITs. Verification. |
There was a problem hiding this comment.
Pull request overview
This PR adds an AQuant (A-only quantized) GEMM TileEngine → Dispatcher bridge path, enabling codegen → per-kernel ctypes .so build → Python ctypes launch for the gemm_aquant block-scale ops (decode + preshufflequant) without requiring user C++.
Changes:
- Introduces AQuant kernel code generator and shared kernel-name construction to keep codegen ↔ Python name parity byte-exact.
- Adds a dedicated ctypes direct-launch library (
QuantGemmHostArgs, registry-bypass) plus a Python utility layer and a self-test script. - Adds CPU-only unit tests to lock down naming/contracts and codegen JSON projection.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/dispatcher/tests/test_aquant_bridge.py | New CPU-only tests for naming, scope, and codegen-config projection. |
| projects/composablekernel/dispatcher/scripts/aquant_selftest.py | New script to validate name parity and optionally build all default configs. |
| projects/composablekernel/dispatcher/python/gemm_aquant_utils.py | New Python configs + build helpers + ctypes wrapper + runner for AQuant GEMM. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py | New AQuant header generator producing single-kernel SelectedKernel headers. |
| projects/composablekernel/dispatcher/codegen/codegen_common.py | Adds shared AQuant kernel-name construction helpers for byte-exact parity. |
| projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp | New per-kernel direct-launch C ABI for AQuant GEMM (host-pointer model). |
| projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt | CMake wiring to build an AQuant ctypes lib when a generated header exists. |
Comments suppressed due to low confidence (1)
projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp:256
- For pk_int4 A, this std::copy uses MK elements, but HostTensor<pk_int4_t> stores only (MK)/PackedSize elements (PackedSize==2). This will overrun the HostTensor backing storage and can corrupt memory before the permute/copy to device.
if constexpr(std::is_same_v<ADataType, ck_tile::pk_int4_t>)
| #ifdef CK_TILE_SINGLE_KERNEL_INCLUDE | ||
| using SelectedKernel = {ns}::{struct}; | ||
| constexpr const char* KERNEL_NAME = {ns}::KERNEL_NAME; | ||
| using ADataType = {ck_a}; | ||
| using BDataType = {ck_b}; | ||
| using CDataType = {ck_c}; | ||
| using QDataType = {ck_q}; | ||
| using AccDataType = {ck_acc}; | ||
| using QuantGroupSize = {ns}::QuantGroupSize; | ||
| using ALayout = {ns}::ALayout; | ||
| using BLayout = {ns}::BLayout; | ||
| using CLayout = {ns}::CLayout; | ||
| constexpr ck_tile::index_t GroupSizeK = {ns}::{struct}::GroupSizeK; | ||
| #endif // CK_TILE_SINGLE_KERNEL_INCLUDE |
| // This implementation only supports packed (contiguous) layouts. The expected | ||
| // leading dimensions depend on the compile-time A/B layouts baked into the kernel. | ||
| // A (ALayout): row-major -> stride_A=K ; col-major -> stride_A=M | ||
| // B (BLayout): row-major -> stride_B=N ; col-major -> stride_B=K | ||
| // AQ: row-major [M, QK_A] -> stride_AQ=QK_A | ||
| // C : row-major [M, N] -> stride_C=N (CLayout is always RowMajor) | ||
| { | ||
| constexpr bool a_row = std::is_same_v<ALayout, ck_tile::tensor_layout::gemm::RowMajor>; | ||
| constexpr bool b_row = std::is_same_v<BLayout, ck_tile::tensor_layout::gemm::RowMajor>; | ||
| const int64_t exp_stride_A = a_row ? K : M; | ||
| const int64_t exp_stride_B = b_row ? N : K; | ||
| const int64_t exp_stride_AQ = QK_A; | ||
| const int64_t exp_stride_C = N; | ||
| if(stride_A != exp_stride_A || stride_B != exp_stride_B || stride_AQ != exp_stride_AQ || | ||
| stride_C != exp_stride_C) |
| # Packed strides derived from the layout tag. | ||
| a_char, b_char, _c_char = self._layout[0], self._layout[1], self._layout[2] | ||
| stride_A = K if a_char == "r" else M # A row-major -> K, col-major -> M | ||
| stride_AQ = QK_A # AQ is row-major [M, QK_A] | ||
| stride_B = N if b_char == "r" else K # B row-major -> N, col-major -> K | ||
| stride_C = N # C is row-major [M, N] | ||
|
|
||
| rc, time_ms = self._lib.run( | ||
| A=A, AQ=AQ, B=B, C=C, | ||
| M=M, N=N, K=K, | ||
| stride_A=stride_A, | ||
| stride_AQ=stride_AQ, | ||
| stride_B=stride_B, | ||
| stride_C=stride_C, | ||
| QK_A=QK_A, | ||
| k_batch=problem.k_batch, | ||
| ) |
GPU validation (aquant, MI300X/gfx942)Validated draft PR #9980 (HEAD Shape for correctness: M=N=256, K=512, quant group 1x1x128, gate = max_rel < 2e-2. Functional correctness
ccr correctness is confirmed ( fp8i4/bf8i4 kernels compile and codegen cleanly, but my standalone harness crashes (heap corruption) in its own Arch trap confirmed (documentation)
Building the default config with the hard-coded Performance A/B parity vs Old-TE (
|
| shape (MxNxK) | bridge ms | Old-TE ms | gap % |
|---|---|---|---|
| 512x512x512 | 0.0043 | 0.0045 | -3.5 |
| 1024x1024x1024 | 0.0119 | 0.0133 | -10.5 |
| 2048x2048x2048 | 0.0675 | 0.0787 | -14.2 |
| 4096x4096x4096 | 0.5132 | 0.5331 | -3.7 |
| 1024x4096x4096 | 0.1281 | 0.1467 | -12.7 |
| 2048x512x8192 | 0.0824 | 0.0919 | -10.3 |
Median gap ~ -10.4% (bridge slightly faster), 6/6 = 100% within ±15%.
Verdict
- fp8 & bf8, all four decode layouts (incl. ccr) + fp8 preshufflequant: functionally AT PARITY with Old-TE (max_rel = 0.0).
- Performance: AT PARITY (100% within ±15%, median -10.4%).
- Blocking round-3 item: the default-config
warp_tile_kis hard-coded to128(a gfx950-only value). On gfx942 this compiles but silently produces all-zeros output.warp_tile_kmust be arch-derived (mirrorget_k_warp_tile<PrecType, M_Warp_Tile>()-> 32 on gfx942, 128 on gfx950) in the default configs / selftest, so users on gfx942 don't hit the silent all-zeros trap. (fp8i4/bf8i4 also remain to be GPU-verified with a packed-aware harness.)
Validated on ctr-cx66-mi300x-13, job 67587849.
The default_*_config helpers hardcoded warp_tile_k=128, which is gfx950-only. On gfx942 (MI300X) the kernel compiles but SILENTLY OUTPUTS ALL-ZEROS (max_rel=1.0), GPU-confirmed for fp8/bf8 across rcr/rrr/crr/ccr. With the correct gfx942 value (warp_tile_k=32) the kernels are bit-exact (max_rel=0.0) and at parity (median -10.4%, 6/6 within +/-15%). Make warp_tile_k arch-derived, mirroring ck_tile::get_k_warp_tile<PrecType,16, IsFlatMM>(). All four AQuant variants (fp8, bf8, fp8i4, bf8i4) instantiate the GEMM config with an 8-bit float PrecType (GemmConfig<fp8/bf8_t>; the pk_int4 A operand does not drive the K warp tile), so: gfx950 -> 128 (decode and preshufflequant) gfx942/other, decode -> 32 gfx942/other, preshufflequant -> 64 The value flows into the byte-exact .name. Mirrors the round-2 fixes on tensor_quant (fp8_warp_tile_k_for_arch) and rowcolquant (_warp_tile_k_for). Also make the standalone codegen default sweep arch-aware via --gfx-arch, and extend the CPU unit tests to assert the per-arch/per-dtype/per-pipeline values.
Round-3 fix: arch-aware warp_tile_kThe GPU tester found a blocking bug on gfx942 (MI300X): the Fix (commit
The arch-derived value flows into the byte-exact Tests: |
GPU validation (aquant fp8i4/bf8i4, MI300X/gfx942)Closes the deferred packed-int4 (pk_int4 A-weight) verification. The earlier GPU pass validated fp8/bf8 bit-exact + at parity but its numpy/ctypes harness crashed on pk_int4 host bookkeeping ( Env: node Functional (bridge vs
|
| variant | rcr | rrr | crr | ccr | rcr-presh | rrr-presh | crr-presh |
|---|---|---|---|---|---|---|---|
| fp8i4 | 0 | 0 | 0 | 0 | 0 | 0 | 6.4e-4 |
| bf8i4 | 0 | 0 | 0 | 0 | 0 | 0 | 0 |
14/14 PASS. 13 bit-exact (max_rel=0); fp8i4/crr/preshufflequant max_rel=6.42e-4 (max_abs=0.125, a single fp16-ULP scaling rounding), well under the 2e-2 gate. Every kernel nz_bridge == nz_ref (65536/65536) — no all-zeros trap, no crash.
Perf A/B (interleaved, warmup50/repeat100, flush + rotating_count=4 both sides; rcr — the layout Old-TE exposes for i4 aquant)
| shape (MxNxK) | fp8i4 bridge ms | fp8i4 Old-TE ms | gap % | bf8i4 bridge ms | bf8i4 Old-TE ms | gap % |
|---|---|---|---|---|---|---|
| 512x512x512 | 0.00474 | 0.01509 | -68.6 | 0.00476 | 0.01513 | -68.6 |
| 1024x1024x1024 | 0.01414 | 0.02683 | -47.3 | 0.01407 | 0.02679 | -47.5 |
| 2048x2048x2048 | 0.08169 | 0.09035 | -9.6 | 0.08357 | 0.08976 | -6.9 |
| 4096x4096x4096 | 0.59875 | 0.59919 | -0.1 | 0.59791 | 0.59597 | +0.3 |
| 1024x4096x4096 | 0.15869 | 0.16773 | -5.4 | 0.15873 | 0.16823 | -5.7 |
| 2048x512x8192 | 0.10336 | 0.11927 | -13.3 | 0.10302 | 0.12036 | -14.4 |
Bridge is at parity or faster on every shape — no bridge-slower |gap|>15% case. The large negative gaps on small shapes (bridge faster) are the known registry-bypass / direct-launch advantage and mirror the fp8/bf8 profile exactly, converging to dead-even at 4096^3.
Verdict
fp8i4 / bf8i4 PASS — functionally bit-exact (mod one sub-ULP rounding) across all 4 decode + 3 preshufflequant layouts, and at-or-better perf vs Old-TE. Combined with the earlier fp8/bf8 pass (all bit-exact + at parity), the aquant bridge is now FULLY validated on gfx942: all 4 dtypes {fp8, bf8, fp8i4, bf8i4} x {decode + preshufflequant} at 100% Old-TE parity. No blocking fix needed.
Motivation
Adds the gemm_aquant block-scale op (A is the quantized operand) to the TileEngine ->
Dispatcher bridge (codegen -> C ABI -> Python), so callers can generate, build, and launch
A-quant GEMM at parity with Old-TE without writing C++. Direct-launch, registry-bypass
(
ck_tile::QuantGemmHostArgs).Test Plan
python3 -m pytest dispatcher/tests/test_aquant_bridge.py -vTest Result
decode + preshufflequant, pipeline-key/preshuffle mapping, codegen-JSON projection, scope).
parity and 28/28 compile (one
host_tensor_descriptorbool->bool_constant issue foundand fixed).
gemm_aquant_ctypes_lib.cpp.Scope / known limitations (exact Old-TE match)
Source:
gemm_aquant_quantgrouped.cpp/..._preshufflequant.cpp. A quantized.= 28 kernels (16 decode + 12 preshufflequant;
ccrexcluded from preshufflequant, asOld-TE rejects it). Decode uses
AQuantGemmPipelineAgBgCrMem+Interwave; preshufflequant usesAQuantGemmPipelineAgBgCrCompV3+Intrawave. pk_int4 A usespermute_vectors_i4x4_b;APreshuffleQuantshuffles AQ.k_batch == 1; packed/contiguous strides enforced(non-packed rejected, never silently transposed).
Related PRs / references (TileEngine -> Dispatcher GEMM bridge series): #8997 (regular), #9000 (grouped), #9028 (stream-K), #8887 (fp8/bf8/int8), #9305 (multi-ABD), #9306 (batched), #9307 (preshuffle), #9308 (multi-D), #9328 (batched-contraction), #9329 (mx_gemm), #9978 (tensor_quant). Sibling block-scale quant bridges (this series): abquant, aquant, bquant, rowcolquant, tensor_quant.