2.1 · 粗粒度选择: 从内容分组到动态稀疏块

动态 Sparse Attention 必须在完整 QK 计算以前决定读取哪些 K/V. 粗粒度选择把 token 先归入桶、簇或二维块, 再以组为单位保留候选. 选择结果未必精确到最终高权重 token, 却能把不规则关系整理成 GPU 可以执行的连续 tile. 这条路线关心的核心量是候选块召回、块内多算、建索引成本与负载分布. 只比较逻辑非零比例, 会同时漏掉索引器和物理扩张.

设输入为 $Q,K,V\in\mathbb{R}^{B\times H\times n\times d_h}$ , 每个物理块包含 $b$ 个 token. 路由器不直接输出 $n\times n$ mask, 而是生成块索引 $J\in\mathbb{N}^{B\times H\times N_q\times k_b}$ , 其中 $N_q=\lceil n/b\rceil$ , $k_b$ 是每个 query 块选择的 key 块上限. kernel 根据 $J$ 读取原始 K/V, 在选中 tile 内计算精确 QK、causal mask、在线 softmax 与 AV. 输出 shape 仍为 $B\times H\times n\times d_h$ .

粗粒度选择保留的是计算区域, 不是提前决定 attention 权重. 块内每条边仍由原 QK 分数竞争. 代价是一个关键 token 会带入同块的其他 token; 好处是 K/V 连续、索引短、矩阵核利用率高. 块尺寸 $b$ 越大, 索引越便宜而多算越多; $b$ 越小, 召回更精细而 gather、top-k 和 kernel 调度更重.

1. 两条进入路线

1.1. 从表示空间形成桶

哈希与聚类从向量本身构造邻域. Reformer 用随机旋转产生 LSH 桶, 相似向量以较高概率碰撞; Routing Transformer 用在线 k-means 中心把 query 与 key 分入内容簇. 两种方法通常在训练阶段已经改变 attention 拓扑. 模型会在局部 head、内容桶和容量约束下学习表示, 推理时沿用同一分桶规则.

哈希的状态是随机投影与多轮桶号, 聚类的状态是持续更新的中心. 前者以概率召回换取无需学习中心, 后者让分区贴合数据, 同时承担中心漂移、空簇和热点簇. 二者都要排序、计算 offset、重排 Q/K/V、执行桶内 attention, 再逆置换输出. 算法里的同桶关系只有经过这些步骤才成为实际加速.

这条路线适合从训练开始共同设计. 若直接替换已训练稠密模型, 共享 QK、中心分布和桶内 softmax 都可能改变原模型行为. 评测需要把路由 recall、桶容量溢出、物理 padding 和任务质量放在一起, 不能用语言模型平均 loss 掩盖唯一证据漏选.

从已有 attention 归纳空间模式。 MInference面向已经训练完成的长上下文模型. 它先离线观察每层每个 head, 在 A-shape、Vertical-Slash 与 Block-Sparse 中选择模式和预算; 在线再根据当前输入生成具体列、斜线或块坐标. 模式类型固定, 索引位置动态. 这种分层把昂贵的模式比较移出请求路径, 在线只解决受限几何中的候选选择.

MInference 的目标是 Prefill. 整段 prompt 同时产生大量 query, 稠密关系为平方级, 专用稀疏 kernel 能获得足够工作量. Prefill 结束后仍保存完整 KV Cache, 原方法并未把未访问 token 删除. Decode 每步只有一个新 query, 索引摊销、cache 布局与访存瓶颈不同, 不能从 Prefill 加速直接推出 Decode 收益.

两条路线的区别不在有没有动态性. 哈希与聚类让动态分组成为模型结构, MInference 从稠密模型已有的 head 形状中寻找可执行近似. 前者可在训练中适应漏边, 后者以免训练部署为主要价值. 比较时必须注明模型是否重训、稀疏发生在哪个阶段以及稠密教师是否只在离线使用.

统一的成本与质量口径。索引器先花了多少。设最终物理 attention 关系数为 $\tilde E$ , 主计算约为 $O(\tilde E d_h)$ . 总时间还包含路由投影、摘要、排序或 top-k、整数索引写入、Q/K/V 重排以及输出 scatter. 若索引器计算完整 $QK^T$ 后才选块, 主平方成本已经发生. 有效索引器必须读取低维摘要、少量采样 query、有限中心或规则几何.

哈希常付出 $O(n\log n)$ 排序, k-means 付出 $O(nCd_h)$ 中心打分, MInference 的在线索引付出采样 QK 与 reduce/top-k. 渐近式仍不足以预测延迟. GPU 上需要记录 HBM 字节、kernel 启动、同步和排序轮数. 短序列时稠密 FlashAttention 往往占优, 长度跨过阈值后稀疏主计算的减少才覆盖索引固定成本.

recall@k 不能只按 token 算。以稠密 attention 的 top- $t$ 集合 $G_i$ 为教师, 候选集合为 $\mathcal N(i)$ , token recall 为 $|G_i\cap\mathcal N(i)|/|G_i|$ . 对块方法还应报告物理预算: 一个命中 token 可能要求读取整个块. 两个方法 token recall 相同, 物理 tile 数可能相差很大. 稠密概率质量覆盖能区分教师第一名与边缘 top- $t$ , 仍需任务质量确认误差传播.

平均 recall 也会隐藏少数致命 query. 应按层、head、query 区间和任务类型观察尾部, 尤其是唯一远程证据、多跳汇总、相似干扰与代码引用. 对训练型路由还要分训练长度内和长度外; 对离线配置要分校准域和部署域. 候选召回回答算法近似, kernel 对同一候选的误差回答实现正确性, 两者要分开测.

块大小连接算法与硬件。逻辑索引可以选择单列或单 token, kernel 往往以 $64\times64$ 等 tile 执行. 物理扩张增加边数, 也换来连续加载和矩阵核利用率. 因此预算搜索最好直接使用物理 tile 成本, 而非只最小化逻辑 mask. 多个逻辑候选落进同一 tile 时应合并, 不同分支重叠时必须去重, 否则 softmax 分母会重复计数.

