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 24fbb4708..21c6f6e16 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 @@ -403,9 +403,7 @@ def forward_dsa_prepare_npu( latent_cache, forward_batch, k_nope, k_pe ) - if m.skip_topk: - topk_indices = prev_topk_indices - else: + if not m.skip_topk or (m.is_nextn and prev_topk_indices is None): topk_indices = m.indexer( hidden_states, q_lora, @@ -415,6 +413,8 @@ def forward_dsa_prepare_npu( layer_scatter_modes, dynamic_scale, ) + else: + topk_indices = prev_topk_indices return ( q_pe,