Prism Attention: Proposal-Refined Index Sharing Mechanism for Efficient LLMs Inference¶

会议: NeurIPS 2026
arXiv: 2610.22985
OpenReview: https://openreview.net/forum?id=L5vrlNwsCo
领域: LLM 其他
关键词: 稀疏注意力, KV Cache 优化, 长上下文推理, 跨层索引共享, INT2 量化

一句话总结¶

Prism Attention 针对现有跨层索引共享在深层注意漂移导致的精度衰减痛点，提出“粗选提案-局部细化”范式，以 INT2 量化打分构建扩展候选池并在后续层做与长度解耦的轻量重排，在 128K 上下文实现 3.4 倍端到端解码加速且精度匹敌全注意力。

研究背景与动机¶

随着大语言模型在数学推理、代码生成和复杂智能体任务中深入探索测试时计算（Test-Time Compute），自回归解码过程中的上下文序列被急剧推长至 64K 甚至 128K 以上。在此阶段，注意力机制的显存带宽与计算复杂度成为主要瓶颈：每个新生成的 token 都必须与历史不断膨胀的 Key-Value (KV) Cache 进行完整交互，导致解码延迟随序列长度呈线性增长。为了打破该瓶颈，Token 级稀疏注意力机制（如 DeepSeek Sparse Attention, DSA）凭借其保留关键 token、避开粗粒度块边界截断的优势受到广泛关注。然而，为了在每一层找出与当前 query 最相关的 Top-\(\kappa\) token，模型必须对整段上下文打分，这种层层重复的索引构建开销往往抵消了稀疏计算带来的实际吞吐收益。

针对这一问题，学术界提出了跨层索引共享方案（如 TidalDecode、LessIsMore、OmniKV 和 IndexCache），仅在少数指定的锚点层（Anchor Layer）计算全局 Top-\(\kappa\) 索引，并在后续多个连续层中直接复用该索引集合。然而，这种静态复用范式隐含着“注意力模式随模型深度保持不变”的理想假设。在实际推理中，注意力分布随着网络深度产生动态漂移；层与层之间距离越远，真实 Top-\(\kappa\) 索引与锚点层索引的重叠率呈指数级锐减，在超长上下文（如 8K 相较于 1K）下该衰减尤为严重。静态强行复用不仅会使深层注意力持续偏离最优支撑集，更会造成关键推导依据丢失，使模型在需要密集逻辑的长程推理与问答基准上遭遇严重的性能退化。

本文深入审视了跨层注意力的动态演化规律，发现了一个关键特性：层间显著 token 的漂移具有极高的局部集中性（Localized Shift）。尽管后续各层真实的 Top-\(\kappa\) 集合与锚点层并不完全重合，但核心高注意力 token 保持高度稳定，发生变动的 token 绝大多数仅仅聚集在先前入选边界的邻近区域，而非散落于全上下文。这意味着只需将锚点层的候选集适度放大为一个宽松的“提案池”，就能以极高召回率包容后续层的漂移。核心 idea：用“粗粒度提案生成 + 细粒度局部重排”的渐进式筛选策略取代僵硬的静态跨层复用，提案层利用 INT2 量化快速划定扩展候选边界，细化层仅在候选池内执行与序列长度解耦的轻量重排，以极低开销兼顾跨层效率与逐层自适应精度。

方法详解¶

整体框架¶

Prism Attention 整体由“离线冗余校准划分”、“提案层量化候选粗选”以及“细化层局部索引重排”三个核心阶段组成。在网络结构上，模型所有 Transformer 层被划分为两类：少数负责全局把关的提案层（Proposal Layer，记作 \(\mathcal{F}\)）和多数负责轻量适配的细化层（Refined Layer，记作 \(\mathcal{S}\)）。在推理自回归生成第 \(t\) 步时，提案层利用超轻量的 INT2 量化键缓存，对全上下文做一次低带宽占用的近似打分，筛选出尺寸扩大为 \(n = \rho \kappa\)（\(\rho \ge 1\) 为扩展系数）的候选池 \(\mathcal{C}_t^{\mathcal{F}}\)；提案层自身使用该候选池的前 \(\kappa\) 项作为稀疏支撑集，并将该候选池传递给后续所属的细化层。细化层不再扫描整个上下文，而是仅在继承的扩展候选池内部执行局部重评分，挑选出各自层专属的最优 Top-\(\kappa\) 索引。整个方法的自回归数据流如下所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Query 与历史 KV Cache"] --> B["离线冗余校准<br/>贪心算法确定提案层与细化层分配 (1:3 比例)"]
    B --> C["候选池量化粗选<br/>提案层 INT2 打分与 GQA 聚合，产出 Top-ρκ 扩展候选池"]
    C -->|传递候选池 C_t| D["候选池约束局部重排<br/>细化层仅在池内重评分，选出层专属 Top-κ 索引"]
    C -->|前 κ 项索引| E["稀疏解码计算<br/>Gather 精确 KV 并执行稀疏 FlashAttention"]
    D --> E
    E --> F["输出当前步解码 Token"]

