【bugfix】The NPU's forward_dsa_prepare_npu also needs special handling for is_nextn (#28118)

This commit is contained in:
littleyellowbicycle
2026-06-15 21:31:39 +08:00
committed by GitHub
parent e985422b2b
commit 09e9c4fde3
@@ -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,