执行摘要

• 一句话：FlashInfer 纯 allreduce 后端，复用融合 workspace 加速 DSV3/V4
• 推荐动作：值得精读，尤其是以下三个设计决策：
  1) can_use_flashinfer_allreduce 与 flashinfer_allreduce 的拆分——如何在 Dynamo 符号执行下保持判定 rank-invariant 且不引入 Dynamo guard 重编译；
  2) workspace 与进程组精确绑定的必要性——hybrid EP+TP 下同 world size 不同 rank 集合会静默产生错误归约，这是分布式正确性中易被忽视的细节；
  3) custom op 包装对 piecewise CUDA graph 的支持方式。建议结合 PR#32330 的 SM120 修复与 PR#32339 的性能对比一起阅读，可形成对多节点 allreduce 路径选型的完整认识。

功能与动机

PR body 指出：SGLang 已支持通过 FlashInfer 的 kARResidualRMSNorm 做 fused allreduce（allreduce + residual + RMSNorm 单 kernel），但并非每个 allreduce 点都能走融合路径——RowParallelLinear、attention 输出 fallback、MoE TP fallback 都是纯（非融合）allreduce，目前即使 FlashInfer mnnvl/trtllm workspace 已初始化可用也总是回退到 NCCL/custom-allreduce。本 PR 将这类纯 allreduce 路由到 FlashInfer AllReduceFusionPattern.kAllReduce，复用现有 workspace，无额外内存开销。nvpohanh 在评论中补充了关键背景：DSV4 因 mHC 无法做 AR+Norm 融合，此前 SGLang 只对 DSV3 支持 FlashInfer MNNVL AR，本 PR 补齐 DSV4 支持，可显著提升 DSV4 在 GB200 TP8 小并发场景的性能。

实现拆解

实现按以下 5 步拆解：

1. 新增纯 allreduce 内核封装：在 python/sglang/srt/layers/flashinfer_comm_fusion.py 中新增 can_use_flashinfer_allreduce() 与 flashinfer_allreduce() 两个函数。前者做全部可用性判定（FlashInfer 可用性、2D 连续形状、workspace 初始化、world size 与 group 精确匹配、buffer 容量），后者只负责调用 allreduce_fusion(pattern=kAllReduce)，且刻意不捕获 kernel 异常——吞异常会让本 rank 回退 NCCL 而其他 rank 仍在 kernel 内，导致集体通信失配挂起。判定与执行分离是为了适配 Dynamo：判定在普通 Python 中于 trace 期完成，执行注册为 opaque custom op 可进入 piecewise CUDA graph。

2. 并行状态层 dispatch 与 group 打标：在 python/sglang/srt/distributed/parallel_state.py 中注册 flashinfer_allreduce(Tensor, group_name) custom op，并在 GroupCoordinator.all_reduce() 的编译与非编译两条路径上、在既有 NCCL/custom-allreduce 链之前插入 dispatch。新增 _tag_groups_for_flashinfer_allreduce_only()：只在真正拥有 fusion workspace 的组上打 _fi_workspace_hint 标记（attention-TP 一个、MoE 一个；MoE workspace 在 moe_ep_size > 1 时 rendezvous 于 EP 组，否则于 MoE-TP 组，二者只取其一），避免 hybrid EP+TP 下归约错 peer。

3. 启动流程接线与参数简化：python/sglang/srt/distributed/bootstrap.py 的 _set_all_reduce_flags() 调用 set_flashinfer_allreduce_only(server_args.flashinfer_allreduce_fusion_backend is not None)，_init_parallel_groups() 在 initialize_model_parallel() 之后调用 _tag_groups_for_flashinfer_allreduce_only()。按 review 意见删除了独立的 --enable-flashinfer-pure-allreduce flag 和架构白名单，kAllReduce 直接复用融合后端 flag，随融合路径默认开启。

4. DSV3/V4 自动启用：python/sglang/srt/arg_groups/overrides.py 将 DeepseekV4ForCausalLM 加入 _FLASHINFER_ALLREDUCE_FUSION_ARCHS，使 DSV4 用户无需手动传任何参数即可获得该加速路径。

