К содержанию
Новости

Сжатие KV‑кэша становится ключевой задачей для LLM: подходы TurboQuant, OSCAR и EpiCache

Сжатие KV‑кэша становится ключевой задачей для LLM: подходы TurboQuant, OSCAR и EpiCache
Фото: MarkTechPost

В больших языковых моделях с длинным контекстом узким местом всё чаще оказывается не размер весов, а память под KV‑кэш, именно она начинает диктовать ограничения. Во время декодирования трансформер сохраняет key‑ и value‑векторы для каждого токена на каждом слое, чтобы не пересчитывать attention. Кэш растёт линейно вместе с длиной последовательности и размером батча, поэтому при длинном контексте его объём легко обгоняет сами веса модели. Например, у Llama‑3.1‑70B в BF16 KV‑кэш занимает около 0.31 MB на токен (80 layers × 8 KV heads × 128 head-dim × 2 tensors × 2 bytes). При 128K токенов это около 40 GB, а при 1M уже более 300 GB, что превышает 140 GB весов модели. Каждый новый токен требует чтения всего кэша из HBM, поэтому декодирование в итоге упирается в пропускную способность памяти.

Сжатие KV‑кэша напрямую влияет и на стоимость, и на задержку генерации. Существующие подходы обычно делят на несколько направлений: token eviction (H2O, SnapKV), quantization (KIVI, GEAR), low-rank projection (Palu), merging (KVMerger) и architectural sharing (MLA). Для квантизации одной из главных проблем остаются outlier channels, каналы с аномально большими значениями, которые задают диапазон квантования. Базовым ориентиром здесь стала работа KIVI. В ней показано, что у key‑векторов между токенами сохраняются фиксированные outlier‑каналы, тогда как у value‑векторов такой структуры нет. Поэтому ключи квантуют per‑channel, а значения per‑token. Такой 2‑битный метод без дополнительной настройки снижает пиковое потребление памяти всей системы примерно в 2.6×.

В более свежих исследованиях заметен интерес к сверхнизкой битности. TurboQuant от Google и NYU предлагает двухэтапную схему без использования данных. Сначала каждый вектор случайно вращают, чтобы координаты стали почти независимыми и примерно гауссовыми; затем применяют предвычисленный скалярный квантователь Lloyd–Max. После этого к остатку добавляется 1‑bit Quantized Johnson–Lindenstrauss (QJL) transform, который даёт несмещённую оценку attention logits и не требует дополнительных затрат на нормализацию. Авторы показывают, что искажение TurboQuant остаётся в пределах небольшого постоянного множителя (≈ 2.7×) от информационно‑теоретической нижней границы. При 4× сжатии достигается практически полноточная recall в тесте Needle‑in‑a‑Haystack, при 3.5 bits качество остаётся нейтральным, а при 2.5 bits на канал наблюдается лишь небольшая деградация.

Есть и другой подход. Метод OSCAR от Together AI использует attention‑aware вращение, которое вычисляется во время одноразовой offline‑калибровки. Keys переводятся в собственный базис ковариации запросов, values, в ковариацию значений, взвешенную по score. Отдельно Apple предложила EpiCache, он нацелен на другую проблему KV‑кэша.

Ключевые факты

  • У Llama-3.1-70B в формате BF16 размер KV‑кэша составляет около 0.31 MB на токен (80 слоёв × 8 KV‑голов × 128 head-dim × 2 тензора × 2 байта).

  • При контексте 128K токенов KV‑кэш занимает около 40 GB, а при 1M токенов превышает 300 GB, что больше чем примерно 140 GB весов модели.

  • Метод KIVI использует квантование ключей по каналам и значений по токенам и снижает пиковое потребление памяти (включая веса) примерно в 2.6 раза.

  • TurboQuant показывает практически полное совпадение с полноточной точностью в тесте Needle-in-a-Haystack при 4× сжатии; в работе указана нейтральность качества при 3.5 бита на канал и лишь небольшая деградация при 2.5 бита.