feat(ck-tile): add AQuant and ABQuant grouped GEMM dispatcher with ctypes bridge - #9253
Merged
Merged
Conversation
…ompv3, gfx950) Implements the three-layer dispatcher path for BQuantGrouped GEMM (C = A @ dequant(B, BQ)) directly in the dispatcher, following the same codegen → C ABI → Python pattern used by GEMM and Conv. Independent of PR #8997; does not depend on gemm_utils.py. New files: - dispatcher/codegen/unified_bquant_gemm_codegen.py Generates one .hpp per kernel config. Kernel name is byte-exact with BQuantKernelConfig.name in bquant_gemm_utils.py. Initial scope: fp8/bf8 variants, compv3 pipeline, rcr layout, configurable QuantGroupShape<sequence<M,N,K>>. - dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp Flat C ABI bridge compiled per-kernel via force-include. Exposes dispatcher_run_bquant_gemm() with host-pointer memory model (hipMalloc/hipMemcpy owned by the C lib). Follows current gemm_ctypes_lib.cpp pattern. - dispatcher/python/bquant_gemm_utils.py Self-contained Python bridge: BQuantKernelConfig, BQuantGemmProblem, BQuantDispatcherLib, BQuantGpuGemmRunner, and setup_multiple_bquant_dispatchers() with inline subprocess helpers (no import from gemm_utils.py). - dispatcher/examples/gemm/python/13_bquant_gemm.py End-to-end example: codegen → hipcc → run → NumPy reference verify. - dispatcher/tests/test_bquant_gemm_utils.py 20 CPU-only unit tests covering name generation, config serialization, and problem dimension helpers. All passing. Modified files: - dispatcher/codegen/codegen_common.py Adds BQUANT_DTYPE_MAP as a new top-level dict for reuse by future AQuant/ABQuant codegens. - dispatcher/bindings/ctypes/CMakeLists.txt Adds dispatcher_bquant_gemm_lib target (same pattern as GEMM/Conv). Co-Authored-By: Claude <noreply@anthropic.com>
Renames all files and kernel name prefix to match the operation name (GroupedGemm BQuant) consistently: unified_bquant_gemm_codegen.py → unified_grouped_gemm_bquant_codegen.py bquant_gemm_ctypes_lib.cpp → grouped_gemm_bquant_ctypes_lib.cpp bquant_gemm_utils.py → grouped_gemm_bquant_utils.py 13_bquant_gemm.py → 13_grouped_gemm_bquant.py test_bquant_gemm_utils.py → test_grouped_gemm_bquant_utils.py Kernel KERNEL_NAME prefix updated from "bquant_gemm_" to "grouped_gemm_bquant_" (byte-exact match maintained between codegen and BQuantKernelConfig.name). CMakeLists.txt target renamed to dispatcher_grouped_gemm_bquant_lib. All 20 CPU-only tests pass. Co-Authored-By: Claude <noreply@anthropic.com>
Example file still imported from bquant_gemm_utils after rename. Updated to grouped_gemm_bquant_utils. Co-Authored-By: Claude <noreply@anthropic.com>
…nt GEMM dispatcher Phase 2: Add fp8i4 and bf8i4 dtype variants (A=fp8/bf8, B=pk_int4, Q=fp8/bf8) with 6 quant groups each (12 new kernels). Phase 3: Add preshuffle infrastructure (double_smem_buffer field, preshuffleb/ microscale pipeline map entries) and 12 convenience config constructors covering all 4 dtypes × 3 preshuffle modes (3a: preshuffle_b only via WPQuantBPipelineAgBgCrV2; 3b: preshuffle_bquant only via compv3; 3c: both flags). Phase 4: Add mx_bf16bf16, mx_bf16bf8, mx_bf16fp4 MX microscale variants (Q=e8m0 block scale) via MicroscaleGemmPipelineAgBgCrCompV3. Corrects base pipeline for microscale to BaseWeightPreshufflePipelineAGmemBGmemCRegV2 (matches run_gemm_quant_example.inc else branch for BQuantGrouped non-preshuffle). Tests: 58 total (32 new across TestPhase3Infrastructure, TestPhase3Configs, TestPhase4MXConfigs), all passing. Co-Authored-By: Claude <noreply@anthropic.com>
Seven issues found in code review and fixed: 1. Epilogue selection: emit PermuteNEpilogue when TiledMMAPermuteN=true and quant_group_n==1, CShuffleEpilogue otherwise. CShuffleEpilogue reads Problem::TiledMMAPermuteN which CShuffleEpilogueProblem does not define, causing a compile error for prefill tiles (N_repeat even). Mirrors run_gemm_quant_example.inc logic exactly. 2. ProcessPoolExecutor → ThreadPoolExecutor in setup_multiple_bquant_dispatchers. Nested closures cannot be pickled for multiprocessing; hipcc work is subprocess I/O-bound so ThreadPoolExecutor is correct. 3. Example input dtype: replace float16 stand-in with real fp8 bytes via _float32_to_fp8 (ml_dtypes when available, clamped int8 fallback). CPU reference now uses decoded fp8 values for an apples-to-apples check. 4. stream_config: time_kernel_ false, cold_niters_ 0 so the kernel fires exactly once per dispatch call instead of twice (1 warmup + 1 timed). 5. QK_B/QN_B validation: assert caller-supplied values match ceil(K/QuantGroupSize::kK) and ceil(N/QuantGroupSize::kN) before allocating the BQ device buffer. 6. Deduplicate kernel name logic: extract make_bquant_kernel_name() to codegen_common so BQuantKernelConfig and BQuantKernelSpec both delegate to a single source of truth, preserving the byte-exact invariant. 7. Remove three dead imports (TileConfig, parallel_generate, BQUANT_DTYPE_MAP) from unified_grouped_gemm_bquant_codegen.py. Co-Authored-By: Claude <noreply@anthropic.com>
… ctypes lib Co-Authored-By: Claude <noreply@anthropic.com>
…lib.cpp Co-Authored-By: Claude <noreply@anthropic.com>
…M bridge
- Fix kernel name/epilogue mismatch: add bquant_effective_epilogue() to
codegen_common.py so make_bquant_kernel_name and generate() both compute
the real epilogue tag (permute_n vs cshuffle) from tile geometry, keeping
KERNEL_NAME and generated C++ in sync.
- Fix cross-arch .so cache collision: append _{arch} to the cached .so
filename in setup_multiple_bquant_dispatchers so gfx-specific builds are
not mistakenly reused across architectures.
- Fix C output dtype for MX variants: add c_dtype parameter to
BQuantGpuGemmRunner.run() (default np.float16); MX callers must pass
np.bfloat16 since those variants have CDataType = bf16.
- Add packed-stride validation in dispatcher_run_bquant_gemm: reject
non-packed strides before device memory is allocated to prevent silent
out-of-bounds access.
- Fix wrong filename in 13_grouped_gemm_bquant.py docstring Usage lines.
- Fix _float32_to_fp8 docstring: fallback clamps to [-2.0, 2.0] not [-448, 448].
- Clarify BQuantDispatcherLib.run() docstring: B must be a packed (K, N)
C-contiguous array; the kernel interprets it as column-major via stride_B=K.
Co-Authored-By: Claude <noreply@anthropic.com>
…t_ctypes_lib.cpp Co-Authored-By: Claude <noreply@anthropic.com>
Extends the BQuant ctypes dispatcher bridge (PR #9166) to cover the two remaining grouped-quantization modes: - AQuantGrouped (QuantType::AQuantGrouped): A-side activation quantization. Scale tensor AQ[ceil(M/gM), ceil(K/gK)] is RowMajor. Non-preshuffle kernels use AQuantGemmPipelineAgBgCrMem; preshuffle variants use AQuantGemmPipelineAgBgCrCompV3. - ABQuantGrouped (QuantType::ABQuantGrouped): both A-side and B-side quantization active simultaneously. Requires AQuantGroupSize::kK == BQuantGroupSize::kK. Supports compv3, eightwaves (gfx950), and preshuffleb pipelines. Each operator follows the same three-layer pattern: codegen (.py → .hpp) → C bridge (.cpp, extern "C") → Python utils (.py) New files per operator: codegen script, ctypes C library, Python utils, 50/44 CPU-only unit tests (all passing), and an end-to-end example. Co-Authored-By: Claude <noreply@anthropic.com>
✅ 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. |
ThruptiRajLakshmanaGowda
changed the title
ThruptiRajLakshmanaGowda
changed the title
Contributor
|
can you share the complete results by running tests on GPU to make sure that dispatcher is picking the right instance? |
- Update bquant test assertions: permute_n epilogue is computed from tile geometry (tile_n / (warp_n * warp_tile_n)), not the user-specified string. preshuffleb/preshufflequant/MX configs with n_repeat%2==0 and quant_group_n==1 now correctly expect permute_n (12 tests fixed). - Remove non-ASCII em-dashes from grouped_gemm_abquant_ctypes_lib.cpp (replaced with ASCII hyphens). - CMakeLists.txt: replace hardcoded GFX_ARCH="gfx950" with CMAKE_HIP_ARCHITECTURES-derived variable for all three quant libs (BQuant, AQuant, ABQuant), matching the pattern from PR #9166. - Resolve Jenkinsfile merge conflicts (keep agent none for TILE_ENGINE_GEMM stages that use ck.runOnHealthyNode). Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Resolved conflicts in 6 files under projects/composablekernel/dispatcher: - CMakeLists.txt: kept AQuant/ABQuant library targets (from HEAD) and improved BQuant GFX_ARCH comment (from develop) - grouped_gemm_bquant_ctypes_lib.cpp: took develop's direct-launch design (no registry, packed type handling via elements_to_bytes, runtime GPU arch check, BQ preshuffle, improved stream_config with timing support) - codegen_common.py: kept HEAD's AQuant/ABQuant kernel name helpers - unified_grouped_gemm_bquant_codegen.py: took develop's GroupSizeK field and warp_tile_k=128 default - grouped_gemm_bquant_utils.py: took develop's two-step hipcc build, expand_bquant_sweep, permute_n unscramble, and warp_tile_k=128 defaults - test_grouped_gemm_bquant_utils.py: kept all tests from both sides Co-Authored-By: Claude <noreply@anthropic.com>
…orrectness
The previous warp_tile_k=16 dispatched to WarpGemmWmma_f32_16x16x16_f8_f8
(WMMA gfx12). Its wmma_intrinsic is guarded by #ifdef __gfx120__; on gfx950
the #else branch silently returns CVecType{0.f}, so every warp GEMM call
produced zeros and max_rel=1.0000 in CI.
Fix: set warp_tile_k=32 in default_fp8_config, default_bf8_config,
default_fp8i4_config, and default_bf8i4_config. This selects
WarpGemmMfma_f32_16x16x32_fp8_fp8 (MFMA), which is a valid instruction
on gfx942/gfx950.
Also add sample-value logging on verification failure to simplify future
debugging.
Co-Authored-By: Claude <noreply@anthropic.com>
…onverter The fallback (no ml_dtypes) used `(clip(x,-2,2)*64).astype(int8)` which produces arbitrary byte patterns — including 0x7F and 0xFF which are NaN in float8_e4m3fn. The GPU's fp8 hardware decodes those bytes as NaN, poisoning the entire GEMM output. Replace with _soft_f32_to_fp8 / _soft_fp8_to_f32 that implement correct IEEE-like fp8 encoding/decoding (e4m3fn and e5m2) in pure Python. The encoder never emits NaN bytes, and the decoder matches what the GPU hardware produces, so the CPU reference agrees with the kernel output. Co-Authored-By: Claude <noreply@anthropic.com>
…fx9 correctness
All six default ABQuant config functions (compv3, eightwaves, preshuffleb
for fp8 and bf8) had warp_tile_k=16, which dispatches to
WarpGemmWmma_f32_16x16x16_f8_f8 (WMMA, gfx12-only). Its wmma_intrinsic
is guarded by #ifdef __gfx120__; on gfx950 the #else branch silently
returns CVecType{0.f}, so every warp GEMM call produced zeros.
Fix: set warp_tile_k=32 in all six defaults. This selects
WarpGemmMfma_f32_16x16x32_fp8/bf8_fp8/bf8 (MFMA), a valid instruction
on gfx942 and gfx950. tile_k=128 % (warp_k=1 * warp_tile_k=32) == 0
holds for all three pipelines.
This is the same class of bug fixed for AQuant in d02213e.
Co-Authored-By: Claude <noreply@anthropic.com>
The fallback (no ml_dtypes) used (clip(x,-2,2)*64).astype(int8) which produces arbitrary byte patterns — including 0x7F and 0xFF which are NaN in float8_e4m3fn. The GPU's fp8 hardware decodes those bytes as NaN, poisoning the entire GEMM output (max_rel=nan). Replace with _soft_f32_to_fp8 / _soft_fp8_to_f32 that implement correct IEEE-like fp8 encoding/decoding (e4m3fn and e5m2) in pure Python. Also add failure diagnostics (C_gpu/C_ref sample values, nonzero count). Mirrors the fix applied to example 14 in 369f0df. Co-Authored-By: Claude <noreply@anthropic.com>
…erence numpy.uint8 arithmetic wraps on subtraction: when exp_bits < bias, (exp_bits - bias) underflows to 255 as uint8, making 2.0**255 overflow to inf then NaN in the CPU reference matmul. Fix: cast exp_bits, mant_bits, and sign to Python int before use so subtraction stays in signed integer arithmetic. Affects both example 14 (AQuant) and example 15 (ABQuant). Co-Authored-By: Claude <noreply@anthropic.com>
…operties Both aquant and abquant ctypes libs were querying hipGetDevice + hipGetDeviceProperties on every dispatcher_run_* call. Move the arch query and unsupported-arch rejection into dispatcher_initialize(), cache the result in g_gfx_arch, and clear it in dispatcher_cleanup(). Co-Authored-By: Claude <noreply@anthropic.com>
… libs Co-Authored-By: Claude <noreply@anthropic.com>
…shuffleb on gfx950 On gfx950 (CK_GFX950_SUPPORT defined), get_k_warp_tile<fp8_t/bf8_t, 16> returns 128 for both IsFlatMM=false (EightWaves) and IsFlatMM=true (PreshuffleB). The default configs hardcoded warp_tile_k=32/64, causing a mismatch with the value the C++ pipeline templates compute at compile time, leading to build failures or wrong results on gfx950. Add _eightwaves_warp_tile_k() and _preshuffleb_warp_tile_k() helpers that mirror the get_k_warp_tile logic and select the arch-correct value at config time. Co-Authored-By: Claude <noreply@anthropic.com>
ABQuant prefill configs have kPadM=false and kPadN=false, so M and N must be exact multiples of TileM and TileN. The previous defaults (128x128) are not multiples of the EightWaves tile (192x256), causing IsSupportedArgument() to return false (rc=-2). Use pipeline-aware defaults: compv3/preshuffleb: 1024x1024x1024 (multiples of 128x128) eightwaves: 1152x1024x1024 (multiples of 192x256) Co-Authored-By: Claude <noreply@anthropic.com>
…n gfx950 GemmConfigABQuantPrefill (compv3 pipeline) also uses get_k_warp_tile<fp8_t/bf8_t, 16, IsFlatMM=false>, which returns 128 on gfx950 (CK_GFX950_SUPPORT) and 32 on gfx942. The previous hardcoded warp_tile_k=32 produced a stale cached .so whose permute_n unshuffling was computed with the wrong _r factor, giving max_rel~1.5 on gfx950. Add _compv3_warp_tile_k() helper and use it in default_fp8/bf8_compv3_config(). Co-Authored-By: Claude <noreply@anthropic.com>
ABQuantGemmPipelineAgBgCrCompV3 uses standard MFMA (mfma_f32_16x16x32), not FlatMM. WarpTileK=128 caused the pipeline to accumulate 4x fewer iterations, producing ~4x underscaled results on gfx950. Only the gfx950-native pipelines (eightwaves, preshuffleb) use FlatMM and require WarpTileK=128 on gfx950. CompV3 must stay at WarpTileK=32 on all architectures. Co-Authored-By: Claude <noreply@anthropic.com>
The PermuteNEpilogue riffles columns within each NPerBlock-wide tile, not across the entire N dimension. The Python runners were computing _half = N // r (global), which is correct only when N == tile_n (single tile). With N=1024 and tile_n=128 (8 tiles), the global formula mapped physical col 1 → logical col 512 instead of logical col 64, corrupting results beyond the first tile boundary. Fix: _half = tile_n // r, with the permutation applied tile-by-tile: logical[c] = tile_base + (c_in_tile % r) * half + (c_in_tile // r) Guard updated from (N % r) == 0 to (N % tile_n) == 0, which is the correct alignment requirement (kPadN=False kernels already enforce this). Applies the same fix to abquant, bquant, and aquant runners, which all had the identical bug. Existing tests were not affected because they use N == tile_n (one tile), where global and per-tile formulas coincide. Co-Authored-By: Claude <noreply@anthropic.com>
The EightWaves pipeline (ABQuantGemmPipelineAgBgCrEightWaves) requires the no-packed-fp32-ops kernel attribute so the compiler does not emit packed FP32 instructions that conflict with matrix operations on gfx950. The C++ test fixtures use kernel_attr<eight_waves> (where eight_waves is true for the 4x2x1 warp / warp_tile_k=128 config on gfx950). The dispatcher codegen was using plain make_kernel<kBlockPerCu> without this attribute, causing incorrect results (NaN/garbage) at runtime. Fix: add a static constexpr IsEightWaves flag to the generated struct and pass kernel_attr<IsEightWaves> to make_kernel. For compv3 and preshuffleb pipelines IsEightWaves=false so kernel_attr<false> is a no-op, keeping those paths unchanged. Because the cached .so does not encode this attribute, the stale eightwaves .so must be deleted and recompiled after this fix. Co-Authored-By: Claude <noreply@anthropic.com>
The EightWaves pipeline sets TransposeC=true which swaps the MFMA accumulator M/N register distribution. The PermuteNEpilogue dst-index formula (dst = n_idx + m_lane * NRepeat) assumes the non-transposed layout and produces incorrect output column ordering with TransposeC. The C++ test fixture (GemmConfigEightWaves) explicitly sets TiledMMAPermuteN=false, forcing CShuffleEpilogue. The Python codegen was applying the general PermuteN heuristic (N_Repeat%2==0 && bqgn==1) without this exception, emitting PermuteNEpilogue for the eightwaves pipeline when bquant_group_n=1. This produced NaN/garbage output. Fix: abquant_effective_epilogue() takes an optional pipeline parameter and returns "cshuffle" unconditionally for "eightwaves". This also changes the kernel name for eightwaves+bqgn=1 from permute_n to cshuffle, so stale cached .so files are bypassed automatically. Co-Authored-By: Claude <noreply@anthropic.com>
…e 15 The EightWaves pipeline is only validated with bquant_group_n=128 (GroupSize2D128N in the C++ tests). With bquant_group_n=1 the BQ block tile has 256 N-groups per K-block (NQPerBlock=256), which is outside the tested range and produces wrong results. Change --bquant-group-n default from a hard-coded 1 to None and derive a pipeline-aware default at runtime: compv3 / preshuffleb: 1 (one scale per column, validated) eightwaves: 128 (matches GroupSize2D128N, validated) Users can still pass --bquant-group-n explicitly to override. Co-Authored-By: Claude <noreply@anthropic.com>
Same root cause as the EightWaves fix: PreshuffleB uses TransposeC=true (GemmConfigPreshuffleB_ABQuant_Prefill), so TiledMMAPermuteN=false in all C++ tests — PermuteNEpilogue is incompatible with the transposed accumulator layout and produces wrong results. C++ tests validate PreshuffleB ABQuant only with bquant_group_n=128 (GroupSize2D128N). With bquant_group_n=1 (NQPerBlock=128, outside the tested range) the kernel produces incorrect output. - abquant_effective_epilogue: treat "preshuffleb" same as "eightwaves" (always return "cshuffle") - example 15: default bquant_group_n=128 for preshuffleb (was 1) Co-Authored-By: Claude <noreply@anthropic.com>
…nel launch
The WPABQuantBPipelineAgBgCrV2 (preshuffleb) kernel requires B to be
pre-shuffled on the host into the [N/wt_n, K/items, wt_n, items]
interleaved layout before copying to device. The C++ test fixture calls
ck_tile::shuffle_b<GemmConfig>(b_k_n) before b_k_n_dev_buf.ToDevice().
The Python runner was missing this step entirely, sending raw col-major
B bytes to a kernel that expects the shuffled layout.
Add _shuffle_b_cdna() mirroring the gfx9/CDNA branch of shuffle_b:
KLane = warp_size / wt_n = 64 / wt_n
items = min(16, wt_k / KLane)
flatten B col-major → view [N/wt_n, wt_n, K/items, items] → permute {0,2,1,3}
Apply it in ABQuantGpuGemmRunner.run() when 'preshuffleb' is in the
kernel name, extracting wt_n and wt_k from the name via the existing
tile-shape regex. Pass the result as a 1-D array so the ctypes
dispatch layer's np.asfortranarray(B) is a no-op.
Co-Authored-By: Claude <noreply@anthropic.com>
…uffleB ABQuant np.asfortranarray(B).flatten() returns row-major bytes regardless of memory layout; the C++ HostTensor[K,N] col-major iterates column-major (K varies fastest). Use B.flatten(order='F') to get the correct column-major byte sequence for _shuffle_b_cdna. Co-Authored-By: Claude <noreply@anthropic.com>
ozturkosu
approved these changes
Aug 3, 2026
ozturkosu
left a comment
ozturkosu
left a comment
Contributor
There was a problem hiding this comment.
It passed all test and my point was applied. LGTM :)
yraparti
approved these changes
Aug 3, 2026
ThruptiRajLakshmanaGowda
enabled auto-merge (squash)
August 4, 2026 18:06
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 4, 2026
feat(ck-tile): add AQuant and ABQuant grouped GEMM dispatcher with ctypes bridge (#9253) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation This PR extends the CK tile dispatcher with full quantization coverage for grouped GEMM operations. Building on the BQuant dispatcher introduced in #9166, this PR adds dispatcher bridges for AQuant (A-side quantization) and ABQuant (both sides quantized), completing the three-mode quantization dispatch surface for grouped GEMM on gfx950 (MI350X). ## Technical Details Stacked on #9166 (BQuant grouped GEMM dispatcher). Adds AQuant (C = dequant(A, AQ) @ B) and ABQuant (C = dequant(A, AQ) @ dequant(B, BQ)) dispatcher bridges following the same three-layer pattern (ctypes C API → Python dispatcher → codegen). Each mode adds: *_ctypes_lib.cpp — C API exposing the kernel to Python via ctypes unified_*_codegen.py — generates kernel instantiations across the full dtype/layout/quant-group matrix *_utils.py — Python dispatcher handling argument validation, dtype selection, and parallel hipcc builds via ThreadPoolExecutor Example script and CPU-only unit tests codegen_common.py is extended with shared epilogue selection helpers reused across all three quant modes. CMakeLists.txt updated with build targets for both new ctypes shared libraries. Supported dtypes: fp8/bf8 activations, fp8i4/bf8i4 packed int4, preshuffle layouts, and MX microscaling variants — targeting gfx950 (MI350X). ## Test Plan CPU-only unit tests covering: Kernel name generation across all dtype/layout/quant-group combinations Config serialization and dimension helpers Epilogue selection logic for AQuant and ABQuant variants End-to-end example scripts (14_grouped_gemm_aquant.py, 15_grouped_gemm_abquant.py) with CPU reference verification for correctness checks. ## Test Result CPU unit tests pass for both AQuant and ABQuant dispatcher utilities. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 18, 2026
feat(ck-tile): add AQuant and ABQuant grouped GEMM dispatcher with ctypes bridge ## Motivation This PR extends the CK tile dispatcher with full quantization coverage for grouped GEMM operations. Building on the BQuant dispatcher introduced in #9166, this PR adds dispatcher bridges for AQuant (A-side quantization) and ABQuant (both sides quantized), completing the three-mode quantization dispatch surface for grouped GEMM on gfx950 (MI350X). ## Technical Details Stacked on #9166 (BQuant grouped GEMM dispatcher). Adds AQuant (C = dequant(A, AQ) @ B) and ABQuant (C = dequant(A, AQ) @ dequant(B, BQ)) dispatcher bridges following the same three-layer pattern (ctypes C API → Python dispatcher → codegen). Each mode adds: *_ctypes_lib.cpp — C API exposing the kernel to Python via ctypes unified_*_codegen.py — generates kernel instantiations across the full dtype/layout/quant-group matrix *_utils.py — Python dispatcher handling argument validation, dtype selection, and parallel hipcc builds via ThreadPoolExecutor Example script and CPU-only unit tests codegen_common.py is extended with shared epilogue selection helpers reused across all three quant modes. CMakeLists.txt updated with build targets for both new ctypes shared libraries. Supported dtypes: fp8/bf8 activations, fp8i4/bf8i4 packed int4, preshuffle layouts, and MX microscaling variants — targeting gfx950 (MI350X). ## Test Plan CPU-only unit tests covering: Kernel name generation across all dtype/layout/quant-group combinations Config serialization and dimension helpers Epilogue selection logic for AQuant and ABQuant variants End-to-end example scripts (14_grouped_gemm_aquant.py, 15_grouped_gemm_abquant.py) with CPU reference verification for correctness checks. ## Test Result CPU unit tests pass for both AQuant and ABQuant dispatcher utilities. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.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.
Motivation
This PR extends the CK tile dispatcher with full quantization coverage for grouped GEMM operations. Building on the BQuant dispatcher introduced in #9166, this PR adds dispatcher bridges for AQuant (A-side quantization) and ABQuant (both sides quantized), completing the three-mode quantization dispatch surface for grouped GEMM on gfx950 (MI350X).
Technical Details
Stacked on #9166 (BQuant grouped GEMM dispatcher). Adds AQuant (C = dequant(A, AQ) @ B) and ABQuant (C = dequant(A, AQ) @ dequant(B, BQ)) dispatcher bridges following the same three-layer pattern (ctypes C API → Python dispatcher → codegen).
Each mode adds:
ctypes_lib.cpp — C API exposing the kernel to Python via ctypes
unified_codegen.py — generates kernel instantiations across the full dtype/layout/quant-group matrix
*_utils.py — Python dispatcher handling argument validation, dtype selection, and parallel hipcc builds via ThreadPoolExecutor
Example script and CPU-only unit tests
codegen_common.py is extended with shared epilogue selection helpers reused across all three quant modes. CMakeLists.txt updated with build targets for both new ctypes shared libraries.
Supported dtypes: fp8/bf8 activations, fp8i4/bf8i4 packed int4, preshuffle layouts, and MX microscaling variants — targeting gfx950 (MI350X).
Test Plan
CPU-only unit tests covering:
Kernel name generation across all dtype/layout/quant-group combinations
Config serialization and dimension helpers
Epilogue selection logic for AQuant and ABQuant variants
End-to-end example scripts (14_grouped_gemm_aquant.py, 15_grouped_gemm_abquant.py) with CPU reference verification for correctness checks.
Test Result
CPU unit tests pass for both AQuant and ABQuant dispatcher utilities.
Submission Checklist