[rocm-libraries] ROCm/rocm-libraries#9166 (commit 61284f3)
feat(ck-tile): add BQuantGrouped GEMM dispatcher with ctypes bridge (#9166) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIRA ID : AICK-1289 ## Motivation The CK tile dispatcher lacked a Python-callable path for BQuantGrouped GEMM, which is the primary quantized GEMM kernel used for weight-only quantization inference on gfx950 (MI350X). This PR adds a complete three-layer bridge (codegen → compile → dispatch) so that Python callers can drive BQuant GEMM through the dispatcher without writing any C++ directly. The implementation is based on the reference kernels in `example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_*` and covers all variants used in production: fp8/bf8 decode and prefill tiles, fp8i4/bf8i4 weight-only quantization, preshuffle variants (preshuffle_b, preshuffle_bquant, both combined), and MX microscale (bf16+bf16, bf16+bf8, bf16+fp4 with e8m0 block scale). ## Technical Details **New files added:** - `dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp` — C API (`dispatcher_initialize`, `dispatcher_run_bquant_gemm`) for Python ctypes integration. Manages host↔device memory transfers internally. Compiled per-kernel via `hipcc -include <kernel.hpp>` force-include pattern. - `dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py` — kernel header generator. Produces one `.hpp` per config, selecting `PermuteNEpilogue` vs `CShuffleEpilogue` based on `TiledMMAPermuteN` (mirrors `run_gemm_quant_example.inc`). - `dispatcher/python/grouped_gemm_bquant_utils.py` — Python bridge: `BQuantKernelConfig` describes a kernel, `BQuantDispatcherLib` wraps the ctypes API, `BQuantGpuGemmRunner` accepts NumPy arrays and returns results. Build pipeline runs codegen + hipcc in parallel via `ThreadPoolExecutor`. - `dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py` — end-to-end example demonstrating codegen → hipcc compile → GPU run → CPU reference verification using real fp8 bit patterns via `ml_dtypes`. - `dispatcher/tests/test_grouped_gemm_bquant_utils.py` — 58 CPU-only unit tests covering kernel name generation, config serialization round-trips, and problem dimension helpers for all variant families. **Modified files:** - `include/ck_tile/ops/gemm_quant/pipeline/gemm_abquant_pipeline_ag_bg_cr_policy.hpp` — adds `wg_attr_num_access` selection logic to `WarpGemmDispatcher` instantiation in `GemmABQuantPipelineAgBgCrDefaultPolicy`, required for ABQuant pipeline correctness on prefill tiles. **Shared infrastructure:** - `dispatcher/codegen/codegen_common.py` — adds `make_bquant_kernel_name()` so both the codegen (`BQuantKernelSpec.name`) and the utils (`BQuantKernelConfig.name`) produce byte-exact identical kernel names from a single implementation. ## Test Plan - Run the 58 CPU-only unit tests (no GPU required): python3 -m pytest dispatcher/tests/test_grouped_gemm_bquant_utils.py -v - Run the end-to-end example on gfx950 (MI350X): python3 dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py --dtype fp8 python3 dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py --dtype bf8 - Build and run `test_gemm_quant_abquant_preshuffle_preshuffleQuant` on gfx950 and gfx120 to verify the ABQuant pipeline policy change. ## Test Result - All 58 CPU-only unit tests pass. - End-to-end GPU tests pending on gfx950 hardware. - `test_gemm_quant_abquant_preshuffle_preshuffleQuant` build and run pending. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
T
Thrupti Raj Lakshmana Gowda committed
5305cddddcb2dee9c3fb8a4f87386b651841b772
Parent: e6d9916
Committed by assistant-librarian[bot] <assistant-librarian[bot]@users.noreply.github.com>
on 7/22/2026, 7:07:07 PM