[NVIDIA] Add flashinfer all-to-all MOE dispatcher (#14668)

2026-01-24 22:59:55 +08:00
parent 458a43d4ac
commit 2c2c4e446b
14 changed files with 723 additions and 16 deletions
@@ -468,6 +468,7 @@ class DeepseekV2MoE(nn.Module):
                    if get_moe_a2a_backend().is_deepep()
                    or get_moe_a2a_backend().is_mooncake()
                    or get_moe_a2a_backend().is_ascend_fuseep()
+                    or get_moe_a2a_backend().is_flashinfer()
                    or should_use_flashinfer_cutlass_moe_fp4_allgather()
                    else {}
                ),
@@ -530,6 +531,7 @@ class DeepseekV2MoE(nn.Module):
            get_moe_a2a_backend().is_deepep()
            or get_moe_a2a_backend().is_mooncake()
            or get_moe_a2a_backend().is_ascend_fuseep()
+            or get_moe_a2a_backend().is_flashinfer()
        )
        self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()