5. hybrid EP+TP 配套修复与测试：python/sglang/srt/layers/communicator.py 的 should_fuse_mlp_allreduce_with_next_layer() 新增守卫：当 moe_ep_size > 1 且 moe_tp_size > 1 时禁止融合 post-experts all-reduce（融合会让下一层 residual+LN 只在一个组上归约，导致激活只归约一半 peer 的静默错误）。测试侧在 test/registered/unit/layers/test_flashinfer_comm_fusion.py 扩展了 TestFlashInferAllReduceOnly（12 个用例，覆盖正确性、形状守卫、容量/dtype 守卫、group 匹配守卫）与 TestTagGroupsForFlashInferAllReduceOnly，并新增 test/registered/unit/layers/test_layer_communicator_fusion_gate.py 覆盖融合门控的 hybrid/pure 场景。

关键文件：

• python/sglang/srt/distributed/parallel_state.py（模块 并行组；类别 source；类型 core-logic；符号 flashinfer_allreduce, _can_use_flashinfer_allreduce, _flashinfer_allreduce, set_flashinfer_allreduce_only）: 核心 dispatch 所在：注册 flashinfer_allreduce custom op，在 GroupCoordinator.all_reduce() 热路径插入 FlashInfer 分支，并通过 _tag_groups_for_flashinfer_allreduce_only() 精确打标避免 hybrid EP+TP 下归约错 peer。
• python/sglang/srt/layers/flashinfer_comm_fusion.py（模块 通信融合；类别 source；类型 core-logic；符号 can_use_flashinfer_allreduce, flashinfer_allreduce）: 新增 can_use_flashinfer_allreduce 与 flashinfer_allreduce 两个核心函数：前者集中所有可用性判定（rank-invariant），后者执行 AllReduceFusionPattern.kAllReduce kernel 调用且不吞异常，是整条路径的正确性基石。
• python/sglang/srt/layers/communicator.py（模块 通信层；类别 source；类型 bugfix；符号 should_fuse_mlp_allreduce_with_next_layer）: 修复 hybrid EP+TP 下融合 post-experts all-reduce 会静默产生未完全归约激活的严重正确性问题（Qwen3-30B-A3B tp=4 ep=2 实测为垃圾输出）。
• python/sglang/srt/distributed/bootstrap.py（模块 启动流程；类别 source；类型 dependency-wiring；符号 _set_all_reduce_flags, _init_parallel_groups, set_flashinfer_allreduce_only, _tag_groups_for_flashinfer_allreduce_only）: 启动接线：在 _set_all_reduce_flags 中根据融合后端 flag 设置 set_flashinfer_allreduce_only，在组初始化后调用 _tag_groups_for_flashinfer_allreduce_only，是功能开关生效的入口。
• python/sglang/srt/arg_groups/overrides.py（模块 参数覆盖；类别 source；类型 configuration；符号 _FLASHINFER_ALLREDUCE_FUSION_ARCHS）: 将 DeepseekV4ForCausalLM 加入 FlashInfer allreduce fusion 自动启用架构集，使 DSV4 用户开箱即用；同时清理了按 review 删除的独立纯 allreduce 架构白名单。
• test/registered/unit/layers/test_flashinfer_comm_fusion.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestFlashInferAllReduceOnly, TestTagGroupsForFlashInferAllReduceOnly, _make_manager, _patched_attn_workspace）: 新增 TestFlashInferAllReduceOnly（12 个用例）与 TestTagGroupsForFlashInferAllReduceOnly，覆盖输出正确性、2D/连续形状守卫、容量/dtype 守卫、group 精确匹配守卫及打标拓扑选择，是判定逻辑的主要回归保障。
• test/registered/unit/layers/test_layer_communicator_fusion_gate.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestFuseMlpAllReduceGate, test_hybrid_ep_tp_does_not_fuse, test_pure_tp_still_fuses, test_pure_ep_still_fuses）: 新增测试验证 hybrid EP+TP 下禁止融合、纯 TP/纯 EP 下仍允许融合，直接守护 communicator.py 的正确性修复。