负载均衡同样受块影响. 热点簇或热门 key 块会拖慢一个执行 wave, 在序列并行中还会形成通信热点. 固定 top-k 给出工作上限, 不能保证实际设备均衡. 应报告每行 tile 数、最大/平均比、padding、溢出和跨设备字节. 粗粒度方法的工程成败经常取决于这些分布量, 而非论文图上的平均稀疏率.

1.2. 从一个 query 块走完数据路径

粗粒度选择最容易被热图迷惑. 热图里一块区域看起来很稀疏, 并不说明执行时真的只读了这些位置. 要判断一个方法有没有节省, 可以固定一个 query 块, 沿着路由、索引、取数、计算和归一化走一遍. 设 query 块大小为 $b_q$ , key 块大小为 $b_k$ , 当前 query 块记作 $Q_r\in\mathbb{R}^{b_q\times d_h}$ . 路由器为它选出 $k_b$ 个 key 块:

J_r=(j_1,j_2,\ldots,j_{k_b}),\qquad K_{J_r},V_{J_r}\in\mathbb{R}^{k_b\times b_k\times d_h}. \tag{1}

式 (1) 只是地址列表. kernel 还要将这些地址变成实际加载. 如果 K/V 按 token 连续存放, 单个 key 块内部可以合并读取; 不同块之间仍是 gather. 若 $J_r$ 已按地址排序, 相邻块有机会合并, L2 命中也更稳定. 若它只按路由分数降序排列, 数学结果不受影响, 访存顺序却可能来回跳. 因此索引既有「选中了谁」这一语义, 也有「以什么顺序访问」这一物理属性.

候选进入 kernel 后, 每个 tile 仍计算原始点积:

S_{r,t}=\frac{Q_rK_{j_t}^{\mathsf T}}{\sqrt{d_h}} \in\mathbb{R}^{b_q\times b_k}. \tag{2}

自回归 mask、末块 padding 和重复候选必须在 softmax 前处理. 多个 key 块不能各自做一次 softmax 后再平均, 因为它们共享一行分母. FlashAttention 的在线归一化可以直接延伸到稀疏块. 处理第 $t$ 个块时维护行最大值 $m^{(t)}$ 、指数和 $\ell^{(t)}$ 与未归一化输出 $u^{(t)}$ . 新块的行最大值为 $\hat m$ , 指数和为 $\hat\ell$ , 加权 Value 为 $\hat u$ , 则:

m^{(t)}=\max(m^{(t-1)},\hat m), \tag{3}

\ell^{(t)}= e^{m^{(t-1)}-m^{(t)}}\ell^{(t-1)}+ e^{\hat m-m^{(t)}}\hat\ell, \tag{4}

u^{(t)}= e^{m^{(t-1)}-m^{(t)}}u^{(t-1)}+ e^{\hat m-m^{(t)}}\hat u. \tag{5}

随后输出 $u^{(k_b)}/\ell^{(k_b)}$ . 式 (3) 到式 (5) 说明稀疏块顺序可以变化, 但一个候选不能被无意计算两次. 局部窗口、全局块和动态 top-k 常会重叠. 若同一 key 块从两个分支进入, 直接拼接会把它的指数项加入两遍, 这已经改变数学定义. 去重可以放在索引生成阶段, 也可以给分支规定互斥范围; 在 kernel 内靠逐项比较去重通常太贵.

粗粒度选择的完整数据路径由此分成五个可测阶段: 生成 query 摘要, 对候选块打分, 选择并整理块号, 加载 K/V tile, 合并稀疏 softmax. 论文若只报告末项 kernel, 还不足以说明端到端收益. 部署时至少要分别计时路由和 attention, 并确认计时包含索引张量从生成位置到消费位置的同步.

2. 路由分数怎样代表一个块

2.1. 中心、池化和最大值回答的问题不同

一个块含 $b_k$ 个 key, 路由器不可能先让每个 query 与块内所有 key 点积, 否则候选阶段已经接近原计算. 常见做法是把一个块压成摘要 $c_j$ . 最简单的是均值:

c_j=\frac{1}{b_k}\sum_{s=1}^{b_k}k_{j,s},\qquad g_{rj}=\operatorname{pool}_{q\in Q_r} q^{\mathsf T}c_j. \tag{6}

均值中心适合块内方向相近的情况. 若一个块包含一个强相关 key 和许多无关 key, 这个信号会被稀释. 最大池化可保留尖峰, 但要么读取更多细粒度特征, 要么依赖额外摘要. 学习得到的 block representative 可以针对召回优化, 代价是训练状态、额外参数以及部署时的一致性要求.

query 侧也要聚合. 用末项 query 代表整个块, 偏向块尾需求; 对所有 query 得分取最大值, 能照顾任意一个 query 的强匹配, 同时容易被异常高分支配; 取均值更稳定, 可能漏掉少数查询的唯一证据. 所以「块分数」没有脱离聚合规则的统一含义. 评测路由器时, 应把块内每个 query 的召回分别保留, 再汇总均值和低分位, 不能只用聚合分数自证准确.

MoBA 采用块级门控, query 对历史 KV 块的代表进行打分并选择少数块, 当前块保留因果范围内的细粒度注意力. 它把 MoE 的 top-k 路由思路移到注意力块, 但 expert 与 KV 块有一处根本差别: expert 是可学习函数, KV 块是当前序列产生的状态. 热门 KV 块不会带来 expert 参数争抢, 却会造成大量 query 重复读取同一段 K/V. 这可能提高缓存复用, 也可能形成某些工作项远重于其他工作项. 判断方向取决于调度和缓存层级, 不能只从 top-k 数字推出.

