tritonBLAS: Triton-based Analytical Approach for GEMM Kernel Parameter Selection

TL;DR

The paper tackles the autotuning overhead in GPU GEMM kernel parameter selection by introducing tritonBLAS, a deterministic analytical model that predicts near-optimal tiling configurations without runtime benchmarking. Built atop Triton, it replaces empirical search with a latency-based, architecture-aware framework that accounts for hierarchical tiling, parallelism, and data locality. Evaluations across 150k GEMM shapes and Llama3 workloads show ~94.7% efficiency compared to exhaustive autotuning, with dramatically reduced configuration overhead and strong cross-architecture portability. The approach offers a practical drop-in alternative for production HPC/ML workloads, enabling fast, predictable GEMM performance without autotuning costs.

Abstract

We present tritonBLAS, a fast and deterministic analytical model that uses architectural parameters like the cache hierarchy, and relative code and data placement to generate performant GPU GEMM kernels. tritonBLAS explicitly models the relationship between architectural topology, matrix shapes, and algorithmic blocking behavior to predict near-optimal configurations without runtime autotuning. Based on this model, we developed and implemented a lightweight GEMM framework entirely within Triton. We evaluate the performance of tritonBLAS across a diverse set of GEMM problem sizes on modern GPUs. tritonBLAS achieves over 95% of the performance of autotuning solutions, while reducing autotuning time to zero. This makes tritonBLAS a practical drop-in replacement for empirical tuning in production HPC and ML workloads.

Figures (6)

• Figure 1: Example of hierarchical tiling of output matrix for AMD's MI300X GPU (see Table \ref{['tab:mi300_taxonomy']}, illustrating device, XCD, CU, SIMD and Thread breakdown of the tile).
• Figure 2: Example of occupancy and CUs across time (timestep$_0$, timestep$_1$ and timestep$_2$) and available CUs ($\text{CU}_0, \text{CU}_1, \text{CU}_2$).
• Figure 3: Efficiency of tritonBLAS relative to Triton's autotune exhaustive search across 150,000 sizes.
• Figure 4: Comparison between tritonBLAS and torch.matmul() on MI300X, showing peak percent performance (top-left), delta percent (top-right), performance vs. arithmetic intensity (bottom-left), and binned peak percent performance (bottom-right).
• Figure 5: tritonBLAS performance characteristics on AMD's Instinct™ MI350X. The tritonBLAS framework demonstrates strong performance on AMD Instinct™ MI350X, requiring only minor constant adjustments from the MI300X implementation.