🔬 Метод

В основе идеи статьи лежит следующее наблюдение - query меняется не очень сильно между соседними токенами, косинусная близость порядка 0.9. Отсюда мысль - давайте будем держать в памяти, то что мы подгрузили на прошлом шаге и менять по мере необходимости.

В качестве кратких summary 📜 по группам токенов, используемых для определения важных токенов на оффлоадинге, берутся ландмарки из QuEST (min-max pooled ключи).

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

Кроме того, подгрузку KV можно перекрыть с вычислениями в attention / mlp, тем самым дополнительно сэкономив время.

🧪 Эксперименты

Предложенный подход валидируют на семействах моделей Llama-3, Qwen-2.5 и дистиллах DeepSeek-R1. Удается достичь качества оффлоадинга, при этом будучи в разы, а то и десяток раз быстрее бейзлайнов. Правда, стоит заметить, что везде инференс поверх HF transformers, так что это гонки на инвалидных колясках.

💡 Выводы

Вроде бы не слишком сложно реализуемый и логичный метод оффлоадинга. Но практическое применение будет упираться в возможность интеграции в эффективные фреймворки инференса vllm / sglang.