NSA 的压缩分支也生成粗粒度表示, 但它不是单纯用摘要决定随后读取哪些原始 token. 压缩分支本身参与输出, 与选择分支、局部分支共同组成 attention. 因此「摘要只用于路由」和「摘要本身也是 Value 信息通道」要分开. 前者的质量依赖原始块是否被召回, 后者即使没有回读原块也能传递压缩信息, 同时会引入压缩误差.

top-k、阈值与容量约束。固定 top- $k_b$ 让每个 query 块拥有相同的工作上限, kernel 调度最直接. 它也把内容难度压成同一预算: 简单局部段落可能多读, 跨文档、多跳或代码依赖可能少读. 阈值选择允许候选数随分数分布变化, 质量更贴合样本, 但产生 ragged 行. 工程上常将阈值结果再裁到 $k_{\max}$ , 同时保留 $k_{\min}$ , 于是规则写成:

k_r=\min\left(k_{\max}, \max\left(k_{\min},\left|\{j:g_{rj}\ge\tau_r\}\right|\right)\right). \tag{7}

$\tau_r$ 可以是绝对阈值、相对最大值的差, 也可以由累计路由概率决定. 绝对阈值对层间尺度敏感; 相对阈值在所有分数都很差时仍会选出一些块; 累计概率依赖路由分数的校准. 无论采用哪种, 物理执行通常仍按若干容量桶分组, 例如把 $k_r$ 向上归入 4、8、16. 逻辑候选变长并不意味着 kernel 真正按每行精确长度运行.

训练时的容量约束还会影响梯度. top-k 之外的块没有主 attention 梯度, 路由器若只从最终语言模型损失学习, 早期错选可能长期收不到纠正信号. 解决办法包括稠密教师产生的路由监督、可微近似、辅助负载损失或由其他分支维持信息通路. 每一种都改变训练成本. 「推理只算稀疏块」与「训练从头到尾都省去稠密 QK」也需要分别核实.

因果边界不能交给块号猜。对第 $r$ 个 query 块, 历史完整 key 块编号满足 $j<r$ ; 当前块 $j=r$ 只能看块内下三角. 若块大小为 64, query 位置 130 位于第 2 块, 它能看到第 0、1 块全部位置, 以及第 2 块中的 128 到 130. 路由器可以只为完整历史块打分, 当前块作为固定局部分支加入. 这样 top-k 不会浪费在必选当前块上, 索引协议也更简单.

变长 batch 会增加一层边界. 同一个物理 batch 中不同样本的有效块数不同, 块号必须相对样本起点解释. packed sequence 还可能把多个文档放进连续 token 区间, 单纯的 $j\le r$ 无法阻止跨文档访问. kernel 需要 sequence ID 或累积长度边界. 稀疏索引正确而文档边界错误, 输出仍会发生信息泄漏.

三类粗粒度路线的真实分叉。内容分组改变模型看到的图。 LSH 与在线聚类先把 token 放进内容邻域, 注意力图由桶成员关系决定. 它们适合在训练期使用, 因为表示可以逐渐适应分组规则. Reformer 的共享 QK、随机旋转和多轮哈希是一整套约束, Routing Transformer 的中心更新与局部 head 也是模型设计的一部分. 把其中一个索引器直接塞进已有模型, 不能假定原 QK 排序仍然成立.

这类方法的好处是分组不限于位置模板. 两个相距很远但表示相似的 token 可以进入同桶. 问题也来自同一点: 桶成员数量不可预知, 相似性空间会随层与训练阶段变化. 排序、容量裁切和跨设备重排不是附属实现, 它们决定算法能否保持预期候选.

空间模式近似已有稠密模型。 MInference 观察已训练模型的 attention 图, 用 A-shape、竖线-斜线和块稀疏三类模板覆盖不同 head. 它没有让任意 token 在表示空间里自由配对, 而是先把候选限制在容易执行的几何族内. 离线阶段决定每个 layer-head 使用哪一类模式及预算,在线阶段才根据输入确定具体列、斜线或块.

这种路线保留原模型权重, 适合超长 Prefill. 它的风险是配置域迁移: 校准数据里的 head 形状不一定覆盖另一语言、代码、表格或多轮上下文. 模式名称相同也不表示索引一致. VS 分支怎样采样 query、怎样聚合列分数、斜线位移如何选择, 都会影响召回. 因此配置文件要与模型权重、长上下文训练版本和 head 编号绑定.

原生块路由让模型适应选择。 MoBA 直接让 query 路由到 KV blocks, NSA 则把压缩、选择和局部通道联合训练. 它们与免训练近似相比, 允许模型围绕稀疏图重新组织表示. 训练适应可以弥补部分漏边, 也意味着结果不能脱离训练配方比较. 一个从头使用稀疏 attention 的模型与一个把稠密模型后装稀疏 kernel 的模型, 即便物理边数相近, 误差来源也不同.

原生训练还把反向传播带入系统设计. forward 选了哪些块, backward 至少要重算或保存相同索引; 路由器本身需要梯度或辅助监督; 分布式训练要处理索引对应的 K/V 所在设备. NSA 明确把训练、前向和解码一并纳入硬件对齐目标, 这比只优化 Prefill 推理覆盖更广. MoBA 允许从 full attention 平滑过渡到 sparse attention, 其价值包含训练过程中的兼容性, 不能缩成一句「块 top-k」.

这三类路线可以放进同一张坐标表:

路线	候选依据	是否依赖重训	主要执行单位	主要风险
LSH/聚类	表示空间的桶或簇	是	排序后的桶与 chunk	碰撞漏选、热点桶、重排成本
MInference	稠密 head 的空间模式	否	列、斜线、二维 tile	校准域迁移、代理索引误差
MoBA	学习到的块门控	是	KV block	块摘要失真、top-k 训练信号
NSA	压缩、选择与局部多分支	是	压缩块、选择块、局部窗口	分支融合、训练与内核复杂度

