免费 Kaggle GPU 跑出 28 TPS：跨 WAN 的张量并行实践

免费 Kaggle GPU 跑出 28 TPS：跨 WAN 的张量并行实践 | 前端进阶之旅

Full translation

完整中文译文

先坦白一件事：整个项目跑在免费的 Kaggle T4 Notebook、两台几乎不要钱的 AWS EC2 t3.micro 中继，以及公共互联网上。没有 A100，没有私有数据中心网络，也没有预算。

即便如此，ShardFlow v2.1 在 Qwen2.5-7B 上跑出了 28.10 TPS 的峰值，跨两个独立的云区域，走的还是 WAN。

这就是故事的全貌，以及 v2.1 中那个我完全没有预料到的关键修复。

一个 FP16 的 7B 参数模型需要大约 15 GB VRAM。一块 Kaggle T4 有 16 GB。技术上勉强能装下，但 KV cache 就没地方放了。

解决方案是张量并行：把模型切分到两台机器上。节点 0（Iowa）负责第 0 到 14 层。节点 1（Oregon）负责第 14 到 28 层，外加 LM head 和最终验证。它们通过俄亥俄州一台 EC2 t3.micro 上的 TCP 中继互相通信。

这个配置下的基准吞吐量（没有任何优化）：4.92 TPS。能跑，但不快。

LLM 推理之所以慢，是因为它是串行的。生成一个 token，等一下，再生成下一个，再等一下。每次跨 WAN 的往返大约耗时 86ms RTT。如果每个往返只出一个 token，那整个过程都在和网络做斗争。

投机解码把这个顺序倒过来了。不再一个 token 一个 token 地发，而是在本地跑一个微型草稿模型，提前猜出接下来 K 个 token。然后把 K 个猜测一次发给验证器。如果大模型同意其中 M 个，那就一个往返内确认了 M 个 token，而不是一个。

ShardFlow 用 Qwen2.5-0.5B 作为草稿模型，跑在节点 0 的 cuda:1 上，而 7B 目标分片跑在 cuda:0 上。零 VRAM 争用。草稿器提出 8 个候选，节点 1 并行验证它们，平均每个往返拿到 4.07 个 token，而不是 1 个。

开启投机解码（eager 模式）：14.3 TPS 峰值。3 倍提升。

我以为 14.3 就是天花板了。网络是明显的瓶颈：两个 Kaggle 实例在不同州，中间夹着 EC2 中继，走公共互联网路由。还能咋整？

然后我更仔细地看了草稿模型实际在干什么。

每一轮，生成 8 个候选 token 意味着对 0.5B 模型跑 8 次独立的前向传播。每次前向传播大约启动 1500 个 CUDA kernel，一个一个从 Python 循环里发出去。

问题来了：每个 CUDA kernel 在 GPU 上执行一次需要 2 到 5 微秒。但 Python 光是发出 launch 调用就需要 8 到 10 微秒。GPU 闲置的时间比实际运算的时间还多。每一轮的草稿生成：112ms。GPU 空闲率：65%。

Python 在悄悄谋杀 GPU 利用率，而我浑然不知。

CUDA Graph 是一种捕获一组 GPU 操作然后作为单次驱动调用重放的方式。

通常，每次模型做一次前向传播，Python 都会发出成百上千个独立的 kernel 启动。每个都是一次单独的 CUDA 驱动调用。这种开销累积得非常快，尤其是在循环里反复做的时候。

有了 CUDA Graph，你把 0.5B 草稿模型的整个前向传播捕获下来：24 层 transformer、LM head、下一个 token 的 argmax。只需做一次。之后，重放整个过程只需要一次驱动调用。热路径上完全不需要 Python。

草稿生成：112ms 降到 25ms。4.5 倍提速。

每次我尝试 CUDA Graphs，模型就开始循环输出："the the the the the"。显然哪里出了问题。

CUDA Graphs 在记录时捕获精确的 GPU 内存地址。如果重放过程中有任何张量被重新分配，graph 就会从一个过时的地址读取，输出就变成垃圾。

HuggingFace 默认的 KV cache（DynamicCache）在每个 token 步骤都调用 torch.cat。这意味着每次都分配一个新的缓冲区。Graph 捕获的是旧地址。重放时从那里读。输出：垃圾。

四处改动修好了这个问题：

用 StaticCache 替代 DynamicCache。StaticCache 为 KV cache 预分配固定大小的缓冲区。生成过程中不会再分配。Graph 捕获的地址始终有效。
原地张量变更。不再为中间值创建新张量，所有操作都原地写入。同一块内存，同一个地址，graph 正常工作。
显式更新 position_ids。Graph 需要知道每个 token 处于哪个位置。在动态分配时这是隐式的。有了静态 graph，每次重放前必须手动更新 position_ids。
原地 KV 回滚。当投机验证器拒绝某些草稿 token 时，KV cache 需要回滚到最后一个被接受的位置。这个回滚必须原地进行，不能通过创建新的 cache 对象来完成。

四处修改全部到位之后：不再循环，输出干净。草稿生成 25ms。

在跨 WAN 连接的两块 Kaggle T4 上跑 Qwen2.5-7B：

同样两块 T4 上测试 Qwen2.5-14B，4-bit NF4 量化：WAN 上平均 14.43 TPS。参数量 14.7B。免费的 GPU。

网络不是瓶颈。我花了很多时间以为 WAN 延迟是硬天花板，以为已经没什么可压榨了。真正的瓶颈是 Python kernel 启动开销，而且在我去看 GPU 空闲时间之前它完全看不见。

性能分析比直觉更重要。"网络很慢"是一个很容易做出的假设。"Python 从一个循环里发出 1500 个 kernel，GPU 有 65% 的时间在闲着"需要真正去测量。

CUDA Graphs 不是魔法。它是非常有针对性的。捕获的地址必须保持有效。任何动态分配都会打破它。StaticCache + 原地变更的组合才是让它们在自回归生成中工作的真正原因。

ShardFlow 是开源的，设计为可以在免费的 Kaggle Notebook 上复现。你需要两个 Kaggle 账号和一台 EC2 t3.micro（或者任何有公网 IP 的机器）。

README里有复现精确基准测试的步骤。已经有 583 人克隆了它。我很想知道你在不同硬件上跑出来的数字是否不一样。

v3 等有人赞助我真正的 GPU再说。在那之前，就用免费 T4 和俄亥俄中继凑合着。