Конец эпохи квадратичной сложности: как LLM научились жульничать с памятью

Почему ваш GPU плачет, когда вы говорите "бесконечный контекст"

Забудьте про маркетинговые обещания. Когда компания заявляет "поддержка 1 миллиона токенов", они обычно умалчивают о том, что для этого потребуется серверная стойка стоимостью с квартиру в Москве. Реальность такова: стандартный механизм внимания в трансформерах требует O(n²) операций и O(n) памяти на каждый слой. Для 100К токенов это уже терабайты VRAM. Чистая фантастика.

Но есть и хорошие новости. За последние два года исследователи придумали десятки способов обмануть математику. Не решить проблему квадратичной сложности (это пока невозможно без изменения фундаментальной архитектуры), а сделать так, чтобы модель думала, что она её решила.

KV-кэш: ваш главный враг и лучший друг

Сначала разберёмся, почему вообще возникает проблема. В обычном трансформере для генерации каждого нового токена модель должна вычислить attention между этим токеном и всеми предыдущими. Ключи (K) и значения (V) — это промежуточные представления, которые пересчитывать каждый раз безумно дорого.

Решение? KV-кэш. Сохраняем вычисленные K и V для всех предыдущих токенов, и при генерации нового токена просто подставляем их в формулу внимания. Гениально. Работает. И убивает всю вашу видеопамять.

Для модели Llama 3 70B с контекстом 128К и форматом bfloat16:

# Примерный расчёт памяти для KV-кэша
layers = 80
heads = 64
dim_per_head = 128
context_length = 128000
bytes_per_param = 2  # bfloat16

memory_per_layer = context_length * heads * dim_per_head * bytes_per_param * 2  # *2 для K и V
memory_per_layer /= (1024**3)  # в GB

total_memory = memory_per_layer * layers
print(f"Требуется VRAM только для KV-кэша: {total_memory:.1f} GB")

Цифры получаются пугающие. Даже для скромных 8К токенов KV-кэш съедает гигабайты. Для 128К — десятки гигабайт. Это только кэш, без учёта весов модели.

Сжатие без потерь: когда точность важнее всего

Первый подход — сохранить полную информацию, но упаковать её плотнее. Не квантование (о нём позже), а именно эффективное хранение.

1 Стратегия: бинарная сериализация KV-кэша

Вспомните нашу статью про Binary KV cache. Суть проста до безобразия: вместо того чтобы пересчитывать контекст каждый раз, сохраняем вычисленный KV-кэш на диск в бинарном формате. При следующем запуске — загружаем обратно.

Почему это работает? Потому что 90% контекста в диалоговых системах — это история разговора, которая не меняется между запросами. Зачем вычислять её заново?

# Псевдокод логики бинарного кэша
class BinaryKVCache:
    def __init__(self, model_name, context_hash):
        self.cache_file = f"cache/{model_name}_{context_hash}.bin"
        
    def save(self, k_cache, v_cache):
        # Сериализуем напрямую из GPU памяти
        with open(self.cache_file, 'wb') as f:
            f.write(k_cache.cpu().numpy().tobytes())
            f.write(v_cache.cpu().numpy().tobytes())
        
    def load(self):
        # Десериализуем прямо в VRAM
        if os.path.exists(self.cache_file):
            with open(self.cache_file, 'rb') as f:
                data = f.read()
            # Восстанавливаем тензоры
            return torch.from_numpy(np.frombuffer(data, dtype=np.float16))
        return None

Преимущество: нулевая потеря точности. KV-кэш загружается в том же формате, в котором был сохранён. Недостаток: требует места на диске и не решает проблему с памятью во время выполнения.

2 Стратегия: shared context compression

Более умный подход. Замечаем, что в длинных документах много повторяющихся паттернов. Системные промпты, шаблоны ответов, форматы данных. Можно вычислить KV-кэш для этих повторяющихся частей один раз и переиспользовать.

Как это выглядит на практике:

Квантование KV-кэша: игра в компромиссы

Когда места не хватает, а контекст нужен длинный, начинаются жертвы. Квантование — это уменьшение битности представления чисел. Вместо float16 (16 бит) используем int8 (8 бит) или даже int4 (4 бита).

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

