[1/N] Introduce Mooncake Backend and Mooncake EP to Support Elastic EP (#10423)

Co-authored-by: Hank Han <hanhan7630@outlook.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Xun Sun
2025-10-15 10:40:54 +08:00
committed by GitHub
parent 74737b2863
commit a40229f6f8
13 changed files with 798 additions and 32 deletions

View File

@@ -228,6 +228,8 @@ class ServerArgs:
# Runtime options
device: Optional[str] = None
elastic_ep_backend: Literal[None, "mooncake"] = None
mooncake_ib_device: Optional[str] = None
tp_size: int = 1
pp_size: int = 1
pp_max_micro_batch_size: Optional[int] = None
@@ -344,7 +346,7 @@ class ServerArgs:
# Expert parallelism
ep_size: int = 1
moe_a2a_backend: Literal["none", "deepep"] = "none"
moe_a2a_backend: Literal["none", "deepep", "mooncake"] = "none"
moe_runner_backend: str = "auto"
flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default"
enable_flashinfer_allreduce_fusion: bool = False
@@ -537,7 +539,7 @@ class ServerArgs:
# Handle MoE configurations.
self._handle_moe_kernel_config()
self._handle_deepep_moe()
self._handle_a2a_moe()
self._handle_eplb_and_dispatch()
self._handle_expert_distribution_metrics()
@@ -1091,7 +1093,7 @@ class ServerArgs:
"FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set."
)
def _handle_deepep_moe(self):
def _handle_a2a_moe(self):
if self.moe_a2a_backend == "deepep":
if self.deepep_mode == "normal":
logger.warning("Cuda graph is disabled because deepep_mode=`normal`")
@@ -1101,6 +1103,12 @@ class ServerArgs:
f"DeepEP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
)
if self.moe_a2a_backend == "mooncake":
self.ep_size = self.tp_size
logger.warning(
f"Mooncake MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
)
def _handle_eplb_and_dispatch(self):
if self.enable_eplb and (self.expert_distribution_recorder_mode is None):
self.expert_distribution_recorder_mode = "stat"
@@ -1712,6 +1720,21 @@ class ServerArgs:
default=ServerArgs.device,
help="The device to use ('cuda', 'xpu', 'hpu', 'npu', 'cpu'). Defaults to auto-detection if not specified.",
)
parser.add_argument(
"--elastic-ep-backend",
type=str,
default=ServerArgs.elastic_ep_backend,
choices=["none", "mooncake"],
help="Specify the collective communication backend for elastic EP. Currently supports 'mooncake'.",
)
parser.add_argument(
"--mooncake-ib-device",
type=str,
default=ServerArgs.mooncake_ib_device,
help="The InfiniBand devices for Mooncake Backend transfer, accepts multiple comma-separated devices "
"(e.g., --mooncake-ib-device mlx5_0,mlx5_1). "
"Default is None, which triggers automatic device detection when Mooncake Backend is enabled.",
)
parser.add_argument(
"--tensor-parallel-size",
"--tp-size",
@@ -2333,7 +2356,7 @@ class ServerArgs:
parser.add_argument(
"--moe-a2a-backend",
type=str,
choices=["none", "deepep"],
choices=["none", "deepep", "mooncake"],
default=ServerArgs.moe_a2a_backend,
help="Choose the backend for MoE A2A.",
)