表里的「是否依赖重训」描述原方法的使用方式, 不表示任何一类绝对优于另一类. 已有权重无法重训时, 免训练近似更现实; 能控制预训练时, 原生稀疏可以把训练成本也纳入优化. 路由自由度越高, 通常越难获得规则内存访问; 几何约束越强, 又越容易遇到分布外模式.

2.2. Prefill、Decode 与训练不是同一道题

Prefill 有足够多的 query 摊薄索引。 Prefill 同时处理整段输入. 一个 query block 的路由决定可以被 $b_q$ 行共享, 块摘要也能被许多 query 复用. 当序列很长时, 稠密下三角 QK/AV 占比迅速上升, 稀疏 tile 的减少容易覆盖索引固定成本. 这正是 MInference 主要瞄准的阶段.

但 chunked prefill 会改变条件. 服务系统可能把长 prompt 切成多个 chunk, 每次只处理新增 query, 历史 K/V 已写入 cache. 路由器要同时面对当前 chunk 和全部历史块; 如果每个 chunk 重建所有历史摘要, 复杂度会悄悄增长. 更合理的做法是摘要随 KV Cache 一起追加, 当前请求只计算新块摘要. 配置与缓存状态必须在抢占、迁移和 prefix cache 复用后保持一致.

Decode 的主成本是读取历史 K/V。 Decode 每步常只有一个 query. 块路由可以减少 K/V 读取, 对内存带宽瓶颈很有吸引力, 但此时没有 $b_q$ 个 query 共同摊销路由. 若路由器为每步扫描全部 $N_k$ 个块摘要, 成本为 $O(N_kd_r)$ ; $d_r$ 是摘要或低维索引维度. 它仍小于读取全部 $N_kb_kd_h$ 个 K/V, 前提是摘要足够小、连续且常驻较快存储.

Decode 还要处理索引随新 token 变化. 固定使用 Prefill 末行选出的块, 会错过后续 query 的新需求; 每步重新 top-k 又带来小算子和同步. 跨层共享、跨 token 复用以及候选缓存因此变得重要. 这些优化要测命中率: 上一步选中的块在下一步是否仍覆盖高概率质量, 浅层索引能否服务深层. 复用不是免费结论, 它用陈旧候选换取较低延迟.

训练还包含反向与激活生命周期。训练使用整段 query, 与 Prefill 在前向形状上相近, 但需要反向传播. 稀疏 attention 若保存索引和 softmax 统计, 会占用激活显存; 若重算路由, 随机哈希、top-k 平局和中心状态必须可复现. 分布式序列并行时, 某个 query 选择的 KV 块可能位于另一设备, 路由结果会转化为不规则通信.

原生稀疏训练的收益应报告完整 step time, 包含 forward、backward、优化器之外的路由和通信. 只测 forward kernel 会高估训练收益. 质量方面也不能只在相同 step 数比较: 若稀疏训练每步处理更多 token, 应同时报告总训练 token、计算量和收敛曲线. NSA 将算术强度与硬件执行共同考虑, 说明训练型稀疏已经不能只用理论边数描述.

一次可复算的块预算。设序列长度 $n=8192$ , $b_q=b_k=64$ , 因而共有 $N=128$ 个块. 单个 head 的稠密因果 attention 约处理:

E_{dense}\approx\frac{n(n+1)}{2}\approx 3.36\times10^7 \tag{8}

个 token 对. 现在让每个 query 块固定保留最近 2 个块、1 个全局起始块, 再动态选择 5 个历史块. 忽略序列开头不足预算以及分支重叠, 每行最多 8 个 key 块, 物理关系数约为:

\tilde E=128\times8\times64\times64=4{,}194{,}304. \tag{9}

相对稠密下三角约为八分之一. 这个数字还不能写成 8 倍加速. 假设索引器先让每个 query 块摘要与全部 128 个 key 块摘要点积, 摘要维度 $d_r=32$ , 路由乘加约为 $128^2\times32$ . top-k 还包含选择与索引写入. 主 attention 的每条边则做 $d_h$ 维 QK 和 AV, 若 $d_h=128$ , 工作量远大于摘要打分, 路由在长序列上可能划算.

再看块内浪费. 若动态分支真正需要的每个块只有 8 个 token, 5 个动态块的逻辑关系为 $128\times64\times5\times8=327{,}680$ ; 物理执行却是 $128\times64\times5\times64=2{,}621{,}440$ , 放大 8 倍. 把块改成 16 能减少多算, tile 数和索引长度会增加 4 倍, 小 tile 的矩阵核效率也可能下降. 块大小不是算法论文末尾的实现参数, 它直接决定召回和吞吐之间的工作点.

如果局部分支与动态分支有两个块重叠, 去重后的总块数是 6 而非 8, 式 (9) 会降为 $3{,}145{,}728$ . 如果没有去重, 重叠块还会污染 softmax. 因此性能计数和正确性检查都会用到同一个指标: 每行唯一物理 tile 数.

这个例子还省略了 causal 三角. 当前块只有一半左右的有效关系, 历史块是完整矩形. 高性能 kernel 可能仍按整 tile 做矩阵操作, 用 mask 去掉上三角, 因而 FLOPs 统计和有效边数不同. 报告结果时应说明采用有效 token 对、物理 tile FLOPs 还是 profiler 实测指令量.

3. 质量评测要追踪被删掉的概率质量

3.1. 三层验证各自回答一件事

第一层固定候选索引, 比较稀疏 kernel 与朴素 masked attention. 这一步检查 causal mask、重复 tile、在线 softmax和 GQA head 映射, 不评价路由器. 给定相同 $J$ , 两者应在 dtype 容差内一致.

第二层将路由候选与稠密 attention 比较. top- $t$ recall 适合观察高分边, 概率质量覆盖更直接:

C_i=\sum_{j\in\mathcal N(i)}\alpha_{ij}^{dense}. \tag{10}

