Fast Matrix Multiplications for Lookup Table-Quantized LLMs (2407.10960v3)

Abstract: The deployment of 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.

Fast Matrix Multiplications for Lookup Table-Quantized LLMs

The paper "Fast Matrix Multiplications for Lookup Table-Quantized LLMs" addresses the computational challenges of deploying LLMs in memory-constrained environments, such as GPUs. The proposed method, FLUTE (Flexible LookUp Table Engine), introduces a novel kernel design that fuses dequantization and matrix multiplication operations to achieve significant speed-ups during LLM inference. The authors focus on weight-only quantization, a technique that compresses the LLM weight parameters to lower precision, thereby reducing the memory bandwidth requirements.

Key Contributions and Implementation

FLUTE is engineered to address several practical challenges in deploying weight-quantized LLMs. Among these are:

1. Offline Matrix Restructuring: The paper underscores the importance of offline restructuring of quantized weight matrices to conform to the optimal layout required by GPU Tensor Cores. This pre-processing step ensures that the matrices are aligned in the most efficient format for subsequent dequantization and matrix multiplication operations.
2. Vectorized Lookup Table Design: To mitigate the inefficiencies associated with dynamic indexing operations during lookup table-based dequantization, the authors implement a vectorized lookup table approach. This method significantly reduces shared memory bandwidth consumption by enabling parallel access to pairs of values in the lookup table.
3. Stream-K Workload Partitioning: Standard matrix multiplication kernels can suffer from workload imbalances, especially in low-bit and low-batch scenarios. FLUTE adopts Stream-K workload partitioning, which decomposes computation tasks more granularly to ensure better distribution across GPU streaming multiprocessors (SMs). This technique minimizes the wave quantization effect and maximizes SM utilization.

Performance and Comparisons

The performance evaluation of FLUTE includes both kernel-level benchmarks and end-to-end LLM inference tests:

1. Kernel Benchmarks: The kernel-level tests compare FLUTE against existing mixed-input matrix multiplication kernels, including bitsandbytes and BitBLAS, which support lookup table-based dequantization. FLUTE demonstrated consistent speed-ups, particularly in memory-bound scenarios where batch sizes are small. It achieves up to 4x speed-up in certain configurations, thereby validating the efficiency of the fused kernel approach.
2. End-to-End LLM Benchmarks: For practical applications, the authors quantized LLaMA3 models using their proposed method and evaluated these on tasks that include WikiText-2 and C4. The results show that FLUTE not only meets but exceeds the performance of existing baselines, achieving competitive perplexity scores while increasing throughput by 1.5 to 2 times. The experimental results on Gemma-2 models further illustrate the robustness and versatility of FLUTE across different LLM architectures.

Implications and Future Directions

From a theoretical standpoint, the paper's contributions extend the boundaries of efficient LLM inference by addressing both algorithmic and architectural challenges in deploying quantized models. The proposed methods can be particularly beneficial in scenarios where memory bandwidth is a critical bottleneck, such as real-time AI applications deployed on embedded or edge devices.

Practically, this work equips researchers and practitioners with a robust framework for implementing efficient inference pipelines. The flexibility in supporting various bit widths and lookup table configurations offers the potential for further optimization in different deployment settings. Moreover, the implementation of FLUTE could inspire hardware manufacturers to consider native support for mixed-type operations in future GPU architectures, enhancing the overall capability for weight-quantized LLM inference.

Conclusion

FLUTE represents a significant advancement in the efficient deployment of LLMs, particularly in memory-constrained environments. By effectively addressing the challenges associated with weight-only lookup table quantization and leveraging offline restructuring, vectorized lookup, and Stream-K workload partitioning, FLUTE delivers substantial performance improvements that are validated through rigorous experimentation. This work lays a solid foundation for further research and development in the area of efficient AI inference.