Быстрые умножения матриц для квантованных LLM с таблицей поиска

Fast Matrix Multiplications for Lookup Table-Quantized LLMs

July 15, 2024

Авторы: Han Guo, William Brandon, Radostin Cholakov, Jonathan Ragan-Kelley, Eric P. Xing, Yoon Kim

cs.AI

Аннотация

Развертывание больших языковых моделей (LLM) часто ограничивается пропускной способностью памяти, где основным узким местом является стоимость передачи параметров модели из глобальной памяти GPU в ее регистры. В сочетании с пользовательскими ядрами, объединяющими операции деквантования и матричного умножения, квантование только весов может обеспечить более быстрое вывода, уменьшая объем перемещения памяти. Однако разработка высокопроизводительных ядер для весово-квантованных LLM представляет существенные вызовы, особенно когда веса сжимаются до неравномерных битовых ширин (например, 3 бита) с неравномерным квантованием через таблицу поиска (LUT). В данной статье описывается FLUTE, гибкий движок таблицы поиска для LUT-квантованных LLM, который использует офлайн-перестройку квантованной матрицы весов для минимизации битовых манипуляций, связанных с распаковкой, а также векторизацию и дублирование таблицы поиска для смягчения ограничений на пропускную способность общей памяти. При размерах партий < 32 и размере группы квантования 128 (типичном для вывода LLM) ядро FLUTE может быть в 2-4 раза быстрее существующих ядер GEMM. В качестве применения FLUTE мы исследуем простое расширение квантования на основе таблицы поиска NormalFloat и применяем его для квантования LLaMA3 в различные конфигурации, достигая конкурентоспособной производительности квантования по сравнению с сильными базовыми значениями, а также увеличивая пропускную способность от начала до конца на 1,5-2 раза.

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.