浏览: 9

概览

NoMAD-Attention 通过**“把 Q·K 乘加改写成寄存器查表＋加法”**，用 CPU SIMD 寄存器的超低延迟随机 shuffle（≈1–2 cycles）来取代传统注意力里的大规模 FMA 运算，从而在 16 k 上下文长度下可为 4-bit LLaMA-7B 带来近 2× 整体吞吐提升，同时几乎不损失精度。(arxiv.org, arxiv.org)

1 设计要点速览

组件

作用

关键数字 / 特性

Product Quantization (PQ)

把每个 K 向量按维度均匀切成 S 段（sub-vector），每段用 16 质心 → 4 bit 编码

常用 d_sub∈{1,2,4}；单 K 仅存 4·S bit (faiss.ai, arxiv.org)

寄存器 LUT

每个解码步中，Query 与 16×S 质心做点积，结果线性量化到 8 bit 并整张装进 128-bit 寄存器

128 = 16 × 8 正好匹配 _mm_shuffle_epi8 随机取数 (stackoverflow.com, news.ycombinator.com)

MAD-free 点积

对任意 K：在每段用其 4-bit 代码索引 LUT，取值后累加

全程 0 乘法，仅寄存器查表＋加法 (arxiv.org, arxiv.org)

Key 编码

新生成的 K 需在线找到最近质心并写入 Key-code cache

只做 O(d) 距离比较，不随序列长度 t 增长 (arxiv.org)

2 「分段」(Sub-quantizer/Segment) 的真正含义

  1. 按特征维度切块：把 d-维向量拆成 S = d/d_sub 个子向量，每段独立量化。

  2. 匹配 SIMD 寄存器：一段 16 质心 × 8 bit = 128 bit；恰好一条 xmm/NEON 可一次随机 shuffle。

  3. 速度–精度权衡：

    • d_sub=1 误差最低，Q·K 阶段加速 ≈ 5–10×；

    • d_sub=4 精度略降，Key 体积仅 4 B，整体吞吐提升可达 2.1×。(proceedings.neurips.cc, arxiv.org)

3 在线推理三步流程

3.1 Key 编码（每新 K 一次）

cs(kt)=arg⁡min⁡c∈[0,15]∥πs(kt)−bs,c∥2,s=1…Sc_s(k_t)=\arg\min_{c\in[0,15]}\|\pi_s(k_t)-b_{s,c}\|^2,\quad s=1\ldots S

得到的 4-bit 代码写入 Key-code Cache。O(d) 乘加，与序列长度无关。(arxiv.org)

3.2 Query 构建 LUT（每步一次）

对每段计算

LUTs[c]=⟨πs(qt), bs,c⟩,  c=0…15\text{LUT}_s[c]=\langle\pi_s(q_t),\,b_{s,c}\rangle,\;c=0\ldots15

再动态映射到 0–255 并存入寄存器。总算量仍 O(d)。(arxiv.org, arxiv.org)

3.3 MAD-free 点积

⟨qt,ki⟩^  =  ∑s=1SLUTs ⁣[codes(ki)]\widehat{\langle q_t,k_i\rangle}\;=\;\sum_{s=1}^{S}\text{LUT}_s\!\bigl[\text{code}_s(k_i)\bigr]

每条 K 只需 S 次寄存器取数＋加法，完全无乘法。(arxiv.org, stackoverflow.com)

4 复杂度与性能

阶段

原生注意力

NoMAD-Attn

Q·K

O(t·d) MAD

查表 + 加法

LUT / Key 编码

–

O(d) MAD

整体解码

重度乘加瓶颈

16 k 序列吞吐 ↑1.8–2.1×

论文 Figure 7 的延迟分解显示，Q·K 被压缩为 1/10 以内，而 Key Caching 仅占 3–6 % 总时延。(neurips.cc, proceedings.neurips.cc)

5 常见疑问 & 澄清

疑问

说明

Q 每步变，岂非还要乘法？

必须为新 Q 重算 LUT，但这只需 O(d) MAD；与原来随 t 爆炸的 O(t·d) 乘加不可同日而语。

“把 K 放进 bucket 后再与 Q 乘”吗？

不完全正确：K 只给出 4-bit 索引，真正参与取值的是 “Query→质心” 的点积，故每条 K 仍有独立权重。

KV Cache 越长开销会回来吗？

Key 编码写 cache 依旧 O(d) 级，查表阶段 0 乘法；累积开销斜率远小于原生注意力。

6 优势与局限概览

优势

  • 零微调、兼容现有权重。

  • 仅依赖 x86 AVX2/AVX-512 或 ARM NEON 的基本 shuffle 指令。

  • Key-cache 更小，减少内存带宽。

局限

  • 质心固定 16 个受 128-bit 限制，表达力受限。

  • MLP 等算子仍需 FMA；整模型上限 ≈ 2×。

  • 老旧 CPU（无 SSSE3/NEON）不支持。

7 术语速查

术语

释义

PQ / Sub-quantizer

产品量化，把向量拆分子空间各自量化。

LUT

Look-Up Table；存 Query × 质心 点积，8 bit 存寄存器。

Key-code Cache

所有 K 的 4-bit 代码按 32 条一块转置存储，方便 SIMD 批查。

d_sub

每段维度；越小误差越低，查表次数越多。

参考文献 / 延伸阅读

  1. NoMAD-Attention 原论文 (2024)(arxiv.org)

  2. NeurIPS 2024 Poster 补充材料 & Figure 7(neurips.cc)

  3. NoMAD-Attention HTML 附录（算法与图示）(arxiv.org)

  4. StackOverflow 对 _mm_shuffle_epi8 用法解释(stackoverflow.com)

  5. FAISS ProductQuantizer API 文档(faiss.ai)

  6. Quicker ADC：SIMD-友好 PQ 查表优化(arxiv.org)

  7. CPU 寄存器访问 1–2 cycle 说明（论文 HTML 第 1 节）(arxiv.org)

  8. SIMD shuffle 1 cycle 延迟讨论（Hacker News 节选）(news.ycombinator.com)

  9. SIMD 性能博客：常见指令延迟(stackoverflow.blog)

  10. Branchfree.org SIMD 分类文章(branchfree.org)

（以上链接均可进一步追踪算法细节与 SIMD 微架构特性）