JSSA：用联合子空间与免聚集内核加速长上下文解码

长上下文模型的推理瓶颈，往往不只来自矩阵计算，也来自每个生成步都要读取大量 KV-cache。随着上下文长度增长，完整 KV 注意力需要搬运的数据持续增加，因此只计算一部分相关 token，成为降低解码成本的重要方向。

Jiwon Song 等人提出的 Joint Subspace Sparse Attention（JSSA），试图同时解决稀疏注意力的两个实际难题：如何选得准，以及选完之后如何真正跑得快。

核心方法

• 在联合子空间中进行筛选。 许多 token 级方法会把查询和键压缩到较低维空间，再据此估计相关性。但如果压缩过程破坏了查询—键之间的结构，排名就可能不准确。JSSA 学习一个共享的低维查询—键子空间，以最小化 token 分数的保持误差。
• 将难优化目标转化为可解形式。 论文原始目标是非凸的。作者构造了一个可处理的边界目标，并将其转化为具有闭式谱解的边际目标，从而可以直接求得子空间。论文还报告，在测试设置中，对精确目标继续做直接细化并没有带来可测量的改进。
• 用免聚集内核执行稀疏读取。 稀疏选择通常会产生分散的 token 访问；如果依赖传统 gather 操作，索引和内存访问成本可能抵消稀疏计算收益。JSSA 设计了 gather-free 的 token 级 indexed-attention kernel，并支持在标准 GQA KV-cache 布局下，为不同 KV head 使用不同的 token 选择结果。

实验观察

作者在标准 GQA、滑动窗口与全局注意力混合架构，以及 Multi-head Latent Attention（MLA）模型上进行评估。摘要显示，JSSA 的结果整体接近 full-KV attention，同时避免了竞争方法在生成长度上的明显膨胀和部分架构相关的失效问题。

在 NVIDIA B200 GPU 上、128K 上下文条件下，JSSA 相比 full-KV attention 的解码吞吐最高提升 6.43 倍。对于大规模模型，论文称其推理准确率与模型原生 DSA 相当，不需要额外训练；即使在 700B 规模模型上，解码吞吐也可达到 DSA 的 96%—99%。这些结果来自论文所述测试设置，不应直接视为所有硬件和模型的普遍表现。

意义与限制

JSSA 的价值在于把“稀疏选择质量”和“稀疏访问效率”放在同一个系统问题中处理。仅优化打分模型，可能无法消除随机访存；仅优化内核，也可能因为选错 token 而损失质量。JSSA 的联合设计为长上下文服务提供了一条更贴近实际部署的路径，尤其适合 KV-cache 流量占主导的解码场景。

不过，方法仍依赖特定的选择策略、注意力架构和 GPU 内核实现。摘要没有提供完整的选择比例、端到端延迟构成或不同硬件上的全面对比，因此其收益需要结合具体模型、上下文长度与部署环境进一步验证。

来源：Hugging Face Daily Papers