[rocm-libraries] ROCm/rocm-libraries#8519 (commit 9637390)

feat(ck-tile): add block-scale GEMM operators (aquant,
 bquant, abquant) (#8519)

JIRA ID - AICK-1289
Motivation
Adds three new block-scale quantized GEMM operators to the CK Tile
Engine for FP8/BF8 inference workloads.

Technical Details
gemm_aquant: A-matrix quantized GEMM with per-row-group scale tensor [M,
K/group_size_k]
gemm_bquant: B-matrix quantized GEMM with per-column-group scale tensor
[K/group_size_k, N]
gemm_abquant: Both A and B quantized with independent group-scale
tensors
Each operator includes CMakeLists, Python instance builder with tier
sampling, C++ benchmark/profiler with host reference verification, and
config JSONs. Supporting changes to gemm_instance_builder.py,
gemm_validation_utils.py, sampling infra, and the operation support
matrix.

Test Plan
 Build and run all three operators with fp8/bf8 on gfx942/gfx950
 Verify correctness against CPU reference
 Verify CI config builds pass
This commit is contained in:
Thrupti Raj Lakshmana Gowda
2026-07-07 18:22:48 +00:00
committed by assistant-librarian[bot]
parent 05615558c6
commit b0f200713a
41 changed files with 8702 additions and 43 deletions

View File

@@ -13,6 +13,9 @@
| GEMM | flatmm<br>example: 18_flatmm/ | ❌ | ❌ | ❌ | ❌ | | ❌ | ❌ | ❌ | | | | ❌ | ❌ | ❌ | ❌ | |
| GEMM | gemm_multi_abd<br>example: 22_gemm_multi_abd/ | ✅ | | | | | | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | 0.0833 |
| GEMM | gemm_quant | | ❌ | | ❌ | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ | |
| GEMM | block_scale_gemm/gemm_aquant<br>engine: block_scale_gemm/gemm_aquant/<br>example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | | | | ❌ | ✅ | ✅ | ❌ | 0.0625 |
| GEMM | block_scale_gemm/gemm_bquant<br>engine: block_scale_gemm/gemm_bquant/<br>example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | 0.0625 |
| GEMM | block_scale_gemm/gemm_abquant<br>engine: block_scale_gemm/gemm_abquant/<br>example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | 0.0625 |
| GEMM | grouped_gemm [10]<br>engine: grouped_gemm/<br>example: 17_grouped_gemm/ | ✅ | ✅ | | | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | 0.0834 |
| GEMM | grouped_gemm_quant/grouped_gemm_rowcolquant<br>engine: grouped_gemm_quant/grouped_gemm_rowcolquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 |
| GEMM | grouped_gemm_quant/grouped_gemm_tensorquant<br>engine: grouped_gemm_quant/grouped_gemm_tensorquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 |

View File

@@ -15,7 +15,7 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
if(_te_budget GREATER 0)
# Detect active ops from their DATATYPE variables
set(_active_ops "")
foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant)
foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant gemm_aquant gemm_bquant gemm_abquant)
string(TOUPPER ${_op} _OP_UPPER)
if(NOT "${${_OP_UPPER}_DATATYPE}" STREQUAL "")
list(APPEND _active_ops ${_op})
@@ -45,7 +45,7 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
message(STATUS "Sampling budget allocation:\n${_alloc_output}")
# Read per-op allocations (only if not already overridden)
foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant)
foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant gemm_aquant gemm_bquant gemm_abquant)
string(TOUPPER ${_op} _OP_UPPER)
if("${${_OP_UPPER}_MAX_INSTANCES}" STREQUAL "")
if(EXISTS "${_alloc_dir}/${_op}_budget.txt")

View File

@@ -1,5 +1,8 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
add_subdirectory(gemm_aquant)
add_subdirectory(gemm_bquant)
add_subdirectory(gemm_abquant)
add_subdirectory(gemm_rowcolquant)
add_subdirectory(gemm_tensor_quant)

View File