$C_i$ 接近 1 表示候选覆盖了稠密 softmax 大部分质量. 它仍不是输出误差的充分条件: 被漏掉的 Value 方向可能互相抵消, 也可能恰好承载关键内容. 可以再计算稠密输出与稀疏输出的相对误差、余弦相似度, 并按 layer-head-query 保存分布.

第三层是语言模型与下游任务. 困惑度观察平均预测, Needle 类任务检查远程定位, KV retrieval 对相似干扰更敏感, 多跳和代码任务检查多次信息搬运. 单一 passkey 很容易被全局列或局部规律解决, 无法证明动态路由已经可靠. MInference 官方资料分别使用 RULER、InfiniteBench、Needle 与 PG-19, 对应的长上下文能力也应分项报告.

平均值会藏住最危险的 query。若 99% query 的覆盖率为 0.99, 1% query 为 0, 平均覆盖仍接近 0.98. 那 1% 可能正是答案位置、函数定义或否定条件. 应至少报告低分位、最差 layer-head、不同位置区间以及不同输入类型. 对粗粒度块还要区分「路由没选中目标块」与「目标 token 在选中块内但 softmax 权重变化」.

长度外推需要单列. 训练或校准在 32K, 部署到 128K 时, 块数量增加, 固定 top-k 对历史覆盖比例下降. 路由分数的极值分布也随候选数变化, 固定阈值可能选出更多或更少块. 只在训练长度内得到的 recall 不能推出更长上下文稳定.

稠密回退怎样使用。回退路径适合三类情况: 序列短到稀疏没有收益, 模型或 shape 不在已验证配置中, 线上监控发现质量或数值异常. 回退应在请求开始前由明确条件选择, 避免某一层候选为空时静默切到稠密, 使延迟分布难以解释. 影子评测可以抽样运行稠密对照, 但主流量双算会抵消节省.

回退也不能掩盖实现错误. 固定索引的单元测试若不过, 应修正 kernel; 路由覆盖不足, 才考虑增加预算、重新训练或重新校准. 把所有异常都切回稠密, 最终只能得到一个偶尔加速且不可预测的系统.

从索引到 kernel 的验收表。索引表示会反过来限制算法。候选集合可以用四种常见形式传给 kernel. 第一种是稠密布尔 mask, shape 为 $N_q\times N_k$ . 它最容易验证, 却仍占平方空间, 极长序列下仅生成 mask 就可能成为瓶颈. 第二种是定长块号表 $J\in\mathbb{N}^{N_q\times k_b}$ , 配合固定预算, 地址计算直接, 也是块稀疏实现最常见的接口. 第三种是 CSR 一类压缩行格式, 用 row_ptr 表示每行起止, 用 col_idx 保存所有候选. 它支持变长预算, kernel 的循环长度与负载更不规则. 第四种是隐式规则, 例如局部窗口、固定 sink 或斜线位移; kernel 根据 query 块号现场算出 key 坐标, 不需要写出完整索引.

混合路线通常同时使用多种表示. 局部窗口采用隐式规则, 动态块放在定长或 CSR 表里, 全局块只保存一份短列表. 这比把所有分支展开成统一 COO 坐标更省控制数据. 代价是 kernel 必须合并多个来源, 并处理重叠. 若先在主机端把分支展开、排序和去重, GPU attention 会等待 CPU; 若在 GPU 上做通用集合运算, 小整数算子又可能占据显著时间. 更常见的选择是从规则上减少重叠, 例如动态 top-k 明确排除局部窗口覆盖的块.

索引 dtype 也有边界. 16 位整数能表示的块数有限, 但在百万 token、64 token 一块时仍只有 15625 个块, 单条序列可容纳; batch 偏移、物理页号或全局地址可能超过范围. 32 位索引更稳妥, 控制数据与带宽翻倍. 采用相对位移可以压缩局部候选, 任意远程块仍要绝对编号. 编码之前应先明确编号对象是序列内逻辑块、batch 展平后的块, 还是 KV 分页器里的物理页.

排序索引不会改变 softmax 结果, 却可能显著改变缓存行为. 同一个 query 行的 key blocks 若按物理地址升序访问, K/V 加载更连续; 按路由分数排序没有数学优势, 因为在线 softmax会正确合并任意顺序. 多个 query 行之间还可以按候选签名分组, 让相似工作项连续执行. 分组本身需要排序或哈希, 只有复用收益大于整理开销时才值得.

块摘要也属于缓存状态。训练时可以从当前 K 临时计算摘要, 推理服务更适合在写 KV Cache 时同步写入. 若每产生一个完整块就计算 $c_j$ , 摘要数量随序列线性增长. 当前未满块需要增量状态. 对均值摘要, 保存向量和与计数即可:

s_j^{(t)}=s_j^{(t-1)}+k_t,\qquad c_j^{(t)}=\frac{s_j^{(t)}}{r_j}, \tag{11}

其中 $r_j$ 是当前块已有 token 数. 学习型压缩器可能需要保留一小段 token 到块完成后统一前向, 或支持流式更新. 摘要更新发生在 KV 量化之前还是之后也会影响一致性. 若路由器在 FP16 key 上训练, 部署却从 INT4 KV 解码近似摘要, 路由误差与主 attention 量化误差会叠加.

prefix cache 复用时, 只复用原始 KV 而重新扫描前缀生成摘要, 会让长前缀命中的延迟收益缩水. 摘要应与 prefix block 使用相同的内容哈希和引用计数. 当多个请求共享物理 KV 页时, 摘要也可以共享, 因为它只由 key 内容决定; query 相关的 top-k 索引不能共享, 除非 query 状态也相同.

