← 返回归档

Scout Before You Attend: Sketch-and-Walk Sparse Attention for Efficient LLM Inference 深度解读

arXiv: 2602.07397
作者：Hoang Anh Duy Le, Sahil Joshi, Zeyu Yang, Zhaozhuo Xu, Anshumali Shrivastava
机构：Rice University, Stevens Institute of Technology
主题：长上下文推理、稀疏注意力、Hadamard sketch、block-sparse attention、LLM serving

TL;DR

这篇论文提出 Sketch&Walk Attention：一种不需要训练、不改模型权重、同时覆盖 prefill 与 decode 的动态稀疏注意力方法。它的核心判断是：现有很多稀疏注意力方法只根据“当前层的一跳 attention score”选择 key block，但深层 Transformer 的真实依赖往往是跨层、多跳组合出来的；如果早期层把中间 block 或未来会变重要的 block 剪掉，后面层就很难恢复。

Sketch&Walk 的解决方案是两步：

  1. Sketch：用 token-space block averaging + feature-space Hadamard projection，低成本估计 block-level attention score；
  2. Walk：在 sketched block attention matrix 上维护跨层 walk state，选择 top-$\tau$ blocks 做真正的 sparse attention。

论文声称在 20% attention density / 80% sparsity 下，Sketch&Walk 在 LongBench 与 RULER 上接近 dense attention，有时还略好；结合自定义 CUDA/Triton kernel，prefill 最高可到约 6× 加速，decode 最高约 1.6× 加速。

一句话评价：这篇不是又一个“把 attention score 近似一下然后 top-k”的稀疏注意力，而是把 sparse selection 的依据从“局部一跳相似度”提升到“跨层组合后的 block influence”。这是它最有价值的地方。

Sketch&Walk 总览

1. 背景：长上下文推理的瓶颈不只是算得多，而是选错了会断链

长上下文 LLM 推理里，self-attention 的成本随序列长度近似二次增长。对于长度 $n$、head dimension $d$ 的注意力，密集 attention 至少要面对 $O(n^2 d)$ 的 score 计算和巨大的 memory traffic。实际部署中它会同时拖慢两个阶段：

  • Prefill：一次性处理 prompt，sequence length 很长，attention matrix 大；
  • Decode：逐 token 生成，单步 query 少，但要访问越来越长的 KV cache，memory bandwidth 成为瓶颈。

因此 block-sparse attention 很自然：只让 query block attend 到少量 key blocks。但问题在于：怎么选这些 key blocks？

许多已有方法（如 MInference、Quest、FlexPrefill、AdaMAs 等）本质上依赖当前层或当前 step 的某种一跳 relevance score。这个思路默认了一个假设：当前层直接 attention 分数大的 block，就是当前层以及后续层真正重要的 block。

论文指出这个假设并不稳。Transformer 的信息传播是跨层叠加的：token/block $i$ 对 $k$ 的依赖可能不是 $i\to k$ 的直接高分，而是 $i\to j\to k$ 这样的多跳路径。更麻烦的是，attention score 是 inner product，不是 metric similarity：

$$ q_i^\top k_j \text{ 高},\quad q_j^\top k_k \text{ 高} \nRightarrow q_i^\top k_k \text{ 不一定高}. $$

所以一跳 score 可能完全看不出多跳依赖。如果 sparse selection 在早期层把中间节点或终点节点剪掉，后续层就失去了信息通路。这就是论文标题里的 “Scout Before You Attend”：先侦察多跳结构，再决定真正 attend 哪里。

一跳与多跳 attention 可视化

图中 $A^1$ 只反映直接 attention；$A^4$、$A^7$ 则展示 repeated attention composition 后显现出的结构。论文的直觉是：真正应该保留的 block，可能在 $A^1$ 里不突出，但在高阶组合里很重要。

2. 方法总览：Sketch to estimate, Walk to aggregate

