SIGN IN SIGN UP

[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