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 689c8c95f..199d3f74f 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 @@ -302,6 +302,7 @@ def forward_dsa_prepare_npu( forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", layer_scatter_modes, + prev_topk_indices: torch.Tensor = None, ): dynamic_scale = None if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode(): @@ -401,15 +402,18 @@ def forward_dsa_prepare_npu( latent_cache, forward_batch, k_nope, k_pe ) - topk_indices = m.indexer( - hidden_states, - q_lora, - positions, - forward_batch, - m.layer_id, - layer_scatter_modes, - dynamic_scale, - ) + if m.skip_topk: + topk_indices = prev_topk_indices + else: + topk_indices = m.indexer( + hidden_states, + q_lora, + positions, + forward_batch, + m.layer_id, + layer_scatter_modes, + dynamic_scale, + ) return ( q_pe, @@ -472,7 +476,10 @@ def forward_dsa_core_npu( attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim) output, _ = m.o_proj(attn_bmm_output) - return output + if not m.next_skip_topk: + return output, None + else: + return output, topk_indices def npu_mla_preprocess( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8f89bfd21..81697d86a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1467,6 +1467,7 @@ class DeepseekV2AttentionMLA( forward_batch, zero_allocator, layer_scatter_modes, + prev_topk_indices, ) else: raise NotImplementedError