minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/qwen2_moe.py
+1
View File
@@ -540,6 +540,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
use_fused_gate = ( use_fused_gate = (
self.shared_expert_gate is not None self.shared_expert_gate is not None
and not use_intel_amx_backend(self.shared_expert_gate) and not use_intel_amx_backend(self.shared_expert_gate)
and not is_npu()
) )
if hidden_states.shape[0] == 0: if hidden_states.shape[0] == 0:
python/sglang/srt/models/qwen3_5.py
+5 -1
View File
@@ -983,7 +983,11 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
attn_output = self.attn(q, k, v, forward_batch) attn_output = self.attn(q, k, v, forward_batch)
if self.attn_output_gate: 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) output, _ = self.o_proj(attn_output)
return output return output