TurboQuant и PolarQuant: математика квантизации KV-кэша в LLM - UPROGER

Автор этого разбора провёл 31 час, разбирая математику TurboQuant — и написал этот разбор специально для тех, кто не хочет делать это самостоятельно. В статье мы шаг за шагом разберём, как работает TurboQuant, действительно ли он отличается от современных методов квантизации вроде NVFP4 от NVIDIA, и что лежит в основе метода PolarQuant.

Проблема KV-кэша в трансформерах

Каждая языковая модель на основе трансформера вычисляет механизм внимания одинаково. Для каждого токена модель создаёт три вектора: query (что я ищу?), key (что я содержу?) и value (какую информацию я несу?). Оценки внимания вычисляются как softmax(Q·Kᵀ/√d)·V: запрос скалярно перемножается со всеми ключами, чтобы определить важные токены, а затем берётся взвешенная сумма значений.

Трюк, ускоряющий авторегрессивную генерацию — это KV-кэширование. Как только ключевой и ценностный векторы токена вычислены, они не меняются на протяжении всей последовательности. Поэтому их сохраняют и используют повторно для каждого последующего токена запроса.

Проблема в том, что этот кэш растёт линейно вместе с длиной последовательности. Каждый новый токен добавляет по одному вектору ключа и значения на каждый слой и каждую голову. Для Llama-3.1-8B с 32 слоями, 8 KV-головами, размерностью головы 128 и контекстом 128K это: 128 000 × 32 × 8 × 128 × 2 (K и V) × 2 байта (FP16) = 16 ГБ KV-кэша для одной пользовательской сессии. Добавьте несколько параллельных сессий на одном GPU — и KV-кэш становится главным узким местом по памяти.

Существующие решения: NVFP4 двухуровневая квантизация

На сегодняшний день наиболее агрессивным решением для квантизации является NVFP4 с двухуровневой схемой. Вот как она работает:

Сканируем всю матрицу и находим глобальный максимум.
Находим максимум для каждого блока из 16 элементов — локальная гранулярность.
Делим числа на локальный максимум, а масштабирующие коэффициенты — на глобальный максимум.
Приводим числа к ближайшему 4-битному значению.
Перемножаем две матрицы (веса и активации) с помощью специализированных ядер, которые внутренне восстанавливают масштабирующие факторы и возвращают результат в полную точность.

Проблема «корзин» при квантизации

Любой метод квантизации — это по сути отображение. Вы берёте непрерывное число и присваиваете его ближайшей «корзине». При 4 битах у вас 16 корзин, при 2 битах — 4. Главная проблема не в самом отображении, а в том, что нужно заранее знать, где расставить корзины — а это требует предварительного анализа данных.

Для каждого блока значений вычисляется масштабирующий коэффициент (максимальное значение) и нулевая точка (смещение), которые хранятся с полной 16-битной точностью вместе с квантизованными значениями. Эти константы нормализации — чистые накладные расходы. Один выброс портит точность всего блока.

Но что если бы мы знали распределение данных заранее? Что если бы можно было гарантировать, что все значения сгруппированы в предсказуемом, плотном диапазоне? Именно это и предлагает PolarQuant.

PolarQuant: ключевая идея

Исследователи Google Research утверждают, что они действительно могут привести данные к плотно сгруппированному распределению. Метод основан на двух свойствах многомерных нормальных случайных величин.

Факт 1: случайное предобусловливание

Если умножить любой фиксированный вектор на случайную матрицу с элементами из нормального распределения (колоколообразная кривая), результат будет представлять собой многомерное гауссово распределение с центром в нуле и дисперсией, равной квадрату длины исходного вектора:

S·x ~ N(0, ‖x‖² · Iₘ)

Факт 2: концентрация нормы

Если каждая координата вектора взята из стандартного нормального распределения, длина этого вектора следует обобщённому гамма-распределению. В многомерном пространстве эта длина концентрируется вблизи √d.

Объединяя оба факта: предобусловливание умножает наш вектор на случайную матрицу S, получая на выходе вектор, координаты которого ведут себя как образцы из стандартного нормального распределения. Это значит, что длина результирующего вектора будет стабильно концентрироваться вблизи √d в высокой размерности.

Полярное преобразование: от декартовых к полярным координатам

После предобусловливания алгоритм рекурсивно преобразует вектор в полярные координаты. Вот как это работает:

Уровень 1. Группируем координаты d-мерного вектора в пары (x₁, x₂). Каждая пара преобразуется в полярные координаты (r, θ), где θ = arctan(x₂/x₁). Получаем d/2 радиусов и d/2 углов. Углы сохраняем.
Уровень 2. Берём все радиусы как новые пары координат и снова применяем полярное преобразование. Поскольку радиусы всегда положительны, углы теперь лежат в диапазоне [0, π/2] — а не [0, 2π], как на уровне 1.
Последующие уровни. Продолжаем рекурсивно log₂(d) раз. На каждом уровне каждый радиус суммирует всё больше координат, и из-за факта 1 их нормы концентрируются всё теснее вокруг √d.

Ключевое наблюдение: после предобусловливания x₁ и x₂ оба взяты из одного распределения N(0, σ²). Операция arctan(x₂/x₁) — это arctan от отношения двух одинаково распределённых величин. Статистически это отношение стремится к 1, и arctan(1) = π/4 = 45°.

Согласно лемме из исходной статьи, угловое распределение описывается формулой:

