vLLM 首日支持 Thinking Machines 的 Inkling 2 Small:594B 总参数、17B 激活,单用户解码约 630 tok/s
核心概要
Thinking Machines 的 Inkling 2 Small 在发布当天获得 vLLM 支持。这是一个文本与图像输入、文本输出的混合专家模型,总参数 594B、激活 17B,上下文最长 1M token,并自带 DSpark 推测解码草稿器。vLLM 通过内核融合、程序化依赖启动、序列并行和针对 6:1 GQA 比例重写的 Gluon 注意力内核完成适配,在 GB200 上 TP4 单用户解码约 630 tok/s、接受长度 3.5,并支持 LoRA 与经由 SkyRL 的强化学习。
深度剖析
vLLM 在 Inkling 2 Small 发布当天就能服务这个模型,包括它自带的 DSpark 推测解码草稿器。 Inkling 2 Small 相对前代 Inkling 改动了骨干:层数从 66 层降为 48 层,滑动窗口从 512 扩大到 1,024,MoE 从 256 个路由专家(top-6)加 2 个共享专家改为 384 个路由专家(top-8)且全部层都是 MoE,图像编码器从 4 层 hMLP 换成 23 层 ViT。 在 GB200 上以 TP4 服务时,单用户解码约 630 tok/s,接受长度 3.5;在 TP8 的智能体负载基准中,32K 会话上下文下每 GPU 约 1,070 tok/s,对应约 84 个并发会话、前缀缓存命中率 91%。
适配的关键是把解码和预填充中重复、零散的内存受限步骤合并,并让相邻内核的启动相互重叠。 块注意力残差让每个子层边界都要做一次 all-reduce、一次短卷积、一次残差更新、一次注意力残差混合和一次 RMSNorm,这五步都读写同一份隐藏状态。 融合内核 allreduce_conv_mix 在 NVLink 多播上完成 all-reduce 并把后续步骤留在寄存器中;GEMM、注意力、路由和激活内核使用程序化依赖启动(PDL),让每个内核的前导段与上一个内核的尾部重叠,这在解码这种由大量短内核串成的长链、小批量场景下尤其重要。
预填充阶段用序列并行去掉张量并行下每张卡重复计算的那部分工作。 在普通张量并行下,每次 all-reduce 之后每张 GPU 都要对全部 token 跑短卷积、注意力残差混合和 RMSNorm,解码时开销不大,预填充时则代价明显。 每张 GPU 先 reduce-scatter,只在自己那片 token 上用一个融合的 conv_mix 内核完成短卷积、残差更新、注意力残差混合和 RMSNorm,再 all-gather 给下一个 GEMM;8K token 的预填充步骤在 TP4 下缩短 22%,TP8 下缩短 26%,每 GPU 释放 1.3–1.9 GiB 给 KV 缓存。
解码路径另写了一个面向相对注意力的 Gluon 内核,因为原有 FlashAttention-4 解码路径在这个模型上会变慢。 Inkling 2 Small 的 6:1 GQA 比例会让 GQA packing 被静默关闭,使 Inkling 的 FlashAttention-4 解码路径效率下降。 新的 warp 专用 Gluon 内核把 6 个 head 打包进一个 64 行 tile、只读一次 KV:TMA warp 流式读取 KV,tcgen05 warp 执行矩阵乘加,四个 softmax warp 施加相对偏置;100K 上下文下 batch-1 解码一步从 15.1 ms 降到 4.8 ms,预填充仍沿用 Inkling 的 sheared-bias FlashAttention-4 路径。
启示与展望
演示范围是 vLLM 在 GB200 上以 TP4 和 TP8 服务 Inkling 2 Small,性能数字来自带 DSpark 概率式草稿和前缀缓存的智能体负载,每轮新增约 3K token、生成 1,024 token、温度 1.0 并开启推理。精度核对覆盖最长 1M token 的 NIAH、MMMU-Pro 和 HLE,且是在开启 DSpark 的情况下与不开推测解码的运行对比。
精度核对只说明开启 DSpark 后的分数与不开推测解码的运行一致,并未给出与参考实现逐项对齐的绝对分数。序列并行的 22% 和 26% 预填充加速分别对应 TP4 和 TP8 的 8K token 预填充步骤,换到其他并行度或序列长度时文中没有给出对应数字。
