mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-09 00:27:01 +00:00
[rocm-libraries] ROCm/rocm-libraries#8519 (commit 9637390)
feat(ck-tile): add block-scale GEMM operators (aquant, bquant, abquant) (#8519) JIRA ID - AICK-1289 Motivation Adds three new block-scale quantized GEMM operators to the CK Tile Engine for FP8/BF8 inference workloads. Technical Details gemm_aquant: A-matrix quantized GEMM with per-row-group scale tensor [M, K/group_size_k] gemm_bquant: B-matrix quantized GEMM with per-column-group scale tensor [K/group_size_k, N] gemm_abquant: Both A and B quantized with independent group-scale tensors Each operator includes CMakeLists, Python instance builder with tier sampling, C++ benchmark/profiler with host reference verification, and config JSONs. Supporting changes to gemm_instance_builder.py, gemm_validation_utils.py, sampling infra, and the operation support matrix. Test Plan Build and run all three operators with fp8/bf8 on gfx942/gfx950 Verify correctness against CPU reference Verify CI config builds pass
This commit is contained in:
committed by
assistant-librarian[bot]
parent
05615558c6
commit
b0f200713a
@@ -13,6 +13,9 @@
|
||||
| GEMM | flatmm<br>example: 18_flatmm/ | ❌ | ❌ | ❌ | ❌ | | ❌ | ❌ | ❌ | | | | ❌ | ❌ | ❌ | ❌ | |
|
||||
| GEMM | gemm_multi_abd<br>example: 22_gemm_multi_abd/ | ✅ | | | | | | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | 0.0833 |
|
||||
| GEMM | gemm_quant | | ❌ | | ❌ | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ | |
|
||||
| GEMM | block_scale_gemm/gemm_aquant<br>engine: block_scale_gemm/gemm_aquant/<br>example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | | | | ❌ | ✅ | ✅ | ❌ | 0.0625 |
|
||||
| GEMM | block_scale_gemm/gemm_bquant<br>engine: block_scale_gemm/gemm_bquant/<br>example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | 0.0625 |
|
||||
| GEMM | block_scale_gemm/gemm_abquant<br>engine: block_scale_gemm/gemm_abquant/<br>example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | 0.0625 |
|
||||
| GEMM | grouped_gemm [10]<br>engine: grouped_gemm/<br>example: 17_grouped_gemm/ | ✅ | ✅ | | | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | 0.0834 |
|
||||
| GEMM | grouped_gemm_quant/grouped_gemm_rowcolquant<br>engine: grouped_gemm_quant/grouped_gemm_rowcolquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 |
|
||||
| GEMM | grouped_gemm_quant/grouped_gemm_tensorquant<br>engine: grouped_gemm_quant/grouped_gemm_tensorquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 |
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "gemm_abquant_common.hpp"
|
||||
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions"
|
||||
|
||||
// Data types and Layouts are defined by the generated kernel headers:
|
||||
// ADataType, BDataType, AQDataType, BQDataType, AccDataType, CDataType
|
||||
// ALayout, BLayout, CLayout, AQLayout, BQLayout
|
||||
|
||||
enum class Metric
|
||||
{
|
||||
LATENCY = 0,
|
||||
TFLOPS = 1,
|
||||
BANDWIDTH = 2
|
||||
};
|
||||
|
||||
inline constexpr auto get_metric_name(Metric m)
|
||||
{
|
||||
switch(m)
|
||||
{
|
||||
case Metric::LATENCY: return "latency";
|
||||
case Metric::TFLOPS: return "tflops";
|
||||
case Metric::BANDWIDTH: return "bandwidth";
|
||||
default: throw std::invalid_argument("Unsupported metric type");
|
||||
}
|
||||
}
|
||||
|
||||
struct ABQuantGemmProblem
|
||||
{
|
||||
int split_k_;
|
||||
int m_, n_, k_;
|
||||
int stride_a_, stride_b_, stride_c_;
|
||||
int stride_aq_;
|
||||
int stride_bq_;
|
||||
int group_size_k_;
|
||||
int group_size_n_;
|
||||
|
||||
std::string dtype_a_, dtype_b_, dtype_aq_, dtype_bq_, dtype_acc_, dtype_c_;
|
||||
std::string layout_a_, layout_b_, layout_c_;
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const ABQuantGemmProblem& problem)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"split_k\":" << problem.split_k_ << ",\n"
|
||||
<< " \"m\":" << problem.m_ << ",\n"
|
||||
<< " \"n\":" << problem.n_ << ",\n"
|
||||
<< " \"k\":" << problem.k_ << ",\n"
|
||||
<< " \"stride_a\":" << problem.stride_a_ << ",\n"
|
||||
<< " \"stride_b\":" << problem.stride_b_ << ",\n"
|
||||
<< " \"stride_c\":" << problem.stride_c_ << ",\n"
|
||||
<< " \"stride_aq\":" << problem.stride_aq_ << ",\n"
|
||||
<< " \"stride_bq\":" << problem.stride_bq_ << ",\n"
|
||||
<< " \"group_size_k\":" << problem.group_size_k_ << ",\n"
|
||||
<< " \"group_size_n\":" << problem.group_size_n_ << ",\n"
|
||||
<< " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n"
|
||||
<< " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n"
|
||||
<< " \"dtype_aq\":\"" << problem.dtype_aq_ << "\",\n"
|
||||
<< " \"dtype_bq\":\"" << problem.dtype_bq_ << "\",\n"
|
||||
<< " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n"
|
||||
<< " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n"
|
||||
<< " \"layout_a\":\"" << problem.layout_a_ << "\",\n"
|
||||
<< " \"layout_b\":\"" << problem.layout_b_ << "\",\n"
|
||||
<< " \"layout_c\":\"" << problem.layout_c_ << "\"\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct PerformanceResult
|
||||
{
|
||||
double latency_;
|
||||
double tflops_;
|
||||
double bandwidth_;
|
||||
|
||||
static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m)
|
||||
{
|
||||
switch(m)
|
||||
{
|
||||
case Metric::LATENCY: return a.latency_ < b.latency_;
|
||||
case Metric::TFLOPS: return a.tflops_ > b.tflops_;
|
||||
case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_;
|
||||
default: throw std::invalid_argument("Unsupported metric type");
|
||||
}
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_
|
||||
<< ",\n"
|
||||
<< " \"tflops(TFlops)\": " << result.tflops_ << ",\n"
|
||||
<< " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct KernelInstance
|
||||
{
|
||||
std::string name_;
|
||||
ABQuantGemmProblem problem_;
|
||||
PerformanceResult perf_result_;
|
||||
|
||||
static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m)
|
||||
{
|
||||
return PerformanceResult::compare(a.perf_result_, b.perf_result_, m);
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"name\": \"" << obj.name_ << "\",\n"
|
||||
<< " \"problem\": " << obj.problem_ << ",\n"
|
||||
<< " \"perf_result\": " << obj.perf_result_ << "\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct Setting
|
||||
{
|
||||
int n_warmup_;
|
||||
int n_repeat_;
|
||||
bool is_gpu_timer_;
|
||||
int verify_;
|
||||
int init_method_;
|
||||
bool log_;
|
||||
std::string csv_filename_;
|
||||
bool flush_cache_;
|
||||
int rotating_count_;
|
||||
bool json_output_;
|
||||
};
|
||||
|
||||
inline std::string get_rocm_version()
|
||||
{
|
||||
std::ifstream version_file("/opt/rocm/.info/version");
|
||||
if(version_file.is_open())
|
||||
{
|
||||
std::string version;
|
||||
std::getline(version_file, version);
|
||||
return version;
|
||||
}
|
||||
return "Unknown";
|
||||
}
|
||||
|
||||
template <typename ADataType_,
|
||||
typename BDataType_,
|
||||
typename AQDataType_,
|
||||
typename BQDataType_,
|
||||
typename AccDataType_,
|
||||
typename CDataType_>
|
||||
auto calculate_rtol_atol_abquant(const ck_tile::index_t K,
|
||||
const ck_tile::index_t kbatch,
|
||||
const float max_accumulated_value)
|
||||
{
|
||||
// Both A and B are the same FP8/BF8 type for abquant; assert so mixed-precision additions are
|
||||
// caught.
|
||||
static_assert(sizeof(ADataType_) == sizeof(BDataType_),
|
||||
"calculate_rtol_atol_abquant assumes equal-width A and B types");
|
||||
using ComputeType = ADataType_;
|
||||
const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType_, AccDataType_>(
|
||||
ck_tile::integer_divide_ceil(K, kbatch));
|
||||
const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType_, AccDataType_>(
|
||||
max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
|
||||
const auto rtol_split_k =
|
||||
ck_tile::get_relative_threshold<CDataType_, CDataType_, CDataType_>(kbatch);
|
||||
const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType_, CDataType_, CDataType_>(
|
||||
max_accumulated_value, kbatch);
|
||||
return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
|
||||
}
|
||||
|
||||
/// @brief Compare device and host results for ABQuant GEMM
|
||||
bool compare_abquant(std::string instanceName,
|
||||
ck_tile::index_t K,
|
||||
ck_tile::index_t kbatch,
|
||||
const ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
const ck_tile::HostTensor<CDataType>& c_m_n_host_result)
|
||||
{
|
||||
const float max_accumulated_value =
|
||||
std::abs(static_cast<float>(*std::max_element(c_m_n_host_result.mData.begin(),
|
||||
c_m_n_host_result.mData.end(),
|
||||
[](const auto& a, const auto& b) {
|
||||
return std::abs(static_cast<float>(a)) <
|
||||
std::abs(static_cast<float>(b));
|
||||
})));
|
||||
const auto rtol_atol = calculate_rtol_atol_abquant<ADataType,
|
||||
BDataType,
|
||||
AQDataType,
|
||||
BQDataType,
|
||||
AccDataType,
|
||||
CDataType>(K, kbatch, max_accumulated_value);
|
||||
bool pass = ck_tile::check_err(c_m_n_dev_result,
|
||||
c_m_n_host_result,
|
||||
"Error: Incorrect results!",
|
||||
rtol_atol.at(ck_tile::number<0>{}),
|
||||
rtol_atol.at(ck_tile::number<1>{}));
|
||||
|
||||
std::cout << "For " << instanceName << " Relative error threshold is "
|
||||
<< rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
|
||||
<< rtol_atol.at(ck_tile::number<1>{}) << std::endl;
|
||||
std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;
|
||||
|
||||
return pass;
|
||||
}
|
||||
|
||||
/// @brief CPU reference implementation for ABQuant GEMM
|
||||
void abquant_gemm_host_reference(int verify,
|
||||
ck_tile::HostTensor<ADataType>& a_m_k,
|
||||
ck_tile::HostTensor<AQDataType>& aq_m_qk,
|
||||
ck_tile::HostTensor<BDataType>& b_k_n,
|
||||
ck_tile::HostTensor<BQDataType>& bq_qk_n,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
|
||||
{
|
||||
if(verify == 1)
|
||||
{
|
||||
c_m_n_host_result.SetZero();
|
||||
ck_tile::reference_gemm_abquant<ADataType,
|
||||
AQDataType,
|
||||
BDataType,
|
||||
BQDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
SelectedKernel::AQuantGroupSize,
|
||||
SelectedKernel::BQuantGroupSize>(
|
||||
a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result);
|
||||
}
|
||||
}
|
||||
#pragma clang diagnostic pop
|
||||
@@ -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())
|
||||
@@ -0,0 +1,178 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#include <iostream>
|
||||
#include <functional>
|
||||
#include <tuple>
|
||||
#include <exception>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "gemm_abquant_profiler.hpp"
|
||||
#include "gemm_abquant_common.hpp"
|
||||
|
||||
// The kernel header is included via the compile command line with -include flag
|
||||
// It defines SelectedKernel struct, KERNEL_NAME, and type aliases:
|
||||
// ADataType, BDataType, AQDataType, BQDataType, AccDataType, CDataType
|
||||
// ALayout, BLayout, CLayout, AQLayout, BQLayout
|
||||
|
||||
inline auto create_args(int argc, char* argv[])
|
||||
{
|
||||
ck_tile::ArgParser arg_parser;
|
||||
arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.")
|
||||
.insert("n", "4096", "The value for n dimension. Default is 4096.")
|
||||
.insert("k", "2048", "The value for k dimension. Default is 2048.")
|
||||
.insert("stride_a", "0", "The stride value for tensor A. Default is 0.")
|
||||
.insert("stride_b", "0", "The stride value for tensor B. Default is 0.")
|
||||
.insert("stride_c", "0", "The stride value for tensor C. Default is 0.")
|
||||
.insert("split_k", "1", "The split value for k dimension. Default is 1.")
|
||||
.insert("group_size_k",
|
||||
std::to_string(SelectedKernel::GroupSizeK),
|
||||
"The quantization group size along K. Default matches kernel config.")
|
||||
.insert("group_size_n",
|
||||
std::to_string(SelectedKernel::GroupSizeN),
|
||||
"The quantization group size along N for BQ. Default matches kernel config.")
|
||||
.insert("verify",
|
||||
"1",
|
||||
"The type of validation. Set to 0 for no validation, 1 for validation on CPU. "
|
||||
"Default is 1, CPU validation.")
|
||||
.insert("log",
|
||||
"false",
|
||||
"Whether output kernel instance information or not. Possible values are true or "
|
||||
"false. Default is false")
|
||||
.insert(
|
||||
"warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.")
|
||||
.insert(
|
||||
"repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.")
|
||||
.insert("timer",
|
||||
"true",
|
||||
"Whether if the timer is gpu timer or not. Possible values are false or true. "
|
||||
"Default is true.")
|
||||
.insert("init",
|
||||
"0",
|
||||
"The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 "
|
||||
"for constant(1). Default is 0, random.")
|
||||
.insert("flush_cache",
|
||||
"true",
|
||||
"To flush cache, possible values are true or false. Default is true.")
|
||||
.insert(
|
||||
"rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.")
|
||||
.insert("metric",
|
||||
"0",
|
||||
"Metric with which to measure kernel performance. Set to 0 for latency, 1 for "
|
||||
"tflops, or 2 for bandwidth. Default is 0, latency.")
|
||||
.insert("csv_filename",
|
||||
"",
|
||||
"The filename of benchmark result. Default is empty (no CSV output).")
|
||||
.insert("json_output",
|
||||
"false",
|
||||
"Whether to output results in JSON format only. Possible values are true or false. "
|
||||
"Default is false");
|
||||
|
||||
bool result = arg_parser.parse(argc, argv);
|
||||
return std::make_tuple(result, arg_parser);
|
||||
}
|
||||
|
||||
void benchmark_single(const ck_tile::ArgParser& arg_parser)
|
||||
{
|
||||
std::string dtype_a = DataTypeTraits<ADataType>::name;
|
||||
std::string dtype_b = DataTypeTraits<BDataType>::name;
|
||||
std::string dtype_aq = DataTypeTraits<AQDataType>::name;
|
||||
std::string dtype_bq = DataTypeTraits<BQDataType>::name;
|
||||
std::string dtype_acc = DataTypeTraits<AccDataType>::name;
|
||||
std::string dtype_c = DataTypeTraits<CDataType>::name;
|
||||
|
||||
std::string layout_a = ALayout::name;
|
||||
std::string layout_b = BLayout::name;
|
||||
std::string layout_c = CLayout::name;
|
||||
|
||||
int M = arg_parser.get_int("m");
|
||||
int N = arg_parser.get_int("n");
|
||||
int K = arg_parser.get_int("k");
|
||||
int group_size_k = arg_parser.get_int("group_size_k");
|
||||
int group_size_n = arg_parser.get_int("group_size_n");
|
||||
|
||||
if(M <= 0 || N <= 0 || K <= 0)
|
||||
{
|
||||
throw std::invalid_argument("m, n, k must be positive integers");
|
||||
}
|
||||
if(group_size_k <= 0 || K % group_size_k != 0)
|
||||
{
|
||||
throw std::invalid_argument(
|
||||
"group_size_k must be positive and k must be divisible by group_size_k");
|
||||
}
|
||||
if(group_size_n > 1 && N % group_size_n != 0)
|
||||
{
|
||||
throw std::invalid_argument("n must be divisible by group_size_n when group_size_n > 1");
|
||||
}
|
||||
|
||||
ABQuantGemmProblem problem{arg_parser.get_int("split_k"),
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
arg_parser.get_int("stride_a"),
|
||||
arg_parser.get_int("stride_b"),
|
||||
arg_parser.get_int("stride_c"),
|
||||
0, // stride_aq computed by profiler
|
||||
0, // stride_bq computed by profiler
|
||||
group_size_k,
|
||||
group_size_n,
|
||||
dtype_a,
|
||||
dtype_b,
|
||||
dtype_aq,
|
||||
dtype_bq,
|
||||
dtype_acc,
|
||||
dtype_c,
|
||||
layout_a,
|
||||
layout_b,
|
||||
layout_c};
|
||||
|
||||
Setting setting{arg_parser.get_int("warmup"),
|
||||
arg_parser.get_int("repeat"),
|
||||
arg_parser.get_bool("timer"),
|
||||
arg_parser.get_int("verify"),
|
||||
arg_parser.get_int("init"),
|
||||
arg_parser.get_bool("log"),
|
||||
arg_parser.get_str("csv_filename"),
|
||||
arg_parser.get_bool("flush_cache"),
|
||||
arg_parser.get_int("rotating_count"),
|
||||
arg_parser.get_bool("json_output")};
|
||||
|
||||
auto& profiler = ABQuantGemmProfiler::instance(setting);
|
||||
|
||||
try
|
||||
{
|
||||
auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args,
|
||||
const ck_tile::stream_config& stream) {
|
||||
return SelectedKernel::launch(args, stream);
|
||||
};
|
||||
|
||||
profiler.benchmark(problem, kernel_func);
|
||||
profiler.select_best_instance(static_cast<Metric>(arg_parser.get_int("metric")));
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Benchmark failed: " << e.what() << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
try
|
||||
{
|
||||
auto [result, parser] = create_args(argc, argv);
|
||||
if(!result)
|
||||
return EXIT_FAILURE;
|
||||
|
||||
benchmark_single(parser);
|
||||
return 0;
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Error: " << e.what() << "\n";
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/core/numeric/integer.hpp"
|
||||
|
||||
// DataTypeTraits for all supported types
|
||||
template <typename T>
|
||||
struct DataTypeTraits;
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<float>
|
||||
{
|
||||
static constexpr const char* name = "fp32";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::half_t>
|
||||
{
|
||||
static constexpr const char* name = "fp16";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::bf16_t>
|
||||
{
|
||||
static constexpr const char* name = "bf16";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::fp8_t>
|
||||
{
|
||||
static constexpr const char* name = "fp8";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::bf8_t>
|
||||
{
|
||||
static constexpr const char* name = "bf8";
|
||||
};
|
||||
|
||||
// Helper function to determine if a layout is row-major
|
||||
template <typename Layout>
|
||||
constexpr auto is_row_major(Layout)
|
||||
{
|
||||
return ck_tile::bool_constant<std::is_same_v<Layout, ck_tile::tensor_layout::gemm::RowMajor>>{};
|
||||
}
|
||||
|
||||
// Structure to hold kernel traits for dispatcher
|
||||
struct ABQuantKernelTraits
|
||||
{
|
||||
std::string pipeline; // compv3
|
||||
std::string scheduler; // intrawave
|
||||
std::string epilogue; // default, cshuffle
|
||||
bool pad_m;
|
||||
bool pad_n;
|
||||
bool pad_k;
|
||||
bool a_preshuffle_quant;
|
||||
bool b_preshuffle_quant;
|
||||
|
||||
ABQuantKernelTraits()
|
||||
: pipeline("compv3"),
|
||||
scheduler("intrawave"),
|
||||
epilogue("default"),
|
||||
pad_m(false),
|
||||
pad_n(false),
|
||||
pad_k(false),
|
||||
a_preshuffle_quant(false),
|
||||
b_preshuffle_quant(false)
|
||||
{
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,904 @@
|
||||
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import os
|
||||
import argparse
|
||||
import importlib.util
|
||||
import multiprocessing
|
||||
import concurrent.futures
|
||||
import itertools
|
||||
import logging
|
||||
|
||||
|
||||
def _import_gemm_kernel_builder():
|
||||
"""Import the base GemmKernelBuilder from the gemm directory."""
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
gemm_dir = os.path.dirname(os.path.dirname(current_dir))
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"gemm_instance_builder",
|
||||
os.path.join(gemm_dir, "gemm_instance_builder.py"),
|
||||
)
|
||||
gemm_builder_module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(gemm_builder_module)
|
||||
|
||||
return gemm_builder_module.GemmKernelBuilder
|
||||
|
||||
|
||||
def _import_validation_utils():
|
||||
"""Import validation utilities."""
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
gemm_dir = os.path.dirname(os.path.dirname(current_dir))
|
||||
parent_dir = os.path.dirname(gemm_dir)
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"validation_utils",
|
||||
os.path.join(parent_dir, "gemm", "gemm_validation_utils.py"),
|
||||
)
|
||||
validation_utils = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(validation_utils)
|
||||
|
||||
return validation_utils
|
||||
|
||||
|
||||
GemmKernelBuilder = _import_gemm_kernel_builder()
|
||||
_validation_utils = _import_validation_utils()
|
||||
is_trait_combination_valid = _validation_utils.is_trait_combination_valid
|
||||
|
||||
|
||||
class GemmABQuantKernelBuilder(GemmKernelBuilder):
|
||||
def __init__(
|
||||
self,
|
||||
kernel_name_prefix,
|
||||
working_path,
|
||||
gpu_target,
|
||||
datatype,
|
||||
layout,
|
||||
config_json=None,
|
||||
max_instances=None,
|
||||
seed=None,
|
||||
tier=None,
|
||||
manifest_path=None,
|
||||
):
|
||||
super().__init__(
|
||||
kernel_name_prefix,
|
||||
working_path,
|
||||
gpu_target,
|
||||
datatype,
|
||||
layout,
|
||||
config_json,
|
||||
max_instances=max_instances,
|
||||
seed=seed,
|
||||
tier=tier,
|
||||
manifest_path=manifest_path,
|
||||
)
|
||||
self.group_size_k = self.config.get("group_size_k", 128)
|
||||
# group_size_n can be an int or a dict with "values" key
|
||||
group_size_n_cfg = self.config.get("group_size_n", 1)
|
||||
if isinstance(group_size_n_cfg, dict):
|
||||
self.group_size_n_values = group_size_n_cfg.get("values", [1])
|
||||
else:
|
||||
self.group_size_n_values = [group_size_n_cfg]
|
||||
|
||||
@staticmethod
|
||||
def _tile_config_to_str(tile_config):
|
||||
return (
|
||||
f"{tile_config['tile_m']}x{tile_config['tile_n']}x{tile_config['tile_k']}_"
|
||||
f"{tile_config['warp_m']}x{tile_config['warp_n']}x{tile_config['warp_k']}_"
|
||||
f"{tile_config['warp_tile_m']}x{tile_config['warp_tile_n']}x{tile_config['warp_tile_k']}"
|
||||
)
|
||||
|
||||
def _apply_sampling(self, kernel_list):
|
||||
"""Apply RFC Sobol+LHS+maximin sampling for ABQuant kernels.
|
||||
|
||||
Overrides base class to handle ABQuant's 8-tuple trait_combo
|
||||
and group_size_n field.
|
||||
"""
|
||||
if self.max_instances is None or len(kernel_list) <= self.max_instances:
|
||||
return kernel_list
|
||||
|
||||
import sys
|
||||
|
||||
sampling_parent = os.path.join(
|
||||
os.path.dirname(__file__), "..", "..", "..", ".."
|
||||
)
|
||||
if sampling_parent not in sys.path:
|
||||
sys.path.insert(0, sampling_parent)
|
||||
|
||||
from sampling.sampler import sample_feasible_set
|
||||
from sampling.seed import make_seed
|
||||
from sampling.feasible_set import GEMM_ABQUANT_AXES
|
||||
|
||||
effective_seed = make_seed(
|
||||
self.seed, self.gpu_target, self.datatype, self.layout
|
||||
)
|
||||
|
||||
flat_items = []
|
||||
for k in kernel_list:
|
||||
flat = dict(k["tile_config"])
|
||||
(
|
||||
pipeline,
|
||||
epilogue,
|
||||
scheduler,
|
||||
pad_m,
|
||||
pad_n,
|
||||
pad_k,
|
||||
a_preshuffle_quant,
|
||||
b_preshuffle_quant,
|
||||
) = k["trait_combo"]
|
||||
flat.update(
|
||||
{
|
||||
"pipeline": pipeline,
|
||||
"epilogue": epilogue,
|
||||
"scheduler": scheduler,
|
||||
"pad_m": pad_m,
|
||||
"pad_n": pad_n,
|
||||
"pad_k": pad_k,
|
||||
"a_preshuffle_quant": a_preshuffle_quant,
|
||||
"b_preshuffle_quant": b_preshuffle_quant,
|
||||
"group_size_n": k.get("group_size_n", 1),
|
||||
}
|
||||
)
|
||||
flat_items.append(flat)
|
||||
|
||||
selected, method, selected_indices = sample_feasible_set(
|
||||
flat_items,
|
||||
self.max_instances,
|
||||
effective_seed,
|
||||
GEMM_ABQUANT_AXES,
|
||||
)
|
||||
|
||||
kernel_list = [kernel_list[i] for i in selected_indices]
|
||||
|
||||
if self.manifest_path:
|
||||
from sampling.manifest import write_manifest
|
||||
|
||||
write_manifest(
|
||||
selected,
|
||||
self.manifest_path,
|
||||
self.kernel_name_prefix,
|
||||
self.datatype,
|
||||
self.layout,
|
||||
self.gpu_target,
|
||||
effective_seed,
|
||||
self.tier or "daily",
|
||||
method,
|
||||
)
|
||||
|
||||
print(
|
||||
f"Sampled {len(kernel_list)} from feasible set "
|
||||
f"(budget={self.max_instances}, seed={effective_seed}, method={method})"
|
||||
)
|
||||
return kernel_list
|
||||
|
||||
def _get_group_size_n_values(self):
|
||||
"""Return the list of group_size_n values to generate kernels for."""
|
||||
return self.group_size_n_values
|
||||
|
||||
def _generate_trait_combinations(self):
|
||||
"""Generate all combinations of traits for ABQuant.
|
||||
|
||||
ABQuant uses 8-tuples:
|
||||
(pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant, b_preshuffle_quant)
|
||||
"""
|
||||
trait_config = self.config["trait_config"]
|
||||
|
||||
pipelines = trait_config.get("pipeline").get("values")
|
||||
epilogues = trait_config.get("epilogue").get("values")
|
||||
schedulers = trait_config.get("scheduler").get("values")
|
||||
pad_m_values = trait_config.get("pad_m").get("values")
|
||||
pad_n_values = trait_config.get("pad_n").get("values")
|
||||
pad_k_values = trait_config.get("pad_k").get("values")
|
||||
a_preshuffle_values = trait_config.get("a_preshuffle_quant").get("values")
|
||||
b_preshuffle_values = trait_config.get("b_preshuffle_quant").get("values")
|
||||
|
||||
all_combinations = list(
|
||||
itertools.product(
|
||||
pipelines,
|
||||
epilogues,
|
||||
schedulers,
|
||||
pad_m_values,
|
||||
pad_n_values,
|
||||
pad_k_values,
|
||||
a_preshuffle_values,
|
||||
b_preshuffle_values,
|
||||
)
|
||||
)
|
||||
|
||||
# Filter out unsupported trait combinations
|
||||
combinations = []
|
||||
for combo in all_combinations:
|
||||
pipeline, epilogue, scheduler = combo[:3]
|
||||
a_preshuffle = combo[6]
|
||||
b_preshuffle = combo[7]
|
||||
if is_trait_combination_valid(
|
||||
pipeline,
|
||||
epilogue,
|
||||
scheduler,
|
||||
b_preshuffle,
|
||||
self.kernel_name_prefix,
|
||||
self.layout,
|
||||
):
|
||||
combinations.append(combo)
|
||||
else:
|
||||
logging.debug(
|
||||
f"Skipping unsupported ABQuant trait combination: "
|
||||
f"{pipeline}-{epilogue}-{scheduler}-a_pre={a_preshuffle}-b_pre={b_preshuffle}"
|
||||
)
|
||||
return combinations
|
||||
|
||||
def _list_kernels(self):
|
||||
"""Write kernel list to file for CMake to read."""
|
||||
tile_configs = self._get_tile_configs()
|
||||
trait_combos = self._generate_trait_combinations()
|
||||
group_size_n_values = self._get_group_size_n_values()
|
||||
|
||||
kernel_list = []
|
||||
for tile_config in tile_configs:
|
||||
for trait_combo in trait_combos:
|
||||
for group_size_n in group_size_n_values:
|
||||
(
|
||||
pipeline,
|
||||
epilogue,
|
||||
scheduler,
|
||||
pad_m,
|
||||
pad_n,
|
||||
pad_k,
|
||||
a_preshuffle,
|
||||
b_preshuffle,
|
||||
) = trait_combo
|
||||
|
||||
# Create kernel name with proper boolean capitalization
|
||||
kernel_name = (
|
||||
f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_"
|
||||
f"{pipeline}_{epilogue}_{scheduler}_"
|
||||
f"{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_"
|
||||
f"{str(a_preshuffle).capitalize()}_{str(b_preshuffle).capitalize()}_"
|
||||
f"gsn{group_size_n}"
|
||||
)
|
||||
|
||||
tile_str = self._tile_config_to_str(tile_config)
|
||||
kernel_name += f"_{tile_str}"
|
||||
|
||||
kernel_list.append(
|
||||
{
|
||||
"name": kernel_name,
|
||||
"tile_config": tile_config,
|
||||
"trait_combo": trait_combo,
|
||||
"group_size_n": group_size_n,
|
||||
}
|
||||
)
|
||||
|
||||
# Apply RFC-compliant sampling (Sobol + LHS + maximin)
|
||||
kernel_list = self._apply_sampling(kernel_list)
|
||||
|
||||
# Write kernel count
|
||||
with open(
|
||||
self.working_path / f"{self.kernel_name_prefix}_kernel_count.txt", "w"
|
||||
) as f:
|
||||
f.write(str(len(kernel_list)))
|
||||
|
||||
# Write kernel list
|
||||
with open(
|
||||
self.working_path / f"{self.kernel_name_prefix}_kernel_list.txt", "w"
|
||||
) as f:
|
||||
for kernel in kernel_list:
|
||||
tile_config = kernel["tile_config"]
|
||||
trait_combo = kernel["trait_combo"]
|
||||
group_size_n = kernel["group_size_n"]
|
||||
|
||||
tile_str = self._tile_config_to_str(tile_config)
|
||||
|
||||
trait_str = (
|
||||
f"{trait_combo[0]}_{trait_combo[1]}_{trait_combo[2]}_"
|
||||
+ "_".join(str(x) for x in trait_combo[3:])
|
||||
+ f"_gsn{group_size_n}"
|
||||
)
|
||||
|
||||
f.write(f"{kernel['name']}|{tile_str}|{trait_str}\n")
|
||||
|
||||
print(f"Listed {len(kernel_list)} kernel configurations")
|
||||
|
||||
def _generate_kernel_instance(self, tile_config, trait_combo, group_size_n=None):
|
||||
"""Generate a single kernel instance for ABQuant.
|
||||
|
||||
trait_combo is an 8-tuple for ABQuant.
|
||||
group_size_n is the BQ group size along N dimension.
|
||||
"""
|
||||
if group_size_n is None:
|
||||
group_size_n = self.group_size_n_values[0]
|
||||
|
||||
k_block_per_cu = self.config.get("k_block_per_cu", 1)
|
||||
|
||||
(
|
||||
pipeline,
|
||||
epilogue,
|
||||
scheduler,
|
||||
pad_m,
|
||||
pad_n,
|
||||
pad_k,
|
||||
a_preshuffle,
|
||||
b_preshuffle,
|
||||
) = trait_combo
|
||||
|
||||
# Create kernel name
|
||||
kernel_name = (
|
||||
f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_"
|
||||
f"{pipeline}_{epilogue}_{scheduler}_"
|
||||
f"{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_"
|
||||
f"{str(a_preshuffle).capitalize()}_{str(b_preshuffle).capitalize()}_"
|
||||
f"gsn{group_size_n}"
|
||||
)
|
||||
|
||||
tile_str = self._tile_config_to_str(tile_config)
|
||||
kernel_name += f"_{tile_str}"
|
||||
|
||||
# Pipeline maps
|
||||
pipeline_impl_map = {
|
||||
"compv3": "ck_tile::ABQuantGemmPipelineAgBgCrCompV3",
|
||||
}
|
||||
base_pipeline_map = {
|
||||
"compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3",
|
||||
}
|
||||
scheduler_type_map = {
|
||||
"intrawave": "ck_tile::GemmPipelineScheduler::Intrawave",
|
||||
"interwave": "ck_tile::GemmPipelineScheduler::Interwave",
|
||||
"default": "ck_tile::GemmPipelineScheduler::Default",
|
||||
}
|
||||
|
||||
instance_code = self.populate_kernel_header(kernel_name)
|
||||
instance_code += self.populate_kernel_dtype_layout()
|
||||
instance_code += self.populate_strut_begin(kernel_name)
|
||||
instance_code += self.populate_tile_config(tile_config)
|
||||
instance_code += self._populate_abquant_trait_config(trait_combo, group_size_n)
|
||||
instance_code += self._populate_abquant_initialization(
|
||||
base_pipeline_map, pipeline
|
||||
)
|
||||
instance_code += self._populate_launch_abquant(
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
)
|
||||
|
||||
# Write into a file
|
||||
simplified_name = kernel_name
|
||||
if simplified_name.startswith(f"{self.kernel_name_prefix}_"):
|
||||
simplified_name = simplified_name[len(self.kernel_name_prefix) + 1 :]
|
||||
|
||||
header_file = (
|
||||
self.working_path
|
||||
/ f"{self.kernel_name_prefix}_single_{simplified_name}.hpp"
|
||||
)
|
||||
with open(header_file, "w") as f:
|
||||
f.write(instance_code)
|
||||
|
||||
print(f"Generated {header_file}")
|
||||
|
||||
return kernel_name, instance_code
|
||||
|
||||
def populate_kernel_header(self, kernel_name):
|
||||
instance_code = f"""// Generated kernel instance for {kernel_name}
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include <utility>
|
||||
#include <tuple>
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host/kernel_launch.hpp"
|
||||
#include "ck_tile/ops/gemm.hpp"
|
||||
#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp"
|
||||
#include "ck_tile/ops/common/tensor_layout.hpp"
|
||||
#include "ck_tile/ops/epilogue/default_2d_epilogue.hpp"
|
||||
#include "ck_tile/ops/epilogue/cshuffle_epilogue.hpp"
|
||||
#include "ck_tile/ops/gemm_quant.hpp"
|
||||
#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp"
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def populate_kernel_dtype_layout(self):
|
||||
a_layout, b_layout, c_layout = _validation_utils.get_abc_layouts(self.layout)
|
||||
dtype_str = _validation_utils.get_dtype_string(self.datatype)
|
||||
c_type_str = _validation_utils.get_dtype_string(
|
||||
"fp16" if self.datatype in ["fp8", "bf8"] else self.datatype
|
||||
)
|
||||
|
||||
instance_code = f"""
|
||||
using ADataType = {dtype_str};
|
||||
using BDataType = {dtype_str};
|
||||
using AccDataType = float;
|
||||
using CDataType = {c_type_str};
|
||||
using AQDataType = float;
|
||||
using BQDataType = float;
|
||||
|
||||
using ALayout = {a_layout};
|
||||
using BLayout = {b_layout};
|
||||
using CLayout = {c_layout};
|
||||
using AQLayout = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
using BQLayout = ck_tile::tensor_layout::gemm::ColumnMajor;
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def _populate_abquant_trait_config(self, trait_combo, group_size_n):
|
||||
(
|
||||
pipeline,
|
||||
epilogue,
|
||||
scheduler,
|
||||
pad_m,
|
||||
pad_n,
|
||||
pad_k,
|
||||
a_preshuffle,
|
||||
b_preshuffle,
|
||||
) = trait_combo
|
||||
|
||||
instance_code = f"""
|
||||
|
||||
// Traits configurations
|
||||
static constexpr bool kPadM = {"true" if pad_m in [True, "true"] else "false"};
|
||||
static constexpr bool kPadN = {"true" if pad_n in [True, "true"] else "false"};
|
||||
static constexpr bool kPadK = {"true" if pad_k in [True, "true"] else "false"};
|
||||
static constexpr bool TransposeC = false;
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr bool APreshuffleQuant = {"true" if a_preshuffle in [True, "true"] else "false"};
|
||||
static constexpr bool BPreshuffleQuant = {"true" if b_preshuffle in [True, "true"] else "false"};
|
||||
static constexpr bool PreshuffleB = false;
|
||||
static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k};
|
||||
static constexpr ck_tile::index_t GroupSizeN = {group_size_n};"""
|
||||
|
||||
return instance_code
|
||||
|
||||
def _populate_abquant_initialization(self, base_pipeline_map, pipeline):
|
||||
instance_code = """
|
||||
|
||||
// Tile shape
|
||||
using TileShape = ck_tile::TileGemmShape<
|
||||
ck_tile::sequence<TileM, TileN, TileK>,
|
||||
ck_tile::sequence<WarpPerBlock_M, WarpPerBlock_N, WarpPerBlock_K>,
|
||||
ck_tile::sequence<WarpTileM, WarpTileN, WarpTileK>>;
|
||||
|
||||
// Tile partitioner
|
||||
using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner<TileShape, 8, 4>;
|
||||
|
||||
// Quantization group sizes
|
||||
using AQuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, GroupSizeK>>;
|
||||
using BQuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, GroupSizeN, GroupSizeK>>;
|
||||
|
||||
// Traits
|
||||
using Traits = ck_tile::TileGemmQuantTraits<
|
||||
kPadM, kPadN, kPadK,
|
||||
APreshuffleQuant,
|
||||
BPreshuffleQuant,
|
||||
PreshuffleB,
|
||||
ALayout, BLayout, CLayout,
|
||||
ck_tile::QuantType::ABQuantGrouped,
|
||||
AQLayout, BQLayout>;
|
||||
|
||||
// Pipeline problem (base, for hot loop detection)
|
||||
using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<
|
||||
ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
BDataType>;"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
// Base pipeline for hot loop detection
|
||||
using BaseGemmPipeline = {base_pipeline_map.get(pipeline)}<GemmPipelineProblem>;"""
|
||||
|
||||
return instance_code
|
||||
|
||||
def _populate_launch_abquant(
|
||||
self,
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
):
|
||||
"""Generate the complete launch function for ABQuant kernels."""
|
||||
instance_code = """
|
||||
|
||||
// Launch function
|
||||
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
|
||||
// Hot loop detection
|
||||
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
constexpr auto tail_number_v = tail_number_.value;
|
||||
|
||||
// Full pipeline problem with hot loop and tail number
|
||||
using PipelineProblem = ck_tile::GemmABQuantPipelineProblem<
|
||||
ADataType,
|
||||
AQDataType,
|
||||
BDataType,
|
||||
BQDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
AQuantGroupSize,
|
||||
BQuantGroupSize,
|
||||
false,
|
||||
ADataType,
|
||||
{scheduler_type_map.get(scheduler)},
|
||||
has_hot_loop_v,
|
||||
tail_number_v>;
|
||||
|
||||
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
|
||||
|
||||
instance_code += """
|
||||
|
||||
// Epilogue"""
|
||||
if epilogue == "cshuffle":
|
||||
instance_code += """
|
||||
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>, // DsDataType
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>, // DsLayout
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TileM, // kM_
|
||||
TileN, // kN_
|
||||
WarpPerBlock_M, // MWave_
|
||||
WarpPerBlock_N, // NWave_
|
||||
WarpTileM, // kMPerXdl_
|
||||
WarpTileN, // kNPerXdl_
|
||||
WarpTileK, // kKPerXdl_
|
||||
TransposeC>; // isCTransposed_
|
||||
|
||||
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;"""
|
||||
else:
|
||||
instance_code += """
|
||||
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>, // DsDataType
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>, // DsLayout
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TileM, // kM_
|
||||
TileN, // kN_
|
||||
kPadM,
|
||||
kPadN,
|
||||
WarpTileM, // kMPerXdl_
|
||||
WarpTileN, // kNPerXdl_
|
||||
WarpTileK, // kKPerXdl_
|
||||
TransposeC>; // isCTransposed_
|
||||
|
||||
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
// Kernel type
|
||||
using Kernel = ck_tile::QuantGemmKernel<
|
||||
TilePartitioner, GemmPipeline, GemmEpilogue,
|
||||
ck_tile::QuantType::ABQuantGrouped>;
|
||||
|
||||
// Kernel arguments
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
|
||||
if (!Kernel::IsSupportedArgument(kargs)) {{
|
||||
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
if(stream.log_level_ > 0) {{
|
||||
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
|
||||
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
|
||||
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
|
||||
<< std::endl;
|
||||
}}
|
||||
|
||||
// Launch kernel
|
||||
constexpr int kBlockPerCu = {k_block_per_cu};
|
||||
float ave_time = ck_tile::launch_kernel(
|
||||
stream,
|
||||
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
|
||||
|
||||
return ave_time;
|
||||
}};
|
||||
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def _generate_all_individual(self, num_workers=None):
|
||||
"""Generate individual kernel files for separate compilation with parallel processing"""
|
||||
if num_workers is None:
|
||||
num_workers = min(multiprocessing.cpu_count(), 8)
|
||||
|
||||
tile_configs = self._get_tile_configs()
|
||||
trait_combos = self._generate_trait_combinations()
|
||||
group_size_n_values = self._get_group_size_n_values()
|
||||
|
||||
work_items = []
|
||||
for tile_config in tile_configs:
|
||||
for trait_combo in trait_combos:
|
||||
for group_size_n in group_size_n_values:
|
||||
work_items.append(
|
||||
(
|
||||
tile_config,
|
||||
trait_combo,
|
||||
group_size_n,
|
||||
self.kernel_name_prefix,
|
||||
self.working_path,
|
||||
self.gpu_target,
|
||||
self.datatype,
|
||||
self.layout,
|
||||
self.config_json,
|
||||
)
|
||||
)
|
||||
|
||||
# Apply RFC-compliant sampling (Sobol + LHS + maximin)
|
||||
if self.max_instances is not None and len(work_items) > self.max_instances:
|
||||
kernel_dicts = [
|
||||
{
|
||||
"tile_config": item[0],
|
||||
"trait_combo": item[1],
|
||||
"group_size_n": item[2],
|
||||
"_work_item": item,
|
||||
}
|
||||
for item in work_items
|
||||
]
|
||||
sampled = self._apply_sampling(kernel_dicts)
|
||||
work_items = [k["_work_item"] for k in sampled]
|
||||
|
||||
print(
|
||||
f"Generating {len(work_items)} individual kernel files using {num_workers} workers..."
|
||||
)
|
||||
print(f" Tile configs: {len(tile_configs)}")
|
||||
print(f" Trait combinations: {len(trait_combos)}")
|
||||
print(f" Group size N values: {group_size_n_values}")
|
||||
print(f" Total kernels: {len(work_items)}")
|
||||
|
||||
if work_items:
|
||||
print(" First work item example:")
|
||||
tile_config, trait_combo = work_items[0][:2]
|
||||
print(f" Tile config: {tile_config}")
|
||||
print(f" Trait combo: {trait_combo[:3]}")
|
||||
|
||||
kernel_list = []
|
||||
completed = 0
|
||||
|
||||
with concurrent.futures.ProcessPoolExecutor(
|
||||
max_workers=num_workers
|
||||
) as executor:
|
||||
print(f" Submitting {len(work_items)} tasks to executor...")
|
||||
future_to_item = {
|
||||
executor.submit(_generate_single_kernel_individual, item): item
|
||||
for item in work_items
|
||||
}
|
||||
print(" All tasks submitted, waiting for completion...")
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_item):
|
||||
completed += 1
|
||||
if completed % 100 == 0 or completed == len(work_items):
|
||||
print(
|
||||
f" Progress: {completed}/{len(work_items)} kernels generated"
|
||||
)
|
||||
try:
|
||||
result = future.result()
|
||||
if result:
|
||||
kernel_list.append(result)
|
||||
except Exception as exc:
|
||||
item = future_to_item[future]
|
||||
print(f"Kernel generation failed for {item}: {exc}")
|
||||
|
||||
kernel_list.sort(key=lambda x: x[0])
|
||||
self._generate_cmake_individual_targets(kernel_list)
|
||||
|
||||
print(
|
||||
f"Generated {len(kernel_list)} individual kernel files in {self.working_path}"
|
||||
)
|
||||
|
||||
|
||||
def _generate_single_kernel_individual(work_item):
|
||||
"""Worker function to generate a single individual kernel file"""
|
||||
(
|
||||
tile_config,
|
||||
trait_combo,
|
||||
group_size_n,
|
||||
kernel_name_prefix,
|
||||
working_path,
|
||||
gpu_target,
|
||||
datatype,
|
||||
layout,
|
||||
config_json,
|
||||
) = work_item
|
||||
|
||||
builder = GemmABQuantKernelBuilder(
|
||||
kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json
|
||||
)
|
||||
|
||||
try:
|
||||
kernel_name, _ = builder._generate_kernel_instance(
|
||||
tile_config, trait_combo, group_size_n
|
||||
)
|
||||
return (kernel_name, trait_combo, tile_config)
|
||||
except Exception as e:
|
||||
print(f"Error generating individual kernel: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="GEMM ABQuant kernel instance builder")
|
||||
parser.add_argument("--working_path", required=True, help="Working directory path")
|
||||
parser.add_argument("--gpu_target", required=True, help="GPU target architecture")
|
||||
parser.add_argument(
|
||||
"--datatype",
|
||||
required=True,
|
||||
choices=["fp8", "bf8"],
|
||||
help="Data type (fp8 or bf8)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layout",
|
||||
required=True,
|
||||
choices=["rcr", "rrr", "ccr", "crr"],
|
||||
help="Matrix layout",
|
||||
)
|
||||
parser.add_argument("--config_json", help="Configuration JSON file")
|
||||
parser.add_argument("--num_workers", type=int, help="Number of parallel workers")
|
||||
parser.add_argument(
|
||||
"--gen_all_individual",
|
||||
action="store_true",
|
||||
help="Generate individual kernel files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gen_single",
|
||||
action="store_true",
|
||||
help="Generate a single kernel file",
|
||||
)
|
||||
parser.add_argument("--kernel_name", help="Kernel name for single generation")
|
||||
parser.add_argument("--tile_config", help="Tile configuration string")
|
||||
parser.add_argument("--trait_combo", help="Trait combination string")
|
||||
parser.add_argument(
|
||||
"--list_kernels",
|
||||
action="store_true",
|
||||
help="List kernel configurations without generating files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-instances",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum number of kernel instances to select via sampling",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="RNG seed for deterministic sampling; if omitted, derived from today's date",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tier",
|
||||
default=None,
|
||||
help="Sampling tier name (e.g., 'daily')",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--manifest-path",
|
||||
default=None,
|
||||
help="Directory to write chosen_instances manifest JSON",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
assert args.datatype in ["fp16", "bf16", "fp8", "bf8"], (
|
||||
f"Invalid datatype string: {args.datatype} (supported datatypes are [fp16, bf16, fp8, and bf8])"
|
||||
)
|
||||
|
||||
layout_parts = args.layout.lower()
|
||||
assert len(layout_parts) == 3, (
|
||||
f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr')"
|
||||
)
|
||||
assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], (
|
||||
f"Invalid matrix_a layout: {layout_parts[0]} or matrix_b layout: {layout_parts[1]}"
|
||||
)
|
||||
assert layout_parts[2] == "r", (
|
||||
f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' for row major)"
|
||||
)
|
||||
|
||||
kernel_name_prefix = "gemm_abquant"
|
||||
builder = GemmABQuantKernelBuilder(
|
||||
kernel_name_prefix,
|
||||
args.working_path,
|
||||
args.gpu_target,
|
||||
args.datatype,
|
||||
args.layout,
|
||||
args.config_json,
|
||||
max_instances=args.max_instances,
|
||||
seed=args.seed,
|
||||
tier=args.tier,
|
||||
manifest_path=args.manifest_path,
|
||||
)
|
||||
|
||||
if args.list_kernels:
|
||||
builder._list_kernels()
|
||||
elif args.gen_single:
|
||||
if not args.kernel_name or not args.tile_config or not args.trait_combo:
|
||||
parser.error(
|
||||
"--gen_single requires --kernel_name, --tile_config, and --trait_combo"
|
||||
)
|
||||
|
||||
# Parse tile config
|
||||
tile_parts = args.tile_config.split("_")
|
||||
tile_dims = tile_parts[0].split("x")
|
||||
warp_dims = tile_parts[1].split("x")
|
||||
warp_tile_dims = tile_parts[2].split("x")
|
||||
|
||||
tile_config = {
|
||||
"tile_m": int(tile_dims[0]),
|
||||
"tile_n": int(tile_dims[1]),
|
||||
"tile_k": int(tile_dims[2]),
|
||||
"warp_m": int(warp_dims[0]),
|
||||
"warp_n": int(warp_dims[1]),
|
||||
"warp_k": int(warp_dims[2]),
|
||||
"warp_tile_m": int(warp_tile_dims[0]),
|
||||
"warp_tile_n": int(warp_tile_dims[1]),
|
||||
"warp_tile_k": int(warp_tile_dims[2]),
|
||||
}
|
||||
|
||||
# Parse trait combo for ABQuant:
|
||||
# pipeline_epilogue_scheduler_padM_padN_padK_aPreshuffle_bPreshuffle_gsnN
|
||||
trait_parts = args.trait_combo.split("_")
|
||||
# Find the gsn part
|
||||
group_size_n = 1
|
||||
gsn_idx = None
|
||||
for i, part in enumerate(trait_parts):
|
||||
if part.startswith("gsn"):
|
||||
group_size_n = int(part[3:])
|
||||
gsn_idx = i
|
||||
break
|
||||
|
||||
if gsn_idx is not None:
|
||||
# Remove gsn from trait_parts for parsing
|
||||
trait_parts_no_gsn = trait_parts[:gsn_idx] + trait_parts[gsn_idx + 1 :]
|
||||
else:
|
||||
trait_parts_no_gsn = trait_parts
|
||||
|
||||
if len(trait_parts_no_gsn) < 8:
|
||||
parser.error(
|
||||
f"--trait_combo must have 8 underscore-separated fields + gsn suffix "
|
||||
f"(e.g. 'compv3_default_intrawave_False_False_False_False_False_gsn1'), "
|
||||
f"got {len(trait_parts_no_gsn)} field(s): '{args.trait_combo}'"
|
||||
)
|
||||
trait_combo = (
|
||||
trait_parts_no_gsn[0], # pipeline
|
||||
trait_parts_no_gsn[1], # epilogue
|
||||
trait_parts_no_gsn[2], # scheduler
|
||||
trait_parts_no_gsn[3] == "True", # pad_m
|
||||
trait_parts_no_gsn[4] == "True", # pad_n
|
||||
trait_parts_no_gsn[5] == "True", # pad_k
|
||||
trait_parts_no_gsn[6] == "True", # a_preshuffle_quant
|
||||
trait_parts_no_gsn[7] == "True", # b_preshuffle_quant
|
||||
)
|
||||
|
||||
builder._generate_kernel_instance(tile_config, trait_combo, group_size_n)
|
||||
elif args.gen_all_individual:
|
||||
builder._generate_all_individual(args.num_workers)
|
||||
else:
|
||||
parser.error(
|
||||
"Must specify one of: --list_kernels, --gen_all_individual, or --gen_single"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,336 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <functional>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <stdexcept>
|
||||
|
||||
#include "ck_tile/host/device_prop.hpp"
|
||||
#include "ck_tile/host/tensor_shuffle_utils.hpp"
|
||||
#include "ck_tile/ops/gemm_quant.hpp"
|
||||
#include "gemm_abquant_benchmark.hpp"
|
||||
|
||||
class ABQuantGemmProfiler
|
||||
{
|
||||
public:
|
||||
static ABQuantGemmProfiler& instance(Setting setting)
|
||||
{
|
||||
static ABQuantGemmProfiler instance{setting};
|
||||
return instance;
|
||||
}
|
||||
|
||||
void reset() { kernel_instances_.clear(); }
|
||||
|
||||
void benchmark(ABQuantGemmProblem& problem,
|
||||
std::function<float(const ck_tile::QuantGemmHostArgs&,
|
||||
const ck_tile::stream_config&)> kernel_func)
|
||||
{
|
||||
std::vector<std::function<std::tuple<std::string, float>(ck_tile::QuantGemmHostArgs&,
|
||||
const ck_tile::stream_config&)>>
|
||||
callables;
|
||||
|
||||
callables.push_back(
|
||||
[kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
float time = kernel_func(args, stream);
|
||||
return std::make_tuple(std::string(KERNEL_NAME), time);
|
||||
});
|
||||
|
||||
benchmark(problem, callables);
|
||||
}
|
||||
|
||||
void benchmark(ABQuantGemmProblem& problem,
|
||||
std::vector<std::function<std::tuple<std::string, float>(
|
||||
ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables)
|
||||
{
|
||||
const ALayout layout_a = ALayout{};
|
||||
const BLayout layout_b = BLayout{};
|
||||
const CLayout layout_c = CLayout{};
|
||||
const AQLayout layout_aq = AQLayout{};
|
||||
const BQLayout layout_bq = BQLayout{};
|
||||
|
||||
problem.stride_a_ = ck_tile::get_default_stride(
|
||||
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a));
|
||||
problem.stride_b_ = ck_tile::get_default_stride(
|
||||
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b));
|
||||
problem.stride_c_ = ck_tile::get_default_stride(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c));
|
||||
|
||||
// Compute AQ scale tensor dimensions: [M, K / group_size_k]
|
||||
if(problem.k_ % problem.group_size_k_ != 0)
|
||||
throw std::runtime_error("k_ must be divisible by group_size_k_");
|
||||
const ck_tile::index_t QK_A = problem.k_ / problem.group_size_k_;
|
||||
problem.stride_aq_ = ck_tile::get_default_stride(
|
||||
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq));
|
||||
|
||||
// Compute BQ scale tensor dimensions: [K / group_size_k, N / group_size_n]
|
||||
const ck_tile::index_t QK_B = problem.k_ / problem.group_size_k_;
|
||||
const ck_tile::index_t QN_B =
|
||||
(problem.group_size_n_ > 1) ? (problem.n_ / problem.group_size_n_) : problem.n_;
|
||||
problem.stride_bq_ =
|
||||
ck_tile::get_default_stride(QK_B, QN_B, problem.stride_bq_, is_row_major(layout_bq));
|
||||
|
||||
// Allocate host tensors
|
||||
ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
|
||||
ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
|
||||
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
|
||||
ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
|
||||
|
||||
// AQ scale tensor: [M, QK_A]
|
||||
ck_tile::HostTensor<AQDataType> aq_m_qk(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq)));
|
||||
|
||||
// BQ scale tensor: [QK_B, QN_B]
|
||||
ck_tile::HostTensor<BQDataType> bq_qk_n(ck_tile::host_tensor_descriptor(
|
||||
QK_B, QN_B, problem.stride_bq_, is_row_major(layout_bq)));
|
||||
|
||||
// Initialize tensors
|
||||
if(setting_.init_method_ == 0)
|
||||
{
|
||||
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
|
||||
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
|
||||
ck_tile::FillUniformDistribution<AQDataType>{0.5f, 1.5f}(aq_m_qk);
|
||||
ck_tile::FillUniformDistribution<BQDataType>{0.5f, 1.5f}(bq_qk_n);
|
||||
}
|
||||
else if(setting_.init_method_ == 1)
|
||||
{
|
||||
ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
|
||||
ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
|
||||
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
|
||||
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
|
||||
}
|
||||
else if(setting_.init_method_ == 2)
|
||||
{
|
||||
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(1)}(a_m_k);
|
||||
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(1)}(b_k_n);
|
||||
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
|
||||
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
|
||||
}
|
||||
else
|
||||
{
|
||||
a_m_k.SetZero();
|
||||
b_k_n.SetZero();
|
||||
aq_m_qk.SetZero();
|
||||
bq_qk_n.SetZero();
|
||||
}
|
||||
|
||||
// Allocate device memory
|
||||
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem aq_dev_buf(aq_m_qk.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem bq_dev_buf(bq_qk_n.get_element_space_size_in_bytes());
|
||||
|
||||
a_m_k_dev_buf.ToDevice(a_m_k.data());
|
||||
b_k_n_dev_buf.ToDevice(b_k_n.data());
|
||||
|
||||
// Shuffle AQ data if preshuffle is enabled
|
||||
if constexpr(SelectedKernel::APreshuffleQuant)
|
||||
{
|
||||
constexpr int block_aq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK;
|
||||
ck_tile::HostTensor<AQDataType> aq_shuffled = ck_tile::shuffle_aq(&aq_m_qk, block_aq_k);
|
||||
aq_dev_buf.ToDevice(aq_shuffled.data());
|
||||
}
|
||||
else
|
||||
{
|
||||
aq_dev_buf.ToDevice(aq_m_qk.data());
|
||||
}
|
||||
|
||||
// Shuffle BQ data if preshuffle is enabled
|
||||
if constexpr(SelectedKernel::BPreshuffleQuant)
|
||||
{
|
||||
constexpr int block_bq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK;
|
||||
ck_tile::HostTensor<BQDataType> bq_shuffled = ck_tile::shuffle_bq(&bq_qk_n, block_bq_k);
|
||||
bq_dev_buf.ToDevice(bq_shuffled.data());
|
||||
}
|
||||
else
|
||||
{
|
||||
bq_dev_buf.ToDevice(bq_qk_n.data());
|
||||
}
|
||||
|
||||
c_m_n_dev_buf.SetZero();
|
||||
c_m_n_dev_result.SetZero();
|
||||
|
||||
// Build QuantGemmHostArgs
|
||||
ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
|
||||
b_k_n_dev_buf.GetDeviceBuffer(),
|
||||
c_m_n_dev_buf.GetDeviceBuffer(),
|
||||
aq_dev_buf.GetDeviceBuffer(),
|
||||
bq_dev_buf.GetDeviceBuffer(),
|
||||
problem.split_k_,
|
||||
problem.m_,
|
||||
problem.n_,
|
||||
problem.k_,
|
||||
QK_A,
|
||||
QK_B,
|
||||
problem.stride_a_,
|
||||
problem.stride_b_,
|
||||
problem.stride_c_,
|
||||
problem.stride_aq_,
|
||||
problem.stride_bq_);
|
||||
|
||||
// Host reference for verification
|
||||
ck_tile::HostTensor<CDataType> c_m_n_host_result(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
|
||||
|
||||
if(setting_.verify_)
|
||||
{
|
||||
abquant_gemm_host_reference(
|
||||
setting_.verify_, a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result);
|
||||
}
|
||||
|
||||
// Run kernel(s)
|
||||
for(auto& callable : callables)
|
||||
{
|
||||
auto kernel_run_result = callable(gemm_args,
|
||||
ck_tile::stream_config{nullptr,
|
||||
true,
|
||||
setting_.log_,
|
||||
setting_.n_warmup_,
|
||||
setting_.n_repeat_,
|
||||
setting_.is_gpu_timer_,
|
||||
setting_.flush_cache_,
|
||||
setting_.rotating_count_});
|
||||
process_result(problem,
|
||||
QK_A,
|
||||
QK_B,
|
||||
c_m_n_dev_buf,
|
||||
c_m_n_host_result,
|
||||
c_m_n_dev_result,
|
||||
kernel_run_result);
|
||||
}
|
||||
}
|
||||
|
||||
void process_result(const ABQuantGemmProblem& problem,
|
||||
ck_tile::index_t QK_A,
|
||||
ck_tile::index_t QK_B,
|
||||
ck_tile::DeviceMem& c_m_n_dev_buf,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
const std::tuple<std::string, float>& kernel_run_result)
|
||||
{
|
||||
auto [name, avg_time] = kernel_run_result;
|
||||
|
||||
KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}};
|
||||
|
||||
// Compute performance metrics
|
||||
std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_;
|
||||
std::size_t num_byte =
|
||||
sizeof(ADataType) * problem.m_ * problem.k_ +
|
||||
sizeof(BDataType) * problem.n_ * problem.k_ + sizeof(AQDataType) * problem.m_ * QK_A +
|
||||
sizeof(BQDataType) * QK_B *
|
||||
(problem.group_size_n_ > 1 ? problem.n_ / problem.group_size_n_ : problem.n_) +
|
||||
sizeof(CDataType) * problem.m_ * problem.n_;
|
||||
|
||||
kernel_instance.perf_result_.latency_ = avg_time;
|
||||
kernel_instance.perf_result_.tflops_ = static_cast<float>(flop) / 1.E9 / avg_time;
|
||||
kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time;
|
||||
|
||||
if(setting_.log_ > 0 && !setting_.json_output_)
|
||||
{
|
||||
std::cout << kernel_instance << std::endl;
|
||||
}
|
||||
|
||||
// Verify result
|
||||
bool verified_correct = true;
|
||||
if(setting_.verify_)
|
||||
{
|
||||
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
|
||||
verified_correct = compare_abquant(
|
||||
name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result);
|
||||
}
|
||||
|
||||
if(verified_correct)
|
||||
{
|
||||
kernel_instances_.emplace_back(kernel_instance);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Verification failed, skip kernel: " << name << std::endl;
|
||||
}
|
||||
|
||||
c_m_n_dev_buf.SetZero();
|
||||
c_m_n_dev_result.SetZero();
|
||||
}
|
||||
|
||||
KernelInstance select_best_instance(Metric metric)
|
||||
{
|
||||
if(kernel_instances_.empty())
|
||||
throw std::runtime_error("Empty instances");
|
||||
|
||||
auto kernel_instance = *std::max_element(kernel_instances_.begin(),
|
||||
kernel_instances_.end(),
|
||||
[metric](const auto& a, const auto& b) {
|
||||
return PerformanceResult::compare(
|
||||
b.perf_result_, a.perf_result_, metric);
|
||||
});
|
||||
|
||||
if(setting_.json_output_)
|
||||
{
|
||||
std::cout << kernel_instance << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "**********************************" << std::endl;
|
||||
std::cout << "According to given metrics: " << get_metric_name(metric) << "\n"
|
||||
<< "Current kernel performance is: " << kernel_instance << std::endl;
|
||||
std::cout << "**********************************" << std::endl;
|
||||
}
|
||||
|
||||
if(!setting_.csv_filename_.empty())
|
||||
{
|
||||
std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app);
|
||||
|
||||
if(!file.is_open())
|
||||
{
|
||||
std::cerr << "Warning: Failed to open CSV file for writing." << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
if(file.tellp() == 0)
|
||||
{
|
||||
file << "rocm_version,device_name,"
|
||||
<< "split_k,m,n,k,stride_a,stride_b,stride_c,stride_aq,stride_bq,"
|
||||
<< "group_size_k,group_size_n,"
|
||||
<< "dtype_a,dtype_b,dtype_aq,dtype_bq,dtype_acc,dtype_c,"
|
||||
<< "layout_a,layout_b,layout_c," << "name,"
|
||||
<< "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n";
|
||||
}
|
||||
|
||||
const auto& prob = kernel_instance.problem_;
|
||||
const auto& perf = kernel_instance.perf_result_;
|
||||
|
||||
file << get_rocm_version() << "," << ck_tile::get_device_name() << ","
|
||||
<< prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << ","
|
||||
<< prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << ","
|
||||
<< prob.stride_aq_ << "," << prob.stride_bq_ << "," << prob.group_size_k_
|
||||
<< "," << prob.group_size_n_ << "," << prob.dtype_a_ << "," << prob.dtype_b_
|
||||
<< "," << prob.dtype_aq_ << "," << prob.dtype_bq_ << "," << prob.dtype_acc_
|
||||
<< "," << prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_
|
||||
<< "," << prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed
|
||||
<< std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << ","
|
||||
<< perf.bandwidth_ << "," << get_metric_name(metric) << "\n";
|
||||
}
|
||||
}
|
||||
|
||||
return kernel_instance;
|
||||
}
|
||||
|
||||
ABQuantGemmProfiler(const ABQuantGemmProfiler&) = delete;
|
||||
ABQuantGemmProfiler& operator=(const ABQuantGemmProfiler&) = delete;
|
||||
|
||||
private:
|
||||
~ABQuantGemmProfiler() { kernel_instances_.clear(); }
|
||||
ABQuantGemmProfiler(Setting setting) : setting_(setting) {}
|
||||
|
||||
Setting setting_;
|
||||
std::vector<KernelInstance> kernel_instances_;
|
||||
};
|
||||
@@ -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()
|
||||
297
tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt
Normal file
297
tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt
Normal file
@@ -0,0 +1,297 @@
|
||||
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
set(GEMM_AQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM AQuant (semicolon-separated)")
|
||||
set(GEMM_AQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM AQuant (semicolon-separated)")
|
||||
set(GEMM_AQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)")
|
||||
option(ENABLE_CCACHE_GEMM_AQUANT "Enable ccache for GEMM AQuant ops compilation" OFF)
|
||||
|
||||
# Store the directory path for use in functions
|
||||
set(GEMM_AQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR})
|
||||
|
||||
# Function to create individual GEMM AQuant targets
|
||||
function(create_individual_gemm_aquant_target datatype layout trait tile_config config_json)
|
||||
if(NOT GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL)
|
||||
message(WARNING "Skipping individual GEMM AQuant target ${datatype}_${layout}_${trait}_${tile_config}: No supported GPU targets")
|
||||
return()
|
||||
endif()
|
||||
|
||||
# Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k
|
||||
string(REPLACE "_" ";" config_groups ${tile_config})
|
||||
list(GET config_groups 0 tile_dims)
|
||||
list(GET config_groups 1 warp_dims)
|
||||
list(GET config_groups 2 warp_tile_dims)
|
||||
|
||||
string(REPLACE "x" ";" tile_parts ${tile_dims})
|
||||
list(GET tile_parts 0 tile_m)
|
||||
list(GET tile_parts 1 tile_n)
|
||||
list(GET tile_parts 2 tile_k)
|
||||
|
||||
string(REPLACE "x" ";" warp_parts ${warp_dims})
|
||||
list(GET warp_parts 0 warp_m)
|
||||
list(GET warp_parts 1 warp_n)
|
||||
list(GET warp_parts 2 warp_k)
|
||||
|
||||
string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims})
|
||||
list(GET warp_tile_parts 0 warp_tile_m)
|
||||
list(GET warp_tile_parts 1 warp_tile_n)
|
||||
list(GET warp_tile_parts 2 warp_tile_k)
|
||||
|
||||
set(target_name "benchmark_gemm_aquant_${datatype}_${layout}_${trait}_${tile_config}")
|
||||
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
|
||||
|
||||
# Generate the single instance header for this kernel
|
||||
set(instance_header "${working_path}/gemm_aquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp")
|
||||
|
||||
# Add custom command to generate the header file at build time
|
||||
add_custom_command(
|
||||
OUTPUT ${instance_header}
|
||||
COMMAND ${Python3_EXECUTABLE} ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_instance_builder.py
|
||||
--working_path ${working_path}
|
||||
--datatype ${datatype}
|
||||
--layout ${layout}
|
||||
--config_json ${config_json}
|
||||
--gen_single
|
||||
--kernel_name "gemm_aquant_${datatype}_${layout}_${trait}_${tile_config}"
|
||||
--tile_config "${tile_config}"
|
||||
--trait_combo "${trait}"
|
||||
--gpu_target "${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}"
|
||||
DEPENDS ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_instance_builder.py ${config_json}
|
||||
COMMENT "Generating ${instance_header}"
|
||||
)
|
||||
|
||||
# Create the executable
|
||||
add_executable(${target_name}
|
||||
EXCLUDE_FROM_ALL
|
||||
${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_benchmark_single.cpp
|
||||
${instance_header}
|
||||
)
|
||||
|
||||
# Set GPU architectures
|
||||
set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL})
|
||||
|
||||
# Set compile definitions
|
||||
target_compile_definitions(${target_name} PRIVATE
|
||||
GEMM_AQUANT_SINGLE_INSTANCE_HPP="${instance_header}"
|
||||
)
|
||||
|
||||
# Include directories
|
||||
target_include_directories(${target_name} PRIVATE
|
||||
${GEMM_AQUANT_SOURCE_DIR}
|
||||
${working_path}
|
||||
)
|
||||
|
||||
# Compile options
|
||||
target_compile_options(${target_name} PRIVATE
|
||||
-Wno-undefined-func-template
|
||||
-Wno-float-equal
|
||||
--offload-compress
|
||||
-include ${instance_header}
|
||||
)
|
||||
|
||||
# Add to collection targets
|
||||
add_dependencies(benchmark_gemm_aquant_all ${target_name})
|
||||
add_dependencies(benchmark_gemm_aquant_${datatype} ${target_name})
|
||||
add_dependencies(benchmark_gemm_aquant_${layout} ${target_name})
|
||||
add_dependencies(benchmark_gemm_aquant_${datatype}_${layout} ${target_name})
|
||||
|
||||
# Add to pipeline-specific and scheduler-specific targets
|
||||
string(REPLACE "_" ";" trait_parts ${trait})
|
||||
list(GET trait_parts 0 pipeline)
|
||||
list(GET trait_parts 1 epilogue)
|
||||
list(GET trait_parts 2 scheduler)
|
||||
|
||||
add_dependencies(benchmark_gemm_aquant_${pipeline}_pipeline ${target_name})
|
||||
add_dependencies(benchmark_gemm_aquant_${epilogue}_epilogue ${target_name})
|
||||
add_dependencies(benchmark_gemm_aquant_${scheduler}_scheduler ${target_name})
|
||||
endfunction()
|
||||
|
||||
# Function to build individual GEMM AQuant targets
|
||||
function(build_individual_gemm_aquant_targets datatype layout)
|
||||
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
|
||||
|
||||
# Choose config file (priority: env var > cmake var > default)
|
||||
if(DEFINED ENV{GEMM_AQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_AQUANT_CONFIG_FILE}" STREQUAL "")
|
||||
set(config_filename "$ENV{GEMM_AQUANT_CONFIG_FILE}")
|
||||
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}")
|
||||
message(VERBOSE " Using config from environment variable: ${config_filename}")
|
||||
elseif(NOT "${GEMM_AQUANT_CONFIG_FILE}" STREQUAL "")
|
||||
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_AQUANT_CONFIG_FILE}")
|
||||
message(VERBOSE " Using custom config: ${GEMM_AQUANT_CONFIG_FILE}")
|
||||
else()
|
||||
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json")
|
||||
message(VERBOSE " Using default config for layout ${layout}")
|
||||
endif()
|
||||
|
||||
if(NOT EXISTS ${json_blob})
|
||||
message(FATAL_ERROR "Config file not found: ${json_blob}")
|
||||
endif()
|
||||
|
||||
# Create working directory
|
||||
file(MAKE_DIRECTORY ${working_path})
|
||||
|
||||
message(VERBOSE "Generating individual AQuant kernels for ${datatype} ${layout}...")
|
||||
|
||||
# Build sampling arguments
|
||||
set(extra_list_args "")
|
||||
if(NOT "${GEMM_AQUANT_MAX_INSTANCES}" STREQUAL "")
|
||||
list(APPEND extra_list_args --max-instances ${GEMM_AQUANT_MAX_INSTANCES})
|
||||
endif()
|
||||
if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
|
||||
list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER})
|
||||
list(APPEND extra_list_args --manifest-path ${working_path})
|
||||
endif()
|
||||
if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "")
|
||||
list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED})
|
||||
endif()
|
||||
|
||||
# List kernels
|
||||
message(VERBOSE " Listing AQuant kernel configurations for ${datatype} ${layout}...")
|
||||
execute_process(
|
||||
COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_aquant_instance_builder.py
|
||||
--working_path ${working_path}
|
||||
--datatype ${datatype}
|
||||
--layout ${layout}
|
||||
--config_json ${json_blob}
|
||||
--gpu_target "${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}"
|
||||
--list_kernels
|
||||
${extra_list_args}
|
||||
WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
|
||||
RESULT_VARIABLE ret
|
||||
OUTPUT_VARIABLE list_output
|
||||
ERROR_VARIABLE list_error
|
||||
)
|
||||
|
||||
if(NOT ret EQUAL 0)
|
||||
message(FATAL_ERROR "Failed to list AQuant kernels for ${datatype} ${layout}: ${list_error}")
|
||||
endif()
|
||||
|
||||
# Read kernel count
|
||||
if(EXISTS ${working_path}/gemm_aquant_kernel_count.txt)
|
||||
file(READ ${working_path}/gemm_aquant_kernel_count.txt kernel_count)
|
||||
string(STRIP "${kernel_count}" kernel_count)
|
||||
message(VERBOSE " Found ${kernel_count} AQuant kernel configurations")
|
||||
else()
|
||||
message(FATAL_ERROR "Kernel count file not found")
|
||||
endif()
|
||||
|
||||
# Read kernel list and create targets
|
||||
if(EXISTS ${working_path}/gemm_aquant_kernel_list.txt)
|
||||
file(STRINGS ${working_path}/gemm_aquant_kernel_list.txt kernel_lines)
|
||||
foreach(line IN LISTS kernel_lines)
|
||||
string(REPLACE "|" ";" parts "${line}")
|
||||
list(GET parts 0 kernel_name)
|
||||
list(GET parts 1 tile_config)
|
||||
list(GET parts 2 trait_combo)
|
||||
|
||||
create_individual_gemm_aquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}")
|
||||
endforeach()
|
||||
else()
|
||||
message(FATAL_ERROR "Kernel list file not found")
|
||||
endif()
|
||||
endfunction()
|
||||
|
||||
# Main build logic
|
||||
message(VERBOSE "=== Starting Tile Engine GEMM AQuant Configuration ===")
|
||||
message(VERBOSE "GEMM_AQUANT_DATATYPE: ${GEMM_AQUANT_DATATYPE}")
|
||||
message(VERBOSE "GEMM_AQUANT_LAYOUT: ${GEMM_AQUANT_LAYOUT}")
|
||||
message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
|
||||
|
||||
# Filter GPU targets to supported architectures
|
||||
set(GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL "")
|
||||
set(DESIRED_TARGETS "gfx90a;gfx942;gfx950")
|
||||
|
||||
foreach(target IN LISTS SUPPORTED_GPU_TARGETS)
|
||||
if(target IN_LIST DESIRED_TARGETS)
|
||||
list(APPEND GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL ${target})
|
||||
message(VERBOSE " Adding GPU target: ${target}")
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
# Skip build if no matching targets found
|
||||
if(NOT GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL)
|
||||
message(WARNING "Skipping Tile Engine GEMM AQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
|
||||
else()
|
||||
message(VERBOSE "Building individual GEMM AQuant targets for GPU targets: ${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}")
|
||||
|
||||
# Enable parallel compilation optimizations (directory-scoped to avoid overriding parent)
|
||||
set_property(DIRECTORY PROPERTY JOB_POOLS
|
||||
compile_heavy=4 # Limit heavy compilations to prevent OOM
|
||||
compile_normal=16 # Allow more parallel normal compilations
|
||||
)
|
||||
|
||||
# Enable compiler cache if requested
|
||||
if(ENABLE_CCACHE_GEMM_AQUANT)
|
||||
find_program(CCACHE_PROGRAM ccache)
|
||||
if(CCACHE_PROGRAM)
|
||||
set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
|
||||
message(VERBOSE "Using ccache for faster compilation")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Create master collection target
|
||||
add_custom_target(benchmark_gemm_aquant_all)
|
||||
|
||||
# Create datatype collection targets
|
||||
foreach(dt IN LISTS GEMM_AQUANT_DATATYPE)
|
||||
add_custom_target(benchmark_gemm_aquant_${dt})
|
||||
endforeach()
|
||||
|
||||
# Create layout collection targets
|
||||
foreach(l IN LISTS GEMM_AQUANT_LAYOUT)
|
||||
add_custom_target(benchmark_gemm_aquant_${l})
|
||||
endforeach()
|
||||
|
||||
# Create combined collection targets
|
||||
foreach(dt IN LISTS GEMM_AQUANT_DATATYPE)
|
||||
foreach(l IN LISTS GEMM_AQUANT_LAYOUT)
|
||||
add_custom_target(benchmark_gemm_aquant_${dt}_${l})
|
||||
endforeach()
|
||||
endforeach()
|
||||
|
||||
# Create pipeline-specific collection targets
|
||||
set(GEMM_AQUANT_PIPELINES "mem;compv3")
|
||||
foreach(pipeline IN LISTS GEMM_AQUANT_PIPELINES)
|
||||
add_custom_target(benchmark_gemm_aquant_${pipeline}_pipeline)
|
||||
endforeach()
|
||||
|
||||
# Create epilogue-specific collection targets
|
||||
set(GEMM_AQUANT_EPILOGUES "default;cshuffle")
|
||||
foreach(epilogue IN LISTS GEMM_AQUANT_EPILOGUES)
|
||||
add_custom_target(benchmark_gemm_aquant_${epilogue}_epilogue)
|
||||
endforeach()
|
||||
|
||||
# Create scheduler-specific collection targets
|
||||
set(GEMM_AQUANT_SCHEDULERS "intrawave;interwave")
|
||||
foreach(scheduler IN LISTS GEMM_AQUANT_SCHEDULERS)
|
||||
add_custom_target(benchmark_gemm_aquant_${scheduler}_scheduler)
|
||||
endforeach()
|
||||
|
||||
# Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that
|
||||
# sampling fires per-combo rather than being a single cap.
|
||||
# CMake integer division truncates; clamp the result to at least 1 so that a
|
||||
# small total budget (e.g. 5 instances across 8 combos) does not silently
|
||||
# produce a per-combo budget of 0 and skip all kernels.
|
||||
if(NOT "${GEMM_AQUANT_MAX_INSTANCES}" STREQUAL "")
|
||||
list(LENGTH GEMM_AQUANT_DATATYPE _aq_n_dt)
|
||||
list(LENGTH GEMM_AQUANT_LAYOUT _aq_n_lay)
|
||||
math(EXPR _aq_n_combos "${_aq_n_dt} * ${_aq_n_lay}")
|
||||
if(_aq_n_combos GREATER 0)
|
||||
math(EXPR GEMM_AQUANT_MAX_INSTANCES
|
||||
"${GEMM_AQUANT_MAX_INSTANCES} / ${_aq_n_combos}")
|
||||
if(GEMM_AQUANT_MAX_INSTANCES EQUAL 0)
|
||||
set(GEMM_AQUANT_MAX_INSTANCES 1)
|
||||
message(WARNING " gemm_aquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_aq_n_combos} combos)")
|
||||
else()
|
||||
message(STATUS " gemm_aquant: per-combo budget = ${GEMM_AQUANT_MAX_INSTANCES} (${_aq_n_combos} combos)")
|
||||
endif()
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Build individual targets for each datatype/layout combination
|
||||
foreach(dt IN LISTS GEMM_AQUANT_DATATYPE)
|
||||
foreach(l IN LISTS GEMM_AQUANT_LAYOUT)
|
||||
build_individual_gemm_aquant_targets(${dt} ${l})
|
||||
endforeach()
|
||||
endforeach()
|
||||
endif()
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <fstream>
|
||||
#include <stdexcept>
|
||||
#include <iomanip>
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "gemm_aquant_common.hpp"
|
||||
|
||||
// Data types and Layouts are defined by the generated kernel headers:
|
||||
// ADataType, BDataType, AQDataType, AccDataType, CDataType
|
||||
// ALayout, BLayout, CLayout, AQLayout
|
||||
|
||||
enum class Metric
|
||||
{
|
||||
LATENCY = 0,
|
||||
TFLOPS = 1,
|
||||
BANDWIDTH = 2
|
||||
};
|
||||
|
||||
inline constexpr auto get_metric_name(Metric m)
|
||||
{
|
||||
switch(m)
|
||||
{
|
||||
case Metric::LATENCY: return "latency";
|
||||
case Metric::TFLOPS: return "tflops";
|
||||
case Metric::BANDWIDTH: return "bandwidth";
|
||||
default: throw std::invalid_argument("Unsupported metric type");
|
||||
}
|
||||
}
|
||||
|
||||
struct AQuantGemmProblem
|
||||
{
|
||||
int split_k_;
|
||||
int m_, n_, k_;
|
||||
int stride_a_, stride_b_, stride_c_;
|
||||
int stride_aq_;
|
||||
int group_size_k_;
|
||||
|
||||
std::string dtype_a_, dtype_b_, dtype_aq_, dtype_acc_, dtype_c_;
|
||||
std::string layout_a_, layout_b_, layout_c_;
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const AQuantGemmProblem& problem)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"split_k\":" << problem.split_k_ << ",\n"
|
||||
<< " \"m\":" << problem.m_ << ",\n"
|
||||
<< " \"n\":" << problem.n_ << ",\n"
|
||||
<< " \"k\":" << problem.k_ << ",\n"
|
||||
<< " \"stride_a\":" << problem.stride_a_ << ",\n"
|
||||
<< " \"stride_b\":" << problem.stride_b_ << ",\n"
|
||||
<< " \"stride_c\":" << problem.stride_c_ << ",\n"
|
||||
<< " \"stride_aq\":" << problem.stride_aq_ << ",\n"
|
||||
<< " \"group_size_k\":" << problem.group_size_k_ << ",\n"
|
||||
<< " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n"
|
||||
<< " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n"
|
||||
<< " \"dtype_aq\":\"" << problem.dtype_aq_ << "\",\n"
|
||||
<< " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n"
|
||||
<< " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n"
|
||||
<< " \"layout_a\":\"" << problem.layout_a_ << "\",\n"
|
||||
<< " \"layout_b\":\"" << problem.layout_b_ << "\",\n"
|
||||
<< " \"layout_c\":\"" << problem.layout_c_ << "\"\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct PerformanceResult
|
||||
{
|
||||
double latency_;
|
||||
double tflops_;
|
||||
double bandwidth_;
|
||||
|
||||
static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m)
|
||||
{
|
||||
switch(m)
|
||||
{
|
||||
case Metric::LATENCY: return a.latency_ < b.latency_;
|
||||
case Metric::TFLOPS: return a.tflops_ > b.tflops_;
|
||||
case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_;
|
||||
default: throw std::invalid_argument("Unsupported metric type");
|
||||
}
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_
|
||||
<< ",\n"
|
||||
<< " \"tflops(TFlops)\": " << result.tflops_ << ",\n"
|
||||
<< " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct KernelInstance
|
||||
{
|
||||
std::string name_;
|
||||
AQuantGemmProblem problem_;
|
||||
PerformanceResult perf_result_;
|
||||
|
||||
static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m)
|
||||
{
|
||||
return PerformanceResult::compare(a.perf_result_, b.perf_result_, m);
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"name\": \"" << obj.name_ << "\",\n"
|
||||
<< " \"problem\": " << obj.problem_ << ",\n"
|
||||
<< " \"perf_result\": " << obj.perf_result_ << "\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct Setting
|
||||
{
|
||||
int n_warmup_;
|
||||
int n_repeat_;
|
||||
bool is_gpu_timer_;
|
||||
int verify_;
|
||||
int init_method_;
|
||||
bool log_;
|
||||
std::string csv_filename_;
|
||||
bool flush_cache_;
|
||||
int rotating_count_;
|
||||
bool json_output_;
|
||||
};
|
||||
|
||||
inline std::string get_rocm_version()
|
||||
{
|
||||
// Try the legacy path first, then the path used by newer ROCm installs.
|
||||
for(const char* path : {"/opt/rocm/.info/version", "/opt/rocm/lib/rocm_version.h"})
|
||||
{
|
||||
std::ifstream version_file(path);
|
||||
if(version_file.is_open())
|
||||
{
|
||||
std::string version;
|
||||
std::getline(version_file, version);
|
||||
if(!version.empty())
|
||||
return version;
|
||||
}
|
||||
}
|
||||
return "Unknown";
|
||||
}
|
||||
|
||||
template <typename ADataType_,
|
||||
typename AQDataType_,
|
||||
typename BDataType_,
|
||||
typename AccDataType_,
|
||||
typename CDataType_>
|
||||
auto calculate_rtol_atol_aquant(const ck_tile::index_t K,
|
||||
const ck_tile::index_t kbatch,
|
||||
const float max_accumulated_value)
|
||||
{
|
||||
using ComputeType =
|
||||
std::conditional_t<sizeof(ADataType_) < sizeof(BDataType_), ADataType_, BDataType_>;
|
||||
const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType_, AccDataType_>(
|
||||
ck_tile::integer_divide_ceil(K, kbatch));
|
||||
const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType_, AccDataType_>(
|
||||
max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
|
||||
const auto rtol_split_k =
|
||||
ck_tile::get_relative_threshold<CDataType_, CDataType_, CDataType_>(kbatch);
|
||||
const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType_, CDataType_, CDataType_>(
|
||||
max_accumulated_value, kbatch);
|
||||
return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
|
||||
}
|
||||
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions"
|
||||
/// @brief Compare device and host results for AQuant GEMM
|
||||
bool compare_aquant(std::string instanceName,
|
||||
ck_tile::index_t K,
|
||||
ck_tile::index_t kbatch,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
|
||||
{
|
||||
const float max_accumulated_value =
|
||||
std::abs(static_cast<float>(*std::max_element(c_m_n_host_result.mData.begin(),
|
||||
c_m_n_host_result.mData.end(),
|
||||
[](const auto& a, const auto& b) {
|
||||
return std::abs(static_cast<float>(a)) <
|
||||
std::abs(static_cast<float>(b));
|
||||
})));
|
||||
const auto rtol_atol =
|
||||
calculate_rtol_atol_aquant<ADataType, AQDataType, BDataType, AccDataType, CDataType>(
|
||||
K, kbatch, max_accumulated_value);
|
||||
bool pass = ck_tile::check_err(c_m_n_dev_result,
|
||||
c_m_n_host_result,
|
||||
"Error: Incorrect results!",
|
||||
rtol_atol.at(ck_tile::number<0>{}),
|
||||
rtol_atol.at(ck_tile::number<1>{}));
|
||||
|
||||
std::cerr << "For " << instanceName << " Relative error threshold is "
|
||||
<< rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
|
||||
<< rtol_atol.at(ck_tile::number<1>{}) << std::endl;
|
||||
std::cerr << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;
|
||||
|
||||
return pass;
|
||||
}
|
||||
|
||||
/// @brief CPU reference implementation for AQuant GEMM
|
||||
void aquant_gemm_host_reference(int verify,
|
||||
ck_tile::HostTensor<ADataType>& a_m_k,
|
||||
ck_tile::HostTensor<AQDataType>& aq_m_qk,
|
||||
ck_tile::HostTensor<BDataType>& b_k_n,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
|
||||
{
|
||||
if(verify == 1)
|
||||
{
|
||||
c_m_n_host_result.SetZero();
|
||||
using QuantGroupSize = typename SelectedKernel::QuantGroupSize;
|
||||
ck_tile::reference_gemm_quant<ADataType,
|
||||
AQDataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
QuantGroupSize,
|
||||
true /* aquant */>(a_m_k, aq_m_qk, b_k_n, c_m_n_host_result);
|
||||
}
|
||||
}
|
||||
#pragma clang diagnostic pop
|
||||
@@ -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())
|
||||
@@ -0,0 +1,153 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#include <iostream>
|
||||
#include <functional>
|
||||
#include <tuple>
|
||||
#include <exception>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "gemm_aquant_profiler.hpp"
|
||||
#include "gemm_aquant_common.hpp"
|
||||
|
||||
// The kernel header is included via the compile command line with -include flag
|
||||
// It defines SelectedKernel struct, KERNEL_NAME, and type aliases:
|
||||
// ADataType, BDataType, AQDataType, AccDataType, CDataType
|
||||
// ALayout, BLayout, CLayout, AQLayout
|
||||
|
||||
inline auto create_args(int argc, char* argv[])
|
||||
{
|
||||
ck_tile::ArgParser arg_parser;
|
||||
arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.")
|
||||
.insert("n", "4096", "The value for n dimension. Default is 4096.")
|
||||
.insert("k", "2048", "The value for k dimension. Default is 2048.")
|
||||
.insert("stride_a", "0", "The stride value for tensor A. Default is 0.")
|
||||
.insert("stride_b", "0", "The stride value for tensor B. Default is 0.")
|
||||
.insert("stride_c", "0", "The stride value for tensor C. Default is 0.")
|
||||
.insert("split_k", "1", "The split value for k dimension. Default is 1.")
|
||||
.insert("group_size_k",
|
||||
std::to_string(SelectedKernel::GroupSizeK),
|
||||
"The quantization group size along K. Default matches kernel config.")
|
||||
.insert("verify",
|
||||
"1",
|
||||
"The type of validation. Set to 0 for no validation, 1 for validation on CPU. "
|
||||
"Default is 1, CPU validation.")
|
||||
.insert("log",
|
||||
"false",
|
||||
"Whether output kernel instance information or not. Possible values are true or "
|
||||
"false. Default is false")
|
||||
.insert(
|
||||
"warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.")
|
||||
.insert(
|
||||
"repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.")
|
||||
.insert("timer",
|
||||
"true",
|
||||
"Whether if the timer is gpu timer or not. Possible values are false or true. "
|
||||
"Default is true.")
|
||||
.insert("init",
|
||||
"0",
|
||||
"The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 "
|
||||
"for constant(1). Default is 0, random.")
|
||||
.insert("flush_cache",
|
||||
"true",
|
||||
"To flush cache, possible values are true or false. Default is true.")
|
||||
.insert(
|
||||
"rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.")
|
||||
.insert("metric",
|
||||
"0",
|
||||
"Metric with which to measure kernel performance. Set to 0 for latency, 1 for "
|
||||
"tflops, or 2 for bandwidth. Default is 0, latency.")
|
||||
.insert("csv_filename",
|
||||
"",
|
||||
"The filename of benchmark result. Default is empty (no CSV output).")
|
||||
.insert("json_output",
|
||||
"false",
|
||||
"Whether to output results in JSON format only. Possible values are true or false. "
|
||||
"Default is false");
|
||||
|
||||
bool result = arg_parser.parse(argc, argv);
|
||||
return std::make_tuple(result, arg_parser);
|
||||
}
|
||||
|
||||
void benchmark_single(const ck_tile::ArgParser& arg_parser)
|
||||
{
|
||||
std::string dtype_a = DataTypeTraits<ADataType>::name;
|
||||
std::string dtype_b = DataTypeTraits<BDataType>::name;
|
||||
std::string dtype_aq = DataTypeTraits<AQDataType>::name;
|
||||
std::string dtype_acc = DataTypeTraits<AccDataType>::name;
|
||||
std::string dtype_c = DataTypeTraits<CDataType>::name;
|
||||
|
||||
std::string layout_a = ALayout::name;
|
||||
std::string layout_b = BLayout::name;
|
||||
std::string layout_c = CLayout::name;
|
||||
|
||||
int group_size_k = arg_parser.get_int("group_size_k");
|
||||
|
||||
AQuantGemmProblem problem{arg_parser.get_int("split_k"),
|
||||
arg_parser.get_int("m"),
|
||||
arg_parser.get_int("n"),
|
||||
arg_parser.get_int("k"),
|
||||
arg_parser.get_int("stride_a"),
|
||||
arg_parser.get_int("stride_b"),
|
||||
arg_parser.get_int("stride_c"),
|
||||
0, // stride_aq computed by profiler
|
||||
group_size_k,
|
||||
dtype_a,
|
||||
dtype_b,
|
||||
dtype_aq,
|
||||
dtype_acc,
|
||||
dtype_c,
|
||||
layout_a,
|
||||
layout_b,
|
||||
layout_c};
|
||||
|
||||
Setting setting{arg_parser.get_int("warmup"),
|
||||
arg_parser.get_int("repeat"),
|
||||
arg_parser.get_bool("timer"),
|
||||
arg_parser.get_int("verify"),
|
||||
arg_parser.get_int("init"),
|
||||
arg_parser.get_bool("log"),
|
||||
arg_parser.get_str("csv_filename"),
|
||||
arg_parser.get_bool("flush_cache"),
|
||||
arg_parser.get_int("rotating_count"),
|
||||
arg_parser.get_bool("json_output")};
|
||||
|
||||
AQuantGemmProfiler profiler{setting};
|
||||
|
||||
try
|
||||
{
|
||||
auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args,
|
||||
const ck_tile::stream_config& stream) {
|
||||
return SelectedKernel::launch(args, stream);
|
||||
};
|
||||
|
||||
profiler.benchmark(problem, kernel_func);
|
||||
profiler.select_best_instance(static_cast<Metric>(arg_parser.get_int("metric")));
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Benchmark failed: " << e.what() << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
try
|
||||
{
|
||||
auto [result, parser] = create_args(argc, argv);
|
||||
if(!result)
|
||||
return EXIT_FAILURE;
|
||||
|
||||
benchmark_single(parser);
|
||||
return 0;
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Error: " << e.what() << "\n";
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/core/numeric/integer.hpp"
|
||||
|
||||
// DataTypeTraits for all supported types
|
||||
template <typename T>
|
||||
struct DataTypeTraits;
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<float>
|
||||
{
|
||||
static constexpr const char* name = "fp32";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::half_t>
|
||||
{
|
||||
static constexpr const char* name = "fp16";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::bf16_t>
|
||||
{
|
||||
static constexpr const char* name = "bf16";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::fp8_t>
|
||||
{
|
||||
static constexpr const char* name = "fp8";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::bf8_t>
|
||||
{
|
||||
static constexpr const char* name = "bf8";
|
||||
};
|
||||
|
||||
// Helper function to determine if a layout is row-major
|
||||
template <typename Layout>
|
||||
constexpr auto is_row_major(Layout)
|
||||
{
|
||||
return ck_tile::bool_constant<std::is_same_v<Layout, ck_tile::tensor_layout::gemm::RowMajor>>{};
|
||||
}
|
||||
|
||||
// Structure to hold kernel traits for dispatcher
|
||||
struct AQuantKernelTraits
|
||||
{
|
||||
std::string pipeline; // mem, compv3
|
||||
std::string scheduler; // intrawave, interwave
|
||||
std::string epilogue; // default
|
||||
bool pad_m;
|
||||
bool pad_n;
|
||||
bool pad_k;
|
||||
bool a_preshuffle_quant;
|
||||
|
||||
AQuantKernelTraits()
|
||||
: pipeline("mem"),
|
||||
scheduler("interwave"),
|
||||
epilogue("default"),
|
||||
pad_m(false),
|
||||
pad_n(false),
|
||||
pad_k(true),
|
||||
a_preshuffle_quant(false)
|
||||
{
|
||||
}
|
||||
};
|
||||
@@ -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()
|
||||
@@ -0,0 +1,276 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <fstream>
|
||||
#include <functional>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
#include "ck_tile/host/device_prop.hpp"
|
||||
#include "ck_tile/ops/gemm_quant.hpp"
|
||||
#include "gemm_aquant_benchmark.hpp"
|
||||
|
||||
class AQuantGemmProfiler
|
||||
{
|
||||
public:
|
||||
explicit AQuantGemmProfiler(Setting setting) : setting_(setting) {}
|
||||
|
||||
AQuantGemmProfiler(const AQuantGemmProfiler&) = delete;
|
||||
AQuantGemmProfiler& operator=(const AQuantGemmProfiler&) = delete;
|
||||
|
||||
void reset() { kernel_instances_.clear(); }
|
||||
|
||||
void benchmark(AQuantGemmProblem& problem,
|
||||
std::function<float(const ck_tile::QuantGemmHostArgs&,
|
||||
const ck_tile::stream_config&)> kernel_func)
|
||||
{
|
||||
std::vector<std::function<std::tuple<std::string, float>(ck_tile::QuantGemmHostArgs&,
|
||||
const ck_tile::stream_config&)>>
|
||||
callables;
|
||||
|
||||
callables.push_back(
|
||||
[kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
float time = kernel_func(args, stream);
|
||||
return std::make_tuple(std::string(KERNEL_NAME), time);
|
||||
});
|
||||
|
||||
benchmark(problem, callables);
|
||||
}
|
||||
|
||||
void benchmark(AQuantGemmProblem& problem,
|
||||
std::vector<std::function<std::tuple<std::string, float>(
|
||||
ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables)
|
||||
{
|
||||
const ALayout layout_a = ALayout{};
|
||||
const BLayout layout_b = BLayout{};
|
||||
const CLayout layout_c = CLayout{};
|
||||
const AQLayout layout_aq = AQLayout{};
|
||||
|
||||
problem.stride_a_ = ck_tile::get_default_stride(
|
||||
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a));
|
||||
problem.stride_b_ = ck_tile::get_default_stride(
|
||||
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b));
|
||||
problem.stride_c_ = ck_tile::get_default_stride(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c));
|
||||
|
||||
// Compute AQ scale tensor dimensions: [M, K / group_size_k]
|
||||
const ck_tile::index_t QK_A = problem.k_ / problem.group_size_k_;
|
||||
problem.stride_aq_ = ck_tile::get_default_stride(
|
||||
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq));
|
||||
|
||||
// Allocate host tensors
|
||||
ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
|
||||
ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
|
||||
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
|
||||
ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
|
||||
|
||||
// AQ scale tensor: [M, QK_A]
|
||||
ck_tile::HostTensor<AQDataType> aq_m_qk(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq)));
|
||||
|
||||
// Initialize tensors
|
||||
if(setting_.init_method_ == 0)
|
||||
{
|
||||
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
|
||||
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
|
||||
ck_tile::FillUniformDistribution<AQDataType>{0.5f, 1.5f}(aq_m_qk);
|
||||
}
|
||||
else if(setting_.init_method_ == 1)
|
||||
{
|
||||
ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
|
||||
ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
|
||||
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
|
||||
}
|
||||
else if(setting_.init_method_ == 2)
|
||||
{
|
||||
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(1)}(a_m_k);
|
||||
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(1)}(b_k_n);
|
||||
ck_tile::FillConstant<AQDataType>{static_cast<AQDataType>(1)}(aq_m_qk);
|
||||
}
|
||||
else
|
||||
{
|
||||
a_m_k.SetZero();
|
||||
b_k_n.SetZero();
|
||||
aq_m_qk.SetZero();
|
||||
}
|
||||
|
||||
// Allocate device memory
|
||||
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem aq_dev_buf(aq_m_qk.get_element_space_size_in_bytes());
|
||||
|
||||
a_m_k_dev_buf.ToDevice(a_m_k.data());
|
||||
b_k_n_dev_buf.ToDevice(b_k_n.data());
|
||||
aq_dev_buf.ToDevice(aq_m_qk.data());
|
||||
c_m_n_dev_buf.SetZero();
|
||||
c_m_n_dev_result.SetZero();
|
||||
|
||||
// Build QuantGemmHostArgs
|
||||
ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
|
||||
b_k_n_dev_buf.GetDeviceBuffer(),
|
||||
c_m_n_dev_buf.GetDeviceBuffer(),
|
||||
aq_dev_buf.GetDeviceBuffer(),
|
||||
nullptr, // bq_ptr not used for AQuant
|
||||
problem.split_k_,
|
||||
problem.m_,
|
||||
problem.n_,
|
||||
problem.k_,
|
||||
QK_A,
|
||||
0, // QK_B not used for AQuant
|
||||
problem.stride_a_,
|
||||
problem.stride_b_,
|
||||
problem.stride_c_,
|
||||
problem.stride_aq_,
|
||||
0 // stride_BQ not used for AQuant
|
||||
);
|
||||
|
||||
// Host reference for verification
|
||||
ck_tile::HostTensor<CDataType> c_m_n_host_result(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
|
||||
|
||||
if(setting_.verify_)
|
||||
{
|
||||
aquant_gemm_host_reference(setting_.verify_, a_m_k, aq_m_qk, b_k_n, c_m_n_host_result);
|
||||
}
|
||||
|
||||
// Run kernel(s)
|
||||
for(auto& callable : callables)
|
||||
{
|
||||
auto kernel_run_result = callable(gemm_args,
|
||||
ck_tile::stream_config{nullptr,
|
||||
true,
|
||||
setting_.log_,
|
||||
setting_.n_warmup_,
|
||||
setting_.n_repeat_,
|
||||
setting_.is_gpu_timer_,
|
||||
setting_.flush_cache_,
|
||||
setting_.rotating_count_});
|
||||
process_result(problem,
|
||||
QK_A,
|
||||
c_m_n_dev_buf,
|
||||
c_m_n_host_result,
|
||||
c_m_n_dev_result,
|
||||
kernel_run_result);
|
||||
}
|
||||
}
|
||||
|
||||
void process_result(const AQuantGemmProblem& problem,
|
||||
ck_tile::index_t QK_A,
|
||||
ck_tile::DeviceMem& c_m_n_dev_buf,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
const std::tuple<std::string, float>& kernel_run_result)
|
||||
{
|
||||
auto [name, avg_time] = kernel_run_result;
|
||||
|
||||
KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}};
|
||||
|
||||
// Compute performance metrics
|
||||
std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_;
|
||||
std::size_t num_byte = sizeof(ADataType) * problem.m_ * problem.k_ +
|
||||
sizeof(BDataType) * problem.n_ * problem.k_ +
|
||||
sizeof(AQDataType) * problem.m_ * QK_A +
|
||||
sizeof(CDataType) * problem.m_ * problem.n_;
|
||||
|
||||
kernel_instance.perf_result_.latency_ = avg_time;
|
||||
kernel_instance.perf_result_.tflops_ = static_cast<float>(flop) / 1.E9 / avg_time;
|
||||
kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time;
|
||||
|
||||
if(setting_.log_ > 0 && !setting_.json_output_)
|
||||
{
|
||||
std::cout << kernel_instance << std::endl;
|
||||
}
|
||||
|
||||
// Verify result
|
||||
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
|
||||
bool verified_correct =
|
||||
!setting_.verify_ ||
|
||||
compare_aquant(name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result);
|
||||
|
||||
if(verified_correct)
|
||||
{
|
||||
kernel_instances_.emplace_back(kernel_instance);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Verification failed, skip kernel: " << name << std::endl;
|
||||
}
|
||||
|
||||
c_m_n_dev_buf.SetZero();
|
||||
c_m_n_dev_result.SetZero();
|
||||
}
|
||||
|
||||
KernelInstance select_best_instance(Metric metric)
|
||||
{
|
||||
if(kernel_instances_.empty())
|
||||
throw std::runtime_error("Empty instances");
|
||||
|
||||
auto kernel_instance = *std::max_element(kernel_instances_.begin(),
|
||||
kernel_instances_.end(),
|
||||
[metric](const auto& a, const auto& b) {
|
||||
return PerformanceResult::compare(
|
||||
b.perf_result_, a.perf_result_, metric);
|
||||
});
|
||||
|
||||
if(setting_.json_output_)
|
||||
{
|
||||
std::cout << kernel_instance << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "**********************************" << std::endl;
|
||||
std::cout << "According to given metrics: " << get_metric_name(metric) << "\n"
|
||||
<< "Current kernel performance is: " << kernel_instance << std::endl;
|
||||
std::cout << "**********************************" << std::endl;
|
||||
}
|
||||
|
||||
if(!setting_.csv_filename_.empty())
|
||||
{
|
||||
std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app);
|
||||
|
||||
if(!file.is_open())
|
||||
{
|
||||
std::cerr << "Warning: Failed to open CSV file for writing." << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
if(file.tellp() == 0)
|
||||
{
|
||||
file << "rocm_version,device_name,"
|
||||
<< "split_k,m,n,k,stride_a,stride_b,stride_c,stride_aq,group_size_k,"
|
||||
<< "dtype_a,dtype_b,dtype_aq,dtype_acc,dtype_c,"
|
||||
<< "layout_a,layout_b,layout_c," << "name,"
|
||||
<< "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n";
|
||||
}
|
||||
|
||||
const auto& prob = kernel_instance.problem_;
|
||||
const auto& perf = kernel_instance.perf_result_;
|
||||
|
||||
file << get_rocm_version() << "," << ck_tile::get_device_name() << ","
|
||||
<< prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << ","
|
||||
<< prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << ","
|
||||
<< prob.stride_aq_ << "," << prob.group_size_k_ << "," << prob.dtype_a_ << ","
|
||||
<< prob.dtype_b_ << "," << prob.dtype_aq_ << "," << prob.dtype_acc_ << ","
|
||||
<< prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_ << ","
|
||||
<< prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed
|
||||
<< std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << ","
|
||||
<< perf.bandwidth_ << "," << get_metric_name(metric) << "\n";
|
||||
}
|
||||
}
|
||||
|
||||
return kernel_instance;
|
||||
}
|
||||
|
||||
private:
|
||||
Setting setting_;
|
||||
std::vector<KernelInstance> kernel_instances_;
|
||||
};
|
||||
@@ -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()
|
||||
289
tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt
Normal file
289
tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt
Normal file
@@ -0,0 +1,289 @@
|
||||
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
set(GEMM_BQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM BQuant (semicolon-separated)")
|
||||
set(GEMM_BQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM BQuant (semicolon-separated)")
|
||||
set(GEMM_BQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)")
|
||||
option(ENABLE_CCACHE_GEMM_BQUANT "Enable ccache for GEMM BQuant ops compilation" OFF)
|
||||
|
||||
# Store the directory path for use in functions
|
||||
set(GEMM_BQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR})
|
||||
|
||||
# Function to create individual GEMM BQuant targets
|
||||
function(create_individual_gemm_bquant_target datatype layout trait tile_config config_json)
|
||||
# GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL is guaranteed non-empty here (caller checks at line 192)
|
||||
if(NOT GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL)
|
||||
message(FATAL_ERROR "BUG: create_individual_gemm_bquant_target called with no GPU targets")
|
||||
endif()
|
||||
|
||||
# Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k
|
||||
string(REPLACE "_" ";" config_groups ${tile_config})
|
||||
list(GET config_groups 0 tile_dims)
|
||||
list(GET config_groups 1 warp_dims)
|
||||
list(GET config_groups 2 warp_tile_dims)
|
||||
|
||||
string(REPLACE "x" ";" tile_parts ${tile_dims})
|
||||
list(GET tile_parts 0 tile_m)
|
||||
list(GET tile_parts 1 tile_n)
|
||||
list(GET tile_parts 2 tile_k)
|
||||
|
||||
string(REPLACE "x" ";" warp_parts ${warp_dims})
|
||||
list(GET warp_parts 0 warp_m)
|
||||
list(GET warp_parts 1 warp_n)
|
||||
list(GET warp_parts 2 warp_k)
|
||||
|
||||
string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims})
|
||||
list(GET warp_tile_parts 0 warp_tile_m)
|
||||
list(GET warp_tile_parts 1 warp_tile_n)
|
||||
list(GET warp_tile_parts 2 warp_tile_k)
|
||||
|
||||
set(target_name "benchmark_gemm_bquant_${datatype}_${layout}_${trait}_${tile_config}")
|
||||
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
|
||||
|
||||
# Generate the single instance header for this kernel
|
||||
set(instance_header "${working_path}/gemm_bquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp")
|
||||
|
||||
# Add custom command to generate the header file at build time
|
||||
add_custom_command(
|
||||
OUTPUT ${instance_header}
|
||||
COMMAND ${Python3_EXECUTABLE} ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_instance_builder.py
|
||||
--working_path ${working_path}
|
||||
--datatype ${datatype}
|
||||
--layout ${layout}
|
||||
--config_json ${config_json}
|
||||
--gen_single
|
||||
--kernel_name "gemm_bquant_${datatype}_${layout}_${trait}_${tile_config}"
|
||||
--tile_config "${tile_config}"
|
||||
--trait_combo "${trait}"
|
||||
--gpu_target "${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}"
|
||||
DEPENDS ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_instance_builder.py ${config_json}
|
||||
COMMENT "Generating ${instance_header}"
|
||||
)
|
||||
|
||||
# Create the executable
|
||||
add_executable(${target_name}
|
||||
EXCLUDE_FROM_ALL
|
||||
${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_benchmark_single.cpp
|
||||
${instance_header}
|
||||
)
|
||||
|
||||
# Set GPU architectures
|
||||
set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL})
|
||||
|
||||
# Set compile definitions
|
||||
target_compile_definitions(${target_name} PRIVATE
|
||||
GEMM_BQUANT_SINGLE_INSTANCE_HPP="${instance_header}"
|
||||
)
|
||||
|
||||
# Include directories
|
||||
target_include_directories(${target_name} PRIVATE
|
||||
${GEMM_BQUANT_SOURCE_DIR}
|
||||
${working_path}
|
||||
)
|
||||
|
||||
# Compile options
|
||||
target_compile_options(${target_name} PRIVATE
|
||||
-Wno-undefined-func-template
|
||||
-Wno-float-equal
|
||||
--offload-compress
|
||||
-include ${instance_header}
|
||||
)
|
||||
|
||||
# Add to collection targets
|
||||
add_dependencies(benchmark_gemm_bquant_all ${target_name})
|
||||
add_dependencies(benchmark_gemm_bquant_${datatype} ${target_name})
|
||||
add_dependencies(benchmark_gemm_bquant_${layout} ${target_name})
|
||||
add_dependencies(benchmark_gemm_bquant_${datatype}_${layout} ${target_name})
|
||||
|
||||
# Add to pipeline-specific, epilogue-specific, and scheduler-specific targets
|
||||
string(REPLACE "_" ";" trait_parts ${trait})
|
||||
list(GET trait_parts 0 pipeline)
|
||||
list(GET trait_parts 1 epilogue)
|
||||
list(GET trait_parts 2 scheduler)
|
||||
|
||||
add_dependencies(benchmark_gemm_bquant_${pipeline}_pipeline ${target_name})
|
||||
add_dependencies(benchmark_gemm_bquant_${epilogue}_epilogue ${target_name})
|
||||
add_dependencies(benchmark_gemm_bquant_${scheduler}_scheduler ${target_name})
|
||||
endfunction()
|
||||
|
||||
# Function to build individual GEMM BQuant targets
|
||||
function(build_individual_gemm_bquant_targets datatype layout)
|
||||
set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
|
||||
|
||||
# Choose config file (priority: env var > cmake var > default)
|
||||
if(DEFINED ENV{GEMM_BQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_BQUANT_CONFIG_FILE}" STREQUAL "")
|
||||
set(config_filename "$ENV{GEMM_BQUANT_CONFIG_FILE}")
|
||||
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}")
|
||||
message(VERBOSE " Using config from environment variable: ${config_filename}")
|
||||
elseif(NOT "${GEMM_BQUANT_CONFIG_FILE}" STREQUAL "")
|
||||
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_BQUANT_CONFIG_FILE}")
|
||||
message(VERBOSE " Using custom config: ${GEMM_BQUANT_CONFIG_FILE}")
|
||||
else()
|
||||
set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json")
|
||||
message(VERBOSE " Using default config for layout ${layout}")
|
||||
endif()
|
||||
|
||||
if(NOT EXISTS ${json_blob})
|
||||
message(FATAL_ERROR "Config file not found: ${json_blob}")
|
||||
endif()
|
||||
|
||||
# Create working directory
|
||||
file(MAKE_DIRECTORY ${working_path})
|
||||
|
||||
# Build sampling arguments
|
||||
set(extra_list_args "")
|
||||
if(NOT "${GEMM_BQUANT_MAX_INSTANCES}" STREQUAL "")
|
||||
list(APPEND extra_list_args --max-instances ${GEMM_BQUANT_MAX_INSTANCES})
|
||||
endif()
|
||||
if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "")
|
||||
list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER})
|
||||
list(APPEND extra_list_args --manifest-path ${working_path})
|
||||
endif()
|
||||
if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "")
|
||||
list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED})
|
||||
endif()
|
||||
|
||||
# List kernels
|
||||
message(VERBOSE " Listing BQuant kernel configurations for ${datatype} ${layout}...")
|
||||
execute_process(
|
||||
COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_bquant_instance_builder.py
|
||||
--working_path ${working_path}
|
||||
--datatype ${datatype}
|
||||
--layout ${layout}
|
||||
--config_json ${json_blob}
|
||||
--gpu_target "${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}"
|
||||
--list_kernels
|
||||
${extra_list_args}
|
||||
WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
|
||||
RESULT_VARIABLE ret
|
||||
OUTPUT_VARIABLE list_output
|
||||
ERROR_VARIABLE list_error
|
||||
)
|
||||
|
||||
if(NOT ret EQUAL 0)
|
||||
message(FATAL_ERROR "Failed to list BQuant kernels for ${datatype} ${layout}: ${list_error}")
|
||||
endif()
|
||||
|
||||
# Read kernel count
|
||||
if(EXISTS ${working_path}/gemm_bquant_kernel_count.txt)
|
||||
file(READ ${working_path}/gemm_bquant_kernel_count.txt kernel_count)
|
||||
string(STRIP "${kernel_count}" kernel_count)
|
||||
message(VERBOSE " Found ${kernel_count} BQuant kernel configurations")
|
||||
else()
|
||||
message(FATAL_ERROR "Kernel count file not found")
|
||||
endif()
|
||||
|
||||
# Read kernel list and create targets
|
||||
if(EXISTS ${working_path}/gemm_bquant_kernel_list.txt)
|
||||
file(STRINGS ${working_path}/gemm_bquant_kernel_list.txt kernel_lines)
|
||||
foreach(line IN LISTS kernel_lines)
|
||||
string(REPLACE "|" ";" parts "${line}")
|
||||
list(GET parts 0 kernel_name)
|
||||
list(GET parts 1 tile_config)
|
||||
list(GET parts 2 trait_combo)
|
||||
|
||||
create_individual_gemm_bquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}")
|
||||
endforeach()
|
||||
else()
|
||||
message(FATAL_ERROR "Kernel list file not found")
|
||||
endif()
|
||||
endfunction()
|
||||
|
||||
# Main build logic
|
||||
message(VERBOSE "=== Starting Tile Engine GEMM BQuant Configuration ===")
|
||||
message(VERBOSE "GEMM_BQUANT_DATATYPE: ${GEMM_BQUANT_DATATYPE}")
|
||||
message(VERBOSE "GEMM_BQUANT_LAYOUT: ${GEMM_BQUANT_LAYOUT}")
|
||||
message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
|
||||
|
||||
# Filter GPU targets to supported architectures
|
||||
set(GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL "")
|
||||
set(DESIRED_TARGETS "gfx90a;gfx942;gfx950")
|
||||
|
||||
foreach(target IN LISTS SUPPORTED_GPU_TARGETS)
|
||||
if(target IN_LIST DESIRED_TARGETS)
|
||||
list(APPEND GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL ${target})
|
||||
message(VERBOSE " Adding GPU target: ${target}")
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
# Skip build if no matching targets found
|
||||
if(NOT GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL)
|
||||
message(WARNING "Skipping Tile Engine GEMM BQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
|
||||
else()
|
||||
message(VERBOSE "Building individual GEMM BQuant targets for GPU targets: ${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}")
|
||||
|
||||
# Enable compiler cache if requested
|
||||
if(ENABLE_CCACHE_GEMM_BQUANT)
|
||||
find_program(CCACHE_PROGRAM ccache)
|
||||
if(CCACHE_PROGRAM)
|
||||
set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
|
||||
message(VERBOSE "Using ccache for faster compilation")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Create master collection target
|
||||
add_custom_target(benchmark_gemm_bquant_all)
|
||||
|
||||
# Create datatype collection targets
|
||||
foreach(dt IN LISTS GEMM_BQUANT_DATATYPE)
|
||||
add_custom_target(benchmark_gemm_bquant_${dt})
|
||||
endforeach()
|
||||
|
||||
# Create layout collection targets
|
||||
foreach(l IN LISTS GEMM_BQUANT_LAYOUT)
|
||||
add_custom_target(benchmark_gemm_bquant_${l})
|
||||
endforeach()
|
||||
|
||||
# Create combined collection targets
|
||||
foreach(dt IN LISTS GEMM_BQUANT_DATATYPE)
|
||||
foreach(l IN LISTS GEMM_BQUANT_LAYOUT)
|
||||
add_custom_target(benchmark_gemm_bquant_${dt}_${l})
|
||||
endforeach()
|
||||
endforeach()
|
||||
|
||||
# Create pipeline-specific collection targets
|
||||
set(GEMM_BQUANT_PIPELINES "compv3")
|
||||
foreach(pipeline IN LISTS GEMM_BQUANT_PIPELINES)
|
||||
add_custom_target(benchmark_gemm_bquant_${pipeline}_pipeline)
|
||||
endforeach()
|
||||
|
||||
# Create epilogue-specific collection targets
|
||||
set(GEMM_BQUANT_EPILOGUES "default;cshuffle")
|
||||
foreach(epilogue IN LISTS GEMM_BQUANT_EPILOGUES)
|
||||
add_custom_target(benchmark_gemm_bquant_${epilogue}_epilogue)
|
||||
endforeach()
|
||||
|
||||
# Create scheduler-specific collection targets
|
||||
set(GEMM_BQUANT_SCHEDULERS "intrawave")
|
||||
foreach(scheduler IN LISTS GEMM_BQUANT_SCHEDULERS)
|
||||
add_custom_target(benchmark_gemm_bquant_${scheduler}_scheduler)
|
||||
endforeach()
|
||||
|
||||
# Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that
|
||||
# sampling fires per-combo rather than being a single cap.
|
||||
# CMake integer division truncates; clamp the result to at least 1 so that a
|
||||
# small total budget (e.g. 5 instances across 8 combos) does not silently
|
||||
# produce a per-combo budget of 0 and skip all kernels.
|
||||
if(NOT "${GEMM_BQUANT_MAX_INSTANCES}" STREQUAL "")
|
||||
list(LENGTH GEMM_BQUANT_DATATYPE _bq_n_dt)
|
||||
list(LENGTH GEMM_BQUANT_LAYOUT _bq_n_lay)
|
||||
math(EXPR _bq_n_combos "${_bq_n_dt} * ${_bq_n_lay}")
|
||||
if(_bq_n_combos GREATER 0)
|
||||
math(EXPR GEMM_BQUANT_MAX_INSTANCES
|
||||
"${GEMM_BQUANT_MAX_INSTANCES} / ${_bq_n_combos}")
|
||||
if(GEMM_BQUANT_MAX_INSTANCES EQUAL 0)
|
||||
set(GEMM_BQUANT_MAX_INSTANCES 1)
|
||||
message(WARNING " gemm_bquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_bq_n_combos} combos)")
|
||||
else()
|
||||
message(STATUS " gemm_bquant: per-combo budget = ${GEMM_BQUANT_MAX_INSTANCES} (${_bq_n_combos} combos)")
|
||||
endif()
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Build individual targets for each datatype/layout combination
|
||||
foreach(dt IN LISTS GEMM_BQUANT_DATATYPE)
|
||||
foreach(l IN LISTS GEMM_BQUANT_LAYOUT)
|
||||
build_individual_gemm_bquant_targets(${dt} ${l})
|
||||
endforeach()
|
||||
endforeach()
|
||||
endif()
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "gemm_bquant_common.hpp"
|
||||
|
||||
// Data types and Layouts are defined by the generated kernel headers:
|
||||
// ADataType, BDataType, BQDataType, AccDataType, CDataType
|
||||
// ALayout, BLayout, CLayout, BQLayout
|
||||
|
||||
enum class Metric
|
||||
{
|
||||
LATENCY = 0,
|
||||
TFLOPS = 1,
|
||||
BANDWIDTH = 2
|
||||
};
|
||||
|
||||
inline constexpr auto get_metric_name(Metric m)
|
||||
{
|
||||
switch(m)
|
||||
{
|
||||
case Metric::LATENCY: return "latency";
|
||||
case Metric::TFLOPS: return "tflops";
|
||||
case Metric::BANDWIDTH: return "bandwidth";
|
||||
default: throw std::invalid_argument("Unsupported metric type");
|
||||
}
|
||||
}
|
||||
|
||||
struct BQuantGemmProblem
|
||||
{
|
||||
int split_k_;
|
||||
int m_, n_, k_;
|
||||
int stride_a_, stride_b_, stride_c_;
|
||||
int stride_bq_;
|
||||
int group_size_k_;
|
||||
|
||||
std::string dtype_a_, dtype_b_, dtype_bq_, dtype_acc_, dtype_c_;
|
||||
std::string layout_a_, layout_b_, layout_c_;
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const BQuantGemmProblem& problem)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"split_k\":" << problem.split_k_ << ",\n"
|
||||
<< " \"m\":" << problem.m_ << ",\n"
|
||||
<< " \"n\":" << problem.n_ << ",\n"
|
||||
<< " \"k\":" << problem.k_ << ",\n"
|
||||
<< " \"stride_a\":" << problem.stride_a_ << ",\n"
|
||||
<< " \"stride_b\":" << problem.stride_b_ << ",\n"
|
||||
<< " \"stride_c\":" << problem.stride_c_ << ",\n"
|
||||
<< " \"stride_bq\":" << problem.stride_bq_ << ",\n"
|
||||
<< " \"group_size_k\":" << problem.group_size_k_ << ",\n"
|
||||
<< " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n"
|
||||
<< " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n"
|
||||
<< " \"dtype_bq\":\"" << problem.dtype_bq_ << "\",\n"
|
||||
<< " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n"
|
||||
<< " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n"
|
||||
<< " \"layout_a\":\"" << problem.layout_a_ << "\",\n"
|
||||
<< " \"layout_b\":\"" << problem.layout_b_ << "\",\n"
|
||||
<< " \"layout_c\":\"" << problem.layout_c_ << "\"\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct PerformanceResult
|
||||
{
|
||||
double latency_;
|
||||
double tflops_;
|
||||
double bandwidth_;
|
||||
|
||||
static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m)
|
||||
{
|
||||
switch(m)
|
||||
{
|
||||
case Metric::LATENCY: return a.latency_ < b.latency_;
|
||||
case Metric::TFLOPS: return a.tflops_ > b.tflops_;
|
||||
case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_;
|
||||
default: throw std::invalid_argument("Unsupported metric type");
|
||||
}
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_
|
||||
<< ",\n"
|
||||
<< " \"tflops(TFlops)\": " << result.tflops_ << ",\n"
|
||||
<< " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct KernelInstance
|
||||
{
|
||||
std::string name_;
|
||||
BQuantGemmProblem problem_;
|
||||
PerformanceResult perf_result_;
|
||||
|
||||
static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m)
|
||||
{
|
||||
return PerformanceResult::compare(a.perf_result_, b.perf_result_, m);
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj)
|
||||
{
|
||||
os << "{\n"
|
||||
<< " \"name\": \"" << obj.name_ << "\",\n"
|
||||
<< " \"problem\": " << obj.problem_ << ",\n"
|
||||
<< " \"perf_result\": " << obj.perf_result_ << "\n"
|
||||
<< "}";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct Setting
|
||||
{
|
||||
int n_warmup_;
|
||||
int n_repeat_;
|
||||
bool is_gpu_timer_;
|
||||
int verify_;
|
||||
int init_method_;
|
||||
bool log_;
|
||||
std::string csv_filename_;
|
||||
bool flush_cache_;
|
||||
int rotating_count_;
|
||||
bool json_output_;
|
||||
};
|
||||
|
||||
inline std::string get_rocm_version()
|
||||
{
|
||||
std::ifstream version_file("/opt/rocm/.info/version");
|
||||
if(version_file.is_open())
|
||||
{
|
||||
std::string version;
|
||||
std::getline(version_file, version);
|
||||
return version;
|
||||
}
|
||||
return "Unknown";
|
||||
}
|
||||
|
||||
template <typename ADataType_,
|
||||
typename BDataType_,
|
||||
typename BQDataType_,
|
||||
typename AccDataType_,
|
||||
typename CDataType_>
|
||||
auto calculate_rtol_atol_bquant(const ck_tile::index_t K,
|
||||
const ck_tile::index_t kbatch,
|
||||
const float max_accumulated_value)
|
||||
{
|
||||
using ComputeType =
|
||||
std::conditional_t<sizeof(ADataType_) < sizeof(BDataType_), ADataType_, BDataType_>;
|
||||
const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType_, AccDataType_>(
|
||||
ck_tile::integer_divide_ceil(K, kbatch));
|
||||
const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType_, AccDataType_>(
|
||||
max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
|
||||
const auto rtol_split_k =
|
||||
ck_tile::get_relative_threshold<CDataType_, CDataType_, CDataType_>(kbatch);
|
||||
const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType_, CDataType_, CDataType_>(
|
||||
max_accumulated_value, kbatch);
|
||||
return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
|
||||
}
|
||||
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions"
|
||||
/// @brief Compare device and host results for BQuant GEMM
|
||||
bool compare_bquant(std::string instanceName,
|
||||
ck_tile::index_t K,
|
||||
ck_tile::index_t kbatch,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
|
||||
{
|
||||
const float max_accumulated_value =
|
||||
std::abs(static_cast<float>(*std::max_element(c_m_n_host_result.mData.begin(),
|
||||
c_m_n_host_result.mData.end(),
|
||||
[](const auto& a, const auto& b) {
|
||||
return std::abs(static_cast<float>(a)) <
|
||||
std::abs(static_cast<float>(b));
|
||||
})));
|
||||
const auto rtol_atol =
|
||||
calculate_rtol_atol_bquant<ADataType, BDataType, BQDataType, AccDataType, CDataType>(
|
||||
K, kbatch, max_accumulated_value);
|
||||
bool pass = ck_tile::check_err(c_m_n_dev_result,
|
||||
c_m_n_host_result,
|
||||
"Error: Incorrect results!",
|
||||
rtol_atol.at(ck_tile::number<0>{}),
|
||||
rtol_atol.at(ck_tile::number<1>{}));
|
||||
|
||||
std::cout << "For " << instanceName << " Relative error threshold is "
|
||||
<< rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
|
||||
<< rtol_atol.at(ck_tile::number<1>{}) << std::endl;
|
||||
std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;
|
||||
|
||||
return pass;
|
||||
}
|
||||
|
||||
/// @brief CPU reference implementation for BQuant GEMM
|
||||
void bquant_gemm_host_reference(int verify,
|
||||
ck_tile::HostTensor<ADataType>& a_m_k,
|
||||
ck_tile::HostTensor<BDataType>& b_k_n,
|
||||
ck_tile::HostTensor<BQDataType>& bq_qk_n,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result)
|
||||
{
|
||||
if(verify == 1)
|
||||
{
|
||||
c_m_n_host_result.SetZero();
|
||||
using QuantGroupSize = typename SelectedKernel::QuantGroupSize;
|
||||
ck_tile::reference_gemm_quant<ADataType,
|
||||
BQDataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
QuantGroupSize,
|
||||
false /* aquant=false, bquant */>(
|
||||
a_m_k, bq_qk_n, b_k_n, c_m_n_host_result);
|
||||
}
|
||||
}
|
||||
#pragma clang diagnostic pop
|
||||
@@ -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())
|
||||
@@ -0,0 +1,166 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#include <iostream>
|
||||
#include <functional>
|
||||
#include <tuple>
|
||||
#include <exception>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "gemm_bquant_profiler.hpp"
|
||||
#include "gemm_bquant_common.hpp"
|
||||
|
||||
// The kernel header is included via the compile command line with -include flag
|
||||
// It defines SelectedKernel struct, KERNEL_NAME, and type aliases:
|
||||
// ADataType, BDataType, BQDataType, AccDataType, CDataType
|
||||
// ALayout, BLayout, CLayout, BQLayout
|
||||
|
||||
inline auto create_args(int argc, char* argv[])
|
||||
{
|
||||
ck_tile::ArgParser arg_parser;
|
||||
arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.")
|
||||
.insert("n", "4096", "The value for n dimension. Default is 4096.")
|
||||
.insert("k", "2048", "The value for k dimension. Default is 2048.")
|
||||
.insert("stride_a", "0", "The stride value for tensor A. Default is 0.")
|
||||
.insert("stride_b", "0", "The stride value for tensor B. Default is 0.")
|
||||
.insert("stride_c", "0", "The stride value for tensor C. Default is 0.")
|
||||
.insert("split_k", "1", "The split value for k dimension. Default is 1.")
|
||||
.insert("group_size_k",
|
||||
std::to_string(SelectedKernel::GroupSizeK),
|
||||
"The quantization group size along K. Default matches kernel config.")
|
||||
.insert("verify",
|
||||
"1",
|
||||
"The type of validation. Set to 0 for no validation, 1 for validation on CPU. "
|
||||
"Default is 1, CPU validation.")
|
||||
.insert("log",
|
||||
"false",
|
||||
"Whether output kernel instance information or not. Possible values are true or "
|
||||
"false. Default is false")
|
||||
.insert(
|
||||
"warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.")
|
||||
.insert(
|
||||
"repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.")
|
||||
.insert("timer",
|
||||
"true",
|
||||
"Whether if the timer is gpu timer or not. Possible values are false or true. "
|
||||
"Default is true.")
|
||||
.insert("init",
|
||||
"0",
|
||||
"The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 "
|
||||
"for constant(1). Default is 0, random.")
|
||||
.insert("flush_cache",
|
||||
"true",
|
||||
"To flush cache, possible values are true or false. Default is true.")
|
||||
.insert(
|
||||
"rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.")
|
||||
.insert("metric",
|
||||
"0",
|
||||
"Metric with which to measure kernel performance. Set to 0 for latency, 1 for "
|
||||
"tflops, or 2 for bandwidth. Default is 0, latency.")
|
||||
.insert("csv_filename",
|
||||
"",
|
||||
"The filename of benchmark result. Default is empty (no CSV output).")
|
||||
.insert("json_output",
|
||||
"false",
|
||||
"Whether to output results in JSON format only. Possible values are true or false. "
|
||||
"Default is false");
|
||||
|
||||
bool result = arg_parser.parse(argc, argv);
|
||||
return std::make_tuple(result, arg_parser);
|
||||
}
|
||||
|
||||
void benchmark_single(const ck_tile::ArgParser& arg_parser)
|
||||
{
|
||||
std::string dtype_a = DataTypeTraits<ADataType>::name;
|
||||
std::string dtype_b = DataTypeTraits<BDataType>::name;
|
||||
std::string dtype_bq = DataTypeTraits<BQDataType>::name;
|
||||
std::string dtype_acc = DataTypeTraits<AccDataType>::name;
|
||||
std::string dtype_c = DataTypeTraits<CDataType>::name;
|
||||
|
||||
std::string layout_a = ALayout::name;
|
||||
std::string layout_b = BLayout::name;
|
||||
std::string layout_c = CLayout::name;
|
||||
|
||||
int M = arg_parser.get_int("m");
|
||||
int N = arg_parser.get_int("n");
|
||||
int K = arg_parser.get_int("k");
|
||||
int group_size_k = arg_parser.get_int("group_size_k");
|
||||
|
||||
if(M <= 0 || N <= 0 || K <= 0)
|
||||
{
|
||||
throw std::invalid_argument("m, n, k must be positive integers");
|
||||
}
|
||||
if(group_size_k <= 0 || K % group_size_k != 0)
|
||||
{
|
||||
throw std::invalid_argument(
|
||||
"group_size_k must be positive and k must be divisible by group_size_k");
|
||||
}
|
||||
|
||||
BQuantGemmProblem problem{arg_parser.get_int("split_k"),
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
arg_parser.get_int("stride_a"),
|
||||
arg_parser.get_int("stride_b"),
|
||||
arg_parser.get_int("stride_c"),
|
||||
0, // stride_bq computed by profiler
|
||||
group_size_k,
|
||||
dtype_a,
|
||||
dtype_b,
|
||||
dtype_bq,
|
||||
dtype_acc,
|
||||
dtype_c,
|
||||
layout_a,
|
||||
layout_b,
|
||||
layout_c};
|
||||
|
||||
Setting setting{arg_parser.get_int("warmup"),
|
||||
arg_parser.get_int("repeat"),
|
||||
arg_parser.get_bool("timer"),
|
||||
arg_parser.get_int("verify"),
|
||||
arg_parser.get_int("init"),
|
||||
arg_parser.get_bool("log"),
|
||||
arg_parser.get_str("csv_filename"),
|
||||
arg_parser.get_bool("flush_cache"),
|
||||
arg_parser.get_int("rotating_count"),
|
||||
arg_parser.get_bool("json_output")};
|
||||
|
||||
BQuantGemmProfiler profiler{setting};
|
||||
|
||||
try
|
||||
{
|
||||
auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args,
|
||||
const ck_tile::stream_config& stream) {
|
||||
return SelectedKernel::launch(args, stream);
|
||||
};
|
||||
|
||||
profiler.benchmark(problem, kernel_func);
|
||||
profiler.select_best_instance(static_cast<Metric>(arg_parser.get_int("metric")));
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Benchmark failed: " << e.what() << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
try
|
||||
{
|
||||
auto [result, parser] = create_args(argc, argv);
|
||||
if(!result)
|
||||
return EXIT_FAILURE;
|
||||
|
||||
benchmark_single(parser);
|
||||
return 0;
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Error: " << e.what() << "\n";
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/core/numeric/integer.hpp"
|
||||
|
||||
// DataTypeTraits for all supported types
|
||||
template <typename T>
|
||||
struct DataTypeTraits;
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<float>
|
||||
{
|
||||
static constexpr const char* name = "fp32";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::half_t>
|
||||
{
|
||||
static constexpr const char* name = "fp16";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::bf16_t>
|
||||
{
|
||||
static constexpr const char* name = "bf16";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::fp8_t>
|
||||
{
|
||||
static constexpr const char* name = "fp8";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeTraits<ck_tile::bf8_t>
|
||||
{
|
||||
static constexpr const char* name = "bf8";
|
||||
};
|
||||
|
||||
// Helper function to determine if a layout is row-major
|
||||
template <typename Layout>
|
||||
constexpr auto is_row_major(Layout)
|
||||
{
|
||||
return ck_tile::bool_constant<std::is_same_v<Layout, ck_tile::tensor_layout::gemm::RowMajor>>{};
|
||||
}
|
||||
|
||||
// Structure to hold kernel traits for dispatcher
|
||||
struct BQuantKernelTraits
|
||||
{
|
||||
std::string pipeline; // compv3
|
||||
std::string scheduler; // intrawave
|
||||
std::string epilogue; // default, cshuffle
|
||||
bool pad_m;
|
||||
bool pad_n;
|
||||
bool pad_k;
|
||||
bool b_preshuffle_quant;
|
||||
|
||||
BQuantKernelTraits()
|
||||
: pipeline("compv3"),
|
||||
scheduler("intrawave"),
|
||||
epilogue("default"),
|
||||
pad_m(false),
|
||||
pad_n(false),
|
||||
pad_k(false),
|
||||
b_preshuffle_quant(false)
|
||||
{
|
||||
}
|
||||
};
|
||||
@@ -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()
|
||||
@@ -0,0 +1,293 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <functional>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <stdexcept>
|
||||
|
||||
#include "ck_tile/host/device_prop.hpp"
|
||||
#include "ck_tile/host/tensor_shuffle_utils.hpp"
|
||||
#include "ck_tile/ops/gemm_quant.hpp"
|
||||
#include "gemm_bquant_benchmark.hpp"
|
||||
|
||||
class BQuantGemmProfiler
|
||||
{
|
||||
public:
|
||||
explicit BQuantGemmProfiler(Setting setting) : setting_(setting) {}
|
||||
|
||||
BQuantGemmProfiler(const BQuantGemmProfiler&) = delete;
|
||||
BQuantGemmProfiler& operator=(const BQuantGemmProfiler&) = delete;
|
||||
|
||||
void reset() { kernel_instances_.clear(); }
|
||||
|
||||
void benchmark(BQuantGemmProblem problem,
|
||||
std::function<float(const ck_tile::QuantGemmHostArgs&,
|
||||
const ck_tile::stream_config&)> kernel_func)
|
||||
{
|
||||
std::vector<std::function<std::tuple<std::string, float>(ck_tile::QuantGemmHostArgs&,
|
||||
const ck_tile::stream_config&)>>
|
||||
callables;
|
||||
|
||||
callables.push_back(
|
||||
[kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
float time = kernel_func(args, stream);
|
||||
return std::make_tuple(std::string(KERNEL_NAME), time);
|
||||
});
|
||||
|
||||
benchmark(problem, callables);
|
||||
}
|
||||
|
||||
void benchmark(BQuantGemmProblem problem,
|
||||
std::vector<std::function<std::tuple<std::string, float>(
|
||||
ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables)
|
||||
{
|
||||
const ALayout layout_a = ALayout{};
|
||||
const BLayout layout_b = BLayout{};
|
||||
const CLayout layout_c = CLayout{};
|
||||
const BQLayout layout_bq = BQLayout{};
|
||||
|
||||
problem.stride_a_ = ck_tile::get_default_stride(
|
||||
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a));
|
||||
problem.stride_b_ = ck_tile::get_default_stride(
|
||||
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b));
|
||||
problem.stride_c_ = ck_tile::get_default_stride(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c));
|
||||
|
||||
// Compute BQ scale tensor dimensions: [K / group_size_k, N]
|
||||
if(problem.k_ % problem.group_size_k_ != 0)
|
||||
throw std::runtime_error("k_ must be divisible by group_size_k_");
|
||||
const ck_tile::index_t QK_B = problem.k_ / problem.group_size_k_;
|
||||
problem.stride_bq_ = ck_tile::get_default_stride(
|
||||
QK_B, problem.n_, problem.stride_bq_, is_row_major(layout_bq));
|
||||
|
||||
// Allocate host tensors
|
||||
ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)));
|
||||
ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
|
||||
problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)));
|
||||
ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
|
||||
|
||||
// BQ scale tensor: [QK_B, N]
|
||||
ck_tile::HostTensor<BQDataType> bq_qk_n(ck_tile::host_tensor_descriptor(
|
||||
QK_B, problem.n_, problem.stride_bq_, is_row_major(layout_bq)));
|
||||
|
||||
// Initialize tensors
|
||||
if(setting_.init_method_ == 0)
|
||||
{
|
||||
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
|
||||
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
|
||||
ck_tile::FillUniformDistribution<BQDataType>{0.5f, 1.5f}(bq_qk_n);
|
||||
}
|
||||
else if(setting_.init_method_ == 1)
|
||||
{
|
||||
ck_tile::FillMonotonicSeq<ADataType>{}(a_m_k);
|
||||
ck_tile::FillMonotonicSeq<BDataType>{}(b_k_n);
|
||||
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
|
||||
}
|
||||
else if(setting_.init_method_ == 2)
|
||||
{
|
||||
ck_tile::FillConstant<ADataType>{static_cast<ADataType>(1)}(a_m_k);
|
||||
ck_tile::FillConstant<BDataType>{static_cast<BDataType>(1)}(b_k_n);
|
||||
ck_tile::FillConstant<BQDataType>{static_cast<BQDataType>(1)}(bq_qk_n);
|
||||
}
|
||||
else
|
||||
{
|
||||
a_m_k.SetZero();
|
||||
b_k_n.SetZero();
|
||||
bq_qk_n.SetZero();
|
||||
}
|
||||
|
||||
// Allocate device memory
|
||||
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem bq_dev_buf(bq_qk_n.get_element_space_size_in_bytes());
|
||||
|
||||
a_m_k_dev_buf.ToDevice(a_m_k.data());
|
||||
b_k_n_dev_buf.ToDevice(b_k_n.data());
|
||||
|
||||
// Shuffle BQ data if preshuffle is enabled
|
||||
if constexpr(SelectedKernel::BPreshuffleQuant)
|
||||
{
|
||||
constexpr int block_bq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK;
|
||||
ck_tile::HostTensor<BQDataType> bq_shuffled = ck_tile::shuffle_bq(&bq_qk_n, block_bq_k);
|
||||
bq_dev_buf.ToDevice(bq_shuffled.data());
|
||||
}
|
||||
else
|
||||
{
|
||||
bq_dev_buf.ToDevice(bq_qk_n.data());
|
||||
}
|
||||
|
||||
c_m_n_dev_buf.SetZero();
|
||||
c_m_n_dev_result.SetZero();
|
||||
|
||||
// Build QuantGemmHostArgs
|
||||
ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
|
||||
b_k_n_dev_buf.GetDeviceBuffer(),
|
||||
c_m_n_dev_buf.GetDeviceBuffer(),
|
||||
nullptr, // aq_ptr not used for BQuant
|
||||
bq_dev_buf.GetDeviceBuffer(),
|
||||
problem.split_k_,
|
||||
problem.m_,
|
||||
problem.n_,
|
||||
problem.k_,
|
||||
0, // QK_A not used for BQuant
|
||||
QK_B,
|
||||
problem.stride_a_,
|
||||
problem.stride_b_,
|
||||
problem.stride_c_,
|
||||
0, // stride_AQ not used for BQuant
|
||||
problem.stride_bq_);
|
||||
|
||||
// Host reference for verification
|
||||
ck_tile::HostTensor<CDataType> c_m_n_host_result(ck_tile::host_tensor_descriptor(
|
||||
problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)));
|
||||
|
||||
if(setting_.verify_)
|
||||
{
|
||||
bquant_gemm_host_reference(setting_.verify_, a_m_k, b_k_n, bq_qk_n, c_m_n_host_result);
|
||||
}
|
||||
|
||||
// Run kernel(s)
|
||||
for(auto& callable : callables)
|
||||
{
|
||||
auto kernel_run_result = callable(gemm_args,
|
||||
ck_tile::stream_config{nullptr,
|
||||
true,
|
||||
setting_.log_,
|
||||
setting_.n_warmup_,
|
||||
setting_.n_repeat_,
|
||||
setting_.is_gpu_timer_,
|
||||
setting_.flush_cache_,
|
||||
setting_.rotating_count_});
|
||||
process_result(problem,
|
||||
QK_B,
|
||||
c_m_n_dev_buf,
|
||||
c_m_n_host_result,
|
||||
c_m_n_dev_result,
|
||||
kernel_run_result);
|
||||
}
|
||||
}
|
||||
|
||||
void process_result(const BQuantGemmProblem& problem,
|
||||
ck_tile::index_t QK_B,
|
||||
ck_tile::DeviceMem& c_m_n_dev_buf,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_host_result,
|
||||
ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
const std::tuple<std::string, float>& kernel_run_result)
|
||||
{
|
||||
auto [name, avg_time] = kernel_run_result;
|
||||
|
||||
KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}};
|
||||
|
||||
// Compute performance metrics
|
||||
std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_;
|
||||
std::size_t num_byte = sizeof(ADataType) * problem.m_ * problem.k_ +
|
||||
sizeof(BDataType) * problem.n_ * problem.k_ +
|
||||
sizeof(BQDataType) * problem.n_ * QK_B +
|
||||
sizeof(CDataType) * problem.m_ * problem.n_;
|
||||
|
||||
kernel_instance.perf_result_.latency_ = avg_time;
|
||||
kernel_instance.perf_result_.tflops_ = static_cast<float>(flop) / 1.E9 / avg_time;
|
||||
kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time;
|
||||
|
||||
if(setting_.log_ > 0 && !setting_.json_output_)
|
||||
{
|
||||
std::cout << kernel_instance << std::endl;
|
||||
}
|
||||
|
||||
// Verify result
|
||||
bool verified_correct = true;
|
||||
if(setting_.verify_)
|
||||
{
|
||||
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
|
||||
verified_correct = compare_bquant(
|
||||
name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result);
|
||||
}
|
||||
|
||||
if(verified_correct)
|
||||
{
|
||||
kernel_instances_.emplace_back(kernel_instance);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Verification failed, skip kernel: " << name << std::endl;
|
||||
}
|
||||
|
||||
c_m_n_dev_buf.SetZero();
|
||||
c_m_n_dev_result.SetZero();
|
||||
}
|
||||
|
||||
KernelInstance select_best_instance(Metric metric)
|
||||
{
|
||||
if(kernel_instances_.empty())
|
||||
throw std::runtime_error("Empty instances");
|
||||
|
||||
auto kernel_instance = *std::max_element(kernel_instances_.begin(),
|
||||
kernel_instances_.end(),
|
||||
[metric](const auto& a, const auto& b) {
|
||||
return PerformanceResult::compare(
|
||||
b.perf_result_, a.perf_result_, metric);
|
||||
});
|
||||
|
||||
if(setting_.json_output_)
|
||||
{
|
||||
std::cout << kernel_instance << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "**********************************" << std::endl;
|
||||
std::cout << "According to given metrics: " << get_metric_name(metric) << "\n"
|
||||
<< "Current kernel performance is: " << kernel_instance << std::endl;
|
||||
std::cout << "**********************************" << std::endl;
|
||||
}
|
||||
|
||||
if(!setting_.csv_filename_.empty())
|
||||
{
|
||||
std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app);
|
||||
|
||||
if(!file.is_open())
|
||||
{
|
||||
std::cerr << "Warning: Failed to open CSV file for writing." << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
if(file.tellp() == 0)
|
||||
{
|
||||
file << "rocm_version,device_name,"
|
||||
<< "split_k,m,n,k,stride_a,stride_b,stride_c,stride_bq,group_size_k,"
|
||||
<< "dtype_a,dtype_b,dtype_bq,dtype_acc,dtype_c,"
|
||||
<< "layout_a,layout_b,layout_c," << "name,"
|
||||
<< "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n";
|
||||
}
|
||||
|
||||
const auto& prob = kernel_instance.problem_;
|
||||
const auto& perf = kernel_instance.perf_result_;
|
||||
|
||||
file << get_rocm_version() << "," << ck_tile::get_device_name() << ","
|
||||
<< prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << ","
|
||||
<< prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << ","
|
||||
<< prob.stride_bq_ << "," << prob.group_size_k_ << "," << prob.dtype_a_ << ","
|
||||
<< prob.dtype_b_ << "," << prob.dtype_bq_ << "," << prob.dtype_acc_ << ","
|
||||
<< prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_ << ","
|
||||
<< prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed
|
||||
<< std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << ","
|
||||
<< perf.bandwidth_ << "," << get_metric_name(metric) << "\n";
|
||||
}
|
||||
}
|
||||
|
||||
return kernel_instance;
|
||||
}
|
||||
|
||||
private:
|
||||
Setting setting_;
|
||||
std::vector<KernelInstance> kernel_instances_;
|
||||
};
|
||||
@@ -0,0 +1,191 @@
|
||||
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Unit tests for GemmBQuantKernelBuilder (gemm_bquant operator)."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.insert(0, _HERE)
|
||||
|
||||
from gemm_bquant_instance_builder import GemmBQuantKernelBuilder # noqa: E402
|
||||
|
||||
_MINIMAL_CONFIG = {
|
||||
"tile_config": {
|
||||
"tile_m": {"values": [64]},
|
||||
"tile_n": {"values": [64]},
|
||||
"tile_k": {"values": [128]},
|
||||
"warp_m": {"values": [2]},
|
||||
"warp_n": {"values": [2]},
|
||||
"warp_k": {"values": [1]},
|
||||
"warp_tile_m": {"values": [16]},
|
||||
"warp_tile_n": {"values": [16]},
|
||||
"warp_tile_k": {"values": [64]},
|
||||
},
|
||||
"trait_config": {
|
||||
"pipeline": {"values": ["compv3"]},
|
||||
"scheduler": {"values": ["intrawave"]},
|
||||
"epilogue": {"values": ["default"]},
|
||||
"pad_m": {"values": [False]},
|
||||
"pad_n": {"values": [False]},
|
||||
"pad_k": {"values": [False]},
|
||||
"b_preshuffle_quant": {"values": [False, True]},
|
||||
},
|
||||
"k_block_per_cu": 1,
|
||||
"group_size_k": 128,
|
||||
}
|
||||
|
||||
|
||||
def _make_builder(tmpdir, config=None, **kwargs):
|
||||
cfg = config if config is not None else _MINIMAL_CONFIG
|
||||
cfg_path = os.path.join(tmpdir, "config.json")
|
||||
with open(cfg_path, "w") as f:
|
||||
json.dump(cfg, f)
|
||||
return GemmBQuantKernelBuilder(
|
||||
kernel_name_prefix="gemm_bquant",
|
||||
working_path=tmpdir,
|
||||
gpu_target=kwargs.get("gpu_target", "gfx942"),
|
||||
datatype=kwargs.get("datatype", "fp8"),
|
||||
layout=kwargs.get("layout", "rcr"),
|
||||
config_json=cfg_path,
|
||||
)
|
||||
|
||||
|
||||
class TestGemmBQuantBuilderInit(unittest.TestCase):
|
||||
def test_group_size_k_from_config(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
builder = _make_builder(tmpdir)
|
||||
self.assertEqual(builder.group_size_k, 128)
|
||||
|
||||
def test_group_size_k_custom(self):
|
||||
cfg = json.loads(json.dumps(_MINIMAL_CONFIG))
|
||||
cfg["group_size_k"] = 32
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
builder = _make_builder(tmpdir, config=cfg)
|
||||
self.assertEqual(builder.group_size_k, 32)
|
||||
|
||||
def test_kernel_name_prefix_stored(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
builder = _make_builder(tmpdir)
|
||||
self.assertEqual(builder.kernel_name_prefix, "gemm_bquant")
|
||||
|
||||
def test_working_path_created(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
sub = os.path.join(tmpdir, "workdir")
|
||||
cfg_path = os.path.join(tmpdir, "config.json")
|
||||
with open(cfg_path, "w") as f:
|
||||
json.dump(_MINIMAL_CONFIG, f)
|
||||
GemmBQuantKernelBuilder("gemm_bquant", sub, "gfx942", "fp8", "rcr", cfg_path)
|
||||
self.assertTrue(os.path.isdir(sub))
|
||||
|
||||
|
||||
class TestGemmBQuantTraitCombinations(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmpdir = tempfile.TemporaryDirectory()
|
||||
self.builder = _make_builder(self._tmpdir.name)
|
||||
|
||||
def tearDown(self):
|
||||
self._tmpdir.cleanup()
|
||||
|
||||
def test_trait_combinations_non_empty(self):
|
||||
combos = self.builder._generate_trait_combinations()
|
||||
self.assertGreater(len(combos), 0)
|
||||
|
||||
def test_trait_combo_is_7_tuple(self):
|
||||
"""BQuant trait tuple: (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, b_preshuffle_quant)."""
|
||||
combos = self.builder._generate_trait_combinations()
|
||||
for combo in combos:
|
||||
self.assertEqual(len(combo), 7, f"Expected 7-tuple, got {len(combo)}: {combo}")
|
||||
|
||||
def test_pipeline_is_compv3(self):
|
||||
combos = self.builder._generate_trait_combinations()
|
||||
for combo in combos:
|
||||
self.assertEqual(combo[0], "compv3")
|
||||
|
||||
def test_preshuffle_quant_values(self):
|
||||
combos = self.builder._generate_trait_combinations()
|
||||
preshuffle_vals = {c[6] for c in combos}
|
||||
self.assertIn(True, preshuffle_vals)
|
||||
self.assertIn(False, preshuffle_vals)
|
||||
|
||||
def test_scheduler_intrawave(self):
|
||||
combos = self.builder._generate_trait_combinations()
|
||||
schedulers = {c[2] for c in combos}
|
||||
self.assertIn("intrawave", schedulers)
|
||||
|
||||
|
||||
class TestGemmBQuantTileConfigs(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmpdir = tempfile.TemporaryDirectory()
|
||||
self.builder = _make_builder(self._tmpdir.name)
|
||||
|
||||
def tearDown(self):
|
||||
self._tmpdir.cleanup()
|
||||
|
||||
def test_tile_configs_non_empty(self):
|
||||
configs = self.builder._get_tile_configs()
|
||||
self.assertGreater(len(configs), 0)
|
||||
|
||||
def test_tile_config_has_required_keys(self):
|
||||
required = {
|
||||
"tile_m", "tile_n", "tile_k",
|
||||
"warp_m", "warp_n", "warp_k",
|
||||
"warp_tile_m", "warp_tile_n", "warp_tile_k",
|
||||
}
|
||||
for cfg in self.builder._get_tile_configs():
|
||||
self.assertTrue(required.issubset(cfg.keys()))
|
||||
|
||||
def test_tile_m_matches_config(self):
|
||||
configs = self.builder._get_tile_configs()
|
||||
tile_m_vals = {c["tile_m"] for c in configs}
|
||||
self.assertIn(64, tile_m_vals)
|
||||
|
||||
|
||||
class TestGemmBQuantLayoutVariants(unittest.TestCase):
|
||||
LAYOUTS = ["rcr", "rrr", "ccr", "crr"]
|
||||
|
||||
def test_all_layouts_construct(self):
|
||||
for layout in self.LAYOUTS:
|
||||
with self.subTest(layout=layout):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
builder = _make_builder(tmpdir, layout=layout)
|
||||
self.assertEqual(builder.layout, layout)
|
||||
|
||||
|
||||
class TestGemmBQuantDataTypes(unittest.TestCase):
|
||||
DTYPES = ["fp8", "bf8"]
|
||||
|
||||
def test_all_dtypes_construct(self):
|
||||
for dtype in self.DTYPES:
|
||||
with self.subTest(dtype=dtype):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
builder = _make_builder(tmpdir, datatype=dtype)
|
||||
self.assertEqual(builder.datatype, dtype)
|
||||
|
||||
|
||||
class TestGemmBQuantMaxInstances(unittest.TestCase):
|
||||
def test_max_instances_none_returns_all(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
cfg_path = os.path.join(tmpdir, "config.json")
|
||||
with open(cfg_path, "w") as f:
|
||||
json.dump(_MINIMAL_CONFIG, f)
|
||||
builder = GemmBQuantKernelBuilder(
|
||||
"gemm_bquant", tmpdir, "gfx942", "fp8", "rcr",
|
||||
config_json=cfg_path, max_instances=None,
|
||||
)
|
||||
kernel_list = [
|
||||
{"tile_config": {"tile_m": 64, "tile_n": 64, "tile_k": 128,
|
||||
"warp_m": 2, "warp_n": 2, "warp_k": 1,
|
||||
"warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 64},
|
||||
"trait_combo": ("compv3", "default", "intrawave", False, False, False, False)}
|
||||
]
|
||||
result = builder._apply_sampling(kernel_list)
|
||||
self.assertEqual(len(result), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -66,6 +66,7 @@ class GemmKernelBuilder:
|
||||
if config_json and os.path.exists(config_json):
|
||||
with open(config_json, "r") as f:
|
||||
self.config = json.load(f)
|
||||
self.group_size_k = self.config.get("group_size_k", 128)
|
||||
|
||||
def _apply_sampling(self, kernel_list):
|
||||
"""Apply RFC Sobol+LHS+maximin sampling. Returns sampled subset."""
|
||||
@@ -319,10 +320,12 @@ class GemmKernelBuilder:
|
||||
pipelines = ["preshufflev2"]
|
||||
elif self.kernel_name_prefix == "mx_gemm":
|
||||
pipelines = ["comp_async"]
|
||||
elif self.kernel_name_prefix in ["grouped_gemm_rowcolquant", "grouped_gemm_tensorquant", "gemm_rowcolquant", "gemm_tensor_quant"]:
|
||||
elif self.kernel_name_prefix in ["grouped_gemm_rowcolquant", "grouped_gemm_tensorquant", "gemm_rowcolquant", "gemm_tensor_quant", "gemm_aquant", "gemm_bquant"]:
|
||||
pipelines = ["compv3"]
|
||||
elif self.kernel_name_prefix in ["gemm_universal", "gemm_multi_d", "gemm_multi_abd", "grouped_gemm", "batched_contraction", "batched_gemm"]:
|
||||
pipelines = ["compv4"]
|
||||
elif self.kernel_name_prefix == "gemm_abquant":
|
||||
pipelines = ["compv3"]
|
||||
|
||||
configs = []
|
||||
for tile_m in tile_m_values:
|
||||
@@ -426,6 +429,7 @@ class GemmKernelBuilder:
|
||||
layout,
|
||||
self.gpu_target,
|
||||
self.kernel_name_prefix,
|
||||
self.config.get("group_size_k", 128),
|
||||
)
|
||||
|
||||
def _generate_trait_combinations(self):
|
||||
@@ -439,7 +443,15 @@ class GemmKernelBuilder:
|
||||
pad_m_values = trait_config.get("pad_m").get("values")
|
||||
pad_n_values = trait_config.get("pad_n").get("values")
|
||||
pad_k_values = trait_config.get("pad_k").get("values")
|
||||
if self.kernel_name_prefix in ["gemm_rowcolquant", "batched_gemm"]:
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
persistent_values = trait_config.get(
|
||||
"a_preshuffle_quant", {}
|
||||
).get("values", [False])
|
||||
elif self.kernel_name_prefix == "gemm_bquant":
|
||||
persistent_values = trait_config.get(
|
||||
"b_preshuffle_quant", {}
|
||||
).get("values", [False])
|
||||
elif self.kernel_name_prefix in ["gemm_rowcolquant", "batched_gemm"]:
|
||||
persistent_values = [
|
||||
False
|
||||
] # Force disable persistent where it is unsupported or not part of the trait key
|
||||
@@ -462,8 +474,14 @@ class GemmKernelBuilder:
|
||||
combinations = []
|
||||
for combo in all_combinations:
|
||||
pipeline, epilogue, scheduler = combo[:3]
|
||||
persistent_or_preshuffle_quant = combo[6] if len(combo) > 6 else False
|
||||
if is_trait_combination_valid(
|
||||
pipeline, epilogue, scheduler, self.kernel_name_prefix
|
||||
pipeline,
|
||||
epilogue,
|
||||
scheduler,
|
||||
persistent_or_preshuffle_quant,
|
||||
self.kernel_name_prefix,
|
||||
self.layout,
|
||||
):
|
||||
combinations.append(combo)
|
||||
else:
|
||||
@@ -484,7 +502,7 @@ class GemmKernelBuilder:
|
||||
pad_m,
|
||||
pad_n,
|
||||
pad_k,
|
||||
persistent,
|
||||
persistent_or_preshuffle_quant,
|
||||
) = self._normalize_trait_combo(trait_combo)
|
||||
|
||||
kernel_name = self._format_kernel_name(trait_combo, tile_config)
|
||||
@@ -523,6 +541,22 @@ class GemmKernelBuilder:
|
||||
"comp_async": "ck_tile::GemmPipelineAgBgCrCompAsync",
|
||||
}
|
||||
base_pipeline_map = {}
|
||||
elif self.kernel_name_prefix == "gemm_aquant":
|
||||
pipeline_impl_map = {
|
||||
"mem": "ck_tile::AQuantGemmPipelineAgBgCrMem",
|
||||
"compv3": "ck_tile::AQuantGemmPipelineAgBgCrCompV3",
|
||||
}
|
||||
base_pipeline_map = {
|
||||
"mem": "ck_tile::BaseGemmPipelineAgBgCrMem",
|
||||
"compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3",
|
||||
}
|
||||
elif self.kernel_name_prefix == "gemm_bquant":
|
||||
pipeline_impl_map = {
|
||||
"compv3": "ck_tile::BQuantGemmPipelineAgBgCrCompV3",
|
||||
}
|
||||
base_pipeline_map = {
|
||||
"compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3",
|
||||
}
|
||||
|
||||
scheduler_type_map = {
|
||||
"intrawave": "ck_tile::GemmPipelineScheduler::Intrawave",
|
||||
@@ -543,7 +577,7 @@ class GemmKernelBuilder:
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
persistent,
|
||||
persistent_or_preshuffle_quant,
|
||||
)
|
||||
|
||||
# Write into a file
|
||||
@@ -581,6 +615,13 @@ class GemmKernelBuilder:
|
||||
instance_code += """#include <vector>
|
||||
#include <hip/hip_runtime.h>
|
||||
#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp"
|
||||
"""
|
||||
elif self.kernel_name_prefix == "mx_gemm":
|
||||
instance_code += """#include "ck_tile/ops/gemm_mx.hpp"
|
||||
"""
|
||||
elif self.kernel_name_prefix in ["gemm_aquant", "gemm_bquant", "gemm_abquant"]:
|
||||
instance_code += """#include "ck_tile/ops/gemm_quant.hpp"
|
||||
#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp"
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
@@ -602,6 +643,9 @@ class GemmKernelBuilder:
|
||||
"grouped_gemm",
|
||||
"mx_gemm",
|
||||
"batched_gemm",
|
||||
"gemm_aquant",
|
||||
"gemm_bquant",
|
||||
"gemm_abquant",
|
||||
]:
|
||||
a_layout, b_layout, c_layout = get_abc_layouts(self.layout)
|
||||
|
||||
@@ -611,11 +655,21 @@ using BDataType = {get_dtype_string(self.datatype)};
|
||||
using AccDataType = {acc_type};
|
||||
using CDataType = {get_dtype_string(c_type)};"""
|
||||
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
# AQ scale factors are always fp32: the CK AQuant kernel expects float scales
|
||||
# regardless of the A/B element type. Changing this requires a matching kernel update.
|
||||
instance_code += """
|
||||
using AQDataType = float;"""
|
||||
|
||||
if self.kernel_name_prefix == "mx_gemm":
|
||||
instance_code += """
|
||||
using ScaleType = ck_tile::e8m0_t;
|
||||
using MxGemmHostArgs = ck_tile::MxGemmHostArgs<1, 1, 0>;"""
|
||||
|
||||
if self.kernel_name_prefix == "gemm_bquant":
|
||||
instance_code += """
|
||||
using BQDataType = float;"""
|
||||
|
||||
if self.kernel_name_prefix == "gemm_multi_d":
|
||||
instance_code += f"""
|
||||
using D0DataType = {get_dtype_string(self.datatype)};
|
||||
@@ -627,6 +681,17 @@ using ALayout = {a_layout};
|
||||
using BLayout = {b_layout};
|
||||
using CLayout = {c_layout};
|
||||
"""
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
instance_code += """using AQLayout = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
"""
|
||||
elif self.kernel_name_prefix == "gemm_bquant":
|
||||
# BQ scale tensor has shape [K/group_size_k, N], so its leading dimension
|
||||
# matches B's K dimension. Setting BQLayout = BLayout is only valid when B
|
||||
# is column-major (layout[1]=='c'); BPreshuffleQuant with row-major B is
|
||||
# blocked in is_trait_combination_valid to enforce this.
|
||||
instance_code += """using BQLayout = BLayout;
|
||||
"""
|
||||
|
||||
if self.kernel_name_prefix == "gemm_multi_d":
|
||||
instance_code += f"""
|
||||
using D0Layout = {ds_layout[0]};
|
||||
@@ -681,6 +746,20 @@ struct SelectedKernel {{
|
||||
static constexpr bool TransposeC = false;
|
||||
static constexpr bool DoubleSmemBuffer = {"true" if pipeline in ["compv4", "preshufflev2", "comp_async"] else "false"};"""
|
||||
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
instance_code += f"""
|
||||
static constexpr bool APreshuffleQuant = {"true" if persistent in [True, "true"] else "false"};
|
||||
static constexpr bool BPreshuffleQuant = false;
|
||||
static constexpr bool PreshuffleB = false;
|
||||
static constexpr ck_tile::index_t GroupSizeK = {self.config.get("group_size_k", 128)};"""
|
||||
|
||||
elif self.kernel_name_prefix == "gemm_bquant":
|
||||
instance_code += f"""
|
||||
static constexpr bool APreshuffleQuant = false;
|
||||
static constexpr bool BPreshuffleQuant = {"true" if persistent in [True, "true"] else "false"};
|
||||
static constexpr bool PreshuffleB = false;
|
||||
static constexpr ck_tile::index_t GroupSizeK = {self.config.get("group_size_k", 128)};"""
|
||||
|
||||
if self.kernel_name_prefix in [
|
||||
"gemm_universal",
|
||||
"gemm_preshuffle",
|
||||
@@ -700,11 +779,26 @@ struct SelectedKernel {{
|
||||
else:
|
||||
instance_code += """
|
||||
static constexpr bool Preshuffle = false;"""
|
||||
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
instance_code += f"""
|
||||
static constexpr bool APreshuffleQuant = {"true" if persistent_or_preshuffle_quant in [True, "true"] else "false"};
|
||||
static constexpr bool BPreshuffleQuant = false;
|
||||
static constexpr bool PreshuffleB = false;
|
||||
static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k};"""
|
||||
|
||||
elif self.kernel_name_prefix == "gemm_bquant":
|
||||
instance_code += f"""
|
||||
static constexpr bool APreshuffleQuant = false;
|
||||
static constexpr bool BPreshuffleQuant = {"true" if persistent_or_preshuffle_quant in [True, "true"] else "false"};
|
||||
static constexpr bool PreshuffleB = false;
|
||||
static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k};"""
|
||||
|
||||
return instance_code
|
||||
|
||||
def populate_initialization(self, base_pipeline_map, pipeline):
|
||||
# Tile Shape
|
||||
if self.kernel_name_prefix in ["gemm_multi_d", "batched_gemm"]:
|
||||
if self.kernel_name_prefix in ["gemm_multi_d", "batched_gemm", "gemm_aquant", "gemm_bquant", "gemm_abquant"]:
|
||||
instance_code = """
|
||||
|
||||
// Tile shape
|
||||
@@ -736,7 +830,39 @@ struct SelectedKernel {{
|
||||
using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner<TileShape, 8, 4>;"""
|
||||
|
||||
# Traits
|
||||
if self.kernel_name_prefix == "gemm_multi_d":
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
instance_code += """
|
||||
|
||||
// Quantization group size
|
||||
using QuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, GroupSizeK>>;
|
||||
|
||||
// Traits
|
||||
using Traits = ck_tile::TileGemmQuantTraits<
|
||||
kPadM, kPadN, kPadK,
|
||||
APreshuffleQuant,
|
||||
BPreshuffleQuant,
|
||||
PreshuffleB,
|
||||
ALayout, BLayout, CLayout,
|
||||
ck_tile::QuantType::AQuantGrouped,
|
||||
AQLayout>;"""
|
||||
|
||||
elif self.kernel_name_prefix == "gemm_bquant":
|
||||
instance_code += """
|
||||
|
||||
// Quantization group size
|
||||
using QuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, GroupSizeK>>;
|
||||
|
||||
// Traits
|
||||
using Traits = ck_tile::TileGemmQuantTraits<
|
||||
kPadM, kPadN, kPadK,
|
||||
APreshuffleQuant,
|
||||
BPreshuffleQuant,
|
||||
PreshuffleB,
|
||||
ALayout, BLayout, CLayout,
|
||||
ck_tile::QuantType::BQuantGrouped,
|
||||
BQLayout>;"""
|
||||
|
||||
elif self.kernel_name_prefix == "gemm_multi_d":
|
||||
instance_code += """
|
||||
|
||||
// Traits
|
||||
@@ -748,7 +874,19 @@ struct SelectedKernel {{
|
||||
using Traits = ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, CLayout, NumWaveGroups>;"""
|
||||
|
||||
# Pipeline problem
|
||||
if self.kernel_name_prefix in ["gemm_preshuffle", "gemm_multi_d"]:
|
||||
if self.kernel_name_prefix in ["gemm_aquant", "gemm_bquant"]:
|
||||
instance_code += """
|
||||
|
||||
// Pipeline problem (base, for hot loop detection)
|
||||
using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<
|
||||
ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
BDataType>;"""
|
||||
|
||||
elif self.kernel_name_prefix in ["gemm_preshuffle", "gemm_multi_d"]:
|
||||
instance_code += """
|
||||
|
||||
// Pipeline problem
|
||||
@@ -766,7 +904,7 @@ struct SelectedKernel {{
|
||||
// Base pipeline for hot loop detection
|
||||
using BaseGemmPipeline = {base_pipeline_map.get(pipeline, "ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2")}<GemmPipelineProblem>;"""
|
||||
|
||||
elif self.kernel_name_prefix == "gemm_multi_d":
|
||||
elif self.kernel_name_prefix in ["gemm_multi_d", "gemm_aquant", "gemm_bquant"]:
|
||||
instance_code += f"""
|
||||
|
||||
// Base pipeline for hot loop detection
|
||||
@@ -782,8 +920,28 @@ struct SelectedKernel {{
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
persistent,
|
||||
persistent_or_preshuffle_quant,
|
||||
):
|
||||
if self.kernel_name_prefix == "gemm_aquant":
|
||||
return self._populate_launch_aquant(
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
)
|
||||
|
||||
if self.kernel_name_prefix == "gemm_bquant":
|
||||
return self._populate_launch_bquant(
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
)
|
||||
|
||||
# Function Signature
|
||||
if self.kernel_name_prefix == "gemm_multi_d":
|
||||
instance_code = """
|
||||
@@ -959,7 +1117,7 @@ struct SelectedKernel {{
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = {"GemmKernel::MaxOccupancyGridSize(stream)" if persistent in [True, "true"] else "GemmKernel::GridSize(args.M, args.N, args.k_batch)"};
|
||||
const dim3 grids = {"GemmKernel::MaxOccupancyGridSize(stream)" if persistent_or_preshuffle_quant in [True, "true"] else "GemmKernel::GridSize(args.M, args.N, args.k_batch)"};
|
||||
const dim3 blocks = GemmKernel::BlockSize();
|
||||
|
||||
if(stream.log_level_ > 0) {{
|
||||
@@ -1032,7 +1190,7 @@ struct SelectedKernel {{
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = {"Kernel::MaxOccupancyGridSize(stream)" if persistent in [True, "true"] else "dim3(kargs.empty() ? 0 : kargs.back().block_end, 1, 1)"};
|
||||
const dim3 grids = {"Kernel::MaxOccupancyGridSize(stream)" if persistent_or_preshuffle_quant in [True, "true"] else "dim3(kargs.empty() ? 0 : kargs.back().block_end, 1, 1)"};
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
|
||||
@@ -1093,6 +1251,187 @@ struct SelectedKernel {{
|
||||
return ave_time;
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def _populate_launch_aquant(
|
||||
self,
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
):
|
||||
"""Generate the complete launch function for AQuant kernels."""
|
||||
instance_code = """
|
||||
|
||||
// Launch function
|
||||
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
|
||||
// Hot loop detection
|
||||
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
constexpr auto tail_number_v = tail_number_.value;
|
||||
|
||||
// Full pipeline problem with hot loop and tail number
|
||||
using PipelineProblem = ck_tile::GemmAQuantPipelineProblem<
|
||||
ADataType,
|
||||
AQDataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
QuantGroupSize,
|
||||
TransposeC,
|
||||
BDataType,
|
||||
{scheduler_type_map.get(scheduler)},
|
||||
has_hot_loop_v,
|
||||
tail_number_v>;
|
||||
|
||||
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
|
||||
|
||||
instance_code += """
|
||||
|
||||
// Epilogue"""
|
||||
if epilogue == "cshuffle":
|
||||
instance_code += self.populate_cshuffle_gemm_aquant()
|
||||
else:
|
||||
instance_code += self.populate_default_gemm_aquant()
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
// Kernel type
|
||||
using Kernel = ck_tile::QuantGemmKernel<
|
||||
TilePartitioner, GemmPipeline, GemmEpilogue,
|
||||
ck_tile::QuantType::AQuantGrouped>;
|
||||
|
||||
// Kernel arguments
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
|
||||
if (!Kernel::IsSupportedArgument(kargs)) {{
|
||||
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
if(stream.log_level_ > 0) {{
|
||||
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
|
||||
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
|
||||
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
|
||||
<< std::endl;
|
||||
}}
|
||||
|
||||
// Launch kernel
|
||||
constexpr int kBlockPerCu = {k_block_per_cu};
|
||||
float ave_time = ck_tile::launch_kernel(
|
||||
stream,
|
||||
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
|
||||
|
||||
return ave_time;
|
||||
}};
|
||||
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def _populate_launch_bquant(
|
||||
self,
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
):
|
||||
"""Generate the complete launch function for BQuant kernels."""
|
||||
instance_code = """
|
||||
|
||||
// Launch function
|
||||
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
|
||||
// Hot loop detection
|
||||
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
constexpr auto tail_number_v = tail_number_.value;
|
||||
|
||||
// Full pipeline problem with hot loop and tail number
|
||||
using PipelineProblem = ck_tile::GemmBQuantPipelineProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
BQDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
QuantGroupSize,
|
||||
ADataType,
|
||||
{scheduler_type_map.get(scheduler)},
|
||||
has_hot_loop_v,
|
||||
tail_number_v>;
|
||||
|
||||
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
|
||||
|
||||
instance_code += """
|
||||
|
||||
// Epilogue"""
|
||||
if epilogue == "cshuffle":
|
||||
instance_code += self.populate_cshuffle_gemm_bquant()
|
||||
else:
|
||||
instance_code += self.populate_default_gemm_bquant()
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
// Kernel type
|
||||
using Kernel = ck_tile::QuantGemmKernel<
|
||||
TilePartitioner, GemmPipeline, GemmEpilogue,
|
||||
ck_tile::QuantType::BQuantGrouped>;
|
||||
|
||||
// Kernel arguments
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
|
||||
if (!Kernel::IsSupportedArgument(kargs)) {{
|
||||
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
if(stream.log_level_ > 0) {{
|
||||
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
|
||||
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
|
||||
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
|
||||
<< std::endl;
|
||||
}}
|
||||
|
||||
// Launch kernel
|
||||
constexpr int kBlockPerCu = {k_block_per_cu};
|
||||
float ave_time = ck_tile::launch_kernel(
|
||||
stream,
|
||||
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
|
||||
|
||||
return ave_time;
|
||||
}};
|
||||
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
@@ -1122,6 +1461,8 @@ struct SelectedKernel {{
|
||||
instance_code += self.populate_default_gemm_preshuffle()
|
||||
elif self.kernel_name_prefix == "mx_gemm":
|
||||
raise ValueError("MX GEMM currently supports only cshuffle epilogue")
|
||||
elif self.kernel_name_prefix == "gemm_aquant":
|
||||
instance_code += self.populate_default_gemm_aquant()
|
||||
|
||||
return instance_code
|
||||
|
||||
@@ -1322,6 +1663,278 @@ struct SelectedKernel {{
|
||||
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
|
||||
return instance_code
|
||||
|
||||
def _populate_launch_aquant(
|
||||
self,
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
):
|
||||
instance_code = """
|
||||
|
||||
// Launch function
|
||||
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
|
||||
// Hot loop detection
|
||||
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
constexpr auto tail_number_v = tail_number_.value;
|
||||
|
||||
// Full pipeline problem with hot loop and tail number
|
||||
using PipelineProblem = ck_tile::GemmAQuantPipelineProblem<
|
||||
ADataType,
|
||||
AQDataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
QuantGroupSize,
|
||||
TransposeC,
|
||||
BDataType,
|
||||
{scheduler_type_map.get(scheduler)},
|
||||
has_hot_loop_v,
|
||||
tail_number_v>;
|
||||
|
||||
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
|
||||
|
||||
instance_code += """
|
||||
|
||||
// Epilogue"""
|
||||
if epilogue == "cshuffle":
|
||||
instance_code += self.populate_cshuffle_gemm_aquant()
|
||||
else:
|
||||
instance_code += self.populate_default_gemm_aquant()
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
// Kernel type
|
||||
using Kernel = ck_tile::QuantGemmKernel<
|
||||
TilePartitioner, GemmPipeline, GemmEpilogue,
|
||||
ck_tile::QuantType::AQuantGrouped>;
|
||||
|
||||
// Kernel arguments
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
|
||||
if (!Kernel::IsSupportedArgument(kargs)) {{
|
||||
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
if(stream.log_level_ > 0) {{
|
||||
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
|
||||
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
|
||||
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
|
||||
<< std::endl;
|
||||
}}
|
||||
|
||||
// Launch kernel
|
||||
constexpr int kBlockPerCu = {k_block_per_cu};
|
||||
float ave_time = ck_tile::launch_kernel(
|
||||
stream,
|
||||
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
|
||||
|
||||
return ave_time;
|
||||
}};
|
||||
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def populate_default_gemm_aquant(self):
|
||||
instance_code = """
|
||||
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>, // DsDataType
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>, // DsLayout
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TileM, // kM_
|
||||
TileN, // kN_
|
||||
kPadM,
|
||||
kPadN,
|
||||
WarpTileM, // kMPerXdl_
|
||||
WarpTileN, // kNPerXdl_
|
||||
WarpTileK, // kKPerXdl_
|
||||
TransposeC>; // isCTransposed_
|
||||
|
||||
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
|
||||
return instance_code
|
||||
|
||||
def populate_cshuffle_gemm_aquant(self):
|
||||
instance_code = """
|
||||
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>, // DsDataType
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>, // DsLayout
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TileM, // kM_
|
||||
TileN, // kN_
|
||||
WarpPerBlock_M, // MWave_
|
||||
WarpPerBlock_N, // NWave_
|
||||
WarpTileM, // kMPerXdl_
|
||||
WarpTileN, // kNPerXdl_
|
||||
WarpTileK, // kKPerXdl_
|
||||
TransposeC>; // isCTransposed_
|
||||
|
||||
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;"""
|
||||
return instance_code
|
||||
|
||||
def _populate_launch_bquant(
|
||||
self,
|
||||
scheduler_type_map,
|
||||
scheduler,
|
||||
pipeline_impl_map,
|
||||
pipeline,
|
||||
epilogue,
|
||||
k_block_per_cu,
|
||||
):
|
||||
instance_code = """
|
||||
|
||||
// Launch function
|
||||
static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) {
|
||||
|
||||
// Hot loop detection
|
||||
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);"""
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
constexpr auto tail_number_v = tail_number_.value;
|
||||
|
||||
// Full pipeline problem with hot loop and tail number
|
||||
using PipelineProblem = ck_tile::GemmBQuantPipelineProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
BQDataType,
|
||||
AccDataType,
|
||||
TileShape,
|
||||
Traits,
|
||||
QuantGroupSize,
|
||||
TransposeC,
|
||||
BDataType,
|
||||
{scheduler_type_map.get(scheduler)},
|
||||
has_hot_loop_v,
|
||||
tail_number_v>;
|
||||
|
||||
using GemmPipeline = {pipeline_impl_map.get(pipeline)}<PipelineProblem>;"""
|
||||
|
||||
instance_code += """
|
||||
|
||||
// Epilogue"""
|
||||
if epilogue == "cshuffle":
|
||||
instance_code += self.populate_cshuffle_gemm_bquant()
|
||||
else:
|
||||
instance_code += self.populate_default_gemm_bquant()
|
||||
|
||||
instance_code += f"""
|
||||
|
||||
// Kernel type
|
||||
using Kernel = ck_tile::QuantGemmKernel<
|
||||
TilePartitioner, GemmPipeline, GemmEpilogue,
|
||||
ck_tile::QuantType::BQuantGrouped>;
|
||||
|
||||
// Kernel arguments
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
|
||||
if (!Kernel::IsSupportedArgument(kargs)) {{
|
||||
throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch.");
|
||||
}}
|
||||
|
||||
// Get grid and block sizes
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
if(stream.log_level_ > 0) {{
|
||||
std::cout << "Launching kernel: " << Kernel::GetName() << '\\n'
|
||||
<< "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}"
|
||||
<< ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}"
|
||||
<< std::endl;
|
||||
}}
|
||||
|
||||
// Launch kernel
|
||||
constexpr int kBlockPerCu = {k_block_per_cu};
|
||||
float ave_time = ck_tile::launch_kernel(
|
||||
stream,
|
||||
ck_tile::make_kernel<kBlockPerCu>(Kernel{{}}, grids, blocks, 0, kargs));
|
||||
|
||||
return ave_time;
|
||||
}};
|
||||
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
return instance_code
|
||||
|
||||
def populate_default_gemm_bquant(self):
|
||||
instance_code = """
|
||||
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>, // DsDataType
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>, // DsLayout
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TileM, // kM_
|
||||
TileN, // kN_
|
||||
kPadM,
|
||||
kPadN,
|
||||
WarpTileM, // kMPerXdl_
|
||||
WarpTileN, // kNPerXdl_
|
||||
WarpTileK, // kKPerXdl_
|
||||
TransposeC>; // isCTransposed_
|
||||
|
||||
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;"""
|
||||
return instance_code
|
||||
|
||||
def populate_cshuffle_gemm_bquant(self):
|
||||
instance_code = """
|
||||
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<
|
||||
ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>, // DsDataType
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>, // DsLayout
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TileM, // kM_
|
||||
TileN, // kN_
|
||||
WarpPerBlock_M, // MWave_
|
||||
WarpPerBlock_N, // NWave_
|
||||
WarpTileM, // kMPerXdl_
|
||||
WarpTileN, // kNPerXdl_
|
||||
WarpTileK, // kKPerXdl_
|
||||
TransposeC>; // isCTransposed_
|
||||
|
||||
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;"""
|
||||
return instance_code
|
||||
|
||||
def _generate_cmake_individual_targets(self, kernel_list):
|
||||
"""Generate CMake include file that creates individual targets"""
|
||||
cmake_code = f"""# Generated CMake file for individual {self.kernel_name_prefix} targets
|
||||
|
||||
@@ -10,6 +10,9 @@ GEMM_PRESHUFFLE_PIPELINES = ["preshufflev2"]
|
||||
|
||||
GEMM_ROWCOLQUANT_PIPELINES = ["compv3"]
|
||||
GEMM_MX_PIPELINES = ["comp_async"]
|
||||
GEMM_BQUANT_PIPELINES = ["compv3"]
|
||||
|
||||
GEMM_ABQUANT_PIPELINES = ["compv3"]
|
||||
|
||||
LAYOUT_MAP = {
|
||||
"r": "ck_tile::tensor_layout::gemm::RowMajor",
|
||||
@@ -231,6 +234,24 @@ TRAIT_UNSUPPORTED_COMBINATIONS = {
|
||||
("comp_async", "cshuffle", "interwave"),
|
||||
}
|
||||
|
||||
AQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
|
||||
("mem", "default", "interwave"),
|
||||
("mem", "cshuffle", "interwave"),
|
||||
("compv3", "default", "interwave"),
|
||||
("compv3", "cshuffle", "interwave"),
|
||||
}
|
||||
|
||||
BQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
|
||||
("compv3", "default", "interwave"),
|
||||
("compv3", "cshuffle", "interwave"),
|
||||
}
|
||||
|
||||
ABQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
|
||||
("compv3", "default", "interwave"),
|
||||
("compv3", "cshuffle", "interwave"),
|
||||
}
|
||||
|
||||
|
||||
def element_size(data_type: str) -> float:
|
||||
"""Calculate the size (in bytes) of a single element for given data type."""
|
||||
data_type = data_type.lower()
|
||||
@@ -240,10 +261,32 @@ def element_size(data_type: str) -> float:
|
||||
|
||||
|
||||
def is_trait_combination_valid(
|
||||
pipeline: str, epilogue: str, scheduler: str, kernel_name_prefix: str = ""
|
||||
pipeline: str,
|
||||
epilogue: str,
|
||||
scheduler: str,
|
||||
persistent_or_preshuffle_quant=None,
|
||||
kernel_name_prefix: str = "",
|
||||
layout: str = "",
|
||||
) -> bool:
|
||||
"""Check if a trait combination is valid."""
|
||||
if (
|
||||
if kernel_name_prefix == "gemm_aquant":
|
||||
if (pipeline, epilogue, scheduler) in AQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
|
||||
return False
|
||||
# mem pipeline does not support preshuffle
|
||||
if pipeline == "mem" and persistent_or_preshuffle_quant is True:
|
||||
return False
|
||||
return True
|
||||
elif kernel_name_prefix == "gemm_bquant":
|
||||
if (pipeline, epilogue, scheduler) in BQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
|
||||
return False
|
||||
# bquant only supports compv3 + intrawave
|
||||
if pipeline != "compv3" or scheduler != "intrawave":
|
||||
return False
|
||||
# BPreshuffleQuant requires ColumnMajor BLayout (second char of layout is 'c')
|
||||
if persistent_or_preshuffle_quant is True and len(layout) >= 2 and layout[1] != "c":
|
||||
return False
|
||||
return True
|
||||
elif (
|
||||
kernel_name_prefix == "gemm_rowcolquant"
|
||||
or kernel_name_prefix == "grouped_gemm_rowcolquant"
|
||||
or kernel_name_prefix == "grouped_gemm_tensorquant"
|
||||
@@ -252,6 +295,16 @@ def is_trait_combination_valid(
|
||||
if pipeline != "compv3" or scheduler != "intrawave" or epilogue != "cshuffle":
|
||||
return False
|
||||
return True
|
||||
elif kernel_name_prefix == "gemm_abquant":
|
||||
if (pipeline, epilogue, scheduler) in ABQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
|
||||
return False
|
||||
# abquant only supports compv3 + intrawave
|
||||
if pipeline != "compv3" or scheduler != "intrawave":
|
||||
return False
|
||||
# BPreshuffleQuant requires ColumnMajor BLayout (second char of layout is 'c')
|
||||
if persistent_or_preshuffle_quant is True and len(layout) >= 2 and layout[1] != "c":
|
||||
return False
|
||||
return True
|
||||
else:
|
||||
return (pipeline, epilogue, scheduler) not in TRAIT_UNSUPPORTED_COMBINATIONS
|
||||
|
||||
@@ -493,6 +546,7 @@ def is_tile_config_valid(
|
||||
layout: str,
|
||||
gpu_target: str,
|
||||
kernel_name_prefix: str = "",
|
||||
group_size_k: int = 128,
|
||||
) -> bool:
|
||||
"""
|
||||
Comprehensive tile configuration validation.
|
||||
@@ -691,6 +745,75 @@ def is_tile_config_valid(
|
||||
)
|
||||
return False
|
||||
|
||||
if kernel_name_prefix == "gemm_aquant":
|
||||
aquant_valid, aquant_valid_error = validate_gemm_aquant(
|
||||
tile_m,
|
||||
tile_n,
|
||||
tile_k,
|
||||
warp_m,
|
||||
warp_n,
|
||||
warp_k,
|
||||
warp_tile_m,
|
||||
warp_tile_n,
|
||||
warp_tile_k,
|
||||
a_datatype,
|
||||
b_datatype,
|
||||
c_datatype,
|
||||
pipeline,
|
||||
layout,
|
||||
gpu_target,
|
||||
group_size_k,
|
||||
)
|
||||
if not aquant_valid:
|
||||
logging.debug(f"GEMM AQuant validation failed: {aquant_valid_error}")
|
||||
return False
|
||||
|
||||
elif kernel_name_prefix == "gemm_bquant":
|
||||
bquant_valid, bquant_valid_error = validate_gemm_bquant(
|
||||
tile_m,
|
||||
tile_n,
|
||||
tile_k,
|
||||
warp_m,
|
||||
warp_n,
|
||||
warp_k,
|
||||
warp_tile_m,
|
||||
warp_tile_n,
|
||||
warp_tile_k,
|
||||
a_datatype,
|
||||
b_datatype,
|
||||
c_datatype,
|
||||
pipeline,
|
||||
layout,
|
||||
gpu_target,
|
||||
group_size_k,
|
||||
)
|
||||
if not bquant_valid:
|
||||
logging.debug(f"GEMM BQuant validation failed: {bquant_valid_error}")
|
||||
return False
|
||||
|
||||
elif kernel_name_prefix == "gemm_abquant":
|
||||
abquant_valid, abquant_valid_error = validate_gemm_abquant(
|
||||
tile_m,
|
||||
tile_n,
|
||||
tile_k,
|
||||
warp_m,
|
||||
warp_n,
|
||||
warp_k,
|
||||
warp_tile_m,
|
||||
warp_tile_n,
|
||||
warp_tile_k,
|
||||
a_datatype,
|
||||
b_datatype,
|
||||
c_datatype,
|
||||
pipeline,
|
||||
layout,
|
||||
gpu_target,
|
||||
group_size_k,
|
||||
)
|
||||
if not abquant_valid:
|
||||
logging.debug(f"GEMM ABQuant validation failed: {abquant_valid_error}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@@ -1259,6 +1382,47 @@ def validate_m0_m1_m2_configuration(
|
||||
return False, f"Error in M0/M1/M2 validation: {str(e)}"
|
||||
|
||||
|
||||
def _validate_fp8_mfma_warp_tile_k(
|
||||
warp_tile_m: int,
|
||||
warp_tile_n: int,
|
||||
warp_tile_k: int,
|
||||
a_datatype: str,
|
||||
gpu_target: str,
|
||||
op_label: str = "",
|
||||
) -> Tuple[bool, str]:
|
||||
"""Validate MFMA warp-tile constraints shared by all fp8/bf8 quantized GEMM ops.
|
||||
|
||||
Checks:
|
||||
- warp_tile_m == warp_tile_n (square MFMA requirement)
|
||||
- warp_tile_k matches the ISA-mandated K-block for the given warp_tile_m and gpu_target
|
||||
"""
|
||||
suffix = f" ({op_label})" if op_label else ""
|
||||
if warp_tile_m != warp_tile_n:
|
||||
return False, (
|
||||
f"warp_tile_m({warp_tile_m}) must equal warp_tile_n({warp_tile_n})"
|
||||
f" — MFMA requires a square warp tile{suffix}"
|
||||
)
|
||||
|
||||
if a_datatype in ["fp8", "bf8"]:
|
||||
# MFMA instruction shapes for fp8/bf8 (from CDNA ISA reference manual):
|
||||
# gfx90a/gfx942: MFMA_F32_16x16x128_F8 (warp_tile_m=16) → warp_tile_k=64
|
||||
# MFMA_F32_32x32x64_F8 (warp_tile_m=32) → warp_tile_k=32
|
||||
# gfx950 doubles the K-block:
|
||||
# MFMA_F32_16x16x256_F8 (warp_tile_m=16) → warp_tile_k=128
|
||||
# MFMA_F32_32x32x128_F8 (warp_tile_m=32) → warp_tile_k=64
|
||||
if gpu_target == "gfx950":
|
||||
expected_k = 64 if warp_tile_m == 32 else 128
|
||||
else:
|
||||
expected_k = 32 if warp_tile_m == 32 else 64
|
||||
if warp_tile_k != expected_k:
|
||||
return False, (
|
||||
f"For {a_datatype} on {gpu_target}, warp_tile_m={warp_tile_m} "
|
||||
f"requires warp_tile_k={expected_k}, got warp_tile_k={warp_tile_k}{suffix}"
|
||||
)
|
||||
|
||||
return True, ""
|
||||
|
||||
|
||||
def validate_gemm_rowcol_tensor_quant(
|
||||
tile_m: int,
|
||||
tile_n: int,
|
||||
@@ -1294,21 +1458,171 @@ def validate_gemm_rowcol_tensor_quant(
|
||||
if not whole_workgroup_cover_valid:
|
||||
return False, whole_workgroup_cover_error
|
||||
|
||||
if warp_tile_m != warp_tile_n:
|
||||
return _validate_fp8_mfma_warp_tile_k(
|
||||
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
|
||||
op_label="RowColQuant / TensorQuant"
|
||||
)
|
||||
|
||||
|
||||
def validate_gemm_aquant(
|
||||
tile_m: int,
|
||||
tile_n: int,
|
||||
tile_k: int,
|
||||
warp_m: int,
|
||||
warp_n: int,
|
||||
warp_k: int,
|
||||
warp_tile_m: int,
|
||||
warp_tile_n: int,
|
||||
warp_tile_k: int,
|
||||
a_datatype: str,
|
||||
b_datatype: str,
|
||||
c_datatype: str,
|
||||
pipeline: str,
|
||||
layout: str,
|
||||
gpu_target: str,
|
||||
group_size_k: int = 128,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Validate AQuant GEMM-specific constraints."""
|
||||
|
||||
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
|
||||
validate_whole_wg_cover_configuration(
|
||||
tile_m,
|
||||
tile_n,
|
||||
tile_k,
|
||||
warp_m,
|
||||
warp_n,
|
||||
warp_k,
|
||||
layout,
|
||||
a_datatype,
|
||||
b_datatype,
|
||||
gpu_target,
|
||||
)
|
||||
)
|
||||
if not whole_workgroup_cover_valid:
|
||||
logging.debug(
|
||||
f"Whole workgroup cover configuration validation failed: {whole_workgroup_cover_error}"
|
||||
)
|
||||
return False, whole_workgroup_cover_error
|
||||
|
||||
if tile_k % group_size_k != 0 or tile_k < group_size_k:
|
||||
return False, (
|
||||
f"warp_tile_m({warp_tile_m}) must be equal to warp_tile_n({warp_tile_n}) "
|
||||
f"(MFMA requirement for RowColQuant / TensorQuant)"
|
||||
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
|
||||
f"and tile_k >= group_size_k"
|
||||
)
|
||||
|
||||
if a_datatype in ["fp8", "bf8"]:
|
||||
if gpu_target == "gfx950":
|
||||
expected_k = 64 if warp_tile_m == 32 else 128
|
||||
else:
|
||||
expected_k = 32 if warp_tile_m == 32 else 64
|
||||
if warp_tile_k != expected_k:
|
||||
return False, (
|
||||
f"For {a_datatype} on {gpu_target}, warp_tile_m={warp_tile_m} "
|
||||
f"requires warp_tile_k={expected_k}, got warp_tile_k={warp_tile_k}"
|
||||
)
|
||||
if group_size_k % warp_tile_k != 0:
|
||||
return False, (
|
||||
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
|
||||
)
|
||||
|
||||
return True, ""
|
||||
return _validate_fp8_mfma_warp_tile_k(
|
||||
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
|
||||
op_label="AQuant"
|
||||
)
|
||||
|
||||
|
||||
def validate_gemm_bquant(
|
||||
tile_m: int,
|
||||
tile_n: int,
|
||||
tile_k: int,
|
||||
warp_m: int,
|
||||
warp_n: int,
|
||||
warp_k: int,
|
||||
warp_tile_m: int,
|
||||
warp_tile_n: int,
|
||||
warp_tile_k: int,
|
||||
a_datatype: str,
|
||||
b_datatype: str,
|
||||
c_datatype: str,
|
||||
pipeline: str,
|
||||
layout: str,
|
||||
gpu_target: str,
|
||||
group_size_k: int = 128,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Validate BQuant GEMM-specific constraints."""
|
||||
|
||||
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
|
||||
validate_whole_wg_cover_configuration(
|
||||
tile_m,
|
||||
tile_n,
|
||||
tile_k,
|
||||
warp_m,
|
||||
warp_n,
|
||||
warp_k,
|
||||
layout,
|
||||
a_datatype,
|
||||
b_datatype,
|
||||
gpu_target,
|
||||
)
|
||||
)
|
||||
if not whole_workgroup_cover_valid:
|
||||
return False, whole_workgroup_cover_error
|
||||
|
||||
if tile_k % group_size_k != 0 or tile_k < group_size_k:
|
||||
return False, (
|
||||
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
|
||||
f"and tile_k >= group_size_k"
|
||||
)
|
||||
|
||||
if group_size_k % warp_tile_k != 0:
|
||||
return False, (
|
||||
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
|
||||
)
|
||||
|
||||
return _validate_fp8_mfma_warp_tile_k(
|
||||
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
|
||||
op_label="BQuant"
|
||||
)
|
||||
|
||||
|
||||
def validate_gemm_abquant(
|
||||
tile_m: int,
|
||||
tile_n: int,
|
||||
tile_k: int,
|
||||
warp_m: int,
|
||||
warp_n: int,
|
||||
warp_k: int,
|
||||
warp_tile_m: int,
|
||||
warp_tile_n: int,
|
||||
warp_tile_k: int,
|
||||
a_datatype: str,
|
||||
b_datatype: str,
|
||||
c_datatype: str,
|
||||
pipeline: str,
|
||||
layout: str,
|
||||
gpu_target: str,
|
||||
group_size_k: int = 128,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Validate ABQuant GEMM-specific constraints."""
|
||||
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
|
||||
validate_whole_wg_cover_configuration(
|
||||
tile_m,
|
||||
tile_n,
|
||||
tile_k,
|
||||
warp_m,
|
||||
warp_n,
|
||||
warp_k,
|
||||
layout,
|
||||
a_datatype,
|
||||
b_datatype,
|
||||
gpu_target,
|
||||
)
|
||||
)
|
||||
if not whole_workgroup_cover_valid:
|
||||
return False, whole_workgroup_cover_error
|
||||
|
||||
if tile_k % group_size_k != 0 or tile_k < group_size_k:
|
||||
return False, (
|
||||
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
|
||||
f"and tile_k >= group_size_k"
|
||||
)
|
||||
|
||||
if group_size_k % warp_tile_k != 0:
|
||||
return False, (
|
||||
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
|
||||
)
|
||||
|
||||
return _validate_fp8_mfma_warp_tile_k(
|
||||
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
|
||||
op_label="ABQuant"
|
||||
)
|
||||
|
||||
@@ -7,4 +7,5 @@ from sampling.budget import allocate_budget as allocate_budget
|
||||
from sampling.budget import load_op_weights as load_op_weights
|
||||
from sampling.manifest import write_manifest as write_manifest
|
||||
from sampling.feasible_set import GEMM_AXES as GEMM_AXES
|
||||
from sampling.feasible_set import GEMM_ABQUANT_AXES as GEMM_ABQUANT_AXES
|
||||
from sampling.feasible_set import normalize_axis_values as normalize_axis_values
|
||||
|
||||
@@ -22,6 +22,27 @@ GEMM_AXES = [
|
||||
|
||||
GEMM_STREAMK_AXES = GEMM_AXES + ["reduction_strategy"]
|
||||
|
||||
GEMM_ABQUANT_AXES = [
|
||||
"tile_m",
|
||||
"tile_n",
|
||||
"tile_k",
|
||||
"warp_m",
|
||||
"warp_n",
|
||||
"warp_k",
|
||||
"warp_tile_m",
|
||||
"warp_tile_n",
|
||||
"warp_tile_k",
|
||||
"pipeline",
|
||||
"epilogue",
|
||||
"scheduler",
|
||||
"pad_m",
|
||||
"pad_n",
|
||||
"pad_k",
|
||||
"a_preshuffle_quant",
|
||||
"b_preshuffle_quant",
|
||||
"group_size_n",
|
||||
]
|
||||
|
||||
CATEGORICAL_AXES = {
|
||||
"pipeline",
|
||||
"epilogue",
|
||||
@@ -31,6 +52,8 @@ CATEGORICAL_AXES = {
|
||||
"pad_n",
|
||||
"pad_k",
|
||||
"persistent",
|
||||
"a_preshuffle_quant",
|
||||
"b_preshuffle_quant",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -2,18 +2,21 @@
|
||||
"version": "1.0",
|
||||
"description": "Op weights for Daily Tier budget allocation (RFC section 3.4)",
|
||||
"weights": {
|
||||
"gemm_universal": 0.0770,
|
||||
"gemm_preshuffle": 0.0770,
|
||||
"gemm_multi_d": 0.0770,
|
||||
"grouped_gemm": 0.0769,
|
||||
"gemm_streamk": 0.0769,
|
||||
"batched_contraction": 0.0769,
|
||||
"batched_gemm": 0.0769,
|
||||
"gemm_multi_abd": 0.0769,
|
||||
"mx_gemm": 0.0769,
|
||||
"gemm_rowcolquant": 0.0769,
|
||||
"gemm_tensor_quant": 0.0769,
|
||||
"grouped_gemm_rowcolquant": 0.0769,
|
||||
"grouped_gemm_tensorquant": 0.0769
|
||||
"gemm_universal": 0.0625,
|
||||
"gemm_preshuffle": 0.0625,
|
||||
"gemm_multi_d": 0.0625,
|
||||
"grouped_gemm": 0.0625,
|
||||
"gemm_streamk": 0.0625,
|
||||
"batched_contraction": 0.0625,
|
||||
"batched_gemm": 0.0625,
|
||||
"gemm_multi_abd": 0.0625,
|
||||
"mx_gemm": 0.0625,
|
||||
"gemm_rowcolquant": 0.0625,
|
||||
"gemm_tensor_quant": 0.0625,
|
||||
"grouped_gemm_rowcolquant": 0.0625,
|
||||
"grouped_gemm_tensorquant": 0.0625,
|
||||
"gemm_aquant": 0.0625,
|
||||
"gemm_bquant": 0.0625,
|
||||
"gemm_abquant": 0.0625
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user