执行摘要
本PR修复了V1中_dummy_run函数将NaN写入KV缓存null block的bug,该bug影响所有使用数据并行(DP)和专家并行(EP)的部署,导致模型精度显著下降(如DeepSeek R1从97%降至55%)。修复通过填充slot mapping为-1使内核跳过KV写入,与V2行为对齐,已通过现场测试验证NaN消除。
功能与动机
问题背景:在DP+EP部署中,空闲的DP rank会执行_dummy_run以保持同步。由于#25954重构后,_dummy_run无条件调用_get_slot_mappings,而slot mapping缓冲区初始化为torch.zeros,导致slot_idx=0映射到KV缓存的null block(block 0),使concat_and_cache_mla内核写入NaN。
影响:所有使用DP+EP的V1部署均受影响,无论量化类型(FP8、NVFP4、BF16),表现为精度严重回归。PR body中提到在DeepSeek R1 NVFP4部署中观察到精度从97%跌至55%。
实现拆解
主要改动集中在vllm/v1/worker/gpu_model_runner.py的_dummy_run函数中:
-
参数更正:将_get_slot_mappings调用中的num_tokens_padded参数从num_tokens(未填充计数)改为num_tokens_padded(填充后计数),确保内部填充逻辑正确执行。
python
slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
num_tokens_padded=num_tokens_padded, # 原为num_tokens
...
)
-
slot mapping填充:在获取slot mapping后,遍历slot_mappings_by_group中的所有张量,使用fill_(-1)设置为PAD_SLOT_ID。
python
if slot_mappings_by_group is not None:
for sm in slot_mappings_by_group.values():
sm.fill_(-1)
这使concat_and_cache内核跳过KV写入,避免污染null block。
评论区精华
review中主要讨论了修复的完整性:
- gemini-code-assist[bot] 指出:“Filling the slot mappings with
-1 is a correct approach... However, this fix appears to be incomplete due to an existing issue... _get_slot_mappings is called with num_tokens_padded=num_tokens, where num_tokens is the unpadded count.”
- elvircrn 回应:“I added this change
num_tokens_padded=num_tokens_padded now as it seemed like a typo. But the reason why this wasn't an issue was because... slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) already handles the rest of the data.”
- tlrmchlsmth 最终批准修复。
风险与影响
风险分析:
- 回归风险低:修复针对性强,且通过现场测试验证(slot mapping从[0]变为[-1],NaN消除)。
- 性能影响可忽略:
fill_(-1)操作开销小,dummy run本身频率较低。
- 兼容性良好:与V2的
get_dummy_slot_mappings行为一致,无breaking change。
影响评估:
- 用户:修复后模型推理准确性恢复,避免因NaN传播导致的错误输出。
- 系统:消除KV缓存污染,提升DP+EP部署的稳定性。
- 团队:揭示了slot mapping初始化与dummy run交互的设计缺陷,提示需加强相关测试覆盖。
关联脉络
- 历史PR #25954:该PR重构了slot mapping逻辑,使
_dummy_run无条件调用_get_slot_mappings,引入了本bug。
- 近期PR 38794:同属v1 worker模块优化,涉及GPU模型运行器性能改进,可对比学习核心路径变更模式。
- 近期PR 39169:同为bugfix,涉及预热逻辑与真实路径对齐,展示了类似“修复不完整”讨论模式。
演进趋势:本PR是V1核心worker模块的持续稳定性改进的一部分,反映了团队对DP+EP等复杂并行场景下边缘案例的深入排查。
参与讨论