作为月之暗面 AI Infra 团队推理侧的一员，眼看着 K2.5 上线不到一周，我们的请求量就已经翻了好几倍。我们正全力应对这甜蜜的烦恼，如果你觉得我们还不够快的话，消消气，我先远程给你磕一个 hhh回归正题，在过去的一年中，整个行业对 RL 的重要性有了更深刻的认识。其中 Rollout 环节已经成为整个流程中最耗时的部分。借着 K2.5 发布，我们想和大家分享一些我们在加速 Sync RL Rollout 过程中积累的经验。去年 8 月我们在 vllm meetup 上公开了我们 CP(Context Parallel) Prefill + DP Decode 的 K2 Rollout 优化方案，这套方案有很多好处，比如可以利用 prefill 节点的 ram cache 来增大 prefix cache 的命中率，节省 prefill 的耗时；还可以通过 PD 分离，来解决 DP 上做 prefill 不均衡的问题。同时这套方案也很适合部署线上大规模服务，比如在 H20 上，以 32k 的输入，15 tokens/s 作为 SLA，DP+EP 相比于 CP 能够节省一倍以上的成本。去年九月我们也开源了我们的 CP 方案集成到了 vllm 中，欢迎大家试用～随着模型架构的不断演进，正如我们在 K2 Thinking 的分享中提到的，我们提出的 INT4 QAT 方案已经可以做到精度无损。目前，该方案已经成为了我们的标准范式。这意味着 K2.5 这样一个 1T 参数的模型现在单机就可以装下，那么在 Rollout 方案的选择上，我们就会产生跟之前训练 K2 FP8 模型不同的答案，本文就来分享一下我们的 CoT。单机 vs 多机一般来说，多机相比于单机最大的问题在于通信慢，虽然 DP+EP 在大 batch 下可以通过 TBO 等技术 overlap 通信，但是由于 RL Rollout 众所周知的长尾问题，我们很长一部分时间都是小 batch，此时 attention/moe 的计算时间短，很难掩盖掉通信。由于 RL Rollout 中并不会使用非常大的并发，且我们的负载长尾占比非常高，我们可以认为单机一定是优于多机。并行方式的选择我们的选择包括 TP/CP/DP，由于 K2 系列模型普遍采用 MLA，使用 TP 会重复存储 kv cache，首先 pass 掉。CP 通过在 context 纬度进行切分，可以将 kv cache 存到不同的 cp_rank 上，相比于 TP 直接多了 dcp_world_size 倍 的 token_budget，显著提升并发以提供更高的吞吐。类似的，DP Attention 也可以达到相同的目的。目前为止，可供我们选择的还剩下四个部署方案，即 Mooncake/Allinone x CP/DP。DP 在做 prefill 时比较难凑出均衡的 batch，导致 prefill 的效率不高。尤其是在 K2.5 这样一个包含 multi-round agent 的 Rollout 场景下，该问题格外严重。我们会通过 PD 分离的方式，在 prefill 节点上使用 CP 来提升 prefill 的效率，故 Allinone DP 方案 pass。现在来比较一下 Allinone CP 和 Mooncake CP。由于 RL Rollout 的负载特性，大量的时间主要在处理长尾请求，此时 prefill 节点在空转，存在浪费行为。而 Mooncake 利用 ram cache 的优势，我们在 Allinone 上可以通过 KVCacheConnector 来抹平，并且我们在实验中发现，CP 由于自身提供了足够多的 token_budget，prefix cache 已经可以维持在一个较高的水准。故 CP 下，Allinone 是优于 Mooncake 的。此时决赛圈仅剩 Allinone CP 和 Mooncake DP 两位选手。DP 原本的一大优势是通过 EP 来省通信，相比于 DP+TP，每个 rank 上需要传输的数据量变少了。然而，得益于 QAT INT4，K2.5 现在可完整部署于单机。在此场景下，DP+TP 由于不需要考虑 expert 负载均衡是优于 DP+EP 的，导致 DP 的通信优势不复存在。结合刚提到的 Mooncake 在 RL Rollout 场景下的若干问题，最后胜出的就是 Allinone CP，在提供足够 token_budget 的同时，也可以结合 MTP 来缓解长尾问题。接下来我们就针对长尾问题，来讲讲我们遇到的问题以及解决方案。CP aware maskSpeculative Decoding 是非常适合加速长尾区的优化方案，我们采用了 MTP 来进行加速。然而，想要将 DCP 跟 MTP 结合起来，需要 MLA kernel 能够正确处理 CP 下的 attention mask。一般情况下，MLA kernel 中使用的是 causal mask，由于 CP 会将 kv cache 存在不同的 cp_rank 上，导致当 query length > 1 时，我们需要的并不是一个标准的 causal mask，如下图所示。MLA kernel 需要确保每个 query 能够正确访问其对应的 kv cache。我们在 inhouse FlashMLA 中，通过额外传入的 cp_rank 以及 cp_world_size 来做到这点Normal:
kv > 0 1 2 3 4 5
query v _____________
2 | 1 1 1
3 | 1 1 1 1
4 | 1 1 1 1 1
5 | 1 1 1 1 1 1
CP_SIZE=2, on CP rank 0:
kv > 0 2 4
query v _______
2 | 1 1
3 | 1 1
4 | 1 1 1
5 | 1 1 1
CP_SIZE=2, On CP rank 1:
kv > 1 3 5
query v ______
2 | 1
3 | 1 1
4 | 1 1
5 | 1 1 1为了使用正确的 attention mask，vllm 也支持了类似的方案，使用 FlashAttn MLA 作为 DCP+MTP 的默认 attention backend，但是我们实测下来，FlashAttn MLA 相比于 FlashMLA 在处理 varlen 请求的性能方面还是有较明显的差距。MTP 支持 varlen在开启投机采样时，由于一个 batch 内每个 request 被接受的 draft token 数量存在差异，在 draft model forward 时输入就是 varlen 的，从而无法调用高性能的 FlashMLA kernel。虽然 vllm 社区也意识到了这一点并提出了相应的优化方案，但该方案仅仅解决了 draft model 可能遇到的变长输入，且无法处理由 short prefill 带来的 varlen。比如在 num_speculative_tokens = 3 时，假如一个 batch 由若干长度为 4 的 decode request 和 一个 长度为 2 的 prefill request 组成，该 short prefill 的存在将再次导致 varlen。所以，在 target model 和 draft model 我们都可能遇到变长输入，并且据我们观察，概率并不低。该问题不仅仅影响 CP+MTP，对所有的并行方式都有影响，因此我们需要新的解决方案。我们采用了一个 naive but effective 的方案，即在 model forward 之前进行 padding，如下图所示。这样我们既可以使用 FlashMLA，也可以开启 full cuda graph。由于 short prefill 的存在，输入长度是 varlen 的
┌─────────────────────────┐
│ Req 0: [S₀][S₁]
│ Req 1: [S₀][S₁][S₂][S₃]
│ Req 2: [S₀][S₁][S₂]
│ Req 3: [S₀]
└─────────────────────────┘
将输入 Pad 成等长，[X] 为 pad token，在 forward 结束之后 unpad 即可
┌─────────────────────────┐
│ Req 0: [S₀][S₁][X ][X ]
│ Req 1: [S₀][S₁][S₂][S₃]
│ Req 2: [S₀][S₁][S₂][X ]
│ Req 3: [S₀][X ][X ][X ]
└─────────────────────────┘结语目前 CP+MTP 已经成为了 K2.5 RL Rollout 的默认方案，CP 通过提供更多的 token_budget，相比于 TP 可获得 2-3 倍的吞吐提升，MTP 进一步缓解长尾问题，将 RL Rollout e2e 提速 40%+。我们将 CP+MTP 这样一个在生产环境中经过充分验证的方案分享了出来，希望大家能够有所收获。可以肯定的是，RL Rollout 系统还有很多待挖掘的宝藏，我们正在积极探索下一代 RL Rollout 系统，比如将 Seer 中的 N-gram 跟 MTP 结合，朝着 Adaptive Speculatvie Decoding 的方向发展。Stay tuned！