Fast Matrix Multiplications for Lookup Table-Quantized LLM s (2024)

Abstract: 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.

## introduction large language model (llm) deployment faces significant latency challenges due to the memory bandwidth constraints inherent in generative (tokenby-token) inference. the primary bottleneck is the cost of transferring model parameters from the gpu's global memory to the registers, i.e., llm inference is memory-bound. to overcome this "memory wall" @xcite , practitioners have increasingly adopted weight-only quantization methods, wherein the parameters of an llm are compressed to lower precision (e.g., 4 or 8 bits) than the precision in which they were trained (typically 16 bits). in addition to latency improvements, weight quantization can also drastically reduce gpu memory required for deployment. realizing practical speed-ups with weight-only quantization requires custom mixed-type matrixmatrix multiply (matmul) kernels which must (1) move a layer's quantized weights from gpu offchip dram to on-chip sram, (2) dequantize the weights to floating-point (fp) format (on chip), (3) perform the fp matmul, and (4) write the results back to dram. existing kernels such as bitsandbytes @xcite , marlin @xcite , and bitblas @xcite demonstrate that this strategy can result in significant matmul speed-ups, e.g. up to four times faster when going from w16a16 to w4a16. however, these kernels are typically specialized to 4-bit quantization, and while some kernels support nonuniform, lookup table (lut) quantization, they are generally slower than the uniform counterparts. given the recent promising results with odd-bit @xcite and nonuniform @xcite quantization methods, there is thus a need to develop flexible kernels that can support mixed-type matmuls with a wider range of settings. this paper describes flute, a flexible lookuptable engine for deploying weight-quantized llms, with a focus on the low-bit and non-uniform quantization setting. this setting raises several challenges. first, going beyond 8-bit quantization involves packing sub-8-bit matrices into supported data types, followed by unpacking during dequantization. structuring the unpacked data to match gpu-native matmul formats is especially challenging when the weights are quantized to non-standard bit-widths. second, while uniformly-quantized models can rely on assembly-level optimizations to convert from int to fp through bit-level manipula-tions, lookup table-based dequantization involves dynamic indexing, and a naïve implementation can lead to substantial overhead. finally, typical matmul implementations which distribute the workload across a grid of parallel thread blocks become inefficient with small batches and low bit-width weights; this necessitates more sophisticated partitioning strategies to optimize hardware resource utilization. flute addresses these challenges through a combination of (1) offline weight restructuring, (2) a shared-memory lookup table for efficient dequantization, and (3) stream-k partitioning for optimized workload distribution. we compare flute against existing kernels on standard llm mixed-precision matmul settings where weights are quantized to 4 bits in groups of 128, and find that it outperforms existing non-uniform quantization kernels, and even matches the simpler uniform-quantization kernels in some cases. as an application of flute, we experiment with quantizing llama3-which has been found to be difficult to quantize @xcite -using a variant of normal float (nf) quantization @xcite which learns the quantization parameters based on calibration data. we find that we can achieve a 1.5 to 2 times increase in end-to-end throughput when integrated with frameworks such as vllm @xcite . ## flute: a fast and flexible kernel for mixed-type matrix multiplications let @xmath0. concretely, given a lookup table our objective is to perform a fast matrix multiplication between a dense input activation matrix @xmath1. a straightforward implementation of such mixed-type matrix multiplication uses separate kernels. the first kernel loads the quantized matrix q from the gpu's off-chip global memory into its on-chip memory, performs dequantization, and writes back the dequantized matrix w back to the dram. the second kernel is a standard fp matmul kernel over x and w. this separate-kernel can introduce substantial overhead since w is moved back and forth. we can achieve faster matmuls by fusing the dequantization and matmul kernels, where we dequantize on chip and immediately use the dequantized values for the matmul. ## experiments our experiments consist of two settings: kernellevel experiments which compare flute matmuls standalone against existing mixed-input matmul kernels ( §4.1), and end-to-end experiments which assess whether practicals speed-ups are obtainable on realistic llm workloads ( §4.2). ## discussion and conclusion early work on llm quantization generally worked with uniform quantization methods @xcite @xcite . na na na na single gpu (8b, 9b) llama-3 (a6000) llama-3 (a100) gemma-2 (a6000) gemma-2 (a100) single gpu (70b, 27b) tensor parallel (70b, 27b) unquantized 3-bit 4-bit more recent work has shown the benefits of lutquantization, both from ptq @xcite and finetuning @xcite perspectives. insofar as lookup tables can represent flexible quantization functions, our hope is that flute can enable researchers and practitioners to explore new quantization algorithms that can learn better lookup tables @xcite @xcite . for example, recent work has found that codebook-based quantization schemes-which generalize lookup tables to vector-valued valuescan enable even lower-bit (e.g., 2-bit) llm quantization without significant performance degradations @xcite . we anticipate that ideas from this work can aid in developing kernels for such methods. algorithmic considerations aside, one of the main challenges in developing fused quantized ma-trix multiplication kernels stems from the lack of hardware support for "mixed-type" instructions, necessitating software-level implementations. existing tensor core instructions support scenarios where the input and output/accumulation data have different types (e.g., compute in fp16 and output/accumulate in fp32). however, they do not support cases where the input operands themselves are of different types (e.g., fp16 inputs and int4 weights). as weight-only quantization becomes increasingly common in llm inference applications, native support for such instructions in future hardware could be beneficial. additionally, the lack of in-register dynamic indexing means that developers must devise software solutions. enhanced hardware acceleration for indexing into small lookup tables could also prove beneficial in the upcoming generations of ai accelerator hardware.