Sketch&Walk 的算法可以拆成两个层次：

  • Small-World Sketching (SWS)：快速估计 block-level attention matrix；
  • Sketch-Determined Walk：把每层估计出的 block matrix 组合起来，维护跨层 influence score。

设 query/key 为 $Q,K\in\mathbb{R}^{n\times d}$，把 token 序列切成 $b=\lceil n/B\rceil$ 个 block，每个 block 大小为 $B$。对 block $i,j$，真实 block score 可写成：

$$ A_{ij}^{\text{true}}=\frac{1}{\sqrt d}\frac{1}{B^2}\sum_{s=1}^{B}\sum_{t=1}^{B}Q^{(i)}[s,:]\cdot K^{(j)}[t,:]^\top. $$

直接算所有 block pair 仍然贵，尤其每个 pair 原始上要聚合 $B^2$ 个 token interaction。SWS 用两个 sketch 降成本。

3. Small-World Sketching：token-space + feature-space 双重压缩

3.1 Token-space sketch：block aggregation

对每个 block 取均值：

$$ \bar q_i=\frac{1}{B}\sum_{t=1}^{B}Q^{(i)}[t,:],\qquad \bar k_i=\frac{1}{B}\sum_{t=1}^{B}K^{(i)}[t,:]. $$

这一步把 $n$ 个 token 压成 $b$ 个 block vectors。它隐含一个假设：局部连续 token 有一定语义一致性，或者至少局部表示的方差相对 block mean 不太大。这个假设在自然语言长上下文里常常成立，但在极端结构化输入、代码跳转、表格、混杂检索上下文里可能变弱。

3.2 Feature-space sketch：Hadamard projection

然后用 Subsampled Randomized Hadamard Transform / Hadamard projection 把 $d$ 维压到 $k$ 维：

$$ \tilde q_i=\bar q_i H_d,\qquad \tilde k_i=\bar k_i H_d, $$

其中 $H_d\in\mathbb{R}^{d\times k}$，$k\ll d$。由此估计 block attention：

$$ \widehat A^{\text{SWS}}=\frac{\tilde Q\tilde K^\top}{\sqrt k}= \frac{(\bar QH_d)(\bar KH_d)^\top}{\sqrt k}. $$

论文给出的复杂度直觉是：block pair score 从 $O(B^2d)$ 级别降到 $O(k)$ 级别。更重要的是，这个估计只用于选择 block，真正 attention 仍在原始 QKV 上对选中的 block 计算，因此 sketch 的误差不会直接替代最终 attention value，只影响稀疏模式。

4. Sketch-Determined Walk：把跨层多跳依赖显式纳入选择依据

拿到第 $l$ 层的 sketched block attention matrix $\widehat A^l_{\text{block}}$ 后，Sketch&Walk 不直接根据它的一跳分数 top-k，而是维护 walk state：

$$ R^0=(\widehat A^0_{\text{block}})^s, $$ $$ R^l=R^{l-1}(\widehat A^l_{\text{block}})^s,\quad l>0. $$

这里 $s$ 是 walk degree / sparsity exponent。直观上，$R^l[i,j]$ 不是“当前层 $i$ 对 $j$ 的直接分数”，而是“截至当前层，block $j$ 对 block $i$ 的累计可达影响”。最终对每个 query block，从 $R^l$ 的对应行选 top-$\tau$ key blocks，然后在原始 QKV 上执行 block-sparse attention。

这个设计有两个关键点：

  1. 跨层记忆：selection 不再是每层独立决策，而带有前面层的 influence trace；
  2. 低维 walk：walk 发生在 $b\times b$ block matrix 上，而不是 $n\times n$ token matrix 上，否则成本会抵消稀疏收益。

论文把这个结构称为 small-world，有点类比社交网络：不是每个 token 都直接连接到每个 token，而是通过少数高价值 block transition 在少数跳内到达重要信息。

5. Prefill 与 decode：同一套选择机制，两套工程 kernel

