Fix DP attention on CPU (#12961)

python/sglang/srt/distributed/parallel_state.py

@@ -1079,7 +1079,8 @@ class GroupCoordinator:
         return output
 
     def reduce_scatter_tensor(self, output: torch.Tensor, input: torch.Tensor):
-        if _is_npu:
+        if _is_npu or _is_cpu:
+            # TODO: add optimized reduce_scatter_tensor kernel for cpu
             self._reduce_scatter_tensor(output, input)
         elif self._maybe_aiter_reduce_scatter(output, input):
             return
@@ -1259,7 +1260,8 @@ class GroupCoordinator:
         return envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
 
     def all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor):
-        if _is_npu:
+        if _is_npu or _is_cpu:
+            # TODO: add optimized all_gather_into_tensor kernel for cpu
             self._all_gather_into_tensor(output, input)
         else:
             # XPU and CUDA both go through reg_all_gather_into_tensor (custom_op) to

python/sglang/srt/layers/attention/intel_amx_backend.py

@@ -8,7 +8,7 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-from sglang.srt.runtime_context import get_spec
+from sglang.srt.runtime_context import get_parallel, get_spec
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
@@ -39,7 +39,7 @@ class IntelAMXAttnBackend(AttentionBackend):
         self.swa_out_cache_loc = None
 
         self.num_head = (
-            model_runner.model_config.num_attention_heads // model_runner.ps.tp_size
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
 
         # [NB]: `layer_id` set to 0 for qwen3-next models, as not all attn layers require kv pool

python/sglang/srt/layers/dp_attention.py

@@ -36,7 +36,7 @@ from sglang.srt.runtime_context import (
     get_flags,
     get_parallel,
 )
-from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils import get_bool_env_var, is_cpu, is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
@@ -75,6 +75,7 @@ def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int):
 
 _is_hip = is_hip()
 _USE_ROCM700A_WA = _is_hip and get_bool_env_var("SGLANG_USE_ROCM700A")
+_is_cpu = is_cpu()
 
 
 class DpPaddingMode(IntEnum):
@@ -451,6 +452,40 @@ def get_dp_local_slice_cpu(
 from sglang.kernels.ops.memory.memcpy_triton import memcpy_triton
 
 
+# TODO: write c++ kernel for cpu
+def memcpy_cpu(dst, src, dim, offset, sz, offset_src):
+    assert dim == 0, "Only dim=0 supported"
+    assert src.shape[1:] == dst.shape[1:], "src and dst must have same trailing shape"
+
+    total_rows_dst, total_rows_src = dst.shape[0], src.shape[0]
+    dst_start, src_start = 0, 0
+
+    if offset_src:
+        # src[offset:] → dst[0:]
+        src_start = offset
+        dst_start = 0
+    else:
+        # src[0:] → dst[offset:]
+        src_start = 0
+        dst_start = offset
+
+    dst_end = min(dst_start + sz, total_rows_dst)
+    src_end = min(src_start + sz, total_rows_src)
+    actual_sz = min(dst_end - dst_start, src_end - src_start)
+
+    if actual_sz <= 0:
+        return
+
+    dst[dst_start : dst_start + actual_sz].copy_(src[src_start : src_start + actual_sz])
+
+
+memcpy_func = memcpy_cpu if _is_cpu else memcpy_triton
+
+
+def memcpy(dst, src, dim, offset, sz, offset_src):
+    memcpy_func(dst, src, dim, offset, sz, offset_src)
+
+
 def _dp_gather_via_all_reduce(
     global_tokens: torch.Tensor,
     local_tokens: torch.Tensor,
@@ -470,9 +505,7 @@ def _dp_gather_via_all_reduce(
             local_tokens.untyped_storage() is not global_tokens.untyped_storage()
         ), "aliasing between global_tokens and local_tokens not allowed"
 
-        memcpy_triton(
-            global_tokens, local_tokens, 0, local_start_pos, local_num_tokens, False
-        )
+        memcpy(global_tokens, local_tokens, 0, local_start_pos, local_num_tokens, False)
 
     # Input IDs are in int 32. We should use inplace_all_reduce for local case because of custom all reduce.
     if world_dp_gather_enabled():
@@ -827,9 +860,7 @@ def dp_scatter(
             local_tokens.untyped_storage() is not global_tokens.untyped_storage()
         ), "aliasing between local_tokens and global_tokens not allowed"
 
-        memcpy_triton(
-            local_tokens, global_tokens, 0, local_start_pos, local_num_tokens, True
-        )
+        memcpy(local_tokens, global_tokens, 0, local_start_pos, local_num_tokens, True)
 
 
 def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):

python/sglang/srt/model_executor/forward_batch_info.py

@@ -53,6 +53,7 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
 )
 from sglang.srt.runtime_context import get_exec, get_parallel
 from sglang.srt.utils import (
+    is_cpu,
     is_cuda,
     is_hip,
     is_npu,
@@ -74,6 +75,7 @@ if TYPE_CHECKING:
 _skip_attn_backend_init_warned = False
 
 _is_npu = is_npu()
+_is_cpu = is_cpu()
 
 
 def _elastic_should_preserve_local_token_counts(
@@ -1457,8 +1459,13 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         # padding
         self._pad_inputs_to_size(model_runner, num_tokens, bs)
         self.global_num_tokens_cpu = global_num_tokens
-        global_num_tokens_pinned = torch.tensor(global_num_tokens, pin_memory=True)
-        self.global_num_tokens_gpu.copy_(global_num_tokens_pinned, non_blocking=True)
+        self.use_pin_memory = not _is_cpu
+        global_num_tokens_pinned = torch.tensor(
+            global_num_tokens, pin_memory=self.use_pin_memory
+        )
+        self.global_num_tokens_gpu.copy_(
+            global_num_tokens_pinned, non_blocking=self.use_pin_memory
+        )
 
         TboForwardBatchPreparer.prepare(
             batch=self, is_draft_worker=model_runner.is_draft_worker
@@ -1484,7 +1491,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             self.lora_ids.extend((bs - len(self.lora_ids)) * [None])
 
         seq_len_fill_value = (
-            model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
+            model_runner.attn_backend.get_cpu_graph_seq_len_fill_value()
+            if _is_cpu
+            else model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
         )
         # Keep gpu_only batches sync-free: leave seq_lens_sum None and let the
         # attention backend over-allocate from an upper bound (see #26738).

python/sglang/srt/models/deepseek_v2.py

@@ -1230,7 +1230,9 @@ class DeepseekV2MoE(nn.Module):
             ),  # block_size
             True,  # is_vnni
         )
-        if self.tp_size > 1 and not get_forward().fuse_mlp_allreduce:
+        if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
+            is_tp_path=True,
+        ):
             final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
         return final_hidden_states