Разобьем последовательность на префикс P и суффикс S.
Пусть N — числитель, а D — знаменатель. Тогда выход внимания имеет следующий вид:
O = (N_P + N_S) / (D_P + D_S)Предположим, что для текущей query нашлась достаточно близкая (по какой-то мере похожести).
Тогда мы можем переиспользовать предпосчитанные числитель и знаменатель от этой близкой query, и остается только посчитать члены между той query и текущей позицией.
⚙️ Важные практические соображения
Хранить целесообразно небольшой скользящий буфер query. Полный query cache (в силу распространения GQA / MLA) будет в разы дороже KV-кэша, да и выигрыш тем больше, чем больше токенов мы переиспользуем (чем ближе по позиции близкая query к текущей).
Хотим, с одной стороны, переиспользовать как можно больше вычислений, но при этом не терять в качестве.
В качестве меры похожести используется L2-расстояние между pre-RoPE ключами. pre-RoPE критично, иначе близкие семантически query разъедутся из-за позиционных эмбеддингов. Если похожая query не нашлась (минимальная ошибка выше порога), то считаем так, как считали бы при стандартном инференсе.
Под это все у них есть самописные кернелы эффективного поиска похожей query и внимания с переиспользованием вычислений.
📊 Эксперименты
Метод валидируют на моделях семейства Llama-3.1 / GLM-4 / Phi-4 и бенчах на длинный контекст LongBench v1/v2 и RULER.
На практике берут окно из 1024 последних токенов для поиска ближайшей query, а порог близости был как-то подобран вручную.
Качество якобы не просаживается, а где-то даже и растет (даже на бюджетах порядка 1%). Но похоже на скам.
Свои кернелы они интегрировали в SGLang и утверждают, что достигают 4× уменьшения latency attention против full attention и работают быстрее всяких прунингов. Подозрительно, что RocketKV такой медленный…
🚀 На больших батчах и с 99% sparsity говорят, что якобы ускоряют внимание до 40 раз.
💾 Сам метод дает примерно 5% дополнительного оверхеда по GPU-памяти (в зависимости от числа групп в GQA) на длине контекста 128k с настройками по умолчанию.
📝 Вывод
Идейка прикольная, даже удивительно, что до нее додумались сравнительно недавно — только в апреле этого года.
Не хватает замеров на более серьезных моделях и бенчах (типа математики, кода и агентских задач), а также end-to-end сравнений по скорости.