[NPU] bugs fix: fix a condition bug when using speculative inference on Qwen3 and Qwen3 moe (#19532)

This commit is contained in:
shengzhaotian
2026-03-03 17:59:25 +08:00
committed by GitHub
parent 666caaf9ce
commit 365ca1edb5
2 changed files with 8 additions and 2 deletions

View File

@@ -181,7 +181,10 @@ class Qwen3Attention(nn.Module):
if get_global_server_args().rl_on_policy_target is not None:
hidden_states = hidden_states.bfloat16()
if not _is_npu or forward_batch.forward_mode.is_extend():
if (
not _is_npu
or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
):
q, k, v = self.forward_prepare_native(
positions=positions,
hidden_states=hidden_states,

View File

@@ -620,7 +620,10 @@ class Qwen3MoeAttention(nn.Module):
):
if hidden_states.shape[0] == 0:
return hidden_states, forward_batch, None
if not _is_npu or forward_batch.forward_mode.is_extend():
if (
not _is_npu
or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
):
return self.forward_prepare_native(
positions=positions,
hidden_states=hidden_states,