【Bug已解决】RuntimeError: shape mismatch during KV cache init with EP DP on MoE model 解决方案一、现象长什么样在一个启用「专家并行EP 数据并行DP」的 MoE 模型上初始化 KV 缓存时引擎启动阶段抛出RuntimeError: shape mismatch并指向 KV cache 初始化那一段。典型日志形如RuntimeError: shape mismatch: KV cache block table shape [dp_rank1] expects num_blocks2048 but allocated 1024 for expert-parallel group 0或者更笼统RuntimeError: shape mismatch during KV cache init with EP DP on MoE model这类问题的几个标志方便你判断是不是同一个坑报错明确卡在KV cache init / block table 构建阶段模型权重可能已经加载完但缓存层一初始化就崩。错误里带有EP、DP、num_blocks、block_table、shape mismatch这些关键字。只要把并行策略退回到「纯 TP无 EP」或者「纯 DP无 EP」KV cache 就能正常初始化——说明问题出在「EP 与 DP 并存时KV cache 的分块视图对不齐」。崩溃往往只在num_experts不能被 EP 分组整除、或 DP 副本数改变每块显存预算时才出现调大/调小--gpu-memory-utilization有时能「蒙」过去但换个 batch 又复发。二、背景理解这个报错得先理清 KV 缓存在 vLLM 里是怎么组织的以及 EP DP 叠加后发生了什么。KV 缓存的基本结构vLLM 用 PagedAttention把每个序列的 KV 缓存切成固定大小的「块」block。全局维护一张block_table记录「序列 i 的第 j 块」存在哪块物理显存。所有序列共享一块连续显存池池子大小由可用显存和block_size决定记为num_blocks。DP 的影响开了 DP 后一个 DP rank数据并行副本独立服务一部分请求因此每个 DP rank 应该有自己的 KV 缓存池与 block_table——它只管自己那批序列不和其他 DP rank 共享。所以num_blocks会按 DP rank 数被切分每个 rank 拿到num_blocks // dp_size的预算。EP 的影响EP 把专家切到不同卡但 KV 缓存是注意力层的状态按理和专家无关——每个 token 的 K/V 张量维度是(num_kv_heads, head_dim)与专家数无关。然而在 vLLM 的实现里MoE 模型的 KV cache 管理有时和「专家并行组」做了耦合不同 EP rank 由于显存预算或分组约定不同可能对num_blocks、block_size产生不同预期。冲突点当 EP 与 DP 同时开启时理想情况下应形成[dp_rank][ep_rank]的二维缓存视图每个(dp, ep)单元内的 KV 池大小应当一致。但实际初始化代码如果只在「EP 维度」或只在「DP 维度」单边做切分就会出现DP rank 0 按num_blocks // dp_size分配EP group 内的某个 rank 因为另一个公式算出了不同的num_blocks两边在block_table的 shape 上不一致 →shape mismatch。本质KV cache 的「逻辑块数」在 EP 与 DP 两个维度上被各算了一遍两个结果没对齐。三、根因根因一句话在 EP DP 并存的 MoE 模型上KV cache 初始化时num_blocks/block_table的 shape被 EP 维度和 DP 维度各自独立计算且没有保证二者乘积与全局显存预算一致导致block_table在跨 rank 拼接/校验时形状对不上抛出RuntimeError: shape mismatch。展开看常见的几个具体成因DP 切分与 EP 切分重复计算全局num_blocks先被dp_size切又在某些 EP rank 内部被ep_size再切一次导致实际每块预算 num_blocks // dp // ep但block_table的 shape 仍按num_blocks // dp建二者差ep倍。EP 组内num_blocks不一致不同 EP rank 因为残差显存、对齐 padding 不同算出的num_blocks略有差异拼接block_table时列数不齐。block_size在 EP/DP 下被错误缩放block_size每块 token 数本应全 rank 统一却因 EP 相关代码路径把它和专家维关联导致某些 rank 的块大小不同KV 张量形状随之错位。KV head 维与 EP 耦合错误理论上num_kv_heads不受 EP 影响但若初始化代码误把 EP rank 数乘进 KV 头维KV 张量 shape 直接翻倍/缩小必然 mismatch。缺少跨 rank 一致性断言初始化时没有「所有 rank 的num_blocks必须相等」这类断言于是形状差被推迟到真正填block_table时才暴露成RuntimeError。总结这是 KV cache 分配器在「多并行维度并存」场景下预算核算逻辑不完整 缺少一致性校验导致的。四、最小可运行复现下面用纯 Python 字典/列表模拟「DP 切分后再被 EP 重复切分」导致block_tableshape 不一致不依赖 GPU 即可运行# reproduce_kvcache.py # 复现EPDP 下 KV 缓存块数被重复切分block_table 形状对不齐 def build_block_tables(num_blocks_global, dp_size, ep_size, wrong_double_splitFalse): tables {} for dp in range(dp_size): for ep in range(ep_size): if wrong_double_split: # bug: 既按 dp 切又按 ep 切 n num_blocks_global // dp_size // ep_size else: # 正确: 只按 dp 切ep 共享同一预算视图 n num_blocks_global // dp_size tables[(dp, ep)] [0] * n # 用列表长度代表 block_table 行数 return tables def check_shape_consistent(tables): shapes {len(v) for v in tables.values()} return len(shapes) 1, shapes if __name__ __main__: GLOBAL 4096 # 错误做法: 双重切分 bad build_block_tables(GLOBAL, dp_size2, ep_size4, wrong_double_splitTrue) ok, shapes check_shape_consistent(bad) print(双重切分 - 一致?, ok, 出现的块数集合:, shapes) # False, {512} # 正确做法: 只对 dp 切分 good build_block_tables(GLOBAL, dp_size2, ep_size4, wrong_double_splitFalse) ok, shapes check_shape_consistent(good) print(仅 dp 切分 - 一致?, ok, 块数:, shapes) # True, {2048}运行python reproduce_kvcache.py会看到「双重切分」让不同(dp, ep)单元的block_table长度即num_blocks不一致正是shape mismatch的成因。五、解决方案第一层最小直接修复最小修复KV cache 的全局块预算只按 DP 切分一次EP 维度共享同一预算视图并在构建block_table前断言所有 rank 的num_blocks完全一致。# fix_layer1_kvcache.py from typing import Dict, Tuple def compute_kv_blocks(num_blocks_global: int, dp_size: int, ep_size: int) - Dict[Tuple[int, int], int]: 只按 dp 切分全局预算ep 共享同一视图禁止重复切分。 assert num_blocks_global % dp_size 0, ( f全局块数 {num_blocks_global} 不能被 dp_size{dp_size} 整除 请调整 gpu-memory-utilization 或 block_size ) per_dp num_blocks_global // dp_size tables {} for dp in range(dp_size): for ep in range(ep_size): # 关键: ep 不变量所有 (dp,ep) 拿到相同 per_dp tables[(dp, ep)] per_dp # 一致性断言任何 rank 的块数都必须 equal per_dp bad [k for k, v in tables.items() if v ! per_dp] assert not bad, fblock_table 块数不一致: {bad} return tables def build_block_table(num_blocks: int, max_blocks_per_seq: int): 返回形状固定的 block_table: [num_seqs, max_blocks_per_seq]占位为 -1。 # 真实场景: torch.full((num_seqs, max_blocks_per_seq), -1) return [[-1] * max_blocks_per_seq] # 单序列示意 if __name__ __main__: t compute_kv_blocks(num_blocks_global4096, dp_size2, ep_size4) print(每 (dp,ep) 单元块数:, set(t.values())) # {2048}这一层修复点很小但关键把「EP 重复切」改成「EP 共享」并加一道一致性断言。改动集中在 KV cache 分配器不动模型、不动注意力核。六、解决方案第二层结构性改进把 KV cache 预算核算做成独立的KVCacheBudget模块明确区分「全局预算」「DP 视图」「EP 视图」三层避免任何代码路径再偷偷二次切分# fix_layer2_budget.py from dataclasses import dataclass dataclass class KVCacheBudget: num_blocks_global: int dp_size: int ep_size: int block_size: int def __post_init__(self): assert self.num_blocks_global % self.dp_size 0, 全局块数必须能被 dp 整除 assert self.block_size 1 # EP 不得改变块预算只校验 ep 与专家维的关系在别处 self.per_dp_blocks self.num_blocks_global // self.dp_size property def per_ep_view_blocks(self) - int: EP 视图共享 DP 预算返回同一数值杜绝二次切分。 return self.per_dp_blocks def block_table_shape(self, max_blocks_per_seq: int): # shape (per_dp_blocks 不直接作为行; 这里返回单序列表形状) return (self.per_dp_blocks, max_blocks_per_seq) def validate_all_ranks(self, ranks: list): 校验传入的每个 (dp,ep) rank 报告的块数都等于 per_dp_blocks。 violations [r for r in ranks if r[num_blocks] ! self.per_dp_blocks] assert not violations, fKV 缓存块数跨 rank 不一致: {violations} return True if __name__ __main__: bud KVCacheBudget(num_blocks_global4096, dp_size2, ep_size4, block_size16) ranks [ {dp: 0, ep: 0, num_blocks: 2048}, {dp: 1, ep: 3, num_blocks: 2048}, ] print(所有 rank 一致:, bud.validate_all_ranks(ranks)) # True这样结构性地保证预算只在KVCacheBudget里算一次任何 EP/DP 相关代码要拿块数都必须通过per_ep_view_blocks/per_dp_blocks这两个只读属性从源头消灭「各算各的」。七、解决方案第三层断言 / CI 守护把 KV cache 预算一致性钉进断言和 CI防止回归# fix_layer3_guard.py # ---- pytest 用例进 CI ---- def test_no_double_split_with_ep_dp(): from fix_layer1_kvcache import compute_kv_blocks t compute_kv_blocks(num_blocks_global4096, dp_size2, ep_size4) # 所有 (dp,ep) 必须相等且等于 4096//2绝不是 4096//2//4 assert set(t.values()) {2048} def test_ep_size_does_not_shrink_blocks(): from fix_layer2_budget import KVCacheBudget base KVCacheBudget(4096, dp_size2, ep_size1, block_size16).per_dp_blocks for ep in (2, 4, 8): b KVCacheBudget(4096, dp_size2, ep_sizeep, block_size16) assert b.per_ep_view_blocks base, fep{ep} 不应改变块预算 def test_global_not_divisible_by_dp_raises(): from fix_layer2_budget import KVCacheBudget try: KVCacheBudget(num_blocks_global4095, dp_size2, ep_size4, block_size16) assert False, 应因 4095 不能被 dp2 整除而报错 except AssertionError: pass再加一个启动期断言拦在真正填block_table之前def assert_kv_consistent_before_init(ranks): from fix_layer2_budget import KVCacheBudget # ranks: 每个 (dp,ep) 上报的 num_blocks / block_size sizes {r[block_size] for r in ranks} assert len(sizes) 1, fblock_size 跨 rank 不一致: {sizes} KVCacheBudget(4096, dp_size2, ep_size4, block_size16).validate_all_ranks(ranks)任何「EP 路径偷偷改了块预算」或「block_size 被错误缩放」的提交都会在 CI 立刻失败。八、排查清单看到shape mismatch during KV cache init with EP DP按顺序查先退到纯 TP关掉 EP 和 DP只tp卡数能初始化说明问题在 EP/DP 的缓存视图而非模型。查块数是否被双重切分全局num_blocks是否同时被dp_size和ep_size各除一次正确做法是只按 dp 切ep 共享。校验所有 rank 块数相等把每个(dp,ep)上报的num_blocks打出来应当完全一致不一致就是 mismatch 来源。block_size是否全 rank 统一每块 token 数必须所有 rank 相同EP 不应影响它。KV 头维是否被 EP 误乘num_kv_heads与 EP 无关确认初始化代码没有把ep_size乘进 KV 头维。全局块数能否被 dp 整除num_blocks % dp_size 0否则先调gpu-memory-utilization/block_size让它能整除。显存预算残差不同 EP rank 显存略有差异时给num_blocks做向下取整 padding 对齐避免差 1 块。加一致性断言在block_table构建前断言「所有 ranknum_blocks、block_size相等」把错误从运行期提前到启动期。看 vLLM 版本某些版本对 EPDP 的 KV cache 支持不完整升级或降级可能直接解决。最后才动注意力核优先在分配器层修预算核算不要为了对齐去改 PagedAttention 内核后者风险高且影响所有模型。九、小结EP DP 下的RuntimeError: shape mismatch during KV cache init根子是KV 缓存的全局块预算被 EP 和 DP 两个维度各自算了一遍、结果没对齐导致block_table跨 rank 形状不一致。修复三层递进第一层在分配器里只按 DP 切一次预算、EP 共享视图并断言所有 rank 块数相等第二层抽出一个只读的KVCacheBudget模块让任何 EP/DP 代码都只能从统一入口拿块数从源头消灭二次切分第三层用 pytest 把「EP 不改变块预算」「块数跨 rank 一致」钉进 CI。记住一条原则——KV 缓存预算是全局资源只应有一个权威核算点多并行维度并存时任何维度都只能「视图」它不能「重新切分」它。
)