3 Динамическое 8-битное квантование

Алгоритм выглядит так:

1. Вычисляем KV-кэш в полной точности (float16)
2. Находим min и max значения в каждом тензоре
3. Масштабируем значения в диапазон int8 (-128..127)
4. Сохраняем квантованные значения и масштабирующие коэффициенты
5. При использовании — деквантуем обратно в float16

Проблема: масштабирующие коэффициенты тоже занимают место. Для каждого слоя, каждой головы, каждого вектора. Начинается игра в баланс между степенью квантования и overhead.

# Упрощённая реализация динамического квантования
def quantize_kv_cache(k_cache, v_cache, bits=8):
    """Квантует KV-кэш на лету"""
    
    # Для ключей
    k_min = k_cache.min()
    k_max = k_cache.max()
    k_scale = (k_max - k_min) / (2**bits - 1)
    
    k_quantized = ((k_cache - k_min) / k_scale).round().to(torch.int8)
    
    # Для значений
    v_min = v_cache.min()
    v_max = v_cache.max()
    v_scale = (v_max - v_min) / (2**bits - 1)
    
    v_quantized = ((v_cache - v_min) / v_scale).round().to(torch.int8)
    
    return {
        'k_q': k_quantized,
        'v_q': v_quantized,
        'k_scale': k_scale,
        'k_min': k_min,
        'v_scale': v_scale,
        'v_min': v_min
    }

А что насчёт архитектурных хаков?

Самые интересные решения приходят не из мира квантования, а из изменения самого механизма внимания. Если нельзя уменьшить размер KV-кэша, может, можно сделать так, чтобы он был не нужен?

4 Sliding Window Attention

Идея проста как двери: токену не нужно видеть весь контекст. Достаточно окна из N предыдущих токенов. Как в старых RNN, только лучше.

Mistral 7B использует sliding window размером 4096 токенов. Это значит, что даже если контекст 32К, модель видит только последние 4К. KV-кэш не растёт бесконечно, а циклически перезаписывается.

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

5 StreamingLLM и его потомки

Более умный подход. Вместо того чтобы хранить весь KV-кэш, храним:

• Первые несколько токенов (системный промпт, важные инструкции)
• Последние N токенов (актуальный контекст)
• Ключевые токены, выбранные через attention score

Получается гибридный подход. Часть контекста хранится полностью, часть — выборочно, остальное выбрасывается. Экономия памяти до 80% при минимальной потере качества.

# Псевдокод StreamingLLM
class StreamingLLMKVCache:
    def __init__(self, keep_first_tokens=10, keep_last_tokens=512, attention_threshold=0.1):
        self.keep_first = keep_first_tokens
        self.keep_last = keep_last_tokens
        self.threshold = attention_threshold
        
    def compress_cache(self, k_cache, v_cache, attention_scores):
        """Сжимает KV-кэш на основе важности токенов"""
        
        total_tokens = k_cache.shape[1]
        
        # Индексы для сохранения
        keep_indices = set()
        
        # Первые токены всегда сохраняем
        keep_indices.update(range(self.keep_first))
        
        # Последние токены
        keep_indices.update(range(total_tokens - self.keep_last, total_tokens))
        
        # Токены с высокой attention score
        for i, score in enumerate(attention_scores.mean(dim=(0, 1))):  # усредняем по головам
            if score > self.threshold and i not in keep_indices:
                keep_indices.add(i)
        
        # Собираем новый кэш
        indices = sorted(keep_indices)
        return k_cache[:, indices, :], v_cache[:, indices, :]

Практика: собираем всё вместе

Теория теорией, но как это выглядит в реальном коде? Возьмём пример из нашей статьи про долговременную память и добавим сжатие.

Допустим, у нас есть чат-бот с историей диалога. Каждый новый запрос добавляется к истории, и модель должна помнить всё. Наивная реализация быстро съест всю память.

6 Гибридная стратегия для продакшена

Вот что работает на практике:

