【bugfix】The NPU's forward_dsa_prepare_npu also needs special handling for is_nextn (#28118)
This commit is contained in:
@@ -403,9 +403,7 @@ def forward_dsa_prepare_npu(
|
|||||||
latent_cache, forward_batch, k_nope, k_pe
|
latent_cache, forward_batch, k_nope, k_pe
|
||||||
)
|
)
|
||||||
|
|
||||||
if m.skip_topk:
|
if not m.skip_topk or (m.is_nextn and prev_topk_indices is None):
|
||||||
topk_indices = prev_topk_indices
|
|
||||||
else:
|
|
||||||
topk_indices = m.indexer(
|
topk_indices = m.indexer(
|
||||||
hidden_states,
|
hidden_states,
|
||||||
q_lora,
|
q_lora,
|
||||||
@@ -415,6 +413,8 @@ def forward_dsa_prepare_npu(
|
|||||||
layer_scatter_modes,
|
layer_scatter_modes,
|
||||||
dynamic_scale,
|
dynamic_scale,
|
||||||
)
|
)
|
||||||
|
else:
|
||||||
|
topk_indices = prev_topk_indices
|
||||||
|
|
||||||
return (
|
return (
|
||||||
q_pe,
|
q_pe,
|
||||||
|
|||||||
Reference in New Issue
Block a user