CodeGEMM: A Codebook-Centric Approach to Efficient GEMM in Quantized LLMs
Gunho Park, Jeongin Bae, Byeongwook Kim, Baeseong Park, Jiwon Ryu, Hoseung Kim, Se Jung Kwon, Dongsoo Lee
Codebook-based quantization delivers superior performance in extreme low-bit environments but is hindered by latency overhead in existing dequantization-based kernels, which require loading the entire codebook into programmable cache memory. We propose CodeGEMM, a matrix multiplication kernel optimized for computing codebook-based quantized models by precomputing and storing inner product results in a Psumbook, thereby eliminating the need for traditional dequantization. The kernel supports a wide range of hyperparameter configurations, enabling exploration of trade-offs between latency, memory usage, and accuracy to achieve optimal performance.     CodeGEMM obtains a $2.27\times$ speed-up compared to the state-of-the-art codebook-based quantization method with comparable accuracy in the 2-bit configuration.