CompactAttention: ускорение чанкового префилла с помощью блочно-объединённого выбора KV

Аннотация

Фрагментированная предварительная загрузка (chunked prefill) стала широко используемой стратегией обслуживания для больших языковых моделей с длинным контекстом, однако эффективное вычисление внимания в этом режиме остаётся сложной задачей. Существующие методы разреженного внимания в основном предназначены для однократной предварительной загрузки и неэффективно переносятся на фрагментированную предварительную загрузку: блочно-разреженные ядра теряют эффективность, когда длина запроса ограничена размером фрагмента, в то время как поиск мелкозернистых шаблонов становится затратным при повторении на накопленном KV-кэше для каждого фрагмента. QUOKA — недавний метод, непосредственно ориентированный на фрагментированную предварительную загрузку, — позволяет избежать накладных расходов разреженных ядер, но полагается на выбор KV на уровне токенов с субдискретизацией запросов, что может пропускать специфичные для запроса KV-записи и вносить явные накладные расходы на копирование KV. Для устранения этих ограничений мы предлагаем CompactAttention — механизм внимания для фрагментированной предварительной загрузки, основанный на блочно-объединённом выборе KV (Block-Union KV Selection). CompactAttention рассматривает двумерные блочно-разреженные маски как сигналы для выбора KV, а не как планы непосредственного выполнения разреженных ядер, и преобразует их в осведомлённые о GQA таблицы KV-блоков для каждой группы с помощью объединения Q-блоков и внутригруппового объединения. Такая конструкция создаёт минимальные таблицы блоков, которые сохраняют все KV-блоки, выбранные входными масками, с учётом ограничений постраничного выполнения, что позволяет обращаться к выбранным KV-блокам на месте без явного уплотнения KV. На модели LLaMA-3.1-8B-Instruct CompactAttention сохраняет точность, близкую к точности плотного внимания на бенчмарке RULER, обеспечивая при этом ускорение внимания до 2,72 раза при длине контекста 128K в режиме фрагментированной предварительной загрузки.

English

Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging. Existing sparse attention methods are primarily designed for one-shot prefill and do not translate efficiently to chunked prefill: block-sparse kernels lose efficiency when the query length is limited by the chunk size, while fine-grained pattern search becomes costly when repeated over the accumulated KV cache at every chunk. QUOKA, a recent method that directly targets chunked prefill, avoids sparse-kernel overhead but relies on query-subsampled, token-level KV selection, which can miss query-specific KV entries and introduce explicit KV-copy overhead. To address these limitations, we propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection. CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction produces the minimal block tables that preserve all KV blocks selected by the input masks under paged execution constraints, enabling selected KV blocks to be accessed in place without explicit KV compaction. On LLaMA-3.1-8B-Instruct, CompactAttention maintains accuracy close to dense attention on the RULER benchmark while delivering up to 2.72times attention speedup at 128K context length under chunked prefill.