PageWeaver improves GPU efficiency for sparse attention in AI models
PageWeaver: KV-Guided Query Unions for Sparse Attention
Distributed, Parallel, and Cluster ComputingMachine Learning
Summary
Models that use sparse attention only look at small parts of data to speed up processing, but this doesn't always use GPUs efficiently. The people behind PageWeaver created a way to group these small parts together so GPUs work better without mixing up original data. They tested this on specific GPUs and found meaningful speed improvements in handling large context lengths. Their approach can reduce latency and boost throughput for some cases.
What this means in practice
- •For ai system developers: Enable faster inference in large transformers by optimizing sparse attention execution on GPUs to reduce latency and increase throughput with minimal data reshaping.
- •For high-performance computing engineers: Improve GPU kernel designs to better manage memory and compute cycles when working with sparse attention workloads, especially for long context processing.
Authors
Zhiyuan Li, Zihan Li, Zefang Yuan, Lei Wang, Hao Wang
Abstract
Dynamic sparse attention limits the KV pages selected by each query, but a small support does not necessarily yield efficient GPU work. Query unions share page loads and populate Tensor Core tiles; their cost depends on which queries are grouped together. We present PageWeaver, an execution design that uses selected-page affinity to assemble query groups while preserving each query's original support and complete output ownership. A bounded GPU search produces query IDs, and an ID-aware two-CTA kernel consumes them without materializing reordered Q tensors or cross-page partial outputs. A direct KV-page union implementation provides a complementary design study of nonlocal reuse and reduction cost. With FP8 KV throughout, the H200 Union8 implementation achieves a 1.70x geometric-mean complete-call speedup over the measured FlashInfer path on six captures. Online regrouping further lowers latency by 3.26-7.66% on five selected 64K-context captures. Whole-model prefill throughput is 7.88-14.36% above the tested native path; the incremental regrouping benefit is smaller, with observed median gains of 0.47-0.73% at 32K/64K and regressions at 8K. A B300 comparison identifies cases where preparation cost and a stronger native kernel remove the advantage. These results separate execution-group reuse from the complete cost of exploiting it online.