Оригинальное название: LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers
Циклические трансформеры многократно вызывают один и тот же набор слоёв, достигая производительности моделей, в несколько раз больших по размеру, при меньшем количестве параметров. Однако на каждом шаге цикла необходимо сохранять собственные ключи и значения (KV), из-за чего размер кэша линейно растёт с увеличением числа шагов цикла, становясь основным узким местом по памяти. Авторы заметили высокую избыточность KV вдоль оси цикла и значительные различия в чувствительности к квантованию на разных шагах цикла. В связи с этим они предложили метод LoopVQ: объединение каналов нескольких последовательных шагов цикла в один подвектор с общим индексом кодовой книги (векторное квантование вдоль оси цикла), а затем дообучение кодовой книги с помощью взвешенного по Фишеру алгоритма k-means, чтобы выходные данные квантованной модели максимально приближались к результатам модели с полной точностью (во время обучения на каждом шаге цикла считывается квантованный кэш). Этот метод ортогонален подходам, ориентированным на обычные LLM, и может использоваться совместно с преобразованием ключей из метода TaSQ.
Лайки 8 · Звёзды в репозитории 3
Обоснование
- Авторы утверждают, что при 0.5 битах на элемент метод превосходит существующие методы векторного квантования KV, ориентированные на обычные LLM, в задачах общего назначения, длинных цепочек рассуждений и поиска в длинном контексте для моделей Ouro и Huginn; при совместном использовании с TaSQ сохраняется большая часть точности BF16 даже при 0.25 битах.
- Системные преимущества подтверждены на практике: после реализации в SGLang авторы добились почти двукратного снижения задержки декодирования для контекста в 32K токенов, а пиковая пропускная способность обслуживания для моделей Ouro 1.4B / 2.6B увеличилась в 4.3 и 6.1 раз соответственно по сравнению с базовой линией BF16.
- Код распространяется под лицензией Apache-2.0: в репозитории доступны процесс калибровки, форк SGLang с упакованным хранением индексов и слитым ядром декодирования, а также скрипты для воспроизведения всех таблиц из статьи.
Использование
- Запустите ./setup_envs.sh для создания трёх окружений (движок loopvq-serve, окружение оценки loopvq-eval, окружение калибровки loopvq-build), установите переменную окружения MODELS_DIR, указывающую на модели Ouro-1.4B, Ouro-2.6B, Ouro-*-Thinking и huginn-0125; для воспроизведения таблиц общего назначения выполните MODEL=ouro_1p4b GPUS="0 1 2 3" bash experiments/loopvq_general/build.sh, затем MODEL=ouro_1p4b bash experiments/loopvq_general/run_all.sh, и наконец python3 experiments/loopvq_general/make_table.py.
- Запуск сервиса: в окружении loopvq-serve выполните PREFIX_TOKENS=8 RECENT_TOKENS=64 PREFILL_BACKEND=triton ./scripts/serve_method.sh $MODELS_DIR/Ouro-1.4B vq <каталог_кодовых_книг> --port 30000 --disable-piecewise-cuda-graph, чтобы получить совместимый с OpenAI интерфейс SGLang; начальные и последние токены сохраняются в формате bf16.
Риски
- Метод применим только к циклическим трансформерам (в статье использованы модели Ouro 1.4B / 2.6B и Huginn-3.5B); его нельзя напрямую применять к обычным LLM. Для калибровки использовалось всего 16 последовательностей из набора GPQA.
- Показатели задержки и пропускной способности получены в среде авторов; количество лайков 8 и звёзд в репозитории 3 отражают лишь уровень внимания сообщества.