diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index 0fd8d9788..16dc169ab 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -12,7 +12,6 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( ) from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full from sglang.srt.layers.attention.nsa.utils import ( - cp_split_and_rebuild_position, nsa_use_prefill_cp, ) from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context @@ -224,12 +223,9 @@ def forward_mla_prepare_npu( q_nope_out = q_nope_out.transpose(0, 1) - if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp): - positions = cp_split_and_rebuild_position(forward_batch, positions) - q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) - if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp): + if nsa_use_prefill_cp(forward_batch): # support allgather+rerrange k_nope, k_pe = m.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe @@ -367,12 +363,9 @@ def forward_dsa_prepare_npu( q_nope_out = q_nope_out.transpose(0, 1) - if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp): - positions = cp_split_and_rebuild_position(forward_batch, positions) - q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) - if nsa_use_prefill_cp(forward_batch, m.nsa_enable_prefill_cp): + if nsa_use_prefill_cp(forward_batch): # support allgather+rerrange k_nope, k_pe = m.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe