Xiaomi HySparse2 cuts 1M-token prefill FLOPs 2.92× with two-level KV sharing

HySparse2: Hybrid Sparse Attention with Two-Level KV Sharing

Jianyu Wei, Yizhao Gao, Qihao Zhang, Shimao Chen, Zhengju Tang, Yu Cheng, Shengjie Zhou, Zihan Jiang, Yifan Song, Hailin Zhang, Liang Zhao, Bo Yang, Gang Wang, Shijie Cao, Fuli Luo

cs.CL

2026-09-22

HySparse2 adds YOCO-style KV Bridging and token-level reuse: +19.81 RULER-v2 vs HySparse after light post-training, 2.92× fewer prefill FLOPs and 2.69 GB KV at 1M.

What problem this solves

Long-horizon, multi-turn agents emit short actions and then swallow long tool returns, traces, and documents. Context grows fast, and so do prefill FLOPs and KV-cache size. HySparse already interleaves full-attention layers with sparse layers and reuses KV plus selection indices, which cuts attention cost and cache. Prefill still runs every layer, and block-level sparsity is too coarse for multi-turn retrieval.

Agent inference is increasingly input-dominated. A tool response can dwarf the action that produced it. Sparse attention does not remove the cost of pushing those tokens through the backbone. YOCO splits the stack into a self-decoder and a cross-decoder so prefill can finish global KV construction halfway. HySparse2 grafts that split onto HySparse, aiming at a shorter prefill path, a smaller KV, and sharper long-context retrieval at once.

Method

The backbone follows YOCO. The self-decoder mixes full attention with sliding-window attention (SWA: each query sees only a recent window) for local modeling. The cross-decoder mixes full attention with sparse attention for global retrieval. Sharing happens at two levels.

Outer level: KV Bridging, only between full-attention layers. Keys and values of a cross-decoder full-attention layer are projected from the input hidden states of a matching self-decoder full-attention layer. Queries still come from the cross-decoder's current hidden states. One self-decoder full-attention layer can feed several cross-decoder full-attention layers, so the two hybrids need not share the same ratio. Each layer keeps its own K/V projections; what is shared is the source hidden state, not one cache tensor.

Inner level: HySparse's KV Reuse, with two changes. Token-level top-k replaces block-level selection: the full-attention layer scores tokens, and later sparse layers reuse those KV entries. The separate SWA branch inside sparse layers is removed. The most recent 128 tokens are forced into the sparse set, then 1,024 high-scoring tokens outside the window. Local and global reads share the same full-attention KV.

Dropping the SWA branch is what makes a true early exit possible. A separate SWA needs a suffix KV from its own hidden states, and that dependency cascades with depth, so the cross-decoder cannot be skipped. On the 49-layer model, the prefill node hosts only the first 25 layers, cutting its weight memory almost in half. With a high SWA ratio, prefill runs full attention in a single layer. HySparse2 uses MQA (1 KV head); both baselines use GQA (4 KV heads). SWA layers use 64-dim partial RoPE; full and sparse layers use NoPE, so long-context extension needs no RoPE retune. Full-attention layers also act as an oracle indexer: exact scores pick tokens for the sparse layers, with no extra indexer and no distillation.

Results

Three 80B-A3B MoE models, 49 layers, hidden size 2,048, same data and schedule: about 500B tokens at 32k, then about 100B tokens of light post-training with context extended to 256k. HySparse and HySparse2 keep 5 full-attention layers; Hybrid SWA (the MiMo-V2 design) keeps 9.

SetupRULERNoLiMaKV @1M (FP8)
Hybrid SWA88.7130.1312.09 GB
HySparse84.8940.276.72 GB
HySparse290.7749.762.69 GB

General skills after pretraining are roughly even. HySparse2 is stronger on BBH and MMLU-Pro, weaker on DROP (58.99 vs 63.78), GSM8K (61.94 vs 64.14), and MATH. Long-context is the gap that matters: RULER +5.88 and NoLiMa +9.49 versus HySparse.

After light post-training, HySparse2 leads on all four long-context metrics at every evaluated length. Mean MRCR-v2 / RULER-v2 rise 11.30 / 19.81 points over HySparse, and 6.44 / 18.65 over Hybrid SWA. At 256k, RULER-v2 is 58.45 vs 32.61 vs 35.74. AgentPPL and LongPPL stay lower throughout. At 1M tokens, prefill FLOPs drop 2.92× vs HySparse and 5.02× vs Hybrid SWA.

Under the same 1,024-token global budget, token-level selection lifts RULER-v2 from 49.56 to 56.13, 2-needle MRCR-v2 from 12.94 to 21.08, and GraphWalks from 29.38 to 34.92 versus 64-token blocks. Forced windows beat gated SWA on RULER, but lose 5.08 GSM8K and 4.99 MRCR-v2. At 290B-A8B on about 1.8T tokens, KV Bridging stays close in quality: DROP falls 71.37 to 68.17, LongPPL 3.6053 to 3.4202. Bridging finishes RULER at 87.65 versus 81.28 for U-shaped KV Mirror.

Why it matters

This is an attention topology for agent serving, not another sparse kernel. Under prefill-decode disaggregation the prefill node holds half the weights and one full-attention layer, which is a real 80B-MoE deployment cut. Token-level selection matches multi-turn traces where the useful bits sit on role delimiters and tool boundaries: keeping a block to save one of those tokens wastes budget on neighbors.

The catch: this is a pretraining change, not a drop-in patch. The comparison also stacks MQA versus GQA and 256/256 versus 192/128 head dims, so the KV shrink is not purely from two-level sharing.

Limitations

General capability does not win across the board. DROP, GSM8K, and MATH all regress after pretraining; Forced SWA buys the early-exit prefill with some reasoning score. AgentPPL and Repo Code PPL are internal. Post-training is a light 100B-token stage, not a full alignment recipe. Speculative decoding from the self/cross boundary (MTP, later DFlash) is sketched, not trained. Full-attention layers still sit in the stack as indexers; the paper says they could be replaced later by a light indexer. The baselines are not "same backbone, sharing on/off."

Terms

Source

What people are saying

Related papers

All paper explainers