From 56767acda486fe93a9bc5fb8ac2eec24428afa85 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Mon, 6 Jul 2026 19:45:04 -0400 Subject: [PATCH 01/35] [CK][Dispatcher] Add BQuantGrouped GEMM dispatcher bridge (fp8/bf8, compv3, gfx950) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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>. - 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 --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 31 + .../ctypes/bquant_gemm_ctypes_lib.cpp | 265 +++++++ .../dispatcher/codegen/codegen_common.py | 19 + .../codegen/unified_bquant_gemm_codegen.py | 560 +++++++++++++++ .../examples/gemm/python/13_bquant_gemm.py | 214 ++++++ .../dispatcher/python/bquant_gemm_utils.py | 648 ++++++++++++++++++ .../tests/test_bquant_gemm_utils.py | 245 +++++++ 7 files changed, 1982 insertions(+) create mode 100644 projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp create mode 100644 projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py create mode 100644 projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py create mode 100644 projects/composablekernel/dispatcher/python/bquant_gemm_utils.py create mode 100644 projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 18314017f225..300c9f46b74c 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -154,6 +154,37 @@ else() ) endif() +# ============================================================================= +# BQuantGrouped GEMM ctypes Library +# ============================================================================= + +file(GLOB BQUANT_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/bquant_gemm_*.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_bquant_gemm_lib + bquant_gemm_ctypes_lib.cpp + KERNEL_HEADER ${BQUANT_GEMM_KERNEL_HEADER} + ) + target_compile_definitions(dispatcher_bquant_gemm_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE + GFX_ARCH="gfx950" + ) +else() + message(STATUS "No BQuant GEMM kernel found for ctypes lib - building without kernel") + add_library(dispatcher_bquant_gemm_lib SHARED bquant_gemm_ctypes_lib.cpp) + target_include_directories(dispatcher_bquant_gemm_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ${PROJECT_SOURCE_DIR}/dispatcher/include + ) + target_link_libraries(dispatcher_bquant_gemm_lib PRIVATE hip::device) + set_target_properties(dispatcher_bquant_gemm_lib PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD 17 + ) +endif() + # ============================================================================= # GPU Helper Executable # ============================================================================= diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp new file mode 100644 index 000000000000..cffeb1d59c33 --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp @@ -0,0 +1,265 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * BQuantGrouped GEMM Dispatcher ctypes Library + * + * Provides C API for Python ctypes integration. + * Kernel header is force-included at compile time via: + * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE bquant_gemm_ctypes_lib.cpp + * + * Force-include defines (from generated kernel header): + * SelectedKernel, KERNEL_NAME + * ADataType, BDataType, CDataType, QDataType, AccDataType, QuantGroupSize + * + * Memory model: host-pointer (this library owns hipMalloc/hipMemcpy/hipFree). + * The Python caller passes host numpy arrays; device memory is managed here. + * + * Pattern: follows current gemm_ctypes_lib.cpp (no GEMM_KEY_* macros). + */ + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/dispatcher/dispatcher.hpp" +#include "ck_tile/dispatcher/registry.hpp" +#include "ck_tile/dispatcher/backends/generated_tile_backend.hpp" + +// Kernel header force-included via -include compiler flag. +// Defines: ADataType, BDataType, CDataType, QDataType, AccDataType, +// QuantGroupSize, SelectedKernel, KERNEL_NAME + +#ifndef GFX_ARCH +#define GFX_ARCH "gfx950" +#endif + +using namespace ck_tile::dispatcher; +using namespace ck_tile::dispatcher::backends; +using Priority = ck_tile::dispatcher::Registry::Priority; + +static std::shared_ptr g_dispatcher = nullptr; +static bool g_initialized = false; + +#define HIP_CHECK(call) \ + { \ + hipError_t _err = (call); \ + if(_err != hipSuccess) \ + { \ + std::cerr << "HIP error: " << hipGetErrorString(_err) \ + << " at " << __FILE__ << ":" << __LINE__ << "\n"; \ + return -1; \ + } \ + } + +extern "C" { + +/** + * Initialize dispatcher — must be called before dispatcher_run_bquant_gemm. + * + * Registers SelectedKernel (from the force-included header) into the Registry. + * Returns 0 on success, -1 on error. + */ +int dispatcher_initialize() +{ + if(g_initialized) + return 0; + + // Build a minimal KernelKey so the kernel is registered and selectable. + // For BQuant the key is used only for registration; actual dispatch goes + // through SelectedKernel::launch() directly via dispatcher_run_bquant_gemm. + KernelKey key; + key.signature.dtype_a = DataType::FP8; + key.signature.dtype_b = DataType::FP8; + key.signature.dtype_c = DataType::FP16; + key.signature.dtype_acc = DataType::FP32; + key.signature.layout_a = LayoutTag::RowMajor; + key.signature.layout_b = LayoutTag::ColMajor; + key.signature.layout_c = LayoutTag::RowMajor; + key.signature.transpose_a = false; + key.signature.transpose_b = false; + key.signature.grouped = false; + key.signature.split_k = 1; + key.signature.elementwise_op = "PassThrough"; + key.signature.num_d_tensors = 0; + key.signature.structured_sparsity = false; + + key.algorithm.tile_shape = {SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; + key.algorithm.wave_shape = {SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; + key.algorithm.warp_tile_shape = {SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; + key.algorithm.pipeline = Pipeline::CompV3; + key.algorithm.scheduler = Scheduler::Intrawave; + key.algorithm.epilogue = Epilogue::CShuffle; + key.algorithm.block_size = SelectedKernel::BlockSize; + key.algorithm.double_buffer = false; + key.algorithm.persistent = false; + key.algorithm.preshuffle = SelectedKernel::PreshuffleB; + key.algorithm.transpose_c = false; + key.algorithm.num_wave_groups = 1; + key.gfx_arch = GFX_ARCH; + + auto kernel = + create_generated_tile_kernel( + key, KERNEL_NAME); + + Registry::instance().clear(); + Registry::instance().register_kernel(kernel, Priority::High); + + g_dispatcher = std::make_shared(); + g_initialized = true; + return 0; +} + +/** + * Run BQuantGrouped GEMM: C[M,N] = A[M,K] @ dequant(B[K,N], BQ[ceil(K/gK), ceil(N/gN)]) + * + * A, B, BQ, C are host pointers. This function manages device memory internally. + * + * Parameters: + * A, B, BQ, C - host data pointers + * M, N, K - matrix dimensions + * stride_A - leading dimension of A (row-major: K; col-major: M) + * stride_B - leading dimension of B (col-major: K; row-major: N) + * stride_BQ - leading dimension of BQ (row-major: ceil(N/gN)) + * stride_C - leading dimension of C (row-major: N) + * QK_B - number of K-groups = ceil(K / quant_group_k) + * QN_B - number of N-groups = ceil(N / quant_group_n) + * k_batch - split-K factor (1 = no split) + * time_ms - output: kernel execution time in ms (may be NULL) + * + * Returns 0 on success, negative on error. + */ +int dispatcher_run_bquant_gemm( + const void* A, + const void* B, + const void* BQ, + void* C, + int64_t M, + int64_t N, + int64_t K, + int64_t stride_A, + int64_t stride_B, + int64_t stride_BQ, + int64_t stride_C, + int64_t QK_B, + int64_t QN_B, + int k_batch, + float* time_ms) +{ + if(!g_initialized) + { + std::cerr << "dispatcher_run_bquant_gemm: not initialized\n"; + return -1; + } + if(!A || !B || !BQ || !C) + { + std::cerr << "dispatcher_run_bquant_gemm: null pointer argument\n"; + return -1; + } + if(M <= 0 || N <= 0 || K <= 0 || QK_B <= 0 || QN_B <= 0) + { + std::cerr << "dispatcher_run_bquant_gemm: invalid dimensions\n"; + return -1; + } + + const ADataType* A_host = static_cast(A); + const BDataType* B_host = static_cast(B); + const QDataType* BQ_host = static_cast(BQ); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + BDataType* B_dev = nullptr; + QDataType* BQ_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup = [&]() { + if(A_dev) (void)hipFree(A_dev); + if(B_dev) (void)hipFree(B_dev); + if(BQ_dev) (void)hipFree(BQ_dev); + if(C_dev) (void)hipFree(C_dev); + }; + + // Allocate device buffers + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + + // Copy inputs to device + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + + // Build QuantGemmHostArgs (aq_ptr = nullptr, QK_A = 0, stride_AQ = 0 for BQuant-only) + ck_tile::QuantGemmHostArgs args; + args.a_ptr = A_dev; + args.b_ptr = B_dev; + args.aq_ptr = nullptr; + args.bq_ptr = BQ_dev; + args.c_ptr = C_dev; + args.k_batch = k_batch; + args.M = static_cast(M); + args.N = static_cast(N); + args.K = static_cast(K); + args.QK_A = 0; + args.QK_B = static_cast(QK_B); + args.stride_A = static_cast(stride_A); + args.stride_B = static_cast(stride_B); + args.stride_C = static_cast(stride_C); + args.stride_AQ = 0; + args.stride_BQ = static_cast(stride_BQ); + + ck_tile::stream_config stream_cfg{nullptr, true, 0, 1, 1, false, false, 1}; + + float exec_time = SelectedKernel::launch(args, stream_cfg); + + if(exec_time < 0.0f) + { + std::cerr << "dispatcher_run_bquant_gemm: kernel reported unsupported args\n"; + cleanup(); + return -2; + } + + // Copy result back + if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + { + cleanup(); + return -1; + } + + if(time_ms) + *time_ms = exec_time; + + cleanup(); + return 0; +} + +/** + * Return the compile-time KERNEL_NAME of the force-included kernel. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Initialize dispatcher (alias kept for consistency with gemm_ctypes_lib). + */ +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Number of kernels registered (always 1 for the single-kernel-per-.so model). + */ +int dispatcher_get_kernel_count() { return static_cast(Registry::instance().size()); } + +/** + * Release dispatcher resources. + */ +void dispatcher_cleanup() +{ + g_dispatcher.reset(); + g_initialized = false; +} + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index a0486da66d74..1b0858128137 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -348,3 +348,22 @@ def needs_warp_expansion(config: dict) -> bool: def needs_pipeline_expansion(config: dict) -> bool: """True if pipeline is a wildcard (\"*\").""" return config.get("pipeline", "compv4") == "*" + + +# ============================================================================ +# BQuant-specific Type Mappings +# ============================================================================ + +# CK qualified type names for BQuant dtype fields (A, B, C, Q). +# Used by unified_bquant_gemm_codegen.py. Kept here so future AQuant/ABQuant +# codegens can reuse the same map without duplication. +BQUANT_DTYPE_MAP = { + "fp8": "ck_tile::fp8_t", + "bf8": "ck_tile::bf8_t", + "pk_int4": "ck_tile::pk_int4_t", + "pk_fp4": "ck_tile::pk_fp4_t", + "half": "ck_tile::half_t", + "bf16": "ck_tile::bf16_t", + "float": "float", + "e8m0": "ck_tile::e8m0_t", +} diff --git a/projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py new file mode 100644 index 000000000000..0d378dbc3874 --- /dev/null +++ b/projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py @@ -0,0 +1,560 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +BQuantGrouped GEMM Code Generator + +Generates one .hpp per kernel config for the dispatcher's ctypes path. +Each header defines a SelectedKernel struct with a static launch() method +taking QuantGemmHostArgs — compiled per-kernel via force-include: + + hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE bquant_gemm_ctypes_lib.cpp + +Initial scope: fp8 and bf8 dtype variants, non-preshuffle, compv3 pipeline, +rcr layout, configurable QuantGroupShape. + +Naming convention (byte-exact with BQuantKernelConfig.name in bquant_gemm_utils.py): + bquant_gemm_{dtype_a}_{layout}_{pipeline}_{epilogue}_{scheduler}_ + {TileM}x{TileN}x{TileK}_{WarpM}x{WarpN}x{WarpK}_{WtM}x{WtN}x{WtK}_ + qg{gM}x{gN}x{gK}[_preshuffleb][_preshufflebq] + +Reference: + example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8.cpp + example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc + example/ck_tile/38_block_scale_gemm/gemm_utils.hpp (GemmConfigQuantDecode) +""" + +import argparse +import itertools +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from codegen_common import TileConfig, parallel_generate, BQUANT_DTYPE_MAP + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# Dtype variant definitions +# Each entry: (dtype_key, ADataType, BDataType, CDataType, QDataType) +# Matches example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_*.cpp +# ============================================================================= + +BQUANT_VARIANTS: Dict[str, Dict[str, str]] = { + "fp8": { + "dtype_a": "fp8", + "dtype_b": "fp8", + "dtype_c": "half", + "dtype_q": "float", + "ck_a": "ck_tile::fp8_t", + "ck_b": "ck_tile::fp8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + }, + "bf8": { + "dtype_a": "bf8", + "dtype_b": "bf8", + "dtype_c": "half", + "dtype_q": "float", + "ck_a": "ck_tile::bf8_t", + "ck_b": "ck_tile::bf8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + }, +} + +# Layout strings supported: only rcr for initial implementation +# (RowMajor A, ColMajor B, RowMajor C) — standard GEMM layout for quant kernels +BQUANT_LAYOUT_TO_CK = { + "r": "ck_tile::tensor_layout::gemm::RowMajor", + "c": "ck_tile::tensor_layout::gemm::ColumnMajor", +} + +# Pipeline map for BQuant (compv3 for initial; preshuffle/microscale are follow-ons) +BQUANT_PIPELINE_MAP = { + "compv3": "ck_tile::BQuantGemmPipelineAgBgCrCompV3", + # preshuffle_b -> WPQuantBPipelineAgBgCrV2 (follow-on) + # mx -> MicroscaleGemmPipelineAgBgCrCompV3 (follow-on) +} + +BQUANT_BASE_PIPELINE_MAP = { + "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", +} + +BQUANT_SCHEDULER_TO_CK = { + "intrawave": "ck_tile::GemmPipelineScheduler::Intrawave", + "interwave": "ck_tile::GemmPipelineScheduler::Interwave", +} + + +# ============================================================================= +# Configuration dataclasses +# ============================================================================= + + +@dataclass +class BQuantTileConfig: + tile_m: int + tile_n: int + tile_k: int + warp_m: int + warp_n: int + warp_k: int + warp_tile_m: int + warp_tile_n: int + warp_tile_k: int + + def is_valid(self) -> bool: + if self.tile_m <= 0 or self.tile_n <= 0 or self.tile_k <= 0: + return False + return ( + self.tile_m % (self.warp_m * self.warp_tile_m) == 0 + and self.tile_n % (self.warp_n * self.warp_tile_n) == 0 + and self.tile_k % (self.warp_k * self.warp_tile_k) == 0 + ) + + +@dataclass +class BQuantKernelSpec: + """Complete specification for one BQuant kernel.""" + + variant_key: str # "fp8" or "bf8" + layout: str # "rcr" + pipeline: str # "compv3" + epilogue: str # "cshuffle" + scheduler: str # "intrawave" + tile: BQuantTileConfig + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + preshuffle_b: bool = False + preshuffle_bquant: bool = False + pad_m: bool = False + pad_n: bool = False + pad_k: bool = True + block_size: int = 256 + k_block_per_cu: int = 1 + + @property + def name(self) -> str: + t = self.tile + parts = [ + "bquant_gemm", + self.variant_key, + self.layout, + self.pipeline, + self.epilogue, + self.scheduler, + f"{t.tile_m}x{t.tile_n}x{t.tile_k}", + f"{t.warp_m}x{t.warp_n}x{t.warp_k}", + f"{t.warp_tile_m}x{t.warp_tile_n}x{t.warp_tile_k}", + f"qg{self.quant_group_m}x{self.quant_group_n}x{self.quant_group_k}", + ] + if self.preshuffle_b: + parts.append("preshuffleb") + if self.preshuffle_bquant: + parts.append("preshufflebq") + return "_".join(parts) + + +# ============================================================================= +# Header generator +# ============================================================================= + + +class BQuantKernelHeaderGenerator: + """Generates a .hpp kernel specialization header for one BQuantKernelSpec.""" + + def generate(self, spec: BQuantKernelSpec) -> str: + variant = BQUANT_VARIANTS[spec.variant_key] + t = spec.tile + ns = "ns_" + spec.name + struct = "Kernel_" + spec.name + + ck_a = variant["ck_a"] + ck_b = variant["ck_b"] + ck_c = variant["ck_c"] + ck_q = variant["ck_q"] + ck_acc = variant["ck_acc"] + + layout_a_ck = BQUANT_LAYOUT_TO_CK[spec.layout[0]] + layout_b_ck = BQUANT_LAYOUT_TO_CK[spec.layout[1]] + layout_c_ck = BQUANT_LAYOUT_TO_CK[spec.layout[2]] + # BQ is always RowMajor: scales are stored [ceil(K/gK), ceil(N/gN)] + layout_bq_ck = BQUANT_LAYOUT_TO_CK["r"] + # AQ layout placeholder (unused for BQuant-only, same as A layout) + layout_aq_ck = layout_a_ck + + pipeline_ck = BQUANT_PIPELINE_MAP[spec.pipeline] + base_pipeline_ck = BQUANT_BASE_PIPELINE_MAP[spec.pipeline] + scheduler_ck = BQUANT_SCHEDULER_TO_CK[spec.scheduler] + + pad_m = str(spec.pad_m).lower() + pad_n = str(spec.pad_n).lower() + pad_k = str(spec.pad_k).lower() + preshuffle_b = str(spec.preshuffle_b).lower() + preshuffle_bquant = str(spec.preshuffle_bquant).lower() + + # CShuffleEpilogue problem parameters + # TiledMMAPermuteN is false when quant_group_n > 1 (see run_gemm_quant_example.inc) + tiled_permute_n = "false" + + return f"""\ +// SPDX-License-Identifier: MIT +// Auto-generated BQuantGrouped GEMM kernel header. +// DO NOT EDIT — regenerate via unified_bquant_gemm_codegen.py +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "ck_tile/ops/epilogue.hpp" + +namespace {ns} {{ + +constexpr const char* KERNEL_NAME = "{spec.name}"; + +using ADataType = {ck_a}; +using BDataType = {ck_b}; +using CDataType = {ck_c}; +using QDataType = {ck_q}; +using AccDataType = {ck_acc}; + +using ALayout = {layout_a_ck}; +using BLayout = {layout_b_ck}; +using CLayout = {layout_c_ck}; +using AQLayout = {layout_aq_ck}; +using BQLayout = {layout_bq_ck}; + +// Single QuantGroupSize alias — same type used for both AQ and BQ slots in the +// pipeline template; AQ is disabled via aq_ptr=nullptr at runtime for BQuant-only. +using QuantGroupSize = ck_tile::QuantGroupShape>; + +struct {struct} {{ + using ADataType = {ns}::ADataType; + using BDataType = {ns}::BDataType; + using CDataType = {ns}::CDataType; + using QDataType = {ns}::QDataType; + using AccDataType = {ns}::AccDataType; + + static constexpr ck_tile::index_t TileM = {t.tile_m}; + static constexpr ck_tile::index_t TileN = {t.tile_n}; + static constexpr ck_tile::index_t TileK = {t.tile_k}; + static constexpr ck_tile::index_t WarpM = {t.warp_m}; + static constexpr ck_tile::index_t WarpN = {t.warp_n}; + static constexpr ck_tile::index_t WarpK = {t.warp_k}; + static constexpr ck_tile::index_t WarpTileM = {t.warp_tile_m}; + static constexpr ck_tile::index_t WarpTileN = {t.warp_tile_n}; + static constexpr ck_tile::index_t WarpTileK = {t.warp_tile_k}; + static constexpr ck_tile::index_t BlockSize = {spec.block_size}; + static constexpr int kBlockPerCu = {spec.k_block_per_cu}; + + static constexpr bool kPadM = {pad_m}; + static constexpr bool kPadN = {pad_n}; + static constexpr bool kPadK = {pad_k}; + static constexpr bool APreshuffleQuant = false; + static constexpr bool BPreshuffleQuant = {preshuffle_bquant}; + static constexpr bool PreshuffleB = {preshuffle_b}; + static constexpr bool TransposeC = false; + static constexpr bool DoubleSmemBuffer = false; + + using TileShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + + using GemmTraits = ck_tile::TileGemmQuantTraits< + kPadM, kPadN, kPadK, + APreshuffleQuant, BPreshuffleQuant, PreshuffleB, + {ns}::ALayout, {ns}::BLayout, {ns}::CLayout, + ck_tile::QuantType::BQuantGrouped, + {ns}::AQLayout, {ns}::BQLayout, + TransposeC, DoubleSmemBuffer>; + + using GemmPipelineProblemBase = ck_tile::GemmPipelineProblemBase< + ADataType, BDataType, AccDataType, TileShape, GemmTraits>; + + using BaseGemmPipeline = {base_pipeline_ck}; + + static float launch(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + {{ + // hot-loop / tail dispatch — mirrors run_gemm_quant_example.inc + const ck_tile::index_t K_split = + (args.k_batch == 1) + ? ck_tile::integer_least_multiple(args.K, TileK) + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, TileK); + + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](auto has_hot_loop_, auto tail_number_) {{ + using PipelineProblem = ck_tile::GemmBQuantPipelineProblem< + ADataType, + BDataType, + QDataType, + AccDataType, + TileShape, + GemmTraits, + QuantGroupSize, + ADataType, // ComputeDataType + {scheduler_ck}, + has_hot_loop_.value, + tail_number_.value>; + + using GemmPipeline = {pipeline_ck}; + + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC>>; + + using Kernel = ck_tile::QuantGemmKernel< + TilePartitioner, GemmPipeline, GemmEpilogue, + ck_tile::QuantType::BQuantGrouped>; + + auto kargs = Kernel::MakeKernelArgs(args); + if(!Kernel::IsSupportedArgument(kargs)) + return -1.0f; + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + return ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{{}}, grids, blocks, 0, kargs)); + }}; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + }} +}}; + +using SelectedKernel = {struct}; + +}} // namespace {ns} + +#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE +using SelectedKernel = {ns}::{struct}; +constexpr const char* KERNEL_NAME = {ns}::KERNEL_NAME; +using ADataType = {ck_a}; +using BDataType = {ck_b}; +using CDataType = {ck_c}; +using QDataType = {ck_q}; +using AccDataType = {ck_acc}; +using QuantGroupSize = {ns}::QuantGroupSize; +#endif // CK_TILE_SINGLE_KERNEL_INCLUDE +""" + + +# ============================================================================= +# Config sweep +# ============================================================================= + + +def _default_config() -> dict: + """Default sweep config matching GemmConfigQuantDecode tile defaults.""" + return { + "variant_keys": ["fp8", "bf8"], + "layouts": ["rcr"], + "pipeline": "compv3", + "epilogue": "cshuffle", + "scheduler": "intrawave", + "tile_configs": [ + # GemmConfigQuantDecode: M=16, N=64, K=256/sizeof(fp8_t)=256 + {"tile_m": 16, "tile_n": 64, "tile_k": 256, + "warp_m": 1, "warp_n": 4, "warp_k": 1, + "warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 16}, + ], + "quant_groups": [ + {"quant_group_m": 1, "quant_group_n": 1, "quant_group_k": 128}, + ], + "pad_m": False, + "pad_n": False, + "pad_k": True, + "block_size": 256, + "k_block_per_cu": 1, + } + + +def _build_specs(config: dict) -> List[BQuantKernelSpec]: + specs = [] + pipeline = config.get("pipeline", "compv3") + epilogue = config.get("epilogue", "cshuffle") + scheduler = config.get("scheduler", "intrawave") + pad_m = config.get("pad_m", False) + pad_n = config.get("pad_n", False) + pad_k = config.get("pad_k", True) + block_size = config.get("block_size", 256) + k_block_per_cu = config.get("k_block_per_cu", 1) + + for variant_key, layout, tile_dict, qg in itertools.product( + config.get("variant_keys", ["fp8"]), + config.get("layouts", ["rcr"]), + config.get("tile_configs", []), + config.get("quant_groups", [{"quant_group_m": 1, "quant_group_n": 1, "quant_group_k": 128}]), + ): + if variant_key not in BQUANT_VARIANTS: + log.warning("Unknown variant_key %s — skipping", variant_key) + continue + if pipeline not in BQUANT_PIPELINE_MAP: + log.warning("Unsupported pipeline %s — skipping", pipeline) + continue + + tile = BQuantTileConfig( + tile_m=tile_dict["tile_m"], + tile_n=tile_dict["tile_n"], + tile_k=tile_dict["tile_k"], + warp_m=tile_dict["warp_m"], + warp_n=tile_dict["warp_n"], + warp_k=tile_dict["warp_k"], + warp_tile_m=tile_dict["warp_tile_m"], + warp_tile_n=tile_dict["warp_tile_n"], + warp_tile_k=tile_dict["warp_tile_k"], + ) + if not tile.is_valid(): + log.debug("Invalid tile config %s — skipping", tile) + continue + + specs.append(BQuantKernelSpec( + variant_key=variant_key, + layout=layout, + pipeline=pipeline, + epilogue=epilogue, + scheduler=scheduler, + tile=tile, + quant_group_m=qg.get("quant_group_m", 1), + quant_group_n=qg.get("quant_group_n", 1), + quant_group_k=qg.get("quant_group_k", 128), + pad_m=pad_m, + pad_n=pad_n, + pad_k=pad_k, + block_size=block_size, + k_block_per_cu=k_block_per_cu, + )) + + return specs + + +# ============================================================================= +# Generation entry point +# ============================================================================= + + +def generate_kernels( + output_dir: Path, + config: Optional[dict] = None, + parallel: bool = True, +) -> List[Path]: + """Generate all BQuant kernel headers into output_dir. + + Returns list of generated .hpp paths. + """ + output_dir.mkdir(parents=True, exist_ok=True) + cfg = config or _default_config() + specs = _build_specs(cfg) + + if not specs: + log.warning("No kernel specs produced from config — check variant_keys and tile_configs") + return [] + + log.info("Generating %d BQuant kernel headers into %s", len(specs), output_dir) + + gen = BQuantKernelHeaderGenerator() + generated: List[Path] = [] + + def _generate_one(spec: BQuantKernelSpec) -> Path: + header = gen.generate(spec) + out_path = output_dir / f"{spec.name}.hpp" + out_path.write_text(header) + log.info(" wrote %s", out_path.name) + return out_path + + if parallel and len(specs) > 1: + import concurrent.futures + with concurrent.futures.ThreadPoolExecutor() as ex: + futures = {ex.submit(_generate_one, s): s for s in specs} + for fut in concurrent.futures.as_completed(futures): + try: + generated.append(fut.result()) + except Exception as e: + log.error("Failed generating %s: %s", futures[fut].name, e) + else: + for spec in specs: + try: + generated.append(_generate_one(spec)) + except Exception as e: + log.error("Failed generating %s: %s", spec.name, e) + + log.info("Generated %d / %d headers", len(generated), len(specs)) + return generated + + +# ============================================================================= +# CLI +# ============================================================================= + + +def main() -> int: + parser = argparse.ArgumentParser( + description="BQuantGrouped GEMM kernel header generator" + ) + parser.add_argument("--output-dir", type=Path, required=True, + help="Directory to write generated .hpp files") + parser.add_argument("--config", type=Path, + help="JSON config file (defaults to built-in sweep)") + parser.add_argument("--config-json", type=str, + help="Inline JSON config string") + parser.add_argument("--no-parallel", action="store_true", + help="Disable parallel generation") + parser.add_argument("--list-names", action="store_true", + help="Print kernel names that would be generated and exit") + args = parser.parse_args() + + cfg: Optional[dict] = None + if args.config_json: + try: + cfg = json.loads(args.config_json) + except json.JSONDecodeError as e: + log.error("Invalid --config-json: %s", e) + return 1 + elif args.config: + with open(args.config) as f: + cfg = json.load(f) + + if args.list_names: + specs = _build_specs(cfg or _default_config()) + for s in specs: + print(s.name) + return 0 + + paths = generate_kernels( + output_dir=args.output_dir, + config=cfg, + parallel=not args.no_parallel, + ) + return 0 if paths else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py b/projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py new file mode 100644 index 000000000000..3d1c6b91308d --- /dev/null +++ b/projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +Example 13 — BQuantGrouped GEMM via the Dispatcher + +Demonstrates the full three-layer path: + 1. Codegen — unified_bquant_gemm_codegen.py → .hpp + 2. Compile — hipcc → .so + 3. Run — BQuantGpuGemmRunner → C = A @ dequant(B, BQ) + +Verifies the GPU result against a NumPy fp32 reference. + +Requirements: + - gfx950 GPU (MI350X) + - hipcc in PATH + - CK include path discoverable relative to this repo + +Usage: + python3 13_bquant_gemm.py # fp8, 1x1x128 groups, M=16 N=64 K=256 + python3 13_bquant_gemm.py --dtype bf8 + python3 13_bquant_gemm.py --dtype fp8 --M 32 --N 128 --K 512 --quant-group-k 128 + python3 13_bquant_gemm.py --no-verify # skip CPU reference check +""" + +import argparse +import logging +import math +import sys +import tempfile +from pathlib import Path + +import numpy as np + +# Add dispatcher/python to path +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "python")) + +from bquant_gemm_utils import ( + BQuantKernelConfig, + BQuantGemmProblem, + BQuantGpuGemmRunner, + setup_multiple_bquant_dispatchers, + default_fp8_config, + default_bf8_config, +) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# NumPy reference: C = A @ dequant(B, BQ) +# ============================================================================= + + +def _fp8_to_float32(arr: np.ndarray) -> np.ndarray: + """Cast fp8 byte array to float32 (treating as float8_e4m3fn).""" + # numpy doesn't have fp8; cast uint8 bit pattern to float via ml_dtypes if available, + # otherwise fall back to float32 direct cast (close enough for reference at low values). + try: + import ml_dtypes + return arr.view(ml_dtypes.float8_e4m3fn).astype(np.float32) + except ImportError: + return arr.astype(np.float32) + + +def reference_bquant_gemm( + A: np.ndarray, + B: np.ndarray, + BQ: np.ndarray, + problem: BQuantGemmProblem, +) -> np.ndarray: + """ + CPU fp32 reference for C = A @ dequant(B, BQ). + + A [M, K] float32 (upcast from fp8/bf8) + B [K, N] float32 (upcast from fp8/bf8) + BQ [QK_B, QN_B] float32 scale factors + + Dequant: B[k, n] *= BQ[k // gK, n // gN] + """ + M, N, K = problem.M, problem.N, problem.K + gK = problem.quant_group_k + gN = problem.quant_group_n + + A_f32 = A.astype(np.float32) + B_f32 = B.astype(np.float32) + + # Apply per-block scales to B + B_dequant = B_f32.copy() + for qi in range(problem.QK_B): + for qj in range(problem.QN_B): + k_start = qi * gK + k_end = min(k_start + gK, K) + n_start = qj * gN + n_end = min(n_start + gN, N) + scale = float(BQ[qi, qj]) + B_dequant[k_start:k_end, n_start:n_end] *= scale + + C_ref = A_f32 @ B_dequant + return C_ref.astype(np.float16) + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + parser = argparse.ArgumentParser(description="BQuantGrouped GEMM dispatcher example") + parser.add_argument("--dtype", choices=["fp8", "bf8"], default="fp8") + parser.add_argument("--M", type=int, default=16) + parser.add_argument("--N", type=int, default=64) + parser.add_argument("--K", type=int, default=256) + parser.add_argument("--quant-group-k", type=int, default=128) + parser.add_argument("--quant-group-n", type=int, default=1) + parser.add_argument("--no-verify", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + parser.add_argument("--gfx-arch", type=str, default="gfx950") + args = parser.parse_args() + + M, N, K = args.M, args.N, args.K + gK = args.quant_group_k + gN = args.quant_group_n + + # ------------------------------------------------------------------------- + # 1. Build kernel config + # ------------------------------------------------------------------------- + if args.dtype == "fp8": + config = default_fp8_config(quant_group_k=gK, quant_group_n=gN, gfx_arch=args.gfx_arch) + else: + config = default_bf8_config(quant_group_k=gK, quant_group_n=gN, gfx_arch=args.gfx_arch) + + log.info("Kernel: %s", config.name) + + # ------------------------------------------------------------------------- + # 2. Codegen + compile + # ------------------------------------------------------------------------- + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="bquant_ex13_")) + log.info("Output dir: %s", out_dir) + + so_paths = setup_multiple_bquant_dispatchers( + configs=[config], + output_dir=out_dir, + gfx_arch=args.gfx_arch, + ) + + if not so_paths or so_paths[0] is None: + log.error("Kernel build failed — see errors above") + return 1 + + so_path = so_paths[0] + log.info("Built: %s", so_path) + + # ------------------------------------------------------------------------- + # 3. Generate inputs + # ------------------------------------------------------------------------- + rng = np.random.default_rng(42) + QK_B = math.ceil(K / gK) + QN_B = math.ceil(N / gN) + + # A and B as uint8 (fp8 byte representation) in [-2, 2] float range + # Using simple values so the reference is easy to verify + A_f32 = rng.uniform(-2.0, 2.0, (M, K)).astype(np.float32) + B_f32 = rng.uniform(-2.0, 2.0, (K, N)).astype(np.float32) + BQ_f32 = rng.uniform(0.5, 2.0, (QK_B, QN_B)).astype(np.float32) + + # Cast to fp8 approximation (uint8 for raw byte passing to C lib) + # Real fp8 encoding would use ml_dtypes; here we use float16 as a stand-in + # that the kernel can safely interpret as fp8 bit patterns for testing. + A_raw = A_f32.astype(np.float16) + B_raw = B_f32.astype(np.float16) + + # ------------------------------------------------------------------------- + # 4. Run on GPU + # ------------------------------------------------------------------------- + problem = BQuantGemmProblem(M=M, N=N, K=K, + quant_group_m=1, + quant_group_n=gN, + quant_group_k=gK) + + runner = BQuantGpuGemmRunner(so_path) + log.info("Running kernel: %s", runner.kernel_name) + + result = runner.run(A=A_raw, B=B_raw, BQ=BQ_f32, problem=problem) + log.info("Kernel time: %.3f ms", result.time_ms) + + # ------------------------------------------------------------------------- + # 5. Verify + # ------------------------------------------------------------------------- + if not args.no_verify: + C_ref = reference_bquant_gemm(A_f32, B_f32, BQ_f32, problem) + C_gpu = result.C + + max_rel = float(np.max(np.abs(C_gpu.astype(np.float32) - C_ref.astype(np.float32))) + / (np.max(np.abs(C_ref.astype(np.float32))) + 1e-6)) + + tolerance = 0.05 # fp8 ~1e-2 to 5e-2 + if max_rel <= tolerance: + log.info("PASSED (max_rel=%.4f, tol=%.4f)", max_rel, tolerance) + else: + log.error("FAILED (max_rel=%.4f > tol=%.4f)", max_rel, tolerance) + return 1 + else: + log.info("Verification skipped (--no-verify)") + + log.info("Example 13 complete.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/python/bquant_gemm_utils.py b/projects/composablekernel/dispatcher/python/bquant_gemm_utils.py new file mode 100644 index 000000000000..27fcbd922439 --- /dev/null +++ b/projects/composablekernel/dispatcher/python/bquant_gemm_utils.py @@ -0,0 +1,648 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +BQuantGrouped GEMM dispatcher utilities. + +Three-layer Python bridge for the dispatcher's BQuantGrouped GEMM path: + + BQuantKernelConfig — describes one kernel; .name is byte-exact with codegen KERNEL_NAME + BQuantDispatcherLib — thin ctypes wrapper around a compiled .so + BQuantGpuGemmRunner — high-level runner that accepts numpy arrays + +Build helpers (self-contained, do not import from gemm_utils.py): + setup_multiple_bquant_dispatchers(configs, ...) + codegen → hipcc → list of .so paths, all in parallel + +Usage (end-to-end): + configs = [BQuantKernelConfig(variant_key="fp8", layout="rcr", ...)] + so_paths = setup_multiple_bquant_dispatchers(configs, output_dir=Path("/tmp/bq")) + runner = BQuantGpuGemmRunner(so_paths[0]) + result = runner.run(A, B, BQ, BQuantGemmProblem(M=16, N=64, K=256)) +""" + +import ctypes +import json +import logging +import os +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Tuple +import concurrent.futures + +log = logging.getLogger(__name__) + +# ============================================================================= +# Constants +# ============================================================================= + +_CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_bquant_gemm_codegen.py" +_CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "bquant_gemm_ctypes_lib.cpp" + +_DEFAULT_HIPCC = "hipcc" +_DEFAULT_GFX_ARCH = "gfx950" + +# Flags that match the tile engine / dispatcher build flags for BQuant kernels +_HIPCC_BASE_FLAGS = [ + "-std=c++17", + "-O3", + "-fPIC", + "-shared", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", + "-w", # suppress warnings during generated-code compilation +] + + +# ============================================================================= +# BQuantKernelConfig — byte-exact naming with codegen +# ============================================================================= + + +@dataclass +class BQuantKernelConfig: + """ + Complete description of one BQuantGrouped GEMM kernel. + + The .name property produces the exact string that unified_bquant_gemm_codegen.py + emits as KERNEL_NAME, ensuring the Python side and compiled .so always agree. + """ + + variant_key: str # "fp8" or "bf8" + layout: str # "rcr" (A=RowMajor, B=ColMajor, C=RowMajor) + pipeline: str # "compv3" + epilogue: str # "cshuffle" + scheduler: str # "intrawave" + + tile_m: int + tile_n: int + tile_k: int + warp_m: int + warp_n: int + warp_k: int + warp_tile_m: int + warp_tile_n: int + warp_tile_k: int + + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + + preshuffle_b: bool = False + preshuffle_bquant: bool = False + + gfx_arch: str = _DEFAULT_GFX_ARCH + + @property + def name(self) -> str: + """Byte-exact match to codegen KERNEL_NAME.""" + parts = [ + "bquant_gemm", + self.variant_key, + self.layout, + self.pipeline, + self.epilogue, + self.scheduler, + f"{self.tile_m}x{self.tile_n}x{self.tile_k}", + f"{self.warp_m}x{self.warp_n}x{self.warp_k}", + f"{self.warp_tile_m}x{self.warp_tile_n}x{self.warp_tile_k}", + f"qg{self.quant_group_m}x{self.quant_group_n}x{self.quant_group_k}", + ] + if self.preshuffle_b: + parts.append("preshuffleb") + if self.preshuffle_bquant: + parts.append("preshufflebq") + return "_".join(parts) + + def to_codegen_config(self) -> dict: + """Produce the JSON config dict consumed by unified_bquant_gemm_codegen.py.""" + return { + "variant_keys": [self.variant_key], + "layouts": [self.layout], + "pipeline": self.pipeline, + "epilogue": self.epilogue, + "scheduler": self.scheduler, + "tile_configs": [{ + "tile_m": self.tile_m, + "tile_n": self.tile_n, + "tile_k": self.tile_k, + "warp_m": self.warp_m, + "warp_n": self.warp_n, + "warp_k": self.warp_k, + "warp_tile_m": self.warp_tile_m, + "warp_tile_n": self.warp_tile_n, + "warp_tile_k": self.warp_tile_k, + }], + "quant_groups": [{ + "quant_group_m": self.quant_group_m, + "quant_group_n": self.quant_group_n, + "quant_group_k": self.quant_group_k, + }], + } + + +# ============================================================================= +# BQuantGemmProblem +# ============================================================================= + + +@dataclass +class BQuantGemmProblem: + M: int + N: int + K: int + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + k_batch: int = 1 + + @property + def QK_B(self) -> int: + """Number of K-groups: ceil(K / quant_group_k).""" + return (self.K + self.quant_group_k - 1) // self.quant_group_k + + @property + def QN_B(self) -> int: + """Number of N-groups: ceil(N / quant_group_n).""" + return (self.N + self.quant_group_n - 1) // self.quant_group_n + + +# ============================================================================= +# BQuantGemmResult +# ============================================================================= + + +@dataclass +class BQuantGemmResult: + C: object # numpy array + time_ms: float + kernel_name: str + + +# ============================================================================= +# BQuantDispatcherLib — thin ctypes wrapper +# ============================================================================= + + +class BQuantDispatcherLib: + """ + Loads a compiled bquant_gemm .so and wraps its C API. + + Expected .so exports: + int dispatcher_initialize() + int dispatcher_run_bquant_gemm(A, B, BQ, C, M, N, K, + stride_A, stride_B, stride_BQ, stride_C, + QK_B, QN_B, k_batch, *time_ms) + char* dispatcher_get_kernel_name() + int dispatcher_get_kernel_count() + void dispatcher_cleanup() + """ + + def __init__(self, so_path: Path): + self.so_path = Path(so_path) + if not self.so_path.exists(): + raise FileNotFoundError(f"BQuant .so not found: {self.so_path}") + self._lib = ctypes.CDLL(str(self.so_path)) + self._setup() + rc = self._lib.dispatcher_initialize() + if rc != 0: + raise RuntimeError(f"dispatcher_initialize() returned {rc}") + + def _setup(self): + lib = self._lib + + lib.dispatcher_initialize.restype = ctypes.c_int + lib.dispatcher_initialize.argtypes = [] + + lib.dispatcher_run_bquant_gemm.restype = ctypes.c_int + lib.dispatcher_run_bquant_gemm.argtypes = [ + ctypes.c_void_p, # A + ctypes.c_void_p, # B + ctypes.c_void_p, # BQ + ctypes.c_void_p, # C + ctypes.c_int64, # M + ctypes.c_int64, # N + ctypes.c_int64, # K + ctypes.c_int64, # stride_A + ctypes.c_int64, # stride_B + ctypes.c_int64, # stride_BQ + ctypes.c_int64, # stride_C + ctypes.c_int64, # QK_B + ctypes.c_int64, # QN_B + ctypes.c_int, # k_batch + ctypes.POINTER(ctypes.c_float), # time_ms + ] + + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + lib.dispatcher_get_kernel_name.argtypes = [] + + lib.dispatcher_get_kernel_count.restype = ctypes.c_int + lib.dispatcher_get_kernel_count.argtypes = [] + + lib.dispatcher_cleanup.restype = None + lib.dispatcher_cleanup.argtypes = [] + + def run( + self, + A, + B, + BQ, + C, + M: int, + N: int, + K: int, + stride_A: int, + stride_B: int, + stride_BQ: int, + stride_C: int, + QK_B: int, + QN_B: int, + k_batch: int = 1, + ) -> Tuple[int, float]: + """ + Call dispatcher_run_bquant_gemm with ctypes-wrapped pointers. + + A, B, BQ, C should be numpy arrays (contiguous). + Returns (status, time_ms). + """ + import numpy as np + + A = np.ascontiguousarray(A) + B = np.ascontiguousarray(B) + BQ = np.ascontiguousarray(BQ) + C = np.ascontiguousarray(C) + + time_ms = ctypes.c_float(0.0) + + rc = self._lib.dispatcher_run_bquant_gemm( + A.ctypes.data_as(ctypes.c_void_p), + B.ctypes.data_as(ctypes.c_void_p), + BQ.ctypes.data_as(ctypes.c_void_p), + C.ctypes.data_as(ctypes.c_void_p), + ctypes.c_int64(M), + ctypes.c_int64(N), + ctypes.c_int64(K), + ctypes.c_int64(stride_A), + ctypes.c_int64(stride_B), + ctypes.c_int64(stride_BQ), + ctypes.c_int64(stride_C), + ctypes.c_int64(QK_B), + ctypes.c_int64(QN_B), + ctypes.c_int(k_batch), + ctypes.byref(time_ms), + ) + return rc, time_ms.value + + def get_kernel_name(self) -> str: + raw = self._lib.dispatcher_get_kernel_name() + return raw.decode("utf-8") if raw else "" + + def get_kernel_count(self) -> int: + return self._lib.dispatcher_get_kernel_count() + + def cleanup(self): + self._lib.dispatcher_cleanup() + + def __del__(self): + try: + self._lib.dispatcher_cleanup() + except Exception: + pass + + +# ============================================================================= +# BQuantGpuGemmRunner — high-level runner +# ============================================================================= + + +class BQuantGpuGemmRunner: + """ + High-level runner that loads a BQuant .so and executes GEMM on the GPU. + + Accepts numpy arrays for A, B, BQ; allocates C; returns BQuantGemmResult. + """ + + def __init__(self, so_path: Path): + self._lib = BQuantDispatcherLib(so_path) + + @property + def kernel_name(self) -> str: + return self._lib.get_kernel_name() + + def run(self, A, B, BQ, problem: BQuantGemmProblem) -> BQuantGemmResult: + """ + Run BQuantGrouped GEMM. + + A shape: (M, K) dtype: fp8/bf8 + B shape: (K, N) col-major dtype: fp8/bf8 + BQ shape: (QK_B, QN_B) dtype: float/fp8 + Returns BQuantGemmResult with C shape (M, N). + """ + import numpy as np + + M, N, K = problem.M, problem.N, problem.K + QK_B = problem.QK_B + QN_B = problem.QN_B + + # Output buffer — dtype matches CDataType (half for fp8/bf8 variants) + C = np.zeros((M, N), dtype=np.float16) + + # Strides (in elements, row-major for A and C; col-major for B means stride = K) + stride_A = K # A is row-major [M, K] + stride_B = K # B is col-major [K, N] → leading dim = K + stride_BQ = QN_B + stride_C = N # C is row-major [M, N] + + rc, time_ms = self._lib.run( + A=A, B=B, BQ=BQ, C=C, + M=M, N=N, K=K, + stride_A=stride_A, + stride_B=stride_B, + stride_BQ=stride_BQ, + stride_C=stride_C, + QK_B=QK_B, + QN_B=QN_B, + k_batch=problem.k_batch, + ) + + if rc != 0: + raise RuntimeError( + f"dispatcher_run_bquant_gemm failed with code {rc} " + f"for kernel {self.kernel_name}" + ) + + return BQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) + + +# ============================================================================= +# Subprocess helpers (self-contained, do not call ctypes_utils.py) +# ============================================================================= + + +def _detect_gpu_arch() -> str: + """Detect current GPU arch via rocm_agent_enumerator. Falls back to gfx950.""" + try: + result = subprocess.run( + ["rocm_agent_enumerator"], + capture_output=True, text=True, timeout=10, + ) + for line in result.stdout.splitlines(): + line = line.strip() + if line.startswith("gfx") and line != "gfx000": + return line + except Exception: + pass + return _DEFAULT_GFX_ARCH + + +def _get_ck_include_dir() -> Optional[Path]: + """Attempt to locate the CK include directory relative to this file.""" + # Walk up from dispatcher/python/ to find project root + here = Path(__file__).resolve().parent + for parent in [here.parent.parent, here.parent.parent.parent]: + candidate = parent / "include" + if (candidate / "ck_tile").is_dir(): + return candidate + return None + + +def _generate_bquant_kernel( + config: BQuantKernelConfig, + output_dir: Path, +) -> Optional[Path]: + """ + Run unified_bquant_gemm_codegen.py for one config; return the .hpp path or None. + """ + config_dict = config.to_codegen_config() + config_json = json.dumps(config_dict) + + cmd = [ + sys.executable, + str(_CODEGEN_SCRIPT), + "--output-dir", str(output_dir), + "--config-json", config_json, + ] + + try: + result = subprocess.run( + cmd, + capture_output=True, text=True, timeout=120, + ) + if result.returncode != 0: + log.error("Codegen failed for %s:\n%s", config.name, result.stderr) + return None + except subprocess.TimeoutExpired: + log.error("Codegen timed out for %s", config.name) + return None + + hpp = output_dir / f"{config.name}.hpp" + if not hpp.exists(): + log.error("Codegen succeeded but %s not found", hpp) + return None + + return hpp + + +def _compile_bquant_kernel( + hpp_path: Path, + so_path: Path, + gfx_arch: str, + hipcc: str = _DEFAULT_HIPCC, + extra_include_dirs: Optional[List[str]] = None, +) -> bool: + """ + Compile a generated .hpp into a .so via hipcc. + Returns True on success. + """ + ck_include = _get_ck_include_dir() + + cmd = [hipcc] + _HIPCC_BASE_FLAGS + [ + f"--offload-arch={gfx_arch}", + f"-DGFX_ARCH=\"{gfx_arch}\"", + f"-include", str(hpp_path), + str(_CTYPES_LIB_SRC), + "-o", str(so_path), + ] + + if ck_include: + cmd += [f"-I{ck_include}"] + + # Dispatcher include + dispatcher_include = _CTYPES_LIB_SRC.parent.parent.parent / "dispatcher" / "include" + if dispatcher_include.is_dir(): + cmd += [f"-I{dispatcher_include}"] + + if extra_include_dirs: + for d in extra_include_dirs: + cmd += [f"-I{d}"] + + log.debug("Compiling %s:\n %s", so_path.name, " ".join(cmd)) + + try: + result = subprocess.run( + cmd, + capture_output=True, text=True, timeout=600, + ) + if result.returncode != 0: + log.error("Compile failed for %s:\n%s", so_path.name, result.stderr[-2000:]) + return False + return True + except subprocess.TimeoutExpired: + log.error("Compile timed out for %s", so_path.name) + return False + + +# ============================================================================= +# setup_multiple_bquant_dispatchers — build pipeline +# ============================================================================= + + +def setup_multiple_bquant_dispatchers( + configs: List[BQuantKernelConfig], + output_dir: Optional[Path] = None, + hipcc: str = _DEFAULT_HIPCC, + gfx_arch: Optional[str] = None, + extra_include_dirs: Optional[List[str]] = None, + parallel: bool = True, + max_workers: Optional[int] = None, +) -> List[Optional[Path]]: + """ + For each BQuantKernelConfig: codegen → hipcc compile → .so path. + + Returns a list parallel to `configs` — each entry is the Path to the + compiled .so, or None if that config failed. + + No GPU is required to call this function. + """ + if not configs: + return [] + + arch = gfx_arch or _detect_gpu_arch() + base_dir = output_dir or Path(tempfile.mkdtemp(prefix="bquant_dispatcher_")) + base_dir.mkdir(parents=True, exist_ok=True) + + headers_dir = base_dir / "generated_kernels" + so_dir = base_dir / "libs" + headers_dir.mkdir(exist_ok=True) + so_dir.mkdir(exist_ok=True) + + log.info( + "Building %d BQuant kernel(s) for %s into %s", + len(configs), arch, base_dir, + ) + + # Deduplicate by name so we don't build the same kernel twice + seen: Dict[str, int] = {} # name → index of first occurrence + deduped: List[Tuple[int, BQuantKernelConfig]] = [] + for i, cfg in enumerate(configs): + if cfg.name not in seen: + seen[cfg.name] = i + deduped.append((i, cfg)) + + # results[i] = Path or None, aligned with input configs + results: List[Optional[Path]] = [None] * len(configs) + + def _build_one(idx: int, cfg: BQuantKernelConfig) -> Tuple[int, Optional[Path]]: + hpp = _generate_bquant_kernel(cfg, headers_dir) + if hpp is None: + return idx, None + + so = so_dir / f"lib{cfg.name}.so" + if so.exists(): + log.info(" [cached] %s", so.name) + return idx, so + + ok = _compile_bquant_kernel( + hpp_path=hpp, + so_path=so, + gfx_arch=arch, + hipcc=hipcc, + extra_include_dirs=extra_include_dirs, + ) + return idx, so if ok else None + + if parallel and len(deduped) > 1: + workers = max_workers or min(len(deduped), os.cpu_count() or 4) + with concurrent.futures.ProcessPoolExecutor(max_workers=workers) as ex: + futures = {ex.submit(_build_one, idx, cfg): (idx, cfg) for idx, cfg in deduped} + for fut in concurrent.futures.as_completed(futures): + try: + idx, so_path = fut.result() + results[idx] = so_path + if so_path: + log.info(" built %s", so_path.name) + else: + _, cfg = futures[fut] + log.error(" FAILED %s", cfg.name) + except Exception as e: + _, cfg = futures[fut] + log.error(" EXCEPTION for %s: %s", cfg.name, e) + else: + for idx, cfg in deduped: + _, so_path = _build_one(idx, cfg) + results[idx] = so_path + + # Fill in duplicates + for i, cfg in enumerate(configs): + if results[i] is None: + first_idx = seen.get(cfg.name) + if first_idx is not None and first_idx != i: + results[i] = results[first_idx] + + built = sum(1 for r in results if r is not None) + log.info("Built %d / %d BQuant kernels", built, len(configs)) + return results + + +# ============================================================================= +# Convenience: default fp8 config (matches GemmConfigQuantDecode) +# ============================================================================= + + +def default_fp8_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """Return the default fp8 BQuant config (tile = 16x64x256, warp = 1x4x1).""" + return BQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) + + +def default_bf8_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """Return the default bf8 BQuant config (tile = 16x64x256, warp = 1x4x1).""" + return BQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) diff --git a/projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py b/projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py new file mode 100644 index 000000000000..97f61f3e7e62 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +CPU-only unit tests for bquant_gemm_utils.py. + +Tests kernel name generation, config serialization, and problem dimension helpers. +No GPU or hipcc required. + +Run: + python3 -m pytest dispatcher/tests/test_bquant_gemm_utils.py -v +""" + +import math +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + +from bquant_gemm_utils import ( + BQuantKernelConfig, + BQuantGemmProblem, + default_fp8_config, + default_bf8_config, +) + + +# ============================================================================= +# BQuantKernelConfig.name — byte-exact match with codegen KERNEL_NAME +# ============================================================================= + + +class TestKernelName: + + def test_fp8_rcr_default_name(self): + cfg = BQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert cfg.name == ( + "bquant_gemm_fp8_rcr_compv3_cshuffle_intrawave_" + "16x64x256_1x4x1_16x16x16_qg1x1x128" + ) + + def test_bf8_rcr_default_name(self): + cfg = BQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert cfg.name == ( + "bquant_gemm_bf8_rcr_compv3_cshuffle_intrawave_" + "16x64x256_1x4x1_16x16x16_qg1x1x128" + ) + + def test_different_quant_groups_produce_different_names(self): + def make(gk, gn): + return BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=gn, quant_group_k=gk, + ).name + + names = [make(64, 1), make(128, 1), make(128, 8), make(128, 128)] + assert len(names) == len(set(names)), "All quant-group variants must have unique names" + + def test_preshuffle_b_suffix(self): + cfg = BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_b=True, + ) + assert cfg.name.endswith("_preshuffleb") + + def test_preshuffle_bquant_suffix(self): + cfg = BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_b=True, preshuffle_bquant=True, + ) + assert "_preshuffleb_preshufflebq" in cfg.name + + def test_preshuffle_bquant_only_suffix(self): + cfg = BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_bquant=True, + ) + # preshuffle_bquant without preshuffle_b: _preshufflebq present but not _preshuffleb_ + # Use word-boundary check: the suffix "_preshuffleb" must not appear as a standalone token + name = cfg.name + assert "_preshufflebq" in name + # "_preshuffleb" is a prefix of "_preshufflebq" — check it's not a standalone suffix + assert not name.endswith("_preshuffleb") + + def test_name_no_spaces(self): + cfg = default_fp8_config() + assert " " not in cfg.name + + def test_name_only_valid_chars(self): + import re + cfg = default_fp8_config() + assert re.match(r'^[a-z0-9_]+$', cfg.name), f"Invalid chars in name: {cfg.name}" + + def test_default_fp8_config_name(self): + cfg = default_fp8_config(quant_group_k=128) + assert "fp8" in cfg.name + assert "qg1x1x128" in cfg.name + + def test_default_bf8_config_name(self): + cfg = default_bf8_config(quant_group_k=128) + assert "bf8" in cfg.name + assert "qg1x1x128" in cfg.name + + +# ============================================================================= +# BQuantKernelConfig.to_codegen_config — round-trip shape +# ============================================================================= + + +class TestCodegenConfig: + + def test_codegen_config_contains_correct_variant(self): + cfg = default_fp8_config() + d = cfg.to_codegen_config() + assert d["variant_keys"] == ["fp8"] + + def test_codegen_config_tile_roundtrip(self): + cfg = BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=32, tile_n=128, tile_k=64, + warp_m=2, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=4, quant_group_k=64, + ) + d = cfg.to_codegen_config() + tc = d["tile_configs"][0] + assert tc["tile_m"] == 32 + assert tc["tile_n"] == 128 + assert tc["tile_k"] == 64 + assert tc["warp_m"] == 2 + qg = d["quant_groups"][0] + assert qg["quant_group_k"] == 64 + assert qg["quant_group_n"] == 4 + + def test_codegen_config_single_layout(self): + cfg = default_fp8_config() + d = cfg.to_codegen_config() + assert d["layouts"] == ["rcr"] + + +# ============================================================================= +# BQuantGemmProblem dimension helpers +# ============================================================================= + + +class TestBQuantGemmProblem: + + def test_QK_B_exact_multiple(self): + p = BQuantGemmProblem(M=16, N=64, K=256, quant_group_k=128) + assert p.QK_B == 2 # 256 / 128 + + def test_QK_B_ceil(self): + p = BQuantGemmProblem(M=16, N=64, K=300, quant_group_k=128) + assert p.QK_B == math.ceil(300 / 128) # == 3 + + def test_QN_B_exact_multiple(self): + p = BQuantGemmProblem(M=16, N=64, K=256, quant_group_n=32) + assert p.QN_B == 2 # 64 / 32 + + def test_QN_B_default_one_group(self): + p = BQuantGemmProblem(M=16, N=64, K=256) + assert p.QN_B == 64 # default quant_group_n=1 → every column is its own group + + def test_QN_B_ceil(self): + p = BQuantGemmProblem(M=16, N=65, K=256, quant_group_n=32) + assert p.QN_B == math.ceil(65 / 32) # == 3 + + def test_default_k_batch(self): + p = BQuantGemmProblem(M=16, N=64, K=256) + assert p.k_batch == 1 + + +# ============================================================================= +# Name uniqueness across a small sweep +# ============================================================================= + + +class TestNameUniqueness: + + def _make_configs(self): + configs = [] + for variant in ("fp8", "bf8"): + for gk in (64, 128): + for gn in (1, 8): + configs.append(BQuantKernelConfig( + variant_key=variant, + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=gn, quant_group_k=gk, + )) + return configs + + def test_all_names_unique(self): + configs = self._make_configs() + names = [c.name for c in configs] + assert len(names) == len(set(names)), f"Duplicate names: {[n for n in names if names.count(n) > 1]}" From eb0dd894cc954a125c1dab9cf935364bd82e6d52 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Mon, 6 Jul 2026 19:54:08 -0400 Subject: [PATCH 02/35] [CK][Dispatcher] Rename bquant_gemm -> grouped_gemm_bquant throughout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 18 +++++++++--------- ....cpp => grouped_gemm_bquant_ctypes_lib.cpp} | 4 ++-- ... => unified_grouped_gemm_bquant_codegen.py} | 12 ++++++------ ...quant_gemm.py => 13_grouped_gemm_bquant.py} | 6 +++--- ...m_utils.py => grouped_gemm_bquant_utils.py} | 14 +++++++------- ...ls.py => test_grouped_gemm_bquant_utils.py} | 10 +++++----- 6 files changed, 32 insertions(+), 32 deletions(-) rename projects/composablekernel/dispatcher/bindings/ctypes/{bquant_gemm_ctypes_lib.cpp => grouped_gemm_bquant_ctypes_lib.cpp} (99%) rename projects/composablekernel/dispatcher/codegen/{unified_bquant_gemm_codegen.py => unified_grouped_gemm_bquant_codegen.py} (98%) rename projects/composablekernel/dispatcher/examples/gemm/python/{13_bquant_gemm.py => 13_grouped_gemm_bquant.py} (97%) rename projects/composablekernel/dispatcher/python/{bquant_gemm_utils.py => grouped_gemm_bquant_utils.py} (97%) rename projects/composablekernel/dispatcher/tests/{test_bquant_gemm_utils.py => test_grouped_gemm_bquant_utils.py} (96%) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 300c9f46b74c..6a9132a5fcfe 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -155,31 +155,31 @@ else() endif() # ============================================================================= -# BQuantGrouped GEMM ctypes Library +# GroupedGemm BQuant ctypes Library # ============================================================================= -file(GLOB BQUANT_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/bquant_gemm_*.hpp") +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_bquant_gemm_lib - bquant_gemm_ctypes_lib.cpp + add_ctypes_library(dispatcher_grouped_gemm_bquant_lib + grouped_gemm_bquant_ctypes_lib.cpp KERNEL_HEADER ${BQUANT_GEMM_KERNEL_HEADER} ) - target_compile_definitions(dispatcher_bquant_gemm_lib PRIVATE + target_compile_definitions(dispatcher_grouped_gemm_bquant_lib PRIVATE CK_TILE_SINGLE_KERNEL_INCLUDE GFX_ARCH="gfx950" ) else() message(STATUS "No BQuant GEMM kernel found for ctypes lib - building without kernel") - add_library(dispatcher_bquant_gemm_lib SHARED bquant_gemm_ctypes_lib.cpp) - target_include_directories(dispatcher_bquant_gemm_lib PRIVATE + 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_bquant_gemm_lib PRIVATE hip::device) - set_target_properties(dispatcher_bquant_gemm_lib PROPERTIES + 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 ) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp similarity index 99% rename from projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp rename to projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp index cffeb1d59c33..366e278e0021 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/bquant_gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp @@ -2,11 +2,11 @@ // SPDX-License-Identifier: MIT /** - * BQuantGrouped GEMM Dispatcher ctypes Library + * GroupedGemm BQuant Dispatcher ctypes Library * * Provides C API for Python ctypes integration. * Kernel header is force-included at compile time via: - * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE bquant_gemm_ctypes_lib.cpp + * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE grouped_gemm_bquant_ctypes_lib.cpp * * Force-include defines (from generated kernel header): * SelectedKernel, KERNEL_NAME diff --git a/projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py similarity index 98% rename from projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py rename to projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py index 0d378dbc3874..3a7f1060b504 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_bquant_gemm_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py @@ -4,19 +4,19 @@ # SPDX-License-Identifier: MIT """ -BQuantGrouped GEMM Code Generator +GroupedGemm BQuant Code Generator Generates one .hpp per kernel config for the dispatcher's ctypes path. Each header defines a SelectedKernel struct with a static launch() method taking QuantGemmHostArgs — compiled per-kernel via force-include: - hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE bquant_gemm_ctypes_lib.cpp + hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE grouped_gemm_bquant_ctypes_lib.cpp Initial scope: fp8 and bf8 dtype variants, non-preshuffle, compv3 pipeline, rcr layout, configurable QuantGroupShape. -Naming convention (byte-exact with BQuantKernelConfig.name in bquant_gemm_utils.py): - bquant_gemm_{dtype_a}_{layout}_{pipeline}_{epilogue}_{scheduler}_ +Naming convention (byte-exact with BQuantKernelConfig.name in grouped_gemm_bquant_utils.py): + grouped_gemm_bquant_{dtype_a}_{layout}_{pipeline}_{epilogue}_{scheduler}_ {TileM}x{TileN}x{TileK}_{WarpM}x{WarpN}x{WarpK}_{WtM}x{WtN}x{WtK}_ qg{gM}x{gN}x{gK}[_preshuffleb][_preshufflebq] @@ -147,7 +147,7 @@ class BQuantKernelSpec: def name(self) -> str: t = self.tile parts = [ - "bquant_gemm", + "grouped_gemm_bquant", self.variant_key, self.layout, self.pipeline, @@ -210,7 +210,7 @@ def generate(self, spec: BQuantKernelSpec) -> str: return f"""\ // SPDX-License-Identifier: MIT // Auto-generated BQuantGrouped GEMM kernel header. -// DO NOT EDIT — regenerate via unified_bquant_gemm_codegen.py +// DO NOT EDIT — regenerate via unified_grouped_gemm_bquant_codegen.py #pragma once #include "ck_tile/core.hpp" diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py similarity index 97% rename from projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py rename to projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py index 3d1c6b91308d..30487c0ee3b8 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/13_bquant_gemm.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py @@ -4,10 +4,10 @@ # SPDX-License-Identifier: MIT """ -Example 13 — BQuantGrouped GEMM via the Dispatcher +Example 13 — GroupedGemm BQuant via the Dispatcher Demonstrates the full three-layer path: - 1. Codegen — unified_bquant_gemm_codegen.py → .hpp + 1. Codegen — unified_grouped_gemm_bquant_codegen.py → .hpp 2. Compile — hipcc → .so 3. Run — BQuantGpuGemmRunner → C = A @ dequant(B, BQ) @@ -109,7 +109,7 @@ def reference_bquant_gemm( def main(): - parser = argparse.ArgumentParser(description="BQuantGrouped GEMM dispatcher example") + parser = argparse.ArgumentParser(description="GroupedGemm BQuant dispatcher example") parser.add_argument("--dtype", choices=["fp8", "bf8"], default="fp8") parser.add_argument("--M", type=int, default=16) parser.add_argument("--N", type=int, default=64) diff --git a/projects/composablekernel/dispatcher/python/bquant_gemm_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py similarity index 97% rename from projects/composablekernel/dispatcher/python/bquant_gemm_utils.py rename to projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py index 27fcbd922439..d420fa050cca 100644 --- a/projects/composablekernel/dispatcher/python/bquant_gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: MIT """ -BQuantGrouped GEMM dispatcher utilities. +GroupedGemm BQuant dispatcher utilities. Three-layer Python bridge for the dispatcher's BQuantGrouped GEMM path: @@ -42,8 +42,8 @@ # Constants # ============================================================================= -_CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_bquant_gemm_codegen.py" -_CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "bquant_gemm_ctypes_lib.cpp" +_CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_grouped_gemm_bquant_codegen.py" +_CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "grouped_gemm_bquant_ctypes_lib.cpp" _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -69,7 +69,7 @@ class BQuantKernelConfig: """ Complete description of one BQuantGrouped GEMM kernel. - The .name property produces the exact string that unified_bquant_gemm_codegen.py + The .name property produces the exact string that unified_grouped_gemm_bquant_codegen.py emits as KERNEL_NAME, ensuring the Python side and compiled .so always agree. """ @@ -102,7 +102,7 @@ class BQuantKernelConfig: def name(self) -> str: """Byte-exact match to codegen KERNEL_NAME.""" parts = [ - "bquant_gemm", + "grouped_gemm_bquant", self.variant_key, self.layout, self.pipeline, @@ -120,7 +120,7 @@ def name(self) -> str: return "_".join(parts) def to_codegen_config(self) -> dict: - """Produce the JSON config dict consumed by unified_bquant_gemm_codegen.py.""" + """Produce the JSON config dict consumed by unified_grouped_gemm_bquant_codegen.py.""" return { "variant_keys": [self.variant_key], "layouts": [self.layout], @@ -416,7 +416,7 @@ def _generate_bquant_kernel( output_dir: Path, ) -> Optional[Path]: """ - Run unified_bquant_gemm_codegen.py for one config; return the .hpp path or None. + Run unified_grouped_gemm_bquant_codegen.py for one config; return the .hpp path or None. """ config_dict = config.to_codegen_config() config_json = json.dumps(config_dict) diff --git a/projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py similarity index 96% rename from projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py rename to projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py index 97f61f3e7e62..0b6ec482a2a3 100644 --- a/projects/composablekernel/dispatcher/tests/test_bquant_gemm_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py @@ -4,13 +4,13 @@ # SPDX-License-Identifier: MIT """ -CPU-only unit tests for bquant_gemm_utils.py. +CPU-only unit tests for grouped_gemm_bquant_utils.py. Tests kernel name generation, config serialization, and problem dimension helpers. No GPU or hipcc required. Run: - python3 -m pytest dispatcher/tests/test_bquant_gemm_utils.py -v + python3 -m pytest dispatcher/tests/test_grouped_gemm_bquant_utils.py -v """ import math @@ -21,7 +21,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) -from bquant_gemm_utils import ( +from grouped_gemm_bquant_utils import ( BQuantKernelConfig, BQuantGemmProblem, default_fp8_config, @@ -49,7 +49,7 @@ def test_fp8_rcr_default_name(self): quant_group_m=1, quant_group_n=1, quant_group_k=128, ) assert cfg.name == ( - "bquant_gemm_fp8_rcr_compv3_cshuffle_intrawave_" + "grouped_gemm_bquant_fp8_rcr_compv3_cshuffle_intrawave_" "16x64x256_1x4x1_16x16x16_qg1x1x128" ) @@ -66,7 +66,7 @@ def test_bf8_rcr_default_name(self): quant_group_m=1, quant_group_n=1, quant_group_k=128, ) assert cfg.name == ( - "bquant_gemm_bf8_rcr_compv3_cshuffle_intrawave_" + "grouped_gemm_bquant_bf8_rcr_compv3_cshuffle_intrawave_" "16x64x256_1x4x1_16x16x16_qg1x1x128" ) From 852997dfb5c4acac0eeb044926d97eb67703943e Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 7 Jul 2026 12:24:44 -0400 Subject: [PATCH 03/35] [CK][Dispatcher] Fix stale import in 13_grouped_gemm_bquant.py Example file still imported from bquant_gemm_utils after rename. Updated to grouped_gemm_bquant_utils. Co-Authored-By: Claude --- .../dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py index 30487c0ee3b8..3fc225941ce0 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py @@ -37,7 +37,7 @@ # Add dispatcher/python to path sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "python")) -from bquant_gemm_utils import ( +from grouped_gemm_bquant_utils import ( BQuantKernelConfig, BQuantGemmProblem, BQuantGpuGemmRunner, From d564c6aa6cce67c13aeccfcfa07d7fa845dc7c6c Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 7 Jul 2026 14:54:05 -0400 Subject: [PATCH 04/35] [CK][Dispatcher] Add fp8i4/bf8i4, preshuffle, and MX variants to BQuant GEMM dispatcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../unified_grouped_gemm_bquant_codegen.py | 89 +++- .../python/grouped_gemm_bquant_utils.py | 384 +++++++++++++++++- .../tests/test_grouped_gemm_bquant_utils.py | 337 ++++++++++++++- 3 files changed, 799 insertions(+), 11 deletions(-) diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py index 3a7f1060b504..67a9fa6050e3 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py @@ -69,6 +69,62 @@ "ck_q": "float", "ck_acc": "float", }, + "fp8i4": { + "dtype_a": "fp8", + "dtype_b": "pk_int4", + "dtype_c": "half", + "dtype_q": "fp8", + "ck_a": "ck_tile::fp8_t", + "ck_b": "ck_tile::pk_int4_t", + "ck_c": "ck_tile::half_t", + "ck_q": "ck_tile::fp8_t", + "ck_acc": "float", + }, + "bf8i4": { + "dtype_a": "bf8", + "dtype_b": "pk_int4", + "dtype_c": "half", + "dtype_q": "bf8", + "ck_a": "ck_tile::bf8_t", + "ck_b": "ck_tile::pk_int4_t", + "ck_c": "ck_tile::half_t", + "ck_q": "ck_tile::bf8_t", + "ck_acc": "float", + }, + # MX microscale variants — Q-type is e8m0 (block scale), pipeline = microscale + "mx_bf16bf16": { + "dtype_a": "bf16", + "dtype_b": "bf16", + "dtype_c": "bf16", + "dtype_q": "e8m0", + "ck_a": "ck_tile::bf16_t", + "ck_b": "ck_tile::bf16_t", + "ck_c": "ck_tile::bf16_t", + "ck_q": "ck_tile::e8m0_t", + "ck_acc": "float", + }, + "mx_bf16bf8": { + "dtype_a": "bf16", + "dtype_b": "bf8", + "dtype_c": "bf16", + "dtype_q": "e8m0", + "ck_a": "ck_tile::bf16_t", + "ck_b": "ck_tile::bf8_t", + "ck_c": "ck_tile::bf16_t", + "ck_q": "ck_tile::e8m0_t", + "ck_acc": "float", + }, + "mx_bf16fp4": { + "dtype_a": "bf16", + "dtype_b": "pk_fp4", + "dtype_c": "bf16", + "dtype_q": "e8m0", + "ck_a": "ck_tile::bf16_t", + "ck_b": "ck_tile::pk_fp4_t", + "ck_c": "ck_tile::bf16_t", + "ck_q": "ck_tile::e8m0_t", + "ck_acc": "float", + }, } # Layout strings supported: only rcr for initial implementation @@ -78,15 +134,21 @@ "c": "ck_tile::tensor_layout::gemm::ColumnMajor", } -# Pipeline map for BQuant (compv3 for initial; preshuffle/microscale are follow-ons) +# Pipeline map for BQuant kernels. +# "preshuffleb" -> WPQuantBPipelineAgBgCrV2 (preshuffle_b=true variants) +# "microscale" -> MicroscaleGemmPipelineAgBgCrCompV3 (MX e8m0 scale variants) BQUANT_PIPELINE_MAP = { - "compv3": "ck_tile::BQuantGemmPipelineAgBgCrCompV3", - # preshuffle_b -> WPQuantBPipelineAgBgCrV2 (follow-on) - # mx -> MicroscaleGemmPipelineAgBgCrCompV3 (follow-on) + "compv3": "ck_tile::BQuantGemmPipelineAgBgCrCompV3", + "preshuffleb": "ck_tile::WPQuantBPipelineAgBgCrV2", + "microscale": "ck_tile::MicroscaleGemmPipelineAgBgCrCompV3", } BQUANT_BASE_PIPELINE_MAP = { - "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", + "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", + "preshuffleb": "ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2", + # MX BQuant (QDataType=e8m0, PreshuffleB=false) falls into the else branch in + # run_gemm_quant_example.inc — same base as preshuffleb. + "microscale": "ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2", } BQUANT_SCHEDULER_TO_CK = { @@ -137,6 +199,7 @@ class BQuantKernelSpec: quant_group_k: int = 128 preshuffle_b: bool = False preshuffle_bquant: bool = False + double_smem_buffer: bool = False pad_m: bool = False pad_n: bool = False pad_k: bool = True @@ -202,6 +265,7 @@ def generate(self, spec: BQuantKernelSpec) -> str: pad_k = str(spec.pad_k).lower() preshuffle_b = str(spec.preshuffle_b).lower() preshuffle_bquant = str(spec.preshuffle_bquant).lower() + double_smem_buffer = str(spec.double_smem_buffer).lower() # CShuffleEpilogue problem parameters # TiledMMAPermuteN is false when quant_group_n > 1 (see run_gemm_quant_example.inc) @@ -265,7 +329,7 @@ def generate(self, spec: BQuantKernelSpec) -> str: static constexpr bool BPreshuffleQuant = {preshuffle_bquant}; static constexpr bool PreshuffleB = {preshuffle_b}; static constexpr bool TransposeC = false; - static constexpr bool DoubleSmemBuffer = false; + static constexpr bool DoubleSmemBuffer = {double_smem_buffer}; using TileShape = ck_tile::TileGemmShape< ck_tile::sequence, @@ -394,6 +458,9 @@ def _default_config() -> dict: "pad_k": True, "block_size": 256, "k_block_per_cu": 1, + "double_smem_buffer": False, + "preshuffle_b": False, + "preshuffle_bquant": False, } @@ -405,8 +472,11 @@ def _build_specs(config: dict) -> List[BQuantKernelSpec]: pad_m = config.get("pad_m", False) pad_n = config.get("pad_n", False) pad_k = config.get("pad_k", True) - block_size = config.get("block_size", 256) - k_block_per_cu = config.get("k_block_per_cu", 1) + block_size = config.get("block_size", 256) + k_block_per_cu = config.get("k_block_per_cu", 1) + double_smem_buffer = config.get("double_smem_buffer", False) + preshuffle_b = config.get("preshuffle_b", False) + preshuffle_bquant = config.get("preshuffle_bquant", False) for variant_key, layout, tile_dict, qg in itertools.product( config.get("variant_keys", ["fp8"]), @@ -446,6 +516,9 @@ def _build_specs(config: dict) -> List[BQuantKernelSpec]: quant_group_m=qg.get("quant_group_m", 1), quant_group_n=qg.get("quant_group_n", 1), quant_group_k=qg.get("quant_group_k", 128), + preshuffle_b=preshuffle_b, + preshuffle_bquant=preshuffle_bquant, + double_smem_buffer=double_smem_buffer, pad_m=pad_m, pad_n=pad_n, pad_k=pad_k, diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py index d420fa050cca..38b40c490c89 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py @@ -93,8 +93,10 @@ class BQuantKernelConfig: quant_group_n: int = 1 quant_group_k: int = 128 - preshuffle_b: bool = False - preshuffle_bquant: bool = False + preshuffle_b: bool = False + preshuffle_bquant: bool = False + double_smem_buffer: bool = False + k_block_per_cu: int = 1 gfx_arch: str = _DEFAULT_GFX_ARCH @@ -143,6 +145,10 @@ def to_codegen_config(self) -> dict: "quant_group_n": self.quant_group_n, "quant_group_k": self.quant_group_k, }], + "preshuffle_b": self.preshuffle_b, + "preshuffle_bquant": self.preshuffle_bquant, + "double_smem_buffer": self.double_smem_buffer, + "k_block_per_cu": self.k_block_per_cu, } @@ -646,3 +652,377 @@ def default_bf8_config( quant_group_k=quant_group_k, gfx_arch=gfx_arch, ) + + +def default_fp8i4_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """Return the default fp8i4 BQuant config (A=fp8, B=pk_int4, Q=fp8; tile = 16x64x256).""" + return BQuantKernelConfig( + variant_key="fp8i4", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) + + +def default_bf8i4_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """Return the default bf8i4 BQuant config (A=bf8, B=pk_int4, Q=bf8; tile = 16x64x256).""" + return BQuantKernelConfig( + variant_key="bf8i4", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) + + +# ============================================================================= +# Phase 3a — preshuffle_b only (WPQuantBPipelineAgBgCrV2, prefill tile 128x128x128) +# +# Tile dims from GemmConfigPreshuffleB_BQuant_Prefill: +# M=128, N=128, K=128/sizeof(PrecType) +# Warp: 1x4x1, WarpTile: 16x16x K_warp +# K_warp from get_k_warp_tile on gfx950: +# fp8/bf8 (8-bit float): K_warp = 128 → tile_k = 128 +# pk_int4 (not 8b float): K_warp = 32 → tile_k = 128 (sizeof=0.5 → 128/0.5=256 packed) +# DoubleSmemBuffer=true, kBlockPerCu=2 +# ============================================================================= + + +def _preshuffleb_config( + variant_key: str, + warp_tile_k: int, + quant_group_k: int, + quant_group_n: int, + gfx_arch: str, +) -> BQuantKernelConfig: + return BQuantKernelConfig( + variant_key=variant_key, + layout="rcr", + pipeline="preshuffleb", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=warp_tile_k, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + preshuffle_b=True, + preshuffle_bquant=False, + double_smem_buffer=True, + k_block_per_cu=2, + gfx_arch=gfx_arch, + ) + + +def default_fp8_preshuffleb_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """fp8 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill).""" + return _preshuffleb_config("fp8", warp_tile_k=128, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_bf8_preshuffleb_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """bf8 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill).""" + return _preshuffleb_config("bf8", warp_tile_k=128, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_fp8i4_preshuffleb_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """fp8i4 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill, B=pk_int4). + + K_warp_tile=32: pk_int4 is not an 8-bit float type so get_k_warp_tile returns 32 on gfx950. + """ + return _preshuffleb_config("fp8i4", warp_tile_k=32, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_bf8i4_preshuffleb_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """bf8i4 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill, B=pk_int4). + + K_warp_tile=32: pk_int4 is not an 8-bit float type so get_k_warp_tile returns 32 on gfx950. + """ + return _preshuffleb_config("bf8i4", warp_tile_k=32, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +# ============================================================================= +# Phase 3b — preshuffle_bquant only (BQuantGemmPipelineAgBgCrCompV3, prefill tile 128x128x128) +# +# Tile dims from GemmConfigPreshuffleBQuantPrefill (inherits GemmConfigQuantPrefill): +# M=128, N=128, K=128, Warp: 1x4x1, WarpTile: 16x16xK_warp +# K_warp from get_k_warp_tile on gfx950: +# fp8/bf8: K_warp = 128 (is_8bit_float=true, M_Warp_Tile != 32) +# pk_int4: K_warp = 32 +# DoubleSmemBuffer=false (default), kBlockPerCu=1 (default) +# Extra quant group: 1x16x128 (not in 3a) +# ============================================================================= + + +def _preshufflequant_config( + variant_key: str, + warp_tile_k: int, + quant_group_k: int, + quant_group_n: int, + gfx_arch: str, +) -> BQuantKernelConfig: + return BQuantKernelConfig( + variant_key=variant_key, + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=warp_tile_k, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + preshuffle_b=False, + preshuffle_bquant=True, + gfx_arch=gfx_arch, + ) + + +def default_fp8_preshufflequant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """fp8 preshuffle_bquant prefill config (GemmConfigPreshuffleBQuantPrefill).""" + return _preshufflequant_config("fp8", warp_tile_k=128, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_bf8_preshufflequant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """bf8 preshuffle_bquant prefill config (GemmConfigPreshuffleBQuantPrefill).""" + return _preshufflequant_config("bf8", warp_tile_k=128, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_fp8i4_preshufflequant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """fp8i4 preshuffle_bquant prefill config.""" + return _preshufflequant_config("fp8i4", warp_tile_k=32, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_bf8i4_preshufflequant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """bf8i4 preshuffle_bquant prefill config.""" + return _preshufflequant_config("bf8i4", warp_tile_k=32, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +# ============================================================================= +# Phase 3c — preshuffle_b + preshuffle_bquant (WPQuantBPipelineAgBgCrV2, prefill tile) +# +# Tile dims from GemmConfigPreshuffleB_PreshuffleBQuant_Prefill +# (inherits GemmConfigPreshuffleB_BQuant_Prefill): identical tile/warp dims to 3a. +# PreshuffleB=true, BPreshuffleQuant=true, DoubleSmemBuffer=true, kBlockPerCu=2 +# ============================================================================= + + +def _preshuffleb_bquant_config( + variant_key: str, + warp_tile_k: int, + quant_group_k: int, + quant_group_n: int, + gfx_arch: str, +) -> BQuantKernelConfig: + return BQuantKernelConfig( + variant_key=variant_key, + layout="rcr", + pipeline="preshuffleb", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=warp_tile_k, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + preshuffle_b=True, + preshuffle_bquant=True, + double_smem_buffer=True, + k_block_per_cu=2, + gfx_arch=gfx_arch, + ) + + +def default_fp8_preshuffleb_bquant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """fp8 preshuffle_b+preshuffle_bquant config (GemmConfigPreshuffleB_PreshuffleBQuant_Prefill).""" + return _preshuffleb_bquant_config("fp8", warp_tile_k=128, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_bf8_preshuffleb_bquant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """bf8 preshuffle_b+preshuffle_bquant config.""" + return _preshuffleb_bquant_config("bf8", warp_tile_k=128, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_fp8i4_preshuffleb_bquant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """fp8i4 preshuffle_b+preshuffle_bquant config.""" + return _preshuffleb_bquant_config("fp8i4", warp_tile_k=32, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +def default_bf8i4_preshuffleb_bquant_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """bf8i4 preshuffle_b+preshuffle_bquant config.""" + return _preshuffleb_bquant_config("bf8i4", warp_tile_k=32, quant_group_k=quant_group_k, + quant_group_n=quant_group_n, gfx_arch=gfx_arch) + + +# ============================================================================= +# Phase 4 — MX microscale variants (A=bf16, Q=e8m0 block scale) +# +# Pipeline for all three: MicroscaleGemmPipelineAgBgCrCompV3 +# (QDataType=e8m0_t, PreshuffleB=false → microscale branch in BQuantPipeline selection) +# Base pipeline: BaseWeightPreshufflePipelineAGmemBGmemCRegV2 +# (BQuantGrouped, PreshuffleB=false, not AQuant/ABQuant → else branch) +# +# Tile dims: +# mx_bf16bf16: GemmConfigQuantPrefill → 128x128x128, warp 1x4x1, warp_tile_k=32 +# mx_bf16bf8: GemmConfigMixedPrecision → 128x128x128, warp 1x4x1, warp_tile_k=64 +# mx_bf16fp4: GemmConfigQuantPrefill → 128x128x128, warp 1x4x1, warp_tile_k=32 +# ============================================================================= + + +def default_mx_bf16bf16_config( + quant_group_k: int = 32, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """MX bf16+bf16 config (A=bf16, B=bf16, Q=e8m0; GemmConfigQuantPrefill). + + Default quant_group_k=32 matches the smallest MX block size for bf16+bf16. + """ + return BQuantKernelConfig( + variant_key="mx_bf16bf16", + layout="rcr", + pipeline="microscale", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) + + +def default_mx_bf16bf8_config( + quant_group_k: int = 128, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """MX bf16+bf8 config (A=bf16, B=bf8, Q=e8m0; GemmConfigMixedPrecision). + + warp_tile_k=64 is hardcoded in GemmConfigMixedPrecision (mixed-precision KPack=16). + """ + return BQuantKernelConfig( + variant_key="mx_bf16bf8", + layout="rcr", + pipeline="microscale", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=64, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) + + +def default_mx_bf16fp4_config( + quant_group_k: int = 32, + quant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> BQuantKernelConfig: + """MX bf16+fp4 config (A=bf16, B=pk_fp4, Q=e8m0; GemmConfigQuantPrefill). + + Default quant_group_k=32 matches the smallest MX block size for fp4. + """ + return BQuantKernelConfig( + variant_key="mx_bf16fp4", + layout="rcr", + pipeline="microscale", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + gfx_arch=gfx_arch, + ) diff --git a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py index 0b6ec482a2a3..5afd3332c293 100644 --- a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py @@ -26,6 +26,23 @@ BQuantGemmProblem, default_fp8_config, default_bf8_config, + default_fp8i4_config, + default_bf8i4_config, + default_fp8_preshuffleb_config, + default_bf8_preshuffleb_config, + default_fp8i4_preshuffleb_config, + default_bf8i4_preshuffleb_config, + default_fp8_preshufflequant_config, + default_bf8_preshufflequant_config, + default_fp8i4_preshufflequant_config, + default_bf8i4_preshufflequant_config, + default_fp8_preshuffleb_bquant_config, + default_bf8_preshuffleb_bquant_config, + default_fp8i4_preshuffleb_bquant_config, + default_bf8i4_preshuffleb_bquant_config, + default_mx_bf16bf16_config, + default_mx_bf16bf8_config, + default_mx_bf16fp4_config, ) @@ -144,6 +161,58 @@ def test_default_bf8_config_name(self): assert "bf8" in cfg.name assert "qg1x1x128" in cfg.name + def test_fp8i4_rcr_default_name(self): + cfg = BQuantKernelConfig( + variant_key="fp8i4", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert cfg.name == ( + "grouped_gemm_bquant_fp8i4_rcr_compv3_cshuffle_intrawave_" + "16x64x256_1x4x1_16x16x16_qg1x1x128" + ) + + def test_bf8i4_rcr_default_name(self): + cfg = BQuantKernelConfig( + variant_key="bf8i4", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert cfg.name == ( + "grouped_gemm_bquant_bf8i4_rcr_compv3_cshuffle_intrawave_" + "16x64x256_1x4x1_16x16x16_qg1x1x128" + ) + + def test_default_fp8i4_config_name(self): + cfg = default_fp8i4_config(quant_group_k=128) + assert "fp8i4" in cfg.name + assert "qg1x1x128" in cfg.name + + def test_default_bf8i4_config_name(self): + cfg = default_bf8i4_config(quant_group_k=128) + assert "bf8i4" in cfg.name + assert "qg1x1x128" in cfg.name + + def test_fp8i4_quant_group64_name(self): + cfg = default_fp8i4_config(quant_group_k=64) + assert "qg1x1x64" in cfg.name + + def test_bf8i4_quant_groupn8_name(self): + cfg = default_bf8i4_config(quant_group_k=128, quant_group_n=8) + assert "qg1x8x128" in cfg.name + # ============================================================================= # BQuantKernelConfig.to_codegen_config — round-trip shape @@ -214,6 +283,272 @@ def test_default_k_batch(self): assert p.k_batch == 1 +# ============================================================================= +# Phase 3 infrastructure — preshuffle fields and new pipeline key +# ============================================================================= + + +class TestPhase3Infrastructure: + + def test_double_smem_buffer_defaults_false(self): + cfg = default_fp8_config() + assert cfg.double_smem_buffer is False + + def test_k_block_per_cu_defaults_one(self): + cfg = default_fp8_config() + assert cfg.k_block_per_cu == 1 + + def test_preshuffleb_pipeline_accepted(self): + cfg = BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="preshuffleb", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_b=True, double_smem_buffer=True, k_block_per_cu=2, + ) + assert "preshuffleb" in cfg.name + assert cfg.double_smem_buffer is True + assert cfg.k_block_per_cu == 2 + + def test_codegen_config_carries_preshuffle_fields(self): + cfg = BQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="preshuffleb", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_b=True, double_smem_buffer=True, k_block_per_cu=2, + ) + d = cfg.to_codegen_config() + assert d["preshuffle_b"] is True + assert d["preshuffle_bquant"] is False + assert d["double_smem_buffer"] is True + assert d["k_block_per_cu"] == 2 + + def test_codegen_config_default_preshuffle_fields_false(self): + d = default_fp8_config().to_codegen_config() + assert d["preshuffle_b"] is False + assert d["preshuffle_bquant"] is False + assert d["double_smem_buffer"] is False + assert d["k_block_per_cu"] == 1 + + +# ============================================================================= +# Phase 3a/b/c — preshuffle convenience configs +# ============================================================================= + + +class TestPhase3Configs: + + # ---- 3a: preshuffle_b only ---- + + def test_fp8_preshuffleb_name(self): + cfg = default_fp8_preshuffleb_config() + assert cfg.name == ( + "grouped_gemm_bquant_fp8_rcr_preshuffleb_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x128_qg1x1x128_preshuffleb" + ) + + def test_bf8_preshuffleb_name(self): + cfg = default_bf8_preshuffleb_config() + assert cfg.name == ( + "grouped_gemm_bquant_bf8_rcr_preshuffleb_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x128_qg1x1x128_preshuffleb" + ) + + def test_fp8i4_preshuffleb_name(self): + cfg = default_fp8i4_preshuffleb_config() + assert cfg.name == ( + "grouped_gemm_bquant_fp8i4_rcr_preshuffleb_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x32_qg1x1x128_preshuffleb" + ) + + def test_bf8i4_preshuffleb_name(self): + cfg = default_bf8i4_preshuffleb_config() + assert cfg.name == ( + "grouped_gemm_bquant_bf8i4_rcr_preshuffleb_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x32_qg1x1x128_preshuffleb" + ) + + def test_preshuffleb_flags(self): + cfg = default_fp8_preshuffleb_config() + assert cfg.preshuffle_b is True + assert cfg.preshuffle_bquant is False + assert cfg.double_smem_buffer is True + assert cfg.k_block_per_cu == 2 + assert cfg.pipeline == "preshuffleb" + + def test_preshuffleb_i4_warp_tile_k(self): + # pk_int4 is not 8-bit float → K_warp_tile=32 on gfx950 + cfg = default_fp8i4_preshuffleb_config() + assert cfg.warp_tile_k == 32 + + def test_preshuffleb_fp8_warp_tile_k(self): + cfg = default_fp8_preshuffleb_config() + assert cfg.warp_tile_k == 128 + + # ---- 3b: preshuffle_bquant only ---- + + def test_fp8_preshufflequant_name(self): + cfg = default_fp8_preshufflequant_config() + assert cfg.name == ( + "grouped_gemm_bquant_fp8_rcr_compv3_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x128_qg1x1x128_preshufflebq" + ) + + def test_bf8_preshufflequant_name(self): + cfg = default_bf8_preshufflequant_config() + assert cfg.name == ( + "grouped_gemm_bquant_bf8_rcr_compv3_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x128_qg1x1x128_preshufflebq" + ) + + def test_fp8i4_preshufflequant_name(self): + cfg = default_fp8i4_preshufflequant_config() + assert cfg.name == ( + "grouped_gemm_bquant_fp8i4_rcr_compv3_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x32_qg1x1x128_preshufflebq" + ) + + def test_preshufflequant_flags(self): + cfg = default_fp8_preshufflequant_config() + assert cfg.preshuffle_b is False + assert cfg.preshuffle_bquant is True + assert cfg.double_smem_buffer is False + assert cfg.k_block_per_cu == 1 + assert cfg.pipeline == "compv3" + + def test_preshufflequant_16n_group(self): + # 1x16x128 is the extra quant group for 3b — verify config accepts it + cfg = default_fp8_preshufflequant_config(quant_group_n=16) + assert "qg1x16x128" in cfg.name + + # ---- 3c: preshuffle_b + preshuffle_bquant ---- + + def test_fp8_preshuffleb_bquant_name(self): + cfg = default_fp8_preshuffleb_bquant_config() + assert cfg.name == ( + "grouped_gemm_bquant_fp8_rcr_preshuffleb_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x128_qg1x1x128_preshuffleb_preshufflebq" + ) + + def test_bf8i4_preshuffleb_bquant_name(self): + cfg = default_bf8i4_preshuffleb_bquant_config() + assert cfg.name == ( + "grouped_gemm_bquant_bf8i4_rcr_preshuffleb_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x32_qg1x1x128_preshuffleb_preshufflebq" + ) + + def test_preshuffleb_bquant_flags(self): + cfg = default_fp8_preshuffleb_bquant_config() + assert cfg.preshuffle_b is True + assert cfg.preshuffle_bquant is True + assert cfg.double_smem_buffer is True + assert cfg.k_block_per_cu == 2 + assert cfg.pipeline == "preshuffleb" + + def test_all_preshuffle_names_unique(self): + configs = [ + default_fp8_preshuffleb_config(), + default_bf8_preshuffleb_config(), + default_fp8i4_preshuffleb_config(), + default_bf8i4_preshuffleb_config(), + default_fp8_preshufflequant_config(), + default_bf8_preshufflequant_config(), + default_fp8i4_preshufflequant_config(), + default_bf8i4_preshufflequant_config(), + default_fp8_preshuffleb_bquant_config(), + default_bf8_preshuffleb_bquant_config(), + default_fp8i4_preshuffleb_bquant_config(), + default_bf8i4_preshuffleb_bquant_config(), + ] + names = [c.name for c in configs] + assert len(names) == len(set(names)), f"Duplicate: {[n for n in names if names.count(n) > 1]}" + + +# ============================================================================= +# Phase 4 — MX microscale variants +# ============================================================================= + + +class TestPhase4MXConfigs: + + def test_mx_bf16bf16_name(self): + cfg = default_mx_bf16bf16_config(quant_group_k=32) + assert cfg.name == ( + "grouped_gemm_bquant_mx_bf16bf16_rcr_microscale_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x32_qg1x1x32" + ) + + def test_mx_bf16bf8_name(self): + cfg = default_mx_bf16bf8_config(quant_group_k=128) + assert cfg.name == ( + "grouped_gemm_bquant_mx_bf16bf8_rcr_microscale_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x64_qg1x1x128" + ) + + def test_mx_bf16fp4_name(self): + cfg = default_mx_bf16fp4_config(quant_group_k=32) + assert cfg.name == ( + "grouped_gemm_bquant_mx_bf16fp4_rcr_microscale_cshuffle_intrawave_" + "128x128x128_1x4x1_16x16x32_qg1x1x32" + ) + + def test_mx_bf16bf16_warp_tile_k(self): + # GemmConfigQuantPrefill: get_k_warp_tile on gfx950 = 32 + assert default_mx_bf16bf16_config().warp_tile_k == 32 + + def test_mx_bf16bf8_warp_tile_k(self): + # GemmConfigMixedPrecision: hardcoded K_Warp_Tile = 64 + assert default_mx_bf16bf8_config().warp_tile_k == 64 + + def test_mx_bf16fp4_warp_tile_k(self): + # GemmConfigQuantPrefill: same as bf16bf16 + assert default_mx_bf16fp4_config().warp_tile_k == 32 + + def test_mx_all_use_microscale_pipeline(self): + for cfg in [ + default_mx_bf16bf16_config(), + default_mx_bf16bf8_config(), + default_mx_bf16fp4_config(), + ]: + assert cfg.pipeline == "microscale", f"{cfg.variant_key} should use microscale" + + def test_mx_no_preshuffle_flags(self): + for cfg in [ + default_mx_bf16bf16_config(), + default_mx_bf16bf8_config(), + default_mx_bf16fp4_config(), + ]: + assert cfg.preshuffle_b is False + assert cfg.preshuffle_bquant is False + assert cfg.double_smem_buffer is False + + def test_mx_bf16bf16_quant_group_k64(self): + cfg = default_mx_bf16bf16_config(quant_group_k=64) + assert "qg1x1x64" in cfg.name + + def test_mx_bf16fp4_quant_group_k128(self): + cfg = default_mx_bf16fp4_config(quant_group_k=128) + assert "qg1x1x128" in cfg.name + + def test_mx_all_names_unique(self): + configs = [ + default_mx_bf16bf16_config(quant_group_k=32), + default_mx_bf16bf16_config(quant_group_k=64), + default_mx_bf16bf8_config(quant_group_k=64), + default_mx_bf16bf8_config(quant_group_k=128), + default_mx_bf16fp4_config(quant_group_k=32), + default_mx_bf16fp4_config(quant_group_k=64), + default_mx_bf16fp4_config(quant_group_k=128), + ] + names = [c.name for c in configs] + assert len(names) == len(set(names)) + + # ============================================================================= # Name uniqueness across a small sweep # ============================================================================= @@ -223,7 +558,7 @@ class TestNameUniqueness: def _make_configs(self): configs = [] - for variant in ("fp8", "bf8"): + for variant in ("fp8", "bf8", "fp8i4", "bf8i4"): for gk in (64, 128): for gn in (1, 8): configs.append(BQuantKernelConfig( From 5b83676197bc63af1e0c4ac071d2576f68841974 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 7 Jul 2026 15:44:45 -0400 Subject: [PATCH 05/35] [CK][Dispatcher] Fix correctness bugs in BQuantGrouped GEMM dispatcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../ctypes/grouped_gemm_bquant_ctypes_lib.cpp | 22 +++- .../dispatcher/codegen/codegen_common.py | 44 ++++++++ .../unified_grouped_gemm_bquant_codegen.py | 100 +++++++++++------- .../gemm/python/13_grouped_gemm_bquant.py | 60 ++++++++--- .../python/grouped_gemm_bquant_utils.py | 43 ++++---- 5 files changed, 197 insertions(+), 72 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp index 366e278e0021..d546a4f526d7 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp @@ -165,6 +165,26 @@ int dispatcher_run_bquant_gemm( return -1; } + // Validate that the caller's QK_B/QN_B match the compile-time quant group sizes + // baked into this .so. A mismatch means the BQ device buffer would be allocated + // with the wrong size while the kernel indexes it with different strides. + { + const int64_t expected_QK_B = + (K + static_cast(QuantGroupSize::kK) - 1) / QuantGroupSize::kK; + const int64_t expected_QN_B = + (N + static_cast(QuantGroupSize::kN) - 1) / QuantGroupSize::kN; + if(QK_B != expected_QK_B || QN_B != expected_QN_B) + { + std::cerr << "dispatcher_run_bquant_gemm: QK_B/QN_B mismatch. " + << "Got (" << QK_B << ", " << QN_B << "), " + << "expected (" << expected_QK_B << ", " << expected_QN_B << ") " + << "for K=" << K << " N=" << N + << " with QuantGroupSize kK=" << QuantGroupSize::kK + << " kN=" << QuantGroupSize::kN << "\n"; + return -1; + } + } + const ADataType* A_host = static_cast(A); const BDataType* B_host = static_cast(B); const QDataType* BQ_host = static_cast(BQ); @@ -213,7 +233,7 @@ int dispatcher_run_bquant_gemm( args.stride_AQ = 0; args.stride_BQ = static_cast(stride_BQ); - ck_tile::stream_config stream_cfg{nullptr, true, 0, 1, 1, false, false, 1}; + ck_tile::stream_config stream_cfg{nullptr, false, 0, 0, 1, false, false, 1}; float exec_time = SelectedKernel::launch(args, stream_cfg); diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index 1b0858128137..d00baeda4159 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -350,6 +350,50 @@ def needs_pipeline_expansion(config: dict) -> bool: return config.get("pipeline", "compv4") == "*" +# ============================================================================ +# BQuant kernel name construction +# ============================================================================ + + +def make_bquant_kernel_name( + variant_key: str, + layout: str, + pipeline: str, + epilogue: str, + scheduler: str, + tile_m: int, tile_n: int, tile_k: int, + warp_m: int, warp_n: int, warp_k: int, + warp_tile_m: int, warp_tile_n: int, warp_tile_k: int, + quant_group_m: int, + quant_group_n: int, + quant_group_k: int, + preshuffle_b: bool = False, + preshuffle_bquant: bool = False, +) -> str: + """Return the canonical BQuant kernel name used as KERNEL_NAME in generated headers. + + Both BQuantKernelConfig (utils) and BQuantKernelSpec (codegen) delegate to this + function so the two sides are guaranteed to stay byte-exact. + """ + parts = [ + "grouped_gemm_bquant", + variant_key, + layout, + pipeline, + epilogue, + scheduler, + f"{tile_m}x{tile_n}x{tile_k}", + f"{warp_m}x{warp_n}x{warp_k}", + f"{warp_tile_m}x{warp_tile_n}x{warp_tile_k}", + f"qg{quant_group_m}x{quant_group_n}x{quant_group_k}", + ] + if preshuffle_b: + parts.append("preshuffleb") + if preshuffle_bquant: + parts.append("preshufflebq") + return "_".join(parts) + + # ============================================================================ # BQuant-specific Type Mappings # ============================================================================ diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py index 67a9fa6050e3..2660eb9790b9 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py @@ -34,7 +34,7 @@ from pathlib import Path from typing import Dict, List, Optional, Tuple -from codegen_common import TileConfig, parallel_generate, BQUANT_DTYPE_MAP +from codegen_common import make_bquant_kernel_name logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") log = logging.getLogger(__name__) @@ -209,23 +209,21 @@ class BQuantKernelSpec: @property def name(self) -> str: t = self.tile - parts = [ - "grouped_gemm_bquant", - self.variant_key, - self.layout, - self.pipeline, - self.epilogue, - self.scheduler, - f"{t.tile_m}x{t.tile_n}x{t.tile_k}", - f"{t.warp_m}x{t.warp_n}x{t.warp_k}", - f"{t.warp_tile_m}x{t.warp_tile_n}x{t.warp_tile_k}", - f"qg{self.quant_group_m}x{self.quant_group_n}x{self.quant_group_k}", - ] - if self.preshuffle_b: - parts.append("preshuffleb") - if self.preshuffle_bquant: - parts.append("preshufflebq") - return "_".join(parts) + return make_bquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline, + epilogue=self.epilogue, + scheduler=self.scheduler, + tile_m=t.tile_m, tile_n=t.tile_n, tile_k=t.tile_k, + warp_m=t.warp_m, warp_n=t.warp_n, warp_k=t.warp_k, + warp_tile_m=t.warp_tile_m, warp_tile_n=t.warp_tile_n, warp_tile_k=t.warp_tile_k, + quant_group_m=self.quant_group_m, + quant_group_n=self.quant_group_n, + quant_group_k=self.quant_group_k, + preshuffle_b=self.preshuffle_b, + preshuffle_bquant=self.preshuffle_bquant, + ) # ============================================================================= @@ -267,9 +265,53 @@ def generate(self, spec: BQuantKernelSpec) -> str: preshuffle_bquant = str(spec.preshuffle_bquant).lower() double_smem_buffer = str(spec.double_smem_buffer).lower() - # CShuffleEpilogue problem parameters - # TiledMMAPermuteN is false when quant_group_n > 1 (see run_gemm_quant_example.inc) - tiled_permute_n = "false" + # TiledMMAPermuteN = (N_Repeat % 2 == 0) where N_Repeat = TileN / (WarpN * WarpTileN) + # Mirrors GemmConfig::TiledMMAPermuteN in gemm_utils.hpp. + n_repeat = t.tile_n // (t.warp_n * t.warp_tile_n) + tiled_mma_permute_n = (n_repeat % 2 == 0) + # PermuteNEpilogue is used when quant_group_n == 1 and TiledMMAPermuteN is true; + # CShuffleEpilogue otherwise. Mirrors run_gemm_quant_example.inc: + # TiledPermuteN = (BQuantGroupSize::kN > 1) ? false : GemmConfig::TiledMMAPermuteN + use_permute_n_epilogue = tiled_mma_permute_n and spec.quant_group_n == 1 + + # Build the epilogue block outside the f-string to keep it readable. + # PermuteNEpilogueProblem takes two extra trailing args (false, 1) vs CShuffleEpilogueProblem. + if use_permute_n_epilogue: + epilogue_block = f"""\ + using GemmEpilogue = ck_tile::PermuteNEpilogue< + ck_tile::PermuteNEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC, + false, + 1>>;""" + else: + epilogue_block = f"""\ + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC>>;""" return f"""\ // SPDX-License-Identifier: MIT @@ -380,21 +422,7 @@ def generate(self, spec: BQuantKernelSpec) -> str: using GemmPipeline = {pipeline_ck}; - using GemmEpilogue = ck_tile::CShuffleEpilogue< - ck_tile::CShuffleEpilogueProblem< - typename PipelineProblem::AComputeDataType, - typename PipelineProblem::BComputeDataType, - ck_tile::tuple<>, - AccDataType, - CDataType, - ck_tile::tuple<>, - {ns}::CLayout, - ck_tile::element_wise::PassThrough, - TilePartitioner::MPerBlock, - TilePartitioner::NPerBlock, - WarpM, WarpN, - WarpTileM, WarpTileN, WarpTileK, - TransposeC>>; +{epilogue_block} using Kernel = ck_tile::QuantGemmKernel< TilePartitioner, GemmPipeline, GemmEpilogue, diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py index 3fc225941ce0..97b76d022721 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py @@ -55,15 +55,38 @@ # ============================================================================= -def _fp8_to_float32(arr: np.ndarray) -> np.ndarray: - """Cast fp8 byte array to float32 (treating as float8_e4m3fn).""" - # numpy doesn't have fp8; cast uint8 bit pattern to float via ml_dtypes if available, - # otherwise fall back to float32 direct cast (close enough for reference at low values). +def _float32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: + """Encode float32 values as fp8 bytes (uint8 view of the fp8 bit pattern). + + dtype: "fp8" -> float8_e4m3fn, "bf8" -> float8_e5m2. + Falls back to saturating the float32 values to [-448, 448] and storing as + uint8 when ml_dtypes is not installed; the bit patterns are not true fp8 in + that case but the buffer has the correct element size (1 byte/element). + """ + try: + import ml_dtypes + ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return arr.astype(ml_t).view(np.uint8) + except ImportError: + # Fallback: clamp to a small range so values fit in fp8, store as uint8. + # The bit patterns are not genuine fp8, but the buffer size is correct. + clamped = np.clip(arr, -2.0, 2.0) + return (clamped * 64).astype(np.int8).view(np.uint8) + + +def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: + """Decode fp8 bytes (uint8 view) back to float32. + + dtype: "fp8" -> float8_e4m3fn, "bf8" -> float8_e5m2. + Must be called on the same array produced by _float32_to_fp8 to get the + values the kernel actually computes on. + """ try: import ml_dtypes - return arr.view(ml_dtypes.float8_e4m3fn).astype(np.float32) + ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return arr.view(ml_t).astype(np.float32) except ImportError: - return arr.astype(np.float32) + return arr.view(np.int8).astype(np.float32) / 64.0 def reference_bquant_gemm( @@ -75,8 +98,8 @@ def reference_bquant_gemm( """ CPU fp32 reference for C = A @ dequant(B, BQ). - A [M, K] float32 (upcast from fp8/bf8) - B [K, N] float32 (upcast from fp8/bf8) + A [M, K] float32 — must be decoded from the same fp8 bytes sent to the GPU + B [K, N] float32 — must be decoded from the same fp8 bytes sent to the GPU BQ [QK_B, QN_B] float32 scale factors Dequant: B[k, n] *= BQ[k // gK, n // gN] @@ -161,17 +184,21 @@ def main(): QK_B = math.ceil(K / gK) QN_B = math.ceil(N / gN) - # A and B as uint8 (fp8 byte representation) in [-2, 2] float range - # Using simple values so the reference is easy to verify + # Generate float32 values in the fp8 representable range, then encode as + # real fp8 bytes (1 byte/element) so the C layer receives the correct size. A_f32 = rng.uniform(-2.0, 2.0, (M, K)).astype(np.float32) B_f32 = rng.uniform(-2.0, 2.0, (K, N)).astype(np.float32) BQ_f32 = rng.uniform(0.5, 2.0, (QK_B, QN_B)).astype(np.float32) - # Cast to fp8 approximation (uint8 for raw byte passing to C lib) - # Real fp8 encoding would use ml_dtypes; here we use float16 as a stand-in - # that the kernel can safely interpret as fp8 bit patterns for testing. - A_raw = A_f32.astype(np.float16) - B_raw = B_f32.astype(np.float16) + # Encode as fp8 bytes (uint8 view). _float32_to_fp8 uses ml_dtypes when + # available for accurate fp8 bit patterns, with a fallback otherwise. + A_raw = _float32_to_fp8(A_f32, args.dtype) # shape (M, K), 1 byte/element + B_raw = _float32_to_fp8(B_f32, args.dtype) # shape (K, N), 1 byte/element + + # Decode back to float32 so the CPU reference sees the same values the + # kernel will compute on (fp8 encoding introduces rounding). + A_dec = _fp8_to_float32(A_raw, args.dtype) + B_dec = _fp8_to_float32(B_raw, args.dtype) # ------------------------------------------------------------------------- # 4. Run on GPU @@ -191,7 +218,8 @@ def main(): # 5. Verify # ------------------------------------------------------------------------- if not args.no_verify: - C_ref = reference_bquant_gemm(A_f32, B_f32, BQ_f32, problem) + # Reference uses decoded fp8 values — the same bit patterns the kernel sees. + C_ref = reference_bquant_gemm(A_dec, B_dec, BQ_f32, problem) C_gpu = result.C max_rel = float(np.max(np.abs(C_gpu.astype(np.float32) - C_ref.astype(np.float32))) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py index 38b40c490c89..5e27a5579b8d 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py @@ -45,6 +45,13 @@ _CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_grouped_gemm_bquant_codegen.py" _CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "grouped_gemm_bquant_ctypes_lib.cpp" +# Import the shared name-construction helper from codegen_common so both sides +# stay byte-exact without duplicating the logic. +_codegen_dir = str(Path(__file__).parent.parent / "codegen") +if _codegen_dir not in sys.path: + sys.path.insert(0, _codegen_dir) +from codegen_common import make_bquant_kernel_name # noqa: E402 + _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -102,24 +109,22 @@ class BQuantKernelConfig: @property def name(self) -> str: - """Byte-exact match to codegen KERNEL_NAME.""" - parts = [ - "grouped_gemm_bquant", - self.variant_key, - self.layout, - self.pipeline, - self.epilogue, - self.scheduler, - f"{self.tile_m}x{self.tile_n}x{self.tile_k}", - f"{self.warp_m}x{self.warp_n}x{self.warp_k}", - f"{self.warp_tile_m}x{self.warp_tile_n}x{self.warp_tile_k}", - f"qg{self.quant_group_m}x{self.quant_group_n}x{self.quant_group_k}", - ] - if self.preshuffle_b: - parts.append("preshuffleb") - if self.preshuffle_bquant: - parts.append("preshufflebq") - return "_".join(parts) + """Byte-exact match to codegen KERNEL_NAME (delegates to make_bquant_kernel_name).""" + return make_bquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline, + epilogue=self.epilogue, + scheduler=self.scheduler, + tile_m=self.tile_m, tile_n=self.tile_n, tile_k=self.tile_k, + warp_m=self.warp_m, warp_n=self.warp_n, warp_k=self.warp_k, + warp_tile_m=self.warp_tile_m, warp_tile_n=self.warp_tile_n, warp_tile_k=self.warp_tile_k, + quant_group_m=self.quant_group_m, + quant_group_n=self.quant_group_n, + quant_group_k=self.quant_group_k, + preshuffle_b=self.preshuffle_b, + preshuffle_bquant=self.preshuffle_bquant, + ) def to_codegen_config(self) -> dict: """Produce the JSON config dict consumed by unified_grouped_gemm_bquant_codegen.py.""" @@ -574,7 +579,7 @@ def _build_one(idx: int, cfg: BQuantKernelConfig) -> Tuple[int, Optional[Path]]: if parallel and len(deduped) > 1: workers = max_workers or min(len(deduped), os.cpu_count() or 4) - with concurrent.futures.ProcessPoolExecutor(max_workers=workers) as ex: + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex: futures = {ex.submit(_build_one, idx, cfg): (idx, cfg) for idx, cfg in deduped} for fut in concurrent.futures.as_completed(futures): try: From 83d07b68b9db2e644d314a623eb2972a284fd2c5 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 7 Jul 2026 16:17:42 -0400 Subject: [PATCH 06/35] [CK][Dispatcher] Fix clang-format-18 formatting in BQuantGrouped GEMM ctypes lib Co-Authored-By: Claude --- .../ctypes/grouped_gemm_bquant_ctypes_lib.cpp | 128 ++++++++++++------ 1 file changed, 83 insertions(+), 45 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp index d546a4f526d7..baaadf66e553 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp @@ -44,15 +44,15 @@ using Priority = ck_tile::dispatcher::Registry::Priority; static std::shared_ptr g_dispatcher = nullptr; static bool g_initialized = false; -#define HIP_CHECK(call) \ - { \ - hipError_t _err = (call); \ - if(_err != hipSuccess) \ - { \ - std::cerr << "HIP error: " << hipGetErrorString(_err) \ - << " at " << __FILE__ << ":" << __LINE__ << "\n"; \ - return -1; \ - } \ +#define HIP_CHECK(call) \ + { \ + hipError_t _err = (call); \ + if(_err != hipSuccess) \ + { \ + std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ + << __LINE__ << "\n"; \ + return -1; \ + } \ } extern "C" { @@ -87,9 +87,12 @@ int dispatcher_initialize() key.signature.num_d_tensors = 0; key.signature.structured_sparsity = false; - key.algorithm.tile_shape = {SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; - key.algorithm.wave_shape = {SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; - key.algorithm.warp_tile_shape = {SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; + key.algorithm.tile_shape = { + SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; + key.algorithm.wave_shape = { + SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; + key.algorithm.warp_tile_shape = { + SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; key.algorithm.pipeline = Pipeline::CompV3; key.algorithm.scheduler = Scheduler::Intrawave; key.algorithm.epilogue = Epilogue::CShuffle; @@ -132,22 +135,21 @@ int dispatcher_initialize() * * Returns 0 on success, negative on error. */ -int dispatcher_run_bquant_gemm( - const void* A, - const void* B, - const void* BQ, - void* C, - int64_t M, - int64_t N, - int64_t K, - int64_t stride_A, - int64_t stride_B, - int64_t stride_BQ, - int64_t stride_C, - int64_t QK_B, - int64_t QN_B, - int k_batch, - float* time_ms) +int dispatcher_run_bquant_gemm(const void* A, + const void* B, + const void* BQ, + void* C, + int64_t M, + int64_t N, + int64_t K, + int64_t stride_A, + int64_t stride_B, + int64_t stride_BQ, + int64_t stride_C, + int64_t QK_B, + int64_t QN_B, + int k_batch, + float* time_ms) { if(!g_initialized) { @@ -175,10 +177,9 @@ int dispatcher_run_bquant_gemm( (N + static_cast(QuantGroupSize::kN) - 1) / QuantGroupSize::kN; if(QK_B != expected_QK_B || QN_B != expected_QN_B) { - std::cerr << "dispatcher_run_bquant_gemm: QK_B/QN_B mismatch. " - << "Got (" << QK_B << ", " << QN_B << "), " - << "expected (" << expected_QK_B << ", " << expected_QN_B << ") " - << "for K=" << K << " N=" << N + std::cerr << "dispatcher_run_bquant_gemm: QK_B/QN_B mismatch. " << "Got (" << QK_B + << ", " << QN_B << "), " << "expected (" << expected_QK_B << ", " + << expected_QN_B << ") " << "for K=" << K << " N=" << N << " with QuantGroupSize kK=" << QuantGroupSize::kK << " kN=" << QuantGroupSize::kN << "\n"; return -1; @@ -188,7 +189,7 @@ int dispatcher_run_bquant_gemm( const ADataType* A_host = static_cast(A); const BDataType* B_host = static_cast(B); const QDataType* BQ_host = static_cast(BQ); - CDataType* C_host = static_cast(C); + CDataType* C_host = static_cast(C); ADataType* A_dev = nullptr; BDataType* B_dev = nullptr; @@ -196,23 +197,60 @@ int dispatcher_run_bquant_gemm( CDataType* C_dev = nullptr; auto cleanup = [&]() { - if(A_dev) (void)hipFree(A_dev); - if(B_dev) (void)hipFree(B_dev); - if(BQ_dev) (void)hipFree(BQ_dev); - if(C_dev) (void)hipFree(C_dev); + if(A_dev) + (void)hipFree(A_dev); + if(B_dev) + (void)hipFree(B_dev); + if(BQ_dev) + (void)hipFree(BQ_dev); + if(C_dev) + (void)hipFree(C_dev); }; // Allocate device buffers - if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup(); + return -1; + } // Copy inputs to device - if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup(); + return -1; + } // Build QuantGemmHostArgs (aq_ptr = nullptr, QK_A = 0, stride_AQ = 0 for BQuant-only) ck_tile::QuantGemmHostArgs args; From cc17718d82819bb3944e99b1424b46827ba698f5 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 7 Jul 2026 18:42:00 -0400 Subject: [PATCH 07/35] [CK][Dispatcher] Fix non-ASCII em-dash in grouped_gemm_bquant_ctypes_lib.cpp Co-Authored-By: Claude --- .../bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp index baaadf66e553..2eda59e15d21 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp @@ -58,7 +58,7 @@ static bool g_initialized = false; extern "C" { /** - * Initialize dispatcher — must be called before dispatcher_run_bquant_gemm. + * Initialize dispatcher -- must be called before dispatcher_run_bquant_gemm. * * Registers SelectedKernel (from the force-included header) into the Registry. * Returns 0 on success, -1 on error. From b3c453012ea86de7b63672bdadefeddeb88533b5 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 8 Jul 2026 01:34:15 -0400 Subject: [PATCH 08/35] [CK][Dispatcher] Address Copilot review comments on BQuantGrouped GEMM 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 --- .../ctypes/grouped_gemm_bquant_ctypes_lib.cpp | 14 +++++++++ .../dispatcher/codegen/codegen_common.py | 30 +++++++++++++++++-- .../unified_grouped_gemm_bquant_codegen.py | 17 +++++------ .../gemm/python/13_grouped_gemm_bquant.py | 10 +++---- .../python/grouped_gemm_bquant_utils.py | 27 ++++++++++++----- 5 files changed, 74 insertions(+), 24 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp index 2eda59e15d21..5db3499c5921 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp @@ -186,6 +186,20 @@ int dispatcher_run_bquant_gemm(const void* A, } } + // This implementation only supports packed (contiguous) layouts. + // Device buffers are allocated and copied as M*K, K*N, QK_B*QN_B, M*N packed arrays. + // Non-packed strides would cause the kernel to index into a differently-sized buffer, + // producing incorrect results or out-of-bounds accesses. + if(stride_A != K || stride_B != K || stride_BQ != QN_B || stride_C != N) + { + std::cerr << "dispatcher_run_bquant_gemm: non-packed strides are not supported. " + << "Expected stride_A=" << K << " stride_B=" << K << " stride_BQ=" << QN_B + << " stride_C=" << N << ", got stride_A=" << stride_A + << " stride_B=" << stride_B << " stride_BQ=" << stride_BQ + << " stride_C=" << stride_C << "\n"; + return -1; + } + const ADataType* A_host = static_cast(A); const BDataType* B_host = static_cast(B); const QDataType* BQ_host = static_cast(BQ); diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index d00baeda4159..f8ab9decf2fb 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -355,11 +355,31 @@ def needs_pipeline_expansion(config: dict) -> bool: # ============================================================================ +def bquant_effective_epilogue( + tile_n: int, + warp_n: int, + warp_tile_n: int, + quant_group_n: int, +) -> str: + """Return the epilogue tag that the codegen will actually emit for the given tile params. + + Mirrors the TiledMMAPermuteN / use_permute_n_epilogue logic in + unified_grouped_gemm_bquant_codegen.py and run_gemm_quant_example.inc: + TiledMMAPermuteN = (N_Repeat % 2 == 0), N_Repeat = TileN / (WarpN * WarpTileN) + use_permute_n_epilogue = TiledMMAPermuteN and quant_group_n == 1 + + Returns "permute_n" when PermuteNEpilogue is selected, "cshuffle" otherwise. + """ + n_repeat = tile_n // (warp_n * warp_tile_n) + use_permute_n = (n_repeat % 2 == 0) and (quant_group_n == 1) + return "permute_n" if use_permute_n else "cshuffle" + + def make_bquant_kernel_name( variant_key: str, layout: str, pipeline: str, - epilogue: str, + epilogue: str, # ignored — actual epilogue is computed from tile params via bquant_effective_epilogue scheduler: str, tile_m: int, tile_n: int, tile_k: int, warp_m: int, warp_n: int, warp_k: int, @@ -374,13 +394,19 @@ def make_bquant_kernel_name( Both BQuantKernelConfig (utils) and BQuantKernelSpec (codegen) delegate to this function so the two sides are guaranteed to stay byte-exact. + + The epilogue segment in the name reflects the epilogue the codegen actually emits + (computed via bquant_effective_epilogue from tile params) rather than the + user-specified epilogue string, so the name always matches the compiled kernel. + The ``epilogue`` parameter is accepted for call-site compatibility but not used. """ + effective_epilogue = bquant_effective_epilogue(tile_n, warp_n, warp_tile_n, quant_group_n) parts = [ "grouped_gemm_bquant", variant_key, layout, pipeline, - epilogue, + effective_epilogue, scheduler, f"{tile_m}x{tile_n}x{tile_k}", f"{warp_m}x{warp_n}x{warp_k}", diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py index 2660eb9790b9..690a2e16b0b6 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_bquant_codegen.py @@ -34,7 +34,7 @@ from pathlib import Path from typing import Dict, List, Optional, Tuple -from codegen_common import make_bquant_kernel_name +from codegen_common import make_bquant_kernel_name, bquant_effective_epilogue logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") log = logging.getLogger(__name__) @@ -265,14 +265,13 @@ def generate(self, spec: BQuantKernelSpec) -> str: preshuffle_bquant = str(spec.preshuffle_bquant).lower() double_smem_buffer = str(spec.double_smem_buffer).lower() - # TiledMMAPermuteN = (N_Repeat % 2 == 0) where N_Repeat = TileN / (WarpN * WarpTileN) - # Mirrors GemmConfig::TiledMMAPermuteN in gemm_utils.hpp. - n_repeat = t.tile_n // (t.warp_n * t.warp_tile_n) - tiled_mma_permute_n = (n_repeat % 2 == 0) - # PermuteNEpilogue is used when quant_group_n == 1 and TiledMMAPermuteN is true; - # CShuffleEpilogue otherwise. Mirrors run_gemm_quant_example.inc: - # TiledPermuteN = (BQuantGroupSize::kN > 1) ? false : GemmConfig::TiledMMAPermuteN - use_permute_n_epilogue = tiled_mma_permute_n and spec.quant_group_n == 1 + # Determine which epilogue the kernel will use, mirroring run_gemm_quant_example.inc. + # Delegates to bquant_effective_epilogue (same logic used by make_bquant_kernel_name) + # so the generated C++ and the kernel name always agree. + use_permute_n_epilogue = ( + bquant_effective_epilogue(t.tile_n, t.warp_n, t.warp_tile_n, spec.quant_group_n) + == "permute_n" + ) # Build the epilogue block outside the f-string to keep it readable. # PermuteNEpilogueProblem takes two extra trailing args (false, 1) vs CShuffleEpilogueProblem. diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py index 97b76d022721..2fe1449e9f74 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/13_grouped_gemm_bquant.py @@ -19,10 +19,10 @@ - CK include path discoverable relative to this repo Usage: - python3 13_bquant_gemm.py # fp8, 1x1x128 groups, M=16 N=64 K=256 - python3 13_bquant_gemm.py --dtype bf8 - python3 13_bquant_gemm.py --dtype fp8 --M 32 --N 128 --K 512 --quant-group-k 128 - python3 13_bquant_gemm.py --no-verify # skip CPU reference check + python3 13_grouped_gemm_bquant.py # fp8, 1x1x128 groups, M=16 N=64 K=256 + python3 13_grouped_gemm_bquant.py --dtype bf8 + python3 13_grouped_gemm_bquant.py --dtype fp8 --M 32 --N 128 --K 512 --quant-group-k 128 + python3 13_grouped_gemm_bquant.py --no-verify # skip CPU reference check """ import argparse @@ -59,7 +59,7 @@ def _float32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: """Encode float32 values as fp8 bytes (uint8 view of the fp8 bit pattern). dtype: "fp8" -> float8_e4m3fn, "bf8" -> float8_e5m2. - Falls back to saturating the float32 values to [-448, 448] and storing as + Falls back to clamping the float32 values to [-2.0, 2.0] and storing as uint8 when ml_dtypes is not installed; the bit patterns are not true fp8 in that case but the buffer has the correct element size (1 byte/element). """ diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py index 5e27a5579b8d..bb3dd03b4e8d 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py @@ -278,7 +278,11 @@ def run( """ Call dispatcher_run_bquant_gemm with ctypes-wrapped pointers. - A, B, BQ, C should be numpy arrays (contiguous). + A, B, BQ, C must be numpy arrays (C-contiguous, packed). + B should be a packed (K, N) C-contiguous array — the kernel interprets + it as column-major via stride_B=K, not via numpy's Fortran-order flag. + C must be the array that will receive output; a non-contiguous C would + produce a temporary copy that is not returned to the caller. Returns (status, time_ms). """ import numpy as np @@ -345,13 +349,17 @@ def __init__(self, so_path: Path): def kernel_name(self) -> str: return self._lib.get_kernel_name() - def run(self, A, B, BQ, problem: BQuantGemmProblem) -> BQuantGemmResult: + def run(self, A, B, BQ, problem: BQuantGemmProblem, c_dtype=None) -> BQuantGemmResult: """ Run BQuantGrouped GEMM. - A shape: (M, K) dtype: fp8/bf8 - B shape: (K, N) col-major dtype: fp8/bf8 - BQ shape: (QK_B, QN_B) dtype: float/fp8 + A shape: (M, K) dtype: fp8/bf8 + B shape: (K, N) col-major dtype: fp8/bf8 + BQ shape: (QK_B, QN_B) dtype: float/fp8 + c_dtype numpy dtype for the output C buffer. Defaults to np.float16 + (correct for fp8/bf8/fp8i4/bf8i4 variants). Pass np.bfloat16 + for MX variants (mx_bf16bf16, mx_bf16bf8, mx_bf16fp4) whose + CDataType is bf16. Returns BQuantGemmResult with C shape (M, N). """ import numpy as np @@ -360,8 +368,11 @@ def run(self, A, B, BQ, problem: BQuantGemmProblem) -> BQuantGemmResult: QK_B = problem.QK_B QN_B = problem.QN_B - # Output buffer — dtype matches CDataType (half for fp8/bf8 variants) - C = np.zeros((M, N), dtype=np.float16) + if c_dtype is None: + c_dtype = np.float16 + + # Output buffer — dtype must match the compiled kernel's CDataType. + C = np.zeros((M, N), dtype=c_dtype) # Strides (in elements, row-major for A and C; col-major for B means stride = K) stride_A = K # A is row-major [M, K] @@ -563,7 +574,7 @@ def _build_one(idx: int, cfg: BQuantKernelConfig) -> Tuple[int, Optional[Path]]: if hpp is None: return idx, None - so = so_dir / f"lib{cfg.name}.so" + so = so_dir / f"lib{cfg.name}_{arch}.so" if so.exists(): log.info(" [cached] %s", so.name) return idx, so From f6a6000324bbf35cbcb5613d9ccfc03fca390c18 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 8 Jul 2026 12:27:52 -0400 Subject: [PATCH 09/35] [CK][Dispatcher] Fix clang-format-18 formatting in grouped_gemm_bquant_ctypes_lib.cpp Co-Authored-By: Claude --- .../bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp index 5db3499c5921..b95b8b2d95d3 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_bquant_ctypes_lib.cpp @@ -194,9 +194,8 @@ int dispatcher_run_bquant_gemm(const void* A, { std::cerr << "dispatcher_run_bquant_gemm: non-packed strides are not supported. " << "Expected stride_A=" << K << " stride_B=" << K << " stride_BQ=" << QN_B - << " stride_C=" << N << ", got stride_A=" << stride_A - << " stride_B=" << stride_B << " stride_BQ=" << stride_BQ - << " stride_C=" << stride_C << "\n"; + << " stride_C=" << N << ", got stride_A=" << stride_A << " stride_B=" << stride_B + << " stride_BQ=" << stride_BQ << " stride_C=" << stride_C << "\n"; return -1; } From fe9e91ed531252b1f28625a93408d45f8a547170 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 8 Jul 2026 14:21:05 -0400 Subject: [PATCH 10/35] feat(ck-tile): add AQuant and ABQuant grouped GEMM dispatcher bridges MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 62 ++ .../grouped_gemm_abquant_ctypes_lib.cpp | 317 ++++++++ .../ctypes/grouped_gemm_aquant_ctypes_lib.cpp | 332 ++++++++ .../dispatcher/codegen/codegen_common.py | 131 +++ .../unified_grouped_gemm_abquant_codegen.py | 657 +++++++++++++++ .../unified_grouped_gemm_aquant_codegen.py | 632 +++++++++++++++ .../gemm/python/14_grouped_gemm_aquant.py | 242 ++++++ .../gemm/python/15_grouped_gemm_abquant.py | 266 ++++++ .../python/grouped_gemm_abquant_utils.py | 762 ++++++++++++++++++ .../python/grouped_gemm_aquant_utils.py | 741 +++++++++++++++++ .../tests/test_grouped_gemm_abquant_utils.py | 363 +++++++++ .../tests/test_grouped_gemm_aquant_utils.py | 410 ++++++++++ 12 files changed, 4915 insertions(+) create mode 100644 projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp create mode 100644 projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp create mode 100644 projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py create mode 100644 projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_aquant_codegen.py create mode 100644 projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py create mode 100644 projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py create mode 100644 projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py create mode 100644 projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py create mode 100644 projects/composablekernel/dispatcher/tests/test_grouped_gemm_abquant_utils.py create mode 100644 projects/composablekernel/dispatcher/tests/test_grouped_gemm_aquant_utils.py diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 6a9132a5fcfe..082c6f33f9fc 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -185,6 +185,68 @@ else() ) endif() +# ============================================================================= +# GroupedGemm AQuant ctypes Library +# ============================================================================= + +file(GLOB AQUANT_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/grouped_gemm_aquant_*.hpp") +if(AQUANT_GEMM_KERNEL_HEADERS) + list(GET AQUANT_GEMM_KERNEL_HEADERS 0 AQUANT_GEMM_KERNEL_HEADER) + message(STATUS "Found AQuant GEMM kernel for ctypes lib: ${AQUANT_GEMM_KERNEL_HEADER}") + + add_ctypes_library(dispatcher_grouped_gemm_aquant_lib + grouped_gemm_aquant_ctypes_lib.cpp + KERNEL_HEADER ${AQUANT_GEMM_KERNEL_HEADER} + ) + target_compile_definitions(dispatcher_grouped_gemm_aquant_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE + GFX_ARCH="gfx950" + ) +else() + message(STATUS "No AQuant GEMM kernel found for ctypes lib - building without kernel") + add_library(dispatcher_grouped_gemm_aquant_lib SHARED grouped_gemm_aquant_ctypes_lib.cpp) + target_include_directories(dispatcher_grouped_gemm_aquant_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ${PROJECT_SOURCE_DIR}/dispatcher/include + ) + target_link_libraries(dispatcher_grouped_gemm_aquant_lib PRIVATE hip::device) + set_target_properties(dispatcher_grouped_gemm_aquant_lib PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD 17 + ) +endif() + +# ============================================================================= +# GroupedGemm ABQuant ctypes Library +# ============================================================================= + +file(GLOB ABQUANT_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/grouped_gemm_abquant_*.hpp") +if(ABQUANT_GEMM_KERNEL_HEADERS) + list(GET ABQUANT_GEMM_KERNEL_HEADERS 0 ABQUANT_GEMM_KERNEL_HEADER) + message(STATUS "Found ABQuant GEMM kernel for ctypes lib: ${ABQUANT_GEMM_KERNEL_HEADER}") + + add_ctypes_library(dispatcher_grouped_gemm_abquant_lib + grouped_gemm_abquant_ctypes_lib.cpp + KERNEL_HEADER ${ABQUANT_GEMM_KERNEL_HEADER} + ) + target_compile_definitions(dispatcher_grouped_gemm_abquant_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE + GFX_ARCH="gfx950" + ) +else() + message(STATUS "No ABQuant GEMM kernel found for ctypes lib - building without kernel") + add_library(dispatcher_grouped_gemm_abquant_lib SHARED grouped_gemm_abquant_ctypes_lib.cpp) + target_include_directories(dispatcher_grouped_gemm_abquant_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ${PROJECT_SOURCE_DIR}/dispatcher/include + ) + target_link_libraries(dispatcher_grouped_gemm_abquant_lib PRIVATE hip::device) + set_target_properties(dispatcher_grouped_gemm_abquant_lib PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD 17 + ) +endif() + # ============================================================================= # GPU Helper Executable # ============================================================================= diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp new file mode 100644 index 000000000000..855ec91b73a5 --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -0,0 +1,317 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * GroupedGemm ABQuant Dispatcher ctypes Library + * + * Provides C API for Python ctypes integration. + * Kernel header is force-included at compile time via: + * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE grouped_gemm_abquant_ctypes_lib.cpp + * + * Force-include defines (from generated kernel header): + * SelectedKernel, KERNEL_NAME + * ADataType, BDataType, CDataType, QDataType, AccDataType + * AQuantGroupSize, BQuantGroupSize + * + * ABQuant: both A-side and B-side quantization active. + * AQ[ceil(M/aM), ceil(K/aK)] — A-side scale tensor (RowMajor) + * BQ[ceil(K/bK), ceil(N/bN)] — B-side scale tensor (RowMajor) + * Constraint: AQuantGroupSize::kK == BQuantGroupSize::kK + * + * Memory model: host-pointer (this library owns hipMalloc/hipMemcpy/hipFree). + */ + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/dispatcher/dispatcher.hpp" +#include "ck_tile/dispatcher/registry.hpp" +#include "ck_tile/dispatcher/backends/generated_tile_backend.hpp" + +// Kernel header force-included via -include compiler flag. +// Defines: ADataType, BDataType, CDataType, QDataType (shared AQ/BQ type), AccDataType, +// AQuantGroupSize, BQuantGroupSize, SelectedKernel, KERNEL_NAME + +#ifndef GFX_ARCH +#define GFX_ARCH "gfx950" +#endif + +using namespace ck_tile::dispatcher; +using namespace ck_tile::dispatcher::backends; +using Priority = ck_tile::dispatcher::Registry::Priority; + +static std::shared_ptr g_dispatcher = nullptr; +static bool g_initialized = false; + +#define HIP_CHECK(call) \ + { \ + hipError_t _err = (call); \ + if(_err != hipSuccess) \ + { \ + std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ + << __LINE__ << "\n"; \ + return -1; \ + } \ + } + +extern "C" { + +/** + * Initialize dispatcher -- must be called before dispatcher_run_abquant_gemm. + * Returns 0 on success, -1 on error. + */ +int dispatcher_initialize() +{ + if(g_initialized) + return 0; + + KernelKey key; + key.signature.dtype_a = DataType::FP8; + key.signature.dtype_b = DataType::FP8; + key.signature.dtype_c = DataType::FP16; + key.signature.dtype_acc = DataType::FP32; + key.signature.layout_a = LayoutTag::RowMajor; + key.signature.layout_b = LayoutTag::ColMajor; + key.signature.layout_c = LayoutTag::RowMajor; + key.signature.transpose_a = false; + key.signature.transpose_b = false; + key.signature.grouped = false; + key.signature.split_k = 1; + key.signature.elementwise_op = "PassThrough"; + key.signature.num_d_tensors = 0; + key.signature.structured_sparsity = false; + + key.algorithm.tile_shape = { + SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; + key.algorithm.wave_shape = { + SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; + key.algorithm.warp_tile_shape = { + SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; + key.algorithm.pipeline = Pipeline::CompV3; + key.algorithm.scheduler = Scheduler::Intrawave; + key.algorithm.epilogue = Epilogue::CShuffle; + key.algorithm.block_size = SelectedKernel::BlockSize; + key.algorithm.double_buffer = SelectedKernel::DoubleSmemBuffer; + key.algorithm.persistent = false; + key.algorithm.preshuffle = SelectedKernel::PreshuffleB; + key.algorithm.transpose_c = SelectedKernel::TransposeC; + key.algorithm.num_wave_groups = 1; + key.gfx_arch = GFX_ARCH; + + auto kernel = + create_generated_tile_kernel( + key, KERNEL_NAME); + + Registry::instance().clear(); + Registry::instance().register_kernel(kernel, Priority::High); + + g_dispatcher = std::make_shared(); + g_initialized = true; + return 0; +} + +/** + * Run ABQuantGrouped GEMM: + * C[M,N] = dequant(A[M,K], AQ[ceil(M/aM), ceil(K/aK)]) @ dequant(B[K,N], BQ[ceil(K/bK), ceil(N/bN)]) + * + * A, B, AQ, BQ, C are host pointers. This function manages device memory internally. + * + * Parameters: + * A, B, AQ, BQ, C — host data pointers + * M, N, K — matrix dimensions + * stride_A — leading dim of A (row-major: K) + * stride_B — leading dim of B (col-major: K) + * stride_AQ — leading dim of AQ (row-major: ceil(K/aK)) + * stride_BQ — leading dim of BQ (row-major: ceil(N/bN)) + * stride_C — leading dim of C (row-major: N) + * QK_A — ceil(K / aquant_group_k) + * QM_A — ceil(M / aquant_group_m) (typically == M when aM=1) + * QK_B — ceil(K / bquant_group_k) + * QN_B — ceil(N / bquant_group_n) + * k_batch — split-K factor (1 = no split) + * time_ms — output: kernel execution time in ms (may be NULL) + * + * Returns 0 on success, negative on error. + */ +int dispatcher_run_abquant_gemm(const void* A, + const void* B, + const void* AQ, + const void* BQ, + void* C, + int64_t M, + int64_t N, + int64_t K, + int64_t stride_A, + int64_t stride_B, + int64_t stride_AQ, + int64_t stride_BQ, + int64_t stride_C, + int64_t QK_A, + int64_t QM_A, + int64_t QK_B, + int64_t QN_B, + int k_batch, + float* time_ms) +{ + if(!g_initialized) + { + std::cerr << "dispatcher_run_abquant_gemm: not initialized\n"; + return -1; + } + if(!A || !B || !AQ || !BQ || !C) + { + std::cerr << "dispatcher_run_abquant_gemm: null pointer argument\n"; + return -1; + } + if(M <= 0 || N <= 0 || K <= 0 || QK_A <= 0 || QM_A <= 0 || QK_B <= 0 || QN_B <= 0) + { + std::cerr << "dispatcher_run_abquant_gemm: invalid dimensions\n"; + return -1; + } + + // Validate AQ dimensions + { + const int64_t exp_QK_A = + (K + static_cast(AQuantGroupSize::kK) - 1) / AQuantGroupSize::kK; + const int64_t exp_QM_A = + (M + static_cast(AQuantGroupSize::kM) - 1) / AQuantGroupSize::kM; + if(QK_A != exp_QK_A || QM_A != exp_QM_A) + { + std::cerr << "dispatcher_run_abquant_gemm: QK_A/QM_A mismatch. " + << "Got (" << QK_A << ", " << QM_A << "), " + << "expected (" << exp_QK_A << ", " << exp_QM_A << ")\n"; + return -1; + } + } + + // Validate BQ dimensions + { + const int64_t exp_QK_B = + (K + static_cast(BQuantGroupSize::kK) - 1) / BQuantGroupSize::kK; + const int64_t exp_QN_B = + (N + static_cast(BQuantGroupSize::kN) - 1) / BQuantGroupSize::kN; + if(QK_B != exp_QK_B || QN_B != exp_QN_B) + { + std::cerr << "dispatcher_run_abquant_gemm: QK_B/QN_B mismatch. " + << "Got (" << QK_B << ", " << QN_B << "), " + << "expected (" << exp_QK_B << ", " << exp_QN_B << ")\n"; + return -1; + } + } + + // Only packed (contiguous) layouts are supported. + // AQ is RowMajor [QM_A, QK_A]: stride_AQ == QK_A + // BQ is RowMajor [QK_B, QN_B]: stride_BQ == QN_B + if(stride_A != K || stride_B != K || stride_AQ != QK_A || stride_BQ != QN_B || stride_C != N) + { + std::cerr << "dispatcher_run_abquant_gemm: non-packed strides not supported. " + << "Expected stride_A=" << K << " stride_B=" << K + << " stride_AQ=" << QK_A << " stride_BQ=" << QN_B << " stride_C=" << N + << "\n"; + return -1; + } + + const ADataType* A_host = static_cast(A); + const BDataType* B_host = static_cast(B); + const QDataType* AQ_host = static_cast(AQ); + const QDataType* BQ_host = static_cast(BQ); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + BDataType* B_dev = nullptr; + QDataType* AQ_dev = nullptr; + QDataType* BQ_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup = [&]() { + if(A_dev) (void)hipFree(A_dev); + if(B_dev) (void)hipFree(B_dev); + if(AQ_dev) (void)hipFree(AQ_dev); + if(BQ_dev) (void)hipFree(BQ_dev); + if(C_dev) (void)hipFree(C_dev); + }; + + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&AQ_dev, QM_A * QK_A * sizeof(QDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(AQ_dev, AQ_host, QM_A * QK_A * sizeof(QDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + + // Build QuantGemmHostArgs: both AQ and BQ active + ck_tile::QuantGemmHostArgs args; + args.a_ptr = A_dev; + args.b_ptr = B_dev; + args.aq_ptr = AQ_dev; // A-side scale: active + args.bq_ptr = BQ_dev; // B-side scale: active + args.c_ptr = C_dev; + args.k_batch = k_batch; + args.M = static_cast(M); + args.N = static_cast(N); + args.K = static_cast(K); + args.QK_A = static_cast(QK_A); + args.QK_B = static_cast(QK_B); + args.stride_A = static_cast(stride_A); + args.stride_B = static_cast(stride_B); + args.stride_C = static_cast(stride_C); + args.stride_AQ = static_cast(stride_AQ); + args.stride_BQ = static_cast(stride_BQ); + + ck_tile::stream_config stream_cfg{nullptr, false, 0, 0, 1, false, false, 1}; + + float exec_time = SelectedKernel::launch(args, stream_cfg); + + if(exec_time < 0.0f) + { + std::cerr << "dispatcher_run_abquant_gemm: kernel reported unsupported args\n"; + cleanup(); + return -2; + } + + if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + { + cleanup(); + return -1; + } + + if(time_ms) + *time_ms = exec_time; + + cleanup(); + return 0; +} + +/** + * Return the compile-time KERNEL_NAME of the force-included kernel. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Initialize dispatcher (alias for consistency). + */ +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Number of kernels registered (always 1 for single-kernel-per-.so model). + */ +int dispatcher_get_kernel_count() { return static_cast(Registry::instance().size()); } + +/** + * Release dispatcher resources. + */ +void dispatcher_cleanup() +{ + g_dispatcher.reset(); + g_initialized = false; +} + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp new file mode 100644 index 000000000000..886bcaafda5b --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp @@ -0,0 +1,332 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * GroupedGemm AQuant Dispatcher ctypes Library + * + * Provides C API for Python ctypes integration. + * Kernel header is force-included at compile time via: + * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE grouped_gemm_aquant_ctypes_lib.cpp + * + * Force-include defines (from generated kernel header): + * SelectedKernel, KERNEL_NAME + * ADataType, BDataType, CDataType, QDataType (=AQDataType), AccDataType, QuantGroupSize + * + * AQuant: A-side activation quantization. + * AQ[ceil(M/gM), ceil(K/gK)] is the A-side scale tensor (RowMajor). + * BQ is unused (bq_ptr=nullptr, QK_B=0, stride_BQ=0 in QuantGemmHostArgs). + * + * Memory model: host-pointer (this library owns hipMalloc/hipMemcpy/hipFree). + */ + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/dispatcher/dispatcher.hpp" +#include "ck_tile/dispatcher/registry.hpp" +#include "ck_tile/dispatcher/backends/generated_tile_backend.hpp" + +// Kernel header force-included via -include compiler flag. +// Defines: ADataType, BDataType, CDataType, QDataType (AQDataType), AccDataType, +// QuantGroupSize, SelectedKernel, KERNEL_NAME + +#ifndef GFX_ARCH +#define GFX_ARCH "gfx950" +#endif + +using namespace ck_tile::dispatcher; +using namespace ck_tile::dispatcher::backends; +using Priority = ck_tile::dispatcher::Registry::Priority; + +static std::shared_ptr g_dispatcher = nullptr; +static bool g_initialized = false; + +#define HIP_CHECK(call) \ + { \ + hipError_t _err = (call); \ + if(_err != hipSuccess) \ + { \ + std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ + << __LINE__ << "\n"; \ + return -1; \ + } \ + } + +extern "C" { + +/** + * Initialize dispatcher -- must be called before dispatcher_run_aquant_gemm. + * + * Registers SelectedKernel (from the force-included header) into the Registry. + * Returns 0 on success, -1 on error. + */ +int dispatcher_initialize() +{ + if(g_initialized) + return 0; + + KernelKey key; + key.signature.dtype_a = DataType::FP8; + key.signature.dtype_b = DataType::FP8; + key.signature.dtype_c = DataType::FP16; + key.signature.dtype_acc = DataType::FP32; + key.signature.layout_a = LayoutTag::RowMajor; + key.signature.layout_b = LayoutTag::ColMajor; + key.signature.layout_c = LayoutTag::RowMajor; + key.signature.transpose_a = false; + key.signature.transpose_b = false; + key.signature.grouped = false; + key.signature.split_k = 1; + key.signature.elementwise_op = "PassThrough"; + key.signature.num_d_tensors = 0; + key.signature.structured_sparsity = false; + + key.algorithm.tile_shape = { + SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; + key.algorithm.wave_shape = { + SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; + key.algorithm.warp_tile_shape = { + SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; + key.algorithm.pipeline = Pipeline::CompV3; + key.algorithm.scheduler = Scheduler::Intrawave; + key.algorithm.epilogue = Epilogue::CShuffle; + key.algorithm.block_size = SelectedKernel::BlockSize; + key.algorithm.double_buffer = false; + key.algorithm.persistent = false; + key.algorithm.preshuffle = SelectedKernel::APreshuffleQuant; + key.algorithm.transpose_c = SelectedKernel::TransposeC; + key.algorithm.num_wave_groups = 1; + key.gfx_arch = GFX_ARCH; + + auto kernel = + create_generated_tile_kernel( + key, KERNEL_NAME); + + Registry::instance().clear(); + Registry::instance().register_kernel(kernel, Priority::High); + + g_dispatcher = std::make_shared(); + g_initialized = true; + return 0; +} + +/** + * Run AQuantGrouped GEMM: C[M,N] = dequant(A[M,K], AQ[ceil(M/gM), ceil(K/gK)]) @ B[K,N] + * + * A, B, AQ, C are host pointers. This function manages device memory internally. + * + * Parameters: + * A, B, AQ, C - host data pointers + * M, N, K - matrix dimensions + * stride_A - leading dimension of A (row-major: K) + * stride_B - leading dimension of B (col-major: K) + * stride_AQ - leading dimension of AQ (row-major: ceil(K/gK)) + * stride_C - leading dimension of C (row-major: N) + * QK_A - number of K-groups for A = ceil(K / quant_group_k) + * QM_A - number of M-groups for A = ceil(M / quant_group_m) (typically == M when gM=1) + * k_batch - split-K factor (1 = no split) + * time_ms - output: kernel execution time in ms (may be NULL) + * + * Returns 0 on success, negative on error. + */ +int dispatcher_run_aquant_gemm(const void* A, + const void* B, + const void* AQ, + void* C, + int64_t M, + int64_t N, + int64_t K, + int64_t stride_A, + int64_t stride_B, + int64_t stride_AQ, + int64_t stride_C, + int64_t QK_A, + int64_t QM_A, + int k_batch, + float* time_ms) +{ + if(!g_initialized) + { + std::cerr << "dispatcher_run_aquant_gemm: not initialized\n"; + return -1; + } + if(!A || !B || !AQ || !C) + { + std::cerr << "dispatcher_run_aquant_gemm: null pointer argument\n"; + return -1; + } + if(M <= 0 || N <= 0 || K <= 0 || QK_A <= 0 || QM_A <= 0) + { + std::cerr << "dispatcher_run_aquant_gemm: invalid dimensions\n"; + return -1; + } + + // Validate that the caller's QK_A/QM_A match the compile-time quant group sizes. + { + const int64_t expected_QK_A = + (K + static_cast(QuantGroupSize::kK) - 1) / QuantGroupSize::kK; + const int64_t expected_QM_A = + (M + static_cast(QuantGroupSize::kM) - 1) / QuantGroupSize::kM; + if(QK_A != expected_QK_A || QM_A != expected_QM_A) + { + std::cerr << "dispatcher_run_aquant_gemm: QK_A/QM_A mismatch. " + << "Got (" << QK_A << ", " << QM_A << "), " + << "expected (" << expected_QK_A << ", " << expected_QM_A << ") " + << "for K=" << K << " M=" << M + << " with QuantGroupSize kK=" << QuantGroupSize::kK + << " kM=" << QuantGroupSize::kM << "\n"; + return -1; + } + } + + // Only packed (contiguous) layouts are supported. + // AQ is RowMajor [QM_A, QK_A]: stride_AQ == QK_A + if(stride_A != K || stride_B != K || stride_AQ != QK_A || stride_C != N) + { + std::cerr << "dispatcher_run_aquant_gemm: non-packed strides are not supported. " + << "Expected stride_A=" << K << " stride_B=" << K << " stride_AQ=" << QK_A + << " stride_C=" << N << ", got stride_A=" << stride_A + << " stride_B=" << stride_B << " stride_AQ=" << stride_AQ + << " stride_C=" << stride_C << "\n"; + return -1; + } + + const ADataType* A_host = static_cast(A); + const BDataType* B_host = static_cast(B); + const QDataType* AQ_host = static_cast(AQ); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + BDataType* B_dev = nullptr; + QDataType* AQ_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup = [&]() { + if(A_dev) + (void)hipFree(A_dev); + if(B_dev) + (void)hipFree(B_dev); + if(AQ_dev) + (void)hipFree(AQ_dev); + if(C_dev) + (void)hipFree(C_dev); + }; + + // Allocate device buffers + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&AQ_dev, QM_A * QK_A * sizeof(QDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + + // Copy inputs to device + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(AQ_dev, AQ_host, QM_A * QK_A * sizeof(QDataType), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + + // Build QuantGemmHostArgs: AQ active, BQ unused + ck_tile::QuantGemmHostArgs args; + args.a_ptr = A_dev; + args.b_ptr = B_dev; + args.aq_ptr = AQ_dev; // A-side scale: active for AQuant + args.bq_ptr = nullptr; // B-side scale: unused for AQuant-only + args.c_ptr = C_dev; + args.k_batch = k_batch; + args.M = static_cast(M); + args.N = static_cast(N); + args.K = static_cast(K); + args.QK_A = static_cast(QK_A); + args.QK_B = 0; // unused + args.stride_A = static_cast(stride_A); + args.stride_B = static_cast(stride_B); + args.stride_C = static_cast(stride_C); + args.stride_AQ = static_cast(stride_AQ); + args.stride_BQ = 0; // unused + + ck_tile::stream_config stream_cfg{nullptr, false, 0, 0, 1, false, false, 1}; + + float exec_time = SelectedKernel::launch(args, stream_cfg); + + if(exec_time < 0.0f) + { + std::cerr << "dispatcher_run_aquant_gemm: kernel reported unsupported args\n"; + cleanup(); + return -2; + } + + // Copy result back + if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + { + cleanup(); + return -1; + } + + if(time_ms) + *time_ms = exec_time; + + cleanup(); + return 0; +} + +/** + * Return the compile-time KERNEL_NAME of the force-included kernel. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Initialize dispatcher (alias kept for consistency). + */ +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Number of kernels registered (always 1 for the single-kernel-per-.so model). + */ +int dispatcher_get_kernel_count() { return static_cast(Registry::instance().size()); } + +/** + * Release dispatcher resources. + */ +void dispatcher_cleanup() +{ + g_dispatcher.reset(); + g_initialized = false; +} + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index f8ab9decf2fb..1facc7dc85a9 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -437,3 +437,134 @@ def make_bquant_kernel_name( "float": "float", "e8m0": "ck_tile::e8m0_t", } + + +# ============================================================================ +# AQuant kernel name construction +# ============================================================================ + + +def aquant_effective_epilogue( + tile_n: int, + warp_n: int, + warp_tile_n: int, + quant_group_n: int, +) -> str: + """Return the epilogue tag the codegen will emit for AQuant kernels. + + Mirrors the same TiledMMAPermuteN / use_permute_n_epilogue logic as BQuant + (the PermuteN condition is driven by B-side tile geometry, regardless of + which side is quantised). Returns "permute_n" when PermuteNEpilogue is + selected, "cshuffle" otherwise. + """ + n_repeat = tile_n // (warp_n * warp_tile_n) + use_permute_n = (n_repeat % 2 == 0) and (quant_group_n == 1) + return "permute_n" if use_permute_n else "cshuffle" + + +def make_aquant_kernel_name( + variant_key: str, + layout: str, + pipeline: str, + epilogue: str, # ignored — actual epilogue is computed from tile params + scheduler: str, + tile_m: int, tile_n: int, tile_k: int, + warp_m: int, warp_n: int, warp_k: int, + warp_tile_m: int, warp_tile_n: int, warp_tile_k: int, + quant_group_m: int, + quant_group_n: int, + quant_group_k: int, + preshuffle_aq: bool = False, +) -> str: + """Return the canonical AQuant kernel name used as KERNEL_NAME in generated headers. + + Both AQuantKernelConfig (utils) and AQuantKernelSpec (codegen) delegate to + this function so the two sides are guaranteed to stay byte-exact. + """ + effective_epilogue = aquant_effective_epilogue(tile_n, warp_n, warp_tile_n, quant_group_n) + parts = [ + "grouped_gemm_aquant", + variant_key, + layout, + pipeline, + effective_epilogue, + scheduler, + f"{tile_m}x{tile_n}x{tile_k}", + f"{warp_m}x{warp_n}x{warp_k}", + f"{warp_tile_m}x{warp_tile_n}x{warp_tile_k}", + f"aqg{quant_group_m}x{quant_group_n}x{quant_group_k}", + ] + if preshuffle_aq: + parts.append("preshuffleaq") + return "_".join(parts) + + +# ============================================================================ +# ABQuant kernel name construction +# ============================================================================ + + +def abquant_effective_epilogue( + tile_n: int, + warp_n: int, + warp_tile_n: int, + bquant_group_n: int, +) -> str: + """Return the epilogue tag the codegen will emit for ABQuant kernels. + + The PermuteN selection is governed by the B-side quant group N, matching + the same logic used by BQuant / AQuant. + """ + n_repeat = tile_n // (warp_n * warp_tile_n) + use_permute_n = (n_repeat % 2 == 0) and (bquant_group_n == 1) + return "permute_n" if use_permute_n else "cshuffle" + + +def make_abquant_kernel_name( + variant_key: str, + layout: str, + pipeline: str, + epilogue: str, # ignored — actual epilogue is computed from tile params + scheduler: str, + tile_m: int, tile_n: int, tile_k: int, + warp_m: int, warp_n: int, warp_k: int, + warp_tile_m: int, warp_tile_n: int, warp_tile_k: int, + aquant_group_m: int, + aquant_group_n: int, + aquant_group_k: int, + bquant_group_m: int, + bquant_group_n: int, + bquant_group_k: int, + preshuffle_b: bool = False, + preshuffle_aq: bool = False, + preshuffle_bq: bool = False, + transpose_c: bool = False, +) -> str: + """Return the canonical ABQuant kernel name used as KERNEL_NAME in generated headers. + + Both ABQuantKernelConfig (utils) and ABQuantKernelSpec (codegen) delegate + to this function so the two sides are guaranteed to stay byte-exact. + """ + effective_epilogue = abquant_effective_epilogue(tile_n, warp_n, warp_tile_n, bquant_group_n) + parts = [ + "grouped_gemm_abquant", + variant_key, + layout, + pipeline, + effective_epilogue, + scheduler, + f"{tile_m}x{tile_n}x{tile_k}", + f"{warp_m}x{warp_n}x{warp_k}", + f"{warp_tile_m}x{warp_tile_n}x{warp_tile_k}", + f"aqg{aquant_group_m}x{aquant_group_n}x{aquant_group_k}", + f"bqg{bquant_group_m}x{bquant_group_n}x{bquant_group_k}", + ] + if preshuffle_b: + parts.append("preshuffleb") + if preshuffle_aq: + parts.append("preshuffleaq") + if preshuffle_bq: + parts.append("preshufflebq") + if transpose_c: + parts.append("transposec") + return "_".join(parts) diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py new file mode 100644 index 000000000000..72ef77acc08c --- /dev/null +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py @@ -0,0 +1,657 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +GroupedGemm ABQuant Code Generator + +Generates one .hpp per kernel config for the dispatcher's ctypes path. +Each header defines a SelectedKernel struct with a static launch() method +taking QuantGemmHostArgs — compiled per-kernel via force-include: + + hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE grouped_gemm_abquant_ctypes_lib.cpp + +ABQuant: A-side and B-side quantization simultaneously (QuantType::ABQuantGrouped). +Covers fp8, bf8 dtype variants with separate AQuantGroupSize and BQuantGroupSize. + +Constraint: AQuantGroupSize::kK == BQuantGroupSize::kK (enforced at codegen time). + +Pipeline selection (mirrors run_gemm_quant_example.inc): + compv3: ABQuantGemmPipelineAgBgCrCompV3 (GemmConfigABQuantPrefill, non-gfx950) + eightwaves: ABQuantGemmPipelineAgBgCrEightWaves (GemmConfigEightWaves, gfx950) + preshuffleb: WPABQuantBPipelineAgBgCrV2 (GemmConfigPreshuffleB_ABQuant_Prefill) + +Naming convention (byte-exact with ABQuantKernelConfig.name in grouped_gemm_abquant_utils.py): + grouped_gemm_abquant_{variant_key}_{layout}_{pipeline}_{epilogue}_{scheduler}_ + {TileM}x{TileN}x{TileK}_{WarpM}x{WarpN}x{WarpK}_{WtM}x{WtN}x{WtK}_ + aqg{aM}x{aN}x{aK}_bqg{bM}x{bN}x{bK} + [_preshuffleb][_preshuffleaq][_preshufflebq][_transposec] + +Reference: + example/ck_tile/38_block_scale_gemm/gemm_abquant_quantgrouped_fp8.cpp + example/ck_tile/38_block_scale_gemm/gemm_abquant_quantgrouped.h + example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc + example/ck_tile/38_block_scale_gemm/gemm_utils.hpp +""" + +import argparse +import itertools +import json +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +from codegen_common import make_abquant_kernel_name, abquant_effective_epilogue + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# Dtype variant definitions +# Both AQ and BQ share the same QDataType for ABQuant. +# ============================================================================= + +ABQUANT_VARIANTS: Dict[str, Dict[str, str]] = { + "fp8": { + "dtype_a": "fp8", + "dtype_b": "fp8", + "dtype_c": "half", + "dtype_q": "float", # both AQDataType and BQDataType + "ck_a": "ck_tile::fp8_t", + "ck_b": "ck_tile::fp8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + }, + "bf8": { + "dtype_a": "bf8", + "dtype_b": "bf8", + "dtype_c": "half", + "dtype_q": "float", + "ck_a": "ck_tile::bf8_t", + "ck_b": "ck_tile::bf8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + }, +} + +ABQUANT_LAYOUT_TO_CK = { + "r": "ck_tile::tensor_layout::gemm::RowMajor", + "c": "ck_tile::tensor_layout::gemm::ColumnMajor", +} + +# Pipeline map for ABQuant kernels. +# compv3 → ABQuantGemmPipelineAgBgCrCompV3 (non-gfx950, GemmConfigABQuantPrefill) +# eightwaves → ABQuantGemmPipelineAgBgCrEightWaves (gfx950, GemmConfigEightWaves) +# preshuffleb→ WPABQuantBPipelineAgBgCrV2 (GemmConfigPreshuffleB_ABQuant_Prefill) +ABQUANT_PIPELINE_MAP = { + "compv3": "ck_tile::ABQuantGemmPipelineAgBgCrCompV3", + "eightwaves": "ck_tile::ABQuantGemmPipelineAgBgCrEightWaves", + "preshuffleb": "ck_tile::WPABQuantBPipelineAgBgCrV2", +} + +ABQUANT_BASE_PIPELINE_MAP = { + "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", + "eightwaves": "ck_tile::BaseGemmPipelineAgBgCrCompV3", + "preshuffleb": "ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2", +} + +ABQUANT_SCHEDULER_TO_CK = { + "intrawave": "ck_tile::GemmPipelineScheduler::Intrawave", + "interwave": "ck_tile::GemmPipelineScheduler::Interwave", +} + + +# ============================================================================= +# Configuration dataclasses +# ============================================================================= + + +@dataclass +class ABQuantTileConfig: + tile_m: int + tile_n: int + tile_k: int + warp_m: int + warp_n: int + warp_k: int + warp_tile_m: int + warp_tile_n: int + warp_tile_k: int + + def is_valid(self) -> bool: + if self.tile_m <= 0 or self.tile_n <= 0 or self.tile_k <= 0: + return False + return ( + self.tile_m % (self.warp_m * self.warp_tile_m) == 0 + and self.tile_n % (self.warp_n * self.warp_tile_n) == 0 + and self.tile_k % (self.warp_k * self.warp_tile_k) == 0 + ) + + +@dataclass +class ABQuantKernelSpec: + """Complete specification for one ABQuant kernel.""" + + variant_key: str # "fp8", "bf8" + layout: str # "rcr" + pipeline: str # "compv3", "eightwaves", "preshuffleb" + epilogue: str # "cshuffle" (effective may be permute_n) + scheduler: str # "intrawave" + tile: ABQuantTileConfig + + # Separate group sizes for A-side and B-side quantization + aquant_group_m: int = 1 + aquant_group_n: int = 1 + aquant_group_k: int = 128 + + bquant_group_m: int = 1 + bquant_group_n: int = 1 + bquant_group_k: int = 128 + + preshuffle_b: bool = False # PreshuffleB (weight preshuffle) + preshuffle_aq: bool = False # APreshuffleQuant + preshuffle_bq: bool = False # BPreshuffleQuant + transpose_c: bool = False # TransposeC (true for eightwaves/gfx950) + double_smem_buffer: bool = False + pad_m: bool = False + pad_n: bool = False + pad_k: bool = False # ABQuant prefill tiles have kPadK=false + block_size: int = 256 + k_block_per_cu: int = 1 + + def __post_init__(self): + if self.aquant_group_k != self.bquant_group_k: + raise ValueError( + f"ABQuant requires AQuantGroupSize::kK == BQuantGroupSize::kK, " + f"got {self.aquant_group_k} != {self.bquant_group_k}" + ) + + @property + def name(self) -> str: + t = self.tile + return make_abquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline, + epilogue=self.epilogue, + scheduler=self.scheduler, + tile_m=t.tile_m, tile_n=t.tile_n, tile_k=t.tile_k, + warp_m=t.warp_m, warp_n=t.warp_n, warp_k=t.warp_k, + warp_tile_m=t.warp_tile_m, warp_tile_n=t.warp_tile_n, warp_tile_k=t.warp_tile_k, + aquant_group_m=self.aquant_group_m, + aquant_group_n=self.aquant_group_n, + aquant_group_k=self.aquant_group_k, + bquant_group_m=self.bquant_group_m, + bquant_group_n=self.bquant_group_n, + bquant_group_k=self.bquant_group_k, + preshuffle_b=self.preshuffle_b, + preshuffle_aq=self.preshuffle_aq, + preshuffle_bq=self.preshuffle_bq, + transpose_c=self.transpose_c, + ) + + +# ============================================================================= +# Header generator +# ============================================================================= + + +class ABQuantKernelHeaderGenerator: + """Generates a .hpp kernel specialization header for one ABQuantKernelSpec.""" + + def generate(self, spec: ABQuantKernelSpec) -> str: + variant = ABQUANT_VARIANTS[spec.variant_key] + t = spec.tile + ns = "ns_" + spec.name + struct = "Kernel_" + spec.name + + ck_a = variant["ck_a"] + ck_b = variant["ck_b"] + ck_c = variant["ck_c"] + ck_q = variant["ck_q"] # same type for both AQ and BQ + ck_acc = variant["ck_acc"] + + layout_a_ck = ABQUANT_LAYOUT_TO_CK[spec.layout[0]] + layout_b_ck = ABQUANT_LAYOUT_TO_CK[spec.layout[1]] + layout_c_ck = ABQUANT_LAYOUT_TO_CK[spec.layout[2]] + # AQ is RowMajor: [ceil(M/gM), ceil(K/gK)] + layout_aq_ck = ABQUANT_LAYOUT_TO_CK["r"] + # BQ is RowMajor: [ceil(K/gK), ceil(N/gN)] + layout_bq_ck = ABQUANT_LAYOUT_TO_CK["r"] + + pipeline_ck = ABQUANT_PIPELINE_MAP[spec.pipeline] + base_pipeline_ck = ABQUANT_BASE_PIPELINE_MAP[spec.pipeline] + scheduler_ck = ABQUANT_SCHEDULER_TO_CK[spec.scheduler] + + pad_m = str(spec.pad_m).lower() + pad_n = str(spec.pad_n).lower() + pad_k = str(spec.pad_k).lower() + preshuffle_b = str(spec.preshuffle_b).lower() + preshuffle_aq = str(spec.preshuffle_aq).lower() + preshuffle_bq = str(spec.preshuffle_bq).lower() + transpose_c = str(spec.transpose_c).lower() + double_smem_buffer = str(spec.double_smem_buffer).lower() + + # Epilogue selection — B-side tile geometry governs PermuteN (same logic as BQuant) + use_permute_n_epilogue = ( + abquant_effective_epilogue(t.tile_n, t.warp_n, t.warp_tile_n, spec.bquant_group_n) + == "permute_n" + ) + + if use_permute_n_epilogue: + epilogue_block = f"""\ + using GemmEpilogue = ck_tile::PermuteNEpilogue< + ck_tile::PermuteNEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC, + false, + 1>>;""" + else: + epilogue_block = f"""\ + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC>>;""" + + return f"""\ +// SPDX-License-Identifier: MIT +// Auto-generated ABQuantGrouped GEMM kernel header. +// DO NOT EDIT — regenerate via unified_grouped_gemm_abquant_codegen.py +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "ck_tile/ops/epilogue.hpp" + +namespace {ns} {{ + +constexpr const char* KERNEL_NAME = "{spec.name}"; + +using ADataType = {ck_a}; +using BDataType = {ck_b}; +using CDataType = {ck_c}; +using QDataType = {ck_q}; // shared by AQDataType and BQDataType +using AccDataType = {ck_acc}; + +using ALayout = {layout_a_ck}; +using BLayout = {layout_b_ck}; +using CLayout = {layout_c_ck}; +using AQLayout = {layout_aq_ck}; // RowMajor: [ceil(M/aM), ceil(K/aK)] +using BQLayout = {layout_bq_ck}; // RowMajor: [ceil(K/bK), ceil(N/bN)] + +// Separate group sizes — ABQuant requires AQuantGroupSize::kK == BQuantGroupSize::kK +using AQuantGroupSize = ck_tile::QuantGroupShape>; +using BQuantGroupSize = ck_tile::QuantGroupShape>; + +struct {struct} {{ + using ADataType = {ns}::ADataType; + using BDataType = {ns}::BDataType; + using CDataType = {ns}::CDataType; + using QDataType = {ns}::QDataType; + using AccDataType = {ns}::AccDataType; + + static constexpr ck_tile::index_t TileM = {t.tile_m}; + static constexpr ck_tile::index_t TileN = {t.tile_n}; + static constexpr ck_tile::index_t TileK = {t.tile_k}; + static constexpr ck_tile::index_t WarpM = {t.warp_m}; + static constexpr ck_tile::index_t WarpN = {t.warp_n}; + static constexpr ck_tile::index_t WarpK = {t.warp_k}; + static constexpr ck_tile::index_t WarpTileM = {t.warp_tile_m}; + static constexpr ck_tile::index_t WarpTileN = {t.warp_tile_n}; + static constexpr ck_tile::index_t WarpTileK = {t.warp_tile_k}; + static constexpr ck_tile::index_t BlockSize = {spec.block_size}; + static constexpr int kBlockPerCu = {spec.k_block_per_cu}; + + static constexpr bool kPadM = {pad_m}; + static constexpr bool kPadN = {pad_n}; + static constexpr bool kPadK = {pad_k}; + static constexpr bool APreshuffleQuant = {preshuffle_aq}; + static constexpr bool BPreshuffleQuant = {preshuffle_bq}; + static constexpr bool PreshuffleB = {preshuffle_b}; + static constexpr bool TransposeC = {transpose_c}; + static constexpr bool DoubleSmemBuffer = {double_smem_buffer}; + + using TileShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + + using GemmTraits = ck_tile::TileGemmQuantTraits< + kPadM, kPadN, kPadK, + APreshuffleQuant, BPreshuffleQuant, PreshuffleB, + {ns}::ALayout, {ns}::BLayout, {ns}::CLayout, + ck_tile::QuantType::ABQuantGrouped, + {ns}::AQLayout, {ns}::BQLayout, + TransposeC, DoubleSmemBuffer>; + + using GemmPipelineProblemBase = ck_tile::GemmPipelineProblemBase< + ADataType, BDataType, AccDataType, TileShape, GemmTraits>; + + using BaseGemmPipeline = {base_pipeline_ck}; + + static float launch(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + {{ + const ck_tile::index_t K_split = + (args.k_batch == 1) + ? ck_tile::integer_least_multiple(args.K, TileK) + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, TileK); + + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](auto has_hot_loop_, auto tail_number_) {{ + using PipelineProblem = ck_tile::GemmABQuantPipelineProblem< + ADataType, + QDataType, // AQDataType (A-side scale, same type as BQ) + BDataType, + QDataType, // BQDataType (B-side scale) + AccDataType, // 5th arg is AccDataType, matching run_gemm_quant_example.inc + TileShape, + GemmTraits, + AQuantGroupSize, + BQuantGroupSize, + TransposeC, + ADataType, // AComputeDataType + {scheduler_ck}, + has_hot_loop_.value, + tail_number_.value>; + + using GemmPipeline = {pipeline_ck}; + +{epilogue_block} + + using Kernel = ck_tile::QuantGemmKernel< + TilePartitioner, GemmPipeline, GemmEpilogue, + ck_tile::QuantType::ABQuantGrouped>; + + auto kargs = Kernel::MakeKernelArgs(args); + if(!Kernel::IsSupportedArgument(kargs)) + return -1.0f; + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + return ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{{}}, grids, blocks, 0, kargs)); + }}; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + }} +}}; + +using SelectedKernel = {struct}; + +}} // namespace {ns} + +#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE +using SelectedKernel = {ns}::{struct}; +constexpr const char* KERNEL_NAME = {ns}::KERNEL_NAME; +using ADataType = {ck_a}; +using BDataType = {ck_b}; +using CDataType = {ck_c}; +using QDataType = {ck_q}; +using AccDataType = {ck_acc}; +using AQuantGroupSize = {ns}::AQuantGroupSize; +using BQuantGroupSize = {ns}::BQuantGroupSize; +#endif // CK_TILE_SINGLE_KERNEL_INCLUDE +""" + + +# ============================================================================= +# Config sweep +# ============================================================================= + + +def _default_config() -> dict: + """Default sweep config: fp8 ABQuant, compv3 pipeline, 128x128x128 tile (non-gfx950).""" + return { + "variant_keys": ["fp8", "bf8"], + "layouts": ["rcr"], + "pipeline": "compv3", + "epilogue": "cshuffle", + "scheduler": "intrawave", + "tile_configs": [ + # GemmConfigABQuantPrefill: M=128, N=128, K=128, Warp 1x4x1 + {"tile_m": 128, "tile_n": 128, "tile_k": 128, + "warp_m": 1, "warp_n": 4, "warp_k": 1, + "warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 16}, + ], + "aquant_groups": [ + {"aquant_group_m": 1, "aquant_group_n": 1, "aquant_group_k": 128}, + ], + "bquant_groups": [ + {"bquant_group_m": 1, "bquant_group_n": 1, "bquant_group_k": 128}, + ], + "preshuffle_b": False, + "preshuffle_aq": False, + "preshuffle_bq": False, + "transpose_c": False, + "double_smem_buffer": False, + "pad_m": False, + "pad_n": False, + "pad_k": False, + "block_size": 256, + "k_block_per_cu": 1, + } + + +def _build_specs(config: dict) -> List[ABQuantKernelSpec]: + specs = [] + pipeline = config.get("pipeline", "compv3") + epilogue = config.get("epilogue", "cshuffle") + scheduler = config.get("scheduler", "intrawave") + pad_m = config.get("pad_m", False) + pad_n = config.get("pad_n", False) + pad_k = config.get("pad_k", False) + block_size = config.get("block_size", 256) + k_block_per_cu = config.get("k_block_per_cu", 1) + double_smem_buffer = config.get("double_smem_buffer", False) + preshuffle_b = config.get("preshuffle_b", False) + preshuffle_aq = config.get("preshuffle_aq", False) + preshuffle_bq = config.get("preshuffle_bq", False) + transpose_c = config.get("transpose_c", False) + + default_aquant = [{"aquant_group_m": 1, "aquant_group_n": 1, "aquant_group_k": 128}] + default_bquant = [{"bquant_group_m": 1, "bquant_group_n": 1, "bquant_group_k": 128}] + + for variant_key, layout, tile_dict, aqg, bqg in itertools.product( + config.get("variant_keys", ["fp8"]), + config.get("layouts", ["rcr"]), + config.get("tile_configs", []), + config.get("aquant_groups", default_aquant), + config.get("bquant_groups", default_bquant), + ): + if variant_key not in ABQUANT_VARIANTS: + log.warning("Unknown variant_key %s — skipping", variant_key) + continue + if pipeline not in ABQUANT_PIPELINE_MAP: + log.warning("Unsupported pipeline %s — skipping", pipeline) + continue + + aqk = aqg.get("aquant_group_k", 128) + bqk = bqg.get("bquant_group_k", 128) + if aqk != bqk: + log.warning( + "ABQuant kK mismatch: aquant_group_k=%d != bquant_group_k=%d — skipping", + aqk, bqk, + ) + continue + + tile = ABQuantTileConfig( + tile_m=tile_dict["tile_m"], + tile_n=tile_dict["tile_n"], + tile_k=tile_dict["tile_k"], + warp_m=tile_dict["warp_m"], + warp_n=tile_dict["warp_n"], + warp_k=tile_dict["warp_k"], + warp_tile_m=tile_dict["warp_tile_m"], + warp_tile_n=tile_dict["warp_tile_n"], + warp_tile_k=tile_dict["warp_tile_k"], + ) + if not tile.is_valid(): + log.debug("Invalid tile config %s — skipping", tile) + continue + + try: + specs.append(ABQuantKernelSpec( + variant_key=variant_key, + layout=layout, + pipeline=pipeline, + epilogue=epilogue, + scheduler=scheduler, + tile=tile, + aquant_group_m=aqg.get("aquant_group_m", 1), + aquant_group_n=aqg.get("aquant_group_n", 1), + aquant_group_k=aqk, + bquant_group_m=bqg.get("bquant_group_m", 1), + bquant_group_n=bqg.get("bquant_group_n", 1), + bquant_group_k=bqk, + preshuffle_b=preshuffle_b, + preshuffle_aq=preshuffle_aq, + preshuffle_bq=preshuffle_bq, + transpose_c=transpose_c, + double_smem_buffer=double_smem_buffer, + pad_m=pad_m, + pad_n=pad_n, + pad_k=pad_k, + block_size=block_size, + k_block_per_cu=k_block_per_cu, + )) + except ValueError as e: + log.warning("Skipping invalid spec: %s", e) + + return specs + + +# ============================================================================= +# Generation entry point +# ============================================================================= + + +def generate_kernels( + output_dir: Path, + config: Optional[dict] = None, + parallel: bool = True, +) -> List[Path]: + """Generate all ABQuant kernel headers into output_dir. Returns list of generated .hpp paths.""" + output_dir.mkdir(parents=True, exist_ok=True) + cfg = config or _default_config() + specs = _build_specs(cfg) + + if not specs: + log.warning("No kernel specs produced — check variant_keys, tile_configs, and kK equality") + return [] + + log.info("Generating %d ABQuant kernel headers into %s", len(specs), output_dir) + + gen = ABQuantKernelHeaderGenerator() + generated: List[Path] = [] + + def _generate_one(spec: ABQuantKernelSpec) -> Path: + header = gen.generate(spec) + out_path = output_dir / f"{spec.name}.hpp" + out_path.write_text(header) + log.info(" wrote %s", out_path.name) + return out_path + + if parallel and len(specs) > 1: + import concurrent.futures + with concurrent.futures.ThreadPoolExecutor() as ex: + futures = {ex.submit(_generate_one, s): s for s in specs} + for fut in concurrent.futures.as_completed(futures): + try: + generated.append(fut.result()) + except Exception as e: + log.error("Failed generating %s: %s", futures[fut].name, e) + else: + for spec in specs: + try: + generated.append(_generate_one(spec)) + except Exception as e: + log.error("Failed generating %s: %s", spec.name, e) + + log.info("Generated %d / %d headers", len(generated), len(specs)) + return generated + + +# ============================================================================= +# CLI +# ============================================================================= + + +def main() -> int: + parser = argparse.ArgumentParser( + description="ABQuantGrouped GEMM kernel header generator" + ) + parser.add_argument("--output-dir", type=Path, required=True, + help="Directory to write generated .hpp files") + parser.add_argument("--config", type=Path, + help="JSON config file") + parser.add_argument("--config-json", type=str, + help="Inline JSON config string") + parser.add_argument("--no-parallel", action="store_true", + help="Disable parallel generation") + parser.add_argument("--list-names", action="store_true", + help="Print kernel names that would be generated and exit") + args = parser.parse_args() + + cfg: Optional[dict] = None + if args.config_json: + try: + cfg = json.loads(args.config_json) + except json.JSONDecodeError as e: + log.error("Invalid --config-json: %s", e) + return 1 + elif args.config: + with open(args.config) as f: + cfg = json.load(f) + + if args.list_names: + specs = _build_specs(cfg or _default_config()) + for s in specs: + print(s.name) + return 0 + + paths = generate_kernels( + output_dir=args.output_dir, + config=cfg, + parallel=not args.no_parallel, + ) + return 0 if paths else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_aquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_aquant_codegen.py new file mode 100644 index 000000000000..17c0230fe07c --- /dev/null +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_aquant_codegen.py @@ -0,0 +1,632 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +GroupedGemm AQuant Code Generator + +Generates one .hpp per kernel config for the dispatcher's ctypes path. +Each header defines a SelectedKernel struct with a static launch() method +taking QuantGemmHostArgs — compiled per-kernel via force-include: + + hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE grouped_gemm_aquant_ctypes_lib.cpp + +Covers fp8, bf8, fp8i4 and bf8i4 dtype variants with configurable QuantGroupShape. +APreshuffleQuant is a trait flag on AQuantGemmPipelineAgBgCrCompV3 (not a separate pipeline). + +Naming convention (byte-exact with AQuantKernelConfig.name in grouped_gemm_aquant_utils.py): + grouped_gemm_aquant_{variant_key}_{layout}_{pipeline}_{epilogue}_{scheduler}_ + {TileM}x{TileN}x{TileK}_{WarpM}x{WarpN}x{WarpK}_{WtM}x{WtN}x{WtK}_ + aqg{gM}x{gN}x{gK}[_preshuffleaq] + +Reference: + example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped.cpp + example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped_preshufflequant.cpp + example/ck_tile/38_block_scale_gemm/gemm_utils.hpp (GemmConfigQuantDecodeInterwave) +""" + +import argparse +import itertools +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional + +from codegen_common import make_aquant_kernel_name, aquant_effective_epilogue + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# Dtype variant definitions +# Each entry: (variant_key, ADataType, BDataType, CDataType, QDataType=AQDataType) +# Matches example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped*.cpp +# ============================================================================= + +AQUANT_VARIANTS: Dict[str, Dict[str, str]] = { + "fp8": { + "dtype_a": "fp8", + "dtype_b": "fp8", + "dtype_c": "half", + "dtype_q": "float", # AQ scale type + "ck_a": "ck_tile::fp8_t", + "ck_b": "ck_tile::fp8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", # AQDataType + "ck_acc": "float", + }, + "bf8": { + "dtype_a": "bf8", + "dtype_b": "bf8", + "dtype_c": "half", + "dtype_q": "float", + "ck_a": "ck_tile::bf8_t", + "ck_b": "ck_tile::bf8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + }, + # fp8i4/bf8i4: A is fp8/bf8 (scale side), B is int4 (weight) + # QDataType (AQDataType) is fp8/bf8 — the A-side scale + "fp8i4": { + "dtype_a": "fp8", + "dtype_b": "pk_int4", + "dtype_c": "half", + "dtype_q": "fp8", + "ck_a": "ck_tile::fp8_t", + "ck_b": "ck_tile::pk_int4_t", + "ck_c": "ck_tile::half_t", + "ck_q": "ck_tile::fp8_t", + "ck_acc": "float", + }, + "bf8i4": { + "dtype_a": "bf8", + "dtype_b": "pk_int4", + "dtype_c": "half", + "dtype_q": "bf8", + "ck_a": "ck_tile::bf8_t", + "ck_b": "ck_tile::pk_int4_t", + "ck_c": "ck_tile::half_t", + "ck_q": "ck_tile::bf8_t", + "ck_acc": "float", + }, +} + +# Standard rcr layout: RowMajor A, ColMajor B, RowMajor C +AQUANT_LAYOUT_TO_CK = { + "r": "ck_tile::tensor_layout::gemm::RowMajor", + "c": "ck_tile::tensor_layout::gemm::ColumnMajor", +} + +# AQuant pipeline selection: +# - non-preshuffle (APreshuffleQuant=false) → AQuantGemmPipelineAgBgCrMem / BaseGemmPipelineAgBgCrMem +# - preshuffle (APreshuffleQuant=true) → AQuantGemmPipelineAgBgCrCompV3 / BaseGemmPipelineAgBgCrCompV3 +# APreshuffleQuant is a traits flag that also switches the pipeline class. +AQUANT_PIPELINE_MAP = { + "mem": "ck_tile::AQuantGemmPipelineAgBgCrMem", # non-preshuffle + "compv3": "ck_tile::AQuantGemmPipelineAgBgCrCompV3", # preshuffle (APreshuffleQuant=true) +} + +AQUANT_BASE_PIPELINE_MAP = { + "mem": "ck_tile::BaseGemmPipelineAgBgCrMem", + "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", +} + +AQUANT_SCHEDULER_TO_CK = { + "intrawave": "ck_tile::GemmPipelineScheduler::Intrawave", + "interwave": "ck_tile::GemmPipelineScheduler::Interwave", +} + + +# ============================================================================= +# Configuration dataclasses +# ============================================================================= + + +@dataclass +class AQuantTileConfig: + tile_m: int + tile_n: int + tile_k: int + warp_m: int + warp_n: int + warp_k: int + warp_tile_m: int + warp_tile_n: int + warp_tile_k: int + + def is_valid(self) -> bool: + if self.tile_m <= 0 or self.tile_n <= 0 or self.tile_k <= 0: + return False + return ( + self.tile_m % (self.warp_m * self.warp_tile_m) == 0 + and self.tile_n % (self.warp_n * self.warp_tile_n) == 0 + and self.tile_k % (self.warp_k * self.warp_tile_k) == 0 + ) + + +@dataclass +class AQuantKernelSpec: + """Complete specification for one AQuant kernel.""" + + variant_key: str # "fp8", "bf8", "fp8i4", "bf8i4" + layout: str # "rcr" + pipeline: str # "mem" (non-preshuffle) or "compv3" (preshuffle) + epilogue: str # "cshuffle" or "permute_n" (effective epilogue overrides) + scheduler: str # "intrawave" + tile: AQuantTileConfig + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + preshuffle_aq: bool = False # APreshuffleQuant trait flag + double_smem_buffer: bool = False + pad_m: bool = False + pad_n: bool = False + pad_k: bool = True + block_size: int = 256 + k_block_per_cu: int = 1 + transpose_c: bool = False + + @property + def name(self) -> str: + t = self.tile + return make_aquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline, + epilogue=self.epilogue, + scheduler=self.scheduler, + tile_m=t.tile_m, tile_n=t.tile_n, tile_k=t.tile_k, + warp_m=t.warp_m, warp_n=t.warp_n, warp_k=t.warp_k, + warp_tile_m=t.warp_tile_m, warp_tile_n=t.warp_tile_n, warp_tile_k=t.warp_tile_k, + quant_group_m=self.quant_group_m, + quant_group_n=self.quant_group_n, + quant_group_k=self.quant_group_k, + preshuffle_aq=self.preshuffle_aq, + ) + + +# ============================================================================= +# Header generator +# ============================================================================= + + +class AQuantKernelHeaderGenerator: + """Generates a .hpp kernel specialization header for one AQuantKernelSpec.""" + + def generate(self, spec: AQuantKernelSpec) -> str: + variant = AQUANT_VARIANTS[spec.variant_key] + t = spec.tile + ns = "ns_" + spec.name + struct = "Kernel_" + spec.name + + ck_a = variant["ck_a"] + ck_b = variant["ck_b"] + ck_c = variant["ck_c"] + ck_q = variant["ck_q"] # AQDataType + ck_acc = variant["ck_acc"] + + layout_a_ck = AQUANT_LAYOUT_TO_CK[spec.layout[0]] + layout_b_ck = AQUANT_LAYOUT_TO_CK[spec.layout[1]] + layout_c_ck = AQUANT_LAYOUT_TO_CK[spec.layout[2]] + # AQ is RowMajor: scale is stored [ceil(M/gM), ceil(K/gK)] + layout_aq_ck = AQUANT_LAYOUT_TO_CK["r"] + # BQ layout placeholder (unused for AQuant-only) + layout_bq_ck = layout_b_ck + + pipeline_ck = AQUANT_PIPELINE_MAP[spec.pipeline] + base_pipeline_ck = AQUANT_BASE_PIPELINE_MAP[spec.pipeline] + scheduler_ck = AQUANT_SCHEDULER_TO_CK[spec.scheduler] + + pad_m = str(spec.pad_m).lower() + pad_n = str(spec.pad_n).lower() + pad_k = str(spec.pad_k).lower() + preshuffle_aq = str(spec.preshuffle_aq).lower() + double_smem_buffer = str(spec.double_smem_buffer).lower() + transpose_c = str(spec.transpose_c).lower() + + # Determine epilogue — same PermuteN logic as BQuant (B-side tile geometry governs) + use_permute_n_epilogue = ( + aquant_effective_epilogue(t.tile_n, t.warp_n, t.warp_tile_n, spec.quant_group_n) + == "permute_n" + ) + + if use_permute_n_epilogue: + epilogue_block = f"""\ + using GemmEpilogue = ck_tile::PermuteNEpilogue< + ck_tile::PermuteNEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC, + false, + 1>>;""" + else: + epilogue_block = f"""\ + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem< + typename PipelineProblem::AComputeDataType, + typename PipelineProblem::BComputeDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + ck_tile::tuple<>, + {ns}::CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + WarpM, WarpN, + WarpTileM, WarpTileN, WarpTileK, + TransposeC>>;""" + + return f"""\ +// SPDX-License-Identifier: MIT +// Auto-generated AQuantGrouped GEMM kernel header. +// DO NOT EDIT — regenerate via unified_grouped_gemm_aquant_codegen.py +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "ck_tile/ops/epilogue.hpp" + +namespace {ns} {{ + +constexpr const char* KERNEL_NAME = "{spec.name}"; + +using ADataType = {ck_a}; +using BDataType = {ck_b}; +using CDataType = {ck_c}; +using QDataType = {ck_q}; // AQDataType (A-side scale) +using AccDataType = {ck_acc}; + +using ALayout = {layout_a_ck}; +using BLayout = {layout_b_ck}; +using CLayout = {layout_c_ck}; +using AQLayout = {layout_aq_ck}; +using BQLayout = {layout_bq_ck}; + +// AQuantGroupSize: scale is [ceil(M/gM), ceil(K/gK)] in RowMajor +using QuantGroupSize = ck_tile::QuantGroupShape>; + +struct {struct} {{ + using ADataType = {ns}::ADataType; + using BDataType = {ns}::BDataType; + using CDataType = {ns}::CDataType; + using QDataType = {ns}::QDataType; + using AccDataType = {ns}::AccDataType; + + static constexpr ck_tile::index_t TileM = {t.tile_m}; + static constexpr ck_tile::index_t TileN = {t.tile_n}; + static constexpr ck_tile::index_t TileK = {t.tile_k}; + static constexpr ck_tile::index_t WarpM = {t.warp_m}; + static constexpr ck_tile::index_t WarpN = {t.warp_n}; + static constexpr ck_tile::index_t WarpK = {t.warp_k}; + static constexpr ck_tile::index_t WarpTileM = {t.warp_tile_m}; + static constexpr ck_tile::index_t WarpTileN = {t.warp_tile_n}; + static constexpr ck_tile::index_t WarpTileK = {t.warp_tile_k}; + static constexpr ck_tile::index_t BlockSize = {spec.block_size}; + static constexpr int kBlockPerCu = {spec.k_block_per_cu}; + + static constexpr bool kPadM = {pad_m}; + static constexpr bool kPadN = {pad_n}; + static constexpr bool kPadK = {pad_k}; + static constexpr bool APreshuffleQuant = {preshuffle_aq}; + static constexpr bool BPreshuffleQuant = false; + static constexpr bool PreshuffleB = false; + static constexpr bool TransposeC = {transpose_c}; + static constexpr bool DoubleSmemBuffer = {double_smem_buffer}; + + using TileShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + + using GemmTraits = ck_tile::TileGemmQuantTraits< + kPadM, kPadN, kPadK, + APreshuffleQuant, BPreshuffleQuant, PreshuffleB, + {ns}::ALayout, {ns}::BLayout, {ns}::CLayout, + ck_tile::QuantType::AQuantGrouped, + {ns}::AQLayout, {ns}::BQLayout, + TransposeC, DoubleSmemBuffer>; + + using GemmPipelineProblemBase = ck_tile::GemmPipelineProblemBase< + ADataType, BDataType, AccDataType, TileShape, GemmTraits>; + + using BaseGemmPipeline = {base_pipeline_ck}; + + static float launch(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + {{ + const ck_tile::index_t K_split = + (args.k_batch == 1) + ? ck_tile::integer_least_multiple(args.K, TileK) + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, TileK); + + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](auto has_hot_loop_, auto tail_number_) {{ + using PipelineProblem = ck_tile::GemmAQuantPipelineProblem< + ADataType, + QDataType, // AQDataType (A-side scale type) + BDataType, + AccDataType, // 4th arg is AccDataType, matching run_gemm_quant_example.inc + TileShape, + GemmTraits, + QuantGroupSize, + TransposeC, + ADataType, // AComputeDataType — A activations compute type + {scheduler_ck}, + has_hot_loop_.value, + tail_number_.value>; + + using GemmPipeline = {pipeline_ck}; + +{epilogue_block} + + using Kernel = ck_tile::QuantGemmKernel< + TilePartitioner, GemmPipeline, GemmEpilogue, + ck_tile::QuantType::AQuantGrouped>; + + auto kargs = Kernel::MakeKernelArgs(args); + if(!Kernel::IsSupportedArgument(kargs)) + return -1.0f; + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + return ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{{}}, grids, blocks, 0, kargs)); + }}; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + }} +}}; + +using SelectedKernel = {struct}; + +}} // namespace {ns} + +#ifdef CK_TILE_SINGLE_KERNEL_INCLUDE +using SelectedKernel = {ns}::{struct}; +constexpr const char* KERNEL_NAME = {ns}::KERNEL_NAME; +using ADataType = {ck_a}; +using BDataType = {ck_b}; +using CDataType = {ck_c}; +using QDataType = {ck_q}; +using AccDataType = {ck_acc}; +using QuantGroupSize = {ns}::QuantGroupSize; +#endif // CK_TILE_SINGLE_KERNEL_INCLUDE +""" + + +# ============================================================================= +# Config sweep +# ============================================================================= + + +def _default_config() -> dict: + """Default sweep config matching GemmConfigQuantDecodeInterwave tile defaults (non-preshuffle). + + Uses the Mem pipeline (AQuantGemmPipelineAgBgCrMem) which is correct for + non-preshuffle AQuant kernels. For APreshuffleQuant=true kernels, set + pipeline="compv3" and preshuffle_aq=true in the config. + """ + return { + "variant_keys": ["fp8", "bf8"], + "layouts": ["rcr"], + "pipeline": "mem", # non-preshuffle AQuant uses Mem pipeline + "epilogue": "cshuffle", + "scheduler": "intrawave", + "tile_configs": [ + # GemmConfigQuantDecodeInterwave: M=16, N=64, K=256/sizeof(fp8_t)=256 + {"tile_m": 16, "tile_n": 64, "tile_k": 256, + "warp_m": 1, "warp_n": 4, "warp_k": 1, + "warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 16}, + ], + "quant_groups": [ + {"quant_group_m": 1, "quant_group_n": 1, "quant_group_k": 128}, + ], + "pad_m": False, + "pad_n": False, + "pad_k": True, + "block_size": 256, + "k_block_per_cu": 1, + "double_smem_buffer": False, + "preshuffle_aq": False, + "transpose_c": False, + } + + +def _build_specs(config: dict) -> List[AQuantKernelSpec]: + specs = [] + pipeline = config.get("pipeline", "compv3") + epilogue = config.get("epilogue", "cshuffle") + scheduler = config.get("scheduler", "intrawave") + pad_m = config.get("pad_m", False) + pad_n = config.get("pad_n", False) + pad_k = config.get("pad_k", True) + block_size = config.get("block_size", 256) + k_block_per_cu = config.get("k_block_per_cu", 1) + double_smem_buffer = config.get("double_smem_buffer", False) + preshuffle_aq = config.get("preshuffle_aq", False) + transpose_c = config.get("transpose_c", False) + + for variant_key, layout, tile_dict, qg in itertools.product( + config.get("variant_keys", ["fp8"]), + config.get("layouts", ["rcr"]), + config.get("tile_configs", []), + config.get("quant_groups", [{"quant_group_m": 1, "quant_group_n": 1, "quant_group_k": 128}]), + ): + if variant_key not in AQUANT_VARIANTS: + log.warning("Unknown variant_key %s — skipping", variant_key) + continue + # Derive the correct pipeline from preshuffle_aq if not explicitly set + effective_pipeline = pipeline + if effective_pipeline not in AQUANT_PIPELINE_MAP: + log.warning("Unsupported pipeline %s — skipping", effective_pipeline) + continue + # Guard: compv3 requires preshuffle_aq=true; mem requires preshuffle_aq=false + if effective_pipeline == "compv3" and not preshuffle_aq: + log.warning("pipeline=compv3 requires preshuffle_aq=true — skipping") + continue + if effective_pipeline == "mem" and preshuffle_aq: + log.warning("pipeline=mem does not support preshuffle_aq=true — skipping") + continue + + tile = AQuantTileConfig( + tile_m=tile_dict["tile_m"], + tile_n=tile_dict["tile_n"], + tile_k=tile_dict["tile_k"], + warp_m=tile_dict["warp_m"], + warp_n=tile_dict["warp_n"], + warp_k=tile_dict["warp_k"], + warp_tile_m=tile_dict["warp_tile_m"], + warp_tile_n=tile_dict["warp_tile_n"], + warp_tile_k=tile_dict["warp_tile_k"], + ) + if not tile.is_valid(): + log.debug("Invalid tile config %s — skipping", tile) + continue + + specs.append(AQuantKernelSpec( + variant_key=variant_key, + layout=layout, + pipeline=pipeline, + epilogue=epilogue, + scheduler=scheduler, + tile=tile, + quant_group_m=qg.get("quant_group_m", 1), + quant_group_n=qg.get("quant_group_n", 1), + quant_group_k=qg.get("quant_group_k", 128), + preshuffle_aq=preshuffle_aq, + double_smem_buffer=double_smem_buffer, + pad_m=pad_m, + pad_n=pad_n, + pad_k=pad_k, + block_size=block_size, + k_block_per_cu=k_block_per_cu, + transpose_c=transpose_c, + )) + + return specs + + +# ============================================================================= +# Generation entry point +# ============================================================================= + + +def generate_kernels( + output_dir: Path, + config: Optional[dict] = None, + parallel: bool = True, +) -> List[Path]: + """Generate all AQuant kernel headers into output_dir. + + Returns list of generated .hpp paths. + """ + output_dir.mkdir(parents=True, exist_ok=True) + cfg = config or _default_config() + specs = _build_specs(cfg) + + if not specs: + log.warning("No kernel specs produced from config — check variant_keys and tile_configs") + return [] + + log.info("Generating %d AQuant kernel headers into %s", len(specs), output_dir) + + gen = AQuantKernelHeaderGenerator() + generated: List[Path] = [] + + def _generate_one(spec: AQuantKernelSpec) -> Path: + header = gen.generate(spec) + out_path = output_dir / f"{spec.name}.hpp" + out_path.write_text(header) + log.info(" wrote %s", out_path.name) + return out_path + + if parallel and len(specs) > 1: + import concurrent.futures + with concurrent.futures.ThreadPoolExecutor() as ex: + futures = {ex.submit(_generate_one, s): s for s in specs} + for fut in concurrent.futures.as_completed(futures): + try: + generated.append(fut.result()) + except Exception as e: + log.error("Failed generating %s: %s", futures[fut].name, e) + else: + for spec in specs: + try: + generated.append(_generate_one(spec)) + except Exception as e: + log.error("Failed generating %s: %s", spec.name, e) + + log.info("Generated %d / %d headers", len(generated), len(specs)) + return generated + + +# ============================================================================= +# CLI +# ============================================================================= + + +def main() -> int: + parser = argparse.ArgumentParser( + description="AQuantGrouped GEMM kernel header generator" + ) + parser.add_argument("--output-dir", type=Path, required=True, + help="Directory to write generated .hpp files") + parser.add_argument("--config", type=Path, + help="JSON config file (defaults to built-in sweep)") + parser.add_argument("--config-json", type=str, + help="Inline JSON config string") + parser.add_argument("--no-parallel", action="store_true", + help="Disable parallel generation") + parser.add_argument("--list-names", action="store_true", + help="Print kernel names that would be generated and exit") + args = parser.parse_args() + + cfg: Optional[dict] = None + if args.config_json: + try: + cfg = json.loads(args.config_json) + except json.JSONDecodeError as e: + log.error("Invalid --config-json: %s", e) + return 1 + elif args.config: + with open(args.config) as f: + cfg = json.load(f) + + if args.list_names: + specs = _build_specs(cfg or _default_config()) + for s in specs: + print(s.name) + return 0 + + paths = generate_kernels( + output_dir=args.output_dir, + config=cfg, + parallel=not args.no_parallel, + ) + return 0 if paths else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py new file mode 100644 index 000000000000..839cbd6a7e2f --- /dev/null +++ b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +Example 14 — GroupedGemm AQuant via the Dispatcher + +Demonstrates the full three-layer path: + 1. Codegen — unified_grouped_gemm_aquant_codegen.py -> .hpp + 2. Compile — hipcc -> .so + 3. Run — AQuantGpuGemmRunner -> C = dequant(A, AQ) @ B + +AQuant: A-side activation quantization. + C[M,N] = dequant(A[M,K], AQ[ceil(M/gM), ceil(K/gK)]) @ B[K,N] + AQ is the per-group scale for A; B is unquantized. + +Verifies the GPU result against a NumPy fp32 reference. + +Requirements: + - gfx950 GPU (MI350X) + - hipcc in PATH + - CK include path discoverable relative to this repo + +Usage: + python3 14_grouped_gemm_aquant.py # fp8, 1x1x128 groups, M=16 N=64 K=256 + python3 14_grouped_gemm_aquant.py --dtype bf8 + python3 14_grouped_gemm_aquant.py --dtype fp8 --M 32 --N 128 --K 512 --quant-group-k 128 + python3 14_grouped_gemm_aquant.py --no-verify # skip CPU reference check +""" + +import argparse +import logging +import math +import sys +import tempfile +from pathlib import Path + +import numpy as np + +# Add dispatcher/python to path +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "python")) + +from grouped_gemm_aquant_utils import ( + AQuantKernelConfig, + AQuantGemmProblem, + AQuantGpuGemmRunner, + setup_multiple_aquant_dispatchers, + default_fp8_config, + default_bf8_config, +) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# fp8 encode/decode helpers +# ============================================================================= + + +def _float32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: + """Encode float32 values as fp8 bytes (uint8 view of the fp8 bit pattern). + + dtype: "fp8" -> float8_e4m3fn, "bf8" -> float8_e5m2. + Falls back to clamping when ml_dtypes is not installed. + """ + try: + import ml_dtypes + ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return arr.astype(ml_t).view(np.uint8) + except ImportError: + clamped = np.clip(arr, -2.0, 2.0) + return (clamped * 64).astype(np.int8).view(np.uint8) + + +def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: + """Decode fp8 bytes (uint8 view) back to float32.""" + try: + import ml_dtypes + ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return arr.view(ml_t).astype(np.float32) + except ImportError: + return arr.view(np.int8).astype(np.float32) / 64.0 + + +# ============================================================================= +# CPU reference: C = dequant(A, AQ) @ B +# ============================================================================= + + +def reference_aquant_gemm( + A: np.ndarray, + B: np.ndarray, + AQ: np.ndarray, + problem: AQuantGemmProblem, +) -> np.ndarray: + """ + CPU fp32 reference for C = dequant(A, AQ) @ B. + + A [M, K] float32 — decoded from fp8 bytes + B [K, N] float32 — decoded from fp8 bytes + AQ [QM_A, QK_A] float32 per-block scales for A + + Dequant: A[m, k] *= AQ[m // gM, k // gK] + """ + M, N, K = problem.M, problem.N, problem.K + gM = problem.quant_group_m + gK = problem.quant_group_k + + A_f32 = A.astype(np.float32) + B_f32 = B.astype(np.float32) + + # Apply per-block scales to A + A_dequant = A_f32.copy() + for qi in range(problem.QM_A): + for qk in range(problem.QK_A): + m_start = qi * gM + m_end = min(m_start + gM, M) + k_start = qk * gK + k_end = min(k_start + gK, K) + scale = float(AQ[qi, qk]) + A_dequant[m_start:m_end, k_start:k_end] *= scale + + C_ref = A_dequant @ B_f32 + return C_ref.astype(np.float16) + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + parser = argparse.ArgumentParser(description="GroupedGemm AQuant dispatcher example") + parser.add_argument("--dtype", choices=["fp8", "bf8"], default="fp8") + parser.add_argument("--M", type=int, default=16) + parser.add_argument("--N", type=int, default=64) + parser.add_argument("--K", type=int, default=256) + parser.add_argument("--quant-group-k", type=int, default=128) + parser.add_argument("--quant-group-m", type=int, default=1) + parser.add_argument("--no-verify", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + parser.add_argument("--gfx-arch", type=str, default="gfx950") + args = parser.parse_args() + + M, N, K = args.M, args.N, args.K + gK = args.quant_group_k + gM = args.quant_group_m + + # ------------------------------------------------------------------------- + # 1. Build kernel config + # ------------------------------------------------------------------------- + if args.dtype == "fp8": + config = default_fp8_config(quant_group_k=gK, quant_group_m=gM, gfx_arch=args.gfx_arch) + else: + config = default_bf8_config(quant_group_k=gK, quant_group_m=gM, gfx_arch=args.gfx_arch) + + log.info("Kernel: %s", config.name) + + # ------------------------------------------------------------------------- + # 2. Codegen + compile + # ------------------------------------------------------------------------- + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="aquant_ex14_")) + log.info("Output dir: %s", out_dir) + + so_paths = setup_multiple_aquant_dispatchers( + configs=[config], + output_dir=out_dir, + gfx_arch=args.gfx_arch, + ) + + if not so_paths or so_paths[0] is None: + log.error("Kernel build failed — see errors above") + return 1 + + so_path = so_paths[0] + log.info("Built: %s", so_path) + + # ------------------------------------------------------------------------- + # 3. Generate inputs + # ------------------------------------------------------------------------- + rng = np.random.default_rng(42) + QK_A = math.ceil(K / gK) + QM_A = math.ceil(M / gM) + + # Generate float32 values in fp8 representable range, then encode as fp8 bytes. + A_f32 = rng.uniform(-2.0, 2.0, (M, K)).astype(np.float32) + B_f32 = rng.uniform(-2.0, 2.0, (K, N)).astype(np.float32) + # AQ scale: per-group scale for A, shape [QM_A, QK_A], dtype float32 + AQ_f32 = rng.uniform(0.5, 2.0, (QM_A, QK_A)).astype(np.float32) + + # Encode A and B as fp8 bytes (uint8 view, 1 byte/element) + A_raw = _float32_to_fp8(A_f32, args.dtype) + B_raw = _float32_to_fp8(B_f32, args.dtype) + + # Decode back so the CPU reference uses the same rounded values the kernel sees + A_dec = _fp8_to_float32(A_raw, args.dtype) + B_dec = _fp8_to_float32(B_raw, args.dtype) + + # ------------------------------------------------------------------------- + # 4. Run on GPU + # ------------------------------------------------------------------------- + problem = AQuantGemmProblem( + M=M, N=N, K=K, + quant_group_m=gM, + quant_group_n=1, + quant_group_k=gK, + ) + + runner = AQuantGpuGemmRunner(so_path) + log.info("Running kernel: %s", runner.kernel_name) + + result = runner.run(A=A_raw, B=B_raw, AQ=AQ_f32, problem=problem) + log.info("Kernel time: %.3f ms", result.time_ms) + + # ------------------------------------------------------------------------- + # 5. Verify + # ------------------------------------------------------------------------- + if not args.no_verify: + C_ref = reference_aquant_gemm(A_dec, B_dec, AQ_f32, problem) + C_gpu = result.C + + max_abs_err = float(np.max(np.abs(C_gpu.astype(np.float32) - C_ref.astype(np.float32)))) + max_ref = float(np.max(np.abs(C_ref.astype(np.float32)))) + 1e-6 + max_rel = max_abs_err / max_ref + + tolerance = 0.05 # fp8 ~1e-2 to 5e-2 + if max_rel <= tolerance: + log.info("PASSED (max_rel=%.4f, tol=%.4f)", max_rel, tolerance) + else: + log.error("FAILED (max_rel=%.4f > tol=%.4f)", max_rel, tolerance) + return 1 + else: + log.info("Verification skipped (--no-verify)") + + log.info("Example 14 complete.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py new file mode 100644 index 000000000000..4f75a5cfda43 --- /dev/null +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -0,0 +1,266 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +Example 15 — GroupedGemm ABQuant via the Dispatcher + +Demonstrates the full three-layer path: + 1. Codegen — unified_grouped_gemm_abquant_codegen.py -> .hpp + 2. Compile — hipcc -> .so + 3. Run — ABQuantGpuGemmRunner -> C = dequant(A, AQ) @ dequant(B, BQ) + +ABQuant: both A-side and B-side quantization active simultaneously. + C[M,N] = dequant(A[M,K], AQ[ceil(M/aM), ceil(K/aK)]) @ dequant(B[K,N], BQ[ceil(K/bK), ceil(N/bN)]) + +Verifies the GPU result against a NumPy fp32 reference. + +Requirements: + - gfx950 GPU (MI350X) + - hipcc in PATH + - CK include path discoverable relative to this repo + +Usage: + python3 15_grouped_gemm_abquant.py # fp8, 128x128x128, no preshuffle + python3 15_grouped_gemm_abquant.py --dtype bf8 + python3 15_grouped_gemm_abquant.py --pipeline eightwaves --M 192 --N 256 --K 128 + python3 15_grouped_gemm_abquant.py --no-verify +""" + +import argparse +import logging +import math +import sys +import tempfile +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "python")) + +from grouped_gemm_abquant_utils import ( + ABQuantKernelConfig, + ABQuantGemmProblem, + ABQuantGpuGemmRunner, + setup_multiple_abquant_dispatchers, + default_fp8_compv3_config, + default_bf8_compv3_config, + default_fp8_eightwaves_config, + default_bf8_eightwaves_config, +) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + + +# ============================================================================= +# fp8 encode/decode helpers +# ============================================================================= + + +def _float32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: + try: + import ml_dtypes + ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return arr.astype(ml_t).view(np.uint8) + except ImportError: + clamped = np.clip(arr, -2.0, 2.0) + return (clamped * 64).astype(np.int8).view(np.uint8) + + +def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: + try: + import ml_dtypes + ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return arr.view(ml_t).astype(np.float32) + except ImportError: + return arr.view(np.int8).astype(np.float32) / 64.0 + + +# ============================================================================= +# CPU reference: C = dequant(A, AQ) @ dequant(B, BQ) +# ============================================================================= + + +def reference_abquant_gemm( + A: np.ndarray, + B: np.ndarray, + AQ: np.ndarray, + BQ: np.ndarray, + problem: ABQuantGemmProblem, +) -> np.ndarray: + """ + CPU fp32 reference for C = dequant(A, AQ) @ dequant(B, BQ). + + A [M, K] float32 — decoded from fp8 bytes + B [K, N] float32 — decoded from fp8 bytes + AQ [QM_A, QK_A] float32 per-block scales for A + BQ [QK_B, QN_B] float32 per-block scales for B + """ + M, N, K = problem.M, problem.N, problem.K + aM = problem.aquant_group_m + aK = problem.aquant_group_k + bK = problem.bquant_group_k + bN = problem.bquant_group_n + + A_f32 = A.astype(np.float32) + B_f32 = B.astype(np.float32) + + # Dequantize A: A[m, k] *= AQ[m // aM, k // aK] + A_dequant = A_f32.copy() + for qi in range(problem.QM_A): + for qk in range(problem.QK_A): + m_start = qi * aM + m_end = min(m_start + aM, M) + k_start = qk * aK + k_end = min(k_start + aK, K) + A_dequant[m_start:m_end, k_start:k_end] *= float(AQ[qi, qk]) + + # Dequantize B: B[k, n] *= BQ[k // bK, n // bN] + B_dequant = B_f32.copy() + for qk in range(problem.QK_B): + for qn in range(problem.QN_B): + k_start = qk * bK + k_end = min(k_start + bK, K) + n_start = qn * bN + n_end = min(n_start + bN, N) + B_dequant[k_start:k_end, n_start:n_end] *= float(BQ[qk, qn]) + + C_ref = A_dequant @ B_dequant + return C_ref.astype(np.float16) + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + parser = argparse.ArgumentParser(description="GroupedGemm ABQuant dispatcher example") + parser.add_argument("--dtype", choices=["fp8", "bf8"], default="fp8") + parser.add_argument("--pipeline", choices=["compv3", "eightwaves", "preshuffleb"], + default="compv3") + parser.add_argument("--M", type=int, default=128) + parser.add_argument("--N", type=int, default=128) + parser.add_argument("--K", type=int, default=128) + parser.add_argument("--quant-group-k", type=int, default=128) + parser.add_argument("--bquant-group-n", type=int, default=1) + parser.add_argument("--no-verify", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + parser.add_argument("--gfx-arch", type=str, default="gfx950") + args = parser.parse_args() + + M, N, K = args.M, args.N, args.K + gK = args.quant_group_k + bN = args.bquant_group_n + + # ------------------------------------------------------------------------- + # 1. Build kernel config + # ------------------------------------------------------------------------- + if args.pipeline == "compv3": + if args.dtype == "fp8": + config = default_fp8_compv3_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=args.gfx_arch) + else: + config = default_bf8_compv3_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=args.gfx_arch) + elif args.pipeline == "eightwaves": + if args.dtype == "fp8": + config = default_fp8_eightwaves_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=args.gfx_arch) + else: + config = default_bf8_eightwaves_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=args.gfx_arch) + else: + from grouped_gemm_abquant_utils import default_fp8_preshuffleb_config, default_bf8_preshuffleb_config + if args.dtype == "fp8": + config = default_fp8_preshuffleb_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=args.gfx_arch) + else: + config = default_bf8_preshuffleb_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=args.gfx_arch) + + log.info("Kernel: %s", config.name) + + # ------------------------------------------------------------------------- + # 2. Codegen + compile + # ------------------------------------------------------------------------- + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="abquant_ex15_")) + log.info("Output dir: %s", out_dir) + + so_paths = setup_multiple_abquant_dispatchers( + configs=[config], + output_dir=out_dir, + gfx_arch=args.gfx_arch, + ) + + if not so_paths or so_paths[0] is None: + log.error("Kernel build failed — see errors above") + return 1 + + so_path = so_paths[0] + log.info("Built: %s", so_path) + + # ------------------------------------------------------------------------- + # 3. Generate inputs + # ------------------------------------------------------------------------- + rng = np.random.default_rng(42) + QK_A = math.ceil(K / gK) + QM_A = 1 # aquant_group_m=1 by default → QM_A == M + QK_B = math.ceil(K / gK) + QN_B = math.ceil(N / bN) + + # Recompute using actual problem to be safe + problem = ABQuantGemmProblem( + M=M, N=N, K=K, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=gK, + bquant_group_m=1, bquant_group_n=bN, bquant_group_k=gK, + ) + + A_f32 = rng.uniform(-2.0, 2.0, (M, K)).astype(np.float32) + B_f32 = rng.uniform(-2.0, 2.0, (K, N)).astype(np.float32) + AQ_f32 = rng.uniform(0.5, 2.0, (problem.QM_A, problem.QK_A)).astype(np.float32) + BQ_f32 = rng.uniform(0.5, 2.0, (problem.QK_B, problem.QN_B)).astype(np.float32) + + A_raw = _float32_to_fp8(A_f32, args.dtype) + B_raw = _float32_to_fp8(B_f32, args.dtype) + + A_dec = _fp8_to_float32(A_raw, args.dtype) + B_dec = _fp8_to_float32(B_raw, args.dtype) + + # ------------------------------------------------------------------------- + # 4. Run on GPU + # ------------------------------------------------------------------------- + runner = ABQuantGpuGemmRunner(so_path) + log.info("Running kernel: %s", runner.kernel_name) + + result = runner.run(A=A_raw, B=B_raw, AQ=AQ_f32, BQ=BQ_f32, problem=problem) + log.info("Kernel time: %.3f ms", result.time_ms) + + # ------------------------------------------------------------------------- + # 5. Verify + # ------------------------------------------------------------------------- + if not args.no_verify: + C_ref = reference_abquant_gemm(A_dec, B_dec, AQ_f32, BQ_f32, problem) + C_gpu = result.C + + max_abs_err = float(np.max(np.abs(C_gpu.astype(np.float32) - C_ref.astype(np.float32)))) + max_ref = float(np.max(np.abs(C_ref.astype(np.float32)))) + 1e-6 + max_rel = max_abs_err / max_ref + + tolerance = 0.05 + if max_rel <= tolerance: + log.info("PASSED (max_rel=%.4f, tol=%.4f)", max_rel, tolerance) + else: + log.error("FAILED (max_rel=%.4f > tol=%.4f)", max_rel, tolerance) + return 1 + else: + log.info("Verification skipped (--no-verify)") + + log.info("Example 15 complete.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py new file mode 100644 index 000000000000..bcf0d2dd466e --- /dev/null +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -0,0 +1,762 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +GroupedGemm ABQuant dispatcher utilities. + +Three-layer Python bridge for the dispatcher's ABQuantGrouped GEMM path: + + ABQuantKernelConfig — describes one kernel; .name is byte-exact with codegen KERNEL_NAME + ABQuantDispatcherLib — thin ctypes wrapper around a compiled .so + ABQuantGpuGemmRunner — high-level runner that accepts numpy arrays + +Build helpers: + setup_multiple_abquant_dispatchers(configs, ...) + codegen -> hipcc -> list of .so paths, all in parallel + +ABQuant: both A-side and B-side quantization active simultaneously. + AQ[ceil(M/aM), ceil(K/aK)] — A-side scale (RowMajor) + BQ[ceil(K/bK), ceil(N/bN)] — B-side scale (RowMajor) + Constraint: aquant_group_k == bquant_group_k + +Usage: + configs = [ABQuantKernelConfig(variant_key="fp8", pipeline="compv3", ...)] + so_paths = setup_multiple_abquant_dispatchers(configs, output_dir=Path("/tmp/abq")) + runner = ABQuantGpuGemmRunner(so_paths[0]) + result = runner.run(A, B, AQ, BQ, ABQuantGemmProblem(M=128, N=128, K=128)) +""" + +import ctypes +import json +import logging +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple +import concurrent.futures + +log = logging.getLogger(__name__) + +# ============================================================================= +# Constants +# ============================================================================= + +_CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_grouped_gemm_abquant_codegen.py" +_CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "grouped_gemm_abquant_ctypes_lib.cpp" + +_codegen_dir = str(Path(__file__).parent.parent / "codegen") +if _codegen_dir not in sys.path: + sys.path.insert(0, _codegen_dir) +from codegen_common import make_abquant_kernel_name # noqa: E402 + +_DEFAULT_HIPCC = "hipcc" +_DEFAULT_GFX_ARCH = "gfx950" + +_HIPCC_BASE_FLAGS = [ + "-std=c++17", + "-O3", + "-fPIC", + "-shared", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", + "-w", +] + + +# ============================================================================= +# ABQuantKernelConfig — byte-exact naming with codegen +# ============================================================================= + + +@dataclass +class ABQuantKernelConfig: + """ + Complete description of one ABQuantGrouped GEMM kernel. + + Constraint: aquant_group_k must equal bquant_group_k. + + Pipelines: + "compv3" — ABQuantGemmPipelineAgBgCrCompV3 (non-gfx950, GemmConfigABQuantPrefill) + "eightwaves" — ABQuantGemmPipelineAgBgCrEightWaves (gfx950, TransposeC=true) + "preshuffleb" — WPABQuantBPipelineAgBgCrV2 (GemmConfigPreshuffleB_ABQuant_Prefill) + """ + + variant_key: str # "fp8", "bf8" + layout: str # "rcr" + pipeline: str # "compv3", "eightwaves", "preshuffleb" + epilogue: str # "cshuffle" (effective epilogue may be permute_n) + scheduler: str # "intrawave" + + tile_m: int + tile_n: int + tile_k: int + warp_m: int + warp_n: int + warp_k: int + warp_tile_m: int + warp_tile_n: int + warp_tile_k: int + + # A-side quantization group + aquant_group_m: int = 1 + aquant_group_n: int = 1 + aquant_group_k: int = 128 + + # B-side quantization group + bquant_group_m: int = 1 + bquant_group_n: int = 1 + bquant_group_k: int = 128 + + preshuffle_b: bool = False # PreshuffleB + preshuffle_aq: bool = False # APreshuffleQuant + preshuffle_bq: bool = False # BPreshuffleQuant + transpose_c: bool = False # true for eightwaves/gfx950 + double_smem_buffer: bool = False + k_block_per_cu: int = 1 + + gfx_arch: str = _DEFAULT_GFX_ARCH + + def __post_init__(self): + if self.aquant_group_k != self.bquant_group_k: + raise ValueError( + f"ABQuant requires aquant_group_k == bquant_group_k, " + f"got {self.aquant_group_k} != {self.bquant_group_k}" + ) + + @property + def name(self) -> str: + """Byte-exact match to codegen KERNEL_NAME.""" + return make_abquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline, + epilogue=self.epilogue, + scheduler=self.scheduler, + tile_m=self.tile_m, tile_n=self.tile_n, tile_k=self.tile_k, + warp_m=self.warp_m, warp_n=self.warp_n, warp_k=self.warp_k, + warp_tile_m=self.warp_tile_m, warp_tile_n=self.warp_tile_n, warp_tile_k=self.warp_tile_k, + aquant_group_m=self.aquant_group_m, + aquant_group_n=self.aquant_group_n, + aquant_group_k=self.aquant_group_k, + bquant_group_m=self.bquant_group_m, + bquant_group_n=self.bquant_group_n, + bquant_group_k=self.bquant_group_k, + preshuffle_b=self.preshuffle_b, + preshuffle_aq=self.preshuffle_aq, + preshuffle_bq=self.preshuffle_bq, + transpose_c=self.transpose_c, + ) + + def to_codegen_config(self) -> dict: + """Produce the JSON config dict consumed by unified_grouped_gemm_abquant_codegen.py.""" + return { + "variant_keys": [self.variant_key], + "layouts": [self.layout], + "pipeline": self.pipeline, + "epilogue": self.epilogue, + "scheduler": self.scheduler, + "tile_configs": [{ + "tile_m": self.tile_m, + "tile_n": self.tile_n, + "tile_k": self.tile_k, + "warp_m": self.warp_m, + "warp_n": self.warp_n, + "warp_k": self.warp_k, + "warp_tile_m": self.warp_tile_m, + "warp_tile_n": self.warp_tile_n, + "warp_tile_k": self.warp_tile_k, + }], + "aquant_groups": [{ + "aquant_group_m": self.aquant_group_m, + "aquant_group_n": self.aquant_group_n, + "aquant_group_k": self.aquant_group_k, + }], + "bquant_groups": [{ + "bquant_group_m": self.bquant_group_m, + "bquant_group_n": self.bquant_group_n, + "bquant_group_k": self.bquant_group_k, + }], + "preshuffle_b": self.preshuffle_b, + "preshuffle_aq": self.preshuffle_aq, + "preshuffle_bq": self.preshuffle_bq, + "transpose_c": self.transpose_c, + "double_smem_buffer": self.double_smem_buffer, + "k_block_per_cu": self.k_block_per_cu, + } + + +# ============================================================================= +# ABQuantGemmProblem +# ============================================================================= + + +@dataclass +class ABQuantGemmProblem: + M: int + N: int + K: int + aquant_group_m: int = 1 + aquant_group_n: int = 1 + aquant_group_k: int = 128 + bquant_group_m: int = 1 + bquant_group_n: int = 1 + bquant_group_k: int = 128 + k_batch: int = 1 + + @property + def QK_A(self) -> int: + """ceil(K / aquant_group_k)""" + return (self.K + self.aquant_group_k - 1) // self.aquant_group_k + + @property + def QM_A(self) -> int: + """ceil(M / aquant_group_m) — typically == M when aM=1""" + return (self.M + self.aquant_group_m - 1) // self.aquant_group_m + + @property + def QK_B(self) -> int: + """ceil(K / bquant_group_k)""" + return (self.K + self.bquant_group_k - 1) // self.bquant_group_k + + @property + def QN_B(self) -> int: + """ceil(N / bquant_group_n)""" + return (self.N + self.bquant_group_n - 1) // self.bquant_group_n + + +# ============================================================================= +# ABQuantGemmResult +# ============================================================================= + + +@dataclass +class ABQuantGemmResult: + C: object # numpy array (M, N) + time_ms: float + kernel_name: str + + +# ============================================================================= +# ABQuantDispatcherLib — thin ctypes wrapper +# ============================================================================= + + +class ABQuantDispatcherLib: + """ + Loads a compiled abquant_gemm .so and wraps its C API. + + Expected .so exports: + int dispatcher_initialize() + int dispatcher_run_abquant_gemm(A, B, AQ, BQ, C, M, N, K, + stride_A, stride_B, stride_AQ, stride_BQ, stride_C, + QK_A, QM_A, QK_B, QN_B, k_batch, *time_ms) + char* dispatcher_get_kernel_name() + int dispatcher_get_kernel_count() + void dispatcher_cleanup() + """ + + def __init__(self, so_path: Path): + self.so_path = Path(so_path) + if not self.so_path.exists(): + raise FileNotFoundError(f"ABQuant .so not found: {self.so_path}") + self._lib = ctypes.CDLL(str(self.so_path)) + self._setup() + rc = self._lib.dispatcher_initialize() + if rc != 0: + raise RuntimeError(f"dispatcher_initialize() returned {rc}") + + def _setup(self): + lib = self._lib + + lib.dispatcher_initialize.restype = ctypes.c_int + lib.dispatcher_initialize.argtypes = [] + + lib.dispatcher_run_abquant_gemm.restype = ctypes.c_int + lib.dispatcher_run_abquant_gemm.argtypes = [ + ctypes.c_void_p, # A + ctypes.c_void_p, # B + ctypes.c_void_p, # AQ + ctypes.c_void_p, # BQ + ctypes.c_void_p, # C + ctypes.c_int64, # M + ctypes.c_int64, # N + ctypes.c_int64, # K + ctypes.c_int64, # stride_A + ctypes.c_int64, # stride_B + ctypes.c_int64, # stride_AQ + ctypes.c_int64, # stride_BQ + ctypes.c_int64, # stride_C + ctypes.c_int64, # QK_A + ctypes.c_int64, # QM_A + ctypes.c_int64, # QK_B + ctypes.c_int64, # QN_B + ctypes.c_int, # k_batch + ctypes.POINTER(ctypes.c_float), # time_ms + ] + + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + lib.dispatcher_get_kernel_name.argtypes = [] + + lib.dispatcher_get_kernel_count.restype = ctypes.c_int + lib.dispatcher_get_kernel_count.argtypes = [] + + lib.dispatcher_cleanup.restype = None + lib.dispatcher_cleanup.argtypes = [] + + def run( + self, + A, B, AQ, BQ, C, + M: int, N: int, K: int, + stride_A: int, stride_B: int, + stride_AQ: int, stride_BQ: int, stride_C: int, + QK_A: int, QM_A: int, QK_B: int, QN_B: int, + k_batch: int = 1, + ) -> Tuple[int, float]: + """Call dispatcher_run_abquant_gemm. Returns (status, time_ms).""" + import numpy as np + + A = np.ascontiguousarray(A) + B = np.ascontiguousarray(B) + AQ = np.ascontiguousarray(AQ) + BQ = np.ascontiguousarray(BQ) + C = np.ascontiguousarray(C) + + time_ms = ctypes.c_float(0.0) + + rc = self._lib.dispatcher_run_abquant_gemm( + A.ctypes.data_as(ctypes.c_void_p), + B.ctypes.data_as(ctypes.c_void_p), + AQ.ctypes.data_as(ctypes.c_void_p), + BQ.ctypes.data_as(ctypes.c_void_p), + C.ctypes.data_as(ctypes.c_void_p), + ctypes.c_int64(M), ctypes.c_int64(N), ctypes.c_int64(K), + ctypes.c_int64(stride_A), ctypes.c_int64(stride_B), + ctypes.c_int64(stride_AQ), ctypes.c_int64(stride_BQ), ctypes.c_int64(stride_C), + ctypes.c_int64(QK_A), ctypes.c_int64(QM_A), + ctypes.c_int64(QK_B), ctypes.c_int64(QN_B), + ctypes.c_int(k_batch), + ctypes.byref(time_ms), + ) + return rc, time_ms.value + + def get_kernel_name(self) -> str: + raw = self._lib.dispatcher_get_kernel_name() + return raw.decode("utf-8") if raw else "" + + def get_kernel_count(self) -> int: + return self._lib.dispatcher_get_kernel_count() + + def cleanup(self): + self._lib.dispatcher_cleanup() + + def __del__(self): + try: + self._lib.dispatcher_cleanup() + except Exception: + pass + + +# ============================================================================= +# ABQuantGpuGemmRunner — high-level runner +# ============================================================================= + + +class ABQuantGpuGemmRunner: + """High-level runner for ABQuantGrouped GEMM.""" + + def __init__(self, so_path: Path): + self._lib = ABQuantDispatcherLib(so_path) + + @property + def kernel_name(self) -> str: + return self._lib.get_kernel_name() + + def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuantGemmResult: + """ + Run ABQuantGrouped GEMM. + + A shape: (M, K) dtype: fp8/bf8 + B shape: (K, N) col-major dtype: fp8/bf8 + AQ shape: (QM_A, QK_A) dtype: float32 (A-side scale) + BQ shape: (QK_B, QN_B) dtype: float32 (B-side scale) + c_dtype: numpy dtype for C output buffer. Defaults to np.float16. + """ + import numpy as np + + M, N, K = problem.M, problem.N, problem.K + QK_A = problem.QK_A + QM_A = problem.QM_A + QK_B = problem.QK_B + QN_B = problem.QN_B + + if c_dtype is None: + c_dtype = np.float16 + + C = np.zeros((M, N), dtype=c_dtype) + + # Stride layout: + # AQ RowMajor [QM_A, QK_A]: stride_AQ = QK_A + # BQ RowMajor [QK_B, QN_B]: stride_BQ = QN_B + stride_A = K + stride_B = K # col-major B: leading dim = K + stride_AQ = QK_A + stride_BQ = QN_B + stride_C = N + + rc, time_ms = self._lib.run( + A=A, B=B, AQ=AQ, BQ=BQ, C=C, + M=M, N=N, K=K, + stride_A=stride_A, stride_B=stride_B, + stride_AQ=stride_AQ, stride_BQ=stride_BQ, stride_C=stride_C, + QK_A=QK_A, QM_A=QM_A, QK_B=QK_B, QN_B=QN_B, + k_batch=problem.k_batch, + ) + + if rc != 0: + raise RuntimeError( + f"dispatcher_run_abquant_gemm failed with code {rc} " + f"for kernel {self.kernel_name}" + ) + + return ABQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) + + +# ============================================================================= +# Subprocess helpers +# ============================================================================= + + +def _detect_gpu_arch() -> str: + try: + result = subprocess.run( + ["rocm_agent_enumerator"], capture_output=True, text=True, timeout=10, + ) + for line in result.stdout.splitlines(): + line = line.strip() + if line.startswith("gfx") and line != "gfx000": + return line + except Exception: + pass + return _DEFAULT_GFX_ARCH + + +def _get_ck_include_dir() -> Optional[Path]: + here = Path(__file__).resolve().parent + for parent in [here.parent.parent, here.parent.parent.parent]: + candidate = parent / "include" + if (candidate / "ck_tile").is_dir(): + return candidate + return None + + +def _generate_abquant_kernel( + config: ABQuantKernelConfig, + output_dir: Path, +) -> Optional[Path]: + config_json = json.dumps(config.to_codegen_config()) + cmd = [ + sys.executable, str(_CODEGEN_SCRIPT), + "--output-dir", str(output_dir), + "--config-json", config_json, + ] + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + log.error("Codegen failed for %s:\n%s", config.name, result.stderr) + return None + except subprocess.TimeoutExpired: + log.error("Codegen timed out for %s", config.name) + return None + + hpp = output_dir / f"{config.name}.hpp" + if not hpp.exists(): + log.error("Codegen succeeded but %s not found", hpp) + return None + return hpp + + +def _compile_abquant_kernel( + hpp_path: Path, + so_path: Path, + gfx_arch: str, + hipcc: str = _DEFAULT_HIPCC, + extra_include_dirs: Optional[List[str]] = None, +) -> bool: + ck_include = _get_ck_include_dir() + + cmd = [hipcc] + _HIPCC_BASE_FLAGS + [ + f"--offload-arch={gfx_arch}", + f"-DGFX_ARCH=\"{gfx_arch}\"", + "-include", str(hpp_path), + str(_CTYPES_LIB_SRC), + "-o", str(so_path), + ] + + if ck_include: + cmd += [f"-I{ck_include}"] + + dispatcher_include = _CTYPES_LIB_SRC.parent.parent.parent / "dispatcher" / "include" + if dispatcher_include.is_dir(): + cmd += [f"-I{dispatcher_include}"] + + if extra_include_dirs: + for d in extra_include_dirs: + cmd += [f"-I{d}"] + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=600) + if result.returncode != 0: + log.error("Compile failed for %s:\n%s", so_path.name, result.stderr[-2000:]) + return False + return True + except subprocess.TimeoutExpired: + log.error("Compile timed out for %s", so_path.name) + return False + + +# ============================================================================= +# setup_multiple_abquant_dispatchers — build pipeline +# ============================================================================= + + +def setup_multiple_abquant_dispatchers( + configs: List[ABQuantKernelConfig], + output_dir: Optional[Path] = None, + hipcc: str = _DEFAULT_HIPCC, + gfx_arch: Optional[str] = None, + extra_include_dirs: Optional[List[str]] = None, + parallel: bool = True, + max_workers: Optional[int] = None, +) -> List[Optional[Path]]: + """For each ABQuantKernelConfig: codegen -> hipcc compile -> .so path.""" + if not configs: + return [] + + arch = gfx_arch or _detect_gpu_arch() + base_dir = output_dir or Path(tempfile.mkdtemp(prefix="abquant_dispatcher_")) + base_dir.mkdir(parents=True, exist_ok=True) + + headers_dir = base_dir / "generated_kernels" + so_dir = base_dir / "libs" + headers_dir.mkdir(exist_ok=True) + so_dir.mkdir(exist_ok=True) + + log.info("Building %d ABQuant kernel(s) for %s into %s", len(configs), arch, base_dir) + + seen: Dict[str, int] = {} + deduped: List[Tuple[int, ABQuantKernelConfig]] = [] + for i, cfg in enumerate(configs): + if cfg.name not in seen: + seen[cfg.name] = i + deduped.append((i, cfg)) + + results: List[Optional[Path]] = [None] * len(configs) + + def _build_one(idx: int, cfg: ABQuantKernelConfig) -> Tuple[int, Optional[Path]]: + hpp = _generate_abquant_kernel(cfg, headers_dir) + if hpp is None: + return idx, None + + so = so_dir / f"lib{cfg.name}_{arch}.so" + if so.exists(): + log.info(" [cached] %s", so.name) + return idx, so + + ok = _compile_abquant_kernel(hpp, so, arch, hipcc, extra_include_dirs) + return idx, so if ok else None + + if parallel and len(deduped) > 1: + workers = max_workers or min(len(deduped), os.cpu_count() or 4) + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex: + futures = {ex.submit(_build_one, idx, cfg): (idx, cfg) for idx, cfg in deduped} + for fut in concurrent.futures.as_completed(futures): + try: + idx, so_path = fut.result() + results[idx] = so_path + if so_path: + log.info(" built %s", so_path.name) + else: + _, cfg = futures[fut] + log.error(" FAILED %s", cfg.name) + except Exception as e: + _, cfg = futures[fut] + log.error(" EXCEPTION for %s: %s", cfg.name, e) + else: + for idx, cfg in deduped: + _, so_path = _build_one(idx, cfg) + results[idx] = so_path + + for i, cfg in enumerate(configs): + if results[i] is None: + first_idx = seen.get(cfg.name) + if first_idx is not None and first_idx != i: + results[i] = results[first_idx] + + built = sum(1 for r in results if r is not None) + log.info("Built %d / %d ABQuant kernels", built, len(configs)) + return results + + +# ============================================================================= +# Default configs (mapped from reference examples) +# ============================================================================= + + +def default_fp8_compv3_config( + quant_group_k: int = 128, + bquant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> ABQuantKernelConfig: + """fp8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False). + + Tile: 128x128x128, warp 1x4x1. kPadK=false (ABQuant prefill tiles). + """ + return ABQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, + bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, + preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, + transpose_c=False, + gfx_arch=gfx_arch, + ) + + +def default_bf8_compv3_config( + quant_group_k: int = 128, + bquant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> ABQuantKernelConfig: + """bf8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False).""" + return ABQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, + bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, + preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, + transpose_c=False, + gfx_arch=gfx_arch, + ) + + +def default_fp8_eightwaves_config( + quant_group_k: int = 128, + bquant_group_n: int = 128, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> ABQuantKernelConfig: + """fp8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True). + + Tile: 192x256x128, warp 4x2x1 — 8-wave configuration for MI350X. + warp_tile_k=16: get_k_warp_tile on gfx950 = 128 for 8-bit floats, + but GemmConfigEightWaves inherits GemmConfigABQuantPrefill whose warp_tile_k uses + get_k_warp_tile without IsFlatMM — use 16 as concrete value for now. + """ + return ABQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="eightwaves", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=192, tile_n=256, tile_k=128, + warp_m=4, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, + bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, + preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, + transpose_c=True, + k_block_per_cu=1, + gfx_arch=gfx_arch, + ) + + +def default_bf8_eightwaves_config( + quant_group_k: int = 128, + bquant_group_n: int = 128, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> ABQuantKernelConfig: + """bf8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True).""" + return ABQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="eightwaves", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=192, tile_n=256, tile_k=128, + warp_m=4, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, + bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, + preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, + transpose_c=True, + k_block_per_cu=1, + gfx_arch=gfx_arch, + ) + + +def default_fp8_preshuffleb_config( + quant_group_k: int = 128, + bquant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> ABQuantKernelConfig: + """fp8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill). + + Tile: 128x128x128, warp 2x2x1 (vs 1x4x1 for compv3). DoubleSmemBuffer=True. + """ + return ABQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="preshuffleb", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=2, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, + bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, + preshuffle_b=True, preshuffle_aq=False, preshuffle_bq=False, + transpose_c=True, + double_smem_buffer=True, + k_block_per_cu=2, + gfx_arch=gfx_arch, + ) + + +def default_bf8_preshuffleb_config( + quant_group_k: int = 128, + bquant_group_n: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> ABQuantKernelConfig: + """bf8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill).""" + return ABQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="preshuffleb", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=2, warp_n=2, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, + bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, + preshuffle_b=True, preshuffle_aq=False, preshuffle_bq=False, + transpose_c=True, + double_smem_buffer=True, + k_block_per_cu=2, + gfx_arch=gfx_arch, + ) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py new file mode 100644 index 000000000000..3dee600f37df --- /dev/null +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -0,0 +1,741 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +GroupedGemm AQuant dispatcher utilities. + +Three-layer Python bridge for the dispatcher's AQuantGrouped GEMM path: + + AQuantKernelConfig — describes one kernel; .name is byte-exact with codegen KERNEL_NAME + AQuantDispatcherLib — thin ctypes wrapper around a compiled .so + AQuantGpuGemmRunner — high-level runner that accepts numpy arrays + +Build helpers: + setup_multiple_aquant_dispatchers(configs, ...) + codegen -> hipcc -> list of .so paths, all in parallel + +AQuant: A-side activation quantization. + AQ[ceil(M/gM), ceil(K/gK)] is the A-side scale tensor (RowMajor). + Non-preshuffle kernels use AQuantGemmPipelineAgBgCrMem (pipeline="mem"). + Preshuffle kernels (APreshuffleQuant=true) use AQuantGemmPipelineAgBgCrCompV3 (pipeline="compv3"). + +Usage: + configs = [AQuantKernelConfig(variant_key="fp8", layout="rcr", pipeline="mem", ...)] + so_paths = setup_multiple_aquant_dispatchers(configs, output_dir=Path("/tmp/aq")) + runner = AQuantGpuGemmRunner(so_paths[0]) + result = runner.run(A, B, AQ, AQuantGemmProblem(M=16, N=64, K=256)) +""" + +import ctypes +import json +import logging +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple +import concurrent.futures + +log = logging.getLogger(__name__) + +# ============================================================================= +# Constants +# ============================================================================= + +_CODEGEN_SCRIPT = Path(__file__).parent.parent / "codegen" / "unified_grouped_gemm_aquant_codegen.py" +_CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "grouped_gemm_aquant_ctypes_lib.cpp" + +_codegen_dir = str(Path(__file__).parent.parent / "codegen") +if _codegen_dir not in sys.path: + sys.path.insert(0, _codegen_dir) +from codegen_common import make_aquant_kernel_name # noqa: E402 + +_DEFAULT_HIPCC = "hipcc" +_DEFAULT_GFX_ARCH = "gfx950" + +_HIPCC_BASE_FLAGS = [ + "-std=c++17", + "-O3", + "-fPIC", + "-shared", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", + "-w", +] + + +# ============================================================================= +# AQuantKernelConfig — byte-exact naming with codegen +# ============================================================================= + + +@dataclass +class AQuantKernelConfig: + """ + Complete description of one AQuantGrouped GEMM kernel. + + The .name property produces the exact string that unified_grouped_gemm_aquant_codegen.py + emits as KERNEL_NAME, ensuring the Python side and compiled .so always agree. + + pipeline: "mem" — non-preshuffle (AQuantGemmPipelineAgBgCrMem) + "compv3" — preshuffle, requires preshuffle_aq=True + """ + + variant_key: str # "fp8", "bf8", "fp8i4", "bf8i4" + layout: str # "rcr" + pipeline: str # "mem" or "compv3" + epilogue: str # "cshuffle" (effective epilogue may differ per tile) + scheduler: str # "intrawave" + + tile_m: int + tile_n: int + tile_k: int + warp_m: int + warp_n: int + warp_k: int + warp_tile_m: int + warp_tile_n: int + warp_tile_k: int + + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + + preshuffle_aq: bool = False # APreshuffleQuant — requires pipeline="compv3" + double_smem_buffer: bool = False + k_block_per_cu: int = 1 + transpose_c: bool = False + + gfx_arch: str = _DEFAULT_GFX_ARCH + + @property + def name(self) -> str: + """Byte-exact match to codegen KERNEL_NAME (delegates to make_aquant_kernel_name).""" + return make_aquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline, + epilogue=self.epilogue, + scheduler=self.scheduler, + tile_m=self.tile_m, tile_n=self.tile_n, tile_k=self.tile_k, + warp_m=self.warp_m, warp_n=self.warp_n, warp_k=self.warp_k, + warp_tile_m=self.warp_tile_m, warp_tile_n=self.warp_tile_n, warp_tile_k=self.warp_tile_k, + quant_group_m=self.quant_group_m, + quant_group_n=self.quant_group_n, + quant_group_k=self.quant_group_k, + preshuffle_aq=self.preshuffle_aq, + ) + + def to_codegen_config(self) -> dict: + """Produce the JSON config dict consumed by unified_grouped_gemm_aquant_codegen.py.""" + return { + "variant_keys": [self.variant_key], + "layouts": [self.layout], + "pipeline": self.pipeline, + "epilogue": self.epilogue, + "scheduler": self.scheduler, + "tile_configs": [{ + "tile_m": self.tile_m, + "tile_n": self.tile_n, + "tile_k": self.tile_k, + "warp_m": self.warp_m, + "warp_n": self.warp_n, + "warp_k": self.warp_k, + "warp_tile_m": self.warp_tile_m, + "warp_tile_n": self.warp_tile_n, + "warp_tile_k": self.warp_tile_k, + }], + "quant_groups": [{ + "quant_group_m": self.quant_group_m, + "quant_group_n": self.quant_group_n, + "quant_group_k": self.quant_group_k, + }], + "preshuffle_aq": self.preshuffle_aq, + "double_smem_buffer": self.double_smem_buffer, + "k_block_per_cu": self.k_block_per_cu, + "transpose_c": self.transpose_c, + } + + +# ============================================================================= +# AQuantGemmProblem +# ============================================================================= + + +@dataclass +class AQuantGemmProblem: + M: int + N: int + K: int + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + k_batch: int = 1 + + @property + def QK_A(self) -> int: + """Number of K-groups for A: ceil(K / quant_group_k).""" + return (self.K + self.quant_group_k - 1) // self.quant_group_k + + @property + def QM_A(self) -> int: + """Number of M-groups for A: ceil(M / quant_group_m). Typically == M when gM=1.""" + return (self.M + self.quant_group_m - 1) // self.quant_group_m + + +# ============================================================================= +# AQuantGemmResult +# ============================================================================= + + +@dataclass +class AQuantGemmResult: + C: object # numpy array (M, N) + time_ms: float + kernel_name: str + + +# ============================================================================= +# AQuantDispatcherLib — thin ctypes wrapper +# ============================================================================= + + +class AQuantDispatcherLib: + """ + Loads a compiled aquant_gemm .so and wraps its C API. + + Expected .so exports: + int dispatcher_initialize() + int dispatcher_run_aquant_gemm(A, B, AQ, C, M, N, K, + stride_A, stride_B, stride_AQ, stride_C, + QK_A, QM_A, k_batch, *time_ms) + char* dispatcher_get_kernel_name() + int dispatcher_get_kernel_count() + void dispatcher_cleanup() + """ + + def __init__(self, so_path: Path): + self.so_path = Path(so_path) + if not self.so_path.exists(): + raise FileNotFoundError(f"AQuant .so not found: {self.so_path}") + self._lib = ctypes.CDLL(str(self.so_path)) + self._setup() + rc = self._lib.dispatcher_initialize() + if rc != 0: + raise RuntimeError(f"dispatcher_initialize() returned {rc}") + + def _setup(self): + lib = self._lib + + lib.dispatcher_initialize.restype = ctypes.c_int + lib.dispatcher_initialize.argtypes = [] + + lib.dispatcher_run_aquant_gemm.restype = ctypes.c_int + lib.dispatcher_run_aquant_gemm.argtypes = [ + ctypes.c_void_p, # A + ctypes.c_void_p, # B + ctypes.c_void_p, # AQ + ctypes.c_void_p, # C + ctypes.c_int64, # M + ctypes.c_int64, # N + ctypes.c_int64, # K + ctypes.c_int64, # stride_A + ctypes.c_int64, # stride_B + ctypes.c_int64, # stride_AQ + ctypes.c_int64, # stride_C + ctypes.c_int64, # QK_A + ctypes.c_int64, # QM_A + ctypes.c_int, # k_batch + ctypes.POINTER(ctypes.c_float), # time_ms + ] + + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + lib.dispatcher_get_kernel_name.argtypes = [] + + lib.dispatcher_get_kernel_count.restype = ctypes.c_int + lib.dispatcher_get_kernel_count.argtypes = [] + + lib.dispatcher_cleanup.restype = None + lib.dispatcher_cleanup.argtypes = [] + + def run( + self, + A, + B, + AQ, + C, + M: int, + N: int, + K: int, + stride_A: int, + stride_B: int, + stride_AQ: int, + stride_C: int, + QK_A: int, + QM_A: int, + k_batch: int = 1, + ) -> Tuple[int, float]: + """ + Call dispatcher_run_aquant_gemm with ctypes-wrapped pointers. + + A, B, AQ, C must be numpy arrays (C-contiguous, packed). + Returns (status, time_ms). + """ + import numpy as np + + A = np.ascontiguousarray(A) + B = np.ascontiguousarray(B) + AQ = np.ascontiguousarray(AQ) + C = np.ascontiguousarray(C) + + time_ms = ctypes.c_float(0.0) + + rc = self._lib.dispatcher_run_aquant_gemm( + A.ctypes.data_as(ctypes.c_void_p), + B.ctypes.data_as(ctypes.c_void_p), + AQ.ctypes.data_as(ctypes.c_void_p), + C.ctypes.data_as(ctypes.c_void_p), + ctypes.c_int64(M), + ctypes.c_int64(N), + ctypes.c_int64(K), + ctypes.c_int64(stride_A), + ctypes.c_int64(stride_B), + ctypes.c_int64(stride_AQ), + ctypes.c_int64(stride_C), + ctypes.c_int64(QK_A), + ctypes.c_int64(QM_A), + ctypes.c_int(k_batch), + ctypes.byref(time_ms), + ) + return rc, time_ms.value + + def get_kernel_name(self) -> str: + raw = self._lib.dispatcher_get_kernel_name() + return raw.decode("utf-8") if raw else "" + + def get_kernel_count(self) -> int: + return self._lib.dispatcher_get_kernel_count() + + def cleanup(self): + self._lib.dispatcher_cleanup() + + def __del__(self): + try: + self._lib.dispatcher_cleanup() + except Exception: + pass + + +# ============================================================================= +# AQuantGpuGemmRunner — high-level runner +# ============================================================================= + + +class AQuantGpuGemmRunner: + """ + High-level runner that loads an AQuant .so and executes GEMM on the GPU. + + Accepts numpy arrays for A, B, AQ; allocates C; returns AQuantGemmResult. + """ + + def __init__(self, so_path: Path): + self._lib = AQuantDispatcherLib(so_path) + + @property + def kernel_name(self) -> str: + return self._lib.get_kernel_name() + + def run(self, A, B, AQ, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmResult: + """ + Run AQuantGrouped GEMM. + + A shape: (M, K) dtype: fp8/bf8 + B shape: (K, N) col-major dtype: fp8/bf8 or pk_int4 + AQ shape: (QM_A, QK_A) dtype: float (for fp8/bf8) or fp8/bf8 (for fp8i4/bf8i4) + c_dtype numpy dtype for the output C buffer. Defaults to np.float16. + Returns AQuantGemmResult with C shape (M, N). + """ + import numpy as np + + M, N, K = problem.M, problem.N, problem.K + QK_A = problem.QK_A + QM_A = problem.QM_A + + if c_dtype is None: + c_dtype = np.float16 + + C = np.zeros((M, N), dtype=c_dtype) + + # AQ is RowMajor [QM_A, QK_A]: stride_AQ == QK_A (leading dim = number of K-groups) + stride_A = K + stride_B = K # B col-major [K, N]: leading dim = K + stride_AQ = QK_A + stride_C = N + + rc, time_ms = self._lib.run( + A=A, B=B, AQ=AQ, C=C, + M=M, N=N, K=K, + stride_A=stride_A, + stride_B=stride_B, + stride_AQ=stride_AQ, + stride_C=stride_C, + QK_A=QK_A, + QM_A=QM_A, + k_batch=problem.k_batch, + ) + + if rc != 0: + raise RuntimeError( + f"dispatcher_run_aquant_gemm failed with code {rc} " + f"for kernel {self.kernel_name}" + ) + + return AQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) + + +# ============================================================================= +# Subprocess helpers +# ============================================================================= + + +def _detect_gpu_arch() -> str: + """Detect current GPU arch via rocm_agent_enumerator. Falls back to gfx950.""" + try: + result = subprocess.run( + ["rocm_agent_enumerator"], + capture_output=True, text=True, timeout=10, + ) + for line in result.stdout.splitlines(): + line = line.strip() + if line.startswith("gfx") and line != "gfx000": + return line + except Exception: + pass + return _DEFAULT_GFX_ARCH + + +def _get_ck_include_dir() -> Optional[Path]: + """Attempt to locate the CK include directory relative to this file.""" + here = Path(__file__).resolve().parent + for parent in [here.parent.parent, here.parent.parent.parent]: + candidate = parent / "include" + if (candidate / "ck_tile").is_dir(): + return candidate + return None + + +def _generate_aquant_kernel( + config: AQuantKernelConfig, + output_dir: Path, +) -> Optional[Path]: + """Run unified_grouped_gemm_aquant_codegen.py for one config; return the .hpp path or None.""" + config_dict = config.to_codegen_config() + config_json = json.dumps(config_dict) + + cmd = [ + sys.executable, + str(_CODEGEN_SCRIPT), + "--output-dir", str(output_dir), + "--config-json", config_json, + ] + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + log.error("Codegen failed for %s:\n%s", config.name, result.stderr) + return None + except subprocess.TimeoutExpired: + log.error("Codegen timed out for %s", config.name) + return None + + hpp = output_dir / f"{config.name}.hpp" + if not hpp.exists(): + log.error("Codegen succeeded but %s not found", hpp) + return None + return hpp + + +def _compile_aquant_kernel( + hpp_path: Path, + so_path: Path, + gfx_arch: str, + hipcc: str = _DEFAULT_HIPCC, + extra_include_dirs: Optional[List[str]] = None, +) -> bool: + """Compile a generated .hpp into a .so via hipcc. Returns True on success.""" + ck_include = _get_ck_include_dir() + + cmd = [hipcc] + _HIPCC_BASE_FLAGS + [ + f"--offload-arch={gfx_arch}", + f"-DGFX_ARCH=\"{gfx_arch}\"", + "-include", str(hpp_path), + str(_CTYPES_LIB_SRC), + "-o", str(so_path), + ] + + if ck_include: + cmd += [f"-I{ck_include}"] + + dispatcher_include = _CTYPES_LIB_SRC.parent.parent.parent / "dispatcher" / "include" + if dispatcher_include.is_dir(): + cmd += [f"-I{dispatcher_include}"] + + if extra_include_dirs: + for d in extra_include_dirs: + cmd += [f"-I{d}"] + + log.debug("Compiling %s:\n %s", so_path.name, " ".join(cmd)) + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=600) + if result.returncode != 0: + log.error("Compile failed for %s:\n%s", so_path.name, result.stderr[-2000:]) + return False + return True + except subprocess.TimeoutExpired: + log.error("Compile timed out for %s", so_path.name) + return False + + +# ============================================================================= +# setup_multiple_aquant_dispatchers — build pipeline +# ============================================================================= + + +def setup_multiple_aquant_dispatchers( + configs: List[AQuantKernelConfig], + output_dir: Optional[Path] = None, + hipcc: str = _DEFAULT_HIPCC, + gfx_arch: Optional[str] = None, + extra_include_dirs: Optional[List[str]] = None, + parallel: bool = True, + max_workers: Optional[int] = None, +) -> List[Optional[Path]]: + """ + For each AQuantKernelConfig: codegen -> hipcc compile -> .so path. + + Returns a list parallel to `configs` — each entry is the Path to the + compiled .so, or None if that config failed. + """ + if not configs: + return [] + + arch = gfx_arch or _detect_gpu_arch() + base_dir = output_dir or Path(tempfile.mkdtemp(prefix="aquant_dispatcher_")) + base_dir.mkdir(parents=True, exist_ok=True) + + headers_dir = base_dir / "generated_kernels" + so_dir = base_dir / "libs" + headers_dir.mkdir(exist_ok=True) + so_dir.mkdir(exist_ok=True) + + log.info("Building %d AQuant kernel(s) for %s into %s", len(configs), arch, base_dir) + + seen: Dict[str, int] = {} + deduped: List[Tuple[int, AQuantKernelConfig]] = [] + for i, cfg in enumerate(configs): + if cfg.name not in seen: + seen[cfg.name] = i + deduped.append((i, cfg)) + + results: List[Optional[Path]] = [None] * len(configs) + + def _build_one(idx: int, cfg: AQuantKernelConfig) -> Tuple[int, Optional[Path]]: + hpp = _generate_aquant_kernel(cfg, headers_dir) + if hpp is None: + return idx, None + + so = so_dir / f"lib{cfg.name}_{arch}.so" + if so.exists(): + log.info(" [cached] %s", so.name) + return idx, so + + ok = _compile_aquant_kernel( + hpp_path=hpp, + so_path=so, + gfx_arch=arch, + hipcc=hipcc, + extra_include_dirs=extra_include_dirs, + ) + return idx, so if ok else None + + if parallel and len(deduped) > 1: + workers = max_workers or min(len(deduped), os.cpu_count() or 4) + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex: + futures = {ex.submit(_build_one, idx, cfg): (idx, cfg) for idx, cfg in deduped} + for fut in concurrent.futures.as_completed(futures): + try: + idx, so_path = fut.result() + results[idx] = so_path + if so_path: + log.info(" built %s", so_path.name) + else: + _, cfg = futures[fut] + log.error(" FAILED %s", cfg.name) + except Exception as e: + _, cfg = futures[fut] + log.error(" EXCEPTION for %s: %s", cfg.name, e) + else: + for idx, cfg in deduped: + _, so_path = _build_one(idx, cfg) + results[idx] = so_path + + for i, cfg in enumerate(configs): + if results[i] is None: + first_idx = seen.get(cfg.name) + if first_idx is not None and first_idx != i: + results[i] = results[first_idx] + + built = sum(1 for r in results if r is not None) + log.info("Built %d / %d AQuant kernels", built, len(configs)) + return results + + +# ============================================================================= +# Default configs (mapped from reference examples) +# ============================================================================= + + +def default_fp8_config( + quant_group_k: int = 128, + quant_group_m: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> AQuantKernelConfig: + """fp8 AQuant decode config (GemmConfigQuantDecodeInterwave, Mem pipeline).""" + return AQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="mem", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=quant_group_m, + quant_group_n=1, + quant_group_k=quant_group_k, + preshuffle_aq=False, + gfx_arch=gfx_arch, + ) + + +def default_bf8_config( + quant_group_k: int = 128, + quant_group_m: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> AQuantKernelConfig: + """bf8 AQuant decode config (GemmConfigQuantDecodeInterwave, Mem pipeline).""" + return AQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="mem", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=quant_group_m, + quant_group_n=1, + quant_group_k=quant_group_k, + preshuffle_aq=False, + gfx_arch=gfx_arch, + ) + + +def default_fp8i4_config( + quant_group_k: int = 128, + quant_group_m: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> AQuantKernelConfig: + """fp8i4 AQuant decode config (A=fp8, B=pk_int4, AQ=fp8; Mem pipeline).""" + return AQuantKernelConfig( + variant_key="fp8i4", + layout="rcr", + pipeline="mem", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=quant_group_m, + quant_group_n=1, + quant_group_k=quant_group_k, + preshuffle_aq=False, + gfx_arch=gfx_arch, + ) + + +def default_bf8i4_config( + quant_group_k: int = 128, + quant_group_m: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> AQuantKernelConfig: + """bf8i4 AQuant decode config (A=bf8, B=pk_int4, AQ=bf8; Mem pipeline).""" + return AQuantKernelConfig( + variant_key="bf8i4", + layout="rcr", + pipeline="mem", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=quant_group_m, + quant_group_n=1, + quant_group_k=quant_group_k, + preshuffle_aq=False, + gfx_arch=gfx_arch, + ) + + +def default_fp8_preshuffleaq_config( + quant_group_k: int = 128, + quant_group_m: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> AQuantKernelConfig: + """fp8 preshuffle-AQ config (GemmConfigPreshuffleQuantDecode, CompV3 pipeline). + + APreshuffleQuant=true, BPreshuffleQuant=true — both scale tensors are preshuffled. + Tile: 16x64x256, warp_tile_k=128 (gfx950 FlatMM, 8-bit float). + """ + return AQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + quant_group_m=quant_group_m, + quant_group_n=1, + quant_group_k=quant_group_k, + preshuffle_aq=True, + gfx_arch=gfx_arch, + ) + + +def default_bf8_preshuffleaq_config( + quant_group_k: int = 128, + quant_group_m: int = 1, + gfx_arch: str = _DEFAULT_GFX_ARCH, +) -> AQuantKernelConfig: + """bf8 preshuffle-AQ config (GemmConfigPreshuffleQuantDecode, CompV3 pipeline).""" + return AQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="compv3", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + quant_group_m=quant_group_m, + quant_group_n=1, + quant_group_k=quant_group_k, + preshuffle_aq=True, + gfx_arch=gfx_arch, + ) diff --git a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_abquant_utils.py new file mode 100644 index 000000000000..1d102e28f00c --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_abquant_utils.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +CPU-only unit tests for grouped_gemm_abquant_utils.py. + +Tests kernel name generation, config serialization, and problem dimension helpers. +No GPU or hipcc required. + +Run: + python3 -m pytest dispatcher/tests/test_grouped_gemm_abquant_utils.py -v +""" + +import json +import re +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + +from grouped_gemm_abquant_utils import ( + ABQuantKernelConfig, + ABQuantGemmProblem, + default_fp8_compv3_config, + default_bf8_compv3_config, + default_fp8_eightwaves_config, + default_bf8_eightwaves_config, + default_fp8_preshuffleb_config, + default_bf8_preshuffleb_config, +) + + +# ============================================================================= +# ABQuantKernelConfig.name — byte-exact match with codegen KERNEL_NAME +# ============================================================================= + + +class TestKernelName: + + def test_fp8_compv3_default_name(self): + # tile_n=128, warp_n=4, warp_tile_n=16 → N_repeat=2 (even) → effective epilogue=permute_n + cfg = ABQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=128, + bquant_group_m=1, bquant_group_n=1, bquant_group_k=128, + ) + assert cfg.name == ( + "grouped_gemm_abquant_fp8_rcr_compv3_permute_n_intrawave_" + "128x128x128_1x4x1_16x16x16_aqg1x1x128_bqg1x1x128" + ) + + def test_bf8_compv3_default_name(self): + # tile_n=128, warp_n=4, warp_tile_n=16 → N_repeat=2 (even) → effective epilogue=permute_n + cfg = ABQuantKernelConfig( + variant_key="bf8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=128, + bquant_group_m=1, bquant_group_n=1, bquant_group_k=128, + ) + assert cfg.name == ( + "grouped_gemm_abquant_bf8_rcr_compv3_permute_n_intrawave_" + "128x128x128_1x4x1_16x16x16_aqg1x1x128_bqg1x1x128" + ) + + def test_prefix_is_grouped_gemm_abquant(self): + for cfg in [default_fp8_compv3_config(), default_bf8_compv3_config(), + default_fp8_eightwaves_config(), default_fp8_preshuffleb_config()]: + assert cfg.name.startswith("grouped_gemm_abquant_") + + def test_aqg_and_bqg_segments_present(self): + cfg = default_fp8_compv3_config(quant_group_k=128) + assert "aqg1x1x128" in cfg.name + assert "bqg1x1x128" in cfg.name + + def test_different_bquant_group_n_unique_names(self): + cfg1 = default_fp8_compv3_config(bquant_group_n=1) + cfg2 = default_fp8_compv3_config(bquant_group_n=128) + assert cfg1.name != cfg2.name + assert "bqg1x1x128" in cfg1.name + assert "bqg1x128x128" in cfg2.name + + def test_transpose_c_suffix(self): + cfg = default_fp8_eightwaves_config() + assert "transposec" in cfg.name + + def test_no_transpose_c_no_suffix(self): + cfg = default_fp8_compv3_config() + assert "transposec" not in cfg.name + + def test_preshuffleb_suffix(self): + cfg = default_fp8_preshuffleb_config() + assert "_preshuffleb" in cfg.name + + def test_preshuffle_aq_suffix(self): + cfg = ABQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=128, + bquant_group_m=1, bquant_group_n=1, bquant_group_k=128, + preshuffle_aq=True, + ) + assert "_preshuffleaq" in cfg.name + + def test_preshuffle_bq_suffix(self): + cfg = ABQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=128, + bquant_group_m=1, bquant_group_n=1, bquant_group_k=128, + preshuffle_bq=True, + ) + assert "_preshufflebq" in cfg.name + + def test_name_no_spaces(self): + for cfg in [default_fp8_compv3_config(), default_fp8_eightwaves_config(), + default_fp8_preshuffleb_config()]: + assert " " not in cfg.name + + def test_name_only_valid_chars(self): + cfg = default_fp8_compv3_config() + assert re.match(r'^[a-z0-9_x]+$', cfg.name), f"Unexpected chars in: {cfg.name}" + + def test_variant_key_in_name(self): + assert "fp8" in default_fp8_compv3_config().name + assert "bf8" in default_bf8_compv3_config().name + + def test_pipeline_in_name(self): + assert "compv3" in default_fp8_compv3_config().name + assert "eightwaves" in default_fp8_eightwaves_config().name + assert "preshuffleb" in default_fp8_preshuffleb_config().name + + def test_all_defaults_unique_names(self): + names = [ + default_fp8_compv3_config().name, + default_bf8_compv3_config().name, + default_fp8_eightwaves_config().name, + default_bf8_eightwaves_config().name, + default_fp8_preshuffleb_config().name, + default_bf8_preshuffleb_config().name, + ] + assert len(names) == len(set(names)), "All defaults must have unique names" + + +# ============================================================================= +# ABQuantKernelConfig — constraint validation +# ============================================================================= + + +class TestConstraintValidation: + + def test_mismatched_kk_raises(self): + with pytest.raises(ValueError, match="aquant_group_k"): + ABQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=64, # different! + bquant_group_m=1, bquant_group_n=1, bquant_group_k=128, + ) + + def test_matching_kk_succeeds(self): + cfg = ABQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + aquant_group_m=1, aquant_group_n=1, aquant_group_k=64, + bquant_group_m=1, bquant_group_n=1, bquant_group_k=64, # same! + ) + assert cfg.aquant_group_k == cfg.bquant_group_k + + +# ============================================================================= +# ABQuantGemmProblem dimension properties +# ============================================================================= + + +class TestABQuantGemmProblem: + + def test_qk_a_basic(self): + p = ABQuantGemmProblem(M=128, N=128, K=256, aquant_group_k=128) + assert p.QK_A == 2 + + def test_qk_b_basic(self): + p = ABQuantGemmProblem(M=128, N=128, K=256, bquant_group_k=128) + assert p.QK_B == 2 + + def test_qm_a_is_m_when_gm1(self): + p = ABQuantGemmProblem(M=128, N=128, K=256, aquant_group_m=1) + assert p.QM_A == 128 + + def test_qn_b_basic(self): + p = ABQuantGemmProblem(M=128, N=128, K=256, bquant_group_n=1) + assert p.QN_B == 128 + + def test_qn_b_grouped(self): + p = ABQuantGemmProblem(M=128, N=128, K=256, bquant_group_n=128) + assert p.QN_B == 1 + + def test_qk_a_ceil(self): + p = ABQuantGemmProblem(M=128, N=128, K=257, aquant_group_k=128) + assert p.QK_A == 3 # ceil(257/128)=3 + + def test_qk_b_ceil(self): + p = ABQuantGemmProblem(M=128, N=128, K=257, bquant_group_k=128) + assert p.QK_B == 3 + + def test_k_batch_default(self): + p = ABQuantGemmProblem(M=128, N=128, K=256) + assert p.k_batch == 1 + + def test_default_group_sizes(self): + p = ABQuantGemmProblem(M=128, N=128, K=256) + assert p.aquant_group_k == 128 + assert p.bquant_group_k == 128 + assert p.aquant_group_m == 1 + assert p.bquant_group_n == 1 + + +# ============================================================================= +# to_codegen_config serialization +# ============================================================================= + + +class TestCodegenConfig: + + def test_round_trip_variant_key(self): + cfg = default_bf8_compv3_config() + d = cfg.to_codegen_config() + assert d["variant_keys"] == ["bf8"] + + def test_round_trip_pipeline(self): + cfg = default_fp8_compv3_config() + d = cfg.to_codegen_config() + assert d["pipeline"] == "compv3" + + def test_round_trip_aquant_groups(self): + cfg = default_fp8_compv3_config(quant_group_k=64, bquant_group_n=4) + d = cfg.to_codegen_config() + aqg = d["aquant_groups"][0] + assert aqg["aquant_group_m"] == 1 + assert aqg["aquant_group_n"] == 1 + assert aqg["aquant_group_k"] == 64 + + def test_round_trip_bquant_groups(self): + cfg = default_fp8_compv3_config(quant_group_k=64, bquant_group_n=4) + d = cfg.to_codegen_config() + bqg = d["bquant_groups"][0] + assert bqg["bquant_group_n"] == 4 + assert bqg["bquant_group_k"] == 64 + + def test_round_trip_preshuffle_b(self): + cfg = default_fp8_preshuffleb_config() + d = cfg.to_codegen_config() + assert d["preshuffle_b"] is True + assert d["preshuffle_aq"] is False + assert d["preshuffle_bq"] is False + + def test_round_trip_transpose_c(self): + cfg = default_fp8_eightwaves_config() + d = cfg.to_codegen_config() + assert d["transpose_c"] is True + + def test_json_serializable(self): + cfg = default_fp8_compv3_config() + json_str = json.dumps(cfg.to_codegen_config()) + parsed = json.loads(json_str) + assert parsed["variant_keys"] == ["fp8"] + + +# ============================================================================= +# Default config factories +# ============================================================================= + + +class TestDefaultConfigs: + + def test_compv3_not_transpose_c(self): + assert default_fp8_compv3_config().transpose_c is False + assert default_bf8_compv3_config().transpose_c is False + + def test_eightwaves_transpose_c(self): + assert default_fp8_eightwaves_config().transpose_c is True + assert default_bf8_eightwaves_config().transpose_c is True + + def test_preshuffleb_preshuffle_b_set(self): + assert default_fp8_preshuffleb_config().preshuffle_b is True + assert default_bf8_preshuffleb_config().preshuffle_b is True + + def test_compv3_preshuffle_flags_unset(self): + cfg = default_fp8_compv3_config() + assert cfg.preshuffle_b is False + assert cfg.preshuffle_aq is False + assert cfg.preshuffle_bq is False + + def test_custom_bquant_group_n(self): + cfg = default_fp8_compv3_config(bquant_group_n=128) + assert cfg.bquant_group_n == 128 + assert "bqg1x128x128" in cfg.name + + def test_gfx_arch_propagated(self): + cfg = default_fp8_compv3_config(gfx_arch="gfx942") + assert cfg.gfx_arch == "gfx942" + + def test_aquant_and_bquant_kk_equal(self): + for cfg in [ + default_fp8_compv3_config(), + default_fp8_eightwaves_config(), + default_fp8_preshuffleb_config(), + ]: + assert cfg.aquant_group_k == cfg.bquant_group_k + + +# ============================================================================= +# Edge cases +# ============================================================================= + + +class TestEdgeCases: + + def test_qk_a_k_equals_group_k(self): + p = ABQuantGemmProblem(M=128, N=128, K=128, aquant_group_k=128) + assert p.QK_A == 1 + + def test_name_consistency_across_instances(self): + assert default_fp8_compv3_config().name == default_fp8_compv3_config().name + + def test_large_dimensions(self): + p = ABQuantGemmProblem( + M=4096, N=4096, K=8192, + aquant_group_m=1, aquant_group_k=128, + bquant_group_n=128, bquant_group_k=128, + ) + assert p.QK_A == 64 # 8192/128=64 + assert p.QK_B == 64 + assert p.QN_B == 32 # 4096/128=32 + assert p.QM_A == 4096 + + def test_transpose_c_in_name_only_when_true(self): + with_tc = default_fp8_eightwaves_config() + without_tc = default_fp8_compv3_config() + assert "transposec" in with_tc.name + assert "transposec" not in without_tc.name diff --git a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_aquant_utils.py new file mode 100644 index 000000000000..be95e88c0822 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_aquant_utils.py @@ -0,0 +1,410 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +CPU-only unit tests for grouped_gemm_aquant_utils.py. + +Tests kernel name generation, config serialization, and problem dimension helpers. +No GPU or hipcc required. + +Run: + python3 -m pytest dispatcher/tests/test_grouped_gemm_aquant_utils.py -v +""" + +import math +import re +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + +from grouped_gemm_aquant_utils import ( + AQuantKernelConfig, + AQuantGemmProblem, + default_fp8_config, + default_bf8_config, + default_fp8i4_config, + default_bf8i4_config, + default_fp8_preshuffleaq_config, + default_bf8_preshuffleaq_config, +) + + +# ============================================================================= +# AQuantKernelConfig.name — byte-exact match with codegen KERNEL_NAME +# ============================================================================= + + +class TestKernelName: + + def test_fp8_rcr_mem_default_name(self): + cfg = AQuantKernelConfig( + variant_key="fp8", + layout="rcr", + pipeline="mem", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert cfg.name == ( + "grouped_gemm_aquant_fp8_rcr_mem_cshuffle_intrawave_" + "16x64x256_1x4x1_16x16x16_aqg1x1x128" + ) + + def test_bf8_rcr_mem_default_name(self): + cfg = AQuantKernelConfig( + variant_key="bf8", + layout="rcr", + pipeline="mem", + epilogue="cshuffle", + scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert cfg.name == ( + "grouped_gemm_aquant_bf8_rcr_mem_cshuffle_intrawave_" + "16x64x256_1x4x1_16x16x16_aqg1x1x128" + ) + + def test_preshuffle_aq_suffix(self): + cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_aq=True, + ) + assert cfg.name.endswith("_preshuffleaq") + + def test_no_preshuffle_no_suffix(self): + cfg = default_fp8_config() + assert "preshuffleaq" not in cfg.name + + def test_prefix_is_grouped_gemm_aquant(self): + for cfg in [default_fp8_config(), default_bf8_config(), + default_fp8i4_config(), default_bf8i4_config()]: + assert cfg.name.startswith("grouped_gemm_aquant_") + + def test_different_variants_unique_names(self): + names = [ + default_fp8_config().name, + default_bf8_config().name, + default_fp8i4_config().name, + default_bf8i4_config().name, + ] + assert len(names) == len(set(names)) + + def test_different_quant_groups_unique_names(self): + def make(gk): + return AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="mem", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=gk, + ).name + + names = [make(64), make(128), make(256)] + assert len(names) == len(set(names)) + + def test_aqg_segment_present(self): + cfg = default_fp8_config(quant_group_k=128, quant_group_m=1) + assert "aqg1x1x128" in cfg.name + + def test_name_no_spaces(self): + cfg = default_fp8_config() + assert " " not in cfg.name + + def test_name_only_valid_chars(self): + cfg = default_fp8_config() + assert re.match(r'^[a-z0-9_x]+$', cfg.name), f"Unexpected chars in: {cfg.name}" + + def test_permute_n_epilogue_selected_when_tile_allows(self): + # tile_n=128, warp_n=4, warp_tile_n=16 → N_repeat=128/(4*16)=2; 2%2==0 → permute_n + cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="mem", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + assert "permute_n" in cfg.name + + def test_cshuffle_epilogue_when_quant_group_n_gt1(self): + # quant_group_n=8 → PermuteN not selected even if N_repeat is even + cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="mem", epilogue="cshuffle", scheduler="intrawave", + tile_m=128, tile_n=128, tile_k=128, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=8, quant_group_k=128, + ) + assert "cshuffle" in cfg.name + assert "permute_n" not in cfg.name + + def test_preshuffle_aq_compv3_name(self): + cfg = default_fp8_preshuffleaq_config() + assert "compv3" in cfg.name + assert cfg.name.endswith("_preshuffleaq") + + def test_fp8i4_in_name(self): + cfg = default_fp8i4_config() + assert "fp8i4" in cfg.name + + def test_bf8i4_in_name(self): + cfg = default_bf8i4_config() + assert "bf8i4" in cfg.name + + +# ============================================================================= +# AQuantGemmProblem dimension properties +# ============================================================================= + + +class TestAQuantGemmProblem: + + def test_qk_a_exact_divisible(self): + p = AQuantGemmProblem(M=16, N=64, K=256, quant_group_k=128) + assert p.QK_A == 2 + + def test_qk_a_ceil_division(self): + p = AQuantGemmProblem(M=16, N=64, K=257, quant_group_k=128) + assert p.QK_A == 3 # ceil(257/128)=3 + + def test_qm_a_is_m_when_gm1(self): + p = AQuantGemmProblem(M=32, N=64, K=256, quant_group_m=1, quant_group_k=128) + assert p.QM_A == 32 # ceil(32/1)=32 + + def test_qm_a_groups_when_gm_gt1(self): + p = AQuantGemmProblem(M=32, N=64, K=256, quant_group_m=8, quant_group_k=128) + assert p.QM_A == 4 # ceil(32/8)=4 + + def test_qm_a_ceil(self): + p = AQuantGemmProblem(M=33, N=64, K=256, quant_group_m=8) + assert p.QM_A == 5 # ceil(33/8)=5 + + def test_qk_a_group64(self): + p = AQuantGemmProblem(M=16, N=64, K=256, quant_group_k=64) + assert p.QK_A == 4 + + def test_default_group_sizes(self): + p = AQuantGemmProblem(M=16, N=64, K=256) + assert p.quant_group_m == 1 + assert p.quant_group_n == 1 + assert p.quant_group_k == 128 + + def test_k_batch_default(self): + p = AQuantGemmProblem(M=16, N=64, K=256) + assert p.k_batch == 1 + + def test_k_batch_splitk(self): + p = AQuantGemmProblem(M=16, N=64, K=256, k_batch=4) + assert p.k_batch == 4 + + +# ============================================================================= +# to_codegen_config serialization +# ============================================================================= + + +class TestCodegenConfig: + + def test_round_trip_variant_key(self): + cfg = default_bf8_config() + d = cfg.to_codegen_config() + assert d["variant_keys"] == ["bf8"] + + def test_round_trip_pipeline(self): + cfg = default_fp8_config() + d = cfg.to_codegen_config() + assert d["pipeline"] == "mem" + + def test_round_trip_preshuffle_aq_false(self): + cfg = default_fp8_config() + d = cfg.to_codegen_config() + assert d["preshuffle_aq"] is False + + def test_round_trip_preshuffle_aq_true(self): + cfg = default_fp8_preshuffleaq_config() + d = cfg.to_codegen_config() + assert d["preshuffle_aq"] is True + assert d["pipeline"] == "compv3" + + def test_round_trip_tile_config(self): + cfg = default_fp8_config() + d = cfg.to_codegen_config() + tc = d["tile_configs"][0] + assert tc["tile_m"] == 16 + assert tc["tile_n"] == 64 + assert tc["tile_k"] == 256 + + def test_round_trip_quant_groups(self): + cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", pipeline="mem", + epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=2, quant_group_n=1, quant_group_k=64, + ) + d = cfg.to_codegen_config() + qg = d["quant_groups"][0] + assert qg["quant_group_m"] == 2 + assert qg["quant_group_n"] == 1 + assert qg["quant_group_k"] == 64 + + def test_round_trip_transpose_c(self): + cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", pipeline="compv3", + epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_aq=True, transpose_c=True, + ) + d = cfg.to_codegen_config() + assert d["transpose_c"] is True + + def test_config_is_serializable_as_json(self): + import json + cfg = default_fp8_config() + json_str = json.dumps(cfg.to_codegen_config()) + parsed = json.loads(json_str) + assert parsed["variant_keys"] == ["fp8"] + + +# ============================================================================= +# Default config factories +# ============================================================================= + + +class TestDefaultConfigs: + + def test_default_fp8_pipeline_is_mem(self): + assert default_fp8_config().pipeline == "mem" + + def test_default_bf8_pipeline_is_mem(self): + assert default_bf8_config().pipeline == "mem" + + def test_default_fp8i4_variant_key(self): + assert default_fp8i4_config().variant_key == "fp8i4" + + def test_default_bf8i4_variant_key(self): + assert default_bf8i4_config().variant_key == "bf8i4" + + def test_preshuffle_aq_pipeline_is_compv3(self): + assert default_fp8_preshuffleaq_config().pipeline == "compv3" + assert default_bf8_preshuffleaq_config().pipeline == "compv3" + + def test_preshuffle_aq_flag_set(self): + assert default_fp8_preshuffleaq_config().preshuffle_aq is True + assert default_bf8_preshuffleaq_config().preshuffle_aq is True + + def test_non_preshuffle_flag_unset(self): + assert default_fp8_config().preshuffle_aq is False + assert default_bf8_config().preshuffle_aq is False + + def test_default_layout_is_rcr(self): + for cfg in [default_fp8_config(), default_bf8_config(), + default_fp8i4_config(), default_bf8i4_config()]: + assert cfg.layout == "rcr" + + def test_custom_quant_group_k(self): + cfg = default_fp8_config(quant_group_k=64) + assert cfg.quant_group_k == 64 + assert "aqg1x1x64" in cfg.name + + def test_custom_quant_group_m(self): + cfg = default_fp8_config(quant_group_m=4) + assert cfg.quant_group_m == 4 + assert "aqg4x1x128" in cfg.name + + def test_gfx_arch_propagated(self): + cfg = default_fp8_config(gfx_arch="gfx942") + assert cfg.gfx_arch == "gfx942" + + def test_all_defaults_have_unique_names(self): + names = [ + default_fp8_config().name, + default_bf8_config().name, + default_fp8i4_config().name, + default_bf8i4_config().name, + default_fp8_preshuffleaq_config().name, + default_bf8_preshuffleaq_config().name, + ] + assert len(names) == len(set(names)), "All default configs must produce unique names" + + +# ============================================================================= +# Edge cases +# ============================================================================= + + +class TestEdgeCases: + + def test_qk_a_k_equals_group_k(self): + p = AQuantGemmProblem(M=16, N=64, K=128, quant_group_k=128) + assert p.QK_A == 1 + + def test_qk_a_k_one(self): + p = AQuantGemmProblem(M=16, N=64, K=1, quant_group_k=128) + assert p.QK_A == 1 + + def test_qm_a_large_m(self): + p = AQuantGemmProblem(M=4096, N=1024, K=8192, quant_group_m=1, quant_group_k=128) + assert p.QM_A == 4096 + assert p.QK_A == 64 # 8192/128=64 + + def test_name_consistency_across_instances(self): + cfg1 = default_fp8_config() + cfg2 = default_fp8_config() + assert cfg1.name == cfg2.name + + def test_preshuffle_aq_false_does_not_append_suffix(self): + cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="mem", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_aq=False, + ) + assert not cfg.name.endswith("_preshuffleaq") + + def test_mem_and_compv3_pipelines_produce_different_names(self): + mem_cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="mem", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + ) + compv3_cfg = AQuantKernelConfig( + variant_key="fp8", layout="rcr", + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + tile_m=16, tile_n=64, tile_k=256, + warp_m=1, warp_n=4, warp_k=1, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + quant_group_m=1, quant_group_n=1, quant_group_k=128, + preshuffle_aq=True, + ) + assert mem_cfg.name != compv3_cfg.name + assert "mem" in mem_cfg.name + assert "compv3" in compv3_cfg.name From d05303d202cac24e01ad2485f7282fa849199ad8 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 22 Jul 2026 12:04:52 -0400 Subject: [PATCH 11/35] fix(ck-tile): sync quant dispatcher with PR #9166 patterns - 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 --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 39 +++++++++++++++++-- .../grouped_gemm_abquant_ctypes_lib.cpp | 30 +++++++------- .../tests/test_grouped_gemm_bquant_utils.py | 24 ++++++------ 3 files changed, 63 insertions(+), 30 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 082c6f33f9fc..3bc153fc7c4e 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -167,9 +167,22 @@ if(BQUANT_GEMM_KERNEL_HEADERS) grouped_gemm_bquant_ctypes_lib.cpp KERNEL_HEADER ${BQUANT_GEMM_KERNEL_HEADER} ) + # Resolve GFX arch from CMAKE_HIP_ARCHITECTURES if available, falling back to + # CK_TILE_BQUANT_GFX_ARCH cache variable (default gfx942). + # The Python build path always passes -DGFX_ARCH=... explicitly via hipcc. + 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="gfx950" + GFX_ARCH="${CK_TILE_BQUANT_GFX_ARCH}" ) else() message(STATUS "No BQuant GEMM kernel found for ctypes lib - building without kernel") @@ -198,9 +211,19 @@ if(AQUANT_GEMM_KERNEL_HEADERS) grouped_gemm_aquant_ctypes_lib.cpp KERNEL_HEADER ${AQUANT_GEMM_KERNEL_HEADER} ) + if(NOT DEFINED CK_TILE_AQUANT_GFX_ARCH OR CK_TILE_AQUANT_GFX_ARCH STREQUAL "") + if(CMAKE_HIP_ARCHITECTURES) + list(GET CMAKE_HIP_ARCHITECTURES 0 _aquant_gfx_arch) + else() + set(_aquant_gfx_arch "gfx942") + endif() + set(CK_TILE_AQUANT_GFX_ARCH "${_aquant_gfx_arch}" CACHE STRING + "GFX arch for the CMake-built AQuant ctypes .so (default: first CMAKE_HIP_ARCHITECTURES or gfx942)") + endif() + message(STATUS "AQuant ctypes lib GFX_ARCH: ${CK_TILE_AQUANT_GFX_ARCH}") target_compile_definitions(dispatcher_grouped_gemm_aquant_lib PRIVATE CK_TILE_SINGLE_KERNEL_INCLUDE - GFX_ARCH="gfx950" + GFX_ARCH="${CK_TILE_AQUANT_GFX_ARCH}" ) else() message(STATUS "No AQuant GEMM kernel found for ctypes lib - building without kernel") @@ -229,9 +252,19 @@ if(ABQUANT_GEMM_KERNEL_HEADERS) grouped_gemm_abquant_ctypes_lib.cpp KERNEL_HEADER ${ABQUANT_GEMM_KERNEL_HEADER} ) + if(NOT DEFINED CK_TILE_ABQUANT_GFX_ARCH OR CK_TILE_ABQUANT_GFX_ARCH STREQUAL "") + if(CMAKE_HIP_ARCHITECTURES) + list(GET CMAKE_HIP_ARCHITECTURES 0 _abquant_gfx_arch) + else() + set(_abquant_gfx_arch "gfx942") + endif() + set(CK_TILE_ABQUANT_GFX_ARCH "${_abquant_gfx_arch}" CACHE STRING + "GFX arch for the CMake-built ABQuant ctypes .so (default: first CMAKE_HIP_ARCHITECTURES or gfx942)") + endif() + message(STATUS "ABQuant ctypes lib GFX_ARCH: ${CK_TILE_ABQUANT_GFX_ARCH}") target_compile_definitions(dispatcher_grouped_gemm_abquant_lib PRIVATE CK_TILE_SINGLE_KERNEL_INCLUDE - GFX_ARCH="gfx950" + GFX_ARCH="${CK_TILE_ABQUANT_GFX_ARCH}" ) else() message(STATUS "No ABQuant GEMM kernel found for ctypes lib - building without kernel") diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp index 855ec91b73a5..bea8625965da 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -14,8 +14,8 @@ * AQuantGroupSize, BQuantGroupSize * * ABQuant: both A-side and B-side quantization active. - * AQ[ceil(M/aM), ceil(K/aK)] — A-side scale tensor (RowMajor) - * BQ[ceil(K/bK), ceil(N/bN)] — B-side scale tensor (RowMajor) + * AQ[ceil(M/aM), ceil(K/aK)] - A-side scale tensor (RowMajor) + * BQ[ceil(K/bK), ceil(N/bN)] - B-side scale tensor (RowMajor) * Constraint: AQuantGroupSize::kK == BQuantGroupSize::kK * * Memory model: host-pointer (this library owns hipMalloc/hipMemcpy/hipFree). @@ -121,19 +121,19 @@ int dispatcher_initialize() * A, B, AQ, BQ, C are host pointers. This function manages device memory internally. * * Parameters: - * A, B, AQ, BQ, C — host data pointers - * M, N, K — matrix dimensions - * stride_A — leading dim of A (row-major: K) - * stride_B — leading dim of B (col-major: K) - * stride_AQ — leading dim of AQ (row-major: ceil(K/aK)) - * stride_BQ — leading dim of BQ (row-major: ceil(N/bN)) - * stride_C — leading dim of C (row-major: N) - * QK_A — ceil(K / aquant_group_k) - * QM_A — ceil(M / aquant_group_m) (typically == M when aM=1) - * QK_B — ceil(K / bquant_group_k) - * QN_B — ceil(N / bquant_group_n) - * k_batch — split-K factor (1 = no split) - * time_ms — output: kernel execution time in ms (may be NULL) + * A, B, AQ, BQ, C - host data pointers + * M, N, K - matrix dimensions + * stride_A - leading dim of A (row-major: K) + * stride_B - leading dim of B (col-major: K) + * stride_AQ - leading dim of AQ (row-major: ceil(K/aK)) + * stride_BQ - leading dim of BQ (row-major: ceil(N/bN)) + * stride_C - leading dim of C (row-major: N) + * QK_A - ceil(K / aquant_group_k) + * QM_A - ceil(M / aquant_group_m) (typically == M when aM=1) + * QK_B - ceil(K / bquant_group_k) + * QN_B - ceil(N / bquant_group_n) + * k_batch - split-K factor (1 = no split) + * time_ms - output: kernel execution time in ms (may be NULL) * * Returns 0 on success, negative on error. */ diff --git a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py index 5afd3332c293..2ee9cf8077e8 100644 --- a/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_grouped_gemm_bquant_utils.py @@ -348,28 +348,28 @@ class TestPhase3Configs: def test_fp8_preshuffleb_name(self): cfg = default_fp8_preshuffleb_config() assert cfg.name == ( - "grouped_gemm_bquant_fp8_rcr_preshuffleb_cshuffle_intrawave_" + "grouped_gemm_bquant_fp8_rcr_preshuffleb_permute_n_intrawave_" "128x128x128_1x4x1_16x16x128_qg1x1x128_preshuffleb" ) def test_bf8_preshuffleb_name(self): cfg = default_bf8_preshuffleb_config() assert cfg.name == ( - "grouped_gemm_bquant_bf8_rcr_preshuffleb_cshuffle_intrawave_" + "grouped_gemm_bquant_bf8_rcr_preshuffleb_permute_n_intrawave_" "128x128x128_1x4x1_16x16x128_qg1x1x128_preshuffleb" ) def test_fp8i4_preshuffleb_name(self): cfg = default_fp8i4_preshuffleb_config() assert cfg.name == ( - "grouped_gemm_bquant_fp8i4_rcr_preshuffleb_cshuffle_intrawave_" + "grouped_gemm_bquant_fp8i4_rcr_preshuffleb_permute_n_intrawave_" "128x128x128_1x4x1_16x16x32_qg1x1x128_preshuffleb" ) def test_bf8i4_preshuffleb_name(self): cfg = default_bf8i4_preshuffleb_config() assert cfg.name == ( - "grouped_gemm_bquant_bf8i4_rcr_preshuffleb_cshuffle_intrawave_" + "grouped_gemm_bquant_bf8i4_rcr_preshuffleb_permute_n_intrawave_" "128x128x128_1x4x1_16x16x32_qg1x1x128_preshuffleb" ) @@ -395,21 +395,21 @@ def test_preshuffleb_fp8_warp_tile_k(self): def test_fp8_preshufflequant_name(self): cfg = default_fp8_preshufflequant_config() assert cfg.name == ( - "grouped_gemm_bquant_fp8_rcr_compv3_cshuffle_intrawave_" + "grouped_gemm_bquant_fp8_rcr_compv3_permute_n_intrawave_" "128x128x128_1x4x1_16x16x128_qg1x1x128_preshufflebq" ) def test_bf8_preshufflequant_name(self): cfg = default_bf8_preshufflequant_config() assert cfg.name == ( - "grouped_gemm_bquant_bf8_rcr_compv3_cshuffle_intrawave_" + "grouped_gemm_bquant_bf8_rcr_compv3_permute_n_intrawave_" "128x128x128_1x4x1_16x16x128_qg1x1x128_preshufflebq" ) def test_fp8i4_preshufflequant_name(self): cfg = default_fp8i4_preshufflequant_config() assert cfg.name == ( - "grouped_gemm_bquant_fp8i4_rcr_compv3_cshuffle_intrawave_" + "grouped_gemm_bquant_fp8i4_rcr_compv3_permute_n_intrawave_" "128x128x128_1x4x1_16x16x32_qg1x1x128_preshufflebq" ) @@ -431,14 +431,14 @@ def test_preshufflequant_16n_group(self): def test_fp8_preshuffleb_bquant_name(self): cfg = default_fp8_preshuffleb_bquant_config() assert cfg.name == ( - "grouped_gemm_bquant_fp8_rcr_preshuffleb_cshuffle_intrawave_" + "grouped_gemm_bquant_fp8_rcr_preshuffleb_permute_n_intrawave_" "128x128x128_1x4x1_16x16x128_qg1x1x128_preshuffleb_preshufflebq" ) def test_bf8i4_preshuffleb_bquant_name(self): cfg = default_bf8i4_preshuffleb_bquant_config() assert cfg.name == ( - "grouped_gemm_bquant_bf8i4_rcr_preshuffleb_cshuffle_intrawave_" + "grouped_gemm_bquant_bf8i4_rcr_preshuffleb_permute_n_intrawave_" "128x128x128_1x4x1_16x16x32_qg1x1x128_preshuffleb_preshufflebq" ) @@ -479,21 +479,21 @@ class TestPhase4MXConfigs: def test_mx_bf16bf16_name(self): cfg = default_mx_bf16bf16_config(quant_group_k=32) assert cfg.name == ( - "grouped_gemm_bquant_mx_bf16bf16_rcr_microscale_cshuffle_intrawave_" + "grouped_gemm_bquant_mx_bf16bf16_rcr_microscale_permute_n_intrawave_" "128x128x128_1x4x1_16x16x32_qg1x1x32" ) def test_mx_bf16bf8_name(self): cfg = default_mx_bf16bf8_config(quant_group_k=128) assert cfg.name == ( - "grouped_gemm_bquant_mx_bf16bf8_rcr_microscale_cshuffle_intrawave_" + "grouped_gemm_bquant_mx_bf16bf8_rcr_microscale_permute_n_intrawave_" "128x128x128_1x4x1_16x16x64_qg1x1x128" ) def test_mx_bf16fp4_name(self): cfg = default_mx_bf16fp4_config(quant_group_k=32) assert cfg.name == ( - "grouped_gemm_bquant_mx_bf16fp4_rcr_microscale_cshuffle_intrawave_" + "grouped_gemm_bquant_mx_bf16fp4_rcr_microscale_permute_n_intrawave_" "128x128x128_1x4x1_16x16x32_qg1x1x32" ) From 398782dc589d50b63a660623e5435a4a7a2fa8ab Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 22 Jul 2026 13:18:30 -0400 Subject: [PATCH 12/35] style(ck-tile): apply clang-format-18 to quant ctypes dispatcher files Co-Authored-By: Claude --- .../grouped_gemm_abquant_ctypes_lib.cpp | 101 +++++++++++++----- .../ctypes/grouped_gemm_aquant_ctypes_lib.cpp | 12 +-- 2 files changed, 79 insertions(+), 34 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp index bea8625965da..02ac58108884 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -116,7 +116,8 @@ int dispatcher_initialize() /** * Run ABQuantGrouped GEMM: - * C[M,N] = dequant(A[M,K], AQ[ceil(M/aM), ceil(K/aK)]) @ dequant(B[K,N], BQ[ceil(K/bK), ceil(N/bN)]) + * C[M,N] = dequant(A[M,K], AQ[ceil(M/aM), ceil(K/aK)]) @ dequant(B[K,N], BQ[ceil(K/bK), + * ceil(N/bN)]) * * A, B, AQ, BQ, C are host pointers. This function manages device memory internally. * @@ -181,9 +182,9 @@ int dispatcher_run_abquant_gemm(const void* A, (M + static_cast(AQuantGroupSize::kM) - 1) / AQuantGroupSize::kM; if(QK_A != exp_QK_A || QM_A != exp_QM_A) { - std::cerr << "dispatcher_run_abquant_gemm: QK_A/QM_A mismatch. " - << "Got (" << QK_A << ", " << QM_A << "), " - << "expected (" << exp_QK_A << ", " << exp_QM_A << ")\n"; + std::cerr << "dispatcher_run_abquant_gemm: QK_A/QM_A mismatch. " << "Got (" << QK_A + << ", " << QM_A << "), " << "expected (" << exp_QK_A << ", " << exp_QM_A + << ")\n"; return -1; } } @@ -196,9 +197,9 @@ int dispatcher_run_abquant_gemm(const void* A, (N + static_cast(BQuantGroupSize::kN) - 1) / BQuantGroupSize::kN; if(QK_B != exp_QK_B || QN_B != exp_QN_B) { - std::cerr << "dispatcher_run_abquant_gemm: QK_B/QN_B mismatch. " - << "Got (" << QK_B << ", " << QN_B << "), " - << "expected (" << exp_QK_B << ", " << exp_QN_B << ")\n"; + std::cerr << "dispatcher_run_abquant_gemm: QK_B/QN_B mismatch. " << "Got (" << QK_B + << ", " << QN_B << "), " << "expected (" << exp_QK_B << ", " << exp_QN_B + << ")\n"; return -1; } } @@ -209,9 +210,8 @@ int dispatcher_run_abquant_gemm(const void* A, if(stride_A != K || stride_B != K || stride_AQ != QK_A || stride_BQ != QN_B || stride_C != N) { std::cerr << "dispatcher_run_abquant_gemm: non-packed strides not supported. " - << "Expected stride_A=" << K << " stride_B=" << K - << " stride_AQ=" << QK_A << " stride_BQ=" << QN_B << " stride_C=" << N - << "\n"; + << "Expected stride_A=" << K << " stride_B=" << K << " stride_AQ=" << QK_A + << " stride_BQ=" << QN_B << " stride_C=" << N << "\n"; return -1; } @@ -228,31 +228,78 @@ int dispatcher_run_abquant_gemm(const void* A, CDataType* C_dev = nullptr; auto cleanup = [&]() { - if(A_dev) (void)hipFree(A_dev); - if(B_dev) (void)hipFree(B_dev); - if(AQ_dev) (void)hipFree(AQ_dev); - if(BQ_dev) (void)hipFree(BQ_dev); - if(C_dev) (void)hipFree(C_dev); + if(A_dev) + (void)hipFree(A_dev); + if(B_dev) + (void)hipFree(B_dev); + if(AQ_dev) + (void)hipFree(AQ_dev); + if(BQ_dev) + (void)hipFree(BQ_dev); + if(C_dev) + (void)hipFree(C_dev); }; - if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&AQ_dev, QM_A * QK_A * sizeof(QDataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&AQ_dev, QM_A * QK_A * sizeof(QDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup(); + return -1; + } - if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(AQ_dev, AQ_host, QM_A * QK_A * sizeof(QDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) { cleanup(); return -1; } + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(AQ_dev, AQ_host, QM_A * QK_A * sizeof(QDataType), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup(); + return -1; + } // Build QuantGemmHostArgs: both AQ and BQ active ck_tile::QuantGemmHostArgs args; args.a_ptr = A_dev; args.b_ptr = B_dev; - args.aq_ptr = AQ_dev; // A-side scale: active - args.bq_ptr = BQ_dev; // B-side scale: active + args.aq_ptr = AQ_dev; // A-side scale: active + args.bq_ptr = BQ_dev; // B-side scale: active args.c_ptr = C_dev; args.k_batch = k_batch; args.M = static_cast(M); diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp index 886bcaafda5b..0e61de049a03 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp @@ -173,10 +173,9 @@ int dispatcher_run_aquant_gemm(const void* A, (M + static_cast(QuantGroupSize::kM) - 1) / QuantGroupSize::kM; if(QK_A != expected_QK_A || QM_A != expected_QM_A) { - std::cerr << "dispatcher_run_aquant_gemm: QK_A/QM_A mismatch. " - << "Got (" << QK_A << ", " << QM_A << "), " - << "expected (" << expected_QK_A << ", " << expected_QM_A << ") " - << "for K=" << K << " M=" << M + std::cerr << "dispatcher_run_aquant_gemm: QK_A/QM_A mismatch. " << "Got (" << QK_A + << ", " << QM_A << "), " << "expected (" << expected_QK_A << ", " + << expected_QM_A << ") " << "for K=" << K << " M=" << M << " with QuantGroupSize kK=" << QuantGroupSize::kK << " kM=" << QuantGroupSize::kM << "\n"; return -1; @@ -189,9 +188,8 @@ int dispatcher_run_aquant_gemm(const void* A, { std::cerr << "dispatcher_run_aquant_gemm: non-packed strides are not supported. " << "Expected stride_A=" << K << " stride_B=" << K << " stride_AQ=" << QK_A - << " stride_C=" << N << ", got stride_A=" << stride_A - << " stride_B=" << stride_B << " stride_AQ=" << stride_AQ - << " stride_C=" << stride_C << "\n"; + << " stride_C=" << N << ", got stride_A=" << stride_A << " stride_B=" << stride_B + << " stride_AQ=" << stride_AQ << " stride_C=" << stride_C << "\n"; return -1; } From 79d5ecebdc8ff49aa678b72b5ab1a54c3f7082a1 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Mon, 27 Jul 2026 23:22:03 -0400 Subject: [PATCH 13/35] fix(ck-tile): address PR #9253 review comments for aquant/abquant bridges - Replace dispatcher registry launch path with direct SelectedKernel::launch() in both aquant and abquant ctypes libs; QuantGemmHostArgs is incompatible with the registry backend's run() signature - Add elements_to_bytes() helper for correct packed-type buffer sizing (pk_int4_t etc.) in hipMalloc/hipMemcpy/hipMemset calls - Remove hardcoded GFX_ARCH; derive arch at runtime via hipGetDeviceProperties and reject unsupported architectures - Make timing conditional on time_ms != nullptr with warmup (cold_niters=3, nrepeat=10) matching BQuant pattern; previously time_kernel_=false returned zero timing data - Fix B and BQ array layout in Python utils: np.asfortranarray instead of np.ascontiguousarray so leading dim = K (col-major) is correct - Fix ABQuant stride_BQ: use QK_B (ColumnMajor leading dim) not QN_B - Add permute_n column-reordering unscramble to aquant and abquant runners (ported from BQuant) so callers receive logically-ordered output - Fix ABQuant codegen to emit ColumnMajor BQLayout; kernel static-asserts against RowMajor BQ causing build failures Co-Authored-By: Claude --- .../grouped_gemm_abquant_ctypes_lib.cpp | 158 +++++++++--------- .../ctypes/grouped_gemm_aquant_ctypes_lib.cpp | 145 ++++++++-------- .../unified_grouped_gemm_abquant_codegen.py | 8 +- .../python/grouped_gemm_abquant_utils.py | 31 +++- .../python/grouped_gemm_aquant_utils.py | 20 ++- 5 files changed, 197 insertions(+), 165 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp index 02ac58108884..4f3220ab1934 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -15,9 +15,14 @@ * * ABQuant: both A-side and B-side quantization active. * AQ[ceil(M/aM), ceil(K/aK)] - A-side scale tensor (RowMajor) - * BQ[ceil(K/bK), ceil(N/bN)] - B-side scale tensor (RowMajor) + * BQ[ceil(K/bK), ceil(N/bN)] - B-side scale tensor (ColumnMajor, as required by kernel) * Constraint: AQuantGroupSize::kK == BQuantGroupSize::kK * + * Design: direct launch -- SelectedKernel::launch(QuantGemmHostArgs, stream_config) is + * called directly. No dispatcher registry is used: ABQuant kernels take QuantGemmHostArgs, + * which is incompatible with the GeneratedTileKernelInstance::run() signature used by + * the dispatcher's registry backend. + * * Memory model: host-pointer (this library owns hipMalloc/hipMemcpy/hipFree). */ @@ -25,27 +30,26 @@ #include #include #include -#include #include - -#include "ck_tile/dispatcher/dispatcher.hpp" -#include "ck_tile/dispatcher/registry.hpp" -#include "ck_tile/dispatcher/backends/generated_tile_backend.hpp" +#include // Kernel header force-included via -include compiler flag. // Defines: ADataType, BDataType, CDataType, QDataType (shared AQ/BQ type), AccDataType, // AQuantGroupSize, BQuantGroupSize, SelectedKernel, KERNEL_NAME -#ifndef GFX_ARCH -#define GFX_ARCH "gfx950" -#endif +// Compute the byte count for N logical elements of type T. +// For packed types (pk_int4_t, pk_fp4_t) PackedSize=2, so N logical values +// occupy N/2 bytes even though sizeof(T)==1. For all other types PackedSize=1. +template +static constexpr std::size_t elements_to_bytes(std::size_t n) +{ + return n * sizeof(T) / ck_tile::numeric_traits::PackedSize; +} -using namespace ck_tile::dispatcher; -using namespace ck_tile::dispatcher::backends; -using Priority = ck_tile::dispatcher::Registry::Priority; +// GPU architecture is derived from the running device at launch time rather than +// assumed at compile time -- do not hardcode a default architecture here. -static std::shared_ptr g_dispatcher = nullptr; -static bool g_initialized = false; +static bool g_initialized = false; #define HIP_CHECK(call) \ { \ @@ -61,55 +65,19 @@ static bool g_initialized = false; extern "C" { /** - * Initialize dispatcher -- must be called before dispatcher_run_abquant_gemm. - * Returns 0 on success, -1 on error. + * Initialize the ctypes lib. Must be called before dispatcher_run_abquant_gemm. + * + * This library uses a single-kernel-per-.so model: SelectedKernel is + * force-included at compile time and invoked directly via SelectedKernel::launch(). + * No dispatcher registry is involved -- ABQuant kernels require QuantGemmHostArgs + * which is incompatible with the GeneratedTileKernelInstance::run() signature. + * + * Returns 0 on success. */ int dispatcher_initialize() { if(g_initialized) return 0; - - KernelKey key; - key.signature.dtype_a = DataType::FP8; - key.signature.dtype_b = DataType::FP8; - key.signature.dtype_c = DataType::FP16; - key.signature.dtype_acc = DataType::FP32; - key.signature.layout_a = LayoutTag::RowMajor; - key.signature.layout_b = LayoutTag::ColMajor; - key.signature.layout_c = LayoutTag::RowMajor; - key.signature.transpose_a = false; - key.signature.transpose_b = false; - key.signature.grouped = false; - key.signature.split_k = 1; - key.signature.elementwise_op = "PassThrough"; - key.signature.num_d_tensors = 0; - key.signature.structured_sparsity = false; - - key.algorithm.tile_shape = { - SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; - key.algorithm.wave_shape = { - SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; - key.algorithm.warp_tile_shape = { - SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; - key.algorithm.pipeline = Pipeline::CompV3; - key.algorithm.scheduler = Scheduler::Intrawave; - key.algorithm.epilogue = Epilogue::CShuffle; - key.algorithm.block_size = SelectedKernel::BlockSize; - key.algorithm.double_buffer = SelectedKernel::DoubleSmemBuffer; - key.algorithm.persistent = false; - key.algorithm.preshuffle = SelectedKernel::PreshuffleB; - key.algorithm.transpose_c = SelectedKernel::TransposeC; - key.algorithm.num_wave_groups = 1; - key.gfx_arch = GFX_ARCH; - - auto kernel = - create_generated_tile_kernel( - key, KERNEL_NAME); - - Registry::instance().clear(); - Registry::instance().register_kernel(kernel, Priority::High); - - g_dispatcher = std::make_shared(); g_initialized = true; return 0; } @@ -127,7 +95,7 @@ int dispatcher_initialize() * stride_A - leading dim of A (row-major: K) * stride_B - leading dim of B (col-major: K) * stride_AQ - leading dim of AQ (row-major: ceil(K/aK)) - * stride_BQ - leading dim of BQ (row-major: ceil(N/bN)) + * stride_BQ - leading dim of BQ (col-major: ceil(K/bK)) * stride_C - leading dim of C (row-major: N) * QK_A - ceil(K / aquant_group_k) * QM_A - ceil(M / aquant_group_m) (typically == M when aM=1) @@ -174,6 +142,25 @@ int dispatcher_run_abquant_gemm(const void* A, return -1; } + // Derive the GPU architecture from the running device and reject unsupported archs. + { + int dev = 0; + hipDeviceProp_t props{}; + if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) + { + std::cerr << "dispatcher_run_abquant_gemm: could not query device architecture\n"; + return -1; + } + const std::string arch(props.gcnArchName); + if(arch.rfind("gfx950", 0) != 0 && arch.rfind("gfx942", 0) != 0 && + arch.rfind("gfx90a", 0) != 0) + { + std::cerr << "dispatcher_run_abquant_gemm: unsupported GPU architecture '" << arch + << "' (supported: gfx90a, gfx942, gfx950)\n"; + return -1; + } + } + // Validate AQ dimensions { const int64_t exp_QK_A = @@ -206,12 +193,12 @@ int dispatcher_run_abquant_gemm(const void* A, // Only packed (contiguous) layouts are supported. // AQ is RowMajor [QM_A, QK_A]: stride_AQ == QK_A - // BQ is RowMajor [QK_B, QN_B]: stride_BQ == QN_B - if(stride_A != K || stride_B != K || stride_AQ != QK_A || stride_BQ != QN_B || stride_C != N) + // BQ is ColumnMajor [QK_B, QN_B]: stride_BQ == QK_B (leading dim = number of K-groups) + if(stride_A != K || stride_B != K || stride_AQ != QK_A || stride_BQ != QK_B || stride_C != N) { std::cerr << "dispatcher_run_abquant_gemm: non-packed strides not supported. " << "Expected stride_A=" << K << " stride_B=" << K << " stride_AQ=" << QK_A - << " stride_BQ=" << QN_B << " stride_C=" << N << "\n"; + << " stride_BQ=" << QK_B << " stride_C=" << N << "\n"; return -1; } @@ -240,55 +227,59 @@ int dispatcher_run_abquant_gemm(const void* A, (void)hipFree(C_dev); }; - if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + // Allocate device buffers. + // elements_to_bytes(n) handles packed types (pk_int4_t etc.) via PackedSize. + if(hipMalloc(&A_dev, elements_to_bytes(M * K)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + if(hipMalloc(&B_dev, elements_to_bytes(K * N)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&AQ_dev, QM_A * QK_A * sizeof(QDataType)) != hipSuccess) + if(hipMalloc(&AQ_dev, elements_to_bytes(QM_A * QK_A)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&BQ_dev, QK_B * QN_B * sizeof(QDataType)) != hipSuccess) + if(hipMalloc(&BQ_dev, elements_to_bytes(QK_B * QN_B)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + if(hipMalloc(&C_dev, elements_to_bytes(M * N)) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(AQ_dev, AQ_host, QM_A * QK_A * sizeof(QDataType), hipMemcpyHostToDevice) != + if(hipMemcpy(AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(BQ_dev, BQ_host, QK_B * QN_B * sizeof(QDataType), hipMemcpyHostToDevice) != + if(hipMemcpy(BQ_dev, BQ_host, elements_to_bytes(QK_B * QN_B), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + if(hipMemset(C_dev, 0, elements_to_bytes(M * N)) != hipSuccess) { cleanup(); return -1; @@ -313,7 +304,17 @@ int dispatcher_run_abquant_gemm(const void* A, args.stride_AQ = static_cast(stride_AQ); args.stride_BQ = static_cast(stride_BQ); - ck_tile::stream_config stream_cfg{nullptr, false, 0, 0, 1, false, false, 1}; + const bool do_time = (time_ms != nullptr); + ck_tile::stream_config stream_cfg{ + nullptr, // stream_id_ + do_time, // time_kernel_ + 0, // log_level_ + do_time ? 3 : 0, // cold_niters_ + do_time ? 10 : 1, // nrepeat_ + do_time, // is_gpu_timer_ + false, // flush_cache_ + 1, // rotating_count_ + }; float exec_time = SelectedKernel::launch(args, stream_cfg); @@ -324,7 +325,8 @@ int dispatcher_run_abquant_gemm(const void* A, return -2; } - if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + if(hipMemcpy(C_host, C_dev, elements_to_bytes(M * N), hipMemcpyDeviceToHost) != + hipSuccess) { cleanup(); return -1; @@ -350,15 +352,11 @@ int dispatcher_init() { return dispatcher_initialize(); } /** * Number of kernels registered (always 1 for single-kernel-per-.so model). */ -int dispatcher_get_kernel_count() { return static_cast(Registry::instance().size()); } +int dispatcher_get_kernel_count() { return 1; } /** * Release dispatcher resources. */ -void dispatcher_cleanup() -{ - g_dispatcher.reset(); - g_initialized = false; -} +void dispatcher_cleanup() { g_initialized = false; } } // extern "C" diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp index 0e61de049a03..ada7cbb8290e 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp @@ -16,6 +16,11 @@ * AQ[ceil(M/gM), ceil(K/gK)] is the A-side scale tensor (RowMajor). * BQ is unused (bq_ptr=nullptr, QK_B=0, stride_BQ=0 in QuantGemmHostArgs). * + * Design: direct launch -- SelectedKernel::launch(QuantGemmHostArgs, stream_config) is + * called directly. No dispatcher registry is used: AQuant kernels take QuantGemmHostArgs, + * which is incompatible with the GeneratedTileKernelInstance::run() signature used by + * the dispatcher's registry backend. + * * Memory model: host-pointer (this library owns hipMalloc/hipMemcpy/hipFree). */ @@ -23,27 +28,26 @@ #include #include #include -#include #include - -#include "ck_tile/dispatcher/dispatcher.hpp" -#include "ck_tile/dispatcher/registry.hpp" -#include "ck_tile/dispatcher/backends/generated_tile_backend.hpp" +#include // Kernel header force-included via -include compiler flag. // Defines: ADataType, BDataType, CDataType, QDataType (AQDataType), AccDataType, // QuantGroupSize, SelectedKernel, KERNEL_NAME -#ifndef GFX_ARCH -#define GFX_ARCH "gfx950" -#endif +// Compute the byte count for N logical elements of type T. +// For packed types (pk_int4_t, pk_fp4_t) PackedSize=2, so N logical values +// occupy N/2 bytes even though sizeof(T)==1. For all other types PackedSize=1. +template +static constexpr std::size_t elements_to_bytes(std::size_t n) +{ + return n * sizeof(T) / ck_tile::numeric_traits::PackedSize; +} -using namespace ck_tile::dispatcher; -using namespace ck_tile::dispatcher::backends; -using Priority = ck_tile::dispatcher::Registry::Priority; +// GPU architecture is derived from the running device at launch time rather than +// assumed at compile time -- do not hardcode a default architecture here. -static std::shared_ptr g_dispatcher = nullptr; -static bool g_initialized = false; +static bool g_initialized = false; #define HIP_CHECK(call) \ { \ @@ -59,57 +63,19 @@ static bool g_initialized = false; extern "C" { /** - * Initialize dispatcher -- must be called before dispatcher_run_aquant_gemm. + * Initialize the ctypes lib. Must be called before dispatcher_run_aquant_gemm. * - * Registers SelectedKernel (from the force-included header) into the Registry. - * Returns 0 on success, -1 on error. + * This library uses a single-kernel-per-.so model: SelectedKernel is + * force-included at compile time and invoked directly via SelectedKernel::launch(). + * No dispatcher registry is involved -- AQuant kernels require QuantGemmHostArgs + * which is incompatible with the GeneratedTileKernelInstance::run() signature. + * + * Returns 0 on success. */ int dispatcher_initialize() { if(g_initialized) return 0; - - KernelKey key; - key.signature.dtype_a = DataType::FP8; - key.signature.dtype_b = DataType::FP8; - key.signature.dtype_c = DataType::FP16; - key.signature.dtype_acc = DataType::FP32; - key.signature.layout_a = LayoutTag::RowMajor; - key.signature.layout_b = LayoutTag::ColMajor; - key.signature.layout_c = LayoutTag::RowMajor; - key.signature.transpose_a = false; - key.signature.transpose_b = false; - key.signature.grouped = false; - key.signature.split_k = 1; - key.signature.elementwise_op = "PassThrough"; - key.signature.num_d_tensors = 0; - key.signature.structured_sparsity = false; - - key.algorithm.tile_shape = { - SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; - key.algorithm.wave_shape = { - SelectedKernel::WarpM, SelectedKernel::WarpN, SelectedKernel::WarpK}; - key.algorithm.warp_tile_shape = { - SelectedKernel::WarpTileM, SelectedKernel::WarpTileN, SelectedKernel::WarpTileK}; - key.algorithm.pipeline = Pipeline::CompV3; - key.algorithm.scheduler = Scheduler::Intrawave; - key.algorithm.epilogue = Epilogue::CShuffle; - key.algorithm.block_size = SelectedKernel::BlockSize; - key.algorithm.double_buffer = false; - key.algorithm.persistent = false; - key.algorithm.preshuffle = SelectedKernel::APreshuffleQuant; - key.algorithm.transpose_c = SelectedKernel::TransposeC; - key.algorithm.num_wave_groups = 1; - key.gfx_arch = GFX_ARCH; - - auto kernel = - create_generated_tile_kernel( - key, KERNEL_NAME); - - Registry::instance().clear(); - Registry::instance().register_kernel(kernel, Priority::High); - - g_dispatcher = std::make_shared(); g_initialized = true; return 0; } @@ -165,6 +131,25 @@ int dispatcher_run_aquant_gemm(const void* A, return -1; } + // Derive the GPU architecture from the running device and reject unsupported archs. + { + int dev = 0; + hipDeviceProp_t props{}; + if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) + { + std::cerr << "dispatcher_run_aquant_gemm: could not query device architecture\n"; + return -1; + } + const std::string arch(props.gcnArchName); + if(arch.rfind("gfx950", 0) != 0 && arch.rfind("gfx942", 0) != 0 && + arch.rfind("gfx90a", 0) != 0) + { + std::cerr << "dispatcher_run_aquant_gemm: unsupported GPU architecture '" << arch + << "' (supported: gfx90a, gfx942, gfx950)\n"; + return -1; + } + } + // Validate that the caller's QK_A/QM_A match the compile-time quant group sizes. { const int64_t expected_QK_A = @@ -214,46 +199,49 @@ int dispatcher_run_aquant_gemm(const void* A, (void)hipFree(C_dev); }; - // Allocate device buffers - if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + // Allocate device buffers. + // B may be a packed type (pk_int4_t): elements_to_bytes(n) handles PackedSize correctly. + if(hipMalloc(&A_dev, elements_to_bytes(M * K)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + if(hipMalloc(&B_dev, elements_to_bytes(K * N)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&AQ_dev, QM_A * QK_A * sizeof(QDataType)) != hipSuccess) + if(hipMalloc(&AQ_dev, elements_to_bytes(QM_A * QK_A)) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + if(hipMalloc(&C_dev, elements_to_bytes(M * N)) != hipSuccess) { cleanup(); return -1; } // Copy inputs to device - if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(AQ_dev, AQ_host, QM_A * QK_A * sizeof(QDataType), hipMemcpyHostToDevice) != + if(hipMemcpy(AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + if(hipMemset(C_dev, 0, elements_to_bytes(M * N)) != hipSuccess) { cleanup(); return -1; @@ -278,7 +266,17 @@ int dispatcher_run_aquant_gemm(const void* A, args.stride_AQ = static_cast(stride_AQ); args.stride_BQ = 0; // unused - ck_tile::stream_config stream_cfg{nullptr, false, 0, 0, 1, false, false, 1}; + const bool do_time = (time_ms != nullptr); + ck_tile::stream_config stream_cfg{ + nullptr, // stream_id_ + do_time, // time_kernel_ + 0, // log_level_ + do_time ? 3 : 0, // cold_niters_ + do_time ? 10 : 1, // nrepeat_ + do_time, // is_gpu_timer_ + false, // flush_cache_ + 1, // rotating_count_ + }; float exec_time = SelectedKernel::launch(args, stream_cfg); @@ -290,7 +288,8 @@ int dispatcher_run_aquant_gemm(const void* A, } // Copy result back - if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + if(hipMemcpy(C_host, C_dev, elements_to_bytes(M * N), hipMemcpyDeviceToHost) != + hipSuccess) { cleanup(); return -1; @@ -316,15 +315,11 @@ int dispatcher_init() { return dispatcher_initialize(); } /** * Number of kernels registered (always 1 for the single-kernel-per-.so model). */ -int dispatcher_get_kernel_count() { return static_cast(Registry::instance().size()); } +int dispatcher_get_kernel_count() { return 1; } /** * Release dispatcher resources. */ -void dispatcher_cleanup() -{ - g_dispatcher.reset(); - g_initialized = false; -} +void dispatcher_cleanup() { g_initialized = false; } } // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py index 72ef77acc08c..7c555f530114 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py @@ -17,6 +17,8 @@ Constraint: AQuantGroupSize::kK == BQuantGroupSize::kK (enforced at codegen time). +BQ layout: ColumnMajor [ceil(K/bK), ceil(N/bN)] — the kernel asserts ColumnMajor for BQ. + Pipeline selection (mirrors run_gemm_quant_example.inc): compv3: ABQuantGemmPipelineAgBgCrCompV3 (GemmConfigABQuantPrefill, non-gfx950) eightwaves: ABQuantGemmPipelineAgBgCrEightWaves (GemmConfigEightWaves, gfx950) @@ -221,8 +223,8 @@ def generate(self, spec: ABQuantKernelSpec) -> str: layout_c_ck = ABQUANT_LAYOUT_TO_CK[spec.layout[2]] # AQ is RowMajor: [ceil(M/gM), ceil(K/gK)] layout_aq_ck = ABQUANT_LAYOUT_TO_CK["r"] - # BQ is RowMajor: [ceil(K/gK), ceil(N/gN)] - layout_bq_ck = ABQUANT_LAYOUT_TO_CK["r"] + # BQ is ColumnMajor: [ceil(K/gK), ceil(N/gN)] — kernel asserts against RowMajor BQ. + layout_bq_ck = ABQUANT_LAYOUT_TO_CK["c"] pipeline_ck = ABQUANT_PIPELINE_MAP[spec.pipeline] base_pipeline_ck = ABQUANT_BASE_PIPELINE_MAP[spec.pipeline] @@ -305,7 +307,7 @@ def generate(self, spec: ABQuantKernelSpec) -> str: using BLayout = {layout_b_ck}; using CLayout = {layout_c_ck}; using AQLayout = {layout_aq_ck}; // RowMajor: [ceil(M/aM), ceil(K/aK)] -using BQLayout = {layout_bq_ck}; // RowMajor: [ceil(K/bK), ceil(N/bN)] +using BQLayout = {layout_bq_ck}; // ColumnMajor: [ceil(K/bK), ceil(N/bN)] // Separate group sizes — ABQuant requires AQuantGroupSize::kK == BQuantGroupSize::kK using AQuantGroupSize = ck_tile::QuantGroupShape ABQuan A shape: (M, K) dtype: fp8/bf8 B shape: (K, N) col-major dtype: fp8/bf8 - AQ shape: (QM_A, QK_A) dtype: float32 (A-side scale) - BQ shape: (QK_B, QN_B) dtype: float32 (B-side scale) + AQ shape: (QM_A, QK_A) dtype: float32 (A-side scale, RowMajor) + BQ shape: (QK_B, QN_B) dtype: float32 (B-side scale, ColumnMajor) c_dtype: numpy dtype for C output buffer. Defaults to np.float16. + Pass np.bfloat16 for MX variants whose CDataType is bf16. """ import numpy as np @@ -400,11 +403,11 @@ def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuan # Stride layout: # AQ RowMajor [QM_A, QK_A]: stride_AQ = QK_A - # BQ RowMajor [QK_B, QN_B]: stride_BQ = QN_B + # BQ ColumnMajor [QK_B, QN_B]: stride_BQ = QK_B (leading dim = K-groups) stride_A = K stride_B = K # col-major B: leading dim = K stride_AQ = QK_A - stride_BQ = QN_B + stride_BQ = QK_B stride_C = N rc, time_ms = self._lib.run( @@ -422,6 +425,22 @@ def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuan f"for kernel {self.kernel_name}" ) + # permute_n epilogue writes C with N-columns riffled into r groups + # (r = tile_n / warp_tile_n / warp_n). Undo it so the caller gets logical C. + _name = self.kernel_name + if 'permute_n' in _name: + import re as _re + _m = _re.search(r'_(\d+)x(\d+)x(\d+)_(\d+)x(\d+)x(\d+)_(\d+)x(\d+)x(\d+)_', _name) + if _m: + _tile_n = int(_m.group(2)); _warp_n = int(_m.group(5)); _wt_n = int(_m.group(8)) + _r = _tile_n // _wt_n // _warp_n + if _r > 1 and (N % _r) == 0: + _half = N // _r + _logical = [(c % _r) * _half + (c // _r) for c in range(N)] + _Cp = np.empty_like(C) + _Cp[:, _logical] = C + C = _Cp + return ABQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py index 3dee600f37df..7d865a4d6561 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -287,7 +287,8 @@ def run( import numpy as np A = np.ascontiguousarray(A) - B = np.ascontiguousarray(B) + # B is col-major [K, N]: Fortran order makes the leading dim = K (stride_B = K). + B = np.asfortranarray(B) AQ = np.ascontiguousarray(AQ) C = np.ascontiguousarray(C) @@ -356,6 +357,7 @@ def run(self, A, B, AQ, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmR B shape: (K, N) col-major dtype: fp8/bf8 or pk_int4 AQ shape: (QM_A, QK_A) dtype: float (for fp8/bf8) or fp8/bf8 (for fp8i4/bf8i4) c_dtype numpy dtype for the output C buffer. Defaults to np.float16. + Pass np.bfloat16 for MX variants whose CDataType is bf16. Returns AQuantGemmResult with C shape (M, N). """ import numpy as np @@ -393,6 +395,22 @@ def run(self, A, B, AQ, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmR f"for kernel {self.kernel_name}" ) + # permute_n epilogue writes C with N-columns riffled into r groups + # (r = tile_n / warp_tile_n / warp_n). Undo it so the caller gets logical C. + _name = self.kernel_name + if 'permute_n' in _name: + import re as _re + _m = _re.search(r'_(\d+)x(\d+)x(\d+)_(\d+)x(\d+)x(\d+)_(\d+)x(\d+)x(\d+)_', _name) + if _m: + _tile_n = int(_m.group(2)); _warp_n = int(_m.group(5)); _wt_n = int(_m.group(8)) + _r = _tile_n // _wt_n // _warp_n + if _r > 1 and (N % _r) == 0: + _half = N // _r + _logical = [(c % _r) * _half + (c // _r) for c in range(N)] + _Cp = np.empty_like(C) + _Cp[:, _logical] = C + C = _Cp + return AQuantGemmResult(C=C, time_ms=time_ms, kernel_name=self.kernel_name) From 84732e27211de5f6a4117d887f6712ee25a608b2 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Mon, 27 Jul 2026 23:29:44 -0400 Subject: [PATCH 14/35] fix(ck-tile): remove dead dispatcher_include from aquant/abquant compile helpers The path _CTYPES_LIB_SRC.parent.parent.parent / "dispatcher" / "include" resolved to dispatcher/dispatcher/include which never exists; the is_dir() guard silently no-opped it rather than failing loudly. Since both aquant and abquant ctypes libs now call SelectedKernel::launch() directly (no dispatcher registry headers), the include is not needed at all. Matches the intentional exclusion already present in the bquant compile helper. Co-Authored-By: Claude --- .../dispatcher/python/grouped_gemm_abquant_utils.py | 5 ++--- .../dispatcher/python/grouped_gemm_aquant_utils.py | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index c927ce4fcd81..fc9ac8c99cfe 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -518,9 +518,8 @@ def _compile_abquant_kernel( if ck_include: cmd += [f"-I{ck_include}"] - dispatcher_include = _CTYPES_LIB_SRC.parent.parent.parent / "dispatcher" / "include" - if dispatcher_include.is_dir(): - cmd += [f"-I{dispatcher_include}"] + # NOTE: dispatcher/include is intentionally excluded here — the abquant ctypes lib + # calls SelectedKernel::launch() directly and does not use any dispatcher headers. if extra_include_dirs: for d in extra_include_dirs: diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py index 7d865a4d6561..1e9ae64133bf 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -497,9 +497,8 @@ def _compile_aquant_kernel( if ck_include: cmd += [f"-I{ck_include}"] - dispatcher_include = _CTYPES_LIB_SRC.parent.parent.parent / "dispatcher" / "include" - if dispatcher_include.is_dir(): - cmd += [f"-I{dispatcher_include}"] + # NOTE: dispatcher/include is intentionally excluded here — the aquant ctypes lib + # calls SelectedKernel::launch() directly and does not use any dispatcher headers. if extra_include_dirs: for d in extra_include_dirs: From 1d389a28b13f860f7612a673baa1c8790025c827 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Mon, 27 Jul 2026 23:31:49 -0400 Subject: [PATCH 15/35] fix(ck-tile): drop unused QK_A/QM_A/QK_B/QN_B pre-calculations in abquant example The four precomputed variables were shadowed immediately by ABQuantGemmProblem whose .QK_A/.QM_A/.QK_B/.QN_B properties are used everywhere. QM_A=1 was also wrong: with aquant_group_m=1 the correct value is ceil(M/1)=M. Create the problem first and derive quant-group counts from it exclusively. Also removes the now-unused math import. Co-Authored-By: Claude --- .../examples/gemm/python/15_grouped_gemm_abquant.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py index 4f75a5cfda43..d229c3dec45a 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -30,7 +30,6 @@ import argparse import logging -import math import sys import tempfile from pathlib import Path @@ -206,12 +205,6 @@ def main(): # 3. Generate inputs # ------------------------------------------------------------------------- rng = np.random.default_rng(42) - QK_A = math.ceil(K / gK) - QM_A = 1 # aquant_group_m=1 by default → QM_A == M - QK_B = math.ceil(K / gK) - QN_B = math.ceil(N / bN) - - # Recompute using actual problem to be safe problem = ABQuantGemmProblem( M=M, N=N, K=K, aquant_group_m=1, aquant_group_n=1, aquant_group_k=gK, From fd4cdf06f0b6f79a76b11a4d65bedef4f1f7454d Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 28 Jul 2026 00:05:49 -0400 Subject: [PATCH 16/35] fix(ck-tile): fix clang-format violations in aquant/abquant ctypes libs Co-Authored-By: Claude --- .../bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp | 6 ++++-- .../bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp | 3 ++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp index 4f3220ab1934..d9619fbd357f 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -267,13 +267,15 @@ int dispatcher_run_abquant_gemm(const void* A, cleanup(); return -1; } - if(hipMemcpy(AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != + if(hipMemcpy( + AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(BQ_dev, BQ_host, elements_to_bytes(QK_B * QN_B), hipMemcpyHostToDevice) != + if(hipMemcpy( + BQ_dev, BQ_host, elements_to_bytes(QK_B * QN_B), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp index ada7cbb8290e..8bb45ae83efd 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp @@ -235,7 +235,8 @@ int dispatcher_run_aquant_gemm(const void* A, cleanup(); return -1; } - if(hipMemcpy(AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != + if(hipMemcpy( + AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); From 6be7241a11e370b6b189f5d10345d58fd6d42742 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 28 Jul 2026 17:49:44 -0400 Subject: [PATCH 17/35] fix(ck-tile): add dispatcher/include to aquant kernel compile command The ctypes lib may include ck_tile/dispatcher/dispatcher.hpp on some checkouts; passing -I ensures the compile succeeds regardless of the cpp version present on the build machine. Co-Authored-By: Claude --- .../dispatcher/python/grouped_gemm_aquant_utils.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py index 1e9ae64133bf..e0e03fe76fb6 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -445,6 +445,15 @@ def _get_ck_include_dir() -> Optional[Path]: return None +def _get_dispatcher_include_dir() -> Optional[Path]: + """Attempt to locate the dispatcher include directory relative to this file.""" + here = Path(__file__).resolve().parent + candidate = here.parent / "include" + if (candidate / "ck_tile" / "dispatcher").is_dir(): + return candidate + return None + + def _generate_aquant_kernel( config: AQuantKernelConfig, output_dir: Path, @@ -497,8 +506,9 @@ def _compile_aquant_kernel( if ck_include: cmd += [f"-I{ck_include}"] - # NOTE: dispatcher/include is intentionally excluded here — the aquant ctypes lib - # calls SelectedKernel::launch() directly and does not use any dispatcher headers. + dispatcher_include = _get_dispatcher_include_dir() + if dispatcher_include: + cmd += [f"-I{dispatcher_include}"] if extra_include_dirs: for d in extra_include_dirs: From d02213e318f4d209331b187eb4e810c26a39a09c Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Tue, 28 Jul 2026 19:13:45 -0400 Subject: [PATCH 18/35] fix(ck-tile): use warp_tile_k=32 in default aquant configs for gfx9 correctness 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 --- .../gemm/python/14_grouped_gemm_aquant.py | 5 +++ .../python/grouped_gemm_aquant_utils.py | 32 ++++++++++++++----- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py index 839cbd6a7e2f..b3f954dddb22 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py @@ -230,6 +230,11 @@ def main(): log.info("PASSED (max_rel=%.4f, tol=%.4f)", max_rel, tolerance) else: log.error("FAILED (max_rel=%.4f > tol=%.4f)", max_rel, tolerance) + C_gpu_f32 = C_gpu.astype(np.float32) + C_ref_f32 = C_ref.astype(np.float32) + log.error("C_gpu[:2,:4] = %s", C_gpu_f32[:2, :4]) + log.error("C_ref[:2,:4] = %s", C_ref_f32[:2, :4]) + log.error("C_gpu nonzero count: %d / %d", int(np.count_nonzero(C_gpu_f32)), C_gpu_f32.size) return 1 else: log.info("Verification skipped (--no-verify)") diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py index e0e03fe76fb6..371eb2b3dfbc 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -631,7 +631,12 @@ def default_fp8_config( quant_group_m: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> AQuantKernelConfig: - """fp8 AQuant decode config (GemmConfigQuantDecodeInterwave, Mem pipeline).""" + """fp8 AQuant decode config (GemmConfigQuantDecodeInterwave, Mem pipeline). + + warp_tile_k=32 maps to the MFMA mfma_f32_16x16x32_fp8_fp8 instruction which is + valid on gfx9 (gfx942, gfx950). warp_tile_k=16 would select the gfx12-only WMMA + instruction and silently returns zeros on gfx9 hardware. + """ return AQuantKernelConfig( variant_key="fp8", layout="rcr", @@ -640,7 +645,7 @@ def default_fp8_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, quant_group_m=quant_group_m, quant_group_n=1, quant_group_k=quant_group_k, @@ -654,7 +659,12 @@ def default_bf8_config( quant_group_m: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> AQuantKernelConfig: - """bf8 AQuant decode config (GemmConfigQuantDecodeInterwave, Mem pipeline).""" + """bf8 AQuant decode config (GemmConfigQuantDecodeInterwave, Mem pipeline). + + warp_tile_k=32 maps to the MFMA mfma_f32_16x16x32_bf8_bf8 instruction which is + valid on gfx9 (gfx942, gfx950). warp_tile_k=16 would select the gfx12-only WMMA + instruction and silently returns zeros on gfx9 hardware. + """ return AQuantKernelConfig( variant_key="bf8", layout="rcr", @@ -663,7 +673,7 @@ def default_bf8_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, quant_group_m=quant_group_m, quant_group_n=1, quant_group_k=quant_group_k, @@ -677,7 +687,10 @@ def default_fp8i4_config( quant_group_m: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> AQuantKernelConfig: - """fp8i4 AQuant decode config (A=fp8, B=pk_int4, AQ=fp8; Mem pipeline).""" + """fp8i4 AQuant decode config (A=fp8, B=pk_int4, AQ=fp8; Mem pipeline). + + warp_tile_k=32 maps to valid MFMA instructions on gfx9 (gfx942, gfx950). + """ return AQuantKernelConfig( variant_key="fp8i4", layout="rcr", @@ -686,7 +699,7 @@ def default_fp8i4_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, quant_group_m=quant_group_m, quant_group_n=1, quant_group_k=quant_group_k, @@ -700,7 +713,10 @@ def default_bf8i4_config( quant_group_m: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> AQuantKernelConfig: - """bf8i4 AQuant decode config (A=bf8, B=pk_int4, AQ=bf8; Mem pipeline).""" + """bf8i4 AQuant decode config (A=bf8, B=pk_int4, AQ=bf8; Mem pipeline). + + warp_tile_k=32 maps to valid MFMA instructions on gfx9 (gfx942, gfx950). + """ return AQuantKernelConfig( variant_key="bf8i4", layout="rcr", @@ -709,7 +725,7 @@ def default_bf8i4_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, quant_group_m=quant_group_m, quant_group_n=1, quant_group_k=quant_group_k, From 369f0df86ad390ba505250cd5ba6cf93224ad0e3 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 29 Jul 2026 00:52:07 -0400 Subject: [PATCH 19/35] fix(ck-tile): replace broken fp8 fallback encoder with software fp8 converter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../gemm/python/14_grouped_gemm_aquant.py | 73 ++++++++++++++++++- 1 file changed, 69 insertions(+), 4 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py index b3f954dddb22..2bea190ecee7 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py @@ -63,15 +63,14 @@ def _float32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: """Encode float32 values as fp8 bytes (uint8 view of the fp8 bit pattern). dtype: "fp8" -> float8_e4m3fn, "bf8" -> float8_e5m2. - Falls back to clamping when ml_dtypes is not installed. + Falls back to software conversion when ml_dtypes is not installed. """ try: import ml_dtypes ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 return arr.astype(ml_t).view(np.uint8) except ImportError: - clamped = np.clip(arr, -2.0, 2.0) - return (clamped * 64).astype(np.int8).view(np.uint8) + return _soft_f32_to_fp8(arr, dtype) def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: @@ -81,7 +80,73 @@ def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 return arr.view(ml_t).astype(np.float32) except ImportError: - return arr.view(np.int8).astype(np.float32) / 64.0 + return _soft_fp8_to_f32(arr, dtype) + + +def _soft_f32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: + """Software float32 → fp8 conversion (no ml_dtypes dependency).""" + flat = arr.astype(np.float32).ravel() + if dtype == "fp8": + ebits, mbits, max_val = 4, 3, 448.0 # e4m3fn: bias=7, no inf, NaN=0x7F + else: + ebits, mbits, max_val = 5, 2, 57344.0 # e5m2: bias=15 + bias = (1 << (ebits - 1)) - 1 + out = np.zeros(flat.shape, dtype=np.uint8) + for i, v in enumerate(flat): + sign = 0 + if v < 0: + sign = 1 + v = -v + v = min(v, max_val) + if v == 0: + out[i] = sign << 7 + continue + exp = int(math.floor(math.log2(v))) if v > 0 else 0 + biased_exp = exp + bias + if biased_exp <= 0: + # subnormal + frac = v / (2.0 ** (1 - bias)) + mant = int(round(frac * (1 << mbits))) + mant = min(mant, (1 << mbits) - 1) + out[i] = (sign << 7) | mant + else: + max_exp = (1 << ebits) - 2 if dtype == "fp8" else (1 << ebits) - 1 + biased_exp = min(biased_exp, max_exp) + frac = v / (2.0 ** (biased_exp - bias)) - 1.0 + mant = int(round(frac * (1 << mbits))) + if mant >= (1 << mbits): + mant = 0 + biased_exp += 1 + biased_exp = min(biased_exp, max_exp) + out[i] = (sign << 7) | (biased_exp << mbits) | mant + return out.reshape(arr.shape) + + +def _soft_fp8_to_f32(arr: np.ndarray, dtype: str) -> np.ndarray: + """Software fp8 → float32 conversion (no ml_dtypes dependency).""" + flat = arr.ravel().astype(np.uint8) + if dtype == "fp8": + ebits, mbits = 4, 3 # e4m3fn: bias=7 + else: + ebits, mbits = 5, 2 # e5m2: bias=15 + bias = (1 << (ebits - 1)) - 1 + max_exp_bits = (1 << ebits) - 1 + out = np.zeros(flat.shape, dtype=np.float32) + for i, byte in enumerate(flat): + sign = (byte >> 7) & 1 + exp_bits = (byte >> mbits) & ((1 << ebits) - 1) + mant_bits = byte & ((1 << mbits) - 1) + if dtype == "fp8" and exp_bits == max_exp_bits and mant_bits == (1 << mbits) - 1: + out[i] = float('nan') + elif dtype == "bf8" and exp_bits == max_exp_bits: + out[i] = float('nan') if mant_bits != 0 else float('inf') * (-1 if sign else 1) + elif exp_bits == 0: + val = (mant_bits / (1 << mbits)) * (2.0 ** (1 - bias)) + out[i] = -val if sign else val + else: + val = (1.0 + mant_bits / (1 << mbits)) * (2.0 ** (exp_bits - bias)) + out[i] = -val if sign else val + return out.reshape(arr.shape) # ============================================================================= From 84d171c47b0fc1212700fac69eb1b68d85495c92 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Thu, 30 Jul 2026 17:20:01 -0400 Subject: [PATCH 20/35] fix(ck-tile): use warp_tile_k=32 in all default abquant configs for gfx9 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 d02213e318f. Co-Authored-By: Claude --- .../python/grouped_gemm_abquant_utils.py | 50 ++++++++++++++----- 1 file changed, 38 insertions(+), 12 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index fc9ac8c99cfe..98a7f400506c 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -632,6 +632,10 @@ def default_fp8_compv3_config( """fp8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False). Tile: 128x128x128, warp 1x4x1. kPadK=false (ABQuant prefill tiles). + + warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8, valid on gfx942 and gfx950. + warp_tile_k=16 would select the gfx12-only WMMA instruction which silently returns + zeros on gfx9 hardware. """ return ABQuantKernelConfig( variant_key="fp8", @@ -641,7 +645,7 @@ def default_fp8_compv3_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -655,7 +659,12 @@ def default_bf8_compv3_config( bquant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> ABQuantKernelConfig: - """bf8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False).""" + """bf8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False). + + warp_tile_k=32 selects mfma_f32_16x16x32_bf8_bf8, valid on gfx942 and gfx950. + warp_tile_k=16 would select the gfx12-only WMMA instruction which silently returns + zeros on gfx9 hardware. + """ return ABQuantKernelConfig( variant_key="bf8", layout="rcr", @@ -664,7 +673,7 @@ def default_bf8_compv3_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -681,9 +690,12 @@ def default_fp8_eightwaves_config( """fp8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True). Tile: 192x256x128, warp 4x2x1 — 8-wave configuration for MI350X. - warp_tile_k=16: get_k_warp_tile on gfx950 = 128 for 8-bit floats, - but GemmConfigEightWaves inherits GemmConfigABQuantPrefill whose warp_tile_k uses - get_k_warp_tile without IsFlatMM — use 16 as concrete value for now. + warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8, valid on gfx942 and gfx950. + GemmConfigEightWaves inherits GemmConfigABQuantPrefill which uses + get_k_warp_tile: gfx950 returns 128, gfx942 returns 32. + We use 32 (= gfx942 value) as the cross-arch safe choice — it is 4 MFMA iterations + on gfx942 and also valid on gfx950 (128 % 32 == 0). warp_tile_k=16 would select + the gfx12-only WMMA instruction which silently returns zeros on gfx9 hardware. """ return ABQuantKernelConfig( variant_key="fp8", @@ -693,7 +705,7 @@ def default_fp8_eightwaves_config( scheduler="intrawave", tile_m=192, tile_n=256, tile_k=128, warp_m=4, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -708,7 +720,11 @@ def default_bf8_eightwaves_config( bquant_group_n: int = 128, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> ABQuantKernelConfig: - """bf8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True).""" + """bf8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True). + + warp_tile_k=32 selects mfma_f32_16x16x32_bf8_bf8, valid on gfx942 and gfx950. + See default_fp8_eightwaves_config for the full rationale. + """ return ABQuantKernelConfig( variant_key="bf8", layout="rcr", @@ -717,7 +733,7 @@ def default_bf8_eightwaves_config( scheduler="intrawave", tile_m=192, tile_n=256, tile_k=128, warp_m=4, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -735,6 +751,12 @@ def default_fp8_preshuffleb_config( """fp8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill). Tile: 128x128x128, warp 2x2x1 (vs 1x4x1 for compv3). DoubleSmemBuffer=True. + + warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8, valid on gfx942 and gfx950. + GemmConfigPreshuffleB_ABQuant_Prefill uses get_k_warp_tile: + gfx950 returns 128, gfx942 returns 64. We use 32 as the cross-arch safe choice + (128 % 32 == 0 and 64 % 32 == 0). warp_tile_k=16 would select the gfx12-only WMMA + instruction which silently returns zeros on gfx9 hardware. """ return ABQuantKernelConfig( variant_key="fp8", @@ -744,7 +766,7 @@ def default_fp8_preshuffleb_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=2, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=True, preshuffle_aq=False, preshuffle_bq=False, @@ -760,7 +782,11 @@ def default_bf8_preshuffleb_config( bquant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> ABQuantKernelConfig: - """bf8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill).""" + """bf8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill). + + warp_tile_k=32 selects mfma_f32_16x16x32_bf8_bf8, valid on gfx942 and gfx950. + See default_fp8_preshuffleb_config for the full rationale. + """ return ABQuantKernelConfig( variant_key="bf8", layout="rcr", @@ -769,7 +795,7 @@ def default_bf8_preshuffleb_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=2, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=16, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=True, preshuffle_aq=False, preshuffle_bq=False, From 04e75f68559383e5de716255fb36a7647c3d2f7b Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Thu, 30 Jul 2026 17:37:56 -0400 Subject: [PATCH 21/35] fix(ck-tile): replace broken fp8 fallback encoder in abquant example 15 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 369f0df86ad. Co-Authored-By: Claude --- .../gemm/python/15_grouped_gemm_abquant.py | 82 ++++++++++++++++++- 1 file changed, 79 insertions(+), 3 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py index d229c3dec45a..c6fc18a0e643 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -30,6 +30,7 @@ import argparse import logging +import math import sys import tempfile from pathlib import Path @@ -59,13 +60,17 @@ def _float32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: + """Encode float32 values as fp8 bytes (uint8 view of the fp8 bit pattern). + + dtype: "fp8" -> float8_e4m3fn, "bf8" -> float8_e5m2. + Falls back to software conversion when ml_dtypes is not installed. + """ try: import ml_dtypes ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 return arr.astype(ml_t).view(np.uint8) except ImportError: - clamped = np.clip(arr, -2.0, 2.0) - return (clamped * 64).astype(np.int8).view(np.uint8) + return _soft_f32_to_fp8(arr, dtype) def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: @@ -74,7 +79,73 @@ def _fp8_to_float32(arr: np.ndarray, dtype: str) -> np.ndarray: ml_t = ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 return arr.view(ml_t).astype(np.float32) except ImportError: - return arr.view(np.int8).astype(np.float32) / 64.0 + return _soft_fp8_to_f32(arr, dtype) + + +def _soft_f32_to_fp8(arr: np.ndarray, dtype: str) -> np.ndarray: + """Software float32 → fp8 conversion (no ml_dtypes dependency).""" + flat = arr.astype(np.float32).ravel() + if dtype == "fp8": + ebits, mbits, max_val = 4, 3, 448.0 # e4m3fn: bias=7, no inf, NaN=0x7F + else: + ebits, mbits, max_val = 5, 2, 57344.0 # e5m2: bias=15 + bias = (1 << (ebits - 1)) - 1 + out = np.zeros(flat.shape, dtype=np.uint8) + for i, v in enumerate(flat): + sign = 0 + if v < 0: + sign = 1 + v = -v + v = min(v, max_val) + if v == 0: + out[i] = sign << 7 + continue + exp = int(math.floor(math.log2(v))) if v > 0 else 0 + biased_exp = exp + bias + if biased_exp <= 0: + # subnormal + frac = v / (2.0 ** (1 - bias)) + mant = int(round(frac * (1 << mbits))) + mant = min(mant, (1 << mbits) - 1) + out[i] = (sign << 7) | mant + else: + max_exp = (1 << ebits) - 2 if dtype == "fp8" else (1 << ebits) - 1 + biased_exp = min(biased_exp, max_exp) + frac = v / (2.0 ** (biased_exp - bias)) - 1.0 + mant = int(round(frac * (1 << mbits))) + if mant >= (1 << mbits): + mant = 0 + biased_exp += 1 + biased_exp = min(biased_exp, max_exp) + out[i] = (sign << 7) | (biased_exp << mbits) | mant + return out.reshape(arr.shape) + + +def _soft_fp8_to_f32(arr: np.ndarray, dtype: str) -> np.ndarray: + """Software fp8 → float32 conversion (no ml_dtypes dependency).""" + flat = arr.ravel().astype(np.uint8) + if dtype == "fp8": + ebits, mbits = 4, 3 # e4m3fn: bias=7 + else: + ebits, mbits = 5, 2 # e5m2: bias=15 + bias = (1 << (ebits - 1)) - 1 + max_exp_bits = (1 << ebits) - 1 + out = np.zeros(flat.shape, dtype=np.float32) + for i, byte in enumerate(flat): + sign = (byte >> 7) & 1 + exp_bits = (byte >> mbits) & ((1 << ebits) - 1) + mant_bits = byte & ((1 << mbits) - 1) + if dtype == "fp8" and exp_bits == max_exp_bits and mant_bits == (1 << mbits) - 1: + out[i] = float('nan') + elif dtype == "bf8" and exp_bits == max_exp_bits: + out[i] = float('nan') if mant_bits != 0 else float('inf') * (-1 if sign else 1) + elif exp_bits == 0: + val = (mant_bits / (1 << mbits)) * (2.0 ** (1 - bias)) + out[i] = -val if sign else val + else: + val = (1.0 + mant_bits / (1 << mbits)) * (2.0 ** (exp_bits - bias)) + out[i] = -val if sign else val + return out.reshape(arr.shape) # ============================================================================= @@ -247,6 +318,11 @@ def main(): log.info("PASSED (max_rel=%.4f, tol=%.4f)", max_rel, tolerance) else: log.error("FAILED (max_rel=%.4f > tol=%.4f)", max_rel, tolerance) + C_gpu_f32 = result.C.astype(np.float32) + C_ref_f32 = C_ref.astype(np.float32) + log.error("C_gpu[:2,:4] = %s", C_gpu_f32[:2, :4]) + log.error("C_ref[:2,:4] = %s", C_ref_f32[:2, :4]) + log.error("C_gpu nonzero count: %d / %d", int(np.count_nonzero(C_gpu_f32)), C_gpu_f32.size) return 1 else: log.info("Verification skipped (--no-verify)") From f8932137d676c67d87757e800592c7fedb3d9325 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Thu, 30 Jul 2026 17:41:36 -0400 Subject: [PATCH 22/35] fix(ck-tile): fix uint8 underflow in soft fp8 decoder causing NaN reference 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 --- .../examples/gemm/python/14_grouped_gemm_aquant.py | 6 +++--- .../examples/gemm/python/15_grouped_gemm_abquant.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py index 2bea190ecee7..38b030636a2f 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/14_grouped_gemm_aquant.py @@ -133,9 +133,9 @@ def _soft_fp8_to_f32(arr: np.ndarray, dtype: str) -> np.ndarray: max_exp_bits = (1 << ebits) - 1 out = np.zeros(flat.shape, dtype=np.float32) for i, byte in enumerate(flat): - sign = (byte >> 7) & 1 - exp_bits = (byte >> mbits) & ((1 << ebits) - 1) - mant_bits = byte & ((1 << mbits) - 1) + sign = int((byte >> 7) & 1) + exp_bits = int((byte >> mbits) & ((1 << ebits) - 1)) + mant_bits = int(byte & ((1 << mbits) - 1)) if dtype == "fp8" and exp_bits == max_exp_bits and mant_bits == (1 << mbits) - 1: out[i] = float('nan') elif dtype == "bf8" and exp_bits == max_exp_bits: diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py index c6fc18a0e643..1ac6300b9c47 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -132,9 +132,9 @@ def _soft_fp8_to_f32(arr: np.ndarray, dtype: str) -> np.ndarray: max_exp_bits = (1 << ebits) - 1 out = np.zeros(flat.shape, dtype=np.float32) for i, byte in enumerate(flat): - sign = (byte >> 7) & 1 - exp_bits = (byte >> mbits) & ((1 << ebits) - 1) - mant_bits = byte & ((1 << mbits) - 1) + sign = int((byte >> 7) & 1) + exp_bits = int((byte >> mbits) & ((1 << ebits) - 1)) + mant_bits = int(byte & ((1 << mbits) - 1)) if dtype == "fp8" and exp_bits == max_exp_bits and mant_bits == (1 << mbits) - 1: out[i] = float('nan') elif dtype == "bf8" and exp_bits == max_exp_bits: From 18b974662d8412acc173ba975d80f9138e57bdb1 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Thu, 30 Jul 2026 18:14:27 -0400 Subject: [PATCH 23/35] fix(ck-tile): cache GPU arch at init to avoid per-call hipGetDeviceProperties 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 --- .../grouped_gemm_abquant_ctypes_lib.cpp | 152 +++++++++++------- .../ctypes/grouped_gemm_aquant_ctypes_lib.cpp | 119 ++++++++------ 2 files changed, 166 insertions(+), 105 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp index d9619fbd357f..080a397e5dc1 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -18,6 +18,10 @@ * BQ[ceil(K/bK), ceil(N/bN)] - B-side scale tensor (ColumnMajor, as required by kernel) * Constraint: AQuantGroupSize::kK == BQuantGroupSize::kK * + * APreshuffleQuant/BPreshuffleQuant: when enabled, AQ and/or BQ are reordered in host memory + * via shuffle_aq()/shuffle_bq() before the device copy, matching the interleaved layout + * the kernel's block tiles expect. + * * Design: direct launch -- SelectedKernel::launch(QuantGemmHostArgs, stream_config) is * called directly. No dispatcher registry is used: ABQuant kernels take QuantGemmHostArgs, * which is incompatible with the GeneratedTileKernelInstance::run() signature used by @@ -27,11 +31,10 @@ */ #include +#include "ck_tile/host/tensor_shuffle_utils.hpp" #include -#include #include #include -#include // Kernel header force-included via -include compiler flag. // Defines: ADataType, BDataType, CDataType, QDataType (shared AQ/BQ type), AccDataType, @@ -50,34 +53,38 @@ static constexpr std::size_t elements_to_bytes(std::size_t n) // assumed at compile time -- do not hardcode a default architecture here. static bool g_initialized = false; - -#define HIP_CHECK(call) \ - { \ - hipError_t _err = (call); \ - if(_err != hipSuccess) \ - { \ - std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ - << __LINE__ << "\n"; \ - return -1; \ - } \ - } +static std::string g_gfx_arch; extern "C" { /** * Initialize the ctypes lib. Must be called before dispatcher_run_abquant_gemm. * - * This library uses a single-kernel-per-.so model: SelectedKernel is - * force-included at compile time and invoked directly via SelectedKernel::launch(). - * No dispatcher registry is involved -- ABQuant kernels require QuantGemmHostArgs - * which is incompatible with the GeneratedTileKernelInstance::run() signature. + * Queries and caches the GPU architecture so subsequent run calls avoid the + * per-call hipGetDeviceProperties overhead. * - * Returns 0 on success. + * Returns 0 on success, -1 if device query fails or arch is unsupported. */ int dispatcher_initialize() { if(g_initialized) return 0; + + int dev = 0; + hipDeviceProp_t props{}; + if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) + { + std::cerr << "dispatcher_initialize: could not query device architecture\n"; + return -1; + } + g_gfx_arch = props.gcnArchName; + if(g_gfx_arch.rfind("gfx950", 0) != 0 && g_gfx_arch.rfind("gfx942", 0) != 0 && + g_gfx_arch.rfind("gfx90a", 0) != 0) + { + std::cerr << "dispatcher_initialize: unsupported GPU architecture '" << g_gfx_arch + << "' (supported: gfx90a, gfx942, gfx950)\n"; + return -1; + } g_initialized = true; return 0; } @@ -142,25 +149,6 @@ int dispatcher_run_abquant_gemm(const void* A, return -1; } - // Derive the GPU architecture from the running device and reject unsupported archs. - { - int dev = 0; - hipDeviceProp_t props{}; - if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) - { - std::cerr << "dispatcher_run_abquant_gemm: could not query device architecture\n"; - return -1; - } - const std::string arch(props.gcnArchName); - if(arch.rfind("gfx950", 0) != 0 && arch.rfind("gfx942", 0) != 0 && - arch.rfind("gfx90a", 0) != 0) - { - std::cerr << "dispatcher_run_abquant_gemm: unsupported GPU architecture '" << arch - << "' (supported: gfx90a, gfx942, gfx950)\n"; - return -1; - } - } - // Validate AQ dimensions { const int64_t exp_QK_A = @@ -229,12 +217,12 @@ int dispatcher_run_abquant_gemm(const void* A, // Allocate device buffers. // elements_to_bytes(n) handles packed types (pk_int4_t etc.) via PackedSize. - if(hipMalloc(&A_dev, elements_to_bytes(M * K)) != hipSuccess) + if(hipMalloc(&A_dev, elements_to_bytes(static_cast(M) * static_cast(K))) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, elements_to_bytes(K * N)) != hipSuccess) + if(hipMalloc(&B_dev, elements_to_bytes(static_cast(K) * static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -249,39 +237,89 @@ int dispatcher_run_abquant_gemm(const void* A, cleanup(); return -1; } - if(hipMalloc(&C_dev, elements_to_bytes(M * N)) != hipSuccess) + if(hipMalloc(&C_dev, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice) != + if(hipMemcpy(A_dev, A_host, elements_to_bytes(static_cast(M) * static_cast(K)), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice) != + if(hipMemcpy(B_dev, B_host, elements_to_bytes(static_cast(K) * static_cast(N)), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy( - AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != - hipSuccess) + // Copy AQ to device; preshuffle when the kernel expects the interleaved layout. + if constexpr(SelectedKernel::APreshuffleQuant) { - cleanup(); - return -1; + constexpr int block_aq_k = + static_cast(SelectedKernel::TileK) / static_cast(AQuantGroupSize::kK); + ck_tile::HostTensor aq_h( + ck_tile::host_tensor_descriptor(static_cast(QM_A), + static_cast(QK_A), + static_cast(QK_A), + ck_tile::bool_constant{} /*row-major*/)); + std::copy(AQ_host, AQ_host + QM_A * QK_A, aq_h.begin()); + auto aq_shuffled = ck_tile::shuffle_aq(&aq_h, block_aq_k); + if(hipMemcpy(AQ_dev, + aq_shuffled.data(), + elements_to_bytes(QM_A * QK_A), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } } - if(hipMemcpy( - BQ_dev, BQ_host, elements_to_bytes(QK_B * QN_B), hipMemcpyHostToDevice) != - hipSuccess) + else { - cleanup(); - return -1; + if(hipMemcpy(AQ_dev, + AQ_host, + elements_to_bytes(QM_A * QK_A), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } } - if(hipMemset(C_dev, 0, elements_to_bytes(M * N)) != hipSuccess) + // Copy BQ to device; preshuffle when the kernel expects the interleaved layout. + if constexpr(SelectedKernel::BPreshuffleQuant) + { + constexpr int block_bq_k = + static_cast(SelectedKernel::TileK) / static_cast(BQuantGroupSize::kK); + ck_tile::HostTensor bq_h( + ck_tile::host_tensor_descriptor(static_cast(QK_B), + static_cast(QN_B), + static_cast(QN_B), + ck_tile::bool_constant{} /*row-major*/)); + std::copy(BQ_host, BQ_host + QK_B * QN_B, bq_h.begin()); + auto bq_shuffled = ck_tile::shuffle_bq(&bq_h, block_bq_k); + if(hipMemcpy(BQ_dev, + bq_shuffled.data(), + elements_to_bytes(QK_B * QN_B), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + else + { + if(hipMemcpy(BQ_dev, + BQ_host, + elements_to_bytes(QK_B * QN_B), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + if(hipMemset(C_dev, 0, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -327,7 +365,7 @@ int dispatcher_run_abquant_gemm(const void* A, return -2; } - if(hipMemcpy(C_host, C_dev, elements_to_bytes(M * N), hipMemcpyDeviceToHost) != + if(hipMemcpy(C_host, C_dev, elements_to_bytes(static_cast(M) * static_cast(N)), hipMemcpyDeviceToHost) != hipSuccess) { cleanup(); @@ -359,6 +397,10 @@ int dispatcher_get_kernel_count() { return 1; } /** * Release dispatcher resources. */ -void dispatcher_cleanup() { g_initialized = false; } +void dispatcher_cleanup() +{ + g_initialized = false; + g_gfx_arch.clear(); +} } // extern "C" diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp index 8bb45ae83efd..6cdd57b6c889 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp @@ -16,6 +16,10 @@ * AQ[ceil(M/gM), ceil(K/gK)] is the A-side scale tensor (RowMajor). * BQ is unused (bq_ptr=nullptr, QK_B=0, stride_BQ=0 in QuantGemmHostArgs). * + * APreshuffleQuant: when SelectedKernel::APreshuffleQuant==true (pipeline=compv3), + * AQ is reordered in host memory via shuffle_aq() before the device copy, matching + * the interleaved layout the kernel's block tile expects. + * * Design: direct launch -- SelectedKernel::launch(QuantGemmHostArgs, stream_config) is * called directly. No dispatcher registry is used: AQuant kernels take QuantGemmHostArgs, * which is incompatible with the GeneratedTileKernelInstance::run() signature used by @@ -25,11 +29,10 @@ */ #include +#include "ck_tile/host/tensor_shuffle_utils.hpp" #include -#include #include #include -#include // Kernel header force-included via -include compiler flag. // Defines: ADataType, BDataType, CDataType, QDataType (AQDataType), AccDataType, @@ -48,34 +51,38 @@ static constexpr std::size_t elements_to_bytes(std::size_t n) // assumed at compile time -- do not hardcode a default architecture here. static bool g_initialized = false; - -#define HIP_CHECK(call) \ - { \ - hipError_t _err = (call); \ - if(_err != hipSuccess) \ - { \ - std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ - << __LINE__ << "\n"; \ - return -1; \ - } \ - } +static std::string g_gfx_arch; extern "C" { /** * Initialize the ctypes lib. Must be called before dispatcher_run_aquant_gemm. * - * This library uses a single-kernel-per-.so model: SelectedKernel is - * force-included at compile time and invoked directly via SelectedKernel::launch(). - * No dispatcher registry is involved -- AQuant kernels require QuantGemmHostArgs - * which is incompatible with the GeneratedTileKernelInstance::run() signature. + * Queries and caches the GPU architecture so subsequent run calls avoid the + * per-call hipGetDeviceProperties overhead. * - * Returns 0 on success. + * Returns 0 on success, -1 if device query fails or arch is unsupported. */ int dispatcher_initialize() { if(g_initialized) return 0; + + int dev = 0; + hipDeviceProp_t props{}; + if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) + { + std::cerr << "dispatcher_initialize: could not query device architecture\n"; + return -1; + } + g_gfx_arch = props.gcnArchName; + if(g_gfx_arch.rfind("gfx950", 0) != 0 && g_gfx_arch.rfind("gfx942", 0) != 0 && + g_gfx_arch.rfind("gfx90a", 0) != 0) + { + std::cerr << "dispatcher_initialize: unsupported GPU architecture '" << g_gfx_arch + << "' (supported: gfx90a, gfx942, gfx950)\n"; + return -1; + } g_initialized = true; return 0; } @@ -131,25 +138,6 @@ int dispatcher_run_aquant_gemm(const void* A, return -1; } - // Derive the GPU architecture from the running device and reject unsupported archs. - { - int dev = 0; - hipDeviceProp_t props{}; - if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&props, dev) != hipSuccess) - { - std::cerr << "dispatcher_run_aquant_gemm: could not query device architecture\n"; - return -1; - } - const std::string arch(props.gcnArchName); - if(arch.rfind("gfx950", 0) != 0 && arch.rfind("gfx942", 0) != 0 && - arch.rfind("gfx90a", 0) != 0) - { - std::cerr << "dispatcher_run_aquant_gemm: unsupported GPU architecture '" << arch - << "' (supported: gfx90a, gfx942, gfx950)\n"; - return -1; - } - } - // Validate that the caller's QK_A/QM_A match the compile-time quant group sizes. { const int64_t expected_QK_A = @@ -201,12 +189,12 @@ int dispatcher_run_aquant_gemm(const void* A, // Allocate device buffers. // B may be a packed type (pk_int4_t): elements_to_bytes(n) handles PackedSize correctly. - if(hipMalloc(&A_dev, elements_to_bytes(M * K)) != hipSuccess) + if(hipMalloc(&A_dev, elements_to_bytes(static_cast(M) * static_cast(K))) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, elements_to_bytes(K * N)) != hipSuccess) + if(hipMalloc(&B_dev, elements_to_bytes(static_cast(K) * static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -216,33 +204,60 @@ int dispatcher_run_aquant_gemm(const void* A, cleanup(); return -1; } - if(hipMalloc(&C_dev, elements_to_bytes(M * N)) != hipSuccess) + if(hipMalloc(&C_dev, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) { cleanup(); return -1; } // Copy inputs to device - if(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice) != + if(hipMemcpy(A_dev, A_host, elements_to_bytes(static_cast(M) * static_cast(K)), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice) != + if(hipMemcpy(B_dev, B_host, elements_to_bytes(static_cast(K) * static_cast(N)), hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy( - AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != - hipSuccess) + // Copy AQ to device, preshuffling when the kernel requires it. + // APreshuffleQuant=true (pipeline=compv3): AQ is reordered so the kernel's block tile + // finds scale values in the interleaved layout it expects. + if constexpr(SelectedKernel::APreshuffleQuant) { - cleanup(); - return -1; + constexpr int block_aq_k = + static_cast(SelectedKernel::TileK) / static_cast(QuantGroupSize::kK); + ck_tile::HostTensor aq_h( + ck_tile::host_tensor_descriptor(static_cast(QM_A), + static_cast(QK_A), + static_cast(QK_A), + ck_tile::bool_constant{} /*row-major*/)); + std::copy(AQ_host, AQ_host + QM_A * QK_A, aq_h.begin()); + auto aq_shuffled = ck_tile::shuffle_aq(&aq_h, block_aq_k); + if(hipMemcpy(AQ_dev, + aq_shuffled.data(), + elements_to_bytes(QM_A * QK_A), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + else + { + if(hipMemcpy(AQ_dev, + AQ_host, + elements_to_bytes(QM_A * QK_A), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } } - if(hipMemset(C_dev, 0, elements_to_bytes(M * N)) != hipSuccess) + if(hipMemset(C_dev, 0, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -289,7 +304,7 @@ int dispatcher_run_aquant_gemm(const void* A, } // Copy result back - if(hipMemcpy(C_host, C_dev, elements_to_bytes(M * N), hipMemcpyDeviceToHost) != + if(hipMemcpy(C_host, C_dev, elements_to_bytes(static_cast(M) * static_cast(N)), hipMemcpyDeviceToHost) != hipSuccess) { cleanup(); @@ -321,6 +336,10 @@ int dispatcher_get_kernel_count() { return 1; } /** * Release dispatcher resources. */ -void dispatcher_cleanup() { g_initialized = false; } +void dispatcher_cleanup() +{ + g_initialized = false; + g_gfx_arch.clear(); +} } // extern "C" From a8bb8f9426783707a76dda575531ed72e78bb5d1 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 15:02:49 -0400 Subject: [PATCH 24/35] style(composablekernel): apply clang-format-18 to grouped_gemm ctypes libs Co-Authored-By: Claude --- .../grouped_gemm_abquant_ctypes_lib.cpp | 52 ++++++++++++------- .../ctypes/grouped_gemm_aquant_ctypes_lib.cpp | 45 +++++++++++----- 2 files changed, 65 insertions(+), 32 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp index 080a397e5dc1..844e54796eae 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_abquant_ctypes_lib.cpp @@ -217,12 +217,16 @@ int dispatcher_run_abquant_gemm(const void* A, // Allocate device buffers. // elements_to_bytes(n) handles packed types (pk_int4_t etc.) via PackedSize. - if(hipMalloc(&A_dev, elements_to_bytes(static_cast(M) * static_cast(K))) != hipSuccess) + if(hipMalloc(&A_dev, + elements_to_bytes(static_cast(M) * + static_cast(K))) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, elements_to_bytes(static_cast(K) * static_cast(N))) != hipSuccess) + if(hipMalloc(&B_dev, + elements_to_bytes(static_cast(K) * + static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -237,20 +241,28 @@ int dispatcher_run_abquant_gemm(const void* A, cleanup(); return -1; } - if(hipMalloc(&C_dev, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) + if(hipMalloc(&C_dev, + elements_to_bytes(static_cast(M) * + static_cast(N))) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(A_dev, A_host, elements_to_bytes(static_cast(M) * static_cast(K)), hipMemcpyHostToDevice) != - hipSuccess) + if(hipMemcpy( + A_dev, + A_host, + elements_to_bytes(static_cast(M) * static_cast(K)), + hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, elements_to_bytes(static_cast(K) * static_cast(N)), hipMemcpyHostToDevice) != - hipSuccess) + if(hipMemcpy( + B_dev, + B_host, + elements_to_bytes(static_cast(K) * static_cast(N)), + hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; @@ -278,10 +290,9 @@ int dispatcher_run_abquant_gemm(const void* A, } else { - if(hipMemcpy(AQ_dev, - AQ_host, - elements_to_bytes(QM_A * QK_A), - hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy( + AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; @@ -310,16 +321,18 @@ int dispatcher_run_abquant_gemm(const void* A, } else { - if(hipMemcpy(BQ_dev, - BQ_host, - elements_to_bytes(QK_B * QN_B), - hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy( + BQ_dev, BQ_host, elements_to_bytes(QK_B * QN_B), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; } } - if(hipMemset(C_dev, 0, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) + if(hipMemset(C_dev, + 0, + elements_to_bytes(static_cast(M) * + static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -365,8 +378,11 @@ int dispatcher_run_abquant_gemm(const void* A, return -2; } - if(hipMemcpy(C_host, C_dev, elements_to_bytes(static_cast(M) * static_cast(N)), hipMemcpyDeviceToHost) != - hipSuccess) + if(hipMemcpy( + C_host, + C_dev, + elements_to_bytes(static_cast(M) * static_cast(N)), + hipMemcpyDeviceToHost) != hipSuccess) { cleanup(); return -1; diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp index 6cdd57b6c889..0597217ff1ec 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/grouped_gemm_aquant_ctypes_lib.cpp @@ -189,12 +189,16 @@ int dispatcher_run_aquant_gemm(const void* A, // Allocate device buffers. // B may be a packed type (pk_int4_t): elements_to_bytes(n) handles PackedSize correctly. - if(hipMalloc(&A_dev, elements_to_bytes(static_cast(M) * static_cast(K))) != hipSuccess) + if(hipMalloc(&A_dev, + elements_to_bytes(static_cast(M) * + static_cast(K))) != hipSuccess) { cleanup(); return -1; } - if(hipMalloc(&B_dev, elements_to_bytes(static_cast(K) * static_cast(N))) != hipSuccess) + if(hipMalloc(&B_dev, + elements_to_bytes(static_cast(K) * + static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -204,21 +208,29 @@ int dispatcher_run_aquant_gemm(const void* A, cleanup(); return -1; } - if(hipMalloc(&C_dev, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) + if(hipMalloc(&C_dev, + elements_to_bytes(static_cast(M) * + static_cast(N))) != hipSuccess) { cleanup(); return -1; } // Copy inputs to device - if(hipMemcpy(A_dev, A_host, elements_to_bytes(static_cast(M) * static_cast(K)), hipMemcpyHostToDevice) != - hipSuccess) + if(hipMemcpy( + A_dev, + A_host, + elements_to_bytes(static_cast(M) * static_cast(K)), + hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; } - if(hipMemcpy(B_dev, B_host, elements_to_bytes(static_cast(K) * static_cast(N)), hipMemcpyHostToDevice) != - hipSuccess) + if(hipMemcpy( + B_dev, + B_host, + elements_to_bytes(static_cast(K) * static_cast(N)), + hipMemcpyHostToDevice) != hipSuccess) { cleanup(); return -1; @@ -248,16 +260,18 @@ int dispatcher_run_aquant_gemm(const void* A, } else { - if(hipMemcpy(AQ_dev, - AQ_host, - elements_to_bytes(QM_A * QK_A), - hipMemcpyHostToDevice) != hipSuccess) + if(hipMemcpy( + AQ_dev, AQ_host, elements_to_bytes(QM_A * QK_A), hipMemcpyHostToDevice) != + hipSuccess) { cleanup(); return -1; } } - if(hipMemset(C_dev, 0, elements_to_bytes(static_cast(M) * static_cast(N))) != hipSuccess) + if(hipMemset(C_dev, + 0, + elements_to_bytes(static_cast(M) * + static_cast(N))) != hipSuccess) { cleanup(); return -1; @@ -304,8 +318,11 @@ int dispatcher_run_aquant_gemm(const void* A, } // Copy result back - if(hipMemcpy(C_host, C_dev, elements_to_bytes(static_cast(M) * static_cast(N)), hipMemcpyDeviceToHost) != - hipSuccess) + if(hipMemcpy( + C_host, + C_dev, + elements_to_bytes(static_cast(M) * static_cast(N)), + hipMemcpyDeviceToHost) != hipSuccess) { cleanup(); return -1; From 4e495f2583671e082817251ebf40d47f8e48be4a Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 17:11:46 -0400 Subject: [PATCH 25/35] fix(ck-tile): match warp_tile_k to get_k_warp_tile for eightwaves/preshuffleb on gfx950 On gfx950 (CK_GFX950_SUPPORT defined), get_k_warp_tile 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 --- .../python/grouped_gemm_abquant_utils.py | 58 ++++++++++++------- 1 file changed, 37 insertions(+), 21 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index 98a7f400506c..02630dd82930 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -682,6 +682,26 @@ def default_bf8_compv3_config( ) +def _eightwaves_warp_tile_k(gfx_arch: str) -> int: + """Return the correct warp_tile_k for the EightWaves pipeline on the given arch. + + Mirrors get_k_warp_tile: + gfx950: CK_GFX950_SUPPORT defined → returns 128 + gfx942: returns 32 + """ + return 128 if gfx_arch.startswith("gfx950") else 32 + + +def _preshuffleb_warp_tile_k(gfx_arch: str) -> int: + """Return the correct warp_tile_k for the PreshuffleB pipeline on the given arch. + + Mirrors get_k_warp_tile: + gfx950: CK_GFX950_SUPPORT defined → returns 128 + gfx942: returns 64 + """ + return 128 if gfx_arch.startswith("gfx950") else 64 + + def default_fp8_eightwaves_config( quant_group_k: int = 128, bquant_group_n: int = 128, @@ -690,12 +710,10 @@ def default_fp8_eightwaves_config( """fp8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True). Tile: 192x256x128, warp 4x2x1 — 8-wave configuration for MI350X. - warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8, valid on gfx942 and gfx950. - GemmConfigEightWaves inherits GemmConfigABQuantPrefill which uses - get_k_warp_tile: gfx950 returns 128, gfx942 returns 32. - We use 32 (= gfx942 value) as the cross-arch safe choice — it is 4 MFMA iterations - on gfx942 and also valid on gfx950 (128 % 32 == 0). warp_tile_k=16 would select - the gfx12-only WMMA instruction which silently returns zeros on gfx9 hardware. + warp_tile_k mirrors get_k_warp_tile: + gfx950 (CK_GFX950_SUPPORT): 128 + gfx942: 32 + The value must match what the C++ kernel header computes at compile time. """ return ABQuantKernelConfig( variant_key="fp8", @@ -705,7 +723,7 @@ def default_fp8_eightwaves_config( scheduler="intrawave", tile_m=192, tile_n=256, tile_k=128, warp_m=4, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_eightwaves_warp_tile_k(gfx_arch), aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -722,8 +740,8 @@ def default_bf8_eightwaves_config( ) -> ABQuantKernelConfig: """bf8 ABQuant gfx950 EightWaves config (GemmConfigEightWaves, TransposeC=True). - warp_tile_k=32 selects mfma_f32_16x16x32_bf8_bf8, valid on gfx942 and gfx950. - See default_fp8_eightwaves_config for the full rationale. + warp_tile_k mirrors get_k_warp_tile: + gfx950: 128, gfx942: 32. See default_fp8_eightwaves_config for rationale. """ return ABQuantKernelConfig( variant_key="bf8", @@ -733,7 +751,7 @@ def default_bf8_eightwaves_config( scheduler="intrawave", tile_m=192, tile_n=256, tile_k=128, warp_m=4, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_eightwaves_warp_tile_k(gfx_arch), aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -750,13 +768,11 @@ def default_fp8_preshuffleb_config( ) -> ABQuantKernelConfig: """fp8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill). - Tile: 128x128x128, warp 2x2x1 (vs 1x4x1 for compv3). DoubleSmemBuffer=True. - - warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8, valid on gfx942 and gfx950. - GemmConfigPreshuffleB_ABQuant_Prefill uses get_k_warp_tile: - gfx950 returns 128, gfx942 returns 64. We use 32 as the cross-arch safe choice - (128 % 32 == 0 and 64 % 32 == 0). warp_tile_k=16 would select the gfx12-only WMMA - instruction which silently returns zeros on gfx9 hardware. + Tile: 128x128x128, warp 2x2x1. DoubleSmemBuffer=True. + warp_tile_k mirrors get_k_warp_tile: + gfx950 (CK_GFX950_SUPPORT): 128 + gfx942: 64 + The value must match what the C++ kernel header computes at compile time. """ return ABQuantKernelConfig( variant_key="fp8", @@ -766,7 +782,7 @@ def default_fp8_preshuffleb_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=2, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=True, preshuffle_aq=False, preshuffle_bq=False, @@ -784,8 +800,8 @@ def default_bf8_preshuffleb_config( ) -> ABQuantKernelConfig: """bf8 ABQuant PreshuffleB config (GemmConfigPreshuffleB_ABQuant_Prefill). - warp_tile_k=32 selects mfma_f32_16x16x32_bf8_bf8, valid on gfx942 and gfx950. - See default_fp8_preshuffleb_config for the full rationale. + warp_tile_k mirrors get_k_warp_tile: + gfx950: 128, gfx942: 64. See default_fp8_preshuffleb_config for rationale. """ return ABQuantKernelConfig( variant_key="bf8", @@ -795,7 +811,7 @@ def default_bf8_preshuffleb_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=2, warp_n=2, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=True, preshuffle_aq=False, preshuffle_bq=False, From eb1a1a70c04d94eb5439e20c8a959484418c0c00 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 17:21:15 -0400 Subject: [PATCH 26/35] fix(ck-tile): use tile-aligned default shapes in example 15 ABQuant 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 --- .../gemm/python/15_grouped_gemm_abquant.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py index 1ac6300b9c47..b99e3d5473f6 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -211,9 +211,9 @@ def main(): parser.add_argument("--dtype", choices=["fp8", "bf8"], default="fp8") parser.add_argument("--pipeline", choices=["compv3", "eightwaves", "preshuffleb"], default="compv3") - parser.add_argument("--M", type=int, default=128) - parser.add_argument("--N", type=int, default=128) - parser.add_argument("--K", type=int, default=128) + parser.add_argument("--M", type=int, default=None) + parser.add_argument("--N", type=int, default=None) + parser.add_argument("--K", type=int, default=1024) parser.add_argument("--quant-group-k", type=int, default=128) parser.add_argument("--bquant-group-n", type=int, default=1) parser.add_argument("--no-verify", action="store_true") @@ -221,7 +221,16 @@ def main(): parser.add_argument("--gfx-arch", type=str, default="gfx950") args = parser.parse_args() - M, N, K = args.M, args.N, args.K + # Default M and N to the smallest valid multiples of the tile for each pipeline. + # kPadM/kPadN are false for all ABQuant prefill configs, so M and N must be exact + # multiples of TileM and TileN respectively. + # compv3/preshuffleb: TileM=128, TileN=128 + # eightwaves: TileM=192, TileN=256 + _default_mn = {"compv3": (1024, 1024), "preshuffleb": (1024, 1024), "eightwaves": (1152, 1024)} + _dm, _dn = _default_mn[args.pipeline] + M = args.M if args.M is not None else _dm + N = args.N if args.N is not None else _dn + K = args.K gK = args.quant_group_k bN = args.bquant_group_n From dd7f9e3d4d7bad9a8707570f6c560c6f78cbc63d Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 17:24:14 -0400 Subject: [PATCH 27/35] fix(ck-tile): apply arch-aware warp_tile_k to compv3 ABQuant config on gfx950 GemmConfigABQuantPrefill (compv3 pipeline) also uses get_k_warp_tile, 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 --- .../python/grouped_gemm_abquant_utils.py | 31 ++++++++++++------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index 02630dd82930..1fd204f151b4 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -629,13 +629,13 @@ def default_fp8_compv3_config( bquant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> ABQuantKernelConfig: - """fp8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False). + """fp8 ABQuant CompV3 config (GemmConfigABQuantPrefill, TransposeC=False). Tile: 128x128x128, warp 1x4x1. kPadK=false (ABQuant prefill tiles). - - warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8, valid on gfx942 and gfx950. - warp_tile_k=16 would select the gfx12-only WMMA instruction which silently returns - zeros on gfx9 hardware. + warp_tile_k mirrors get_k_warp_tile: + gfx950 (CK_GFX950_SUPPORT): 128 + gfx942: 32 + The value must match what the C++ kernel header computes at compile time. """ return ABQuantKernelConfig( variant_key="fp8", @@ -645,7 +645,7 @@ def default_fp8_compv3_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_compv3_warp_tile_k(gfx_arch), aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -659,11 +659,10 @@ def default_bf8_compv3_config( bquant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> ABQuantKernelConfig: - """bf8 ABQuant non-gfx950 config (GemmConfigABQuantPrefill, TransposeC=False). + """bf8 ABQuant CompV3 config (GemmConfigABQuantPrefill, TransposeC=False). - warp_tile_k=32 selects mfma_f32_16x16x32_bf8_bf8, valid on gfx942 and gfx950. - warp_tile_k=16 would select the gfx12-only WMMA instruction which silently returns - zeros on gfx9 hardware. + warp_tile_k mirrors get_k_warp_tile: + gfx950: 128, gfx942: 32. See default_fp8_compv3_config for rationale. """ return ABQuantKernelConfig( variant_key="bf8", @@ -673,7 +672,7 @@ def default_bf8_compv3_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_compv3_warp_tile_k(gfx_arch), aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -682,6 +681,16 @@ def default_bf8_compv3_config( ) +def _compv3_warp_tile_k(gfx_arch: str) -> int: + """Return the correct warp_tile_k for the CompV3 (ABQuantPrefill) pipeline. + + Mirrors get_k_warp_tile: + gfx950 (CK_GFX950_SUPPORT): 128 + gfx942: 32 + """ + return 128 if gfx_arch.startswith("gfx950") else 32 + + def _eightwaves_warp_tile_k(gfx_arch: str) -> int: """Return the correct warp_tile_k for the EightWaves pipeline on the given arch. From 0cf210c8d91899a4422ff8ccd6500a15756020f7 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 17:49:50 -0400 Subject: [PATCH 28/35] =?UTF-8?q?fix(ck-tile):=20revert=20compv3=20warp=5F?= =?UTF-8?q?tile=5Fk=20to=2032=20=E2=80=94=20not=20a=20FlatMM=20pipeline?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../python/grouped_gemm_abquant_utils.py | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index 1fd204f151b4..75098a42a199 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -631,11 +631,11 @@ def default_fp8_compv3_config( ) -> ABQuantKernelConfig: """fp8 ABQuant CompV3 config (GemmConfigABQuantPrefill, TransposeC=False). - Tile: 128x128x128, warp 1x4x1. kPadK=false (ABQuant prefill tiles). - warp_tile_k mirrors get_k_warp_tile: - gfx950 (CK_GFX950_SUPPORT): 128 - gfx942: 32 - The value must match what the C++ kernel header computes at compile time. + Tile: 128x128x128, warp 1x4x1. kPadK=false. + warp_tile_k=32 selects mfma_f32_16x16x32_fp8_fp8 (standard MFMA, valid on gfx942 + and gfx950). CompV3 is a standard compute pipeline — it does NOT use FlatMM, so + WarpTileK must stay at 32 regardless of arch. Only the gfx950-native pipelines + (eightwaves, preshuffleb) use FlatMM and require WarpTileK=128. """ return ABQuantKernelConfig( variant_key="fp8", @@ -645,7 +645,7 @@ def default_fp8_compv3_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=_compv3_warp_tile_k(gfx_arch), + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -661,8 +661,7 @@ def default_bf8_compv3_config( ) -> ABQuantKernelConfig: """bf8 ABQuant CompV3 config (GemmConfigABQuantPrefill, TransposeC=False). - warp_tile_k mirrors get_k_warp_tile: - gfx950: 128, gfx942: 32. See default_fp8_compv3_config for rationale. + warp_tile_k=32 — CompV3 uses standard MFMA, not FlatMM. See default_fp8_compv3_config. """ return ABQuantKernelConfig( variant_key="bf8", @@ -672,7 +671,7 @@ def default_bf8_compv3_config( scheduler="intrawave", tile_m=128, tile_n=128, tile_k=128, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=_compv3_warp_tile_k(gfx_arch), + warp_tile_m=16, warp_tile_n=16, warp_tile_k=32, aquant_group_m=1, aquant_group_n=1, aquant_group_k=quant_group_k, bquant_group_m=1, bquant_group_n=bquant_group_n, bquant_group_k=quant_group_k, preshuffle_b=False, preshuffle_aq=False, preshuffle_bq=False, @@ -681,16 +680,6 @@ def default_bf8_compv3_config( ) -def _compv3_warp_tile_k(gfx_arch: str) -> int: - """Return the correct warp_tile_k for the CompV3 (ABQuantPrefill) pipeline. - - Mirrors get_k_warp_tile: - gfx950 (CK_GFX950_SUPPORT): 128 - gfx942: 32 - """ - return 128 if gfx_arch.startswith("gfx950") else 32 - - def _eightwaves_warp_tile_k(gfx_arch: str) -> int: """Return the correct warp_tile_k for the EightWaves pipeline on the given arch. From 4ec2097754a13d0fd6976a2739e95bb76ea676d5 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 18:07:51 -0400 Subject: [PATCH 29/35] fix(ck-tile): fix permute_n unscramble to operate per-tile, not globally MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../dispatcher/python/grouped_gemm_abquant_utils.py | 13 ++++++++----- .../dispatcher/python/grouped_gemm_aquant_utils.py | 13 ++++++++----- .../dispatcher/python/grouped_gemm_bquant_utils.py | 13 ++++++++----- 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index 75098a42a199..9c3fca274538 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -425,8 +425,8 @@ def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuan f"for kernel {self.kernel_name}" ) - # permute_n epilogue writes C with N-columns riffled into r groups - # (r = tile_n / warp_tile_n / warp_n). Undo it so the caller gets logical C. + # permute_n epilogue riffles N-columns within each tile of width tile_n. + # Undo it per-tile so the caller gets logical (row-major) C. _name = self.kernel_name if 'permute_n' in _name: import re as _re @@ -434,9 +434,12 @@ def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuan if _m: _tile_n = int(_m.group(2)); _warp_n = int(_m.group(5)); _wt_n = int(_m.group(8)) _r = _tile_n // _wt_n // _warp_n - if _r > 1 and (N % _r) == 0: - _half = N // _r - _logical = [(c % _r) * _half + (c // _r) for c in range(N)] + if _r > 1 and (N % _tile_n) == 0: + _half = _tile_n // _r + _logical = [ + (c // _tile_n) * _tile_n + (c % _tile_n % _r) * _half + (c % _tile_n // _r) + for c in range(N) + ] _Cp = np.empty_like(C) _Cp[:, _logical] = C C = _Cp diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py index 371eb2b3dfbc..60e529446e2c 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -395,8 +395,8 @@ def run(self, A, B, AQ, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmR f"for kernel {self.kernel_name}" ) - # permute_n epilogue writes C with N-columns riffled into r groups - # (r = tile_n / warp_tile_n / warp_n). Undo it so the caller gets logical C. + # permute_n epilogue riffles N-columns within each tile of width tile_n. + # Undo it per-tile so the caller gets logical (row-major) C. _name = self.kernel_name if 'permute_n' in _name: import re as _re @@ -404,9 +404,12 @@ def run(self, A, B, AQ, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmR if _m: _tile_n = int(_m.group(2)); _warp_n = int(_m.group(5)); _wt_n = int(_m.group(8)) _r = _tile_n // _wt_n // _warp_n - if _r > 1 and (N % _r) == 0: - _half = N // _r - _logical = [(c % _r) * _half + (c // _r) for c in range(N)] + if _r > 1 and (N % _tile_n) == 0: + _half = _tile_n // _r + _logical = [ + (c // _tile_n) * _tile_n + (c % _tile_n % _r) * _half + (c % _tile_n // _r) + for c in range(N) + ] _Cp = np.empty_like(C) _Cp[:, _logical] = C C = _Cp diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py index a24ab8b6a3b6..92fb0038c8f7 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py @@ -401,8 +401,8 @@ def run(self, A, B, BQ, problem: BQuantGemmProblem, c_dtype=None) -> BQuantGemmR f"for kernel {self.kernel_name}" ) - # permute_n epilogue writes C with N-columns riffled into r groups - # (r = tile_n / warp_tile_n / warp_n). Undo it so the caller gets logical C. + # permute_n epilogue riffles N-columns within each tile of width tile_n. + # Undo it per-tile so the caller gets logical (row-major) C. _name = self.kernel_name if 'permute_n' in _name: import re as _re @@ -410,9 +410,12 @@ def run(self, A, B, BQ, problem: BQuantGemmProblem, c_dtype=None) -> BQuantGemmR if _m: _tile_n = int(_m.group(2)); _warp_n = int(_m.group(5)); _wt_n = int(_m.group(8)) _r = _tile_n // _wt_n // _warp_n - if _r > 1 and (N % _r) == 0: - _half = N // _r - _logical = [(c % _r) * _half + (c // _r) for c in range(N)] + if _r > 1 and (N % _tile_n) == 0: + _half = _tile_n // _r + _logical = [ + (c // _tile_n) * _tile_n + (c % _tile_n % _r) * _half + (c % _tile_n // _r) + for c in range(N) + ] _Cp = np.empty_like(C) _Cp[:, _logical] = C C = _Cp From 1e699eb14fcd3cc0b07fb7fd6db3302c2169094d Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 18:18:50 -0400 Subject: [PATCH 30/35] fix(ck-tile): add kernel_attr for EightWaves ABQuant dispatcher 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 (where eight_waves is true for the 4x2x1 warp / warp_tile_k=128 config on gfx950). The dispatcher codegen was using plain make_kernel 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 to make_kernel. For compv3 and preshuffleb pipelines IsEightWaves=false so kernel_attr 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 --- .../codegen/unified_grouped_gemm_abquant_codegen.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py index 7c555f530114..0857cc8ec8b4 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py @@ -236,8 +236,9 @@ def generate(self, spec: ABQuantKernelSpec) -> str: preshuffle_b = str(spec.preshuffle_b).lower() preshuffle_aq = str(spec.preshuffle_aq).lower() preshuffle_bq = str(spec.preshuffle_bq).lower() - transpose_c = str(spec.transpose_c).lower() + transpose_c = str(spec.transpose_c).lower() double_smem_buffer = str(spec.double_smem_buffer).lower() + is_eight_waves = str(spec.pipeline == "eightwaves").lower() # Epilogue selection — B-side tile geometry governs PermuteN (same logic as BQuant) use_permute_n_epilogue = ( @@ -342,6 +343,7 @@ def generate(self, spec: ABQuantKernelSpec) -> str: static constexpr bool PreshuffleB = {preshuffle_b}; static constexpr bool TransposeC = {transpose_c}; static constexpr bool DoubleSmemBuffer = {double_smem_buffer}; + static constexpr bool IsEightWaves = {is_eight_waves}; using TileShape = ck_tile::TileGemmShape< ck_tile::sequence, @@ -406,8 +408,10 @@ def generate(self, spec: ABQuantKernelSpec) -> str: const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); const dim3 blocks = Kernel::BlockSize(); + // EightWaves pipeline requires no-packed-fp32-ops to co-execute matrix ops. + using KAttr = ck_tile::kernel_attr; return ck_tile::launch_kernel( - s, ck_tile::make_kernel(Kernel{{}}, grids, blocks, 0, kargs)); + s, ck_tile::make_kernel(Kernel{{}}, grids, blocks, 0, kargs)); }}; return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); From 9cef769345c7bfefbe4cb83236d47a92afd3dc8f Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 18:27:33 -0400 Subject: [PATCH 31/35] fix(ck-tile): EightWaves ABQuant must use CShuffleEpilogue, not PermuteN 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 --- .../dispatcher/codegen/codegen_common.py | 9 +++++++-- .../codegen/unified_grouped_gemm_abquant_codegen.py | 7 ++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index 294c1f5062b7..7eac854129ba 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -530,12 +530,17 @@ def abquant_effective_epilogue( warp_n: int, warp_tile_n: int, bquant_group_n: int, + pipeline: str = "compv3", ) -> str: """Return the epilogue tag the codegen will emit for ABQuant kernels. The PermuteN selection is governed by the B-side quant group N, matching - the same logic used by BQuant / AQuant. + the same logic used by BQuant / AQuant. The EightWaves pipeline uses a + transposed accumulator layout (TransposeC=true) that is incompatible with + PermuteNEpilogue — it must always use CShuffleEpilogue. """ + if pipeline == "eightwaves": + return "cshuffle" n_repeat = tile_n // (warp_n * warp_tile_n) use_permute_n = (n_repeat % 2 == 0) and (bquant_group_n == 1) return "permute_n" if use_permute_n else "cshuffle" @@ -566,7 +571,7 @@ def make_abquant_kernel_name( Both ABQuantKernelConfig (utils) and ABQuantKernelSpec (codegen) delegate to this function so the two sides are guaranteed to stay byte-exact. """ - effective_epilogue = abquant_effective_epilogue(tile_n, warp_n, warp_tile_n, bquant_group_n) + effective_epilogue = abquant_effective_epilogue(tile_n, warp_n, warp_tile_n, bquant_group_n, pipeline) parts = [ "grouped_gemm_abquant", variant_key, diff --git a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py index 0857cc8ec8b4..93e0883adb30 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_grouped_gemm_abquant_codegen.py @@ -240,10 +240,11 @@ def generate(self, spec: ABQuantKernelSpec) -> str: double_smem_buffer = str(spec.double_smem_buffer).lower() is_eight_waves = str(spec.pipeline == "eightwaves").lower() - # Epilogue selection — B-side tile geometry governs PermuteN (same logic as BQuant) + # Epilogue selection — B-side tile geometry governs PermuteN (except EightWaves). use_permute_n_epilogue = ( - abquant_effective_epilogue(t.tile_n, t.warp_n, t.warp_tile_n, spec.bquant_group_n) - == "permute_n" + abquant_effective_epilogue( + t.tile_n, t.warp_n, t.warp_tile_n, spec.bquant_group_n, spec.pipeline + ) == "permute_n" ) if use_permute_n_epilogue: From 2330ac7a1fd55a7c42aad2b0b00c056f502d71bb Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 18:36:54 -0400 Subject: [PATCH 32/35] fix(ck-tile): use bquant_group_n=128 default for eightwaves in example 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 --- .../examples/gemm/python/15_grouped_gemm_abquant.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py index b99e3d5473f6..8b3e9fedf876 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -24,7 +24,7 @@ Usage: python3 15_grouped_gemm_abquant.py # fp8, 128x128x128, no preshuffle python3 15_grouped_gemm_abquant.py --dtype bf8 - python3 15_grouped_gemm_abquant.py --pipeline eightwaves --M 192 --N 256 --K 128 + python3 15_grouped_gemm_abquant.py --pipeline eightwaves --M 192 --N 256 --K 128 # bqgn=128 default python3 15_grouped_gemm_abquant.py --no-verify """ @@ -215,7 +215,7 @@ def main(): parser.add_argument("--N", type=int, default=None) parser.add_argument("--K", type=int, default=1024) parser.add_argument("--quant-group-k", type=int, default=128) - parser.add_argument("--bquant-group-n", type=int, default=1) + parser.add_argument("--bquant-group-n", type=int, default=None) parser.add_argument("--no-verify", action="store_true") parser.add_argument("--output-dir", type=Path, default=None) parser.add_argument("--gfx-arch", type=str, default="gfx950") @@ -232,7 +232,10 @@ def main(): N = args.N if args.N is not None else _dn K = args.K gK = args.quant_group_k - bN = args.bquant_group_n + # EightWaves requires bquant_group_n >= warp_tile_n (16); tested config uses 128. + # CompV3/PreshuffleB support bquant_group_n=1 (one scale per column). + _default_bqn = {"compv3": 1, "preshuffleb": 1, "eightwaves": 128} + bN = args.bquant_group_n if args.bquant_group_n is not None else _default_bqn[args.pipeline] # ------------------------------------------------------------------------- # 1. Build kernel config From dae5945246023ba77c29612abd7ce0bc24c08cfd Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 18:51:49 -0400 Subject: [PATCH 33/35] fix(ck-tile): PreshuffleB ABQuant must use CShuffleEpilogue + bqgn=128 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../dispatcher/codegen/codegen_common.py | 9 +++++---- .../examples/gemm/python/15_grouped_gemm_abquant.py | 9 +++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index 7eac854129ba..58483f3597f2 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -535,11 +535,12 @@ def abquant_effective_epilogue( """Return the epilogue tag the codegen will emit for ABQuant kernels. The PermuteN selection is governed by the B-side quant group N, matching - the same logic used by BQuant / AQuant. The EightWaves pipeline uses a - transposed accumulator layout (TransposeC=true) that is incompatible with - PermuteNEpilogue — it must always use CShuffleEpilogue. + the same logic used by BQuant / AQuant. EightWaves and PreshuffleB use + TransposeC=true (transposed accumulator layout) which is incompatible with + PermuteNEpilogue — both must always use CShuffleEpilogue (TiledMMAPermuteN=false + in the C++ test fixtures for both GemmConfigEightWaves and GemmConfigPreshuffleB_ABQuant_Prefill). """ - if pipeline == "eightwaves": + if pipeline in ("eightwaves", "preshuffleb"): return "cshuffle" n_repeat = tile_n // (warp_n * warp_tile_n) use_permute_n = (n_repeat % 2 == 0) and (bquant_group_n == 1) diff --git a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py index 8b3e9fedf876..9c28ae514118 100644 --- a/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py +++ b/projects/composablekernel/dispatcher/examples/gemm/python/15_grouped_gemm_abquant.py @@ -24,7 +24,8 @@ Usage: python3 15_grouped_gemm_abquant.py # fp8, 128x128x128, no preshuffle python3 15_grouped_gemm_abquant.py --dtype bf8 - python3 15_grouped_gemm_abquant.py --pipeline eightwaves --M 192 --N 256 --K 128 # bqgn=128 default + python3 15_grouped_gemm_abquant.py --pipeline eightwaves --M 192 --N 256 --K 128 # bqgn=128 default + python3 15_grouped_gemm_abquant.py --pipeline preshuffleb # bqgn=128 default python3 15_grouped_gemm_abquant.py --no-verify """ @@ -232,9 +233,9 @@ def main(): N = args.N if args.N is not None else _dn K = args.K gK = args.quant_group_k - # EightWaves requires bquant_group_n >= warp_tile_n (16); tested config uses 128. - # CompV3/PreshuffleB support bquant_group_n=1 (one scale per column). - _default_bqn = {"compv3": 1, "preshuffleb": 1, "eightwaves": 128} + # EightWaves and PreshuffleB are only validated with bquant_group_n=128 in C++ tests. + # CompV3 supports bquant_group_n=1 (one scale per column). + _default_bqn = {"compv3": 1, "preshuffleb": 128, "eightwaves": 128} bN = args.bquant_group_n if args.bquant_group_n is not None else _default_bqn[args.pipeline] # ------------------------------------------------------------------------- From fe637493c3d55f357644bc322ec039d3bb7dd1df Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Fri, 31 Jul 2026 18:59:17 -0400 Subject: [PATCH 34/35] fix(ck-tile): apply host shuffle_b for PreshuffleB ABQuant before kernel launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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(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 --- .../python/grouped_gemm_abquant_utils.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index 9c3fca274538..e1183ed6f0b7 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -367,6 +367,30 @@ def __del__(self): # ============================================================================= +def _shuffle_b_cdna(B: "np.ndarray", wt_n: int, wt_k: int) -> "np.ndarray": + """Apply ck_tile::shuffle_b (gfx9/CDNA path) to B in-Python. + + Mirrors the gfx9 branch of tensor_shuffle_utils.hpp::shuffle_b: + KLane = warp_size / wt_n (64 / wt_n) + ItemsPerAccess = min(16 / sizeof(fp8) = 16, wt_k / KLane) + view shape = [N/wt_n, wt_n, K/items, items] (C++ row-major) + permute = {0, 2, 1, 3} + + B must be a numpy array with shape (K, N) and uint8 dtype (fp8/bf8 bytes). + The shuffle is applied before uploading to the device; the kernel reads + the shuffled layout as [N/wt_n, K/items, wt_n, items] (after permute). + """ + import numpy as np + K, N = B.shape + kLane = 64 // wt_n + items = min(16, wt_k // kLane) + # C++ copies B col-major (HostTensor[K,N] with stride=[1,K]) into + # a row-major view [N/wt_n, wt_n, K/items, items], then permutes {0,2,1,3}. + flat = np.asfortranarray(B).flatten() # col-major byte order + view = flat.reshape(N // wt_n, wt_n, K // items, items) + return np.ascontiguousarray(view.transpose(0, 2, 1, 3)).reshape(B.shape) + + class ABQuantGpuGemmRunner: """High-level runner for ABQuantGrouped GEMM.""" @@ -389,6 +413,7 @@ def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuan Pass np.bfloat16 for MX variants whose CDataType is bf16. """ import numpy as np + import re as _re M, N, K = problem.M, problem.N, problem.K QK_A = problem.QK_A @@ -401,6 +426,17 @@ def run(self, A, B, AQ, BQ, problem: ABQuantGemmProblem, c_dtype=None) -> ABQuan C = np.zeros((M, N), dtype=c_dtype) + # PreshuffleB kernels require B to be host-shuffled before upload. + # Mirrors ck_tile::shuffle_b (gfx9/CDNA path) used by the C++ test fixture. + # After shuffling, pass B as 1-D so the ctypes layer's asfortranarray(B) + # is a no-op (1-D arrays have only one memory order). + _name = self.kernel_name + if 'preshuffleb' in _name: + _pm = _re.search(r'_(\d+)x(\d+)x(\d+)_(\d+)x(\d+)x(\d+)_(\d+)x(\d+)x(\d+)_', _name) + if _pm: + _wt_n = int(_pm.group(8)); _wt_k = int(_pm.group(9)) + B = _shuffle_b_cdna(B, wt_n=_wt_n, wt_k=_wt_k).ravel() + # Stride layout: # AQ RowMajor [QM_A, QK_A]: stride_AQ = QK_A # BQ ColumnMajor [QK_B, QN_B]: stride_BQ = QK_B (leading dim = K-groups) From 44008386f028ac11a35925c538ae8f378fce5160 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Mon, 3 Aug 2026 16:06:19 -0400 Subject: [PATCH 35/35] fix(ck-tile): use flatten(order='F') for col-major B shuffle in PreshuffleB 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 --- .../dispatcher/python/grouped_gemm_abquant_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index e1183ed6f0b7..a55458104fac 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -386,7 +386,7 @@ def _shuffle_b_cdna(B: "np.ndarray", wt_n: int, wt_k: int) -> "np.ndarray": items = min(16, wt_k // kLane) # C++ copies B col-major (HostTensor[K,N] with stride=[1,K]) into # a row-major view [N/wt_n, wt_n, K/items, items], then permutes {0,2,1,3}. - flat = np.asfortranarray(B).flatten() # col-major byte order + flat = B.flatten(order='F') # col-major byte order view = flat.reshape(N // wt_n, wt_n, K // items, items) return np.ascontiguousarray(view.transpose(0, 2, 1, 3)).reshape(B.shape)