关键设计¶

1. 离线冗余校准：自动化确定提案层与细化层的网络拓扑

传统的跨层索引共享要么依赖人工经验固定步长，要么忽略了不同层间注意力演化速率差异。Prism Attention 继承并优化了无须重新微调训练的离线贪心校准机制。在较小的校准集（如 OpenR1-MATH-220k 中的 64 条样本）上，将目标序列长度设为 32K，以 token 级交叉熵损失为优化目标，为每一层决策其归属为提案层还是细化层。以典型设置 group-4（提案层与细化层比例为 1:3）为例，校准算法自动感知深浅层注意力模式的剧烈跳跃点（例如浅层特征提取与深层输出映射通常被选为提案层），将深层中模式相近的层组织成依赖簇，确保每个细化层与对应提案层之间的语义漂移始终处于可控范围。

2. 候选池量化粗选：以 INT2 低比特键缓存支撑极低带宽的全局筛选

在长上下文解码阶段，提案层即便仅需为少数层构建候选池，直接读取 FP16/BF16 全精度历史 Key 也会受到显存读取带宽的严重挤压。Prism 在提案层引入了分组非对称 2-bit（INT2）整型量化方案。对于每个注意力头的键向量 \(k_i^{\mathcal{F}} \in \mathbb{R}^d\)，将其划分为 \(M = d / g\) 个互不重叠的子向量（组大小 \(g\)），各子组独立计算缩放系数与零点：

\[\alpha_{i,m}^{\mathcal{F}} = \frac{\max(k_{i,m}^{\mathcal{F}}) - \min(k_{i,m}^{\mathcal{F}})}{2^b - 1}, \quad \beta_{i,m}^{\mathcal{F}} = \min(k_{i,m}^{\mathcal{F}})\]

量化后的整数键 \(\tilde{k}_{i,m}^{\mathcal{F}}\) 表示为：

\[\tilde{k}_{i,m}^{\mathcal{F}} = \operatorname{clip}\left(\left\lfloor \frac{k_{i,m}^{\mathcal{F}} - \beta_{i,m}^{\mathcal{F}}}{\alpha_{i,m}^{\mathcal{F}}} \right\rceil, 0, 2^b - 1\right)\]

在分组查询注意力（Grouped Query Attention, GQA）架构下，多个 query head 共享一个 KV head。为了避免对每个 query head 重复访存加载同一份量化键，Prism 采用组内 Query 最大化降维（Max-reduction）操作：

\[\tilde{s}_t^{\mathcal{F}}(i) = \max_{j \in \mathcal{G}} \big\langle q_{t,j}^{\mathcal{F}}, \tilde{k}_i^{\mathcal{F}} \big\rangle, \quad i \in [L_t]\]

这一设计使得每个 KV head 在解码步只需被读取一次，显存访存量削减了数倍。通过对评分 \(\tilde{s}_t^{\mathcal{F}}\) 取 Top-\(n\)（其中 \(n = \rho \kappa\)，\(\rho \in [2, 4]\)），便得到了既紧凑又高召回的扩展候选池 \(\mathcal{C}_t^{\mathcal{F}}\)。

3. 候选池约束局部重排：与上下文长度解耦的高精度层级自适应

针对静态复用导致深层误差累积的痼疾，细化层打破“全盘照搬”的教条，实施细粒度的局部重排机制。在接收到提案层传递的候选池 \(\mathcal{C}_t^{f(\mathcal{S})}\) 之后，细化层 \(\mathcal{S}\) 拒绝在整个序列长度 \(L_t\) 内进行昂贵的全局搜索，而仅将计算严格限定在候选池大小 \(n = \rho \kappa\) 内：

