Multiplications matricielles rapides pour les LLM quantifiés par table de correspondance

Fast Matrix Multiplications for Lookup Table-Quantized LLMs

July 15, 2024

Auteurs: Han Guo, William Brandon, Radostin Cholakov, Jonathan Ragan-Kelley, Eric P. Xing, Yoon Kim

cs.AI

Résumé

Le déploiement des grands modèles de langage (LLMs) est souvent limité par la bande passante mémoire, où le principal goulot d'étranglement réside dans le coût de transfert des paramètres du modèle depuis la mémoire globale du GPU vers ses registres. Lorsqu'elle est couplée à des noyaux personnalisés qui fusionnent les opérations de déquantification et de multiplication matricielle (matmul), la quantification des poids uniquement peut ainsi permettre une inférence plus rapide en réduisant la quantité de mouvements de mémoire. Cependant, le développement de noyaux hautes performances pour les LLMs à poids quantifiés présente des défis substantiels, en particulier lorsque les poids sont compressés à des largeurs de bits non uniformément divisibles (par exemple, 3 bits) avec une quantification non uniforme basée sur des tables de correspondance (LUT). Cet article décrit FLUTE, un moteur de table de correspondance flexible pour les LLMs quantifiés par LUT, qui utilise une restructuration hors ligne de la matrice de poids quantifiée pour minimiser les manipulations de bits associées au dépaquetage, ainsi qu'une vectorisation et duplication de la table de correspondance pour atténuer les contraintes de bande passante de la mémoire partagée. Pour des tailles de lot inférieures à 32 et une taille de groupe de quantification de 128 (typique dans l'inférence des LLMs), le noyau FLUTE peut être 2 à 4 fois plus rapide que les noyaux GEMM existants. En tant qu'application de FLUTE, nous explorons une extension simple à la quantification NormalFloat basée sur des tables de correspondance et l'appliquons pour quantifier LLaMA3 à diverses configurations, obtenant des performances de quantification compétitives par rapport à des bases de référence solides tout en obtenant une augmentation du débit de bout en bout de 1,5 à 2 fois.

English

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.