Одно кэш-состояние вместо N: LinearKV вдвое ускоряет вывод в гибридных LLM

Исследователи предложили новый подход LinearKV для ускорения инференса гибридных больших языковых моделей (LLM). Такие архитектуры сочетают слои attention с линейными рекуррентными слоями, что создаёт сложности для стандартных методов кэширования. Новая техника опирается на позиционно-независимое кэширование (PIC) и использует всего одно кэш-состояние для линейных слоёв вместо набора состояний.

В гибридных LLM, таких как GDN и Mamba-2, линейные слои выдают только фиксированное по размеру состояние, в отличие от attention, где каждый токен имеет свой ключ и значение. Это ломает классические механизмы PIC, рассчитанные на полное внимание. Авторы LinearKV нашли решение: линейные слои инициализируются из одного сопоставленного локального состояния, а attention-слои работают как раньше.

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

Эксперименты на трёх гибридных моделях и трёх PIC-селекторах показали, что одного состояния достаточно для восстановления качества. На двух GDN-моделях LinearKV сравнялся с более сложным методом HYPIC, восстанавливая до 92% полного качества. На Mamba-2 результат оказался ещё показательнее: точное составление всех состояний (как в HYPIC) приводило к коллапсу — всего 46,6% качества при использовании селектора EPIC, тогда как одиночный инициализатор давал 86,8%.

Помимо качества, LinearKV выигрывает по скорости. Одиночное кэш-состояние сокращает время до первого токена до 0,46 от полного prefill, то есть более чем вдвое ускоряет выдачу первого ответа. В то же время точное составление всех состояний добавляет 5–17% накладных расходов. Результаты стабильны на LongBench QA и RULER при длинах контекста от 8K до 32K.

Авторы подчёркивают, что их метод не требует дообучения модели и может быть легко интегрирован в существующие системы обслуживания LLM. Это практическое решение для гибридных архитектур, которые становятся всё более популярными благодаря балансу качества и эффективности.