PageWeaver: KV-Guided Query Unions for Sparse Attention

The gist The authors 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.

How it works

PAGEWEAVER is an execution design that uses selected-page affinity to assemble query groups while preserving each query’s original support and complete output ownership Query unions share page loads and populate Tensor Core tiles; their cost depends on which queries are grouped together The design uses a bounded GPU search to produce query IDs, and an ID-aware two-CTA kernel consumes them without materializing reordered Q tensors or cross-page partial outputs The selector is unchanged; membership and causal masks follow original query IDs Within a bounded window, a GPU kernel greedily collects queries whose selected sets add few new pages The method rebuilds IDs from current Top-K sets on every call; it does not reuse page-pair decisions or previous requests’ routes This design trades an online grouping cost for a smaller union without introducing crosspage output reduction The execution size is fixed at g = 8, independently of the search window W Each CTA applies its own four-query membership and causal masks, and retains FP32 running softmax statistics through the complete traversal The H200 path accepts BF16 Q and FP8 E4M3 KV Each CTA finishes its queries locally and stores directly to their original rows, eliminating a separate reorder-output pass and page-partial combine The design is useful when Snet > 0

Key Findings

The H200 Union8 implementation achieves a 1.70× 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 These results separate execution-group reuse from the complete cost of exploiting it online

Design Alternatives and Comparisons

The paper compares PAGEWEAVER against a direct KV-page union implementation, which provides a complementary design study of nonlocal reuse and reduction cost The direct KV-page implementation inverts incidence, proposes disjoint page pairs, and forms common/exclusive query lists using full membership This organization partitions a query’s output across page groups The B300 port maps the same IDs into a CuTe Union8 path with direct FP8 Q and physical page size 64 The auxiliary B300 experiment shows that neither online window beats the faster existing backend on any of the eight tested inputs

Limitations and Scope

The main evidence is an H200 MiniMax-M3 case study; selected operator captures and limited service repetitions constrain the generality of its conclusions The fixed 64K bitmap limits active regrouping on longer contexts The present greedy objective minimizes the eight-query union, not the subgroup compute proxy in Equation 4 A future selector could consider both, but no such cost-aware layer policy is evaluated here The operator evidence emphasizes selected inputs; independent natural-text validation should precede a general scheduling rule The paired task-quality evaluation of the final policy remains necessary before a quality-qualified deployment claim

Conclusion

PAGEWEAVER uses the relation between selected queries and KV pages to improve execution groups while retaining each query’s full attention output The central design principle is to optimize the cost of executing a fixed sparse relation, including the cost of constructing its execution groups PAGEWEAVER uses the relation between selected queries and KV pages to improve execution groups while retaining each query’s full attention output The direct KV-page execution and the B300 comparison identify why reuse must be evaluated together with preparation and output costs PAGEWEAVER uses the relation between selected queries and KV pages to improve execution groups while retaining each query’s full attention output The central design principle is to optimize the cost of executing a fixed sparse relation, including the cost of constructing its execution groups The central design principle is to optimize the cost of executing a fixed sparse relation, including the cost of constructing its execution groups

How it works

The authors connect query-group size, page-set affinity, and output ownership in a common execution model, with direct KV-page union as a measured design alternative The authors implement bounded GPU regrouping and an ID-aware FP8 sparse-attention path on Hopper, retaining a complete query output within one CTA The execution slot r reads qπ(r), uses membership Mπ(r,p) and causal position cπ(r), and writes oπ(r) This is a supportpreservation argument in exact arithmetic

Key Findings

The authors show that controlled interleaved supports make the page-side opportunity explicit: adjacent queries select different page communities, while distant queries share them Direct KV-page union lowers complete-call time relative to FlashInfer by 15.6% at 32K context (1.340 versus 1.588 ms) and 10.1% at 64K (2.852 versus 3.172 ms) Online W128 lowers complete-call latency by 3.26–7.66% relative to original Union8 The net saving is 58.41 µs

Design Alternatives and Comparisons

The authors compare the base backends against the same Triton reference, and Union8’s differences are close to FlashInfer’s on this set, while native differs more The B300 port preserves support and has at most 1.55 × 10−5 full-output NRMS relative to its parent Native algorithms, input sets, Q entry formats, page sizes, and descriptor costs differ between the two platforms

Conclusion

The operator evidence emphasizes selected inputs; independent natural-text validation should precede a general scheduling rule The paired task-quality evaluation of the final policy remains necessary before a quality-qualified deployment claim The central design principle is to optimize the cost of executing a fixed sparse relation, including the cost of constructing its execution groups The central design principle is to optimize the cost of executing a fixed sparse relation, including the cost of constructing its execution groups The central design principle is to optimize the cost of executing a fixed sparse relation, including the cost of constructing its execution groups<ref:2610.

1. Bold query grouping to share page loads and populate Tensor Core tiles, as Query unions share page loads and populate Tensor Core tiles; their cost depends on which queries are grouped together. This allows for efficient execution of sparse attention by grouping queries whose selected pages overlap, amortizing KV loads.

2. Implement an ID-aware two-CTA kernel that consumes IDs without materializing reordered Q tensors or cross-page partial outputs, as this design trades an online grouping cost for a smaller union without introducing crosspage output reduction. This preserves the original query ownership while exploiting page reuse.

3. Utilize bounded online regrouping to lower latency by 3.26–7.66% on five selected 64K-context captures, as Online regrouping further lowers latency by 3.26–7.66% on five selected 64K-context captures. This provides a practical speedup for dynamic workloads where query sets can be grouped based on page affinity.

4. Develop a direct KV-page union implementation that retains each query’s original support and complete output ownership, as this design trades an online grouping cost for a smaller union without introducing crosspage output reduction. This enables the system to exploit nonlocal reuse while ensuring every query receives its full, correct output.