A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 15:00 UTC
[Bugfix] Size FlashInfer NVLink MNNVL workspace to EP group (#40893)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Dao007forever committed
Apr 26, 2026 at 01:26 UTC
9558f43903faa1b6db08ac98802bf88111196345
1 file changed
+13
-7
vllm/distributed/device_communicators/all2all.py
+13
-7
@@ -492,15 +492,18 @@ class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
492
CustomCommunicator,
493
)
494
495
- dp_config = MnnvlConfig(
496
- comm_backend=CustomCommunicator(get_dp_group().cpu_group),
495
+ # MNNVL workspace is allocated per rank in the comm_backend's group; the
496
+ # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
497
+ # must span the EP group (= DP*PCP*TP), not the DP group.
498
+ ep_config = MnnvlConfig(
499
+ comm_backend=CustomCommunicator(self.cpu_group),
500
fabric_page_size=1 << 29, # 512MB
501
allocation_granularity=0, # Auto-detect
502
)
503
501
- self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, dp_config)
504
+ self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, ep_config)
505
self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace(
503
- self.mapping, dp_config
506
+ self.mapping, ep_config
507
)
508
509
self.world_size = world_size
@@ -605,8 +608,11 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
608
CustomCommunicator,
609
)
610
608
- dp_config = MnnvlConfig(
609
- comm_backend=CustomCommunicator(get_dp_group().cpu_group),
611
+ # MNNVL workspace is allocated per rank in the comm_backend's group; the
612
+ # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
613
+ # must span the EP group (= DP*PCP*TP), not the DP group.
614
+ ep_config = MnnvlConfig(
615
+ comm_backend=CustomCommunicator(self.cpu_group),
616
)
617
total_dispatch_payload_size_per_token = (
618
hidden_size // 2 # nvfp4 hidden states
@@ -628,7 +634,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
634
top_k=top_k,
635
num_experts=num_experts,
636
workspace_size_per_rank=self.workspace_size,
631
- mnnvl_config=dp_config,
637
+ mnnvl_config=ep_config,
638
)
639
640
self.gpus_per_node = gpus_per_node