feat(ck-tile): rowcolquant GEMM TE to dispatcher bridge#9979
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. |
| rng = np.random.default_rng(0) | ||
| # fp8/bf8 inputs are represented on the host as float32 here for the | ||
| # reference; the kernel's ADataType/BDataType handle narrowing on device. | ||
| # For a genuine numeric check, use small integer-ish values. | ||
| A = rng.uniform(-2.0, 2.0, size=(M, K)).astype(np.float32) | ||
| B = rng.uniform(-2.0, 2.0, size=(K, N)).astype(np.float32) | ||
| AQ = rng.uniform(0.5, 1.5, size=(M,)).astype(np.float32) | ||
| BQ = rng.uniform(0.5, 1.5, size=(N,)).astype(np.float32) |
| A = np.ascontiguousarray(A) | ||
| # Kernel BLayout is ColumnMajor (rcr): B[k,n] lives at offset n*K+k. | ||
| # Supply column-major bytes for 2-D B; ascontiguousarray would force | ||
| # row-major and silently transpose. | ||
| B = np.asfortranarray(B) if B.ndim == 2 else np.ascontiguousarray(B) | ||
| AQ = np.ascontiguousarray(AQ) | ||
| BQ = np.ascontiguousarray(BQ) | ||
| C = np.ascontiguousarray(C) |
| else() | ||
| message(STATUS "No RowColQuant GEMM kernel found for ctypes lib - building without kernel") | ||
| add_library(dispatcher_gemm_rowcolquant_lib SHARED gemm_rowcolquant_ctypes_lib.cpp) | ||
| target_include_directories(dispatcher_gemm_rowcolquant_lib PRIVATE | ||
| ${PROJECT_SOURCE_DIR}/include | ||
| ${PROJECT_SOURCE_DIR}/dispatcher/include | ||
| ) | ||
| target_link_libraries(dispatcher_gemm_rowcolquant_lib PRIVATE hip::device) | ||
| set_target_properties(dispatcher_gemm_rowcolquant_lib PROPERTIES | ||
| POSITION_INDEPENDENT_CODE ON | ||
| CXX_STANDARD 17 | ||
| ) | ||
| endif() |
| if variant_key not in ROWCOLQUANT_VARIANTS: | ||
| log.warning("Unknown variant_key %s -- skipping", variant_key) | ||
| continue | ||
| if pipeline not in ROWCOLQUANT_PIPELINE_MAP: | ||
| log.warning("Unsupported pipeline %s -- skipping", pipeline) | ||
| continue | ||
|
|
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 |
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.