diff --git a/benchmark/ops/bench_hipblaslt_algo_tuning.py b/benchmark/ops/bench_hipblaslt_algo_tuning.py new file mode 100644 index 000000000..6d371b124 --- /dev/null +++ b/benchmark/ops/bench_hipblaslt_algo_tuning.py @@ -0,0 +1,191 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""hipBLASLt algorithm tuning benchmark. + +Runs each dense-model GEMM shape (MBS=1) with level-2 AutoTune so +that the framework benchmarks all hipBLASLt heuristic algorithms and +prints the winner for each shape. The tuning logs come from the +AutoKernelDispatcher logger. + +By default all dtype/granularity combinations are tested. Use --dtype +and --granularity to restrict to a single combination. + +Usage: + python benchmark/ops/bench_hipblaslt_algo_tuning.py # all combos + python benchmark/ops/bench_hipblaslt_algo_tuning.py --dtype bf16 # bf16 only + python benchmark/ops/bench_hipblaslt_algo_tuning.py --dtype fp8 --granularity tensorwise + python benchmark/ops/bench_hipblaslt_algo_tuning.py --max-algos 5 +""" + +import argparse +import os +import sys +from datetime import datetime + +os.environ.setdefault("PRIMUS_TURBO_LOG_LEVEL", "INFO") + +import torch # noqa: E402 +from config import ( # noqa: E402 + DenseModelConfigs, + gen_gemm_test_cases, + get_platform_info, +) + +import primus_turbo.pytorch as turbo # noqa: E402 +from primus_turbo.pytorch.core.backend import GlobalBackendManager # noqa: E402 +from primus_turbo.pytorch.core.low_precision import ( # noqa: E402 + Float8QuantConfig, + Format, + ScaleDtype, + ScalingGranularity, +) + +FP8_GRANULARITY_CONFIGS = { + "tensorwise": Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE), + "rowwise": Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.ROWWISE), + "blockwise": Float8QuantConfig( + format=Format.E4M3, + granularity=ScalingGranularity.BLOCKWISE, + block_size=128, + ), +} + + +def run_shapes(dtype_label, device, fp8_config=None): + """Run all dense-model GEMM shapes for a given dtype/granularity.""" + for model_name, model_config in DenseModelConfigs.items(): + test_cases = gen_gemm_test_cases(model_config) + for shape in test_cases: + M, N, K = shape[0], shape[1], shape[2] + + print(f"{'='*60}") + print(f"Case: {model_name}, M={M}, N={N}, K={K}, dtype={dtype_label}") + print(f"{'='*60}") + + a = torch.randn((M, K), dtype=torch.bfloat16, device=device) + b = torch.randn((N, K), dtype=torch.bfloat16, device=device) + + if fp8_config is not None: + turbo.ops.gemm_fp8(a, b, trans_b=True, config=fp8_config) + turbo.ops.gemm_fp8(a, b, trans_b=True, config=fp8_config) + else: + turbo.ops.gemm(a, b, trans_b=True) + turbo.ops.gemm(a, b, trans_b=True) + + torch.cuda.synchronize() + print() + + +def main(): + parser = argparse.ArgumentParser(description="hipBLASLt Algorithm Tuning Benchmark") + parser.add_argument( + "--dtype", + type=str, + choices=["bf16", "fp8", "all"], + default="all", + help="Data type to test (default: all)", + ) + parser.add_argument( + "--granularity", + type=str, + choices=["tensorwise", "rowwise", "blockwise", "all"], + default="all", + help="FP8 scaling granularity (only used when dtype includes fp8, default: all)", + ) + parser.add_argument( + "--max-algos", + type=int, + default=100, + help="Max hipBLASLt algorithms to try per shape (default: 10)", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="Output log file (default: auto-generated name). Use 'stdout' to skip file output.", + ) + args = parser.parse_args() + + os.environ["PRIMUS_TURBO_HIPBLASLT_TUNE_MAX_ALGOS"] = str(args.max_algos) + GlobalBackendManager.set_auto_tune(2) + + platform, gpu_name = get_platform_info() + device = "cuda" + + log_file = None + if args.output != "stdout": + filename = args.output or ( + f"hipblaslt_algo_tuning_{args.dtype}_{args.granularity}" + f"_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{gpu_name}.log" + ) + log_file = open(filename, "w") + + # Tee output to both stdout and file + class Tee: + def __init__(self, *streams): + self.streams = streams + + def write(self, data): + for s in self.streams: + s.write(data) + s.flush() + + def flush(self): + for s in self.streams: + s.flush() + + sys.stdout = Tee(sys.__stdout__, log_file) + sys.stderr = Tee(sys.__stderr__, log_file) + + # Point existing logger handlers at the new stderr so their + # output is captured in the log file too. + import logging + + for handler in logging.getLogger("primus_turbo").handlers: + if isinstance(handler, logging.StreamHandler): + handler.stream = sys.stdout + + print(f"Platform: {platform}, GPU: {gpu_name}") + print(f"AutoTune level: 2, hipBLASLt algo cap: {args.max_algos}") + print(f"Dtype: {args.dtype}, Granularity: {args.granularity}\n") + + run_bf16 = args.dtype in ("bf16", "all") + run_fp8 = args.dtype in ("fp8", "all") + + if run_bf16: + print(f"\n{'#'*60}") + print(f"# BF16 GEMM") + print(f"{'#'*60}\n") + run_shapes("bf16", device) + GlobalBackendManager.reset() + GlobalBackendManager.set_auto_tune(2) + + if run_fp8: + granularities = ( + list(FP8_GRANULARITY_CONFIGS.keys()) if args.granularity == "all" else [args.granularity] + ) + for gran in granularities: + print(f"\n{'#'*60}") + print(f"# FP8 GEMM ({gran})") + print(f"{'#'*60}\n") + run_shapes(f"fp8-{gran}", device, fp8_config=FP8_GRANULARITY_CONFIGS[gran]) + GlobalBackendManager.reset() + GlobalBackendManager.set_auto_tune(2) + + GlobalBackendManager.reset() + print("Done.") + + if log_file is not None: + sys.stdout = sys.__stdout__ + sys.stderr = sys.__stderr__ + log_file.close() + print(f"Log saved to {filename}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/ops/benchmark_suite.yaml b/benchmark/ops/benchmark_suite.yaml index 89b3d3c3d..0c1a617a7 100644 --- a/benchmark/ops/benchmark_suite.yaml +++ b/benchmark/ops/benchmark_suite.yaml @@ -281,6 +281,13 @@ tasks: env: {PRIMUS_TURBO_GROUPED_GEMM_BACKEND: TRITON} output: grouped_gemm_fp8_blockwise_triton_benchmark.csv + # ── Online Tuning (level 2) ────────────────────────────── + - label: online_tuning_all + group: online_tuning + script: bench_hipblaslt_algo_tuning.py + args: [--max-algos, 100] + output: online_tuning_all_benchmark.log + # ── DeepEP (multi-GPU, placed last) ────────────────────── - label: deep_ep_intranode group: deepep diff --git a/csrc/include/primus_turbo/gemm.h b/csrc/include/primus_turbo/gemm.h index a02b4395e..dc93e9a6e 100644 --- a/csrc/include/primus_turbo/gemm.h +++ b/csrc/include/primus_turbo/gemm.h @@ -42,6 +42,7 @@ inline size_t hipblaslt_dtype_bytes(hipDataType dtype) { int64_t get_hipblaslt_workspace_size_in_byte(); +// Run a GEMM using the heuristic algorithm at algo_index (0 = top-ranked). void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t rows_a, const int64_t cols_a, const int64_t lda, const void *scaleA_inv, hipblasOperation_t transA, const void *B, const hipDataType B_type, @@ -51,7 +52,19 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t const int64_t ldd, void *workspace, const int64_t workspace_size, const bool use_low_precision, const hipblasLtMatmulMatrixScale_t scale_mode, hipblasLtHandle_t handle, - hipStream_t stream); + hipStream_t stream, const int algo_index = 0); + +// Query how many heuristic algorithms are available for this problem. +int hipblaslt_gemm_get_algo_count(const void *A, const hipDataType A_type, const int64_t rows_a, + const int64_t cols_a, const int64_t lda, const void *scaleA_inv, + hipblasOperation_t transA, const void *B, + const hipDataType B_type, const int64_t rows_b, + const int64_t cols_b, const int64_t ldb, const void *scaleB_inv, + hipblasOperation_t transB, const hipDataType D_type, + const int64_t rows_d, const int64_t cols_d, const int64_t ldd, + const int64_t workspace_size, const bool use_low_precision, + hipblasLtMatmulMatrixScale_t scale_mode, + hipblasLtHandle_t handle); //================================================================== // CK GEMM diff --git a/csrc/kernels/gemm/hipblaslt_gemm.cu b/csrc/kernels/gemm/hipblaslt_gemm.cu index 36cd0c3bb..02f41c2ab 100644 --- a/csrc/kernels/gemm/hipblaslt_gemm.cu +++ b/csrc/kernels/gemm/hipblaslt_gemm.cu @@ -26,7 +26,7 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t const hipDataType D_type, const int64_t rows_d, const int64_t cols_d, const int64_t ldd, void *workspace, const int64_t workspace_size, const bool use_low_precision, hipblasLtMatmulMatrixScale_t scale_mode, - hipblasLtHandle_t handle, hipStream_t stream) { + hipblasLtHandle_t handle, hipStream_t stream, const int algo_index) { hipblasLtMatmulDesc_t operation_desc = nullptr; hipblasLtMatrixLayout_t A_desc = nullptr, B_desc = nullptr, D_desc = nullptr; hipblasLtMatmulPreference_t preference = nullptr; @@ -69,7 +69,7 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t operation_desc, scaleB_inv_ptr_desc, &scaleB_inv, sizeof(scaleB_inv))); } - const int request_solutions = 1; + const int request_solutions = algo_index + 1; std::vector algos(request_solutions); int returnedAlgoCount = 0; @@ -81,8 +81,8 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulAlgoGetHeuristic( handle, operation_desc, A_desc, B_desc, D_desc, D_desc, preference, request_solutions, algos.data(), &returnedAlgoCount)); - PRIMUS_TURBO_CHECK(returnedAlgoCount > 0, - "hipBLASLt: no valid algorithm found for current matmul config"); + PRIMUS_TURBO_CHECK(returnedAlgoCount > algo_index, + "hipBLASLt: requested algo_index exceeds available algorithms"); const float alpha = 1.0; const float beta = 0.0; @@ -96,7 +96,7 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t &beta, D, D_desc, D, D_desc, - &algos[0].algo, + &algos[algo_index].algo, workspace, workspace_size, stream)); // clang-format on @@ -108,4 +108,73 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulPreferenceDestroy(preference)); } +int hipblaslt_gemm_get_algo_count(const void *A, const hipDataType A_type, const int64_t rows_a, + const int64_t cols_a, const int64_t lda, const void *scaleA_inv, + hipblasOperation_t transA, const void *B, + const hipDataType B_type, const int64_t rows_b, + const int64_t cols_b, const int64_t ldb, const void *scaleB_inv, + hipblasOperation_t transB, const hipDataType D_type, + const int64_t rows_d, const int64_t cols_d, const int64_t ldd, + const int64_t workspace_size, const bool use_low_precision, + hipblasLtMatmulMatrixScale_t scale_mode, + hipblasLtHandle_t handle) { + hipblasLtMatmulDesc_t operation_desc = nullptr; + hipblasLtMatrixLayout_t A_desc = nullptr, B_desc = nullptr, D_desc = nullptr; + hipblasLtMatmulPreference_t preference = nullptr; + hipblasLtEpilogue_t epilogue = HIPBLASLT_EPILOGUE_DEFAULT; + hipblasComputeType_t gemm_compute_type = HIPBLAS_COMPUTE_32F; + + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatrixLayoutCreate(&A_desc, A_type, rows_a, cols_a, lda)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatrixLayoutCreate(&B_desc, B_type, rows_b, cols_b, ldb)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatrixLayoutCreate(&D_desc, D_type, rows_d, cols_d, ldd)); + + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatmulDescCreate(&operation_desc, gemm_compute_type, HIP_R_32F)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, HIPBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, HIPBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); + + if (use_low_precision) { + PRIMUS_TURBO_CHECK(scaleA_inv != nullptr); + PRIMUS_TURBO_CHECK(scaleB_inv != nullptr); + + hipblasLtMatmulDescAttributes_t scaleA_inv_ptr_desc = HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER; + hipblasLtMatmulDescAttributes_t scaleB_inv_ptr_desc = HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER; + + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, HIPBLASLT_MATMUL_DESC_A_SCALE_MODE, &scale_mode, sizeof(scale_mode))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, HIPBLASLT_MATMUL_DESC_B_SCALE_MODE, &scale_mode, sizeof(scale_mode))); + + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, scaleA_inv_ptr_desc, &scaleA_inv, sizeof(scaleA_inv))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + operation_desc, scaleB_inv_ptr_desc, &scaleB_inv, sizeof(scaleB_inv))); + } + + const int max_algos = 32768; + std::vector algos(max_algos); + int returnedAlgoCount = 0; + + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulPreferenceCreate(&preference)); + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatmulPreferenceSetAttribute(preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, sizeof(workspace_size))); + + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatmulAlgoGetHeuristic(handle, operation_desc, A_desc, B_desc, D_desc, D_desc, + preference, max_algos, algos.data(), &returnedAlgoCount)); + + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatrixLayoutDestroy(D_desc)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatrixLayoutDestroy(B_desc)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatrixLayoutDestroy(A_desc)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescDestroy(operation_desc)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulPreferenceDestroy(preference)); + + return returnedAlgoCount; +} + } // namespace primus_turbo diff --git a/csrc/pytorch/bindings_pytorch.cpp b/csrc/pytorch/bindings_pytorch.cpp index 2fa3f33f4..2549d663d 100644 --- a/csrc/pytorch/bindings_pytorch.cpp +++ b/csrc/pytorch/bindings_pytorch.cpp @@ -12,14 +12,21 @@ namespace primus_turbo::pytorch { TORCH_LIBRARY(primus_turbo_cpp_extension, m) { // ********* Gemm ********* - m.def("hipblaslt_gemm(Tensor A, Tensor B, " - "ScalarType out_dtype, bool transA, bool transB, bool transC) -> Tensor"); m.def( - "hipblaslt_gemm_fp8(Tensor A, Tensor scaleA_inv, Tensor B, Tensor scaleB_inv," - "ScalarType out_dtype, bool transA, bool transB, bool transC, str granularity) -> Tensor"); - m.def( - "hipblaslt_gemm_fp4(Tensor A, Tensor scaleA_inv, Tensor B, Tensor scaleB_inv," - "ScalarType out_dtype, bool transA, bool transB, bool transC, str granularity) -> Tensor"); + "hipblaslt_gemm(Tensor A, Tensor B, " + "ScalarType out_dtype, bool transA, bool transB, bool transC, int algo_index=0) -> Tensor"); + m.def("hipblaslt_gemm_fp8(Tensor A, Tensor scaleA_inv, Tensor B, Tensor scaleB_inv," + "ScalarType out_dtype, bool transA, bool transB, bool transC, str granularity, int " + "algo_index=0) -> Tensor"); + m.def("hipblaslt_gemm_fp4(Tensor A, Tensor scaleA_inv, Tensor B, Tensor scaleB_inv," + "ScalarType out_dtype, bool transA, bool transB, bool transC, str granularity, int " + "algo_index=0) -> Tensor"); + m.def("hipblaslt_gemm_algo_count(Tensor A, Tensor B, " + "ScalarType out_dtype, bool transA, bool transB, bool transC) -> int"); + m.def("hipblaslt_gemm_fp8_algo_count(Tensor A, Tensor scaleA_inv, Tensor B, Tensor scaleB_inv," + "ScalarType out_dtype, bool transA, bool transB, bool transC, str granularity) -> int"); + m.def("hipblaslt_gemm_fp4_algo_count(Tensor A, Tensor scaleA_inv, Tensor B, Tensor scaleB_inv," + "ScalarType out_dtype, bool transA, bool transB, bool transC, str granularity) -> int"); m.def("ck_gemm_fp8(Tensor a, Tensor b, Tensor a_scales, Tensor b_scales, bool transA," "bool transB, ScalarType out_dtype, str granularity) -> Tensor"); @@ -86,6 +93,9 @@ TORCH_LIBRARY_IMPL(primus_turbo_cpp_extension, CUDA, m) { m.impl("hipblaslt_gemm", hipblaslt_gemm); m.impl("hipblaslt_gemm_fp8", hipblaslt_gemm_fp8); m.impl("hipblaslt_gemm_fp4", hipblaslt_gemm_fp4); + m.impl("hipblaslt_gemm_algo_count", hipblaslt_gemm_algo_count); + m.impl("hipblaslt_gemm_fp8_algo_count", hipblaslt_gemm_fp8_algo_count); + m.impl("hipblaslt_gemm_fp4_algo_count", hipblaslt_gemm_fp4_algo_count); m.impl("ck_gemm_fp8", ck_gemm_fp8); m.impl("turbo_gemm_fp8", turbo_gemm_fp8); // ********* Quantization ********* diff --git a/csrc/pytorch/extensions.h b/csrc/pytorch/extensions.h index b7c09f85e..8856a4fcd 100644 --- a/csrc/pytorch/extensions.h +++ b/csrc/pytorch/extensions.h @@ -114,28 +114,43 @@ at::Tensor shuffle_weight_impl_meta(const at::Tensor weight, at::IntArrayRef lay //================================================================== at::Tensor hipblaslt_gemm(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, bool transA, - bool transB, bool transC); + bool transB, bool transC, int64_t algo_index = 0); at::Tensor hipblaslt_gemm_meta(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, - bool transA, bool transB, bool transC); + bool transA, bool transB, bool transC, int64_t algo_index = 0); at::Tensor hipblaslt_gemm_fp8(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, - bool transB, bool transC, const std::string &granularity); + bool transB, bool transC, const std::string &granularity, + int64_t algo_index = 0); at::Tensor hipblaslt_gemm_fp8_meta(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, bool transB, bool transC, - const std::string &granularity); + const std::string &granularity, int64_t algo_index = 0); at::Tensor hipblaslt_gemm_fp4(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, - bool transB, bool transC, const std::string &granularity); + bool transB, bool transC, const std::string &granularity, + int64_t algo_index = 0); at::Tensor hipblaslt_gemm_fp4_meta(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, bool transB, bool transC, - const std::string &granularity); + const std::string &granularity, int64_t algo_index = 0); + +int64_t hipblaslt_gemm_algo_count(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC); + +int64_t hipblaslt_gemm_fp8_algo_count(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, + const std::string &granularity); + +int64_t hipblaslt_gemm_fp4_algo_count(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, + const std::string &granularity); at::Tensor ck_gemm_fp8(at::Tensor &a, at::Tensor &b, at::Tensor &a_scales, at::Tensor &b_scales, const bool transA, const bool transB, at::ScalarType out_dtype, diff --git a/csrc/pytorch/gemm/gemm_meta.cpp b/csrc/pytorch/gemm/gemm_meta.cpp index be84a32a6..9d623227d 100644 --- a/csrc/pytorch/gemm/gemm_meta.cpp +++ b/csrc/pytorch/gemm/gemm_meta.cpp @@ -9,7 +9,7 @@ namespace primus_turbo::pytorch { at::Tensor hipblaslt_gemm_meta(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, - bool transA, bool transB, bool transC) { + bool transA, bool transB, bool transC, int64_t algo_index) { const int64_t m = transA ? A.size(1) : A.size(0); const int64_t n = transB ? B.size(0) : B.size(1); @@ -21,8 +21,8 @@ at::Tensor hipblaslt_gemm_meta(at::Tensor A, at::Tensor B, const at::ScalarType at::Tensor hipblaslt_gemm_fp8_meta(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, bool transB, bool transC, - const std::string &granularity) { - return hipblaslt_gemm_meta(A, B, out_dtype, transA, transB, transC); + const std::string &granularity, int64_t algo_index) { + return hipblaslt_gemm_meta(A, B, out_dtype, transA, transB, transC, algo_index); } at::Tensor ck_gemm_fp8_meta(at::Tensor &a, at::Tensor &b, at::Tensor &a_scales, diff --git a/csrc/pytorch/gemm/hipblaslt_gemm.cpp b/csrc/pytorch/gemm/hipblaslt_gemm.cpp index 51de4bdac..e066fe133 100644 --- a/csrc/pytorch/gemm/hipblaslt_gemm.cpp +++ b/csrc/pytorch/gemm/hipblaslt_gemm.cpp @@ -9,8 +9,9 @@ namespace primus_turbo::pytorch { -at::Tensor hipblaslt_gemm(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, bool transA, - bool transB, bool transC) { +static at::Tensor hipblaslt_gemm_core(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, int64_t algo_index = 0, + int *out_algo_count = nullptr) { PRIMUS_TURBO_CHECK(is_floating_point_dtype(A.scalar_type())); PRIMUS_TURBO_CHECK(is_floating_point_dtype(B.scalar_type())); PRIMUS_TURBO_CHECK(A.scalar_type() == B.scalar_type(), "A and B dtype mismatch"); @@ -61,18 +62,30 @@ at::Tensor hipblaslt_gemm(at::Tensor A, at::Tensor B, const at::ScalarType out_d PRIMUS_TURBO_ERROR("Not support layout."); } - at::Tensor C = at::empty({m, n}, torch::dtype(out_dtype).device(at::kCUDA)); - - auto stream = at::cuda::getCurrentCUDAStream(); - auto handle = at::cuda::getCurrentCUDABlasLtHandle(); - + auto handle = at::cuda::getCurrentCUDABlasLtHandle(); hipblasOperation_t trans_operation_A = transA ? HIPBLAS_OP_T : HIPBLAS_OP_N; hipblasOperation_t trans_operation_B = transB ? HIPBLAS_OP_T : HIPBLAS_OP_N; const hipDataType A_type = get_hipblaslt_dtype(A.scalar_type()); const hipDataType B_type = get_hipblaslt_dtype(B.scalar_type()); - const hipDataType C_type = get_hipblaslt_dtype(C.scalar_type()); + const hipDataType D_type = get_hipblaslt_dtype(out_dtype); + const int64_t workspace_size = get_hipblaslt_workspace_size_in_byte(); + + if (out_algo_count) { + // clang-format off + *out_algo_count = hipblaslt_gemm_get_algo_count( + static_cast(B.data_ptr()), B_type, + rows_b, cols_b, ldb, nullptr, trans_operation_B, + static_cast(A.data_ptr()), A_type, + rows_a, cols_a, lda, nullptr, trans_operation_A, + D_type, rows_d, cols_d, ldd, + workspace_size, false, HIPBLASLT_MATMUL_MATRIX_SCALE_END, handle); + // clang-format on + return {}; + } - const int64_t workspace_size = get_hipblaslt_workspace_size_in_byte(); + at::Tensor C = at::empty({m, n}, torch::dtype(out_dtype).device(at::kCUDA)); + auto stream = at::hip::getCurrentHIPStream(); + const hipDataType C_type = get_hipblaslt_dtype(C.scalar_type()); at::Tensor workspace = at::empty({workspace_size}, torch::dtype(at::kByte).device(at::kCUDA)); // clang-format off @@ -92,15 +105,17 @@ at::Tensor hipblaslt_gemm(at::Tensor A, at::Tensor B, const at::ScalarType out_d static_cast(workspace.data_ptr()), workspace_size, false, HIPBLASLT_MATMUL_MATRIX_SCALE_END, - handle, stream); + handle, stream, static_cast(algo_index)); // clang-format on return C; } -at::Tensor hipblaslt_gemm_fp8(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, - at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, - bool transB, bool transC, const std::string &granularity) { +static at::Tensor hipblaslt_gemm_fp8_core(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, + const std::string &granularity, int64_t algo_index = 0, + int *out_algo_count = nullptr) { const bool use_fp8 = is_8bit_floating_point_dtype(A.scalar_type()) && is_8bit_floating_point_dtype(B.scalar_type()); @@ -187,18 +202,30 @@ at::Tensor hipblaslt_gemm_fp8(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, PRIMUS_TURBO_CHECK(scaleB_inv.dim() == 2, "Scale B must be a 2-D tensor."); } - at::Tensor C = at::empty({m, n}, torch::dtype(out_dtype).device(at::kCUDA)); - - auto stream = at::cuda::getCurrentCUDAStream(); - auto handle = at::cuda::getCurrentCUDABlasLtHandle(); - + auto handle = at::cuda::getCurrentCUDABlasLtHandle(); hipblasOperation_t trans_operation_A = transA ? HIPBLAS_OP_T : HIPBLAS_OP_N; hipblasOperation_t trans_operation_B = transB ? HIPBLAS_OP_T : HIPBLAS_OP_N; const hipDataType A_type = get_hipblaslt_dtype(A.scalar_type()); const hipDataType B_type = get_hipblaslt_dtype(B.scalar_type()); - const hipDataType C_type = get_hipblaslt_dtype(C.scalar_type()); + const hipDataType D_type = get_hipblaslt_dtype(out_dtype); + const int64_t workspace_size = get_hipblaslt_workspace_size_in_byte(); + + if (out_algo_count) { + // clang-format off + *out_algo_count = hipblaslt_gemm_get_algo_count( + static_cast(B.data_ptr()), B_type, + rows_b, cols_b, ldb, nullptr, trans_operation_B, + static_cast(A.data_ptr()), A_type, + rows_a, cols_a, lda, nullptr, trans_operation_A, + D_type, rows_d, cols_d, ldd, + workspace_size, false, HIPBLASLT_MATMUL_MATRIX_SCALE_END, handle); + // clang-format on + return {}; + } - const int64_t workspace_size = get_hipblaslt_workspace_size_in_byte(); + at::Tensor C = at::empty({m, n}, torch::dtype(out_dtype).device(at::kCUDA)); + auto stream = at::hip::getCurrentHIPStream(); + const hipDataType C_type = get_hipblaslt_dtype(C.scalar_type()); at::Tensor workspace = at::empty({workspace_size}, torch::dtype(at::kByte).device(at::kCUDA)); // clang-format off @@ -224,9 +251,11 @@ at::Tensor hipblaslt_gemm_fp8(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, return C; } -at::Tensor hipblaslt_gemm_fp4(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, - at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, - bool transB, bool transC, const std::string &granularity) { +static at::Tensor hipblaslt_gemm_fp4_core(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, + const std::string &granularity, int64_t algo_index = 0, + int *out_algo_count = nullptr) { const bool use_fp4 = is_4bit_floating_point_dtype(A.scalar_type()) && is_4bit_floating_point_dtype(B.scalar_type()); @@ -310,18 +339,30 @@ at::Tensor hipblaslt_gemm_fp4(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, PRIMUS_TURBO_CHECK(scaleB_inv.dim() == 2, "Scale B must be a 2-D tensor."); } - at::Tensor C = at::empty({m, n}, torch::dtype(out_dtype).device(at::kCUDA)); - - auto stream = at::cuda::getCurrentCUDAStream(); - auto handle = at::cuda::getCurrentCUDABlasLtHandle(); - + auto handle = at::cuda::getCurrentCUDABlasLtHandle(); hipblasOperation_t trans_operation_A = transA ? HIPBLAS_OP_T : HIPBLAS_OP_N; hipblasOperation_t trans_operation_B = transB ? HIPBLAS_OP_T : HIPBLAS_OP_N; const hipDataType A_type = get_hipblaslt_dtype(A.scalar_type()); const hipDataType B_type = get_hipblaslt_dtype(B.scalar_type()); - const hipDataType C_type = get_hipblaslt_dtype(C.scalar_type()); + const hipDataType D_type = get_hipblaslt_dtype(out_dtype); + const int64_t workspace_size = get_hipblaslt_workspace_size_in_byte(); + + if (out_algo_count) { + // clang-format off + *out_algo_count = hipblaslt_gemm_get_algo_count( + static_cast(B.data_ptr()), B_type, + rows_b, cols_b, ldb, nullptr, trans_operation_B, + static_cast(A.data_ptr()), A_type, + rows_a, cols_a, lda, nullptr, trans_operation_A, + D_type, rows_d, cols_d, ldd, + workspace_size, false, HIPBLASLT_MATMUL_MATRIX_SCALE_END, handle); + // clang-format on + return {}; + } - const int64_t workspace_size = get_hipblaslt_workspace_size_in_byte(); + at::Tensor C = at::empty({m, n}, torch::dtype(out_dtype).device(at::kCUDA)); + auto stream = at::hip::getCurrentHIPStream(); + const hipDataType C_type = get_hipblaslt_dtype(C.scalar_type()); at::Tensor workspace = at::empty({workspace_size}, torch::dtype(at::kByte).device(at::kCUDA)); // clang-format off @@ -341,10 +382,60 @@ at::Tensor hipblaslt_gemm_fp4(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, static_cast(workspace.data_ptr()), workspace_size, use_fp4, scale_mode, - handle, stream); + handle, stream, static_cast(algo_index)); // clang-format on return C; } +// ---- Public API: thin wrappers around *_core ---- + +at::Tensor hipblaslt_gemm(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, bool transA, + bool transB, bool transC, int64_t algo_index) { + return hipblaslt_gemm_core(A, B, out_dtype, transA, transB, transC, algo_index); +} + +int64_t hipblaslt_gemm_algo_count(at::Tensor A, at::Tensor B, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC) { + int count = 0; + hipblaslt_gemm_core(A, B, out_dtype, transA, transB, transC, 0, &count); + return count; +} + +at::Tensor hipblaslt_gemm_fp8(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, + bool transB, bool transC, const std::string &granularity, + int64_t algo_index) { + return hipblaslt_gemm_fp8_core(A, scaleA_inv, B, scaleB_inv, out_dtype, transA, transB, transC, + granularity, algo_index); +} + +int64_t hipblaslt_gemm_fp8_algo_count(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, + const std::string &granularity) { + int count = 0; + hipblaslt_gemm_fp8_core(A, scaleA_inv, B, scaleB_inv, out_dtype, transA, transB, transC, + granularity, 0, &count); + return count; +} + +at::Tensor hipblaslt_gemm_fp4(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA, + bool transB, bool transC, const std::string &granularity, + int64_t algo_index) { + return hipblaslt_gemm_fp4_core(A, scaleA_inv, B, scaleB_inv, out_dtype, transA, transB, transC, + granularity, algo_index); +} + +int64_t hipblaslt_gemm_fp4_algo_count(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, + at::Tensor scaleB_inv, const at::ScalarType out_dtype, + bool transA, bool transB, bool transC, + const std::string &granularity) { + int count = 0; + hipblaslt_gemm_fp4_core(A, scaleA_inv, B, scaleB_inv, out_dtype, transA, transB, transC, + granularity, 0, &count); + return count; +} + } // namespace primus_turbo::pytorch diff --git a/docs/examples.md b/docs/examples.md index dac810863..461dfe906 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -490,16 +490,31 @@ GlobalBackendManager.set_gemm_backend(None) ### 5.2 AutoTune -AutoTune profiles compatible backends on the first call (per input "key", usually shape/dtype/layout) and caches the fastest choice. This is useful for stable shapes, but may add overhead if shapes change frequently. +AutoTune profiles compatible backends on the first call (per input "key", usually shape/dtype/layout) and caches the fastest choice. Two levels are available: + +| Level | Description | +|-------|-------------| +| **0** | Disabled (default). | +| **1** | Backend selection — profiles each backend (hipBLASLt, Triton, CK, ...) with its default algorithm and picks the fastest. | +| **2** | Backend + algorithm tuning — additionally benchmarks multiple hipBLASLt heuristic algorithms per problem shape, then picks the best backend *and* algorithm. | NOTE: AutoTune is skipped during CUDA graph capture. +NOTE: Setting an explicit backend (via `set_gemm_backend()` or `PRIMUS_TURBO_GEMM_BACKEND`) takes priority over AutoTune at any level. If you want level-2 algo tuning, do **not** set an explicit backend — let AutoTune select the backend and algorithm together. + ```python import torch import primus_turbo.pytorch as turbo from primus_turbo.pytorch.core.backend import GlobalBackendManager -GlobalBackendManager.set_auto_tune(True) # or set PRIMUS_TURBO_AUTO_TUNE=1 +# Level 1: backend selection only (same as the old True / PRIMUS_TURBO_AUTO_TUNE=1) +GlobalBackendManager.set_auto_tune(1) + +# Level 2: backend selection + hipBLASLt multi-algo tuning +GlobalBackendManager.set_auto_tune(2) + +# True maps to level 1. +GlobalBackendManager.set_auto_tune(True) # First call may be slower due to profiling; later calls hit the cache. # Use fixed shapes for best results. @@ -521,8 +536,13 @@ GlobalBackendManager.reset() You can also control backend selection and AutoTune via environment variables: ```bash +# AutoTune level: 0=off, 1=backend selection, 2=backend+algo tuning export PRIMUS_TURBO_AUTO_TUNE=1 + export PRIMUS_TURBO_GEMM_BACKEND=HIPBLASLT export PRIMUS_TURBO_GROUPED_GEMM_BACKEND=CK export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP + +# Cap the number of hipBLASLt heuristic algorithms tried during level-2 tuning (default: 50) +export PRIMUS_TURBO_HIPBLASLT_TUNE_MAX_ALGOS=50 ``` diff --git a/primus_turbo/pytorch/core/backend.py b/primus_turbo/pytorch/core/backend.py index 9a4ed25ae..58fdfc3ca 100644 --- a/primus_turbo/pytorch/core/backend.py +++ b/primus_turbo/pytorch/core/backend.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from enum import Enum, auto from functools import lru_cache -from typing import Any, Dict, Hashable, List, Optional, Type +from typing import Any, Dict, Hashable, List, Optional, Tuple, Type import torch @@ -33,6 +33,7 @@ ] _ENV_GEMM_BACKEND_KEY = "PRIMUS_TURBO_GEMM_BACKEND" +_ENV_HIPBLASLT_MAX_ALGOS_KEY = "PRIMUS_TURBO_HIPBLASLT_TUNE_MAX_ALGOS" _ENV_GROUPED_GEMM_BACKEND_KEY = "PRIMUS_TURBO_GROUPED_GEMM_BACKEND" _ENV_MOE_DISPATCH_COMBINE_BACKEND_KEY = "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND" _ENV_AUTO_TUNE_KEY = "PRIMUS_TURBO_AUTO_TUNE" @@ -71,7 +72,7 @@ class GlobalBackendManager: Priority (high to low): 1. Code settings - set_gemm_backend(), etc. 2. Environment variables - PRIMUS_TURBO_GEMM_BACKEND, etc. - 3. Auto-tune - PRIMUS_TURBO_AUTO_TUNE=1 + 3. Auto-tune - PRIMUS_TURBO_AUTO_TUNE=1 (backend selection) or =2 (backend + algo tuning) 4. Code defaults 5. Fallback: try all backends """ @@ -79,7 +80,7 @@ class GlobalBackendManager: _gemm_backend: Dict[PrecisionType, Optional[BackendType]] = None _grouped_gemm_backend: Dict[PrecisionType, Optional[BackendType]] = None _moe_dispatch_combine_backend: Dict[PrecisionType, Optional[BackendType]] = None - _auto_tune: Optional[bool] = None + _auto_tune: Optional[int] = None @classmethod @lru_cache(maxsize=32) @@ -161,9 +162,21 @@ def set_grouped_gemm_backend( cls._grouped_gemm_backend[precision] = backend @classmethod - def set_auto_tune(cls, enabled: Optional[bool]) -> None: - """Set whether auto-tune is enabled in code.""" - cls._auto_tune = enabled + def set_auto_tune(cls, level: Optional[object]) -> None: + """Set the auto-tune level. + + Args: + level: ``None`` to defer to the env var, ``False``/``0`` to disable, + ``True``/``1`` for backend selection only, + ``2`` for backend selection **plus** intra-backend algo tuning + (e.g. hipBLASLt multi-algorithm benchmark). + """ + if level is None or level is False: + cls._auto_tune = 0 if level is False else None + elif level is True: + cls._auto_tune = 1 + else: + cls._auto_tune = int(level) @classmethod def get_gemm_backend(cls, precision: PrecisionType) -> Optional[BackendType]: @@ -226,11 +239,16 @@ def get_moe_dispatch_combine_backend(cls, precision: PrecisionType) -> Optional[ return None @classmethod - def auto_tune_enabled(cls) -> bool: - """Check if auto-tune is enabled.""" + def auto_tune_level(cls) -> int: + """Return the auto-tune level (0=off, 1=backend selection, 2=backend+algo tuning).""" if cls._auto_tune is not None: return cls._auto_tune - return os.environ.get(_ENV_AUTO_TUNE_KEY, "0") == "1" + return int(os.environ.get(_ENV_AUTO_TUNE_KEY, "0")) + + @classmethod + def auto_tune_enabled(cls) -> bool: + """Check if auto-tune is enabled (level >= 1).""" + return cls.auto_tune_level() >= 1 @classmethod def reset(cls) -> None: @@ -252,6 +270,15 @@ def can_handle(**kwargs) -> bool: def execute(**kwargs): raise NotImplementedError("execute is not implemented") + @staticmethod + def get_algo_count(**kwargs) -> int: + """Number of heuristic algorithms available for the given problem. + + Returns 1 by default (single algorithm, no intra-backend tuning). + Override in backends that support multi-algo selection (e.g. hipBLASLt). + """ + return 1 + @dataclass(frozen=True) class BackendEntry: @@ -269,19 +296,27 @@ class BackendEntry: class TuneCache: - """LRU cache for storing tuned backend results.""" + """LRU cache for storing tuned backend + algo results.""" def __init__(self, capacity: int = 1024): self._capacity = capacity - self._cache: OrderedDict[Hashable, Type[KernelBackend]] = OrderedDict() + self._cache: OrderedDict[Hashable, Tuple[Type[KernelBackend], int]] = OrderedDict() def get(self, key: Hashable) -> Optional[Type[KernelBackend]]: + """Get cached backend class (ignores algo_index).""" + result = self.get_with_algo(key) + if result is not None: + return result[0] + return None + + def get_with_algo(self, key: Hashable) -> Optional[Tuple[Type[KernelBackend], int]]: + """Get cached (backend_class, algo_index) tuple.""" if key in self._cache: self._cache.move_to_end(key) return self._cache[key] return None - def put(self, key: Hashable, value: Type[KernelBackend]) -> None: + def put(self, key: Hashable, value: Type[KernelBackend], algo_index: int = -1) -> None: if key in self._cache: self._cache.move_to_end(key) elif len(self._cache) >= self._capacity: @@ -292,7 +327,7 @@ def put(self, key: Hashable, value: Type[KernelBackend]) -> None: stacklevel=2, ) self._cache.popitem(last=False) - self._cache[key] = value + self._cache[key] = (value, algo_index) def clear(self) -> None: self._cache.clear() @@ -351,7 +386,6 @@ def profile(cls, backend: Type[KernelBackend], **kwargs) -> float: start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) - # warm-up for _ in range(cls._warmup_iters): backend.execute(**kwargs) torch.cuda.synchronize() @@ -365,34 +399,86 @@ def profile(cls, backend: Type[KernelBackend], **kwargs) -> float: return start.elapsed_time(end) / cls._profile_iters + @classmethod + @torch.no_grad() + def _tune_algo_for_backend(cls, backend: Type[KernelBackend], **kwargs) -> Tuple[float, int]: + """Benchmark each heuristic algorithm for a backend (level-2 tuning). + + Returns (best_time, best_algo_index). The number of algorithms + tried is capped by the ``PRIMUS_TURBO_HIPBLASLT_TUNE_MAX_ALGOS`` + environment variable (default 10). + """ + max_algos = int(os.environ.get(_ENV_HIPBLASLT_MAX_ALGOS_KEY, "50")) + algo_count = min(backend.get_algo_count(**kwargs), max_algos) + if algo_count <= 1: + return cls.profile(backend, **kwargs), -1 + + best_time = float("inf") + best_algo = -1 + for idx in range(algo_count): + algo_kwargs = {**kwargs, "algo_index": idx} + try: + t = cls.profile(backend, **algo_kwargs) + except Exception: + t = float("inf") + if t < best_time: + best_time = t + best_algo = idx + logger.info( + f"[AutoTune L2] {backend.__name__} algo {idx}/{algo_count}: {t:.3f} ms" + f"{' (best so far)' if idx == best_algo else ''}" + ) + return best_time, best_algo + @classmethod def tune(cls, **kwargs) -> Optional[Type[KernelBackend]]: - """Profile all compatible backends and cache the fastest one.""" + """Profile compatible backends and cache the fastest. + + Level 1: pick the best *backend* (each backend uses its default algo). + Level 2: additionally benchmark multiple hipBLASLt heuristic algorithms. + """ key = cls.make_key(**kwargs) - cached_backend = cls._cache.get(key) - if cached_backend is not None: - return cached_backend + cached = cls._cache.get_with_algo(key) + if cached is not None: + backend_cls, algo_index = cached + if algo_index >= 0: + kwargs["algo_index"] = algo_index + return backend_cls + level = GlobalBackendManager.auto_tune_level() best_backend = None best_time = float("inf") + best_algo = -1 for entry in cls._backends.values(): if not entry.autotune: continue if entry.impl.can_handle(**kwargs): torch.cuda.synchronize() try: - cur_time = cls.profile(entry.impl, **kwargs) + if level >= 2: + cur_time, cur_algo = cls._tune_algo_for_backend(entry.impl, **kwargs) + else: + cur_time, cur_algo = cls.profile(entry.impl, **kwargs), -1 except Exception: - cur_time = float("inf") + cur_time, cur_algo = float("inf"), -1 finally: torch.cuda.synchronize() + logger.info( + f"[AutoTune] {entry.impl.__name__}: {cur_time:.3f} ms" + f"{f' (algo {cur_algo})' if cur_algo >= 0 else ''}" + ) if cur_time < best_time: best_time = cur_time best_backend = entry.impl + best_algo = cur_algo if best_backend is not None: - cls._cache.put(key, best_backend) + cls._cache.put(key, best_backend, best_algo) + logger.info( + f"[AutoTune] Winner for key={key}: {best_backend.__name__} " + f"({best_time:.3f} ms{f', algo {best_algo}' if best_algo >= 0 else ''})" + ) return best_backend @classmethod @@ -417,8 +503,21 @@ def dispatch( # 2. Auto tune # NOTE: Skip autotune during cuda graph capture. if GlobalBackendManager.auto_tune_enabled() and not cls._is_graph_capturing(): + key = cls.make_key(**kwargs) + cached = cls._cache.get_with_algo(key) + if cached is not None: + backend_cls, algo_index = cached + if algo_index >= 0: + return backend_cls.execute(**kwargs, algo_index=algo_index) + return backend_cls.execute(**kwargs) + backend_cls = cls.tune(**kwargs) if backend_cls is not None: + cached = cls._cache.get_with_algo(key) + if cached is not None: + _, algo_index = cached + if algo_index >= 0: + return backend_cls.execute(**kwargs, algo_index=algo_index) return backend_cls.execute(**kwargs) # 3. Default backend diff --git a/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py b/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py index 5145513be..16da50f1c 100644 --- a/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py +++ b/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py @@ -105,10 +105,35 @@ def execute( out_dtype: torch.dtype, trans_c: bool, granularity: ScalingGranularity, + **kwargs, ): + algo_index = kwargs.get("algo_index", 0) # TODO(ruibin): Add padding return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm_fp4( - a, a_scale_inv, b, b_scale_inv, out_dtype, trans_a, trans_b, trans_c, granularity.name + a, + a_scale_inv, + b, + b_scale_inv, + out_dtype, + trans_a, + trans_b, + trans_c, + granularity.name, + algo_index, + ) + + @staticmethod + def get_algo_count(**kwargs) -> int: + return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm_fp4_algo_count( + kwargs["a"], + kwargs["a_scale_inv"], + kwargs["b"], + kwargs["b_scale_inv"], + kwargs["out_dtype"], + kwargs["trans_a"], + kwargs["trans_b"], + kwargs["trans_c"], + kwargs["granularity"].name, ) diff --git a/primus_turbo/pytorch/kernels/gemm/gemm_fp8_impl.py b/primus_turbo/pytorch/kernels/gemm/gemm_fp8_impl.py index 15e13a799..17c2e8f5b 100644 --- a/primus_turbo/pytorch/kernels/gemm/gemm_fp8_impl.py +++ b/primus_turbo/pytorch/kernels/gemm/gemm_fp8_impl.py @@ -113,9 +113,34 @@ def execute( out_dtype: torch.dtype, trans_c: bool, granularity: ScalingGranularity, + **kwargs, ): + algo_index = kwargs.get("algo_index", 0) return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm_fp8( - a, a_scale_inv, b, b_scale_inv, out_dtype, trans_a, trans_b, trans_c, granularity.name + a, + a_scale_inv, + b, + b_scale_inv, + out_dtype, + trans_a, + trans_b, + trans_c, + granularity.name, + algo_index, + ) + + @staticmethod + def get_algo_count(**kwargs) -> int: + return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm_fp8_algo_count( + kwargs["a"], + kwargs["a_scale_inv"], + kwargs["b"], + kwargs["b_scale_inv"], + kwargs["out_dtype"], + kwargs["trans_a"], + kwargs["trans_b"], + kwargs["trans_c"], + kwargs["granularity"].name, ) diff --git a/primus_turbo/pytorch/kernels/gemm/gemm_impl.py b/primus_turbo/pytorch/kernels/gemm/gemm_impl.py index a43f14ae2..e90a1b341 100644 --- a/primus_turbo/pytorch/kernels/gemm/gemm_impl.py +++ b/primus_turbo/pytorch/kernels/gemm/gemm_impl.py @@ -50,7 +50,21 @@ def execute( trans_c: bool, **kwargs, ) -> torch.Tensor: - return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm(a, b, out_dtype, trans_a, trans_b, trans_c) + algo_index = kwargs.get("algo_index", 0) + return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm( + a, b, out_dtype, trans_a, trans_b, trans_c, algo_index + ) + + @staticmethod + def get_algo_count(**kwargs) -> int: + return torch.ops.primus_turbo_cpp_extension.hipblaslt_gemm_algo_count( + kwargs["a"], + kwargs["b"], + kwargs["out_dtype"], + kwargs["trans_a"], + kwargs["trans_b"], + kwargs["trans_c"], + ) class GEMMTritonBackend(KernelBackend): diff --git a/tests/pytorch/core/test_global_backend_manager.py b/tests/pytorch/core/test_global_backend_manager.py index a3b626bef..a62b52ce5 100644 --- a/tests/pytorch/core/test_global_backend_manager.py +++ b/tests/pytorch/core/test_global_backend_manager.py @@ -23,6 +23,7 @@ def clean_backend_state(monkeypatch): "PRIMUS_TURBO_GROUPED_GEMM_BACKEND", "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "PRIMUS_TURBO_AUTO_TUNE", + "PRIMUS_TURBO_HIPBLASLT_TUNE_MAX_ALGOS", ): monkeypatch.delenv(key, raising=False) yield @@ -118,6 +119,36 @@ def test_set_auto_tune(self): GlobalBackendManager.set_auto_tune(False) assert GlobalBackendManager.auto_tune_enabled() is False + def test_set_auto_tune_levels(self): + """Level 0=off, 1=backend selection, 2=backend+algo tuning.""" + GlobalBackendManager.set_auto_tune(0) + assert GlobalBackendManager.auto_tune_level() == 0 + assert GlobalBackendManager.auto_tune_enabled() is False + + GlobalBackendManager.set_auto_tune(1) + assert GlobalBackendManager.auto_tune_level() == 1 + assert GlobalBackendManager.auto_tune_enabled() is True + + GlobalBackendManager.set_auto_tune(2) + assert GlobalBackendManager.auto_tune_level() == 2 + assert GlobalBackendManager.auto_tune_enabled() is True + + def test_set_auto_tune_backward_compat(self): + """True maps to level 1, False to 0, None defers to env.""" + GlobalBackendManager.set_auto_tune(True) + assert GlobalBackendManager.auto_tune_level() == 1 + + GlobalBackendManager.set_auto_tune(False) + assert GlobalBackendManager.auto_tune_level() == 0 + + GlobalBackendManager.set_auto_tune(None) + assert GlobalBackendManager.auto_tune_level() == 0 # env default is "0" + + def test_auto_tune_env_level_2(self, monkeypatch): + monkeypatch.setenv("PRIMUS_TURBO_AUTO_TUNE", "2") + assert GlobalBackendManager.auto_tune_level() == 2 + assert GlobalBackendManager.auto_tune_enabled() is True + def test_reset_clears_code_settings(self): self._init_gemm_backend() GlobalBackendManager.set_gemm_backend(BackendType.CK, PrecisionType.FP8) diff --git a/tests/pytorch/modules/test_grouped_linear.py b/tests/pytorch/modules/test_grouped_linear.py index 17052fe9e..546304a8a 100644 --- a/tests/pytorch/modules/test_grouped_linear.py +++ b/tests/pytorch/modules/test_grouped_linear.py @@ -21,7 +21,7 @@ @pytest.mark.parametrize("N_K", [(1024, 2048), (512, 1024)]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("balance", [True, False]) -@pytest.mark.parametrize("auto_tune", [False, True]) +@pytest.mark.parametrize("auto_tune", [False, True, 2]) @pytest.mark.parametrize("enable_torch_compile", [True, False]) def test_grouped_linear(B, M, N_K, dtype, balance, auto_tune, enable_torch_compile): GlobalBackendManager.set_auto_tune(auto_tune) diff --git a/tests/pytorch/ops/test_gemm_fp4.py b/tests/pytorch/ops/test_gemm_fp4.py index ce3ad3eed..78aaed923 100644 --- a/tests/pytorch/ops/test_gemm_fp4.py +++ b/tests/pytorch/ops/test_gemm_fp4.py @@ -39,7 +39,7 @@ ) @pytest.mark.parametrize("granularity", [ScalingGranularity.MX_BLOCKWISE]) @pytest.mark.parametrize("backend", [None, BackendType.HIPBLASLT, BackendType.AITER]) -@pytest.mark.parametrize("auto_tune", [False, True]) +@pytest.mark.parametrize("auto_tune", [False, True, 2]) def test_gemm_fp4_mx_blockwise(m, n, k, layout, format, dtype, granularity, backend, auto_tune): if backend == BackendType.AITER: if dtype != torch.bfloat16: diff --git a/tests/pytorch/ops/test_gemm_fp8.py b/tests/pytorch/ops/test_gemm_fp8.py index 19920258a..492fefa59 100644 --- a/tests/pytorch/ops/test_gemm_fp8.py +++ b/tests/pytorch/ops/test_gemm_fp8.py @@ -207,7 +207,7 @@ def _run_once(): @pytest.mark.parametrize("format", [Format.E4M3, Format.E5M2, Format.HYBRID]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("backend", [None, BackendType.TRITON, BackendType.CK, BackendType.HIPBLASLT]) -@pytest.mark.parametrize("auto_tune", [False, True]) +@pytest.mark.parametrize("auto_tune", [False, True, 2]) def test_gemm_fp8_tensorwise(m, n, k, layout, format, dtype, backend, auto_tune): if backend == BackendType.TRITON and format == Format.HYBRID: pytest.skip("TRITON backend does not support HYBRID format currently.") @@ -232,7 +232,7 @@ def test_gemm_fp8_tensorwise(m, n, k, layout, format, dtype, backend, auto_tune) @pytest.mark.parametrize("format", [Format.E4M3, Format.E5M2]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("backend", [None, BackendType.TRITON, BackendType.CK]) -@pytest.mark.parametrize("auto_tune", [False, True]) +@pytest.mark.parametrize("auto_tune", [False, True, 2]) def test_gemm_fp8_rowwise(m, n, k, layout, format, dtype, backend, auto_tune): _run_gemm_fp8_test( m=m, @@ -255,7 +255,7 @@ def test_gemm_fp8_rowwise(m, n, k, layout, format, dtype, backend, auto_tune): @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("block_size", [128]) @pytest.mark.parametrize("backend", [None, BackendType.TRITON, BackendType.CK]) -@pytest.mark.parametrize("auto_tune", [False, True]) +@pytest.mark.parametrize("auto_tune", [False, True, 2]) def test_gemm_fp8_blockwise(m, n, k, layout, format, dtype, block_size, backend, auto_tune): _run_gemm_fp8_test( m=m, @@ -278,7 +278,7 @@ def test_gemm_fp8_blockwise(m, n, k, layout, format, dtype, block_size, backend, @pytest.mark.parametrize("format", [Format.E4M3, Format.E5M2, Format.HYBRID]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("backend", [None, BackendType.HIPBLASLT, BackendType.TURBO]) -@pytest.mark.parametrize("auto_tune", [False, True]) +@pytest.mark.parametrize("auto_tune", [False, True, 2]) def test_gemm_fp8_mx_blockwise(m, n, k, layout, format, dtype, backend, auto_tune): # NOTE: m, n and k must be multiples of 16 for MX_BLOCKWISE. assert m % 16 == 0 and n % 16 == 0 and k % 16 == 0, "m, n and k must be multiples of 16"