论文强调 Sketch&Walk 是统一方法：prefill 和 decode 都用 SWS + walk state 决定 sparse blocks。

  • 在 prefill 中，输入 prompt 已知，可以按 query blocks 批量生成 block-level sparse pattern，再调用 sparse attention kernel；
  • 在 decode 中，每步新 query 对历史 KV cache 的访问也可用同样的 block-level walk score 来决定保留哪些 key blocks。

这点比很多 stage-specific 方法更干净：有些方法只优化 prefill，有些只适配 decode，而 Sketch&Walk 试图用同一个“先 sketch 再 walk”的选择准则覆盖完整推理链路。工程上则仍需要分别实现高效 prefill sparse kernel 与 decode sparse kernel，因为两阶段的访存形态不同。

6. 理论部分：保证有用，但依赖假设要看清

论文理论部分主要想说明两件事：

  1. SWS 可以高概率保留 block-level inner product / top-$\tau$ selection；
  2. 在 heavy-tailed attention 分布下，只保留 top-$\tau$ blocks 的 attention output 接近 dense attention。

6.1 Block coherence

论文假设同一 block 内 token 表示可分解为：

$$ q_t^{(i)}=\mu_i+\varepsilon_t,\quad \mathbb{E}[\varepsilon_t]=0,\quad \text{Var}(\varepsilon_t)\le \sigma^2I. $$

于是 block mean 满足类似：

$$ \|\bar q_i-\mu_i\|_2\le \frac{\sigma\sqrt{d\log(1/\delta)}}{\sqrt B} $$

的集中界。换句话说，block 越大，均值估计越稳定；但 block 太大又会降低稀疏选择粒度。这是工程调参里的核心 trade-off。

6.2 Hadamard sketch 保 inner product

对固定向量 $u,v$，如果 sketch dimension 足够大：

$$ k=\Omega(\varepsilon^{-2}\log(b/\delta)), $$

则 Hadamard sketch 后的 inner product 近似原始 inner product：

$$ |(u^\top H_d)(H_d^\top v)-u^\top v|\le \varepsilon\|u\|_2\|v\|_2. $$

这解释了为什么 $k=64$ 这类小 sketch dimension 在实验中也可能够用：selection 只要稳定分出头部 block，不需要精确恢复所有 score。

6.3 Heavy-tailed attention 与 output approximation

论文还假设 block attention 是 heavy-tailed：每个 query block 的大部分 attention mass 集中在少量 $\tau$ 个 block 上：

$$ \sum_{j\in S_i^*}\text{softmax}(A_{ij}^{\text{true}})\ge 1-\eta. $$

在这个假设下，sparse output 与 dense output 的误差可被 tail mass 与 block aggregation noise 控制：

$$ \|O^{\text{Sketch\&Walk}}-O^{\text{full}}\|_F \le \eta\|V\|_F+O\left(\frac{\tau\sigma}{\sqrt B}\right)\|V\|_2. $$

我的判断：理论证明给的是“合理性背书”，不是无条件保证。真正关键的假设是 block coherence 和 heavy-tailed attention 是否在目标 workload 上成立。对于自然语言 QA、摘要、检索型长上下文，这通常还行；对于需要精确全局结构、代码跨文件依赖、密集表格推理，风险更大。

7. 实验设置

论文评测覆盖：

  • 模型：Llama-3.1-8B-Instruct、Llama-3.2-1B-Instruct、Qwen2-7B-Instruct；上下文最高 128K；
  • 任务：LongBench（QA、摘要、代码、多文档等真实任务）与 RULER（长上下文检索/位置敏感诊断）；
  • 硬件：单张 NVIDIA H100 94GB；
  • 实现：PyTorch inference pipeline + Triton sparse attention kernels + custom CUDA kernel 做 Hadamard sketch；
  • 默认参数：block size 64，sketch dimension 64，walk degree / exponent 8；前两层通常不稀疏化，因为早期层可稀疏性低；
  • 稀疏水平：accuracy 多在 80% sparsity / 20% density 下评估；加速图中也使用 90% sparsity。

