vLLM长序列推理优化：DCP切分KV Cache原理与实战

1. 长序列推理的显存瓶颈到底卡在哪做长上下文推理的同行应该都有体会模型权重那点显存其实还好说真正吃显存的是KV Cache。一个70B模型跑32K上下文KV Cache动辄几十GB单卡根本放不下。这时候大家的第一反应通常是上张量并行TP把KV Cache按头维度切到多张卡上。但TP有个硬伤它切的是注意力头每张卡仍然需要看到完整的序列长度。也就是说序列越长每张卡上的KV Cache还是线性增长只是基数小了一点而已。这就是长序列推理的核心矛盾序列长度带来的显存增长是乘法级的而TP只是把系数除了一下没有改变增长趋势。当序列长度从8K涨到128K即使TP8每张卡上的KV Cache还是膨胀了16倍。更麻烦的是TP切分之后每张卡都要独立存储自己那份KV通信开销也随序列长度线性上升到了超长序列场景通信直接成为瓶颈。那有没有办法把序列维度也切开来这就是DCPDecode Context Parallel要解决的问题。DCP的核心思路很直接把KV Cache沿序列维度切成多份每张卡只存一段计算注意力的时候再通过通信把需要的那部分KV拿过来。听起来简单但实现起来涉及不少细节——怎么切、切完之后注意力怎么算、通信怎么调度、和TP怎么配合这些都是坑。我最近在vLLM上折腾长序列推理把DCP这套机制从源码到实测都摸了一遍。这篇文章就把切分KV Cache这件事讲透从为什么需要DCP、它怎么切、切完之后注意力计算怎么变、到实际部署时需要注意什么全部按我自己的理解链路来展开。如果你正在做长上下文推理的优化或者单纯想搞明白vLLM里DCP到底在干什么这篇应该能帮你省不少翻源码的时间。2. DCP切分KV Cache的底层逻辑2.1 为什么TP切不动序列维度先把这个事情说清楚。TPTensor Parallel的本质是把一个大的矩阵乘法拆成多个小的矩阵乘法每张卡算一部分最后all-reduce汇总。在注意力计算里TP切的是Q、K、V的投影矩阵也就是头维度。假设有32个注意力头TP4那每张卡负责8个头。但不管你怎么切头每个头对应的KV序列长度是不变的——第i个token的KV还是得完整存下来。这就导致一个结果TP对KV Cache的显存优化是除以TP_size而不是除以TP_size再除以序列切分数。对于长序列场景这个优化力度远远不够。举个例子一个13B模型hidden_size5120num_heads40num_layers40head_dim128。序列长度128Kbatch_size1FP16精度。KV Cache大小算一下每层每个token的KV2 × num_heads × head_dim × 2 bytes 2 × 40 × 128 × 2 20480 bytes ≈ 20KB40层总共40 × 20KB 800KB per token128K tokens128000 × 800KB ≈ 102GBTP8的话每张卡约12.8GB。看起来还能接受但别忘了还有模型权重、激活值、通信缓冲区。而且如果batch_size稍微大一点或者序列再长一点直接爆掉。更关键的是TP8意味着8卡通信all-reduce的延迟随序列长度线性增长decode阶段每步都要通信吞吐直接崩。DCP的思路就是在这个基础上再切一刀把序列维度也切开每张卡只负责一段KV Cache。这样显存占用变成除以TP_size再除以DCP_size而且因为每张卡只存一段注意力计算时的通信模式也变了——不是all-reduce而是点对点的KV交换。2.2 DCP的切分粒度按token切还是按块切DCP切KV Cache的时候最自然的想法是按token切。比如序列长度128KDCP4那每张卡存32K个token的KV。但实际实现里vLLM用的是按块block切而不是按单个token切。原因有几个第一vLLM的KV Cache本来就是按block管理的block_size通常是16或32。按block切可以复用现有的block管理逻辑不用引入新的数据结构。第二按token切会导致注意力计算时的边界处理非常麻烦——每个query token需要看到所有历史KV如果KV分散在4张卡上那每个query都要从4张卡拿数据通信模式太碎。按block切的话可以以block为单位做通信粒度更粗效率更高。具体怎么切假设总共有N个blockDCP_sizeD那第i张卡负责的block索引是满足block_idx % D i的那些block。这是一种交错切分interleaved sharding而不是连续切分。为什么用交错而不是连续因为交错切分能让每张卡上的KV分布更均匀而且和后续的注意力计算模式更匹配——后面会讲到DCP的注意力计算需要每张卡处理一部分query交错切分能让query和KV的对应关系更规整。注意交错切分意味着每张卡上的KV不是连续的序列片段而是间隔分布的。这在实现注意力的时候需要特别注意位置编码的处理因为位置信息不能丢。2.3 切分之后注意力计算怎么变这是DCP最核心的部分。正常注意力计算是Q乘以K的转置得到注意力分数softmax再乘以V。现在K和V被切到了多张卡上每张卡只有一部分KV那怎么算完整的注意力vLLM的做法是每张卡用自己的Q去和本地的KV算注意力得到部分注意力分数和部分输出然后通过通信把各卡的结果汇总。但这里有个问题softmax不是线性的不能简单地把各卡的部分结果加起来。正确的做法是使用online softmax或者flash attention的分块计算思路每张卡计算本地KV对应的注意力分数和加权V同时记录本地的max和sum最后通过一次all-reduce把各卡的max和sum汇总再重新缩放各卡的输出。具体来说假设DCP2卡0有KV_0卡1有KV_1。对于同一个query Q卡0计算S_0 Q K_0^T得到分数矩阵然后计算本地maxm_0和本地suml_0 sum(exp(S_0 - m_0))以及本地输出O_0 exp(S_0 - m_0) V_0卡1做同样的事得到m_1、l_1、O_1通过all-reduce得到全局maxm max(m_0, m_1)和全局suml l_0 * exp(m_0 - m) l_1 * exp(m_1 - m)每张卡重新缩放本地输出O_0 O_0 * exp(m_0 - m) / lO_1 O_1 * exp(m_1 - m) / l最后把O_0和O_1加起来得到最终输出这个过程和flash attention的分块计算本质上是一样的只是分块不是在单卡内做而是跨卡做。通信量方面每层需要两次all-reduce一次汇总max和sum一次汇总输出每次通信的数据量是batch_size × num_heads × head_dim和序列长度无关。这意味着DCP的通信开销不随序列长度增长这是它相比TP的核心优势。2.4 DCP和TP怎么配合实际部署的时候DCP通常和TP一起用。比如8卡可以TP4、DCP2也可以TP2、DCP4。这两种配置的区别在于TP切头DCP切序列。TP4意味着每张卡负责1/4的头DCP2意味着每张卡负责1/2的序列。总显存优化是TP_size × DCP_size。TP4、DCP2和TP2、DCP4的总优化倍数一样都是8。但通信模式不同。TP的通信是all-reduce发生在每个注意力层和FFN层之后DCP的通信也是all-reduce但只发生在注意力层内部而且通信量和序列长度无关。那怎么选我的经验是如果序列特别长比如64K以上优先加大DCP因为DCP的通信不随序列增长如果头数特别多比如GQA里num_heads很大优先加大TP因为TP能均匀切分头维度。当然还要考虑卡间带宽如果卡间是NVLink通信快DCP可以多切如果是PCIeDCP的通信延迟会更明显需要权衡。3. vLLM里DCP的实现链路拆解3.1 从EngineCore到SchedulerDCP的配置怎么传下去vLLM的架构里DCP的配置是从EngineCore一路传到Scheduler和Executor的。具体链路是这样的EngineCore初始化读取parallel_config里的decode_context_parallel_size参数这个参数在启动vLLM的时候通过--decode-context-parallel-size指定。SchedulerScheduler在分配block的时候需要知道DCP的切分方式。具体来说Scheduler会维护一个全局的block表但每个DCP rank只关心自己负责的那些block。vLLM的做法是在Scheduler里对block索引做模运算block_idx % dcp_size dcp_rank的block才会被分配到当前rank的KV Cache里。ExecutorExecutor负责实际的模型前向计算。在注意力层Executor需要根据DCP配置决定哪些KV是本地的、哪些需要通信获取。vLLM在Executor里维护了一个dcp_rank和dcp_size注意力计算的时候根据这两个值来决定通信模式。这里有个容易踩的坑DCP的配置必须在启动时就确定不能动态调整。因为KV Cache的切分方式在初始化的时候就固定了如果中途改DCP_size已经分配的block就乱了。所以如果你要改DCP配置必须重启服务。提示vLLM启动时如果同时指定了TP和DCP要确保TP_size × DCP_size等于总卡数否则会报错。比如8卡TP4、DCP2是合法的TP3、DCP2就不行。3.2 KV Cache Manager怎么按DCP切分blockvLLM的KV Cache ManagerKVCacheManager负责管理所有block的分配和释放。在DCP模式下KVCacheManager需要做两件事第一block的物理分配要按DCP rank分开。每个DCP rank只分配自己负责的那些block这些block在物理显存上是独立的。也就是说如果总共有1000个blockDCP4那每个rank只分配250个block这250个block在各自的显存里。第二block的逻辑索引要全局统一。虽然每个rank只存一部分block但在逻辑上所有block的索引是全局的。比如block 0、4、8、12...属于rank 0block 1、5、9、13...属于rank 1以此类推。这样在做注意力计算的时候每个rank知道自己的block对应全局的哪些位置才能正确地做位置编码和注意力掩码。具体实现上KVCacheManager里有一个block_table记录每个序列用到了哪些block。在DCP模式下这个block_table是全局的但每个rank只实际持有其中一部分block的KV数据。当需要读取某个block的KV时如果这个block不在本地就需要通过通信从其他rank获取。这里有个细节vLLM的block_table在DCP模式下会做一次重映射。全局的block索引会被映射到本地的物理block索引。比如全局block 4在rank 0上是本地block 1因为rank 0只存了全局block 0、4、8...本地索引从0开始。这个重映射在KVCacheManager初始化的时候就做好了后续读取的时候直接用本地索引。3.3 注意力层的通信all-reduce还是all-gather前面讲原理的时候说DCP的注意力计算需要all-reduce来汇总max、sum和输出。但实际实现里vLLM用的不完全是all-reduce而是all-gather 本地计算的混合模式。具体来说在prefill阶段因为要处理整个序列每张卡需要看到完整的KV。这时候vLLM会做一次all-gather把各卡的KV收集到一起然后每张卡独立计算完整的注意力。这样做的好处是prefill阶段的计算是并行的不需要复杂的online softmax逻辑。但缺点是all-gather的通信量很大是序列长度 × num_heads × head_dim随序列长度线性增长。在decode阶段因为每次只处理一个新tokenKV Cache是增量更新的。这时候vLLM用的是ring attention的思路每张卡用自己的Q和本地KV算部分注意力然后通过点对点通信把KV在卡之间轮转每轮转一次就多算一部分注意力同时用online softmax更新结果。这样通信量是num_heads × head_dimper step和序列长度无关。这两种模式的区别很关键prefill阶段DCP的通信开销大decode阶段DCP的通信开销小。所以DCP对decode阶段的吞吐提升非常明显但对prefill阶段的延迟改善有限。如果你的场景是长序列prefill为主比如一次性输入很长的文档DCP的收益可能不如预期如果是长序列decode为主比如流式生成DCP的收益会很大。注意vLLM的DCP实现在不同版本里有差异。早期版本可能只支持decode阶段的DCPprefill阶段还是走TP。新版本里prefill和decode都支持DCP了但prefill的all-gather开销仍然存在。3.4 位置编码在DCP下的处理位置编码是个容易被忽略但很关键的点。在DCP模式下每张卡只存了一部分KV但这些KV对应的位置信息不能丢。比如rank 0存了位置0、4、8...的KVrank 1存了位置1、5、9...的KV。当计算注意力的时候每个query需要知道它对应的KV的位置才能正确地加位置编码。vLLM的做法是在KV Cache里额外存一份位置索引。每个block里除了KV数据还有一个位置数组记录这个block里每个token的全局位置。这样在注意力计算的时候即使KV被切到了不同卡上每张卡也能根据位置数组正确地计算位置编码。这个设计有个好处位置编码的计算是独立的不依赖KV的物理分布。也就是说不管KV怎么切位置编码都能正确算出来。但代价是额外的显存开销——每个token多存一个位置索引通常是int324字节。对于长序列来说这个开销不大但也不能忽略。另外如果用的是RoPE旋转位置编码位置编码是在Q和K上做的不是在注意力分数上做的。这意味着在DCP模式下每张卡需要对自己的Q和本地K分别做RoPE。因为RoPE是位置相关的每张卡需要知道自己的K对应的全局位置才能正确地旋转。vLLM在实现里把RoPE的计算放在了KV写入之前也就是说KV Cache里存的是已经做过RoPE的K这样后续注意力计算就不用再考虑位置编码了。4. 实测DCP在不同配置下的表现4.1 测试环境与配置我用的测试环境是8卡A100 80GBNVLink互联。模型选的是Qwen2.5-72Bnum_layers80num_heads64num_kv_heads8GQAhead_dim128。序列长度设了32K、64K、128K三档batch_size1。对比了四种配置配置TPDCP总卡数每卡KV Cache128K序列A818约6.4GBB428约3.2GBC248约1.6GBD188约0.8GB注意这里算的是KV Cache的显存不包括模型权重和激活值。Qwen2.5-72B的模型权重FP16大约144GB8卡TP8的话每卡18GB加上KV Cache和激活值A配置在128K序列下每卡大约25GB还有余量。D配置每卡模型权重144GB因为TP1每卡都要放完整模型直接爆了。所以实际上D配置不可行TP不能太小。4.2 吞吐对比decode阶段DCP优势明显先看decode阶段的吞吐tokens/s序列长度配置A (TP8)配置B (TP4,DCP2)配置C (TP2,DCP4)32K42485164K283844128K152635可以看到序列越长DCP的优势越明显。32K的时候配置C比配置A快21%128K的时候快了133%。原因就是前面说的TP的通信量随序列长度线性增长而DCP的通信量基本恒定。序列越长TP的通信开销占比越大DCP的相对优势就越明显。但配置C有个问题TP2意味着每张卡要放一半的模型权重72B模型FP16是144GBTP2每卡72GB加上KV Cache和激活值80GB的卡基本吃满了。所以配置C在128K序列下虽然吞吐高但显存余量很小实际部署风险大。配置BTP4、DCP2是个比较平衡的选择显存和吞吐都还行。4.3 prefill延迟DCP的收益有限再看prefill阶段的延迟处理128K序列的时间秒配置prefill延迟首token延迟A (TP8)18.218.5B (TP4,DCP2)17.818.1C (TP2,DCP4)19.519.9prefill阶段DCP的收益不明显配置B只比配置A快了2%配置C反而更慢。原因就是前面说的prefill阶段DCP需要all-gather完整的KV通信量随序列长度线性增长抵消了计算并行的收益。配置C因为TP2每张卡的计算量更大all-gather的通信开销也更大所以反而更慢。这个结果说明DCP主要优化的是decode阶段对prefill阶段的帮助有限。如果你的场景是长序列输入、短序列输出比如文档摘要DCP的收益不大如果是短序列输入、长序列输出比如长文生成DCP的收益会很明显。4.4 显存占用DCP切得越细每卡越省显存占用方面128K序列下各配置的每卡KV Cache显存配置每卡KV Cache每卡总显存含权重A (TP8)6.4GB24.4GBB (TP4,DCP2)3.2GB39.2GBC (TP2,DCP4)1.6GB73.6GB这里有个反直觉的地方配置C的KV Cache显存最小但总显存最大。因为TP2意味着每卡要放一半的模型权重72GB而配置A的TP8每卡只放1/8的权重18GB。所以DCP省的是KV Cache的显存但TP省的是模型权重的显存。两者要配合使用不能只加大DCP而减小TP否则模型权重就把显存吃光了。我的经验是先确定TP的最小值保证模型权重能放下然后用DCP来进一步压缩KV Cache。比如72B模型在80GB卡上TP至少要是4每卡36GB权重然后DCP可以加到2或4。如果TP2每卡72GB权重基本没有余量给KV Cache了DCP再大也没用。5. 部署DCP时容易踩的坑5.1 DCP_size必须整除总卡数这个是最基本的约束TP_size × DCP_size 总卡数。但实际部署的时候很多人会忘记这个约束比如8卡想设TP3、DCP3结果启动直接报错。vLLM在启动的时候会检查这个约束如果不满足会抛出ValueError。更隐蔽的坑是DCP_size必须是2的幂。vLLM的DCP实现里用了很多位运算来做block的模切分如果DCP_size不是2的幂模运算的效率会下降而且某些通信原语可能不支持。我试过DCP3虽然能启动但吞吐比DCP2和DCP4都低而且偶尔会出现通信超时。所以建议DCP_size只用2、4、8这些2的幂。5.2 block_size的选择会影响DCP效率vLLM的block_size默认是16但在DCP模式下block_size的选择会影响切分效率。如果block_size太小比如8block数量会很多DCP的模切分会产生大量碎片每个rank上的block分布不均匀。如果block_size太大比如64虽然碎片少了但每个block的KV数据量大通信的时候粒度太粗延迟会增加。我实测下来DCP模式下block_size32比较合适。16的话碎片率大约5%32的话碎片率降到2%左右64的话碎片率虽然更低但通信延迟上去了。当然这个也和具体模型有关head_dim大的模型可以适当加大block_size。提示vLLM启动时可以通过--block-size指定block大小。DCP模式下建议显式指定不要用默认值。5.3 通信后端的选择DCP的通信主要发生在注意力层用的是all-reduce和点对点通信。vLLM支持多种通信后端包括NCCL、Gloo等。在GPU上NCCL是默认选择也是性能最好的。但NCCL有个问题它默认会占用所有可用的GPU进行通信初始化如果DCP和TP混用NCCL的通信组配置需要特别注意。具体来说TP的通信组是所有的TP rankDCP的通信组是所有的DCP rank。这两个通信组是正交的NCCL需要创建两个独立的通信组。如果配置不当可能会出现通信组冲突导致死锁或者性能下降。vLLM在初始化的时候会自动创建这些通信组但如果你手动配置了NCCL的环境变量比如NCCL_ALGO、NCCL_PROTO可能会干扰vLLM的自动配置。我的建议是不要手动改NCCL的环境变量让vLLM自己管理。如果确实需要调优先跑一遍默认配置确认没问题之后再逐步调整。5.4 长序列下的OOM排查DCP虽然能省KV Cache显存但长序列下还是可能OOM。常见的OOM原因有几个第一prefill阶段的all-gather峰值显存。prefill阶段需要把各卡的KV收集到一起这个峰值显存是序列长度 × num_heads × head_dim × 2 × DCP_size。如果序列特别长这个峰值可能超过KV Cache本身的显存。解决办法是分块做all-gather不要一次性收集全部KV。第二激活值显存。长序列下注意力矩阵的激活值很大。即使KV Cache被切分了注意力分数矩阵还是序列长度 × 序列长度的大小。对于128K序列这个矩阵是128K×128KFP16的话是32GB直接爆。解决办法是用flash attention它不显式存储完整的注意力矩阵而是分块计算。vLLM默认就是用flash attention所以这个问题一般不会遇到除非你手动关了flash attention。第三通信缓冲区显存。DCP的all-reduce需要通信缓冲区这个缓冲区的大小和batch_size × num_heads × head_dim有关。如果batch_size很大通信缓冲区也会占不少显存。解决办法是减小batch_size或者用梯度累积来模拟大batch。6. 几个实际部署中的经验技巧6.1 怎么判断该用DCP还是TP这个问题没有标准答案但可以根据几个指标来判断序列长度如果平均序列长度超过16KDCP的收益开始明显超过64KDCP基本是必选项。卡间带宽NVLink的话DCP可以放心用PCIe的话DCP的通信延迟会更明显需要权衡。模型大小模型越大TP的最小值越高留给DCP的空间越小。比如72B模型在80GB卡上TP至少4那DCP最多28卡的话。吞吐vs延迟DCP优化吞吐对首token延迟帮助有限。如果场景对首token延迟敏感DCP可能不是最优选择。我的经验是先用TP把模型放下然后用DCP把KV Cache压到显存余量的50%以下。这样既有吞吐提升又留了足够的余量应对峰值。6.2 DCP和PagedAttention的配合vLLM的PagedAttention是它KV Cache管理的核心DCP和PagedAttention是配合使用的。PagedAttention负责block的分配和回收DCP负责block的切分和通信。两者配合的时候有个细节PagedAttention的block_table在DCP模式下需要做一次重映射把全局block索引映射到本地物理block索引。这个重映射是在KVCacheManager初始化的时候做的后续运行时不需要重复计算。但如果你动态调整了序列长度比如从32K扩展到64Kblock_table会重新分配重映射也需要更新。vLLM在每次序列长度变化的时候都会检查block_table如果需要扩展就重新做重映射。注意DCP模式下block_table的重映射是O(block数量)的操作。如果block数量很多比如长序列下有几万个block重映射的开销不能忽略。vLLM的做法是增量更新只重映射新增的block不重映射已有的block。6.3 监控DCP的通信开销DCP的通信开销是它性能的关键。如果通信开销太大DCP的收益会被抵消。监控DCP通信开销的方法有几个第一看NCCL的日志。vLLM启动的时候会打印NCCL的初始化信息包括通信组的大小和拓扑。如果通信组配置不对NCCL会打印警告。第二用nsys或者nvprof做profiling。这两个工具能看到每个kernel的执行时间和通信时间。如果通信时间占比超过30%说明DCP的通信开销偏大需要调整配置。第三看vLLM的metrics。vLLM暴露了一些metrics包括每层的计算时间和通信时间。如果通信时间随序列长度增长说明DCP的通信模式有问题比如prefill阶段的all-gather。我一般会先跑一遍默认配置用nsys看一下通信时间占比。如果占比在20%以下说明DCP配置合理如果超过30%就要考虑减小DCP_size或者换通信后端。6.4 一个容易忽略的细节DCP和CUDA Graph的兼容性vLLM支持CUDA Graph来加速decode阶段。但DCP和CUDA Graph的兼容性有个坑CUDA Graph要求每次执行的kernel序列完全一样但DCP的通信模式在decode阶段是动态的因为KV Cache的长度在变。所以vLLM在DCP模式下默认关闭了CUDA Graph或者只对固定长度的序列启用。如果你手动开了CUDA Graph可能会遇到通信死锁或者结果错误。解决办法是在DCP模式下不要手动开CUDA Graph让vLLM自己决定。如果确实需要CUDA Graph来加速可以考虑用固定长度的序列或者用piecewise CUDA GraphvLLM新版本支持。这个坑我踩过一次当时为了加速decode手动开了CUDA Graph结果跑了几百步之后突然hang住查了半天才发现是DCP的通信和CUDA Graph冲突。后来关了CUDA Graph就正常了。所以DCP模式下CUDA Graph要慎用。7. 从DCP延伸出去长序列推理还有哪些优化方向DCP解决了KV Cache的显存问题但长序列推理还有其他瓶颈。比如注意力计算本身是O(n²)的复杂度序列越长计算量越大。这方面有flash attention、sparse attention、linear attention等优化方向。另外KV Cache的量化比如FP8 KV Cache也能进一步压缩显存和DCP是互补的。还有一个方向是分层DCP不同层的KV Cache用不同的DCP_size。因为浅层的KV Cache对精度要求低可以切得更细深层的KV Cache对精度要求高切得粗一点。这样能在显存和精度之间做更好的权衡。vLLM目前还不支持分层DCP但社区里有相关的讨论未来可能会加。最后说一个我自己的体会DCP不是银弹它解决的是显存问题不是计算问题。长序列推理的瓶颈是多方面的显存、计算、通信、调度每个环节都可能成为瓶颈。DCP只是把显存这个瓶颈缓解了如果计算或者通信成为新瓶颈还需要其他优化手段。所以部署的时候要全面评估不要指望一个DCP就能解决所有问题。我在实际使用中发现DCP配合FP8 KV Cache量化能把128K序列的显存占用压到原来的1/4左右吞吐提升2倍以上。这个组合目前是我做长序列推理的默认配置。当然具体参数还要根据模型和硬件调没有一套配置能通吃所有场景。