vLLM speeds up Gemma 4 E2B 2.85x on A100, Google credits Gemini 4 Argon
The author breaks down the vLLM PRs that raised Gemma 4 E2B throughput on A100 from about 37,900 to over 108,000 tok/s, covering KV-sharing prefill shortcuts, kernel fusion, attention tuning and speculative decoding draft layers.
vLLM 跑 Gemma 4 E2B，A100 上的吞吐从约 37,900 tok/s 涨到 108,000 以上，2.85 倍，生成的 token 和原来完全一致。 Google 说这是用 Gemini 4 Argon 几天做出来的，发布页还附了 PR 列表。我把这组 PR 从头翻了一遍，改动都很具体，按顺序拆一下。 一、先看 E2B 长什么样 按 Hugging Face 上的 config：35 层，其中 28 层是滑动窗口注意力（窗口 512），7 层是全局注意力。 后 20 层不自己存 KV，直接用前面层算好的 KV。隐藏维度只有 1536，词表却有 262,144。 模型这么小，每一层的矩阵乘本身很快。时间大量花在 kernel 启动、每一步的元数据准备、来回读写显存这些固定开销上。 下面几个 PR，基本都在砍这部分。 二、最大的一刀：后 20 层在 prefill 时少算 一段 4K 的提示词进来，前面的层要给每个 token 算 KV，后面生成时要用。 后 20 层不写 KV，提示词里绝大多数 token 在这 20 层算出来的结果都用不上，只有要出下一个字的那个位置需要。所以 prefill 时，这 20 层只算那一个位置就够了。 vLLM 里本来就有这个开关（--kv-sharing-fast-prefill，Gemma 3n 时加的），但要手动打开。PR 60445 做了三件事： 1. 检测到 KV 共享层就默认打开，开了投机解码或流水线并行时不开 2. 把这条路径接进 CUDA Graph 3. 删掉每一步重复的元数据计算。比如每个请求只出 1 个 token 时，起始位置直接用 arange 生成，省掉 5 次 kernel 启动和 1 次显存到内存的拷贝 PR 里的数字：E2B 偏 prefill 的负载快 1.34 到 2.06 倍，4K 长提示快 1.95 到 2.10 倍。 12B、26B、31B 没有 KV 共享层，不受影响。 三、把小 kernel 合成一个 每个注意力块里，Q 和 K 要先做 RMSNorm 再做 RoPE，原来是几次分开的 kernel 启动。 PR 60674 用 Triton 把它们合成一个，FFN 前后的 norm、PLE 投影、MoE 路由也一样处理。 单看这一段，滑动窗口层的 QKV norm 加 RoPE 快了 1.88 到 3.85 倍。 端到端看，E2B 在 BS=4、1K 提示时，prefill 从 35,480.8 到 67,488.9 tok/s，1.90 倍。decode 只快了 1.06 倍。 到 12B 以上，只剩 4% 到 5%。 模型越小，固定开销占的比例越大，合并 kernel 的收益也就越大。 四、注意力 kernel 重新调参 PR 60317 给 head_dim 512 的全局层和窗口 512、1024 的滑动层重新选 tile 大小。端到端 E2B 快 1.06 到 1.12 倍。 这个 PR 被 reviewer 抓到了问题。 有人在 RTX 5090 上复测，滑动窗口层单序列 decode 的延迟变成原来的 2.3 到 3.3 倍。原因是 PR 在小 batch 时关掉了原来的 3D kernel 路径。 作者后来改成小 batch 保留原路径，等 GPU 吃满了再切换。 五、投机解码的草稿层 Gemma 4 的 MTP 草稿层要给候选词打分。原来先把选中的词表行拷出来，再做矩阵乘，E2B 这一步中间要分配 1,541 MB。 PR 60689 改成按索引直接读权重做 GEMV，中间分配降到 5 MB。 5 步草稿生成快 1.37 到 3.23 倍。 六、还有两处被 reviewer 改掉 1. 视觉塔的 RMSNorm，原来新写了一个 Triton 版。reviewer 测出 vLLM 自带的 CUDA 版在小 batch 上更快，PR 换成了自带的 2. 12B 的 2D 位置编码，原来也是新写的 Triton kernel。review 之后改成普通的向量化 PyTorch，结果逐位一致，还少一个要维护的 kernel 另外，26B-A4B 的 MoE 调参文件一共 269 个。PR 自己写明只有 A100 是在真机上调出来的，其余显卡是按模板推出来的。 七、我怎么看 1. 这 8 个性能 PR 还在等合并，2.85 倍是 Google 在 A100 上测的口径 2. 收益集中在 E2B、E4B 这种小模型，大模型只有几个百分点 3. 写 kernel、跑基准、写很长的 PR 说明，AI 已经能做了。换一张卡会不会变慢，目前还是靠手里有那张卡的人去测 我的 RTX PRO 6000 和 5090 一样是 sm_120。 这组 PR 合进去以后，我更想看消费级 Blackwell 上的复测数字。