Repository navigation
Conversation
Add a registry-bypass, single-kernel-per-.so, direct-launch ctypes bridge for the TensorQuant (single per-tensor scale for A and B) block-scale GEMM operator, mirroring the grouped_gemm_bquant bridge pattern. Behavioral parity with Old-TE example/ck_tile/38_block_scale_gemm/ gemm_quant_tensor.cpp: fp8/bf8 dtypes, rcr layout (RowMajor A, ColumnMajor B, RowMajor C), QuantType::TensorQuant via GemmRowColTensorQuantPipelineProblem + GemmPipelineAgBgCrCompV3, with the two scalar scales applied in the epilogue. Files: - dispatcher/codegen/unified_gemm_tensor_quant_codegen.py - dispatcher/bindings/ctypes/gemm_tensor_quant_ctypes_lib.cpp - dispatcher/python/gemm_tensor_quant_utils.py - dispatcher/bindings/ctypes/CMakeLists.txt wiring Arch handling is generic (gfx942 + gfx950, runtime device guard; no hardcoded default). Codegen/utils kernel names are byte-exact. Both fp8 and bf8 kernels compile+link via hipcc on gfx950; end-to-end Python codegen->hipcc->.so path verified.
Reviewer pass (tensor_quant bridge)Reviewed Verified correct
Should-fix
Nit
Nice work — scope, the AccDataType slot, arch guard, and name contract are all correct. LGTM pending the optional k_batch guard. |
✅ 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. |
There was a problem hiding this comment.
Pull request overview
Adds a complete TileEngine → Dispatcher bridge for TensorQuant GEMM (per-tensor A/B scalar scales) following the existing direct-launch, registry-bypass ctypes pattern, so kernels can be codegen’d, built into per-kernel .sos, and launched from Python without C++.
Changes:
- Introduces TensorQuant kernel codegen that emits a single-kernel
SelectedKernel::launch(QuantGemmHostArgs, …)header per config. - Adds a TensorQuant ctypes C ABI library and a self-contained Python build+ctypes runner (
setup_multiple_tensor_quant_dispatchers→TensorQuantGpuGemmRunner). - Adds CPU-only unit tests that lock the name contract and JSON projection behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/dispatcher/tests/test_tensor_quant_bridge.py | CPU-only tests covering name contract and config→JSON projection for tensor-quant GEMM. |
| projects/composablekernel/dispatcher/python/gemm_tensor_quant_utils.py | Python utilities to define configs/problems, run codegen+hipcc builds, and launch via ctypes. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_tensor_quant_codegen.py | Code generator for tensor-quant GEMM headers + shared kernel-name builder. |
| projects/composablekernel/dispatcher/bindings/ctypes/gemm_tensor_quant_ctypes_lib.cpp | Registry-bypass C ABI that allocates/copies host buffers and directly launches SelectedKernel. |
| projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt | Adds a CMake target to build the TensorQuant ctypes library from generated headers. |
GPU validation (tensor_quant, MI300X/gfx942)Ran on a single MI300X (gfx942), ROCm 7.2. Built the fp8 + bf8 Functional correctness (fp8, bf8; rcr; per-tensor scale)Inputs are real fp8/bf8-encoded bytes (encoded via ck_tile
Gate 2e-2. Both bit-exact vs reference.
Performance A/B parity vs Old-TE (interleaved per (kernel,shape), median of 3 rounds)Bridge: its GPU timer. Old-TE:
Median gap +5.18%, 80% within +/-15%. Verdict: AT PARITY. The two +37% outliers are the tiny 1024^3 case only: the bridge's ctypes timing path hardcodes Verdict
Follow-up (round-2 coder)
|
Round-2 fix for the GPU-confirmed blocking bug: default_fp8_config / default_bf8_config hardcoded warp_tile_k=128, which is valid only on gfx950. On gfx942 (MI300X) the 16x16x128 fp8/bf8 warp-gemm has no valid instruction, so the kernel compiled fine but silently produced all-zeros output. Derive warp_tile_k from the config's gfx_arch, mirroring ck_tile get_k_warp_tile<fp8_t/bf8_t, M_Warp_Tile=16>(): 32 on gfx942, 128 on gfx950. With warp_tile_k=32 the gfx942 kernel is bit-exact and at parity with Old-TE (which launches ...16x16x32 on gfx942). - python/gemm_tensor_quant_utils.py: new fp8_warp_tile_k_for_arch() helper; default_fp8_config/default_bf8_config now pick warp_tile_k from gfx_arch. - codegen: mirror the same arch-derived default; add --gfx-arch flag so the standalone built-in default config is also correct. warp_tile_k flows into the byte-exact KERNEL_NAME (16x16x32 vs 16x16x128). - tests: assert warp_tile_k==32 on gfx942 and ==128 on gfx950 (helper, fp8/bf8 defaults, and the encoded kernel name). gfx942 fp8 and bf8 kernels smoke-recompiled+linked via hipcc (warp_tile_k=32).
Round-2 fix: arch-aware warp_tile_kThe GPU tester (gfx942 / MI300X) found a blocking functional bug: Fix —
Changes:
Verification:
No other behavior changed (fp8/bf8 rcr per-tensor scale, accumulator = AccDataType, etc. remain as confirmed correct). |
GPU re-validation (tensor_quant, shipped default, MI300X/gfx942)Re-test of the shipped default config at HEAD Default kernel name resolves correctly on gfx942:
Functional (fp8 + bf8, rcr, per-tensor scale) — max_rel vs FNUZ quant→dequant numpy reference, gate 2e-2
Bit-exact (the ~5e-4 at large K is half-precision accumulation rounding, far under the 2e-2 gate). 8/8 PASS. Performance A/B vs Old-TE (
|
| variant | M | N | K | oldTE ms | bridge ms | gap % |
|---|---|---|---|---|---|---|
| fp8 | 16 | 64 | 256 | 0.01314 | 0.01310 | -0.27 |
| fp8 | 64 | 128 | 512 | 0.01389 | 0.01398 | +0.67 |
| fp8 | 256 | 256 | 256 | 0.01344 | 0.01356 | +0.93 |
| fp8 | 512 | 512 | 512 | 0.01544 | 0.01468 | -4.86 |
| fp8 | 1024 | 1024 | 1024 | 0.02493 | 0.02495 | +0.08 |
| fp8 | 2048 | 2048 | 2048 | 0.08623 | 0.07593 | -11.94 |
| fp8 | 3840 | 4096 | 2048 | 0.26529 | 0.22523 | -15.10 |
| bf8 | 16 | 64 | 256 | 0.01309 | 0.01333 | +1.82 |
| bf8 | 64 | 128 | 512 | 0.01394 | 0.01397 | +0.18 |
| bf8 | 256 | 256 | 256 | 0.01457 | 0.01357 | -6.90 |
| bf8 | 512 | 512 | 512 | 0.01554 | 0.01470 | -5.45 |
| bf8 | 1024 | 1024 | 1024 | 0.02485 | 0.02496 | +0.45 |
| bf8 | 2048 | 2048 | 2048 | 0.08592 | 0.07596 | -11.59 |
| bf8 | 3840 | 4096 | 2048 | 0.26488 | 0.22367 | -15.56 |
median gap = -2.57%, within ±15% = 12/14 (86%).
The earlier +37% tiny-shape outlier is GONE with symmetric flush+rotating timing (tiny shapes now -0.27% … +1.82%). The two off-gate rows are the largest shape (3840x4096x2048), where the bridge is ~15% faster than Old-TE.
Diagnosis of the two |gap|>15% rows (bridge faster at large shapes)
This is a real compiled-kernel difference, not a measurement artifact — confirmed by measuring both sides with each side's own internal gpu-timer (flush off, no external loop): fp8 3840x4096x2048 = Old-TE 0.2537 ms vs bridge 0.2149 ms (-15.3%); bf8 similar. rocprof kernel-trace shows both dispatch the same QuantGemmMultiDKernel<GemmPipelineAgBgCrCompV3, CShuffleEpilogue> at identical tile 16x64x256 / 1x4x1 / 16x16x32, identical grid (3932160). The only delta: Old-TE passes kernel_attr<eight_waves=false> into make_kernel (routing through the attributed kentry overload, VGPR=132), whereas the bridge's generated header uses plain make_kernel<kBlockPerCu> (VGPR=136). Net effect on gfx942 at large shapes: the bridge kernel is marginally more efficient (~15%); tiny/medium shapes are unaffected (at parity). The divergence favors the bridge and is benign — no correctness or performance loss.
Verdict
- Default kernel name: confirmed
...16x16x32on gfx942 (no 128, no all-zeros). - Functional: 100% PASS, bit-exact for fp8 + bf8.
- Performance: 100% at-or-better than Old-TE; median -2.57%, 12/14 within ±15%, and the 2 out-of-band rows are the bridge running ~15% faster at the largest shape (diagnosed as a benign kernel-codegen delta, not a regression). The prior tiny-shape +37% artifact is resolved by symmetric external flush+rotating timing.
Shipped default = functionally correct and at-or-better performance parity on gfx942.
…overload At large shapes the tensor_quant bridge kernel was ~15% faster than Old-TE because the generated launch used the plain make_kernel<kBlockPerCu> path (kentry<MinBlockPerCu, ...> overload -> VGPR 136), while Old-TE launches through the kernel_attr<...> / kentry<Attr, MinBlockPerCu, ...> overload (VGPR 132). Same QuantGemmMultiDKernel, tile and grid -- only the launch overload differed, producing a non-identical compiled kernel. Mirror run_gemm_quant_example.inc: compute eight_waves (always false for TensorQuant since IS_FP8BLOCKSCALE=false, guarded under CK_GFX950_SUPPORT) and launch via make_kernel<kBlockPerCu, kernel_attr<eight_waves>>. The emitted device symbol is now kentry<kernel_attr<false>, 1, QuantGemmMultiDKernel<...>>, matching Old-TE's instantiation so the compiled kernel is identical. Functional behavior is unchanged (still bit-exact fp8/bf8 rcr, per-tensor scale, arch-aware warp_tile_k). CPU unit tests unchanged (kernel name unaffected).
Round-3: align launch to Old-TE kernel_attr/kentry (perf parity)At large shapes (e.g.
Fix: the generated Verified on gfx942 via Functional behavior is unchanged (still bit-exact fp8/bf8 rcr per-tensor scale, arch-aware Note: a GPU re-test on gfx942 with the real bench harness is still needed to confirm the large-shape gap is now ≤15% (the absolute VGPR number in the untuned standalone compile differs from the harness's 132/136, but the load-bearing |
GPU re-validation round-3 (tensor_quant kernel_attr, MI300X/gfx942)Re-tested the round-3 launch fix ( Launch overload now matches Old-TE (kernels identical)The emitted device symbol is now the (host stub: VGPR before/after (same TE flags):
The bridge and Old-TE tensor_quant kernels now have identical resource usage (144/138) across all 4 hot-loop/tail specializations for both fp8 and bf8. Kernels are proven identical (symbol + VGPR + grid + tile). Functional (fp8 & bf8, rcr, per-tensor scale vs FNUZ quant→dequant reference)
All shapes bit-exact (max_rel ~4.9e-4 is fp16 output rounding). Full set 10/10 PASS. Performance A/B vs Old-TE (interleaved, warmup50/repeat100, gpu-timer, same flags, matched random inputs)
Note on the earlier ~15% "gap"The round-2 report's ~15% large-shape gap was reproducible only when the bench filled inputs with Verdicttensor_quant is at 100% parity (function + performance) on the shipped default (gfx942). Launch now uses Old-TE's |
Parity results — tensor_quant bridge vs Old-TE (gfx942 / MI300X)Full Verdict: AT PARITY.
(negative = bridge faster). Coverage: 56 total cases = 48 comparable + 8 N/A (M=1 unsupported on both sides — the default_config TileM=16 rejects M=1, Old-TE throws "Arguments not supported"). Fairness (all enforced): warmup=50/repeat=100 on both sides (bridge stream_config patched + rebuilt; Old-TE 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 (tensor_quant): 48 comparable cases, median +1.67%, 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 explain packed
types explicitly.
Co-Authored-By: Claude <noreply@anthropic.com>
…types_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>
- codegen: coerce non-"cshuffle" epilogue configs to "cshuffle" with a warning in _build_specs. TensorQuant codegen only ever emits the CShuffle epilogue and the kernel name ignores the field, so an arbitrary epilogue string could otherwise mislead callers. - CMake: when no generated TensorQuant kernel header exists, skip dispatcher_gemm_tensor_quant_lib instead of building it kernel-less. The ctypes .cpp requires a force-included generated header under CK_TILE_SINGLE_KERNEL_INCLUDE, so the old fallback could not compile. Now matches the sibling bridge targets (multi_abd/batched/bc).
|
LGTM! |
Bridge vs Old-TE parity resultsFull data (secret gist): https://gist.github.com/ozturkosu/3ffa718281ff6f50b454b029cc0b467c AT PARITY. gfx950 (MI355X), 48 cases, median +1.67%, 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
The TileEngine -> Dispatcher bridge had no path for the gemm_tensor_quant block-scale
op (single per-tensor scale on A and B). This PR adds a complete bridge (codegen -> C ABI ->
Python) so callers can generate, build, and launch tensor-quant GEMM at parity with the
legacy Tile Engine version without writing any C++.
It follows the direct-launch, registry-bypass pattern used by the other block-scale
quant bridges, because the launch takes
ck_tile::QuantGemmHostArgs(with quant scalepointers) that the single-pointer registry backend cannot express.
Test Plan
python3 -m pytest dispatcher/tests/test_tensor_quant_bridge.py -vTest Result
codegen-JSON projection round-trip, fp8/bf8-rcr scope, problem defaults).
.sos; end-to-end codegen -> hipcc -> .so -> ctypes-load succeeds;.so-reported kernelname matches config
.name.KERNEL_NAME(byte-exact),e.g.
gemm_tensor_quant_fp8_rcr_compv3_cshuffle_intrawave_16x64x256_1x4x1_16x16x128.gemm_tensor_quant_ctypes_lib.cpp.serialized GPU run (draft until then).
Scope / known limitations (exact Old-TE match)
fp8/bf8only and layoutrcronly (RowMajor A / ColMajor B / RowMajor C)-- exactly the two LUT entries
gemm_quant_tensor.cppregisters ({fp8,tensor},{bf8,tensor}); the only comborun_gemm_quant_example.incdispatches for these types.QK_A=QK_B=1), applied inthe CShuffle epilogue; matches
QuantType::TensorQuant.k_batch == 1launch semantics (split-K not registered for tensor-quant in Old-TE).IsSupportedArguments.Performance parity
Pending -- serialized A/B perf-parity vs Old-TE
gemm_tensor_quantwill be added as acomment after the GPU run (fair 50/100/flush/rotating both sides, interleaved).
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). Sibling
block-scale quant bridges (this series): abquant, aquant, bquant, rowcolquant, tensor_quant.