diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index c40c0186c..ca5c45cd7 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/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, diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index f56ff48e4..06fa25d6a 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/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) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index e6ce967ca..6733e56c8 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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)), diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index 0879d3184..7ad2c26a3 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/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() diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 87f6e4326..3c6f2b736 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/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.cuda() + weight = weight.to(weight_device) 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() - w_scales = w_scales.cuda() + device = get_device() + 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() diff --git a/python/sglang/srt/models/llava.py b/python/sglang/srt/models/llava.py index 1f07f8a41..c5f62f6b2 100644 --- a/python/sglang/srt/models/llava.py +++ b/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() diff --git a/python/sglang/srt/models/llavavid.py b/python/sglang/srt/models/llavavid.py index f21c74485..61b6bd482 100644 --- a/python/sglang/srt/models/llavavid.py +++ b/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 diff --git a/python/sglang/srt/models/phi4mm_audio.py b/python/sglang/srt/models/phi4mm_audio.py index fd199836e..10ccbbc09 100644 --- a/python/sglang/srt/models/phi4mm_audio.py +++ b/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.cuda() - xs_pad = xs_pad.cuda() + enc_streaming_mask = enc_streaming_mask.to(xs_pad.device) 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.cuda() + enc_streaming_mask_nc = enc_streaming_mask_nc.to(xs_pad.device) if masks is not None: hs_mask_nc = masks & enc_streaming_mask_nc else: diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 2012068db..439c4f6ea 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/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() - torch.cuda.synchronize() + current_platform.empty_cache() + current_platform.synchronize() def set_lm_head_from_target(self, target_lm_head): if self.config.tie_word_embeddings: diff --git a/python/sglang/srt/models/qwen3_5_text.py b/python/sglang/srt/models/qwen3_5_text.py index 95a2d0844..814d6ff8a 100644 --- a/python/sglang/srt/models/qwen3_5_text.py +++ b/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() - torch.cuda.synchronize() + current_platform.empty_cache() + current_platform.synchronize() def set_dflash_layers_to_capture(self, layers_to_capture: list[int]): if not self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/yivl.py b/python/sglang/srt/models/yivl.py index efd7eb52e..f051f9101 100644 --- a/python/sglang/srt/models/yivl.py +++ b/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() diff --git a/test/registered/xpu/test_layernorm_xpu.py b/test/registered/xpu/test_layernorm_xpu.py new file mode 100644 index 000000000..ae63aa21d --- /dev/null +++ b/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()