Why Decode Is Memory Bound

为什么值得回答

它决定 decode 优化应优先关注计算量、内存带宽、批大小还是数据复用。

验证路径

  • 计算单步 decode 的参数与 KV Cache 读取量
  • 计算对应 FLOPs 和算术强度
  • 与目标 GPU 的 roofline 边界比较
  • 改变 batch size 并观察瓶颈迁移

机制草图

Decode 每步只生成少量新 token,却要从显存读取大量层权重,并访问随历史长度增长的 KV;相较大矩阵乘,算术强度更容易落在带宽/访存一侧。增大 batch 能复用权重、提高算术强度,但会增加 KV、排队和 p99。

1
2
3
4
batch / context ↑
→ 权重读取摊薄、吞吐可能 ↑
→ KV 与调度压力 ↑
→ 某个拐点后延迟/OOM 主导

最小验证要同时测 single-request 与并发下的 TPOT、显存读写、KV 占用、GPU kernel 时间和 p99;再比较 MQA/GQA、量化、paged KV 或 speculative decoding 的影响。相关节点:DecodeKV Cache

从一个解码步看 roofline

以每层 hidden size 为 d 的模型为例,一个 decode step 的新输入只有一个 token,但每层仍要使用完整权重,并读取该序列已有的 K/V。权重读取规模大致跟参数量和 dtype 成正比,KV 读取则跟层数、KV heads、head dimension、历史长度和 batch 成正比。新增 token 的计算不会随历史长度线性增加到同等程度,因此上下文变长时,带宽压力和 cache 访问更快成为主导。

1
2
3
4
5
one-token dependency
→ no long sequence GEMM within one request
→ weight/KV reads dominate useful FLOPs
→ arithmetic intensity low
→ HBM bandwidth and cache layout set TPOT

“memory-bound”不是说 GPU 没有计算,而是说继续增加峰值 FLOPs 不一定降低单 token 时间。量化若只压缩权重但引入昂贵反量化,或分页导致访问不连续,实际收益可能小于理论字节数下降。

batch 会改变而不是消除瓶颈

并发 decode 时,同一批请求可以复用一次权重加载,算术强度上升;但每个请求的历史长度不同,KV 变成不规则 gather,且 batch 增大带来显存占用、排队和尾延迟。于是会出现一个拐点:吞吐继续增加,但单请求 TPOT/P99 恶化,或者新请求因 KV 不够被拒绝。Continuous Batching 和 PagedAttention 解决的是利用率和碎片化,不会改变自回归依赖本身。

实验如何证伪误判

固定模型和 kernel,做 batch × history length 网格,至少记录每 token 的 HBM bytes、KV block 命中/访问、有效 FLOPs、TPOT、P99 和 OOM。再逐个打开 GQA、权重量化、KV 量化、speculative decoding,观察改善来自读写减少、batch 复用还是串行步数减少。如果 GPU utilization 提升但 TPOT 不变,可能只是把更多请求排队进 GPU;如果单请求变快而并发吞吐下降,则优化改变了服务目标,不能只报其中一个数字。