feat: support data parallel (DP) for DeepSeek-V3.2 Python model executor.

Author: yingxudeng

xllm/models/llm/py_causal_lm.cpp

@@ -87,8 +87,6 @@ PyCausalLM::PyCausalLM(const ModelContext& context)
   dp_size_ = (dp_group != nullptr) ? dp_group->world_size() : 1;
   dp_rank_ = (dp_group != nullptr) ? dp_group->rank() : 0;
   ep_size_ = parallel_args.ep_size();
-  CHECK(ep_size_ == 1 || ep_size_ == parallel_args.world_size())
-      << "Python models support only ep_size=1 or ep_size=world_size.";
 
   CHECK(parallel_args.moe_tp_group_ != nullptr);
   ProcessGroup* moe_tp_group = parallel_args.moe_tp_group_;

xllm/python/distributed/__init__.py

@@ -20,6 +20,7 @@
     all_gather,
     all_gather_variable,
     all_reduce_,
+    barrier,
     cp_rank,
     cp_world_size,
     init_process_group,
@@ -36,4 +37,5 @@
     "all_reduce_",
     "all_gather",
     "all_gather_variable",
+    "barrier",
 ]

xllm/python/distributed/collectives.py

@@ -383,6 +383,24 @@ def _(
     return x.new_empty(shape)
 
 
+def barrier(device: torch.device | str, group_name: str = "dp") -> None:
+    """Block until all ranks in the group have reached this point.
+
+    Drains the local device queue first (synchronize), then coordinates
+    across ranks via dist.barrier. Unlike an all_reduce approach, no
+    data-plane traffic is injected into the work queue — this matters
+    before NPUGraph capture which expects a clean queue.
+    """
+    device_obj = torch.device(device) if isinstance(device, str) else device
+    group = _groups.get((group_name, str(device_obj)))
+    if group is None or group.size() <= 1:
+        return
+    device_module = getattr(torch, device_obj.type, None)
+    if device_module is not None and hasattr(device_module, "synchronize"):
+        device_module.synchronize(device_obj)
+    dist.barrier(group=group)
+
+
 __all__ = [
     "init_process_group",
     "init_tp_group",
@@ -392,4 +410,5 @@ def _(
     "all_reduce_",
     "all_gather",
     "all_gather_variable",
+    "barrier",
 ]

xllm/python/layers/fused_moe.py

@@ -103,18 +103,18 @@ def __init__(
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        local_hidden_states = hidden_states
-        token_counts: list[int] | None = None
+        local_tokens: int = 0
+        padded_tokens: int = 0
         if self.dp_size > 1:
             token_counts = list(get_forward_context().metadata.dp_token_counts)
             if len(token_counts) != self.dp_size:
                 raise RuntimeError(f"expected {self.dp_size} DP token counts, got {token_counts}")
-            hidden_states = distributed.all_gather_variable(
-                hidden_states,
-                token_counts,
-                self.dp_rank,
-                "dp",
-            )
+            padded_tokens = max(token_counts)
+            local_tokens = hidden_states.shape[0]
+            pad_size = padded_tokens - local_tokens
+            if pad_size > 0:
+                hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, pad_size))
+            hidden_states = distributed.all_gather(hidden_states, dim=0, world_size=self.dp_size, group_name="dp")
 
         router_logits = self.gate(hidden_states)
         topk_weights, topk_ids = kernels.moe_fused_topk(
@@ -150,15 +150,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
                 distributed.all_reduce_(output, "moe_tp")
             if self.ep_size > 1:
                 distributed.all_reduce_(output, "moe_ep")
-        if token_counts is not None:
-            local_tokens = token_counts[self.dp_rank]
-            if local_tokens == 0:
-                return torch.zeros_like(local_hidden_states)
-            start = sum(token_counts[: self.dp_rank])
-            local_output = output.narrow(0, start, local_tokens)
-            if local_tokens == local_hidden_states.shape[0]:
-                return local_output
-            padding_shape = list(local_hidden_states.shape)
-            padding_shape[0] -= local_tokens
-            return torch.cat([local_output, local_hidden_states.new_zeros(padding_shape)], dim=0)
+        if padded_tokens > 0:
+            start = self.dp_rank * padded_tokens
+            output = output.narrow(0, start, local_tokens)
         return output

xllm/python/model_executor/runners/decode_acl_graph.py

@@ -35,7 +35,7 @@
 import torch
 import torch.nn as nn
 
-from xllm.python import kernels
+from xllm.python import distributed, kernels
 from xllm.python.attention.backend import AttentionBackend, AttentionMetadata
 from xllm.python.attention.expanded_decode_metadata import (
     ExpandedDecodeMetadata,
@@ -138,14 +138,20 @@ def can_execute(
         if input_ids.dim() != 1:
             return False
         batch_size = input_ids.numel()
-        bucket_size = _decode_bucket(batch_size)
         is_expanded_spec_verify = resolve_expanded_decode_metadata(metadata) is not None
-        return (
+        if not (
             ((not metadata.is_prefill and not metadata.is_chunked_prefill) or is_expanded_spec_verify)
             and self._has_compatible_decode_metadata(input_ids, metadata)
             and (input_embedding is None or input_embedding.shape[0] == batch_size)
-            and bucket_size <= self.max_batch
-        )
+        ):
+            return False
+        if self.dp_size > 1:
+            dp_token_counts = getattr(metadata, "dp_token_counts", None)
+            if dp_token_counts is None or len(dp_token_counts) != self.dp_size:
+                return False
+            global_batch = max(max(int(c) for c in dp_token_counts), batch_size)
+            return _decode_bucket(global_batch) <= self.max_batch
+        return _decode_bucket(batch_size) <= self.max_batch
 
     def _decode_metadata(
         self, metadata: AttentionMetadata
@@ -385,7 +391,14 @@ def execute(
         input_embedding: torch.Tensor | None = None,
     ) -> torch.Tensor:
         batch_size = input_ids.shape[0]
-        padded_batch_size = _decode_bucket(batch_size)
+
+        if self.dp_size > 1:
+            dp_token_counts = tuple(int(c) for c in metadata.dp_token_counts)
+            global_batch = max(max(dp_token_counts, default=0), batch_size)
+            padded_batch_size = _decode_bucket(global_batch)
+        else:
+            padded_batch_size = _decode_bucket(batch_size)
+
         if padded_batch_size > self.max_batch:
             raise ValueError("decode batch exceeds ACL graph capacity")
 
@@ -418,6 +431,9 @@ def execute(
             self.attention_backend.prepare(entry.static_metadata, graph_mode=True)
 
         if first_capture:
+            self._stream.wait_stream(torch.npu.current_stream())
+            if self.dp_size > 1:
+                distributed.barrier(self.device, "dp")
             self._capture(entry)
 
         self._stream.wait_stream(torch.npu.current_stream())
@@ -524,6 +540,7 @@ def _allocate_entry(
             paged_kv_last_page_len_host=torch.ones(padded_batch_size, dtype=torch.int32, device="cpu"),
             kv_seq_lens_host_values=[1] * padded_batch_size,
             block_table=static_block_table,
+            dp_token_counts=tuple([padded_batch_size] * self.dp_size) if self.dp_size > 1 else (),
         )
         is_expanded = resolve_expanded_decode_metadata(metadata) is not None
         entry.kv_seq_lens_delta = torch.empty(padded_batch_size, dtype=torch.int32, device=device)
@@ -663,7 +680,7 @@ def _capture(self, entry: _DecodeGraphEntry) -> None:
             self.layer_caches,
             execution_state=entry.execution_state,
         )
-        with forward_context(context):
+        with forward_context(context), torch.npu.stream(self._stream):
             for _ in range(_CAPTURE_WARMUP_STEPS):
                 self._forward_static(entry)
         torch.npu.synchronize()

xllm/python/models/deepseek_v32.py

@@ -829,16 +829,16 @@ def process_weights_after_loading(self) -> None:
         self.shared_experts.process_weights_after_loading()
 
     def forward(self, hidden: torch.Tensor) -> torch.Tensor:
-        local_hidden = hidden
-        token_counts: list[int] | None = None
+        local_tokens: int = 0
+        padded_tokens: int = 0
         if self.dp_size > 1:
             token_counts = list(get_forward_context().metadata.dp_token_counts)
-            hidden = distributed.all_gather_variable(
-                hidden,
-                token_counts,
-                self.dp_rank,
-                "dp",
-            )
+            padded_tokens = max(token_counts)
+            local_tokens = hidden.shape[0]
+            pad_size = padded_tokens - local_tokens
+            if pad_size > 0:
+                hidden = torch.nn.functional.pad(hidden, (0, 0, 0, pad_size))
+            hidden = distributed.all_gather(hidden, dim=0, world_size=self.dp_size, group_name="dp")
 
         logits = self.gate(hidden)
         routed = kernels.grouped_moe(
@@ -867,11 +867,8 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor:
         elif self.cfg.tp_size > 1 and self.ep_size == 1:
             distributed.all_reduce_(final)
 
-        if token_counts is not None:
-            local_tokens = token_counts[self.dp_rank]
-            if local_tokens == 0:
-                return torch.zeros_like(local_hidden)
-            start = sum(token_counts[: self.dp_rank])
+        if padded_tokens > 0:
+            start = self.dp_rank * padded_tokens
             final = final.narrow(0, start, local_tokens)
 
         return final
@@ -967,6 +964,12 @@ def __init__(self, config: dict, build_model: bool = True) -> None:
         self.cfg.moe_tp_size = int(config.get("moe_tp_size", 1))
         self.cfg.moe_tp_rank = int(config.get("moe_tp_rank", 0))
         self.cfg.world_size = int(config.get("world_size", self.cfg.tp_size))
+        # C++ computes moe_tp_size = world_size / ep_size, which conflates DP
+        # replicas with TP shards. Correct to exclude DP ranks: each DP replica
+        # holds the full expert set independently — no MoE TP reduce needed
+        # across DP boundaries.
+        if self.cfg.dp_size > 1 and self.cfg.moe_tp_size > 1:
+            self.cfg.moe_tp_size = max(1, self.cfg.moe_tp_size // self.cfg.dp_size)
         if hasattr(self.cfg, "validate"):
             self.cfg.validate()
         dtype = self.resolve_dtype(config.get("dtype") or config.get("torch_dtype"))