关键符号：flashinfer_allreduce, can_use_flashinfer_allreduce, _can_use_flashinfer_allreduce, _flashinfer_allreduce, _tag_groups_for_flashinfer_allreduce_only, set_flashinfer_allreduce_only, should_fuse_mlp_allreduce_with_next_layer

关键源码片段

python/sglang/srt/distributed/parallel_state.py

核心 dispatch 所在：注册 flashinfer_allreduce custom op，在 GroupCoordinator.all_reduce() 热路径插入 FlashInfer 分支，并通过 _tag_groups_for_flashinfer_allreduce_only() 精确打标避免 hybrid EP+TP 下归约错 peer。

# python/sglang/srt/distributed/parallel_state.py
# 注册为 custom op，使 Dynamo 将其视为不透明调用，可进入 piecewise CUDA graph；
# 判定逻辑在 trace 期由 _can_use_flashinfer_allreduce 完成，op 自身无回退。
@register_custom_op(out_shape="tensor")
def flashinfer_allreduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
    assert group_name in _groups, f"Group {group_name} is not found."
    group = _groups[group_name]()
    if group is None:
        raise ValueError(f"Group {group_name} is destroyed.")
    return group._flashinfer_allreduce(tensor)
​
​
class GroupCoordinator:
    def _can_use_flashinfer_allreduce(self, input_: torch.Tensor) -> bool:
        # 只有被 _tag_groups_for_flashinfer_allreduce_only 标记过的组
        # （即真正拥有某个 fusion workspace 的组）才允许走该路径。
        if self._fi_workspace_hint is None:
            return False
        from sglang.srt.layers.flashinfer_comm_fusion import (
            can_use_flashinfer_allreduce,
        )
​
        return can_use_flashinfer_allreduce(
            input_,
            use_attn_tp_group=(self._fi_workspace_hint == "attn_tp"),
            expected_world_size=self.world_size,
            expected_group_key=(self.device_group, self.cpu_group),
        )
​
    def _flashinfer_allreduce(self, input_: torch.Tensor) -> torch.Tensor:
        from sglang.srt.layers.flashinfer_comm_fusion import (
            flashinfer_allreduce as _impl,
        )
​
        return _impl(input_, use_attn_tp_group=(self._fi_workspace_hint == "attn_tp"))
​
​
def _tag_groups_for_flashinfer_allreduce_only():
    """给真正拥有 fusion workspace 的组打上 _fi_workspace_hint 标记。
​
    只有两个 workspace：attention-TP 一个、MoE 一个；MoE workspace 在
    moe_ep_size > 1 时于 EP 组上 rendezvous，否则于 MoE-TP 组上。
    因此 _MOE_EP 与 _MOE_TP 中只有一个是合法的——两个都打标会让
    MoE-TP 的 allreduce 在 hybrid EP+TP（如 tp=4, ep=2）下错误归约 EP peer。
    """
    if not _ENABLE_FLASHINFER_ALLREDUCE_ONLY:
        return
​
    moe_group = _MOE_EP if (_MOE_EP is not None and _MOE_EP.world_size > 1) else _MOE_TP
    for group, hint in ((moe_group, "moe"), (_ATTN_TP, "attn_tp")):
        if group is not None:
            group._fi_workspace_hint = hint

python/sglang/srt/layers/flashinfer_comm_fusion.py

新增 can_use_flashinfer_allreduce 与 flashinfer_allreduce 两个核心函数：前者集中所有可用性判定（rank-invariant），后者执行 AllReduceFusionPattern.kAllReduce kernel 调用且不吞异常，是整条路径的正确性基石。