\[\widehat{s}_t^{\mathcal{S}}(i) = \phi\big(q_t^{\mathcal{S}}, k_i^{\mathcal{S}}\big), \quad i \in \mathcal{C}_t^{f(\mathcal{S})}\]

最终通过局部排序筛选出专属于当前层的稀疏支撑集：

\[\widehat{\mathcal{T}}_t^{\mathcal{S}} = \operatorname{TopK}\big(\widehat{s}_t^{\mathcal{S}}, \kappa\big)\]

由于候选池规模 \(n\) 是固定超参数（通常为 4096 或 8192），无论输入上下文是从 8K 延展到 128K 乃至更大，细化层的索引重排延迟始终恒定在 0.011 ms 左右。该机制以可以忽略不计的微小开销，成功恢复了静态复用遗漏的高价值 token，彻底终结了长程解码中注意力质量与计算效率不可兼得的矛盾。

损失函数 / 训练策略¶

Prism Attention 为完全免训练（Training-Free）推理加速框架，不需要对底座模型进行任何参数微调或梯度反向传播。离线阶段仅运行一层贪心校准搜索，以校准样本的自回归交叉熵损失为标尺，在数十秒内确定提案层与细化层的拓扑对应关系。在在线部署阶段，采用基于 Triton 深度定制的 CUDA Graph 算子链，将 INT2 量化、Top-K 粗选、局部重排及稀疏 Gather 注意力高效融合。

实验关键数据¶

主实验¶

评估在 Qwen3-4B 与 Qwen3-8B 模型上展开，涵盖数学与复杂推理基准（AIME24、AIME25、GPQA-Diamond、MATH500，以 Pass@1 为指标，最大生成长度 32768）以及长上下文理解基准 LongBench（9 个平均输入长度超过 8192 的任务子集）。对比基线包括标准 Full Attention 及四大代表性稀疏方法：Quest、SeerAttention-R (Seer-R)、LessIsMore (LIM) 和 IndexCache (IDC)。

表 1：推理基准与长文本评测性能对比 (Qwen3-4B / Qwen3-8B)

模型	方法	AIME24	AIME25	GPQA	MATH500	推理均分	LongBench 均分
Qwen3-4B	Full Attention	71.25	66.41	56.19	93.93	71.95	55.07
	Quest	46.67	39.79	24.74	73.40	46.15	-
	Seer-R	69.32	58.59	55.40	94.10	69.35	54.37
	LessIsMore	71.48	62.87	56.23	94.12	71.18	52.36
	IndexCache	70.99	61.30	54.29	94.05	70.16	53.27
	Prism Attention	72.60	64.27	54.83	94.00	71.43	54.98
Qwen3-8B	Full Attention	74.48	67.86	60.54	94.43	74.33	54.58
	Quest	59.79	57.71	47.22	87.80	63.13	-
	Seer-R	71.35	57.81	60.48	94.00	70.91	53.64
	LessIsMore	73.00	65.24	60.65	94.55	73.36	50.98
	IndexCache	74.84	64.47	60.39	94.40	73.53	52.65
	Prism Attention	75.05	65.52	60.45	94.30	73.83	54.97

消融实验¶

消融实验在 Qwen3-4B 上系统分析了逐层索引独立计算、静态索引复用、INT2 量化、分组结构、离线校准以及“提案-细化（Proposal-Refined）”机制对模型精度的独立与协同贡献（稀疏预算 \(\kappa=1024\)，扩展系数 \(\rho=4\)）。

表 2：索引构建与提案-细化机制消融实验 (Qwen3-4B)

