From f1b41a5b3df70639677f98038e70e74864dd81ab Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Sun, 2 Aug 2026 04:50:10 +0800 Subject: [PATCH] [CP]: FIx some issue for glm5.2 cp v2 (#33100) --- .../sglang/srt/layers/attention/dsa/utils.py | 7 +++++++ .../srt/layers/attention/dsa_backend.py | 11 +++++++--- python/sglang/srt/server_args.py | 20 +++++++------------ 3 files changed, 22 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 5caf46cc9..7df042c37 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -170,6 +170,13 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"): # Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch, # calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode. from sglang.srt.layers.cp.padding import get_cp_padding_align_size + from sglang.srt.layers.cp.utils import is_cp_v2_active + + # CP-v2 already pads each rank-local shard to its physical size + if is_cp_v2_active(forward_batch): + return forward_batch.attn_cp_metadata.per_rank_actual_token[ + get_parallel().attn_cp_rank + ] global_num_tokens = forward_batch.global_num_tokens_cpu.copy() sync_group_size = len(global_num_tokens) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index b9094679b..c029f4496 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -3218,9 +3218,14 @@ class DeepseekSparseAttnBackend( rope_positions = forward_batch.positions if dsa_use_prefill_cp(forward_batch): - rope_positions = cp_split_and_rebuild_position( - forward_batch, rope_positions - ) + if is_cp_v2_active(forward_batch): + rope_positions = get_cp_strategy().shard_position_ids( + rope_positions, forward_batch + ) + else: + rope_positions = cp_split_and_rebuild_position( + forward_batch, rope_positions + ) q, k, k_rope = mla_quantize_and_rope_for_fp8( q, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d7d1b1156..c297e8e33 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -6203,19 +6203,13 @@ class ServerArgs: hf_config = model_config.hf_config model_arch = hf_config.architectures[0] if model_arch in CP_V2_DEFAULT_MODEL_CLASSES: - if getattr(hf_config, "index_share_for_mtp_iteration", False): - # GLM 5.2 (DSA index-share MTP): CP-v2 is not ready for it - # yet, so default the env to off and keep the legacy CP path. - if not envs.SGLANG_ENABLE_CP_V2.is_set(): - envs.SGLANG_ENABLE_CP_V2.set(False) - else: - is_dsa_default_model = is_deepseek_dsa(hf_config) - # DSA CP-v2 currently supports only the interleave strategy. - enable_default_cp_v2 = not is_dsa_default_model or ( - self.enable_prefill_cp and self.cp_strategy == "interleave" - ) - if enable_default_cp_v2 and not envs.SGLANG_ENABLE_CP_V2.is_set(): - envs.SGLANG_ENABLE_CP_V2.set(True) + is_dsa_default_model = is_deepseek_dsa(hf_config) + # DSA CP-v2 currently supports only the interleave strategy. + enable_default_cp_v2 = not is_dsa_default_model or ( + self.enable_prefill_cp and self.cp_strategy == "interleave" + ) + if enable_default_cp_v2 and not envs.SGLANG_ENABLE_CP_V2.is_set(): + envs.SGLANG_ENABLE_CP_V2.set(True) if ( self.enable_prefill_cp