diff --git a/tile_engine/operation_support_matrix.md b/tile_engine/operation_support_matrix.md
index 6444956ad9..3fb3379ad7 100644
--- a/tile_engine/operation_support_matrix.md
+++ b/tile_engine/operation_support_matrix.md
@@ -13,6 +13,9 @@
| GEMM | flatmm
example: 18_flatmm/ | ❌ | ❌ | ❌ | ❌ | | ❌ | ❌ | ❌ | | | | ❌ | ❌ | ❌ | ❌ | |
| GEMM | gemm_multi_abd
example: 22_gemm_multi_abd/ | ✅ | | | | | | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | 0.0833 |
| GEMM | gemm_quant | | ❌ | | ❌ | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ | |
+| GEMM | block_scale_gemm/gemm_aquant
engine: block_scale_gemm/gemm_aquant/
example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | | | | ❌ | ✅ | ✅ | ❌ | 0.0625 |
+| GEMM | block_scale_gemm/gemm_bquant
engine: block_scale_gemm/gemm_bquant/
example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | 0.0625 |
+| GEMM | block_scale_gemm/gemm_abquant
engine: block_scale_gemm/gemm_abquant/
example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | 0.0625 |
| GEMM | grouped_gemm [10]
engine: grouped_gemm/
example: 17_grouped_gemm/ | ✅ | ✅ | | | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | 0.0834 |
| GEMM | grouped_gemm_quant/grouped_gemm_rowcolquant
engine: grouped_gemm_quant/grouped_gemm_rowcolquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 |
| GEMM | grouped_gemm_quant/grouped_gemm_tensorquant
engine: grouped_gemm_quant/grouped_gemm_tensorquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 |
diff --git a/tile_engine/ops/gemm/CMakeLists.txt b/tile_engine/ops/gemm/CMakeLists.txt
index b50a679010..c2967f1f57 100644
--- a/tile_engine/ops/gemm/CMakeLists.txt
+++ b/tile_engine/ops/gemm/CMakeLists.txt
@@ -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")
diff --git a/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt
index 83803dabc3..035ac7008b 100644
--- a/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt
+++ b/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt
@@ -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)
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt
new file mode 100644
index 0000000000..980b3005a4
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json
new file mode 100644
index 0000000000..ef77633ddb
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json
@@ -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
+ ]
+ }
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json
new file mode 100644
index 0000000000..70c307ba10
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json
@@ -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
+ ]
+ }
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json
new file mode 100644
index 0000000000..a20dca95f7
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json
@@ -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
+ ]
+ }
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp
new file mode 100644
index 0000000000..18fbbc2c0a
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp
@@ -0,0 +1,243 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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
+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(
+ ck_tile::integer_divide_ceil(K, kbatch));
+ const auto atol = ck_tile::get_absolute_threshold(
+ max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+ const auto rtol_split_k =
+ ck_tile::get_relative_threshold(kbatch);
+ const auto atol_split_k = ck_tile::get_absolute_threshold(
+ 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& c_m_n_dev_result,
+ const ck_tile::HostTensor& c_m_n_host_result)
+{
+ const float max_accumulated_value =
+ std::abs(static_cast(*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(a)) <
+ std::abs(static_cast(b));
+ })));
+ const auto rtol_atol = calculate_rtol_atol_abquant(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& a_m_k,
+ ck_tile::HostTensor& aq_m_qk,
+ ck_tile::HostTensor& b_k_n,
+ ck_tile::HostTensor& bq_qk_n,
+ ck_tile::HostTensor& c_m_n_host_result)
+{
+ if(verify == 1)
+ {
+ c_m_n_host_result.SetZero();
+ ck_tile::reference_gemm_abquant(
+ a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result);
+ }
+}
+#pragma clang diagnostic pop
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py
new file mode 100644
index 0000000000..8642298646
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py
@@ -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())
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp
new file mode 100644
index 0000000000..960473cbb4
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp
@@ -0,0 +1,178 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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::name;
+ std::string dtype_b = DataTypeTraits::name;
+ std::string dtype_aq = DataTypeTraits::name;
+ std::string dtype_bq = DataTypeTraits::name;
+ std::string dtype_acc = DataTypeTraits::name;
+ std::string dtype_c = DataTypeTraits::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(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;
+ }
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp
new file mode 100644
index 0000000000..c8726993c8
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp
@@ -0,0 +1,76 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/core/numeric/integer.hpp"
+
+// DataTypeTraits for all supported types
+template
+struct DataTypeTraits;
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp32";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp16";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "bf16";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp8";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "bf8";
+};
+
+// Helper function to determine if a layout is row-major
+template
+constexpr auto is_row_major(Layout)
+{
+ return ck_tile::bool_constant>{};
+}
+
+// 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)
+ {
+ }
+};
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py
new file mode 100644
index 0000000000..93d51ee09e
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py
@@ -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
+#include
+#include
+#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,
+ ck_tile::sequence,
+ ck_tile::sequence>;
+
+ // Tile partitioner
+ using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner;
+
+ // Quantization group sizes
+ using AQuantGroupSize = ck_tile::QuantGroupShape>;
+ using BQuantGroupSize = ck_tile::QuantGroupShape>;
+
+ // 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)};"""
+
+ 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)};"""
+
+ 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;"""
+ 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;"""
+
+ 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(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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp
new file mode 100644
index 0000000000..7efb4dab55
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp
@@ -0,0 +1,336 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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 kernel_func)
+ {
+ std::vector(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(
+ 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 a_m_k(ck_tile::host_tensor_descriptor(
+ problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
+ ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor(
+ problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
+ ck_tile::HostTensor 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 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 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{-1.f, 1.f}(a_m_k);
+ ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n);
+ ck_tile::FillUniformDistribution{0.5f, 1.5f}(aq_m_qk);
+ ck_tile::FillUniformDistribution{0.5f, 1.5f}(bq_qk_n);
+ }
+ else if(setting_.init_method_ == 1)
+ {
+ ck_tile::FillMonotonicSeq{}(a_m_k);
+ ck_tile::FillMonotonicSeq{}(b_k_n);
+ ck_tile::FillConstant{static_cast(1)}(aq_m_qk);
+ ck_tile::FillConstant{static_cast(1)}(bq_qk_n);
+ }
+ else if(setting_.init_method_ == 2)
+ {
+ ck_tile::FillConstant{static_cast(1)}(a_m_k);
+ ck_tile::FillConstant{static_cast(1)}(b_k_n);
+ ck_tile::FillConstant{static_cast(1)}(aq_m_qk);
+ ck_tile::FillConstant{static_cast(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 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 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 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& c_m_n_host_result,
+ ck_tile::HostTensor& c_m_n_dev_result,
+ const std::tuple& 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(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 kernel_instances_;
+};
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py
new file mode 100644
index 0000000000..320599f115
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt
new file mode 100644
index 0000000000..49de66cfec
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json
new file mode 100644
index 0000000000..a998050ebe
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json
@@ -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
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json
new file mode 100644
index 0000000000..7e6c8581ec
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json
@@ -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
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json
new file mode 100644
index 0000000000..37b077e1a2
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json
@@ -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
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp
new file mode 100644
index 0000000000..3fc07b2793
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp
@@ -0,0 +1,231 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+
+#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
+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;
+ const auto rtol = ck_tile::get_relative_threshold(
+ ck_tile::integer_divide_ceil(K, kbatch));
+ const auto atol = ck_tile::get_absolute_threshold(
+ max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+ const auto rtol_split_k =
+ ck_tile::get_relative_threshold(kbatch);
+ const auto atol_split_k = ck_tile::get_absolute_threshold(
+ 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& c_m_n_dev_result,
+ ck_tile::HostTensor& c_m_n_host_result)
+{
+ const float max_accumulated_value =
+ std::abs(static_cast(*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(a)) <
+ std::abs(static_cast(b));
+ })));
+ const auto rtol_atol =
+ calculate_rtol_atol_aquant(
+ 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& a_m_k,
+ ck_tile::HostTensor& aq_m_qk,
+ ck_tile::HostTensor& b_k_n,
+ ck_tile::HostTensor& c_m_n_host_result)
+{
+ if(verify == 1)
+ {
+ c_m_n_host_result.SetZero();
+ using QuantGroupSize = typename SelectedKernel::QuantGroupSize;
+ ck_tile::reference_gemm_quant(a_m_k, aq_m_qk, b_k_n, c_m_n_host_result);
+ }
+}
+#pragma clang diagnostic pop
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py
new file mode 100644
index 0000000000..2f55340643
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py
@@ -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())
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp
new file mode 100644
index 0000000000..637cef5c54
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp
@@ -0,0 +1,153 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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::name;
+ std::string dtype_b = DataTypeTraits::name;
+ std::string dtype_aq = DataTypeTraits::name;
+ std::string dtype_acc = DataTypeTraits::name;
+ std::string dtype_c = DataTypeTraits::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(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;
+ }
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp
new file mode 100644
index 0000000000..00fb4e1d7a
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp
@@ -0,0 +1,73 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/core/numeric/integer.hpp"
+
+// DataTypeTraits for all supported types
+template
+struct DataTypeTraits;
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp32";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp16";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "bf16";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp8";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "bf8";
+};
+
+// Helper function to determine if a layout is row-major
+template
+constexpr auto is_row_major(Layout)
+{
+ return ck_tile::bool_constant>{};
+}
+
+// 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)
+ {
+ }
+};
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py
new file mode 100644
index 0000000000..a57a3bb2e8
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp
new file mode 100644
index 0000000000..00a19b0b44
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp
@@ -0,0 +1,276 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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 kernel_func)
+ {
+ std::vector(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(
+ 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 a_m_k(ck_tile::host_tensor_descriptor(
+ problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
+ ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor(
+ problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
+ ck_tile::HostTensor 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 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{-1.f, 1.f}(a_m_k);
+ ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n);
+ ck_tile::FillUniformDistribution{0.5f, 1.5f}(aq_m_qk);
+ }
+ else if(setting_.init_method_ == 1)
+ {
+ ck_tile::FillMonotonicSeq{}(a_m_k);
+ ck_tile::FillMonotonicSeq{}(b_k_n);
+ ck_tile::FillConstant{static_cast(1)}(aq_m_qk);
+ }
+ else if(setting_.init_method_ == 2)
+ {
+ ck_tile::FillConstant{static_cast(1)}(a_m_k);
+ ck_tile::FillConstant{static_cast(1)}(b_k_n);
+ ck_tile::FillConstant{static_cast(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 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& c_m_n_host_result,
+ ck_tile::HostTensor& c_m_n_dev_result,
+ const std::tuple& 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(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 kernel_instances_;
+};
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py
new file mode 100644
index 0000000000..4cbfc024c8
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt
new file mode 100644
index 0000000000..d0c0785df6
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json
new file mode 100644
index 0000000000..79ecb47a1f
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json
@@ -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
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json
new file mode 100644
index 0000000000..9ed8302878
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json
@@ -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
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json
new file mode 100644
index 0000000000..c2ea50a65c
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json
@@ -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
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp
new file mode 100644
index 0000000000..d50ce46755
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp
@@ -0,0 +1,229 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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
+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;
+ const auto rtol = ck_tile::get_relative_threshold(
+ ck_tile::integer_divide_ceil(K, kbatch));
+ const auto atol = ck_tile::get_absolute_threshold(
+ max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+ const auto rtol_split_k =
+ ck_tile::get_relative_threshold(kbatch);
+ const auto atol_split_k = ck_tile::get_absolute_threshold(
+ 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& c_m_n_dev_result,
+ ck_tile::HostTensor& c_m_n_host_result)
+{
+ const float max_accumulated_value =
+ std::abs(static_cast(*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(a)) <
+ std::abs(static_cast(b));
+ })));
+ const auto rtol_atol =
+ calculate_rtol_atol_bquant(
+ 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& a_m_k,
+ ck_tile::HostTensor& b_k_n,
+ ck_tile::HostTensor& bq_qk_n,
+ ck_tile::HostTensor& c_m_n_host_result)
+{
+ if(verify == 1)
+ {
+ c_m_n_host_result.SetZero();
+ using QuantGroupSize = typename SelectedKernel::QuantGroupSize;
+ ck_tile::reference_gemm_quant(
+ a_m_k, bq_qk_n, b_k_n, c_m_n_host_result);
+ }
+}
+#pragma clang diagnostic pop
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py
new file mode 100644
index 0000000000..b19d3702e2
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py
@@ -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())
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp
new file mode 100644
index 0000000000..746ced532f
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp
@@ -0,0 +1,166 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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::name;
+ std::string dtype_b = DataTypeTraits::name;
+ std::string dtype_bq = DataTypeTraits::name;
+ std::string dtype_acc = DataTypeTraits::name;
+ std::string dtype_c = DataTypeTraits::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(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;
+ }
+}
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp
new file mode 100644
index 0000000000..1408f4473d
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp
@@ -0,0 +1,74 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/core/numeric/integer.hpp"
+
+// DataTypeTraits for all supported types
+template
+struct DataTypeTraits;
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp32";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp16";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "bf16";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "fp8";
+};
+
+template <>
+struct DataTypeTraits
+{
+ static constexpr const char* name = "bf8";
+};
+
+// Helper function to determine if a layout is row-major
+template
+constexpr auto is_row_major(Layout)
+{
+ return ck_tile::bool_constant>{};
+}
+
+// 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)
+ {
+ }
+};
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py
new file mode 100644
index 0000000000..770ec9d6f1
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py
@@ -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()
diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp
new file mode 100644
index 0000000000..7f562ef0ad
--- /dev/null
+++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp
@@ -0,0 +1,293 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#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 kernel_func)
+ {
+ std::vector(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(
+ 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 a_m_k(ck_tile::host_tensor_descriptor(
+ problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
+ ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor(
+ problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
+ ck_tile::HostTensor 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 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{-1.f, 1.f}(a_m_k);
+ ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n);
+ ck_tile::FillUniformDistribution{0.5f, 1.5f}(bq_qk_n);
+ }
+ else if(setting_.init_method_ == 1)
+ {
+ ck_tile::FillMonotonicSeq{}(a_m_k);
+ ck_tile::FillMonotonicSeq{}(b_k_n);
+ ck_tile::FillConstant{static_cast(1)}(bq_qk_n);
+ }
+ else if(setting_.init_method_ == 2)
+ {
+ ck_tile::FillConstant{static_cast(1)}(a_m_k);
+ ck_tile::FillConstant{static_cast(1)}(b_k_n);
+ ck_tile::FillConstant