滑动窗口删除旧页后, 逻辑位置可能继续增长而物理页循环复用. 路由索引不能把物理槽号当成永久块 ID. 一种稳妥的协议是让路由器输出请求内单调逻辑块号, KV 管理器在 attention 启动前读取当前 block table 完成映射. 这样分页器可以搬迁页面, 路由语义保持不变. 映射表若在 CPU, 每步传输会成为新开销; 生产实现通常要求它常驻设备.

正确性。先构造小尺寸张量, 手写每个 query 块的索引. 覆盖第一个块、末项不足整块的尾部、重复索引、空动态分支、当前块因果三角以及 packed sequence 边界. 对同一索引运行朴素实现与目标 kernel, 逐层比较 logits 行最大值、指数和和最终输出. 只比较最终生成 token, 会让误差经过几十层后才暴露.

GQA 要验证 $h_q\mapsto h_{kv}$ 映射. tensor parallel 要验证本地 query head 使用正确的全局路由配置. 如果索引按全局块号生成而 K/V 在各 rank 分片, 还要明确块号到本地地址的转换. 一个错误配置通常 shape 完全合法, 结果也未必出现 NaN, 因而需要特定断言.

性能。至少记录 $T_{route}$ 、 $T_{select}$ 、 $T_{attn}$ 和整层时间. CUDA 异步执行要求在计时边界同步或使用事件; 首次编译与稳态分开. 同时记录序列长度、batch、query/KV head、 $d_h$ 、dtype、块大小、每行唯一 tile 数、候选长度分布和硬件版本. 缺少这些条件, 倍数无法复现.

与稠密 FlashAttention 比较时使用同一 QKV 布局、同一 dtype 和同一 causal 语义. 稀疏实现若省略 padding 而稠密基线处理大量 padding, 比较会把变长打包收益算进注意力方法. 反过来, 稀疏实现若额外转置和复制, 也要计入端到端路径, 不能只截取最内层 kernel.

服务状态。 KV block 摘要、索引缓存和原始 KV 的生命周期要一致. prefix cache 命中时, 摘要也应复用或可由缓存重建; 请求被抢占并恢复时, block 数和位置偏移不能变化; sliding window 丢弃旧 KV 时, 相应摘要与块号也要更新. 索引状态若仍指向已释放页面, 错误会表现为静默读取另一请求的数据, 远比普通精度下降严重.

Paged KV Cache 下, 逻辑块与物理页未必相同. 路由器输出逻辑 token block, block table 再映射到物理页. 若注意力块跨两个页面, kernel 需要两次加载或调整块对齐. 因此算法块大小最好与分页粒度协同, 但两者服务的目标不同: 算法块控制召回, 物理页控制内存分配. 强行设成同一个值可能牺牲其中一侧.

阅读顺序与路线导航。哈希与聚类从向量生成内容桶, $O(n^{3/2})$ 来自中心打分与桶内配对的平衡, 训练、Prefill、Decode 则各自维护索引. LSH 碰撞与在线中心更新的手算还能直接说明一件事: 索引本身也是算子, 它同样消耗计算、状态和同步.

再读MInference, 可以把同一套判断用于免训练 Prefill. 重点是离线按 head 选模式与在线索引的分工, 三类形状如何变成不同 tile, 以及在线 softmax怎样跨稀疏 tile 合并. 如果只关心已有模型的百万 token prompt 加速, 可以直接从这一篇进入, 但仍应回看前篇的 recall、重复候选与索引成本口径.

MoBA 与 NSA 也属于粗粒度动态选择的重要节点. 前者强调块级门控, 后者把压缩、选择与局部分支共同纳入训练. NSA已有独立解析; MoBA可沿「块摘要—top-k—当前块因果注意力—稀疏 kernel」这一链条阅读. 路线首页提供共同坐标, 具体论文页负责参数、实验与代码路径.

下一层的细粒度 indexer 会把候选生成压到 token 或小组级, 常用低维代理分数先召回、再由原 QK 精排. 它能减少块内多算, 同时更容易被 top-k、随机访存和 KV 反量化限制. 跨层共享则进一步复用 main KV、indexer K 或 top-k 索引, 用较旧的路由结果换取更低开销. 三条路线应使用同一质量和性能口径, 才能比较节省来自稀疏本身还是来自共享.

阅读任何动态路由方案时, 先沿数据路径逐项核对: 路由输入是什么, 索引 shape 是什么, 候选何时生成, 完整 KV 何时读取, softmax 在哪里归一化, 缓存保存什么, 未命中怎样影响后续层. 这些问题能把相似的热图还原为不同的训练与执行机制.

3.2. 粗粒度路由的统一数学对象

把长度为 $N$ 的序列切成 $B=\lceil N/b\rceil$ 个块，每块最多 $b$ 个 token。第 $u$ 个 query 块记为 $Q_u\in\mathbb R^{b\times d}$ ，第 $v$ 个历史块为 $(K_v,V_v)$ 。路由器先为块生成摘要 $c_v=f(K_v,V_v)$ ，再计算

s_{uv}=g(Q_u,c_v),\qquad \mathcal I_u=\operatorname{TopK}_{v\le u}(s_{uv},k).

稀疏注意力只展开候选集合中的原始 K/V：

O_u=\operatorname{Attn}\left(Q_u, \operatorname{Concat}_{v\in\mathcal I_u}K_v, \operatorname{Concat}_{v\in\mathcal I_u}V_v\right).

这三步分别对应「压缩」「检索」「精算」。摘要并不直接替代原始 Value，它只决定哪些原始块值得读取；因此质量取决于 $f,g$ 的召回，性能取决于摘要生成、全块打分、top-k 和稀疏核的总成本。只计算末尾一项，会把路由器当成免费组件。

若每个 query 块与全部 $B$ 个摘要打分，索引复杂度约为 $O(B^2d_r)$ ，其中 $d_r$ 是路由维度；精算部分约为 $O(Bkb^2d)$ 。代入 $B=N/b$ ：

C_{route}=O\left(\frac{N^2d_r}{b^2}\right),\qquad C_{attn}=O(Nkbd).

