diff --git a/benchmark/ops/bench_grouped_gemm_gmm.py b/benchmark/ops/bench_grouped_gemm_gmm.py new file mode 100644 index 000000000..856ef858d --- /dev/null +++ b/benchmark/ops/bench_grouped_gemm_gmm.py @@ -0,0 +1,205 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""GMM (grouped_gemm) BF16 Grouped GEMM Baseline Benchmark. + +Reference: https://github.com/tgale96/grouped_gemm +""" + +import argparse +from datetime import datetime + +import grouped_gemm.ops as gmm_ops +import pandas as pd +import torch +import torch.utils.benchmark as benchmark +from config import ( + check_allclose, + gen_grouped_gemm_group_lens, + gen_grouped_gemm_test_cases, + get_platform_info, + grouped_gemm_ref, +) +from tabulate import tabulate + + +def check_grouped_gemm_gmm_correctness(a, b, batch_sizes, grad_out, dtype): + """Check correctness of gmm forward and backward against reference.""" + out = gmm_ops.gmm(a, b, batch_sizes, trans_b=True) + + group_lens = batch_sizes.to(device=a.device, dtype=torch.int64) + out_ref = grouped_gemm_ref(a.detach(), b.detach(), group_lens, trans_b=True) + fwd_correct = check_allclose(out.detach(), out_ref, dtype) + + a_ref = a.detach().clone().requires_grad_() + b_ref = b.detach().clone().requires_grad_() + out_ref = grouped_gemm_ref(a_ref, b_ref, group_lens, trans_b=True) + out_ref.backward(grad_out) + out.backward(grad_out, retain_graph=True) + bwd_correct = check_allclose(a.grad, a_ref.grad, dtype) + + a.grad = None + + correct = fwd_correct and bwd_correct + status = "PASS" if correct else "FAIL" + print(f"Correctness Check: {status} (fwd={fwd_correct}, bwd={bwd_correct})") + + return correct + + +def profile_grouped_gemm_gmm(B, M, N, K, dtype, balance=True, num_topk=None): + """Profile BF16 Grouped GEMM using grouped_gemm (gmm) library.""" + device = "cuda" + # batch_sizes must be on CPU for gmm + batch_sizes = gen_grouped_gemm_group_lens(B, M, balance=balance, num_topk=num_topk) + total_m = batch_sizes.sum().item() + + a = torch.randn((total_m, K), dtype=dtype, device=device, requires_grad=True) + b = torch.randn((B, N, K), dtype=dtype, device=device, requires_grad=True) + + out = gmm_ops.gmm(a, b, batch_sizes, trans_b=True) + grad_out = torch.randn_like(out) + + correct = check_grouped_gemm_gmm_correctness(a, b, batch_sizes, grad_out, dtype) + + def fwd_func(): + return gmm_ops.gmm(a, b, batch_sizes, trans_b=True) + + def fwd_bwd_func(): + out = gmm_ops.gmm(a, b, batch_sizes, trans_b=True) + out.backward(grad_out) + a.grad = None + b.grad = None + + fwd_total_flops = 2 * B * M * N * K + bwd_total_flops = 2 * fwd_total_flops + + for _ in range(100): + fwd_bwd_func() + torch.cuda.synchronize() + + fwd_timer = benchmark.Timer(stmt="fn()", globals={"fn": fwd_func}) + fwd_measurement = fwd_timer.timeit(200) + + fwd_bwd_timer = benchmark.Timer(stmt="fn()", globals={"fn": fwd_bwd_func}) + fwd_bwd_measurement = fwd_bwd_timer.timeit(200) + + fwd_mean_time_ms = fwd_measurement.mean * 1e3 + fwd_bwd_mean_time_ms = fwd_bwd_measurement.mean * 1e3 + bwd_mean_time_ms = fwd_bwd_mean_time_ms - fwd_mean_time_ms + + fwd_tflops = fwd_total_flops / (fwd_mean_time_ms * 1e-3) / 1e12 + bwd_tflops = bwd_total_flops / (bwd_mean_time_ms * 1e-3) / 1e12 + print(f"Forward Mean time: {fwd_mean_time_ms:.3f} ms | TFLOPS: {fwd_tflops:.2f}") + print(f"Backward Mean time: {bwd_mean_time_ms:.3f} ms | TFLOPS: {bwd_tflops:.2f}") + + return fwd_mean_time_ms, fwd_tflops, bwd_mean_time_ms, bwd_tflops, correct + + +def benchmark_grouped_gemm_gmm(output_csv=None): + platform, gpu_name = get_platform_info() + + test_cases = gen_grouped_gemm_test_cases() + + rows = [] + test_id = 0 + dtype = torch.bfloat16 + + for case in test_cases: + test_id += 1 + B, M, N, K = case["B"], case["M"], case["N"], case["K"] + + balance = case.get("balance", True) + balance_str = "balanced" if balance else "unbalanced" + print(f"\n{'='*60}") + print( + f"TestID: {test_id}, Case: {case['Case']}, B: {B}, M: {M}, N: {N}, K: {K}, dtype: bf16, {balance_str}" + ) + print(f"{'='*60}") + + try: + fwd_time_ms, fwd_tflops, bwd_time_ms, bwd_tflops, correct = profile_grouped_gemm_gmm( + B=B, + M=M, + N=N, + K=K, + dtype=dtype, + balance=balance, + num_topk=case.get("num_topk"), + ) + + row = { + "TestID": test_id, + "Platform": platform, + "GPU": gpu_name, + "Case": case["Case"], + "B": B, + "M": M, + "N": N, + "K": K, + "Dtype": "bf16", + "Balance": "Y" if balance else "N", + "Check": "PASS" if correct else "FAIL", + "Forward Time (ms)": f"{fwd_time_ms:.2f}", + "Forward TFLOPS": f"{fwd_tflops:.2f}", + "Backward Time (ms)": f"{bwd_time_ms:.2f}", + "Backward TFLOPS": f"{bwd_tflops:.2f}", + } + rows.append(row) + + except Exception as e: + import traceback + + print(f"Failed: {str(e)}") + traceback.print_exc() + row = { + "TestID": test_id, + "Platform": platform, + "GPU": gpu_name, + "Case": case["Case"], + "B": B, + "M": M, + "N": N, + "K": K, + "Dtype": "bf16", + "Balance": "Y" if balance else "N", + "Check": "ERROR", + "Forward Time (ms)": "ERROR", + "Forward TFLOPS": "0.00", + "Backward Time (ms)": "ERROR", + "Backward TFLOPS": "0.00", + } + rows.append(row) + + results = pd.DataFrame(rows) + print("\nFinal Results:") + print(tabulate(results, headers="keys", tablefmt="grid", showindex=False)) + + avg_fwd_tflops = results["Forward TFLOPS"].astype(float).mean() + avg_bwd_tflops = results["Backward TFLOPS"].astype(float).mean() + print(f"\nAverage Forward TFLOPS: {avg_fwd_tflops:.2f}") + print(f"Average Backward TFLOPS: {avg_bwd_tflops:.2f}") + + if output_csv: + filename = output_csv + else: + timestamp = datetime.now().strftime("%Y%m%d") + filename = f"grouped_gemm_gmm_bf16_{timestamp}_{gpu_name}.csv" + results.to_csv(filename, index=False) + print(f"Results saved to {filename}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Benchmark grouped_gemm (gmm) BF16 Grouped GEMM") + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="Output CSV filename", + ) + args = parser.parse_args() + benchmark_grouped_gemm_gmm(output_csv=args.output) diff --git a/benchmark/ops/bench_grouped_gemm_turbo.py b/benchmark/ops/bench_grouped_gemm_turbo.py index 5692dc3ac..6f7833c7f 100644 --- a/benchmark/ops/bench_grouped_gemm_turbo.py +++ b/benchmark/ops/bench_grouped_gemm_turbo.py @@ -90,12 +90,13 @@ def check_grouped_gemm_fp8_correctness(a, b, out, grad_out, group_lens, fp8_form return correct -def profile_grouped_gemm(B, M, N, K, dtype): +def profile_grouped_gemm(B, M, N, K, dtype, balance=True): """Profile BF16 Grouped GEMM.""" device = "cuda" - x = torch.randn((B * M, K), dtype=dtype, device=device, requires_grad=True) + group_lens = gen_grouped_gemm_group_lens(B, M, balance=balance).to(device) + total_m = group_lens.sum().item() + x = torch.randn((total_m, K), dtype=dtype, device=device, requires_grad=True) w = torch.randn((B, N, K), dtype=dtype, device=device, requires_grad=True) - group_lens = gen_grouped_gemm_group_lens(B, M, balance=True).to(device) out = turbo.ops.grouped_gemm(x, w, group_lens, trans_b=True) grad_out = torch.randn_like(out) @@ -109,15 +110,15 @@ def profile_grouped_gemm(B, M, N, K, dtype): fwd_total_flops = 2 * B * M * N * K bwd_total_flops = 2 * fwd_total_flops - for _ in range(20): + for _ in range(100): fwd_func() bwd_func() torch.cuda.synchronize() fwd_timer = benchmark.Timer(stmt="fn()", globals={"fn": fwd_func}) bwd_timer = benchmark.Timer(stmt="fn()", globals={"fn": bwd_func}) - fwd_measurement = fwd_timer.timeit(100) - bwd_measurement = bwd_timer.timeit(100) + fwd_measurement = fwd_timer.timeit(200) + bwd_measurement = bwd_timer.timeit(200) fwd_mean_time_ms = fwd_measurement.mean * 1e3 bwd_mean_time_ms = bwd_measurement.mean * 1e3 @@ -148,15 +149,15 @@ def profile_grouped_gemm_fp8(B, M, N, K, dtype, config): fwd_total_flops = 2 * B * M * N * K bwd_total_flops = 2 * fwd_total_flops - for _ in range(20): + for _ in range(100): fwd_func() bwd_func() torch.cuda.synchronize() fwd_timer = benchmark.Timer(stmt="fn()", globals={"fn": fwd_func}) bwd_timer = benchmark.Timer(stmt="fn()", globals={"fn": bwd_func}) - fwd_measurement = fwd_timer.timeit(100) - bwd_measurement = bwd_timer.timeit(100) + fwd_measurement = fwd_timer.timeit(200) + bwd_measurement = bwd_timer.timeit(200) fwd_mean_time_ms = fwd_measurement.mean * 1e3 bwd_mean_time_ms = bwd_measurement.mean * 1e3 @@ -168,7 +169,9 @@ def profile_grouped_gemm_fp8(B, M, N, K, dtype, config): return fwd_mean_time_ms, fwd_tflops, bwd_mean_time_ms, bwd_tflops, correct -def benchmark_grouped_gemm_turbo(dtype_name="bf16", granularity_name="tensorwise", output_csv=None): +def benchmark_grouped_gemm_turbo( + dtype_name="bf16", granularity_name="tensorwise", output_csv=None, balance=True +): platform, gpu_name = get_platform_info() is_fp8 = dtype_name == "fp8" @@ -183,14 +186,17 @@ def benchmark_grouped_gemm_turbo(dtype_name="bf16", granularity_name="tensorwise B, M, N, K = case["B"], case["M"], case["N"], case["K"] dtype = case["dtype"] + balance_str = "balanced" if balance else "unbalanced" print(f"\n{'='*60}") if is_fp8: print( f"TestID: {test_id}, Case: {case['Case']}, B: {B}, M: {M}, N: {N}, K: {K}, " - f"dtype: fp8, granularity: {granularity_name}" + f"dtype: fp8, granularity: {granularity_name}, {balance_str}" ) else: - print(f"TestID: {test_id}, Case: {case['Case']}, B: {B}, M: {M}, N: {N}, K: {K}, dtype: bf16") + print( + f"TestID: {test_id}, Case: {case['Case']}, B: {B}, M: {M}, N: {N}, K: {K}, dtype: bf16, {balance_str}" + ) print(f"{'='*60}") try: @@ -200,7 +206,12 @@ def benchmark_grouped_gemm_turbo(dtype_name="bf16", granularity_name="tensorwise ) else: fwd_time_ms, fwd_tflops, bwd_time_ms, bwd_tflops, correct = profile_grouped_gemm( - B=B, M=M, N=N, K=K, dtype=dtype + B=B, + M=M, + N=N, + K=K, + dtype=dtype, + balance=balance, ) row = { @@ -213,6 +224,7 @@ def benchmark_grouped_gemm_turbo(dtype_name="bf16", granularity_name="tensorwise "N": N, "K": K, "Dtype": dtype_name, + "Balance": "Y" if balance else "N", } if is_fp8: row["Granularity"] = granularity_name @@ -239,6 +251,7 @@ def benchmark_grouped_gemm_turbo(dtype_name="bf16", granularity_name="tensorwise "N": N, "K": K, "Dtype": dtype_name, + "Balance": "Y" if balance else "N", } if is_fp8: row["Granularity"] = granularity_name @@ -297,7 +310,17 @@ def benchmark_grouped_gemm_turbo(dtype_name="bf16", granularity_name="tensorwise default=None, help="Output CSV filename", ) + parser.add_argument( + "--balance", + type=str, + choices=["balanced", "unbalanced"], + default="balanced", + help="Routing balance mode (default: balanced)", + ) args = parser.parse_args() benchmark_grouped_gemm_turbo( - dtype_name=args.dtype, granularity_name=args.granularity, output_csv=args.output + dtype_name=args.dtype, + granularity_name=args.granularity, + output_csv=args.output, + balance=(args.balance == "balanced"), ) diff --git a/benchmark/ops/config.py b/benchmark/ops/config.py index 58c113c99..985b0eaec 100644 --- a/benchmark/ops/config.py +++ b/benchmark/ops/config.py @@ -295,6 +295,29 @@ def grouped_gemm_ref(a, b, group_lens, trans_b=True): # "head_dim": 128, # "seqlen": 4096, # }, + # # https://huggingface.co/LiquidAI/LFM2-8B-A1B + # "LFM2-8B-A1B": { + # "n_routed_experts": 32, + # "moe_intermediate_size": 1792, + # "hidden_size": 2048, + # # GQA attention config + # "num_attention_heads": 32, + # "num_key_value_heads": 8, + # "head_dim": 128, # kv_channels + # "seqlen": 4096, + # }, + # # /shared_nfs/kyle/test/Primus/examples/megatron/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml + # "gpt_oss_20B": { + # "n_routed_experts": 32, + # "moe_intermediate_size": 2880, + # "hidden_size": 2880, + # # GQA attention config + # "num_attention_heads": 64, + # "num_key_value_heads": 8, + # "head_dim_qk": 128, # qk_head_dim + # "head_dim_v": 64, # kv_channels + # "seqlen": 4096, + # }, } ############################################################################### diff --git a/csrc/include/primus_turbo/gemm.h b/csrc/include/primus_turbo/gemm.h index a02b4395e..54df37dff 100644 --- a/csrc/include/primus_turbo/gemm.h +++ b/csrc/include/primus_turbo/gemm.h @@ -42,6 +42,8 @@ inline size_t hipblaslt_dtype_bytes(hipDataType dtype) { int64_t get_hipblaslt_workspace_size_in_byte(); +void clear_hipblaslt_gemm_runtime_caches(); + 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, diff --git a/csrc/include/primus_turbo/grouped_gemm.h b/csrc/include/primus_turbo/grouped_gemm.h index 0a864315a..1c11a3677 100644 --- a/csrc/include/primus_turbo/grouped_gemm.h +++ b/csrc/include/primus_turbo/grouped_gemm.h @@ -17,6 +17,8 @@ std::int64_t get_ck_grouped_gemm_fp8_args_sizes(const int group_num); std::int64_t get_hipblaslt_grouped_gemm_workspace_size(); +void clear_hipblaslt_grouped_gemm_runtime_state(); + //================================================================== // Grouped GEMM Params //================================================================== @@ -67,13 +69,14 @@ struct HipblasltGroupedGemmParams { hipDataType c_type; std::vector c_shape; - const int64_t *group_lens_ptr = nullptr; - const int64_t *group_offs_ptr = nullptr; - bool transA = false; - bool transB = false; - int32_t group_num = 0; - hipStream_t stream = nullptr; - void *workspace = nullptr; + const int64_t *group_lens_ptr = nullptr; + const int64_t *group_offs_ptr = nullptr; + bool group_lens_on_host = false; + bool transA = false; + bool transB = false; + int32_t group_num = 0; + hipStream_t stream = nullptr; + void *workspace = nullptr; bool use_low_precision = false; diff --git a/csrc/kernels/gemm/hipblaslt_gemm.cu b/csrc/kernels/gemm/hipblaslt_gemm.cu index 36cd0c3bb..66ebd8248 100644 --- a/csrc/kernels/gemm/hipblaslt_gemm.cu +++ b/csrc/kernels/gemm/hipblaslt_gemm.cu @@ -2,6 +2,8 @@ // // See LICENSE for license information. +#include + #include "primus_turbo/common.h" #include "primus_turbo/gemm.h" @@ -18,6 +20,176 @@ int64_t get_hipblaslt_workspace_size_in_byte() { } } +// ── algo cache ─────────────────────────────────────────────────────────────── +// hipblasLtMatmulAlgoGetHeuristic() can cost 100 ms+ on first encounter of a +// new shape. Cache the winning algo per (shape, dtype, trans, handle) tuple +// so subsequent calls for the same problem are free. + +// ── Algo cache key: optimistic M-invariant fast path ───────────────────────── +// Reusing an algo across different M values avoids re-autotuning every expert +// shape in grouped GEMM. Some BF16 hipBLASLt configs still require an exact-M +// override, so an exact-shape cache is checked first and populated on fallback. +struct AlgoKey { + int64_t cols_a; // K (M-invariant) + int64_t rows_b, cols_b, ldb; // weight matrix shape (fixed per layer) + int64_t cols_d; // N (M-invariant) + hipDataType A_type, B_type, D_type; + hipblasOperation_t transA, transB; + bool use_low_precision; + hipblasLtMatmulMatrixScale_t scale_mode; + hipblasLtHandle_t handle; + + bool operator==(const AlgoKey &o) const { + return cols_a == o.cols_a && rows_b == o.rows_b && cols_b == o.cols_b && ldb == o.ldb && + cols_d == o.cols_d && A_type == o.A_type && B_type == o.B_type && + D_type == o.D_type && transA == o.transA && transB == o.transB && + use_low_precision == o.use_low_precision && scale_mode == o.scale_mode && + handle == o.handle; + } +}; + +struct AlgoKeyHash { + size_t operator()(const AlgoKey &k) const { + size_t s = 0; + auto hc = [&s](auto v) { + s ^= std::hash{}(v) + 0x9e3779b9u + (s << 6) + (s >> 2); + }; + hc(k.cols_a); + hc(k.rows_b); + hc(k.cols_b); + hc(k.ldb); + hc(k.cols_d); + hc(static_cast(k.A_type)); + hc(static_cast(k.B_type)); + hc(static_cast(k.D_type)); + hc(static_cast(k.transA)); + hc(static_cast(k.transB)); + hc(k.use_low_precision); + hc(static_cast(k.scale_mode)); + hc(reinterpret_cast(k.handle)); + return s; + } +}; + +static thread_local std::unordered_map + algo_cache; + +// ── Descriptor cache key: includes full shape (M-dependent) ───────────────── +// Matrix layouts must match exact dimensions, so descriptors are cached per +// full shape. Same-shape experts (balanced routing) still get cache hits. +struct DescKey { + int64_t rows_a, cols_a, lda; + int64_t rows_b, cols_b, ldb; + int64_t rows_d, cols_d, ldd; + hipDataType A_type, B_type, D_type; + hipblasOperation_t transA, transB; + bool use_low_precision; + hipblasLtMatmulMatrixScale_t scale_mode; + hipblasLtHandle_t handle; + + bool operator==(const DescKey &o) const { + return rows_a == o.rows_a && cols_a == o.cols_a && lda == o.lda && rows_b == o.rows_b && + cols_b == o.cols_b && ldb == o.ldb && rows_d == o.rows_d && cols_d == o.cols_d && + ldd == o.ldd && A_type == o.A_type && B_type == o.B_type && D_type == o.D_type && + transA == o.transA && transB == o.transB && + use_low_precision == o.use_low_precision && scale_mode == o.scale_mode && + handle == o.handle; + } +}; + +struct DescKeyHash { + size_t operator()(const DescKey &k) const { + size_t s = 0; + auto hc = [&s](auto v) { + s ^= std::hash{}(v) + 0x9e3779b9u + (s << 6) + (s >> 2); + }; + hc(k.rows_a); + hc(k.cols_a); + hc(k.lda); + hc(k.rows_b); + hc(k.cols_b); + hc(k.ldb); + hc(k.rows_d); + hc(k.cols_d); + hc(k.ldd); + hc(static_cast(k.A_type)); + hc(static_cast(k.B_type)); + hc(static_cast(k.D_type)); + hc(static_cast(k.transA)); + hc(static_cast(k.transB)); + hc(k.use_low_precision); + hc(static_cast(k.scale_mode)); + hc(reinterpret_cast(k.handle)); + return s; + } +}; + +struct DescCache { + hipblasLtMatrixLayout_t A_desc = nullptr; + hipblasLtMatrixLayout_t B_desc = nullptr; + hipblasLtMatrixLayout_t D_desc = nullptr; + hipblasLtMatmulDesc_t op_desc = nullptr; + + DescCache() = default; + DescCache(const DescCache &) = delete; + DescCache &operator=(const DescCache &) = delete; + + DescCache(DescCache &&other) noexcept + : A_desc(other.A_desc), B_desc(other.B_desc), D_desc(other.D_desc), op_desc(other.op_desc) { + other.A_desc = nullptr; + other.B_desc = nullptr; + other.D_desc = nullptr; + other.op_desc = nullptr; + } + + DescCache &operator=(DescCache &&other) noexcept { + if (this != &other) { + reset(); + A_desc = other.A_desc; + B_desc = other.B_desc; + D_desc = other.D_desc; + op_desc = other.op_desc; + other.A_desc = nullptr; + other.B_desc = nullptr; + other.D_desc = nullptr; + other.op_desc = nullptr; + } + return *this; + } + + ~DescCache() { reset(); } + + void reset() { + if (A_desc != nullptr) { + (void) hipblasLtMatrixLayoutDestroy(A_desc); + A_desc = nullptr; + } + if (B_desc != nullptr) { + (void) hipblasLtMatrixLayoutDestroy(B_desc); + B_desc = nullptr; + } + if (D_desc != nullptr) { + (void) hipblasLtMatrixLayoutDestroy(D_desc); + D_desc = nullptr; + } + if (op_desc != nullptr) { + (void) hipblasLtMatmulDescDestroy(op_desc); + op_desc = nullptr; + } + } +}; + +static thread_local std::unordered_map desc_cache; +static thread_local std::unordered_map + exact_algo_cache; +// ───────────────────────────────────────────────────────────────────────────── + +void clear_hipblaslt_gemm_runtime_caches() { + algo_cache.clear(); + exact_algo_cache.clear(); + desc_cache.clear(); +} + 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, @@ -27,85 +199,200 @@ 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, hipblasLtMatmulMatrixScale_t scale_mode, hipblasLtHandle_t handle, hipStream_t stream) { - 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))); + // Look up or create cached descriptors for this exact shape. + DescKey dkey{rows_a, cols_a, lda, rows_b, cols_b, + ldb, rows_d, cols_d, ldd, A_type, + B_type, D_type, transA, transB, use_low_precision, + scale_mode, handle}; - if (use_low_precision) { - if (scale_mode == HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0) { - PRIMUS_TURBO_CHECK( - is_gfx950(), - "The HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0 only support on gfx950."); + auto &dc = desc_cache[dkey]; + if (dc.op_desc == nullptr) { + dc.reset(); + hipblasLtEpilogue_t epilogue = HIPBLASLT_EPILOGUE_DEFAULT; + hipblasComputeType_t gemm_compute_type = HIPBLAS_COMPUTE_32F; + + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatrixLayoutCreate(&dc.A_desc, A_type, rows_a, cols_a, lda)); + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatrixLayoutCreate(&dc.B_desc, B_type, rows_b, cols_b, ldb)); + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatrixLayoutCreate(&dc.D_desc, D_type, rows_d, cols_d, ldd)); + + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatmulDescCreate(&dc.op_desc, gemm_compute_type, HIP_R_32F)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + dc.op_desc, HIPBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + dc.op_desc, HIPBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + dc.op_desc, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); + + if (use_low_precision) { + if (scale_mode == HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0) { + PRIMUS_TURBO_CHECK( + is_gfx950(), + "The HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0 only support on gfx950."); + } + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + dc.op_desc, HIPBLASLT_MATMUL_DESC_A_SCALE_MODE, &scale_mode, sizeof(scale_mode))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( + dc.op_desc, HIPBLASLT_MATMUL_DESC_B_SCALE_MODE, &scale_mode, sizeof(scale_mode))); } + } + + hipblasLtMatmulDesc_t operation_desc = dc.op_desc; + hipblasLtMatrixLayout_t A_desc = dc.A_desc, B_desc = dc.B_desc, D_desc = dc.D_desc; + + // Scale pointers change per call — must update every time for FP8. + if (use_low_precision) { PRIMUS_TURBO_CHECK(scaleA_inv != nullptr); PRIMUS_TURBO_CHECK(scaleB_inv != nullptr); + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatmulDescSetAttribute(operation_desc, HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &scaleA_inv, sizeof(scaleA_inv))); + PRIMUS_TURBO_CHECK_HIPBLAS( + hipblasLtMatmulDescSetAttribute(operation_desc, HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &scaleB_inv, sizeof(scaleB_inv))); + } - hipblasLtMatmulDescAttributes_t scaleA_inv_ptr_desc = HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER; - hipblasLtMatmulDescAttributes_t scaleB_inv_ptr_desc = HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER; + auto autotune_algo = [&]() -> hipblasLtMatmulHeuristicResult_t { + static constexpr int kMaxCandidates = 8; + static constexpr int kWarmupIters = 10; + static constexpr int kBenchIters = 30; + hipblasLtMatmulHeuristicResult_t candidates[kMaxCandidates]; + hipblasLtMatmulHeuristicResult_t tuned_algo{}; + int returnedAlgoCount = 0; + hipblasLtMatmulPreference_t preference = nullptr; - 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(hipblasLtMatmulPreferenceCreate(&preference)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulPreferenceSetAttribute( + preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspace_size, + sizeof(workspace_size))); - 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))); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulAlgoGetHeuristic( + handle, operation_desc, A_desc, B_desc, D_desc, D_desc, preference, kMaxCandidates, + candidates, &returnedAlgoCount)); + PRIMUS_TURBO_CHECK(returnedAlgoCount > 0, + "hipBLASLt: no valid algorithm found for current matmul config"); + + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmulPreferenceDestroy(preference)); + + if (returnedAlgoCount == 1 || D_type == HIP_R_32F) { + // Float32 coverage in CI is correctness-oriented rather than performance-critical. + // Avoid timing-based autotune here because mixed shared-GPU workloads can skew the + // ranking and make the chosen algo flaky for tiny shapes. + tuned_algo = candidates[0]; + } else { + // Autotune: benchmark each candidate, pick fastest. + const float a1 = 1.0f, b0 = 0.0f; + hipEvent_t ev_start, ev_stop; + PRIMUS_TURBO_CHECK_HIP(hipEventCreate(&ev_start)); + PRIMUS_TURBO_CHECK_HIP(hipEventCreate(&ev_stop)); + + float best_ms = 1e30f; + int best_idx = -1; + for (int c = 0; c < returnedAlgoCount; ++c) { + bool candidate_ok = true; + hipblasStatus_t status = HIPBLAS_STATUS_SUCCESS; + + // warm-up + for (int w = 0; w < kWarmupIters; ++w) { + status = hipblasLtMatmul(handle, operation_desc, &a1, A, A_desc, B, B_desc, &b0, + D, D_desc, D, D_desc, &candidates[c].algo, workspace, + workspace_size, stream); + if (status != HIPBLAS_STATUS_SUCCESS) { + candidate_ok = false; + break; + } + } + if (!candidate_ok) + continue; + + PRIMUS_TURBO_CHECK_HIP(hipEventRecord(ev_start, stream)); + for (int i = 0; i < kBenchIters; ++i) { + status = hipblasLtMatmul(handle, operation_desc, &a1, A, A_desc, B, B_desc, &b0, + D, D_desc, D, D_desc, &candidates[c].algo, workspace, + workspace_size, stream); + if (status != HIPBLAS_STATUS_SUCCESS) { + candidate_ok = false; + break; + } + } + if (!candidate_ok) + continue; + + PRIMUS_TURBO_CHECK_HIP(hipEventRecord(ev_stop, stream)); + PRIMUS_TURBO_CHECK_HIP(hipEventSynchronize(ev_stop)); + + float ms = 0; + PRIMUS_TURBO_CHECK_HIP(hipEventElapsedTime(&ms, ev_start, ev_stop)); + if (ms < best_ms) { + best_ms = ms; + best_idx = c; + } + } + PRIMUS_TURBO_CHECK_HIP(hipEventDestroy(ev_start)); + PRIMUS_TURBO_CHECK_HIP(hipEventDestroy(ev_stop)); + PRIMUS_TURBO_CHECK(best_idx >= 0, + "hipBLASLt: all autotune candidates failed during benchmarking"); + tuned_algo = candidates[best_idx]; + } + + return tuned_algo; + }; + + // Check exact-shape overrides first, then fall back to the optimistic + // M-invariant cache to avoid repeated autotune on similar expert shapes. + AlgoKey akey{cols_a, + rows_b, + cols_b, + ldb, + cols_d, + A_type, + B_type, + D_type, + transA, + transB, + use_low_precision, + scale_mode, + handle}; + + hipblasLtMatmulHeuristicResult_t algo_result{}; + bool used_shared_algo_cache = false; + auto exact_it = exact_algo_cache.find(dkey); + if (exact_it != exact_algo_cache.end()) { + algo_result = exact_it->second; + } else { + auto it = algo_cache.find(akey); + if (it != algo_cache.end()) { + algo_result = it->second; + used_shared_algo_cache = true; + } else { + algo_result = autotune_algo(); + algo_cache[akey] = algo_result; + } } - const int request_solutions = 1; - std::vector algos(request_solutions); - 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, request_solutions, - algos.data(), &returnedAlgoCount)); - PRIMUS_TURBO_CHECK(returnedAlgoCount > 0, - "hipBLASLt: no valid algorithm found for current matmul config"); - - const float alpha = 1.0; - const float beta = 0.0; - // clang-format off - PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtMatmul( - handle, - operation_desc, - &alpha, - A, A_desc, - B, B_desc, - &beta, - D, D_desc, - D, D_desc, - &algos[0].algo, - workspace, workspace_size, - stream)); - // clang-format on - - 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)); + const float alpha = 1.0; + const float beta = 0.0; + hipblasStatus_t status = + hipblasLtMatmul(handle, operation_desc, &alpha, A, A_desc, B, B_desc, &beta, D, D_desc, D, + D_desc, &algo_result.algo, workspace, workspace_size, stream); + + if (status == HIPBLAS_STATUS_INVALID_VALUE && used_shared_algo_cache) { + // Some configs still need an exact-shape algo even when N/K/trans/dtype + // match. Retune for this descriptor and remember the exact override. + algo_result = autotune_algo(); + exact_algo_cache[dkey] = algo_result; + status = + hipblasLtMatmul(handle, operation_desc, &alpha, A, A_desc, B, B_desc, &beta, D, D_desc, + D, D_desc, &algo_result.algo, workspace, workspace_size, stream); + } + + PRIMUS_TURBO_CHECK_HIPBLAS(status); + + // Descriptors live in thread-local caches and are released when cache + // entries are destroyed, so no per-call destroy is needed here. } } // namespace primus_turbo diff --git a/csrc/kernels/grouped_gemm/hipblaslt_grouped_gemm.cu b/csrc/kernels/grouped_gemm/hipblaslt_grouped_gemm.cu index 91a57f72e..97b33bef6 100644 --- a/csrc/kernels/grouped_gemm/hipblaslt_grouped_gemm.cu +++ b/csrc/kernels/grouped_gemm/hipblaslt_grouped_gemm.cu @@ -5,58 +5,64 @@ #include #include #include +#include #include #include +#include +#include #include "primus_turbo/gemm.h" #include "primus_turbo/grouped_gemm.h" namespace primus_turbo { -static constexpr size_t kMaxNumStreams = 4; +static constexpr size_t kMaxNumStreams = 8; + +// Per-expert token threshold: above this, individual GEMMs are large enough +// to saturate the GPU, so serial dispatch avoids resource contention. +// Below this, GEMMs are small and benefit from concurrent multi-stream +// dispatch to overlap kernel launch overhead. +static constexpr size_t kSerialThreshold = 512; std::int64_t get_hipblaslt_grouped_gemm_workspace_size() { + // Baseline requirement for the multi-stream path: one workspace per stream. + // The PyTorch wrapper may over-allocate for the serial path so each expert + // gets a dedicated workspace slice as well. return kMaxNumStreams * get_hipblaslt_workspace_size_in_byte(); } class HipblasltGroupedGemm { public: HipblasltGroupedGemm() { + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtCreate(&handle_)); PRIMUS_TURBO_CHECK_HIP(hipEventCreateWithFlags(&sync_event_, hipEventDisableTiming)); for (size_t i = 0; i < kMaxNumStreams; ++i) { - PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtCreate(&handles_[i])); - PRIMUS_TURBO_CHECK_HIP( - hipStreamCreateWithPriority(&compute_streams_[i], hipStreamNonBlocking, -1)); + PRIMUS_TURBO_CHECK_HIPBLAS(hipblasLtCreate(&par_handles_[i])); PRIMUS_TURBO_CHECK_HIP( - hipEventCreateWithFlags(&hipblaslt_events_[i], hipEventDisableTiming)); + hipStreamCreateWithPriority(&par_streams_[i], hipStreamNonBlocking, -1)); + PRIMUS_TURBO_CHECK_HIP(hipEventCreateWithFlags(&par_events_[i], hipEventDisableTiming)); } - workspaces_.resize(kMaxNumStreams); } ~HipblasltGroupedGemm() { - if (sync_event_ != nullptr) { + if (sync_event_) (void) hipEventDestroy(sync_event_); - } - for (size_t i = 0; i < kMaxNumStreams; ++i) { - if (compute_streams_[i] != nullptr) { - (void) hipStreamDestroy(compute_streams_[i]); - } - if (handles_[i] != nullptr) { - (void) hipblasLtDestroy(handles_[i]); - } - if (hipblaslt_events_[i] != nullptr) { - (void) hipEventDestroy(hipblaslt_events_[i]); - } + if (par_streams_[i]) + (void) hipStreamDestroy(par_streams_[i]); + if (par_handles_[i]) + (void) hipblasLtDestroy(par_handles_[i]); + if (par_events_[i]) + (void) hipEventDestroy(par_events_[i]); } + if (handle_) + (void) hipblasLtDestroy(handle_); } void check(const HipblasltGroupedGemmParams ¶ms) { PRIMUS_TURBO_CHECK(params.a_shape.size() == 2); if (params.transA) { - // For a * grad_c = grad_b - // [m, k]^T * [m, n] = [b, k, n] PRIMUS_TURBO_CHECK(params.b_shape.size() == 2); PRIMUS_TURBO_CHECK(params.c_shape.size() == 3); PRIMUS_TURBO_CHECK(params.a_shape[0] == params.b_shape[0]); @@ -64,7 +70,6 @@ public: PRIMUS_TURBO_CHECK(params.c_shape[1] == params.a_shape[1]); PRIMUS_TURBO_CHECK(params.c_shape[2] == params.b_shape[1]); } else { - // For a * b = c and grad_c * b = grad_a PRIMUS_TURBO_CHECK(params.b_shape.size() == 3); PRIMUS_TURBO_CHECK(params.c_shape.size() == 2); PRIMUS_TURBO_CHECK(params.b_shape[0] == params.group_num); @@ -76,58 +81,139 @@ public: PRIMUS_TURBO_CHECK_HIP(hipStreamSynchronize(params.stream)); } - // Check check(params); - // Compute arguments compute_args(params); - const size_t num_gemms{gemm_ptrs_.size()}; - const size_t num_stream_used{std::min(kMaxNumStreams, num_gemms)}; + const size_t num_gemms = gemm_ptrs_.size(); + if (num_gemms == 0) + return; + + // Warm the hipBLASLt algo cache with balanced M (total_M / group_num) + // so the tuned algo is representative of the average expert size, + // not biased by whichever expert happens to be dispatched first. + warm_algo_cache(params, num_gemms); + + // Determine max per-expert token count to choose dispatch strategy. + size_t max_tokens = 0; + for (size_t i = 0; i < num_gemms; ++i) { + // fwd/dgrad: cols_c = token count; wgrad: cols_a = token count + size_t tokens = params.transA ? cols_b_[i] : cols_c_[i]; + max_tokens = std::max(max_tokens, tokens); + } - if (num_gemms > 0) { - // wait for current stream to finish - PRIMUS_TURBO_CHECK_HIP(hipEventRecord(sync_event_, params.stream)); + if (max_tokens >= kSerialThreshold || num_gemms <= kMaxNumStreams) { + dispatch_serial(params, num_gemms); + } else { + dispatch_parallel(params, num_gemms); + } + } - for (size_t s = 0; s < num_stream_used; ++s) { - PRIMUS_TURBO_CHECK_HIP(hipStreamWaitEvent(compute_streams_[s], sync_event_, 0)); - } + void clear_runtime_state() { warmed_shapes_.clear(); } - for (size_t idx = 0; idx < num_gemms; ++idx) { - const auto stream_idx = idx % kMaxNumStreams; - auto stream = compute_streams_[stream_idx]; - auto handle = handles_[stream_idx]; - auto workspace = workspaces_[stream_idx]; - // clang-format off - hipblaslt_gemm_impl( - gemm_ptrs_[idx].b_ptr, params.b_type, rows_b_[idx], cols_b_[idx], ld_b_[idx], - gemm_ptrs_[idx].b_scale_ptr, - params.transB ? HIPBLAS_OP_T : HIPBLAS_OP_N, - gemm_ptrs_[idx].a_ptr, params.a_type, rows_a_[idx], cols_a_[idx], ld_a_[idx], - gemm_ptrs_[idx].a_scale_ptr, - params.transA ? HIPBLAS_OP_T : HIPBLAS_OP_N, - gemm_ptrs_[idx].c_ptr, params.c_type, rows_c_[idx], cols_c_[idx], ld_c_[idx], - workspace, get_hipblaslt_workspace_size_in_byte(), - params.use_low_precision, - params.scale_mode, - handle, - stream - ); - // clang-format on - } +private: + void warm_algo_cache(const HipblasltGroupedGemmParams ¶ms, size_t num_gemms) { + // Use (rows_b, cols_a, transA, transB) as a compact shape key. + // This matches the M-invariant part of AlgoKey in hipblaslt_gemm.cu. + auto shape_key = std::make_tuple(rows_b_[0], cols_a_[0], params.transA, params.transB); + if (warmed_shapes_.count(shape_key)) + return; + + // Compute balanced M = total_tokens / group_num. + int64_t total_tokens = 0; + for (size_t i = 0; i < num_gemms; ++i) { + total_tokens += params.transA ? cols_b_[i] : cols_c_[i]; + } + const int64_t balanced_m = total_tokens / params.group_num; + if (balanced_m <= 0) + return; + + // Call hipblaslt_gemm_impl with balanced M to trigger algo tuning + // with a representative workload size. + const int64_t bal_cols_a = params.transA ? rows_a_[0] : balanced_m; + const int64_t bal_cols_b = params.transA ? balanced_m : cols_b_[0]; + const int64_t bal_cols_d = params.transA ? cols_c_[0] : balanced_m; + // clang-format off + hipblaslt_gemm_impl( + gemm_ptrs_[0].b_ptr, params.b_type, rows_b_[0], bal_cols_b, ld_b_[0], + gemm_ptrs_[0].b_scale_ptr, + params.transB ? HIPBLAS_OP_T : HIPBLAS_OP_N, + gemm_ptrs_[0].a_ptr, params.a_type, rows_a_[0], bal_cols_a, ld_a_[0], + gemm_ptrs_[0].a_scale_ptr, + params.transA ? HIPBLAS_OP_T : HIPBLAS_OP_N, + gemm_ptrs_[0].c_ptr, params.c_type, rows_c_[0], bal_cols_d, ld_c_[0], + params.workspace, get_hipblaslt_workspace_size_in_byte(), + params.use_low_precision, + params.scale_mode, + handle_, + params.stream + ); + // clang-format on + PRIMUS_TURBO_CHECK_HIP(hipStreamSynchronize(params.stream)); + warmed_shapes_.insert(shape_key); + } - // record events on compute streams - for (size_t s = 0; s < num_stream_used; ++s) { - PRIMUS_TURBO_CHECK_HIP(hipEventRecord(hipblaslt_events_[s], compute_streams_[s])); - } + void dispatch_serial(const HipblasltGroupedGemmParams ¶ms, size_t num_gemms) { + char *workspace_base = static_cast(params.workspace); + for (size_t idx = 0; idx < num_gemms; ++idx) { + void *ws = workspace_base + idx * get_hipblaslt_workspace_size_in_byte(); + // clang-format off + hipblaslt_gemm_impl( + gemm_ptrs_[idx].b_ptr, params.b_type, rows_b_[idx], cols_b_[idx], ld_b_[idx], + gemm_ptrs_[idx].b_scale_ptr, + params.transB ? HIPBLAS_OP_T : HIPBLAS_OP_N, + gemm_ptrs_[idx].a_ptr, params.a_type, rows_a_[idx], cols_a_[idx], ld_a_[idx], + gemm_ptrs_[idx].a_scale_ptr, + params.transA ? HIPBLAS_OP_T : HIPBLAS_OP_N, + gemm_ptrs_[idx].c_ptr, params.c_type, rows_c_[idx], cols_c_[idx], ld_c_[idx], + ws, get_hipblaslt_workspace_size_in_byte(), + params.use_low_precision, + params.scale_mode, + handle_, + params.stream + ); + // clang-format on + } + } - // wait for all compute streams to finish - for (size_t s = 0; s < num_stream_used; ++s) { - PRIMUS_TURBO_CHECK_HIP(hipStreamWaitEvent(params.stream, hipblaslt_events_[s], 0)); - } + void dispatch_parallel(const HipblasltGroupedGemmParams ¶ms, size_t num_gemms) { + const size_t num_streams = std::min(kMaxNumStreams, num_gemms); + char *workspace_base = static_cast(params.workspace); + + // Fork: make compute streams wait for the caller's stream. + PRIMUS_TURBO_CHECK_HIP(hipEventRecord(sync_event_, params.stream)); + for (size_t s = 0; s < num_streams; ++s) { + PRIMUS_TURBO_CHECK_HIP(hipStreamWaitEvent(par_streams_[s], sync_event_, 0)); + } + + // Round-robin dispatch across streams. + for (size_t idx = 0; idx < num_gemms; ++idx) { + const size_t s = idx % kMaxNumStreams; + void *ws = workspace_base + s * get_hipblaslt_workspace_size_in_byte(); + // clang-format off + hipblaslt_gemm_impl( + gemm_ptrs_[idx].b_ptr, params.b_type, rows_b_[idx], cols_b_[idx], ld_b_[idx], + gemm_ptrs_[idx].b_scale_ptr, + params.transB ? HIPBLAS_OP_T : HIPBLAS_OP_N, + gemm_ptrs_[idx].a_ptr, params.a_type, rows_a_[idx], cols_a_[idx], ld_a_[idx], + gemm_ptrs_[idx].a_scale_ptr, + params.transA ? HIPBLAS_OP_T : HIPBLAS_OP_N, + gemm_ptrs_[idx].c_ptr, params.c_type, rows_c_[idx], cols_c_[idx], ld_c_[idx], + ws, get_hipblaslt_workspace_size_in_byte(), + params.use_low_precision, + params.scale_mode, + par_handles_[s], + par_streams_[s] + ); + // clang-format on + } + + // Join: caller's stream waits for all compute streams. + for (size_t s = 0; s < num_streams; ++s) { + PRIMUS_TURBO_CHECK_HIP(hipEventRecord(par_events_[s], par_streams_[s])); + PRIMUS_TURBO_CHECK_HIP(hipStreamWaitEvent(params.stream, par_events_[s], 0)); } } -private: struct GemmPtr { const void *a_ptr = nullptr; const void *a_scale_ptr = nullptr; @@ -138,16 +224,30 @@ private: void compute_args(const HipblasltGroupedGemmParams ¶ms) { - // TODO(xiaobochen): group_lens_ptr is device pointer, but host can access - // it directly on MI300/MI350. Need to investigate why this works. + group_lens_host_.resize(params.group_num); + if (params.group_lens_on_host) { + std::memcpy(group_lens_host_.data(), params.group_lens_ptr, + params.group_num * sizeof(int64_t)); + } else { + PRIMUS_TURBO_CHECK_HIP(hipMemcpy(group_lens_host_.data(), params.group_lens_ptr, + params.group_num * sizeof(int64_t), + hipMemcpyDeviceToHost)); + } + int valid_group_num = 0; for (size_t i = 0; i < params.group_num; ++i) { - valid_group_num += params.group_lens_ptr[i] > 0 ? 1 : 0; + valid_group_num += group_lens_host_[i] > 0 ? 1 : 0; } const char *a_ptr = static_cast(params.a_ptr); const char *b_ptr = static_cast(params.b_ptr); char *c_ptr = static_cast(params.c_ptr); + + const int64_t b_expert_stride = params.transA ? 0 + : get_dim(params.b_shape, -1) * + get_dim(params.b_shape, -2) * + hipblaslt_dtype_bytes(params.b_type); + gemm_ptrs_.resize(valid_group_num); ld_a_.resize(valid_group_num); ld_b_.resize(valid_group_num); @@ -161,10 +261,8 @@ private: int write_idx = 0; for (size_t i = 0; i < params.group_num; ++i) { - int64_t len = params.group_lens_ptr[i]; + int64_t len = group_lens_host_[i]; - // For grad_b (transA), if the group len is 0, set the output memory to 0 - // c shape is [group_num, k, n], so each group's c size is k * n if (params.transA && len == 0) { int64_t c_rows_t = get_dim(params.c_shape, -1); int64_t c_cols_t = get_dim(params.c_shape, -2); @@ -176,38 +274,35 @@ private: if (len == 0) continue; - // pointers gemm_ptrs_[write_idx].a_ptr = a_ptr; - gemm_ptrs_[write_idx].b_ptr = b_ptr; + gemm_ptrs_[write_idx].b_ptr = params.transA + ? b_ptr + : static_cast(params.b_ptr) + + static_cast(i) * b_expert_stride; gemm_ptrs_[write_idx].c_ptr = c_ptr; if (params.use_low_precision) { - // TODO(xiaobochen): support variable scale mode gemm_ptrs_[write_idx].a_scale_ptr = params.a_scale_ptr; gemm_ptrs_[write_idx].b_scale_ptr = params.b_scale_ptr; } - // leading dimension - ld_a_[write_idx] = get_dim(params.a_shape, -1); - ld_b_[write_idx] = get_dim(params.b_shape, -1); - ld_c_[write_idx] = get_dim(params.c_shape, -1); - // rows and cols of matrices + ld_a_[write_idx] = get_dim(params.a_shape, -1); + ld_b_[write_idx] = get_dim(params.b_shape, -1); + ld_c_[write_idx] = get_dim(params.c_shape, -1); rows_a_[write_idx] = get_dim(params.a_shape, -1); cols_a_[write_idx] = len; rows_b_[write_idx] = get_dim(params.b_shape, -1); cols_b_[write_idx] = params.transA ? len : get_dim(params.b_shape, -2); rows_c_[write_idx] = get_dim(params.c_shape, -1); cols_c_[write_idx] = params.transA ? get_dim(params.c_shape, -2) : len; - // advance the pointers + a_ptr += rows_a_[write_idx] * cols_a_[write_idx] * hipblaslt_dtype_bytes(params.a_type); - b_ptr += rows_b_[write_idx] * cols_b_[write_idx] * hipblaslt_dtype_bytes(params.b_type); + if (params.transA) { + b_ptr += + rows_b_[write_idx] * cols_b_[write_idx] * hipblaslt_dtype_bytes(params.b_type); + } c_ptr += rows_c_[write_idx] * cols_c_[write_idx] * hipblaslt_dtype_bytes(params.c_type); write_idx++; } - - char *workspace_ptr = static_cast(params.workspace); - for (size_t i = 0; i < kMaxNumStreams; ++i) { - workspaces_[i] = workspace_ptr + i * get_hipblaslt_workspace_size_in_byte(); - } } template T get_dim(const std::vector &shape, int idx) { @@ -217,18 +312,20 @@ private: return shape.at(idx); } - // Handles, events, streams, heuristic, epilogue - hipblasLtHandle_t handles_[kMaxNumStreams]{}; - hipEvent_t sync_event_{nullptr}; - hipStream_t compute_streams_[kMaxNumStreams]{}; - hipEvent_t hipblaslt_events_[kMaxNumStreams]{}; - hipblasLtEpilogue_t epilogue_{HIPBLASLT_EPILOGUE_DEFAULT}; + // Primary handle for serial dispatch + hipblasLtHandle_t handle_{}; + + // Parallel dispatch infrastructure + hipblasLtHandle_t par_handles_[kMaxNumStreams]{}; + hipStream_t par_streams_[kMaxNumStreams]{}; + hipEvent_t par_events_[kMaxNumStreams]{}; + hipEvent_t sync_event_{nullptr}; // Gemm Pointers std::vector gemm_ptrs_; - // workspace - std::vector workspaces_; + // host-side copy of group_lens (avoids GPU cache coherence issues) + std::vector group_lens_host_; // Leading dimensions std::vector ld_a_; @@ -242,11 +339,27 @@ private: std::vector cols_b_; std::vector rows_c_; std::vector cols_c_; + + // Track which (N,K,transA,transB) shapes have been warmed for algo tuning + using ShapeKey = std::tuple; + std::set warmed_shapes_; }; -void hipblaslt_grouped_gemm(const HipblasltGroupedGemmParams ¶ms, const bool pre_sync) { +namespace { + +HipblasltGroupedGemm &get_thread_local_grouped_gemm_instance() { static thread_local HipblasltGroupedGemm instance; - instance.run(params, pre_sync); + return instance; +} + +} // namespace + +void hipblaslt_grouped_gemm(const HipblasltGroupedGemmParams ¶ms, const bool pre_sync) { + get_thread_local_grouped_gemm_instance().run(params, pre_sync); +} + +void clear_hipblaslt_grouped_gemm_runtime_state() { + get_thread_local_grouped_gemm_instance().clear_runtime_state(); } } // namespace primus_turbo diff --git a/csrc/pytorch/bindings_pytorch.cpp b/csrc/pytorch/bindings_pytorch.cpp index 2fa3f33f4..daf5dbe8a 100644 --- a/csrc/pytorch/bindings_pytorch.cpp +++ b/csrc/pytorch/bindings_pytorch.cpp @@ -5,6 +5,8 @@ #include #include "extensions.h" +#include "primus_turbo/gemm.h" +#include "primus_turbo/grouped_gemm.h" namespace primus_turbo::pytorch { @@ -206,6 +208,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { auto runtime_module = m.def_submodule("runtime", "Runtime utilities"); runtime_module.def("create_stream_with_cu_masks", &create_stream_with_cu_masks); runtime_module.def("destroy_stream", &destroy_stream); + runtime_module.def("clear_hipblaslt_gemm_runtime_caches", + &primus_turbo::clear_hipblaslt_gemm_runtime_caches); + runtime_module.def("clear_hipblaslt_grouped_gemm_runtime_state", + &primus_turbo::clear_hipblaslt_grouped_gemm_runtime_state); } /********************************************/ diff --git a/csrc/pytorch/grouped_gemm/hipblaslt_grouped_gemm.cpp b/csrc/pytorch/grouped_gemm/hipblaslt_grouped_gemm.cpp index aee2af7ec..f9f52adb3 100644 --- a/csrc/pytorch/grouped_gemm/hipblaslt_grouped_gemm.cpp +++ b/csrc/pytorch/grouped_gemm/hipblaslt_grouped_gemm.cpp @@ -34,14 +34,15 @@ make_hipblaslt_grouped_gemm_params(const at::Tensor &a, const at::Tensor &b, at: params.c_type = get_hipblaslt_dtype(c.scalar_type()); params.c_shape = c.sizes().vec(); - params.group_lens_ptr = reinterpret_cast(group_lens.data_ptr()); - params.group_offs_ptr = reinterpret_cast(group_offs.data_ptr()); - params.transA = transA; - params.transB = transB; - params.group_num = group_lens.numel(); - params.stream = at::cuda::getCurrentCUDAStream(); - params.workspace = workspace.data_ptr(); - params.handle = at::cuda::getCurrentCUDABlasLtHandle(); + params.group_lens_ptr = reinterpret_cast(group_lens.data_ptr()); + params.group_offs_ptr = reinterpret_cast(group_offs.data_ptr()); + params.group_lens_on_host = group_lens.is_cpu(); + params.transA = transA; + params.transB = transB; + params.group_num = group_lens.numel(); + params.stream = at::cuda::getCurrentCUDAStream(); + params.workspace = workspace.data_ptr(); + params.handle = at::cuda::getCurrentCUDABlasLtHandle(); return params; } @@ -65,13 +66,14 @@ inline HipblasltGroupedGemmParams make_hipblaslt_grouped_gemm_fp8_params( params.c_type = get_hipblaslt_dtype(c.scalar_type()); params.c_shape = c.sizes().vec(); - params.group_lens_ptr = reinterpret_cast(group_lens.data_ptr()); - params.group_offs_ptr = reinterpret_cast(group_offs.data_ptr()); - params.transA = transA; - params.transB = transB; - params.group_num = group_lens.numel(); - params.stream = at::cuda::getCurrentCUDAStream(); - params.workspace = workspace.data_ptr(); + params.group_lens_ptr = reinterpret_cast(group_lens.data_ptr()); + params.group_offs_ptr = reinterpret_cast(group_offs.data_ptr()); + params.group_lens_on_host = group_lens.is_cpu(); + params.transA = transA; + params.transB = transB; + params.group_num = group_lens.numel(); + params.stream = at::cuda::getCurrentCUDAStream(); + params.workspace = workspace.data_ptr(); params.use_low_precision = true; @@ -106,8 +108,13 @@ at::Tensor hipblaslt_grouped_gemm(at::Tensor &a, at::Tensor &b, at::Tensor &grou c = at::empty({m, n}, a.options()); } - const int64_t workspace_size = primus_turbo::get_hipblaslt_grouped_gemm_workspace_size(); - at::Tensor workspace = + const int64_t single_workspace_size = primus_turbo::get_hipblaslt_workspace_size_in_byte(); + const int64_t grouped_workspace_size = group_lens.numel() * single_workspace_size; + const int64_t workspace_size = + grouped_workspace_size > primus_turbo::get_hipblaslt_grouped_gemm_workspace_size() + ? grouped_workspace_size + : primus_turbo::get_hipblaslt_grouped_gemm_workspace_size(); + at::Tensor workspace = at::empty({workspace_size}, at::TensorOptions().dtype(at::kByte).device(a.device())); auto params = make_hipblaslt_grouped_gemm_params(a, b, c, group_lens, group_offs, transA, @@ -151,8 +158,13 @@ at::Tensor hipblaslt_grouped_gemm_fp8(at::Tensor &a, at::Tensor &b, at::Tensor & c = at::empty({m, n}, a.options().dtype(out_dtype)); } - const int64_t workspace_size = primus_turbo::get_hipblaslt_grouped_gemm_workspace_size(); - at::Tensor workspace = + const int64_t single_workspace_size = primus_turbo::get_hipblaslt_workspace_size_in_byte(); + const int64_t grouped_workspace_size = group_lens.numel() * single_workspace_size; + const int64_t workspace_size = + grouped_workspace_size > primus_turbo::get_hipblaslt_grouped_gemm_workspace_size() + ? grouped_workspace_size + : primus_turbo::get_hipblaslt_grouped_gemm_workspace_size(); + at::Tensor workspace = at::empty({workspace_size}, at::TensorOptions().dtype(at::kByte).device(a.device())); auto params = make_hipblaslt_grouped_gemm_fp8_params( diff --git a/primus_turbo/pytorch/core/backend.py b/primus_turbo/pytorch/core/backend.py index 7ab040359..0b91c1ad1 100644 --- a/primus_turbo/pytorch/core/backend.py +++ b/primus_turbo/pytorch/core/backend.py @@ -3,6 +3,7 @@ # # See LICENSE for license information. ############################################################################### +import importlib import os import warnings from abc import ABC, abstractmethod @@ -240,6 +241,10 @@ def reset(cls) -> None: cls._grouped_gemm_backend = None cls._auto_tune = None AutoKernelDispatcher.clear_all_caches() + runtime = importlib.import_module("primus_turbo.pytorch._C.runtime") + + runtime.clear_hipblaslt_gemm_runtime_caches() + runtime.clear_hipblaslt_grouped_gemm_runtime_state() clear_origami_caches() diff --git a/primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_impl.py b/primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_impl.py index 3b7a7c5fd..76b4393c6 100644 --- a/primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_impl.py +++ b/primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_impl.py @@ -216,7 +216,7 @@ def execute( _GROUPED_GEMM_BACKENDS = { BackendType.CK: BackendEntry(GroupedGEMMCKBackend), - BackendType.HIPBLASLT: BackendEntry(GroupedGEMMHipblasltBackend, autotune=False), + BackendType.HIPBLASLT: BackendEntry(GroupedGEMMHipblasltBackend), BackendType.TRITON: BackendEntry(GroupedGEMMTritonBackend), } @@ -263,7 +263,7 @@ def execute( _GROUPED_GEMM_VARIABLE_K_BACKENDS = { BackendType.CK: BackendEntry(GroupedGEMMVariableKCKBackend), - BackendType.HIPBLASLT: BackendEntry(GroupedGEMMVariableKHipblasltBackend, autotune=False), + BackendType.HIPBLASLT: BackendEntry(GroupedGEMMVariableKHipblasltBackend), BackendType.TRITON: BackendEntry(GroupedGEMMVariableKTritonBackend), } @@ -278,7 +278,6 @@ def make_key(cls, a, b, group_lens, group_offs, trans_a, trans_b, num_cu, **kwar m = a.shape[1] if trans_a else a.shape[0] n = b.shape[-2] if trans_b else b.shape[-1] k = a.shape[0] if trans_a else a.shape[1] - # bs, m, n, k, a.dtype, b.dtype, out_dtype, trans_a, trans_b, trans_c return (bs, m, n, k, a.dtype, b.dtype, a.dtype, trans_a, trans_b, False) diff --git a/requirements.txt b/requirements.txt index df0df151a..9873e05bb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ clang-format==18.1.6 hipify_torch @ git+https://github.com/ROCm/hipify_torch.git build scikit-build-core -triton==3.6.0 +triton expecttest pytest pytest-xdist