Year: 2026
Type: article
Abstract: Batched general matrix multiplication (GEMM) operations in latency-sensitive inference workloads often expose limited cooperative thread array (CTA)-level parallelism within each GEMM, making efficient GPU utilization challenging. Independent GEMMs can run concurrently, but this creates a tradeoff: each GEMM should expose more CTA-level parallelism when the batch is small and limit its CTA count when the batch already provides enough work to occupy the GPU. This paper presents KABUMI, a communication-avoiding (CA) batched GEMM framework for GPUs. KABUMI treats streaming multiprocessor (SM) participation as a budgeted resource rather than a byproduct of fixed tiling. For each single GEMM, KABUMI selects a CA algorithm according to the reduction length: small- K GEMMs use 1D CA output-space expansion, where output-tile candidates partition the wide column dimension, retain the tiny A operand near its owning CTA, stream the corresponding columns of B , and directly write disjoint C tiles; large- K GEMMs use 3D CA reduction-slice expansion to turn the long reduction dimension into schedulable partial CTAs and combine partial outputs. KABUMI’s batch-aware scheduler then maps the current batch size to a per-GEMM CTA-wave budget, exposing more intra-GEMM parallelism for underfilled batches and capping expansion when batch-level concurrency can already occupy the GPU. We implement KABUMI in CUDA and evaluate it on NVIDIA RTX 5090 and B200 GPUs. Across 1,440 batched-GEMM cases on RTX 5090, KABUMI achieves 1.60 ×, 2.45 ×, 2.64 ×, and 7.00 × geomean speedup over cuBLAS batched GEMM, CUTLASS batched GEMM, CUTLASS grouped GEMM, and MAGMA batched GEMM, respectively. Across the same 1,440 batched-GEMM cases on B200, the corresponding geomean speedups are 1.32 ×, 1.98 ×, 2.09 ×, and 5.03 ×. On model-derived grouped GEMMs, KABUMI reaches 1.30 × –1.56 × geomean speedup over the evaluated BF16 library paths and 1.11 × over FlashInfer in FP8.
DOI: 10.1145/3860333
Authors: Hemeng Wang, Weifeng Liu
Institution: China University of Petroleum, Beijing
Language: English
Subfield: Hardware and Architecture
Field: Computer Science
Domain: Physical Sciences
Open Access status: diamond
Award: G3000813168
OpenAlex ID: W7225401825