From 9558f43903faa1b6db08ac98802bf88111196345 Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Sun, 26 Apr 2026 01:26:34 -0700 Subject: [PATCH] [Bugfix] Size FlashInfer NVLink MNNVL workspace to EP group (#40893) Signed-off-by: Dao Le --- .../device_communicators/all2all.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 5ccea4d50cd..6a15d3f6168 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -492,15 +492,18 @@ class FlashInferNVLinkTwoSidedManager(All2AllManagerBase): CustomCommunicator, ) - dp_config = MnnvlConfig( - comm_backend=CustomCommunicator(get_dp_group().cpu_group), + # MNNVL workspace is allocated per rank in the comm_backend's group; the + # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend + # must span the EP group (= DP*PCP*TP), not the DP group. + ep_config = MnnvlConfig( + comm_backend=CustomCommunicator(self.cpu_group), fabric_page_size=1 << 29, # 512MB allocation_granularity=0, # Auto-detect ) - self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, dp_config) + self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, ep_config) self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace( - self.mapping, dp_config + self.mapping, ep_config ) self.world_size = world_size @@ -605,8 +608,11 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): CustomCommunicator, ) - dp_config = MnnvlConfig( - comm_backend=CustomCommunicator(get_dp_group().cpu_group), + # MNNVL workspace is allocated per rank in the comm_backend's group; the + # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend + # must span the EP group (= DP*PCP*TP), not the DP group. + ep_config = MnnvlConfig( + comm_backend=CustomCommunicator(self.cpu_group), ) total_dispatch_payload_size_per_token = ( hidden_size // 2 # nvfp4 hidden states @@ -628,7 +634,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): top_k=top_k, num_experts=num_experts, workspace_size_per_rank=self.workspace_size, - mnnvl_config=dp_config, + mnnvl_config=ep_config, ) self.gpus_per_node = gpus_per_node