Fast Gemma挑战赛:我们验证过的SOTA配方全解析

0 阅读

在Fast Gemma挑战赛中,我们VIDRAFT团队(参赛代号vidraft-darwin)以510.58 TPS的成绩,在通过验证的提交中取得了最高推理速度,同时保持了2.3930的困惑度(PPL),远低于2.42的质量门槛。这一成绩并非偶然,而是基于对推理栈的深度优化和社区协作的成果。本文将完整公开我们的配置清单,并深入解析每个优化项背后的原理,希望能为在受限硬件上追求极致推理性能的开发者提供可复现的参考。

挑战的本质:在质量约束下榨干硬件潜力

Fast Gemma挑战赛的核心规则很简单:在完全相同的硬件(单张NVIDIA A10G)上,仅通过软件优化提升Google的gemma-4-E4B-it模型的推理速度(TPS)。参赛者不能更换模型或禁用任何功能,最关键的是,不能牺牲输出质量。如果PPL超过约2.42的阈值,该次运行即视为失败。最终排名仅依据组织者在私有提示集上重新运行并标记为“已验证”的结果。

这意味着,任何优化都必须在不损害质量的前提下进行。我们的策略是:只采用质量中立的加速技术,任何导致PPL上升的优化,无论速度提升多大,都坚决放弃。

我们的成绩与定位

我们的最终提交配置为vidraft-fw188-ctk49-n64-patchbridge-v1,在单流A10G上实现了510.58 TPS,PPL为2.3930,并成功通过了重新验证。坦白说,在原始TPS上,有更快的运行(例如535.91 TPS),但那些方案的PPL超过了2.44,未能通过质量验证。我们引以为傲的不是“最快”,而是“在验证结果中最快,且未牺牲质量(PPL 2.39)”。

完整配置:manifest.json 全公开

以下是我们提交的完整配置,该文件可复现整个优化栈。

{
  "name": "vidraft-fw188-ctk49-n64-patchbridge-v1",
  "description": "VIDRAFT W188 CTK49 N64 patch-bridge reproduction: public patch-style warmup bridge, sliding_window=188, CENTROID_TOP_K=49.",
  "dependencies": [
    "https://wheels.vllm.ai/.../vllm-0.22.1rc1.dev307+g3e8afdf78.cu129-...whl",
    "transformers==5.9.0", "jinja2==3.1.6", "MarkupSafe==3.0.3",
    "orjson==3.10.18", "safetensors", "torch"
  ],
  "model_id": "google/gemma-4-E4B-it",
  "served_model_name": "gemma-4-e4b-it",
  "port": 8000,
  "serve": ["python", "serve.py"],
  "env": {
    "WEIGHTS_BUCKET": "hf://buckets/gemma-challenge/gemma-chiku-inu/weights/osoi5-v0-baked",
    "MAX_MODEL_LEN": "4096",
    "GPU_MEMORY_UTILIZATION": "0.90",
    "MAX_NUM_BATCHED_TOKENS": "512",
    "MAX_NUM_SEQS": "1",
    "PERFORMANCE_MODE": "interactivity",
    "SLIDING_WINDOW": "188",
    "HF_OVERRIDES": "{\"text_config\": {\"sliding_window\": 188}}",
    "FA_SLIDING": "1",
    "CENTROID_TOP_K": "49",
    "SPECULATIVE_CONFIG": "{\"method\":\"mtp\",\"model\":\"/tmp/qat-assistant\",\"num_speculative_tokens\":7}",
    "DRAFTER_BUCKET": "hf://buckets/gemma-challenge/gemma-kenyan-duma/weights/drafter-ft/ft-v1-epoch_001",
    "LM_HEAD_PRUNE": "1",
    "LM_HEAD_KEEPSET_BUCKET": "hf://buckets/gemma-challenge/gemma-dixie-flatline/weights/int4-pck04c-12k",
    "PCK04_KEEPSET": "/tmp/osoi5-v0-baked/pck04_keepset.json",
    "WARMUP_BRIDGE": "1",
    "WARMUP_NUM_PROMPTS": "64",
    "WARMUP_MAX_TOKENS": "1",
    "WARMUP_SEED": "42",
    "PRECACHE_BENCH": "0",
    "ONEGRAPH": "1",
    "LOOPGRAPH_REQUIRE_CAPTURE": "1",
    "LOOPGRAPH_WARMUP_CALLS": "20",
    "LOOPGRAPH_PINGPONG_SLOTS": "3",
    "FUSED_SPARSE_ARGMAX": "1",
    "FUSED_SPARSE_ARGMAX_BLOCK": "64",
    "SPLITKV_VERIFY": "1",
    "SPLITKV_VERIFY_MAX_Q": "64",
    "DETOK_ENDONLY": "1",
    "FASTRENDER": "1",
    "OVERRIDE_GENERATION_CONFIG": "{\"temperature\":0.0,\"top_p\":1.0,\"top_k\":0}",
    "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512,expandable_segments:True",
    "LD_PRELOAD": "/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4"
  }
}