# python/sglang/srt/layers/flashinfer_comm_fusion.py
# 纯 allreduce 的判定与执行分离：判定在 trace 期普通 Python 中完成，
# 执行被包裹为不透明 custom op；这样 Dynamo 不会拆开 kernel 调用，
# 而数据相关的守卫都留在 trace 期，避免重编译。
​
def can_use_flashinfer_allreduce(
    input_: torch.Tensor,
    *,
    use_attn_tp_group: bool,
    expected_world_size: int,
    expected_group_key: Tuple[Optional[ProcessGroup], Optional[ProcessGroup]],
) -> bool:
    """判定本次 all-reduce 能否走 FlashInfer kAllReduce 路径。
​
    所有检查必须保持 rank 无关（rank-invariant）：如果某个 rank 悄悄
    回退到 NCCL 而其余 rank 进入 kernel，集体通信会失配并挂起。
    """
    if _flashinfer_allreduce_unavailable or _flashinfer_comm is None:
        return False
​
    # kAllReduce 只支持 2D 连续张量（如 3D 的 vocab embedding 输出则回退）。
    if input_.ndim != 2 or not input_.is_contiguous():
        return False
​
    workspace_manager = _get_workspace_manager(use_attn_tp_group)
    if not workspace_manager.initialized or workspace_manager.workspace is None:
        return False
​
    # workspace 必须是在完全相同的 peer 集合上 rendezvous 的：
    # hybrid EP+TP 下 MoE workspace 可能对应 EP 或 MoE-TP 组，两者 world size
    # 相同但 rank 集合不同，混用会静默归约错 peer（错误输出而非崩溃）。
    if (
        workspace_manager.world_size != expected_world_size
        or workspace_manager.group != expected_group_key
    ):
        return False
​
    # 尺寸检查放在最后：token 维在 Dynamo 下是符号值，静态不可用的配置
    # 必须先短路，避免 trace 期访问 workspace 对象（与融合路径同规则）。
    token_num, hidden_dim = input_.shape
    if torch.compiler.is_compiling():
        return (
            workspace_manager.max_token_num is not None
            and workspace_manager.hidden_dim is not None
            and workspace_manager.dtype is not None
            and token_num <= workspace_manager.max_token_num
            and hidden_dim <= workspace_manager.hidden_dim
            and workspace_manager.dtype == input_.dtype
        )
​
    return workspace_manager.is_buffer_size_sufficient(
        token_num=token_num,
        hidden_dim=hidden_dim,
        dtype=input_.dtype,
    )
​
​
def flashinfer_allreduce(
    input_: torch.Tensor,
    *,
    use_attn_tp_group: bool,
) -> torch.Tensor:
    """执行 FlashInfer kAllReduce（仅 allreduce，无融合）。
​
    假定调用前 can_use_flashinfer_allreduce 已返回 True；这里刻意不捕获
    kernel 异常——吞掉异常会让本 rank 走 NCCL 而其他 rank 留在 kernel 内，
    失配挂起而不是报错。
    """
    workspace_manager = _get_workspace_manager(use_attn_tp_group)
​
    output = torch.empty_like(input_)
    kwargs = dict(
        input=input_,
        workspace=workspace_manager.workspace,
        pattern=_flashinfer_comm.AllReduceFusionPattern.kAllReduce,
        launch_with_pdl=True,
        fp32_acc=False, # 与融合路径保持一致的数值语义
        output=output,
    )
    if _flashinfer_allreduce_supports_trigger_completion:
        kwargs["trigger_completion_at_end"] = False # 与融合路径保持一致
    _flashinfer_comm.allreduce_fusion(**kwargs)
    return output

python/sglang/srt/layers/communicator.py

修复 hybrid EP+TP 下融合 post-experts all-reduce 会静默产生未完全归约激活的严重正确性问题（Qwen3-30B-A3B tp=4 ep=2 实测为垃圾输出）。

# python/sglang/srt/layers/communicator.py
# 决定 MLP 的 all-reduce 是否可以与下一层的 residual+LN 融合。
def should_fuse_mlp_allreduce_with_next_layer(
    self, forward_batch: ForwardBatch
) -> bool:
    # MOE_FULL 时无法融合：融合路径会跳过 postprocess_layer（含 moe_cp scatter），
    # 导致 hidden_states 与 residual 形状不匹配。
    if is_enable_moe_cp_allgather():
        return False
