[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.

Author: ThruptiRajLakshmanaGowda

dispatcher/bindings/ctypes/CMakeLists.txt

@@ -154,6 +154,52 @@ else()
     )
 endif()
 
+# =============================================================================
+# GroupedGemm BQuant ctypes Library
+# =============================================================================
+
+file(GLOB BQUANT_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/grouped_gemm_bquant_*.hpp")
+if(BQUANT_GEMM_KERNEL_HEADERS)
+    list(GET BQUANT_GEMM_KERNEL_HEADERS 0 BQUANT_GEMM_KERNEL_HEADER)
+    message(STATUS "Found BQuant GEMM kernel for ctypes lib: ${BQUANT_GEMM_KERNEL_HEADER}")
+
+    add_ctypes_library(dispatcher_grouped_gemm_bquant_lib
+        grouped_gemm_bquant_ctypes_lib.cpp
+        KERNEL_HEADER ${BQUANT_GEMM_KERNEL_HEADER}
+    )
+    # Resolve a single GFX arch string from CMAKE_HIP_ARCHITECTURES if available,
+    # falling back to the cache variable CK_TILE_BQUANT_GFX_ARCH (default gfx942).
+    # The Python build path always passes -DGFX_ARCH=... explicitly via hipcc, so
+    # this default only affects builds driven through CMake directly.
+    if(NOT DEFINED CK_TILE_BQUANT_GFX_ARCH OR CK_TILE_BQUANT_GFX_ARCH STREQUAL "")
+        if(CMAKE_HIP_ARCHITECTURES)
+            list(GET CMAKE_HIP_ARCHITECTURES 0 _bquant_gfx_arch)
+        else()
+            set(_bquant_gfx_arch "gfx942")
+        endif()
+        set(CK_TILE_BQUANT_GFX_ARCH "${_bquant_gfx_arch}" CACHE STRING
+            "GFX arch for the CMake-built BQuant ctypes .so (default: first CMAKE_HIP_ARCHITECTURES or gfx942)")
+    endif()
+    message(STATUS "BQuant ctypes lib GFX_ARCH: ${CK_TILE_BQUANT_GFX_ARCH}")
+
+    target_compile_definitions(dispatcher_grouped_gemm_bquant_lib PRIVATE
+        CK_TILE_SINGLE_KERNEL_INCLUDE
+        GFX_ARCH="${CK_TILE_BQUANT_GFX_ARCH}"
+    )
+else()
+    message(STATUS "No BQuant GEMM kernel found for ctypes lib - building without kernel")
+    add_library(dispatcher_grouped_gemm_bquant_lib SHARED grouped_gemm_bquant_ctypes_lib.cpp)
+    target_include_directories(dispatcher_grouped_gemm_bquant_lib PRIVATE
+        ${PROJECT_SOURCE_DIR}/include
+        ${PROJECT_SOURCE_DIR}/dispatcher/include
+    )
+    target_link_libraries(dispatcher_grouped_gemm_bquant_lib PRIVATE hip::device)
+    set_target_properties(dispatcher_grouped_gemm_bquant_lib PROPERTIES
+        POSITION_INDEPENDENT_CODE ON
+        CXX_STANDARD 17
+    )
+endif()
+
 # =============================================================================
 # GPU Helper Executable
 # =============================================================================