arXiv · 2407.00044
Stream-K Optimization and Exploration
Abstract
We explore optimization options for the Stream-K algorithm, a work-centric parallelization of matrix multiplication (GEMM). In our study, we investigated differences between the theoretical and practical implementations, particularly noting the impact of padding. Our debugging efforts revealed a persistent bug related to block mapping, which we could not fully resolve, but we managed to implement some optimizations. Setting the padding to zero for the M, N, and K dimensions resulted in an average 0.6% improvement in performance, achieving 1.44 ms, 89.37 TFlops, and 66.91 GB/s. However, adjusting the block size and parameters led to the process getting stuck, indicating a need for further tuning. Additionally, exploring the potential of Block2Time highlighted its promise in enhancing runtime predictions and optimizing load balancing.
Explore related subjects
Keep this discovery
Nick Rackley, Bryan Gonzalez, Casey Morrison. 2024-06-01. Stream-K Optimization and Exploration. https://arxiv.org/abs/2407.00044
Cite the original work for its findings. Save a collection to share your selection of sources.