From 58c5bee3acb815ecf7be4055f69737c020c799d0 Mon Sep 17 00:00:00 2001 From: Chunyuan WU Date: Wed, 19 Aug 2026 09:56:21 +0800 Subject: [PATCH] Fix DP attention on CPU (#12961) --- .../sglang/srt/distributed/parallel_state.py | 6 ++- .../srt/layers/attention/intel_amx_backend.py | 4 +- python/sglang/srt/layers/dp_attention.py | 45 +++++++++++++++--- .../srt/model_executor/forward_batch_info.py | 15 ++++-- python/sglang/srt/models/deepseek_v2.py | 4 +- .../cpu/test_intel_amx_attention_backend_a.py | 47 +++++++++++++++++++ 6 files changed, 106 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 75dfdd0ce..6a0406e27 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py index 7a3f138a1..3ab03afdf 100644 --- a/python/sglang/srt/layers/attention/intel_amx_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 8205287b0..ea7810c0f 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/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): diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 0212e809a..04883e376 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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). diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index f2d76f735..ca8014a7c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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 diff --git a/test/registered/cpu/test_intel_amx_attention_backend_a.py b/test/registered/cpu/test_intel_amx_attention_backend_a.py index f15afacce..11214de23 100644 --- a/test/registered/cpu/test_intel_amx_attention_backend_a.py +++ b/test/registered/cpu/test_intel_amx_attention_backend_a.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( DEFAULT_MLA_MODEL_NAME_FOR_TEST, @@ -72,5 +73,51 @@ class TestIntelAMXAttnBackend(CustomTestCase): kill_process_tree(process.pid) +class TestDPAttention(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--trust-remote-code", + "--disable-radix-cache", + "--attention-backend", + "intel_amx", + "--mem-fraction-static", + "0.1", + "--disable-overlap-schedule", + "--tp", + "2", + "--enable-dp-attention", + "--dp", + "2", + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_dp_attention_DP2TP2(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=32, + parallel=32, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"Eval accuracy of GSM8K: {metrics=}") + + self.assertGreater(metrics["accuracy"], 0.7) + + if __name__ == "__main__": unittest.main()