From a1c30701aade4dc853ff890e3fe194ffb87f4f04 Mon Sep 17 00:00:00 2001 From: Trang Do <200224632+trangdough@users.noreply.github.com> Date: Fri, 31 Jul 2026 05:33:19 +0700 Subject: [PATCH] Integrate pplx a2a backend (#30756) Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> --- .../advanced_features/expert_parallelism.mdx | 9 +- .../advanced_features/server_arguments.mdx | 4 +- .../docs/references/environment_variables.mdx | 5 + python/sglang/srt/arg_groups/overrides.py | 11 +- .../srt/batch_overlap/two_batch_overlap.py | 5 + python/sglang/srt/environ.py | 3 + python/sglang/srt/layers/dp_attention.py | 9 + python/sglang/srt/layers/moe/ep_moe/layer.py | 3 +- .../srt/layers/moe/fused_moe_triton/layer.py | 1 + .../layers/moe/token_dispatcher/__init__.py | 8 + .../srt/layers/moe/token_dispatcher/pplx.py | 527 ++++++++++++++++++ python/sglang/srt/layers/moe/utils.py | 12 +- .../sglang/srt/layers/quantization/unquant.py | 2 +- python/sglang/srt/models/deepseek_v2.py | 1 + python/sglang/srt/server_args.py | 53 +- test/manual/ep/test_pplx_small.py | 145 +++++ 16 files changed, 788 insertions(+), 10 deletions(-) create mode 100644 python/sglang/srt/layers/moe/token_dispatcher/pplx.py create mode 100644 test/manual/ep/test_pplx_small.py diff --git a/docs_new/docs/advanced_features/expert_parallelism.mdx b/docs_new/docs/advanced_features/expert_parallelism.mdx index 15d03f412..26e38cf01 100644 --- a/docs_new/docs/advanced_features/expert_parallelism.mdx +++ b/docs_new/docs/advanced_features/expert_parallelism.mdx @@ -62,12 +62,17 @@ SGLang's EP integrates diverse, highly efficient backends for different use case Ascend NPU native fused all-to-all communication. Ascend NPU deployments. + + `pplx` + pplx-kernels, Perplexity's NVSHMEM-based all-to-all dispatch/combine kernels. Low-latency (masked) only; targets FP8 (DeepGEMM) MoE models on Hopper. Requires NVSHMEM 3.2.5, nvshmem4py, cuda-python, and a prebuilt libpplx_kernels.so (sm_90a). + Low-latency decode EP on Hopper. + -DeepEP and Mooncake backends support two modes for token dispatch: `normal` mode (optimized for prefill workloads with high throughput) and `low_latency` mode (optimized for decode workloads with low latency and CUDA Graph compatibility). MORI backend only supports `normal` mode now. NIXL-EP currently operates in low-latency mode with CUDA Graph support. Users are recommended to set `--deepep-mode auto` to enable automatic dispatch mode switching during runtime. Setting `--deepep-mode normal` or `--deepep-mode low_latency` is useful for debugging or development purposes. +DeepEP and Mooncake backends support two modes for token dispatch: `normal` mode (optimized for prefill workloads with high throughput) and `low_latency` mode (optimized for decode workloads with low latency and CUDA Graph compatibility). MORI backend only supports `normal` mode now. NIXL-EP and PPLX currently operate in low-latency mode with CUDA Graph support (PPLX reuses the DeepEP low-latency masked expert-compute path). Users are recommended to set `--deepep-mode auto` to enable automatic dispatch mode switching during runtime. Setting `--deepep-mode normal` or `--deepep-mode low_latency` is useful for debugging or development purposes. -Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support cases where `ep_size = tp_size`. For hybrid EP and TP (i.e., `ep_size < tp_size`), only the `none` backend (All-Reduce or All-Gather-based dispatching) is supported. +Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep`, `pplx` and MORI only support cases where `ep_size = tp_size`. For hybrid EP and TP (i.e., `ep_size < tp_size`), only the `none` backend (All-Reduce or All-Gather-based dispatching) is supported. Note that `pplx` additionally requires `--enable-dp-attention` with at least 2 DP groups (i.e., `tp_size / attention_tp_size > 1`); otherwise pplx-kernels' AllToAll cannot be constructed. ### Backends for MoE Computation diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index b54b850f5..7117a7173 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1578,7 +1578,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--speculative-moe-a2a-backend` MOE A2A backend for EAGLE speculative decoding, see `--moe-a2a-backend` for options. Same as moe a2a backend if unset. `None` - none, deepep, mooncake, nixl, mori, ascend_fuseep, flashinfer, megamoe + none, deepep, mooncake, nixl, mori, ascend_fuseep, flashinfer, megamoe, pplx `--speculative-draft-model-quantization` @@ -1746,7 +1746,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--moe-a2a-backend` Select the backend for all-to-all communication for expert parallelism. `none` - none, deepep, mooncake, nixl, mori, ascend_fuseep, flashinfer, megamoe + none, deepep, mooncake, nixl, mori, ascend_fuseep, flashinfer, megamoe, pplx `--moe-runner-backend` diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 67750563c..6fe524634 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -475,6 +475,11 @@ SGLang supports various environment variables that can be used to configure its Truncate the dispatch buffer to this many rows before MoE computation, reducing kernel overhead on padding tokens. The value must be >= the actual number of received tokens (totalRecvTokenNum); setting it too small causes incorrect results. 0 disables truncation (use full buffer). 0 + + SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK + Maximum number of dispatched tokens per rank for the PPLX (--moe-a2a-backend pplx) NVSHMEM buffer allocation. Must be ≥ the worst-case per-rank forward tokens (the per-rank prefill chunk chunked_prefill_size / dp_size, or the decode CUDA-graph batch size); the server refuses to start otherwise. + 128 + SGLANG_MORI_QP_PER_TRANSFER Number of RDMA Queue Pairs (QPs) used per transfer operation diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 95998c34b..3d7e40662 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -2151,7 +2151,16 @@ def _cutlass_moe_env_override(view: Any) -> dict: # Every A2A backend that forces expert parallelism to span the TP group. _A2A_EP_SPANNING_BACKENDS = frozenset( - {"megamoe", "deepep", "mooncake", "nixl", "ascend_fuseep", "flashinfer", "mori"} + { + "megamoe", + "deepep", + "mooncake", + "nixl", + "ascend_fuseep", + "flashinfer", + "mori", + "pplx", + } ) diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 326d313df..2d447081c 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -30,6 +30,7 @@ from sglang.srt.layers.moe.token_dispatcher import ( MooncakeEPDispatcher, MoriEPDispatcher, NixlEPDispatcher, + PplxDispatcher, ) from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher from sglang.srt.managers.schedule_batch import ScheduleBatch @@ -1091,6 +1092,10 @@ class MaybeTboDeepEPDispatcher(BaseDispatcher): self._inners = [ NixlEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) ] + elif get_moe_a2a_backend().is_pplx(): + self._inners = [ + PplxDispatcher(**kwargs) for _ in range(num_inner_dispatchers) + ] @property def expert_mask_gpu(self): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 901a6dbf5..0787565a5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -737,6 +737,9 @@ class Envs: SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False) SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # PPLX-EP (Perplexity pplx-kernels NVSHMEM all-to-all) + SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # DSA Backend (canonical names; fall back to SGLANG_NSA_* with deprecation warning) SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias( True, deprecated_name="SGLANG_NSA_FUSE_TOPK" diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index fe51cca90..c1ea7893f 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -88,6 +88,15 @@ class DpPaddingMode(IntEnum): ) -> DpPaddingMode: dp_size = get_attention_dp_size() + # (trangdough) pplx-kernels a2a is a symmetric collective: every EP rank + # must dispatch the same number of tokens or the device-side handshake + # deadlocks (idle DP ranks with 0 tokens never signal their peers). + # Force MAX_LEN so all ranks are padded to equal token counts. + from sglang.srt.layers.moe.utils import get_moe_a2a_backend + + if get_moe_a2a_backend().is_pplx(): + return DpPaddingMode.MAX_LEN + # When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding # overhead from uneven token distribution. # For dp_size=1, max_len equals sum_len, so prefer MAX_LEN mode diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 1d1f40c77..13dba0388 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -114,7 +114,7 @@ class DeepEPMoE(FusedMoE): quant_config is None and self.w13_weight.dtype == torch.bfloat16 and get_moe_runner_backend().is_deep_gemm() - and get_moe_a2a_backend().is_deepep() + and (get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_pplx()) and not _is_npu and not _is_hip ): @@ -282,6 +282,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): or get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() + or get_moe_a2a_backend().is_pplx() ): return DeepEPMoE return FusedMoE diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index ee756e992..1160b8637 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -123,6 +123,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: or a2a_backend.is_mooncake() or a2a_backend.is_mori() or a2a_backend.is_nixl() + or a2a_backend.is_pplx() ): return MaybeTboDeepEPDispatcher( group=_get_deepep_comm_group(a2a_backend), diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py index b605f74c0..7f2c0942f 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -42,6 +42,11 @@ from sglang.srt.layers.moe.token_dispatcher.nixl import ( NixlEPDispatcher, NixlEPDispatchOutput, ) +from sglang.srt.layers.moe.token_dispatcher.pplx import ( + PplxCombineInput, + PplxDispatcher, + PplxDispatchOutput, +) from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardCombineInput, StandardDispatcher, @@ -70,6 +75,9 @@ __all__ = [ "NixlEPCombineInput", "NixlEPDispatchOutput", "NixlEPDispatcher", + "PplxCombineInput", + "PplxDispatchOutput", + "PplxDispatcher", "StandardDispatcher", "StandardDispatchOutput", "StandardCombineInput", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/pplx.py b/python/sglang/srt/layers/moe/token_dispatcher/pplx.py new file mode 100644 index 000000000..ecb0d8194 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/pplx.py @@ -0,0 +1,527 @@ +from __future__ import annotations + +from enum import Enum, auto +from typing import NamedTuple, Optional, Tuple + +import torch +import torch.distributed as dist + +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.dp_attention import get_is_extend_in_batch +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.utils import ( + DeepEPMode, + DispatcherOutputDtype, + get_deepep_output_dtype, +) +from sglang.srt.runtime_context import get_parallel, get_server_args + +# Block size used by pplx-kernels for FP8 block-wise scales, matching the +# DeepSeek / DeepGEMM block quantization convention. +_FP8_BLOCK_SIZE = 128 + +try: + from pplx_kernels import AllToAll, nvshmem_init + from pplx_kernels.nvshmem import PyTorchStreamWrapper # noqa: F401 + + use_pplx = True +except ImportError: + use_pplx = False + + +class PplxDispatchOutput(NamedTuple): + """PPLX EP dispatch output (masked / per-expert batched).""" + + hidden_states: torch.Tensor + hidden_states_scale: Optional[torch.Tensor] + topk_ids: torch.Tensor + topk_weights: torch.Tensor + masked_m: torch.Tensor + expected_m: int + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.DEEPEP_LL + + +assert isinstance(PplxDispatchOutput, DispatchOutput) + + +class PplxCombineInput(NamedTuple): + """PPLX EP combine input.""" + + hidden_states: torch.Tensor + topk_ids: torch.Tensor + topk_weights: torch.Tensor + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.DEEPEP_LL + + +assert isinstance(PplxCombineInput, CombineInput) + + +class PplxAllToAllManager: + _nvshmem_initialized = False + _all_to_all: Optional[AllToAll] = None + _key: Optional[tuple] = None + _group_name: Optional[str] = None + + # Name under which the EP process group is registered with c10d so the + # pplx intranode kernel can resolve it via resolve_process_group(). + _GROUP_NAME = "pplx_ep" + + @classmethod + def _ensure_nvshmem(cls, group: dist.ProcessGroup) -> None: + if cls._nvshmem_initialized: + return + + assert group.size() == dist.get_world_size(), ( + "moe_a2a_backend='pplx' requires the EP group to span the whole " + f"world (got ep_size={group.size()}, world_size={dist.get_world_size()}); " + "pipeline parallelism and EP-subset layouts are not supported." + ) + + global_rank = group.rank() + world_size = group.size() + device = torch.device("cuda", torch.cuda.current_device()) + local_rank = torch.cuda.current_device() + nvshmem_init( + global_rank=global_rank, + local_rank=local_rank, + world_size=world_size, + device=device, + ) + cls._nvshmem_initialized = True + + @classmethod + def _register_group(cls, group: dist.ProcessGroup) -> str: + if cls._group_name is not None: + return cls._group_name + ranks = dist.get_process_group_ranks(group) + combined = dist.new_group(ranks=ranks, backend="cpu:gloo,cuda:nccl") + torch._C._distributed_c10d._register_process_group(cls._GROUP_NAME, combined) + cls._group_name = cls._GROUP_NAME + return cls._group_name + + @classmethod + def get_all_to_all( + cls, + group: dist.ProcessGroup, + max_num_tokens: int, + num_experts: int, + experts_per_token: int, + hidden_dim: int, + hidden_dim_bytes: int, + hidden_dim_scale_bytes: int, + ) -> AllToAll: + world_size = group.size() + rank = group.rank() + # pplx dpSize == number of ranks per DP group == attention TP size. + # numDPGroups == worldSize / dpSize == attention DP size (must be > 1). + dp_size = get_parallel().attn_tp_size + + key = ( + max_num_tokens, + num_experts, + experts_per_token, + hidden_dim, + hidden_dim_bytes, + hidden_dim_scale_bytes, + world_size, + dp_size, + ) + if cls._all_to_all is not None: + assert cls._key == key, ( + "PplxAllToAllManager already initialized with a different " + f"configuration: {cls._key} != {key}" + ) + return cls._all_to_all + + cls._ensure_nvshmem(group) + + # Use the single-node NVLink path when the EP group fits on one node, + # otherwise the NVSHMEM internode path. + + # pplx forces ep_size == world_size + # with pp_size == 1 (enforced in _ensure_nvshmem), so the EP group spans + # a single node iff the whole job runs on one node. + is_internode = get_server_args().nnodes > 1 + + if is_internode: + cls._all_to_all = AllToAll.internode( + max_num_tokens=max_num_tokens, + num_experts=num_experts, + experts_per_token=experts_per_token, + rank=rank, + world_size=world_size, + dp_size=dp_size, + hidden_dim=hidden_dim, + hidden_dim_bytes=hidden_dim_bytes, + hidden_dim_scale_bytes=hidden_dim_scale_bytes, + ) + else: + group_name = cls._register_group(group) + cls._all_to_all = AllToAll.intranode( + max_num_tokens=max_num_tokens, + num_experts=num_experts, + experts_per_token=experts_per_token, + rank=rank, + world_size=world_size, + dp_size=dp_size, + hidden_dim=hidden_dim, + hidden_dim_bytes=hidden_dim_bytes, + hidden_dim_scale_bytes=hidden_dim_scale_bytes, + group_name=group_name, + ) + cls._key = key + return cls._all_to_all + + +class _PplxDispatcherImpl: + def __init__( + self, + group: torch.distributed.ProcessGroup, + router_topk: int, + permute_fusion: bool, + num_experts: int, + num_local_experts: int, + hidden_size: int, + params_dtype: torch.dtype, + deepep_mode: DeepEPMode, + ): + if not use_pplx: + raise ImportError( + "pplx-kernels is not installed. Please build and install it " + "from https://github.com/perplexityai/pplx-kernels (e.g. " + "`TORCH_CUDA_ARCH_LIST=9.0a+PTX python3 setup.py bdist_wheel && " + "pip install dist/*.whl`) to run SGLang with the pplx MoE A2A " + "backend." + ) + + self.group = group + self.router_topk = router_topk + self.permute_fusion = permute_fusion + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.params_dtype = params_dtype + self.params_bytes = torch.tensor([], dtype=params_dtype).element_size() + self.deepep_mode = deepep_mode + + self.num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + + self.quant_config: dict = {} + self.use_fp8 = False + self.set_dispatch_dtype() + + def set_dispatch_dtype(self) -> None: + output_dtype = get_deepep_output_dtype(self) + if output_dtype == DispatcherOutputDtype.BF16: + self.use_fp8 = False + elif output_dtype == DispatcherOutputDtype.FP8: + self.use_fp8 = True + else: + raise NotImplementedError( + f"pplx MoE A2A backend does not support dispatch dtype " + f"{output_dtype}; use bf16 or fp8." + ) + + def _hidden_dim_scale_bytes(self) -> int: + if not self.use_fp8: + return 0 + return ( + (self.hidden_size + _FP8_BLOCK_SIZE - 1) + // _FP8_BLOCK_SIZE + * torch.float32.itemsize + ) + + def _get_all_to_all(self) -> AllToAll: + itemsize = 1 if self.use_fp8 else self.params_bytes + return PplxAllToAllManager.get_all_to_all( + group=self.group, + max_num_tokens=self.num_max_dispatch_tokens_per_rank, + num_experts=self.num_experts, + experts_per_token=self.router_topk, + hidden_dim=self.hidden_size, + hidden_dim_bytes=self.hidden_size * itemsize, + hidden_dim_scale_bytes=self._hidden_dim_scale_bytes(), + ) + + def _quantize( + self, hidden_states: torch.Tensor + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Return (dp_x, dp_x_scale) matching the pplx dispatch contract.""" + if not self.use_fp8: + return hidden_states, None + + from sglang.srt.layers.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + + x_q, x_s = sglang_per_token_group_quant_fp8( + hidden_states, + group_size=_FP8_BLOCK_SIZE, + ) + # pplx expects float32 scales. + return x_q, x_s.to(torch.float32) + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ): + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids + ata = self._get_all_to_all() + + num_tokens = hidden_states.shape[0] + assert num_tokens <= self.num_max_dispatch_tokens_per_rank, ( + f"num_tokens ({num_tokens}) exceeds num_max_dispatch_tokens_per_rank " + f"({self.num_max_dispatch_tokens_per_rank}); raise " + f"SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK or lower the per-rank " + f"decode batch / chunked-prefill size." + ) + num_dp_groups = get_parallel().attn_dp_size + max_batch_tokens = self.num_max_dispatch_tokens_per_rank * num_dp_groups + device = hidden_states.device + + dp_x, dp_x_scale = self._quantize(hidden_states) + out_expert_num_tokens = torch.zeros( + self.num_local_experts, dtype=torch.int32, device=device + ) + out_expert_x = torch.zeros( + (self.num_local_experts, max_batch_tokens, self.hidden_size), + dtype=dp_x.dtype, + device=device, + ) + out_expert_x_scale = None + if self.use_fp8: + scale_dim = self._hidden_dim_scale_bytes() // torch.float32.itemsize + # Zero-init like out_expert_x: padding scale rows beyond masked_m + # must not feed uninitialized floats into FP8 dequant (-> NaNs). + out_expert_x_scale = torch.zeros( + (self.num_local_experts, max_batch_tokens, scale_dim), + dtype=torch.float32, + device=device, + ) + + bound_m = torch.full((1,), num_tokens, dtype=torch.uint32, device=device) + indices = topk_ids.to(torch.uint32) + + ata.dispatch( + out_expert_num_tokens=out_expert_num_tokens, + out_expert_x=out_expert_x, + out_expert_x_scale=out_expert_x_scale, + dp_x=dp_x, + dp_x_scale=dp_x_scale, + indices=indices, + bound_m=bound_m, + ) + + expected_m = ( + num_tokens * num_dp_groups * self.router_topk + self.num_experts + ) // self.num_experts + + return ( + out_expert_x, + out_expert_x_scale, + topk_ids, + topk_weights, + out_expert_num_tokens, + expected_m, + ) + + def dispatch_b( + self, + out_expert_x, + out_expert_x_scale, + topk_ids, + topk_weights, + out_expert_num_tokens, + expected_m, + ): + get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency( + out_expert_num_tokens + ) + return PplxDispatchOutput( + out_expert_x, + out_expert_x_scale, + topk_ids, + topk_weights, + out_expert_num_tokens, + expected_m, + ) + + def combine_a( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ): + ata = self._get_all_to_all() + num_tokens = topk_ids.shape[0] + assert num_tokens <= self.num_max_dispatch_tokens_per_rank, ( + f"num_tokens ({num_tokens}) exceeds num_max_dispatch_tokens_per_rank " + f"({self.num_max_dispatch_tokens_per_rank}); raise " + f"SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK or lower the per-rank " + f"decode batch / chunked-prefill size." + ) + device = topk_ids.device + + out_tokens = torch.zeros( + (self.num_max_dispatch_tokens_per_rank, self.hidden_size), + dtype=self.params_dtype, + device=device, + ) + + bound_m = torch.full((1,), num_tokens, dtype=torch.uint32, device=device) + + ata.combine( + out_tokens=out_tokens, + indices=topk_ids.to(torch.uint32), + weights=topk_weights.to(torch.float32), + expert_y=hidden_states, + bound_m=bound_m, + ) + return (out_tokens[:num_tokens],) + + def combine_b(self, hidden_states): + return hidden_states + + def set_quant_config(self, quant_config: dict) -> None: + self.quant_config = quant_config + self.set_dispatch_dtype() + self._get_all_to_all() + + +class _Stage(Enum): + INITIAL = auto() + AFTER_DISPATCH_A = auto() + AFTER_DISPATCH_B = auto() + AFTER_COMBINE_A = auto() + + +class PplxDispatcher(BaseDispatcher): + """MoE all-to-all dispatcher backed by Perplexity's pplx-kernels. + + Reuse the DEEPEP_LL dispatch/combine format so the existing masked + expert-compute path is unchanged. + """ + + def __init__( + self, + group: torch.distributed.ProcessGroup, + router_topk: int, + permute_fusion: bool = False, + num_experts: int = None, + num_local_experts: int = None, + hidden_size: int = None, + params_dtype: torch.dtype = None, + deepep_mode: DeepEPMode = DeepEPMode.AUTO, + async_finish: bool = False, + return_recv_hook: bool = False, + ): + super().__init__() + + self.deepep_mode = deepep_mode + + if self.deepep_mode.enable_normal(): + raise NotImplementedError( + "pplx MoE A2A backend supports low-latency mode only." + ) + + self._low_latency_dispatcher = _PplxDispatcherImpl( + group=group, + router_topk=router_topk, + permute_fusion=permute_fusion, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + params_dtype=params_dtype, + deepep_mode=deepep_mode, + ) + + self._stage = _Stage.INITIAL + + def dispatch( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ) -> DispatchOutput: + self.dispatch_a(hidden_states, topk_output) + return self.dispatch_b() + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ): + self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A) + inner_state = self._get_impl().dispatch_a( + hidden_states=hidden_states, + topk_output=topk_output, + ) + self._dispatch_intermediate_state = inner_state + + def dispatch_b(self): + self._update_stage(_Stage.AFTER_DISPATCH_A, _Stage.AFTER_DISPATCH_B) + inner_state = self._dispatch_intermediate_state + del self._dispatch_intermediate_state + return self._get_impl().dispatch_b(*inner_state) + + def combine( + self, + combine_input: CombineInput, + ) -> torch.Tensor: + self.combine_a(combine_input) + return self.combine_b() + + def combine_a( + self, + combine_input: CombineInput, + ): + hidden_states, topk_ids, topk_weights = combine_input + self._update_stage(_Stage.AFTER_DISPATCH_B, _Stage.AFTER_COMBINE_A) + inner_state = self._get_impl().combine_a( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + self._combine_intermediate_state = inner_state + + def combine_b(self): + self._update_stage(_Stage.AFTER_COMBINE_A, _Stage.INITIAL) + inner_state = self._combine_intermediate_state + del self._combine_intermediate_state + return self._get_impl().combine_b(*inner_state) + + def set_quant_config(self, quant_config: dict) -> None: + self.quant_config = quant_config + self._low_latency_dispatcher.set_quant_config(quant_config) + + def _get_impl(self) -> _PplxDispatcherImpl: + is_extend_in_batch = get_is_extend_in_batch() + resolved_deepep_mode = self.deepep_mode.resolve(is_extend_in_batch) + if resolved_deepep_mode == DeepEPMode.NORMAL: + raise NotImplementedError( + "pplx MoE A2A backend supports low-latency mode only." + ) + elif resolved_deepep_mode == DeepEPMode.LOW_LATENCY: + return self._low_latency_dispatcher + else: + raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}") + + def _update_stage(self, old_stage, new_stage): + assert self._stage == old_stage + self._stage = new_stage diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 424b685dc..09bc8834f 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -36,6 +36,7 @@ class MoeA2ABackend(Enum): ASCEND_TP = "ascend_tp" FLASHINFER = "flashinfer" MEGAMOE = "megamoe" + PPLX = "pplx" CUSTOMIZED = "customized" @classmethod @@ -74,6 +75,9 @@ class MoeA2ABackend(Enum): def is_megamoe(self): return self == MoeA2ABackend.MEGAMOE + def is_pplx(self): + return self == MoeA2ABackend.PPLX + def is_customized(self): return self == MoeA2ABackend.CUSTOMIZED @@ -384,9 +388,9 @@ def is_sbo_enabled() -> bool: def is_deepep_class_backend() -> bool: - """Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, or Mori).""" + """Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, Mori, or PPLX).""" b = get_moe_a2a_backend() - return b.is_deepep() or b.is_mooncake() or b.is_mori() + return b.is_deepep() or b.is_mooncake() or b.is_mori() or b.is_pplx() def uses_per_rank_fused_shared_slots() -> bool: @@ -510,6 +514,10 @@ def should_skip_post_experts_all_reduce(*, is_tp_path: bool) -> bool: return True if get_moe_a2a_backend().is_flashinfer(): return True + if get_moe_a2a_backend().is_pplx(): + # pplx's AllToAll.combine already sums each token's expert outputs back + # to the source rank + return True return False diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 1834b038c..f32d68f17 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -365,7 +365,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): if ( self.use_deep_gemm and layer.w13_weight.dtype == torch.bfloat16 - and get_moe_a2a_backend().is_deepep() + and (get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_pplx()) and not _is_npu and not _is_hip and hasattr(layer, "dispatcher") diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index e00c411df..e6ffbc034 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -723,6 +723,7 @@ class DeepseekV2MoE(nn.Module): # not divisible by the global TP size. _shared_expert_use_tp1 = ( get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_pplx() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_mori() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1182ccbb2..097cb55db 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -272,6 +272,7 @@ MOE_A2A_BACKEND_CHOICES = [ "ascend_fuseep", "flashinfer", "megamoe", + "pplx", "ascend_tp", ] @@ -2253,7 +2254,7 @@ class ServerArgs: "ascend_fuseep", "flashinfer", "megamoe", - "ascend_tp", + "pplx", ], Arg( help="Choose the backend for MoE A2A.", @@ -6669,10 +6670,60 @@ class ServerArgs: "(chunked_prefill_size by default)" ) + if a2a_backend == "pplx": + if self.deepep_mode == "normal": + raise ValueError( + "moe_a2a_backend='pplx' only supports low-latency mode; " + "set --deepep-mode to 'low_latency' or 'auto'." + ) + if self.deepep_mode == "auto": + self.deepep_mode = "low_latency" + logger.warning("auto set deepep_mode=`low_latency` for PPLX EP") + # pplx-kernels' AllToAll needs numDPGroups (== attention dp_size) > 1; + # without DP attention numDPGroups == 1 and construction fails deep in + # the kernel. This also implies ep_size >= 2. + assert resolved_view(self).enable_dp_attention and self.dp_size >= 2, ( + "moe_a2a_backend='pplx' requires --enable-dp-attention with at " + "least 2 DP groups (--dp-size >= 2)." + ) + # pplx runs the masked DeepGEMM expert path (sm_90a): reject other + # runners and resolve auto -> deep_gemm. Unquantized bf16 pplx needs + # an explicit deep_gemm backend, otherwise the expert layer falls + # through to the deprecated masked path and asserts at runtime. + assert resolved_view(self).moe_runner_backend in ("deep_gemm", "auto"), ( + "moe_a2a_backend='pplx' is only supported with --moe-runner-backend " + "deep_gemm (or auto)." + ) + if self.moe_runner_backend == "auto": + self.moe_runner_backend = "deep_gemm" + logger.warning("auto set moe_runner_backend=`deep_gemm` for PPLX EP") + logger.warning( + f"PPLX MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." + ) + + # Check per-rank dispatch tokens for pplx + # Skip validation if chunked prefill is disabled (i.e., size <= 0) + # Skip validation if disaggregation mode is decode + if self.chunked_prefill_size > 0 and self.disaggregation_mode != "decode": + assert ( + self._required_pplx_dispatch_tokens_per_rank() + ) <= envs.SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get(), ( + "SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK (default 128) " + "must be >= the per-rank pplx dispatch tokens " + "(chunked_prefill_size, or the decode cuda-graph batch size)" + ) + def _required_mori_dispatch_tokens_per_rank(self) -> int: """Max tokens a single rank dispatches through MoRI in one forward.""" return self.chunked_prefill_size + def _required_pplx_dispatch_tokens_per_rank(self) -> int: + """Max tokens a single rank dispatches through pplx in one forward.""" + required = self.chunked_prefill_size + if self.cuda_graph_max_bs_decode is not None: + required = max(required, self.cuda_graph_max_bs_decode) + return required + def _handle_eplb_and_dispatch(self): if self.enable_eplb and (self.expert_distribution_recorder_mode is None): self.expert_distribution_recorder_mode = "stat" diff --git a/test/manual/ep/test_pplx_small.py b/test/manual/ep/test_pplx_small.py new file mode 100644 index 000000000..382103a64 --- /dev/null +++ b/test/manual/ep/test_pplx_small.py @@ -0,0 +1,145 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_MLA, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +# Manual test: pplx-kernels are not available in the CI environment, so this is +# not CI-registered. Run locally on a 4x H100 node with pplx-kernels installed. + + +class TestPureDP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp", + "4", + "--enable-dp-attention", + "--dp", + "4", + "--moe-a2a-backend", + "pplx", + "--deepep-mode", + "low_latency", + "--cuda-graph-max-bs-decode", + "128", + "--max-running-requests", + "512", + "--mem-fraction-static", + "0.5", + ], + # Per-rank dispatch cap must cover the per-rank prefill chunk + # (chunked_prefill_size // dp_size = 8192 // 4 = 2048 on H100). + env={"SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "4096"}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.60) + + def test_gsm8k_single_stream(self): + # Regression guard for the post-experts all-reduce double-count: with + # num_threads=1 only one DP rank has a real request at a time, leaving + # the others idle. If pplx does not skip the post-experts all-reduce, + # those idle ranks' outputs corrupt the answer and the score collapses + # (~0). Keep this serial + low example count so it stays cheap. + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=40, + num_threads=1, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.50) + + +class TestHybridDPTP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp", + "4", + "--enable-dp-attention", + "--dp", + "2", + "--moe-a2a-backend", + "pplx", + "--deepep-mode", + "low_latency", + "--cuda-graph-max-bs-decode", + "128", + "--max-running-requests", + "256", + "--mem-fraction-static", + "0.5", + ], + # Per-rank dispatch cap must cover the per-rank prefill chunk + # (chunked_prefill_size // dp_size = 8192 // 2 = 4096 on H100). + env={"SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "4096"}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.60) + + +if __name__ == "__main__": + unittest.main()