f_Θ(θ) = Γ(d) / (2^(d-2) · Γ(d/2)²) · sin^(d-1)(2θ),
E[Θ] = π/4,  Var(Θ) = O(1/√d)

Почему каждый следующий уровень точнее

На уровне 3 каждый радиус суммирует 4 координаты, на уровне 4 — 8, а на уровне 7 — 64. Чем длиннее подвектор, тем теснее его норма сконцентрирована вблизи своего ожидаемого значения (Факт 1). Когда вы делите две нормы 64-мерных подвекторов, отношение чрезвычайно близко к 1, и arctangent практически равен 45 градусам. На высоких уровнях рекурсии углы настолько предсказуемы, что их можно квантизовать буквально одним битом.

Алгоритм PolarQuant на практике

Шаг 1: предобусловливание

Строим матрицу вращения: генерируем случайную матрицу и применяем к ней алгоритм Грама–Шмидта для ортогонализации. Эта матрица сохраняется и переиспользуется.

Шаг 2: полярное преобразование

Применяем матрицу вращения к входным данным, а затем рекурсивно выполняем полярное преобразование, сохраняя углы на каждом уровне:

for level in range(n_levels):
    a = r[0::2]  # нечётные элементы
    b = r[1::2]  # чётные элементы
    level_angles = torch.atan2(b, a)
    if level == 0:
        level_angles = level_angles % (2 * pi)  # уровень 1: [0, 2π)
    # уровни 2+: радиусы всегда положительны → [0, π/2]
    new_r = torch.sqrt(a**2 + b**2)
    angles.append(level_angles)
    r = new_r

Шаг 3: построение кодовых книг

Весь смысл математики в том, что мы заранее знаем точное распределение углов на каждом уровне — до начала инференса. Это позволяет заранее вычислить оптимальные центроиды корзин и сохранить их в таблицах поиска (SMEM). При квантизации просто находим ближайший центроид:

codebooks = [
    build_codebook(n_bits=4, lo=0, hi=2*pi, level=0),   # уровень 1: 16 корзин
    build_codebook(n_bits=2, lo=0, hi=pi/2, level=1),   # уровень 2: 4 корзины
    build_codebook(n_bits=2, lo=0, hi=pi/2, level=2),   # уровень 3: 4 корзины
    build_codebook(n_bits=2, lo=0, hi=pi/2, level=3),   # уровень 4: 4 корзины
]

Для построения каждой кодовой книги запускается алгоритм Ллойда: границы корзин устанавливаются посередине между соседними центроидами, затем центроиды пересчитываются как средневзвешенные значения всех точек в корзине (взвешенные по вероятности распределения). Этот процесс повторяется до сходимости. Центроиды смещаются к пику распределения (π/4), а не к геометрическому центру корзины.

Степень сжатия

Рассмотрим конкретный пример: 128-мерный токен после квантизации хранится как:

64 индекса угла по 4 бита (уровень 1, углы охватывают полный круг): 256 бит
32 индекса по 2 бита (уровень 2): 64 бита
16 индексов по 2 бита (уровень 3): 32 бита
8 индексов по 2 бита (уровень 4): 16 бит
8 оставшихся радиусов в полной 16-битной точности: 128 бит

Итого: 256 + 64 + 32 + 16 + 128 = 496 бит на токен. Исходный FP16-вектор занимал 128 × 16 = 2048 бит. Степень сжатия: 4,13×.

TurboQuant vs. NVFP4: три ключевых отличия

Теперь ответим на исходный вопрос: чем TurboQuant/PolarQuant принципиально отличается от NVFP4?

Нет масштабирующих коэффициентов на блок. NVFP4 хранит отдельные масштабные факторы для каждого блока — это чистые накладные расходы. PolarQuant не делает этого вообще.
Неравномерные корзины, учитывающие распределение. NVFP4 использует равномерно расставленные корзины — это всё, что он может. PolarQuant заранее знает точное распределение углов и расставляет корзины оптимально, с концентрацией вблизи π/4.
Полностью аналитическое предвычисление. NVFP4 требует калибровочных данных для PTQ (с потерей качества) или онлайн-сканирования максимумов во время инференса (с потерей скорости). PolarQuant предвычисляет всё аналитически — бесплатно, без компромиссов.

Производительность ядер

Ядра для PolarQuant были созданы с помощью агентов, использующих авто-тюнинг фреймворк Андрея Карпатого. Результаты оказались неоднозначными. Наиболее оптимизированное ядро достигло ~75% производительности cuBLAS при длинах последовательностей 65K и 512K. При любой длине последовательности менее 8K время упиралось в запуск ядра. CuBLAS стабильно превосходит на ~50% во всех остальных случаях.

75% от скорости cuBLAS пока недостаточно для промышленного использования. Однако авторы планируют продолжить работу над более оптимизированными ядрами.

Итог

PolarQuant — это элегантный метод квантизации KV-кэша, который использует математическую гарантию: после случайного предобусловливания углы полярного преобразования следуют точно известному распределению с пиком у π/4. Это позволяет строить оптимальные кодовые книги аналитически, без калибровочных данных, без накладных расходов на масштабирующие коэффициенты, и достигать степени сжатия 4,13× при хорошем сохранении точности. Главным нерешённым вопросом остаётся производительность CUDA-ядер — но это уже вопрос инженерии, а не математики.

Исходный код ядер доступен на GitHub: github.com/AliesTaha/polar_quant

Просмотры: 785