Integrate pplx a2a backend (#30756)

Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>

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",
+    }
 )
 
 

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):

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"

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

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

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),

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",

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

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
 
 

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")

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()

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"