​
    # 关键修复：hybrid EP+TP（moe_ep_size > 1 且 moe_tp_size > 1）下，
    # post-experts 的归约要跨两个不相交的组（_MOE_EP 再 _MOE_TP）。
    # 融合让下一层 residual+LN 只在一个组上归约，而
    # should_skip_post_experts_all_reduce() 会把两个都跳过——
    # 结果激活只归约了一半 peer，静默错误而非崩溃。
    parallel = get_parallel()
    if parallel.moe_ep_size > 1 and parallel.moe_tp_size > 1:
        return False
​
    # 其余原有守卫（EAGLE、input_scattered、SCATTERED 模式等）保持不变。
    if (
        is_dp_attention_enabled()
        and self._speculative_algo is not None
        and self._speculative_algo.is_eagle()
    ):
        return False
​
    if get_attn_tp_context().input_scattered:
        return False
​
    batch_size = (
        forward_batch.input_ids.shape[0]
        if hasattr(forward_batch, "input_ids")
        else 0
    )
    if self.layer_scatter_modes.mlp_mode == ScatterMode.SCATTERED:
        return False
​
    return apply_flashinfer_allreduce_fusion(batch_size) or ...

评论区精华

Review 中围绕「如何安全地把 NCCL 热路径换成 FlashInfer」展开了高质量讨论：

• b8zhong 认为独立 flag 与架构白名单冗余，建议直接复用 --flashinfer-allreduce-fusion-backend 控制，作者采纳并删除 _FLASHINFER_ALLREDUCE_ONLY_ARCHS。

• mmangkad 连续指出三个正确性问题：flashinfer_allreduce() 的 try-except 回退不是跨 rank 同步的，divergent rank 会 NCCL/FlashInfer 失配挂起；fp32_acc=True 与融合路径默认 False 不一致；trigger_completion_at_end 未显式设置而继承默认 True，与融合路径 False 不一致。作者逐一修正：移除 try-except、改为 fp32_acc=False、显式传 trigger_completion_at_end=False（按 capability 门控）。

• mmangkad 指出 _MOE_TP 与 _MOE_EP 同时打标会在 hybrid EP+TP 下归约错 peer，并建议用 custom op 包装以支持 piecewise CUDA graph；作者据此实现 can_use_/execute 分离与精确打标。

• b8zhong 与 shyeh25 将本 PR 与 PR#32339（custom-allreduce 多节点升级）对比实测，二者性能相当（1P1D 约 51.4 TPS，1P4D 约 929 TPS）；结论是本 PR 对 TP16/TP32 仍有独立价值。

• ormandj 在 SM120 上实测发现本 PR fails closed（arch 校验拒绝 + 非 multicast TRT-LLM workspace 被 preflight 拒绝），需叠加两个 commit 才能启用（+5.2% median），并开 PR#32330 跟进。

• 是否保留独立的纯 allreduce 开关与架构白名单 (design): 作者删除独立 flag 与白名单，kAllReduce 随融合后端 flag 默认开启，仅保留 DSV3/V4 自动启用逻辑（DSV4 加入 _FLASHINFER_ALLREDUCE_FUSION_ARCHS）。

• fallback 路径的跨 rank 一致性 (correctness): 作者移除 try-except，改为「判定（rank-invariant）与执行（不吞异常）分离」，并强调所有判定必须保持 rank 无关。
• fp32_acc 与 trigger_completion_at_end 与融合路径不一致 (correctness): 作者修正为 fp32_acc=False，并显式传 trigger_completion_at_end=False（按 _flashinfer_allreduce_supports_trigger_completion 门控），与融合路径对齐。
• workspace 与进程组的精确绑定（hybrid EP+TP） (correctness): 作者改为只打标真正持有 workspace 的组（EP 或 MoE-TP 二选一），_fi_workspace_hint 在 __init__ 初始化；同时将 kernel 调用注册为 custom op，并在 is_compiling() 路径也接入 dispatch。
• 与 PR#32339（custom-allreduce 多节点升级）的性能对比 (performance): 两者性能 on par；本 PR 作为不依赖 NCCL ring 的替代路径保留，对超大 TP 场景有价值。
• SM120 上 fails closed (performance): 合入前未在 SM120 上启用，由 #32330 以叠加 commit 方式跟进；作者未在本 PR 内处理。
• 测试类与注释的规范性 (style): 作者已将新增测试类改为 CustomTestCase，并精简注释。

