Repository navigation
Conversation
Registry-bypass, single-kernel-per-.so, direct-launch ctypes bridge for the RowColQuant block-scale GEMM (per-row scale of A + per-column scale of B), mirroring the grouped_gemm_bquant bridge pattern. Scope matches Old-TE gemm_quant_rowcol.cpp exactly: fp8/bf8, rcr layout (RowMajor A / ColMajor B / RowMajor C), compv3 pipeline, intrawave scheduler. Uses QuantGemmHostArgs with AQ/BQ float scale vectors (QK_A=QK_B=1, no quant groups). Arch is derived at runtime and throws on unknown (gfx942+gfx950). Files: codegen/unified_gemm_rowcolquant_codegen.py bindings/ctypes/gemm_rowcolquant_ctypes_lib.cpp python/gemm_rowcolquant_utils.py python/rowcolquant_selftest.py codegen/codegen_common.py (make_rowcolquant_kernel_name helper) bindings/ctypes/CMakeLists.txt (dispatcher_gemm_rowcolquant_lib target)
Reviewer pass (rowcolquant bridge)Reviewed against Old-TE source of truth Overall this is in good shape and closely mirrors Old-TE. The core correctness traps are handled right. Findings below. BlockingNone. Parity, correctness, arch handling, registry-bypass, name contract, and CI test are all met (see notes). Should-fix
Nit
Confirmed correct
|
✅ 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. |
Round-2 review fixes for the gemm_rowcolquant dispatcher bridge: - Self-test now feeds REAL 1-byte-per-element fp8/bf8 bytes to the kernel (via ml_dtypes e4m3/e5m2 encode) and rounds the numpy reference through the same quantize->dequantize, so "max relative error" is a genuine numeric check with a pass/fail tolerance. When ml_dtypes is unavailable the run is clearly marked SMOKE-ONLY (no correctness claim). - _detect_gpu_arch now raises instead of silently falling back to gfx950, so a flaky enumerator on a gfx942 box can never build gfx950 objects. Explicit gfx_arch= still short-circuits detection. - Add encode_fp8_bytes / quantize_dequantize_fp8 / fp8_encoding_available helpers plus CPU unit tests locking their byte-width and rounding.
Round-2 fixesAddressed both should-fix review items on the gemm_rowcolquant bridge (commit #1 — Self-test now actually validates.
#2 — No more silent arch fallback. Also added 3 CPU unit tests locking the encoder byte-width (1 byte/elem) and fp8-vs-bf8 rounding. Tests: |
There was a problem hiding this comment.
Pull request overview
This PR adds a RowColQuant (per-row A scale + per-col B scale) GEMM path to the CK TileEngine → Dispatcher bridge, enabling codegen → per-kernel .so build → Python ctypes launch for the fp8/bf8 rcr scope using the direct-launch (registry-bypass) pattern.
Changes:
- Adds a dedicated RowColQuant codegen script and shared kernel-name constructor to keep codegen↔Python naming byte-exact.
- Introduces a per-kernel ctypes
.so(directSelectedKernel::launch(QuantGemmHostArgs)) plus Python utilities/runner for building and launching. - Adds CPU-only unit tests and a Python self-test script for build/run validation.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/dispatcher/tests/test_rowcolquant_bridge.py | New CPU-only unit tests that lock naming contract and config projection for RowColQuant. |
| projects/composablekernel/dispatcher/python/rowcolquant_selftest.py | New self-test CLI to build and optionally run/verify default fp8/bf8 RowColQuant kernels. |
| projects/composablekernel/dispatcher/python/gemm_rowcolquant_utils.py | New Python build utilities + ctypes wrapper + high-level runner for RowColQuant. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_rowcolquant_codegen.py | New code generator producing per-kernel headers defining SelectedKernel + KERNEL_NAME. |
| projects/composablekernel/dispatcher/codegen/codegen_common.py | Adds make_rowcolquant_kernel_name() shared helper to enforce byte-exact naming across layers. |
| projects/composablekernel/dispatcher/bindings/ctypes/gemm_rowcolquant_ctypes_lib.cpp | New ctypes C ABI implementation doing host-pointer allocation/copy and direct kernel launch. |
| projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt | Adds a CMake target for the RowColQuant ctypes library when a generated kernel header is present. |
Round-2 fix: default_fp8_config/default_bf8_config now derive warp_tile_k from the target arch via _warp_tile_k_for(), mirroring ck_tile get_k_warp_tile<PrecType,16>(): fp8/bf8 -> 128 on gfx950, 32 on gfx942. The previous hardcoded warp_tile_k=128 is valid only on gfx950; on gfx942 there is no 16x16x128 fp8/bf8 warp-gemm, so the kernel compiles but silently outputs all-zeros. This exact bug was confirmed on the sibling tensor_quant bridge GPU tester, where warp_tile_k=32 (Old-TE's gfx942 value) is bit-exact. Also updates the byte-exact kernel .name, the codegen default-config note, and adds CPU tests asserting 32 on gfx942 / 128 on gfx950.
Round-2 fix: arch-aware warp_tile_k
This exact bug was just confirmed on the sibling tensor_quant bridge GPU tester: with Fix:
Updated: codegen default-config note, the byte-exact Verification:
The rest of the op (fp8/bf8 rcr, per-row A + per-col B scale, |
GPU validation (rowcolquant, MI300X/gfx942)Node Arch trap confirmed & round-2 fix works. The gfx942 default correctly resolves to Functional correctness (quant→dequant numpy ref, gate 2e-2)Encoded with FNUZ fp8/bf8 (gfx942 has no
Bridge output is bit-identical to Old-TE ( Performance A/B parity vs Old-TEIdentical external timing BOTH sides: same standalone harness, same CK
median gap = -0.02%, 100% within ±15% (14/14), max |gap| = 3.40% (small shapes only — measurement noise). No >15% rows. VerdictPASS — functional + performance at 100% Old-TE parity on gfx942 for fp8 and bf8. Ready from a validation standpoint. Non-blocking note for a coder
|
Round-3 GPU-tester findings (self-test/reference path only; shipping kernel is confirmed Old-TE bit-exact + perf parity): 1. Arch not threaded: rowcolquant_selftest main() called default_fp8_config()/default_bf8_config() with no arch, so --arch gfx942 still built the gfx950 16x16x128 tile -> all-zeros. Resolve arch once (arg or autodetect) and pass it to both configs and the encoders. 2. FP8 encoding was OCP-only: encode_fp8_bytes/quantize_dequantize_fp8 hardcoded OCP float8_e4m3/e5m2, so the numpy reference silently NaN'd on gfx942 (which uses FNUZ fp8_t). Select FNUZ vs OCP by arch, mirroring CK_USE_OCP_FP8 (gfx942->FNUZ, gfx950/gfx12->OCP), via ml_dtypes e4m3fnuz/e5m2fnuz. SMOKE-ONLY fallback when ml_dtypes is absent kept. Add tests asserting the arch threads into the default config and that the encoder picks FNUZ on gfx942 / OCP on gfx950. 23/23 pass.
Round-3: self-test arch + FNUZ fixesAddresses the two non-blocking self-test/reference bugs the GPU tester flagged (the shipping kernel itself is confirmed Old-TE bit-exact + perf within ±15% — these were reference-path-only issues). Fixed in
Tests: added |
Parity results — rowcolquant bridge vs Old-TE (gfx942 / MI300X)Full Verdict: AT PARITY.
(negative = bridge faster). Coverage: 56 total = 48 comparable + 8 N/A (M=1 unsupported on both sides). Fairness (all enforced): warmup=50/repeat=100 both sides; Old-TE via develop CMake (real TE Note: Old-TE's own CPU reference-verify reports FAIL on all 48 rows here — a pre-existing Old-TE reference/tolerance quirk for rowcol (identical device kernel runs on both sides; bridge output is non-zero). It is not a bridge issue and does not affect timing parity. Correctness gate: arch-aware |
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 (rowcolquant): 48 comparable cases, median +2.09%, 100% within ±15%. Arch-aware Fairness enforced: warmup/repeat 50/100 both sides; Old-TE built via develop CMake real TE Coverage: 8 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), and expand the elements_to_bytes comment to the canonical
3-line form explaining packed types.
Co-Authored-By: Claude <noreply@anthropic.com>
…ypes_lib Remove trailing alignment spaces before backslash continuations in HIP_CHECK macro and remove manual column-alignment spaces in hipMalloc and hipMemcpy call sites. Co-Authored-By: Claude <noreply@anthropic.com>
- run(): reject A/B whose dtype is not 1 byte (fp8/bf8) with a clear TypeError, and coerce AQ/BQ to float32. Prevents float32/float16 A/B from being silently reinterpreted as fp8 on the device. - self-test: the smoke-only path (no ml_dtypes) now feeds correctly sized 1-byte buffers instead of float32, so it satisfies the new dtype guard and no longer implies a bogus fp8 encoding. The verify path already used real fp8 bytes via encode_fp8_bytes(). - codegen: skip unsupported layouts (only rcr) with a warning in _build_specs; downstream indexes layout[0/1/2] + ROWCOLQUANT_LAYOUT_TO_CK and would otherwise raise IndexError/KeyError. - CMake: skip dispatcher_gemm_rowcolquant_lib when no generated kernel header exists instead of building it kernel-less (the ctypes .cpp requires a force-included header under CK_TILE_SINGLE_KERNEL_INCLUDE). Matches the sibling bridge targets.
|
LGTM! |
Bridge vs Old-TE parity resultsFull data (secret gist): https://gist.github.com/ozturkosu/0387019a2f18cdb43c494286bc755d1f AT PARITY. gfx950 (MI355X), 48 cases, median +2.09%, 100% within +/-15%. 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 |
#10439) JIRA ID : AICK-2090 ## Summary This PR lands the five block-scale quantized GEMM TileEngine→Dispatcher bridges as a single branch on top of `develop`, so the dispatcher can run every block-scale quant GEMM that the Tile-Engine supports. Each op differs in *how* the scales are laid out on the A and B operands: - **tensor_quant** (#9978) — **per-tensor** quantization: a single scalar scale for the whole A matrix and a single scalar scale for the whole B matrix. The simplest scheme; the kernel dequantizes with two scalars in the epilogue. - **rowcolquant** (#9979) — **per-row / per-column** quantization: one scale per row of A and one scale per column of B (a length-M vector for A, a length-N vector for B). Common for activation×weight GEMMs where each output row/column has its own dynamic range. - **aquant** (#9980) — **A block-scale** quantization: A is split into blocks along K (`quant_group_k`) and each block carries its own scale; B is the plain low-precision operand. Supports fp8/bf8 and the packed i4 variants, and both the decode (`mem`) and preshuffle (`compv3`) pipelines across the rcr/rrr/crr/ccr layouts. - **abquant** (#9981) — **A+B block-scale** quantization: A is block-scaled along K (`aquant_group_k`) and B is block-scaled along N and K (`bquant_group_n`, `bquant_group_k`). Includes the gfx950 **EightWaves** path and the B-preshuffle path. - **bquant** (#9982) — **B block-scale** quantization: B is split into blocks and each block carries its own scale (A is the plain operand). Includes the **MX / microscaling** variants (e8m0 block scales: bf16×bf16, bf16×bf8, bf16×fp4) and the **grouped**-GEMM form. Each op is a registry-bypass, single-kernel-per-`.so` ctypes bridge that builds its kernel at runtime through the Python codegen path. Supersedes / combines #9978–#9982. ## Motivation The five bridges started as sibling draft PRs that all touch the same three shared files (`bindings/ctypes/CMakeLists.txt`, `tests/CMakeLists.txt`, `codegen/codegen_common.py`). Landing them together resolves the cross-PR merge conflicts once and gives a single, fully-tested branch instead of five that keep colliding. ## Design - **Registry-bypass ctypes bridges** — each op appends its own target/test/helper to the shared files, so the merges are purely additive. - **Arch safety** — get-arch-or-`FATAL_ERROR` (no silent gfx942 default). - **abquant clang guard** — `gemm_abquant_utils.py` gates `-amdgpu-coerce-illegal-types=1` on the local clang accepting it. clang≥22 (ROCm 7.2) removed the flag and aborts the abquant `.so` compile; it only tightens EightWaves register allocation on older toolchains and the kernel is bit-accurate without it. This lets abquant actually build+run on ROCm 7.2 instead of self-skipping. - **aquant pipeline/preshuffle decoupling** (latest commit) — the aquant codegen used to hard-couple the pipeline to the preshuffle flag (`compv3 ⇔ preshuffle`), so it could not emit the `compv3 + APreshuffleQuant=false` family: 616 valid kernels the Tile-Engine builds. I added an independent `pipeline` field so `compv3` is selectable without preshuffle, giving the bridge full coverage of the Tile-Engine's aquant kernel space. This is a Python-codegen-only change (the C++ template already exposes `APreshuffleQuant` as a separate compile-time Traits flag). ## How this was tested **1. CPU unit / bridge tests** (host-only, run in CI): `test_{tensor_quant,rowcolquant,aquant,abquant,bquant}_bridge.py`, `test_quant_bridge_shared.py`, `test_codegen_common.py`, and the per-op self-tests. These cover config construction, kernel-name generation, codegen output, and the dtype/layout guards. **2. On-device GPU correctness** (MI350X / gfx950, ROCm 7.2) — bridge output checked against an independent CPU reference, all passing: - [x] tensor_quant — 3/3 (fp8, bf8, not-all-zeros) - [x] rowcolquant — 3/3 - [x] aquant — 3/3 (incl. i4) - [x] abquant — 3/3 (fp8, bf8, not-all-zeros) after the clang guard above - [x] bquant — 6/6 (fp8/bf8, i4, MX bf16bf16 / bf16bf8 / bf16fp4) **3. CI regression** — build + tests green across gfx942, gfx950, gfx90a, gfx11, gfx1201, gfx1250, plus the static-analysis / clang-tidy / pre-commit gates. **4. Comprehensive same-kernel performance parity vs the Tile-Engine.** This is the main validation. For every kernel the Tile-Engine builds from its **full `default_config`** (the whole tuning space, not the small CI subset), I construct the **byte-identical** kernel on the dispatcher bridge and run the two interleaved A/B with the **same warmup=50 / repeat=100 and the same TE `-mllvm` flags**, on **both MI300 (gfx942) and MI350 (gfx950)**. Every `>15%` outlier is re-measured standalone (median-of-3 on an isolated GPU) before it is trusted, to rule out contention artifacts. Result: the bridges are **at performance parity with the Tile-Engine** — per-op medians land within a few percent and the large majority of kernels are within ±15%; the handful of survivors are narrow, mostly bridge-*faster* differences at awkward shapes, not regressions. The aquant coverage fix above was validated at full scale on gfx942: **1232/1232 build, 1200/1232 GPU-correct, performance median +0.47% (95.8% within ±15%).** The 32 correctness failures are a **pre-existing shared CK kernel bug** — a `192-tile + 2x2x1 + 32x32x32 + row-major-A` cluster that the **Tile-Engine also fails its own verification on** (byte-identical kernel) — so it is not introduced by this PR; I'll track it as a separate upstream issue. **5. bquant correctness.** The Tile-Engine's bquant benchmark baseline does not compile at this head (a pre-existing upstream codegen defect, unrelated to the bridge), so there is no Tile-Engine performance baseline for bquant — it is validated by **correctness only**: every buildable bquant kernel passes on-device against a CPU reference (MX variants are correctly gfx950-only). --- Supersedes / combines #9978, #9979, #9980, #9981, #9982. --- > **Note — abquant layout coverage:** the abquant bridge in this PR currently supports the `rcr` layout. Full `ccr`/`crr`/`rrr` layout support is added in the follow-up PR #11477 (stacked on this branch, to merge after this one). --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: ThruptiRajLakshmanaGowda <Thruptiraj.LakshmanaGowda@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Motivation
Adds the gemm_rowcolquant block-scale op (per-row scale on A, per-column scale on B) to
the TileEngine -> Dispatcher bridge (codegen -> C ABI -> Python), so callers can generate,
build, and launch row/col-quant GEMM at parity with Old-TE without writing C++. Uses the
direct-launch, registry-bypass pattern (launch takes
ck_tile::QuantGemmHostArgswithscale pointers the registry backend can't express).
Test Plan
python3 -m pytest dispatcher/tests/test_rowcolquant_bridge.py -vTest Result
codegen-JSON projection, fp8/bf8-rcr scope, problem defaults).
.soexporting all 6
dispatcher_*symbols. Name parity byte-exact(
gemm_rowcolquant_{fp8,bf8}_rcr_compv3_cshuffle_intrawave_16x64x256_1x4x1_16x16x128).gemm_rowcolquant_ctypes_lib.cpp.Scope / known limitations (exact Old-TE match)
fp8/bf8, layoutrcronly — the exact setgemm_quant_rowcol.cppregisters.QK_A=QK_B=1), applied inthe CShuffle epilogue;
QuantType::RowColQuant. Accumulator slot =AccDataType(float),narrowed to half only in the epilogue.
k_batch == 1(split-K not registered for rowcol in Old-TE). Non-tile-multiple M/N/Krejected by
IsSupportedArguments.Performance parity
Pending — serialized A/B vs Old-TE
gemm_rowcolquant, added as a comment after the GPU run.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.