[DSv32] Fix MTP and CP compatability (#19062)

Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
This commit is contained in:
Vladislav Nosivskoy
2026-02-21 06:10:34 +03:00
committed by GitHub
parent 463baafe10
commit afd91e8782

View File

@@ -34,8 +34,8 @@ from sglang.srt.layers.attention.nsa.utils import (
prepare_input_dp_with_cp_dsa,
)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
get_attention_cp_rank,
get_attention_cp_size,
is_dp_attention_enabled,
)
from sglang.srt.layers.layernorm import RMSNorm
@@ -123,7 +123,7 @@ class DeepseekModelNextN(nn.Module):
self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.cp_size = get_attention_tp_size()
self.cp_size = get_attention_cp_size()
else:
self.cp_size = None
@@ -206,8 +206,8 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
self.use_nsa = is_deepseek_nsa(config)
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.cp_rank = get_attention_tp_rank()
self.cp_size = get_attention_tp_size()
self.cp_rank = get_attention_cp_rank()
self.cp_size = get_attention_cp_size()
else:
self.cp_rank = None
self.cp_size = None