Xpu/weekly simple model enablement 2026 08 30 (#37193)

Co-authored-by: dayanandav <dayananda.vasantha.kumar@intel.com>
Co-authored-by: Girijala, Pavan Sivaram <pavan.sivaram.girijala@intel.com>
Co-authored-by: Cui, Lily <lily.cui@intel.com>
Co-authored-by: Juan Muneton <juan.muneton.gallego@intel.com>
Co-authored-by: Gao, Pengfei <pengfei.gao@intel.com>

python/sglang/srt/layers/rotary_embedding/base.py

@@ -454,8 +454,6 @@ class RotaryEmbedding(BaseFusedOp):
         ), "fused_set_kv_buffer_arg is not supported for xpu implementation"
         positions = torch.add(positions, offsets) if offsets is not None else positions
 
-        self._match_cos_sin_cache_dtype(query)
-
         # Fused_qk_rope only supports aligned head_size
         if self.head_size in [128, 256, 512]:
             num_tokens = positions.size(0)
@@ -475,6 +473,7 @@ class RotaryEmbedding(BaseFusedOp):
             return query, key
         else:
             # Use fallback kernel of 'rotary_embedding'
+            self._match_cos_sin_cache_dtype(query)
             return torch.ops.sgl_kernel.rotary_embedding(
                 positions,
                 query,

python/sglang/srt/layers/rotary_embedding/mrope.py

@@ -610,6 +610,32 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
 
         return self.forward_native(positions, query, key)
 
+    def forward_xpu(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor = None,
+    ):
+        assert key is not None
+        assert positions.ndim in (1, 2)
+        self._match_cos_sin_cache_dtype(query)
+
+        if positions.ndim == 2:
+            assert self.mrope_section is not None
+            triton_ernie45_rope_fused_inplace(
+                q=query,
+                k=key,
+                cos_sin_cache=self.cos_sin_cache,
+                positions=positions,
+                mrope_section=self.mrope_section,
+                head_size=self.head_size,
+                rotary_dim=self.rotary_dim,
+                is_neox_style=self.is_neox_style,
+            )
+            return query, key
+
+        return self.forward_native(positions, query, key)
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -618,4 +644,6 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
         fused_set_kv_buffer_arg=None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         assert positions.ndim == 1 or positions.ndim == 2
+        if _is_xpu:
+            return self.forward_xpu(positions, query, key)
         return self.forward_cuda(positions, query, key)

python/sglang/srt/models/deepseek_v4.py

@@ -194,6 +194,9 @@ class MhcOps(NamedTuple):
     hc_split_sinkhorn: Callable[..., Any]
     mhc_fused_post_pre: Optional[Callable[..., Any]]
     npu_hc_pre: Optional[Callable[..., Any]]
+    mhc_pre: Optional[Callable[..., Any]]
+    mhc_post: Optional[Callable[..., Any]]
+    fused_hc_head: Optional[Callable[..., Any]]
 
 
 @functools.cache
@@ -207,9 +210,22 @@ def _get_mhc_ops() -> MhcOps:
     their communication workspaces.  DeepSeek-V4 is the sole consumer here.
     """
     if _is_xpu:
-        from sgl_kernel import hc_split_sinkhorn
+        from sgl_kernel import (
+            fused_hc_head,
+            hc_post,
+            hc_split_sinkhorn,
+            mhc_fused_post_pre,
+            mhc_pre,
+        )
 
-        return MhcOps(hc_split_sinkhorn, None, None)
+        return MhcOps(
+            hc_split_sinkhorn=hc_split_sinkhorn,
+            mhc_fused_post_pre=mhc_fused_post_pre,
+            npu_hc_pre=None,
+            mhc_pre=mhc_pre,
+            mhc_post=hc_post,
+            fused_hc_head=fused_hc_head,
+        )
 
     from sglang.kernels.ops.layernorm.mhc import (
         hc_split_sinkhorn,
@@ -217,7 +233,14 @@ def _get_mhc_ops() -> MhcOps:
         npu_hc_pre,
     )
 
-    return MhcOps(hc_split_sinkhorn, mhc_fused_post_pre, npu_hc_pre)
+    return MhcOps(
+        hc_split_sinkhorn=hc_split_sinkhorn,
+        mhc_fused_post_pre=mhc_fused_post_pre,
+        npu_hc_pre=npu_hc_pre,
+        mhc_pre=None,
+        mhc_post=None,
+        fused_hc_head=None,
+    )
 
 
 logger = logging.getLogger(__name__)
@@ -232,6 +255,13 @@ DEEPSEEK_V4_STACKED_PARAMS_MAPPING: List[Tuple[str, str, int]] = [
 ]
 
 
+def _is_fused_mhc_post_pre_enabled_xpu() -> bool:
+    if _is_xpu:
+        return envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get()
+
+    return False
+
+
 # FlashInfer's mhc_pre_big_fuse only accepts these split-K counts.
 _FLASHINFER_MHC_PRE_SPLITS = (1, 2, 4, 8, 16)
 
@@ -1854,7 +1884,9 @@ class DeepseekV4DecoderLayer(nn.Module):
         ) = make_hc_mixing_params(hc_mult, config.hidden_size)
         self.rms_norm_eps = config.rms_norm_eps
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
-        self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled()
+        self.use_fused_mhc_post_pre = (
+            is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu()
+        )
         self._input_layernorm_weight_bf16 = None
         self._post_attention_layernorm_weight_bf16 = None
 
@@ -1930,6 +1962,26 @@ class DeepseekV4DecoderLayer(nn.Module):
             )
             return y, post, comb, False
 
+        if _is_xpu:
+            norm_kwargs = {}
+            if norm is not None:
+                norm_kwargs["norm_weight"] = norm.weight.data
+                norm_kwargs["norm_eps"] = norm.variance_epsilon
+
+            post, comb, y = _get_mhc_ops().mhc_pre(
+                residual=x,
+                fn=hc_fn,
+                hc_scale=hc_scale,
+                hc_base=hc_base,
+                rms_eps=self.rms_norm_eps,
+                hc_pre_eps=self.hc_eps,
+                hc_sinkhorn_eps=self.hc_eps,
+                hc_post_mult_value=_MHC_POST_MULT_VALUE,
+                sinkhorn_repeat=self.hc_sinkhorn_iters,
+                **norm_kwargs,
+            )
+            return y, post, comb, norm is not None
+
         if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get():
             y, post, comb = _flashinfer_hc_pre(
                 x,
@@ -2042,6 +2094,9 @@ class DeepseekV4DecoderLayer(nn.Module):
         if _is_npu:
             return torch.ops.custom.npu_hc_post(x, residual, post, comb)
 
+        if _is_xpu:
+            return _get_mhc_ops().mhc_post(x, residual, post, comb)
+
         if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get():
             from flashinfer.mhc import mhc_post
 
@@ -2806,7 +2861,9 @@ class DeepseekV4Model(nn.Module):
             ) = make_hc_head_params(hc_mult, config.hidden_size)
 
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
-        self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled()
+        self.use_fused_mhc_post_pre = (
+            is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu()
+        )
         if self.dsa_enable_prefill_cp:
             self.cp_size = get_parallel().attn_cp_size
 
@@ -2823,6 +2880,15 @@ class DeepseekV4Model(nn.Module):
         hc_base: torch.Tensor,
     ):
         if x.numel() > 0:
+            if _is_xpu:
+                return _get_mhc_ops().fused_hc_head(
+                    x.contiguous(),
+                    hc_fn,
+                    hc_scale,
+                    hc_base,
+                    norm_eps=self.norm_eps,
+                    hc_eps=self.hc_eps,
+                )
             from sglang.kernels.ops.layernorm.mhc_head import fused_hc_head
 
             return fused_hc_head(
@@ -3524,7 +3590,7 @@ class DeepseekV4ForCausalLM(nn.Module):
         if self._mhc_prewarmed_at_load:
             return
         self._mhc_prewarmed_at_load = True
-        if _is_npu or not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
+        if _is_npu or _is_xpu or not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
             return
         layer = next(
             (m for m in self.model.layers if isinstance(m, DeepseekV4DecoderLayer)),

python/sglang/srt/models/ernie45_moe_vl.py

@@ -41,6 +41,7 @@ from sglang.srt.layers.rotary_embedding import Ernie4_5_VLRotaryEmbedding
 from sglang.srt.layers.utils import PPMissingLayer
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, make_layers
@@ -282,7 +283,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
         hidden_dim = hidden_states.shape[-1]
         hidden_states = hidden_states.view(-1, hidden_dim)
 
-        capturing = torch.cuda.is_current_stream_capturing()
+        capturing = get_is_capture_mode()
 
         if visual_token_mask is not None and not capturing:
             all_visual = visual_token_mask.all()

python/sglang/srt/models/gpt_oss.py

@@ -77,6 +77,7 @@ from sglang.srt.runtime_context import (
 from sglang.srt.utils import (
     LazyValue,
     add_prefix,
+    get_device,
     is_cpu,
     is_cuda,
     is_flashinfer_available,
@@ -1000,9 +1001,10 @@ class GptOssForCausalLM(nn.Module):
         moe_ep_rank_start = moe_ep_rank * moe_num_local_experts
         moe_ep_rank_end = (moe_ep_rank + 1) * moe_num_local_experts
 
+        weight_device = next(iter(params_dict.values())).device
+
         for name, weight in weights:
-            if _is_cuda:
+            weight = weight.to(weight_device)
-                weight = weight.cuda()
 
             if "gate_up_proj_blocks" in name:
                 # Handle MLP gate and up projection weights
@@ -1392,8 +1394,9 @@ def _dequant_mlp_weight(debug_name, w_blocks, w_scales):
 
     original_device = w_blocks.device
 
-    w_blocks = w_blocks.cuda()
+    device = get_device()
-    w_scales = w_scales.cuda()
+    w_blocks = w_blocks.to(device)
+    w_scales = w_scales.to(device)
 
     w_bf16 = dequant_mxfp4(w_block=w_blocks, w_scale=w_scales, out_dtype=torch.bfloat16)
     w_bf16 = w_bf16.transpose(-2, -1).contiguous()

python/sglang/srt/models/llava.py

@@ -473,14 +473,15 @@ class LlavaBaseForCausalLM(nn.Module):
         # huggingface_name or path_of_clip_relative_to_llava_model_dir
         # We put the initialization here instead of __init__ to allow it being reused by other subclasses.
         vision_path = self.config.mm_vision_tower
+        device = next(self.language_model.parameters()).device
         if "clip" in vision_path:
             self.vision_tower = CLIPVisionModel.from_pretrained(
                 vision_path, torch_dtype=torch.float16
-            ).cuda()
+            ).to(device)
         elif "siglip" in vision_path:
             self.vision_tower = SiglipVisionModel.from_pretrained(
                 vision_path, torch_dtype=torch.float16
-            ).cuda()
+            ).to(device)
             # Siglip needs all feature tokens
             self.config.mm_vision_select_feature = "full"
         self.vision_tower.eval()

python/sglang/srt/models/llavavid.py

@@ -228,9 +228,10 @@ class LlavaVidForCausalLM(nn.Module):
         # huggingface_name or path_of_clip_relative_to_llava_model_dir
         # We put the initialization here instead of __init__ to allow it being reused by other subclasses.
         vision_path = self.config.mm_vision_tower
+        device = next(self.language_model.parameters()).device
         self.vision_tower = CLIPVisionModel.from_pretrained(
             vision_path, torch_dtype=torch.float16
-        ).cuda()
+        ).to(device)
         self.vision_tower.eval()
 
         self.vision_feature_layer = self.config.mm_vision_select_layer

python/sglang/srt/models/phi4mm_audio.py

@@ -561,9 +561,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
             seq_len, batch_size, self.chunk_size, self.left_chunk
         )
 
-        if xs_pad.is_cuda:
+        enc_streaming_mask = enc_streaming_mask.to(xs_pad.device)
-            enc_streaming_mask = enc_streaming_mask.cuda()
-            xs_pad = xs_pad.cuda()
 
         input_tensor = xs_pad
         input_tensor, masks = self._forward_embeddings_core(input_tensor, masks)
@@ -580,8 +578,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
             enc_streaming_mask_nc = self._streaming_mask(
                 seq_len, batch_size, chunk_size_nc, left_chunk_nc
             )
-            if xs_pad.is_cuda:
+            enc_streaming_mask_nc = enc_streaming_mask_nc.to(xs_pad.device)
-                enc_streaming_mask_nc = enc_streaming_mask_nc.cuda()
             if masks is not None:
                 hs_mask_nc = masks & enc_streaming_mask_nc
             else:

python/sglang/srt/models/qwen3_5_mtp.py

@@ -34,6 +34,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen3_5 import QWEN3_5_KV_SCALE_MAPPER, Qwen3_5ForCausalLM
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_model,
     get_parallel,
@@ -162,8 +163,8 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
         if head is not None and not self.config.tie_word_embeddings:
             del self.lm_head.weight
             self.lm_head.weight = head
-        torch.cuda.empty_cache()
+        current_platform.empty_cache()
-        torch.cuda.synchronize()
+        current_platform.synchronize()
 
     def set_lm_head_from_target(self, target_lm_head):
         if self.config.tie_word_embeddings:

python/sglang/srt/models/qwen3_5_text.py

@@ -29,6 +29,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models import qwen3_5
 from sglang.srt.models.qwen2_moe import Qwen2MoeSparseMoeBlock
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import LazyValue, add_prefix
 
@@ -143,8 +144,8 @@ class Qwen3_5ForCausalLM(nn.Module):
         del self.lm_head.weight
         self.model.embed_tokens.weight = embed
         self.lm_head.weight = head
-        torch.cuda.empty_cache()
+        current_platform.empty_cache()
-        torch.cuda.synchronize()
+        current_platform.synchronize()
 
     def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):
         if not self.pp_group.is_last_rank:

python/sglang/srt/models/yivl.py

@@ -40,11 +40,12 @@ class YiVLForCausalLM(LlavaLlamaForCausalLM):
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         # We have to use the subfolder of the main model directory (e.g. 01-ai/Yi-VL-6B)
+        device = next(self.language_model.parameters()).device
         self.vision_tower = CLIPVisionModel.from_pretrained(
             self.config._name_or_path,
             torch_dtype=torch.float16,
             subfolder=self.vision_tower_subfolder,
-        ).to("cuda")
+        ).to(device)
 
         self.vision_tower.eval()
 

test/registered/xpu/test_layernorm_xpu.py

@@ -0,0 +1,52 @@
+"""Tests for Gemma4RMSNorm forward_xpu dispatch."""
+
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
+
+
+class TestGemma4RMSNormXPU(CustomTestCase):
+    def setUp(self):
+        if not torch.xpu.is_available():
+            self.skipTest("XPU not available")
+        torch.manual_seed(42)
+        from sglang.srt.layers.layernorm import Gemma4RMSNorm
+
+        self.norm = Gemma4RMSNorm(128, eps=1e-6, scale_shift=1.0).to("xpu")
+
+    def test_2d_input(self):
+        x = torch.randn(4, 128, dtype=torch.bfloat16, device="xpu")
+        out = self.norm.forward_xpu(x)
+        self.assertEqual(out.shape, (4, 128))
+        ref = self.norm.forward_native(x)
+        torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-3)
+
+    def test_3d_input(self):
+        x = torch.randn(4, 8, 128, dtype=torch.bfloat16, device="xpu")
+        out = self.norm.forward_xpu(x)
+        self.assertEqual(out.shape, (4, 8, 128))
+        ref = self.norm.forward_native(x)
+        torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-3)
+
+    def test_scale_shift_zero(self):
+        from sglang.srt.layers.layernorm import Gemma4RMSNorm
+
+        norm0 = Gemma4RMSNorm(128, eps=1e-6, scale_shift=0.0).to("xpu")
+        x = torch.randn(4, 128, dtype=torch.bfloat16, device="xpu")
+        out = norm0.forward_xpu(x)
+        ref = norm0.forward_native(x)
+        torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-3)
+
+    def test_empty_tensor(self):
+        x = torch.empty(0, 128, dtype=torch.bfloat16, device="xpu")
+        out = self.norm.forward_xpu(x)
+        self.assertEqual(out.shape, (0, 128))
+
+
+if __name__ == "__main__":
+    unittest.main()