Monarch v3: Вдохновленная NES система постраничной организации KV-памяти для ускорения вывода LLM на 78%

Что делает Monarch v3

Monarch v3 — это реализация с открытым исходным кодом NES-вдохновленного разбиения памяти на страницы для вывода трансформеров, которая решает проблему линейного роста KV-кэша с увеличением длины последовательности. К 4K токенам большая часть KV-кэша остаётся неиспользуемой, но при этом занимает VRAM в полной точности.

Как это работает

Система разделяет KV-кэш на две области:

• Горячая область: Недавние токены хранятся в полной точности
• Холодная область: Старые токены сжимаются до ~20 байт каждый (против 64-128 байт горячих)

Четыре компонента работают вместе:

• Сжатие TurboQuant: Квантует KV в 4-битные целые числа с полярным кодированием и коррекцией остатков, достигая ~97% уменьшения размера с потерей перплексии ~0.3%
• Скользящее окно вытеснения: Последние N токенов остаются горячими по умолчанию, старые токены сжимаются в холодное хранилище
• Продвижение на основе внимания: Токены с высоким вниманием возвращаются в горячую область с липким механизмом для предотвращения трэшинга
• Свопинг страниц: Небольшие пакеты холодных токенов материализуются при доступе с локальным циклом декодирования вместо пакетного matmul

Результаты тестирования

Настройка: TinyLlama-1.1B fp16, 50 сгенерированных токенов

• Стандартный: 17.01 токенов/с, 2112 МБ VRAM
• Monarch-v3: 30.42 токенов/с, 2131 МБ VRAM, 512 горячих токенов, 1024 холодных токенов
• Прирост: +78.7% пропускной способности, +0.9% VRAM

Упрощённый цикл декодирования

for step in 1..100:
    q = project_query(next_token)
    # Вычисление внимания: только горячие (быстро)
    scores_hot = q @ kv_hot.T
    # Доступ к холодным при высоком внимании (редко)
    if max(scores_hot) < threshold:
        kv_cold_promoted = decompress(cold_pages)
        scores_cold = q @ kv_cold_promoted.T
        # Переместить в горячие для следующего шага
    # Агрегация, softmax, применение внимания ...
    # Вытеснение старых токенов из горячих → холодные
    if len(kv_hot) > window_size:
        evict_oldest_to_cold()

Текущий статус

• Реализация: Работает на Hugging Face Transformers с пользовательским бэкендом кэша
• Лицензия: Apache 2.0
• Статья: Полная техническая спецификация доступна
• Далее: Запланировано слияние CUDA ядер для холодной декомпрессии

Как попробовать

git clone https://github.com/JohannaWeb/Monarch.git
cd Monarch
pip install -r requirements.txt
python train_tinyllama_fp16.py
python src/benchmark_monarch.py \
    --model models/tinyllama_fp16 \
    --mode both \
    --max-new-tokens 100 \
    --promotion-threshold 0.15 \
    --sticky-threshold 3 \
    --json

Ограничения

Подход основан на недавности (недавние токены = высокое внимание), что работает для большинства задач, но может не подходить для рабочих нагрузок с интенсивным поиском. Извлечение внимания доступно в базовых моделях, но не в чат-вариантах; запасной вариант использует только оконное разбиение на страницы.