运行方法很简单:设置上述环境变量,然后执行python serve.py(端口8000)。所有文件均可从下方存储桶下载。

发布文件hf://buckets/gemma-challenge/gemma-vidraft-darwin/submissions/vidraft-darwin/break-fw188-ctk49-n64-patchbridge-v1/):

文件 作用
serve.py 主服务入口
serve_patch_warmup_bridge.py 合成预热桥(N64)
fa_sliding_patch.py FlashAttention滑动窗口补丁
serve_patch_precache.py 预缓存路径(本配置中关闭)
splitkv_verify_patch.py split-KV验证内核
serve_patch_pck04.py · detok_endonly.py · lsk_patch.py · steptime_patch.py 分词/解码/计时优化
manifest.json 上述完整配置

浏览:huggingface.co/buckets/gemma-challenge/gemma-vidraft-darwin/tree/submissions/vidraft-darwin/break-fw188-ctk49-n64-patchbridge-v1

每个优化项的作用与原理

滑动窗口(SLIDING_WINDOW=188)

解码阶段的瓶颈是KV缓存的内存带宽。通过将注意力限制在最近的188个token,我们显著减少了KV缓存的大小,从而提升了内存带宽利用率。窗口太窄(如128)会破坏PPL,太宽则会拖慢速度。经过实验,188-192是性能与质量的最佳平衡点。HF_OVERRIDES用于修改模型配置。

质心Top-K(CENTROID_TOP_K=49)

这是一个内核级别的调优参数。吞吐量和PPL会随该值变化,因此我们的目标是找到在PPL预算内能达到最高吞吐量的值。我们扫描了44、48、49等值,最终确定49为最优。

合成预热桥(WARMUP_BRIDGE=1, WARMUP_NUM_PROMPTS=64)

在计时运行开始前,我们发送64个合成提示(每个仅1个token),以彻底完成CUDA图的捕获和JIT编译。这样,在测量窗口内就不会残留任何编译或捕获开销,从而缩小了公开与私有(已验证)TPS之间的差距,并稳定了高速记录。没有这个预热桥,我们大约会损失15 TPS。

关闭预缓存(PRECACHE_BENCH=0)

预缓存路径可能会虚增自测TPS,但该数字在私有重新运行时无法复现,并会被标记为无效。关闭后,自测TPS与验证TPS基本一致,确保我们报告的数字就是最终验证的数字。

推测解码(SPECULATIVE_CONFIG)

我们使用多token预测(MTP)草稿模型,设置推测token数为7,以提高每步的吞吐量。

其他内核优化

ONEGRAPHFUSED_SPARSE_ARGMAXSPLITKV_VERIFYDETOK_ENDONLY等选项分别用于消除内核启动、采样和解码的开销。

所有这些优化的核心原则是:只堆叠质量中立的加速技术。任何导致PPL上升的优化,无论速度提升多大,我们都坚决放弃。

社区协作:成绩属于大家

细心的读者可能已经注意到,我们的配置中引用了多个社区共享的资产:

  • WEIGHTS_BUCKET:来自@chiku-inu的INT4量化权重(osoi5-v0-baked
  • DRAFTER_BUCKET:来自@kenyan-duma的推测解码草稿模型(drafter-ft
  • LM_HEAD_KEEPSET_BUCKET:来自@dixie-flatline的12k lm-head保留集
  • 预热桥的基础来自@firfir-cast,前沿配置由@gemma-slayer共享、复现和验证

正是因为大家即使失败也分享经验,社区的整体记录在六天内显著提升。我们的510.58 TPS只是这个共享基础之上的一块砖。

结语

我们的大部分工作集中在如何在受限硬件上高效服务模型,希望这份配置和文件能为进行类似实验的开发者提供一个起点。我们非常期待看到复现结果、改进方案,尤其是将其移植到其他GPU上的尝试。再次感谢Google Gemma团队、Hugging Face以及所有参与者。

—— vidraft-darwin (VIDRAFT)

排行榜:huggingface.co/spaces/gemma-challenge/gemma-dashboard