Fast Matrix Multiplications for Lookup Table-Quantized LLMs
作者: Han Guo, William Brandon, Radostin Cholakov, Jonathan Ragan-Kelley, Eric P. Xing, Yoon Kim
分类: cs.LG, cs.CL, cs.DC
发布日期: 2024-07-15 (更新: 2025-01-17)
备注: EMNLP 2024 (Findings)
💡 一句话要点
FLUTE：加速查找表量化LLM的矩阵乘法，提升推理速度
🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 查找表量化 大型语言模型 矩阵乘法 推理加速 CUDA内核
📋 核心要点
- 现有LLM推理受限于内存带宽，权重量化虽可加速，但非均匀量化位宽和查找表量化给高性能内核开发带来挑战。
- FLUTE通过离线重构量化权重矩阵，减少位操作，并向量化查找表，缓解共享内存带宽瓶颈，从而加速推理。
- 实验表明，FLUTE内核在典型LLM推理设置下比现有GEMM内核快2-4倍，并使LLaMA3的端到端吞吐量提高1.5-2倍。
📝 摘要（中文）
大型语言模型（LLM）的部署通常受限于内存带宽，主要瓶颈在于将模型参数从GPU全局内存传输到寄存器的成本。结合融合反量化和矩阵乘法操作的自定义内核，权重量化可以通过减少内存移动来加速推理。然而，为权重已量化的LLM开发高性能内核面临巨大挑战，特别是当权重被压缩到非均匀可分的位宽（例如3位）并采用非均匀查找表（LUT）量化时。本文介绍了FLUTE，一种用于LUT量化LLM的灵活查找表引擎，它使用量化权重矩阵的离线重构来最小化与解包相关的位操作，并对查找表进行向量化和复制以缓解共享内存带宽约束。在批大小<32且量化组大小为128（LLM推理中的典型值）时，FLUTE内核比现有GEMM内核快2-4倍。作为FLUTE的应用，我们探索了基于查找表的NormalFloat量化的简单扩展，并将其应用于量化LLaMA3到各种配置，在获得与强基线相比具有竞争力的量化性能的同时，获得了1.5到2倍的端到端吞吐量提升。
🔬 方法详解
问题定义：论文旨在解决查找表（LUT）量化的大型语言模型（LLM）在推理过程中矩阵乘法速度慢的问题。现有的通用矩阵乘法（GEMM）内核无法有效处理非均匀位宽的量化权重，导致大量的时间消耗在解包和位操作上，成为性能瓶颈。
核心思路：论文的核心思路是通过离线重构量化后的权重矩阵，减少在线推理时所需的位操作。同时，通过向量化和复制查找表，缓解共享内存带宽的限制，从而提高矩阵乘法的效率。这种方法旨在优化数据访问模式和计算流程，以适应LUT量化的特点。
技术框架：FLUTE框架主要包含两个阶段：离线预处理和在线推理。在离线预处理阶段，对量化后的权重矩阵进行重构，优化数据存储格式。在在线推理阶段，使用优化的内核进行矩阵乘法运算，该内核利用了预处理后的数据格式和向量化的查找表。整体流程是先加载输入激活，然后使用重构后的权重和查找表进行计算，最后输出结果。
关键创新：FLUTE的关键创新在于针对LUT量化权重的特定优化策略。与通用的GEMM内核不同，FLUTE专门设计用于处理非均匀位宽的量化权重，并通过离线重构和查找表优化，显著减少了在线推理时的计算量和内存访问量。这种针对性优化是FLUTE能够获得显著性能提升的关键。
关键设计：FLUTE的关键设计包括：1) 权重矩阵的离线重构策略，旨在最小化解包所需的位操作；2) 查找表的向量化和复制，以提高共享内存的带宽利用率；3) 自定义的CUDA内核，针对特定的硬件架构进行优化。具体的参数设置和网络结构取决于所量化的LLM模型，但FLUTE的设计目标是提供一个通用的加速框架，可以应用于不同的LUT量化配置。
🖼️ 关键图片
📊 实验亮点
FLUTE内核在批大小小于32且量化组大小为128的典型LLM推理设置下，比现有GEMM内核快2-4倍。此外，将FLUTE应用于量化LLaMA3，在保持与强基线相当的量化性能的同时，实现了1.5到2倍的端到端吞吐量提升。这些实验结果表明，FLUTE是一种高效的LUT量化LLM加速方法。
🎯 应用场景
FLUTE具有广泛的应用前景，可以加速各种基于查找表量化的LLM的推理过程。这对于需要在资源受限的设备上部署LLM，例如移动设备、边缘服务器等，具有重要的实际价值。通过提高推理速度和降低内存占用，FLUTE可以使LLM在更多场景中得到应用，并推动人工智能技术的普及。
📄 摘要（原文）
The deployment of large language models (LLMs) is often constrained by memory bandwidth, where the primary bottleneck is the cost of transferring model parameters from the GPU's global memory to its registers. When coupled with custom kernels that fuse the dequantization and matmul operations, weight-only quantization can thus enable faster inference by reducing the amount of memory movement. However, developing high-performance kernels for weight-quantized LLMs presents substantial challenges, especially when the weights are compressed to non-evenly-divisible bit widths (e.g., 3 bits) with non-uniform, lookup table (LUT) quantization. This paper describes FLUTE, a flexible lookup table engine for LUT-quantized LLMs, which uses offline restructuring of the quantized weight matrix to minimize bit manipulations associated with unpacking, and vectorization and duplication of the lookup table to mitigate shared memory bandwidth constraints. At batch sizes < 32 and quantization group size of 128 (typical in LLM inference), the FLUTE kernel can be 2-4x faster than existing GEMM kernels. As an application of FLUTE, we explore a simple extension to lookup table-based NormalFloat quantization and apply it to quantize LLaMA3 to various configurations, obtaining competitive quantization performance against strong baselines while obtaining an end-to-end throughput increase of 1.5 to 2 times.