VSPrefill: Vertical-Slash Sparse Attention with Lightweight Indexing for Long-Context Prefilling
Guanzhong Chen ⋅ Tianhao Gao ⋅ Kai Yang
Abstract
The computational complexity of standard self-attention grows quadratically with sequence length, making long-context prefilling a major latency bottleneck in modern large language model (LLM) serving. Existing sparse prefill methods often model each attention head independently, overlooking shared sparse patterns within key-value (KV) groups, or rely on coarse block-level selection that can miss token-level salient positions in vertical-slash patterns. Through empirical analysis of long-context LLMs with grouped-query attention, we are, to the best of our knowledge, the first to empirically discover group-structured vertical-slash sparsity: vertical-slash patterns exhibit group-level structure, with similar patterns within each KV group and distinct patterns across groups, while salient attention remains concentrated at fine-grained token-level granularity rather than coarse block-level. Motivated by this observation, we propose the Vertical-Slash Prefill (VSPrefill), a group-aware fine-grained sparse attention mechanism equipped with VSIndexer, a lightweight trainable index module that predicts context-aware vertical and slash importance scores per KV group from KV representations. The VSIndexer is trained to approximate ground-truth attention distributions, requiring no alterations to the original model weights. During inference, we introduce an adaptive cumulative score threshold strategy to allocate per-group sparsity budgets, coupled with a fused kernel to execute sparse attention via on-the-fly index merging. At the 128K extreme context length, VSPrefill achieves 1.32$\times$ higher prefill speedup than the strong sparse prefill baseline FlexPrefill at comparable accuracy. Across long-context benchmarks, VSPrefill preserves near-FlashAttention accuracy while achieving up to a 4.95$\times$ prefill speedup over the FlashAttention full-attention baseline.
Chat is not available.
Successful Page Load