Разреженное блочно-разреженное внимание через перестановку токенов
Sparser Block-Sparse Attention via Token Permutation
October 24, 2025
Авторы: Xinghao Wang, Pengyu Wang, Dong Zhang, Chenkun Tan, Shaojun Zhou, Zhaoxiang Liu, Shiguo Lian, Fangxu Liu, Kai Song, Xipeng Qiu
cs.AI
Аннотация
Увеличение длины контекста в больших языковых моделях (LLM) дает значительные преимущества, но вычислительно затратно. Основная причина затратности — механизм самовнимания, чья сложность O(N^2) по отношению к длине последовательности создает серьезное узкое место для памяти и задержек. К счастью, матрица внимания часто является разреженной, особенно для длинных последовательностей, что указывает на возможность оптимизации. Блочно-разреженное внимание стало перспективным решением, которое разбивает последовательности на блоки и пропускает вычисления для части этих блоков. Однако эффективность этого метода сильно зависит от базовых паттернов внимания, что может приводить к неоптимальной разреженности на уровне блоков. Например, важные ключевые токены для запросов в пределах одного блока могут быть распределены по множеству других блоков, что приводит к вычислительной избыточности. В данной работе мы предлагаем Permuted Block-Sparse Attention (PBS-Attn) — модульный метод, который использует свойства перестановки внимания для увеличения разреженности на уровне блоков и повышения вычислительной эффективности предварительного заполнения (prefilling) в LLM. Мы провели всесторонние эксперименты на сложных реалистичных наборах данных с длинным контекстом, показав, что PBS-Attn стабильно превосходит существующие методы блочно-разреженного внимания по точности модели и близко соответствует базовому уровню полного внимания. Благодаря нашим специализированным ядрам permuted-FlashAttention, PBS-Attn обеспечивает сквозное ускорение до 2.75 раз при предварительном заполнении длинного контекста, подтверждая свою практическую жизнеспособность. Код доступен по адресу https://github.com/xinghaow99/pbs-attn.
English
Scaling the context length of large language models (LLMs) offers significant
benefits but is computationally expensive. This expense stems primarily from
the self-attention mechanism, whose O(N^2) complexity with respect to
sequence length presents a major bottleneck for both memory and latency.
Fortunately, the attention matrix is often sparse, particularly for long
sequences, suggesting an opportunity for optimization. Block-sparse attention
has emerged as a promising solution that partitions sequences into blocks and
skips computation for a subset of these blocks. However, the effectiveness of
this method is highly dependent on the underlying attention patterns, which can
lead to sub-optimal block-level sparsity. For instance, important key tokens
for queries within a single block may be scattered across numerous other
blocks, leading to computational redundancy. In this work, we propose Permuted
Block-Sparse Attention (PBS-Attn), a plug-and-play method that
leverages the permutation properties of attention to increase block-level
sparsity and enhance the computational efficiency of LLM prefilling. We conduct
comprehensive experiments on challenging real-world long-context datasets,
demonstrating that PBS-Attn consistently outperforms existing block-sparse
attention methods in model accuracy and closely matches the full attention
baseline. Powered by our custom permuted-FlashAttention kernels, PBS-Attn
achieves an end-to-end speedup of up to 2.75times in long-context
prefilling, confirming its practical viability. Code available at
https://github.com/xinghaow99/pbs-attn