[Bugfix] Size FlashInfer NVLink MNNVL workspace to EP group (#40893)

Signed-off-by: Dao Le <Dao007forever@gmail.com>
This commit is contained in:
Dao007forever
2026-04-26 01:26:34 -07:00
committed by GitHub
parent 8cd174fa35
commit 9558f43903
@@ -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