@@ -0,0 +1,272 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
set(GEMM_ABQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM ABQuant (semicolon-separated)")
set(GEMM_ABQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM ABQuant (semicolon-separated)")
set(GEMM_ABQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)")
option(ENABLE_CCACHE_GEMM_ABQUANT "Enable ccache for GEMM ABQuant ops compilation" OFF)
# Store the directory path for use in functions
set(GEMM_ABQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR})
# Function to create individual GEMM ABQuant targets
function(create_individual_gemm_abquant_target datatype layout trait tile_config config_json)
# GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL is guaranteed non-empty here (caller checks at line 192)
if(NOT GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL)
message(FATAL_ERROR "BUG: create_individual_gemm_abquant_target called with no GPU targets")
endif()
# Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k
string(REPLACE "_" ";" config_groups ${tile_config})
list(GET config_groups 0 tile_dims)
list(GET config_groups 1 warp_dims)
list(GET config_groups 2 warp_tile_dims)
string(REPLACE "x" ";" tile_parts ${tile_dims})
list(GET tile_parts 0 tile_m)
list(GET tile_parts 1 tile_n)
list(GET tile_parts 2 tile_k)
string(REPLACE "x" ";" warp_parts ${warp_dims})
list(GET warp_parts 0 warp_m)
list(GET warp_parts 1 warp_n)
list(GET warp_parts 2 warp_k)
string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims})
list(GET warp_tile_parts 0 warp_tile_m)
list(GET warp_tile_parts 1 warp_tile_n)
list(GET warp_tile_parts 2 warp_tile_k)
set(target_name "benchmark_gemm_abquant_${datatype}_${layout}_${trait}_${tile_config}")
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
# Generate the single instance header for this kernel
set(instance_header "${working_path}/gemm_abquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp")
# Add custom command to generate the header file at build time
add_custom_command(
OUTPUT ${instance_header}
COMMAND ${Python3_EXECUTABLE} ${GEMM_ABQUANT_SOURCE_DIR}/gemm_abquant_instance_builder.py
--working_path ${working_path}
--datatype ${datatype}
--layout ${layout}
--config_json ${config_json}
--gen_single
--kernel_name "gemm_abquant_${datatype}_${layout}_${trait}_${tile_config}"
--tile_config "${tile_config}"
--trait_combo "${trait}"
--gpu_target "${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}"
DEPENDS ${GEMM_ABQUANT_SOURCE_DIR}/gemm_abquant_instance_builder.py ${config_json}
COMMENT "Generating ${instance_header}"
)
# Create the executable
add_executable(${target_name}
EXCLUDE_FROM_ALL
${GEMM_ABQUANT_SOURCE_DIR}/gemm_abquant_benchmark_single.cpp
${instance_header}
)
# Set GPU architectures
set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL})
# Set compile definitions
target_compile_definitions(${target_name} PRIVATE
GEMM_ABQUANT_SINGLE_INSTANCE_HPP="${instance_header}"
)
# Include directories
target_include_directories(${target_name} PRIVATE
${GEMM_ABQUANT_SOURCE_DIR}
${working_path}
)
# Compile options
target_compile_options(${target_name} PRIVATE
-Wno-undefined-func-template
-Wno-float-equal
--offload-compress
-include ${instance_header}
)
# Add to collection targets
add_dependencies(benchmark_gemm_abquant_all ${target_name})
add_dependencies(benchmark_gemm_abquant_${datatype} ${target_name})
add_dependencies(benchmark_gemm_abquant_${layout} ${target_name})
add_dependencies(benchmark_gemm_abquant_${datatype}_${layout} ${target_name})
# Add to pipeline-specific targets
string(REPLACE "_" ";" trait_parts ${trait})
list(GET trait_parts 0 pipeline)
add_dependencies(benchmark_gemm_abquant_${pipeline}_pipeline ${target_name})
endfunction()
# Function to build individual GEMM ABQuant targets
function(build_individual_gemm_abquant_targets datatype layout)
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
# Choose config file (priority: env var > cmake var > default)
if(DEFINED ENV{GEMM_ABQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_ABQUANT_CONFIG_FILE}" STREQUAL "")
set(config_filename "$ENV{GEMM_ABQUANT_CONFIG_FILE}")
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}")
message(VERBOSE " Using config from environment variable: ${config_filename}")
elseif(NOT "${GEMM_ABQUANT_CONFIG_FILE}" STREQUAL "")
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_ABQUANT_CONFIG_FILE}")
message(VERBOSE " Using custom config: ${GEMM_ABQUANT_CONFIG_FILE}")
else()
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json")
message(VERBOSE " Using default config for layout ${layout}")
endif()
if(NOT EXISTS ${json_blob})
message(FATAL_ERROR "Config file not found: ${json_blob}")
endif()
# Create working directory
file(MAKE_DIRECTORY ${working_path})
# Build sampling arguments
set(extra_list_args "")
if(NOT "${GEMM_ABQUANT_MAX_INSTANCES}" STREQUAL "")
list(APPEND extra_list_args --max-instances ${GEMM_ABQUANT_MAX_INSTANCES})
endif()
if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER})
list(APPEND extra_list_args --manifest-path ${working_path})
endif()
if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "")
list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED})
endif()
# List kernels
message(VERBOSE " Listing ABQuant kernel configurations for ${datatype} ${layout}...")
execute_process(
COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_abquant_instance_builder.py
--working_path ${working_path}
--datatype ${datatype}
--layout ${layout}
--config_json ${json_blob}
--gpu_target "${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}"
--list_kernels
${extra_list_args}
WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
RESULT_VARIABLE ret
OUTPUT_VARIABLE list_output
ERROR_VARIABLE list_error
)
if(NOT ret EQUAL 0)
message(FATAL_ERROR "Failed to list ABQuant kernels for ${datatype} ${layout}: ${list_error}")
endif()
# Read kernel count
if(EXISTS ${working_path}/gemm_abquant_kernel_count.txt)
file(READ ${working_path}/gemm_abquant_kernel_count.txt kernel_count)
string(STRIP "${kernel_count}" kernel_count)
message(VERBOSE " Found ${kernel_count} ABQuant kernel configurations")
else()
message(FATAL_ERROR "Kernel count file not found")
endif()
# Read kernel list and create targets
if(EXISTS ${working_path}/gemm_abquant_kernel_list.txt)
file(STRINGS ${working_path}/gemm_abquant_kernel_list.txt kernel_lines)
foreach(line IN LISTS kernel_lines)
string(REPLACE "|" ";" parts "${line}")
list(GET parts 0 kernel_name)
list(GET parts 1 tile_config)
list(GET parts 2 trait_combo)
create_individual_gemm_abquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}")
endforeach()
else()
message(FATAL_ERROR "Kernel list file not found")
endif()
endfunction()
# Main build logic
message(VERBOSE "=== Starting Tile Engine GEMM ABQuant Configuration ===")
message(VERBOSE "GEMM_ABQUANT_DATATYPE: ${GEMM_ABQUANT_DATATYPE}")
message(VERBOSE "GEMM_ABQUANT_LAYOUT: ${GEMM_ABQUANT_LAYOUT}")
message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
# Filter GPU targets to supported architectures
set(GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL "")
set(DESIRED_TARGETS "gfx90a;gfx942;gfx950")
foreach(target IN LISTS SUPPORTED_GPU_TARGETS)
if(target IN_LIST DESIRED_TARGETS)
list(APPEND GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL ${target})
message(VERBOSE " Adding GPU target: ${target}")
endif()
endforeach()
# Skip build if no matching targets found
if(NOT GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL)
message(WARNING "Skipping Tile Engine GEMM ABQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
else()
message(VERBOSE "Building individual GEMM ABQuant targets for GPU targets: ${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}")
# Enable compiler cache if requested
if(ENABLE_CCACHE_GEMM_ABQUANT)
find_program(CCACHE_PROGRAM ccache)
if(CCACHE_PROGRAM)
set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
message(VERBOSE "Using ccache for faster compilation")
endif()
endif()
# Create master collection target
add_custom_target(benchmark_gemm_abquant_all)
# Create datatype collection targets
foreach(dt IN LISTS GEMM_ABQUANT_DATATYPE)
add_custom_target(benchmark_gemm_abquant_${dt})
endforeach()
# Create layout collection targets
foreach(l IN LISTS GEMM_ABQUANT_LAYOUT)
add_custom_target(benchmark_gemm_abquant_${l})
endforeach()
# Create combined collection targets
foreach(dt IN LISTS GEMM_ABQUANT_DATATYPE)
foreach(l IN LISTS GEMM_ABQUANT_LAYOUT)
add_custom_target(benchmark_gemm_abquant_${dt}_${l})
endforeach()
endforeach()
# Create pipeline-specific collection targets.
# Must stay in sync with trait_config.pipeline.values in the config JSONs:
# if a new pipeline is added there, add it here too or add_dependencies() will fail at build time.
set(GEMM_ABQUANT_PIPELINES "compv3")
foreach(pipeline IN LISTS GEMM_ABQUANT_PIPELINES)
add_custom_target(benchmark_gemm_abquant_${pipeline}_pipeline)
endforeach()
# Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that
# sampling fires per-combo rather than being a single cap.
if(NOT "${GEMM_ABQUANT_MAX_INSTANCES}" STREQUAL "")
list(LENGTH GEMM_ABQUANT_DATATYPE _abq_n_dt)
list(LENGTH GEMM_ABQUANT_LAYOUT _abq_n_lay)
math(EXPR _abq_n_combos "${_abq_n_dt} * ${_abq_n_lay}")
if(_abq_n_combos GREATER 0)
math(EXPR GEMM_ABQUANT_MAX_INSTANCES
"${GEMM_ABQUANT_MAX_INSTANCES} / ${_abq_n_combos}")
if(GEMM_ABQUANT_MAX_INSTANCES EQUAL 0)
set(GEMM_ABQUANT_MAX_INSTANCES 1)
message(WARNING " gemm_abquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_abq_n_combos} combos)")
else()
message(STATUS " gemm_abquant: per-combo budget = ${GEMM_ABQUANT_MAX_INSTANCES} (${_abq_n_combos} combos)")
endif()
endif()
endif()
# Build individual targets for each datatype/layout combination
foreach(dt IN LISTS GEMM_ABQUANT_DATATYPE)
foreach(l IN LISTS GEMM_ABQUANT_LAYOUT)
build_individual_gemm_abquant_targets(${dt} ${l})
endforeach()
endforeach()
endif()

View File

@@ -0,0 +1,104 @@
{
"tile_config": {
"tile_m": {
"values": [
16
]
},
"tile_n": {
"values": [
64
]
},
"tile_k": {
"values": [
256
]
},
"warp_m": {
"values": [
1
]
},
"warp_n": {
"values": [
4
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16
]
},
"warp_tile_n": {
"values": [
16
]
},
"warp_tile_k": {
"values": [
16,
32,
64,
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3"
]
},
"scheduler": {
"values": [
"intrawave"
]
},
"epilogue": {
"values": [
"default",
"cshuffle"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"a_preshuffle_quant": {
"values": [
false
]
},
"b_preshuffle_quant": {
"values": [
false,
true
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128,
"group_size_n": {
"values": [
1,
128
]
}
}

View File

@@ -0,0 +1,113 @@
{
"tile_config": {
"tile_m": {
"max": 256,
"min": 64,
"step": 64
},
"tile_n": {
"max": 256,
"min": 64,
"step": 64
},
"tile_k": {
"max": 256,
"min": 64,
"step": 64
},
"warp_m": {
"values": [
4,
2,
1
]
},
"warp_n": {
"values": [
4,
2,
1
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
4,
16,
32
]
},
"warp_tile_n": {
"values": [
16,
32,
64
]
},
"warp_tile_k": {
"values": [
8,
16,
32,
64,
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3"
]
},
"scheduler": {
"values": [
"intrawave"
]
},
"epilogue": {
"values": [
"cshuffle",
"default"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"a_preshuffle_quant": {
"values": [
false
]
},
"b_preshuffle_quant": {
"values": [
false,
true
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128,
"group_size_n": {
"values": [
1,
128
]
}
}

View File

@@ -0,0 +1,98 @@
{
"tile_config": {
"tile_m": {
"values": [
16
]
},
"tile_n": {
"values": [
64
]
},
"tile_k": {
"values": [
256
]
},
"warp_m": {
"values": [
1
]
},
"warp_n": {
"values": [
4
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16
]
},
"warp_tile_n": {
"values": [
16
]
},
"warp_tile_k": {
"values": [
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3"
]
},
"scheduler": {
"values": [
"intrawave"
]
},
"epilogue": {
"values": [
"default"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"a_preshuffle_quant": {
"values": [
false
]
},
"b_preshuffle_quant": {
"values": [
false
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128,
"group_size_n": {
"values": [
1
]
}
}

View File

@@ -0,0 +1,243 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "gemm_abquant_common.hpp"
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions"
// Data types and Layouts are defined by the generated kernel headers:
// ADataType, BDataType, AQDataType, BQDataType, AccDataType, CDataType
// ALayout, BLayout, CLayout, AQLayout, BQLayout
enum class Metric
{
LATENCY = 0,
TFLOPS = 1,
BANDWIDTH = 2
};
inline constexpr auto get_metric_name(Metric m)
{
switch(m)
{
case Metric::LATENCY: return "latency";
case Metric::TFLOPS: return "tflops";
case Metric::BANDWIDTH: return "bandwidth";
default: throw std::invalid_argument("Unsupported metric type");
}
}
struct ABQuantGemmProblem
{
int split_k_;
int m_, n_, k_;
int stride_a_, stride_b_, stride_c_;
int stride_aq_;
int stride_bq_;
int group_size_k_;
int group_size_n_;
std::string dtype_a_, dtype_b_, dtype_aq_, dtype_bq_, dtype_acc_, dtype_c_;
std::string layout_a_, layout_b_, layout_c_;
friend std::ostream& operator<<(std::ostream& os, const ABQuantGemmProblem& problem)
{
os << "{\n"
<< " \"split_k\":" << problem.split_k_ << ",\n"
<< " \"m\":" << problem.m_ << ",\n"
<< " \"n\":" << problem.n_ << ",\n"
<< " \"k\":" << problem.k_ << ",\n"
<< " \"stride_a\":" << problem.stride_a_ << ",\n"
<< " \"stride_b\":" << problem.stride_b_ << ",\n"
<< " \"stride_c\":" << problem.stride_c_ << ",\n"
<< " \"stride_aq\":" << problem.stride_aq_ << ",\n"
<< " \"stride_bq\":" << problem.stride_bq_ << ",\n"
<< " \"group_size_k\":" << problem.group_size_k_ << ",\n"
<< " \"group_size_n\":" << problem.group_size_n_ << ",\n"
<< " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n"
<< " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n"
<< " \"dtype_aq\":\"" << problem.dtype_aq_ << "\",\n"
<< " \"dtype_bq\":\"" << problem.dtype_bq_ << "\",\n"
<< " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n"
<< " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n"
<< " \"layout_a\":\"" << problem.layout_a_ << "\",\n"
<< " \"layout_b\":\"" << problem.layout_b_ << "\",\n"
<< " \"layout_c\":\"" << problem.layout_c_ << "\"\n"
<< "}";
return os;
}
};
struct PerformanceResult
{
double latency_;
double tflops_;
double bandwidth_;
static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m)
{
switch(m)
{
case Metric::LATENCY: return a.latency_ < b.latency_;
case Metric::TFLOPS: return a.tflops_ > b.tflops_;
case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_;
default: throw std::invalid_argument("Unsupported metric type");
}
}
friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result)
{
os << "{\n"
<< " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_
<< ",\n"
<< " \"tflops(TFlops)\": " << result.tflops_ << ",\n"
<< " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n"
<< "}";
return os;
}
};
struct KernelInstance
{
std::string name_;
ABQuantGemmProblem problem_;
PerformanceResult perf_result_;
static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m)
{
return PerformanceResult::compare(a.perf_result_, b.perf_result_, m);
}
friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj)
{
os << "{\n"
<< " \"name\": \"" << obj.name_ << "\",\n"
<< " \"problem\": " << obj.problem_ << ",\n"
<< " \"perf_result\": " << obj.perf_result_ << "\n"
<< "}";
return os;
}
};
struct Setting
{
int n_warmup_;
int n_repeat_;
bool is_gpu_timer_;
int verify_;
int init_method_;
bool log_;
std::string csv_filename_;
bool flush_cache_;
int rotating_count_;
bool json_output_;
};
inline std::string get_rocm_version()
{
std::ifstream version_file("/opt/rocm/.info/version");
if(version_file.is_open())
{
std::string version;
std::getline(version_file, version);
return version;
}
return "Unknown";
}
template <typename ADataType_,
typename BDataType_,
typename AQDataType_,
typename BQDataType_,
typename AccDataType_,
typename CDataType_>
auto calculate_rtol_atol_abquant(const ck_tile::index_t K,
const ck_tile::index_t kbatch,
const float max_accumulated_value)
{
// Both A and B are the same FP8/BF8 type for abquant; assert so mixed-precision additions are
// caught.
static_assert(sizeof(ADataType_) == sizeof(BDataType_),
"calculate_rtol_atol_abquant assumes equal-width A and B types");
using ComputeType = ADataType_;
const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType_, AccDataType_>(
ck_tile::integer_divide_ceil(K, kbatch));
const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType_, AccDataType_>(
max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
const auto rtol_split_k =
ck_tile::get_relative_threshold<CDataType_, CDataType_, CDataType_>(kbatch);
const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType_, CDataType_, CDataType_>(
max_accumulated_value, kbatch);
return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
}
/// @brief Compare device and host results for ABQuant GEMM
bool compare_abquant(std::string instanceName,
ck_tile::index_t K,
ck_tile::index_t kbatch,
const ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
const ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
const float max_accumulated_value =
std::abs(static_cast<float>(*std::max_element(c_m_n_host_result.mData.begin(),
c_m_n_host_result.mData.end(),
[](const auto& a, const auto& b) {
return std::abs(static_cast<float>(a)) <
std::abs(static_cast<float>(b));
})));
const auto rtol_atol = calculate_rtol_atol_abquant<ADataType,
BDataType,
AQDataType,
BQDataType,
AccDataType,
CDataType>(K, kbatch, max_accumulated_value);
bool pass = ck_tile::check_err(c_m_n_dev_result,
c_m_n_host_result,
"Error: Incorrect results!",
rtol_atol.at(ck_tile::number<0>{}),
rtol_atol.at(ck_tile::number<1>{}));
std::cout << "For " << instanceName << " Relative error threshold is "
<< rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
<< rtol_atol.at(ck_tile::number<1>{}) << std::endl;
std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;
return pass;
}
/// @brief CPU reference implementation for ABQuant GEMM
void abquant_gemm_host_reference(int verify,
ck_tile::HostTensor<ADataType>& a_m_k,
ck_tile::HostTensor<AQDataType>& aq_m_qk,
ck_tile::HostTensor<BDataType>& b_k_n,
ck_tile::HostTensor<BQDataType>& bq_qk_n,
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
if(verify == 1)
{
c_m_n_host_result.SetZero();
ck_tile::reference_gemm_abquant<ADataType,
AQDataType,
BDataType,
BQDataType,
AccDataType,
CDataType,
SelectedKernel::AQuantGroupSize,
SelectedKernel::BQuantGroupSize>(
a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result);
}
}
#pragma clang diagnostic pop

View File

@@ -0,0 +1,477 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import sys
import json
import subprocess
import argparse
import time
from pathlib import Path
from typing import List, Dict, Tuple, Optional
class ABQuantGemmBenchmark:
def __init__(self, build_dir: str, verbose: bool = False):
self.build_dir = Path(build_dir)
self.verbose = verbose
self.results = []
def discover_kernels(self) -> List[Path]:
"""Find all benchmark_gemm_abquant_* executables in the build directory."""
bin_dir = self.build_dir / "bin"
if not bin_dir.exists():
print(f"Error: Binary directory {bin_dir} does not exist")
return []
kernels = list(bin_dir.glob("benchmark_gemm_abquant_*"))
if self.verbose:
print(f"Found {len(kernels)} ABQuant kernel executables")
for k in kernels:
print(f" - {k.name}")
return kernels
def extract_kernel_info(self, kernel_path: Path) -> Dict[str, str]:
"""Extract kernel information from filename."""
name = kernel_path.stem
info = {
"executable": str(kernel_path),
"name": name,
"data_type": "unknown",
"layout": "unknown",
"pipeline": "unknown",
"scheduler": "unknown",
"epilogue": "unknown",
"a_preshuffle_quant": False,
"b_preshuffle_quant": False,
"group_size_n": 1,
}
# Parse: benchmark_gemm_abquant_fp8_rcr_compv3_default_intrawave_False_False_False_False_False_gsn1_128x128x128_...
parts = name.split("_")
if len(parts) >= 3:
# Skip "benchmark_gemm_abquant" prefix (3 parts)
idx = 3
if idx < len(parts):
info["data_type"] = parts[idx]
idx += 1
if idx < len(parts):
info["layout"] = parts[idx]
idx += 1
if idx < len(parts):
info["pipeline"] = parts[idx]
idx += 1
if idx < len(parts):
info["epilogue"] = parts[idx]
idx += 1
if idx < len(parts):
info["scheduler"] = parts[idx]
idx += 1
# Skip pad_m, pad_n, pad_k
idx += 3
if idx < len(parts):
info["a_preshuffle_quant"] = parts[idx] == "True"
idx += 1
if idx < len(parts):
info["b_preshuffle_quant"] = parts[idx] == "True"
idx += 1
# Parse gsn value (e.g. gsn1, gsn128)
if idx < len(parts) and parts[idx].startswith("gsn"):
info["group_size_n"] = int(parts[idx][3:])
config_info = self.parse_detailed_config(name)
info.update(config_info)
# Warn on fields that failed to parse — silent "unknown"/zero values hide name format drift.
unknown_fields = [k for k in ("data_type", "layout", "pipeline", "scheduler", "epilogue")
if info.get(k) == "unknown"]
if unknown_fields:
print(f"Warning: could not parse {unknown_fields} from kernel name: {name}")
tile_sizes = info.get("tile_sizes", {})
if any(v == 0 for v in tile_sizes.values()):
print(f"Warning: zero tile dimension parsed from kernel name: {name} -> {tile_sizes}")
info["config_id"] = self.generate_config_id(info)
return info
def parse_detailed_config(self, kernel_name: str) -> Dict:
"""Parse tile dimensions and boolean flags from kernel name."""
config = {
"tile_sizes": {"tile_m": 0, "tile_n": 0, "tile_k": 0},
"warp_config": {"warp_m": 0, "warp_n": 0, "warp_k": 0},
"warp_tile": {"warp_tile_m": 0, "warp_tile_n": 0, "warp_tile_k": 0},
"optimization_flags": {
"pad_m": False,
"pad_n": False,
"pad_k": False,
"a_preshuffle_quant": False,
"b_preshuffle_quant": False,
},
}
parts = kernel_name.split("_")
# Extract boolean flags (5 consecutive True/False values)
bool_sequence = []
for i, part in enumerate(parts):
if part in ["True", "False"]:
bool_sequence.append(part == "True")
j = i + 1
while j < len(parts) and parts[j] in ["True", "False"]:
bool_sequence.append(parts[j] == "True")
j += 1
break
if len(bool_sequence) >= 5:
config["optimization_flags"]["pad_m"] = bool_sequence[0]
config["optimization_flags"]["pad_n"] = bool_sequence[1]
config["optimization_flags"]["pad_k"] = bool_sequence[2]
config["optimization_flags"]["a_preshuffle_quant"] = bool_sequence[3]
config["optimization_flags"]["b_preshuffle_quant"] = bool_sequence[4]
# Extract dimension groups (e.g., 128x128x128) in positional order.
# Kernel names encode groups as: tile_MxNxK_warp_MxNxK_warp_tile_MxNxK
dimension_groups = []
for part in parts:
if "x" in part and len(part.split("x")) == 3:
try:
dims = [int(x) for x in part.split("x")]
if all(d > 0 for d in dims):
dimension_groups.append(dims)
except ValueError:
continue
if len(dimension_groups) >= 3:
config["tile_sizes"]["tile_m"] = dimension_groups[0][0]
config["tile_sizes"]["tile_n"] = dimension_groups[0][1]
config["tile_sizes"]["tile_k"] = dimension_groups[0][2]
config["warp_config"]["warp_m"] = dimension_groups[1][0]
config["warp_config"]["warp_n"] = dimension_groups[1][1]
config["warp_config"]["warp_k"] = dimension_groups[1][2]
config["warp_tile"]["warp_tile_m"] = dimension_groups[2][0]
config["warp_tile"]["warp_tile_n"] = dimension_groups[2][1]
config["warp_tile"]["warp_tile_k"] = dimension_groups[2][2]
return config
def generate_config_id(self, info: Dict) -> str:
"""Generate a compact config ID from kernel info."""
parts = [
info.get("data_type", "unk"),
info.get("layout", "unk"),
info.get("pipeline", "unk"),
info.get("scheduler", "unk"),
]
tile_sizes = info.get("tile_sizes", {})
if tile_sizes.get("tile_m", 0) > 0:
parts.append(
f"{tile_sizes['tile_m']}x{tile_sizes['tile_n']}x{tile_sizes['tile_k']}"
)
gsn = info.get("group_size_n", 1)
parts.append(f"gsn{gsn}")
return "_".join(parts)
def run_kernel(self, kernel_path: Path, params: Dict[str, str]) -> Optional[Dict]:
"""Run a single kernel with given parameters."""
results_dir = self.build_dir / "results"
results_dir.mkdir(exist_ok=True)
json_file = results_dir / f"{kernel_path.stem}.json"
cmd = [str(kernel_path)]
for key, value in params.items():
cmd.append(f"-{key}={value}")
cmd.append("-json_output=true")
if self.verbose:
print(f"Running: {' '.join(cmd)}")
try:
result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
if result.returncode != 0:
print(f"Error running {kernel_path.name}: {result.stderr}")
return None
output = result.stdout.strip()
if output:
with open(json_file, "w") as f:
f.write(output)
return self.parse_json_file(json_file)
else:
print(f"No output from {kernel_path.name}")
return None
except subprocess.TimeoutExpired:
print(f"Timeout running {kernel_path.name}")
return None
except Exception as e:
print(f"Error running {kernel_path.name}: {e}")
return None
def parse_json_file(self, json_file: Path) -> Optional[Dict]:
"""Parse JSON data from kernel output file."""
try:
with open(json_file, "r") as f:
content = f.read().strip()
data = json.loads(content)
result = data.copy()
if "perf_result" in data:
perf = data["perf_result"]
result["time_ms"] = perf.get("latency(ms)", 0)
result["tflops"] = perf.get("tflops(TFlops)", 0)
result["bandwidth_gb_s"] = perf.get("bandwidth(GB/s)", 0)
return result
except (json.JSONDecodeError, Exception) as e:
if self.verbose:
print(f"Failed to parse JSON from {json_file}: {e}")
return None
def benchmark_problem_size(
self,
kernels: List[Path],
m: int,
n: int,
k: int,
group_size_k: int = 128,
verify: int = 0,
warmup: int = 50,
repeat: int = 100,
flush_cache: bool = True,
rotating_count: int = 1000,
) -> List[Dict]:
"""Benchmark all kernels for a specific problem size."""
results = []
print(f"\nBenchmarking M={m}, N={n}, K={k}, group_size_k={group_size_k}")
for kernel_path in kernels:
kernel_info = self.extract_kernel_info(kernel_path)
group_size_n = kernel_info.get("group_size_n", 1)
params = {
"m": m,
"n": n,
"k": k,
"group_size_k": group_size_k,
"group_size_n": group_size_n,
"verify": verify,
"warmup": warmup,
"repeat": repeat,
"flush_cache": str(flush_cache).lower(),
"rotating_count": rotating_count,
}
result = self.run_kernel(kernel_path, params)
if result:
structured_result = {
"name": kernel_info["name"],
"config_id": kernel_info["config_id"],
"problem": result.get("problem", {}),
"perf_result": result.get("perf_result", {}),
"config": {
"data_type": kernel_info["data_type"],
"layout": kernel_info["layout"],
"pipeline": kernel_info["pipeline"],
"scheduler": kernel_info["scheduler"],
"epilogue": kernel_info["epilogue"],
"tile_sizes": kernel_info.get("tile_sizes", {}),
"warp_config": kernel_info.get("warp_config", {}),
"warp_tile": kernel_info.get("warp_tile", {}),
"optimization_flags": kernel_info.get("optimization_flags", {}),
},
"executable": kernel_info["executable"],
"time_ms": result.get("time_ms", 0),
"tflops": result.get("tflops", 0),
"bandwidth_gb_s": result.get("bandwidth_gb_s", 0),
}
results.append(structured_result)
if self.verbose:
print(
f" {kernel_info['config_id']}: "
f"{structured_result['tflops']:.2f} TFLOPS, "
f"{structured_result['bandwidth_gb_s']:.2f} GB/s, "
f"{structured_result['time_ms']:.2f}ms"
)
return results
def find_best_kernel(
self, results: List[Dict], metric: str = "tflops"
) -> Optional[Dict]:
"""Find the best performing kernel based on metric."""
if not results:
return None
if metric == "tflops":
return max(results, key=lambda x: x.get("tflops", 0))
elif metric == "time_ms":
return min(results, key=lambda x: x.get("time_ms", float("inf")))
elif metric == "bandwidth_gb_s":
return max(results, key=lambda x: x.get("bandwidth_gb_s", 0))
else:
raise ValueError(f"Unknown metric: {metric}")
def benchmark_sweep(
self,
problem_sizes: List[Tuple[int, int, int]],
group_size_k: int = 128,
verify: bool = False,
warmup: int = 50,
repeat: int = 100,
flush_cache: bool = True,
rotating_count: int = 1000,
) -> Dict:
"""Run comprehensive benchmark sweep."""
kernels = self.discover_kernels()
if not kernels:
print("No kernels found!")
return {}
all_results = []
best_kernels = {}
for m, n, k in problem_sizes:
results = self.benchmark_problem_size(
kernels,
m,
n,
k,
group_size_k=group_size_k,
verify=1 if verify else 0,
warmup=warmup,
repeat=repeat,
flush_cache=flush_cache,
rotating_count=rotating_count,
)
all_results.extend(results)
best = self.find_best_kernel(results)
if best:
key = f"m{m}_n{n}_k{k}"
best_kernels[key] = best
print(
f"Best for {key}: {best['name']} "
f"({best['tflops']:.2f} TFLOPS, "
f"{best['bandwidth_gb_s']:.2f} GB/s, "
f"{best['time_ms']:.2f}ms)"
)
self.results = all_results
return best_kernels
def export_json(self, filename: str, best_kernels: Dict = None):
"""Export results to JSON."""
from datetime import datetime
output_data = {
"benchmark_metadata": {
"timestamp": datetime.now().isoformat(),
"operator": "gemm_abquant",
"total_kernels_tested": len(self.results),
"successful_runs": len(
[r for r in self.results if r.get("tflops", 0) > 0]
),
},
"kernel_results": self.results,
"best_kernels_by_problem": best_kernels or {},
}
with open(filename, "w") as f:
json.dump(output_data, f, indent=2)
print(f"JSON results exported to {filename}")
def main():
parser = argparse.ArgumentParser(
description="ABQuant GEMM Kernel Benchmarking Tool"
)
parser.add_argument(
"build_dir", help="Build directory containing kernel executables"
)
parser.add_argument(
"--problem-sizes",
nargs="+",
default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"],
help="Problem sizes as M,N,K tuples",
)
parser.add_argument(
"--group-size-k",
type=int,
default=128,
help="Quantization group size along K (default: 128)",
)
parser.add_argument("--verify", action="store_true", help="Enable verification")
parser.add_argument("--verbose", action="store_true", help="Verbose output")
parser.add_argument(
"--warmup",
type=int,
default=50,
help="Number of warmup iterations",
)
parser.add_argument(
"--repeat",
type=int,
default=100,
help="Number of benchmark iterations",
)
parser.add_argument(
"--no-flush-cache",
action="store_true",
help="Disable cache flushing between iterations",
)
parser.add_argument(
"--rotating-count",
type=int,
default=1000,
help="Number of iterations to rotate the cache (default: 1000)",
)
parser.add_argument("--json", help="JSON output filename (optional)")
args = parser.parse_args()
problem_sizes = []
for size_str in args.problem_sizes:
try:
m, n, k = map(int, size_str.split(","))
problem_sizes.append((m, n, k))
except ValueError:
print(f"Invalid problem size: {size_str}")
return 1
benchmark = ABQuantGemmBenchmark(args.build_dir, verbose=args.verbose)
print("Starting ABQuant GEMM kernel benchmark sweep...")
start_time = time.time()
best_kernels = benchmark.benchmark_sweep(
problem_sizes=problem_sizes,
group_size_k=args.group_size_k,
verify=args.verify,
warmup=args.warmup,
repeat=args.repeat,
flush_cache=not args.no_flush_cache,
rotating_count=args.rotating_count,
)
elapsed_time = time.time() - start_time
print(f"\nBenchmark completed in {elapsed_time:.2f} seconds")
if args.json:
benchmark.export_json(args.json, best_kernels)
return 0
if __name__ == "__main__":
sys.exit(main())

View File

@@ -0,0 +1,178 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <iostream>
#include <functional>
#include <tuple>
#include <exception>
#include <sstream>
#include <vector>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "gemm_abquant_profiler.hpp"
#include "gemm_abquant_common.hpp"
// The kernel header is included via the compile command line with -include flag
// It defines SelectedKernel struct, KERNEL_NAME, and type aliases:
// ADataType, BDataType, AQDataType, BQDataType, AccDataType, CDataType
// ALayout, BLayout, CLayout, AQLayout, BQLayout
inline auto create_args(int argc, char* argv[])
{
ck_tile::ArgParser arg_parser;
arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.")
.insert("n", "4096", "The value for n dimension. Default is 4096.")
.insert("k", "2048", "The value for k dimension. Default is 2048.")
.insert("stride_a", "0", "The stride value for tensor A. Default is 0.")
.insert("stride_b", "0", "The stride value for tensor B. Default is 0.")
.insert("stride_c", "0", "The stride value for tensor C. Default is 0.")
.insert("split_k", "1", "The split value for k dimension. Default is 1.")
.insert("group_size_k",
std::to_string(SelectedKernel::GroupSizeK),
"The quantization group size along K. Default matches kernel config.")
.insert("group_size_n",
std::to_string(SelectedKernel::GroupSizeN),
"The quantization group size along N for BQ. Default matches kernel config.")
.insert("verify",
"1",
"The type of validation. Set to 0 for no validation, 1 for validation on CPU. "
"Default is 1, CPU validation.")
.insert("log",
"false",
"Whether output kernel instance information or not. Possible values are true or "
"false. Default is false")
.insert(
"warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.")
.insert(
"repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.")
.insert("timer",
"true",
"Whether if the timer is gpu timer or not. Possible values are false or true. "
"Default is true.")
.insert("init",
"0",
"The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 "
"for constant(1). Default is 0, random.")
.insert("flush_cache",
"true",
"To flush cache, possible values are true or false. Default is true.")
.insert(
"rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.")
.insert("metric",
"0",
"Metric with which to measure kernel performance. Set to 0 for latency, 1 for "
"tflops, or 2 for bandwidth. Default is 0, latency.")
.insert("csv_filename",
"",
"The filename of benchmark result. Default is empty (no CSV output).")
.insert("json_output",
"false",
"Whether to output results in JSON format only. Possible values are true or false. "
"Default is false");
bool result = arg_parser.parse(argc, argv);
return std::make_tuple(result, arg_parser);
}
void benchmark_single(const ck_tile::ArgParser& arg_parser)
{
std::string dtype_a = DataTypeTraits<ADataType>::name;
std::string dtype_b = DataTypeTraits<BDataType>::name;
std::string dtype_aq = DataTypeTraits<AQDataType>::name;
std::string dtype_bq = DataTypeTraits<BQDataType>::name;
std::string dtype_acc = DataTypeTraits<AccDataType>::name;
std::string dtype_c = DataTypeTraits<CDataType>::name;
std::string layout_a = ALayout::name;
std::string layout_b = BLayout::name;
std::string layout_c = CLayout::name;
int M = arg_parser.get_int("m");
int N = arg_parser.get_int("n");
int K = arg_parser.get_int("k");
int group_size_k = arg_parser.get_int("group_size_k");
int group_size_n = arg_parser.get_int("group_size_n");
if(M <= 0 || N <= 0 || K <= 0)
{
throw std::invalid_argument("m, n, k must be positive integers");
}
if(group_size_k <= 0 || K % group_size_k != 0)
{
throw std::invalid_argument(
"group_size_k must be positive and k must be divisible by group_size_k");
}
if(group_size_n > 1 && N % group_size_n != 0)
{
throw std::invalid_argument("n must be divisible by group_size_n when group_size_n > 1");
}
ABQuantGemmProblem problem{arg_parser.get_int("split_k"),
M,
N,
K,
arg_parser.get_int("stride_a"),
arg_parser.get_int("stride_b"),
arg_parser.get_int("stride_c"),
0, // stride_aq computed by profiler
0, // stride_bq computed by profiler
group_size_k,
group_size_n,
dtype_a,
dtype_b,
dtype_aq,
dtype_bq,
dtype_acc,
dtype_c,
layout_a,
layout_b,
layout_c};
Setting setting{arg_parser.get_int("warmup"),
arg_parser.get_int("repeat"),
arg_parser.get_bool("timer"),
arg_parser.get_int("verify"),
arg_parser.get_int("init"),
arg_parser.get_bool("log"),
arg_parser.get_str("csv_filename"),
arg_parser.get_bool("flush_cache"),
arg_parser.get_int("rotating_count"),
arg_parser.get_bool("json_output")};
auto& profiler = ABQuantGemmProfiler::instance(setting);
try
{
auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args,
const ck_tile::stream_config& stream) {
return SelectedKernel::launch(args, stream);
};
profiler.benchmark(problem, kernel_func);
profiler.select_best_instance(static_cast<Metric>(arg_parser.get_int("metric")));
}
catch(const std::exception& e)
{
std::cerr << "Benchmark failed: " << e.what() << std::endl;
}
}
int main(int argc, char* argv[])
{
try
{
auto [result, parser] = create_args(argc, argv);
if(!result)
return EXIT_FAILURE;
benchmark_single(parser);
return 0;
}
catch(const std::exception& e)
{
std::cerr << "Error: " << e.what() << "\n";
return EXIT_FAILURE;
}
}

View File

@@ -0,0 +1,76 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <string>
#include <type_traits>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/core/numeric/integer.hpp"
// DataTypeTraits for all supported types
template <typename T>
struct DataTypeTraits;
template <>
struct DataTypeTraits<float>
{
static constexpr const char* name = "fp32";
};
template <>
struct DataTypeTraits<ck_tile::half_t>
{
static constexpr const char* name = "fp16";
};
template <>
struct DataTypeTraits<ck_tile::bf16_t>
{
static constexpr const char* name = "bf16";
};
template <>
struct DataTypeTraits<ck_tile::fp8_t>
{
static constexpr const char* name = "fp8";
};
template <>
struct DataTypeTraits<ck_tile::bf8_t>
{
static constexpr const char* name = "bf8";
};
// Helper function to determine if a layout is row-major
template <typename Layout>
constexpr auto is_row_major(Layout)
{
return ck_tile::bool_constant<std::is_same_v<Layout, ck_tile::tensor_layout::gemm::RowMajor>>{};
}
// Structure to hold kernel traits for dispatcher
struct ABQuantKernelTraits
{
std::string pipeline; // compv3
std::string scheduler; // intrawave
std::string epilogue; // default, cshuffle
bool pad_m;
bool pad_n;
bool pad_k;
bool a_preshuffle_quant;
bool b_preshuffle_quant;
ABQuantKernelTraits()
: pipeline("compv3"),
scheduler("intrawave"),
epilogue("default"),
pad_m(false),
pad_n(false),
pad_k(false),
a_preshuffle_quant(false),
b_preshuffle_quant(false)
{
}
};

View File

@@ -0,0 +1,904 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import os
import argparse
import importlib.util
import multiprocessing
import concurrent.futures
import itertools
import logging
def _import_gemm_kernel_builder():
"""Import the base GemmKernelBuilder from the gemm directory."""
current_dir = os.path.dirname(os.path.abspath(__file__))
gemm_dir = os.path.dirname(os.path.dirname(current_dir))
spec = importlib.util.spec_from_file_location(
"gemm_instance_builder",
os.path.join(gemm_dir, "gemm_instance_builder.py"),
)
gemm_builder_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(gemm_builder_module)
return gemm_builder_module.GemmKernelBuilder
def _import_validation_utils():
"""Import validation utilities."""
current_dir = os.path.dirname(os.path.abspath(__file__))
gemm_dir = os.path.dirname(os.path.dirname(current_dir))
parent_dir = os.path.dirname(gemm_dir)
spec = importlib.util.spec_from_file_location(
"validation_utils",
os.path.join(parent_dir, "gemm", "gemm_validation_utils.py"),
)
validation_utils = importlib.util.module_from_spec(spec)
spec.loader.exec_module(validation_utils)
return validation_utils
GemmKernelBuilder = _import_gemm_kernel_builder()
_validation_utils = _import_validation_utils()
is_trait_combination_valid = _validation_utils.is_trait_combination_valid
class GemmABQuantKernelBuilder(GemmKernelBuilder):
def __init__(
self,
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json=None,
max_instances=None,
seed=None,
tier=None,
manifest_path=None,
):
super().__init__(
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json,
max_instances=max_instances,
seed=seed,
tier=tier,
manifest_path=manifest_path,
)
self.group_size_k = self.config.get("group_size_k", 128)
# group_size_n can be an int or a dict with "values" key
group_size_n_cfg = self.config.get("group_size_n", 1)
if isinstance(group_size_n_cfg, dict):
self.group_size_n_values = group_size_n_cfg.get("values", [1])
else:
self.group_size_n_values = [group_size_n_cfg]
@staticmethod
def _tile_config_to_str(tile_config):
return (
f"{tile_config['tile_m']}x{tile_config['tile_n']}x{tile_config['tile_k']}_"
f"{tile_config['warp_m']}x{tile_config['warp_n']}x{tile_config['warp_k']}_"
f"{tile_config['warp_tile_m']}x{tile_config['warp_tile_n']}x{tile_config['warp_tile_k']}"
)
def _apply_sampling(self, kernel_list):
"""Apply RFC Sobol+LHS+maximin sampling for ABQuant kernels.
Overrides base class to handle ABQuant's 8-tuple trait_combo
and group_size_n field.
"""
if self.max_instances is None or len(kernel_list) <= self.max_instances:
return kernel_list
import sys
sampling_parent = os.path.join(
os.path.dirname(__file__), "..", "..", "..", ".."
)
if sampling_parent not in sys.path:
sys.path.insert(0, sampling_parent)
from sampling.sampler import sample_feasible_set
from sampling.seed import make_seed
from sampling.feasible_set import GEMM_ABQUANT_AXES
effective_seed = make_seed(
self.seed, self.gpu_target, self.datatype, self.layout
)
flat_items = []
for k in kernel_list:
flat = dict(k["tile_config"])
(
pipeline,
epilogue,
scheduler,
pad_m,
pad_n,
pad_k,
a_preshuffle_quant,
b_preshuffle_quant,
) = k["trait_combo"]
flat.update(
{
"pipeline": pipeline,
"epilogue": epilogue,
"scheduler": scheduler,
"pad_m": pad_m,
"pad_n": pad_n,
"pad_k": pad_k,
"a_preshuffle_quant": a_preshuffle_quant,
"b_preshuffle_quant": b_preshuffle_quant,
"group_size_n": k.get("group_size_n", 1),
}
)
flat_items.append(flat)
selected, method, selected_indices = sample_feasible_set(
flat_items,
self.max_instances,
effective_seed,
GEMM_ABQUANT_AXES,
)
kernel_list = [kernel_list[i] for i in selected_indices]
if self.manifest_path:
from sampling.manifest import write_manifest
write_manifest(
selected,
self.manifest_path,
self.kernel_name_prefix,
self.datatype,
self.layout,
self.gpu_target,
effective_seed,
self.tier or "daily",
method,
)
print(
f"Sampled {len(kernel_list)} from feasible set "
f"(budget={self.max_instances}, seed={effective_seed}, method={method})"
)
return kernel_list
def _get_group_size_n_values(self):
"""Return the list of group_size_n values to generate kernels for."""
return self.group_size_n_values
def _generate_trait_combinations(self):
"""Generate all combinations of traits for ABQuant.
ABQuant uses 8-tuples:
(pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant, b_preshuffle_quant)
"""
trait_config = self.config["trait_config"]
pipelines = trait_config.get("pipeline").get("values")
epilogues = trait_config.get("epilogue").get("values")
schedulers = trait_config.get("scheduler").get("values")
pad_m_values = trait_config.get("pad_m").get("values")
pad_n_values = trait_config.get("pad_n").get("values")
pad_k_values = trait_config.get("pad_k").get("values")
a_preshuffle_values = trait_config.get("a_preshuffle_quant").get("values")
b_preshuffle_values = trait_config.get("b_preshuffle_quant").get("values")
all_combinations = list(
itertools.product(
pipelines,
epilogues,
schedulers,
pad_m_values,
pad_n_values,
pad_k_values,
a_preshuffle_values,
b_preshuffle_values,
)
)
# Filter out unsupported trait combinations
combinations = []
for combo in all_combinations:
pipeline, epilogue, scheduler = combo[:3]
a_preshuffle = combo[6]
b_preshuffle = combo[7]
if is_trait_combination_valid(
pipeline,
epilogue,
scheduler,
b_preshuffle,
self.kernel_name_prefix,
self.layout,
):
combinations.append(combo)
else:
logging.debug(
f"Skipping unsupported ABQuant trait combination: "
f"{pipeline}-{epilogue}-{scheduler}-a_pre={a_preshuffle}-b_pre={b_preshuffle}"
)
return combinations
def _list_kernels(self):
"""Write kernel list to file for CMake to read."""
tile_configs = self._get_tile_configs()
trait_combos = self._generate_trait_combinations()
group_size_n_values = self._get_group_size_n_values()
kernel_list = []
for tile_config in tile_configs:
for trait_combo in trait_combos:
for group_size_n in group_size_n_values:
(
pipeline,
epilogue,
scheduler,
pad_m,
pad_n,
pad_k,
a_preshuffle,
b_preshuffle,
) = trait_combo
# Create kernel name with proper boolean capitalization
kernel_name = (
f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_"
f"{pipeline}_{epilogue}_{scheduler}_"
f"{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_"
f"{str(a_preshuffle).capitalize()}_{str(b_preshuffle).capitalize()}_"
f"gsn{group_size_n}"
)
tile_str = self._tile_config_to_str(tile_config)
kernel_name += f"_{tile_str}"
kernel_list.append(
{
"name": kernel_name,
"tile_config": tile_config,
"trait_combo": trait_combo,
"group_size_n": group_size_n,
}
)
# Apply RFC-compliant sampling (Sobol + LHS + maximin)
kernel_list = self._apply_sampling(kernel_list)
# Write kernel count
with open(
self.working_path / f"{self.kernel_name_prefix}_kernel_count.txt", "w"
) as f:
f.write(str(len(kernel_list)))
# Write kernel list
with open(
self.working_path / f"{self.kernel_name_prefix}_kernel_list.txt", "w"
) as f:
for kernel in kernel_list:
tile_config = kernel["tile_config"]
trait_combo = kernel["trait_combo"]
group_size_n = kernel["group_size_n"]
tile_str = self._tile_config_to_str(tile_config)
trait_str = (
f"{trait_combo[0]}_{trait_combo[1]}_{trait_combo[2]}_"
+ "_".join(str(x) for x in trait_combo[3:])
+ f"_gsn{group_size_n}"
)
f.write(f"{kernel['name']}|{tile_str}|{trait_str}\n")
print(f"Listed {len(kernel_list)} kernel configurations")
def _generate_kernel_instance(self, tile_config, trait_combo, group_size_n=None):
"""Generate a single kernel instance for ABQuant.
trait_combo is an 8-tuple for ABQuant.
group_size_n is the BQ group size along N dimension.
"""
if group_size_n is None:
group_size_n = self.group_size_n_values[0]
k_block_per_cu = self.config.get("k_block_per_cu", 1)
(
pipeline,
epilogue,
scheduler,
pad_m,
pad_n,
pad_k,
a_preshuffle,
b_preshuffle,
) = trait_combo
# Create kernel name
kernel_name = (
f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_"
f"{pipeline}_{epilogue}_{scheduler}_"
f"{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_"
f"{str(a_preshuffle).capitalize()}_{str(b_preshuffle).capitalize()}_"
f"gsn{group_size_n}"
)
tile_str = self._tile_config_to_str(tile_config)
kernel_name += f"_{tile_str}"
# Pipeline maps
pipeline_impl_map = {
"compv3": "ck_tile::ABQuantGemmPipelineAgBgCrCompV3",
}
base_pipeline_map = {
"compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3",
}
scheduler_type_map = {
"intrawave": "ck_tile::GemmPipelineScheduler::Intrawave",
"interwave": "ck_tile::GemmPipelineScheduler::Interwave",
"default": "ck_tile::GemmPipelineScheduler::Default",
}
instance_code = self.populate_kernel_header(kernel_name)
instance_code += self.populate_kernel_dtype_layout()
instance_code += self.populate_strut_begin(kernel_name)
instance_code += self.populate_tile_config(tile_config)
instance_code += self._populate_abquant_trait_config(trait_combo, group_size_n)
instance_code += self._populate_abquant_initialization(
base_pipeline_map, pipeline
)
instance_code += self._populate_launch_abquant(
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
)
# Write into a file
simplified_name = kernel_name
if simplified_name.startswith(f"{self.kernel_name_prefix}_"):
simplified_name = simplified_name[len(self.kernel_name_prefix) + 1 :]
header_file = (
self.working_path
/ f"{self.kernel_name_prefix}_single_{simplified_name}.hpp"
)
with open(header_file, "w") as f:
f.write(instance_code)
print(f"Generated {header_file}")
return kernel_name, instance_code
def populate_kernel_header(self, kernel_name):
instance_code = f"""// Generated kernel instance for {kernel_name}
#pragma once
#include <cstdint>
#include <utility>
#include <tuple>
#include "ck_tile/core.hpp"
#include "ck_tile/host/kernel_launch.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp"
#include "ck_tile/ops/common/tensor_layout.hpp"
#include "ck_tile/ops/epilogue/default_2d_epilogue.hpp"
#include "ck_tile/ops/epilogue/cshuffle_epilogue.hpp"
#include "ck_tile/ops/gemm_quant.hpp"
#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp"
"""
return instance_code
def populate_kernel_dtype_layout(self):
a_layout, b_layout, c_layout = _validation_utils.get_abc_layouts(self.layout)
dtype_str = _validation_utils.get_dtype_string(self.datatype)
c_type_str = _validation_utils.get_dtype_string(
"fp16" if self.datatype in ["fp8", "bf8"] else self.datatype
)
instance_code = f"""
using ADataType = {dtype_str};
using BDataType = {dtype_str};
using AccDataType = float;
using CDataType = {c_type_str};
using AQDataType = float;
using BQDataType = float;
using ALayout = {a_layout};
using BLayout = {b_layout};
using CLayout = {c_layout};
using AQLayout = ck_tile::tensor_layout::gemm::RowMajor;
using BQLayout = ck_tile::tensor_layout::gemm::ColumnMajor;
"""
return instance_code
def _populate_abquant_trait_config(self, trait_combo, group_size_n):
(
pipeline,
epilogue,
scheduler,
pad_m,
pad_n,
pad_k,
a_preshuffle,
b_preshuffle,
) = trait_combo
instance_code = f"""
// Traits configurations
static constexpr bool kPadM = {"true" if pad_m in [True, "true"] else "false"};
static constexpr bool kPadN = {"true" if pad_n in [True, "true"] else "false"};
static constexpr bool kPadK = {"true" if pad_k in [True, "true"] else "false"};
static constexpr bool TransposeC = false;
static constexpr bool DoubleSmemBuffer = false;
static constexpr bool APreshuffleQuant = {"true" if a_preshuffle in [True, "true"] else "false"};
static constexpr bool BPreshuffleQuant = {"true" if b_preshuffle in [True, "true"] else "false"};
static constexpr bool PreshuffleB = false;
static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k};
static constexpr ck_tile::index_t GroupSizeN = {group_size_n};"""
return instance_code
def _populate_abquant_initialization(self, base_pipeline_map, pipeline):
instance_code = """
// Tile shape
using TileShape = ck_tile::TileGemmShape<
ck_tile::sequence<TileM, TileN, TileK>,
ck_tile::sequence<WarpPerBlock_M, WarpPerBlock_N, WarpPerBlock_K>,
ck_tile::sequence<WarpTileM, WarpTileN, WarpTileK>>;
// Tile partitioner
using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner<TileShape, 8, 4>;
// Quantization group sizes
using AQuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, GroupSizeK>>;
using BQuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, GroupSizeN, GroupSizeK>>;
// Traits
using Traits = ck_tile::TileGemmQuantTraits<
kPadM, kPadN, kPadK,
APreshuffleQuant,
BPreshuffleQuant,
PreshuffleB,
ALayout, BLayout, CLayout,
ck_tile::QuantType::ABQuantGrouped,
AQLayout, BQLayout>;
// Pipeline problem (base, for hot loop detection)
using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<
ADataType,
BDataType,
AccDataType,
TileShape,
Traits,
BDataType>;"""
instance_code += f"""
// Base pipeline for hot loop detection
using BaseGemmPipeline = {base_pipeline_map.get(pipeline)}<GemmPipelineProblem>;"""
return instance_code
def _populate_launch_abquant(
self,
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
):
"""Generate the complete launch function for ABQuant kernels."""
instance_code = """
// Launch function
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
// Hot loop detection
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
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);"""
instance_code += f"""
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
// Full pipeline problem with hot loop and tail number
using PipelineProblem = ck_tile::GemmABQuantPipelineProblem<
ADataType,
AQDataType,
BDataType,
BQDataType,
AccDataType,
TileShape,
Traits,
AQuantGroupSize,
BQuantGroupSize,
false,
ADataType,
{scheduler_type_map.get(scheduler)},
has_hot_loop_v,
tail_number_v>;
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
instance_code += """
// Epilogue"""
if epilogue == "cshuffle":
instance_code += """
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<
ADataType,
BDataType,
ck_tile::tuple<>, // DsDataType
AccDataType,
CDataType,
ck_tile::tuple<>, // DsLayout
CLayout,
ck_tile::element_wise::PassThrough,
TileM, // kM_
TileN, // kN_
WarpPerBlock_M, // MWave_
WarpPerBlock_N, // NWave_
WarpTileM, // kMPerXdl_
WarpTileN, // kNPerXdl_
WarpTileK, // kKPerXdl_
TransposeC>; // isCTransposed_
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;"""
else:
instance_code += """
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<
ADataType,
BDataType,
ck_tile::tuple<>, // DsDataType
AccDataType,
CDataType,
ck_tile::tuple<>, // DsLayout
CLayout,
ck_tile::element_wise::PassThrough,
TileM, // kM_
TileN, // kN_
kPadM,
kPadN,
WarpTileM, // kMPerXdl_
WarpTileN, // kNPerXdl_
WarpTileK, // kKPerXdl_
TransposeC>; // isCTransposed_
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
instance_code += f"""
// Kernel type
using Kernel = ck_tile::QuantGemmKernel<
TilePartitioner, GemmPipeline, GemmEpilogue,
ck_tile::QuantType::ABQuantGrouped>;
// Kernel arguments
auto kargs = Kernel::MakeKernelArgs(args);
if (!Kernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
}}
// Get grid and block sizes
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
const dim3 blocks = Kernel::BlockSize();
if(stream.log_level_ > 0) {{
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
<< std::endl;
}}
// Launch kernel
constexpr int kBlockPerCu = {k_block_per_cu};
float ave_time = ck_tile::launch_kernel(
stream,
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}};
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
}}
}};
"""
return instance_code
def _generate_all_individual(self, num_workers=None):
"""Generate individual kernel files for separate compilation with parallel processing"""
if num_workers is None:
num_workers = min(multiprocessing.cpu_count(), 8)
tile_configs = self._get_tile_configs()
trait_combos = self._generate_trait_combinations()
group_size_n_values = self._get_group_size_n_values()
work_items = []
for tile_config in tile_configs:
for trait_combo in trait_combos:
for group_size_n in group_size_n_values:
work_items.append(
(
tile_config,
trait_combo,
group_size_n,
self.kernel_name_prefix,
self.working_path,
self.gpu_target,
self.datatype,
self.layout,
self.config_json,
)
)
# Apply RFC-compliant sampling (Sobol + LHS + maximin)
if self.max_instances is not None and len(work_items) > self.max_instances:
kernel_dicts = [
{
"tile_config": item[0],
"trait_combo": item[1],
"group_size_n": item[2],
"_work_item": item,
}
for item in work_items
]
sampled = self._apply_sampling(kernel_dicts)
work_items = [k["_work_item"] for k in sampled]
print(
f"Generating {len(work_items)} individual kernel files using {num_workers} workers..."
)
print(f" Tile configs: {len(tile_configs)}")
print(f" Trait combinations: {len(trait_combos)}")
print(f" Group size N values: {group_size_n_values}")
print(f" Total kernels: {len(work_items)}")
if work_items:
print(" First work item example:")
tile_config, trait_combo = work_items[0][:2]
print(f" Tile config: {tile_config}")
print(f" Trait combo: {trait_combo[:3]}")
kernel_list = []
completed = 0
with concurrent.futures.ProcessPoolExecutor(
max_workers=num_workers
) as executor:
print(f" Submitting {len(work_items)} tasks to executor...")
future_to_item = {
executor.submit(_generate_single_kernel_individual, item): item
for item in work_items
}
print(" All tasks submitted, waiting for completion...")
for future in concurrent.futures.as_completed(future_to_item):
completed += 1
if completed % 100 == 0 or completed == len(work_items):
print(
f" Progress: {completed}/{len(work_items)} kernels generated"
)
try:
result = future.result()
if result:
kernel_list.append(result)
except Exception as exc:
item = future_to_item[future]
print(f"Kernel generation failed for {item}: {exc}")
kernel_list.sort(key=lambda x: x[0])
self._generate_cmake_individual_targets(kernel_list)
print(
f"Generated {len(kernel_list)} individual kernel files in {self.working_path}"
)
def _generate_single_kernel_individual(work_item):
"""Worker function to generate a single individual kernel file"""
(
tile_config,
trait_combo,
group_size_n,
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json,
) = work_item
builder = GemmABQuantKernelBuilder(
kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json
)
try:
kernel_name, _ = builder._generate_kernel_instance(
tile_config, trait_combo, group_size_n
)
return (kernel_name, trait_combo, tile_config)
except Exception as e:
print(f"Error generating individual kernel: {e}")
return None
def main():
parser = argparse.ArgumentParser(description="GEMM ABQuant kernel instance builder")
parser.add_argument("--working_path", required=True, help="Working directory path")
parser.add_argument("--gpu_target", required=True, help="GPU target architecture")
parser.add_argument(
"--datatype",
required=True,
choices=["fp8", "bf8"],
help="Data type (fp8 or bf8)",
)
parser.add_argument(
"--layout",
required=True,
choices=["rcr", "rrr", "ccr", "crr"],
help="Matrix layout",
)
parser.add_argument("--config_json", help="Configuration JSON file")
parser.add_argument("--num_workers", type=int, help="Number of parallel workers")
parser.add_argument(
"--gen_all_individual",
action="store_true",
help="Generate individual kernel files",
)
parser.add_argument(
"--gen_single",
action="store_true",
help="Generate a single kernel file",
)
parser.add_argument("--kernel_name", help="Kernel name for single generation")
parser.add_argument("--tile_config", help="Tile configuration string")
parser.add_argument("--trait_combo", help="Trait combination string")
parser.add_argument(
"--list_kernels",
action="store_true",
help="List kernel configurations without generating files",
)
parser.add_argument(
"--max-instances",
type=int,
default=None,
help="Maximum number of kernel instances to select via sampling",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="RNG seed for deterministic sampling; if omitted, derived from today's date",
)
parser.add_argument(
"--tier",
default=None,
help="Sampling tier name (e.g., 'daily')",
)
parser.add_argument(
"--manifest-path",
default=None,
help="Directory to write chosen_instances manifest JSON",
)
args = parser.parse_args()
assert args.datatype in ["fp16", "bf16", "fp8", "bf8"], (
f"Invalid datatype string: {args.datatype} (supported datatypes are [fp16, bf16, fp8, and bf8])"
)
layout_parts = args.layout.lower()
assert len(layout_parts) == 3, (
f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr')"
)
assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], (
f"Invalid matrix_a layout: {layout_parts[0]} or matrix_b layout: {layout_parts[1]}"
)
assert layout_parts[2] == "r", (
f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' for row major)"
)
kernel_name_prefix = "gemm_abquant"
builder = GemmABQuantKernelBuilder(
kernel_name_prefix,
args.working_path,
args.gpu_target,
args.datatype,
args.layout,
args.config_json,
max_instances=args.max_instances,
seed=args.seed,
tier=args.tier,
manifest_path=args.manifest_path,
)
if args.list_kernels:
builder._list_kernels()
elif args.gen_single:
if not args.kernel_name or not args.tile_config or not args.trait_combo:
parser.error(
"--gen_single requires --kernel_name, --tile_config, and --trait_combo"
)
# Parse tile config
tile_parts = args.tile_config.split("_")
tile_dims = tile_parts[0].split("x")
warp_dims = tile_parts[1].split("x")
warp_tile_dims = tile_parts[2].split("x")
tile_config = {
"tile_m": int(tile_dims[0]),
"tile_n": int(tile_dims[1]),
"tile_k": int(tile_dims[2]),
"warp_m": int(warp_dims[0]),
"warp_n": int(warp_dims[1]),
"warp_k": int(warp_dims[2]),
"warp_tile_m": int(warp_tile_dims[0]),
"warp_tile_n": int(warp_tile_dims[1]),
"warp_tile_k": int(warp_tile_dims[2]),
}
# Parse trait combo for ABQuant:
# pipeline_epilogue_scheduler_padM_padN_padK_aPreshuffle_bPreshuffle_gsnN
trait_parts = args.trait_combo.split("_")
# Find the gsn part
group_size_n = 1
gsn_idx = None
for i, part in enumerate(trait_parts):
if part.startswith("gsn"):
group_size_n = int(part[3:])
gsn_idx = i
break
if gsn_idx is not None:
# Remove gsn from trait_parts for parsing
trait_parts_no_gsn = trait_parts[:gsn_idx] + trait_parts[gsn_idx + 1 :]
else:
trait_parts_no_gsn = trait_parts
if len(trait_parts_no_gsn) < 8:
parser.error(
f"--trait_combo must have 8 underscore-separated fields + gsn suffix "
f"(e.g. 'compv3_default_intrawave_False_False_False_False_False_gsn1'), "
f"got {len(trait_parts_no_gsn)} field(s): '{args.trait_combo}'"
)
trait_combo = (
trait_parts_no_gsn[0], # pipeline
trait_parts_no_gsn[1], # epilogue
trait_parts_no_gsn[2], # scheduler
trait_parts_no_gsn[3] == "True", # pad_m
trait_parts_no_gsn[4] == "True", # pad_n
trait_parts_no_gsn[5] == "True", # pad_k
trait_parts_no_gsn[6] == "True", # a_preshuffle_quant
trait_parts_no_gsn[7] == "True", # b_preshuffle_quant
)
builder._generate_kernel_instance(tile_config, trait_combo, group_size_n)
elif args.gen_all_individual:
builder._generate_all_individual(args.num_workers)
else:
parser.error(
"Must specify one of: --list_kernels, --gen_all_individual, or --gen_single"
)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,336 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <iostream>
#include <fstream>
#include <iomanip>
#include <functional>
#include <tuple>
#include <vector>
#include <algorithm>
#include <stdexcept>
#include "ck_tile/host/device_prop.hpp"
#include "ck_tile/host/tensor_shuffle_utils.hpp"
#include "ck_tile/ops/gemm_quant.hpp"
#include "gemm_abquant_benchmark.hpp"
class ABQuantGemmProfiler
{
public:
static ABQuantGemmProfiler& instance(Setting setting)
{
static ABQuantGemmProfiler instance{setting};
return instance;
}
void reset() { kernel_instances_.clear(); }
void benchmark(ABQuantGemmProblem& problem,
std::function<float(const ck_tile::QuantGemmHostArgs&,
const ck_tile::stream_config&)> kernel_func)
{
std::vector<std::function<std::tuple<std::string, float>(ck_tile::QuantGemmHostArgs&,
const ck_tile::stream_config&)>>
callables;
callables.push_back(
[kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
float time = kernel_func(args, stream);
return std::make_tuple(std::string(KERNEL_NAME), time);
});
benchmark(problem, callables);
}
void benchmark(ABQuantGemmProblem& problem,
std::vector<std::function<std::tuple<std::string, float>(
ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables)
{
const ALayout layout_a = ALayout{};
const BLayout layout_b = BLayout{};
const CLayout layout_c = CLayout{};
const AQLayout layout_aq = AQLayout{};
const BQLayout layout_bq = BQLayout{};
problem.stride_a_ = ck_tile::get_default_stride(
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a));
problem.stride_b_ = ck_tile::get_default_stride(
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b));
problem.stride_c_ = ck_tile::get_default_stride(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c));
// Compute AQ scale tensor dimensions: [M, K / group_size_k]
if(problem.k_ % problem.group_size_k_ != 0)
throw std::runtime_error("k_ must be divisible by group_size_k_");
const ck_tile::index_t QK_A = problem.k_ / problem.group_size_k_;
problem.stride_aq_ = ck_tile::get_default_stride(
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq));
// Compute BQ scale tensor dimensions: [K / group_size_k, N / group_size_n]
const ck_tile::index_t QK_B = problem.k_ / problem.group_size_k_;
const ck_tile::index_t QN_B =
(problem.group_size_n_ > 1) ? (problem.n_ / problem.group_size_n_) : problem.n_;
problem.stride_bq_ =
ck_tile::get_default_stride(QK_B, QN_B, problem.stride_bq_, is_row_major(layout_bq));
// Allocate host tensors
ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
// AQ scale tensor: [M, QK_A]
ck_tile::HostTensor<AQDataType> aq_m_qk(ck_tile::host_tensor_descriptor(
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq)));
// BQ scale tensor: [QK_B, QN_B]
ck_tile::HostTensor<BQDataType> bq_qk_n(ck_tile::host_tensor_descriptor(
QK_B, QN_B, problem.stride_bq_, is_row_major(layout_bq)));
// Initialize tensors
if(setting_.init_method_ == 0)
{
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
ck_tile::FillUniformDistribution<AQDataType>{0.5f, 1.5f}(aq_m_qk);
ck_tile::FillUniformDistribution<BQDataType>{0.5f, 1.5f}(bq_qk_n);
}
else if(setting_.init_method_ == 1)
{
ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
}
else if(setting_.init_method_ == 2)
{
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(1)}(a_m_k);
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(1)}(b_k_n);
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
}
else
{
a_m_k.SetZero();
b_k_n.SetZero();
aq_m_qk.SetZero();
bq_qk_n.SetZero();
}
// Allocate device memory
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
ck_tile::DeviceMem aq_dev_buf(aq_m_qk.get_element_space_size_in_bytes());
ck_tile::DeviceMem bq_dev_buf(bq_qk_n.get_element_space_size_in_bytes());
a_m_k_dev_buf.ToDevice(a_m_k.data());
b_k_n_dev_buf.ToDevice(b_k_n.data());
// Shuffle AQ data if preshuffle is enabled
if constexpr(SelectedKernel::APreshuffleQuant)
{
constexpr int block_aq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK;
ck_tile::HostTensor<AQDataType> aq_shuffled = ck_tile::shuffle_aq(&aq_m_qk, block_aq_k);
aq_dev_buf.ToDevice(aq_shuffled.data());
}
else
{
aq_dev_buf.ToDevice(aq_m_qk.data());
}
// Shuffle BQ data if preshuffle is enabled
if constexpr(SelectedKernel::BPreshuffleQuant)
{
constexpr int block_bq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK;
ck_tile::HostTensor<BQDataType> bq_shuffled = ck_tile::shuffle_bq(&bq_qk_n, block_bq_k);
bq_dev_buf.ToDevice(bq_shuffled.data());
}
else
{
bq_dev_buf.ToDevice(bq_qk_n.data());
}
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
// Build QuantGemmHostArgs
ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
b_k_n_dev_buf.GetDeviceBuffer(),
c_m_n_dev_buf.GetDeviceBuffer(),
aq_dev_buf.GetDeviceBuffer(),
bq_dev_buf.GetDeviceBuffer(),
problem.split_k_,
problem.m_,
problem.n_,
problem.k_,
QK_A,
QK_B,
problem.stride_a_,
problem.stride_b_,
problem.stride_c_,
problem.stride_aq_,
problem.stride_bq_);
// Host reference for verification
ck_tile::HostTensor<CDataType> c_m_n_host_result(ck_tile::host_tensor_descriptor(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
if(setting_.verify_)
{
abquant_gemm_host_reference(
setting_.verify_, a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result);
}
// Run kernel(s)
for(auto& callable : callables)
{
auto kernel_run_result = callable(gemm_args,
ck_tile::stream_config{nullptr,
true,
setting_.log_,
setting_.n_warmup_,
setting_.n_repeat_,
setting_.is_gpu_timer_,
setting_.flush_cache_,
setting_.rotating_count_});
process_result(problem,
QK_A,
QK_B,
c_m_n_dev_buf,
c_m_n_host_result,
c_m_n_dev_result,
kernel_run_result);
}
}
void process_result(const ABQuantGemmProblem& problem,
ck_tile::index_t QK_A,
ck_tile::index_t QK_B,
ck_tile::DeviceMem& c_m_n_dev_buf,
ck_tile::HostTensor<CDataType>& c_m_n_host_result,
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
const std::tuple<std::string, float>& kernel_run_result)
{
auto [name, avg_time] = kernel_run_result;
KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}};
// Compute performance metrics
std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_;
std::size_t num_byte =
sizeof(ADataType) * problem.m_ * problem.k_ +
sizeof(BDataType) * problem.n_ * problem.k_ + sizeof(AQDataType) * problem.m_ * QK_A +
sizeof(BQDataType) * QK_B *
(problem.group_size_n_ > 1 ? problem.n_ / problem.group_size_n_ : problem.n_) +
sizeof(CDataType) * problem.m_ * problem.n_;
kernel_instance.perf_result_.latency_ = avg_time;
kernel_instance.perf_result_.tflops_ = static_cast<float>(flop) / 1.E9 / avg_time;
kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time;
if(setting_.log_ > 0 && !setting_.json_output_)
{
std::cout << kernel_instance << std::endl;
}
// Verify result
bool verified_correct = true;
if(setting_.verify_)
{
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
verified_correct = compare_abquant(
name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result);
}
if(verified_correct)
{
kernel_instances_.emplace_back(kernel_instance);
}
else
{
std::cout << "Verification failed, skip kernel: " << name << std::endl;
}
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
}
KernelInstance select_best_instance(Metric metric)
{
if(kernel_instances_.empty())
throw std::runtime_error("Empty instances");
auto kernel_instance = *std::max_element(kernel_instances_.begin(),
kernel_instances_.end(),
[metric](const auto& a, const auto& b) {
return PerformanceResult::compare(
b.perf_result_, a.perf_result_, metric);
});
if(setting_.json_output_)
{
std::cout << kernel_instance << std::endl;
}
else
{
std::cout << "**********************************" << std::endl;
std::cout << "According to given metrics: " << get_metric_name(metric) << "\n"
<< "Current kernel performance is: " << kernel_instance << std::endl;
std::cout << "**********************************" << std::endl;
}
if(!setting_.csv_filename_.empty())
{
std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app);
if(!file.is_open())
{
std::cerr << "Warning: Failed to open CSV file for writing." << std::endl;
}
else
{
if(file.tellp() == 0)
{
file << "rocm_version,device_name,"
<< "split_k,m,n,k,stride_a,stride_b,stride_c,stride_aq,stride_bq,"
<< "group_size_k,group_size_n,"
<< "dtype_a,dtype_b,dtype_aq,dtype_bq,dtype_acc,dtype_c,"
<< "layout_a,layout_b,layout_c," << "name,"
<< "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n";
}
const auto& prob = kernel_instance.problem_;
const auto& perf = kernel_instance.perf_result_;
file << get_rocm_version() << "," << ck_tile::get_device_name() << ","
<< prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << ","
<< prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << ","
<< prob.stride_aq_ << "," << prob.stride_bq_ << "," << prob.group_size_k_
<< "," << prob.group_size_n_ << "," << prob.dtype_a_ << "," << prob.dtype_b_
<< "," << prob.dtype_aq_ << "," << prob.dtype_bq_ << "," << prob.dtype_acc_
<< "," << prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_
<< "," << prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed
<< std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << ","
<< perf.bandwidth_ << "," << get_metric_name(metric) << "\n";
}
}
return kernel_instance;
}
ABQuantGemmProfiler(const ABQuantGemmProfiler&) = delete;
ABQuantGemmProfiler& operator=(const ABQuantGemmProfiler&) = delete;
private:
~ABQuantGemmProfiler() { kernel_instances_.clear(); }
ABQuantGemmProfiler(Setting setting) : setting_(setting) {}
Setting setting_;
std::vector<KernelInstance> kernel_instances_;
};

View File

@@ -0,0 +1,236 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""Unit tests for GemmABQuantKernelBuilder (gemm_abquant operator)."""
import json
import os
import sys
import tempfile
import unittest
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
from gemm_abquant_instance_builder import GemmABQuantKernelBuilder # noqa: E402
_MINIMAL_CONFIG = {
"tile_config": {
"tile_m": {"values": [64]},
"tile_n": {"values": [64]},
"tile_k": {"values": [128]},
"warp_m": {"values": [2]},
"warp_n": {"values": [2]},
"warp_k": {"values": [1]},
"warp_tile_m": {"values": [16]},
"warp_tile_n": {"values": [16]},
"warp_tile_k": {"values": [64]},
},
"trait_config": {
"pipeline": {"values": ["compv3"]},
"scheduler": {"values": ["intrawave"]},
"epilogue": {"values": ["default"]},
"pad_m": {"values": [False]},
"pad_n": {"values": [False]},
"pad_k": {"values": [False]},
"a_preshuffle_quant": {"values": [False]},
"b_preshuffle_quant": {"values": [False, True]},
},
"k_block_per_cu": 1,
"group_size_k": 128,
"group_size_n": {"values": [1, 128]},
}
def _make_builder(tmpdir, config=None, **kwargs):
cfg = config if config is not None else _MINIMAL_CONFIG
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(cfg, f)
return GemmABQuantKernelBuilder(
kernel_name_prefix="gemm_abquant",
working_path=tmpdir,
gpu_target=kwargs.get("gpu_target", "gfx942"),
datatype=kwargs.get("datatype", "fp8"),
layout=kwargs.get("layout", "rcr"),
config_json=cfg_path,
)
class TestGemmABQuantBuilderInit(unittest.TestCase):
def test_group_size_k_from_config(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertEqual(builder.group_size_k, 128)
def test_group_size_n_dict_parsed(self):
"""group_size_n specified as {"values": [...]} should expand to a list."""
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertIsInstance(builder.group_size_n_values, list)
self.assertIn(1, builder.group_size_n_values)
self.assertIn(128, builder.group_size_n_values)
def test_group_size_n_scalar_parsed(self):
"""group_size_n specified as a plain int should become a single-element list."""
cfg = json.loads(json.dumps(_MINIMAL_CONFIG))
cfg["group_size_n"] = 64
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, config=cfg)
self.assertEqual(builder.group_size_n_values, [64])
def test_kernel_name_prefix_stored(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertEqual(builder.kernel_name_prefix, "gemm_abquant")
def test_working_path_created(self):
with tempfile.TemporaryDirectory() as tmpdir:
sub = os.path.join(tmpdir, "workdir")
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(_MINIMAL_CONFIG, f)
GemmABQuantKernelBuilder("gemm_abquant", sub, "gfx942", "fp8", "rcr", cfg_path)
self.assertTrue(os.path.isdir(sub))
class TestGemmABQuantTraitCombinations(unittest.TestCase):
def setUp(self):
self._tmpdir = tempfile.TemporaryDirectory()
self.builder = _make_builder(self._tmpdir.name)
def tearDown(self):
self._tmpdir.cleanup()
def test_trait_combinations_non_empty(self):
combos = self.builder._generate_trait_combinations()
self.assertGreater(len(combos), 0)
def test_trait_combo_is_8_tuple(self):
"""ABQuant uses 8-tuples:
(pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant, b_preshuffle_quant)
"""
combos = self.builder._generate_trait_combinations()
for combo in combos:
self.assertEqual(
len(combo), 8, f"Expected 8-tuple, got {len(combo)}: {combo}"
)
def test_pipeline_is_compv3(self):
combos = self.builder._generate_trait_combinations()
for combo in combos:
self.assertEqual(combo[0], "compv3")
def test_b_preshuffle_quant_values(self):
combos = self.builder._generate_trait_combinations()
b_preshuffle_vals = {c[7] for c in combos}
self.assertIn(True, b_preshuffle_vals)
self.assertIn(False, b_preshuffle_vals)
def test_a_preshuffle_quant_default_false(self):
"""Config only lists False for a_preshuffle_quant, so all combos should have False."""
combos = self.builder._generate_trait_combinations()
for combo in combos:
self.assertFalse(combo[6], "a_preshuffle_quant should be False per minimal config")
class TestGemmABQuantTileConfigs(unittest.TestCase):
def setUp(self):
self._tmpdir = tempfile.TemporaryDirectory()
self.builder = _make_builder(self._tmpdir.name)
def tearDown(self):
self._tmpdir.cleanup()
def test_tile_configs_non_empty(self):
configs = self.builder._get_tile_configs()
self.assertGreater(len(configs), 0)
def test_tile_config_has_required_keys(self):
required = {
"tile_m", "tile_n", "tile_k",
"warp_m", "warp_n", "warp_k",
"warp_tile_m", "warp_tile_n", "warp_tile_k",
}
for cfg in self.builder._get_tile_configs():
self.assertTrue(required.issubset(cfg.keys()))
def test_tile_config_values_are_positive(self):
for cfg in self.builder._get_tile_configs():
for key, val in cfg.items():
self.assertGreater(val, 0, f"{key}={val} must be positive")
class TestGemmABQuantTileConfigStr(unittest.TestCase):
def test_tile_config_to_str_format(self):
tile_config = {
"tile_m": 128, "tile_n": 128, "tile_k": 64,
"warp_m": 2, "warp_n": 2, "warp_k": 1,
"warp_tile_m": 32, "warp_tile_n": 32, "warp_tile_k": 16,
}
result = GemmABQuantKernelBuilder._tile_config_to_str(tile_config)
self.assertEqual(result, "128x128x64_2x2x1_32x32x16")
class TestGemmABQuantGroupSizeN(unittest.TestCase):
def test_group_size_n_values_returned(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
gsn = builder._get_group_size_n_values()
self.assertIsInstance(gsn, list)
self.assertGreater(len(gsn), 0)
def test_group_size_n_values_match_config(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
gsn = builder._get_group_size_n_values()
self.assertIn(1, gsn)
self.assertIn(128, gsn)
class TestGemmABQuantLayoutVariants(unittest.TestCase):
LAYOUTS = ["rcr", "rrr", "ccr", "crr"]
def test_all_layouts_construct(self):
for layout in self.LAYOUTS:
with self.subTest(layout=layout):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, layout=layout)
self.assertEqual(builder.layout, layout)
class TestGemmABQuantDataTypes(unittest.TestCase):
DTYPES = ["fp8", "bf8"]
def test_all_dtypes_construct(self):
for dtype in self.DTYPES:
with self.subTest(dtype=dtype):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, datatype=dtype)
self.assertEqual(builder.datatype, dtype)
class TestGemmABQuantMaxInstances(unittest.TestCase):
def test_max_instances_none_returns_all(self):
with tempfile.TemporaryDirectory() as tmpdir:
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(_MINIMAL_CONFIG, f)
builder = GemmABQuantKernelBuilder(
"gemm_abquant", tmpdir, "gfx942", "fp8", "rcr",
config_json=cfg_path, max_instances=None,
)
kernel_list = [
{"tile_config": {"tile_m": 64, "tile_n": 64, "tile_k": 128,
"warp_m": 2, "warp_n": 2, "warp_k": 1,
"warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 64},
"trait_combo": ("compv3", "default", "intrawave",
False, False, False, False, False),
"group_size_n": 1}
]
result = builder._apply_sampling(kernel_list)
self.assertEqual(len(result), 1)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,297 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
set(GEMM_AQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM AQuant (semicolon-separated)")
set(GEMM_AQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM AQuant (semicolon-separated)")
set(GEMM_AQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)")
option(ENABLE_CCACHE_GEMM_AQUANT "Enable ccache for GEMM AQuant ops compilation" OFF)
# Store the directory path for use in functions
set(GEMM_AQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR})
# Function to create individual GEMM AQuant targets
function(create_individual_gemm_aquant_target datatype layout trait tile_config config_json)
if(NOT GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL)
message(WARNING "Skipping individual GEMM AQuant target ${datatype}_${layout}_${trait}_${tile_config}: No supported GPU targets")
return()
endif()
# Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k
string(REPLACE "_" ";" config_groups ${tile_config})
list(GET config_groups 0 tile_dims)
list(GET config_groups 1 warp_dims)
list(GET config_groups 2 warp_tile_dims)
string(REPLACE "x" ";" tile_parts ${tile_dims})
list(GET tile_parts 0 tile_m)
list(GET tile_parts 1 tile_n)
list(GET tile_parts 2 tile_k)
string(REPLACE "x" ";" warp_parts ${warp_dims})
list(GET warp_parts 0 warp_m)
list(GET warp_parts 1 warp_n)
list(GET warp_parts 2 warp_k)
string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims})
list(GET warp_tile_parts 0 warp_tile_m)
list(GET warp_tile_parts 1 warp_tile_n)
list(GET warp_tile_parts 2 warp_tile_k)
set(target_name "benchmark_gemm_aquant_${datatype}_${layout}_${trait}_${tile_config}")
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
# Generate the single instance header for this kernel
set(instance_header "${working_path}/gemm_aquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp")
# Add custom command to generate the header file at build time
add_custom_command(
OUTPUT ${instance_header}
COMMAND ${Python3_EXECUTABLE} ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_instance_builder.py
--working_path ${working_path}
--datatype ${datatype}
--layout ${layout}
--config_json ${config_json}
--gen_single
--kernel_name "gemm_aquant_${datatype}_${layout}_${trait}_${tile_config}"
--tile_config "${tile_config}"
--trait_combo "${trait}"
--gpu_target "${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}"
DEPENDS ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_instance_builder.py ${config_json}
COMMENT "Generating ${instance_header}"
)
# Create the executable
add_executable(${target_name}
EXCLUDE_FROM_ALL
${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_benchmark_single.cpp
${instance_header}
)
# Set GPU architectures
set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL})
# Set compile definitions
target_compile_definitions(${target_name} PRIVATE
GEMM_AQUANT_SINGLE_INSTANCE_HPP="${instance_header}"
)
# Include directories
target_include_directories(${target_name} PRIVATE
${GEMM_AQUANT_SOURCE_DIR}
${working_path}
)
# Compile options
target_compile_options(${target_name} PRIVATE
-Wno-undefined-func-template
-Wno-float-equal
--offload-compress
-include ${instance_header}
)
# Add to collection targets
add_dependencies(benchmark_gemm_aquant_all ${target_name})
add_dependencies(benchmark_gemm_aquant_${datatype} ${target_name})
add_dependencies(benchmark_gemm_aquant_${layout} ${target_name})
add_dependencies(benchmark_gemm_aquant_${datatype}_${layout} ${target_name})
# Add to pipeline-specific and scheduler-specific targets
string(REPLACE "_" ";" trait_parts ${trait})
list(GET trait_parts 0 pipeline)
list(GET trait_parts 1 epilogue)
list(GET trait_parts 2 scheduler)
add_dependencies(benchmark_gemm_aquant_${pipeline}_pipeline ${target_name})
add_dependencies(benchmark_gemm_aquant_${epilogue}_epilogue ${target_name})
add_dependencies(benchmark_gemm_aquant_${scheduler}_scheduler ${target_name})
endfunction()
# Function to build individual GEMM AQuant targets
function(build_individual_gemm_aquant_targets datatype layout)
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
# Choose config file (priority: env var > cmake var > default)
if(DEFINED ENV{GEMM_AQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_AQUANT_CONFIG_FILE}" STREQUAL "")
set(config_filename "$ENV{GEMM_AQUANT_CONFIG_FILE}")
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}")
message(VERBOSE " Using config from environment variable: ${config_filename}")
elseif(NOT "${GEMM_AQUANT_CONFIG_FILE}" STREQUAL "")
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_AQUANT_CONFIG_FILE}")
message(VERBOSE " Using custom config: ${GEMM_AQUANT_CONFIG_FILE}")
else()
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json")
message(VERBOSE " Using default config for layout ${layout}")
endif()
if(NOT EXISTS ${json_blob})
message(FATAL_ERROR "Config file not found: ${json_blob}")
endif()
# Create working directory
file(MAKE_DIRECTORY ${working_path})
message(VERBOSE "Generating individual AQuant kernels for ${datatype} ${layout}...")
# Build sampling arguments
set(extra_list_args "")
if(NOT "${GEMM_AQUANT_MAX_INSTANCES}" STREQUAL "")
list(APPEND extra_list_args --max-instances ${GEMM_AQUANT_MAX_INSTANCES})
endif()
if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER})
list(APPEND extra_list_args --manifest-path ${working_path})
endif()
if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "")
list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED})
endif()
# List kernels
message(VERBOSE " Listing AQuant kernel configurations for ${datatype} ${layout}...")
execute_process(
COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_aquant_instance_builder.py
--working_path ${working_path}
--datatype ${datatype}
--layout ${layout}
--config_json ${json_blob}
--gpu_target "${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}"
--list_kernels
${extra_list_args}
WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
RESULT_VARIABLE ret
OUTPUT_VARIABLE list_output
ERROR_VARIABLE list_error
)
if(NOT ret EQUAL 0)
message(FATAL_ERROR "Failed to list AQuant kernels for ${datatype} ${layout}: ${list_error}")
endif()
# Read kernel count
if(EXISTS ${working_path}/gemm_aquant_kernel_count.txt)
file(READ ${working_path}/gemm_aquant_kernel_count.txt kernel_count)
string(STRIP "${kernel_count}" kernel_count)
message(VERBOSE " Found ${kernel_count} AQuant kernel configurations")
else()
message(FATAL_ERROR "Kernel count file not found")
endif()
# Read kernel list and create targets
if(EXISTS ${working_path}/gemm_aquant_kernel_list.txt)
file(STRINGS ${working_path}/gemm_aquant_kernel_list.txt kernel_lines)
foreach(line IN LISTS kernel_lines)
string(REPLACE "|" ";" parts "${line}")
list(GET parts 0 kernel_name)
list(GET parts 1 tile_config)
list(GET parts 2 trait_combo)
create_individual_gemm_aquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}")
endforeach()
else()
message(FATAL_ERROR "Kernel list file not found")
endif()
endfunction()
# Main build logic
message(VERBOSE "=== Starting Tile Engine GEMM AQuant Configuration ===")
message(VERBOSE "GEMM_AQUANT_DATATYPE: ${GEMM_AQUANT_DATATYPE}")
message(VERBOSE "GEMM_AQUANT_LAYOUT: ${GEMM_AQUANT_LAYOUT}")
message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
# Filter GPU targets to supported architectures
set(GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL "")
set(DESIRED_TARGETS "gfx90a;gfx942;gfx950")
foreach(target IN LISTS SUPPORTED_GPU_TARGETS)
if(target IN_LIST DESIRED_TARGETS)
list(APPEND GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL ${target})
message(VERBOSE " Adding GPU target: ${target}")
endif()
endforeach()
# Skip build if no matching targets found
if(NOT GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL)
message(WARNING "Skipping Tile Engine GEMM AQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
else()
message(VERBOSE "Building individual GEMM AQuant targets for GPU targets: ${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}")
# Enable parallel compilation optimizations (directory-scoped to avoid overriding parent)
set_property(DIRECTORY PROPERTY JOB_POOLS
compile_heavy=4 # Limit heavy compilations to prevent OOM
compile_normal=16 # Allow more parallel normal compilations
)
# Enable compiler cache if requested
if(ENABLE_CCACHE_GEMM_AQUANT)
find_program(CCACHE_PROGRAM ccache)
if(CCACHE_PROGRAM)
set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
message(VERBOSE "Using ccache for faster compilation")
endif()
endif()
# Create master collection target
add_custom_target(benchmark_gemm_aquant_all)
# Create datatype collection targets
foreach(dt IN LISTS GEMM_AQUANT_DATATYPE)
add_custom_target(benchmark_gemm_aquant_${dt})
endforeach()
# Create layout collection targets
foreach(l IN LISTS GEMM_AQUANT_LAYOUT)
add_custom_target(benchmark_gemm_aquant_${l})
endforeach()
# Create combined collection targets
foreach(dt IN LISTS GEMM_AQUANT_DATATYPE)
foreach(l IN LISTS GEMM_AQUANT_LAYOUT)
add_custom_target(benchmark_gemm_aquant_${dt}_${l})
endforeach()
endforeach()
# Create pipeline-specific collection targets
set(GEMM_AQUANT_PIPELINES "mem;compv3")
foreach(pipeline IN LISTS GEMM_AQUANT_PIPELINES)
add_custom_target(benchmark_gemm_aquant_${pipeline}_pipeline)
endforeach()
# Create epilogue-specific collection targets
set(GEMM_AQUANT_EPILOGUES "default;cshuffle")
foreach(epilogue IN LISTS GEMM_AQUANT_EPILOGUES)
add_custom_target(benchmark_gemm_aquant_${epilogue}_epilogue)
endforeach()
# Create scheduler-specific collection targets
set(GEMM_AQUANT_SCHEDULERS "intrawave;interwave")
foreach(scheduler IN LISTS GEMM_AQUANT_SCHEDULERS)
add_custom_target(benchmark_gemm_aquant_${scheduler}_scheduler)
endforeach()
# Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that
# sampling fires per-combo rather than being a single cap.
# CMake integer division truncates; clamp the result to at least 1 so that a
# small total budget (e.g. 5 instances across 8 combos) does not silently
# produce a per-combo budget of 0 and skip all kernels.
if(NOT "${GEMM_AQUANT_MAX_INSTANCES}" STREQUAL "")
list(LENGTH GEMM_AQUANT_DATATYPE _aq_n_dt)
list(LENGTH GEMM_AQUANT_LAYOUT _aq_n_lay)
math(EXPR _aq_n_combos "${_aq_n_dt} * ${_aq_n_lay}")
if(_aq_n_combos GREATER 0)
math(EXPR GEMM_AQUANT_MAX_INSTANCES
"${GEMM_AQUANT_MAX_INSTANCES} / ${_aq_n_combos}")
if(GEMM_AQUANT_MAX_INSTANCES EQUAL 0)
set(GEMM_AQUANT_MAX_INSTANCES 1)
message(WARNING " gemm_aquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_aq_n_combos} combos)")
else()
message(STATUS " gemm_aquant: per-combo budget = ${GEMM_AQUANT_MAX_INSTANCES} (${_aq_n_combos} combos)")
endif()
endif()
endif()
# Build individual targets for each datatype/layout combination
foreach(dt IN LISTS GEMM_AQUANT_DATATYPE)
foreach(l IN LISTS GEMM_AQUANT_LAYOUT)
build_individual_gemm_aquant_targets(${dt} ${l})
endforeach()
endforeach()
endif()

View File

@@ -0,0 +1,92 @@
{
"tile_config": {
"tile_m": {
"values": [
16
]
},
"tile_n": {
"values": [
64
]
},
"tile_k": {
"values": [
256
]
},
"warp_m": {
"values": [
1
]
},
"warp_n": {
"values": [
4
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16
]
},
"warp_tile_n": {
"values": [
16
]
},
"warp_tile_k": {
"values": [
16, 32, 64, 128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"mem",
"compv3"
]
},
"scheduler": {
"values": [
"interwave",
"intrawave"
]
},
"epilogue": {
"values": [
"default",
"cshuffle"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"a_preshuffle_quant": {
"values": [
false,
true
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128
}

View File

@@ -0,0 +1,102 @@
{
"tile_config": {
"tile_m": {
"max": 256,
"min": 64,
"step": 64
},
"tile_n": {
"max": 256,
"min": 64,
"step": 64
},
"tile_k": {
"max": 256,
"min": 64,
"step": 64
},
"warp_m": {
"values": [
4,
2,
1
]
},
"warp_n": {
"values": [
4,
2,
1
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16,
32
]
},
"warp_tile_n": {
"values": [
16,
32
]
},
"warp_tile_k": {
"values": [
8,
16,
32,
64,
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3",
"mem"
]
},
"scheduler": {
"values": [
"intrawave",
"interwave"
]
},
"epilogue": {
"values": [
"cshuffle",
"default"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"a_preshuffle_quant": {
"values": [
false,
true
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128
}

View File

@@ -0,0 +1,88 @@
{
"tile_config": {
"tile_m": {
"values": [
16
]
},
"tile_n": {
"values": [
64
]
},
"tile_k": {
"values": [
256
]
},
"warp_m": {
"values": [
1
]
},
"warp_n": {
"values": [
4
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16
]
},
"warp_tile_n": {
"values": [
16
]
},
"warp_tile_k": {
"values": [
16, 32, 64, 128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"mem"
]
},
"scheduler": {
"values": [
"interwave"
]
},
"epilogue": {
"values": [
"cshuffle"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
true
]
},
"a_preshuffle_quant": {
"values": [
false
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128
}

View File

@@ -0,0 +1,231 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <iostream>
#include <string>
#include <fstream>
#include <stdexcept>
#include <iomanip>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "gemm_aquant_common.hpp"
// Data types and Layouts are defined by the generated kernel headers:
// ADataType, BDataType, AQDataType, AccDataType, CDataType
// ALayout, BLayout, CLayout, AQLayout
enum class Metric
{
LATENCY = 0,
TFLOPS = 1,
BANDWIDTH = 2
};
inline constexpr auto get_metric_name(Metric m)
{
switch(m)
{
case Metric::LATENCY: return "latency";
case Metric::TFLOPS: return "tflops";
case Metric::BANDWIDTH: return "bandwidth";
default: throw std::invalid_argument("Unsupported metric type");
}
}
struct AQuantGemmProblem
{
int split_k_;
int m_, n_, k_;
int stride_a_, stride_b_, stride_c_;
int stride_aq_;
int group_size_k_;
std::string dtype_a_, dtype_b_, dtype_aq_, dtype_acc_, dtype_c_;
std::string layout_a_, layout_b_, layout_c_;
friend std::ostream& operator<<(std::ostream& os, const AQuantGemmProblem& problem)
{
os << "{\n"
<< " \"split_k\":" << problem.split_k_ << ",\n"
<< " \"m\":" << problem.m_ << ",\n"
<< " \"n\":" << problem.n_ << ",\n"
<< " \"k\":" << problem.k_ << ",\n"
<< " \"stride_a\":" << problem.stride_a_ << ",\n"
<< " \"stride_b\":" << problem.stride_b_ << ",\n"
<< " \"stride_c\":" << problem.stride_c_ << ",\n"
<< " \"stride_aq\":" << problem.stride_aq_ << ",\n"
<< " \"group_size_k\":" << problem.group_size_k_ << ",\n"
<< " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n"
<< " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n"
<< " \"dtype_aq\":\"" << problem.dtype_aq_ << "\",\n"
<< " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n"
<< " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n"
<< " \"layout_a\":\"" << problem.layout_a_ << "\",\n"
<< " \"layout_b\":\"" << problem.layout_b_ << "\",\n"
<< " \"layout_c\":\"" << problem.layout_c_ << "\"\n"
<< "}";
return os;
}
};
struct PerformanceResult
{
double latency_;
double tflops_;
double bandwidth_;
static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m)
{
switch(m)
{
case Metric::LATENCY: return a.latency_ < b.latency_;
case Metric::TFLOPS: return a.tflops_ > b.tflops_;
case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_;
default: throw std::invalid_argument("Unsupported metric type");
}
}
friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result)
{
os << "{\n"
<< " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_
<< ",\n"
<< " \"tflops(TFlops)\": " << result.tflops_ << ",\n"
<< " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n"
<< "}";
return os;
}
};
struct KernelInstance
{
std::string name_;
AQuantGemmProblem problem_;
PerformanceResult perf_result_;
static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m)
{
return PerformanceResult::compare(a.perf_result_, b.perf_result_, m);
}
friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj)
{
os << "{\n"
<< " \"name\": \"" << obj.name_ << "\",\n"
<< " \"problem\": " << obj.problem_ << ",\n"
<< " \"perf_result\": " << obj.perf_result_ << "\n"
<< "}";
return os;
}
};
struct Setting
{
int n_warmup_;
int n_repeat_;
bool is_gpu_timer_;
int verify_;
int init_method_;
bool log_;
std::string csv_filename_;
bool flush_cache_;
int rotating_count_;
bool json_output_;
};
inline std::string get_rocm_version()
{
// Try the legacy path first, then the path used by newer ROCm installs.
for(const char* path : {"/opt/rocm/.info/version", "/opt/rocm/lib/rocm_version.h"})
{
std::ifstream version_file(path);
if(version_file.is_open())
{
std::string version;
std::getline(version_file, version);
if(!version.empty())
return version;
}
}
return "Unknown";
}
template <typename ADataType_,
typename AQDataType_,
typename BDataType_,
typename AccDataType_,
typename CDataType_>
auto calculate_rtol_atol_aquant(const ck_tile::index_t K,
const ck_tile::index_t kbatch,
const float max_accumulated_value)
{
using ComputeType =
std::conditional_t<sizeof(ADataType_) < sizeof(BDataType_), ADataType_, BDataType_>;
const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType_, AccDataType_>(
ck_tile::integer_divide_ceil(K, kbatch));
const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType_, AccDataType_>(
max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
const auto rtol_split_k =
ck_tile::get_relative_threshold<CDataType_, CDataType_, CDataType_>(kbatch);
const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType_, CDataType_, CDataType_>(
max_accumulated_value, kbatch);
return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
}
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions"
/// @brief Compare device and host results for AQuant GEMM
bool compare_aquant(std::string instanceName,
ck_tile::index_t K,
ck_tile::index_t kbatch,
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
const float max_accumulated_value =
std::abs(static_cast<float>(*std::max_element(c_m_n_host_result.mData.begin(),
c_m_n_host_result.mData.end(),
[](const auto& a, const auto& b) {
return std::abs(static_cast<float>(a)) <
std::abs(static_cast<float>(b));
})));
const auto rtol_atol =
calculate_rtol_atol_aquant<ADataType, AQDataType, BDataType, AccDataType, CDataType>(
K, kbatch, max_accumulated_value);
bool pass = ck_tile::check_err(c_m_n_dev_result,
c_m_n_host_result,
"Error: Incorrect results!",
rtol_atol.at(ck_tile::number<0>{}),
rtol_atol.at(ck_tile::number<1>{}));
std::cerr << "For " << instanceName << " Relative error threshold is "
<< rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
<< rtol_atol.at(ck_tile::number<1>{}) << std::endl;
std::cerr << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;
return pass;
}
/// @brief CPU reference implementation for AQuant GEMM
void aquant_gemm_host_reference(int verify,
ck_tile::HostTensor<ADataType>& a_m_k,
ck_tile::HostTensor<AQDataType>& aq_m_qk,
ck_tile::HostTensor<BDataType>& b_k_n,
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
if(verify == 1)
{
c_m_n_host_result.SetZero();
using QuantGroupSize = typename SelectedKernel::QuantGroupSize;
ck_tile::reference_gemm_quant<ADataType,
AQDataType,
BDataType,
AccDataType,
CDataType,
QuantGroupSize,
true /* aquant */>(a_m_k, aq_m_qk, b_k_n, c_m_n_host_result);
}
}
#pragma clang diagnostic pop

View File

@@ -0,0 +1,547 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import sys
import json
import csv
import subprocess
import argparse
import time
from pathlib import Path
from typing import List, Dict, Tuple, Optional
class AQuantGemmBenchmark:
def __init__(self, build_dir: str, verbose: bool = False):
self.build_dir = Path(build_dir)
self.verbose = verbose
self.results = []
def discover_kernels(self) -> List[Path]:
"""Find all benchmark_gemm_aquant_* executables in the build directory."""
bin_dir = self.build_dir / "bin"
if not bin_dir.exists():
print(f"Error: Binary directory {bin_dir} does not exist")
return []
kernels = list(bin_dir.glob("benchmark_gemm_aquant_*"))
if self.verbose:
print(f"Found {len(kernels)} AQuant kernel executables")
for k in kernels:
print(f" - {k.name}")
return kernels
def extract_kernel_info(self, kernel_path: Path) -> Dict[str, str]:
"""Extract kernel information from filename."""
name = kernel_path.stem
info = {
"executable": str(kernel_path),
"name": name,
"data_type": "unknown",
"layout": "unknown",
"pipeline": "unknown",
"scheduler": "unknown",
"epilogue": "unknown",
"a_preshuffle_quant": False,
}
# Parse: benchmark_gemm_aquant_fp8_rcr_mem_default_interwave_False_False_True_False_128x128x128_...
parts = name.split("_")
if len(parts) >= 3:
# Skip "benchmark_gemm_aquant" prefix (3 parts)
idx = 3
if idx < len(parts):
info["data_type"] = parts[idx]
idx += 1
if idx < len(parts):
info["layout"] = parts[idx]
idx += 1
if idx < len(parts):
info["pipeline"] = parts[idx]
idx += 1
if idx < len(parts):
info["epilogue"] = parts[idx]
idx += 1
if idx < len(parts):
info["scheduler"] = parts[idx]
# Parse tile dimensions
config_info = self.parse_detailed_config(name)
info.update(config_info)
info["config_id"] = self.generate_config_id(info)
return info
def parse_detailed_config(self, kernel_name: str) -> Dict:
"""Parse tile dimensions and boolean flags from kernel name.
Kernel name format (after "benchmark_gemm_aquant_"):
{dtype}_{layout}_{pipeline}_{epilogue}_{scheduler}
_{padM}_{padN}_{padK}_{aPreshuffle}
_{TileMxTileNxTileK}_{WarpMxWarpNxWarpK}_{WarpTileMxWarpTileNxWarpTileK}
Dimensions are extracted positionally (in order of appearance) rather than
by sorting, which would silently misassign groups that share a max value.
"""
config = {
"tile_sizes": {"tile_m": 0, "tile_n": 0, "tile_k": 0},
"warp_config": {"warp_m": 0, "warp_n": 0, "warp_k": 0},
"warp_tile": {"warp_tile_m": 0, "warp_tile_n": 0, "warp_tile_k": 0},
"optimization_flags": {
"pad_m": False,
"pad_n": False,
"pad_k": False,
"a_preshuffle_quant": False,
},
}
parts = kernel_name.split("_")
# Locate the first boolean token; collect the contiguous boolean run
bool_start = -1
bool_sequence = []
for i, part in enumerate(parts):
if part in ("True", "False"):
if bool_start == -1:
bool_start = i
bool_sequence.append(part == "True")
elif bool_start != -1:
# End of the boolean run
break
if len(bool_sequence) >= 4:
config["optimization_flags"]["pad_m"] = bool_sequence[0]
config["optimization_flags"]["pad_n"] = bool_sequence[1]
config["optimization_flags"]["pad_k"] = bool_sequence[2]
config["optimization_flags"]["a_preshuffle_quant"] = bool_sequence[3]
# Extract dimension groups (NxNxN tokens) in positional order.
# Appearance order in the name is: tile, warp, warp_tile.
dimension_groups = []
for part in parts:
sub = part.split("x")
if len(sub) == 3:
try:
dims = [int(v) for v in sub]
if all(d > 0 for d in dims):
dimension_groups.append(dims)
except ValueError:
continue
if len(dimension_groups) >= 3:
config["tile_sizes"]["tile_m"] = dimension_groups[0][0]
config["tile_sizes"]["tile_n"] = dimension_groups[0][1]
config["tile_sizes"]["tile_k"] = dimension_groups[0][2]
config["warp_config"]["warp_m"] = dimension_groups[1][0]
config["warp_config"]["warp_n"] = dimension_groups[1][1]
config["warp_config"]["warp_k"] = dimension_groups[1][2]
config["warp_tile"]["warp_tile_m"] = dimension_groups[2][0]
config["warp_tile"]["warp_tile_n"] = dimension_groups[2][1]
config["warp_tile"]["warp_tile_k"] = dimension_groups[2][2]
elif len(dimension_groups) == 2:
config["tile_sizes"]["tile_m"] = dimension_groups[0][0]
config["tile_sizes"]["tile_n"] = dimension_groups[0][1]
config["tile_sizes"]["tile_k"] = dimension_groups[0][2]
config["warp_config"]["warp_m"] = dimension_groups[1][0]
config["warp_config"]["warp_n"] = dimension_groups[1][1]
config["warp_config"]["warp_k"] = dimension_groups[1][2]
elif len(dimension_groups) == 1:
config["tile_sizes"]["tile_m"] = dimension_groups[0][0]
config["tile_sizes"]["tile_n"] = dimension_groups[0][1]
config["tile_sizes"]["tile_k"] = dimension_groups[0][2]
return config
def generate_config_id(self, info: Dict) -> str:
"""Generate a compact config ID from kernel info."""
parts = [
info.get("data_type", "unk"),
info.get("layout", "unk"),
info.get("pipeline", "unk"),
info.get("scheduler", "unk"),
]
tile_sizes = info.get("tile_sizes", {})
if tile_sizes.get("tile_m", 0) > 0:
parts.append(
f"{tile_sizes['tile_m']}x{tile_sizes['tile_n']}x{tile_sizes['tile_k']}"
)
warp_config = info.get("warp_config", {})
if warp_config.get("warp_m", 0) > 0:
parts.append(
f"w{warp_config['warp_m']}x{warp_config['warp_n']}x{warp_config['warp_k']}"
)
warp_tile = info.get("warp_tile", {})
if warp_tile.get("warp_tile_m", 0) > 0:
parts.append(
f"wt{warp_tile['warp_tile_m']}x{warp_tile['warp_tile_n']}x{warp_tile['warp_tile_k']}"
)
return "_".join(parts)
def run_kernel(self, kernel_path: Path, params: Dict[str, str]) -> Optional[Dict]:
"""Run a single kernel with given parameters."""
results_dir = self.build_dir / "results"
results_dir.mkdir(exist_ok=True)
json_file = results_dir / f"{kernel_path.stem}.json"
cmd = [str(kernel_path)]
for key, value in params.items():
cmd.append(f"-{key}={value}")
cmd.append("-json_output=true")
if self.verbose:
print(f"Running: {' '.join(cmd)}")
try:
result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
if result.returncode != 0:
print(f"Error running {kernel_path.name}: {result.stderr}")
return None
output = result.stdout.strip()
if output:
with open(json_file, "w") as f:
f.write(output)
return self.parse_json_file(json_file)
else:
print(f"No output from {kernel_path.name}")
return None
except subprocess.TimeoutExpired:
print(f"Timeout running {kernel_path.name}")
return None
except Exception as e:
print(f"Error running {kernel_path.name}: {e}")
return None
def parse_json_file(self, json_file: Path) -> Optional[Dict]:
"""Parse JSON data from kernel output file."""
try:
with open(json_file, "r") as f:
content = f.read().strip()
data = json.loads(content)
result = data.copy()
if "perf_result" in data:
perf = data["perf_result"]
result["time_ms"] = perf.get("latency(ms)", 0)
result["tflops"] = perf.get("tflops(TFlops)", 0)
result["bandwidth_gb_s"] = perf.get("bandwidth(GB/s)", 0)
return result
except (json.JSONDecodeError, Exception) as e:
if self.verbose:
print(f"Failed to parse JSON from {json_file}: {e}")
return None
def benchmark_problem_size(
self,
kernels: List[Path],
m: int,
n: int,
k: int,
group_size_k: int = 128,
split_k: int = 1,
verify: int = 0,
warmup: int = 50,
repeat: int = 100,
flush_cache: bool = True,
rotating_count: int = 1000,
) -> List[Dict]:
"""Benchmark all kernels for a specific problem size."""
results = []
params = {
"m": m,
"n": n,
"k": k,
"group_size_k": group_size_k,
"split_k": split_k,
"verify": verify,
"warmup": warmup,
"repeat": repeat,
"flush_cache": str(flush_cache).lower(),
"rotating_count": rotating_count,
}
print(f"\nBenchmarking M={m}, N={n}, K={k}, group_size_k={group_size_k}")
for kernel_path in kernels:
kernel_info = self.extract_kernel_info(kernel_path)
result = self.run_kernel(kernel_path, params)
if result:
structured_result = {
"name": kernel_info["name"],
"config_id": kernel_info["config_id"],
"problem": result.get("problem", {}),
"perf_result": result.get("perf_result", {}),
"config": {
"data_type": kernel_info["data_type"],
"layout": kernel_info["layout"],
"pipeline": kernel_info["pipeline"],
"scheduler": kernel_info["scheduler"],
"epilogue": kernel_info["epilogue"],
"tile_sizes": kernel_info.get("tile_sizes", {}),
"warp_config": kernel_info.get("warp_config", {}),
"warp_tile": kernel_info.get("warp_tile", {}),
"optimization_flags": kernel_info.get("optimization_flags", {}),
},
"executable": kernel_info["executable"],
"time_ms": result.get("time_ms", 0),
"tflops": result.get("tflops", 0),
"bandwidth_gb_s": result.get("bandwidth_gb_s", 0),
}
results.append(structured_result)
if self.verbose:
print(
f" {kernel_info['config_id']}: "
f"{structured_result['tflops']:.2f} TFLOPS, "
f"{structured_result['bandwidth_gb_s']:.2f} GB/s, "
f"{structured_result['time_ms']:.2f}ms"
)
return results
def find_best_kernel(
self, results: List[Dict], metric: str = "tflops"
) -> Optional[Dict]:
"""Find the best performing kernel based on metric."""
if not results:
return None
if metric == "tflops":
return max(results, key=lambda x: x.get("tflops", 0))
elif metric == "time_ms":
return min(results, key=lambda x: x.get("time_ms", float("inf")))
elif metric == "bandwidth_gb_s":
return max(results, key=lambda x: x.get("bandwidth_gb_s", 0))
else:
raise ValueError(f"Unknown metric: {metric}")
def benchmark_sweep(
self,
problem_sizes: List[Tuple[int, int, int]],
group_size_k: int = 128,
split_k_values: Optional[List[int]] = None,
verify: bool = False,
warmup: int = 50,
repeat: int = 100,
flush_cache: bool = True,
rotating_count: int = 1000,
) -> Dict:
"""Run comprehensive benchmark sweep."""
if split_k_values is None:
split_k_values = [1]
kernels = self.discover_kernels()
if not kernels:
print("No kernels found!")
return {}
all_results = []
best_kernels = {}
for m, n, k in problem_sizes:
for split_k in split_k_values:
results = self.benchmark_problem_size(
kernels,
m,
n,
k,
group_size_k=group_size_k,
split_k=split_k,
verify=1 if verify else 0,
warmup=warmup,
repeat=repeat,
flush_cache=flush_cache,
rotating_count=rotating_count,
)
all_results.extend(results)
best = self.find_best_kernel(results)
if best:
key = f"m{m}_n{n}_k{k}_splitk{split_k}"
best_kernels[key] = best
print(
f"Best for {key}: {best['name']} "
f"({best['tflops']:.2f} TFLOPS, "
f"{best['bandwidth_gb_s']:.2f} GB/s, "
f"{best['time_ms']:.2f}ms)"
)
self.results = all_results
return best_kernels
def export_csv(self, filename: str):
"""Export all results to CSV."""
if not self.results:
print("No results to export")
return
all_keys = set()
for result in self.results:
all_keys.update(result.keys())
fieldnames = sorted(all_keys)
with open(filename, "w", newline="") as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(self.results)
print(f"Results exported to {filename}")
def export_best_kernels(self, best_kernels: Dict, filename: str):
"""Export best kernel selections to file."""
with open(filename, "w") as f:
f.write("# Best AQuant kernel selections\n")
f.write(
"# Format: problem_size -> kernel_name (TFLOPS, bandwidth, latency)\n\n"
)
for key, kernel in sorted(best_kernels.items()):
f.write(
f"{key}: {kernel['name']} "
f"({kernel['tflops']:.2f} TFLOPS, "
f"{kernel['bandwidth_gb_s']:.2f} GB/s, "
f"{kernel['time_ms']:.2f}ms)\n"
)
print(f"Best kernels exported to {filename}")
def export_json(self, filename: str, best_kernels: Dict = None):
"""Export results to JSON."""
from datetime import datetime
output_data = {
"benchmark_metadata": {
"timestamp": datetime.now().isoformat(),
"operator": "gemm_aquant",
"total_kernels_tested": len(self.results),
"successful_runs": len(
[r for r in self.results if r.get("tflops", 0) > 0]
),
},
"kernel_results": self.results,
"best_kernels_by_problem": best_kernels or {},
}
with open(filename, "w") as f:
json.dump(output_data, f, indent=2)
print(f"JSON results exported to {filename}")
def main():
parser = argparse.ArgumentParser(description="AQuant GEMM Kernel Benchmarking Tool")
parser.add_argument(
"build_dir", help="Build directory containing kernel executables"
)
parser.add_argument(
"--problem-sizes",
nargs="+",
default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"],
help="Problem sizes as M,N,K tuples",
)
parser.add_argument(
"--group-size-k",
type=int,
default=128,
help="Quantization group size along K (default: 128)",
)
parser.add_argument(
"--split-k",
nargs="+",
type=int,
default=[1],
help="Split-K values to test",
)
parser.add_argument("--verify", action="store_true", help="Enable verification")
parser.add_argument("--verbose", action="store_true", help="Verbose output")
parser.add_argument(
"--warmup",
type=int,
default=50,
help="Number of warmup iterations",
)
parser.add_argument(
"--repeat",
type=int,
default=100,
help="Number of benchmark iterations",
)
parser.add_argument(
"--csv",
default="gemm_aquant_benchmark_results.csv",
help="CSV output filename (default: gemm_aquant_benchmark_results.csv)",
)
parser.add_argument(
"--best",
default="best_aquant_kernels.txt",
help="Best kernels output filename (default: best_aquant_kernels.txt)",
)
parser.add_argument("--json", help="JSON output filename (optional)")
parser.add_argument(
"--flush-cache",
action="store_true",
default=True,
help="Flush cache between kernel runs (default: true)",
)
parser.add_argument(
"--rotating-count",
type=int,
default=1000,
help="Number of iterations to rotate the cache (default: 1000)",
)
args = parser.parse_args()
problem_sizes = []
for size_str in args.problem_sizes:
try:
m, n, k = map(int, size_str.split(","))
problem_sizes.append((m, n, k))
except ValueError:
print(f"Invalid problem size: {size_str}")
return 1
benchmark = AQuantGemmBenchmark(args.build_dir, verbose=args.verbose)
print("Starting AQuant GEMM kernel benchmark sweep...")
start_time = time.time()
best_kernels = benchmark.benchmark_sweep(
problem_sizes=problem_sizes,
group_size_k=args.group_size_k,
split_k_values=args.split_k,
verify=args.verify,
warmup=args.warmup,
repeat=args.repeat,
flush_cache=args.flush_cache,
rotating_count=args.rotating_count,
)
elapsed_time = time.time() - start_time
print(f"\nBenchmark completed in {elapsed_time:.2f} seconds")
# Export results
benchmark.export_csv(args.csv)
benchmark.export_best_kernels(best_kernels, args.best)
if args.json:
benchmark.export_json(args.json, best_kernels)
return 0
if __name__ == "__main__":
sys.exit(main())

View File

@@ -0,0 +1,153 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <iostream>
#include <functional>
#include <tuple>
#include <exception>
#include <sstream>
#include <vector>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "gemm_aquant_profiler.hpp"
#include "gemm_aquant_common.hpp"
// The kernel header is included via the compile command line with -include flag
// It defines SelectedKernel struct, KERNEL_NAME, and type aliases:
// ADataType, BDataType, AQDataType, AccDataType, CDataType
// ALayout, BLayout, CLayout, AQLayout
inline auto create_args(int argc, char* argv[])
{
ck_tile::ArgParser arg_parser;
arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.")
.insert("n", "4096", "The value for n dimension. Default is 4096.")
.insert("k", "2048", "The value for k dimension. Default is 2048.")
.insert("stride_a", "0", "The stride value for tensor A. Default is 0.")
.insert("stride_b", "0", "The stride value for tensor B. Default is 0.")
.insert("stride_c", "0", "The stride value for tensor C. Default is 0.")
.insert("split_k", "1", "The split value for k dimension. Default is 1.")
.insert("group_size_k",
std::to_string(SelectedKernel::GroupSizeK),
"The quantization group size along K. Default matches kernel config.")
.insert("verify",
"1",
"The type of validation. Set to 0 for no validation, 1 for validation on CPU. "
"Default is 1, CPU validation.")
.insert("log",
"false",
"Whether output kernel instance information or not. Possible values are true or "
"false. Default is false")
.insert(
"warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.")
.insert(
"repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.")
.insert("timer",
"true",
"Whether if the timer is gpu timer or not. Possible values are false or true. "
"Default is true.")
.insert("init",
"0",
"The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 "
"for constant(1). Default is 0, random.")
.insert("flush_cache",
"true",
"To flush cache, possible values are true or false. Default is true.")
.insert(
"rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.")
.insert("metric",
"0",
"Metric with which to measure kernel performance. Set to 0 for latency, 1 for "
"tflops, or 2 for bandwidth. Default is 0, latency.")
.insert("csv_filename",
"",
"The filename of benchmark result. Default is empty (no CSV output).")
.insert("json_output",
"false",
"Whether to output results in JSON format only. Possible values are true or false. "
"Default is false");
bool result = arg_parser.parse(argc, argv);
return std::make_tuple(result, arg_parser);
}
void benchmark_single(const ck_tile::ArgParser& arg_parser)
{
std::string dtype_a = DataTypeTraits<ADataType>::name;
std::string dtype_b = DataTypeTraits<BDataType>::name;
std::string dtype_aq = DataTypeTraits<AQDataType>::name;
std::string dtype_acc = DataTypeTraits<AccDataType>::name;
std::string dtype_c = DataTypeTraits<CDataType>::name;
std::string layout_a = ALayout::name;
std::string layout_b = BLayout::name;
std::string layout_c = CLayout::name;
int group_size_k = arg_parser.get_int("group_size_k");
AQuantGemmProblem problem{arg_parser.get_int("split_k"),
arg_parser.get_int("m"),
arg_parser.get_int("n"),
arg_parser.get_int("k"),
arg_parser.get_int("stride_a"),
arg_parser.get_int("stride_b"),
arg_parser.get_int("stride_c"),
0, // stride_aq computed by profiler
group_size_k,
dtype_a,
dtype_b,
dtype_aq,
dtype_acc,
dtype_c,
layout_a,
layout_b,
layout_c};
Setting setting{arg_parser.get_int("warmup"),
arg_parser.get_int("repeat"),
arg_parser.get_bool("timer"),
arg_parser.get_int("verify"),
arg_parser.get_int("init"),
arg_parser.get_bool("log"),
arg_parser.get_str("csv_filename"),
arg_parser.get_bool("flush_cache"),
arg_parser.get_int("rotating_count"),
arg_parser.get_bool("json_output")};
AQuantGemmProfiler profiler{setting};
try
{
auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args,
const ck_tile::stream_config& stream) {
return SelectedKernel::launch(args, stream);
};
profiler.benchmark(problem, kernel_func);
profiler.select_best_instance(static_cast<Metric>(arg_parser.get_int("metric")));
}
catch(const std::exception& e)
{
std::cerr << "Benchmark failed: " << e.what() << std::endl;
}
}
int main(int argc, char* argv[])
{
try
{
auto [result, parser] = create_args(argc, argv);
if(!result)
return EXIT_FAILURE;
benchmark_single(parser);
return 0;
}
catch(const std::exception& e)
{
std::cerr << "Error: " << e.what() << "\n";
return EXIT_FAILURE;
}
}

View File

@@ -0,0 +1,73 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/core/numeric/integer.hpp"
// DataTypeTraits for all supported types
template <typename T>
struct DataTypeTraits;
template <>
struct DataTypeTraits<float>
{
static constexpr const char* name = "fp32";
};
template <>
struct DataTypeTraits<ck_tile::half_t>
{
static constexpr const char* name = "fp16";
};
template <>
struct DataTypeTraits<ck_tile::bf16_t>
{
static constexpr const char* name = "bf16";
};
template <>
struct DataTypeTraits<ck_tile::fp8_t>
{
static constexpr const char* name = "fp8";
};
template <>
struct DataTypeTraits<ck_tile::bf8_t>
{
static constexpr const char* name = "bf8";
};
// Helper function to determine if a layout is row-major
template <typename Layout>
constexpr auto is_row_major(Layout)
{
return ck_tile::bool_constant<std::is_same_v<Layout, ck_tile::tensor_layout::gemm::RowMajor>>{};
}
// Structure to hold kernel traits for dispatcher
struct AQuantKernelTraits
{
std::string pipeline; // mem, compv3
std::string scheduler; // intrawave, interwave
std::string epilogue; // default
bool pad_m;
bool pad_n;
bool pad_k;
bool a_preshuffle_quant;
AQuantKernelTraits()
: pipeline("mem"),
scheduler("interwave"),
epilogue("default"),
pad_m(false),
pad_n(false),
pad_k(true),
a_preshuffle_quant(false)
{
}
};

View File

@@ -0,0 +1,322 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import os
import argparse
import importlib.util
import multiprocessing
import concurrent.futures
def _import_gemm_kernel_builder():
"""Import the base GemmKernelBuilder from the gemm directory."""
current_dir = os.path.dirname(os.path.abspath(__file__))
gemm_dir = os.path.dirname(os.path.dirname(current_dir))
spec = importlib.util.spec_from_file_location(
"gemm_instance_builder",
os.path.join(gemm_dir, "gemm_instance_builder.py"),
)
gemm_builder_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(gemm_builder_module)
return gemm_builder_module.GemmKernelBuilder
GemmKernelBuilder = _import_gemm_kernel_builder()
class GemmAQuantKernelBuilder(GemmKernelBuilder):
def __init__(
self,
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json=None,
max_instances=None,
seed=None,
tier=None,
manifest_path=None,
):
super().__init__(
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json,
max_instances=max_instances,
seed=seed,
tier=tier,
manifest_path=manifest_path,
)
self.group_size_k = self.config.get("group_size_k", 128)
def _uses_persistent_trait(self):
# The 7th trait slot carries a_preshuffle_quant, not a persistent-kernel flag.
# Return True so the base class includes it in kernel names and sampling features.
return True
def _generate_all_individual(self, num_workers=None):
"""Generate individual kernel files for separate compilation with parallel processing"""
if num_workers is None:
num_workers = min(
multiprocessing.cpu_count(), 8
) # Limit to avoid memory issues
tile_configs = self._get_tile_configs()
trait_combos = self._generate_trait_combinations()
# Prepare work items for parallel processing
work_items = []
for tile_config in tile_configs:
for trait_combo in trait_combos:
work_items.append(
(
tile_config,
trait_combo,
self.kernel_name_prefix,
self.working_path,
self.gpu_target,
self.datatype,
self.layout,
self.config_json,
)
)
# Apply RFC-compliant sampling (Sobol + LHS + maximin)
if self.max_instances is not None and len(work_items) > self.max_instances:
kernel_dicts = [
{
"tile_config": item[0],
"trait_combo": item[1],
"_work_item": item,
}
for item in work_items
]
sampled = self._apply_sampling(kernel_dicts)
work_items = [k["_work_item"] for k in sampled]
print(
f"Generating {len(work_items)} individual kernel files using {num_workers} workers..."
)
print(f" Tile configs: {len(tile_configs)}")
print(f" Trait combinations: {len(trait_combos)}")
print(f" Total kernels: {len(work_items)}")
# Process work items in parallel
kernel_list = []
completed = 0
with concurrent.futures.ProcessPoolExecutor(
max_workers=num_workers
) as executor:
# Submit all work items
print(f" Submitting {len(work_items)} tasks to executor...")
future_to_item = {
executor.submit(_generate_single_kernel_individual, item): item
for item in work_items
}
print(" All tasks submitted, waiting for completion...")
# Collect results with progress reporting
for future in concurrent.futures.as_completed(future_to_item):
completed += 1
if completed % 100 == 0 or completed == len(work_items):
print(
f" Progress: {completed}/{len(work_items)} kernels generated"
)
try:
result = future.result()
if result:
kernel_list.append(result)
except Exception as exc:
item = future_to_item[future]
print(f"Kernel generation failed for {item}: {exc}")
# Sort kernel list for consistent ordering
kernel_list.sort(key=lambda x: x[0]) # Sort by kernel name
# Generate CMake include file for individual targets
self._generate_cmake_individual_targets(kernel_list)
print(
f"Generated {len(kernel_list)} individual kernel files in {self.working_path}"
)
def _generate_single_kernel_individual(work_item):
"""Worker function to generate a single individual kernel file"""
(
tile_config,
trait_combo,
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json,
) = work_item
# Create a temporary builder instance for this worker
builder = GemmAQuantKernelBuilder(
kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json
)
try:
kernel_name, instance_code = builder._generate_kernel_instance(
tile_config, trait_combo
)
# Strip the "gemm_aquant_" prefix so the header filename does not duplicate it
simplified_name = kernel_name.removeprefix(kernel_name_prefix + "_")
# Write individual header file
header_file = working_path / f"gemm_aquant_single_{simplified_name}.hpp"
with open(header_file, "w") as f:
f.write(instance_code)
return (kernel_name, trait_combo, tile_config)
except Exception as e:
print(f"Error generating individual kernel: {e}")
return None
def main():
parser = argparse.ArgumentParser(description="GEMM AQuant kernel instance builder")
parser.add_argument("--working_path", required=True, help="Working directory path")
parser.add_argument("--gpu_target", required=True, help="GPU target architecture")
parser.add_argument(
"--datatype",
required=True,
choices=["fp8", "bf8"],
help="Data type (fp8 or bf8)",
)
parser.add_argument(
"--layout",
required=True,
choices=["rcr", "rrr", "ccr", "crr"],
help="Matrix layout",
)
parser.add_argument("--config_json", help="Configuration JSON file")
parser.add_argument("--num_workers", type=int, help="Number of parallel workers")
parser.add_argument(
"--gen_all_individual",
action="store_true",
help="Generate individual kernel files",
)
parser.add_argument(
"--gen_single",
action="store_true",
help="Generate a single kernel file",
)
parser.add_argument("--kernel_name", help="Kernel name for single generation")
parser.add_argument("--tile_config", help="Tile configuration string")
parser.add_argument("--trait_combo", help="Trait combination string")
parser.add_argument(
"--list_kernels",
action="store_true",
help="List kernel configurations without generating files",
)
parser.add_argument(
"--max-instances",
type=int,
default=None,
help="Maximum number of kernel instances to select via sampling",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="RNG seed for deterministic sampling; if omitted, derived from today's date",
)
parser.add_argument(
"--tier",
default=None,
help="Sampling tier name (e.g., 'daily')",
)
parser.add_argument(
"--manifest-path",
default=None,
help="Directory to write chosen_instances manifest JSON",
)
args = parser.parse_args()
layout_parts = args.layout.lower()
assert len(layout_parts) == 3, (
f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr' where r stands for row major and c stands for column major)"
)
assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], (
f"Invalid matrix_a layout : {layout_parts[0]} or matrix_b layout: {layout_parts[1]} (matrix_a and matrix_b must be either 'r' for row major or 'c' for column major)"
)
assert layout_parts[2] == "r", (
f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' only as currently we are supporting only row major)"
)
kernel_name_prefix = "gemm_aquant"
builder = GemmAQuantKernelBuilder(
kernel_name_prefix,
args.working_path,
args.gpu_target,
args.datatype,
args.layout,
args.config_json,
max_instances=args.max_instances,
seed=args.seed,
tier=args.tier,
manifest_path=args.manifest_path,
)
if args.list_kernels:
builder._list_kernels()
elif args.gen_single:
if not args.kernel_name or not args.tile_config or not args.trait_combo:
parser.error(
"--gen_single requires --kernel_name, --tile_config, and --trait_combo"
)
# Parse tile config
tile_parts = args.tile_config.split("_")
tile_dims = tile_parts[0].split("x")
warp_dims = tile_parts[1].split("x")
warp_tile_dims = tile_parts[2].split("x")
tile_config = {
"tile_m": int(tile_dims[0]),
"tile_n": int(tile_dims[1]),
"tile_k": int(tile_dims[2]),
"warp_m": int(warp_dims[0]),
"warp_n": int(warp_dims[1]),
"warp_k": int(warp_dims[2]),
"warp_tile_m": int(warp_tile_dims[0]),
"warp_tile_n": int(warp_tile_dims[1]),
"warp_tile_k": int(warp_tile_dims[2]),
}
# Parse trait combo:
# pipeline_epilogue_scheduler_padM_padN_padK_aPreshuffle
trait_parts = args.trait_combo.split("_")
trait_combo = (
trait_parts[0], # pipeline
trait_parts[1], # epilogue
trait_parts[2], # scheduler
trait_parts[3] == "True", # pad_m
trait_parts[4] == "True", # pad_n
trait_parts[5] == "True", # pad_k
trait_parts[6] == "True", # a_preshuffle_quant
)
builder._generate_kernel_instance(tile_config, trait_combo)
elif args.gen_all_individual:
builder._generate_all_individual(args.num_workers)
else:
parser.error(
"Must specify one of: --list_kernels, --gen_all_individual, or --gen_single"
)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,276 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <algorithm>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <tuple>
#include <vector>
#include "ck_tile/host/device_prop.hpp"
#include "ck_tile/ops/gemm_quant.hpp"
#include "gemm_aquant_benchmark.hpp"
class AQuantGemmProfiler
{
public:
explicit AQuantGemmProfiler(Setting setting) : setting_(setting) {}
AQuantGemmProfiler(const AQuantGemmProfiler&) = delete;
AQuantGemmProfiler& operator=(const AQuantGemmProfiler&) = delete;
void reset() { kernel_instances_.clear(); }
void benchmark(AQuantGemmProblem& problem,
std::function<float(const ck_tile::QuantGemmHostArgs&,
const ck_tile::stream_config&)> kernel_func)
{
std::vector<std::function<std::tuple<std::string, float>(ck_tile::QuantGemmHostArgs&,
const ck_tile::stream_config&)>>
callables;
callables.push_back(
[kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
float time = kernel_func(args, stream);
return std::make_tuple(std::string(KERNEL_NAME), time);
});
benchmark(problem, callables);
}
void benchmark(AQuantGemmProblem& problem,
std::vector<std::function<std::tuple<std::string, float>(
ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables)
{
const ALayout layout_a = ALayout{};
const BLayout layout_b = BLayout{};
const CLayout layout_c = CLayout{};
const AQLayout layout_aq = AQLayout{};
problem.stride_a_ = ck_tile::get_default_stride(
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a));
problem.stride_b_ = ck_tile::get_default_stride(
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b));
problem.stride_c_ = ck_tile::get_default_stride(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c));
// Compute AQ scale tensor dimensions: [M, K / group_size_k]
const ck_tile::index_t QK_A = problem.k_ / problem.group_size_k_;
problem.stride_aq_ = ck_tile::get_default_stride(
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq));
// Allocate host tensors
ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
// AQ scale tensor: [M, QK_A]
ck_tile::HostTensor<AQDataType> aq_m_qk(ck_tile::host_tensor_descriptor(
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq)));
// Initialize tensors
if(setting_.init_method_ == 0)
{
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
ck_tile::FillUniformDistribution<AQDataType>{0.5f, 1.5f}(aq_m_qk);
}
else if(setting_.init_method_ == 1)
{
ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
}
else if(setting_.init_method_ == 2)
{
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(1)}(a_m_k);
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(1)}(b_k_n);
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
}
else
{
a_m_k.SetZero();
b_k_n.SetZero();
aq_m_qk.SetZero();
}
// Allocate device memory
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
ck_tile::DeviceMem aq_dev_buf(aq_m_qk.get_element_space_size_in_bytes());
a_m_k_dev_buf.ToDevice(a_m_k.data());
b_k_n_dev_buf.ToDevice(b_k_n.data());
aq_dev_buf.ToDevice(aq_m_qk.data());
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
// Build QuantGemmHostArgs
ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
b_k_n_dev_buf.GetDeviceBuffer(),
c_m_n_dev_buf.GetDeviceBuffer(),
aq_dev_buf.GetDeviceBuffer(),
nullptr, // bq_ptr not used for AQuant
problem.split_k_,
problem.m_,
problem.n_,
problem.k_,
QK_A,
0, // QK_B not used for AQuant
problem.stride_a_,
problem.stride_b_,
problem.stride_c_,
problem.stride_aq_,
0 // stride_BQ not used for AQuant
);
// Host reference for verification
ck_tile::HostTensor<CDataType> c_m_n_host_result(ck_tile::host_tensor_descriptor(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
if(setting_.verify_)
{
aquant_gemm_host_reference(setting_.verify_, a_m_k, aq_m_qk, b_k_n, c_m_n_host_result);
}
// Run kernel(s)
for(auto& callable : callables)
{
auto kernel_run_result = callable(gemm_args,
ck_tile::stream_config{nullptr,
true,
setting_.log_,
setting_.n_warmup_,
setting_.n_repeat_,
setting_.is_gpu_timer_,
setting_.flush_cache_,
setting_.rotating_count_});
process_result(problem,
QK_A,
c_m_n_dev_buf,
c_m_n_host_result,
c_m_n_dev_result,
kernel_run_result);
}
}
void process_result(const AQuantGemmProblem& problem,
ck_tile::index_t QK_A,
ck_tile::DeviceMem& c_m_n_dev_buf,
ck_tile::HostTensor<CDataType>& c_m_n_host_result,
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
const std::tuple<std::string, float>& kernel_run_result)
{
auto [name, avg_time] = kernel_run_result;
KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}};
// Compute performance metrics
std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_;
std::size_t num_byte = sizeof(ADataType) * problem.m_ * problem.k_ +
sizeof(BDataType) * problem.n_ * problem.k_ +
sizeof(AQDataType) * problem.m_ * QK_A +
sizeof(CDataType) * problem.m_ * problem.n_;
kernel_instance.perf_result_.latency_ = avg_time;
kernel_instance.perf_result_.tflops_ = static_cast<float>(flop) / 1.E9 / avg_time;
kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time;
if(setting_.log_ > 0 && !setting_.json_output_)
{
std::cout << kernel_instance << std::endl;
}
// Verify result
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
bool verified_correct =
!setting_.verify_ ||
compare_aquant(name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result);
if(verified_correct)
{
kernel_instances_.emplace_back(kernel_instance);
}
else
{
std::cout << "Verification failed, skip kernel: " << name << std::endl;
}
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
}
KernelInstance select_best_instance(Metric metric)
{
if(kernel_instances_.empty())
throw std::runtime_error("Empty instances");
auto kernel_instance = *std::max_element(kernel_instances_.begin(),
kernel_instances_.end(),
[metric](const auto& a, const auto& b) {
return PerformanceResult::compare(
b.perf_result_, a.perf_result_, metric);
});
if(setting_.json_output_)
{
std::cout << kernel_instance << std::endl;
}
else
{
std::cout << "**********************************" << std::endl;
std::cout << "According to given metrics: " << get_metric_name(metric) << "\n"
<< "Current kernel performance is: " << kernel_instance << std::endl;
std::cout << "**********************************" << std::endl;
}
if(!setting_.csv_filename_.empty())
{
std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app);
if(!file.is_open())
{
std::cerr << "Warning: Failed to open CSV file for writing." << std::endl;
}
else
{
if(file.tellp() == 0)
{
file << "rocm_version,device_name,"
<< "split_k,m,n,k,stride_a,stride_b,stride_c,stride_aq,group_size_k,"
<< "dtype_a,dtype_b,dtype_aq,dtype_acc,dtype_c,"
<< "layout_a,layout_b,layout_c," << "name,"
<< "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n";
}
const auto& prob = kernel_instance.problem_;
const auto& perf = kernel_instance.perf_result_;
file << get_rocm_version() << "," << ck_tile::get_device_name() << ","
<< prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << ","
<< prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << ","
<< prob.stride_aq_ << "," << prob.group_size_k_ << "," << prob.dtype_a_ << ","
<< prob.dtype_b_ << "," << prob.dtype_aq_ << "," << prob.dtype_acc_ << ","
<< prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_ << ","
<< prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed
<< std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << ","
<< perf.bandwidth_ << "," << get_metric_name(metric) << "\n";
}
}
return kernel_instance;
}
private:
Setting setting_;
std::vector<KernelInstance> kernel_instances_;
};

View File

@@ -0,0 +1,204 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""Unit tests for GemmAQuantKernelBuilder (gemm_aquant operator)."""
import json
import os
import sys
import tempfile
import unittest
# Make the instance builder importable by inserting the directory into sys.path
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
from gemm_aquant_instance_builder import GemmAQuantKernelBuilder # noqa: E402
# ---------------------------------------------------------------------------
# Minimal default config that mirrors default_config.json, kept small so that
# _get_tile_configs() / _generate_trait_combinations() run quickly in tests.
# ---------------------------------------------------------------------------
_MINIMAL_CONFIG = {
"tile_config": {
"tile_m": {"values": [64]},
"tile_n": {"values": [64]},
"tile_k": {"values": [128]},
"warp_m": {"values": [2]},
"warp_n": {"values": [2]},
"warp_k": {"values": [1]},
"warp_tile_m": {"values": [16]},
"warp_tile_n": {"values": [16]},
"warp_tile_k": {"values": [64]},
},
"trait_config": {
"pipeline": {"values": ["compv3"]},
"scheduler": {"values": ["intrawave"]},
"epilogue": {"values": ["default"]},
"pad_m": {"values": [False]},
"pad_n": {"values": [False]},
"pad_k": {"values": [False]},
"a_preshuffle_quant": {"values": [False, True]},
},
"k_block_per_cu": 1,
"group_size_k": 128,
}
def _make_builder(tmpdir, config=None, **kwargs):
"""Helper: write config to a temp JSON file and create a builder."""
cfg = config if config is not None else _MINIMAL_CONFIG
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(cfg, f)
return GemmAQuantKernelBuilder(
kernel_name_prefix="gemm_aquant",
working_path=tmpdir,
gpu_target=kwargs.get("gpu_target", "gfx942"),
datatype=kwargs.get("datatype", "fp8"),
layout=kwargs.get("layout", "rcr"),
config_json=cfg_path,
)
class TestGemmAQuantBuilderInit(unittest.TestCase):
def test_group_size_k_from_config(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertEqual(builder.group_size_k, 128)
def test_group_size_k_custom(self):
cfg = json.loads(json.dumps(_MINIMAL_CONFIG))
cfg["group_size_k"] = 64
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, config=cfg)
self.assertEqual(builder.group_size_k, 64)
def test_kernel_name_prefix_stored(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertEqual(builder.kernel_name_prefix, "gemm_aquant")
def test_working_path_created(self):
with tempfile.TemporaryDirectory() as tmpdir:
sub = os.path.join(tmpdir, "workdir")
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(_MINIMAL_CONFIG, f)
GemmAQuantKernelBuilder(
"gemm_aquant", sub, "gfx942", "fp8", "rcr", cfg_path
)
self.assertTrue(os.path.isdir(sub))
class TestGemmAQuantTraitCombinations(unittest.TestCase):
def setUp(self):
self._tmpdir = tempfile.TemporaryDirectory()
self.builder = _make_builder(self._tmpdir.name)
def tearDown(self):
self._tmpdir.cleanup()
def test_trait_combinations_non_empty(self):
combos = self.builder._generate_trait_combinations()
self.assertGreater(len(combos), 0)
def test_trait_combo_is_7_tuple(self):
"""AQuant trait tuple: (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant)."""
combos = self.builder._generate_trait_combinations()
for combo in combos:
self.assertEqual(len(combo), 7, f"Expected 7-tuple, got {len(combo)}: {combo}")
def test_pipeline_values_present(self):
combos = self.builder._generate_trait_combinations()
pipelines = {c[0] for c in combos}
self.assertIn("compv3", pipelines)
def test_preshuffle_quant_values(self):
"""Both True and False should appear when config lists both."""
combos = self.builder._generate_trait_combinations()
preshuffle_vals = {c[6] for c in combos}
self.assertIn(True, preshuffle_vals)
self.assertIn(False, preshuffle_vals)
def test_uses_persistent_trait_returns_true(self):
"""AQuant overrides _uses_persistent_trait to accommodate a_preshuffle_quant slot."""
self.assertTrue(self.builder._uses_persistent_trait())
class TestGemmAQuantTileConfigs(unittest.TestCase):
def setUp(self):
self._tmpdir = tempfile.TemporaryDirectory()
self.builder = _make_builder(self._tmpdir.name)
def tearDown(self):
self._tmpdir.cleanup()
def test_tile_configs_non_empty(self):
configs = self.builder._get_tile_configs()
self.assertGreater(len(configs), 0)
def test_tile_config_has_required_keys(self):
configs = self.builder._get_tile_configs()
required = {
"tile_m", "tile_n", "tile_k",
"warp_m", "warp_n", "warp_k",
"warp_tile_m", "warp_tile_n", "warp_tile_k",
}
for cfg in configs:
self.assertTrue(required.issubset(cfg.keys()), f"Missing keys in {cfg}")
def test_tile_config_values_are_positive(self):
configs = self.builder._get_tile_configs()
for cfg in configs:
for key, val in cfg.items():
self.assertGreater(val, 0, f"{key}={val} must be positive")
class TestGemmAQuantLayoutVariants(unittest.TestCase):
"""Verify the builder initialises correctly for all supported layouts."""
LAYOUTS = ["rcr", "rrr", "ccr", "crr"]
def test_all_layouts_construct(self):
for layout in self.LAYOUTS:
with self.subTest(layout=layout):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, layout=layout)
self.assertEqual(builder.layout, layout)
class TestGemmAQuantDataTypes(unittest.TestCase):
DTYPES = ["fp8", "bf8"]
def test_all_dtypes_construct(self):
for dtype in self.DTYPES:
with self.subTest(dtype=dtype):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, datatype=dtype)
self.assertEqual(builder.datatype, dtype)
class TestGemmAQuantMaxInstances(unittest.TestCase):
def test_max_instances_none_returns_all(self):
"""Without a max_instances cap, sampling is a no-op."""
with tempfile.TemporaryDirectory() as tmpdir:
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(_MINIMAL_CONFIG, f)
builder = GemmAQuantKernelBuilder(
"gemm_aquant", tmpdir, "gfx942", "fp8", "rcr",
config_json=cfg_path, max_instances=None,
)
kernel_list = [
{"tile_config": {"tile_m": 64, "tile_n": 64, "tile_k": 128,
"warp_m": 2, "warp_n": 2, "warp_k": 1,
"warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 64},
"trait_combo": ("compv3", "default", "intrawave", False, False, False, False)}
]
result = builder._apply_sampling(kernel_list)
self.assertEqual(len(result), 1)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,289 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
set(GEMM_BQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM BQuant (semicolon-separated)")
set(GEMM_BQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM BQuant (semicolon-separated)")
set(GEMM_BQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)")
option(ENABLE_CCACHE_GEMM_BQUANT "Enable ccache for GEMM BQuant ops compilation" OFF)
# Store the directory path for use in functions
set(GEMM_BQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR})
# Function to create individual GEMM BQuant targets
function(create_individual_gemm_bquant_target datatype layout trait tile_config config_json)
# GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL is guaranteed non-empty here (caller checks at line 192)
if(NOT GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL)
message(FATAL_ERROR "BUG: create_individual_gemm_bquant_target called with no GPU targets")
endif()
# Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k
string(REPLACE "_" ";" config_groups ${tile_config})
list(GET config_groups 0 tile_dims)
list(GET config_groups 1 warp_dims)
list(GET config_groups 2 warp_tile_dims)
string(REPLACE "x" ";" tile_parts ${tile_dims})
list(GET tile_parts 0 tile_m)
list(GET tile_parts 1 tile_n)
list(GET tile_parts 2 tile_k)
string(REPLACE "x" ";" warp_parts ${warp_dims})
list(GET warp_parts 0 warp_m)
list(GET warp_parts 1 warp_n)
list(GET warp_parts 2 warp_k)
string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims})
list(GET warp_tile_parts 0 warp_tile_m)
list(GET warp_tile_parts 1 warp_tile_n)
list(GET warp_tile_parts 2 warp_tile_k)
set(target_name "benchmark_gemm_bquant_${datatype}_${layout}_${trait}_${tile_config}")
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
# Generate the single instance header for this kernel
set(instance_header "${working_path}/gemm_bquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp")
# Add custom command to generate the header file at build time
add_custom_command(
OUTPUT ${instance_header}
COMMAND ${Python3_EXECUTABLE} ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_instance_builder.py
--working_path ${working_path}
--datatype ${datatype}
--layout ${layout}
--config_json ${config_json}
--gen_single
--kernel_name "gemm_bquant_${datatype}_${layout}_${trait}_${tile_config}"
--tile_config "${tile_config}"
--trait_combo "${trait}"
--gpu_target "${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}"
DEPENDS ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_instance_builder.py ${config_json}
COMMENT "Generating ${instance_header}"
)
# Create the executable
add_executable(${target_name}
EXCLUDE_FROM_ALL
${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_benchmark_single.cpp
${instance_header}
)
# Set GPU architectures
set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL})
# Set compile definitions
target_compile_definitions(${target_name} PRIVATE
GEMM_BQUANT_SINGLE_INSTANCE_HPP="${instance_header}"
)
# Include directories
target_include_directories(${target_name} PRIVATE
${GEMM_BQUANT_SOURCE_DIR}
${working_path}
)
# Compile options
target_compile_options(${target_name} PRIVATE
-Wno-undefined-func-template
-Wno-float-equal
--offload-compress
-include ${instance_header}
)
# Add to collection targets
add_dependencies(benchmark_gemm_bquant_all ${target_name})
add_dependencies(benchmark_gemm_bquant_${datatype} ${target_name})
add_dependencies(benchmark_gemm_bquant_${layout} ${target_name})
add_dependencies(benchmark_gemm_bquant_${datatype}_${layout} ${target_name})
# Add to pipeline-specific, epilogue-specific, and scheduler-specific targets
string(REPLACE "_" ";" trait_parts ${trait})
list(GET trait_parts 0 pipeline)
list(GET trait_parts 1 epilogue)
list(GET trait_parts 2 scheduler)
add_dependencies(benchmark_gemm_bquant_${pipeline}_pipeline ${target_name})
add_dependencies(benchmark_gemm_bquant_${epilogue}_epilogue ${target_name})
add_dependencies(benchmark_gemm_bquant_${scheduler}_scheduler ${target_name})
endfunction()
# Function to build individual GEMM BQuant targets
function(build_individual_gemm_bquant_targets datatype layout)
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
# Choose config file (priority: env var > cmake var > default)
if(DEFINED ENV{GEMM_BQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_BQUANT_CONFIG_FILE}" STREQUAL "")
set(config_filename "$ENV{GEMM_BQUANT_CONFIG_FILE}")
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}")
message(VERBOSE " Using config from environment variable: ${config_filename}")
elseif(NOT "${GEMM_BQUANT_CONFIG_FILE}" STREQUAL "")
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_BQUANT_CONFIG_FILE}")
message(VERBOSE " Using custom config: ${GEMM_BQUANT_CONFIG_FILE}")
else()
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json")
message(VERBOSE " Using default config for layout ${layout}")
endif()
if(NOT EXISTS ${json_blob})
message(FATAL_ERROR "Config file not found: ${json_blob}")
endif()
# Create working directory
file(MAKE_DIRECTORY ${working_path})
# Build sampling arguments
set(extra_list_args "")
if(NOT "${GEMM_BQUANT_MAX_INSTANCES}" STREQUAL "")
list(APPEND extra_list_args --max-instances ${GEMM_BQUANT_MAX_INSTANCES})
endif()
if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER})
list(APPEND extra_list_args --manifest-path ${working_path})
endif()
if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "")
list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED})
endif()
# List kernels
message(VERBOSE " Listing BQuant kernel configurations for ${datatype} ${layout}...")
execute_process(
COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_bquant_instance_builder.py
--working_path ${working_path}
--datatype ${datatype}
--layout ${layout}
--config_json ${json_blob}
--gpu_target "${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}"
--list_kernels
${extra_list_args}
WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
RESULT_VARIABLE ret
OUTPUT_VARIABLE list_output
ERROR_VARIABLE list_error
)
if(NOT ret EQUAL 0)
message(FATAL_ERROR "Failed to list BQuant kernels for ${datatype} ${layout}: ${list_error}")
endif()
# Read kernel count
if(EXISTS ${working_path}/gemm_bquant_kernel_count.txt)
file(READ ${working_path}/gemm_bquant_kernel_count.txt kernel_count)
string(STRIP "${kernel_count}" kernel_count)
message(VERBOSE " Found ${kernel_count} BQuant kernel configurations")
else()
message(FATAL_ERROR "Kernel count file not found")
endif()
# Read kernel list and create targets
if(EXISTS ${working_path}/gemm_bquant_kernel_list.txt)
file(STRINGS ${working_path}/gemm_bquant_kernel_list.txt kernel_lines)
foreach(line IN LISTS kernel_lines)
string(REPLACE "|" ";" parts "${line}")
list(GET parts 0 kernel_name)
list(GET parts 1 tile_config)
list(GET parts 2 trait_combo)
create_individual_gemm_bquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}")
endforeach()
else()
message(FATAL_ERROR "Kernel list file not found")
endif()
endfunction()
# Main build logic
message(VERBOSE "=== Starting Tile Engine GEMM BQuant Configuration ===")
message(VERBOSE "GEMM_BQUANT_DATATYPE: ${GEMM_BQUANT_DATATYPE}")
message(VERBOSE "GEMM_BQUANT_LAYOUT: ${GEMM_BQUANT_LAYOUT}")
message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
# Filter GPU targets to supported architectures
set(GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL "")
set(DESIRED_TARGETS "gfx90a;gfx942;gfx950")
foreach(target IN LISTS SUPPORTED_GPU_TARGETS)
if(target IN_LIST DESIRED_TARGETS)
list(APPEND GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL ${target})
message(VERBOSE " Adding GPU target: ${target}")
endif()
endforeach()
# Skip build if no matching targets found
if(NOT GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL)
message(WARNING "Skipping Tile Engine GEMM BQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
else()
message(VERBOSE "Building individual GEMM BQuant targets for GPU targets: ${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}")
# Enable compiler cache if requested
if(ENABLE_CCACHE_GEMM_BQUANT)
find_program(CCACHE_PROGRAM ccache)
if(CCACHE_PROGRAM)
set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
message(VERBOSE "Using ccache for faster compilation")
endif()
endif()
# Create master collection target
add_custom_target(benchmark_gemm_bquant_all)
# Create datatype collection targets
foreach(dt IN LISTS GEMM_BQUANT_DATATYPE)
add_custom_target(benchmark_gemm_bquant_${dt})
endforeach()
# Create layout collection targets
foreach(l IN LISTS GEMM_BQUANT_LAYOUT)
add_custom_target(benchmark_gemm_bquant_${l})
endforeach()
# Create combined collection targets
foreach(dt IN LISTS GEMM_BQUANT_DATATYPE)
foreach(l IN LISTS GEMM_BQUANT_LAYOUT)
add_custom_target(benchmark_gemm_bquant_${dt}_${l})
endforeach()
endforeach()
# Create pipeline-specific collection targets
set(GEMM_BQUANT_PIPELINES "compv3")
foreach(pipeline IN LISTS GEMM_BQUANT_PIPELINES)
add_custom_target(benchmark_gemm_bquant_${pipeline}_pipeline)
endforeach()
# Create epilogue-specific collection targets
set(GEMM_BQUANT_EPILOGUES "default;cshuffle")
foreach(epilogue IN LISTS GEMM_BQUANT_EPILOGUES)
add_custom_target(benchmark_gemm_bquant_${epilogue}_epilogue)
endforeach()
# Create scheduler-specific collection targets
set(GEMM_BQUANT_SCHEDULERS "intrawave")
foreach(scheduler IN LISTS GEMM_BQUANT_SCHEDULERS)
add_custom_target(benchmark_gemm_bquant_${scheduler}_scheduler)
endforeach()
# Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that
# sampling fires per-combo rather than being a single cap.
# CMake integer division truncates; clamp the result to at least 1 so that a
# small total budget (e.g. 5 instances across 8 combos) does not silently
# produce a per-combo budget of 0 and skip all kernels.
if(NOT "${GEMM_BQUANT_MAX_INSTANCES}" STREQUAL "")
list(LENGTH GEMM_BQUANT_DATATYPE _bq_n_dt)
list(LENGTH GEMM_BQUANT_LAYOUT _bq_n_lay)
math(EXPR _bq_n_combos "${_bq_n_dt} * ${_bq_n_lay}")
if(_bq_n_combos GREATER 0)
math(EXPR GEMM_BQUANT_MAX_INSTANCES
"${GEMM_BQUANT_MAX_INSTANCES} / ${_bq_n_combos}")
if(GEMM_BQUANT_MAX_INSTANCES EQUAL 0)
set(GEMM_BQUANT_MAX_INSTANCES 1)
message(WARNING " gemm_bquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_bq_n_combos} combos)")
else()
message(STATUS " gemm_bquant: per-combo budget = ${GEMM_BQUANT_MAX_INSTANCES} (${_bq_n_combos} combos)")
endif()
endif()
endif()
# Build individual targets for each datatype/layout combination
foreach(dt IN LISTS GEMM_BQUANT_DATATYPE)
foreach(l IN LISTS GEMM_BQUANT_LAYOUT)
build_individual_gemm_bquant_targets(${dt} ${l})
endforeach()
endforeach()
endif()

View File

@@ -0,0 +1,93 @@
{
"tile_config": {
"tile_m": {
"values": [
16
]
},
"tile_n": {
"values": [
64
]
},
"tile_k": {
"values": [
256
]
},
"warp_m": {
"values": [
1
]
},
"warp_n": {
"values": [
4
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16
]
},
"warp_tile_n": {
"values": [
16
]
},
"warp_tile_k": {
"values": [
16,
32,
64,
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3"
]
},
"scheduler": {
"values": [
"intrawave"
]
},
"epilogue": {
"values": [
"default",
"cshuffle"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"b_preshuffle_quant": {
"values": [
false,
true
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128
}

View File

@@ -0,0 +1,102 @@
{
"tile_config": {
"tile_m": {
"max": 256,
"min": 64,
"step": 64
},
"tile_n": {
"max": 256,
"min": 64,
"step": 64
},
"tile_k": {
"max": 256,
"min": 64,
"step": 64
},
"warp_m": {
"values": [
4,
2,
1
]
},
"warp_n": {
"values": [
4,
2,
1
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
4,
16,
32
]
},
"warp_tile_n": {
"values": [
16,
32,
64
]
},
"warp_tile_k": {
"values": [
8,
16,
32,
64,
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3"
]
},
"scheduler": {
"values": [
"intrawave"
]
},
"epilogue": {
"values": [
"cshuffle",
"default"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"b_preshuffle_quant": {
"values": [
false,
true
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128
}

View File

@@ -0,0 +1,88 @@
{
"tile_config": {
"tile_m": {
"values": [
16
]
},
"tile_n": {
"values": [
64
]
},
"tile_k": {
"values": [
256
]
},
"warp_m": {
"values": [
1
]
},
"warp_n": {
"values": [
4
]
},
"warp_k": {
"values": [
1
]
},
"warp_tile_m": {
"values": [
16
]
},
"warp_tile_n": {
"values": [
16
]
},
"warp_tile_k": {
"values": [
128
]
}
},
"trait_config": {
"pipeline": {
"values": [
"compv3"
]
},
"scheduler": {
"values": [
"intrawave"
]
},
"epilogue": {
"values": [
"default"
]
},
"pad_m": {
"values": [
false
]
},
"pad_n": {
"values": [
false
]
},
"pad_k": {
"values": [
false
]
},
"b_preshuffle_quant": {
"values": [
false
]
}
},
"k_block_per_cu": 1,
"group_size_k": 128
}

View File

@@ -0,0 +1,229 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "gemm_bquant_common.hpp"
// Data types and Layouts are defined by the generated kernel headers:
// ADataType, BDataType, BQDataType, AccDataType, CDataType
// ALayout, BLayout, CLayout, BQLayout
enum class Metric
{
LATENCY = 0,
TFLOPS = 1,
BANDWIDTH = 2
};
inline constexpr auto get_metric_name(Metric m)
{
switch(m)
{
case Metric::LATENCY: return "latency";
case Metric::TFLOPS: return "tflops";
case Metric::BANDWIDTH: return "bandwidth";
default: throw std::invalid_argument("Unsupported metric type");
}
}
struct BQuantGemmProblem
{
int split_k_;
int m_, n_, k_;
int stride_a_, stride_b_, stride_c_;
int stride_bq_;
int group_size_k_;
std::string dtype_a_, dtype_b_, dtype_bq_, dtype_acc_, dtype_c_;
std::string layout_a_, layout_b_, layout_c_;
friend std::ostream& operator<<(std::ostream& os, const BQuantGemmProblem& problem)
{
os << "{\n"
<< " \"split_k\":" << problem.split_k_ << ",\n"
<< " \"m\":" << problem.m_ << ",\n"
<< " \"n\":" << problem.n_ << ",\n"
<< " \"k\":" << problem.k_ << ",\n"
<< " \"stride_a\":" << problem.stride_a_ << ",\n"
<< " \"stride_b\":" << problem.stride_b_ << ",\n"
<< " \"stride_c\":" << problem.stride_c_ << ",\n"
<< " \"stride_bq\":" << problem.stride_bq_ << ",\n"
<< " \"group_size_k\":" << problem.group_size_k_ << ",\n"
<< " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n"
<< " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n"
<< " \"dtype_bq\":\"" << problem.dtype_bq_ << "\",\n"
<< " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n"
<< " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n"
<< " \"layout_a\":\"" << problem.layout_a_ << "\",\n"
<< " \"layout_b\":\"" << problem.layout_b_ << "\",\n"
<< " \"layout_c\":\"" << problem.layout_c_ << "\"\n"
<< "}";
return os;
}
};
struct PerformanceResult
{
double latency_;
double tflops_;
double bandwidth_;
static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m)
{
switch(m)
{
case Metric::LATENCY: return a.latency_ < b.latency_;
case Metric::TFLOPS: return a.tflops_ > b.tflops_;
case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_;
default: throw std::invalid_argument("Unsupported metric type");
}
}
friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result)
{
os << "{\n"
<< " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_
<< ",\n"
<< " \"tflops(TFlops)\": " << result.tflops_ << ",\n"
<< " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n"
<< "}";
return os;
}
};
struct KernelInstance
{
std::string name_;
BQuantGemmProblem problem_;
PerformanceResult perf_result_;
static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m)
{
return PerformanceResult::compare(a.perf_result_, b.perf_result_, m);
}
friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj)
{
os << "{\n"
<< " \"name\": \"" << obj.name_ << "\",\n"
<< " \"problem\": " << obj.problem_ << ",\n"
<< " \"perf_result\": " << obj.perf_result_ << "\n"
<< "}";
return os;
}
};
struct Setting
{
int n_warmup_;
int n_repeat_;
bool is_gpu_timer_;
int verify_;
int init_method_;
bool log_;
std::string csv_filename_;
bool flush_cache_;
int rotating_count_;
bool json_output_;
};
inline std::string get_rocm_version()
{
std::ifstream version_file("/opt/rocm/.info/version");
if(version_file.is_open())
{
std::string version;
std::getline(version_file, version);
return version;
}
return "Unknown";
}
template <typename ADataType_,
typename BDataType_,
typename BQDataType_,
typename AccDataType_,
typename CDataType_>
auto calculate_rtol_atol_bquant(const ck_tile::index_t K,
const ck_tile::index_t kbatch,
const float max_accumulated_value)
{
using ComputeType =
std::conditional_t<sizeof(ADataType_) < sizeof(BDataType_), ADataType_, BDataType_>;
const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType_, AccDataType_>(
ck_tile::integer_divide_ceil(K, kbatch));
const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType_, AccDataType_>(
max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
const auto rtol_split_k =
ck_tile::get_relative_threshold<CDataType_, CDataType_, CDataType_>(kbatch);
const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType_, CDataType_, CDataType_>(
max_accumulated_value, kbatch);
return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
}
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions"
/// @brief Compare device and host results for BQuant GEMM
bool compare_bquant(std::string instanceName,
ck_tile::index_t K,
ck_tile::index_t kbatch,
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
const float max_accumulated_value =
std::abs(static_cast<float>(*std::max_element(c_m_n_host_result.mData.begin(),
c_m_n_host_result.mData.end(),
[](const auto& a, const auto& b) {
return std::abs(static_cast<float>(a)) <
std::abs(static_cast<float>(b));
})));
const auto rtol_atol =
calculate_rtol_atol_bquant<ADataType, BDataType, BQDataType, AccDataType, CDataType>(
K, kbatch, max_accumulated_value);
bool pass = ck_tile::check_err(c_m_n_dev_result,
c_m_n_host_result,
"Error: Incorrect results!",
rtol_atol.at(ck_tile::number<0>{}),
rtol_atol.at(ck_tile::number<1>{}));
std::cout << "For " << instanceName << " Relative error threshold is "
<< rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
<< rtol_atol.at(ck_tile::number<1>{}) << std::endl;
std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;
return pass;
}
/// @brief CPU reference implementation for BQuant GEMM
void bquant_gemm_host_reference(int verify,
ck_tile::HostTensor<ADataType>& a_m_k,
ck_tile::HostTensor<BDataType>& b_k_n,
ck_tile::HostTensor<BQDataType>& bq_qk_n,
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
if(verify == 1)
{
c_m_n_host_result.SetZero();
using QuantGroupSize = typename SelectedKernel::QuantGroupSize;
ck_tile::reference_gemm_quant<ADataType,
BQDataType,
BDataType,
AccDataType,
CDataType,
QuantGroupSize,
false /* aquant=false, bquant */>(
a_m_k, bq_qk_n, b_k_n, c_m_n_host_result);
}
}
#pragma clang diagnostic pop

View File

@@ -0,0 +1,464 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import sys
import json
import subprocess
import argparse
import time
from pathlib import Path
from typing import List, Dict, Tuple, Optional
class BQuantGemmBenchmark:
def __init__(self, build_dir: str, verbose: bool = False):
self.build_dir = Path(build_dir)
self.verbose = verbose
self.results = []
def discover_kernels(self) -> List[Path]:
"""Find all benchmark_gemm_bquant_* executables in the build directory."""
bin_dir = self.build_dir / "bin"
if not bin_dir.exists():
print(f"Error: Binary directory {bin_dir} does not exist")
return []
kernels = list(bin_dir.glob("benchmark_gemm_bquant_*"))
if self.verbose:
print(f"Found {len(kernels)} BQuant kernel executables")
for k in kernels:
print(f" - {k.name}")
return kernels
def extract_kernel_info(self, kernel_path: Path) -> Dict[str, str]:
"""Extract kernel information from filename."""
name = kernel_path.stem
info = {
"executable": str(kernel_path),
"name": name,
"data_type": "unknown",
"layout": "unknown",
"pipeline": "unknown",
"scheduler": "unknown",
"epilogue": "unknown",
"b_preshuffle_quant": False,
"preshuffle_b": False,
}
# Parse: benchmark_gemm_bquant_fp8_rcr_compv3_default_intrawave_False_False_True_False_False_128x128x128_...
parts = name.split("_")
if len(parts) >= 3:
# Skip "benchmark_gemm_bquant" prefix (3 parts)
idx = 3
if idx < len(parts):
info["data_type"] = parts[idx]
idx += 1
if idx < len(parts):
info["layout"] = parts[idx]
idx += 1
if idx < len(parts):
info["pipeline"] = parts[idx]
idx += 1
if idx < len(parts):
info["epilogue"] = parts[idx]
idx += 1
if idx < len(parts):
info["scheduler"] = parts[idx]
# Parse tile dimensions
config_info = self.parse_detailed_config(name)
info.update(config_info)
info["config_id"] = self.generate_config_id(info)
return info
def parse_detailed_config(self, kernel_name: str) -> Dict:
"""Parse tile dimensions and boolean flags from kernel name."""
config = {
"tile_sizes": {"tile_m": 0, "tile_n": 0, "tile_k": 0},
"warp_config": {"warp_m": 0, "warp_n": 0, "warp_k": 0},
"warp_tile": {"warp_tile_m": 0, "warp_tile_n": 0, "warp_tile_k": 0},
"optimization_flags": {
"pad_m": False,
"pad_n": False,
"pad_k": False,
"b_preshuffle_quant": False,
"preshuffle_b": False,
},
}
parts = kernel_name.split("_")
# Locate the first boolean token; collect the contiguous boolean run
bool_start = -1
bool_sequence = []
for part in parts:
if part in ("True", "False"):
if bool_start == -1:
bool_start = 0
bool_sequence.append(part == "True")
elif bool_start != -1:
break
if len(bool_sequence) >= 5:
config["optimization_flags"]["pad_m"] = bool_sequence[0]
config["optimization_flags"]["pad_n"] = bool_sequence[1]
config["optimization_flags"]["pad_k"] = bool_sequence[2]
config["optimization_flags"]["b_preshuffle_quant"] = bool_sequence[3]
config["optimization_flags"]["preshuffle_b"] = bool_sequence[4]
# Extract dimension groups (e.g., 128x128x128) in positional order.
# Kernel names encode groups as: tile_MxNxK_warp_MxNxK_warp_tile_MxNxK
dimension_groups = []
for part in parts:
if "x" in part and len(part.split("x")) == 3:
try:
dims = [int(x) for x in part.split("x")]
if all(d > 0 for d in dims):
dimension_groups.append(dims)
except ValueError:
continue
if len(dimension_groups) >= 3:
# Use positional order: tile, warp, warp_tile
config["tile_sizes"]["tile_m"] = dimension_groups[0][0]
config["tile_sizes"]["tile_n"] = dimension_groups[0][1]
config["tile_sizes"]["tile_k"] = dimension_groups[0][2]
config["warp_config"]["warp_m"] = dimension_groups[1][0]
config["warp_config"]["warp_n"] = dimension_groups[1][1]
config["warp_config"]["warp_k"] = dimension_groups[1][2]
config["warp_tile"]["warp_tile_m"] = dimension_groups[2][0]
config["warp_tile"]["warp_tile_n"] = dimension_groups[2][1]
config["warp_tile"]["warp_tile_k"] = dimension_groups[2][2]
return config
def generate_config_id(self, info: Dict) -> str:
"""Generate a compact config ID from kernel info."""
parts = [
info.get("data_type", "unk"),
info.get("layout", "unk"),
info.get("pipeline", "unk"),
info.get("scheduler", "unk"),
]
tile_sizes = info.get("tile_sizes", {})
if tile_sizes.get("tile_m", 0) > 0:
parts.append(
f"{tile_sizes['tile_m']}x{tile_sizes['tile_n']}x{tile_sizes['tile_k']}"
)
return "_".join(parts)
def run_kernel(self, kernel_path: Path, params: Dict[str, str]) -> Optional[Dict]:
"""Run a single kernel with given parameters."""
results_dir = self.build_dir / "results"
results_dir.mkdir(exist_ok=True)
json_file = results_dir / f"{kernel_path.stem}.json"
cmd = [str(kernel_path)]
for key, value in params.items():
cmd.append(f"-{key}={value}")
cmd.append("-json_output=true")
if self.verbose:
print(f"Running: {' '.join(cmd)}")
try:
result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
if result.returncode != 0:
print(f"Error running {kernel_path.name}: {result.stderr}")
return None
output = result.stdout.strip()
if output:
with open(json_file, "w") as f:
f.write(output)
return self.parse_json_file(json_file)
else:
print(f"No output from {kernel_path.name}")
return None
except subprocess.TimeoutExpired:
print(f"Timeout running {kernel_path.name}")
return None
except Exception as e:
print(f"Error running {kernel_path.name}: {e}")
return None
def parse_json_file(self, json_file: Path) -> Optional[Dict]:
"""Parse JSON data from kernel output file."""
try:
with open(json_file, "r") as f:
content = f.read().strip()
data = json.loads(content)
result = data.copy()
if "perf_result" in data:
perf = data["perf_result"]
result["time_ms"] = perf.get("latency(ms)", 0)
result["tflops"] = perf.get("tflops(TFlops)", 0)
result["bandwidth_gb_s"] = perf.get("bandwidth(GB/s)", 0)
return result
except (json.JSONDecodeError, Exception) as e:
if self.verbose:
print(f"Failed to parse JSON from {json_file}: {e}")
return None
def benchmark_problem_size(
self,
kernels: List[Path],
m: int,
n: int,
k: int,
group_size_k: int = 128,
split_k: int = 1,
verify: int = 0,
warmup: int = 50,
repeat: int = 100,
flush_cache: bool = True,
rotating_count: int = 1000,
) -> List[Dict]:
"""Benchmark all kernels for a specific problem size."""
results = []
params = {
"m": m,
"n": n,
"k": k,
"group_size_k": group_size_k,
"split_k": split_k,
"verify": verify,
"warmup": warmup,
"repeat": repeat,
"flush_cache": str(flush_cache).lower(),
"rotating_count": rotating_count,
}
print(f"\nBenchmarking M={m}, N={n}, K={k}, group_size_k={group_size_k}")
for kernel_path in kernels:
kernel_info = self.extract_kernel_info(kernel_path)
result = self.run_kernel(kernel_path, params)
if result:
structured_result = {
"name": kernel_info["name"],
"config_id": kernel_info["config_id"],
"problem": result.get("problem", {}),
"perf_result": result.get("perf_result", {}),
"config": {
"data_type": kernel_info["data_type"],
"layout": kernel_info["layout"],
"pipeline": kernel_info["pipeline"],
"scheduler": kernel_info["scheduler"],
"epilogue": kernel_info["epilogue"],
"tile_sizes": kernel_info.get("tile_sizes", {}),
"warp_config": kernel_info.get("warp_config", {}),
"warp_tile": kernel_info.get("warp_tile", {}),
"optimization_flags": kernel_info.get("optimization_flags", {}),
},
"executable": kernel_info["executable"],
"time_ms": result.get("time_ms", 0),
"tflops": result.get("tflops", 0),
"bandwidth_gb_s": result.get("bandwidth_gb_s", 0),
}
results.append(structured_result)
if self.verbose:
print(
f" {kernel_info['config_id']}: "
f"{structured_result['tflops']:.2f} TFLOPS, "
f"{structured_result['bandwidth_gb_s']:.2f} GB/s, "
f"{structured_result['time_ms']:.2f}ms"
)
return results
def find_best_kernel(
self, results: List[Dict], metric: str = "tflops"
) -> Optional[Dict]:
"""Find the best performing kernel based on metric."""
if not results:
return None
if metric == "tflops":
return max(results, key=lambda x: x.get("tflops", 0))
elif metric == "time_ms":
return min(results, key=lambda x: x.get("time_ms", float("inf")))
elif metric == "bandwidth_gb_s":
return max(results, key=lambda x: x.get("bandwidth_gb_s", 0))
else:
raise ValueError(f"Unknown metric: {metric}")
def benchmark_sweep(
self,
problem_sizes: List[Tuple[int, int, int]],
group_size_k: int = 128,
split_k_values: Optional[List[int]] = None,
verify: bool = False,
warmup: int = 50,
repeat: int = 100,
flush_cache: bool = True,
rotating_count: int = 1000,
) -> Dict:
"""Run comprehensive benchmark sweep."""
if split_k_values is None:
split_k_values = [1]
kernels = self.discover_kernels()
if not kernels:
print("No kernels found!")
return {}
all_results = []
best_kernels = {}
for m, n, k in problem_sizes:
for split_k in split_k_values:
results = self.benchmark_problem_size(
kernels,
m,
n,
k,
group_size_k=group_size_k,
split_k=split_k,
verify=1 if verify else 0,
warmup=warmup,
repeat=repeat,
flush_cache=flush_cache,
rotating_count=rotating_count,
)
all_results.extend(results)
best = self.find_best_kernel(results)
if best:
key = f"m{m}_n{n}_k{k}_splitk{split_k}"
best_kernels[key] = best
print(
f"Best for {key}: {best['name']} "
f"({best['tflops']:.2f} TFLOPS, "
f"{best['bandwidth_gb_s']:.2f} GB/s, "
f"{best['time_ms']:.2f}ms)"
)
self.results = all_results
return best_kernels
def export_json(self, filename: str, best_kernels: Dict = None):
"""Export results to JSON."""
from datetime import datetime
output_data = {
"benchmark_metadata": {
"timestamp": datetime.now().isoformat(),
"operator": "gemm_bquant",
"total_kernels_tested": len(self.results),
"successful_runs": len(
[r for r in self.results if r.get("tflops", 0) > 0]
),
},
"kernel_results": self.results,
"best_kernels_by_problem": best_kernels or {},
}
with open(filename, "w") as f:
json.dump(output_data, f, indent=2)
print(f"JSON results exported to {filename}")
def main():
parser = argparse.ArgumentParser(description="BQuant GEMM Kernel Benchmarking Tool")
parser.add_argument(
"build_dir", help="Build directory containing kernel executables"
)
parser.add_argument(
"--problem-sizes",
nargs="+",
default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"],
help="Problem sizes as M,N,K tuples",
)
parser.add_argument(
"--group-size-k",
type=int,
default=128,
help="Quantization group size along K (default: 128)",
)
parser.add_argument(
"--split-k",
nargs="+",
type=int,
default=[1],
help="Split-K values to test",
)
parser.add_argument("--verify", action="store_true", help="Enable verification")
parser.add_argument("--verbose", action="store_true", help="Verbose output")
parser.add_argument(
"--warmup",
type=int,
default=50,
help="Number of warmup iterations",
)
parser.add_argument(
"--repeat",
type=int,
default=100,
help="Number of benchmark iterations",
)
parser.add_argument(
"--no-flush-cache",
action="store_true",
help="Disable cache flushing between iterations",
)
parser.add_argument(
"--rotating-count",
type=int,
default=1000,
help="Number of iterations to rotate the cache (default: 1000)",
)
parser.add_argument("--json", help="JSON output filename (optional)")
args = parser.parse_args()
problem_sizes = []
for size_str in args.problem_sizes:
try:
m, n, k = map(int, size_str.split(","))
problem_sizes.append((m, n, k))
except ValueError:
print(f"Invalid problem size: {size_str}")
return 1
benchmark = BQuantGemmBenchmark(args.build_dir, verbose=args.verbose)
print("Starting BQuant GEMM kernel benchmark sweep...")
start_time = time.time()
best_kernels = benchmark.benchmark_sweep(
problem_sizes=problem_sizes,
group_size_k=args.group_size_k,
split_k_values=args.split_k,
verify=args.verify,
warmup=args.warmup,
repeat=args.repeat,
flush_cache=not args.no_flush_cache,
rotating_count=args.rotating_count,
)
elapsed_time = time.time() - start_time
print(f"\nBenchmark completed in {elapsed_time:.2f} seconds")
if args.json:
benchmark.export_json(args.json, best_kernels)
return 0
if __name__ == "__main__":
sys.exit(main())

View File

@@ -0,0 +1,166 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <iostream>
#include <functional>
#include <tuple>
#include <exception>
#include <sstream>
#include <vector>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "gemm_bquant_profiler.hpp"
#include "gemm_bquant_common.hpp"
// The kernel header is included via the compile command line with -include flag
// It defines SelectedKernel struct, KERNEL_NAME, and type aliases:
// ADataType, BDataType, BQDataType, AccDataType, CDataType
// ALayout, BLayout, CLayout, BQLayout
inline auto create_args(int argc, char* argv[])
{
ck_tile::ArgParser arg_parser;
arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.")
.insert("n", "4096", "The value for n dimension. Default is 4096.")
.insert("k", "2048", "The value for k dimension. Default is 2048.")
.insert("stride_a", "0", "The stride value for tensor A. Default is 0.")
.insert("stride_b", "0", "The stride value for tensor B. Default is 0.")
.insert("stride_c", "0", "The stride value for tensor C. Default is 0.")
.insert("split_k", "1", "The split value for k dimension. Default is 1.")
.insert("group_size_k",
std::to_string(SelectedKernel::GroupSizeK),
"The quantization group size along K. Default matches kernel config.")
.insert("verify",
"1",
"The type of validation. Set to 0 for no validation, 1 for validation on CPU. "
"Default is 1, CPU validation.")
.insert("log",
"false",
"Whether output kernel instance information or not. Possible values are true or "
"false. Default is false")
.insert(
"warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.")
.insert(
"repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.")
.insert("timer",
"true",
"Whether if the timer is gpu timer or not. Possible values are false or true. "
"Default is true.")
.insert("init",
"0",
"The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 "
"for constant(1). Default is 0, random.")
.insert("flush_cache",
"true",
"To flush cache, possible values are true or false. Default is true.")
.insert(
"rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.")
.insert("metric",
"0",
"Metric with which to measure kernel performance. Set to 0 for latency, 1 for "
"tflops, or 2 for bandwidth. Default is 0, latency.")
.insert("csv_filename",
"",
"The filename of benchmark result. Default is empty (no CSV output).")
.insert("json_output",
"false",
"Whether to output results in JSON format only. Possible values are true or false. "
"Default is false");
bool result = arg_parser.parse(argc, argv);
return std::make_tuple(result, arg_parser);
}
void benchmark_single(const ck_tile::ArgParser& arg_parser)
{
std::string dtype_a = DataTypeTraits<ADataType>::name;
std::string dtype_b = DataTypeTraits<BDataType>::name;
std::string dtype_bq = DataTypeTraits<BQDataType>::name;
std::string dtype_acc = DataTypeTraits<AccDataType>::name;
std::string dtype_c = DataTypeTraits<CDataType>::name;
std::string layout_a = ALayout::name;
std::string layout_b = BLayout::name;
std::string layout_c = CLayout::name;
int M = arg_parser.get_int("m");
int N = arg_parser.get_int("n");
int K = arg_parser.get_int("k");
int group_size_k = arg_parser.get_int("group_size_k");
if(M <= 0 || N <= 0 || K <= 0)
{
throw std::invalid_argument("m, n, k must be positive integers");
}
if(group_size_k <= 0 || K % group_size_k != 0)
{
throw std::invalid_argument(
"group_size_k must be positive and k must be divisible by group_size_k");
}
BQuantGemmProblem problem{arg_parser.get_int("split_k"),
M,
N,
K,
arg_parser.get_int("stride_a"),
arg_parser.get_int("stride_b"),
arg_parser.get_int("stride_c"),
0, // stride_bq computed by profiler
group_size_k,
dtype_a,
dtype_b,
dtype_bq,
dtype_acc,
dtype_c,
layout_a,
layout_b,
layout_c};
Setting setting{arg_parser.get_int("warmup"),
arg_parser.get_int("repeat"),
arg_parser.get_bool("timer"),
arg_parser.get_int("verify"),
arg_parser.get_int("init"),
arg_parser.get_bool("log"),
arg_parser.get_str("csv_filename"),
arg_parser.get_bool("flush_cache"),
arg_parser.get_int("rotating_count"),
arg_parser.get_bool("json_output")};
BQuantGemmProfiler profiler{setting};
try
{
auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args,
const ck_tile::stream_config& stream) {
return SelectedKernel::launch(args, stream);
};
profiler.benchmark(problem, kernel_func);
profiler.select_best_instance(static_cast<Metric>(arg_parser.get_int("metric")));
}
catch(const std::exception& e)
{
std::cerr << "Benchmark failed: " << e.what() << std::endl;
}
}
int main(int argc, char* argv[])
{
try
{
auto [result, parser] = create_args(argc, argv);
if(!result)
return EXIT_FAILURE;
benchmark_single(parser);
return 0;
}
catch(const std::exception& e)
{
std::cerr << "Error: " << e.what() << "\n";
return EXIT_FAILURE;
}
}

View File

@@ -0,0 +1,74 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <string>
#include <type_traits>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/core/numeric/integer.hpp"
// DataTypeTraits for all supported types
template <typename T>
struct DataTypeTraits;
template <>
struct DataTypeTraits<float>
{
static constexpr const char* name = "fp32";
};
template <>
struct DataTypeTraits<ck_tile::half_t>
{
static constexpr const char* name = "fp16";
};
template <>
struct DataTypeTraits<ck_tile::bf16_t>
{
static constexpr const char* name = "bf16";
};
template <>
struct DataTypeTraits<ck_tile::fp8_t>
{
static constexpr const char* name = "fp8";
};
template <>
struct DataTypeTraits<ck_tile::bf8_t>
{
static constexpr const char* name = "bf8";
};
// Helper function to determine if a layout is row-major
template <typename Layout>
constexpr auto is_row_major(Layout)
{
return ck_tile::bool_constant<std::is_same_v<Layout, ck_tile::tensor_layout::gemm::RowMajor>>{};
}
// Structure to hold kernel traits for dispatcher
struct BQuantKernelTraits
{
std::string pipeline; // compv3
std::string scheduler; // intrawave
std::string epilogue; // default, cshuffle
bool pad_m;
bool pad_n;
bool pad_k;
bool b_preshuffle_quant;
BQuantKernelTraits()
: pipeline("compv3"),
scheduler("intrawave"),
epilogue("default"),
pad_m(false),
pad_n(false),
pad_k(false),
b_preshuffle_quant(false)
{
}
};

View File

@@ -0,0 +1,288 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import os
import argparse
import importlib.util
import multiprocessing
import concurrent.futures
def _import_gemm_kernel_builder():
"""Import the base GemmKernelBuilder from the gemm directory."""
current_dir = os.path.dirname(os.path.abspath(__file__))
gemm_dir = os.path.dirname(os.path.dirname(current_dir))
spec = importlib.util.spec_from_file_location(
"gemm_instance_builder",
os.path.join(gemm_dir, "gemm_instance_builder.py"),
)
gemm_builder_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(gemm_builder_module)
return gemm_builder_module.GemmKernelBuilder
GemmKernelBuilder = _import_gemm_kernel_builder()
class GemmBQuantKernelBuilder(GemmKernelBuilder):
def __init__(
self,
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json=None,
max_instances=None,
seed=None,
tier=None,
manifest_path=None,
):
super().__init__(
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json,
max_instances=max_instances,
seed=seed,
tier=tier,
manifest_path=manifest_path,
)
self.group_size_k = self.config.get("group_size_k", 128)
def _generate_all_individual(self, num_workers=None):
"""Generate individual kernel files for separate compilation with parallel processing"""
if num_workers is None:
num_workers = min(multiprocessing.cpu_count(), 8)
tile_configs = self._get_tile_configs()
trait_combos = self._generate_trait_combinations()
work_items = []
for tile_config in tile_configs:
for trait_combo in trait_combos:
work_items.append(
(
tile_config,
trait_combo,
self.kernel_name_prefix,
self.working_path,
self.gpu_target,
self.datatype,
self.layout,
self.config_json,
)
)
kernel_list = []
completed = 0
with concurrent.futures.ProcessPoolExecutor(
max_workers=num_workers
) as executor:
future_to_item = {
executor.submit(_generate_single_kernel_individual, item): item
for item in work_items
}
for future in concurrent.futures.as_completed(future_to_item):
completed += 1
if completed % 100 == 0 or completed == len(work_items):
print(
f" Progress: {completed}/{len(work_items)} kernels generated"
)
try:
result = future.result()
if result:
kernel_list.append(result)
except Exception as exc:
item = future_to_item[future]
print(f"Kernel generation failed for {item}: {exc}")
kernel_list.sort(key=lambda x: x[0])
self._generate_cmake_individual_targets(kernel_list)
print(
f"Generated {len(kernel_list)} individual kernel files in {self.working_path}"
)
def _generate_single_kernel_individual(work_item):
"""Worker function to generate a single individual kernel file"""
(
tile_config,
trait_combo,
kernel_name_prefix,
working_path,
gpu_target,
datatype,
layout,
config_json,
) = work_item
builder = GemmBQuantKernelBuilder(
kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json
)
try:
kernel_name, instance_code = builder._generate_kernel_instance(
tile_config, trait_combo
)
simplified_name = kernel_name.removeprefix(kernel_name_prefix + "_")
header_file = working_path / f"gemm_bquant_single_{simplified_name}.hpp"
with open(header_file, "w") as f:
f.write(instance_code)
return (kernel_name, trait_combo, tile_config)
except Exception as e:
print(f"Error generating individual kernel: {e}")
return None
def main():
parser = argparse.ArgumentParser(description="GEMM BQuant kernel instance builder")
parser.add_argument("--working_path", required=True, help="Working directory path")
parser.add_argument("--gpu_target", required=True, help="GPU target architecture")
parser.add_argument(
"--datatype",
required=True,
choices=["fp8", "bf8"],
help="Data type (fp8 or bf8)",
)
parser.add_argument(
"--layout",
required=True,
choices=["rcr", "rrr", "ccr", "crr"],
help="Matrix layout",
)
parser.add_argument("--config_json", help="Configuration JSON file")
parser.add_argument("--num_workers", type=int, help="Number of parallel workers")
parser.add_argument(
"--gen_all_individual",
action="store_true",
help="Generate individual kernel files",
)
parser.add_argument(
"--gen_single",
action="store_true",
help="Generate a single kernel file",
)
parser.add_argument("--kernel_name", help="Kernel name for single generation")
parser.add_argument("--tile_config", help="Tile configuration string")
parser.add_argument("--trait_combo", help="Trait combination string")
parser.add_argument(
"--list_kernels",
action="store_true",
help="List kernel configurations without generating files",
)
parser.add_argument(
"--max-instances",
type=int,
default=None,
help="Maximum number of kernel instances to select via sampling",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="RNG seed for deterministic sampling; if omitted, derived from today's date",
)
parser.add_argument(
"--tier",
default=None,
help="Sampling tier name (e.g., 'daily')",
)
parser.add_argument(
"--manifest-path",
default=None,
help="Directory to write chosen_instances manifest JSON",
)
args = parser.parse_args()
layout_parts = args.layout.lower()
assert len(layout_parts) == 3, (
f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr')"
)
assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], (
f"Invalid matrix_a layout: {layout_parts[0]} or matrix_b layout: {layout_parts[1]}"
)
assert layout_parts[2] == "r", (
f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' for row major)"
)
kernel_name_prefix = "gemm_bquant"
builder = GemmBQuantKernelBuilder(
kernel_name_prefix,
args.working_path,
args.gpu_target,
args.datatype,
args.layout,
args.config_json,
max_instances=args.max_instances,
seed=args.seed,
tier=args.tier,
manifest_path=args.manifest_path,
)
if args.list_kernels:
builder._list_kernels()
elif args.gen_single:
if not args.kernel_name or not args.tile_config or not args.trait_combo:
parser.error(
"--gen_single requires --kernel_name, --tile_config, and --trait_combo"
)
# Parse tile config
tile_parts = args.tile_config.split("_")
tile_dims = tile_parts[0].split("x")
warp_dims = tile_parts[1].split("x")
warp_tile_dims = tile_parts[2].split("x")
tile_config = {
"tile_m": int(tile_dims[0]),
"tile_n": int(tile_dims[1]),
"tile_k": int(tile_dims[2]),
"warp_m": int(warp_dims[0]),
"warp_n": int(warp_dims[1]),
"warp_k": int(warp_dims[2]),
"warp_tile_m": int(warp_tile_dims[0]),
"warp_tile_n": int(warp_tile_dims[1]),
"warp_tile_k": int(warp_tile_dims[2]),
}
# Parse trait combo:
# pipeline_epilogue_scheduler_padM_padN_padK_bPreshuffle
trait_parts = args.trait_combo.split("_")
if len(trait_parts) < 7:
parser.error(
f"--trait_combo must have 7 underscore-separated fields "
f"(e.g. 'compv3_default_intrawave_True_True_True_False'), "
f"got {len(trait_parts)} field(s): '{args.trait_combo}'"
)
trait_combo = (
trait_parts[0], # pipeline
trait_parts[1], # epilogue
trait_parts[2], # scheduler
trait_parts[3] == "True", # pad_m
trait_parts[4] == "True", # pad_n
trait_parts[5] == "True", # pad_k
trait_parts[6] == "True", # b_preshuffle_quant
)
builder._generate_kernel_instance(tile_config, trait_combo)
elif args.gen_all_individual:
builder._generate_all_individual(args.num_workers)
else:
parser.error(
"Must specify one of: --list_kernels, --gen_all_individual, or --gen_single"
)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,293 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <iostream>
#include <fstream>
#include <iomanip>
#include <functional>
#include <tuple>
#include <vector>
#include <algorithm>
#include <stdexcept>
#include "ck_tile/host/device_prop.hpp"
#include "ck_tile/host/tensor_shuffle_utils.hpp"
#include "ck_tile/ops/gemm_quant.hpp"
#include "gemm_bquant_benchmark.hpp"
class BQuantGemmProfiler
{
public:
explicit BQuantGemmProfiler(Setting setting) : setting_(setting) {}
BQuantGemmProfiler(const BQuantGemmProfiler&) = delete;
BQuantGemmProfiler& operator=(const BQuantGemmProfiler&) = delete;
void reset() { kernel_instances_.clear(); }
void benchmark(BQuantGemmProblem problem,
std::function<float(const ck_tile::QuantGemmHostArgs&,
const ck_tile::stream_config&)> kernel_func)
{
std::vector<std::function<std::tuple<std::string, float>(ck_tile::QuantGemmHostArgs&,
const ck_tile::stream_config&)>>
callables;
callables.push_back(
[kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
float time = kernel_func(args, stream);
return std::make_tuple(std::string(KERNEL_NAME), time);
});
benchmark(problem, callables);
}
void benchmark(BQuantGemmProblem problem,
std::vector<std::function<std::tuple<std::string, float>(
ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables)
{
const ALayout layout_a = ALayout{};
const BLayout layout_b = BLayout{};
const CLayout layout_c = CLayout{};
const BQLayout layout_bq = BQLayout{};
problem.stride_a_ = ck_tile::get_default_stride(
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a));
problem.stride_b_ = ck_tile::get_default_stride(
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b));
problem.stride_c_ = ck_tile::get_default_stride(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c));
// Compute BQ scale tensor dimensions: [K / group_size_k, N]
if(problem.k_ % problem.group_size_k_ != 0)
throw std::runtime_error("k_ must be divisible by group_size_k_");
const ck_tile::index_t QK_B = problem.k_ / problem.group_size_k_;
problem.stride_bq_ = ck_tile::get_default_stride(
QK_B, problem.n_, problem.stride_bq_, is_row_major(layout_bq));
// Allocate host tensors
ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
// BQ scale tensor: [QK_B, N]
ck_tile::HostTensor<BQDataType> bq_qk_n(ck_tile::host_tensor_descriptor(
QK_B, problem.n_, problem.stride_bq_, is_row_major(layout_bq)));
// Initialize tensors
if(setting_.init_method_ == 0)
{
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
ck_tile::FillUniformDistribution<BQDataType>{0.5f, 1.5f}(bq_qk_n);
}
else if(setting_.init_method_ == 1)
{
ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
}
else if(setting_.init_method_ == 2)
{
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(1)}(a_m_k);
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(1)}(b_k_n);
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
}
else
{
a_m_k.SetZero();
b_k_n.SetZero();
bq_qk_n.SetZero();
}
// Allocate device memory
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
ck_tile::DeviceMem bq_dev_buf(bq_qk_n.get_element_space_size_in_bytes());
a_m_k_dev_buf.ToDevice(a_m_k.data());
b_k_n_dev_buf.ToDevice(b_k_n.data());
// Shuffle BQ data if preshuffle is enabled
if constexpr(SelectedKernel::BPreshuffleQuant)
{
constexpr int block_bq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK;
ck_tile::HostTensor<BQDataType> bq_shuffled = ck_tile::shuffle_bq(&bq_qk_n, block_bq_k);
bq_dev_buf.ToDevice(bq_shuffled.data());
}
else
{
bq_dev_buf.ToDevice(bq_qk_n.data());
}
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
// Build QuantGemmHostArgs
ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
b_k_n_dev_buf.GetDeviceBuffer(),
c_m_n_dev_buf.GetDeviceBuffer(),
nullptr, // aq_ptr not used for BQuant
bq_dev_buf.GetDeviceBuffer(),
problem.split_k_,
problem.m_,
problem.n_,
problem.k_,
0, // QK_A not used for BQuant
QK_B,
problem.stride_a_,
problem.stride_b_,
problem.stride_c_,
0, // stride_AQ not used for BQuant
problem.stride_bq_);
// Host reference for verification
ck_tile::HostTensor<CDataType> c_m_n_host_result(ck_tile::host_tensor_descriptor(
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
if(setting_.verify_)
{
bquant_gemm_host_reference(setting_.verify_, a_m_k, b_k_n, bq_qk_n, c_m_n_host_result);
}
// Run kernel(s)
for(auto& callable : callables)
{
auto kernel_run_result = callable(gemm_args,
ck_tile::stream_config{nullptr,
true,
setting_.log_,
setting_.n_warmup_,
setting_.n_repeat_,
setting_.is_gpu_timer_,
setting_.flush_cache_,
setting_.rotating_count_});
process_result(problem,
QK_B,
c_m_n_dev_buf,
c_m_n_host_result,
c_m_n_dev_result,
kernel_run_result);
}
}
void process_result(const BQuantGemmProblem& problem,
ck_tile::index_t QK_B,
ck_tile::DeviceMem& c_m_n_dev_buf,
ck_tile::HostTensor<CDataType>& c_m_n_host_result,
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
const std::tuple<std::string, float>& kernel_run_result)
{
auto [name, avg_time] = kernel_run_result;
KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}};
// Compute performance metrics
std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_;
std::size_t num_byte = sizeof(ADataType) * problem.m_ * problem.k_ +
sizeof(BDataType) * problem.n_ * problem.k_ +
sizeof(BQDataType) * problem.n_ * QK_B +
sizeof(CDataType) * problem.m_ * problem.n_;
kernel_instance.perf_result_.latency_ = avg_time;
kernel_instance.perf_result_.tflops_ = static_cast<float>(flop) / 1.E9 / avg_time;
kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time;
if(setting_.log_ > 0 && !setting_.json_output_)
{
std::cout << kernel_instance << std::endl;
}
// Verify result
bool verified_correct = true;
if(setting_.verify_)
{
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
verified_correct = compare_bquant(
name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result);
}
if(verified_correct)
{
kernel_instances_.emplace_back(kernel_instance);
}
else
{
std::cout << "Verification failed, skip kernel: " << name << std::endl;
}
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
}
KernelInstance select_best_instance(Metric metric)
{
if(kernel_instances_.empty())
throw std::runtime_error("Empty instances");
auto kernel_instance = *std::max_element(kernel_instances_.begin(),
kernel_instances_.end(),
[metric](const auto& a, const auto& b) {
return PerformanceResult::compare(
b.perf_result_, a.perf_result_, metric);
});
if(setting_.json_output_)
{
std::cout << kernel_instance << std::endl;
}
else
{
std::cout << "**********************************" << std::endl;
std::cout << "According to given metrics: " << get_metric_name(metric) << "\n"
<< "Current kernel performance is: " << kernel_instance << std::endl;
std::cout << "**********************************" << std::endl;
}
if(!setting_.csv_filename_.empty())
{
std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app);
if(!file.is_open())
{
std::cerr << "Warning: Failed to open CSV file for writing." << std::endl;
}
else
{
if(file.tellp() == 0)
{
file << "rocm_version,device_name,"
<< "split_k,m,n,k,stride_a,stride_b,stride_c,stride_bq,group_size_k,"
<< "dtype_a,dtype_b,dtype_bq,dtype_acc,dtype_c,"
<< "layout_a,layout_b,layout_c," << "name,"
<< "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n";
}
const auto& prob = kernel_instance.problem_;
const auto& perf = kernel_instance.perf_result_;
file << get_rocm_version() << "," << ck_tile::get_device_name() << ","
<< prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << ","
<< prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << ","
<< prob.stride_bq_ << "," << prob.group_size_k_ << "," << prob.dtype_a_ << ","
<< prob.dtype_b_ << "," << prob.dtype_bq_ << "," << prob.dtype_acc_ << ","
<< prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_ << ","
<< prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed
<< std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << ","
<< perf.bandwidth_ << "," << get_metric_name(metric) << "\n";
}
}
return kernel_instance;
}
private:
Setting setting_;
std::vector<KernelInstance> kernel_instances_;
};

View File

@@ -0,0 +1,191 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""Unit tests for GemmBQuantKernelBuilder (gemm_bquant operator)."""
import json
import os
import sys
import tempfile
import unittest
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
from gemm_bquant_instance_builder import GemmBQuantKernelBuilder # noqa: E402
_MINIMAL_CONFIG = {
"tile_config": {
"tile_m": {"values": [64]},
"tile_n": {"values": [64]},
"tile_k": {"values": [128]},
"warp_m": {"values": [2]},
"warp_n": {"values": [2]},
"warp_k": {"values": [1]},
"warp_tile_m": {"values": [16]},
"warp_tile_n": {"values": [16]},
"warp_tile_k": {"values": [64]},
},
"trait_config": {
"pipeline": {"values": ["compv3"]},
"scheduler": {"values": ["intrawave"]},
"epilogue": {"values": ["default"]},
"pad_m": {"values": [False]},
"pad_n": {"values": [False]},
"pad_k": {"values": [False]},
"b_preshuffle_quant": {"values": [False, True]},
},
"k_block_per_cu": 1,
"group_size_k": 128,
}
def _make_builder(tmpdir, config=None, **kwargs):
cfg = config if config is not None else _MINIMAL_CONFIG
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(cfg, f)
return GemmBQuantKernelBuilder(
kernel_name_prefix="gemm_bquant",
working_path=tmpdir,
gpu_target=kwargs.get("gpu_target", "gfx942"),
datatype=kwargs.get("datatype", "fp8"),
layout=kwargs.get("layout", "rcr"),
config_json=cfg_path,
)
class TestGemmBQuantBuilderInit(unittest.TestCase):
def test_group_size_k_from_config(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertEqual(builder.group_size_k, 128)
def test_group_size_k_custom(self):
cfg = json.loads(json.dumps(_MINIMAL_CONFIG))
cfg["group_size_k"] = 32
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, config=cfg)
self.assertEqual(builder.group_size_k, 32)
def test_kernel_name_prefix_stored(self):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir)
self.assertEqual(builder.kernel_name_prefix, "gemm_bquant")
def test_working_path_created(self):
with tempfile.TemporaryDirectory() as tmpdir:
sub = os.path.join(tmpdir, "workdir")
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(_MINIMAL_CONFIG, f)
GemmBQuantKernelBuilder("gemm_bquant", sub, "gfx942", "fp8", "rcr", cfg_path)
self.assertTrue(os.path.isdir(sub))
class TestGemmBQuantTraitCombinations(unittest.TestCase):
def setUp(self):
self._tmpdir = tempfile.TemporaryDirectory()
self.builder = _make_builder(self._tmpdir.name)
def tearDown(self):
self._tmpdir.cleanup()
def test_trait_combinations_non_empty(self):
combos = self.builder._generate_trait_combinations()
self.assertGreater(len(combos), 0)
def test_trait_combo_is_7_tuple(self):
"""BQuant trait tuple: (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, b_preshuffle_quant)."""
combos = self.builder._generate_trait_combinations()
for combo in combos:
self.assertEqual(len(combo), 7, f"Expected 7-tuple, got {len(combo)}: {combo}")
def test_pipeline_is_compv3(self):
combos = self.builder._generate_trait_combinations()
for combo in combos:
self.assertEqual(combo[0], "compv3")
def test_preshuffle_quant_values(self):
combos = self.builder._generate_trait_combinations()
preshuffle_vals = {c[6] for c in combos}
self.assertIn(True, preshuffle_vals)
self.assertIn(False, preshuffle_vals)
def test_scheduler_intrawave(self):
combos = self.builder._generate_trait_combinations()
schedulers = {c[2] for c in combos}
self.assertIn("intrawave", schedulers)
class TestGemmBQuantTileConfigs(unittest.TestCase):
def setUp(self):
self._tmpdir = tempfile.TemporaryDirectory()
self.builder = _make_builder(self._tmpdir.name)
def tearDown(self):
self._tmpdir.cleanup()
def test_tile_configs_non_empty(self):
configs = self.builder._get_tile_configs()
self.assertGreater(len(configs), 0)
def test_tile_config_has_required_keys(self):
required = {
"tile_m", "tile_n", "tile_k",
"warp_m", "warp_n", "warp_k",
"warp_tile_m", "warp_tile_n", "warp_tile_k",
}
for cfg in self.builder._get_tile_configs():
self.assertTrue(required.issubset(cfg.keys()))
def test_tile_m_matches_config(self):
configs = self.builder._get_tile_configs()
tile_m_vals = {c["tile_m"] for c in configs}
self.assertIn(64, tile_m_vals)
class TestGemmBQuantLayoutVariants(unittest.TestCase):
LAYOUTS = ["rcr", "rrr", "ccr", "crr"]
def test_all_layouts_construct(self):
for layout in self.LAYOUTS:
with self.subTest(layout=layout):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, layout=layout)
self.assertEqual(builder.layout, layout)
class TestGemmBQuantDataTypes(unittest.TestCase):
DTYPES = ["fp8", "bf8"]
def test_all_dtypes_construct(self):
for dtype in self.DTYPES:
with self.subTest(dtype=dtype):
with tempfile.TemporaryDirectory() as tmpdir:
builder = _make_builder(tmpdir, datatype=dtype)
self.assertEqual(builder.datatype, dtype)
class TestGemmBQuantMaxInstances(unittest.TestCase):
def test_max_instances_none_returns_all(self):
with tempfile.TemporaryDirectory() as tmpdir:
cfg_path = os.path.join(tmpdir, "config.json")
with open(cfg_path, "w") as f:
json.dump(_MINIMAL_CONFIG, f)
builder = GemmBQuantKernelBuilder(
"gemm_bquant", tmpdir, "gfx942", "fp8", "rcr",
config_json=cfg_path, max_instances=None,
)
kernel_list = [
{"tile_config": {"tile_m": 64, "tile_n": 64, "tile_k": 128,
"warp_m": 2, "warp_n": 2, "warp_k": 1,
"warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 64},
"trait_combo": ("compv3", "default", "intrawave", False, False, False, False)}
]
result = builder._apply_sampling(kernel_list)
self.assertEqual(len(result), 1)
if __name__ == "__main__":
unittest.main()

View File

@@ -66,6 +66,7 @@ class GemmKernelBuilder:
if config_json and os.path.exists(config_json):
with open(config_json, "r") as f:
self.config = json.load(f)
self.group_size_k = self.config.get("group_size_k", 128)
def _apply_sampling(self, kernel_list):
"""Apply RFC Sobol+LHS+maximin sampling. Returns sampled subset."""
@@ -319,10 +320,12 @@ class GemmKernelBuilder:
pipelines = ["preshufflev2"]
elif self.kernel_name_prefix == "mx_gemm":
pipelines = ["comp_async"]
elif self.kernel_name_prefix in ["grouped_gemm_rowcolquant", "grouped_gemm_tensorquant", "gemm_rowcolquant", "gemm_tensor_quant"]:
elif self.kernel_name_prefix in ["grouped_gemm_rowcolquant", "grouped_gemm_tensorquant", "gemm_rowcolquant", "gemm_tensor_quant", "gemm_aquant", "gemm_bquant"]:
pipelines = ["compv3"]
elif self.kernel_name_prefix in ["gemm_universal", "gemm_multi_d", "gemm_multi_abd", "grouped_gemm", "batched_contraction", "batched_gemm"]:
pipelines = ["compv4"]
elif self.kernel_name_prefix == "gemm_abquant":
pipelines = ["compv3"]
configs = []
for tile_m in tile_m_values:
@@ -426,6 +429,7 @@ class GemmKernelBuilder:
layout,
self.gpu_target,
self.kernel_name_prefix,
self.config.get("group_size_k", 128),
)
def _generate_trait_combinations(self):
@@ -439,7 +443,15 @@ class GemmKernelBuilder:
pad_m_values = trait_config.get("pad_m").get("values")
pad_n_values = trait_config.get("pad_n").get("values")
pad_k_values = trait_config.get("pad_k").get("values")
if self.kernel_name_prefix in ["gemm_rowcolquant", "batched_gemm"]:
if self.kernel_name_prefix == "gemm_aquant":
persistent_values = trait_config.get(
"a_preshuffle_quant", {}
).get("values", [False])
elif self.kernel_name_prefix == "gemm_bquant":
persistent_values = trait_config.get(
"b_preshuffle_quant", {}
).get("values", [False])
elif self.kernel_name_prefix in ["gemm_rowcolquant", "batched_gemm"]:
persistent_values = [
False
] # Force disable persistent where it is unsupported or not part of the trait key
@@ -462,8 +474,14 @@ class GemmKernelBuilder:
combinations = []
for combo in all_combinations:
pipeline, epilogue, scheduler = combo[:3]
persistent_or_preshuffle_quant = combo[6] if len(combo) > 6 else False
if is_trait_combination_valid(
pipeline, epilogue, scheduler, self.kernel_name_prefix
pipeline,
epilogue,
scheduler,
persistent_or_preshuffle_quant,
self.kernel_name_prefix,
self.layout,
):
combinations.append(combo)
else:
@@ -484,7 +502,7 @@ class GemmKernelBuilder:
pad_m,
pad_n,
pad_k,
persistent,
persistent_or_preshuffle_quant,
) = self._normalize_trait_combo(trait_combo)
kernel_name = self._format_kernel_name(trait_combo, tile_config)
@@ -523,6 +541,22 @@ class GemmKernelBuilder:
"comp_async": "ck_tile::GemmPipelineAgBgCrCompAsync",
}
base_pipeline_map = {}
elif self.kernel_name_prefix == "gemm_aquant":
pipeline_impl_map = {
"mem": "ck_tile::AQuantGemmPipelineAgBgCrMem",
"compv3": "ck_tile::AQuantGemmPipelineAgBgCrCompV3",
}
base_pipeline_map = {
"mem": "ck_tile::BaseGemmPipelineAgBgCrMem",
"compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3",
}
elif self.kernel_name_prefix == "gemm_bquant":
pipeline_impl_map = {
"compv3": "ck_tile::BQuantGemmPipelineAgBgCrCompV3",
}
base_pipeline_map = {
"compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3",
}
scheduler_type_map = {
"intrawave": "ck_tile::GemmPipelineScheduler::Intrawave",
@@ -543,7 +577,7 @@ class GemmKernelBuilder:
pipeline,
epilogue,
k_block_per_cu,
persistent,
persistent_or_preshuffle_quant,
)
# Write into a file
@@ -581,6 +615,13 @@ class GemmKernelBuilder:
instance_code += """#include <vector>
#include <hip/hip_runtime.h>
#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp"
"""
elif self.kernel_name_prefix == "mx_gemm":
instance_code += """#include "ck_tile/ops/gemm_mx.hpp"
"""
elif self.kernel_name_prefix in ["gemm_aquant", "gemm_bquant", "gemm_abquant"]:
instance_code += """#include "ck_tile/ops/gemm_quant.hpp"
#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp"
"""
return instance_code
@@ -602,6 +643,9 @@ class GemmKernelBuilder:
"grouped_gemm",
"mx_gemm",
"batched_gemm",
"gemm_aquant",
"gemm_bquant",
"gemm_abquant",
]:
a_layout, b_layout, c_layout = get_abc_layouts(self.layout)
@@ -611,11 +655,21 @@ using BDataType = {get_dtype_string(self.datatype)};
using AccDataType = {acc_type};
using CDataType = {get_dtype_string(c_type)};"""
if self.kernel_name_prefix == "gemm_aquant":
# AQ scale factors are always fp32: the CK AQuant kernel expects float scales
# regardless of the A/B element type. Changing this requires a matching kernel update.
instance_code += """
using AQDataType = float;"""
if self.kernel_name_prefix == "mx_gemm":
instance_code += """
using ScaleType = ck_tile::e8m0_t;
using MxGemmHostArgs = ck_tile::MxGemmHostArgs<1, 1, 0>;"""
if self.kernel_name_prefix == "gemm_bquant":
instance_code += """
using BQDataType = float;"""
if self.kernel_name_prefix == "gemm_multi_d":
instance_code += f"""
using D0DataType = {get_dtype_string(self.datatype)};
@@ -627,6 +681,17 @@ using ALayout = {a_layout};
using BLayout = {b_layout};
using CLayout = {c_layout};
"""
if self.kernel_name_prefix == "gemm_aquant":
instance_code += """using AQLayout = ck_tile::tensor_layout::gemm::RowMajor;
"""
elif self.kernel_name_prefix == "gemm_bquant":
# BQ scale tensor has shape [K/group_size_k, N], so its leading dimension
# matches B's K dimension. Setting BQLayout = BLayout is only valid when B
# is column-major (layout[1]=='c'); BPreshuffleQuant with row-major B is
# blocked in is_trait_combination_valid to enforce this.
instance_code += """using BQLayout = BLayout;
"""
if self.kernel_name_prefix == "gemm_multi_d":
instance_code += f"""
using D0Layout = {ds_layout[0]};
@@ -681,6 +746,20 @@ struct SelectedKernel {{
static constexpr bool TransposeC = false;
static constexpr bool DoubleSmemBuffer = {"true" if pipeline in ["compv4", "preshufflev2", "comp_async"] else "false"};"""
if self.kernel_name_prefix == "gemm_aquant":
instance_code += f"""
static constexpr bool APreshuffleQuant = {"true" if persistent in [True, "true"] else "false"};
static constexpr bool BPreshuffleQuant = false;
static constexpr bool PreshuffleB = false;
static constexpr ck_tile::index_t GroupSizeK = {self.config.get("group_size_k", 128)};"""
elif self.kernel_name_prefix == "gemm_bquant":
instance_code += f"""
static constexpr bool APreshuffleQuant = false;
static constexpr bool BPreshuffleQuant = {"true" if persistent in [True, "true"] else "false"};
static constexpr bool PreshuffleB = false;
static constexpr ck_tile::index_t GroupSizeK = {self.config.get("group_size_k", 128)};"""
if self.kernel_name_prefix in [
"gemm_universal",
"gemm_preshuffle",
@@ -700,11 +779,26 @@ struct SelectedKernel {{
else:
instance_code += """
static constexpr bool Preshuffle = false;"""
if self.kernel_name_prefix == "gemm_aquant":
instance_code += f"""
static constexpr bool APreshuffleQuant = {"true" if persistent_or_preshuffle_quant in [True, "true"] else "false"};
static constexpr bool BPreshuffleQuant = false;
static constexpr bool PreshuffleB = false;
static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k};"""
elif self.kernel_name_prefix == "gemm_bquant":
instance_code += f"""
static constexpr bool APreshuffleQuant = false;
static constexpr bool BPreshuffleQuant = {"true" if persistent_or_preshuffle_quant in [True, "true"] else "false"};
static constexpr bool PreshuffleB = false;
static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k};"""
return instance_code
def populate_initialization(self, base_pipeline_map, pipeline):
# Tile Shape
if self.kernel_name_prefix in ["gemm_multi_d", "batched_gemm"]:
if self.kernel_name_prefix in ["gemm_multi_d", "batched_gemm", "gemm_aquant", "gemm_bquant", "gemm_abquant"]:
instance_code = """
// Tile shape
@@ -736,7 +830,39 @@ struct SelectedKernel {{
using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner<TileShape, 8, 4>;"""
# Traits
if self.kernel_name_prefix == "gemm_multi_d":
if self.kernel_name_prefix == "gemm_aquant":
instance_code += """
// Quantization group size
using QuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, GroupSizeK>>;
// Traits
using Traits = ck_tile::TileGemmQuantTraits<
kPadM, kPadN, kPadK,
APreshuffleQuant,
BPreshuffleQuant,
PreshuffleB,
ALayout, BLayout, CLayout,
ck_tile::QuantType::AQuantGrouped,
AQLayout>;"""
elif self.kernel_name_prefix == "gemm_bquant":
instance_code += """
// Quantization group size
using QuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, GroupSizeK>>;
// Traits
using Traits = ck_tile::TileGemmQuantTraits<
kPadM, kPadN, kPadK,
APreshuffleQuant,
BPreshuffleQuant,
PreshuffleB,
ALayout, BLayout, CLayout,
ck_tile::QuantType::BQuantGrouped,
BQLayout>;"""
elif self.kernel_name_prefix == "gemm_multi_d":
instance_code += """
// Traits
@@ -748,7 +874,19 @@ struct SelectedKernel {{
using Traits = ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, CLayout, NumWaveGroups>;"""
# Pipeline problem
if self.kernel_name_prefix in ["gemm_preshuffle", "gemm_multi_d"]:
if self.kernel_name_prefix in ["gemm_aquant", "gemm_bquant"]:
instance_code += """
// Pipeline problem (base, for hot loop detection)
using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<
ADataType,
BDataType,
AccDataType,
TileShape,
Traits,
BDataType>;"""
elif self.kernel_name_prefix in ["gemm_preshuffle", "gemm_multi_d"]:
instance_code += """
// Pipeline problem
@@ -766,7 +904,7 @@ struct SelectedKernel {{
// Base pipeline for hot loop detection
using BaseGemmPipeline = {base_pipeline_map.get(pipeline, "ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2")}<GemmPipelineProblem>;"""
elif self.kernel_name_prefix == "gemm_multi_d":
elif self.kernel_name_prefix in ["gemm_multi_d", "gemm_aquant", "gemm_bquant"]:
instance_code += f"""
// Base pipeline for hot loop detection
@@ -782,8 +920,28 @@ struct SelectedKernel {{
pipeline,
epilogue,
k_block_per_cu,
persistent,
persistent_or_preshuffle_quant,
):
if self.kernel_name_prefix == "gemm_aquant":
return self._populate_launch_aquant(
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
)
if self.kernel_name_prefix == "gemm_bquant":
return self._populate_launch_bquant(
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
)
# Function Signature
if self.kernel_name_prefix == "gemm_multi_d":
instance_code = """
@@ -959,7 +1117,7 @@ struct SelectedKernel {{
}}
// Get grid and block sizes
const dim3 grids = {"GemmKernel::MaxOccupancyGridSize(stream)" if persistent in [True, "true"] else "GemmKernel::GridSize(args.M, args.N, args.k_batch)"};
const dim3 grids = {"GemmKernel::MaxOccupancyGridSize(stream)" if persistent_or_preshuffle_quant in [True, "true"] else "GemmKernel::GridSize(args.M, args.N, args.k_batch)"};
const dim3 blocks = GemmKernel::BlockSize();
if(stream.log_level_ > 0) {{
@@ -1032,7 +1190,7 @@ struct SelectedKernel {{
}}
// Get grid and block sizes
const dim3 grids = {"Kernel::MaxOccupancyGridSize(stream)" if persistent in [True, "true"] else "dim3(kargs.empty() ? 0 : kargs.back().block_end, 1, 1)"};
const dim3 grids = {"Kernel::MaxOccupancyGridSize(stream)" if persistent_or_preshuffle_quant in [True, "true"] else "dim3(kargs.empty() ? 0 : kargs.back().block_end, 1, 1)"};
const dim3 blocks = Kernel::BlockSize();
HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
@@ -1093,6 +1251,187 @@ struct SelectedKernel {{
return ave_time;
}}
}};
"""
return instance_code
def _populate_launch_aquant(
self,
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
):
"""Generate the complete launch function for AQuant kernels."""
instance_code = """
// Launch function
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
// Hot loop detection
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
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);"""
instance_code += f"""
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
// Full pipeline problem with hot loop and tail number
using PipelineProblem = ck_tile::GemmAQuantPipelineProblem<
ADataType,
AQDataType,
BDataType,
AccDataType,
TileShape,
Traits,
QuantGroupSize,
TransposeC,
BDataType,
{scheduler_type_map.get(scheduler)},
has_hot_loop_v,
tail_number_v>;
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
instance_code += """
// Epilogue"""
if epilogue == "cshuffle":
instance_code += self.populate_cshuffle_gemm_aquant()
else:
instance_code += self.populate_default_gemm_aquant()
instance_code += f"""
// Kernel type
using Kernel = ck_tile::QuantGemmKernel<
TilePartitioner, GemmPipeline, GemmEpilogue,
ck_tile::QuantType::AQuantGrouped>;
// Kernel arguments
auto kargs = Kernel::MakeKernelArgs(args);
if (!Kernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
}}
// Get grid and block sizes
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
const dim3 blocks = Kernel::BlockSize();
if(stream.log_level_ > 0) {{
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
<< std::endl;
}}
// Launch kernel
constexpr int kBlockPerCu = {k_block_per_cu};
float ave_time = ck_tile::launch_kernel(
stream,
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}};
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
}}
}};
"""
return instance_code
def _populate_launch_bquant(
self,
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
):
"""Generate the complete launch function for BQuant kernels."""
instance_code = """
// Launch function
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
// Hot loop detection
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
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);"""
instance_code += f"""
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
// Full pipeline problem with hot loop and tail number
using PipelineProblem = ck_tile::GemmBQuantPipelineProblem<
ADataType,
BDataType,
BQDataType,
AccDataType,
TileShape,
Traits,
QuantGroupSize,
ADataType,
{scheduler_type_map.get(scheduler)},
has_hot_loop_v,
tail_number_v>;
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
instance_code += """
// Epilogue"""
if epilogue == "cshuffle":
instance_code += self.populate_cshuffle_gemm_bquant()
else:
instance_code += self.populate_default_gemm_bquant()
instance_code += f"""
// Kernel type
using Kernel = ck_tile::QuantGemmKernel<
TilePartitioner, GemmPipeline, GemmEpilogue,
ck_tile::QuantType::BQuantGrouped>;
// Kernel arguments
auto kargs = Kernel::MakeKernelArgs(args);
if (!Kernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
}}
// Get grid and block sizes
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
const dim3 blocks = Kernel::BlockSize();
if(stream.log_level_ > 0) {{
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
<< std::endl;
}}
// Launch kernel
constexpr int kBlockPerCu = {k_block_per_cu};
float ave_time = ck_tile::launch_kernel(
stream,
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}};
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
}}
}};
"""
return instance_code
@@ -1122,6 +1461,8 @@ struct SelectedKernel {{
instance_code += self.populate_default_gemm_preshuffle()
elif self.kernel_name_prefix == "mx_gemm":
raise ValueError("MX GEMM currently supports only cshuffle epilogue")
elif self.kernel_name_prefix == "gemm_aquant":
instance_code += self.populate_default_gemm_aquant()
return instance_code
@@ -1322,6 +1663,278 @@ struct SelectedKernel {{
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
return instance_code
def _populate_launch_aquant(
self,
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
):
instance_code = """
// Launch function
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
// Hot loop detection
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
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);"""
instance_code += f"""
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
// Full pipeline problem with hot loop and tail number
using PipelineProblem = ck_tile::GemmAQuantPipelineProblem<
ADataType,
AQDataType,
BDataType,
AccDataType,
TileShape,
Traits,
QuantGroupSize,
TransposeC,
BDataType,
{scheduler_type_map.get(scheduler)},
has_hot_loop_v,
tail_number_v>;
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
instance_code += """
// Epilogue"""
if epilogue == "cshuffle":
instance_code += self.populate_cshuffle_gemm_aquant()
else:
instance_code += self.populate_default_gemm_aquant()
instance_code += f"""
// Kernel type
using Kernel = ck_tile::QuantGemmKernel<
TilePartitioner, GemmPipeline, GemmEpilogue,
ck_tile::QuantType::AQuantGrouped>;
// Kernel arguments
auto kargs = Kernel::MakeKernelArgs(args);
if (!Kernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
}}
// Get grid and block sizes
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
const dim3 blocks = Kernel::BlockSize();
if(stream.log_level_ > 0) {{
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
<< std::endl;
}}
// Launch kernel
constexpr int kBlockPerCu = {k_block_per_cu};
float ave_time = ck_tile::launch_kernel(
stream,
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}};
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
}}
}};
"""
return instance_code
def populate_default_gemm_aquant(self):
instance_code = """
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<
ADataType,
BDataType,
ck_tile::tuple<>, // DsDataType
AccDataType,
CDataType,
ck_tile::tuple<>, // DsLayout
CLayout,
ck_tile::element_wise::PassThrough,
TileM, // kM_
TileN, // kN_
kPadM,
kPadN,
WarpTileM, // kMPerXdl_
WarpTileN, // kNPerXdl_
WarpTileK, // kKPerXdl_
TransposeC>; // isCTransposed_
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
return instance_code
def populate_cshuffle_gemm_aquant(self):
instance_code = """
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<
ADataType,
BDataType,
ck_tile::tuple<>, // DsDataType
AccDataType,
CDataType,
ck_tile::tuple<>, // DsLayout
CLayout,
ck_tile::element_wise::PassThrough,
TileM, // kM_
TileN, // kN_
WarpPerBlock_M, // MWave_
WarpPerBlock_N, // NWave_
WarpTileM, // kMPerXdl_
WarpTileN, // kNPerXdl_
WarpTileK, // kKPerXdl_
TransposeC>; // isCTransposed_
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;"""
return instance_code
def _populate_launch_bquant(
self,
scheduler_type_map,
scheduler,
pipeline_impl_map,
pipeline,
epilogue,
k_block_per_cu,
):
instance_code = """
// Launch function
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
// Hot loop detection
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
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);"""
instance_code += f"""
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
// Full pipeline problem with hot loop and tail number
using PipelineProblem = ck_tile::GemmBQuantPipelineProblem<
ADataType,
BDataType,
BQDataType,
AccDataType,
TileShape,
Traits,
QuantGroupSize,
TransposeC,
BDataType,
{scheduler_type_map.get(scheduler)},
has_hot_loop_v,
tail_number_v>;
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
instance_code += """
// Epilogue"""
if epilogue == "cshuffle":
instance_code += self.populate_cshuffle_gemm_bquant()
else:
instance_code += self.populate_default_gemm_bquant()
instance_code += f"""
// Kernel type
using Kernel = ck_tile::QuantGemmKernel<
TilePartitioner, GemmPipeline, GemmEpilogue,
ck_tile::QuantType::BQuantGrouped>;
// Kernel arguments
auto kargs = Kernel::MakeKernelArgs(args);
if (!Kernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
}}
// Get grid and block sizes
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
const dim3 blocks = Kernel::BlockSize();
if(stream.log_level_ > 0) {{
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
<< std::endl;
}}
// Launch kernel
constexpr int kBlockPerCu = {k_block_per_cu};
float ave_time = ck_tile::launch_kernel(
stream,
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}};
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
}}
}};
"""
return instance_code
def populate_default_gemm_bquant(self):
instance_code = """
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<
ADataType,
BDataType,
ck_tile::tuple<>, // DsDataType
AccDataType,
CDataType,
ck_tile::tuple<>, // DsLayout
CLayout,
ck_tile::element_wise::PassThrough,
TileM, // kM_
TileN, // kN_
kPadM,
kPadN,
WarpTileM, // kMPerXdl_
WarpTileN, // kNPerXdl_
WarpTileK, // kKPerXdl_
TransposeC>; // isCTransposed_
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
return instance_code
def populate_cshuffle_gemm_bquant(self):
instance_code = """
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<
ADataType,
BDataType,
ck_tile::tuple<>, // DsDataType
AccDataType,
CDataType,
ck_tile::tuple<>, // DsLayout
CLayout,
ck_tile::element_wise::PassThrough,
TileM, // kM_
TileN, // kN_
WarpPerBlock_M, // MWave_
WarpPerBlock_N, // NWave_
WarpTileM, // kMPerXdl_
WarpTileN, // kNPerXdl_
WarpTileK, // kKPerXdl_
TransposeC>; // isCTransposed_
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;"""
return instance_code
def _generate_cmake_individual_targets(self, kernel_list):
"""Generate CMake include file that creates individual targets"""
cmake_code = f"""# Generated CMake file for individual {self.kernel_name_prefix} targets

View File

@@ -10,6 +10,9 @@ GEMM_PRESHUFFLE_PIPELINES = ["preshufflev2"]
GEMM_ROWCOLQUANT_PIPELINES = ["compv3"]
GEMM_MX_PIPELINES = ["comp_async"]
GEMM_BQUANT_PIPELINES = ["compv3"]
GEMM_ABQUANT_PIPELINES = ["compv3"]
LAYOUT_MAP = {
"r": "ck_tile::tensor_layout::gemm::RowMajor",
@@ -231,6 +234,24 @@ TRAIT_UNSUPPORTED_COMBINATIONS = {
("comp_async", "cshuffle", "interwave"),
}
AQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
("mem", "default", "interwave"),
("mem", "cshuffle", "interwave"),
("compv3", "default", "interwave"),
("compv3", "cshuffle", "interwave"),
}
BQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
("compv3", "default", "interwave"),
("compv3", "cshuffle", "interwave"),
}
ABQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
("compv3", "default", "interwave"),
("compv3", "cshuffle", "interwave"),
}
def element_size(data_type: str) -> float:
"""Calculate the size (in bytes) of a single element for given data type."""
data_type = data_type.lower()
@@ -240,10 +261,32 @@ def element_size(data_type: str) -> float:
def is_trait_combination_valid(
pipeline: str, epilogue: str, scheduler: str, kernel_name_prefix: str = ""
pipeline: str,
epilogue: str,
scheduler: str,
persistent_or_preshuffle_quant=None,
kernel_name_prefix: str = "",
layout: str = "",
) -> bool:
"""Check if a trait combination is valid."""
if (
if kernel_name_prefix == "gemm_aquant":
if (pipeline, epilogue, scheduler) in AQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
return False
# mem pipeline does not support preshuffle
if pipeline == "mem" and persistent_or_preshuffle_quant is True:
return False
return True
elif kernel_name_prefix == "gemm_bquant":
if (pipeline, epilogue, scheduler) in BQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
return False
# bquant only supports compv3 + intrawave
if pipeline != "compv3" or scheduler != "intrawave":
return False
# BPreshuffleQuant requires ColumnMajor BLayout (second char of layout is 'c')
if persistent_or_preshuffle_quant is True and len(layout) >= 2 and layout[1] != "c":
return False
return True
elif (
kernel_name_prefix == "gemm_rowcolquant"
or kernel_name_prefix == "grouped_gemm_rowcolquant"
or kernel_name_prefix == "grouped_gemm_tensorquant"
@@ -252,6 +295,16 @@ def is_trait_combination_valid(
if pipeline != "compv3" or scheduler != "intrawave" or epilogue != "cshuffle":
return False
return True
elif kernel_name_prefix == "gemm_abquant":
if (pipeline, epilogue, scheduler) in ABQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
return False
# abquant only supports compv3 + intrawave
if pipeline != "compv3" or scheduler != "intrawave":
return False
# BPreshuffleQuant requires ColumnMajor BLayout (second char of layout is 'c')
if persistent_or_preshuffle_quant is True and len(layout) >= 2 and layout[1] != "c":
return False
return True
else:
return (pipeline, epilogue, scheduler) not in TRAIT_UNSUPPORTED_COMBINATIONS
@@ -493,6 +546,7 @@ def is_tile_config_valid(
layout: str,
gpu_target: str,
kernel_name_prefix: str = "",
group_size_k: int = 128,
) -> bool:
"""
Comprehensive tile configuration validation.
@@ -691,6 +745,75 @@ def is_tile_config_valid(
)
return False
if kernel_name_prefix == "gemm_aquant":
aquant_valid, aquant_valid_error = validate_gemm_aquant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
group_size_k,
)
if not aquant_valid:
logging.debug(f"GEMM AQuant validation failed: {aquant_valid_error}")
return False
elif kernel_name_prefix == "gemm_bquant":
bquant_valid, bquant_valid_error = validate_gemm_bquant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
group_size_k,
)
if not bquant_valid:
logging.debug(f"GEMM BQuant validation failed: {bquant_valid_error}")
return False
elif kernel_name_prefix == "gemm_abquant":
abquant_valid, abquant_valid_error = validate_gemm_abquant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
group_size_k,
)
if not abquant_valid:
logging.debug(f"GEMM ABQuant validation failed: {abquant_valid_error}")
return False
return True
@@ -1259,6 +1382,47 @@ def validate_m0_m1_m2_configuration(
return False, f"Error in M0/M1/M2 validation: {str(e)}"
def _validate_fp8_mfma_warp_tile_k(
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
gpu_target: str,
op_label: str = "",
) -> Tuple[bool, str]:
"""Validate MFMA warp-tile constraints shared by all fp8/bf8 quantized GEMM ops.
Checks:
- warp_tile_m == warp_tile_n (square MFMA requirement)
- warp_tile_k matches the ISA-mandated K-block for the given warp_tile_m and gpu_target
"""
suffix = f" ({op_label})" if op_label else ""
if warp_tile_m != warp_tile_n:
return False, (
f"warp_tile_m({warp_tile_m}) must equal warp_tile_n({warp_tile_n})"
f" — MFMA requires a square warp tile{suffix}"
)
if a_datatype in ["fp8", "bf8"]:
# MFMA instruction shapes for fp8/bf8 (from CDNA ISA reference manual):
# gfx90a/gfx942: MFMA_F32_16x16x128_F8 (warp_tile_m=16) → warp_tile_k=64
# MFMA_F32_32x32x64_F8 (warp_tile_m=32) → warp_tile_k=32
# gfx950 doubles the K-block:
# MFMA_F32_16x16x256_F8 (warp_tile_m=16) → warp_tile_k=128
# MFMA_F32_32x32x128_F8 (warp_tile_m=32) → warp_tile_k=64
if gpu_target == "gfx950":
expected_k = 64 if warp_tile_m == 32 else 128
else:
expected_k = 32 if warp_tile_m == 32 else 64
if warp_tile_k != expected_k:
return False, (
f"For {a_datatype} on {gpu_target}, warp_tile_m={warp_tile_m} "
f"requires warp_tile_k={expected_k}, got warp_tile_k={warp_tile_k}{suffix}"
)
return True, ""
def validate_gemm_rowcol_tensor_quant(
tile_m: int,
tile_n: int,
@@ -1294,21 +1458,171 @@ def validate_gemm_rowcol_tensor_quant(
if not whole_workgroup_cover_valid:
return False, whole_workgroup_cover_error
if warp_tile_m != warp_tile_n:
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="RowColQuant / TensorQuant"
)
def validate_gemm_aquant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
group_size_k: int = 128,
) -> Tuple[bool, str]:
"""Validate AQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
logging.debug(
f"Whole workgroup cover configuration validation failed: {whole_workgroup_cover_error}"
)
return False, whole_workgroup_cover_error
if tile_k % group_size_k != 0 or tile_k < group_size_k:
return False, (
f"warp_tile_m({warp_tile_m}) must be equal to warp_tile_n({warp_tile_n}) "
f"(MFMA requirement for RowColQuant / TensorQuant)"
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
f"and tile_k >= group_size_k"
)
if a_datatype in ["fp8", "bf8"]:
if gpu_target == "gfx950":
expected_k = 64 if warp_tile_m == 32 else 128
else:
expected_k = 32 if warp_tile_m == 32 else 64
if warp_tile_k != expected_k:
return False, (
f"For {a_datatype} on {gpu_target}, warp_tile_m={warp_tile_m} "
f"requires warp_tile_k={expected_k}, got warp_tile_k={warp_tile_k}"
)
if group_size_k % warp_tile_k != 0:
return False, (
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
)
return True, ""
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="AQuant"
)
def validate_gemm_bquant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
group_size_k: int = 128,
) -> Tuple[bool, str]:
"""Validate BQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
return False, whole_workgroup_cover_error
if tile_k % group_size_k != 0 or tile_k < group_size_k:
return False, (
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
f"and tile_k >= group_size_k"
)
if group_size_k % warp_tile_k != 0:
return False, (
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
)
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="BQuant"
)
def validate_gemm_abquant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
group_size_k: int = 128,
) -> Tuple[bool, str]:
"""Validate ABQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
return False, whole_workgroup_cover_error
if tile_k % group_size_k != 0 or tile_k < group_size_k:
return False, (
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
f"and tile_k >= group_size_k"
)
if group_size_k % warp_tile_k != 0:
return False, (
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
)
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="ABQuant"
)

View File

@@ -7,4 +7,5 @@ from sampling.budget import allocate_budget as allocate_budget
from sampling.budget import load_op_weights as load_op_weights
from sampling.manifest import write_manifest as write_manifest
from sampling.feasible_set import GEMM_AXES as GEMM_AXES
from sampling.feasible_set import GEMM_ABQUANT_AXES as GEMM_ABQUANT_AXES
from sampling.feasible_set import normalize_axis_values as normalize_axis_values

View File

@@ -22,6 +22,27 @@ GEMM_AXES = [
GEMM_STREAMK_AXES = GEMM_AXES + ["reduction_strategy"]
GEMM_ABQUANT_AXES = [
"tile_m",
"tile_n",
"tile_k",
"warp_m",
"warp_n",
"warp_k",
"warp_tile_m",
"warp_tile_n",
"warp_tile_k",
"pipeline",
"epilogue",
"scheduler",
"pad_m",
"pad_n",
"pad_k",
"a_preshuffle_quant",
"b_preshuffle_quant",
"group_size_n",
]
CATEGORICAL_AXES = {
"pipeline",
"epilogue",
@@ -31,6 +52,8 @@ CATEGORICAL_AXES = {
"pad_n",
"pad_k",
"persistent",
"a_preshuffle_quant",
"b_preshuffle_quant",
}

View File

@@ -2,18 +2,21 @@
"version": "1.0",
"description": "Op weights for Daily Tier budget allocation (RFC section 3.4)",
"weights": {
"gemm_universal": 0.0770,
"gemm_preshuffle": 0.0770,
"gemm_multi_d": 0.0770,
"grouped_gemm": 0.0769,
"gemm_streamk": 0.0769,
"batched_contraction": 0.0769,
"batched_gemm": 0.0769,
"gemm_multi_abd": 0.0769,
"mx_gemm": 0.0769,
"gemm_rowcolquant": 0.0769,
"gemm_tensor_quant": 0.0769,
"grouped_gemm_rowcolquant": 0.0769,
"grouped_gemm_tensorquant": 0.0769
"gemm_universal": 0.0625,
"gemm_preshuffle": 0.0625,
"gemm_multi_d": 0.0625,
"grouped_gemm": 0.0625,
"gemm_streamk": 0.0625,
"batched_contraction": 0.0625,
"batched_gemm": 0.0625,
"gemm_multi_abd": 0.0625,
"mx_gemm": 0.0625,
"gemm_rowcolquant": 0.0625,
"gemm_tensor_quant": 0.0625,
"grouped_gemm_rowcolquant": 0.0625,
"grouped_gemm_tensorquant": 0.0625,
"gemm_aquant": 0.0625,
"gemm_bquant": 0.0625,
"gemm_abquant": 0.0625
}
}