From d2e37bd693f96930f75f59b63dc08555fbea9644 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Fri, 24 Jul 2026 18:59:38 -0400 Subject: [PATCH 01/11] feat(ck-tile): aquant GEMM TE to dispatcher bridge Add a registry-bypass, single-kernel-per-.so TE->Dispatcher bridge for the gemm_aquant (A-only quantized) block-scale GEMM operator, mirroring the merged grouped_gemm_bquant bridge. - codegen/unified_gemm_aquant_codegen.py: per-kernel .hpp generator using GemmAQuantPipelineProblem + AQuantGemmPipelineAgBgCrMem/CompV3 and QuantType::AQuantGrouped (QuantGemmHostArgs, aq_ptr set / bq_ptr null). - bindings/ctypes/gemm_aquant_ctypes_lib.cpp: direct-launch C API (dispatcher_run_aquant_gemm) with pk_int4 A permute + AQ shuffle for APreshuffleQuant; runtime arch check throws on unknown arch (no gfx942 default). - python/gemm_aquant_utils.py: byte-exact config/runner/build helpers; arch validated, never silently defaulted. - codegen_common.py: shared make_aquant_kernel_name / aquant_effective_epilogue. - CMake: guarded dispatcher_gemm_aquant_lib target. - scripts/aquant_selftest.py: codegen + name-parity + optional hipcc build runner. dtype x layout matrix matches Old-TE gemm_aquant_quantgrouped*.cpp exactly: fp8/bf8/fp8i4/bf8i4 x {rcr,rrr,crr,ccr} decode (16) and {rcr,rrr,crr} preshufflequant (12). All 28 kernels name-byte-exact and hipcc-compile on gfx950. --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 44 + .../ctypes/gemm_aquant_ctypes_lib.cpp | 398 +++++++++ .../dispatcher/codegen/codegen_common.py | 71 ++ .../codegen/unified_gemm_aquant_codegen.py | 637 ++++++++++++++ .../dispatcher/python/gemm_aquant_utils.py | 827 ++++++++++++++++++ .../dispatcher/scripts/aquant_selftest.py | 149 ++++ 6 files changed, 2126 insertions(+) create mode 100644 projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp create mode 100644 projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py create mode 100644 projects/composablekernel/dispatcher/python/gemm_aquant_utils.py create mode 100644 projects/composablekernel/dispatcher/scripts/aquant_selftest.py diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index e6d637833fe7..47caf221abe5 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -276,6 +276,50 @@ else() ) endif() +# ============================================================================= +# AQuant (A-only quantized) GEMM ctypes Library +# ============================================================================= + +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/gemm_aquant_ctypes_lib.cpp) + file(GLOB AQUANT_GEMM_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/gemm_aquant_*.hpp") + if(AQUANT_GEMM_KERNEL_HEADERS) + list(SORT 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_gemm_aquant_lib + gemm_aquant_ctypes_lib.cpp + KERNEL_HEADER ${AQUANT_GEMM_KERNEL_HEADER} + ) + # Resolve a single GFX arch string from CMAKE_HIP_ARCHITECTURES if available, + # falling back to the cache variable CK_TILE_AQUANT_GFX_ARCH. Never hardcode a + # default arch: if none can be resolved, error out (the ctypes source #errors + # on an unset GFX_ARCH too). The Python build path always passes -DGFX_ARCH=... + # explicitly via hipcc, so this only affects CMake-driven builds. + 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) + set(CK_TILE_AQUANT_GFX_ARCH "${_aquant_gfx_arch}" CACHE STRING + "GFX arch for the CMake-built AQuant ctypes .so (from CMAKE_HIP_ARCHITECTURES)") + else() + message(FATAL_ERROR + "CK_TILE_AQUANT_GFX_ARCH is unset and CMAKE_HIP_ARCHITECTURES is empty. " + "Set -DCK_TILE_AQUANT_GFX_ARCH=gfx942|gfx950 to build dispatcher_gemm_aquant_lib.") + endif() + endif() + message(STATUS "AQuant ctypes lib GFX_ARCH: ${CK_TILE_AQUANT_GFX_ARCH}") + + target_compile_definitions(dispatcher_gemm_aquant_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE + GFX_ARCH="${CK_TILE_AQUANT_GFX_ARCH}" + ) + else() + message(STATUS + "No AQuant GEMM kernel found for ctypes lib - skipping dispatcher_gemm_aquant_lib " + "(built at runtime by the Python bridge once a kernel header exists)") + endif() +endif() + # ============================================================================= # TileEngine -> Dispatcher Bridge ctypes Libraries (append-only shared block) # ============================================================================= diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp new file mode 100644 index 000000000000..2229f13dbb01 --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp @@ -0,0 +1,398 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * AQuant (A-only quantized) GEMM ctypes Library + * + * Provides a C API for Python ctypes integration. One .so is compiled per + * kernel variant; the kernel is force-included at compile time: + * hipcc -include -DCK_TILE_SINGLE_KERNEL_INCLUDE gemm_aquant_ctypes_lib.cpp + * + * Force-include defines (from generated kernel header): + * SelectedKernel, KERNEL_NAME + * ADataType, BDataType, CDataType, QDataType, AccDataType, QuantGroupSize + * + * 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). + * + * AQuant vs BQuant: here the *A* matrix is the quantized operand. The scale + * tensor AQ has shape [M, QK_A] (QK_A = ceil(K/gK)), aq_ptr is set and bq_ptr is + * nullptr. For pk_int4 A the raw values are permuted via permute_vectors_i4x4_b + * before the device copy (mirrors run_gemm_quant_example.inc), and APreshuffleQuant + * kernels shuffle AQ via shuffle_aq. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/host/tensor_shuffle_utils.hpp" +#include "ck_tile/host/permute_pk_int4.hpp" + +// Kernel header force-included via -include compiler flag. +// Defines: ADataType, BDataType, CDataType, QDataType, AccDataType, +// QuantGroupSize, SelectedKernel, KERNEL_NAME + +// 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; +} + +// GPU architecture is derived from the running device at launch time (see the +// runtime check in dispatcher_run_aquant_gemm) rather than 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; \ + } \ + } + +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 that + * the dispatcher's registry backend uses. + * + * Returns 0 on success. + */ +int dispatcher_initialize() +{ + if(g_initialized) + return 0; + g_initialized = true; + return 0; +} + +/** + * Run AQuantGrouped GEMM: C[M,N] = dequant(A[M,K], AQ[M, ceil(K/gK)]) @ B[K,N] + * + * A, AQ, B, C are host pointers. This function manages device memory internally. + * + * Parameters: + * A, AQ, B, C - host data pointers + * M, N, K - matrix dimensions + * stride_A - leading dimension of A (row-major: K; col-major: M) + * stride_AQ - leading dimension of AQ (row-major: QK_A) + * stride_B - leading dimension of B (col-major: K; row-major: N) + * stride_C - leading dimension of C (row-major: N) + * QK_A - number of K-groups = ceil(K / quant_group_k) + * 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* AQ, + const void* B, + void* C, + int64_t M, + int64_t N, + int64_t K, + int64_t stride_A, + int64_t stride_AQ, + int64_t stride_B, + int64_t stride_C, + int64_t QK_A, + int k_batch, + float* time_ms) +{ + if(!g_initialized) + { + std::cerr << "dispatcher_run_aquant_gemm: not initialized\n"; + return -1; + } + if(!A || !AQ || !B || !C) + { + std::cerr << "dispatcher_run_aquant_gemm: null pointer argument\n"; + return -1; + } + if(M <= 0 || N <= 0 || K <= 0 || QK_A <= 0) + { + std::cerr << "dispatcher_run_aquant_gemm: invalid dimensions\n"; + return -1; + } + + // Derive the GPU architecture from the running device (do not assume one at + // compile time) and reject unsupported archs, per review feedback. + { + 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 matches the compile-time quant group size baked + // into this .so. A mismatch means the AQ device buffer would be allocated with the + // wrong size while the kernel indexes it with different strides. + { + const int64_t expected_QK_A = + (K + static_cast(QuantGroupSize::kK) - 1) / QuantGroupSize::kK; + if(QK_A != expected_QK_A) + { + std::cerr << "dispatcher_run_aquant_gemm: QK_A mismatch. Got " << QK_A << ", expected " + << expected_QK_A << " for K=" << K + << " with QuantGroupSize kK=" << QuantGroupSize::kK << "\n"; + return -1; + } + } + + // This implementation only supports packed (contiguous) layouts. The expected + // leading dimensions depend on the compile-time A/B layouts baked into the kernel. + // A (ALayout): row-major -> stride_A=K ; col-major -> stride_A=M + // B (BLayout): row-major -> stride_B=N ; col-major -> stride_B=K + // AQ: row-major [M, QK_A] -> stride_AQ=QK_A + // C : row-major [M, N] -> stride_C=N (CLayout is always RowMajor) + { + constexpr bool a_row = std::is_same_v; + constexpr bool b_row = std::is_same_v; + const int64_t exp_stride_A = a_row ? K : M; + const int64_t exp_stride_B = b_row ? N : K; + const int64_t exp_stride_AQ = QK_A; + const int64_t exp_stride_C = N; + if(stride_A != exp_stride_A || stride_B != exp_stride_B || stride_AQ != exp_stride_AQ || + stride_C != exp_stride_C) + { + std::cerr << "dispatcher_run_aquant_gemm: non-packed strides are not supported. " + << "Expected stride_A=" << exp_stride_A << " stride_AQ=" << exp_stride_AQ + << " stride_B=" << exp_stride_B << " stride_C=" << exp_stride_C + << ", got stride_A=" << stride_A << " stride_AQ=" << stride_AQ + << " stride_B=" << stride_B << " stride_C=" << stride_C << "\n"; + return -1; + } + } + + const ADataType* A_host = static_cast(A); + const QDataType* AQ_host = static_cast(AQ); + const BDataType* B_host = static_cast(B); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + QDataType* AQ_dev = nullptr; + BDataType* B_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup = [&]() { + if(A_dev) + (void)hipFree(A_dev); + if(AQ_dev) + (void)hipFree(AQ_dev); + if(B_dev) + (void)hipFree(B_dev); + if(C_dev) + (void)hipFree(C_dev); + }; + + // Allocate device buffers. + // A may be a packed type (pk_int4_t): 2 logical values per byte. + // elements_to_bytes(n) handles the packed case via numeric_traits::PackedSize. + if(hipMalloc(&A_dev, elements_to_bytes(M * K)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&AQ_dev, elements_to_bytes(M * QK_A)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&B_dev, elements_to_bytes(K * N)) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMalloc(&C_dev, elements_to_bytes(M * N)) != hipSuccess) + { + cleanup(); + return -1; + } + + // Copy A to device. For pk_int4 A the raw i4x4 values must be permuted for the + // device implementation -- mirrors run_gemm_quant_example.inc:758-763. + if constexpr(std::is_same_v) + { + constexpr bool a_is_row = std::is_same_v; + ck_tile::HostTensor a_h( + ck_tile::host_tensor_descriptor(static_cast(M), + static_cast(K), + static_cast(stride_A), + ck_tile::bool_constant{})); + std::copy(A_host, A_host + M * K, a_h.begin()); + ck_tile::permute_vectors_i4x4_b(a_h); + if(hipMemcpy( + A_dev, a_h.data(), elements_to_bytes(M * K), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + } + else + { + if(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + } + + // Apply AQ preshuffle when required -- mirrors run_gemm_quant_example.inc:746-751. + // APreshuffleQuant reorders AQ in host memory (shuffle_aq) before the device copy + // so the kernel finds the scale values in the interleaved layout it expects. + if constexpr(SelectedKernel::APreshuffleQuant) + { + 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(M), + static_cast(QK_A), + static_cast(QK_A), + ck_tile::bool_constant{} /*row-major*/)); + std::copy(AQ_host, AQ_host + M * 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(M * QK_A), + hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + else + { + if(hipMemcpy( + AQ_dev, AQ_host, elements_to_bytes(M * QK_A), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + } + if(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemset(C_dev, 0, elements_to_bytes(M * N)) != hipSuccess) + { + cleanup(); + return -1; + } + + // Build QuantGemmHostArgs (bq_ptr = nullptr, QK_B = 0, stride_BQ = 0 for AQuant-only) + ck_tile::QuantGemmHostArgs args; + args.a_ptr = A_dev; + args.b_ptr = B_dev; + args.aq_ptr = AQ_dev; + args.bq_ptr = nullptr; + 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; + 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; + + const bool do_time = (time_ms != nullptr); + // When timing is requested use GPU timer with warmup (cold_niters=3, nrepeat=10). + // Otherwise run once with no overhead. + 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); + + 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, elements_to_bytes(M * N), 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 in this .so (always 1: the force-included SelectedKernel). + */ +int dispatcher_get_kernel_count() { return 1; } + +/** + * Release resources. + */ +void dispatcher_cleanup() { g_initialized = false; } + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index d99e0977729d..5360e14d5675 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -441,6 +441,77 @@ def make_bquant_kernel_name( return "_".join(parts) +# ============================================================================ +# 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 aquant codegen will actually emit. + + Mirrors the TiledPermuteN logic in run_gemm_quant_example.inc for AQuant: + TiledPermuteN = (BQuantGroupSize::kN > 1) ? false : GemmConfig::TiledMMAPermuteN + For the AQuant decode configs (GemmConfigQuantDecodeInterwave / + GemmConfigPreshuffleQuantDecode) TiledMMAPermuteN is false (it is only set on + the PreshuffleB configs, which AQuant does not use — PreshuffleB is rejected + for AQuant in run_gemm_example_prec_type). So AQuant always uses the CShuffle + epilogue. The parameters are accepted for signature symmetry with the bquant + helper and to keep the door open for future permute-N aquant configs. + """ + # AQuant configs never enable TiledMMAPermuteN, so the epilogue is always cshuffle. + _ = (tile_n, warp_n, warp_tile_n, quant_group_n) + return "cshuffle" + + +def make_aquant_kernel_name( + variant_key: str, + layout: str, + pipeline: str, + epilogue: str, # ignored — actual epilogue computed via aquant_effective_epilogue + 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_aquant: 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 stay byte-exact. + + Naming convention: + gemm_aquant_{variant}_{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}[_preshufflequant] + + The ``epilogue`` parameter is accepted for call-site compatibility but not used. + """ + effective_epilogue = aquant_effective_epilogue(tile_n, warp_n, warp_tile_n, quant_group_n) + parts = [ + "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"qg{quant_group_m}x{quant_group_n}x{quant_group_k}", + ] + if preshuffle_aquant: + parts.append("preshufflequant") + return "_".join(parts) + + # ============================================================================ # BQuant-specific Type Mappings # ============================================================================ diff --git a/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py new file mode 100644 index 000000000000..4a638983dc9b --- /dev/null +++ b/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py @@ -0,0 +1,637 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +AQuant (A-only quantized) 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 gemm_aquant_ctypes_lib.cpp + +Scope (matches Old-TE gemm_aquant_quantgrouped*.cpp): + dtypes : fp8, bf8, fp8i4 (A=pk_int4), bf8i4 (A=pk_int4) + layouts: rcr, rrr, crr, ccr (non-preshufflequant) + rcr, rrr, crr (preshufflequant -- ccr rejected by Old-TE) + pipeline: compv3 -> AQuantGemmPipelineAgBgCrMem (non-preshufflequant) + AQuantGemmPipelineAgBgCrCompV3 (preshufflequant) + host args = ck_tile::QuantGemmHostArgs (aq_ptr set, bq_ptr = nullptr) + +Naming convention (byte-exact with AQuantKernelConfig.name in gemm_aquant_utils.py): + gemm_aquant_{variant}_{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}[_preshufflequant] + +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/run_gemm_quant_example.inc + example/ck_tile/38_block_scale_gemm/gemm_utils.hpp + (GemmConfigQuantDecodeInterwave, GemmConfigPreshuffleQuantDecode) +""" + +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_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 maps a variant key to the (A, B, C, Q) CK types. +# For AQuant the *A* matrix is the quantized operand: +# fp8/bf8 : A and B are the same 8-bit float, Q (A-scale) is float +# fp8i4 : A = pk_int4 (quantized weight), B = fp8, Q (A-scale) = fp8 +# bf8i4 : A = pk_int4 (quantized weight), B = bf8, Q (A-scale) = bf8 +# Matches GemmQuantTypeConfig in the Old-TE aquant .cpp files. +# ============================================================================= + +AQUANT_VARIANTS: Dict[str, Dict[str, str]] = { + "fp8": { + "ck_a": "ck_tile::fp8_t", + "ck_b": "ck_tile::fp8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + # PrecType passed to GemmConfig drives K_Tile = 256/sizeof(PrecType). + "prec_bytes": 1, + }, + "bf8": { + "ck_a": "ck_tile::bf8_t", + "ck_b": "ck_tile::bf8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "float", + "ck_acc": "float", + "prec_bytes": 1, + }, + "fp8i4": { + "ck_a": "ck_tile::pk_int4_t", + "ck_b": "ck_tile::fp8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "ck_tile::fp8_t", + "ck_acc": "float", + # GemmConfig -> K_Tile = 256/sizeof(fp8_t) = 256. + "prec_bytes": 1, + }, + "bf8i4": { + "ck_a": "ck_tile::pk_int4_t", + "ck_b": "ck_tile::bf8_t", + "ck_c": "ck_tile::half_t", + "ck_q": "ck_tile::bf8_t", + "ck_acc": "float", + "prec_bytes": 1, + }, +} + +# Layout characters -> CK layout type. +AQUANT_LAYOUT_TO_CK = { + "r": "ck_tile::tensor_layout::gemm::RowMajor", + "c": "ck_tile::tensor_layout::gemm::ColumnMajor", +} + +# The 3-char layout tag encodes (A, B, C). C is always RowMajor for quant kernels +# (static_assert in gemm_calc_quant). The AQ layout is derived from the Old-TE +# run_gemm_example_prec_type dispatch table: +# rcr : A=R B=C C=R -> AQ=R BQ=C +# rrr : A=R B=R C=R -> AQ=R BQ=C +# crr : A=C B=R C=R -> AQ=R BQ=C +# ccr : A=C B=C C=R -> AQ=C BQ=C (non-preshufflequant only) +AQUANT_AQ_LAYOUT = { + "rcr": "r", + "rrr": "r", + "crr": "r", + "ccr": "c", +} + +# Pipeline map for AQuant kernels. +# non-preshufflequant -> AQuantGemmPipelineAgBgCrMem (base: BaseGemmPipelineAgBgCrMem) +# preshufflequant -> AQuantGemmPipelineAgBgCrCompV3(base: BaseGemmPipelineAgBgCrCompV3) +AQUANT_PIPELINE_MAP = { + "mem": "ck_tile::AQuantGemmPipelineAgBgCrMem", + "compv3": "ck_tile::AQuantGemmPipelineAgBgCrCompV3", +} + +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", "rrr", "crr", "ccr" + scheduler: str # "interwave" (decode) or "intrawave" (preshufflequant) + tile: AQuantTileConfig + quant_group_m: int = 1 + quant_group_n: int = 1 + quant_group_k: int = 128 + preshuffle_aquant: bool = False + 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 + + @property + def pipeline_key(self) -> str: + """Pipeline map key: preshufflequant -> compv3, else mem.""" + return "compv3" if self.preshuffle_aquant else "mem" + + @property + def name(self) -> str: + t = self.tile + return make_aquant_kernel_name( + variant_key=self.variant_key, + layout=self.layout, + pipeline=self.pipeline_key, + epilogue="cshuffle", + 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_aquant=self.preshuffle_aquant, + ) + + +# ============================================================================= +# 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"] + 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]] + layout_aq_ck = AQUANT_LAYOUT_TO_CK[AQUANT_AQ_LAYOUT[spec.layout]] + # BQ layout is unused for AQuant-only (bq_ptr=nullptr); Old-TE passes ColumnMajor. + layout_bq_ck = AQUANT_LAYOUT_TO_CK["c"] + + pipeline_key = spec.pipeline_key + pipeline_ck = AQUANT_PIPELINE_MAP[pipeline_key] + base_pipeline_ck = AQUANT_BASE_PIPELINE_MAP[pipeline_key] + 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_aquant = str(spec.preshuffle_aquant).lower() + double_smem_buffer = str(spec.double_smem_buffer).lower() + + # AQuant configs never enable TiledMMAPermuteN (see aquant_effective_epilogue), + # so the epilogue is always CShuffle. Kept as a computed value for parity with + # the bquant codegen and to fail loudly if the assumption ever changes. + use_permute_n_epilogue = ( + aquant_effective_epilogue(t.tile_n, t.warp_n, t.warp_tile_n, spec.quant_group_n) + == "permute_n" + ) + assert not use_permute_n_epilogue, "AQuant does not support PermuteN epilogue" + + 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 AQuant (A-only quantized) GEMM kernel header. +// DO NOT EDIT -- regenerate via unified_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}; +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}; + +// QuantGroupShape> -- same type used for the AQ slot in the +// pipeline template; BQ is disabled via bq_ptr=nullptr at runtime for AQuant-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 ck_tile::index_t GroupSizeK = {spec.quant_group_k}; + + static constexpr bool kPadM = {pad_m}; + static constexpr bool kPadN = {pad_n}; + static constexpr bool kPadK = {pad_k}; + static constexpr bool APreshuffleQuant = {preshuffle_aquant}; + static constexpr bool BPreshuffleQuant = false; + static constexpr bool PreshuffleB = false; + static constexpr bool TransposeC = false; + 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) + {{ + // 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, WarpTileK); + + 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_) {{ + // GemmAQuantPipelineProblem + using PipelineProblem = ck_tile::GemmAQuantPipelineProblem< + ADataType, + QDataType, + BDataType, + AccDataType, + TileShape, + GemmTraits, + QuantGroupSize, + TransposeC, + ADataType, // ComputeDataType + {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; +using ALayout = {ns}::ALayout; +using BLayout = {ns}::BLayout; +using CLayout = {ns}::CLayout; +constexpr ck_tile::index_t GroupSizeK = {ns}::{struct}::GroupSizeK; +#endif // CK_TILE_SINGLE_KERNEL_INCLUDE +""" + + +# ============================================================================= +# Config sweep +# ============================================================================= + + +def _k_tile_for(variant_key: str) -> int: + """K_Tile = 256 / sizeof(PrecType) for the decode config. + + PrecType is the 8-bit float weight type (fp8_t / bf8_t), so sizeof==1 and + K_Tile==256 for every supported variant (fp8, bf8, fp8i4, bf8i4). + """ + return 256 // AQUANT_VARIANTS[variant_key]["prec_bytes"] + + +def _default_config() -> dict: + """Default sweep config matching GemmConfigQuantDecodeInterwave tile defaults. + + Non-preshufflequant decode kernels for every dtype x layout Old-TE supports. + """ + return { + "variant_keys": ["fp8", "bf8", "fp8i4", "bf8i4"], + "layouts": ["rcr", "rrr", "crr", "ccr"], + "scheduler": "interwave", + "tile_configs": [ + # GemmConfigQuantDecodeInterwave: M=16, N=64, K=256/sizeof(PrecType)=256 + # WarpTileK=128: get_k_warp_tile() on gfx950 = 128 + {"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_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_aquant": False, + } + + +def _build_specs(config: dict) -> List[AQuantKernelSpec]: + specs = [] + preshuffle_aquant = config.get("preshuffle_aquant", False) + # Decode config uses Interwave; preshufflequant uses Intrawave (GemmConfigBase default). + default_scheduler = "intrawave" if preshuffle_aquant else "interwave" + scheduler = config.get("scheduler", default_scheduler) + 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) + + 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 + if layout not in AQUANT_AQ_LAYOUT: + log.warning("Unsupported layout %s -- skipping", layout) + continue + # Old-TE rejects the ccr layout for the preshufflequant path. + if preshuffle_aquant and layout == "ccr": + log.warning("ccr layout is unsupported for preshufflequant -- 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, + 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_aquant=preshuffle_aquant, + 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, + )) + + 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="AQuant (A-only quantized) GEMM kernel header generator" + ) + parser.add_argument("--output-dir", type=Path, + help="Directory to write generated .hpp files " + "(required unless --list-names)") + 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 + + if args.output_dir is None: + parser.error("--output-dir is required unless --list-names is given") + + 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/python/gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py new file mode 100644 index 000000000000..626367ecea0d --- /dev/null +++ b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py @@ -0,0 +1,827 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +AQuant (A-only quantized) GEMM 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 (self-contained, do not import from gemm_utils.py): + setup_multiple_aquant_dispatchers(configs, ...) + codegen -> hipcc -> list of .so paths, all in parallel + +Usage (end-to-end): + configs = [AQuantKernelConfig(variant_key="fp8", layout="rcr", ...)] + so_paths = setup_multiple_aquant_dispatchers(configs, output_dir=Path("/tmp/aq")) + runner = AQuantGpuGemmRunner(so_paths[0]) + result = runner.run(A, AQ, B, 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_gemm_aquant_codegen.py" +_CTYPES_LIB_SRC = Path(__file__).parent.parent / "bindings" / "ctypes" / "gemm_aquant_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_aquant_kernel_name # noqa: E402 + +# NEVER default to gfx942 -- arch must be detected or explicitly supplied. +_DEFAULT_HIPCC = "hipcc" +_SUPPORTED_ARCHS = ("gfx90a", "gfx942", "gfx950") + +# Layout tags AQuant supports, and whether preshufflequant is allowed for each. +# Derived from run_gemm_example_prec_type in run_gemm_quant_example.inc: +# rcr, rrr, crr : both decode and preshufflequant +# ccr : decode only (rejected for APreshuffleQuant) +_LAYOUTS_DECODE = ("rcr", "rrr", "crr", "ccr") +_LAYOUTS_PRESHUFFLEQUANT = ("rcr", "rrr", "crr") + +# fp8/bf8 A/B/Q dtype meta for the four variants (A is the quantized operand). +_VARIANT_META: Dict[str, Dict[str, str]] = { + "fp8": {"a": "fp8", "b": "fp8", "q": "float"}, + "bf8": {"a": "bf8", "b": "bf8", "q": "float"}, + "fp8i4": {"a": "pk_int4", "b": "fp8", "q": "fp8"}, + "bf8i4": {"a": "pk_int4", "b": "bf8", "q": "bf8"}, +} + + +def _validate_arch(arch: str) -> str: + """Return arch if supported; raise on unknown arch (never silently default).""" + if not any(arch.startswith(a) for a in _SUPPORTED_ARCHS): + raise ValueError( + f"Unsupported GPU architecture '{arch}' for AQuant " + f"(supported: {', '.join(_SUPPORTED_ARCHS)})" + ) + return arch + + +# ============================================================================= +# 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_gemm_aquant_codegen.py + emits as KERNEL_NAME, ensuring the Python side and compiled .so always agree. + """ + + variant_key: str # "fp8", "bf8", "fp8i4", "bf8i4" + layout: str # "rcr", "rrr", "crr", "ccr" + scheduler: str # "interwave" (decode) or "intrawave" (preshufflequant) + + 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_aquant: bool = False + double_smem_buffer: bool = False + k_block_per_cu: int = 1 + + # No default arch: caller must pass a valid one (or use _detect_gpu_arch()). + gfx_arch: str = "gfx950" + + @property + def pipeline_key(self) -> str: + """Pipeline map key echoed in the kernel name: preshufflequant -> compv3, else mem.""" + return "compv3" if self.preshuffle_aquant else "mem" + + @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_key, + epilogue="cshuffle", + 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_aquant=self.preshuffle_aquant, + ) + + def to_codegen_config(self) -> dict: + """Produce the JSON config dict consumed by unified_gemm_aquant_codegen.py.""" + return { + "variant_keys": [self.variant_key], + "layouts": [self.layout], + "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_aquant": self.preshuffle_aquant, + "double_smem_buffer": self.double_smem_buffer, + "k_block_per_cu": self.k_block_per_cu, + } + + +# ============================================================================= +# 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: ceil(K / quant_group_k).""" + return (self.K + self.quant_group_k - 1) // self.quant_group_k + + +# ============================================================================= +# AQuantGemmResult +# ============================================================================= + + +@dataclass +class AQuantGemmResult: + C: object # numpy array + 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, AQ, B, C, M, N, K, + stride_A, stride_AQ, stride_B, stride_C, + QK_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, # AQ + ctypes.c_void_p, # B + 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_AQ + ctypes.c_int64, # stride_B + ctypes.c_int64, # stride_C + ctypes.c_int64, # QK_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, + AQ, + B, + C, + M: int, + N: int, + K: int, + stride_A: int, + stride_AQ: int, + stride_B: int, + stride_C: int, + QK_A: int, + k_batch: int = 1, + ) -> Tuple[int, float]: + """ + Call dispatcher_run_aquant_gemm with ctypes-wrapped pointers. + + A, AQ, B, C must be numpy arrays (C-contiguous, packed). + Returns (status, time_ms). + """ + import numpy as np + + A = np.ascontiguousarray(A) + AQ = np.ascontiguousarray(AQ) + B = np.ascontiguousarray(B) + 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), + AQ.ctypes.data_as(ctypes.c_void_p), + B.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_AQ), + ctypes.c_int64(stride_B), + ctypes.c_int64(stride_C), + ctypes.c_int64(QK_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, AQ, B; allocates C; returns AQuantGemmResult. + """ + + def __init__(self, so_path: Path, layout: str = "rcr"): + self._lib = AQuantDispatcherLib(so_path) + self._layout = layout + + @property + def kernel_name(self) -> str: + return self._lib.get_kernel_name() + + def run(self, A, AQ, B, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmResult: + """ + Run AQuantGrouped GEMM. + + A shape: (M, K) dtype: fp8/bf8/pk_int4 + AQ shape: (M, QK_A) dtype: float/fp8/bf8 + B shape: (K, N) dtype: fp8/bf8 + c_dtype numpy dtype for the output C buffer. Defaults to np.float16 + (correct for all supported aquant variants, whose CDataType is half). + Returns AQuantGemmResult with C shape (M, N). + + Stride conventions follow the compiled kernel's compile-time A/B layouts, + which are encoded in the layout tag (A, B, C). C is always RowMajor. + """ + import numpy as np + + M, N, K = problem.M, problem.N, problem.K + QK_A = problem.QK_A + + if c_dtype is None: + c_dtype = np.float16 + + # Output buffer -- dtype must match the compiled kernel's CDataType (half). + C = np.zeros((M, N), dtype=c_dtype) + + # Packed strides derived from the layout tag. + a_char, b_char, _c_char = self._layout[0], self._layout[1], self._layout[2] + stride_A = K if a_char == "r" else M # A row-major -> K, col-major -> M + stride_AQ = QK_A # AQ is row-major [M, QK_A] + stride_B = N if b_char == "r" else K # B row-major -> N, col-major -> K + stride_C = N # C is row-major [M, N] + + rc, time_ms = self._lib.run( + A=A, AQ=AQ, B=B, C=C, + M=M, N=N, K=K, + stride_A=stride_A, + stride_AQ=stride_AQ, + stride_B=stride_B, + stride_C=stride_C, + QK_A=QK_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 (self-contained, do not call ctypes_utils.py) +# ============================================================================= + + +def _detect_gpu_arch() -> str: + """Detect current GPU arch via rocm_agent_enumerator. Raises if none found.""" + 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 _validate_arch(line) + except Exception as e: + raise RuntimeError(f"Could not detect GPU arch via rocm_agent_enumerator: {e}") + raise RuntimeError( + "No supported GPU architecture detected; pass gfx_arch explicitly " + f"(supported: {', '.join(_SUPPORTED_ARCHS)})" + ) + + +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_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 (compile then link). + + Two-step build: + 1. Compile to a .o object file. + 2. Link the .o into a shared .so (no dispatcher static lib needed; + the AQuant ctypes lib does not use the registry or dispatcher). + + Returns True on success. + """ + ck_include = _get_ck_include_dir() + + obj_path = so_path.with_suffix(".o") + + # Arch-specific defines: gfx950 uses OCP fp8 (not FNUZ) and native MX support. + # These mirror the CMakeLists.txt definitions normally injected by CMake but + # absent in the standalone hipcc build path. + arch_defines = [] + if "gfx12" in gfx_arch or "gfx950" in gfx_arch: + arch_defines += ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] + if "gfx950" in gfx_arch: + arch_defines += ["-DCK_USE_NATIVE_MX_SUPPORT", "-DCK_GFX950_SUPPORT"] + + compile_cmd = [hipcc, "-c", "-fPIC", "-O3", "-std=c++17", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", "-w", + f"--offload-arch={gfx_arch}", + f"-DGFX_ARCH=\"{gfx_arch}\"", + *arch_defines, + "-include", str(hpp_path), + str(_CTYPES_LIB_SRC), + "-o", str(obj_path)] + + if ck_include: + compile_cmd += [f"-I{ck_include}"] + + # NOTE: dispatcher/include is intentionally excluded here. It pulls in + # generated_tile_backend.hpp which instantiates SelectedKernel::launch(GemmHostArgs&), + # conflicting with the AQuant kernel's launch(QuantGemmHostArgs&). The AQuant ctypes + # lib only needs the main CK include path. + + if extra_include_dirs: + for d in extra_include_dirs: + compile_cmd += [f"-I{d}"] + + log.debug("Compiling %s:\n %s", so_path.name, " ".join(compile_cmd)) + + try: + result = subprocess.run(compile_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 + except subprocess.TimeoutExpired: + log.error("Compile timed out for %s", so_path.name) + return False + + link_cmd = [hipcc, "-shared", "-fPIC", + f"--offload-arch={gfx_arch}", "--hip-link", + str(obj_path), "-o", str(so_path)] + + log.debug("Linking %s:\n %s", so_path.name, " ".join(link_cmd)) + + try: + result = subprocess.run(link_cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + log.error("Link failed for %s:\n%s", so_path.name, result.stderr[-2000:]) + obj_path.unlink(missing_ok=True) + return False + except subprocess.TimeoutExpired: + log.error("Link timed out for %s", so_path.name) + obj_path.unlink(missing_ok=True) + return False + + obj_path.unlink(missing_ok=True) + return True + + +# ============================================================================= +# 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. No GPU is required to call this function. + """ + if not configs: + return [] + + arch = _validate_arch(gfx_arch) if gfx_arch else _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) + + # Deduplicate by name so we don't build the same kernel twice. + 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 + + # 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 AQuant kernels", built, len(configs)) + return results + + +# ============================================================================= +# Sweep expansion: JSON config -> list of AQuantKernelConfig +# ============================================================================= + + +def expand_aquant_sweep( + config_path: str, + gfx_arch: str = "gfx950", +) -> List["AQuantKernelConfig"]: + """Expand an AQuant JSON sweep config into a list of AQuantKernelConfig objects. + + The JSON format mirrors unified_gemm_aquant_codegen.py's _build_specs so the same + config files work for both codegen and Python utils. Every valid + (variant, layout, tile, quant_group) combination produces one AQuantKernelConfig; + duplicates (by .name) are collapsed. + """ + import itertools + + with open(config_path) as f: + cfg = json.load(f) + + preshuffle_aquant = cfg.get("preshuffle_aquant", False) + default_scheduler = "intrawave" if preshuffle_aquant else "interwave" + scheduler = cfg.get("scheduler", default_scheduler) + double_smem_buffer = cfg.get("double_smem_buffer", False) + k_block_per_cu = cfg.get("k_block_per_cu", 1) + + allowed_layouts = _LAYOUTS_PRESHUFFLEQUANT if preshuffle_aquant else _LAYOUTS_DECODE + + configs: List[AQuantKernelConfig] = [] + seen: set = set() + + for variant_key, layout, tile_dict, qg in itertools.product( + cfg.get("variant_keys", ["fp8"]), + cfg.get("layouts", ["rcr"]), + cfg.get("tile_configs", []), + cfg.get("quant_groups", [{"quant_group_m": 1, "quant_group_n": 1, "quant_group_k": 128}]), + ): + if layout not in allowed_layouts: + log.warning("Skipping unsupported layout %s (preshufflequant=%s)", + layout, preshuffle_aquant) + continue + c = AQuantKernelConfig( + variant_key=variant_key, + layout=layout, + scheduler=scheduler, + 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"], + 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_aquant=preshuffle_aquant, + double_smem_buffer=double_smem_buffer, + k_block_per_cu=k_block_per_cu, + gfx_arch=gfx_arch, + ) + if c.name not in seen: + seen.add(c.name) + configs.append(c) + + return configs + + +# ============================================================================= +# Convenience: default configs (match GemmConfigQuantDecodeInterwave tile defaults) +# +# tile = 16x64x256, warp = 1x4x1, warp_tile = 16x16x{K_warp} +# fp8/bf8 : get_k_warp_tile() on gfx950 = 128 -> warp_tile_k=128 +# fp8i4/bf8i4: PrecType is the 8-bit float weight -> warp_tile_k=128 as well +# (A being pk_int4 does not change the weight-driven K warp tile) +# ============================================================================= + + +def _decode_config( + variant_key: str, + warp_tile_k: int, + quant_group_k: int, + quant_group_n: int, + layout: str, + gfx_arch: str, +) -> AQuantKernelConfig: + return AQuantKernelConfig( + variant_key=variant_key, + layout=layout, + scheduler="interwave", + 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=warp_tile_k, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + preshuffle_aquant=False, + gfx_arch=gfx_arch, + ) + + +def default_fp8_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """Default fp8 AQuant decode config (GemmConfigQuantDecodeInterwave).""" + return _decode_config("fp8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def default_bf8_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """Default bf8 AQuant decode config (GemmConfigQuantDecodeInterwave).""" + return _decode_config("bf8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def default_fp8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """Default fp8i4 AQuant decode config (A=pk_int4, B=fp8, Q=fp8).""" + return _decode_config("fp8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def default_bf8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """Default bf8i4 AQuant decode config (A=pk_int4, B=bf8, Q=bf8).""" + return _decode_config("bf8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def _preshufflequant_config( + variant_key: str, + warp_tile_k: int, + quant_group_k: int, + quant_group_n: int, + layout: str, + gfx_arch: str, +) -> AQuantKernelConfig: + return AQuantKernelConfig( + variant_key=variant_key, + layout=layout, + 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=warp_tile_k, + quant_group_m=1, + quant_group_n=quant_group_n, + quant_group_k=quant_group_k, + preshuffle_aquant=True, + gfx_arch=gfx_arch, + ) + + +def default_fp8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", + gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """fp8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode).""" + return _preshufflequant_config("fp8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def default_bf8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", + gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """bf8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode).""" + return _preshufflequant_config("bf8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def default_fp8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", + gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """fp8i4 AQuant preshufflequant config (A=pk_int4, B=fp8, Q=fp8).""" + return _preshufflequant_config("fp8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + + +def default_bf8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, + layout: str = "rcr", + gfx_arch: str = "gfx950") -> AQuantKernelConfig: + """bf8i4 AQuant preshufflequant config (A=pk_int4, B=bf8, Q=bf8).""" + return _preshufflequant_config("bf8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) diff --git a/projects/composablekernel/dispatcher/scripts/aquant_selftest.py b/projects/composablekernel/dispatcher/scripts/aquant_selftest.py new file mode 100644 index 000000000000..187f76491f48 --- /dev/null +++ b/projects/composablekernel/dispatcher/scripts/aquant_selftest.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +AQuant (A-only quantized) GEMM bridge self-test / default-config runner. + +No GPU required for the default (--codegen-only) mode: + 1. Verify AQuantKernelConfig.name (utils) == codegen KERNEL_NAME (byte-exact). + 2. Generate every default-config kernel header and assert the emitted + `constexpr const char* KERNEL_NAME` matches the config .name. + 3. Print the dtype x layout coverage matrix. + +With --build (GPU/hipcc host required) it additionally compiles each default +config into a .so via the Python bridge and reports pass/fail per kernel. + +Usage: + python3 aquant_selftest.py # codegen + name-parity only + python3 aquant_selftest.py --build # also hipcc-compile each .so + python3 aquant_selftest.py --arch gfx942 # target arch for --build + python3 aquant_selftest.py --preshufflequant +""" + +import argparse +import re +import sys +import tempfile +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(_HERE.parent / "python")) +sys.path.insert(0, str(_HERE.parent / "codegen")) + +from gemm_aquant_utils import ( # noqa: E402 + AQuantKernelConfig, + default_fp8_config, + default_bf8_config, + default_fp8i4_config, + default_bf8i4_config, + default_fp8_preshufflequant_config, + default_bf8_preshufflequant_config, + default_fp8i4_preshufflequant_config, + default_bf8i4_preshufflequant_config, + _generate_aquant_kernel, + setup_multiple_aquant_dispatchers, + _LAYOUTS_DECODE, + _LAYOUTS_PRESHUFFLEQUANT, +) + +_VARIANTS = ["fp8", "bf8", "fp8i4", "bf8i4"] + + +def _decode_configs(arch: str): + fns = { + "fp8": default_fp8_config, + "bf8": default_bf8_config, + "fp8i4": default_fp8i4_config, + "bf8i4": default_bf8i4_config, + } + return [fns[v](layout=lay, gfx_arch=arch) + for v in _VARIANTS for lay in _LAYOUTS_DECODE] + + +def _preshufflequant_configs(arch: str): + fns = { + "fp8": default_fp8_preshufflequant_config, + "bf8": default_bf8_preshufflequant_config, + "fp8i4": default_fp8i4_preshufflequant_config, + "bf8i4": default_bf8i4_preshufflequant_config, + } + return [fns[v](layout=lay, gfx_arch=arch) + for v in _VARIANTS for lay in _LAYOUTS_PRESHUFFLEQUANT] + + +_KERNEL_NAME_RE = re.compile(r'constexpr const char\* KERNEL_NAME = "([^"]+)"') + + +def _check_name_parity(configs, out_dir: Path) -> int: + failures = 0 + for cfg in configs: + hpp = _generate_aquant_kernel(cfg, out_dir) + if hpp is None: + print(f" [FAIL] codegen produced no header for {cfg.name}") + failures += 1 + continue + text = hpp.read_text() + m = _KERNEL_NAME_RE.search(text) + emitted = m.group(1) if m else "" + if emitted != cfg.name: + print(f" [FAIL] name mismatch:\n" + f" utils : {cfg.name}\n" + f" header: {emitted}") + failures += 1 + else: + print(f" [ok] {cfg.name}") + return failures + + +def main() -> int: + ap = argparse.ArgumentParser(description="AQuant GEMM bridge self-test") + ap.add_argument("--build", action="store_true", + help="Also hipcc-compile each default config into a .so") + ap.add_argument("--arch", default="gfx950", help="Target GFX arch (default gfx950)") + ap.add_argument("--preshufflequant", action="store_true", + help="Test the preshufflequant configs instead of decode") + args = ap.parse_args() + + configs = (_preshufflequant_configs(args.arch) if args.preshufflequant + else _decode_configs(args.arch)) + + mode = "preshufflequant" if args.preshufflequant else "decode" + print(f"=== AQuant {mode} self-test (arch={args.arch}) ===") + print(f"dtype x layout coverage: {len(configs)} kernels") + print() + + with tempfile.TemporaryDirectory(prefix="aquant_selftest_") as td: + out_dir = Path(td) + print("--- name parity (codegen KERNEL_NAME vs utils .name) ---") + failures = _check_name_parity(configs, out_dir) + + print() + if failures: + print(f"NAME PARITY FAILED: {failures} mismatch(es)") + return 1 + print(f"NAME PARITY OK: {len(configs)} kernels byte-exact") + + if args.build: + print() + print("--- hipcc build ---") + so_paths = setup_multiple_aquant_dispatchers( + configs, output_dir=out_dir, gfx_arch=args.arch, + ) + built = sum(1 for p in so_paths if p is not None) + for cfg, p in zip(configs, so_paths): + status = "ok" if p is not None else "FAIL" + print(f" [{status}] {cfg.name}") + print() + print(f"BUILD: {built}/{len(configs)} kernels compiled") + if built != len(configs): + return 2 + + print() + print("SELF-TEST PASSED") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 244536926aa856591561e2c1e644c05f8a729701 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:03:10 -0400 Subject: [PATCH 02/11] test(ck-tile): CPU unit tests for aquant dispatcher bridge --- .../dispatcher/tests/test_aquant_bridge.py | 115 ++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 projects/composablekernel/dispatcher/tests/test_aquant_bridge.py diff --git a/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py b/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py new file mode 100644 index 000000000000..8b5a7b49a01a --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU-only unit tests for the aquant (A-only quantized) GEMM TileEngine -> Dispatcher bridge. + +Locks the config name format, the byte-exact codegen<->utils kernel-name contract, the +codegen-JSON projection, and the dtype/pipeline scope (fp8/bf8/fp8i4/bf8i4, decode via the +mem/interwave pipeline and preshufflequant via the compv3/intrawave pipeline) that Old-TE +gemm_aquant_quantgrouped{,_preshufflequant}.cpp register. No GPU, no hipcc required. +""" + +import sys +import unittest +from pathlib import Path + +_DISP = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_DISP / "python")) +sys.path.insert(0, str(_DISP / "codegen")) + +from gemm_aquant_utils import ( # noqa: E402 + AQuantGemmProblem, + default_fp8_config, + default_bf8_config, + default_fp8i4_config, + default_bf8i4_config, + default_fp8_preshufflequant_config, + default_bf8i4_preshufflequant_config, +) +from codegen_common import make_aquant_kernel_name # noqa: E402 + +_DECODE = [ + ("fp8", default_fp8_config), + ("bf8", default_bf8_config), + ("fp8i4", default_fp8i4_config), + ("bf8i4", default_bf8i4_config), +] + + +class TestConfigName(unittest.TestCase): + def test_decode_name_prefix(self): + for variant, ctor in _DECODE: + cfg = ctor() + self.assertTrue(cfg.name.startswith(f"gemm_aquant_{variant}_"), cfg.name) + + def test_name_encodes_tiles(self): + cfg = default_fp8_config() + self.assertIn(f"{cfg.tile_m}x{cfg.tile_n}x{cfg.tile_k}", cfg.name) + self.assertIn(f"{cfg.warp_tile_m}x{cfg.warp_tile_n}x{cfg.warp_tile_k}", cfg.name) + + def test_pipeline_key_reflects_preshuffle(self): + self.assertEqual(default_fp8_config().pipeline_key, "mem") + self.assertEqual(default_fp8_preshufflequant_config().pipeline_key, "compv3") + + +class TestNameContract(unittest.TestCase): + def _assert_contract(self, cfg): + expected = make_aquant_kernel_name( + variant_key=cfg.variant_key, + layout=cfg.layout, + pipeline=cfg.pipeline_key, + epilogue="cshuffle", + scheduler=cfg.scheduler, + tile_m=cfg.tile_m, tile_n=cfg.tile_n, tile_k=cfg.tile_k, + warp_m=cfg.warp_m, warp_n=cfg.warp_n, warp_k=cfg.warp_k, + warp_tile_m=cfg.warp_tile_m, warp_tile_n=cfg.warp_tile_n, + warp_tile_k=cfg.warp_tile_k, + quant_group_m=cfg.quant_group_m, + quant_group_n=cfg.quant_group_n, + quant_group_k=cfg.quant_group_k, + preshuffle_aquant=cfg.preshuffle_aquant, + ) + self.assertEqual(cfg.name, expected) + + def test_decode_contracts(self): + for _, ctor in _DECODE: + self._assert_contract(ctor()) + + def test_preshufflequant_contracts(self): + self._assert_contract(default_fp8_preshufflequant_config()) + self._assert_contract(default_bf8i4_preshufflequant_config()) + + +class TestScope(unittest.TestCase): + def test_decode_variants(self): + self.assertEqual([v for v, _ in _DECODE], + [ctor().variant_key for _, ctor in _DECODE]) + + def test_decode_uses_mem_pipeline(self): + for _, ctor in _DECODE: + self.assertFalse(ctor().preshuffle_aquant) + + def test_preshufflequant_flag(self): + self.assertTrue(default_fp8_preshufflequant_config().preshuffle_aquant) + + +class TestCodegenProjection(unittest.TestCase): + def test_to_codegen_config_roundtrip(self): + cfg = default_fp8_config() + d = cfg.to_codegen_config() + self.assertEqual(d["variant_keys"], [cfg.variant_key]) + self.assertEqual(d["layouts"], [cfg.layout]) + self.assertEqual(d["quant_groups"][0]["quant_group_k"], cfg.quant_group_k) + self.assertEqual(d["preshuffle_aquant"], cfg.preshuffle_aquant) + + +class TestProblem(unittest.TestCase): + def test_problem_constructs(self): + p = AQuantGemmProblem(M=128, N=256, K=512) + self.assertEqual((p.M, p.N, p.K), (128, 256, 512)) + + +if __name__ == "__main__": + unittest.main() From 20623a7d8effdcc9e045af1e828bcf5adf452a2d Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:15:35 -0400 Subject: [PATCH 03/11] fix(ck-tile): round-2 aquant bridge fixes (ccr AQ stride, CMake includes) Reviewer feedback on draft PR #9980: - BLOCKING: derive stride_AQ from the AQ layout on both C++ and Python sides. The ccr layout makes the A-scale tensor column-major, so its leading dim is M (row count), not QK_A. Previously stride_AQ was hardcoded to QK_A, silently mis-indexing the ccr A-scale. Codegen now also emits a global `using AQLayout` so the ctypes lib can key off it. - SHOULD-FIX: build dispatcher_gemm_aquant_lib with a main-include-only include set instead of add_ctypes_library, which pulled in dispatcher/include and the generated_tile_backend.hpp run()/launch() conflict the Python build path deliberately avoids. - SHOULD-FIX: pass warp_tile_k through the default_*_config helpers instead of hardcoding 128 (dead named param). - NITs: static_assert ties the shuffle_aq row-major assumption to the ccr preshufflequant exclusion; CPU tests assert ccr column-major AQ stride, the layout scope, the 28-kernel count, and gfx90a validation. --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 22 ++++++- .../ctypes/gemm_aquant_ctypes_lib.cpp | 41 ++++++++---- .../codegen/unified_gemm_aquant_codegen.py | 3 + .../dispatcher/python/gemm_aquant_utils.py | 64 ++++++++++++++----- .../dispatcher/tests/test_aquant_bridge.py | 59 +++++++++++++++++ 5 files changed, 157 insertions(+), 32 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 47caf221abe5..a12f7ee54d8a 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -287,9 +287,25 @@ if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/gemm_aquant_ctypes_lib.cpp) 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_gemm_aquant_lib - gemm_aquant_ctypes_lib.cpp - KERNEL_HEADER ${AQUANT_GEMM_KERNEL_HEADER} + # NOTE: unlike the other ctypes targets, this one does NOT use + # add_ctypes_library(), because that helper adds ${PROJECT_SOURCE_DIR}/dispatcher/include + # to the include path. That directory pulls in generated_tile_backend.hpp, which + # instantiates SelectedKernel::run(GemmHostArgs&) and conflicts with the AQuant + # kernel's launch(QuantGemmHostArgs&) signature. The Python build path + # (gemm_aquant_utils.py:_compile_aquant_kernel) deliberately excludes + # dispatcher/include for the same reason. Build the target by hand with a + # MAIN-include-only include set so the CMake path matches the Python path. + add_library(dispatcher_gemm_aquant_lib SHARED gemm_aquant_ctypes_lib.cpp) + target_include_directories(dispatcher_gemm_aquant_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ) + target_link_libraries(dispatcher_gemm_aquant_lib PRIVATE hip::device) + target_compile_options(dispatcher_gemm_aquant_lib PRIVATE + -include ${AQUANT_GEMM_KERNEL_HEADER} + ) + set_target_properties(dispatcher_gemm_aquant_lib PROPERTIES + POSITION_INDEPENDENT_CODE ON + CXX_STANDARD 17 ) # Resolve a single GFX arch string from CMAKE_HIP_ARCHITECTURES if available, # falling back to the cache variable CK_TILE_AQUANT_GFX_ARCH. Never hardcode a diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp index 2229f13dbb01..ec82f5444c92 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp @@ -21,9 +21,12 @@ * * AQuant vs BQuant: here the *A* matrix is the quantized operand. The scale * tensor AQ has shape [M, QK_A] (QK_A = ceil(K/gK)), aq_ptr is set and bq_ptr is - * nullptr. For pk_int4 A the raw values are permuted via permute_vectors_i4x4_b - * before the device copy (mirrors run_gemm_quant_example.inc), and APreshuffleQuant - * kernels shuffle AQ via shuffle_aq. + * nullptr. Its leading dimension follows AQLayout: row-major (rcr/rrr/crr) uses + * stride_AQ=QK_A, while the ccr layout is column-major and uses stride_AQ=M + * (see run_gemm_quant_example.inc get_default_stride, ~line 528). For pk_int4 A the + * raw values are permuted via permute_vectors_i4x4_b before the device copy (mirrors + * run_gemm_quant_example.inc), and APreshuffleQuant kernels shuffle AQ via shuffle_aq + * (row-major only -- ccr is excluded from the preshufflequant path by Old-TE). */ #include @@ -98,7 +101,7 @@ int dispatcher_initialize() * A, AQ, B, C - host data pointers * M, N, K - matrix dimensions * stride_A - leading dimension of A (row-major: K; col-major: M) - * stride_AQ - leading dimension of AQ (row-major: QK_A) + * stride_AQ - leading dimension of AQ (row-major: QK_A; col-major/ccr: M) * stride_B - leading dimension of B (col-major: K; row-major: N) * stride_C - leading dimension of C (row-major: N) * QK_A - number of K-groups = ceil(K / quant_group_k) @@ -174,17 +177,22 @@ int dispatcher_run_aquant_gemm(const void* A, } // This implementation only supports packed (contiguous) layouts. The expected - // leading dimensions depend on the compile-time A/B layouts baked into the kernel. - // A (ALayout): row-major -> stride_A=K ; col-major -> stride_A=M - // B (BLayout): row-major -> stride_B=N ; col-major -> stride_B=K - // AQ: row-major [M, QK_A] -> stride_AQ=QK_A - // C : row-major [M, N] -> stride_C=N (CLayout is always RowMajor) + // leading dimensions depend on the compile-time A/B/AQ layouts baked into the kernel. + // A (ALayout) : row-major -> stride_A=K ; col-major -> stride_A=M + // B (BLayout) : row-major -> stride_B=N ; col-major -> stride_B=K + // AQ (AQLayout): scale tensor [M, QK_A]. row-major -> stride_AQ=QK_A ; + // col-major (ccr) -> stride_AQ=M. Mirrors Old-TE + // get_default_stride(M, QK_A, 0, is_row_major(aq_layout)) in + // run_gemm_quant_example.inc (~line 528): the column-major branch + // returns the row count M, not the K-group count QK_A. + // C (CLayout) : row-major [M, N] -> stride_C=N (CLayout is always RowMajor) { - constexpr bool a_row = std::is_same_v; - constexpr bool b_row = std::is_same_v; + constexpr bool a_row = std::is_same_v; + constexpr bool b_row = std::is_same_v; + constexpr bool aq_row = std::is_same_v; const int64_t exp_stride_A = a_row ? K : M; const int64_t exp_stride_B = b_row ? N : K; - const int64_t exp_stride_AQ = QK_A; + const int64_t exp_stride_AQ = aq_row ? QK_A : M; const int64_t exp_stride_C = N; if(stride_A != exp_stride_A || stride_B != exp_stride_B || stride_AQ != exp_stride_AQ || stride_C != exp_stride_C) @@ -278,6 +286,15 @@ int dispatcher_run_aquant_gemm(const void* A, // so the kernel finds the scale values in the interleaved layout it expects. if constexpr(SelectedKernel::APreshuffleQuant) { + // shuffle_aq assumes a row-major AQ descriptor. This holds because Old-TE + // rejects the ccr (column-major AQ) layout for the preshufflequant path, so + // every APreshuffleQuant kernel this .so can host has AQLayout == RowMajor. + // The assert is gated on APreshuffleQuant so it is only evaluated for kernels + // that actually take this branch (decode-path ccr kernels are unaffected). + static_assert(!SelectedKernel::APreshuffleQuant || + std::is_same_v, + "APreshuffleQuant requires a row-major AQ layout (ccr is excluded " + "from the preshufflequant path); shuffle_aq below assumes row-major"); constexpr int block_aq_k = static_cast(SelectedKernel::TileK) / static_cast(QuantGroupSize::kK); ck_tile::HostTensor aq_h( diff --git a/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py index 4a638983dc9b..2ea9ebf9db20 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py @@ -416,6 +416,9 @@ def generate(self, spec: AQuantKernelSpec) -> str: using ALayout = {ns}::ALayout; using BLayout = {ns}::BLayout; using CLayout = {ns}::CLayout; +// AQ scale-tensor layout: RowMajor for rcr/rrr/crr, ColumnMajor for ccr. +// The ctypes lib derives stride_AQ from this (RowMajor -> QK_A, ColumnMajor -> M). +using AQLayout = {ns}::AQLayout; constexpr ck_tile::index_t GroupSizeK = {ns}::{struct}::GroupSizeK; #endif // CK_TILE_SINGLE_KERNEL_INCLUDE """ diff --git a/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py index 626367ecea0d..f1ed5696bec9 100644 --- a/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py @@ -62,6 +62,23 @@ _LAYOUTS_DECODE = ("rcr", "rrr", "crr", "ccr") _LAYOUTS_PRESHUFFLEQUANT = ("rcr", "rrr", "crr") +# AQ (A-scale) tensor layout per 3-char layout tag, mirroring AQUANT_AQ_LAYOUT in +# unified_gemm_aquant_codegen.py. The scale tensor is column-major only for ccr +# (A=C B=C), row-major for rcr/rrr/crr. A column-major AQ of shape [M, QK_A] has +# leading dimension M (not QK_A) -- see Old-TE get_default_stride in +# run_gemm_quant_example.inc (~line 528). +_LAYOUTS_AQ_COLMAJOR = frozenset({"ccr"}) + + +def _aq_stride(layout: str, M: int, QK_A: int) -> int: + """Leading dimension of the AQ scale tensor for a given layout tag. + + Column-major AQ (ccr) -> M; row-major AQ (rcr/rrr/crr) -> QK_A. Consistent with + the ColumnMajor/RowMajor AQLayout the codegen emits and the exp_stride_AQ the + ctypes lib validates. + """ + return M if layout in _LAYOUTS_AQ_COLMAJOR else QK_A + # fp8/bf8 A/B/Q dtype meta for the four variants (A is the quantized operand). _VARIANT_META: Dict[str, Dict[str, str]] = { "fp8": {"a": "fp8", "b": "fp8", "q": "float"}, @@ -380,7 +397,8 @@ def run(self, A, AQ, B, problem: AQuantGemmProblem, c_dtype=None) -> AQuantGemmR # Packed strides derived from the layout tag. a_char, b_char, _c_char = self._layout[0], self._layout[1], self._layout[2] stride_A = K if a_char == "r" else M # A row-major -> K, col-major -> M - stride_AQ = QK_A # AQ is row-major [M, QK_A] + # AQ [M, QK_A]: row-major -> QK_A, column-major (ccr) -> M. + stride_AQ = _aq_stride(self._layout, M, QK_A) stride_B = N if b_char == "r" else K # B row-major -> N, col-major -> K stride_C = N # C is row-major [M, N] @@ -753,27 +771,31 @@ def _decode_config( def default_fp8_config(quant_group_k: int = 128, quant_group_n: int = 1, - layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + layout: str = "rcr", gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """Default fp8 AQuant decode config (GemmConfigQuantDecodeInterwave).""" - return _decode_config("fp8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _decode_config("fp8", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def default_bf8_config(quant_group_k: int = 128, quant_group_n: int = 1, - layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + layout: str = "rcr", gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """Default bf8 AQuant decode config (GemmConfigQuantDecodeInterwave).""" - return _decode_config("bf8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _decode_config("bf8", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def default_fp8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, - layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + layout: str = "rcr", gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """Default fp8i4 AQuant decode config (A=pk_int4, B=fp8, Q=fp8).""" - return _decode_config("fp8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _decode_config("fp8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def default_bf8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, - layout: str = "rcr", gfx_arch: str = "gfx950") -> AQuantKernelConfig: + layout: str = "rcr", gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """Default bf8i4 AQuant decode config (A=pk_int4, B=bf8, Q=bf8).""" - return _decode_config("bf8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _decode_config("bf8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def _preshufflequant_config( @@ -801,27 +823,35 @@ def _preshufflequant_config( def default_fp8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", - gfx_arch: str = "gfx950") -> AQuantKernelConfig: + gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """fp8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode).""" - return _preshufflequant_config("fp8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _preshufflequant_config("fp8", warp_tile_k, quant_group_k, quant_group_n, + layout, gfx_arch) def default_bf8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", - gfx_arch: str = "gfx950") -> AQuantKernelConfig: + gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """bf8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode).""" - return _preshufflequant_config("bf8", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _preshufflequant_config("bf8", warp_tile_k, quant_group_k, quant_group_n, + layout, gfx_arch) def default_fp8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", - gfx_arch: str = "gfx950") -> AQuantKernelConfig: + gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """fp8i4 AQuant preshufflequant config (A=pk_int4, B=fp8, Q=fp8).""" - return _preshufflequant_config("fp8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _preshufflequant_config("fp8i4", warp_tile_k, quant_group_k, quant_group_n, + layout, gfx_arch) def default_bf8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", - gfx_arch: str = "gfx950") -> AQuantKernelConfig: + gfx_arch: str = "gfx950", + warp_tile_k: int = 128) -> AQuantKernelConfig: """bf8i4 AQuant preshufflequant config (A=pk_int4, B=bf8, Q=bf8).""" - return _preshufflequant_config("bf8i4", 128, quant_group_k, quant_group_n, layout, gfx_arch) + return _preshufflequant_config("bf8i4", warp_tile_k, quant_group_k, quant_group_n, + layout, gfx_arch) diff --git a/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py b/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py index 8b5a7b49a01a..abb62e5f3455 100644 --- a/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py +++ b/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py @@ -21,6 +21,13 @@ from gemm_aquant_utils import ( # noqa: E402 AQuantGemmProblem, + _aq_stride, + _validate_arch, + _LAYOUTS_DECODE, + _LAYOUTS_PRESHUFFLEQUANT, + _LAYOUTS_AQ_COLMAJOR, + _SUPPORTED_ARCHS, + _VARIANT_META, default_fp8_config, default_bf8_config, default_fp8i4_config, @@ -111,5 +118,57 @@ def test_problem_constructs(self): self.assertEqual((p.M, p.N, p.K), (128, 256, 512)) +class TestAQStride(unittest.TestCase): + """ccr must carry the column-major AQ stride (M), all others row-major (QK_A).""" + + M, QK_A = 96, 4 + + def test_ccr_is_column_major(self): + # ccr: A=C B=C -> AQ column-major, leading dim = M (row count), not QK_A. + self.assertIn("ccr", _LAYOUTS_AQ_COLMAJOR) + self.assertEqual(_aq_stride("ccr", self.M, self.QK_A), self.M) + + def test_row_major_layouts_use_qk_a(self): + for layout in ("rcr", "rrr", "crr"): + self.assertNotIn(layout, _LAYOUTS_AQ_COLMAJOR) + self.assertEqual(_aq_stride(layout, self.M, self.QK_A), self.QK_A) + + def test_ccr_stride_differs_from_row_major(self): + # Guards against the round-1 bug where stride_AQ was always QK_A. + self.assertNotEqual( + _aq_stride("ccr", self.M, self.QK_A), + _aq_stride("rcr", self.M, self.QK_A), + ) + + +class TestLayoutScope(unittest.TestCase): + """Lock the layout scope and the 28-kernel (4 variant x 7 layout) count.""" + + def test_decode_layouts(self): + self.assertEqual(_LAYOUTS_DECODE, ("rcr", "rrr", "crr", "ccr")) + + def test_preshufflequant_excludes_ccr(self): + self.assertEqual(_LAYOUTS_PRESHUFFLEQUANT, ("rcr", "rrr", "crr")) + self.assertNotIn("ccr", _LAYOUTS_PRESHUFFLEQUANT) + + def test_full_kernel_count_is_28(self): + # 4 variants x (4 decode layouts + 3 preshufflequant layouts) = 28. + variants = len(_VARIANT_META) + total = variants * (len(_LAYOUTS_DECODE) + len(_LAYOUTS_PRESHUFFLEQUANT)) + self.assertEqual(variants, 4) + self.assertEqual(total, 28) + + +class TestSupportedArchs(unittest.TestCase): + def test_gfx90a_supported_and_validated(self): + self.assertIn("gfx90a", _SUPPORTED_ARCHS) + self.assertEqual(_validate_arch("gfx90a"), "gfx90a") + self.assertEqual(_validate_arch("gfx90a:sramecc+:xnack-"), "gfx90a:sramecc+:xnack-") + + def test_unsupported_arch_raises(self): + with self.assertRaises(ValueError): + _validate_arch("gfx1030") + + if __name__ == "__main__": unittest.main() From cd5e9261f726ec03c183641357e94934c945671d Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Fri, 24 Jul 2026 20:07:38 -0400 Subject: [PATCH 04/11] fix(ck-tile): round-3 aquant bridge arch-aware warp_tile_k The default_*_config helpers hardcoded warp_tile_k=128, which is gfx950-only. On gfx942 (MI300X) the kernel compiles but SILENTLY OUTPUTS ALL-ZEROS (max_rel=1.0), GPU-confirmed for fp8/bf8 across rcr/rrr/crr/ccr. With the correct gfx942 value (warp_tile_k=32) the kernels are bit-exact (max_rel=0.0) and at parity (median -10.4%, 6/6 within +/-15%). Make warp_tile_k arch-derived, mirroring ck_tile::get_k_warp_tile(). All four AQuant variants (fp8, bf8, fp8i4, bf8i4) instantiate the GEMM config with an 8-bit float PrecType (GemmConfig; the pk_int4 A operand does not drive the K warp tile), so: gfx950 -> 128 (decode and preshufflequant) gfx942/other, decode -> 32 gfx942/other, preshufflequant -> 64 The value flows into the byte-exact .name. Mirrors the round-2 fixes on tensor_quant (fp8_warp_tile_k_for_arch) and rowcolquant (_warp_tile_k_for). Also make the standalone codegen default sweep arch-aware via --gfx-arch, and extend the CPU unit tests to assert the per-arch/per-dtype/per-pipeline values. --- .../codegen/unified_gemm_aquant_codegen.py | 33 +++++- .../dispatcher/python/gemm_aquant_utils.py | 106 ++++++++++++++---- .../dispatcher/tests/test_aquant_bridge.py | 68 +++++++++++ 3 files changed, 184 insertions(+), 23 deletions(-) diff --git a/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py index 2ea9ebf9db20..68b96ad9b82d 100644 --- a/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_gemm_aquant_codegen.py @@ -438,10 +438,31 @@ def _k_tile_for(variant_key: str) -> int: return 256 // AQUANT_VARIANTS[variant_key]["prec_bytes"] -def _default_config() -> dict: +def _warp_tile_k_for_arch(gfx_arch: str, preshuffle_aquant: bool = False) -> int: + """Arch-derived WarpTileK for AQuant with M_Warp_Tile=16. + + Every AQuant variant (fp8, bf8, fp8i4, bf8i4) instantiates the GEMM config with + an 8-bit float PrecType (fp8_t/bf8_t; the pk_int4 A operand does not drive the K + warp tile -- see gemm_aquant_quantgrouped{,_preshufflequant}.cpp GemmConfig). + Mirrors ck_tile::get_k_warp_tile() + (include/ck_tile/ops/gemm/pipeline/tile_gemm_shape.hpp): + gfx950 -> 128 (decode and preshufflequant) + gfx942/other, decode -> 32 + gfx942/other, preshufflequant -> 64 + Using 128 on gfx942 compiles but produces all-zeros output (no valid 16x16x128 + fp8/bf8 warp-gemm on gfx942). + """ + if "gfx950" in gfx_arch: + return 128 + return 64 if preshuffle_aquant else 32 + + +def _default_config(gfx_arch: str = "gfx950") -> dict: """Default sweep config matching GemmConfigQuantDecodeInterwave tile defaults. Non-preshufflequant decode kernels for every dtype x layout Old-TE supports. + WarpTileK is arch-derived (get_k_warp_tile() = 128 on gfx950, + 32 on gfx942 for the decode path). """ return { "variant_keys": ["fp8", "bf8", "fp8i4", "bf8i4"], @@ -449,10 +470,10 @@ def _default_config() -> dict: "scheduler": "interwave", "tile_configs": [ # GemmConfigQuantDecodeInterwave: M=16, N=64, K=256/sizeof(PrecType)=256 - # WarpTileK=128: get_k_warp_tile() on gfx950 = 128 {"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}, + "warp_tile_m": 16, "warp_tile_n": 16, + "warp_tile_k": _warp_tile_k_for_arch(gfx_arch, preshuffle_aquant=False)}, ], "quant_groups": [ {"quant_group_m": 1, "quant_group_n": 1, "quant_group_k": 128}, @@ -604,6 +625,10 @@ def main() -> int: help="Inline JSON config string") parser.add_argument("--no-parallel", action="store_true", help="Disable parallel generation") + parser.add_argument("--gfx-arch", type=str, default="gfx950", + help="Target GPU arch for the built-in default sweep's " + "arch-derived WarpTileK (128 on gfx950, 32/64 on gfx942). " + "Ignored when --config/--config-json is given.") parser.add_argument("--list-names", action="store_true", help="Print kernel names that would be generated and exit") args = parser.parse_args() @@ -620,7 +645,7 @@ def main() -> int: cfg = json.load(f) if args.list_names: - specs = _build_specs(cfg or _default_config()) + specs = _build_specs(cfg or _default_config(args.gfx_arch)) for s in specs: print(s.name) return 0 diff --git a/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py index f1ed5696bec9..45fbf4ac0711 100644 --- a/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py @@ -741,12 +741,36 @@ def expand_aquant_sweep( # Convenience: default configs (match GemmConfigQuantDecodeInterwave tile defaults) # # tile = 16x64x256, warp = 1x4x1, warp_tile = 16x16x{K_warp} -# fp8/bf8 : get_k_warp_tile() on gfx950 = 128 -> warp_tile_k=128 -# fp8i4/bf8i4: PrecType is the 8-bit float weight -> warp_tile_k=128 as well -# (A being pk_int4 does not change the weight-driven K warp tile) +# warp_tile_k is ARCH-DERIVED (never hardcoded); see _warp_tile_k_for below. # ============================================================================= +def _warp_tile_k_for(gfx_arch: str, preshuffle_aquant: bool = False) -> int: + """Arch-derived K warp-tile, mirroring ck_tile::get_k_warp_tile(). + + (tile_gemm_shape.hpp:104-136, M_Warp_Tile=16, non-WMMA path.) For AQuant every + variant -- fp8, bf8, fp8i4, bf8i4 -- instantiates the GEMM config with an 8-bit + float PrecType (fp8_t or bf8_t; the pk_int4 A operand does not drive the K warp + tile -- see gemm_aquant_quantgrouped{,_preshufflequant}.cpp GemmConfig). + So is_8bit_float is always True and warp_tile_k depends only on the arch and the + pipeline (decode = IsFlatMM false, preshufflequant = IsFlatMM true): + + gfx950 (CK_GFX950_SUPPORT): 128 (both decode and preshufflequant) + gfx942/other, decode (IsFlatMM=false): 32 + gfx942/other, preshuf (IsFlatMM=true) : 64 + + This is a BLOCKING correctness constraint, not just a naming detail: a + warp_tile_k=128 fp8/bf8 kernel *compiles* on gfx942 but silently produces + all-zeros output (GPU-confirmed on gfx942 MI300X for this bridge, and earlier + on the sibling tensor_quant/rowcolquant bridges). Old-TE uses 16x16x32 on + gfx942 for decode and is bit-exact there with warp_tile_k=32. + """ + if "gfx950" in gfx_arch: + return 128 + # gfx942 / gfx90a / other: 8-bit-float PrecType, M_Warp_Tile=16 non-WMMA path. + return 64 if preshuffle_aquant else 32 + + def _decode_config( variant_key: str, warp_tile_k: int, @@ -772,29 +796,52 @@ def _decode_config( def default_fp8_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """Default fp8 AQuant decode config (GemmConfigQuantDecodeInterwave).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """Default fp8 AQuant decode config (GemmConfigQuantDecodeInterwave). + + warp_tile_k is arch-derived (get_k_warp_tile()): 128 on gfx950, + 32 on gfx942 (128 silently outputs all-zeros on gfx942). + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=False) return _decode_config("fp8", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def default_bf8_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """Default bf8 AQuant decode config (GemmConfigQuantDecodeInterwave).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """Default bf8 AQuant decode config (GemmConfigQuantDecodeInterwave). + + warp_tile_k is arch-derived: 128 on gfx950, 32 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=False) return _decode_config("bf8", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def default_fp8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """Default fp8i4 AQuant decode config (A=pk_int4, B=fp8, Q=fp8).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """Default fp8i4 AQuant decode config (A=pk_int4, B=fp8, Q=fp8). + + PrecType is fp8_t (GemmConfig), so warp_tile_k is arch-derived like fp8: + 128 on gfx950, 32 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=False) return _decode_config("fp8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) def default_bf8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """Default bf8i4 AQuant decode config (A=pk_int4, B=bf8, Q=bf8).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """Default bf8i4 AQuant decode config (A=pk_int4, B=bf8, Q=bf8). + + PrecType is bf8_t (GemmConfig), so warp_tile_k is arch-derived like bf8: + 128 on gfx950, 32 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=False) return _decode_config("bf8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) @@ -824,8 +871,14 @@ def _preshufflequant_config( def default_fp8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """fp8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """fp8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode). + + warp_tile_k is arch-derived (get_k_warp_tile()): + 128 on gfx950, 64 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=True) return _preshufflequant_config("fp8", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) @@ -833,8 +886,13 @@ def default_fp8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: def default_bf8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """bf8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """bf8 AQuant preshufflequant config (GemmConfigPreshuffleQuantDecode). + + warp_tile_k is arch-derived: 128 on gfx950, 64 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=True) return _preshufflequant_config("bf8", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) @@ -842,8 +900,13 @@ def default_bf8_preshufflequant_config(quant_group_k: int = 128, quant_group_n: def default_fp8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """fp8i4 AQuant preshufflequant config (A=pk_int4, B=fp8, Q=fp8).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """fp8i4 AQuant preshufflequant config (A=pk_int4, B=fp8, Q=fp8). + + PrecType is fp8_t, so warp_tile_k is arch-derived: 128 on gfx950, 64 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=True) return _preshufflequant_config("fp8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) @@ -851,7 +914,12 @@ def default_fp8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n def default_bf8i4_preshufflequant_config(quant_group_k: int = 128, quant_group_n: int = 1, layout: str = "rcr", gfx_arch: str = "gfx950", - warp_tile_k: int = 128) -> AQuantKernelConfig: - """bf8i4 AQuant preshufflequant config (A=pk_int4, B=bf8, Q=bf8).""" + warp_tile_k: Optional[int] = None) -> AQuantKernelConfig: + """bf8i4 AQuant preshufflequant config (A=pk_int4, B=bf8, Q=bf8). + + PrecType is bf8_t, so warp_tile_k is arch-derived: 128 on gfx950, 64 on gfx942. + """ + if warp_tile_k is None: + warp_tile_k = _warp_tile_k_for(gfx_arch, preshuffle_aquant=True) return _preshufflequant_config("bf8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) diff --git a/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py b/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py index abb62e5f3455..9de42c625fde 100644 --- a/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py +++ b/projects/composablekernel/dispatcher/tests/test_aquant_bridge.py @@ -23,6 +23,7 @@ AQuantGemmProblem, _aq_stride, _validate_arch, + _warp_tile_k_for, _LAYOUTS_DECODE, _LAYOUTS_PRESHUFFLEQUANT, _LAYOUTS_AQ_COLMAJOR, @@ -33,6 +34,8 @@ default_fp8i4_config, default_bf8i4_config, default_fp8_preshufflequant_config, + default_bf8_preshufflequant_config, + default_fp8i4_preshufflequant_config, default_bf8i4_preshufflequant_config, ) from codegen_common import make_aquant_kernel_name # noqa: E402 @@ -159,6 +162,71 @@ def test_full_kernel_count_is_28(self): self.assertEqual(total, 28) +class TestArchWarpTileK(unittest.TestCase): + """warp_tile_k must be arch-derived (get_k_warp_tile<8bit_float,16,IsFlatMM>()). + + Every AQuant variant (fp8/bf8/fp8i4/bf8i4) instantiates the GEMM config with an + 8-bit float PrecType (fp8_t/bf8_t; pk_int4 A does not drive the K warp tile), so + warp_tile_k depends only on arch and pipeline. gfx942 has NO valid 16x16x128 + fp8/bf8 warp-gemm: warp_tile_k=128 compiles but SILENTLY OUTPUTS ALL-ZEROS there + (GPU-confirmed on gfx942 MI300X). Old-TE uses 16x16x32 (decode) / 16x16x64 + (preshufflequant) on gfx942. Only gfx950 gets 128. + """ + + _DECODE_CTORS = [ + default_fp8_config, + default_bf8_config, + default_fp8i4_config, + default_bf8i4_config, + ] + _PRESHUF_CTORS = [ + default_fp8_preshufflequant_config, + default_bf8_preshufflequant_config, + default_fp8i4_preshufflequant_config, + default_bf8i4_preshufflequant_config, + ] + + def test_helper_decode(self): + self.assertEqual(_warp_tile_k_for("gfx942", preshuffle_aquant=False), 32) + self.assertEqual(_warp_tile_k_for("gfx950", preshuffle_aquant=False), 128) + + def test_helper_preshufflequant(self): + self.assertEqual(_warp_tile_k_for("gfx942", preshuffle_aquant=True), 64) + self.assertEqual(_warp_tile_k_for("gfx950", preshuffle_aquant=True), 128) + + def test_helper_arch_suffix_tolerant(self): + # Real rocm_agent_enumerator output carries feature suffixes. + self.assertEqual(_warp_tile_k_for("gfx942:sramecc+:xnack-"), 32) + self.assertEqual(_warp_tile_k_for("gfx950:sramecc+:xnack-"), 128) + + def test_decode_configs_gfx942(self): + # BLOCKING: all four decode variants must be 32 on gfx942, NOT 128. + for ctor in self._DECODE_CTORS: + self.assertEqual(ctor(gfx_arch="gfx942").warp_tile_k, 32, ctor.__name__) + + def test_decode_configs_gfx950(self): + for ctor in self._DECODE_CTORS: + self.assertEqual(ctor(gfx_arch="gfx950").warp_tile_k, 128, ctor.__name__) + + def test_preshufflequant_configs_gfx942(self): + for ctor in self._PRESHUF_CTORS: + self.assertEqual(ctor(gfx_arch="gfx942").warp_tile_k, 64, ctor.__name__) + + def test_preshufflequant_configs_gfx950(self): + for ctor in self._PRESHUF_CTORS: + self.assertEqual(ctor(gfx_arch="gfx950").warp_tile_k, 128, ctor.__name__) + + def test_name_reflects_arch_warp_tile_k(self): + # The arch-derived value must flow into the byte-exact .name. + self.assertIn("16x16x32", default_fp8_config(gfx_arch="gfx942").name) + self.assertIn("16x16x128", default_fp8_config(gfx_arch="gfx950").name) + self.assertIn("16x16x64", default_bf8_preshufflequant_config(gfx_arch="gfx942").name) + + def test_explicit_override_respected(self): + # An explicit warp_tile_k still wins (for sweeps / experimentation). + self.assertEqual(default_fp8_config(gfx_arch="gfx942", warp_tile_k=128).warp_tile_k, 128) + + class TestSupportedArchs(unittest.TestCase): def test_gfx90a_supported_and_validated(self): self.assertIn("gfx90a", _SUPPORTED_ARCHS) From da3ae8da3c7c0b23220bfe3827b14e2cbfbfc733 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Sun, 26 Jul 2026 03:54:45 -0400 Subject: [PATCH 05/11] fix(ck-tile): aquant pk_int4 A copy overran packed HostTensor The pk_int4 A path copied M*K logical elements into a HostTensor that (PackedSize=2) only allocates M*K/2 elements, overrunning the heap before permute_vectors_i4x4_b ran. This corrupted the allocator and crashed every fp8i4/bf8i4 config (malloc abort / segfault), leaving all i4 cases N/A in the default_config parity sweep. Copy a_h.size() (the packed element count) instead, mirroring the already -fixed B path in gemm_bquant_ctypes_lib.cpp. A_host is the packed representation (2 nibbles/byte), so the packed count is the correct amount. Verified on gfx942/MI300X: fp8i4 and bf8i4 now build, run, and are bit-exact (maxrel 0.0) vs an empirically-anchored reference; fp8/bf8 non-i4 paths unaffected (change is inside the pk_int4 branch). --- .../dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp index ec82f5444c92..09949a85994b 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp @@ -261,7 +261,13 @@ int dispatcher_run_aquant_gemm(const void* A, static_cast(K), static_cast(stride_A), ck_tile::bool_constant{})); - std::copy(A_host, A_host + M * K, a_h.begin()); + // For packed A (pk_int4_t; PackedSize=2) the HostTensor holds only + // M*K/PackedSize elements and A_host already contains the packed + // representation, so copy a_h.size() elements (== elements_to_bytes + // count) -- copying M*K here overran the buffer and corrupted the heap + // before permute_vectors_i4x4_b ran, crashing all fp8i4/bf8i4 configs. + // Mirrors the B path in gemm_bquant_ctypes_lib.cpp. + std::copy(A_host, A_host + a_h.size(), a_h.begin()); ck_tile::permute_vectors_i4x4_b(a_h); if(hipMemcpy( A_dev, a_h.data(), elements_to_bytes(M * K), hipMemcpyHostToDevice) != From bb5e01dd790b79fa5c908d5db9577b3a4f62e3c5 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk Date: Mon, 27 Jul 2026 03:36:51 -0400 Subject: [PATCH 06/11] ci: re-trigger CI (flaky gfx11 ctest timeout on unrelated test_batched_gemm_b_scale_wmma; i4 fix unchanged) From f881aa3a1d037a60ab06c2a1a3b6dac49da1cf7c Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 29 Jul 2026 13:32:43 -0400 Subject: [PATCH 07/11] fix(ck-tile): activate HIP_CHECK macro in gemm_aquant_ctypes_lib Replace verbose inline if(!= hipSuccess){cleanup();return -1;} guards with HIP_CHECK(), which was defined but never invoked. Also update the HIP_CHECK definition to call cleanup() on error (matching the grouped variant). Preshuffle/permute conditional branches each get their own HIP_CHECK call so error handling is uniform across all paths. Co-Authored-By: Claude --- .../ctypes/gemm_aquant_ctypes_lib.cpp | 93 ++++--------------- 1 file changed, 20 insertions(+), 73 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp index 09949a85994b..cf56ecafe2d5 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp @@ -60,15 +60,16 @@ static constexpr std::size_t elements_to_bytes(std::size_t n) static bool g_initialized = false; -#define HIP_CHECK(call) \ - { \ - hipError_t _err = (call); \ - if(_err != hipSuccess) \ - { \ +#define HIP_CHECK(call) \ + { \ + hipError_t _err = (call); \ + if(_err != hipSuccess) \ + { \ std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ - << __LINE__ << "\n"; \ - return -1; \ - } \ + << __LINE__ << "\n"; \ + cleanup(); \ + return -1; \ + } \ } extern "C" { @@ -230,26 +231,10 @@ int dispatcher_run_aquant_gemm(const void* A, // Allocate device buffers. // A may be a packed type (pk_int4_t): 2 logical values per byte. // elements_to_bytes(n) handles the packed case via numeric_traits::PackedSize. - if(hipMalloc(&A_dev, elements_to_bytes(M * K)) != hipSuccess) - { - cleanup(); - return -1; - } - if(hipMalloc(&AQ_dev, elements_to_bytes(M * QK_A)) != hipSuccess) - { - cleanup(); - return -1; - } - if(hipMalloc(&B_dev, elements_to_bytes(K * N)) != hipSuccess) - { - cleanup(); - return -1; - } - if(hipMalloc(&C_dev, elements_to_bytes(M * N)) != hipSuccess) - { - cleanup(); - return -1; - } + HIP_CHECK(hipMalloc(&A_dev, elements_to_bytes(M * K))); + HIP_CHECK(hipMalloc(&AQ_dev, elements_to_bytes(M * QK_A))); + HIP_CHECK(hipMalloc(&B_dev, elements_to_bytes(K * N))); + HIP_CHECK(hipMalloc(&C_dev, elements_to_bytes(M * N))); // Copy A to device. For pk_int4 A the raw i4x4 values must be permuted for the // device implementation -- mirrors run_gemm_quant_example.inc:758-763. @@ -269,22 +254,11 @@ int dispatcher_run_aquant_gemm(const void* A, // Mirrors the B path in gemm_bquant_ctypes_lib.cpp. std::copy(A_host, A_host + a_h.size(), a_h.begin()); ck_tile::permute_vectors_i4x4_b(a_h); - if(hipMemcpy( - A_dev, a_h.data(), elements_to_bytes(M * K), hipMemcpyHostToDevice) != - hipSuccess) - { - cleanup(); - return -1; - } + HIP_CHECK(hipMemcpy(A_dev, a_h.data(), elements_to_bytes(M * K), hipMemcpyHostToDevice)); } else { - if(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice) != - hipSuccess) - { - cleanup(); - return -1; - } + HIP_CHECK(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice)); } // Apply AQ preshuffle when required -- mirrors run_gemm_quant_example.inc:746-751. @@ -310,36 +284,14 @@ int dispatcher_run_aquant_gemm(const void* A, ck_tile::bool_constant{} /*row-major*/)); std::copy(AQ_host, AQ_host + M * 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(M * QK_A), - hipMemcpyHostToDevice) != hipSuccess) - { - cleanup(); - return -1; - } + HIP_CHECK(hipMemcpy(AQ_dev, aq_shuffled.data(), elements_to_bytes(M * QK_A), hipMemcpyHostToDevice)); } else { - if(hipMemcpy( - AQ_dev, AQ_host, elements_to_bytes(M * QK_A), hipMemcpyHostToDevice) != - hipSuccess) - { - cleanup(); - return -1; - } - } - if(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice) != - hipSuccess) - { - cleanup(); - return -1; - } - if(hipMemset(C_dev, 0, elements_to_bytes(M * N)) != hipSuccess) - { - cleanup(); - return -1; + HIP_CHECK(hipMemcpy(AQ_dev, AQ_host, elements_to_bytes(M * QK_A), hipMemcpyHostToDevice)); } + HIP_CHECK(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemset(C_dev, 0, elements_to_bytes(M * N))); // Build QuantGemmHostArgs (bq_ptr = nullptr, QK_B = 0, stride_BQ = 0 for AQuant-only) ck_tile::QuantGemmHostArgs args; @@ -384,12 +336,7 @@ int dispatcher_run_aquant_gemm(const void* A, } // Copy result back - if(hipMemcpy(C_host, C_dev, elements_to_bytes(M * N), hipMemcpyDeviceToHost) != - hipSuccess) - { - cleanup(); - return -1; - } + HIP_CHECK(hipMemcpy(C_host, C_dev, elements_to_bytes(M * N), hipMemcpyDeviceToHost)); if(time_ms) *time_ms = exec_time; From 09305492b78b09c9e7a031b94facf16e885bf0a7 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 29 Jul 2026 13:48:08 -0400 Subject: [PATCH 08/11] docs(ck-tile): add section headers for decode/preshufflequant in gemm_aquant_utils Add # === section headers before the decode-family and preshufflequant-family config factory functions, mirroring the grouping style from the sibling gemm_bquant_utils.py file. Co-Authored-By: Claude --- .../dispatcher/python/gemm_aquant_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py index 45fbf4ac0711..f1aecb111c43 100644 --- a/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_aquant_utils.py @@ -771,6 +771,12 @@ def _warp_tile_k_for(gfx_arch: str, preshuffle_aquant: bool = False) -> int: return 64 if preshuffle_aquant else 32 +# ============================================================================= +# Decode family (GemmConfigQuantDecodeInterwave, tile 16x64x256, IsFlatMM=false) +# fp8/bf8/fp8i4/bf8i4: K_warp = 128 on gfx950, 32 on gfx942. +# ============================================================================= + + def _decode_config( variant_key: str, warp_tile_k: int, @@ -845,6 +851,12 @@ def default_bf8i4_config(quant_group_k: int = 128, quant_group_n: int = 1, return _decode_config("bf8i4", warp_tile_k, quant_group_k, quant_group_n, layout, gfx_arch) +# ============================================================================= +# Preshufflequant family (GemmConfigPreshuffleQuantDecode, tile 16x64x256, IsFlatMM=true) +# fp8/bf8/fp8i4/bf8i4: K_warp = 128 on gfx950, 64 on gfx942. +# ============================================================================= + + def _preshufflequant_config( variant_key: str, warp_tile_k: int, From d40877fd8d218b9ea9d45cebd7dc572d9d55dc53 Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 29 Jul 2026 15:33:56 -0400 Subject: [PATCH 09/11] style(ck-tile): fix clang-format-18 violations in gemm_aquant_ctypes_lib Remove trailing alignment spaces before backslash continuations in HIP_CHECK macro and remove manual column-alignment spaces in hipMalloc call sites. Co-Authored-By: Claude --- .../ctypes/gemm_aquant_ctypes_lib.cpp | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp index cf56ecafe2d5..166beb6271f9 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp @@ -60,16 +60,16 @@ static constexpr std::size_t elements_to_bytes(std::size_t n) static bool g_initialized = false; -#define HIP_CHECK(call) \ - { \ - hipError_t _err = (call); \ - if(_err != hipSuccess) \ - { \ +#define HIP_CHECK(call) \ + { \ + hipError_t _err = (call); \ + if(_err != hipSuccess) \ + { \ std::cerr << "HIP error: " << hipGetErrorString(_err) << " at " << __FILE__ << ":" \ - << __LINE__ << "\n"; \ - cleanup(); \ - return -1; \ - } \ + << __LINE__ << "\n"; \ + cleanup(); \ + return -1; \ + } \ } extern "C" { @@ -231,10 +231,10 @@ int dispatcher_run_aquant_gemm(const void* A, // Allocate device buffers. // A may be a packed type (pk_int4_t): 2 logical values per byte. // elements_to_bytes(n) handles the packed case via numeric_traits::PackedSize. - HIP_CHECK(hipMalloc(&A_dev, elements_to_bytes(M * K))); + HIP_CHECK(hipMalloc(&A_dev, elements_to_bytes(M * K))); HIP_CHECK(hipMalloc(&AQ_dev, elements_to_bytes(M * QK_A))); - HIP_CHECK(hipMalloc(&B_dev, elements_to_bytes(K * N))); - HIP_CHECK(hipMalloc(&C_dev, elements_to_bytes(M * N))); + HIP_CHECK(hipMalloc(&B_dev, elements_to_bytes(K * N))); + HIP_CHECK(hipMalloc(&C_dev, elements_to_bytes(M * N))); // Copy A to device. For pk_int4 A the raw i4x4 values must be permuted for the // device implementation -- mirrors run_gemm_quant_example.inc:758-763. From c2c0bc170eac3d2f2b660285fe0909dcee506474 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:11:47 -0400 Subject: [PATCH 10/11] test(ck-tile): add on-device GPU correctness test for aquant (registered in CI) --- .../dispatcher/tests/CMakeLists.txt | 12 + .../tests/test_aquant_gpu_correctness.py | 214 ++++++++++++++++++ 2 files changed, 226 insertions(+) create mode 100644 projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py diff --git a/projects/composablekernel/dispatcher/tests/CMakeLists.txt b/projects/composablekernel/dispatcher/tests/CMakeLists.txt index bd19859e5f2d..052d2d59eaaf 100644 --- a/projects/composablekernel/dispatcher/tests/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/tests/CMakeLists.txt @@ -10,6 +10,18 @@ cmake_minimum_required(VERSION 3.16) # Find Python find_package(Python3 COMPONENTS Interpreter REQUIRED) +# Bridge on-device GPU correctness test (SKIPs cleanly without a GPU/built lib). +add_test( + NAME dispatcher_test_aquant_gpu_correctness + COMMAND ${Python3_EXECUTABLE} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR} -p test_aquant_gpu_correctness.py -v + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) +set_tests_properties(dispatcher_test_aquant_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" + TIMEOUT 600 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # ============================================================================= # Python Tests # ============================================================================= diff --git a/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py new file mode 100644 index 000000000000..cb6c776112c1 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +On-device GPU correctness test for the AQuant (A-scaled block-scale) GEMM +dispatcher bridge (PR #9980). + +Mirrors the structure of test_bquant_gpu_correctness.py: build the op's +block-scale dispatcher .so, run it on the GPU via AQuantGpuGemmRunner, and +compare against a NumPy fp32 reference + + C = dequant(A, AQ) @ B + +where AQ applies a per-(row, K-group) scale to A: + + A_deq[m, k] = A[m, k] * AQ[m, k // quant_group_k] + +Two block-scale correctness traps are handled explicitly (both were needed to +get a bit-accurate match on gfx942): + + 1. fp8 encoding flavour. gfx942 uses FNUZ fp8/bf8 (e4m3fnuz / e5m2fnuz); + gfx950/gfx12 use OCP (e4m3 / e5m2). Encoding A/B with the wrong flavour + makes the reference NaN or uncorrelated. Selected from the detected arch. + 2. warp_tile_k. fp8/bf8 needs warp_tile_k=32 on gfx942 (128 silently outputs + all-zeros). default_fp8_config/default_bf8_config derive this from the + arch; the test asserts it matches to guard against regressions. + +Also note the rcr layout uses ColMajor B, so the [K, N] logical B is stored +transposed (a C-contiguous [N, K] buffer == [K, N] col-major) before it is +handed to the kernel. + +Run: + python3 -m pytest test_aquant_gpu_correctness.py -v + python3 test_aquant_gpu_correctness.py # standalone +""" + +import math +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +_HERE = Path(__file__).resolve() + + +def _find_python_dir() -> Path: + for c in (_HERE.parent.parent / "python", _HERE.parent / "python"): + if (c / "gemm_aquant_utils.py").is_file(): + return c + for parent in _HERE.parents: + cand = parent / "dispatcher" / "python" + if (cand / "gemm_aquant_utils.py").is_file(): + return cand + raise RuntimeError("could not locate dispatcher/python/gemm_aquant_utils.py") + + +sys.path.insert(0, str(_find_python_dir())) + +from gemm_aquant_utils import ( # noqa: E402 + AQuantGemmProblem, + AQuantGpuGemmRunner, + setup_multiple_aquant_dispatchers, + default_fp8_config, + default_bf8_config, + _detect_gpu_arch, + _warp_tile_k_for, +) + + +# ============================================================================= +# arch-aware fp8/bf8 codecs (OCP on gfx950/gfx12, FNUZ on gfx942) +# ============================================================================= + + +def _uses_ocp(arch: str) -> bool: + return ("gfx950" in arch) or ("gfx12" in arch) + + +def _ml_type(dtype: str, arch: str): + import ml_dtypes + if _uses_ocp(arch): + return ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return ml_dtypes.float8_e4m3fnuz if dtype == "fp8" else ml_dtypes.float8_e5m2fnuz + + +def _encode(arr, dtype: str, arch: str) -> np.ndarray: + t = _ml_type(dtype, arch) + return np.ascontiguousarray(arr.astype(np.float32).astype(t)).view(np.uint8) + + +def _qdq(arr, dtype: str, arch: str) -> np.ndarray: + """Round through fp8/bf8 and back to fp32 (reference-side counterpart).""" + t = _ml_type(dtype, arch) + return arr.astype(np.float32).astype(t).astype(np.float32) + + +def _have_ml_dtypes() -> bool: + try: + import ml_dtypes # noqa: F401 + return True + except ImportError: + return False + + +def _have_gpu() -> bool: + try: + out = subprocess.run(["rocminfo"], capture_output=True, text=True, timeout=30) + return "gfx" in out.stdout + except Exception: + return False + + +# ============================================================================= +# NumPy reference: C = dequant(A, AQ) @ B +# ============================================================================= + + +def reference_aquant_gemm(A_dec, B_dec, AQ, problem: AQuantGemmProblem) -> np.ndarray: + """A_dec [M,K], B_dec [K,N] are already fp8-rounded fp32 (kernel-identical).""" + M, N, K = problem.M, problem.N, problem.K + gK = problem.quant_group_k + A_deq = A_dec.astype(np.float32).copy() + for qi in range(problem.QK_A): + k0 = qi * gK + k1 = min(k0 + gK, K) + A_deq[:, k0:k1] *= AQ.astype(np.float32)[:, qi][:, None] + return (A_deq @ B_dec.astype(np.float32)).astype(np.float32) + + +# ============================================================================= +# build + run + verify +# ============================================================================= + + +def _run_case(dtype: str, M: int, N: int, K: int, gK: int, out_dir: Path): + arch = _detect_gpu_arch() + make_cfg = default_fp8_config if dtype == "fp8" else default_bf8_config + config = make_cfg(quant_group_k=gK, quant_group_n=1, layout="rcr", gfx_arch=arch) + + expected_wtk = _warp_tile_k_for(arch, preshuffle_aquant=False) + assert config.warp_tile_k == expected_wtk, ( + f"warp_tile_k arch trap: got {config.warp_tile_k}, expected {expected_wtk} for {arch}" + ) + + so_paths = setup_multiple_aquant_dispatchers( + configs=[config], output_dir=out_dir, gfx_arch=arch + ) + assert so_paths and so_paths[0] is not None, "aquant kernel build failed" + so_path = so_paths[0] + + rng = np.random.default_rng(1234) + QK_A = math.ceil(K / gK) + A_f = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) + B_f = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) # logical [K, N] + AQ_f = rng.uniform(0.5, 1.5, (M, QK_A)).astype(np.float32) + + A_raw = _encode(A_f, dtype, arch) # [M, K] row-major fp8 bytes + B_raw = _encode(B_f.T.copy(), dtype, arch) # [N, K] C-contig == [K, N] col-major + A_dec = _qdq(A_f, dtype, arch) + B_dec = _qdq(B_f, dtype, arch) + + problem = AQuantGemmProblem(M=M, N=N, K=K, quant_group_k=gK) + runner = AQuantGpuGemmRunner(so_path, layout="rcr") + result = runner.run(A=A_raw, AQ=AQ_f, B=B_raw, problem=problem) + + C_gpu = np.asarray(result.C, dtype=np.float32) + assert np.max(np.abs(C_gpu)) > 1e-3, "GPU output all-zeros (warp_tile_k arch trap?)" + + C_ref = reference_aquant_gemm(A_dec, B_dec, AQ_f, problem) + max_rel = float(np.max(np.abs(C_gpu - C_ref)) / (np.max(np.abs(C_ref)) + 1e-6)) + + tol = 0.05 # fp8/bf8 block-scale ~1e-2 .. 5e-2 + assert max_rel <= tol, f"max_rel={max_rel:.4f} > tol={tol} ({dtype} {M}x{N}x{K} gK={gK})" + return max_rel, result.time_ms + + +_SKIP_NO_GPU = pytest.mark.skipif(not _have_gpu(), reason="no ROCm GPU detected") +_SKIP_NO_MLD = pytest.mark.skipif(not _have_ml_dtypes(), reason="ml_dtypes not installed") + + +@_SKIP_NO_GPU +@_SKIP_NO_MLD +@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) +def test_aquant_gpu_matches_reference(dtype, tmp_path): + max_rel, _ = _run_case(dtype, M=256, N=256, K=512, gK=128, out_dir=tmp_path) + assert max_rel <= 0.05 + + +@_SKIP_NO_GPU +@_SKIP_NO_MLD +def test_aquant_gpu_not_all_zeros(tmp_path): + _run_case("fp8", M=256, N=256, K=512, gK=128, out_dir=tmp_path) + + +if __name__ == "__main__": + if not _have_gpu(): + print("SKIP: no GPU"); raise SystemExit(0) + if not _have_ml_dtypes(): + print("SKIP: ml_dtypes not installed"); raise SystemExit(0) + d = Path(tempfile.mkdtemp(prefix="aquant_gputest_")) + ok = True + for dt in ("fp8", "bf8"): + try: + mr, t = _run_case(dt, 256, 256, 512, 128, d) + print(f"PASS aquant {dt}: max_rel={mr:.4f} time_ms={t:.3f}") + except Exception as e: + ok = False + print(f"FAIL aquant {dt}: {e}") + raise SystemExit(0 if ok else 1) From 379d91cc9d1d0d4ef77d5cdcb43bcd0616cf5de6 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:57:41 -0400 Subject: [PATCH 11/11] style(ck-tile): fix clang-format-18 violations in gemm_aquant_ctypes_lib --- .../bindings/ctypes/gemm_aquant_ctypes_lib.cpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp index 166beb6271f9..7c2f81db6a36 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_aquant_ctypes_lib.cpp @@ -254,11 +254,13 @@ int dispatcher_run_aquant_gemm(const void* A, // Mirrors the B path in gemm_bquant_ctypes_lib.cpp. std::copy(A_host, A_host + a_h.size(), a_h.begin()); ck_tile::permute_vectors_i4x4_b(a_h); - HIP_CHECK(hipMemcpy(A_dev, a_h.data(), elements_to_bytes(M * K), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy( + A_dev, a_h.data(), elements_to_bytes(M * K), hipMemcpyHostToDevice)); } else { - HIP_CHECK(hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice)); + HIP_CHECK( + hipMemcpy(A_dev, A_host, elements_to_bytes(M * K), hipMemcpyHostToDevice)); } // Apply AQ preshuffle when required -- mirrors run_gemm_quant_example.inc:746-751. @@ -284,11 +286,15 @@ int dispatcher_run_aquant_gemm(const void* A, ck_tile::bool_constant{} /*row-major*/)); std::copy(AQ_host, AQ_host + M * QK_A, aq_h.begin()); auto aq_shuffled = ck_tile::shuffle_aq(&aq_h, block_aq_k); - HIP_CHECK(hipMemcpy(AQ_dev, aq_shuffled.data(), elements_to_bytes(M * QK_A), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(AQ_dev, + aq_shuffled.data(), + elements_to_bytes(M * QK_A), + hipMemcpyHostToDevice)); } else { - HIP_CHECK(hipMemcpy(AQ_dev, AQ_host, elements_to_bytes(M * QK_A), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy( + AQ_dev, AQ_host, elements_to_bytes(M * QK_A), hipMemcpyHostToDevice)); } HIP_CHECK(hipMemcpy(B_dev, B_host, elements_to_bytes(K * N), hipMemcpyHostToDevice)); HIP_CHECK(hipMemset(C_dev, 0, elements_to_bytes(M * N)));