Support V4.1 DP attention in DP-only DSpark PD

(cherry picked from commit 6ef0ebba730b3fffde4b59b953c5c9b58a35c4ed)
This commit is contained in:
SYChen123
2026-09-20 22:11:48 +08:00
committed by minke.yu
parent 4f22146e51
commit 2580c24d1b
2 changed files with 22 additions and 7 deletions
@@ -249,15 +249,13 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
if (
read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
or cfg.disaggregation_transfer_backend != "mooncake"
or cfg.dp_size != 1
or cfg.enable_dp_attention
or cfg.attn_cp_size != 1
or cfg.dcp_size != 1
):
raise ValueError(
"DeepSeek-V4.1 DSpark PD requires static verify, Mooncake, "
"DP=1 and CP=1. Both servers must enable DSpark with the same "
"block size and TP size."
"and CP=1 on both servers. DP attention is supported when "
"both servers use the same block size and target/draft KV layout."
)
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase
@@ -940,10 +940,27 @@ class CommonKVManager(BaseKVManager):
and self.attn_cp_size == 1
and info.attn_tp_size * info.attn_cp_size == self.attn_tp_size
)
if info.attn_tp_size != self.attn_tp_size and not same_tp_with_prefill_cp:
# Combined branch (40323-series + 40177): prefill CP can also pair
# with a DP-attention decode server. MLA KV is replicated across
# prefill CP ranks, so per-rank layouts match when attn_tp matches.
dp_decode_with_prefill_cp = (
info.attn_cp_size > 1
and self.attn_cp_size == 1
and (self.is_mla_backend or self.is_hybrid_mla_backend)
and info.attn_tp_size == self.attn_tp_size
)
non_cp_mla_layout = info.attn_cp_size == self.attn_cp_size == 1 and (
self.is_mla_backend or self.is_hybrid_mla_backend
)
if info.attn_tp_size != self.attn_tp_size and not (
same_tp_with_prefill_cp
or dp_decode_with_prefill_cp
or non_cp_mla_layout
):
raise RuntimeError(
"DeepSeek-V4.1 DSpark PD requires the same TP size on both "
"servers (including prefill CP ranks for MLA)"
"DeepSeek-V4.1 DSpark PD requires matching attention TP "
"unless both servers use CP=1 with an MLA KV layout, "
"or prefill runs CP with an MLA KV layout"
)
if self.dcp_size > 1: