[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:
@@ -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.",
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user