1. Статическую часть (системный промпт, шаблоны) кэшируем в бинарном формате на диск. Загружаем один раз при старте.
2. Историю диалога храним в сжатом виде через StreamingLLM. Держим только важные фрагменты.
3. Текущий запрос обрабатываем в полной точности.
4. Каждые 100 токенов применяем динамическое 8-битное квантование к старой части истории.

Реализация:

class HybridKVCacheManager:
    """Гибридный менеджер KV-кэша для длинных диалогов"""
    
    def __init__(self, model, cache_dir="kv_cache"):
        self.model = model
        self.cache_dir = cache_dir
        self.static_cache = None  # Кэш для статичных частей
        self.dynamic_cache = {}   # Кэш для истории диалога
        self.current_context = [] # Текущий контекст
        
    def add_to_context(self, text):
        """Добавляет текст в контекст, обновляя кэш"""
        tokens = self.tokenize(text)
        
        # Если контекст слишком длинный, сжимаем старую часть
        if len(self.current_context) + len(tokens) > MAX_CONTEXT:
            self._compress_old_context()
        
        # Добавляем новые токены
        self.current_context.extend(tokens)
        
        # Обновляем динамический кэш
        self._update_dynamic_cache(tokens)
        
    def _compress_old_context(self):
        """Сжимает старую часть контекста"""
        # Оставляем первые 10% и последние 50%
        keep_first = len(self.current_context) // 10
        keep_last = len(self.current_context) // 2
        
        # Вычисляем attention scores для оставшихся токенов
        scores = self._compute_attention_scores()
        
        # Выбираем самые важные токены из середины
        middle_indices = self._select_important_tokens(
            self.current_context[keep_first:-keep_last],
            scores[keep_first:-keep_last]
        )
        
        # Собираем новый контекст
        new_context = (
            self.current_context[:keep_first] +
            [self.current_context[i] for i in middle_indices] +
            self.current_context[-keep_last:]
        )
        
        self.current_context = new_context
        
        # Перестраиваем динамический кэш
        self._rebuild_dynamic_cache()
        
    def _update_dynamic_cache(self, new_tokens):
        """Обновляет динамический кэш с квантованием старых частей"""
        # Для новых токенов — полная точность
        # Для токенов старше 100 шагов — 8-битное квантование
        for i, (k, v) in enumerate(self.dynamic_cache.items()):
            if i < len(self.dynamic_cache) - 100:
                self.dynamic_cache[i] = self._quantize_cache(k, v, bits=8)

Ошибки, которые все совершают (и как их избежать)

Я видел десятки реализаций сжатия KV-кэша. 90% из них содержат одни и те же ошибки.

Ошибка 1: Слепое квантование всего подряд

Квантовать ключи и значения с одинаковой агрессивностью — грубая ошибка. Значения (V) гораздо более чувствительны к потере точности. Ключи (K) можно квантовать сильнее.

Ошибка 2: Игнорирование паттернов доступа

Внимание в трансформерах неравномерное. Первые слои смотрят на локальный контекст, последние — на глобальный. Можно агрессивнее сжимать кэш в первых слоях.

Ошибка 3: Забыть про continuous batching

В продакшене запросы приходят пачками. Если у каждого запроса свой KV-кэш, память умрёт. Нужно учитывать батчинг при проектировании системы сжатия.

Что будет дальше? Прогноз на 2025

Нынешние подходы — это костыли. Элегантное решение придёт с изменением архитектуры внимания. Уже сейчас появляются модели с линейной сложностью (Linformer, Perceiver, LongNet).

Мой прогноз: через год мы увидим модели, которые по умолчанию работают с 1M токенов на потребительском железе. Не через сжатие KV-кэша, а через фундаментально другие механизмы внимания.

А пока — используйте гибридные подходы. Кэшируйте статичное, сжимайте динамичное, квантуйте старое. И не верьте маркетинговым заявлениям про "бесконечный контекст без компромиссов". Их не существует.

P.S. Если хотите глубоко разобраться в оптимизации памяти для локальных моделей, посмотрите нашу статью про стратегии масштабирования локальных LLM. Там есть конкретные цифры и примеры для разного железа.