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 给长程 Agent 做两级 KV 共享:外层只桥接全注意力层、内层 token 级复用。80B-A3B 轻后训后 RULER-v2 均值比 HySparse 高 19.81 分,1M token 下 prefill 算力降 2.92 倍、KV 仅 2.69GB。
长程、多轮 Agent 的节奏很固定:动作或工具调用很短,工具返回、执行日志、文档却很长。上下文涨得快,prefill 算力和 KV cache 一起涨。HySparse 已经用全注意力层与稀疏层交错,并把 KV 和选块索引交给后续稀疏层复用,注意力成本和 cache 都压下去了。但它 prefill 仍要把每一层跑完,块级稀疏在多轮检索上也不够准。
Agent 推理越来越偏输入主导。工具返回往往比动作本身长得多,只稀疏注意力,解决不了「这些 token 还得一层层往前传」的成本。YOCO 把骨干拆成 self-decoder 和 cross-decoder,让 prefill 在前半段就把全局 KV 建完。HySparse2 把这套接到 HySparse 上,同时打三个目标:更短的 prefill 路径、更小的 KV、更准的长上下文检索。
骨干按 YOCO 拆成两半。self-decoder 混全注意力和滑动窗口注意力(SWA,每个 query 只看最近一段 token),负责本地建模;cross-decoder 混全注意力和稀疏注意力,负责全局检索。共享分两层。
外层叫 KV Bridging,只桥两边的全注意力层。cross-decoder 全注意力的 K/V,用对应 self-decoder 全注意力层的输入隐状态再投影一次;Q 仍用本层当前隐状态。一层 self-decoder 全注意力可以给多层 cross-decoder 全注意力供源,两边混合比例可以不同。每层仍有自己的 K/V 投影,共享的是源隐状态,不是同一份 cache。
内层沿用 HySparse 的 KV Reuse,改了两处。一是把块级 top-k 换成 token 级:全注意力层按 token 打分,后面稀疏层复用被选中的 KV。二是拆掉稀疏层里单独的 SWA 分支,强制把最近 128 个 token 塞进稀疏选择,再补 1024 个窗口外高分 token。本地和全局读同一份全注意力 KV。
拆独立 SWA 是为了让 prefill 能在 self-decoder 之后直接退出。独立 SWA 要从本层隐状态建一段后缀 KV,依赖会顺着深度往下传,cross-decoder 就跳不过去。49 层模型里 prefill 节点只部署前 25 层,内存几乎砍半;高 SWA 比例下 self-decoder 的 prefill 只需跑一层全注意力。HySparse2 用 MQA(每组 1 个 KV 头),两条基线用 GQA(4 个 KV 头)。SWA 层用 64 维 partial RoPE,全注意力和稀疏层用 NoPE,长上下文扩展不用再调 RoPE。全注意力层还当 oracle 选择器:用精确注意力分数给后续稀疏层挑 token,不必另训 indexer,也不必蒸馏。
三套 80B-A3B MoE,49 层、hidden 2048,数据和日程相同:约 500B token、32k 预训练,再约 100B token 轻后训并把上下文拉到 256k。HySparse 和 HySparse2 各留 5 层全注意力,Hybrid SWA(MiMo-V2 系列同款)留 9 层。
| 设置 | RULER | NoLiMa | KV @1M (FP8) |
| Hybrid SWA | 88.71 | 30.13 | 12.09 GB |
| HySparse | 84.89 | 40.27 | 6.72 GB |
| HySparse2 | 90.77 | 49.76 | 2.69 GB |
预训练通用能力大体打平。HySparse2 在 BBH、MMLU-Pro 上更强,DROP 58.99 对 HySparse 的 63.78,GSM8K 61.94 对 64.14,MATH 也略低。长上下文这边差距清楚:RULER 比 HySparse 高 5.88,NoLiMa 高 9.49。
轻后训后四项长上下文指标全程领先。相对 HySparse,MRCR-v2 均值 +11.30,RULER-v2 均值 +19.81;相对 Hybrid SWA 分别 +6.44 和 +18.65。256k 上 RULER-v2:58.45 / 32.61 / 35.74。AgentPPL、LongPPL 在评估长度内都更低。1M token 时 prefill FLOPs 相对 HySparse 降 2.92 倍、相对 Hybrid SWA 降 5.02 倍。
同样 1024 全局 token 预算,token 级相对块级(块长 64)把 RULER-v2 从 49.56 拉到 56.13,两针 MRCR-v2 从 12.94 拉到 21.08,GraphWalks 从 29.38 拉到 34.92。强制窗口相对独立门控 SWA:RULER 更好,GSM8K 低 5.08 分、MRCR-v2 低 4.99 分。290B-A8B、约 1.8T token 上,有无 KV Bridging 质量接近,DROP 从 71.37 掉到 68.17,LongPPL 从 3.6053 降到 3.4202。KV Bridging 的 RULER 收在 87.65,U 型 KV Mirror 是 81.28。
这是给长程 Agent 推理栈改注意力拓扑,不是又一个稀疏核。prefill–decode 分离时,prefill 节点只扛一半权重、一层全注意力,对 80B 级 MoE 的部署账是真的。token 级选择对准的是多轮轨迹里关键信息散落在角色分隔符和工具边界这种模式:块级要留一个边界 token,预算会砸在邻居上。
能用的前提是愿意从预训练改注意力,不是给现成模型打补丁。MQA 对 GQA、头维 256/256 对 192/128 也叠在对比里,KV 变小不全是两级共享贡献的。
通用能力没有全面赢。预训练 DROP、GSM8K、MATH 都有回退;Forced SWA 换来的 prefill 早退,是用部分推理分换的。AgentPPL、Repo Code PPL 是内部集,外部不好复现。后训只有约 100B token 的轻量,还不是完整对齐。投机解码用 self/cross 边界隐状态做 MTP 或 DFlash 只写了方向,没训更大 drafter。全注意力层仍当选择器留着,作者自己说后训里可以再换成轻量 indexer。基线对比并不「只换共享」。