对比基线：

  • Prefill：MInference、FlexPrefill；
  • Decode：QUEST、AdaMAs；
  • Dense：FlashAttention dense baseline。

8. 实验结果：质量基本贴近 dense，Qwen 上甚至有明显收益

8.1 LongBench / RULER accuracy

论文结果可以概括为：Sketch&Walk 在 20% density 下大体接近 dense attention，并且通常优于同稀疏水平下的现有稀疏方法。

几个值得注意的点：

  • Prefill on RULER：Llama-3.1-8B 平均分 Dense 91.44，Sketch&Walk 91.35，几乎无损；Llama-3.2-1B Dense 68.31，Sketch&Walk 67.59，也接近。
  • Decode on RULER：Llama-3.1-8B Dense 91.44，Sketch&Walk 91.08，接近 dense，并优于 Quest 的 89.06、AdaMAs 的 90.74。
  • LongBench end-to-end：Llama-3.1-8B Dense AVG 47.95，Sketch&Walk 47.63；Llama-3.2-1B Dense 29.75，Sketch&Walk 30.10；Qwen2-7B Dense 40.85，Sketch&Walk 43.22。

Qwen2-7B 上 Sketch&Walk 反而超过 dense 的结果很有意思，但要谨慎解读：稀疏化可能起到了某种 regularization / noise filtering 作用，也可能和 benchmark variance、prompt template、实现细节有关。它说明 sparse 不一定只是在“牺牲精度换速度”，但不能简单外推为普遍提升。

8.2 Prefill 加速

Prefill 加速

prefill 加速图显示，随着 context length 从 16K 增长到 128K，Sketch&Walk 的 speedup 变得更明显。论文声称最高可达约 6× end-to-end / prefill 相关加速。这个趋势符合预期：context 越长，dense attention 的 $n^2$ 成本越高，block-sparse 的收益越容易压过 sketch/walk 的额外开销。

论文还指出，相比 MInference 等方法，Sketch&Walk 在较短长上下文区间也更稳定，不是只有极长上下文才开始有效。

8.3 Decode 加速

Decode 加速

decode 阶段的最高 speedup 约 1.6×。这个数字比 prefill 小是合理的：decode 单步 query 少，主要瓶颈常常是 KV cache 读带宽和 kernel launch/访存效率，而不是完整 $n^2$ 矩阵计算。要在 decode 中获得收益，sparse pattern 不仅要少算，还要让 memory access 真的更规整、更少搬数据。

8.4 Kernel overhead

Kernel 分析

kernel analysis 的重点是：Hadamard sketch 和 walk estimation 的 overhead 相对 sparse attention 总成本较小，尤其长上下文下占比会被摊薄。这个结论对方法能否落地很关键，因为很多稀疏注意力方案在 paper math 上省计算，但 selection overhead 和不规则访存会吃掉收益。

Sketch&Walk 的优势在于 selection 在 block level 做，且 sketch dimension 小，能够把“选哪里”的成本控制住。

8.5 Ablation：sketch dimension、walk degree、sparsity

消融实验

消融实验给出三个信号：

  • sketch dimension 16–32 也能保持不错效果，64 通常稳定；
  • walk degree 增大后性能改善并趋于平台，说明不需要无限高阶 walk；
  • 在高 sparsity 下仍较稳，支持它作为受限 serving 场景的 accuracy-efficiency trade-off。

这说明方法的有效性不是完全靠某个极端参数撑起来的，而是有一定鲁棒性。

9. 与已有方法的差异

