From b0f200713a0b75af6e09d17f97274830fabbe0e9 Mon Sep 17 00:00:00 2001 From: Thrupti Raj Lakshmana Gowda <170563851+ThruptiRajLakshmanaGowda@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:22:48 +0000 Subject: [PATCH] [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 --- tile_engine/operation_support_matrix.md | 3 + tile_engine/ops/gemm/CMakeLists.txt | 4 +- .../ops/gemm/block_scale_gemm/CMakeLists.txt | 3 + .../gemm_abquant/CMakeLists.txt | 272 ++++++ .../configs/default_ci_config.json | 104 ++ .../gemm_abquant/configs/default_config.json | 113 +++ .../gemm_abquant/configs/example_config.json | 98 ++ .../gemm_abquant/gemm_abquant_benchmark.hpp | 243 +++++ .../gemm_abquant/gemm_abquant_benchmark.py | 477 +++++++++ .../gemm_abquant_benchmark_single.cpp | 178 ++++ .../gemm_abquant/gemm_abquant_common.hpp | 76 ++ .../gemm_abquant_instance_builder.py | 904 ++++++++++++++++++ .../gemm_abquant/gemm_abquant_profiler.hpp | 336 +++++++ .../test_gemm_abquant_instance_builder.py | 236 +++++ .../gemm_aquant/CMakeLists.txt | 297 ++++++ .../configs/default_ci_config.json | 92 ++ .../gemm_aquant/configs/default_config.json | 102 ++ .../gemm_aquant/configs/example_config.json | 88 ++ .../gemm_aquant/gemm_aquant_benchmark.hpp | 231 +++++ .../gemm_aquant/gemm_aquant_benchmark.py | 547 +++++++++++ .../gemm_aquant_benchmark_single.cpp | 153 +++ .../gemm_aquant/gemm_aquant_common.hpp | 73 ++ .../gemm_aquant_instance_builder.py | 322 +++++++ .../gemm_aquant/gemm_aquant_profiler.hpp | 276 ++++++ .../test_gemm_aquant_instance_builder.py | 204 ++++ .../gemm_bquant/CMakeLists.txt | 289 ++++++ .../configs/default_ci_config.json | 93 ++ .../gemm_bquant/configs/default_config.json | 102 ++ .../gemm_bquant/configs/example_config.json | 88 ++ .../gemm_bquant/gemm_bquant_benchmark.hpp | 229 +++++ .../gemm_bquant/gemm_bquant_benchmark.py | 464 +++++++++ .../gemm_bquant_benchmark_single.cpp | 166 ++++ .../gemm_bquant/gemm_bquant_common.hpp | 74 ++ .../gemm_bquant_instance_builder.py | 288 ++++++ .../gemm_bquant/gemm_bquant_profiler.hpp | 293 ++++++ .../test_gemm_bquant_instance_builder.py | 191 ++++ tile_engine/ops/gemm/gemm_instance_builder.py | 637 +++++++++++- tile_engine/ops/gemm/gemm_validation_utils.py | 346 ++++++- tile_engine/sampling/__init__.py | 1 + tile_engine/sampling/feasible_set.py | 23 + tile_engine/sampling/op_weights.json | 29 +- 41 files changed, 8702 insertions(+), 43 deletions(-) create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp create mode 100644 tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/test_gemm_bquant_instance_builder.py diff --git a/tile_engine/operation_support_matrix.md b/tile_engine/operation_support_matrix.md index 6444956ad9..3fb3379ad7 100644 --- a/tile_engine/operation_support_matrix.md +++ b/tile_engine/operation_support_matrix.md @@ -13,6 +13,9 @@ | GEMM | flatmm
example: 18_flatmm/ | ❌ | ❌ | ❌ | ❌ | | ❌ | ❌ | ❌ | | | | ❌ | ❌ | ❌ | ❌ | | | GEMM | gemm_multi_abd
example: 22_gemm_multi_abd/ | ✅ | | | | | | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | 0.0833 | | GEMM | gemm_quant | | ❌ | | ❌ | | | | ❌ | | | | ❌ | ❌ | ❌ | ❌ | | +| GEMM | block_scale_gemm/gemm_aquant
engine: block_scale_gemm/gemm_aquant/
example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | | | | ❌ | ✅ | ✅ | ❌ | 0.0625 | +| GEMM | block_scale_gemm/gemm_bquant
engine: block_scale_gemm/gemm_bquant/
example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | 0.0625 | +| GEMM | block_scale_gemm/gemm_abquant
engine: block_scale_gemm/gemm_abquant/
example: 38_block_scale_gemm/ | | ✅ | | ✅ | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | 0.0625 | | GEMM | grouped_gemm [10]
engine: grouped_gemm/
example: 17_grouped_gemm/ | ✅ | ✅ | | | | | | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | 0.0834 | | GEMM | grouped_gemm_quant/grouped_gemm_rowcolquant
engine: grouped_gemm_quant/grouped_gemm_rowcolquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 | | GEMM | grouped_gemm_quant/grouped_gemm_tensorquant
engine: grouped_gemm_quant/grouped_gemm_tensorquant/ | | ✅ | | ✅ | | | | ✅ | | | | | ✅ | ✅ | ✅ | 0.0833 | diff --git a/tile_engine/ops/gemm/CMakeLists.txt b/tile_engine/ops/gemm/CMakeLists.txt index b50a679010..c2967f1f57 100644 --- a/tile_engine/ops/gemm/CMakeLists.txt +++ b/tile_engine/ops/gemm/CMakeLists.txt @@ -15,7 +15,7 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") if(_te_budget GREATER 0) # Detect active ops from their DATATYPE variables set(_active_ops "") - foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant) + foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant gemm_aquant gemm_bquant gemm_abquant) string(TOUPPER ${_op} _OP_UPPER) if(NOT "${${_OP_UPPER}_DATATYPE}" STREQUAL "") list(APPEND _active_ops ${_op}) @@ -45,7 +45,7 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") message(STATUS "Sampling budget allocation:\n${_alloc_output}") # Read per-op allocations (only if not already overridden) - foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant) + foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant gemm_aquant gemm_bquant gemm_abquant) string(TOUPPER ${_op} _OP_UPPER) if("${${_OP_UPPER}_MAX_INSTANCES}" STREQUAL "") if(EXISTS "${_alloc_dir}/${_op}_budget.txt") diff --git a/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt index 83803dabc3..035ac7008b 100644 --- a/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt +++ b/tile_engine/ops/gemm/block_scale_gemm/CMakeLists.txt @@ -1,5 +1,8 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT +add_subdirectory(gemm_aquant) +add_subdirectory(gemm_bquant) +add_subdirectory(gemm_abquant) add_subdirectory(gemm_rowcolquant) add_subdirectory(gemm_tensor_quant) diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt new file mode 100644 index 0000000000..980b3005a4 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/CMakeLists.txt @@ -0,0 +1,272 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +set(GEMM_ABQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM ABQuant (semicolon-separated)") +set(GEMM_ABQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM ABQuant (semicolon-separated)") +set(GEMM_ABQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)") +option(ENABLE_CCACHE_GEMM_ABQUANT "Enable ccache for GEMM ABQuant ops compilation" OFF) + +# Store the directory path for use in functions +set(GEMM_ABQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR}) + +# Function to create individual GEMM ABQuant targets +function(create_individual_gemm_abquant_target datatype layout trait tile_config config_json) + # GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL is guaranteed non-empty here (caller checks at line 192) + if(NOT GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL) + message(FATAL_ERROR "BUG: create_individual_gemm_abquant_target called with no GPU targets") + endif() + + # Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k + string(REPLACE "_" ";" config_groups ${tile_config}) + list(GET config_groups 0 tile_dims) + list(GET config_groups 1 warp_dims) + list(GET config_groups 2 warp_tile_dims) + + string(REPLACE "x" ";" tile_parts ${tile_dims}) + list(GET tile_parts 0 tile_m) + list(GET tile_parts 1 tile_n) + list(GET tile_parts 2 tile_k) + + string(REPLACE "x" ";" warp_parts ${warp_dims}) + list(GET warp_parts 0 warp_m) + list(GET warp_parts 1 warp_n) + list(GET warp_parts 2 warp_k) + + string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims}) + list(GET warp_tile_parts 0 warp_tile_m) + list(GET warp_tile_parts 1 warp_tile_n) + list(GET warp_tile_parts 2 warp_tile_k) + + set(target_name "benchmark_gemm_abquant_${datatype}_${layout}_${trait}_${tile_config}") + set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") + + # Generate the single instance header for this kernel + set(instance_header "${working_path}/gemm_abquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp") + + # Add custom command to generate the header file at build time + add_custom_command( + OUTPUT ${instance_header} + COMMAND ${Python3_EXECUTABLE} ${GEMM_ABQUANT_SOURCE_DIR}/gemm_abquant_instance_builder.py + --working_path ${working_path} + --datatype ${datatype} + --layout ${layout} + --config_json ${config_json} + --gen_single + --kernel_name "gemm_abquant_${datatype}_${layout}_${trait}_${tile_config}" + --tile_config "${tile_config}" + --trait_combo "${trait}" + --gpu_target "${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}" + DEPENDS ${GEMM_ABQUANT_SOURCE_DIR}/gemm_abquant_instance_builder.py ${config_json} + COMMENT "Generating ${instance_header}" + ) + + # Create the executable + add_executable(${target_name} + EXCLUDE_FROM_ALL + ${GEMM_ABQUANT_SOURCE_DIR}/gemm_abquant_benchmark_single.cpp + ${instance_header} + ) + + # Set GPU architectures + set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}) + + # Set compile definitions + target_compile_definitions(${target_name} PRIVATE + GEMM_ABQUANT_SINGLE_INSTANCE_HPP="${instance_header}" + ) + + # Include directories + target_include_directories(${target_name} PRIVATE + ${GEMM_ABQUANT_SOURCE_DIR} + ${working_path} + ) + + # Compile options + target_compile_options(${target_name} PRIVATE + -Wno-undefined-func-template + -Wno-float-equal + --offload-compress + -include ${instance_header} + ) + + # Add to collection targets + add_dependencies(benchmark_gemm_abquant_all ${target_name}) + add_dependencies(benchmark_gemm_abquant_${datatype} ${target_name}) + add_dependencies(benchmark_gemm_abquant_${layout} ${target_name}) + add_dependencies(benchmark_gemm_abquant_${datatype}_${layout} ${target_name}) + + # Add to pipeline-specific targets + string(REPLACE "_" ";" trait_parts ${trait}) + list(GET trait_parts 0 pipeline) + + add_dependencies(benchmark_gemm_abquant_${pipeline}_pipeline ${target_name}) +endfunction() + +# Function to build individual GEMM ABQuant targets +function(build_individual_gemm_abquant_targets datatype layout) + set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") + + # Choose config file (priority: env var > cmake var > default) + if(DEFINED ENV{GEMM_ABQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_ABQUANT_CONFIG_FILE}" STREQUAL "") + set(config_filename "$ENV{GEMM_ABQUANT_CONFIG_FILE}") + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}") + message(VERBOSE " Using config from environment variable: ${config_filename}") + elseif(NOT "${GEMM_ABQUANT_CONFIG_FILE}" STREQUAL "") + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_ABQUANT_CONFIG_FILE}") + message(VERBOSE " Using custom config: ${GEMM_ABQUANT_CONFIG_FILE}") + else() + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json") + message(VERBOSE " Using default config for layout ${layout}") + endif() + + if(NOT EXISTS ${json_blob}) + message(FATAL_ERROR "Config file not found: ${json_blob}") + endif() + + # Create working directory + file(MAKE_DIRECTORY ${working_path}) + + # Build sampling arguments + set(extra_list_args "") + if(NOT "${GEMM_ABQUANT_MAX_INSTANCES}" STREQUAL "") + list(APPEND extra_list_args --max-instances ${GEMM_ABQUANT_MAX_INSTANCES}) + endif() + if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") + list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER}) + list(APPEND extra_list_args --manifest-path ${working_path}) + endif() + if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "") + list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED}) + endif() + + # List kernels + message(VERBOSE " Listing ABQuant kernel configurations for ${datatype} ${layout}...") + execute_process( + COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_abquant_instance_builder.py + --working_path ${working_path} + --datatype ${datatype} + --layout ${layout} + --config_json ${json_blob} + --gpu_target "${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}" + --list_kernels + ${extra_list_args} + WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR} + RESULT_VARIABLE ret + OUTPUT_VARIABLE list_output + ERROR_VARIABLE list_error + ) + + if(NOT ret EQUAL 0) + message(FATAL_ERROR "Failed to list ABQuant kernels for ${datatype} ${layout}: ${list_error}") + endif() + + # Read kernel count + if(EXISTS ${working_path}/gemm_abquant_kernel_count.txt) + file(READ ${working_path}/gemm_abquant_kernel_count.txt kernel_count) + string(STRIP "${kernel_count}" kernel_count) + message(VERBOSE " Found ${kernel_count} ABQuant kernel configurations") + else() + message(FATAL_ERROR "Kernel count file not found") + endif() + + # Read kernel list and create targets + if(EXISTS ${working_path}/gemm_abquant_kernel_list.txt) + file(STRINGS ${working_path}/gemm_abquant_kernel_list.txt kernel_lines) + foreach(line IN LISTS kernel_lines) + string(REPLACE "|" ";" parts "${line}") + list(GET parts 0 kernel_name) + list(GET parts 1 tile_config) + list(GET parts 2 trait_combo) + + create_individual_gemm_abquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}") + endforeach() + else() + message(FATAL_ERROR "Kernel list file not found") + endif() +endfunction() + +# Main build logic +message(VERBOSE "=== Starting Tile Engine GEMM ABQuant Configuration ===") +message(VERBOSE "GEMM_ABQUANT_DATATYPE: ${GEMM_ABQUANT_DATATYPE}") +message(VERBOSE "GEMM_ABQUANT_LAYOUT: ${GEMM_ABQUANT_LAYOUT}") +message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") + +# Filter GPU targets to supported architectures +set(GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL "") +set(DESIRED_TARGETS "gfx90a;gfx942;gfx950") + +foreach(target IN LISTS SUPPORTED_GPU_TARGETS) + if(target IN_LIST DESIRED_TARGETS) + list(APPEND GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL ${target}) + message(VERBOSE " Adding GPU target: ${target}") + endif() +endforeach() + +# Skip build if no matching targets found +if(NOT GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL) + message(WARNING "Skipping Tile Engine GEMM ABQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") +else() + message(VERBOSE "Building individual GEMM ABQuant targets for GPU targets: ${GEMM_ABQUANT_GPU_TARGETS_INDIVIDUAL}") + + # Enable compiler cache if requested + if(ENABLE_CCACHE_GEMM_ABQUANT) + find_program(CCACHE_PROGRAM ccache) + if(CCACHE_PROGRAM) + set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM}) + message(VERBOSE "Using ccache for faster compilation") + endif() + endif() + + # Create master collection target + add_custom_target(benchmark_gemm_abquant_all) + + # Create datatype collection targets + foreach(dt IN LISTS GEMM_ABQUANT_DATATYPE) + add_custom_target(benchmark_gemm_abquant_${dt}) + endforeach() + + # Create layout collection targets + foreach(l IN LISTS GEMM_ABQUANT_LAYOUT) + add_custom_target(benchmark_gemm_abquant_${l}) + endforeach() + + # Create combined collection targets + foreach(dt IN LISTS GEMM_ABQUANT_DATATYPE) + foreach(l IN LISTS GEMM_ABQUANT_LAYOUT) + add_custom_target(benchmark_gemm_abquant_${dt}_${l}) + endforeach() + endforeach() + + # Create pipeline-specific collection targets. + # Must stay in sync with trait_config.pipeline.values in the config JSONs: + # if a new pipeline is added there, add it here too or add_dependencies() will fail at build time. + set(GEMM_ABQUANT_PIPELINES "compv3") + foreach(pipeline IN LISTS GEMM_ABQUANT_PIPELINES) + add_custom_target(benchmark_gemm_abquant_${pipeline}_pipeline) + endforeach() + + # Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that + # sampling fires per-combo rather than being a single cap. + if(NOT "${GEMM_ABQUANT_MAX_INSTANCES}" STREQUAL "") + list(LENGTH GEMM_ABQUANT_DATATYPE _abq_n_dt) + list(LENGTH GEMM_ABQUANT_LAYOUT _abq_n_lay) + math(EXPR _abq_n_combos "${_abq_n_dt} * ${_abq_n_lay}") + if(_abq_n_combos GREATER 0) + math(EXPR GEMM_ABQUANT_MAX_INSTANCES + "${GEMM_ABQUANT_MAX_INSTANCES} / ${_abq_n_combos}") + if(GEMM_ABQUANT_MAX_INSTANCES EQUAL 0) + set(GEMM_ABQUANT_MAX_INSTANCES 1) + message(WARNING " gemm_abquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_abq_n_combos} combos)") + else() + message(STATUS " gemm_abquant: per-combo budget = ${GEMM_ABQUANT_MAX_INSTANCES} (${_abq_n_combos} combos)") + endif() + endif() + endif() + + # Build individual targets for each datatype/layout combination + foreach(dt IN LISTS GEMM_ABQUANT_DATATYPE) + foreach(l IN LISTS GEMM_ABQUANT_LAYOUT) + build_individual_gemm_abquant_targets(${dt} ${l}) + endforeach() + endforeach() +endif() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json new file mode 100644 index 0000000000..ef77633ddb --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_ci_config.json @@ -0,0 +1,104 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 16 + ] + }, + "tile_n": { + "values": [ + 64 + ] + }, + "tile_k": { + "values": [ + 256 + ] + }, + "warp_m": { + "values": [ + 1 + ] + }, + "warp_n": { + "values": [ + 4 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16 + ] + }, + "warp_tile_n": { + "values": [ + 16 + ] + }, + "warp_tile_k": { + "values": [ + 16, + 32, + 64, + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "default", + "cshuffle" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "a_preshuffle_quant": { + "values": [ + false + ] + }, + "b_preshuffle_quant": { + "values": [ + false, + true + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128, + "group_size_n": { + "values": [ + 1, + 128 + ] + } +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json new file mode 100644 index 0000000000..70c307ba10 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/default_config.json @@ -0,0 +1,113 @@ +{ + "tile_config": { + "tile_m": { + "max": 256, + "min": 64, + "step": 64 + }, + "tile_n": { + "max": 256, + "min": 64, + "step": 64 + }, + "tile_k": { + "max": 256, + "min": 64, + "step": 64 + }, + "warp_m": { + "values": [ + 4, + 2, + 1 + ] + }, + "warp_n": { + "values": [ + 4, + 2, + 1 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 4, + 16, + 32 + ] + }, + "warp_tile_n": { + "values": [ + 16, + 32, + 64 + ] + }, + "warp_tile_k": { + "values": [ + 8, + 16, + 32, + 64, + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "cshuffle", + "default" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "a_preshuffle_quant": { + "values": [ + false + ] + }, + "b_preshuffle_quant": { + "values": [ + false, + true + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128, + "group_size_n": { + "values": [ + 1, + 128 + ] + } +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json new file mode 100644 index 0000000000..a20dca95f7 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/configs/example_config.json @@ -0,0 +1,98 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 16 + ] + }, + "tile_n": { + "values": [ + 64 + ] + }, + "tile_k": { + "values": [ + 256 + ] + }, + "warp_m": { + "values": [ + 1 + ] + }, + "warp_n": { + "values": [ + 4 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16 + ] + }, + "warp_tile_n": { + "values": [ + 16 + ] + }, + "warp_tile_k": { + "values": [ + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "default" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "a_preshuffle_quant": { + "values": [ + false + ] + }, + "b_preshuffle_quant": { + "values": [ + false + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128, + "group_size_n": { + "values": [ + 1 + ] + } +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp new file mode 100644 index 0000000000..18fbbc2c0a --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.hpp @@ -0,0 +1,243 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "gemm_abquant_common.hpp" + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" + +// Data types and Layouts are defined by the generated kernel headers: +// ADataType, BDataType, AQDataType, BQDataType, AccDataType, CDataType +// ALayout, BLayout, CLayout, AQLayout, BQLayout + +enum class Metric +{ + LATENCY = 0, + TFLOPS = 1, + BANDWIDTH = 2 +}; + +inline constexpr auto get_metric_name(Metric m) +{ + switch(m) + { + case Metric::LATENCY: return "latency"; + case Metric::TFLOPS: return "tflops"; + case Metric::BANDWIDTH: return "bandwidth"; + default: throw std::invalid_argument("Unsupported metric type"); + } +} + +struct ABQuantGemmProblem +{ + int split_k_; + int m_, n_, k_; + int stride_a_, stride_b_, stride_c_; + int stride_aq_; + int stride_bq_; + int group_size_k_; + int group_size_n_; + + std::string dtype_a_, dtype_b_, dtype_aq_, dtype_bq_, dtype_acc_, dtype_c_; + std::string layout_a_, layout_b_, layout_c_; + + friend std::ostream& operator<<(std::ostream& os, const ABQuantGemmProblem& problem) + { + os << "{\n" + << " \"split_k\":" << problem.split_k_ << ",\n" + << " \"m\":" << problem.m_ << ",\n" + << " \"n\":" << problem.n_ << ",\n" + << " \"k\":" << problem.k_ << ",\n" + << " \"stride_a\":" << problem.stride_a_ << ",\n" + << " \"stride_b\":" << problem.stride_b_ << ",\n" + << " \"stride_c\":" << problem.stride_c_ << ",\n" + << " \"stride_aq\":" << problem.stride_aq_ << ",\n" + << " \"stride_bq\":" << problem.stride_bq_ << ",\n" + << " \"group_size_k\":" << problem.group_size_k_ << ",\n" + << " \"group_size_n\":" << problem.group_size_n_ << ",\n" + << " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n" + << " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n" + << " \"dtype_aq\":\"" << problem.dtype_aq_ << "\",\n" + << " \"dtype_bq\":\"" << problem.dtype_bq_ << "\",\n" + << " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n" + << " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n" + << " \"layout_a\":\"" << problem.layout_a_ << "\",\n" + << " \"layout_b\":\"" << problem.layout_b_ << "\",\n" + << " \"layout_c\":\"" << problem.layout_c_ << "\"\n" + << "}"; + return os; + } +}; + +struct PerformanceResult +{ + double latency_; + double tflops_; + double bandwidth_; + + static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m) + { + switch(m) + { + case Metric::LATENCY: return a.latency_ < b.latency_; + case Metric::TFLOPS: return a.tflops_ > b.tflops_; + case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_; + default: throw std::invalid_argument("Unsupported metric type"); + } + } + + friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result) + { + os << "{\n" + << " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_ + << ",\n" + << " \"tflops(TFlops)\": " << result.tflops_ << ",\n" + << " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n" + << "}"; + return os; + } +}; + +struct KernelInstance +{ + std::string name_; + ABQuantGemmProblem problem_; + PerformanceResult perf_result_; + + static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m) + { + return PerformanceResult::compare(a.perf_result_, b.perf_result_, m); + } + + friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj) + { + os << "{\n" + << " \"name\": \"" << obj.name_ << "\",\n" + << " \"problem\": " << obj.problem_ << ",\n" + << " \"perf_result\": " << obj.perf_result_ << "\n" + << "}"; + return os; + } +}; + +struct Setting +{ + int n_warmup_; + int n_repeat_; + bool is_gpu_timer_; + int verify_; + int init_method_; + bool log_; + std::string csv_filename_; + bool flush_cache_; + int rotating_count_; + bool json_output_; +}; + +inline std::string get_rocm_version() +{ + std::ifstream version_file("/opt/rocm/.info/version"); + if(version_file.is_open()) + { + std::string version; + std::getline(version_file, version); + return version; + } + return "Unknown"; +} + +template +auto calculate_rtol_atol_abquant(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + // Both A and B are the same FP8/BF8 type for abquant; assert so mixed-precision additions are + // caught. + static_assert(sizeof(ADataType_) == sizeof(BDataType_), + "calculate_rtol_atol_abquant assumes equal-width A and B types"); + using ComputeType = ADataType_; + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} + +/// @brief Compare device and host results for ABQuant GEMM +bool compare_abquant(std::string instanceName, + ck_tile::index_t K, + ck_tile::index_t kbatch, + const ck_tile::HostTensor& c_m_n_dev_result, + const ck_tile::HostTensor& c_m_n_host_result) +{ + const float max_accumulated_value = + std::abs(static_cast(*std::max_element(c_m_n_host_result.mData.begin(), + c_m_n_host_result.mData.end(), + [](const auto& a, const auto& b) { + return std::abs(static_cast(a)) < + std::abs(static_cast(b)); + }))); + const auto rtol_atol = calculate_rtol_atol_abquant(K, kbatch, max_accumulated_value); + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_result, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + std::cout << "For " << instanceName << " Relative error threshold is " + << rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is " + << rtol_atol.at(ck_tile::number<1>{}) << std::endl; + std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl; + + return pass; +} + +/// @brief CPU reference implementation for ABQuant GEMM +void abquant_gemm_host_reference(int verify, + ck_tile::HostTensor& a_m_k, + ck_tile::HostTensor& aq_m_qk, + ck_tile::HostTensor& b_k_n, + ck_tile::HostTensor& bq_qk_n, + ck_tile::HostTensor& c_m_n_host_result) +{ + if(verify == 1) + { + c_m_n_host_result.SetZero(); + ck_tile::reference_gemm_abquant( + a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result); + } +} +#pragma clang diagnostic pop diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py new file mode 100644 index 0000000000..8642298646 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py @@ -0,0 +1,477 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import sys +import json +import subprocess +import argparse +import time +from pathlib import Path +from typing import List, Dict, Tuple, Optional + + +class ABQuantGemmBenchmark: + def __init__(self, build_dir: str, verbose: bool = False): + self.build_dir = Path(build_dir) + self.verbose = verbose + self.results = [] + + def discover_kernels(self) -> List[Path]: + """Find all benchmark_gemm_abquant_* executables in the build directory.""" + bin_dir = self.build_dir / "bin" + if not bin_dir.exists(): + print(f"Error: Binary directory {bin_dir} does not exist") + return [] + + kernels = list(bin_dir.glob("benchmark_gemm_abquant_*")) + if self.verbose: + print(f"Found {len(kernels)} ABQuant kernel executables") + for k in kernels: + print(f" - {k.name}") + return kernels + + def extract_kernel_info(self, kernel_path: Path) -> Dict[str, str]: + """Extract kernel information from filename.""" + name = kernel_path.stem + + info = { + "executable": str(kernel_path), + "name": name, + "data_type": "unknown", + "layout": "unknown", + "pipeline": "unknown", + "scheduler": "unknown", + "epilogue": "unknown", + "a_preshuffle_quant": False, + "b_preshuffle_quant": False, + "group_size_n": 1, + } + + # Parse: benchmark_gemm_abquant_fp8_rcr_compv3_default_intrawave_False_False_False_False_False_gsn1_128x128x128_... + parts = name.split("_") + + if len(parts) >= 3: + # Skip "benchmark_gemm_abquant" prefix (3 parts) + idx = 3 + if idx < len(parts): + info["data_type"] = parts[idx] + idx += 1 + if idx < len(parts): + info["layout"] = parts[idx] + idx += 1 + if idx < len(parts): + info["pipeline"] = parts[idx] + idx += 1 + if idx < len(parts): + info["epilogue"] = parts[idx] + idx += 1 + if idx < len(parts): + info["scheduler"] = parts[idx] + idx += 1 + # Skip pad_m, pad_n, pad_k + idx += 3 + if idx < len(parts): + info["a_preshuffle_quant"] = parts[idx] == "True" + idx += 1 + if idx < len(parts): + info["b_preshuffle_quant"] = parts[idx] == "True" + idx += 1 + # Parse gsn value (e.g. gsn1, gsn128) + if idx < len(parts) and parts[idx].startswith("gsn"): + info["group_size_n"] = int(parts[idx][3:]) + + config_info = self.parse_detailed_config(name) + info.update(config_info) + + # Warn on fields that failed to parse — silent "unknown"/zero values hide name format drift. + unknown_fields = [k for k in ("data_type", "layout", "pipeline", "scheduler", "epilogue") + if info.get(k) == "unknown"] + if unknown_fields: + print(f"Warning: could not parse {unknown_fields} from kernel name: {name}") + tile_sizes = info.get("tile_sizes", {}) + if any(v == 0 for v in tile_sizes.values()): + print(f"Warning: zero tile dimension parsed from kernel name: {name} -> {tile_sizes}") + + info["config_id"] = self.generate_config_id(info) + return info + + def parse_detailed_config(self, kernel_name: str) -> Dict: + """Parse tile dimensions and boolean flags from kernel name.""" + config = { + "tile_sizes": {"tile_m": 0, "tile_n": 0, "tile_k": 0}, + "warp_config": {"warp_m": 0, "warp_n": 0, "warp_k": 0}, + "warp_tile": {"warp_tile_m": 0, "warp_tile_n": 0, "warp_tile_k": 0}, + "optimization_flags": { + "pad_m": False, + "pad_n": False, + "pad_k": False, + "a_preshuffle_quant": False, + "b_preshuffle_quant": False, + }, + } + + parts = kernel_name.split("_") + + # Extract boolean flags (5 consecutive True/False values) + bool_sequence = [] + for i, part in enumerate(parts): + if part in ["True", "False"]: + bool_sequence.append(part == "True") + j = i + 1 + while j < len(parts) and parts[j] in ["True", "False"]: + bool_sequence.append(parts[j] == "True") + j += 1 + break + + if len(bool_sequence) >= 5: + config["optimization_flags"]["pad_m"] = bool_sequence[0] + config["optimization_flags"]["pad_n"] = bool_sequence[1] + config["optimization_flags"]["pad_k"] = bool_sequence[2] + config["optimization_flags"]["a_preshuffle_quant"] = bool_sequence[3] + config["optimization_flags"]["b_preshuffle_quant"] = bool_sequence[4] + + # Extract dimension groups (e.g., 128x128x128) in positional order. + # Kernel names encode groups as: tile_MxNxK_warp_MxNxK_warp_tile_MxNxK + dimension_groups = [] + for part in parts: + if "x" in part and len(part.split("x")) == 3: + try: + dims = [int(x) for x in part.split("x")] + if all(d > 0 for d in dims): + dimension_groups.append(dims) + except ValueError: + continue + + if len(dimension_groups) >= 3: + config["tile_sizes"]["tile_m"] = dimension_groups[0][0] + config["tile_sizes"]["tile_n"] = dimension_groups[0][1] + config["tile_sizes"]["tile_k"] = dimension_groups[0][2] + config["warp_config"]["warp_m"] = dimension_groups[1][0] + config["warp_config"]["warp_n"] = dimension_groups[1][1] + config["warp_config"]["warp_k"] = dimension_groups[1][2] + config["warp_tile"]["warp_tile_m"] = dimension_groups[2][0] + config["warp_tile"]["warp_tile_n"] = dimension_groups[2][1] + config["warp_tile"]["warp_tile_k"] = dimension_groups[2][2] + + return config + + def generate_config_id(self, info: Dict) -> str: + """Generate a compact config ID from kernel info.""" + parts = [ + info.get("data_type", "unk"), + info.get("layout", "unk"), + info.get("pipeline", "unk"), + info.get("scheduler", "unk"), + ] + + tile_sizes = info.get("tile_sizes", {}) + if tile_sizes.get("tile_m", 0) > 0: + parts.append( + f"{tile_sizes['tile_m']}x{tile_sizes['tile_n']}x{tile_sizes['tile_k']}" + ) + + gsn = info.get("group_size_n", 1) + parts.append(f"gsn{gsn}") + + return "_".join(parts) + + def run_kernel(self, kernel_path: Path, params: Dict[str, str]) -> Optional[Dict]: + """Run a single kernel with given parameters.""" + results_dir = self.build_dir / "results" + results_dir.mkdir(exist_ok=True) + + json_file = results_dir / f"{kernel_path.stem}.json" + + cmd = [str(kernel_path)] + for key, value in params.items(): + cmd.append(f"-{key}={value}") + cmd.append("-json_output=true") + + if self.verbose: + print(f"Running: {' '.join(cmd)}") + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=60) + + if result.returncode != 0: + print(f"Error running {kernel_path.name}: {result.stderr}") + return None + + output = result.stdout.strip() + if output: + with open(json_file, "w") as f: + f.write(output) + return self.parse_json_file(json_file) + else: + print(f"No output from {kernel_path.name}") + return None + + except subprocess.TimeoutExpired: + print(f"Timeout running {kernel_path.name}") + return None + except Exception as e: + print(f"Error running {kernel_path.name}: {e}") + return None + + def parse_json_file(self, json_file: Path) -> Optional[Dict]: + """Parse JSON data from kernel output file.""" + try: + with open(json_file, "r") as f: + content = f.read().strip() + + data = json.loads(content) + result = data.copy() + if "perf_result" in data: + perf = data["perf_result"] + result["time_ms"] = perf.get("latency(ms)", 0) + result["tflops"] = perf.get("tflops(TFlops)", 0) + result["bandwidth_gb_s"] = perf.get("bandwidth(GB/s)", 0) + return result + + except (json.JSONDecodeError, Exception) as e: + if self.verbose: + print(f"Failed to parse JSON from {json_file}: {e}") + return None + + def benchmark_problem_size( + self, + kernels: List[Path], + m: int, + n: int, + k: int, + group_size_k: int = 128, + verify: int = 0, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> List[Dict]: + """Benchmark all kernels for a specific problem size.""" + results = [] + + print(f"\nBenchmarking M={m}, N={n}, K={k}, group_size_k={group_size_k}") + + for kernel_path in kernels: + kernel_info = self.extract_kernel_info(kernel_path) + group_size_n = kernel_info.get("group_size_n", 1) + + params = { + "m": m, + "n": n, + "k": k, + "group_size_k": group_size_k, + "group_size_n": group_size_n, + "verify": verify, + "warmup": warmup, + "repeat": repeat, + "flush_cache": str(flush_cache).lower(), + "rotating_count": rotating_count, + } + + result = self.run_kernel(kernel_path, params) + + if result: + structured_result = { + "name": kernel_info["name"], + "config_id": kernel_info["config_id"], + "problem": result.get("problem", {}), + "perf_result": result.get("perf_result", {}), + "config": { + "data_type": kernel_info["data_type"], + "layout": kernel_info["layout"], + "pipeline": kernel_info["pipeline"], + "scheduler": kernel_info["scheduler"], + "epilogue": kernel_info["epilogue"], + "tile_sizes": kernel_info.get("tile_sizes", {}), + "warp_config": kernel_info.get("warp_config", {}), + "warp_tile": kernel_info.get("warp_tile", {}), + "optimization_flags": kernel_info.get("optimization_flags", {}), + }, + "executable": kernel_info["executable"], + "time_ms": result.get("time_ms", 0), + "tflops": result.get("tflops", 0), + "bandwidth_gb_s": result.get("bandwidth_gb_s", 0), + } + + results.append(structured_result) + + if self.verbose: + print( + f" {kernel_info['config_id']}: " + f"{structured_result['tflops']:.2f} TFLOPS, " + f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " + f"{structured_result['time_ms']:.2f}ms" + ) + + return results + + def find_best_kernel( + self, results: List[Dict], metric: str = "tflops" + ) -> Optional[Dict]: + """Find the best performing kernel based on metric.""" + if not results: + return None + + if metric == "tflops": + return max(results, key=lambda x: x.get("tflops", 0)) + elif metric == "time_ms": + return min(results, key=lambda x: x.get("time_ms", float("inf"))) + elif metric == "bandwidth_gb_s": + return max(results, key=lambda x: x.get("bandwidth_gb_s", 0)) + else: + raise ValueError(f"Unknown metric: {metric}") + + def benchmark_sweep( + self, + problem_sizes: List[Tuple[int, int, int]], + group_size_k: int = 128, + verify: bool = False, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> Dict: + """Run comprehensive benchmark sweep.""" + kernels = self.discover_kernels() + if not kernels: + print("No kernels found!") + return {} + + all_results = [] + best_kernels = {} + + for m, n, k in problem_sizes: + results = self.benchmark_problem_size( + kernels, + m, + n, + k, + group_size_k=group_size_k, + verify=1 if verify else 0, + warmup=warmup, + repeat=repeat, + flush_cache=flush_cache, + rotating_count=rotating_count, + ) + + all_results.extend(results) + + best = self.find_best_kernel(results) + if best: + key = f"m{m}_n{n}_k{k}" + best_kernels[key] = best + print( + f"Best for {key}: {best['name']} " + f"({best['tflops']:.2f} TFLOPS, " + f"{best['bandwidth_gb_s']:.2f} GB/s, " + f"{best['time_ms']:.2f}ms)" + ) + + self.results = all_results + return best_kernels + + def export_json(self, filename: str, best_kernels: Dict = None): + """Export results to JSON.""" + from datetime import datetime + + output_data = { + "benchmark_metadata": { + "timestamp": datetime.now().isoformat(), + "operator": "gemm_abquant", + "total_kernels_tested": len(self.results), + "successful_runs": len( + [r for r in self.results if r.get("tflops", 0) > 0] + ), + }, + "kernel_results": self.results, + "best_kernels_by_problem": best_kernels or {}, + } + + with open(filename, "w") as f: + json.dump(output_data, f, indent=2) + print(f"JSON results exported to {filename}") + + +def main(): + parser = argparse.ArgumentParser( + description="ABQuant GEMM Kernel Benchmarking Tool" + ) + parser.add_argument( + "build_dir", help="Build directory containing kernel executables" + ) + parser.add_argument( + "--problem-sizes", + nargs="+", + default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"], + help="Problem sizes as M,N,K tuples", + ) + parser.add_argument( + "--group-size-k", + type=int, + default=128, + help="Quantization group size along K (default: 128)", + ) + parser.add_argument("--verify", action="store_true", help="Enable verification") + parser.add_argument("--verbose", action="store_true", help="Verbose output") + parser.add_argument( + "--warmup", + type=int, + default=50, + help="Number of warmup iterations", + ) + parser.add_argument( + "--repeat", + type=int, + default=100, + help="Number of benchmark iterations", + ) + parser.add_argument( + "--no-flush-cache", + action="store_true", + help="Disable cache flushing between iterations", + ) + parser.add_argument( + "--rotating-count", + type=int, + default=1000, + help="Number of iterations to rotate the cache (default: 1000)", + ) + parser.add_argument("--json", help="JSON output filename (optional)") + + args = parser.parse_args() + + problem_sizes = [] + for size_str in args.problem_sizes: + try: + m, n, k = map(int, size_str.split(",")) + problem_sizes.append((m, n, k)) + except ValueError: + print(f"Invalid problem size: {size_str}") + return 1 + + benchmark = ABQuantGemmBenchmark(args.build_dir, verbose=args.verbose) + + print("Starting ABQuant GEMM kernel benchmark sweep...") + start_time = time.time() + + best_kernels = benchmark.benchmark_sweep( + problem_sizes=problem_sizes, + group_size_k=args.group_size_k, + verify=args.verify, + warmup=args.warmup, + repeat=args.repeat, + flush_cache=not args.no_flush_cache, + rotating_count=args.rotating_count, + ) + + elapsed_time = time.time() - start_time + print(f"\nBenchmark completed in {elapsed_time:.2f} seconds") + + if args.json: + benchmark.export_json(args.json, best_kernels) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp new file mode 100644 index 0000000000..960473cbb4 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark_single.cpp @@ -0,0 +1,178 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "gemm_abquant_profiler.hpp" +#include "gemm_abquant_common.hpp" + +// The kernel header is included via the compile command line with -include flag +// It defines SelectedKernel struct, KERNEL_NAME, and type aliases: +// ADataType, BDataType, AQDataType, BQDataType, AccDataType, CDataType +// ALayout, BLayout, CLayout, AQLayout, BQLayout + +inline auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.") + .insert("n", "4096", "The value for n dimension. Default is 4096.") + .insert("k", "2048", "The value for k dimension. Default is 2048.") + .insert("stride_a", "0", "The stride value for tensor A. Default is 0.") + .insert("stride_b", "0", "The stride value for tensor B. Default is 0.") + .insert("stride_c", "0", "The stride value for tensor C. Default is 0.") + .insert("split_k", "1", "The split value for k dimension. Default is 1.") + .insert("group_size_k", + std::to_string(SelectedKernel::GroupSizeK), + "The quantization group size along K. Default matches kernel config.") + .insert("group_size_n", + std::to_string(SelectedKernel::GroupSizeN), + "The quantization group size along N for BQ. Default matches kernel config.") + .insert("verify", + "1", + "The type of validation. Set to 0 for no validation, 1 for validation on CPU. " + "Default is 1, CPU validation.") + .insert("log", + "false", + "Whether output kernel instance information or not. Possible values are true or " + "false. Default is false") + .insert( + "warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.") + .insert( + "repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.") + .insert("timer", + "true", + "Whether if the timer is gpu timer or not. Possible values are false or true. " + "Default is true.") + .insert("init", + "0", + "The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 " + "for constant(1). Default is 0, random.") + .insert("flush_cache", + "true", + "To flush cache, possible values are true or false. Default is true.") + .insert( + "rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.") + .insert("metric", + "0", + "Metric with which to measure kernel performance. Set to 0 for latency, 1 for " + "tflops, or 2 for bandwidth. Default is 0, latency.") + .insert("csv_filename", + "", + "The filename of benchmark result. Default is empty (no CSV output).") + .insert("json_output", + "false", + "Whether to output results in JSON format only. Possible values are true or false. " + "Default is false"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} + +void benchmark_single(const ck_tile::ArgParser& arg_parser) +{ + std::string dtype_a = DataTypeTraits::name; + std::string dtype_b = DataTypeTraits::name; + std::string dtype_aq = DataTypeTraits::name; + std::string dtype_bq = DataTypeTraits::name; + std::string dtype_acc = DataTypeTraits::name; + std::string dtype_c = DataTypeTraits::name; + + std::string layout_a = ALayout::name; + std::string layout_b = BLayout::name; + std::string layout_c = CLayout::name; + + int M = arg_parser.get_int("m"); + int N = arg_parser.get_int("n"); + int K = arg_parser.get_int("k"); + int group_size_k = arg_parser.get_int("group_size_k"); + int group_size_n = arg_parser.get_int("group_size_n"); + + if(M <= 0 || N <= 0 || K <= 0) + { + throw std::invalid_argument("m, n, k must be positive integers"); + } + if(group_size_k <= 0 || K % group_size_k != 0) + { + throw std::invalid_argument( + "group_size_k must be positive and k must be divisible by group_size_k"); + } + if(group_size_n > 1 && N % group_size_n != 0) + { + throw std::invalid_argument("n must be divisible by group_size_n when group_size_n > 1"); + } + + ABQuantGemmProblem problem{arg_parser.get_int("split_k"), + M, + N, + K, + arg_parser.get_int("stride_a"), + arg_parser.get_int("stride_b"), + arg_parser.get_int("stride_c"), + 0, // stride_aq computed by profiler + 0, // stride_bq computed by profiler + group_size_k, + group_size_n, + dtype_a, + dtype_b, + dtype_aq, + dtype_bq, + dtype_acc, + dtype_c, + layout_a, + layout_b, + layout_c}; + + Setting setting{arg_parser.get_int("warmup"), + arg_parser.get_int("repeat"), + arg_parser.get_bool("timer"), + arg_parser.get_int("verify"), + arg_parser.get_int("init"), + arg_parser.get_bool("log"), + arg_parser.get_str("csv_filename"), + arg_parser.get_bool("flush_cache"), + arg_parser.get_int("rotating_count"), + arg_parser.get_bool("json_output")}; + + auto& profiler = ABQuantGemmProfiler::instance(setting); + + try + { + auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& stream) { + return SelectedKernel::launch(args, stream); + }; + + profiler.benchmark(problem, kernel_func); + profiler.select_best_instance(static_cast(arg_parser.get_int("metric"))); + } + catch(const std::exception& e) + { + std::cerr << "Benchmark failed: " << e.what() << std::endl; + } +} + +int main(int argc, char* argv[]) +{ + try + { + auto [result, parser] = create_args(argc, argv); + if(!result) + return EXIT_FAILURE; + + benchmark_single(parser); + return 0; + } + catch(const std::exception& e) + { + std::cerr << "Error: " << e.what() << "\n"; + return EXIT_FAILURE; + } +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp new file mode 100644 index 0000000000..c8726993c8 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_common.hpp @@ -0,0 +1,76 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/core/numeric/integer.hpp" + +// DataTypeTraits for all supported types +template +struct DataTypeTraits; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp32"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp16"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "bf16"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp8"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "bf8"; +}; + +// Helper function to determine if a layout is row-major +template +constexpr auto is_row_major(Layout) +{ + return ck_tile::bool_constant>{}; +} + +// Structure to hold kernel traits for dispatcher +struct ABQuantKernelTraits +{ + std::string pipeline; // compv3 + std::string scheduler; // intrawave + std::string epilogue; // default, cshuffle + bool pad_m; + bool pad_n; + bool pad_k; + bool a_preshuffle_quant; + bool b_preshuffle_quant; + + ABQuantKernelTraits() + : pipeline("compv3"), + scheduler("intrawave"), + epilogue("default"), + pad_m(false), + pad_n(false), + pad_k(false), + a_preshuffle_quant(false), + b_preshuffle_quant(false) + { + } +}; diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py new file mode 100644 index 0000000000..93d51ee09e --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_instance_builder.py @@ -0,0 +1,904 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import os +import argparse +import importlib.util +import multiprocessing +import concurrent.futures +import itertools +import logging + + +def _import_gemm_kernel_builder(): + """Import the base GemmKernelBuilder from the gemm directory.""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + gemm_dir = os.path.dirname(os.path.dirname(current_dir)) + + spec = importlib.util.spec_from_file_location( + "gemm_instance_builder", + os.path.join(gemm_dir, "gemm_instance_builder.py"), + ) + gemm_builder_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(gemm_builder_module) + + return gemm_builder_module.GemmKernelBuilder + + +def _import_validation_utils(): + """Import validation utilities.""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + gemm_dir = os.path.dirname(os.path.dirname(current_dir)) + parent_dir = os.path.dirname(gemm_dir) + + spec = importlib.util.spec_from_file_location( + "validation_utils", + os.path.join(parent_dir, "gemm", "gemm_validation_utils.py"), + ) + validation_utils = importlib.util.module_from_spec(spec) + spec.loader.exec_module(validation_utils) + + return validation_utils + + +GemmKernelBuilder = _import_gemm_kernel_builder() +_validation_utils = _import_validation_utils() +is_trait_combination_valid = _validation_utils.is_trait_combination_valid + + +class GemmABQuantKernelBuilder(GemmKernelBuilder): + def __init__( + self, + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json=None, + max_instances=None, + seed=None, + tier=None, + manifest_path=None, + ): + super().__init__( + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json, + max_instances=max_instances, + seed=seed, + tier=tier, + manifest_path=manifest_path, + ) + self.group_size_k = self.config.get("group_size_k", 128) + # group_size_n can be an int or a dict with "values" key + group_size_n_cfg = self.config.get("group_size_n", 1) + if isinstance(group_size_n_cfg, dict): + self.group_size_n_values = group_size_n_cfg.get("values", [1]) + else: + self.group_size_n_values = [group_size_n_cfg] + + @staticmethod + def _tile_config_to_str(tile_config): + return ( + f"{tile_config['tile_m']}x{tile_config['tile_n']}x{tile_config['tile_k']}_" + f"{tile_config['warp_m']}x{tile_config['warp_n']}x{tile_config['warp_k']}_" + f"{tile_config['warp_tile_m']}x{tile_config['warp_tile_n']}x{tile_config['warp_tile_k']}" + ) + + def _apply_sampling(self, kernel_list): + """Apply RFC Sobol+LHS+maximin sampling for ABQuant kernels. + + Overrides base class to handle ABQuant's 8-tuple trait_combo + and group_size_n field. + """ + if self.max_instances is None or len(kernel_list) <= self.max_instances: + return kernel_list + + import sys + + sampling_parent = os.path.join( + os.path.dirname(__file__), "..", "..", "..", ".." + ) + if sampling_parent not in sys.path: + sys.path.insert(0, sampling_parent) + + from sampling.sampler import sample_feasible_set + from sampling.seed import make_seed + from sampling.feasible_set import GEMM_ABQUANT_AXES + + effective_seed = make_seed( + self.seed, self.gpu_target, self.datatype, self.layout + ) + + flat_items = [] + for k in kernel_list: + flat = dict(k["tile_config"]) + ( + pipeline, + epilogue, + scheduler, + pad_m, + pad_n, + pad_k, + a_preshuffle_quant, + b_preshuffle_quant, + ) = k["trait_combo"] + flat.update( + { + "pipeline": pipeline, + "epilogue": epilogue, + "scheduler": scheduler, + "pad_m": pad_m, + "pad_n": pad_n, + "pad_k": pad_k, + "a_preshuffle_quant": a_preshuffle_quant, + "b_preshuffle_quant": b_preshuffle_quant, + "group_size_n": k.get("group_size_n", 1), + } + ) + flat_items.append(flat) + + selected, method, selected_indices = sample_feasible_set( + flat_items, + self.max_instances, + effective_seed, + GEMM_ABQUANT_AXES, + ) + + kernel_list = [kernel_list[i] for i in selected_indices] + + if self.manifest_path: + from sampling.manifest import write_manifest + + write_manifest( + selected, + self.manifest_path, + self.kernel_name_prefix, + self.datatype, + self.layout, + self.gpu_target, + effective_seed, + self.tier or "daily", + method, + ) + + print( + f"Sampled {len(kernel_list)} from feasible set " + f"(budget={self.max_instances}, seed={effective_seed}, method={method})" + ) + return kernel_list + + def _get_group_size_n_values(self): + """Return the list of group_size_n values to generate kernels for.""" + return self.group_size_n_values + + def _generate_trait_combinations(self): + """Generate all combinations of traits for ABQuant. + + ABQuant uses 8-tuples: + (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant, b_preshuffle_quant) + """ + trait_config = self.config["trait_config"] + + pipelines = trait_config.get("pipeline").get("values") + epilogues = trait_config.get("epilogue").get("values") + schedulers = trait_config.get("scheduler").get("values") + pad_m_values = trait_config.get("pad_m").get("values") + pad_n_values = trait_config.get("pad_n").get("values") + pad_k_values = trait_config.get("pad_k").get("values") + a_preshuffle_values = trait_config.get("a_preshuffle_quant").get("values") + b_preshuffle_values = trait_config.get("b_preshuffle_quant").get("values") + + all_combinations = list( + itertools.product( + pipelines, + epilogues, + schedulers, + pad_m_values, + pad_n_values, + pad_k_values, + a_preshuffle_values, + b_preshuffle_values, + ) + ) + + # Filter out unsupported trait combinations + combinations = [] + for combo in all_combinations: + pipeline, epilogue, scheduler = combo[:3] + a_preshuffle = combo[6] + b_preshuffle = combo[7] + if is_trait_combination_valid( + pipeline, + epilogue, + scheduler, + b_preshuffle, + self.kernel_name_prefix, + self.layout, + ): + combinations.append(combo) + else: + logging.debug( + f"Skipping unsupported ABQuant trait combination: " + f"{pipeline}-{epilogue}-{scheduler}-a_pre={a_preshuffle}-b_pre={b_preshuffle}" + ) + return combinations + + def _list_kernels(self): + """Write kernel list to file for CMake to read.""" + tile_configs = self._get_tile_configs() + trait_combos = self._generate_trait_combinations() + group_size_n_values = self._get_group_size_n_values() + + kernel_list = [] + for tile_config in tile_configs: + for trait_combo in trait_combos: + for group_size_n in group_size_n_values: + ( + pipeline, + epilogue, + scheduler, + pad_m, + pad_n, + pad_k, + a_preshuffle, + b_preshuffle, + ) = trait_combo + + # Create kernel name with proper boolean capitalization + kernel_name = ( + f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_" + f"{pipeline}_{epilogue}_{scheduler}_" + f"{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_" + f"{str(a_preshuffle).capitalize()}_{str(b_preshuffle).capitalize()}_" + f"gsn{group_size_n}" + ) + + tile_str = self._tile_config_to_str(tile_config) + kernel_name += f"_{tile_str}" + + kernel_list.append( + { + "name": kernel_name, + "tile_config": tile_config, + "trait_combo": trait_combo, + "group_size_n": group_size_n, + } + ) + + # Apply RFC-compliant sampling (Sobol + LHS + maximin) + kernel_list = self._apply_sampling(kernel_list) + + # Write kernel count + with open( + self.working_path / f"{self.kernel_name_prefix}_kernel_count.txt", "w" + ) as f: + f.write(str(len(kernel_list))) + + # Write kernel list + with open( + self.working_path / f"{self.kernel_name_prefix}_kernel_list.txt", "w" + ) as f: + for kernel in kernel_list: + tile_config = kernel["tile_config"] + trait_combo = kernel["trait_combo"] + group_size_n = kernel["group_size_n"] + + tile_str = self._tile_config_to_str(tile_config) + + trait_str = ( + f"{trait_combo[0]}_{trait_combo[1]}_{trait_combo[2]}_" + + "_".join(str(x) for x in trait_combo[3:]) + + f"_gsn{group_size_n}" + ) + + f.write(f"{kernel['name']}|{tile_str}|{trait_str}\n") + + print(f"Listed {len(kernel_list)} kernel configurations") + + def _generate_kernel_instance(self, tile_config, trait_combo, group_size_n=None): + """Generate a single kernel instance for ABQuant. + + trait_combo is an 8-tuple for ABQuant. + group_size_n is the BQ group size along N dimension. + """ + if group_size_n is None: + group_size_n = self.group_size_n_values[0] + + k_block_per_cu = self.config.get("k_block_per_cu", 1) + + ( + pipeline, + epilogue, + scheduler, + pad_m, + pad_n, + pad_k, + a_preshuffle, + b_preshuffle, + ) = trait_combo + + # Create kernel name + kernel_name = ( + f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_" + f"{pipeline}_{epilogue}_{scheduler}_" + f"{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_" + f"{str(a_preshuffle).capitalize()}_{str(b_preshuffle).capitalize()}_" + f"gsn{group_size_n}" + ) + + tile_str = self._tile_config_to_str(tile_config) + kernel_name += f"_{tile_str}" + + # Pipeline maps + pipeline_impl_map = { + "compv3": "ck_tile::ABQuantGemmPipelineAgBgCrCompV3", + } + base_pipeline_map = { + "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", + } + scheduler_type_map = { + "intrawave": "ck_tile::GemmPipelineScheduler::Intrawave", + "interwave": "ck_tile::GemmPipelineScheduler::Interwave", + "default": "ck_tile::GemmPipelineScheduler::Default", + } + + instance_code = self.populate_kernel_header(kernel_name) + instance_code += self.populate_kernel_dtype_layout() + instance_code += self.populate_strut_begin(kernel_name) + instance_code += self.populate_tile_config(tile_config) + instance_code += self._populate_abquant_trait_config(trait_combo, group_size_n) + instance_code += self._populate_abquant_initialization( + base_pipeline_map, pipeline + ) + instance_code += self._populate_launch_abquant( + scheduler_type_map, + scheduler, + pipeline_impl_map, + pipeline, + epilogue, + k_block_per_cu, + ) + + # Write into a file + simplified_name = kernel_name + if simplified_name.startswith(f"{self.kernel_name_prefix}_"): + simplified_name = simplified_name[len(self.kernel_name_prefix) + 1 :] + + header_file = ( + self.working_path + / f"{self.kernel_name_prefix}_single_{simplified_name}.hpp" + ) + with open(header_file, "w") as f: + f.write(instance_code) + + print(f"Generated {header_file}") + + return kernel_name, instance_code + + def populate_kernel_header(self, kernel_name): + instance_code = f"""// Generated kernel instance for {kernel_name} +#pragma once + +#include +#include +#include +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp" +#include "ck_tile/ops/common/tensor_layout.hpp" +#include "ck_tile/ops/epilogue/default_2d_epilogue.hpp" +#include "ck_tile/ops/epilogue/cshuffle_epilogue.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp" +""" + return instance_code + + def populate_kernel_dtype_layout(self): + a_layout, b_layout, c_layout = _validation_utils.get_abc_layouts(self.layout) + dtype_str = _validation_utils.get_dtype_string(self.datatype) + c_type_str = _validation_utils.get_dtype_string( + "fp16" if self.datatype in ["fp8", "bf8"] else self.datatype + ) + + instance_code = f""" +using ADataType = {dtype_str}; +using BDataType = {dtype_str}; +using AccDataType = float; +using CDataType = {c_type_str}; +using AQDataType = float; +using BQDataType = float; + +using ALayout = {a_layout}; +using BLayout = {b_layout}; +using CLayout = {c_layout}; +using AQLayout = ck_tile::tensor_layout::gemm::RowMajor; +using BQLayout = ck_tile::tensor_layout::gemm::ColumnMajor; +""" + return instance_code + + def _populate_abquant_trait_config(self, trait_combo, group_size_n): + ( + pipeline, + epilogue, + scheduler, + pad_m, + pad_n, + pad_k, + a_preshuffle, + b_preshuffle, + ) = trait_combo + + instance_code = f""" + + // Traits configurations + static constexpr bool kPadM = {"true" if pad_m in [True, "true"] else "false"}; + static constexpr bool kPadN = {"true" if pad_n in [True, "true"] else "false"}; + static constexpr bool kPadK = {"true" if pad_k in [True, "true"] else "false"}; + static constexpr bool TransposeC = false; + static constexpr bool DoubleSmemBuffer = false; + static constexpr bool APreshuffleQuant = {"true" if a_preshuffle in [True, "true"] else "false"}; + static constexpr bool BPreshuffleQuant = {"true" if b_preshuffle in [True, "true"] else "false"}; + static constexpr bool PreshuffleB = false; + static constexpr ck_tile::index_t GroupSizeK = {self.group_size_k}; + static constexpr ck_tile::index_t GroupSizeN = {group_size_n};""" + + return instance_code + + def _populate_abquant_initialization(self, base_pipeline_map, pipeline): + instance_code = """ + + // Tile shape + using TileShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + + // Tile partitioner + using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner; + + // Quantization group sizes + using AQuantGroupSize = ck_tile::QuantGroupShape>; + using BQuantGroupSize = ck_tile::QuantGroupShape>; + + // Traits + using Traits = ck_tile::TileGemmQuantTraits< + kPadM, kPadN, kPadK, + APreshuffleQuant, + BPreshuffleQuant, + PreshuffleB, + ALayout, BLayout, CLayout, + ck_tile::QuantType::ABQuantGrouped, + AQLayout, BQLayout>; + + // Pipeline problem (base, for hot loop detection) + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase< + ADataType, + BDataType, + AccDataType, + TileShape, + Traits, + BDataType>;""" + + instance_code += f""" + + // Base pipeline for hot loop detection + using BaseGemmPipeline = {base_pipeline_map.get(pipeline)};""" + + return instance_code + + def _populate_launch_abquant( + self, + scheduler_type_map, + scheduler, + pipeline_impl_map, + pipeline, + epilogue, + k_block_per_cu, + ): + """Generate the complete launch function for ABQuant kernels.""" + instance_code = """ + + // Launch function + static float launch(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) { + + // Hot loop detection + const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, TileShape::kK); + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);""" + + instance_code += f""" + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{ + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + + // Full pipeline problem with hot loop and tail number + using PipelineProblem = ck_tile::GemmABQuantPipelineProblem< + ADataType, + AQDataType, + BDataType, + BQDataType, + AccDataType, + TileShape, + Traits, + AQuantGroupSize, + BQuantGroupSize, + false, + ADataType, + {scheduler_type_map.get(scheduler)}, + has_hot_loop_v, + tail_number_v>; + + using GemmPipeline = {pipeline_impl_map.get(pipeline)};""" + + instance_code += """ + + // Epilogue""" + if epilogue == "cshuffle": + instance_code += """ + using EpilogueProblem = ck_tile::CShuffleEpilogueProblem< + ADataType, + BDataType, + ck_tile::tuple<>, // DsDataType + AccDataType, + CDataType, + ck_tile::tuple<>, // DsLayout + CLayout, + ck_tile::element_wise::PassThrough, + TileM, // kM_ + TileN, // kN_ + WarpPerBlock_M, // MWave_ + WarpPerBlock_N, // NWave_ + WarpTileM, // kMPerXdl_ + WarpTileN, // kNPerXdl_ + WarpTileK, // kKPerXdl_ + TransposeC>; // isCTransposed_ + + using GemmEpilogue = ck_tile::CShuffleEpilogue;""" + else: + instance_code += """ + using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem< + ADataType, + BDataType, + ck_tile::tuple<>, // DsDataType + AccDataType, + CDataType, + ck_tile::tuple<>, // DsLayout + CLayout, + ck_tile::element_wise::PassThrough, + TileM, // kM_ + TileN, // kN_ + kPadM, + kPadN, + WarpTileM, // kMPerXdl_ + WarpTileN, // kNPerXdl_ + WarpTileK, // kKPerXdl_ + TransposeC>; // isCTransposed_ + + using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue;""" + + instance_code += f""" + + // Kernel type + using Kernel = ck_tile::QuantGemmKernel< + TilePartitioner, GemmPipeline, GemmEpilogue, + ck_tile::QuantType::ABQuantGrouped>; + + // Kernel arguments + auto kargs = Kernel::MakeKernelArgs(args); + + if (!Kernel::IsSupportedArgument(kargs)) {{ + throw std::runtime_error("Unsupported kernel arguments; skipping GEMM launch."); + }} + + // Get grid and block sizes + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(stream.log_level_ > 0) {{ + std::cout << "Launching kernel: " << Kernel::GetName() << '\\n' + << "grid: {{" << grids.x << ", " << grids.y << ", " << grids.z << "}}" + << ", blocks: {{" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}}" + << std::endl; + }} + + // Launch kernel + constexpr int kBlockPerCu = {k_block_per_cu}; + float ave_time = ck_tile::launch_kernel( + stream, + ck_tile::make_kernel(Kernel{{}}, grids, blocks, 0, kargs)); + + return ave_time; + }}; + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + }} +}}; +""" + return instance_code + + def _generate_all_individual(self, num_workers=None): + """Generate individual kernel files for separate compilation with parallel processing""" + if num_workers is None: + num_workers = min(multiprocessing.cpu_count(), 8) + + tile_configs = self._get_tile_configs() + trait_combos = self._generate_trait_combinations() + group_size_n_values = self._get_group_size_n_values() + + work_items = [] + for tile_config in tile_configs: + for trait_combo in trait_combos: + for group_size_n in group_size_n_values: + work_items.append( + ( + tile_config, + trait_combo, + group_size_n, + self.kernel_name_prefix, + self.working_path, + self.gpu_target, + self.datatype, + self.layout, + self.config_json, + ) + ) + + # Apply RFC-compliant sampling (Sobol + LHS + maximin) + if self.max_instances is not None and len(work_items) > self.max_instances: + kernel_dicts = [ + { + "tile_config": item[0], + "trait_combo": item[1], + "group_size_n": item[2], + "_work_item": item, + } + for item in work_items + ] + sampled = self._apply_sampling(kernel_dicts) + work_items = [k["_work_item"] for k in sampled] + + print( + f"Generating {len(work_items)} individual kernel files using {num_workers} workers..." + ) + print(f" Tile configs: {len(tile_configs)}") + print(f" Trait combinations: {len(trait_combos)}") + print(f" Group size N values: {group_size_n_values}") + print(f" Total kernels: {len(work_items)}") + + if work_items: + print(" First work item example:") + tile_config, trait_combo = work_items[0][:2] + print(f" Tile config: {tile_config}") + print(f" Trait combo: {trait_combo[:3]}") + + kernel_list = [] + completed = 0 + + with concurrent.futures.ProcessPoolExecutor( + max_workers=num_workers + ) as executor: + print(f" Submitting {len(work_items)} tasks to executor...") + future_to_item = { + executor.submit(_generate_single_kernel_individual, item): item + for item in work_items + } + print(" All tasks submitted, waiting for completion...") + + for future in concurrent.futures.as_completed(future_to_item): + completed += 1 + if completed % 100 == 0 or completed == len(work_items): + print( + f" Progress: {completed}/{len(work_items)} kernels generated" + ) + try: + result = future.result() + if result: + kernel_list.append(result) + except Exception as exc: + item = future_to_item[future] + print(f"Kernel generation failed for {item}: {exc}") + + kernel_list.sort(key=lambda x: x[0]) + self._generate_cmake_individual_targets(kernel_list) + + print( + f"Generated {len(kernel_list)} individual kernel files in {self.working_path}" + ) + + +def _generate_single_kernel_individual(work_item): + """Worker function to generate a single individual kernel file""" + ( + tile_config, + trait_combo, + group_size_n, + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json, + ) = work_item + + builder = GemmABQuantKernelBuilder( + kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json + ) + + try: + kernel_name, _ = builder._generate_kernel_instance( + tile_config, trait_combo, group_size_n + ) + return (kernel_name, trait_combo, tile_config) + except Exception as e: + print(f"Error generating individual kernel: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="GEMM ABQuant kernel instance builder") + parser.add_argument("--working_path", required=True, help="Working directory path") + parser.add_argument("--gpu_target", required=True, help="GPU target architecture") + parser.add_argument( + "--datatype", + required=True, + choices=["fp8", "bf8"], + help="Data type (fp8 or bf8)", + ) + parser.add_argument( + "--layout", + required=True, + choices=["rcr", "rrr", "ccr", "crr"], + help="Matrix layout", + ) + parser.add_argument("--config_json", help="Configuration JSON file") + parser.add_argument("--num_workers", type=int, help="Number of parallel workers") + parser.add_argument( + "--gen_all_individual", + action="store_true", + help="Generate individual kernel files", + ) + parser.add_argument( + "--gen_single", + action="store_true", + help="Generate a single kernel file", + ) + parser.add_argument("--kernel_name", help="Kernel name for single generation") + parser.add_argument("--tile_config", help="Tile configuration string") + parser.add_argument("--trait_combo", help="Trait combination string") + parser.add_argument( + "--list_kernels", + action="store_true", + help="List kernel configurations without generating files", + ) + parser.add_argument( + "--max-instances", + type=int, + default=None, + help="Maximum number of kernel instances to select via sampling", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="RNG seed for deterministic sampling; if omitted, derived from today's date", + ) + parser.add_argument( + "--tier", + default=None, + help="Sampling tier name (e.g., 'daily')", + ) + parser.add_argument( + "--manifest-path", + default=None, + help="Directory to write chosen_instances manifest JSON", + ) + + args = parser.parse_args() + + assert args.datatype in ["fp16", "bf16", "fp8", "bf8"], ( + f"Invalid datatype string: {args.datatype} (supported datatypes are [fp16, bf16, fp8, and bf8])" + ) + + layout_parts = args.layout.lower() + assert len(layout_parts) == 3, ( + f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr')" + ) + assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], ( + f"Invalid matrix_a layout: {layout_parts[0]} or matrix_b layout: {layout_parts[1]}" + ) + assert layout_parts[2] == "r", ( + f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' for row major)" + ) + + kernel_name_prefix = "gemm_abquant" + builder = GemmABQuantKernelBuilder( + kernel_name_prefix, + args.working_path, + args.gpu_target, + args.datatype, + args.layout, + args.config_json, + max_instances=args.max_instances, + seed=args.seed, + tier=args.tier, + manifest_path=args.manifest_path, + ) + + if args.list_kernels: + builder._list_kernels() + elif args.gen_single: + if not args.kernel_name or not args.tile_config or not args.trait_combo: + parser.error( + "--gen_single requires --kernel_name, --tile_config, and --trait_combo" + ) + + # Parse tile config + tile_parts = args.tile_config.split("_") + tile_dims = tile_parts[0].split("x") + warp_dims = tile_parts[1].split("x") + warp_tile_dims = tile_parts[2].split("x") + + tile_config = { + "tile_m": int(tile_dims[0]), + "tile_n": int(tile_dims[1]), + "tile_k": int(tile_dims[2]), + "warp_m": int(warp_dims[0]), + "warp_n": int(warp_dims[1]), + "warp_k": int(warp_dims[2]), + "warp_tile_m": int(warp_tile_dims[0]), + "warp_tile_n": int(warp_tile_dims[1]), + "warp_tile_k": int(warp_tile_dims[2]), + } + + # Parse trait combo for ABQuant: + # pipeline_epilogue_scheduler_padM_padN_padK_aPreshuffle_bPreshuffle_gsnN + trait_parts = args.trait_combo.split("_") + # Find the gsn part + group_size_n = 1 + gsn_idx = None + for i, part in enumerate(trait_parts): + if part.startswith("gsn"): + group_size_n = int(part[3:]) + gsn_idx = i + break + + if gsn_idx is not None: + # Remove gsn from trait_parts for parsing + trait_parts_no_gsn = trait_parts[:gsn_idx] + trait_parts[gsn_idx + 1 :] + else: + trait_parts_no_gsn = trait_parts + + if len(trait_parts_no_gsn) < 8: + parser.error( + f"--trait_combo must have 8 underscore-separated fields + gsn suffix " + f"(e.g. 'compv3_default_intrawave_False_False_False_False_False_gsn1'), " + f"got {len(trait_parts_no_gsn)} field(s): '{args.trait_combo}'" + ) + trait_combo = ( + trait_parts_no_gsn[0], # pipeline + trait_parts_no_gsn[1], # epilogue + trait_parts_no_gsn[2], # scheduler + trait_parts_no_gsn[3] == "True", # pad_m + trait_parts_no_gsn[4] == "True", # pad_n + trait_parts_no_gsn[5] == "True", # pad_k + trait_parts_no_gsn[6] == "True", # a_preshuffle_quant + trait_parts_no_gsn[7] == "True", # b_preshuffle_quant + ) + + builder._generate_kernel_instance(tile_config, trait_combo, group_size_n) + elif args.gen_all_individual: + builder._generate_all_individual(args.num_workers) + else: + parser.error( + "Must specify one of: --list_kernels, --gen_all_individual, or --gen_single" + ) + + +if __name__ == "__main__": + main() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp new file mode 100644 index 0000000000..7efb4dab55 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_profiler.hpp @@ -0,0 +1,336 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/host/tensor_shuffle_utils.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "gemm_abquant_benchmark.hpp" + +class ABQuantGemmProfiler +{ + public: + static ABQuantGemmProfiler& instance(Setting setting) + { + static ABQuantGemmProfiler instance{setting}; + return instance; + } + + void reset() { kernel_instances_.clear(); } + + void benchmark(ABQuantGemmProblem& problem, + std::function kernel_func) + { + std::vector(ck_tile::QuantGemmHostArgs&, + const ck_tile::stream_config&)>> + callables; + + callables.push_back( + [kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) { + float time = kernel_func(args, stream); + return std::make_tuple(std::string(KERNEL_NAME), time); + }); + + benchmark(problem, callables); + } + + void benchmark(ABQuantGemmProblem& problem, + std::vector( + ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables) + { + const ALayout layout_a = ALayout{}; + const BLayout layout_b = BLayout{}; + const CLayout layout_c = CLayout{}; + const AQLayout layout_aq = AQLayout{}; + const BQLayout layout_bq = BQLayout{}; + + problem.stride_a_ = ck_tile::get_default_stride( + problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)); + problem.stride_b_ = ck_tile::get_default_stride( + problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)); + problem.stride_c_ = ck_tile::get_default_stride( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)); + + // Compute AQ scale tensor dimensions: [M, K / group_size_k] + if(problem.k_ % problem.group_size_k_ != 0) + throw std::runtime_error("k_ must be divisible by group_size_k_"); + const ck_tile::index_t QK_A = problem.k_ / problem.group_size_k_; + problem.stride_aq_ = ck_tile::get_default_stride( + problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq)); + + // Compute BQ scale tensor dimensions: [K / group_size_k, N / group_size_n] + const ck_tile::index_t QK_B = problem.k_ / problem.group_size_k_; + const ck_tile::index_t QN_B = + (problem.group_size_n_ > 1) ? (problem.n_ / problem.group_size_n_) : problem.n_; + problem.stride_bq_ = + ck_tile::get_default_stride(QK_B, QN_B, problem.stride_bq_, is_row_major(layout_bq)); + + // Allocate host tensors + ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( + problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a))); + ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( + problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b))); + ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c))); + + // AQ scale tensor: [M, QK_A] + ck_tile::HostTensor aq_m_qk(ck_tile::host_tensor_descriptor( + problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq))); + + // BQ scale tensor: [QK_B, QN_B] + ck_tile::HostTensor bq_qk_n(ck_tile::host_tensor_descriptor( + QK_B, QN_B, problem.stride_bq_, is_row_major(layout_bq))); + + // Initialize tensors + if(setting_.init_method_ == 0) + { + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); + ck_tile::FillUniformDistribution{0.5f, 1.5f}(aq_m_qk); + ck_tile::FillUniformDistribution{0.5f, 1.5f}(bq_qk_n); + } + else if(setting_.init_method_ == 1) + { + ck_tile::FillMonotonicSeq{}(a_m_k); + ck_tile::FillMonotonicSeq{}(b_k_n); + ck_tile::FillConstant{static_cast(1)}(aq_m_qk); + ck_tile::FillConstant{static_cast(1)}(bq_qk_n); + } + else if(setting_.init_method_ == 2) + { + ck_tile::FillConstant{static_cast(1)}(a_m_k); + ck_tile::FillConstant{static_cast(1)}(b_k_n); + ck_tile::FillConstant{static_cast(1)}(aq_m_qk); + ck_tile::FillConstant{static_cast(1)}(bq_qk_n); + } + else + { + a_m_k.SetZero(); + b_k_n.SetZero(); + aq_m_qk.SetZero(); + bq_qk_n.SetZero(); + } + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); + ck_tile::DeviceMem aq_dev_buf(aq_m_qk.get_element_space_size_in_bytes()); + ck_tile::DeviceMem bq_dev_buf(bq_qk_n.get_element_space_size_in_bytes()); + + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + + // Shuffle AQ data if preshuffle is enabled + if constexpr(SelectedKernel::APreshuffleQuant) + { + constexpr int block_aq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK; + ck_tile::HostTensor aq_shuffled = ck_tile::shuffle_aq(&aq_m_qk, block_aq_k); + aq_dev_buf.ToDevice(aq_shuffled.data()); + } + else + { + aq_dev_buf.ToDevice(aq_m_qk.data()); + } + + // Shuffle BQ data if preshuffle is enabled + if constexpr(SelectedKernel::BPreshuffleQuant) + { + constexpr int block_bq_k = SelectedKernel::TileK / SelectedKernel::GroupSizeK; + ck_tile::HostTensor bq_shuffled = ck_tile::shuffle_bq(&bq_qk_n, block_bq_k); + bq_dev_buf.ToDevice(bq_shuffled.data()); + } + else + { + bq_dev_buf.ToDevice(bq_qk_n.data()); + } + + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + + // Build QuantGemmHostArgs + ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(), + b_k_n_dev_buf.GetDeviceBuffer(), + c_m_n_dev_buf.GetDeviceBuffer(), + aq_dev_buf.GetDeviceBuffer(), + bq_dev_buf.GetDeviceBuffer(), + problem.split_k_, + problem.m_, + problem.n_, + problem.k_, + QK_A, + QK_B, + problem.stride_a_, + problem.stride_b_, + problem.stride_c_, + problem.stride_aq_, + problem.stride_bq_); + + // Host reference for verification + ck_tile::HostTensor c_m_n_host_result(ck_tile::host_tensor_descriptor( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c))); + + if(setting_.verify_) + { + abquant_gemm_host_reference( + setting_.verify_, a_m_k, aq_m_qk, b_k_n, bq_qk_n, c_m_n_host_result); + } + + // Run kernel(s) + for(auto& callable : callables) + { + auto kernel_run_result = callable(gemm_args, + ck_tile::stream_config{nullptr, + true, + setting_.log_, + setting_.n_warmup_, + setting_.n_repeat_, + setting_.is_gpu_timer_, + setting_.flush_cache_, + setting_.rotating_count_}); + process_result(problem, + QK_A, + QK_B, + c_m_n_dev_buf, + c_m_n_host_result, + c_m_n_dev_result, + kernel_run_result); + } + } + + void process_result(const ABQuantGemmProblem& problem, + ck_tile::index_t QK_A, + ck_tile::index_t QK_B, + ck_tile::DeviceMem& c_m_n_dev_buf, + ck_tile::HostTensor& c_m_n_host_result, + ck_tile::HostTensor& c_m_n_dev_result, + const std::tuple& kernel_run_result) + { + auto [name, avg_time] = kernel_run_result; + + KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}}; + + // Compute performance metrics + std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_; + std::size_t num_byte = + sizeof(ADataType) * problem.m_ * problem.k_ + + sizeof(BDataType) * problem.n_ * problem.k_ + sizeof(AQDataType) * problem.m_ * QK_A + + sizeof(BQDataType) * QK_B * + (problem.group_size_n_ > 1 ? problem.n_ / problem.group_size_n_ : problem.n_) + + sizeof(CDataType) * problem.m_ * problem.n_; + + kernel_instance.perf_result_.latency_ = avg_time; + kernel_instance.perf_result_.tflops_ = static_cast(flop) / 1.E9 / avg_time; + kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time; + + if(setting_.log_ > 0 && !setting_.json_output_) + { + std::cout << kernel_instance << std::endl; + } + + // Verify result + bool verified_correct = true; + if(setting_.verify_) + { + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); + verified_correct = compare_abquant( + name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result); + } + + if(verified_correct) + { + kernel_instances_.emplace_back(kernel_instance); + } + else + { + std::cout << "Verification failed, skip kernel: " << name << std::endl; + } + + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + } + + KernelInstance select_best_instance(Metric metric) + { + if(kernel_instances_.empty()) + throw std::runtime_error("Empty instances"); + + auto kernel_instance = *std::max_element(kernel_instances_.begin(), + kernel_instances_.end(), + [metric](const auto& a, const auto& b) { + return PerformanceResult::compare( + b.perf_result_, a.perf_result_, metric); + }); + + if(setting_.json_output_) + { + std::cout << kernel_instance << std::endl; + } + else + { + std::cout << "**********************************" << std::endl; + std::cout << "According to given metrics: " << get_metric_name(metric) << "\n" + << "Current kernel performance is: " << kernel_instance << std::endl; + std::cout << "**********************************" << std::endl; + } + + if(!setting_.csv_filename_.empty()) + { + std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app); + + if(!file.is_open()) + { + std::cerr << "Warning: Failed to open CSV file for writing." << std::endl; + } + else + { + if(file.tellp() == 0) + { + file << "rocm_version,device_name," + << "split_k,m,n,k,stride_a,stride_b,stride_c,stride_aq,stride_bq," + << "group_size_k,group_size_n," + << "dtype_a,dtype_b,dtype_aq,dtype_bq,dtype_acc,dtype_c," + << "layout_a,layout_b,layout_c," << "name," + << "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n"; + } + + const auto& prob = kernel_instance.problem_; + const auto& perf = kernel_instance.perf_result_; + + file << get_rocm_version() << "," << ck_tile::get_device_name() << "," + << prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << "," + << prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << "," + << prob.stride_aq_ << "," << prob.stride_bq_ << "," << prob.group_size_k_ + << "," << prob.group_size_n_ << "," << prob.dtype_a_ << "," << prob.dtype_b_ + << "," << prob.dtype_aq_ << "," << prob.dtype_bq_ << "," << prob.dtype_acc_ + << "," << prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_ + << "," << prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed + << std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << "," + << perf.bandwidth_ << "," << get_metric_name(metric) << "\n"; + } + } + + return kernel_instance; + } + + ABQuantGemmProfiler(const ABQuantGemmProfiler&) = delete; + ABQuantGemmProfiler& operator=(const ABQuantGemmProfiler&) = delete; + + private: + ~ABQuantGemmProfiler() { kernel_instances_.clear(); } + ABQuantGemmProfiler(Setting setting) : setting_(setting) {} + + Setting setting_; + std::vector kernel_instances_; +}; diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py new file mode 100644 index 0000000000..320599f115 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/test_gemm_abquant_instance_builder.py @@ -0,0 +1,236 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Unit tests for GemmABQuantKernelBuilder (gemm_abquant operator).""" + +import json +import os +import sys +import tempfile +import unittest + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) + +from gemm_abquant_instance_builder import GemmABQuantKernelBuilder # noqa: E402 + +_MINIMAL_CONFIG = { + "tile_config": { + "tile_m": {"values": [64]}, + "tile_n": {"values": [64]}, + "tile_k": {"values": [128]}, + "warp_m": {"values": [2]}, + "warp_n": {"values": [2]}, + "warp_k": {"values": [1]}, + "warp_tile_m": {"values": [16]}, + "warp_tile_n": {"values": [16]}, + "warp_tile_k": {"values": [64]}, + }, + "trait_config": { + "pipeline": {"values": ["compv3"]}, + "scheduler": {"values": ["intrawave"]}, + "epilogue": {"values": ["default"]}, + "pad_m": {"values": [False]}, + "pad_n": {"values": [False]}, + "pad_k": {"values": [False]}, + "a_preshuffle_quant": {"values": [False]}, + "b_preshuffle_quant": {"values": [False, True]}, + }, + "k_block_per_cu": 1, + "group_size_k": 128, + "group_size_n": {"values": [1, 128]}, +} + + +def _make_builder(tmpdir, config=None, **kwargs): + cfg = config if config is not None else _MINIMAL_CONFIG + cfg_path = os.path.join(tmpdir, "config.json") + with open(cfg_path, "w") as f: + json.dump(cfg, f) + return GemmABQuantKernelBuilder( + kernel_name_prefix="gemm_abquant", + working_path=tmpdir, + gpu_target=kwargs.get("gpu_target", "gfx942"), + datatype=kwargs.get("datatype", "fp8"), + layout=kwargs.get("layout", "rcr"), + config_json=cfg_path, + ) + + +class TestGemmABQuantBuilderInit(unittest.TestCase): + def test_group_size_k_from_config(self): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + self.assertEqual(builder.group_size_k, 128) + + def test_group_size_n_dict_parsed(self): + """group_size_n specified as {"values": [...]} should expand to a list.""" + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + self.assertIsInstance(builder.group_size_n_values, list) + self.assertIn(1, builder.group_size_n_values) + self.assertIn(128, builder.group_size_n_values) + + def test_group_size_n_scalar_parsed(self): + """group_size_n specified as a plain int should become a single-element list.""" + cfg = json.loads(json.dumps(_MINIMAL_CONFIG)) + cfg["group_size_n"] = 64 + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir, config=cfg) + self.assertEqual(builder.group_size_n_values, [64]) + + def test_kernel_name_prefix_stored(self): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + self.assertEqual(builder.kernel_name_prefix, "gemm_abquant") + + def test_working_path_created(self): + with tempfile.TemporaryDirectory() as tmpdir: + sub = os.path.join(tmpdir, "workdir") + cfg_path = os.path.join(tmpdir, "config.json") + with open(cfg_path, "w") as f: + json.dump(_MINIMAL_CONFIG, f) + GemmABQuantKernelBuilder("gemm_abquant", sub, "gfx942", "fp8", "rcr", cfg_path) + self.assertTrue(os.path.isdir(sub)) + + +class TestGemmABQuantTraitCombinations(unittest.TestCase): + def setUp(self): + self._tmpdir = tempfile.TemporaryDirectory() + self.builder = _make_builder(self._tmpdir.name) + + def tearDown(self): + self._tmpdir.cleanup() + + def test_trait_combinations_non_empty(self): + combos = self.builder._generate_trait_combinations() + self.assertGreater(len(combos), 0) + + def test_trait_combo_is_8_tuple(self): + """ABQuant uses 8-tuples: + (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant, b_preshuffle_quant) + """ + combos = self.builder._generate_trait_combinations() + for combo in combos: + self.assertEqual( + len(combo), 8, f"Expected 8-tuple, got {len(combo)}: {combo}" + ) + + def test_pipeline_is_compv3(self): + combos = self.builder._generate_trait_combinations() + for combo in combos: + self.assertEqual(combo[0], "compv3") + + def test_b_preshuffle_quant_values(self): + combos = self.builder._generate_trait_combinations() + b_preshuffle_vals = {c[7] for c in combos} + self.assertIn(True, b_preshuffle_vals) + self.assertIn(False, b_preshuffle_vals) + + def test_a_preshuffle_quant_default_false(self): + """Config only lists False for a_preshuffle_quant, so all combos should have False.""" + combos = self.builder._generate_trait_combinations() + for combo in combos: + self.assertFalse(combo[6], "a_preshuffle_quant should be False per minimal config") + + +class TestGemmABQuantTileConfigs(unittest.TestCase): + def setUp(self): + self._tmpdir = tempfile.TemporaryDirectory() + self.builder = _make_builder(self._tmpdir.name) + + def tearDown(self): + self._tmpdir.cleanup() + + def test_tile_configs_non_empty(self): + configs = self.builder._get_tile_configs() + self.assertGreater(len(configs), 0) + + def test_tile_config_has_required_keys(self): + required = { + "tile_m", "tile_n", "tile_k", + "warp_m", "warp_n", "warp_k", + "warp_tile_m", "warp_tile_n", "warp_tile_k", + } + for cfg in self.builder._get_tile_configs(): + self.assertTrue(required.issubset(cfg.keys())) + + def test_tile_config_values_are_positive(self): + for cfg in self.builder._get_tile_configs(): + for key, val in cfg.items(): + self.assertGreater(val, 0, f"{key}={val} must be positive") + + +class TestGemmABQuantTileConfigStr(unittest.TestCase): + def test_tile_config_to_str_format(self): + tile_config = { + "tile_m": 128, "tile_n": 128, "tile_k": 64, + "warp_m": 2, "warp_n": 2, "warp_k": 1, + "warp_tile_m": 32, "warp_tile_n": 32, "warp_tile_k": 16, + } + result = GemmABQuantKernelBuilder._tile_config_to_str(tile_config) + self.assertEqual(result, "128x128x64_2x2x1_32x32x16") + + +class TestGemmABQuantGroupSizeN(unittest.TestCase): + def test_group_size_n_values_returned(self): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + gsn = builder._get_group_size_n_values() + self.assertIsInstance(gsn, list) + self.assertGreater(len(gsn), 0) + + def test_group_size_n_values_match_config(self): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + gsn = builder._get_group_size_n_values() + self.assertIn(1, gsn) + self.assertIn(128, gsn) + + +class TestGemmABQuantLayoutVariants(unittest.TestCase): + LAYOUTS = ["rcr", "rrr", "ccr", "crr"] + + def test_all_layouts_construct(self): + for layout in self.LAYOUTS: + with self.subTest(layout=layout): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir, layout=layout) + self.assertEqual(builder.layout, layout) + + +class TestGemmABQuantDataTypes(unittest.TestCase): + DTYPES = ["fp8", "bf8"] + + def test_all_dtypes_construct(self): + for dtype in self.DTYPES: + with self.subTest(dtype=dtype): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir, datatype=dtype) + self.assertEqual(builder.datatype, dtype) + + +class TestGemmABQuantMaxInstances(unittest.TestCase): + def test_max_instances_none_returns_all(self): + with tempfile.TemporaryDirectory() as tmpdir: + cfg_path = os.path.join(tmpdir, "config.json") + with open(cfg_path, "w") as f: + json.dump(_MINIMAL_CONFIG, f) + builder = GemmABQuantKernelBuilder( + "gemm_abquant", tmpdir, "gfx942", "fp8", "rcr", + config_json=cfg_path, max_instances=None, + ) + kernel_list = [ + {"tile_config": {"tile_m": 64, "tile_n": 64, "tile_k": 128, + "warp_m": 2, "warp_n": 2, "warp_k": 1, + "warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 64}, + "trait_combo": ("compv3", "default", "intrawave", + False, False, False, False, False), + "group_size_n": 1} + ] + result = builder._apply_sampling(kernel_list) + self.assertEqual(len(result), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt new file mode 100644 index 0000000000..49de66cfec --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/CMakeLists.txt @@ -0,0 +1,297 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +set(GEMM_AQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM AQuant (semicolon-separated)") +set(GEMM_AQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM AQuant (semicolon-separated)") +set(GEMM_AQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)") +option(ENABLE_CCACHE_GEMM_AQUANT "Enable ccache for GEMM AQuant ops compilation" OFF) + +# Store the directory path for use in functions +set(GEMM_AQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR}) + +# Function to create individual GEMM AQuant targets +function(create_individual_gemm_aquant_target datatype layout trait tile_config config_json) + if(NOT GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL) + message(WARNING "Skipping individual GEMM AQuant target ${datatype}_${layout}_${trait}_${tile_config}: No supported GPU targets") + return() + endif() + + # Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k + string(REPLACE "_" ";" config_groups ${tile_config}) + list(GET config_groups 0 tile_dims) + list(GET config_groups 1 warp_dims) + list(GET config_groups 2 warp_tile_dims) + + string(REPLACE "x" ";" tile_parts ${tile_dims}) + list(GET tile_parts 0 tile_m) + list(GET tile_parts 1 tile_n) + list(GET tile_parts 2 tile_k) + + string(REPLACE "x" ";" warp_parts ${warp_dims}) + list(GET warp_parts 0 warp_m) + list(GET warp_parts 1 warp_n) + list(GET warp_parts 2 warp_k) + + string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims}) + list(GET warp_tile_parts 0 warp_tile_m) + list(GET warp_tile_parts 1 warp_tile_n) + list(GET warp_tile_parts 2 warp_tile_k) + + set(target_name "benchmark_gemm_aquant_${datatype}_${layout}_${trait}_${tile_config}") + set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") + + # Generate the single instance header for this kernel + set(instance_header "${working_path}/gemm_aquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp") + + # Add custom command to generate the header file at build time + add_custom_command( + OUTPUT ${instance_header} + COMMAND ${Python3_EXECUTABLE} ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_instance_builder.py + --working_path ${working_path} + --datatype ${datatype} + --layout ${layout} + --config_json ${config_json} + --gen_single + --kernel_name "gemm_aquant_${datatype}_${layout}_${trait}_${tile_config}" + --tile_config "${tile_config}" + --trait_combo "${trait}" + --gpu_target "${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}" + DEPENDS ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_instance_builder.py ${config_json} + COMMENT "Generating ${instance_header}" + ) + + # Create the executable + add_executable(${target_name} + EXCLUDE_FROM_ALL + ${GEMM_AQUANT_SOURCE_DIR}/gemm_aquant_benchmark_single.cpp + ${instance_header} + ) + + # Set GPU architectures + set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}) + + # Set compile definitions + target_compile_definitions(${target_name} PRIVATE + GEMM_AQUANT_SINGLE_INSTANCE_HPP="${instance_header}" + ) + + # Include directories + target_include_directories(${target_name} PRIVATE + ${GEMM_AQUANT_SOURCE_DIR} + ${working_path} + ) + + # Compile options + target_compile_options(${target_name} PRIVATE + -Wno-undefined-func-template + -Wno-float-equal + --offload-compress + -include ${instance_header} + ) + + # Add to collection targets + add_dependencies(benchmark_gemm_aquant_all ${target_name}) + add_dependencies(benchmark_gemm_aquant_${datatype} ${target_name}) + add_dependencies(benchmark_gemm_aquant_${layout} ${target_name}) + add_dependencies(benchmark_gemm_aquant_${datatype}_${layout} ${target_name}) + + # Add to pipeline-specific and scheduler-specific targets + string(REPLACE "_" ";" trait_parts ${trait}) + list(GET trait_parts 0 pipeline) + list(GET trait_parts 1 epilogue) + list(GET trait_parts 2 scheduler) + + add_dependencies(benchmark_gemm_aquant_${pipeline}_pipeline ${target_name}) + add_dependencies(benchmark_gemm_aquant_${epilogue}_epilogue ${target_name}) + add_dependencies(benchmark_gemm_aquant_${scheduler}_scheduler ${target_name}) +endfunction() + +# Function to build individual GEMM AQuant targets +function(build_individual_gemm_aquant_targets datatype layout) + set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") + + # Choose config file (priority: env var > cmake var > default) + if(DEFINED ENV{GEMM_AQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_AQUANT_CONFIG_FILE}" STREQUAL "") + set(config_filename "$ENV{GEMM_AQUANT_CONFIG_FILE}") + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}") + message(VERBOSE " Using config from environment variable: ${config_filename}") + elseif(NOT "${GEMM_AQUANT_CONFIG_FILE}" STREQUAL "") + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_AQUANT_CONFIG_FILE}") + message(VERBOSE " Using custom config: ${GEMM_AQUANT_CONFIG_FILE}") + else() + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json") + message(VERBOSE " Using default config for layout ${layout}") + endif() + + if(NOT EXISTS ${json_blob}) + message(FATAL_ERROR "Config file not found: ${json_blob}") + endif() + + # Create working directory + file(MAKE_DIRECTORY ${working_path}) + + message(VERBOSE "Generating individual AQuant kernels for ${datatype} ${layout}...") + + # Build sampling arguments + set(extra_list_args "") + if(NOT "${GEMM_AQUANT_MAX_INSTANCES}" STREQUAL "") + list(APPEND extra_list_args --max-instances ${GEMM_AQUANT_MAX_INSTANCES}) + endif() + if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") + list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER}) + list(APPEND extra_list_args --manifest-path ${working_path}) + endif() + if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "") + list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED}) + endif() + + # List kernels + message(VERBOSE " Listing AQuant kernel configurations for ${datatype} ${layout}...") + execute_process( + COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_aquant_instance_builder.py + --working_path ${working_path} + --datatype ${datatype} + --layout ${layout} + --config_json ${json_blob} + --gpu_target "${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}" + --list_kernels + ${extra_list_args} + WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR} + RESULT_VARIABLE ret + OUTPUT_VARIABLE list_output + ERROR_VARIABLE list_error + ) + + if(NOT ret EQUAL 0) + message(FATAL_ERROR "Failed to list AQuant kernels for ${datatype} ${layout}: ${list_error}") + endif() + + # Read kernel count + if(EXISTS ${working_path}/gemm_aquant_kernel_count.txt) + file(READ ${working_path}/gemm_aquant_kernel_count.txt kernel_count) + string(STRIP "${kernel_count}" kernel_count) + message(VERBOSE " Found ${kernel_count} AQuant kernel configurations") + else() + message(FATAL_ERROR "Kernel count file not found") + endif() + + # Read kernel list and create targets + if(EXISTS ${working_path}/gemm_aquant_kernel_list.txt) + file(STRINGS ${working_path}/gemm_aquant_kernel_list.txt kernel_lines) + foreach(line IN LISTS kernel_lines) + string(REPLACE "|" ";" parts "${line}") + list(GET parts 0 kernel_name) + list(GET parts 1 tile_config) + list(GET parts 2 trait_combo) + + create_individual_gemm_aquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}") + endforeach() + else() + message(FATAL_ERROR "Kernel list file not found") + endif() +endfunction() + +# Main build logic +message(VERBOSE "=== Starting Tile Engine GEMM AQuant Configuration ===") +message(VERBOSE "GEMM_AQUANT_DATATYPE: ${GEMM_AQUANT_DATATYPE}") +message(VERBOSE "GEMM_AQUANT_LAYOUT: ${GEMM_AQUANT_LAYOUT}") +message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") + +# Filter GPU targets to supported architectures +set(GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL "") +set(DESIRED_TARGETS "gfx90a;gfx942;gfx950") + +foreach(target IN LISTS SUPPORTED_GPU_TARGETS) + if(target IN_LIST DESIRED_TARGETS) + list(APPEND GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL ${target}) + message(VERBOSE " Adding GPU target: ${target}") + endif() +endforeach() + +# Skip build if no matching targets found +if(NOT GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL) + message(WARNING "Skipping Tile Engine GEMM AQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") +else() + message(VERBOSE "Building individual GEMM AQuant targets for GPU targets: ${GEMM_AQUANT_GPU_TARGETS_INDIVIDUAL}") + + # Enable parallel compilation optimizations (directory-scoped to avoid overriding parent) + set_property(DIRECTORY PROPERTY JOB_POOLS + compile_heavy=4 # Limit heavy compilations to prevent OOM + compile_normal=16 # Allow more parallel normal compilations + ) + + # Enable compiler cache if requested + if(ENABLE_CCACHE_GEMM_AQUANT) + find_program(CCACHE_PROGRAM ccache) + if(CCACHE_PROGRAM) + set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM}) + message(VERBOSE "Using ccache for faster compilation") + endif() + endif() + + # Create master collection target + add_custom_target(benchmark_gemm_aquant_all) + + # Create datatype collection targets + foreach(dt IN LISTS GEMM_AQUANT_DATATYPE) + add_custom_target(benchmark_gemm_aquant_${dt}) + endforeach() + + # Create layout collection targets + foreach(l IN LISTS GEMM_AQUANT_LAYOUT) + add_custom_target(benchmark_gemm_aquant_${l}) + endforeach() + + # Create combined collection targets + foreach(dt IN LISTS GEMM_AQUANT_DATATYPE) + foreach(l IN LISTS GEMM_AQUANT_LAYOUT) + add_custom_target(benchmark_gemm_aquant_${dt}_${l}) + endforeach() + endforeach() + + # Create pipeline-specific collection targets + set(GEMM_AQUANT_PIPELINES "mem;compv3") + foreach(pipeline IN LISTS GEMM_AQUANT_PIPELINES) + add_custom_target(benchmark_gemm_aquant_${pipeline}_pipeline) + endforeach() + + # Create epilogue-specific collection targets + set(GEMM_AQUANT_EPILOGUES "default;cshuffle") + foreach(epilogue IN LISTS GEMM_AQUANT_EPILOGUES) + add_custom_target(benchmark_gemm_aquant_${epilogue}_epilogue) + endforeach() + + # Create scheduler-specific collection targets + set(GEMM_AQUANT_SCHEDULERS "intrawave;interwave") + foreach(scheduler IN LISTS GEMM_AQUANT_SCHEDULERS) + add_custom_target(benchmark_gemm_aquant_${scheduler}_scheduler) + endforeach() + + # Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that + # sampling fires per-combo rather than being a single cap. + # CMake integer division truncates; clamp the result to at least 1 so that a + # small total budget (e.g. 5 instances across 8 combos) does not silently + # produce a per-combo budget of 0 and skip all kernels. + if(NOT "${GEMM_AQUANT_MAX_INSTANCES}" STREQUAL "") + list(LENGTH GEMM_AQUANT_DATATYPE _aq_n_dt) + list(LENGTH GEMM_AQUANT_LAYOUT _aq_n_lay) + math(EXPR _aq_n_combos "${_aq_n_dt} * ${_aq_n_lay}") + if(_aq_n_combos GREATER 0) + math(EXPR GEMM_AQUANT_MAX_INSTANCES + "${GEMM_AQUANT_MAX_INSTANCES} / ${_aq_n_combos}") + if(GEMM_AQUANT_MAX_INSTANCES EQUAL 0) + set(GEMM_AQUANT_MAX_INSTANCES 1) + message(WARNING " gemm_aquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_aq_n_combos} combos)") + else() + message(STATUS " gemm_aquant: per-combo budget = ${GEMM_AQUANT_MAX_INSTANCES} (${_aq_n_combos} combos)") + endif() + endif() + endif() + + # Build individual targets for each datatype/layout combination + foreach(dt IN LISTS GEMM_AQUANT_DATATYPE) + foreach(l IN LISTS GEMM_AQUANT_LAYOUT) + build_individual_gemm_aquant_targets(${dt} ${l}) + endforeach() + endforeach() +endif() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json new file mode 100644 index 0000000000..a998050ebe --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_ci_config.json @@ -0,0 +1,92 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 16 + ] + }, + "tile_n": { + "values": [ + 64 + ] + }, + "tile_k": { + "values": [ + 256 + ] + }, + "warp_m": { + "values": [ + 1 + ] + }, + "warp_n": { + "values": [ + 4 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16 + ] + }, + "warp_tile_n": { + "values": [ + 16 + ] + }, + "warp_tile_k": { + "values": [ + 16, 32, 64, 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "mem", + "compv3" + ] + }, + "scheduler": { + "values": [ + "interwave", + "intrawave" + ] + }, + "epilogue": { + "values": [ + "default", + "cshuffle" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "a_preshuffle_quant": { + "values": [ + false, + true + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128 +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json new file mode 100644 index 0000000000..7e6c8581ec --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/default_config.json @@ -0,0 +1,102 @@ +{ + "tile_config": { + "tile_m": { + "max": 256, + "min": 64, + "step": 64 + }, + "tile_n": { + "max": 256, + "min": 64, + "step": 64 + }, + "tile_k": { + "max": 256, + "min": 64, + "step": 64 + }, + "warp_m": { + "values": [ + 4, + 2, + 1 + ] + }, + "warp_n": { + "values": [ + 4, + 2, + 1 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16, + 32 + ] + }, + "warp_tile_n": { + "values": [ + 16, + 32 + ] + }, + "warp_tile_k": { + "values": [ + 8, + 16, + 32, + 64, + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3", + "mem" + ] + }, + "scheduler": { + "values": [ + "intrawave", + "interwave" + ] + }, + "epilogue": { + "values": [ + "cshuffle", + "default" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "a_preshuffle_quant": { + "values": [ + false, + true + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128 +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json new file mode 100644 index 0000000000..37b077e1a2 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/configs/example_config.json @@ -0,0 +1,88 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 16 + ] + }, + "tile_n": { + "values": [ + 64 + ] + }, + "tile_k": { + "values": [ + 256 + ] + }, + "warp_m": { + "values": [ + 1 + ] + }, + "warp_n": { + "values": [ + 4 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16 + ] + }, + "warp_tile_n": { + "values": [ + 16 + ] + }, + "warp_tile_k": { + "values": [ + 16, 32, 64, 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "mem" + ] + }, + "scheduler": { + "values": [ + "interwave" + ] + }, + "epilogue": { + "values": [ + "cshuffle" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + true + ] + }, + "a_preshuffle_quant": { + "values": [ + false + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128 +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp new file mode 100644 index 0000000000..3fc07b2793 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.hpp @@ -0,0 +1,231 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "gemm_aquant_common.hpp" + +// Data types and Layouts are defined by the generated kernel headers: +// ADataType, BDataType, AQDataType, AccDataType, CDataType +// ALayout, BLayout, CLayout, AQLayout + +enum class Metric +{ + LATENCY = 0, + TFLOPS = 1, + BANDWIDTH = 2 +}; + +inline constexpr auto get_metric_name(Metric m) +{ + switch(m) + { + case Metric::LATENCY: return "latency"; + case Metric::TFLOPS: return "tflops"; + case Metric::BANDWIDTH: return "bandwidth"; + default: throw std::invalid_argument("Unsupported metric type"); + } +} + +struct AQuantGemmProblem +{ + int split_k_; + int m_, n_, k_; + int stride_a_, stride_b_, stride_c_; + int stride_aq_; + int group_size_k_; + + std::string dtype_a_, dtype_b_, dtype_aq_, dtype_acc_, dtype_c_; + std::string layout_a_, layout_b_, layout_c_; + + friend std::ostream& operator<<(std::ostream& os, const AQuantGemmProblem& problem) + { + os << "{\n" + << " \"split_k\":" << problem.split_k_ << ",\n" + << " \"m\":" << problem.m_ << ",\n" + << " \"n\":" << problem.n_ << ",\n" + << " \"k\":" << problem.k_ << ",\n" + << " \"stride_a\":" << problem.stride_a_ << ",\n" + << " \"stride_b\":" << problem.stride_b_ << ",\n" + << " \"stride_c\":" << problem.stride_c_ << ",\n" + << " \"stride_aq\":" << problem.stride_aq_ << ",\n" + << " \"group_size_k\":" << problem.group_size_k_ << ",\n" + << " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n" + << " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n" + << " \"dtype_aq\":\"" << problem.dtype_aq_ << "\",\n" + << " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n" + << " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n" + << " \"layout_a\":\"" << problem.layout_a_ << "\",\n" + << " \"layout_b\":\"" << problem.layout_b_ << "\",\n" + << " \"layout_c\":\"" << problem.layout_c_ << "\"\n" + << "}"; + return os; + } +}; + +struct PerformanceResult +{ + double latency_; + double tflops_; + double bandwidth_; + + static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m) + { + switch(m) + { + case Metric::LATENCY: return a.latency_ < b.latency_; + case Metric::TFLOPS: return a.tflops_ > b.tflops_; + case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_; + default: throw std::invalid_argument("Unsupported metric type"); + } + } + + friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result) + { + os << "{\n" + << " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_ + << ",\n" + << " \"tflops(TFlops)\": " << result.tflops_ << ",\n" + << " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n" + << "}"; + return os; + } +}; + +struct KernelInstance +{ + std::string name_; + AQuantGemmProblem problem_; + PerformanceResult perf_result_; + + static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m) + { + return PerformanceResult::compare(a.perf_result_, b.perf_result_, m); + } + + friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj) + { + os << "{\n" + << " \"name\": \"" << obj.name_ << "\",\n" + << " \"problem\": " << obj.problem_ << ",\n" + << " \"perf_result\": " << obj.perf_result_ << "\n" + << "}"; + return os; + } +}; + +struct Setting +{ + int n_warmup_; + int n_repeat_; + bool is_gpu_timer_; + int verify_; + int init_method_; + bool log_; + std::string csv_filename_; + bool flush_cache_; + int rotating_count_; + bool json_output_; +}; + +inline std::string get_rocm_version() +{ + // Try the legacy path first, then the path used by newer ROCm installs. + for(const char* path : {"/opt/rocm/.info/version", "/opt/rocm/lib/rocm_version.h"}) + { + std::ifstream version_file(path); + if(version_file.is_open()) + { + std::string version; + std::getline(version_file, version); + if(!version.empty()) + return version; + } + } + return "Unknown"; +} + +template +auto calculate_rtol_atol_aquant(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeType = + std::conditional_t; + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +/// @brief Compare device and host results for AQuant GEMM +bool compare_aquant(std::string instanceName, + ck_tile::index_t K, + ck_tile::index_t kbatch, + ck_tile::HostTensor& c_m_n_dev_result, + ck_tile::HostTensor& c_m_n_host_result) +{ + const float max_accumulated_value = + std::abs(static_cast(*std::max_element(c_m_n_host_result.mData.begin(), + c_m_n_host_result.mData.end(), + [](const auto& a, const auto& b) { + return std::abs(static_cast(a)) < + std::abs(static_cast(b)); + }))); + const auto rtol_atol = + calculate_rtol_atol_aquant( + K, kbatch, max_accumulated_value); + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_result, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + std::cerr << "For " << instanceName << " Relative error threshold is " + << rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is " + << rtol_atol.at(ck_tile::number<1>{}) << std::endl; + std::cerr << "The verification result is:" << (pass ? "correct" : "fail") << std::endl; + + return pass; +} + +/// @brief CPU reference implementation for AQuant GEMM +void aquant_gemm_host_reference(int verify, + ck_tile::HostTensor& a_m_k, + ck_tile::HostTensor& aq_m_qk, + ck_tile::HostTensor& b_k_n, + ck_tile::HostTensor& c_m_n_host_result) +{ + if(verify == 1) + { + c_m_n_host_result.SetZero(); + using QuantGroupSize = typename SelectedKernel::QuantGroupSize; + ck_tile::reference_gemm_quant(a_m_k, aq_m_qk, b_k_n, c_m_n_host_result); + } +} +#pragma clang diagnostic pop diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py new file mode 100644 index 0000000000..2f55340643 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py @@ -0,0 +1,547 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import sys +import json +import csv +import subprocess +import argparse +import time +from pathlib import Path +from typing import List, Dict, Tuple, Optional + + +class AQuantGemmBenchmark: + def __init__(self, build_dir: str, verbose: bool = False): + self.build_dir = Path(build_dir) + self.verbose = verbose + self.results = [] + + def discover_kernels(self) -> List[Path]: + """Find all benchmark_gemm_aquant_* executables in the build directory.""" + bin_dir = self.build_dir / "bin" + if not bin_dir.exists(): + print(f"Error: Binary directory {bin_dir} does not exist") + return [] + + kernels = list(bin_dir.glob("benchmark_gemm_aquant_*")) + if self.verbose: + print(f"Found {len(kernels)} AQuant kernel executables") + for k in kernels: + print(f" - {k.name}") + return kernels + + def extract_kernel_info(self, kernel_path: Path) -> Dict[str, str]: + """Extract kernel information from filename.""" + name = kernel_path.stem + + info = { + "executable": str(kernel_path), + "name": name, + "data_type": "unknown", + "layout": "unknown", + "pipeline": "unknown", + "scheduler": "unknown", + "epilogue": "unknown", + "a_preshuffle_quant": False, + } + + # Parse: benchmark_gemm_aquant_fp8_rcr_mem_default_interwave_False_False_True_False_128x128x128_... + parts = name.split("_") + + if len(parts) >= 3: + # Skip "benchmark_gemm_aquant" prefix (3 parts) + idx = 3 + if idx < len(parts): + info["data_type"] = parts[idx] + idx += 1 + if idx < len(parts): + info["layout"] = parts[idx] + idx += 1 + if idx < len(parts): + info["pipeline"] = parts[idx] + idx += 1 + if idx < len(parts): + info["epilogue"] = parts[idx] + idx += 1 + if idx < len(parts): + info["scheduler"] = parts[idx] + + # Parse tile dimensions + config_info = self.parse_detailed_config(name) + info.update(config_info) + + info["config_id"] = self.generate_config_id(info) + return info + + def parse_detailed_config(self, kernel_name: str) -> Dict: + """Parse tile dimensions and boolean flags from kernel name. + + Kernel name format (after "benchmark_gemm_aquant_"): + {dtype}_{layout}_{pipeline}_{epilogue}_{scheduler} + _{padM}_{padN}_{padK}_{aPreshuffle} + _{TileMxTileNxTileK}_{WarpMxWarpNxWarpK}_{WarpTileMxWarpTileNxWarpTileK} + + Dimensions are extracted positionally (in order of appearance) rather than + by sorting, which would silently misassign groups that share a max value. + """ + config = { + "tile_sizes": {"tile_m": 0, "tile_n": 0, "tile_k": 0}, + "warp_config": {"warp_m": 0, "warp_n": 0, "warp_k": 0}, + "warp_tile": {"warp_tile_m": 0, "warp_tile_n": 0, "warp_tile_k": 0}, + "optimization_flags": { + "pad_m": False, + "pad_n": False, + "pad_k": False, + "a_preshuffle_quant": False, + }, + } + + parts = kernel_name.split("_") + + # Locate the first boolean token; collect the contiguous boolean run + bool_start = -1 + bool_sequence = [] + for i, part in enumerate(parts): + if part in ("True", "False"): + if bool_start == -1: + bool_start = i + bool_sequence.append(part == "True") + elif bool_start != -1: + # End of the boolean run + break + + if len(bool_sequence) >= 4: + config["optimization_flags"]["pad_m"] = bool_sequence[0] + config["optimization_flags"]["pad_n"] = bool_sequence[1] + config["optimization_flags"]["pad_k"] = bool_sequence[2] + config["optimization_flags"]["a_preshuffle_quant"] = bool_sequence[3] + + # Extract dimension groups (NxNxN tokens) in positional order. + # Appearance order in the name is: tile, warp, warp_tile. + dimension_groups = [] + for part in parts: + sub = part.split("x") + if len(sub) == 3: + try: + dims = [int(v) for v in sub] + if all(d > 0 for d in dims): + dimension_groups.append(dims) + except ValueError: + continue + + if len(dimension_groups) >= 3: + config["tile_sizes"]["tile_m"] = dimension_groups[0][0] + config["tile_sizes"]["tile_n"] = dimension_groups[0][1] + config["tile_sizes"]["tile_k"] = dimension_groups[0][2] + config["warp_config"]["warp_m"] = dimension_groups[1][0] + config["warp_config"]["warp_n"] = dimension_groups[1][1] + config["warp_config"]["warp_k"] = dimension_groups[1][2] + config["warp_tile"]["warp_tile_m"] = dimension_groups[2][0] + config["warp_tile"]["warp_tile_n"] = dimension_groups[2][1] + config["warp_tile"]["warp_tile_k"] = dimension_groups[2][2] + elif len(dimension_groups) == 2: + config["tile_sizes"]["tile_m"] = dimension_groups[0][0] + config["tile_sizes"]["tile_n"] = dimension_groups[0][1] + config["tile_sizes"]["tile_k"] = dimension_groups[0][2] + config["warp_config"]["warp_m"] = dimension_groups[1][0] + config["warp_config"]["warp_n"] = dimension_groups[1][1] + config["warp_config"]["warp_k"] = dimension_groups[1][2] + elif len(dimension_groups) == 1: + config["tile_sizes"]["tile_m"] = dimension_groups[0][0] + config["tile_sizes"]["tile_n"] = dimension_groups[0][1] + config["tile_sizes"]["tile_k"] = dimension_groups[0][2] + + return config + + def generate_config_id(self, info: Dict) -> str: + """Generate a compact config ID from kernel info.""" + parts = [ + info.get("data_type", "unk"), + info.get("layout", "unk"), + info.get("pipeline", "unk"), + info.get("scheduler", "unk"), + ] + + tile_sizes = info.get("tile_sizes", {}) + if tile_sizes.get("tile_m", 0) > 0: + parts.append( + f"{tile_sizes['tile_m']}x{tile_sizes['tile_n']}x{tile_sizes['tile_k']}" + ) + + warp_config = info.get("warp_config", {}) + if warp_config.get("warp_m", 0) > 0: + parts.append( + f"w{warp_config['warp_m']}x{warp_config['warp_n']}x{warp_config['warp_k']}" + ) + + warp_tile = info.get("warp_tile", {}) + if warp_tile.get("warp_tile_m", 0) > 0: + parts.append( + f"wt{warp_tile['warp_tile_m']}x{warp_tile['warp_tile_n']}x{warp_tile['warp_tile_k']}" + ) + + return "_".join(parts) + + def run_kernel(self, kernel_path: Path, params: Dict[str, str]) -> Optional[Dict]: + """Run a single kernel with given parameters.""" + results_dir = self.build_dir / "results" + results_dir.mkdir(exist_ok=True) + + json_file = results_dir / f"{kernel_path.stem}.json" + + cmd = [str(kernel_path)] + for key, value in params.items(): + cmd.append(f"-{key}={value}") + cmd.append("-json_output=true") + + if self.verbose: + print(f"Running: {' '.join(cmd)}") + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=60) + + if result.returncode != 0: + print(f"Error running {kernel_path.name}: {result.stderr}") + return None + + output = result.stdout.strip() + if output: + with open(json_file, "w") as f: + f.write(output) + return self.parse_json_file(json_file) + else: + print(f"No output from {kernel_path.name}") + return None + + except subprocess.TimeoutExpired: + print(f"Timeout running {kernel_path.name}") + return None + except Exception as e: + print(f"Error running {kernel_path.name}: {e}") + return None + + def parse_json_file(self, json_file: Path) -> Optional[Dict]: + """Parse JSON data from kernel output file.""" + try: + with open(json_file, "r") as f: + content = f.read().strip() + + data = json.loads(content) + result = data.copy() + if "perf_result" in data: + perf = data["perf_result"] + result["time_ms"] = perf.get("latency(ms)", 0) + result["tflops"] = perf.get("tflops(TFlops)", 0) + result["bandwidth_gb_s"] = perf.get("bandwidth(GB/s)", 0) + return result + + except (json.JSONDecodeError, Exception) as e: + if self.verbose: + print(f"Failed to parse JSON from {json_file}: {e}") + return None + + def benchmark_problem_size( + self, + kernels: List[Path], + m: int, + n: int, + k: int, + group_size_k: int = 128, + split_k: int = 1, + verify: int = 0, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> List[Dict]: + """Benchmark all kernels for a specific problem size.""" + results = [] + + params = { + "m": m, + "n": n, + "k": k, + "group_size_k": group_size_k, + "split_k": split_k, + "verify": verify, + "warmup": warmup, + "repeat": repeat, + "flush_cache": str(flush_cache).lower(), + "rotating_count": rotating_count, + } + + print(f"\nBenchmarking M={m}, N={n}, K={k}, group_size_k={group_size_k}") + + for kernel_path in kernels: + kernel_info = self.extract_kernel_info(kernel_path) + result = self.run_kernel(kernel_path, params) + + if result: + structured_result = { + "name": kernel_info["name"], + "config_id": kernel_info["config_id"], + "problem": result.get("problem", {}), + "perf_result": result.get("perf_result", {}), + "config": { + "data_type": kernel_info["data_type"], + "layout": kernel_info["layout"], + "pipeline": kernel_info["pipeline"], + "scheduler": kernel_info["scheduler"], + "epilogue": kernel_info["epilogue"], + "tile_sizes": kernel_info.get("tile_sizes", {}), + "warp_config": kernel_info.get("warp_config", {}), + "warp_tile": kernel_info.get("warp_tile", {}), + "optimization_flags": kernel_info.get("optimization_flags", {}), + }, + "executable": kernel_info["executable"], + "time_ms": result.get("time_ms", 0), + "tflops": result.get("tflops", 0), + "bandwidth_gb_s": result.get("bandwidth_gb_s", 0), + } + + results.append(structured_result) + + if self.verbose: + print( + f" {kernel_info['config_id']}: " + f"{structured_result['tflops']:.2f} TFLOPS, " + f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " + f"{structured_result['time_ms']:.2f}ms" + ) + + return results + + def find_best_kernel( + self, results: List[Dict], metric: str = "tflops" + ) -> Optional[Dict]: + """Find the best performing kernel based on metric.""" + if not results: + return None + + if metric == "tflops": + return max(results, key=lambda x: x.get("tflops", 0)) + elif metric == "time_ms": + return min(results, key=lambda x: x.get("time_ms", float("inf"))) + elif metric == "bandwidth_gb_s": + return max(results, key=lambda x: x.get("bandwidth_gb_s", 0)) + else: + raise ValueError(f"Unknown metric: {metric}") + + def benchmark_sweep( + self, + problem_sizes: List[Tuple[int, int, int]], + group_size_k: int = 128, + split_k_values: Optional[List[int]] = None, + verify: bool = False, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> Dict: + """Run comprehensive benchmark sweep.""" + if split_k_values is None: + split_k_values = [1] + kernels = self.discover_kernels() + if not kernels: + print("No kernels found!") + return {} + + all_results = [] + best_kernels = {} + + for m, n, k in problem_sizes: + for split_k in split_k_values: + results = self.benchmark_problem_size( + kernels, + m, + n, + k, + group_size_k=group_size_k, + split_k=split_k, + verify=1 if verify else 0, + warmup=warmup, + repeat=repeat, + flush_cache=flush_cache, + rotating_count=rotating_count, + ) + + all_results.extend(results) + + best = self.find_best_kernel(results) + if best: + key = f"m{m}_n{n}_k{k}_splitk{split_k}" + best_kernels[key] = best + print( + f"Best for {key}: {best['name']} " + f"({best['tflops']:.2f} TFLOPS, " + f"{best['bandwidth_gb_s']:.2f} GB/s, " + f"{best['time_ms']:.2f}ms)" + ) + + self.results = all_results + return best_kernels + + def export_csv(self, filename: str): + """Export all results to CSV.""" + if not self.results: + print("No results to export") + return + + all_keys = set() + for result in self.results: + all_keys.update(result.keys()) + + fieldnames = sorted(all_keys) + + with open(filename, "w", newline="") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(self.results) + + print(f"Results exported to {filename}") + + def export_best_kernels(self, best_kernels: Dict, filename: str): + """Export best kernel selections to file.""" + with open(filename, "w") as f: + f.write("# Best AQuant kernel selections\n") + f.write( + "# Format: problem_size -> kernel_name (TFLOPS, bandwidth, latency)\n\n" + ) + + for key, kernel in sorted(best_kernels.items()): + f.write( + f"{key}: {kernel['name']} " + f"({kernel['tflops']:.2f} TFLOPS, " + f"{kernel['bandwidth_gb_s']:.2f} GB/s, " + f"{kernel['time_ms']:.2f}ms)\n" + ) + + print(f"Best kernels exported to {filename}") + + def export_json(self, filename: str, best_kernels: Dict = None): + """Export results to JSON.""" + from datetime import datetime + + output_data = { + "benchmark_metadata": { + "timestamp": datetime.now().isoformat(), + "operator": "gemm_aquant", + "total_kernels_tested": len(self.results), + "successful_runs": len( + [r for r in self.results if r.get("tflops", 0) > 0] + ), + }, + "kernel_results": self.results, + "best_kernels_by_problem": best_kernels or {}, + } + + with open(filename, "w") as f: + json.dump(output_data, f, indent=2) + print(f"JSON results exported to {filename}") + + +def main(): + parser = argparse.ArgumentParser(description="AQuant GEMM Kernel Benchmarking Tool") + parser.add_argument( + "build_dir", help="Build directory containing kernel executables" + ) + parser.add_argument( + "--problem-sizes", + nargs="+", + default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"], + help="Problem sizes as M,N,K tuples", + ) + parser.add_argument( + "--group-size-k", + type=int, + default=128, + help="Quantization group size along K (default: 128)", + ) + parser.add_argument( + "--split-k", + nargs="+", + type=int, + default=[1], + help="Split-K values to test", + ) + parser.add_argument("--verify", action="store_true", help="Enable verification") + parser.add_argument("--verbose", action="store_true", help="Verbose output") + parser.add_argument( + "--warmup", + type=int, + default=50, + help="Number of warmup iterations", + ) + parser.add_argument( + "--repeat", + type=int, + default=100, + help="Number of benchmark iterations", + ) + parser.add_argument( + "--csv", + default="gemm_aquant_benchmark_results.csv", + help="CSV output filename (default: gemm_aquant_benchmark_results.csv)", + ) + parser.add_argument( + "--best", + default="best_aquant_kernels.txt", + help="Best kernels output filename (default: best_aquant_kernels.txt)", + ) + parser.add_argument("--json", help="JSON output filename (optional)") + parser.add_argument( + "--flush-cache", + action="store_true", + default=True, + help="Flush cache between kernel runs (default: true)", + ) + parser.add_argument( + "--rotating-count", + type=int, + default=1000, + help="Number of iterations to rotate the cache (default: 1000)", + ) + + args = parser.parse_args() + + problem_sizes = [] + for size_str in args.problem_sizes: + try: + m, n, k = map(int, size_str.split(",")) + problem_sizes.append((m, n, k)) + except ValueError: + print(f"Invalid problem size: {size_str}") + return 1 + + benchmark = AQuantGemmBenchmark(args.build_dir, verbose=args.verbose) + + print("Starting AQuant GEMM kernel benchmark sweep...") + start_time = time.time() + + best_kernels = benchmark.benchmark_sweep( + problem_sizes=problem_sizes, + group_size_k=args.group_size_k, + split_k_values=args.split_k, + verify=args.verify, + warmup=args.warmup, + repeat=args.repeat, + flush_cache=args.flush_cache, + rotating_count=args.rotating_count, + ) + + elapsed_time = time.time() - start_time + print(f"\nBenchmark completed in {elapsed_time:.2f} seconds") + + # Export results + benchmark.export_csv(args.csv) + benchmark.export_best_kernels(best_kernels, args.best) + + if args.json: + benchmark.export_json(args.json, best_kernels) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp new file mode 100644 index 0000000000..637cef5c54 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark_single.cpp @@ -0,0 +1,153 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "gemm_aquant_profiler.hpp" +#include "gemm_aquant_common.hpp" + +// The kernel header is included via the compile command line with -include flag +// It defines SelectedKernel struct, KERNEL_NAME, and type aliases: +// ADataType, BDataType, AQDataType, AccDataType, CDataType +// ALayout, BLayout, CLayout, AQLayout + +inline auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.") + .insert("n", "4096", "The value for n dimension. Default is 4096.") + .insert("k", "2048", "The value for k dimension. Default is 2048.") + .insert("stride_a", "0", "The stride value for tensor A. Default is 0.") + .insert("stride_b", "0", "The stride value for tensor B. Default is 0.") + .insert("stride_c", "0", "The stride value for tensor C. Default is 0.") + .insert("split_k", "1", "The split value for k dimension. Default is 1.") + .insert("group_size_k", + std::to_string(SelectedKernel::GroupSizeK), + "The quantization group size along K. Default matches kernel config.") + .insert("verify", + "1", + "The type of validation. Set to 0 for no validation, 1 for validation on CPU. " + "Default is 1, CPU validation.") + .insert("log", + "false", + "Whether output kernel instance information or not. Possible values are true or " + "false. Default is false") + .insert( + "warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.") + .insert( + "repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.") + .insert("timer", + "true", + "Whether if the timer is gpu timer or not. Possible values are false or true. " + "Default is true.") + .insert("init", + "0", + "The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 " + "for constant(1). Default is 0, random.") + .insert("flush_cache", + "true", + "To flush cache, possible values are true or false. Default is true.") + .insert( + "rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.") + .insert("metric", + "0", + "Metric with which to measure kernel performance. Set to 0 for latency, 1 for " + "tflops, or 2 for bandwidth. Default is 0, latency.") + .insert("csv_filename", + "", + "The filename of benchmark result. Default is empty (no CSV output).") + .insert("json_output", + "false", + "Whether to output results in JSON format only. Possible values are true or false. " + "Default is false"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} + +void benchmark_single(const ck_tile::ArgParser& arg_parser) +{ + std::string dtype_a = DataTypeTraits::name; + std::string dtype_b = DataTypeTraits::name; + std::string dtype_aq = DataTypeTraits::name; + std::string dtype_acc = DataTypeTraits::name; + std::string dtype_c = DataTypeTraits::name; + + std::string layout_a = ALayout::name; + std::string layout_b = BLayout::name; + std::string layout_c = CLayout::name; + + int group_size_k = arg_parser.get_int("group_size_k"); + + AQuantGemmProblem problem{arg_parser.get_int("split_k"), + arg_parser.get_int("m"), + arg_parser.get_int("n"), + arg_parser.get_int("k"), + arg_parser.get_int("stride_a"), + arg_parser.get_int("stride_b"), + arg_parser.get_int("stride_c"), + 0, // stride_aq computed by profiler + group_size_k, + dtype_a, + dtype_b, + dtype_aq, + dtype_acc, + dtype_c, + layout_a, + layout_b, + layout_c}; + + Setting setting{arg_parser.get_int("warmup"), + arg_parser.get_int("repeat"), + arg_parser.get_bool("timer"), + arg_parser.get_int("verify"), + arg_parser.get_int("init"), + arg_parser.get_bool("log"), + arg_parser.get_str("csv_filename"), + arg_parser.get_bool("flush_cache"), + arg_parser.get_int("rotating_count"), + arg_parser.get_bool("json_output")}; + + AQuantGemmProfiler profiler{setting}; + + try + { + auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& stream) { + return SelectedKernel::launch(args, stream); + }; + + profiler.benchmark(problem, kernel_func); + profiler.select_best_instance(static_cast(arg_parser.get_int("metric"))); + } + catch(const std::exception& e) + { + std::cerr << "Benchmark failed: " << e.what() << std::endl; + } +} + +int main(int argc, char* argv[]) +{ + try + { + auto [result, parser] = create_args(argc, argv); + if(!result) + return EXIT_FAILURE; + + benchmark_single(parser); + return 0; + } + catch(const std::exception& e) + { + std::cerr << "Error: " << e.what() << "\n"; + return EXIT_FAILURE; + } +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp new file mode 100644 index 0000000000..00fb4e1d7a --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_common.hpp @@ -0,0 +1,73 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/core/numeric/integer.hpp" + +// DataTypeTraits for all supported types +template +struct DataTypeTraits; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp32"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp16"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "bf16"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp8"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "bf8"; +}; + +// Helper function to determine if a layout is row-major +template +constexpr auto is_row_major(Layout) +{ + return ck_tile::bool_constant>{}; +} + +// Structure to hold kernel traits for dispatcher +struct AQuantKernelTraits +{ + std::string pipeline; // mem, compv3 + std::string scheduler; // intrawave, interwave + std::string epilogue; // default + bool pad_m; + bool pad_n; + bool pad_k; + bool a_preshuffle_quant; + + AQuantKernelTraits() + : pipeline("mem"), + scheduler("interwave"), + epilogue("default"), + pad_m(false), + pad_n(false), + pad_k(true), + a_preshuffle_quant(false) + { + } +}; diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py new file mode 100644 index 0000000000..a57a3bb2e8 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_instance_builder.py @@ -0,0 +1,322 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import os +import argparse +import importlib.util +import multiprocessing +import concurrent.futures + + +def _import_gemm_kernel_builder(): + """Import the base GemmKernelBuilder from the gemm directory.""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + gemm_dir = os.path.dirname(os.path.dirname(current_dir)) + + spec = importlib.util.spec_from_file_location( + "gemm_instance_builder", + os.path.join(gemm_dir, "gemm_instance_builder.py"), + ) + gemm_builder_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(gemm_builder_module) + + return gemm_builder_module.GemmKernelBuilder + + +GemmKernelBuilder = _import_gemm_kernel_builder() + + +class GemmAQuantKernelBuilder(GemmKernelBuilder): + def __init__( + self, + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json=None, + max_instances=None, + seed=None, + tier=None, + manifest_path=None, + ): + super().__init__( + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json, + max_instances=max_instances, + seed=seed, + tier=tier, + manifest_path=manifest_path, + ) + self.group_size_k = self.config.get("group_size_k", 128) + + def _uses_persistent_trait(self): + # The 7th trait slot carries a_preshuffle_quant, not a persistent-kernel flag. + # Return True so the base class includes it in kernel names and sampling features. + return True + + def _generate_all_individual(self, num_workers=None): + """Generate individual kernel files for separate compilation with parallel processing""" + if num_workers is None: + num_workers = min( + multiprocessing.cpu_count(), 8 + ) # Limit to avoid memory issues + + tile_configs = self._get_tile_configs() + trait_combos = self._generate_trait_combinations() + + # Prepare work items for parallel processing + work_items = [] + for tile_config in tile_configs: + for trait_combo in trait_combos: + work_items.append( + ( + tile_config, + trait_combo, + self.kernel_name_prefix, + self.working_path, + self.gpu_target, + self.datatype, + self.layout, + self.config_json, + ) + ) + + # Apply RFC-compliant sampling (Sobol + LHS + maximin) + if self.max_instances is not None and len(work_items) > self.max_instances: + kernel_dicts = [ + { + "tile_config": item[0], + "trait_combo": item[1], + "_work_item": item, + } + for item in work_items + ] + sampled = self._apply_sampling(kernel_dicts) + work_items = [k["_work_item"] for k in sampled] + + print( + f"Generating {len(work_items)} individual kernel files using {num_workers} workers..." + ) + print(f" Tile configs: {len(tile_configs)}") + print(f" Trait combinations: {len(trait_combos)}") + print(f" Total kernels: {len(work_items)}") + + # Process work items in parallel + kernel_list = [] + completed = 0 + + with concurrent.futures.ProcessPoolExecutor( + max_workers=num_workers + ) as executor: + # Submit all work items + print(f" Submitting {len(work_items)} tasks to executor...") + future_to_item = { + executor.submit(_generate_single_kernel_individual, item): item + for item in work_items + } + print(" All tasks submitted, waiting for completion...") + + # Collect results with progress reporting + for future in concurrent.futures.as_completed(future_to_item): + completed += 1 + if completed % 100 == 0 or completed == len(work_items): + print( + f" Progress: {completed}/{len(work_items)} kernels generated" + ) + try: + result = future.result() + if result: + kernel_list.append(result) + except Exception as exc: + item = future_to_item[future] + print(f"Kernel generation failed for {item}: {exc}") + + # Sort kernel list for consistent ordering + kernel_list.sort(key=lambda x: x[0]) # Sort by kernel name + + # Generate CMake include file for individual targets + self._generate_cmake_individual_targets(kernel_list) + + print( + f"Generated {len(kernel_list)} individual kernel files in {self.working_path}" + ) + + +def _generate_single_kernel_individual(work_item): + """Worker function to generate a single individual kernel file""" + ( + tile_config, + trait_combo, + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json, + ) = work_item + + # Create a temporary builder instance for this worker + builder = GemmAQuantKernelBuilder( + kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json + ) + + try: + kernel_name, instance_code = builder._generate_kernel_instance( + tile_config, trait_combo + ) + + # Strip the "gemm_aquant_" prefix so the header filename does not duplicate it + simplified_name = kernel_name.removeprefix(kernel_name_prefix + "_") + + # Write individual header file + header_file = working_path / f"gemm_aquant_single_{simplified_name}.hpp" + with open(header_file, "w") as f: + f.write(instance_code) + + return (kernel_name, trait_combo, tile_config) + except Exception as e: + print(f"Error generating individual kernel: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="GEMM AQuant kernel instance builder") + parser.add_argument("--working_path", required=True, help="Working directory path") + parser.add_argument("--gpu_target", required=True, help="GPU target architecture") + parser.add_argument( + "--datatype", + required=True, + choices=["fp8", "bf8"], + help="Data type (fp8 or bf8)", + ) + parser.add_argument( + "--layout", + required=True, + choices=["rcr", "rrr", "ccr", "crr"], + help="Matrix layout", + ) + parser.add_argument("--config_json", help="Configuration JSON file") + parser.add_argument("--num_workers", type=int, help="Number of parallel workers") + parser.add_argument( + "--gen_all_individual", + action="store_true", + help="Generate individual kernel files", + ) + parser.add_argument( + "--gen_single", + action="store_true", + help="Generate a single kernel file", + ) + parser.add_argument("--kernel_name", help="Kernel name for single generation") + parser.add_argument("--tile_config", help="Tile configuration string") + parser.add_argument("--trait_combo", help="Trait combination string") + parser.add_argument( + "--list_kernels", + action="store_true", + help="List kernel configurations without generating files", + ) + parser.add_argument( + "--max-instances", + type=int, + default=None, + help="Maximum number of kernel instances to select via sampling", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="RNG seed for deterministic sampling; if omitted, derived from today's date", + ) + parser.add_argument( + "--tier", + default=None, + help="Sampling tier name (e.g., 'daily')", + ) + parser.add_argument( + "--manifest-path", + default=None, + help="Directory to write chosen_instances manifest JSON", + ) + + args = parser.parse_args() + + layout_parts = args.layout.lower() + assert len(layout_parts) == 3, ( + f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr' where r stands for row major and c stands for column major)" + ) + assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], ( + f"Invalid matrix_a layout : {layout_parts[0]} or matrix_b layout: {layout_parts[1]} (matrix_a and matrix_b must be either 'r' for row major or 'c' for column major)" + ) + assert layout_parts[2] == "r", ( + f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' only as currently we are supporting only row major)" + ) + + kernel_name_prefix = "gemm_aquant" + builder = GemmAQuantKernelBuilder( + kernel_name_prefix, + args.working_path, + args.gpu_target, + args.datatype, + args.layout, + args.config_json, + max_instances=args.max_instances, + seed=args.seed, + tier=args.tier, + manifest_path=args.manifest_path, + ) + + if args.list_kernels: + builder._list_kernels() + elif args.gen_single: + if not args.kernel_name or not args.tile_config or not args.trait_combo: + parser.error( + "--gen_single requires --kernel_name, --tile_config, and --trait_combo" + ) + + # Parse tile config + tile_parts = args.tile_config.split("_") + tile_dims = tile_parts[0].split("x") + warp_dims = tile_parts[1].split("x") + warp_tile_dims = tile_parts[2].split("x") + + tile_config = { + "tile_m": int(tile_dims[0]), + "tile_n": int(tile_dims[1]), + "tile_k": int(tile_dims[2]), + "warp_m": int(warp_dims[0]), + "warp_n": int(warp_dims[1]), + "warp_k": int(warp_dims[2]), + "warp_tile_m": int(warp_tile_dims[0]), + "warp_tile_n": int(warp_tile_dims[1]), + "warp_tile_k": int(warp_tile_dims[2]), + } + + # Parse trait combo: + # pipeline_epilogue_scheduler_padM_padN_padK_aPreshuffle + trait_parts = args.trait_combo.split("_") + trait_combo = ( + trait_parts[0], # pipeline + trait_parts[1], # epilogue + trait_parts[2], # scheduler + trait_parts[3] == "True", # pad_m + trait_parts[4] == "True", # pad_n + trait_parts[5] == "True", # pad_k + trait_parts[6] == "True", # a_preshuffle_quant + ) + + builder._generate_kernel_instance(tile_config, trait_combo) + elif args.gen_all_individual: + builder._generate_all_individual(args.num_workers) + else: + parser.error( + "Must specify one of: --list_kernels, --gen_all_individual, or --gen_single" + ) + + +if __name__ == "__main__": + main() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp new file mode 100644 index 0000000000..00a19b0b44 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_profiler.hpp @@ -0,0 +1,276 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "gemm_aquant_benchmark.hpp" + +class AQuantGemmProfiler +{ + public: + explicit AQuantGemmProfiler(Setting setting) : setting_(setting) {} + + AQuantGemmProfiler(const AQuantGemmProfiler&) = delete; + AQuantGemmProfiler& operator=(const AQuantGemmProfiler&) = delete; + + void reset() { kernel_instances_.clear(); } + + void benchmark(AQuantGemmProblem& problem, + std::function kernel_func) + { + std::vector(ck_tile::QuantGemmHostArgs&, + const ck_tile::stream_config&)>> + callables; + + callables.push_back( + [kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) { + float time = kernel_func(args, stream); + return std::make_tuple(std::string(KERNEL_NAME), time); + }); + + benchmark(problem, callables); + } + + void benchmark(AQuantGemmProblem& problem, + std::vector( + ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables) + { + const ALayout layout_a = ALayout{}; + const BLayout layout_b = BLayout{}; + const CLayout layout_c = CLayout{}; + const AQLayout layout_aq = AQLayout{}; + + problem.stride_a_ = ck_tile::get_default_stride( + problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)); + problem.stride_b_ = ck_tile::get_default_stride( + problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)); + problem.stride_c_ = ck_tile::get_default_stride( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)); + + // Compute AQ scale tensor dimensions: [M, K / group_size_k] + const ck_tile::index_t QK_A = problem.k_ / problem.group_size_k_; + problem.stride_aq_ = ck_tile::get_default_stride( + problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq)); + + // Allocate host tensors + ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( + problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a))); + ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( + problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b))); + ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c))); + + // AQ scale tensor: [M, QK_A] + ck_tile::HostTensor aq_m_qk(ck_tile::host_tensor_descriptor( + problem.m_, QK_A, problem.stride_aq_, is_row_major(layout_aq))); + + // Initialize tensors + if(setting_.init_method_ == 0) + { + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); + ck_tile::FillUniformDistribution{0.5f, 1.5f}(aq_m_qk); + } + else if(setting_.init_method_ == 1) + { + ck_tile::FillMonotonicSeq{}(a_m_k); + ck_tile::FillMonotonicSeq{}(b_k_n); + ck_tile::FillConstant{static_cast(1)}(aq_m_qk); + } + else if(setting_.init_method_ == 2) + { + ck_tile::FillConstant{static_cast(1)}(a_m_k); + ck_tile::FillConstant{static_cast(1)}(b_k_n); + ck_tile::FillConstant{static_cast(1)}(aq_m_qk); + } + else + { + a_m_k.SetZero(); + b_k_n.SetZero(); + aq_m_qk.SetZero(); + } + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); + ck_tile::DeviceMem aq_dev_buf(aq_m_qk.get_element_space_size_in_bytes()); + + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + aq_dev_buf.ToDevice(aq_m_qk.data()); + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + + // Build QuantGemmHostArgs + ck_tile::QuantGemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(), + b_k_n_dev_buf.GetDeviceBuffer(), + c_m_n_dev_buf.GetDeviceBuffer(), + aq_dev_buf.GetDeviceBuffer(), + nullptr, // bq_ptr not used for AQuant + problem.split_k_, + problem.m_, + problem.n_, + problem.k_, + QK_A, + 0, // QK_B not used for AQuant + problem.stride_a_, + problem.stride_b_, + problem.stride_c_, + problem.stride_aq_, + 0 // stride_BQ not used for AQuant + ); + + // Host reference for verification + ck_tile::HostTensor c_m_n_host_result(ck_tile::host_tensor_descriptor( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c))); + + if(setting_.verify_) + { + aquant_gemm_host_reference(setting_.verify_, a_m_k, aq_m_qk, b_k_n, c_m_n_host_result); + } + + // Run kernel(s) + for(auto& callable : callables) + { + auto kernel_run_result = callable(gemm_args, + ck_tile::stream_config{nullptr, + true, + setting_.log_, + setting_.n_warmup_, + setting_.n_repeat_, + setting_.is_gpu_timer_, + setting_.flush_cache_, + setting_.rotating_count_}); + process_result(problem, + QK_A, + c_m_n_dev_buf, + c_m_n_host_result, + c_m_n_dev_result, + kernel_run_result); + } + } + + void process_result(const AQuantGemmProblem& problem, + ck_tile::index_t QK_A, + ck_tile::DeviceMem& c_m_n_dev_buf, + ck_tile::HostTensor& c_m_n_host_result, + ck_tile::HostTensor& c_m_n_dev_result, + const std::tuple& kernel_run_result) + { + auto [name, avg_time] = kernel_run_result; + + KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}}; + + // Compute performance metrics + std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_; + std::size_t num_byte = sizeof(ADataType) * problem.m_ * problem.k_ + + sizeof(BDataType) * problem.n_ * problem.k_ + + sizeof(AQDataType) * problem.m_ * QK_A + + sizeof(CDataType) * problem.m_ * problem.n_; + + kernel_instance.perf_result_.latency_ = avg_time; + kernel_instance.perf_result_.tflops_ = static_cast(flop) / 1.E9 / avg_time; + kernel_instance.perf_result_.bandwidth_ = num_byte / 1.E6 / avg_time; + + if(setting_.log_ > 0 && !setting_.json_output_) + { + std::cout << kernel_instance << std::endl; + } + + // Verify result + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); + bool verified_correct = + !setting_.verify_ || + compare_aquant(name, problem.k_, problem.split_k_, c_m_n_dev_result, c_m_n_host_result); + + if(verified_correct) + { + kernel_instances_.emplace_back(kernel_instance); + } + else + { + std::cout << "Verification failed, skip kernel: " << name << std::endl; + } + + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + } + + KernelInstance select_best_instance(Metric metric) + { + if(kernel_instances_.empty()) + throw std::runtime_error("Empty instances"); + + auto kernel_instance = *std::max_element(kernel_instances_.begin(), + kernel_instances_.end(), + [metric](const auto& a, const auto& b) { + return PerformanceResult::compare( + b.perf_result_, a.perf_result_, metric); + }); + + if(setting_.json_output_) + { + std::cout << kernel_instance << std::endl; + } + else + { + std::cout << "**********************************" << std::endl; + std::cout << "According to given metrics: " << get_metric_name(metric) << "\n" + << "Current kernel performance is: " << kernel_instance << std::endl; + std::cout << "**********************************" << std::endl; + } + + if(!setting_.csv_filename_.empty()) + { + std::ofstream file(setting_.csv_filename_ + ".csv", std::ios::app); + + if(!file.is_open()) + { + std::cerr << "Warning: Failed to open CSV file for writing." << std::endl; + } + else + { + if(file.tellp() == 0) + { + file << "rocm_version,device_name," + << "split_k,m,n,k,stride_a,stride_b,stride_c,stride_aq,group_size_k," + << "dtype_a,dtype_b,dtype_aq,dtype_acc,dtype_c," + << "layout_a,layout_b,layout_c," << "name," + << "latency(ms),tflops(TFlops),bandwidth(GB/s),metric\n"; + } + + const auto& prob = kernel_instance.problem_; + const auto& perf = kernel_instance.perf_result_; + + file << get_rocm_version() << "," << ck_tile::get_device_name() << "," + << prob.split_k_ << "," << prob.m_ << "," << prob.n_ << "," << prob.k_ << "," + << prob.stride_a_ << "," << prob.stride_b_ << "," << prob.stride_c_ << "," + << prob.stride_aq_ << "," << prob.group_size_k_ << "," << prob.dtype_a_ << "," + << prob.dtype_b_ << "," << prob.dtype_aq_ << "," << prob.dtype_acc_ << "," + << prob.dtype_c_ << "," << prob.layout_a_ << "," << prob.layout_b_ << "," + << prob.layout_c_ << "," << kernel_instance.name_ << "," << std::fixed + << std::setprecision(4) << perf.latency_ << "," << perf.tflops_ << "," + << perf.bandwidth_ << "," << get_metric_name(metric) << "\n"; + } + } + + return kernel_instance; + } + + private: + Setting setting_; + std::vector kernel_instances_; +}; diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py new file mode 100644 index 0000000000..4cbfc024c8 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/test_gemm_aquant_instance_builder.py @@ -0,0 +1,204 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Unit tests for GemmAQuantKernelBuilder (gemm_aquant operator).""" + +import json +import os +import sys +import tempfile +import unittest + +# Make the instance builder importable by inserting the directory into sys.path +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) + +from gemm_aquant_instance_builder import GemmAQuantKernelBuilder # noqa: E402 + +# --------------------------------------------------------------------------- +# Minimal default config that mirrors default_config.json, kept small so that +# _get_tile_configs() / _generate_trait_combinations() run quickly in tests. +# --------------------------------------------------------------------------- +_MINIMAL_CONFIG = { + "tile_config": { + "tile_m": {"values": [64]}, + "tile_n": {"values": [64]}, + "tile_k": {"values": [128]}, + "warp_m": {"values": [2]}, + "warp_n": {"values": [2]}, + "warp_k": {"values": [1]}, + "warp_tile_m": {"values": [16]}, + "warp_tile_n": {"values": [16]}, + "warp_tile_k": {"values": [64]}, + }, + "trait_config": { + "pipeline": {"values": ["compv3"]}, + "scheduler": {"values": ["intrawave"]}, + "epilogue": {"values": ["default"]}, + "pad_m": {"values": [False]}, + "pad_n": {"values": [False]}, + "pad_k": {"values": [False]}, + "a_preshuffle_quant": {"values": [False, True]}, + }, + "k_block_per_cu": 1, + "group_size_k": 128, +} + + +def _make_builder(tmpdir, config=None, **kwargs): + """Helper: write config to a temp JSON file and create a builder.""" + cfg = config if config is not None else _MINIMAL_CONFIG + cfg_path = os.path.join(tmpdir, "config.json") + with open(cfg_path, "w") as f: + json.dump(cfg, f) + return GemmAQuantKernelBuilder( + kernel_name_prefix="gemm_aquant", + working_path=tmpdir, + gpu_target=kwargs.get("gpu_target", "gfx942"), + datatype=kwargs.get("datatype", "fp8"), + layout=kwargs.get("layout", "rcr"), + config_json=cfg_path, + ) + + +class TestGemmAQuantBuilderInit(unittest.TestCase): + def test_group_size_k_from_config(self): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + self.assertEqual(builder.group_size_k, 128) + + def test_group_size_k_custom(self): + cfg = json.loads(json.dumps(_MINIMAL_CONFIG)) + cfg["group_size_k"] = 64 + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir, config=cfg) + self.assertEqual(builder.group_size_k, 64) + + def test_kernel_name_prefix_stored(self): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir) + self.assertEqual(builder.kernel_name_prefix, "gemm_aquant") + + def test_working_path_created(self): + with tempfile.TemporaryDirectory() as tmpdir: + sub = os.path.join(tmpdir, "workdir") + cfg_path = os.path.join(tmpdir, "config.json") + with open(cfg_path, "w") as f: + json.dump(_MINIMAL_CONFIG, f) + GemmAQuantKernelBuilder( + "gemm_aquant", sub, "gfx942", "fp8", "rcr", cfg_path + ) + self.assertTrue(os.path.isdir(sub)) + + +class TestGemmAQuantTraitCombinations(unittest.TestCase): + def setUp(self): + self._tmpdir = tempfile.TemporaryDirectory() + self.builder = _make_builder(self._tmpdir.name) + + def tearDown(self): + self._tmpdir.cleanup() + + def test_trait_combinations_non_empty(self): + combos = self.builder._generate_trait_combinations() + self.assertGreater(len(combos), 0) + + def test_trait_combo_is_7_tuple(self): + """AQuant trait tuple: (pipeline, epilogue, scheduler, pad_m, pad_n, pad_k, a_preshuffle_quant).""" + combos = self.builder._generate_trait_combinations() + for combo in combos: + self.assertEqual(len(combo), 7, f"Expected 7-tuple, got {len(combo)}: {combo}") + + def test_pipeline_values_present(self): + combos = self.builder._generate_trait_combinations() + pipelines = {c[0] for c in combos} + self.assertIn("compv3", pipelines) + + def test_preshuffle_quant_values(self): + """Both True and False should appear when config lists both.""" + combos = self.builder._generate_trait_combinations() + preshuffle_vals = {c[6] for c in combos} + self.assertIn(True, preshuffle_vals) + self.assertIn(False, preshuffle_vals) + + def test_uses_persistent_trait_returns_true(self): + """AQuant overrides _uses_persistent_trait to accommodate a_preshuffle_quant slot.""" + self.assertTrue(self.builder._uses_persistent_trait()) + + +class TestGemmAQuantTileConfigs(unittest.TestCase): + def setUp(self): + self._tmpdir = tempfile.TemporaryDirectory() + self.builder = _make_builder(self._tmpdir.name) + + def tearDown(self): + self._tmpdir.cleanup() + + def test_tile_configs_non_empty(self): + configs = self.builder._get_tile_configs() + self.assertGreater(len(configs), 0) + + def test_tile_config_has_required_keys(self): + configs = self.builder._get_tile_configs() + required = { + "tile_m", "tile_n", "tile_k", + "warp_m", "warp_n", "warp_k", + "warp_tile_m", "warp_tile_n", "warp_tile_k", + } + for cfg in configs: + self.assertTrue(required.issubset(cfg.keys()), f"Missing keys in {cfg}") + + def test_tile_config_values_are_positive(self): + configs = self.builder._get_tile_configs() + for cfg in configs: + for key, val in cfg.items(): + self.assertGreater(val, 0, f"{key}={val} must be positive") + + +class TestGemmAQuantLayoutVariants(unittest.TestCase): + """Verify the builder initialises correctly for all supported layouts.""" + + LAYOUTS = ["rcr", "rrr", "ccr", "crr"] + + def test_all_layouts_construct(self): + for layout in self.LAYOUTS: + with self.subTest(layout=layout): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir, layout=layout) + self.assertEqual(builder.layout, layout) + + +class TestGemmAQuantDataTypes(unittest.TestCase): + DTYPES = ["fp8", "bf8"] + + def test_all_dtypes_construct(self): + for dtype in self.DTYPES: + with self.subTest(dtype=dtype): + with tempfile.TemporaryDirectory() as tmpdir: + builder = _make_builder(tmpdir, datatype=dtype) + self.assertEqual(builder.datatype, dtype) + + +class TestGemmAQuantMaxInstances(unittest.TestCase): + def test_max_instances_none_returns_all(self): + """Without a max_instances cap, sampling is a no-op.""" + with tempfile.TemporaryDirectory() as tmpdir: + cfg_path = os.path.join(tmpdir, "config.json") + with open(cfg_path, "w") as f: + json.dump(_MINIMAL_CONFIG, f) + builder = GemmAQuantKernelBuilder( + "gemm_aquant", tmpdir, "gfx942", "fp8", "rcr", + config_json=cfg_path, max_instances=None, + ) + kernel_list = [ + {"tile_config": {"tile_m": 64, "tile_n": 64, "tile_k": 128, + "warp_m": 2, "warp_n": 2, "warp_k": 1, + "warp_tile_m": 16, "warp_tile_n": 16, "warp_tile_k": 64}, + "trait_combo": ("compv3", "default", "intrawave", False, False, False, False)} + ] + result = builder._apply_sampling(kernel_list) + self.assertEqual(len(result), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt new file mode 100644 index 0000000000..d0c0785df6 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/CMakeLists.txt @@ -0,0 +1,289 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +set(GEMM_BQUANT_DATATYPE "fp8;bf8" CACHE STRING "List of datatypes for GEMM BQuant (semicolon-separated)") +set(GEMM_BQUANT_LAYOUT "rcr;rrr;ccr;crr" CACHE STRING "List of layout for GEMM BQuant (semicolon-separated)") +set(GEMM_BQUANT_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)") +option(ENABLE_CCACHE_GEMM_BQUANT "Enable ccache for GEMM BQuant ops compilation" OFF) + +# Store the directory path for use in functions +set(GEMM_BQUANT_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR}) + +# Function to create individual GEMM BQuant targets +function(create_individual_gemm_bquant_target datatype layout trait tile_config config_json) + # GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL is guaranteed non-empty here (caller checks at line 192) + if(NOT GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL) + message(FATAL_ERROR "BUG: create_individual_gemm_bquant_target called with no GPU targets") + endif() + + # Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k + string(REPLACE "_" ";" config_groups ${tile_config}) + list(GET config_groups 0 tile_dims) + list(GET config_groups 1 warp_dims) + list(GET config_groups 2 warp_tile_dims) + + string(REPLACE "x" ";" tile_parts ${tile_dims}) + list(GET tile_parts 0 tile_m) + list(GET tile_parts 1 tile_n) + list(GET tile_parts 2 tile_k) + + string(REPLACE "x" ";" warp_parts ${warp_dims}) + list(GET warp_parts 0 warp_m) + list(GET warp_parts 1 warp_n) + list(GET warp_parts 2 warp_k) + + string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims}) + list(GET warp_tile_parts 0 warp_tile_m) + list(GET warp_tile_parts 1 warp_tile_n) + list(GET warp_tile_parts 2 warp_tile_k) + + set(target_name "benchmark_gemm_bquant_${datatype}_${layout}_${trait}_${tile_config}") + set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") + + # Generate the single instance header for this kernel + set(instance_header "${working_path}/gemm_bquant_single_${datatype}_${layout}_${trait}_${tile_config}.hpp") + + # Add custom command to generate the header file at build time + add_custom_command( + OUTPUT ${instance_header} + COMMAND ${Python3_EXECUTABLE} ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_instance_builder.py + --working_path ${working_path} + --datatype ${datatype} + --layout ${layout} + --config_json ${config_json} + --gen_single + --kernel_name "gemm_bquant_${datatype}_${layout}_${trait}_${tile_config}" + --tile_config "${tile_config}" + --trait_combo "${trait}" + --gpu_target "${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}" + DEPENDS ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_instance_builder.py ${config_json} + COMMENT "Generating ${instance_header}" + ) + + # Create the executable + add_executable(${target_name} + EXCLUDE_FROM_ALL + ${GEMM_BQUANT_SOURCE_DIR}/gemm_bquant_benchmark_single.cpp + ${instance_header} + ) + + # Set GPU architectures + set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}) + + # Set compile definitions + target_compile_definitions(${target_name} PRIVATE + GEMM_BQUANT_SINGLE_INSTANCE_HPP="${instance_header}" + ) + + # Include directories + target_include_directories(${target_name} PRIVATE + ${GEMM_BQUANT_SOURCE_DIR} + ${working_path} + ) + + # Compile options + target_compile_options(${target_name} PRIVATE + -Wno-undefined-func-template + -Wno-float-equal + --offload-compress + -include ${instance_header} + ) + + # Add to collection targets + add_dependencies(benchmark_gemm_bquant_all ${target_name}) + add_dependencies(benchmark_gemm_bquant_${datatype} ${target_name}) + add_dependencies(benchmark_gemm_bquant_${layout} ${target_name}) + add_dependencies(benchmark_gemm_bquant_${datatype}_${layout} ${target_name}) + + # Add to pipeline-specific, epilogue-specific, and scheduler-specific targets + string(REPLACE "_" ";" trait_parts ${trait}) + list(GET trait_parts 0 pipeline) + list(GET trait_parts 1 epilogue) + list(GET trait_parts 2 scheduler) + + add_dependencies(benchmark_gemm_bquant_${pipeline}_pipeline ${target_name}) + add_dependencies(benchmark_gemm_bquant_${epilogue}_epilogue ${target_name}) + add_dependencies(benchmark_gemm_bquant_${scheduler}_scheduler ${target_name}) +endfunction() + +# Function to build individual GEMM BQuant targets +function(build_individual_gemm_bquant_targets datatype layout) + set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") + + # Choose config file (priority: env var > cmake var > default) + if(DEFINED ENV{GEMM_BQUANT_CONFIG_FILE} AND NOT "$ENV{GEMM_BQUANT_CONFIG_FILE}" STREQUAL "") + set(config_filename "$ENV{GEMM_BQUANT_CONFIG_FILE}") + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}") + message(VERBOSE " Using config from environment variable: ${config_filename}") + elseif(NOT "${GEMM_BQUANT_CONFIG_FILE}" STREQUAL "") + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_BQUANT_CONFIG_FILE}") + message(VERBOSE " Using custom config: ${GEMM_BQUANT_CONFIG_FILE}") + else() + set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json") + message(VERBOSE " Using default config for layout ${layout}") + endif() + + if(NOT EXISTS ${json_blob}) + message(FATAL_ERROR "Config file not found: ${json_blob}") + endif() + + # Create working directory + file(MAKE_DIRECTORY ${working_path}) + + # Build sampling arguments + set(extra_list_args "") + if(NOT "${GEMM_BQUANT_MAX_INSTANCES}" STREQUAL "") + list(APPEND extra_list_args --max-instances ${GEMM_BQUANT_MAX_INSTANCES}) + endif() + if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") + list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER}) + list(APPEND extra_list_args --manifest-path ${working_path}) + endif() + if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "") + list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED}) + endif() + + # List kernels + message(VERBOSE " Listing BQuant kernel configurations for ${datatype} ${layout}...") + execute_process( + COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_bquant_instance_builder.py + --working_path ${working_path} + --datatype ${datatype} + --layout ${layout} + --config_json ${json_blob} + --gpu_target "${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}" + --list_kernels + ${extra_list_args} + WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR} + RESULT_VARIABLE ret + OUTPUT_VARIABLE list_output + ERROR_VARIABLE list_error + ) + + if(NOT ret EQUAL 0) + message(FATAL_ERROR "Failed to list BQuant kernels for ${datatype} ${layout}: ${list_error}") + endif() + + # Read kernel count + if(EXISTS ${working_path}/gemm_bquant_kernel_count.txt) + file(READ ${working_path}/gemm_bquant_kernel_count.txt kernel_count) + string(STRIP "${kernel_count}" kernel_count) + message(VERBOSE " Found ${kernel_count} BQuant kernel configurations") + else() + message(FATAL_ERROR "Kernel count file not found") + endif() + + # Read kernel list and create targets + if(EXISTS ${working_path}/gemm_bquant_kernel_list.txt) + file(STRINGS ${working_path}/gemm_bquant_kernel_list.txt kernel_lines) + foreach(line IN LISTS kernel_lines) + string(REPLACE "|" ";" parts "${line}") + list(GET parts 0 kernel_name) + list(GET parts 1 tile_config) + list(GET parts 2 trait_combo) + + create_individual_gemm_bquant_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}") + endforeach() + else() + message(FATAL_ERROR "Kernel list file not found") + endif() +endfunction() + +# Main build logic +message(VERBOSE "=== Starting Tile Engine GEMM BQuant Configuration ===") +message(VERBOSE "GEMM_BQUANT_DATATYPE: ${GEMM_BQUANT_DATATYPE}") +message(VERBOSE "GEMM_BQUANT_LAYOUT: ${GEMM_BQUANT_LAYOUT}") +message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") + +# Filter GPU targets to supported architectures +set(GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL "") +set(DESIRED_TARGETS "gfx90a;gfx942;gfx950") + +foreach(target IN LISTS SUPPORTED_GPU_TARGETS) + if(target IN_LIST DESIRED_TARGETS) + list(APPEND GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL ${target}) + message(VERBOSE " Adding GPU target: ${target}") + endif() +endforeach() + +# Skip build if no matching targets found +if(NOT GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL) + message(WARNING "Skipping Tile Engine GEMM BQuant build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") +else() + message(VERBOSE "Building individual GEMM BQuant targets for GPU targets: ${GEMM_BQUANT_GPU_TARGETS_INDIVIDUAL}") + + # Enable compiler cache if requested + if(ENABLE_CCACHE_GEMM_BQUANT) + find_program(CCACHE_PROGRAM ccache) + if(CCACHE_PROGRAM) + set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM}) + message(VERBOSE "Using ccache for faster compilation") + endif() + endif() + + # Create master collection target + add_custom_target(benchmark_gemm_bquant_all) + + # Create datatype collection targets + foreach(dt IN LISTS GEMM_BQUANT_DATATYPE) + add_custom_target(benchmark_gemm_bquant_${dt}) + endforeach() + + # Create layout collection targets + foreach(l IN LISTS GEMM_BQUANT_LAYOUT) + add_custom_target(benchmark_gemm_bquant_${l}) + endforeach() + + # Create combined collection targets + foreach(dt IN LISTS GEMM_BQUANT_DATATYPE) + foreach(l IN LISTS GEMM_BQUANT_LAYOUT) + add_custom_target(benchmark_gemm_bquant_${dt}_${l}) + endforeach() + endforeach() + + # Create pipeline-specific collection targets + set(GEMM_BQUANT_PIPELINES "compv3") + foreach(pipeline IN LISTS GEMM_BQUANT_PIPELINES) + add_custom_target(benchmark_gemm_bquant_${pipeline}_pipeline) + endforeach() + + # Create epilogue-specific collection targets + set(GEMM_BQUANT_EPILOGUES "default;cshuffle") + foreach(epilogue IN LISTS GEMM_BQUANT_EPILOGUES) + add_custom_target(benchmark_gemm_bquant_${epilogue}_epilogue) + endforeach() + + # Create scheduler-specific collection targets + set(GEMM_BQUANT_SCHEDULERS "intrawave") + foreach(scheduler IN LISTS GEMM_BQUANT_SCHEDULERS) + add_custom_target(benchmark_gemm_bquant_${scheduler}_scheduler) + endforeach() + + # Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that + # sampling fires per-combo rather than being a single cap. + # CMake integer division truncates; clamp the result to at least 1 so that a + # small total budget (e.g. 5 instances across 8 combos) does not silently + # produce a per-combo budget of 0 and skip all kernels. + if(NOT "${GEMM_BQUANT_MAX_INSTANCES}" STREQUAL "") + list(LENGTH GEMM_BQUANT_DATATYPE _bq_n_dt) + list(LENGTH GEMM_BQUANT_LAYOUT _bq_n_lay) + math(EXPR _bq_n_combos "${_bq_n_dt} * ${_bq_n_lay}") + if(_bq_n_combos GREATER 0) + math(EXPR GEMM_BQUANT_MAX_INSTANCES + "${GEMM_BQUANT_MAX_INSTANCES} / ${_bq_n_combos}") + if(GEMM_BQUANT_MAX_INSTANCES EQUAL 0) + set(GEMM_BQUANT_MAX_INSTANCES 1) + message(WARNING " gemm_bquant: per-combo budget rounded to 0; clamped to 1 (total budget too small for ${_bq_n_combos} combos)") + else() + message(STATUS " gemm_bquant: per-combo budget = ${GEMM_BQUANT_MAX_INSTANCES} (${_bq_n_combos} combos)") + endif() + endif() + endif() + + # Build individual targets for each datatype/layout combination + foreach(dt IN LISTS GEMM_BQUANT_DATATYPE) + foreach(l IN LISTS GEMM_BQUANT_LAYOUT) + build_individual_gemm_bquant_targets(${dt} ${l}) + endforeach() + endforeach() +endif() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json new file mode 100644 index 0000000000..79ecb47a1f --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_ci_config.json @@ -0,0 +1,93 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 16 + ] + }, + "tile_n": { + "values": [ + 64 + ] + }, + "tile_k": { + "values": [ + 256 + ] + }, + "warp_m": { + "values": [ + 1 + ] + }, + "warp_n": { + "values": [ + 4 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16 + ] + }, + "warp_tile_n": { + "values": [ + 16 + ] + }, + "warp_tile_k": { + "values": [ + 16, + 32, + 64, + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "default", + "cshuffle" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "b_preshuffle_quant": { + "values": [ + false, + true + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128 +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json new file mode 100644 index 0000000000..9ed8302878 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/default_config.json @@ -0,0 +1,102 @@ +{ + "tile_config": { + "tile_m": { + "max": 256, + "min": 64, + "step": 64 + }, + "tile_n": { + "max": 256, + "min": 64, + "step": 64 + }, + "tile_k": { + "max": 256, + "min": 64, + "step": 64 + }, + "warp_m": { + "values": [ + 4, + 2, + 1 + ] + }, + "warp_n": { + "values": [ + 4, + 2, + 1 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 4, + 16, + 32 + ] + }, + "warp_tile_n": { + "values": [ + 16, + 32, + 64 + ] + }, + "warp_tile_k": { + "values": [ + 8, + 16, + 32, + 64, + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "cshuffle", + "default" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "b_preshuffle_quant": { + "values": [ + false, + true + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128 +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json new file mode 100644 index 0000000000..c2ea50a65c --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/configs/example_config.json @@ -0,0 +1,88 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 16 + ] + }, + "tile_n": { + "values": [ + 64 + ] + }, + "tile_k": { + "values": [ + 256 + ] + }, + "warp_m": { + "values": [ + 1 + ] + }, + "warp_n": { + "values": [ + 4 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 16 + ] + }, + "warp_tile_n": { + "values": [ + 16 + ] + }, + "warp_tile_k": { + "values": [ + 128 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "default" + ] + }, + "pad_m": { + "values": [ + false + ] + }, + "pad_n": { + "values": [ + false + ] + }, + "pad_k": { + "values": [ + false + ] + }, + "b_preshuffle_quant": { + "values": [ + false + ] + } + }, + "k_block_per_cu": 1, + "group_size_k": 128 +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp new file mode 100644 index 0000000000..d50ce46755 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.hpp @@ -0,0 +1,229 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "gemm_bquant_common.hpp" + +// Data types and Layouts are defined by the generated kernel headers: +// ADataType, BDataType, BQDataType, AccDataType, CDataType +// ALayout, BLayout, CLayout, BQLayout + +enum class Metric +{ + LATENCY = 0, + TFLOPS = 1, + BANDWIDTH = 2 +}; + +inline constexpr auto get_metric_name(Metric m) +{ + switch(m) + { + case Metric::LATENCY: return "latency"; + case Metric::TFLOPS: return "tflops"; + case Metric::BANDWIDTH: return "bandwidth"; + default: throw std::invalid_argument("Unsupported metric type"); + } +} + +struct BQuantGemmProblem +{ + int split_k_; + int m_, n_, k_; + int stride_a_, stride_b_, stride_c_; + int stride_bq_; + int group_size_k_; + + std::string dtype_a_, dtype_b_, dtype_bq_, dtype_acc_, dtype_c_; + std::string layout_a_, layout_b_, layout_c_; + + friend std::ostream& operator<<(std::ostream& os, const BQuantGemmProblem& problem) + { + os << "{\n" + << " \"split_k\":" << problem.split_k_ << ",\n" + << " \"m\":" << problem.m_ << ",\n" + << " \"n\":" << problem.n_ << ",\n" + << " \"k\":" << problem.k_ << ",\n" + << " \"stride_a\":" << problem.stride_a_ << ",\n" + << " \"stride_b\":" << problem.stride_b_ << ",\n" + << " \"stride_c\":" << problem.stride_c_ << ",\n" + << " \"stride_bq\":" << problem.stride_bq_ << ",\n" + << " \"group_size_k\":" << problem.group_size_k_ << ",\n" + << " \"dtype_a\":\"" << problem.dtype_a_ << "\",\n" + << " \"dtype_b\":\"" << problem.dtype_b_ << "\",\n" + << " \"dtype_bq\":\"" << problem.dtype_bq_ << "\",\n" + << " \"dtype_acc\":\"" << problem.dtype_acc_ << "\",\n" + << " \"dtype_c\":\"" << problem.dtype_c_ << "\",\n" + << " \"layout_a\":\"" << problem.layout_a_ << "\",\n" + << " \"layout_b\":\"" << problem.layout_b_ << "\",\n" + << " \"layout_c\":\"" << problem.layout_c_ << "\"\n" + << "}"; + return os; + } +}; + +struct PerformanceResult +{ + double latency_; + double tflops_; + double bandwidth_; + + static bool compare(const PerformanceResult& a, const PerformanceResult& b, Metric m) + { + switch(m) + { + case Metric::LATENCY: return a.latency_ < b.latency_; + case Metric::TFLOPS: return a.tflops_ > b.tflops_; + case Metric::BANDWIDTH: return a.bandwidth_ > b.bandwidth_; + default: throw std::invalid_argument("Unsupported metric type"); + } + } + + friend std::ostream& operator<<(std::ostream& os, const PerformanceResult& result) + { + os << "{\n" + << " \"latency(ms)\": " << std::fixed << std::setprecision(2) << result.latency_ + << ",\n" + << " \"tflops(TFlops)\": " << result.tflops_ << ",\n" + << " \"bandwidth(GB/s)\": " << result.bandwidth_ << "\n" + << "}"; + return os; + } +}; + +struct KernelInstance +{ + std::string name_; + BQuantGemmProblem problem_; + PerformanceResult perf_result_; + + static bool compare(const KernelInstance& a, const KernelInstance& b, Metric m) + { + return PerformanceResult::compare(a.perf_result_, b.perf_result_, m); + } + + friend std::ostream& operator<<(std::ostream& os, const KernelInstance& obj) + { + os << "{\n" + << " \"name\": \"" << obj.name_ << "\",\n" + << " \"problem\": " << obj.problem_ << ",\n" + << " \"perf_result\": " << obj.perf_result_ << "\n" + << "}"; + return os; + } +}; + +struct Setting +{ + int n_warmup_; + int n_repeat_; + bool is_gpu_timer_; + int verify_; + int init_method_; + bool log_; + std::string csv_filename_; + bool flush_cache_; + int rotating_count_; + bool json_output_; +}; + +inline std::string get_rocm_version() +{ + std::ifstream version_file("/opt/rocm/.info/version"); + if(version_file.is_open()) + { + std::string version; + std::getline(version_file, version); + return version; + } + return "Unknown"; +} + +template +auto calculate_rtol_atol_bquant(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeType = + std::conditional_t; + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +/// @brief Compare device and host results for BQuant GEMM +bool compare_bquant(std::string instanceName, + ck_tile::index_t K, + ck_tile::index_t kbatch, + ck_tile::HostTensor& c_m_n_dev_result, + ck_tile::HostTensor& c_m_n_host_result) +{ + const float max_accumulated_value = + std::abs(static_cast(*std::max_element(c_m_n_host_result.mData.begin(), + c_m_n_host_result.mData.end(), + [](const auto& a, const auto& b) { + return std::abs(static_cast(a)) < + std::abs(static_cast(b)); + }))); + const auto rtol_atol = + calculate_rtol_atol_bquant( + K, kbatch, max_accumulated_value); + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_result, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + std::cout << "For " << instanceName << " Relative error threshold is " + << rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is " + << rtol_atol.at(ck_tile::number<1>{}) << std::endl; + std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl; + + return pass; +} + +/// @brief CPU reference implementation for BQuant GEMM +void bquant_gemm_host_reference(int verify, + ck_tile::HostTensor& a_m_k, + ck_tile::HostTensor& b_k_n, + ck_tile::HostTensor& bq_qk_n, + ck_tile::HostTensor& c_m_n_host_result) +{ + if(verify == 1) + { + c_m_n_host_result.SetZero(); + using QuantGroupSize = typename SelectedKernel::QuantGroupSize; + ck_tile::reference_gemm_quant( + a_m_k, bq_qk_n, b_k_n, c_m_n_host_result); + } +} +#pragma clang diagnostic pop diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py new file mode 100644 index 0000000000..b19d3702e2 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py @@ -0,0 +1,464 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import sys +import json +import subprocess +import argparse +import time +from pathlib import Path +from typing import List, Dict, Tuple, Optional + + +class BQuantGemmBenchmark: + def __init__(self, build_dir: str, verbose: bool = False): + self.build_dir = Path(build_dir) + self.verbose = verbose + self.results = [] + + def discover_kernels(self) -> List[Path]: + """Find all benchmark_gemm_bquant_* executables in the build directory.""" + bin_dir = self.build_dir / "bin" + if not bin_dir.exists(): + print(f"Error: Binary directory {bin_dir} does not exist") + return [] + + kernels = list(bin_dir.glob("benchmark_gemm_bquant_*")) + if self.verbose: + print(f"Found {len(kernels)} BQuant kernel executables") + for k in kernels: + print(f" - {k.name}") + return kernels + + def extract_kernel_info(self, kernel_path: Path) -> Dict[str, str]: + """Extract kernel information from filename.""" + name = kernel_path.stem + + info = { + "executable": str(kernel_path), + "name": name, + "data_type": "unknown", + "layout": "unknown", + "pipeline": "unknown", + "scheduler": "unknown", + "epilogue": "unknown", + "b_preshuffle_quant": False, + "preshuffle_b": False, + } + + # Parse: benchmark_gemm_bquant_fp8_rcr_compv3_default_intrawave_False_False_True_False_False_128x128x128_... + parts = name.split("_") + + if len(parts) >= 3: + # Skip "benchmark_gemm_bquant" prefix (3 parts) + idx = 3 + if idx < len(parts): + info["data_type"] = parts[idx] + idx += 1 + if idx < len(parts): + info["layout"] = parts[idx] + idx += 1 + if idx < len(parts): + info["pipeline"] = parts[idx] + idx += 1 + if idx < len(parts): + info["epilogue"] = parts[idx] + idx += 1 + if idx < len(parts): + info["scheduler"] = parts[idx] + + # Parse tile dimensions + config_info = self.parse_detailed_config(name) + info.update(config_info) + + info["config_id"] = self.generate_config_id(info) + return info + + def parse_detailed_config(self, kernel_name: str) -> Dict: + """Parse tile dimensions and boolean flags from kernel name.""" + config = { + "tile_sizes": {"tile_m": 0, "tile_n": 0, "tile_k": 0}, + "warp_config": {"warp_m": 0, "warp_n": 0, "warp_k": 0}, + "warp_tile": {"warp_tile_m": 0, "warp_tile_n": 0, "warp_tile_k": 0}, + "optimization_flags": { + "pad_m": False, + "pad_n": False, + "pad_k": False, + "b_preshuffle_quant": False, + "preshuffle_b": False, + }, + } + + parts = kernel_name.split("_") + + # Locate the first boolean token; collect the contiguous boolean run + bool_start = -1 + bool_sequence = [] + for part in parts: + if part in ("True", "False"): + if bool_start == -1: + bool_start = 0 + bool_sequence.append(part == "True") + elif bool_start != -1: + break + + if len(bool_sequence) >= 5: + config["optimization_flags"]["pad_m"] = bool_sequence[0] + config["optimization_flags"]["pad_n"] = bool_sequence[1] + config["optimization_flags"]["pad_k"] = bool_sequence[2] + config["optimization_flags"]["b_preshuffle_quant"] = bool_sequence[3] + config["optimization_flags"]["preshuffle_b"] = bool_sequence[4] + + # Extract dimension groups (e.g., 128x128x128) in positional order. + # Kernel names encode groups as: tile_MxNxK_warp_MxNxK_warp_tile_MxNxK + dimension_groups = [] + for part in parts: + if "x" in part and len(part.split("x")) == 3: + try: + dims = [int(x) for x in part.split("x")] + if all(d > 0 for d in dims): + dimension_groups.append(dims) + except ValueError: + continue + + if len(dimension_groups) >= 3: + # Use positional order: tile, warp, warp_tile + config["tile_sizes"]["tile_m"] = dimension_groups[0][0] + config["tile_sizes"]["tile_n"] = dimension_groups[0][1] + config["tile_sizes"]["tile_k"] = dimension_groups[0][2] + config["warp_config"]["warp_m"] = dimension_groups[1][0] + config["warp_config"]["warp_n"] = dimension_groups[1][1] + config["warp_config"]["warp_k"] = dimension_groups[1][2] + config["warp_tile"]["warp_tile_m"] = dimension_groups[2][0] + config["warp_tile"]["warp_tile_n"] = dimension_groups[2][1] + config["warp_tile"]["warp_tile_k"] = dimension_groups[2][2] + + return config + + def generate_config_id(self, info: Dict) -> str: + """Generate a compact config ID from kernel info.""" + parts = [ + info.get("data_type", "unk"), + info.get("layout", "unk"), + info.get("pipeline", "unk"), + info.get("scheduler", "unk"), + ] + + tile_sizes = info.get("tile_sizes", {}) + if tile_sizes.get("tile_m", 0) > 0: + parts.append( + f"{tile_sizes['tile_m']}x{tile_sizes['tile_n']}x{tile_sizes['tile_k']}" + ) + + return "_".join(parts) + + def run_kernel(self, kernel_path: Path, params: Dict[str, str]) -> Optional[Dict]: + """Run a single kernel with given parameters.""" + results_dir = self.build_dir / "results" + results_dir.mkdir(exist_ok=True) + + json_file = results_dir / f"{kernel_path.stem}.json" + + cmd = [str(kernel_path)] + for key, value in params.items(): + cmd.append(f"-{key}={value}") + cmd.append("-json_output=true") + + if self.verbose: + print(f"Running: {' '.join(cmd)}") + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=60) + + if result.returncode != 0: + print(f"Error running {kernel_path.name}: {result.stderr}") + return None + + output = result.stdout.strip() + if output: + with open(json_file, "w") as f: + f.write(output) + return self.parse_json_file(json_file) + else: + print(f"No output from {kernel_path.name}") + return None + + except subprocess.TimeoutExpired: + print(f"Timeout running {kernel_path.name}") + return None + except Exception as e: + print(f"Error running {kernel_path.name}: {e}") + return None + + def parse_json_file(self, json_file: Path) -> Optional[Dict]: + """Parse JSON data from kernel output file.""" + try: + with open(json_file, "r") as f: + content = f.read().strip() + + data = json.loads(content) + result = data.copy() + if "perf_result" in data: + perf = data["perf_result"] + result["time_ms"] = perf.get("latency(ms)", 0) + result["tflops"] = perf.get("tflops(TFlops)", 0) + result["bandwidth_gb_s"] = perf.get("bandwidth(GB/s)", 0) + return result + + except (json.JSONDecodeError, Exception) as e: + if self.verbose: + print(f"Failed to parse JSON from {json_file}: {e}") + return None + + def benchmark_problem_size( + self, + kernels: List[Path], + m: int, + n: int, + k: int, + group_size_k: int = 128, + split_k: int = 1, + verify: int = 0, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> List[Dict]: + """Benchmark all kernels for a specific problem size.""" + results = [] + + params = { + "m": m, + "n": n, + "k": k, + "group_size_k": group_size_k, + "split_k": split_k, + "verify": verify, + "warmup": warmup, + "repeat": repeat, + "flush_cache": str(flush_cache).lower(), + "rotating_count": rotating_count, + } + + print(f"\nBenchmarking M={m}, N={n}, K={k}, group_size_k={group_size_k}") + + for kernel_path in kernels: + kernel_info = self.extract_kernel_info(kernel_path) + result = self.run_kernel(kernel_path, params) + + if result: + structured_result = { + "name": kernel_info["name"], + "config_id": kernel_info["config_id"], + "problem": result.get("problem", {}), + "perf_result": result.get("perf_result", {}), + "config": { + "data_type": kernel_info["data_type"], + "layout": kernel_info["layout"], + "pipeline": kernel_info["pipeline"], + "scheduler": kernel_info["scheduler"], + "epilogue": kernel_info["epilogue"], + "tile_sizes": kernel_info.get("tile_sizes", {}), + "warp_config": kernel_info.get("warp_config", {}), + "warp_tile": kernel_info.get("warp_tile", {}), + "optimization_flags": kernel_info.get("optimization_flags", {}), + }, + "executable": kernel_info["executable"], + "time_ms": result.get("time_ms", 0), + "tflops": result.get("tflops", 0), + "bandwidth_gb_s": result.get("bandwidth_gb_s", 0), + } + + results.append(structured_result) + + if self.verbose: + print( + f" {kernel_info['config_id']}: " + f"{structured_result['tflops']:.2f} TFLOPS, " + f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " + f"{structured_result['time_ms']:.2f}ms" + ) + + return results + + def find_best_kernel( + self, results: List[Dict], metric: str = "tflops" + ) -> Optional[Dict]: + """Find the best performing kernel based on metric.""" + if not results: + return None + + if metric == "tflops": + return max(results, key=lambda x: x.get("tflops", 0)) + elif metric == "time_ms": + return min(results, key=lambda x: x.get("time_ms", float("inf"))) + elif metric == "bandwidth_gb_s": + return max(results, key=lambda x: x.get("bandwidth_gb_s", 0)) + else: + raise ValueError(f"Unknown metric: {metric}") + + def benchmark_sweep( + self, + problem_sizes: List[Tuple[int, int, int]], + group_size_k: int = 128, + split_k_values: Optional[List[int]] = None, + verify: bool = False, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> Dict: + """Run comprehensive benchmark sweep.""" + if split_k_values is None: + split_k_values = [1] + kernels = self.discover_kernels() + if not kernels: + print("No kernels found!") + return {} + + all_results = [] + best_kernels = {} + + for m, n, k in problem_sizes: + for split_k in split_k_values: + results = self.benchmark_problem_size( + kernels, + m, + n, + k, + group_size_k=group_size_k, + split_k=split_k, + verify=1 if verify else 0, + warmup=warmup, + repeat=repeat, + flush_cache=flush_cache, + rotating_count=rotating_count, + ) + + all_results.extend(results) + + best = self.find_best_kernel(results) + if best: + key = f"m{m}_n{n}_k{k}_splitk{split_k}" + best_kernels[key] = best + print( + f"Best for {key}: {best['name']} " + f"({best['tflops']:.2f} TFLOPS, " + f"{best['bandwidth_gb_s']:.2f} GB/s, " + f"{best['time_ms']:.2f}ms)" + ) + + self.results = all_results + return best_kernels + + def export_json(self, filename: str, best_kernels: Dict = None): + """Export results to JSON.""" + from datetime import datetime + + output_data = { + "benchmark_metadata": { + "timestamp": datetime.now().isoformat(), + "operator": "gemm_bquant", + "total_kernels_tested": len(self.results), + "successful_runs": len( + [r for r in self.results if r.get("tflops", 0) > 0] + ), + }, + "kernel_results": self.results, + "best_kernels_by_problem": best_kernels or {}, + } + + with open(filename, "w") as f: + json.dump(output_data, f, indent=2) + print(f"JSON results exported to {filename}") + + +def main(): + parser = argparse.ArgumentParser(description="BQuant GEMM Kernel Benchmarking Tool") + parser.add_argument( + "build_dir", help="Build directory containing kernel executables" + ) + parser.add_argument( + "--problem-sizes", + nargs="+", + default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"], + help="Problem sizes as M,N,K tuples", + ) + parser.add_argument( + "--group-size-k", + type=int, + default=128, + help="Quantization group size along K (default: 128)", + ) + parser.add_argument( + "--split-k", + nargs="+", + type=int, + default=[1], + help="Split-K values to test", + ) + parser.add_argument("--verify", action="store_true", help="Enable verification") + parser.add_argument("--verbose", action="store_true", help="Verbose output") + parser.add_argument( + "--warmup", + type=int, + default=50, + help="Number of warmup iterations", + ) + parser.add_argument( + "--repeat", + type=int, + default=100, + help="Number of benchmark iterations", + ) + parser.add_argument( + "--no-flush-cache", + action="store_true", + help="Disable cache flushing between iterations", + ) + parser.add_argument( + "--rotating-count", + type=int, + default=1000, + help="Number of iterations to rotate the cache (default: 1000)", + ) + parser.add_argument("--json", help="JSON output filename (optional)") + + args = parser.parse_args() + + problem_sizes = [] + for size_str in args.problem_sizes: + try: + m, n, k = map(int, size_str.split(",")) + problem_sizes.append((m, n, k)) + except ValueError: + print(f"Invalid problem size: {size_str}") + return 1 + + benchmark = BQuantGemmBenchmark(args.build_dir, verbose=args.verbose) + + print("Starting BQuant GEMM kernel benchmark sweep...") + start_time = time.time() + + best_kernels = benchmark.benchmark_sweep( + problem_sizes=problem_sizes, + group_size_k=args.group_size_k, + split_k_values=args.split_k, + verify=args.verify, + warmup=args.warmup, + repeat=args.repeat, + flush_cache=not args.no_flush_cache, + rotating_count=args.rotating_count, + ) + + elapsed_time = time.time() - start_time + print(f"\nBenchmark completed in {elapsed_time:.2f} seconds") + + if args.json: + benchmark.export_json(args.json, best_kernels) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp new file mode 100644 index 0000000000..746ced532f --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark_single.cpp @@ -0,0 +1,166 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "gemm_bquant_profiler.hpp" +#include "gemm_bquant_common.hpp" + +// The kernel header is included via the compile command line with -include flag +// It defines SelectedKernel struct, KERNEL_NAME, and type aliases: +// ADataType, BDataType, BQDataType, AccDataType, CDataType +// ALayout, BLayout, CLayout, BQLayout + +inline auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "3840", "The value for m dimension. Default is 3840.") + .insert("n", "4096", "The value for n dimension. Default is 4096.") + .insert("k", "2048", "The value for k dimension. Default is 2048.") + .insert("stride_a", "0", "The stride value for tensor A. Default is 0.") + .insert("stride_b", "0", "The stride value for tensor B. Default is 0.") + .insert("stride_c", "0", "The stride value for tensor C. Default is 0.") + .insert("split_k", "1", "The split value for k dimension. Default is 1.") + .insert("group_size_k", + std::to_string(SelectedKernel::GroupSizeK), + "The quantization group size along K. Default matches kernel config.") + .insert("verify", + "1", + "The type of validation. Set to 0 for no validation, 1 for validation on CPU. " + "Default is 1, CPU validation.") + .insert("log", + "false", + "Whether output kernel instance information or not. Possible values are true or " + "false. Default is false") + .insert( + "warmup", "50", "The number of iterations before benchmark the kernel. Default is 50.") + .insert( + "repeat", "100", "The number of iterations to benchmark the kernel. Default is 100.") + .insert("timer", + "true", + "Whether if the timer is gpu timer or not. Possible values are false or true. " + "Default is true.") + .insert("init", + "0", + "The method of tensor initialization. Set to 0 for random, to 1 for linear, or 2 " + "for constant(1). Default is 0, random.") + .insert("flush_cache", + "true", + "To flush cache, possible values are true or false. Default is true.") + .insert( + "rotating_count", "1000", "number of iterations to rotate the cache. default is 1000.") + .insert("metric", + "0", + "Metric with which to measure kernel performance. Set to 0 for latency, 1 for " + "tflops, or 2 for bandwidth. Default is 0, latency.") + .insert("csv_filename", + "", + "The filename of benchmark result. Default is empty (no CSV output).") + .insert("json_output", + "false", + "Whether to output results in JSON format only. Possible values are true or false. " + "Default is false"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} + +void benchmark_single(const ck_tile::ArgParser& arg_parser) +{ + std::string dtype_a = DataTypeTraits::name; + std::string dtype_b = DataTypeTraits::name; + std::string dtype_bq = DataTypeTraits::name; + std::string dtype_acc = DataTypeTraits::name; + std::string dtype_c = DataTypeTraits::name; + + std::string layout_a = ALayout::name; + std::string layout_b = BLayout::name; + std::string layout_c = CLayout::name; + + int M = arg_parser.get_int("m"); + int N = arg_parser.get_int("n"); + int K = arg_parser.get_int("k"); + int group_size_k = arg_parser.get_int("group_size_k"); + + if(M <= 0 || N <= 0 || K <= 0) + { + throw std::invalid_argument("m, n, k must be positive integers"); + } + if(group_size_k <= 0 || K % group_size_k != 0) + { + throw std::invalid_argument( + "group_size_k must be positive and k must be divisible by group_size_k"); + } + + BQuantGemmProblem problem{arg_parser.get_int("split_k"), + M, + N, + K, + arg_parser.get_int("stride_a"), + arg_parser.get_int("stride_b"), + arg_parser.get_int("stride_c"), + 0, // stride_bq computed by profiler + group_size_k, + dtype_a, + dtype_b, + dtype_bq, + dtype_acc, + dtype_c, + layout_a, + layout_b, + layout_c}; + + Setting setting{arg_parser.get_int("warmup"), + arg_parser.get_int("repeat"), + arg_parser.get_bool("timer"), + arg_parser.get_int("verify"), + arg_parser.get_int("init"), + arg_parser.get_bool("log"), + arg_parser.get_str("csv_filename"), + arg_parser.get_bool("flush_cache"), + arg_parser.get_int("rotating_count"), + arg_parser.get_bool("json_output")}; + + BQuantGemmProfiler profiler{setting}; + + try + { + auto kernel_func = [](const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& stream) { + return SelectedKernel::launch(args, stream); + }; + + profiler.benchmark(problem, kernel_func); + profiler.select_best_instance(static_cast(arg_parser.get_int("metric"))); + } + catch(const std::exception& e) + { + std::cerr << "Benchmark failed: " << e.what() << std::endl; + } +} + +int main(int argc, char* argv[]) +{ + try + { + auto [result, parser] = create_args(argc, argv); + if(!result) + return EXIT_FAILURE; + + benchmark_single(parser); + return 0; + } + catch(const std::exception& e) + { + std::cerr << "Error: " << e.what() << "\n"; + return EXIT_FAILURE; + } +} diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp new file mode 100644 index 0000000000..1408f4473d --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_common.hpp @@ -0,0 +1,74 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/core/numeric/integer.hpp" + +// DataTypeTraits for all supported types +template +struct DataTypeTraits; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp32"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp16"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "bf16"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "fp8"; +}; + +template <> +struct DataTypeTraits +{ + static constexpr const char* name = "bf8"; +}; + +// Helper function to determine if a layout is row-major +template +constexpr auto is_row_major(Layout) +{ + return ck_tile::bool_constant>{}; +} + +// Structure to hold kernel traits for dispatcher +struct BQuantKernelTraits +{ + std::string pipeline; // compv3 + std::string scheduler; // intrawave + std::string epilogue; // default, cshuffle + bool pad_m; + bool pad_n; + bool pad_k; + bool b_preshuffle_quant; + + BQuantKernelTraits() + : pipeline("compv3"), + scheduler("intrawave"), + epilogue("default"), + pad_m(false), + pad_n(false), + pad_k(false), + b_preshuffle_quant(false) + { + } +}; diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py new file mode 100644 index 0000000000..770ec9d6f1 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_instance_builder.py @@ -0,0 +1,288 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import os +import argparse +import importlib.util +import multiprocessing +import concurrent.futures + + +def _import_gemm_kernel_builder(): + """Import the base GemmKernelBuilder from the gemm directory.""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + gemm_dir = os.path.dirname(os.path.dirname(current_dir)) + + spec = importlib.util.spec_from_file_location( + "gemm_instance_builder", + os.path.join(gemm_dir, "gemm_instance_builder.py"), + ) + gemm_builder_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(gemm_builder_module) + + return gemm_builder_module.GemmKernelBuilder + + +GemmKernelBuilder = _import_gemm_kernel_builder() + + +class GemmBQuantKernelBuilder(GemmKernelBuilder): + def __init__( + self, + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json=None, + max_instances=None, + seed=None, + tier=None, + manifest_path=None, + ): + super().__init__( + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json, + max_instances=max_instances, + seed=seed, + tier=tier, + manifest_path=manifest_path, + ) + self.group_size_k = self.config.get("group_size_k", 128) + + def _generate_all_individual(self, num_workers=None): + """Generate individual kernel files for separate compilation with parallel processing""" + if num_workers is None: + num_workers = min(multiprocessing.cpu_count(), 8) + + tile_configs = self._get_tile_configs() + trait_combos = self._generate_trait_combinations() + + work_items = [] + for tile_config in tile_configs: + for trait_combo in trait_combos: + work_items.append( + ( + tile_config, + trait_combo, + self.kernel_name_prefix, + self.working_path, + self.gpu_target, + self.datatype, + self.layout, + self.config_json, + ) + ) + kernel_list = [] + completed = 0 + + with concurrent.futures.ProcessPoolExecutor( + max_workers=num_workers + ) as executor: + future_to_item = { + executor.submit(_generate_single_kernel_individual, item): item + for item in work_items + } + + for future in concurrent.futures.as_completed(future_to_item): + completed += 1 + if completed % 100 == 0 or completed == len(work_items): + print( + f" Progress: {completed}/{len(work_items)} kernels generated" + ) + try: + result = future.result() + if result: + kernel_list.append(result) + except Exception as exc: + item = future_to_item[future] + print(f"Kernel generation failed for {item}: {exc}") + + kernel_list.sort(key=lambda x: x[0]) + self._generate_cmake_individual_targets(kernel_list) + + print( + f"Generated {len(kernel_list)} individual kernel files in {self.working_path}" + ) + + +def _generate_single_kernel_individual(work_item): + """Worker function to generate a single individual kernel file""" + ( + tile_config, + trait_combo, + kernel_name_prefix, + working_path, + gpu_target, + datatype, + layout, + config_json, + ) = work_item + + builder = GemmBQuantKernelBuilder( + kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json + ) + + try: + kernel_name, instance_code = builder._generate_kernel_instance( + tile_config, trait_combo + ) + + simplified_name = kernel_name.removeprefix(kernel_name_prefix + "_") + + header_file = working_path / f"gemm_bquant_single_{simplified_name}.hpp" + with open(header_file, "w") as f: + f.write(instance_code) + + return (kernel_name, trait_combo, tile_config) + except Exception as e: + print(f"Error generating individual kernel: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="GEMM BQuant kernel instance builder") + parser.add_argument("--working_path", required=True, help="Working directory path") + parser.add_argument("--gpu_target", required=True, help="GPU target architecture") + parser.add_argument( + "--datatype", + required=True, + choices=["fp8", "bf8"], + help="Data type (fp8 or bf8)", + ) + parser.add_argument( + "--layout", + required=True, + choices=["rcr", "rrr", "ccr", "crr"], + help="Matrix layout", + ) + parser.add_argument("--config_json", help="Configuration JSON file") + parser.add_argument("--num_workers", type=int, help="Number of parallel workers") + parser.add_argument( + "--gen_all_individual", + action="store_true", + help="Generate individual kernel files", + ) + parser.add_argument( + "--gen_single", + action="store_true", + help="Generate a single kernel file", + ) + parser.add_argument("--kernel_name", help="Kernel name for single generation") + parser.add_argument("--tile_config", help="Tile configuration string") + parser.add_argument("--trait_combo", help="Trait combination string") + parser.add_argument( + "--list_kernels", + action="store_true", + help="List kernel configurations without generating files", + ) + parser.add_argument( + "--max-instances", + type=int, + default=None, + help="Maximum number of kernel instances to select via sampling", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="RNG seed for deterministic sampling; if omitted, derived from today's date", + ) + parser.add_argument( + "--tier", + default=None, + help="Sampling tier name (e.g., 'daily')", + ) + parser.add_argument( + "--manifest-path", + default=None, + help="Directory to write chosen_instances manifest JSON", + ) + + args = parser.parse_args() + + layout_parts = args.layout.lower() + assert len(layout_parts) == 3, ( + f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr')" + ) + assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], ( + f"Invalid matrix_a layout: {layout_parts[0]} or matrix_b layout: {layout_parts[1]}" + ) + assert layout_parts[2] == "r", ( + f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' for row major)" + ) + + kernel_name_prefix = "gemm_bquant" + builder = GemmBQuantKernelBuilder( + kernel_name_prefix, + args.working_path, + args.gpu_target, + args.datatype, + args.layout, + args.config_json, + max_instances=args.max_instances, + seed=args.seed, + tier=args.tier, + manifest_path=args.manifest_path, + ) + + if args.list_kernels: + builder._list_kernels() + elif args.gen_single: + if not args.kernel_name or not args.tile_config or not args.trait_combo: + parser.error( + "--gen_single requires --kernel_name, --tile_config, and --trait_combo" + ) + + # Parse tile config + tile_parts = args.tile_config.split("_") + tile_dims = tile_parts[0].split("x") + warp_dims = tile_parts[1].split("x") + warp_tile_dims = tile_parts[2].split("x") + + tile_config = { + "tile_m": int(tile_dims[0]), + "tile_n": int(tile_dims[1]), + "tile_k": int(tile_dims[2]), + "warp_m": int(warp_dims[0]), + "warp_n": int(warp_dims[1]), + "warp_k": int(warp_dims[2]), + "warp_tile_m": int(warp_tile_dims[0]), + "warp_tile_n": int(warp_tile_dims[1]), + "warp_tile_k": int(warp_tile_dims[2]), + } + + # Parse trait combo: + # pipeline_epilogue_scheduler_padM_padN_padK_bPreshuffle + trait_parts = args.trait_combo.split("_") + if len(trait_parts) < 7: + parser.error( + f"--trait_combo must have 7 underscore-separated fields " + f"(e.g. 'compv3_default_intrawave_True_True_True_False'), " + f"got {len(trait_parts)} field(s): '{args.trait_combo}'" + ) + trait_combo = ( + trait_parts[0], # pipeline + trait_parts[1], # epilogue + trait_parts[2], # scheduler + trait_parts[3] == "True", # pad_m + trait_parts[4] == "True", # pad_n + trait_parts[5] == "True", # pad_k + trait_parts[6] == "True", # b_preshuffle_quant + ) + + builder._generate_kernel_instance(tile_config, trait_combo) + elif args.gen_all_individual: + builder._generate_all_individual(args.num_workers) + else: + parser.error( + "Must specify one of: --list_kernels, --gen_all_individual, or --gen_single" + ) + + +if __name__ == "__main__": + main() diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp new file mode 100644 index 0000000000..7f562ef0ad --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_profiler.hpp @@ -0,0 +1,293 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/host/tensor_shuffle_utils.hpp" +#include "ck_tile/ops/gemm_quant.hpp" +#include "gemm_bquant_benchmark.hpp" + +class BQuantGemmProfiler +{ + public: + explicit BQuantGemmProfiler(Setting setting) : setting_(setting) {} + + BQuantGemmProfiler(const BQuantGemmProfiler&) = delete; + BQuantGemmProfiler& operator=(const BQuantGemmProfiler&) = delete; + + void reset() { kernel_instances_.clear(); } + + void benchmark(BQuantGemmProblem problem, + std::function kernel_func) + { + std::vector(ck_tile::QuantGemmHostArgs&, + const ck_tile::stream_config&)>> + callables; + + callables.push_back( + [kernel_func](ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& stream) { + float time = kernel_func(args, stream); + return std::make_tuple(std::string(KERNEL_NAME), time); + }); + + benchmark(problem, callables); + } + + void benchmark(BQuantGemmProblem problem, + std::vector( + ck_tile::QuantGemmHostArgs&, const ck_tile::stream_config&)>>& callables) + { + const ALayout layout_a = ALayout{}; + const BLayout layout_b = BLayout{}; + const CLayout layout_c = CLayout{}; + const BQLayout layout_bq = BQLayout{}; + + problem.stride_a_ = ck_tile::get_default_stride( + problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a)); + problem.stride_b_ = ck_tile::get_default_stride( + problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b)); + problem.stride_c_ = ck_tile::get_default_stride( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c)); + + // Compute BQ scale tensor dimensions: [K / group_size_k, N] + if(problem.k_ % problem.group_size_k_ != 0) + throw std::runtime_error("k_ must be divisible by group_size_k_"); + const ck_tile::index_t QK_B = problem.k_ / problem.group_size_k_; + problem.stride_bq_ = ck_tile::get_default_stride( + QK_B, problem.n_, problem.stride_bq_, is_row_major(layout_bq)); + + // Allocate host tensors + ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( + problem.m_, problem.k_, problem.stride_a_, is_row_major(layout_a))); + ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( + problem.k_, problem.n_, problem.stride_b_, is_row_major(layout_b))); + ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( + problem.m_, problem.n_, problem.stride_c_, is_row_major(layout_c))); + + // BQ scale tensor: [QK_B, N] + ck_tile::HostTensor bq_qk_n(ck_tile::host_tensor_descriptor( + QK_B, problem.n_, problem.stride_bq_, is_row_major(layout_bq))); + + // Initialize tensors + if(setting_.init_method_ == 0) + { + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); + ck_tile::FillUniformDistribution{0.5f, 1.5f}(bq_qk_n); + } + else if(setting_.init_method_ == 1) + { + ck_tile::FillMonotonicSeq{}(a_m_k); + ck_tile::FillMonotonicSeq{}(b_k_n); + ck_tile::FillConstant{static_cast(1)}(bq_qk_n); + } + else if(setting_.init_method_ == 2) + { + ck_tile::FillConstant{static_cast(1)}(a_m_k); + ck_tile::FillConstant{static_cast(1)}(b_k_n); + ck_tile::FillConstant{static_cast(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 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 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& c_m_n_host_result, + ck_tile::HostTensor& c_m_n_dev_result, + const std::tuple& kernel_run_result) + { + auto [name, avg_time] = kernel_run_result; + + KernelInstance kernel_instance{name, problem, {-1.0f, -1.0f, -1.0f}}; + + // Compute performance metrics + std::size_t flop = std::size_t(2) * problem.m_ * problem.n_ * problem.k_; + std::size_t num_byte = sizeof(ADataType) * problem.m_ * problem.k_ + + sizeof(BDataType) * problem.n_ * problem.k_ + + sizeof(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(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 kernel_instances_; +}; diff --git a/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/test_gemm_bquant_instance_builder.py b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/test_gemm_bquant_instance_builder.py new file mode 100644 index 0000000000..f1be2dd2d2 --- /dev/null +++ b/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/test_gemm_bquant_instance_builder.py @@ -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() diff --git a/tile_engine/ops/gemm/gemm_instance_builder.py b/tile_engine/ops/gemm/gemm_instance_builder.py index 31148661f2..1026990fdf 100644 --- a/tile_engine/ops/gemm/gemm_instance_builder.py +++ b/tile_engine/ops/gemm/gemm_instance_builder.py @@ -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 #include #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;""" # 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>; + + // 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>; + + // 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;""" # 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")};""" - 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)};""" + + 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(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)};""" + + 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(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;""" 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)};""" + + 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(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;""" + 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;""" + 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)};""" + + 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(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;""" + 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;""" + 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 diff --git a/tile_engine/ops/gemm/gemm_validation_utils.py b/tile_engine/ops/gemm/gemm_validation_utils.py index 79de3178f8..183a582458 100644 --- a/tile_engine/ops/gemm/gemm_validation_utils.py +++ b/tile_engine/ops/gemm/gemm_validation_utils.py @@ -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" + ) diff --git a/tile_engine/sampling/__init__.py b/tile_engine/sampling/__init__.py index 9cffef5ac2..de5171dcb4 100644 --- a/tile_engine/sampling/__init__.py +++ b/tile_engine/sampling/__init__.py @@ -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 diff --git a/tile_engine/sampling/feasible_set.py b/tile_engine/sampling/feasible_set.py index 1f45e2908a..3d09647515 100644 --- a/tile_engine/sampling/feasible_set.py +++ b/tile_engine/sampling/feasible_set.py @@ -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", } diff --git a/tile_engine/sampling/op_weights.json b/tile_engine/sampling/op_weights.json index 226a856f91..031d6f3570 100644 --- a/tile_engine/sampling/op_weights.json +++ b/tile_engine/sampling/op_weights.json @@ -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 } }