块越大，路由矩阵越小，块内多算越多；块越小，候选更精确，top-k 与不规则访存更重。这不是单调的「块越细越好」。设相关 token 在命中块内平均只占比例 $\rho_b$ ，每读取一个有效 token 需要搬运约 $1/\rho_b$ 个 token； $\rho_b$ 随块变小上升，kernel 效率却可能下降。最佳块大小由召回曲线、显存带宽和 tile 形状共同决定。

局部窗口、全局 sink 与动态块最好写成集合并集：

\mathcal N_u=\mathcal L_u\cup\mathcal G\cup\mathcal I_u.

实现前要去重。若同一块出现在局部和动态集合里，数学上只能进入 softmax 一次；重复加载后把两份 logit 当作两个元素，会把该块概率质量近似翻倍。正确的 kernel 可以分别遍历多个分支，但在线 softmax必须以唯一逻辑位置合并。性能报告也应使用唯一 tile 数，否则名义稀疏率和实际读取量不一致。

摘要是检索键，不是块的完整语义。均值池化摘要 $c_v=b^{-1}\sum_iK_{v,i}$ 隐含假设：同一块内与未来 query 相关的方向能够在线性平均后保留。若块内大多数 token 属于主题 A，只有一个 token 含关键实体 B，均值会把 B 稀释。最大池化保留各维极值，却可能把来自不同 token 的维度拼成一个不存在的向量。学习型压缩器能针对训练分布调整，仍受固定维度的信息瓶颈限制。

一个块同时含互不相关的两段事实时，单摘要尤其危险。对查询 $q_1,q_2$ ，理想打分分别由不同 token 决定：

s_v(q_1)=\max_{i\in v}q_1^Tk_i,\qquad s_v(q_2)=\max_{i\in v}q_2^Tk_i.

单个中心 $c_v$ 未必同时逼近两个最大值。多摘要可把块切成若干子中心，召回提高，控制数据和打分又相应增加。粗粒度路线的上限因此与块内语义熵有关：内容越混杂，单摘要越难代表；若输入天然按文档、段落或函数边界分块，摘要会更可靠。

摘要训练还存在目标错位。用语言模型 loss 端到端学习时，梯度偏向高频、容易影响下一个 token 的关系；真正的长程关键事实出现稀少，可能被平均 loss 淹没。额外的检索监督可以用稠密 attention top mass、已知引用边或任务答案位置构造，但每种标签只代表一种相关性。稠密 attention 本身也可能分散或错误，不能自动视为真值。

Query 块共享候选的代价。为了让 kernel 规整，常让一个 query 块的所有 token 共享 $\mathcal I_u$ 。若块中不同 token 需要不同历史，候选取并集才能维持召回。设每个 query token 各需 $k_0$ 个块，彼此重叠率低，块级并集可能接近 $bk_0$ ，稀疏收益迅速消失；若只取固定 $k$ ，则部分 token 必然漏召回。

可以用每个 token 的理想候选集合 $S_i$ 定义共享损失：

L_{share}=1-\frac1b\sum_{i\in Q_u} \frac{|S_i\cap\mathcal I_u|}{|S_i|}.

相邻 token 的语义连续时， $S_i$ 重叠较高；代码跳转、表格列和多问题拼接会让重叠下降。评测应按输入类型报告，而非只给平均 recall。自适应 query 块可在语义边界切分，但动态 shape 会增加调度难度；固定物理 tile 加逻辑子组是一种折中。

路由训练的梯度与不可微选择。 top-k 是离散操作。若只有被选块进入 attention，未选块无法从主 attention 输出获得直接梯度，路由器容易延续早期的错误选择。常见缓解包括连续松弛、辅助蒸馏、共享可微分支和探索候选。四种方法改变的假设不同。

连续松弛在训练时对所有或更多摘要做 softmax，使用概率加权近似硬选择：

\tilde c_u=\sum_{v\le u}p_{uv}c_v,\qquad p_{uv}=\frac{e^{s_{uv}/\tau}}{\sum_je^{s_{uj}/\tau}}.

温度 $\tau$ 降低时分布更尖，接近硬路由；训练仍计算全部摘要贡献，省算只在推理出现。松弛目标优化的是摘要混合，不完全等于选择原始 K/V。若多个摘要的线性组合能表达答案，而任何单个 top-k 集合都不够，训练—推理会出现离散化缺口。

蒸馏让路由分数拟合稠密 attention 聚合质量。令块质量

m_{uv}=\sum_{i\in Q_u}\sum_{j\in K_v}\alpha_{ij}^{dense},

可用 KL、排序或 top-k recall loss 训练 $s_{uv}$ 。它给所有块监督，代价是训练需要稠密 teacher，长序列上成本高；并且 $m_{uv}$ 偏爱总质量大的块，一个只对单个关键 query 有高权重的块可能仍被平均淹没。使用最大值或高分位聚合能保护尖峰，又会让路由对噪声异常值敏感。

共享可微分支让摘要注意力始终参与输出，NSA 一类结构正利用这一点。即使某块未进入精细选择，它仍通过压缩分支获得梯度。隐含假设是适合低分辨率压缩注意的块，也应在高分辨率阶段被选中。若摘要足以提供主题，却丢掉块内精确数值，压缩分支可能认为它不必精读，选择分支反而漏掉关键细节。

探索候选在训练中给 top-k 外块少量随机机会，类似路由器的探索。随机率太低无法纠正早期偏差，太高又增加计算并扰动优化。可以按路由不确定性分配：top-k 边界分差小时增加探索，分数稳定时减少。探索样本的梯度要区分「被路由器主动选中」和「随机插入」，否则随机命中会被误当成当前排序正确。

负载均衡也不等于语义召回。若把历史块视作专家，均衡损失会鼓励不同块被相近频率选择；真实文本中近期块、系统提示或关键定义本就更常用。强行均衡可能牺牲质量。容量约束适合限制极端热点和硬件拥塞，目标应是允许符合数据分布的偏斜，同时避免少数块让所有 query 集中读取。

