同一个大模型,换一种芯片,生成速度可能差很多。原因往往不在算力,而在一份叫 KV 缓存的"笔记"怎么被搬来搬去。IBM Technology 的这节课用 GPU、TPU 和 IBM Spyre 三种芯片讲清了这件事。
模型的"笔记":KV 缓存
大模型一个词元一个词元地生成回答,每生成一个都要回看整段对话。如果每次都从头计算,会非常浪费。所以模型把读过的内容记成"笔记":在每个注意力层里,先前词元留下的键(keys)和值(values)可以复用,新词元只需做新的计算。
核心矛盾:大容量 vs. 快速本地存储
上下文越长,笔记越大,需要大容量存储;但计算最适合在小而快的本地工作集上进行。大的慢,快的小。所以真正的难题是在各级存储和计算单元之间高效搬运注意力所需的 KV 数据,也就是:谁来编排搬运?
三种芯片,三种答案
- 数据中心 GPU:运行时加内核。KV 缓存放在设备内存(常为 HBM)。服务运行时决定每个请求的 KV 怎么分配,例如 vLLM 按需发放固定大小的小块并记录位置表,思路类似操作系统的虚拟内存;注意力内核决定访问模式。
- TPU:编译器加内核。同样有大容量 HBM 和片上快速的向量内存(VMEM),但由编译器提前规划分块(tiling),并在计算当前块时预取下一块(流水线)。
- IBM Spyre:显式软件管理。数据流加速器,每卡 128 GB LPDDR5、32 个活跃核心,每核 2 MB 可编程 SRAM 暂存器。它不像硬件缓存那样自动决定保留什么,而是由编译器划分工作、调度加载与存储。
我的理解
同样的 KV 缓存状态,在不同芯片上放的位置、到达计算单元的路径、负责编排的层次都不同,这就是速度差异的来源。比较芯片时,不要只看算力,也要问它如何搬运模型的记忆状态。
术语表
- KV cache KV 缓存:保存已处理词元的键和值,避免重复计算历史。
- HBM 高带宽内存:加速器上容量大、相对较慢的设备内存。
- Scratchpad 暂存器:由软件显式管理的片上 SRAM。
- Tiling 分块:把数据切成小块并预取下一块。
- Kernel 内核:决定注意力计算的访问模式。
自测一题
TPU 上主要由谁编排 KV 数据搬运?
A. 硬件缓存自动处理 B. 用户手动调度 C. 编译器提前规划分块并流水线预取 D. 只靠运行时分页
和求职面试的关系
做 AI 推理、MLOps 或后端岗位,面试常问:为什么长上下文会变慢?KV 缓存占多少显存、怎么优化?vLLM 的分页思路是什么?能从"容量与带宽的矛盾"和"谁来编排"两个角度回答,会比只背名词更有说服力。
答案:C。TPU 由编译器加内核,通过 tile 与 pipeline 在 HBM 和 VMEM 间搬运。
来源
内容整理自 IBM Technology 视频:Why AI Models Run Faster on Some Chips: The KV Cache Effect。完整学习卡与更多自测题:课程页。
没有评论:
发表评论