[Bug]: Batch invariance is broken when sequence parallelism / async TP is enabled (`VLLM_BATCH_INVARIANT=1` + `pass_config.enable_sp`)

问题场景

用户使用 vLLM 以张量并行（TP=2/4/8）加载 Qwen/Qwen3-1.7B，开启 VLLM_BATCH_INVARIANT=1 希望输出与 batch 组成无关，同时通过 compilation_config.pass_config 设置了 enable_sp=True（或设置 fuse_gemm_comms=True，该选项会隐式开启 SP）。在单条请求（bs=1）与 64 条请求同批（bs=64）的对比中，采样 token 的 logprobs 甚至贪心解码出的 token 出现差异。默认编译在 SP 下会强制 full-graph 且 cudagraph_mode=FULL。

Issue 中的测量设置：64 条 prompt，贪心（temperature=0、max_tokens=24、logprobs=1、ignore_eos=True）、seed=0、max_model_len=512、max_num_seqs=64、gpu_memory_utilization=0.5，逐 token 按位比较。sm_120 上 SP 阈值启发式返回 None，因此需显式设置 sp_min_token_num；H100/B200 上启用 enable_sp=True 后启发式会自动开启 SP。

报错原文

[Bug]: Batch invariance is broken when sequence parallelism / async TP is enabled (`VLLM_BATCH_INVARIANT=1` + `pass_config.enable_sp`)

原因分析

最可能的原因是 sequence parallelism（async TP）路径本身不具备批量不变性。启用 enable_sp 后，SequenceParallelismPass 会把 all_reduce + rms_norm 重写为 reduce_scatter + rms_norm + all_gather。在 TP=4 下，即使 SP 作用于所有 compile range（单一 range，不存在跨阈值），结果仍随 batch 组成变化；而 TP=4 关闭 SP 时完全一致，说明问题不在通用 TP 路径，而在 SP 路径。

较新的复现（关闭前缀缓存、sp_min_token_num=1 单一 compile range）进一步显示：TP=4 + SP 下批量执行连 run-to-run 都不确定（同一进程内两次相同的 bs=64 在 1536 个 logprob 中有 1175/865/1030 个不同，9–12 条贪心序列变化）；而 bs=1 在同一进程内与跨进程按位一致，SP 关闭时 bs=1/bs=64 均按位一致。NCCL_DEBUG_SUBSYS=TUNING 显示 TP=4 下所有 ReduceScatter/AllGather 均为 RING / SIMPLE / channel{0..0}，NCCL 侧已被固定，因此可以排除“归约顺序随 chunk owner 变化”作为 TP=4 的完整解释。当前更可能指向 SP 的 reduce-scatter → rms_norm → all-gather 路径（torch.ops.vllm.reduce_scatter / all_gather → _reduce_scatter_out_place / _all_gather_out_place → PyNccl on current_stream()）存在同步/顺序问题，但 Issue 作者尚未在代码中定位。

同时，enable_sp / fuse_gemm_comms 从未与 VLLM_BATCH_INVARIANT 做互斥校验（对比 vllm/config/vllm.py 中 cascade attention 与 fused allreduce+RMSNorm 的既有守卫），文档 docs/features/batch_invariance.md 也未提及 SP / async TP，因此没有配置层拦截。

环境排查

• 确认 vLLM 版本/提交：main @ 7470082f5（2026-09-10）或 v0.29.0 wheel，是否使用 VLLM_USE_PRECOMPILED=1 预编译 kernel。
• 确认是否设置了 VLLM_BATCH_INVARIANT=1。
• 确认 compilation_config.pass_config 中是否 enable_sp=True，或是否设置了 fuse_gemm_comms=True（会隐式开启 SP）。
• 确认 sp_min_token_num 取值（sm_120 上启发式返回 None，需显式设置；H100/B200 会自动开启 SP）。
• 确认 TP 规模（Issue 中 TP=2 与 TP=4/8 表现不同）。
• 确认 attention backend（TRITON_ATTN / FLASH_ATTN）、cudagraph_mode（FULL / NONE）、prefix caching 是否开启。
• 确认 torch 2.13.0+cu130、triton 3.7.1、NCCL 2.27.3。
• 确认显卡与驱动：4x NVIDIA RTX PRO 6000 Blackwell（sm_120, 96 GB），驱动 580.159.04，PCIe 无 NVLink。

解决步骤

1. 先复现并确认触发条件：按 Issue 测量设置运行对比脚本，分别比较 bs=1 与 bs=64 的 token ids 与采样 token logprobs，逐位比对。可用 VLLM_BATCH_INVARIANT=1 VLLM_ATTENTION_BACKEND=TRITON_ATTN python bi_sp_probe.py <tp> <off|threshold> 形式运行 Issue 提供的探针脚本。
2. 做对照实验：在相同环境下重复一次并关闭 SP（enable_sp=False），确认 TP 下 bs=1 与 bs=64 是否恢复按位一致（Issue 中 TP=4 关闭 SP 为 0/1536 差异）。
3. 若确认是 SP 引起的组合问题，可优先尝试的规避方案是在需要 batch invariance 时不要启用 enable_sp / fuse_gemm_comms，避免二者同时生效。
4. 若必须使用 SP，可尝试关闭 prefix caching 以减少交互：Issue 中 TP=2 + SP 在关闭前缀缓存后达到按位不变（0/1536），说明该项能消除 TP=2 的残留差异；但该手段对 TP=4/8 无效，TP=4 关闭前缀缓存后批量执行仍不确定。
5. 若用于缩小范围，可尝试 cudagraph_mode=NONE（仍编译、仍应用 SP）：Issue 中重复运行差异从约 1000 降到 118/1536，但并未消失，说明问题对时序敏感，不是 CUDA graph 捕获产物。
6. 用 NCCL_DEBUG_SUBSYS=TUNING 确认集合通信算法是否被固定（Issue 中 TP=2 与 TP=4 均为 RING/SIMPLE/channel{0..0}），以排除 NCCL 算法切换因素。
7. 关注上游是否有针对 SP reduce-scatter → rms_norm → all-gather 路径的修复；在确认前不要假设该组合可用。

验证方法

用与 Issue 相同的对比方式验证：对同一批 prompt 分别以 bs=1 和 bs=64 生成，逐位比较 token ids 与采样 token logprobs，期望差异为 0/总数；再在同一进程内重复两次相同的 bs=64 生成，确认 run-to-run 也按位一致。若要专门验证 TP=2 的残留，可在关闭 prefix caching 后重跑，Issue 中该配置下达到 0/1536。注意验证需在多 TP 规模下分别进行，TP=2 与 TP=4/8 的结果并不相同。

参考来源

vllm-project/vllm #56370

AI 工具推荐

想把多个 AI 模型放在一个入口？

GamsGo AI 集成 ChatGPT、DeepSeek、Gemini、Claude、Midjourney、Veo 等常用模型，适合写作、绘图、视频和日常 AI 工作流。

了解 GamsGo AI

推广链接：通过此链接购买，我可能获得佣金，不影响你的价格。