路由器坍缩的三种形态。第一种是位置坍缩：无论内容都选最近若干块，动态分支退化成更宽局部窗口。第二种是 sink 坍缩：总选开头或特殊 token 所在块。第三种是主题坍缩：高频领域块长期占据候选，稀有事实很少被检索。只看负载直方图能发现前两类的一部分，无法判断选择是否正确。

诊断需要条件化统计。比较固定相对位置下的内容扰动，检验路由是否随语义变化；打乱块顺序但保留位置编码，观察选择是否跟内容移动；插入相似干扰块，测量目标块的名次变化；按领域和频率分桶报告召回。若交换内容后候选仍固定在原位置，路由器主要学了位置先验。

位置先验并非全无价值。语言本来具有强局部性，稳定局部分支可以直接编码它，让动态分支专注非局部关系。把两者混在一个 top-k 路由器中，容易让便宜的近期命中占满预算。候选集合先预留局部窗口，再从其余历史中动态选择，能把职责分开；训练损失也应单独评估动态分支是否提供新增覆盖。

4. 选择误差怎样传到注意力输出

设稠密注意力概率为 $\alpha_j$ ，候选集合为 $S$ ，漏掉的概率质量为 $\delta=\sum_{j\notin S}\alpha_j$ 。稀疏注意力在 $S$ 内重新归一化： $\hat\alpha_j=\alpha_j/(1-\delta)$ 。若 $|v_j|\le V$ ，输出误差满足

\left\|\sum_j\alpha_jv_j- \sum_{j\in S}\hat\alpha_jv_j\right\|\le2V\delta.

这个上界说明应该测概率质量覆盖，而非只数 top token 命中率。它仍很松：被漏 Value 方向互相抵消时实际误差小；一个低概率但 Value 范数极大的 token 会违反统一 $V$ 的温和假设。多层网络还会经残差、归一化和非线性放大或吸收误差，单层覆盖率不能完全预测最终任务。

块级召回有两层误差。路由器先决定目标 token 所在块是否入选；块入选后，原始 QK softmax 再决定目标 token 的权重。若块摘要与具体 token 相关性不一致，第一层漏掉；若候选包含很多干扰 token，第二层被重新归一化稀释。扩大块可以减少第一层漏召回，却增加第二层竞争和读取成本。

多头注意力还会放大稀有失败。平均所有 head 的覆盖率很高，某个专门负责引用、分隔或位置的 head 仍可能持续漏召回。若各 query head 独立路由，索引量和分歧增加；若 GQA 组共享候选，组内 head 需求冲突。可以取组内分数最大、均值或并集：最大保护任一 head 的尖峰，均值偏向共识，并集最稳但预算膨胀。选择应与实际 kernel 的 head 分组一起评测。

长度增长会改变极值。候选块从 512 增至 8192，即使无关分数分布不变，最大噪声分数也会上升，目标块名次可能下降。固定 top-k 的覆盖比例同时缩小。训练在 32K、推理到 1M 时，路由器面对的并非同一分类问题。可以按长度校准分数、分层检索或让 $k$ 随候选数缓慢增长，仍要实测额外成本。

分层检索先在大区段中选若干，再在区段内选细块，复杂度从全量打分降下来。它引入级联召回：目标若在第一层被删，第二层无从恢复。总召回近似受各层召回乘积限制。例如第一层召回为 95%, 第二层在通过第一层的目标上条件召回为 95%, 联合召回为 90.25%。层数增加前要确保每层非常高的召回，或保留跨层旁路候选。

4.1. 因果掩码、位置与块边界

自回归注意力要求 query 只能访问当前位置之前。完整历史块满足矩形因果关系，当前 query 块对应的 key 块需要块内三角掩码。若路由器以整块摘要打分，摘要不能包含 query token 尚不可见的未来位置。训练时一次并行计算全序列，最容易在这里发生泄漏。

对第 $u$ 块第 $r$ 个 query，当前块摘要只能由 $0\ldots r$ 的 key 构造。若预先计算整块中心再供块内所有 query 使用，前部 token 已间接看见后部 key。安全做法是当前块始终走局部三角分支，不参与动态历史摘要；完整块结束后才写入索引。若需要块内动态选择，则使用前缀可更新摘要，并验证每个位置的可见范围。

packed sequence 把多条样本放在同一张量中，块可能跨样本边界。仅按全局位置 $v\le u$ 过滤不够，候选还需匹配 sequence ID。末尾不足整块的 padding 不能进入摘要和 softmax。一个越界块可能读取另一用户或另一训练样本的 KV，既是正确性错误也是数据隔离问题。

RoPE 等位置编码通常在 K 上携带绝对或相对位置信息。摘要若直接平均已旋转 K，不同位置相位可能抵消；先聚合未旋转表示再选一个代表位置，又会扭曲查询相关性。可学习路由器常使用独立低维投影与位置特征，主 attention 仍读取原始带位置 K。路由分数与最终 QK 因此属于两个空间，需要专门测量排序一致性。

块边界人为切断连续片段。重叠块能让边界 token 出现在两个摘要中，提高召回，却产生重复候选和更多状态。按语义边界切分可减少混杂，但块长不规则，kernel 仍需落到固定 tile。可采用逻辑段落映射到若干物理 tile：路由按逻辑段落打分，执行展开为物理块；超长段落再分片，短段落合并时保留子段标识。

位置外推还会影响路由先验。训练长度内「很远」可能是 16K，百万上下文中它只是相对近处。若路由分数直接使用绝对块号嵌入，未见位置缺乏校准。相对距离桶、对数距离或内容主导的打分更容易外推，但模型仍可能依赖训练中的长度分布。评测必须覆盖超过训练长度的多个点，观察召回和候选距离分布怎样变化。