Repository navigation
Conversation
Extend the abquant (A+B block-scale quant) TE->dispatcher bridge runner to
cover all four GEMM layouts. The initial abquant bridge wires up the rcr
layout end-to-end; the ctypes runner assumed a single (RowMajor-A,
ColumnMajor-B) mapping, so ccr/crr/rrr were not yet plumbed through. The
device kernels are byte-identical across layouts -- the only per-layout
difference is the host-side leading dimension (stride) and buffer byte order.
- python runner: derive A/B column-major from the byte-exact kernel-name
layout token (or problem.layout); set stride_A = M|K and stride_B = K|N and
supply matching Fortran/C-order bytes, mirroring the existing AQ handling.
Add kernel_uses_column_major_a/_b helpers.
- ctypes guard: compute the expected stride_A/stride_B from the compiled
kernel's ALayout/BLayout (std::is_same_v) instead of a hard-coded value.
- codegen: re-export ALayout/BLayout as host-side aliases so the guard can
name them (mirrors the AQuantGroupSize re-export). Device kernel/struct
unchanged.
Validated on gfx950/MI350: ccr/crr/rrr across K in {1024, 4096, 14336} pass
numeric correctness vs an independent CPU reference (max_rel ~3.8e-4); rcr
unchanged.
…non-eight_waves The abquant codegen emitted make_kernel<kBlockPerCu, kernel_attr<false>> for the non-eight_waves families, which selects a distinct kentry specialization with a different register allocation (VGPR 128 -> 132) that is materially slower on large compv3 cshuffle tiles -- rocprof measured +64..90% device time at 4096^3 and the 2048x4096x14336 MoE shape on gfx950, across all A/B layouts. Old-TE's abquant instance builder launches every config through the plain make_kernel<kBlockPerCu> (Attr=void) path (it has no kernel_attr overload). Mirror it: plain launch for the non-eight_waves families, keep kernel_attr only for the eight_waves fast path. Host-only launch-wrapper selection; the device kernel and algorithm are unchanged.
✅ 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. |
…of the base compile _compile_abquant_kernel injected `-mllvm -greedy-reverse-local-assignment=1` (and -enable-noalias-to-md-conversion=1) into EVERY abquant kernel build. These two flags are only needed by the gfx950 EightWaves fast path; Old-TE's tile_engine abquant build ships only the 5 canonical base flags for every config. Applying the greedy flag to the non-EightWaves families raised those kernels' register usage (VGPR 128 -> 136 on gfx950), a pure build-flag asymmetry vs the byte-identical Old-TE kernel. Gate the two -mllvm flags to the EightWaves family (`is_eight_waves` from the generated header name); non-EightWaves configs now build flag-identical to Old-TE, and the EightWaves fast path keeps its flags. Build-only change; device kernel and codegen unchanged. Validated on gfx950: non-EightWaves compile no longer carries the greedy/noalias flags, EightWaves still does, no correctness or fast-path regression.
ozturkosu
marked this pull request as ready for review
August 31, 2026 02:19
ozturkosu
merged commit Aug 31, 2026
d00c4fe
into
users/muozturk/ck/blockscale_all_bridges
109 of 110 checks passed
ozturkosu
added a commit
that referenced
this pull request
Sep 1, 2026
#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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds full A/B GEMM layout support to the abquant (A+B block-scale quant) TileEngine → dispatcher bridge. The abquant bridge introduced in #10439 wires up the
rcrlayout; this PR addsccr,crr, andrrrso the bridge covers the same layout set the tile-engine op exposes.Motivation
abquant kernels are generated for all four layouts, but the ctypes runner wrapper assumed a single (RowMajor-A, ColumnMajor-B) mapping, so only
rcrwas plumbed through end-to-end. Full layout coverage is needed for parity with the tile-engine op.Design note
The device kernels are byte-identical across layouts — the only per-layout difference is the host-side leading dimension (stride) and the buffer byte order. This PR makes the runner layout-aware:
problem.layout), setsstride_A = M|Kandstride_B = K|N, and supplies matching Fortran/C-order bytes — mirroring the existing AQ column-major handling.stride_A/stride_Bfrom the compiled kernel'sALayout/BLayout(std::is_same_v) instead of a hard-coded value.ALayout/BLayoutas host-side type aliases so the guard can name them (mirrors the existingAQuantGroupSizere-export). No device-side change.Changes
dispatcher/python/gemm_abquant_utils.py— layout-aware strides + byte order;kernel_uses_column_major_a/_bhelpers.dispatcher/bindings/ctypes/gemm_abquant_ctypes_lib.cpp— per-layout expected strides viaALayout/BLayout.dispatcher/codegen/unified_gemm_abquant_codegen.py— host-sideALayout/BLayoutre-export.Testing
Validated on gfx950/MI350:
ccr,crr,rrracross K ∈ {1024, 4096, 14336} pass numeric correctness vs an independent CPU reference (max_rel ~3.8e-4);rcrunchanged (no regression). A full-layout parity sweep vs the tile-engine baseline follows once the comprehensive run completes.Stacking / merge order
Stacked on #10439 (base =
users/muozturk/ck/blockscale_all_bridges). Merge order: #10439 first, then this PR. Will retarget todevelopafter #10439 merges.