arXiv：cs.LG· Zhiyuan Li, Zihan Li, Zefang Yuan, Lei Wang, Hao Wang·· 4 小时前
PageWeaver：用 KV 引导的查询合并优化稀疏注意力
PageWeaver: KV-Guided Query Unions for Sparse Attention
AI 导读
PageWeaver 是一种利用选中页亲和性组装查询组的执行设计，在不改变每个查询原有支持集的前提下共享 KV 页加载。在 H200 上采用 FP8 KV，其 Union8 实现相较实测 FlashInfer 路径在六组抓取上取得 1.70 倍几何平均完整调用加速，在线重组在五个 64K 上下文抓取上再降 3.26-7.66% 延迟。
正文
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.
| Comments: | 11 pages, 9 figures, 2 tables. Zhiyuan Li and Zihan Li contributed equally |
| Subjects: | Distributed, Parallel, and Cluster Computing (cs.DC); Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.11201 [cs.DC] |
| (or arXiv:2610.11201v1 [cs.DC] for this version) | |
| https://doi.org/10.48550/arXiv.2610.11201 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Zihan Li [view email]
[v1]
Thu, 8 Oct 2026 04:04:52 UTC (204 KB)
来源：arXiv：cs.LG · arxiv.org