2比特量化下，用预计算加速大模型推理，速度提升近9倍。
CodeGEMM: A Codebook-Centric Approach to Efficient GEMM in Quantized LLMs
- 用预计算的局部和表替代逐元素反量化，减少内存访问
- 在Llama-3上实现8B模型1.83倍、70B模型8.93倍加速
- 适合追求极致推理效率的低比特量化部署场景
权重量化广泛用于缓解大模型推理中的内存瓶颈。基于码本的方法在极低比特（如2比特）下仍能保持良好精度。然而现有内核依赖反量化，需反复读取中心点并重构权重，导致显著延迟和缓存压力。我们提出CodeGEMM，一种以码本为中心的矩阵乘法内核，将反量化替换为预先计算的中心点与激活值之间的内积，并存储于轻量级Psumbook中。推理时，仅通过码索引直接获取部分和，避免逐元素查找，显著降低片上存储开销。该内核支持在统一实现下系统探索延迟-内存-精度的权衡。在Llama-3模型上，CodeGEMM在2比特配置下相比最先进码本量化方法，在相近精度下分别实现8B模型1.83倍、70B模型8.93倍的加速，进一步提升计算效率与内存子系统利用率。
原文摘要 · Abstract (English)
Weight-only quantization is widely used to mitigate the memory-bound nature of LLM inference. Codebook-based methods extend this trend by achieving strong accuracy in the extremely low-bit regime (e.g., 2-bit). However, current kernels rely on dequantization, which repeatedly fetches centroids and reconstructs weights, incurring substantial latency and cache pressure. We present CodeGEMM, a codebook-centric GEMM kernel that replaces dequantization with precomputed inner products between centroids and activations stored in a lightweight Psumbook. At inference, code indices directly gather these partial sums, eliminating per-element lookups and reducing the on-chip footprint. The kernel supports the systematic exploration of latency-memory-accuracy trade-offs under a unified implementation. On Llama-3 models, CodeGEMM delivers 1.83x (8B) and 8.93x (70B) speedups in the 2-bit configuration compared to state-of-the-art codebook-based quantization at comparable accuracy and further improves computing efficiency and memory subsystem utilization.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。