逐层索引独立	静态复用索引	INT2量化	提案-细化比 (Group)	离线校准	提案-细化重排	AIME24	AIME25	说明
✓			逐层独立			71.35	63.87	全精度独立检索（上限参考）
✓		✓	逐层独立			72.08	63.26	逐层独立 + INT2 量化打分
	✓	✓	group-4 (1:3)			65.68	50.05	朴素静态复用（无校准、无重排，精度剧跌）
	✓	✓	group-4 (1:3)	✓		72.19	59.22	静态复用 + 离线校准（IndexCache 复现）
	✓	✓	group-4 (1:3)		✓	71.56	62.60	静态复用 + 提案细化（无校准，已大幅提升）
	✓	✓	group-4 (1:3)	✓	✓	72.03	64.43	完整 Prism 配置（校准 + 提案细化）
	✓	✓	group-9 (1:8)			61.88	43.33	极疏复用（精度严重损失）
	✓	✓	group-9 (1:8)	✓		67.66	52.81	极疏复用 + 离线校准
	✓	✓	group-9 (1:8)		✓	69.58	59.01	极疏复用 + 提案细化
	✓	✓	group-9 (1:8)	✓	✓	70.83	60.89	极疏 Prism 配置（较纯复用 +17.56）

关键发现¶

局部重排挽回关键损失：从表 2 可见，朴素静态复用在 AIME25 上精度跌至 50.05；加入提案-细化重排后迅速拉升至 62.60，配合离线校准进一步达到 64.43，甚至超越了逐层全精度独立检索的 63.87。在更为激进的 group-9（1:8 极疏共享）下，提案-细化重排将 AIME25 精度从 43.33 挽救至 60.89，证明了跨层跳跃越大，局部自适应重排的价值越显著。
扩展倍率 \(\rho\) 的极值敏感性：随候选池扩展系数 \(\rho\) 从 1 增至 4，AIME25 得分从 59.22 攀升至 64.43；但当 \(\rho\) 进一步过度扩大至 8 时，精度轻微回落，表明候选池适度包容局部漂移即可，盲目扩大反而引入过多背景噪声。
长序列端到端显著加速：在 SGLang 框架与 H800 环境下，序列达到 128K（BSZ=16）时，Prism 相较于 FlashAttention 取得 3.36 倍端到端加速（步延迟从 15.03 ms 降至 4.47 ms）；在内核层面达到 7.50 倍加速。细化层引入的恒定 0.011 ms 重排开销在超长序列中几乎被完全摊薄，与 IndexCache 的延迟差距从 16K 时的 15.7% 单调收敛至 128K 时的 0.7%。

亮点与洞察¶

将注意力动态学与稀疏检索有机融合：论文并非盲目叠加量化或搜索工程，而是敏锐抓住“跨层注意力漂移呈局部聚集”这一第一性原理，将传统“单点精确复用”升级为“窗口包容+内部重排”，构思清晰且极具洞察力。
GQA Max-Reduction 极度契合现代硬件访存模型：在提案层巧妙运用同一组内所有 Query 头与量化 Key 的内积最大值作为综合重要性，将多头多次访存合并为单次读取，从算子底层彻底解放了显存带宽。
无状态恒定重排的系统可扩展性：细化层的重排范围严格由超参数 \(n\) 锚定，不随上下文增长而膨胀。这种时间复杂度为 \(O(1)\) 的局部细化设计，对未来面向百万级上下文的推理系统具备极高的借鉴与落地价值。

局限与展望¶

作者承认的局限：提案层与细化层的拓扑切分仍需依赖前置的离线贪心搜索；虽然该过程耗时仅数十秒且无须梯度更新，但在遇到未曾校准的新模型架构或异构配置时仍需单次准备成本。
实测与实验边界局限：实验评测主要集中在 Qwen3 系列与 GLM-5.2 等模型，尽管在附录展示了 DeepSeek-R1-Distill-Llama-8B 与 GLM-4.7-Flash 的注意力模式重叠率，但尚未将全套加速内核直接下沉至极低比特量化大模型中评估复合误差。
未来改进方向：可进一步探索运行时动态自适应扩展因子 \(\rho\)，根据当前生成 token 的注意力熵动态调节候选池大小，在平缓生成阶段采用紧凑候选池，在复杂逻辑转折点扩大局部搜索窗口。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从注意力动态学角度提出跨层提案-局部重排范式，突破了静态索引共享的理论局限。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖推理与长程双重基准、细致的消融与内核级延迟剖析，并在百K长文本上给出详实数据。
写作质量: ⭐⭐⭐⭐⭐ 动机推导层层递进，图表规范工整，原理与硬件感知设计叙述透彻。
价值: ⭐⭐⭐⭐⭐ 兼具算法创新与极高的工业级推理部署落地价值，SGLang 适配表现亮眼。