方法 主要阶段 selection 依据 Sketch&Walk 相对差异
MInference Prefill 预设/估计 sparse pattern，偏一跳结构 Sketch&Walk 显式累计跨层多跳 influence
FlexPrefill Prefill 动态 prefill 稀疏选择 Sketch&Walk 同时覆盖 prefill/decode，且引入 walk state
QUEST Decode decode-time KV/cache relevance Sketch&Walk 不只看当前 query 的局部分数，而看累计 block walk
AdaMAs Decode adaptive sparse decode Sketch&Walk 的 novelty 在 sketch + cross-layer walk，而非单步自适应

最核心差异不是“用了 Hadamard sketch”——这只是工程上让估计便宜；真正差异是 selection criterion：从 one-hop per-layer score 换成 multi-hop cross-layer accumulated score。

10. 贡献总结

我认为这篇论文的贡献可以归纳为四点：

  1. 问题定义清楚：指出一跳稀疏选择会断掉深层多跳依赖，这是很多 sparse attention 工作没有正面处理的问题；
  2. 方法组合合理：用 cheap sketch 得到 block score，再用 walk state 做跨层组合，结构简洁；
  3. 统一 prefill/decode：同一 sparse selection 思路覆盖两个推理阶段，工程叙事完整；
  4. 有 kernel 结果支撑：不仅报告 accuracy，也展示了 prefill/decode speedup 和 kernel overhead。

11. 局限与我会关注的风险

11.1 Block coherence 不总成立

block averaging 默认连续 token 是可聚合的。自然语言段落可以，代码、表格、混合检索片段、工具调用日志未必稳定。如果一个 block 内混合了多个语义对象，block mean 可能抹掉关键 token。

11.2 Walk state 可能引入历史偏置

跨层累计 influence 是优势，也可能是风险。如果早期 sketch 估计错了，walk state 会把错误传播到后续层。论文消融显示总体鲁棒，但极端任务上仍要测。

11.3 Sparse pattern 的硬选择对 tail information 仍敏感

heavy-tailed attention 是很多长上下文任务的经验事实，但不是所有任务都符合。某些需要分散证据聚合的任务，tail mass 可能不是噪声，而是答案所需。

11.4 Serving 集成成本

论文有自定义 CUDA/Triton kernel，但真实 serving 系统还要考虑：batching、paged KV cache、continuous batching、tensor parallel、不同模型结构、GQA/MQA、混合精度、prefix cache 复用。paper kernel 到生产系统仍有距离。

12. 对 LLM serving 的启发

这篇对 serving 的启发挺直接：

  • 稀疏注意力的核心不只是降低 density，而是选对 block；
  • “当前层局部分数”可能不是好 selection signal，尤其长上下文依赖常常跨层形成；
  • block-level low-cost metadata 可以作为 serving runtime 的调度信号；
  • prefill 与 decode 最好共享一套语义一致的 sparse policy，否则端到端行为会割裂；
  • sketch/walk 这类 cheap routing 可能成为未来长上下文推理系统里的常驻 sidecar。

如果把它和 KV cache 压缩、retrieval memory、speculative decoding 放在一起看，Sketch&Walk 更像“attention routing layer”：不压缩模型，不训练 adapter，而是在 runtime 决定注意力计算预算流向哪里。

13. 结论

Sketch&Walk 的价值在于把稀疏注意力从“局部 top-k score”推进到“跨层多跳 influence-aware selection”。它不是完全解决长上下文推理问题，但抓住了一个关键痛点：稀疏化如果只看一跳，很容易省掉后面才会变重要的信息路径。

在实验上，它以 20% density 保持接近 dense 的精度，并在长上下文 prefill 中展示最高约 6× 的加速，decode 也有约 1.6×。更重要的是，论文没有停留在算法层，而是给了 kernel overhead 分析，这让它比纯理论 sparse attention 更接近可落地系统。

我的总体判断：这是一篇值得收进“长上下文推理加速 / sparse attention runtime”谱系的论文。它的思想可以和 KV cache compression、paged attention、prefix cache、speculative decoding 组合，尤其适合探索下一代长上下文 serving runtime 的动态 attention budget 分配。