diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index 1f310d82a..ebe4e3a8a 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -543,7 +543,9 @@ def _compute_num_split_for_mhc_pre(num_tokens: int, hc_hidden_size: int) -> int: block_m, block_k = 64, 64 grid_size = (num_tokens + block_m - 1) // block_m num_block_k = (hc_hidden_size + block_k - 1) // block_k + n_sms = torch.cuda.get_device_properties(0).multi_processor_count + return max(1, min(n_sms // max(grid_size, 1), num_block_k // 4)) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 1a5589160..495d932d1 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -182,22 +182,38 @@ 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 def _get_mhc_ops() -> MhcOps: """Load MHC kernels only when a DeepSeek-V4 layer needs them. - Model modules are imported eagerly by the registry. Importing + Model modules are imported eagerly by the registry. Importing ``sglang.kernels.ops.layernorm.mhc`` owns TileLang-backed MHC kernels. Import it only when a DeepSeek-V4 layer executes so registry discovery cannot initialize an optional CUDA runtime before unrelated models set up - their communication workspaces. DeepSeek-V4 is the sole consumer here. + 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, @@ -205,7 +221,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__) @@ -219,6 +242,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) @@ -1781,7 +1811,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 @@ -1857,6 +1889,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, @@ -1962,6 +2014,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 @@ -2740,6 +2795,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( @@ -3434,7 +3498,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)), @@ -3508,7 +3572,7 @@ class DeepseekV4ForCausalLM(nn.Module): else: raise ValueError("num_nextn_predict_layers is not in the config") - if not envs.SGLANG_OPT_FP8_WO_A_GEMM.get(): + if not _FP8_WO_A_GEMM: weights = _dequant_fp8_wo_a_streaming(weights) stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING