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

[qwen3.5][XPU]Add XPU support for set_embed_and_head and fused QK RMSNorm kernel (#27870)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
jiayisunx
2026-06-24 13:03:49 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent e3f1fa9d8e
commit bf231f01a3
1 changed files with 18 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/qwen3_5.py
+14 -4
View File
@@ -105,6 +105,7 @@ from sglang.srt.utils import (
is_gfx95_supported, is_gfx95_supported,
is_hip, is_hip,
is_npu, is_npu,
is_xpu,
make_layers, make_layers,
set_weight_attrs, set_weight_attrs,
) )
@@ -116,6 +117,7 @@ _is_npu = is_npu()
_is_cpu = is_cpu() _is_cpu = is_cpu()
_is_gfx95 = is_gfx95_supported() _is_gfx95 = is_gfx95_supported()
_is_hip = is_hip() _is_hip = is_hip()
_is_xpu = is_xpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_hip_use_alt_stream = get_bool_env_var("SGLANG_ALT_STREAM") and _is_hip _hip_use_alt_stream = get_bool_env_var("SGLANG_ALT_STREAM") and _is_hip
_gdn_use_alt_stream = _is_cuda or ( _gdn_use_alt_stream = _is_cuda or (
@@ -867,7 +869,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
k_by_head = k.reshape(-1, self.head_dim) k_by_head = k.reshape(-1, self.head_dim)
k_by_head = self.k_norm(k_by_head) k_by_head = self.k_norm(k_by_head)
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
elif _is_hip: elif _is_hip or _is_xpu:
q_by_head, k_by_head = fused_qk_gemma_rmsnorm( q_by_head, k_by_head = fused_qk_gemma_rmsnorm(
q, q,
k, k,
@@ -904,7 +906,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
q, k = self.rotary_emb(positions, q, k) q, k = self.rotary_emb(positions, q, k)
return q, k, v, gate return q, k, v, gate
def forward_prepare_hip(self, positions, hidden_states): def forward_prepare_fused_gate(self, positions, hidden_states):
qkv, _ = self.qkv_proj(hidden_states) qkv, _ = self.qkv_proj(hidden_states)
if self.attn_output_gate: if self.attn_output_gate:
q_gate, k, v = qkv.split( q_gate, k, v = qkv.split(
@@ -958,8 +960,8 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
) -> torch.Tensor: ) -> torch.Tensor:
"""Full attention forward pass.""" """Full attention forward pass."""
if _is_hip and self.attn_output_gate: if (_is_hip or _is_xpu) and self.attn_output_gate:
q, k, v, gate = self.forward_prepare_hip( q, k, v, gate = self.forward_prepare_fused_gate(
positions=positions, positions=positions,
hidden_states=hidden_states, hidden_states=hidden_states,
) )
@@ -1632,6 +1634,10 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration):
if self.pp_group.is_last_rank and head is not None: if self.pp_group.is_last_rank and head is not None:
del self.lm_head.weight del self.lm_head.weight
self.lm_head.weight = head self.lm_head.weight = head
if _is_xpu:
torch.xpu.empty_cache()
torch.xpu.synchronize()
else:
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.synchronize() torch.cuda.synchronize()
@@ -1787,6 +1793,10 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
if self.pp_group.is_last_rank and head is not None: if self.pp_group.is_last_rank and head is not None:
del self.lm_head.weight del self.lm_head.weight
self.lm_head.weight = head self.lm_head.weight = head
if _is_xpu:
torch.xpu.empty_cache()
torch.xpu.synchronize()
else:
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.synchronize() torch.cuda.synchronize()