风险与影响

• 风险：
  1. 核心路径变更回归风险：GroupCoordinator.all_reduce() 是所有 TP 通信的热路径，dispatch 插入位置在既有 pynccl/custom-allreduce 之前，任何守卫的 rank 不一致都会导致集体通信失配挂起而非报错。代码通过「所有判定 rank-invariant + kernel 异常不捕获」缓解，但覆盖多拓扑（hybrid EP+TP、DP attention、PP）的端到端验证仍有限。
  2. CUDA graph 兼容性：作者在评论中承认 PCG capture 未完全端到端验证（"PCG capture not yet validated end to end"），虽然 custom op 设计上可进入 piecewise graph，但真实图捕获场景仍需回归。
  3. SM120 兼容性缺口：ormandj 实测在 SM120（RTX PRO 6000 Blackwell）上该 PR fails closed，需额外 commit；这意味着一部分 Blackwell 用户（非 GB200/GB300 MNNVL 环境）在合入后无法立即受益。
  4. 数值语义变化：fp32_acc=False 与融合路径对齐后，纯 allreduce 的数值行为与 NCCL sum 存在潜在差异；GSM8K 精度持平（0.958），但其他模型/量化组合未覆盖。
  5. hybrid EP+TP 禁用融合的影响：communicator.py 的守卫会让 hybrid EP+TP 场景失去 fuse_mlp_allreduce 优化，这是为避免静默错误的有意取舍，但可能影响该拓扑下的性能预期。
     - 影响：用户影响：DeepSeek-V3/V4 用户在没有额外参数的情况下自动获得 FlashInfer kAllReduce 加速路径；Blackwell（SM100 系）小并发场景收益显著（shyeh25 实测 GB200 TP8 单并发 TPS per GPU 从 78.77 提升到 97.12，约 +23%；BS 1-16 的 decode tok/s 提升 +2.6%~+6.9%），大并发基本持平（BS=64 约 -1.4%）；SM90 单机仍可用 TRT-LLM 后端，但不能用于多节点。系统影响：零额外显存开销，复用现有融合 workspace；与 PR#32339 的 custom-allreduce v2 多节点路径性能相当，但本 PR 不依赖 NCCL ring，在 TP16/TP32 等大 TP 场景有独立价值。团队影响：确立了「可用性判定（trace 期）与 kernel 执行（图内）分离 + 精确 group 绑定」的模式，为后续其他 FlashInfer 融合原语接入提供了可复用的设计范式。
     - 风险标记：核心路径变更, 跨 rank 一致性风险, CUDA graph 兼容性待验证, SM120 兼容性缺口, hybrid EP+TP 边界

关联脉络

• PR #32339 Upgrade custom-allreduce kernels for multi-node GB200/300（评论中提及）: b8zhong 提出其可能替代本 PR；shyeh25 实测两者性能相当，本 PR 对 TP16/TP32 仍有独立价值。
• PR #32963 Fix hybrid EP+TP post-experts all-reduce fusion（评论中提及，作者声明本 PR 基于它）: 作者在 review 回复中说明该 PR 修复的 hybrid EP+TP 归约问题被 #32963 正式修复，本 PR rebase 在其上。
• PR #32330 Enable flashinfer allreduce-only on SM120（ormandj 所开，叠加本 PR 的附加 commit）: ormandj 实测 SM120 fails closed 后提供两个修复 commit（arch 选择器与 multicast preflight 移除），并建议折回本 PR。