[NPU] Add NPU fallback for fused Triton gating kernels (#28293)

This commit is contained in:
iridiumine
2026-06-16 11:37:05 +08:00
committed by GitHub
parent ed9024e09f
commit 486ec150d0
2 changed files with 6 additions and 1 deletions
+1
View File
@@ -540,6 +540,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
use_fused_gate = (
self.shared_expert_gate is not None
and not use_intel_amx_backend(self.shared_expert_gate)
and not is_npu()
)
if hidden_states.shape[0] == 0:
+5 -1
View File
@@ -983,7 +983,11 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
attn_output = self.attn(q, k, v, forward_batch)
if self.attn_output_gate:
attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True)
if not _is_npu:
attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True)
else:
gate_val = gate.reshape(gate.shape[0], -1) if gate.ndim == 3 else gate
attn_output.mul_(torch.sigmoid(gate_val))
output, _ = self.o_proj(attn_output)
return output