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>)
| // 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) |
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.
The pk_int4 A path copied M*K logical elements into a HostTensor<pk_int4_t> that (PackedSize=2) only allocates M*K/2 elements, overrunning the heap before permute_vectors_i4x4_b ran. This corrupted the allocator and crashed every fp8i4/bf8i4 config (malloc abort / segfault), leaving all i4 cases N/A in the default_config parity sweep. Copy a_h.size() (the packed element count) instead, mirroring the already -fixed B path in gemm_bquant_ctypes_lib.cpp. A_host is the packed representation (2 nibbles/byte), so the packed count is the correct amount. Verified on gfx942/MI300X: fp8i4 and bf8i4 now build, run, and are bit-exact (maxrel 0.0) vs an empirically-anchored reference; fp8/bf8 non-i4 paths unaffected (change is inside the pk_int4 branch).
Fix: pk_int4 A-operand heap overflow (all fp8i4/bf8i4 configs crashed)Bug. Every Root cause. In std::copy(A_host, A_host + M * K, a_h.begin()); // overruns a_h by M*K/2 elementsThis wrote past the end of the Fix. Copy the packed element count std::copy(A_host, A_host + a_h.size(), a_h.begin());The change is confined to the Verification (gfx942 / MI300X).
Commit: |
Parity results — aquant bridge vs Old-TE (gfx942 / MI300X)Full Verdict: AT PARITY (incl. i4).
Fairness (all enforced): warmup=50/repeat=100 both sides; Old-TE via develop CMake (real TE Correctness gate: arch-aware |
…d_gemm_b_scale_wmma; i4 fix unchanged)
Fresh MI350X / gfx950 fair parity re-verify — 2026-07-27Verdict: AT PARITY. Full apples-to-apples bridge-vs-Old-TE sweep on 1× MI350X (gfx950), ROCm 7.2.0 / HIP 7.2.26015, 28-shape suite × all dtype×layout.
gap% = (bridge − oldTE)/oldTE × 100 (+ = bridge slower). 870 comparable timing rows across all 5 bridges, 0 gaps >15% (nothing required standalone re-measure), 0 all-zeros/silent-fail rows. This PR (aquant): 384 comparable cases, median −1.65%, 100% within ±15%. The pk_int4 A-copy over-read crash (fixed da3ae8d) is confirmed resolved — all fp8i4/bf8i4 rows ran clean, 0 crashes, i4 gaps within ±5%. 0 all-zeros. Fairness enforced: warmup/repeat 50/100 both sides; Old-TE built via develop CMake real TE Coverage: 64 N/A rows (small-M M≤64 unsupported — Old-TE also rejects). N/A = genuine coverage limits identical on both sides, not perf gaps. Per-bridge CSVs + consolidated report: https://gist.github.com/ozturkosu/f33d8f9bb5b503395334c5051d3c39c9 ( |
Replace verbose inline if(!= hipSuccess){cleanup();return -1;} guards
with HIP_CHECK(), which was defined but never invoked. Also update the
HIP_CHECK definition to call cleanup() on error (matching the grouped
variant). Preshuffle/permute conditional branches each get their own
HIP_CHECK call so error handling is uniform across all paths.
Co-Authored-By: Claude <noreply@anthropic.com>
…_aquant_utils Add # === section headers before the decode-family and preshufflequant-family config factory functions, mirroring the grouping style from the sibling gemm_bquant_utils.py file. Co-Authored-By: Claude <noreply@anthropic.com>
Remove trailing alignment spaces before backslash continuations in HIP_CHECK macro and remove manual column-alignment spaces in hipMalloc call sites. Co-Authored-By: Claude <noreply@anthropic.com>
|
LGTM! |
Bridge vs Old-TE parity resultsFull data (secret gist): https://gist.github.com/ozturkosu/7c9196ac2cc3eb2b25925d286f9d643a AT PARITY. gfx950 (MI355X), 384 cases, median -1.65%, 100% within +/-15% (i4 heap-overflow crash fixed; 0 crashes). Contains: consolidated parity summary + per-PR CSV result file(s). Fair A/B methodology: warmup=50/repeat=100 both sides, interleaved per (stem,shape), Old-TE built via CK CMake with real TE |
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.