前阵子跑一个长文档问答任务7B模型上下文给到128K。数据灌进去还没来得及看效果先瞄了一眼显存KV Cache直接吃掉60多个GB。单卡H100 80GB被压到红线再塞一个batch就得OOM。换成70B模型更夸张光KV部分就要300多GB一张卡根本放不下。当时第一反应是加卡——但加卡并不等于加显存。DP、TP、PP各管一段真正卡脖子的“长序列KV分片”问题这三板斧都打不到点上。后来把Context Parallelism上下文并行简称CP打开把序列在token维度上切到多卡才算把这个题解了。这篇文章就把这套东西从原理到落地完整捋一遍适合正在被长上下文显存问题折磨的推理或训练工程师也适合想搞懂大模型并行策略该怎么选的算法同学。1. 长上下文场景下的显存瓶颈到底卡在哪1.1 先算一笔账KV Cache才是大头输入长度一旦上去显存占用格局会彻底反转。很多人直觉上觉得“模型越大越占显存”这句话在长序列场景下只说对了一半。模型权重是一次性加载、大小固定的KV Cache却是跟着序列长度动态增长的每生成一个token所有层的K和V都要新存一份。计算公式很直接以FP16精度为例KV Cache大小字节 2K和V两份 × 层数 × hidden_size × 序列长度 × batch数 × 2字节拿7B模型来算普通配置是32层、hidden_size 4096序列长度128Kbatch为12 × 32 × 4096 × 131072 × 1 × 2 68,719,476,736 字节约64 GiB70B模型按80层、hidden_size 8192来算同样128K序列2 × 80 × 8192 × 131072 × 1 × 2约320 GiB这个数字意味着什么7B权重才14GB左右但在128K上下文下KV Cache是权重的四倍多。70B权重140GBKV Cache又是权重的两倍多。很多人在长上下文任务里遇到的OOM根本不是模型太大装不下而是KV Cache把显存挤爆了。所以排查长上下文OOM第一步永远是先算这笔账别急着怪模型。1.2 常规并行三板斧DP、TP、PP为什么都不够用数据并行DP解决的是“多个请求怎么分摊”的问题它把模型复制到多张卡上每张卡处理不同的数据。但每张卡上依然有一份完整的模型和完整的KV Cache单条长序列照样放不下。张量并行TP是把权重、注意力头切到多卡确实能把权重摊薄但注意力计算时的all-reduce通信量随模型宽度增长。序列越长中间激活越大TP的通信开销跟着涨而且TP特别吃卡间带宽通常要求NVLink这种高速互联跨节点效果很差。流水线并行PP按层切分每张卡只放一部分层权重是摊下来了。但问题在于每个注意力层在计算时它负责的那一层依然要处理整条序列的KV Cache。哪怕层数被切到8张卡上负责Attention层的某张卡在那一刻还是要吃掉完整的序列KV。PP没有解决“单层KV过大”的问题。这三板斧的共同盲区是它们都在切模型的结构维度没有动“序列本身”这个最暴力的维度。1.3 Context Parallelism的基本思想把序列也当成可切分的资源Context Parallelism的思路就一句话既然KV Cache按序列长度线性增长那就把序列直接切成C块分给C张卡。每张卡只持有序列的一个连续片段也只负责这个片段对应的Attention输出。听起来很简单但实现时有两个核心问题绕不开第一注意力计算是全局的。一个token的Attention要看到序列里所有其他token的K和V每张卡只持有片段怎么让它在不持有全部KV的情况下算出正确结果第二Softmax的归一化是跨整个序列做的。分块之后每块内部的局部softmax和全局softmax不是一回事怎么保证最终结果和未切分时完全一致这两个问题分别对应了CP实现里的两种主流路线All-to-All形式和Ring形式。下一章展开讲。2. Context Parallelism的核心原理拆解2.1 序列切分与KV的归属先把“切”这件事说清楚。假设有C张卡原序列长度是L理想情况下每张卡分到 L/C 个token的连续区间。输入在进入Transformer之前先按token维度切成C份第i份送进第i张卡。经过QKV投影后每张卡上有自己那一段的 Q_i、K_i、V_i。在Prefill阶段这是完整的一块在Decode阶段KV Cache本来就是分片存放的每张卡只为自己负责的那一段积累KV新增量也落在对应的卡上。这样的直接收益是单卡KV Cache峰值从原来的“全长”降到“全长/C”。7B模型128K序列4路CP每卡KV约16GiB70B模型128K序列配4路CP每卡降到约80GiB再叠加TP切权重就能塞进80GB卡了。但这里有个反直觉的点每张卡算自己那段的输出却需要全局的KV才能算。看你怎么组织数据交换这就出现了两种路线。2.2 路线一All-to-All交换DeepSpeed-Ulysses的思路All-to-All路线的想法是与其让每张卡缺东少西不如直接交换一次数据让每张卡拿到自己需要的完整视图。具体来说在QKV投影之后做一次All-to-All通信原来每张卡持有全序列的一小段、但拥有所有注意力头交换机后每张卡持有完整的序列位置、但只保留一部分注意力头。这样每张卡只需要对自己名下的那几个头跑完整的全局序列AttentionSoftmax也是本地完整归一化不需要跨卡合并。这个方案的优点是实现简单、数学上和单卡完全等价每个Attention输出的计算结果都在卡内闭合。代表实现是DeepSpeed-Ulysses。缺点是All-to-All是稠密通信所有卡同时互相收发数据量随序列长度和头数上涨。卡数一多对集合通信带宽的压力相当大跨节点场景尤其吃紧。2.3 路线二Ring Attention环形传KV另一条路线不在开局做全局交换而是把KV块在设备环形拓扑上逐轮传递。初始状态每张卡持有自己的 Q_i 和对应的 K_i/V_i。第一轮每张卡用自己的 Q_i 去计算与本地 KV 块的Attention然后把KV块传给下一张卡同时从上一张卡收到新的KV块开始算下一轮。这样经过C轮每张卡都能看到全部C个KV块Q_i始终留在本地最终得到完整的全局Attention输出。Ring形式的通信从“一次稠密全交换”变成了“C次点对点环形传递”通信压力分散了很多而且可以在传输第k1个KV块的同时计算第k个KV块做到通信和计算重叠。代表实现是Ring Attention。两种路线怎么选我的经验是卡间带宽很好NVLink/IB、卡数不多时All-to-All实现简单、调试方便卡数多、或者跨机、或者追求极致吞吐时Ring的异步重叠和分散通信优势更明显。具体选型还得看框架支持程度别为了理论优雅去手撸一个没人维护的实现。2.4 分块Softmax数学上必须做的rescale不管哪种路线只要Attention不是在单卡全局一次算完就会遇到Softmax分块合并的问题。普通Softmax需要两个全局量所有logits的最大值和指数和。分块时每块只能看到自己那部分logits局部最大值和局部指数和跟全局不是一回事。直接合并会错。业界通行做法是FlashAttention里那套在线Softmaxonline softmax。维护两个运行状态全局最大 m、全局指数和 l。当一个新块到达时先取合并后的新最大值m_new max(m_old, 本块logits最大值)旧块的指数和需要重新缩放l_old × exp(m_old - m_new)新块的贡献累加l_new 缩放后的旧值 本块 exp(scores - m_new)之前的输出块也要rescaleO_old × exp(m_old - m_new)再加上本块的加权贡献最后在全部块处理完后统一除以 l_new 得到最终输出这套过程保证最终结果和单卡一次算完在数学上等价不是近似。FP16下极端场景可能有尾数误差但一般影响很小真遇到精度敏感的任务换BF16基本能压住。理解这个细节很重要。因为一旦你要手写CP实现或者要排查“开了CP之后结果和单卡不完全一样”的问题90%的坑都出在这个rescale逻辑上。2.5 理论收益与代价用带宽换峰值显存CP最核心的收益是KV Cache峰值显存近似降到原来的 1/C。在长序列场景下这往往是决定能不能跑起来的关键。但它不是免费的。总计算量基本不变Attention本身的计算量没少还多了一点点通信相关的开销新增的代价是通信量。Ring形式每层每轮要传一份KV块All-to-All形式每层要做两次全局交换。所以CP本质上是“用通信带宽换单卡峰值显存”。这决定了它的适用边界序列越长收益越明显序列短到一定程度通信固定开销会吞掉全部收益甚至比单卡更慢。后面会专门讲这个阈值怎么估。3. 主流框架里的CP落地方式3.1 Megatron-LM训练侧的CP配置Megatron-LM是最早把Context Parallelism做成生产级能力的框架之一。训练场景下启动参数里加上--context-parallel-size即可。一个典型的长序列训练配置大致长这样伪代码参数以你的版本为准python pretrain_gpt.py \ --tensor-model-parallel-size 4 \ --pipeline-model-parallel-size 2 \ --context-parallel-size 4 \ --sequence-parallel \ --num-layers 48 \ --hidden-size 8192 \ --seq-length 131072 \ ...注意几个点一是CP通常需要和--sequence-parallel配合使用LayerNorm、Dropout这些非张量并行操作也按序列切分避免频繁转换数据排布二是总卡数等于 TP × PP × CP对不上的话初始化阶段就会报错三是CP组最好落在NVLink域内因为每层都有KV交换跨节点走以太网会很痛苦。Megatron的实现里CP和TP是强耦合的。TP切多头、CP切序列两者组合后通信模式更复杂但也正是这种组合才能在超长序列下同时压住权重和KV两座大山。3.2 vLLM推理侧的CP支持如果是做推理vLLM比较新版本已经原生支持CP。启动时可以通过LLM构造参数直接指定cp_size或者命令行加--cp-size。from vllm import LLM llm LLM( modelyour-7b-model, tensor_parallel_size2, cp_size2, max_model_len131072, )这里tensor_parallel_size × cp_size才是占用的总卡数。我实际测下来7B模型128K上下文单卡必爆TP2也只是把权重摊了KV还是要64GiB照样OOMTP2 CP2每卡KV降到32GiB左右权重也摊了一半就能稳定跑了。vLLM里开CP之后KV Cache Manager的分片逻辑会变日志里会打印context_parallel_size之类的信息显存池的分配也会按CP切分。如果开完发现显存占用和没开一样先检查是不是参数没生效很多版本对参数名大小写、位置很敏感。3.3 DeepSpeed-UlyssesAll-to-All路线的代表DeepSpeed-Ulysses走的是前面说的All-to-All路线主要面向训练。它的核心贡献是把“序列并行”和“注意力头并行”统一到一套通信框架里先按序列切分QKV投影后All-to-All交换把序列分片转换成注意力头分片每个设备独立算全局序列的部分头最后再交换回来。它的优点是扩展性好理论上一旦通信域够大CP可以撑到非常大的规模缺点是All-to-All通信对网络质量要求高。实测中单机多卡NVLink环境表现不错跨机节点带宽不足时吞吐会明显下滑。3.4 并行策略编排一张表看懂怎么选策略切分维度通信模式显存收益适合场景DP数据梯度同步无多请求吞吐TP权重/注意力头AllReduce权重均摊大模型权重放不下PP层点对点权重均摊跨节点大模型CP序列All-to-All/RingKV均摊长上下文实际编排时我一般按这个原则单机内优先TPCP把最重的通信留在NVLink里跨节点用DP或PP把通信频率降下来总卡数吃紧时优先保CP因为长上下文场景里KV才是主要矛盾。4. 实操把一个长上下文推理任务切到多卡上4.1 先判断要不要开CP怎么决定CP大小实操第一步不是开参数而是做决策。还是上面那个例子7B模型128K上下文目标是把推理稳定跑在80GB单卡上。先算KV64GiB。再看权重7B FP16约14GiB。两项相加快80GiB再加激活和推理框架自身的开销单卡必然爆。这种情况就必须切序列。CP大小怎么定目标是让“单卡KV 单卡权重 激活”低于可用显存并留出余量。4路CP时KV降到16GiBTP2时权重降到7GiB加起来不到30GiB余量非常充足。如果换成70B模型权重140GiBKV 320GiB那至少要TP8 CP8这种组合才能压到单卡可承受范围成本就上去了。我的习惯是留20%左右的显存余量给激活、临时缓冲和推理框架开销别卡着上限配置否则一旦序列里出现长尾请求就直接OOM。4.2 一个可复现的vLLM配置过程我实际跑通的配置大概是这样的供参考python -m vllm.entrypoints.openai.api_server \ --model /path/to/7b-model \ --tensor-parallel-size 2 \ --cp-size 2 \ --max-model-len 131072 \ --gpu-memory-utilization 0.9 \ --enforce-eager几点说明gpu-memory-utilization我设到0.9给KV Cache池留足空间。CP切分后每张卡的KV池是自己那一份利用率过低头会浪费显存过高容易OOM。第一次验证建议加--enforce-eager关掉CUDA Graph因为有些版本的CP路径和Graph捕获有兼容问题先跑通再开优化。启动后看日志里的context_parallel_size确认等于你设的值。有的版本里cp_size和tensor_parallel_size是分开算的总卡数必须等于两者乘积。验证是否生效最直观的办法是看每张卡的显存占用。CP生效时每张卡的显存占用应该比较接近而且明显比单卡跑同样的模型低。如果某张卡飙高、其他卡闲着多半是负载分配出了问题。4.3 训练侧一个可参考的Megatron启动形态训练场景大同小异核心是配置对齐。除了3.1里的启动参数还要注意数据加载时的序列切分逻辑要和CP大小匹配确保每个数据并行rank拿到的数据块位置正确。这块容易出隐性问题模型并行没问题但loss曲线和单卡对不上查了半天发现是数据切分和attention切分的语义没对齐。Megatron里CP通常要求--sequence-parallel一并打开还要保证--global-batch-size能被DP大小整除--seq-length能被CP大小整除。不能整除时尽量padding到整数倍宁可多算几个token也别让分块不均引发边界错误。4.4 负载均衡一个永远不能忽视的细节CP的理想假设是每张卡分到等长的序列块。但真实场景里序列长度不一定能被CP大小整除多出来的几个token会让某一张卡多算一点。训练时padding就好推理时问题更隐蔽——请求长度天然不齐。我踩过一个典型的坑一批请求里混了4K和128K两种长度开CP后长请求的KV在那张卡上占了大量显存其他卡早早算完空等整体吞吐被最慢的那张卡拖死。后来在接入层按长度分桶把相近长度的请求放到同一个Batch里CP的收益才真正出来。长Prefill阶段也建议配合Chunked Prefill使用。CP已经按序列切分了Chunked Prefill再按时间切分两者叠加能进一步压峰值激活显存尤其适合超长Prompt第一次解析的场景。5. 性能调优与踩坑记录5.1 通信开销是CP的头号敌人CP省的是显存花的是通信。同一个任务CP从2加到4显存确实更省了但端到端延迟未必变快甚至可能变慢。原因就是通信占比上去了。我测过一轮对比7B模型128K序列80GB A100CP2时每卡KV约32GiB吞吐比单卡略低但能跑起来CP4时每卡KV约16GiB显存充裕但吞吐又比CP2掉了一截。这很典型——CP不是开得越大越好够用就行。影响最大的是卡间带宽。NVLink域内开CP通信开销基本可控跨节点走以太网的话All-to-All方案会非常难受Ring方案稍好但仍有限。所以我的经验是CP组尽量限制在单机内跨机的长序列任务优先考虑PPDP而不是CP扩到跨机。5.2 负载不均的木桶效应前面说过负载不均这里再说深一点。CP是强同步模式——所有卡要等最慢的卡算完才能进入下一轮通信或下一层计算。只要有一张卡多分到了几个token或者一张卡因为KV热点导致显存不足触发重算整体性能立刻被打到最低点。推理场景里动态请求调度会让这个问题更难办。vLLM这类框架有调度器可以根据序列长度和显存预算做一定程度的均衡但不会像你想象中那么完美。实测经验是长尾请求多的时候与其依赖框架自动调度不如在业务层做长短请求分离分别跑不同的CP配置甚至不同的实例。5.3 数值精度开CP后结果和单卡不完全一致严格实现下CP在数学上应该和单卡等价。但FP16的舍入误差、在线Softmax的rescale顺序差异可能导致最终输出token级别的微小差异极端情况下会累积成不同的采样结果。排查方法很简单同一个Prompt分别用CP1和CP4跑一遍对比logits。差异在一个很小的epsilon范围内就是正常的。如果差异明显先看是不是误差累积再看是不是实现里softmax的中间统计用了低精度。生产中我用BF16基本没遇到问题FP16在超长序列下偶尔会有漂移敏感场景回退BF16即可。5.4 什么时候不该用CPCP不是银弹。序列只有4K、8K的时候单卡KV可能就几GiB开CP纯属给自己找通信负担。启动通信的内核延迟、All-to-All的带宽占用在短序列下是净亏损。我个人的经验阈值KV Cache占用超过单卡可用显存的一半以上才值得考虑CP不到这个量级优先加Batch提升吞吐或者用TP分摊权重都比CP划算。另外如果任务是短上下文但QPS很高CP基本帮不上忙那是另一个优化方向。6. 常见问题速查表现象可能原因排查与处理开了CP仍然OOMCP参数未生效总卡数不等于TP×CP序列长度未被正确切分检查启动日志的cp_size确认显卡占用接近均衡确认max_model_len正确每张卡显存严重不均请求长度不齐负载均衡没开padding逻辑不对接入层按长度分桶开启/调整chunked prefill检查pad token处理吞吐反而下降通信占比过高CP开得过大跨节点网络瓶颈缩小CP保证CP在NVLink域内对比CP1/2/4的吞吐曲线结果和单卡不一致FP16在线Softmax误差累积rescale逻辑bug换BF16对比logits差异量级关闭CP交叉验证启动报rank不匹配TP、PP、CP乘积不等于总卡数核对总卡数配置检查分布式init参数Ring Attention模式下卡死环形拓扑配置错误死锁检查拓扑生成逻辑加上通信超时和重试机制这些坑我几乎都踩过一遍。最浪费时间的不是配置本身而是“看起来都对了但性能不对”的隐性负载不均问题。我的排查顺序固定成这样先确认CP生效再看显存分布然后看通信占比最后才怀疑精度和代码bug。7. 一点实际经验把CP用顺之后我遇到长任务的第一反应变了不是急着加卡而是先算KV再算权重然后决定CP开几路、TP开几路。这个顺序看着简单但能省下大量试错时间。另一个心得是CP和批量推理是互补的不要二选一——长序列任务用CP切显存短序列任务用大Batch吃满算力两个手段配合才能把卡的价值榨干。最后再分享一个实操小技巧在正式跑长任务之前先用一个短序列、小模型把CP链路整体验证一遍包括日志、显存、精度对齐全部正常后再切到目标模型。这能帮你把“参数没生效”和“配置方式不对”这类低级问题隔离在真正耗时的跑批之前。