feat(ck-tile): enable grouped AQuant/ABQuant/BQuant GEMM bridges on gfx1250 (MI400) - #10927
feat(ck-tile): enable grouped AQuant/ABQuant/BQuant GEMM bridges on gfx1250 (MI400)#10927ozturkosu wants to merge 4 commits into
Conversation
The regular-GEMM TE->Dispatcher bridge (#8997 fp16/bf16, #8998 fp8/bf8/int8) was CDNA-only (gfx90a/gfx942/gfx950, MFMA). This enables it on gfx1250 (MI400, RDNA4/WMMA): - gemm_utils: add gfx1250 to _SUPPORTED_ARCHES + wave-combo table; add an OCP fp8/bf8 host codec (float8_e4m3fn / float8_e5m2) selected by arch (gfx950/gfx12xx -> OCP, gfx942 -> FNUZ), since gfx1250's device fp8_t is OCP. - gemm_validation_utils: gfx1250 warp-configuration combos. - arch_specs_generated: gfx1250 fp8/bf8 WMMA warp tile (16x16x64). - gemm_full_benchmark: allow --dtype fp8/bf8. - configs: gfx1250 CI configs (fp16/bf16 warp_tile 16x16x32; fp8 16x16x64). - tests: CPU-only gfx1250/OCP-fp8 codec + arch tests. Validated on MI400/gfx1250 (ROCm 10.1): fp16/bf16 and fp8 kernels build and pass --verify (fp8 max_rel ~3.6e-2 vs fp32 ref, i.e. E4M3 quantization only).
Fold the grouped (#9000), multi-D (#9308), and multi-ABD (#9305) variant enablement onto the shared gfx1250 foundation so gemm_universal + these three "full" variants ship in one PR: - grouped: gfx1250 CI config (WMMA 16x16x32) + CPU tests - multi-D: gfx1250 CI config (WMMA 16x16x32) + CPU tests + --arch help - multi-ABD: gfx1250 CI config (WMMA 16x16x32, fp16-only) + CPU tests Each merged variant pinned the CDNA MFMA 32x32x16 tile in its CI config, which does not run on gfx1250 WMMA; the new configs use 16x16x32. No variant codegen change was needed beyond the shared arch enablement. Validated on MI400/gfx1250 (ROCm 10.1): grouped 64/64 fp16 OK; multi-D 80 measurements VERIFY (max_rel ~5e-4); multi-ABD 32/32 fp16 VERIFY. 63 CPU tests pass.
…fx1250 (MI400) Enables the three merged grouped quant GEMM TE->dispatcher bridges on gfx1250, stacked on the shared gfx1250 foundation (PR #10921). C++ (small, required): the three ctypes-lib runtime arch allow-lists rejected gfx1250; add gfx1250 to each gate and to the "(supported: ...)" message. clang-format-18 applied. Python: add gfx1250 default-config helpers that pin warp_tile_m/n/k=16 so the codegen selects the gfx12 WMMA instruction (warp_tile_k=32/128 select the gfx9-only MFMA/FlatMM paths, which silently return zeros on gfx12). ABQuant uses the standard CompV3 path (not the gfx950-native eightwaves/FlatMM path). Tests: CPU unit tests assert the gfx1250 configs use WMMA warp_tile_k=16, propagate gfx_arch, and (abquant) use compv3 rather than eightwaves. CPU + codegen validated (176/176 quant utils tests pass); on-GPU validation pending gfx1250 node recovery.
On-GPU validation on MI400/gfx1250 (a07-1)Node: Build (C++ arch-gate + WMMA codegen) — primary goal
All three arch gates accept gfx1250 and the gfx12 WMMA (16x16x16) kernels compile cleanly. On-GPU runEach On-GPU numerical run needs a proper reference harness — deferred. The build result above (C++ gate + WMMA codegen OK for all three variants) is the load-bearing validation for this PR, which changed the ctypes arch gates. No code changes made. |
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
This is PR 3 of 3 in the recommended review order (review after #10921 and #10928; config correction in progress). Review order & dependencies — gfx1250 (MI400) GEMM-bridge enablement seriesThese PRs enable the merged CK GEMM bridges on gfx1250 (MI400, RDNA4/WMMA). They all target 1️⃣ #10921 — foundation: gemm_universal + grouped + multi-D + multi-ABD. 2️⃣ #10928 — batched GEMM + batched contraction. 3️⃣ #10927 — grouped AQuant/ABQuant/BQuant. Merge mechanicsAfter #10921 merges to Separate, not in this order (need C++/kernel follow-ups, kept as drafts): #10916 (mx — GPU blocked by a gfx950-only |
…le_k=128
On-GPU validation on MI400/gfx1250 (fp32-dequant reference) shows the previous
warp_tile_k=16 (WMMA) gfx1250 overrides return all-zeros for every grouped-quant
variant. The correct gfx1250 config for fp8/bf8 is the FlatMM tile
(warp_tile_k=128):
- aquant fp8/bf8: 16x16x16 -> all-zero; stock 16x16x32 -> WRONG (max_rel 1.0);
16x16x128 -> PASS (max_rel ~3e-4 / ~1e-4)
- abquant compv3 fp8/bf8: 16x16x16 and 16x16x32 -> all-zero;
16x16x128 -> PASS (max_rel ~3e-4)
- bquant fp8/bf8: stock is already 16x16x128 -> PASS (max_rel ~5e-4);
16x16x16 -> all-zero. gfx1250 helper now returns the stock config.
Changes:
- aquant/abquant gfx1250 helpers now pin warp_tile_k=128 (FlatMM).
- bquant gfx1250 helpers delegate to the stock fp8/bf8 config.
- Remove fp8i4/bf8i4 gfx1250 helpers: the packed-int4 quant path does not
compile on gfx1250 at any warp_tile_k (no gfx12 instruction). The C++ ctypes
arch-gate stays open for i4, but on-GPU correctness is deferred.
- CPU tests updated to assert warp_tile_k=128 / 16x16x128 names.
C++ ctypes arch-gates (add gfx1250) are unchanged and remain correct.
Corrected gfx1250 config + on-GPU validationThis supersedes my earlier all-zero comment. That comment concluded the all-zero output was a test-harness/scale gap in a lightweight standalone runner. That conclusion was wrong: using a proper fp32-dequant reference harness on MI400/gfx1250 shows the all-zeros were caused by the config ( Node: Per-variant on-GPU result (corrected)
What was wrong before, and why 128 is correctDirect A/B comparison on gfx1250 for each variant (same shapes, fp32 reference):
Changes pushed (
|
|
🎉 All checks passed! This PR is ready for review. |
Enables the 5 block-scale quant GEMM bridges (aquant/abquant/bquant/ rowcolquant/tensor_quant) on gfx1250: - quant_bridge_common.hpp: accept gfx1250 in the shared runtime arch check (+ updated the unsupported-arch message). - aquant/abquant: add gfx1250 to _SUPPORTED_ARCHS (bquant/rowcolquant/ tensor_quant already accept gfx12 via _uses_ocp_fp8). - All 5 arch-derived warp_tile_k selectors (_warp_tile_k_for / fp8_warp_tile_k_for_arch): treat gfx1250 like gfx950 -> 128 (FlatMM) for the 8-bit-float quant path. warp_tile_k=32 is the gfx9 MFMA path and silently outputs zeros on gfx1250 (same trap found in the grouped-quant work #10927); fp8/bf8 quant needs the FlatMM 16x16x128 tile on gfx1250. OCP fp8 for gfx12 was already wired in the utils build path. The CI configs already sweep warp_tile_k in [16,32,64,128], so 128 is exercised.
📋 Merge order for the gfx1250 (MI400) bridge-enablement seriesThis is the coordination note for landing the gfx1250 enablement PRs. There are two independent tracks plus follow-ups. Track A — core GEMM bridges (must be ordered; they share files)1️⃣ #10921 — FIRST. Foundation: gemm_universal + grouped + multi-D + multi-ABD plus the shared gfx1250 enablement ( 2️⃣ #10928 — batched GEMM + batched contraction. After #10921. GPU-validated.
Track B — block-scale quant bridges (independent)#10439 — block-scale AQuant/ABQuant/BQuant/RowColQuant/TensorQuant. ✅ Independent of Track A (verified: 0 file overlap with #10921). It introduces its own bridges and its gfx1250 enablement is self-contained ( Follow-ups — not ready (need C++/kernel work)
TL;DR |
…250 (MI400) Enables the #10010 grouped rowcolquant + tensorquant quant bridges on gfx1250: - grouped_gemm_{rowcolquant,tensorquant}_ctypes_lib.cpp: accept gfx1250 in the runtime arch gate (+ message). clang-format-18 clean. - {rowcolquant,tensorquant}_utils.py default_fp8/bf8_config: gfx12 branch using the FlatMM 8-bit tile (16x64x256, warp 1x4x1, warp_tile 16x16x128). The stock configs use MFMA 32x32x16 / CI 16x16x32, which zero-out on gfx1250 WMMA (same trap as #10439/#10927); the FlatMM warp_tile_k=128 path is what runs there. - test_{rowcolquant,tensorquant}_gpu_correctness.py: add gfx1250 to the harness _SUPPORTED_ARCHES so the on-device correctness tests run on MI400. Validated on MI400 / gfx1250 (a07-1, ROCm 10.1, gpu0): both correctness harnesses PASS 4/4 (fp8 + bf8, vs fp32 dequant reference). NOTE: local commit (not pushed). #10010 is @tlakshma's open PR and also has a pre-existing A@B^T correctness bug independent of this gfx1250 work.
2 participants
Summary
Enables three merged grouped quant GEMM TE->dispatcher bridges on gfx1250 (MI400):
grouped_gemm_aquant+grouped_gemm_abquant(originally from feat(ck-tile): add AQuant and ABQuant grouped GEMM dispatcher with ctypes bridge #9253)grouped_gemm_bquant(originally from feat(ck-tile): add BQuantGrouped GEMM dispatcher with ctypes bridge #9166)Stacked on the shared gfx1250 foundation in #10921 (base branch
users/muozturk/ck/gemm-bridges-gfx1250-enable).Requires C++? YES (small)
The 3 ctypes-lib runtime arch allow-lists rejected gfx1250 and had to be extended — that is the only C++ change. Each of
grouped_gemm_{aquant,abquant,bquant}_ctypes_lib.cpphad a runtime gate of the formif(arch != gfx950 && arch != gfx942 && arch != gfx90a) return -1;; added&& ... != gfx1250and updated the(supported: ...)message. Each file is clang-format-18 clean.Correct gfx1250 config — FlatMM
warp_tile_k=128(on-GPU determined)The original revision of this PR pinned the gfx12 WMMA tile (
warp_tile_k=16) on the theory that gfx1250 needs WMMA. On-GPU validation on MI400/gfx1250 (fp32-dequant reference) proved that wrong:warp_tile_k=16returns all-zeros for every grouped-quant variant. The empirically-correct gfx1250 config for the fp8/bf8 quant path is the FlatMM tile,warp_tile_k=128— the same tile the gfx9 fp8/bf8 FlatMM decode uses.warp_tile_k=32(plain MFMA), which is WRONG on gfx1250 (aquant max_rel=1.0; abquant all-zero). The gfx1250 helpers now pinwarp_tile_k=128.warp_tile_k=128and is correct on gfx1250 as-is, so the gfx1250 helper simply delegates to the stock config.warp_tile_k(no gfx12 instruction for the packed-int4 quant path). The C++ arch-gate stays open, but on-GPU correctness is deferred pending kernel work. No gfx1250 config helper is shipped for i4 (it would emit a non-buildable kernel).Changes
gfx1250to the runtime allow-list +(supported: ...)string ingrouped_gemm_{aquant,abquant,bquant}_ctypes_lib.cpp; clang-format-18 applied.warp_tile_k=128:default_{fp8,bf8}_config_gfx1250default_{fp8,bf8}_config_gfx1250(delegate to stock)default_{fp8,bf8}_compv3_config_gfx1250(compv3,transpose_c=false)TestGfx1250Configsnow assertswarp_tile_k=128/16x16x128names and the correct pipeline (aquantmem; bquant/abquantcompv3; abquant not eightwaves/transpose_c).Validation
pytest dispatcher/tests/test_grouped_gemm_{aquant,bquant,abquant}_utils.py→ 176/176 pass.HIP_VISIBLE_DEVICES=0): all 6 fp8/bf8 configs PASS against an fp32-dequant reference (see table + validation comment). Uses the real corrected helper functions.Dependencies / follow-ups
Note: no
Co-Authored-Bytrailer per team policy.