diff --git a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py index 26a17f4a4..81ba655a9 100644 --- a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py +++ b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py @@ -4,13 +4,55 @@ from typing import Optional, Tuple import torch import triton +from sglang.srt.environ import envs +from sglang.srt.utils import get_bool_env_var, is_gfx95_supported, is_hip +from sglang.srt.utils.common import is_sm120_supported + logger = logging.getLogger(__name__) +_is_hip = is_hip() +_is_gfx95_supported = is_gfx95_supported() + _FUSED_HC_POST_PRE_M_THRESHOLD = 64 _FUSED_HC_POST_PRE_CACHE: dict[tuple, dict[str, torch.Tensor]] = {} _TRITON_MHC_POST_PRE_OPS = None _TRITON_MHC_POST_PRE_RUNTIME_DISABLED = False +# aiter fused mHC state. The kernel self-gates the fused-vs-unfused decision per +# arch internally (see aiter.ops.mhc.mhc_fused_post_pre), so this wrapper only +# tracks import/runtime availability and never re-implements the token threshold. +_AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED = False +_AITER_MHC_IMPORT_WARNED = False + + +def _is_fused_mhc_post_pre_enabled() -> bool: + # SM120 disables the standalone TileLang pre path. mhc_fused_post_pre does + # not read that flag and dispatches independently for both small and large + # token batches, so the standalone pre flag must not veto the fused opt-in. + return ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get() + and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get() + and (envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() or is_sm120_supported()) + ) + + +def _is_aiter_gfx95_mhc_available() -> bool: + return _is_hip and get_bool_env_var("SGLANG_USE_AITER") and is_gfx95_supported() + + +def _is_production_mhc_enabled() -> bool: + return _is_fused_mhc_post_pre_enabled() or _is_aiter_gfx95_mhc_available() + + +def is_cross_layer_mhc_fusion_enabled() -> bool: + """Whether DeepSeek-V4 may defer mHC post across the attn/MoE boundary. + + Cross-layer fusion requires a fused post+pre kernel to be available: either + the TileLang path (``SGLANG_OPT_FUSE_MHC_POST_PRE`` + TileLang pre/post) or + the aiter HIP path on a supported gfx95 device. + """ + return _is_production_mhc_enabled() + def _get_triton_mhc_post_pre_ops(): global _TRITON_MHC_POST_PRE_OPS @@ -153,3 +195,220 @@ def try_fused_hc_post_pre( return None return new_residual, layer_input_out, bufs["h_post"], bufs["h_res"], False + + +def try_aiter_fused_mhc_post_pre( + layer_input: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + hc_post_mult: float, + sinkhorn_iters: int, + norm_weight: Optional[torch.Tensor], + norm_eps: Optional[float], +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]: + """Fused mhc_post + next-layer mhc_pre via the aiter HIP kernel. + + Returns ``(next_residual, layer_input, post_mix, comb_mix, norm_applied)`` or + ``None`` to let the caller fall back. The aiter kernel internally chooses + between its fused and unfused (mhc_post + mhc_pre) implementations based on the + token count and detected arch, so no token threshold is applied here. + """ + global _AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED, _AITER_MHC_IMPORT_WARNED + + if ( + _AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED + or not _is_aiter_gfx95_mhc_available() + or layer_input.shape[0] == 0 + or layer_input.dim() != 2 + or residual.dim() != 3 + ): + return None + + try: + from aiter.ops.mhc import mhc_fused_post_pre + except Exception as err: + if not _AITER_MHC_IMPORT_WARNED: + logger.warning("aiter fused mHC is unavailable, falling back: %s", err) + _AITER_MHC_IMPORT_WARNED = True + _AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED = True + return None + + norm_kwargs = {} + if norm_weight is not None: + norm_kwargs["norm_weight"] = norm_weight + norm_kwargs["norm_eps"] = norm_eps if norm_eps is not None else rms_eps + + try: + post_mix, comb_mix, layer_input_out, next_residual = mhc_fused_post_pre( + layer_input=layer_input, + residual_in=residual, + post_layer_mix=post, + comb_res_mix=comb, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=rms_eps, + hc_pre_eps=hc_eps, + hc_sinkhorn_eps=hc_eps, + hc_post_mult_value=hc_post_mult, + sinkhorn_repeat=sinkhorn_iters, + **norm_kwargs, + ) + except Exception as err: + logger.warning( + "aiter fused mHC kernel failed, disabling fallback path: %s", err + ) + _AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED = True + return None + + post_out = post_mix.squeeze(-1) if post_mix.ndim == 3 else post_mix + return ( + next_residual, + layer_input_out, + post_out, + comb_mix, + norm_weight is not None, + ) + + +def try_mhc_fused_post_pre_boundary( + layer_input: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + rms_eps: float, + hc_eps: float, + hc_post_mult: float, + sinkhorn_iters: int, + norm_weight: Optional[torch.Tensor], + norm_eps: Optional[float], + fn_transpose: bool, + is_gfx95_supported_flag: bool, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]: + """Dispatch the fused mHC post+pre across the attn/MoE boundary. + + Preference order (first available wins): aiter HIP kernel, then the Triton + kernel. Returns ``None`` when neither fires so the caller can fall back to the + TileLang path or the unfused ``hc_post`` + ``hc_pre`` sequence. + + The aiter and Triton kernels expect opposite ``fn`` orientations: the Triton + kernel takes ``hc_fn`` transposed (``fn_transpose``), while aiter consumes it + in the native ``mhc_pre`` layout. + """ + aiter_result = try_aiter_fused_mhc_post_pre( + layer_input, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + hc_post_mult, + sinkhorn_iters, + norm_weight, + norm_eps, + ) + if aiter_result is not None: + return aiter_result + + triton_fn = hc_fn.T if fn_transpose else hc_fn + return try_fused_hc_post_pre( + layer_input, + residual, + post, + comb, + triton_fn, + hc_scale, + hc_base, + hc_mult, + rms_eps, + hc_eps, + hc_post_mult, + sinkhorn_iters, + is_gfx95_supported_flag, + ) + + +def apply_mhc_post_pre_boundary( + layer_input: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + rms_eps: float, + hc_eps: float, + hc_post_mult: float, + sinkhorn_iters: int, + norm_weight: Optional[torch.Tensor], + norm_eps: Optional[float], + *, + fn_transpose: bool, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]: + # Try the aiter/Triton fused post+pre kernels first; if neither fires, + # fall back to the TileLang fused kernel, else return None so the caller + # runs the unfused hc_post + hc_pre sequence. + fused = try_mhc_fused_post_pre_boundary( + layer_input, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + hc_mult, + rms_eps, + hc_eps, + hc_post_mult, + sinkhorn_iters, + norm_weight, + norm_eps, + fn_transpose, + _is_gfx95_supported, + ) + if fused is not None: + return fused + + if not _is_fused_mhc_post_pre_enabled(): + return None + + from sglang.srt.models.deepseek_v4 import _get_mhc_ops + + post_in = post.unsqueeze(-1) if post.ndim == 2 else post + ( + residual, + post_out, + comb_out, + layer_input_out, + ) = _get_mhc_ops().mhc_fused_post_pre( + layer_input, + residual, + post_in, + comb, + hc_fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + hc_eps, + hc_post_mult, + sinkhorn_iters, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + post_out = post_out.squeeze(-1) if post_out.ndim == 3 else post_out + return residual, layer_input_out, post_out, comb_out, True diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index e088f6b05..1a5589160 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -138,7 +138,8 @@ from sglang.srt.model_loader.weight_utils import ( ) from sglang.srt.models.dbrx import ReplicatedLinear from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import ( - try_fused_hc_post_pre, + apply_mhc_post_pre_boundary, + is_cross_layer_mhc_fusion_enabled, ) from sglang.srt.models.deepseek_common.utils import ( _use_aiter_bpreshuffle_gfx95, @@ -168,7 +169,6 @@ from sglang.srt.utils import ( log_info_on_rank0, make_layers, ) -from sglang.srt.utils.common import is_sm120_supported from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -219,17 +219,6 @@ DEEPSEEK_V4_STACKED_PARAMS_MAPPING: List[Tuple[str, str, int]] = [ ] -def _is_fused_mhc_post_pre_enabled() -> bool: - # SM120 disables the standalone TileLang pre path. mhc_fused_post_pre does - # not read that flag and dispatches independently for both small and large - # token batches, so the standalone pre flag must not veto the fused opt-in. - return ( - envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get() - and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get() - and (envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() or is_sm120_supported()) - ) - - # FlashInfer's mhc_pre_big_fuse only accepts these split-K counts. _FLASHINFER_MHC_PRE_SPLITS = (1, 2, 4, 8, 16) @@ -1792,7 +1781,7 @@ 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_fused_mhc_post_pre_enabled() + self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled() self._input_layernorm_weight_bf16 = None self._post_attention_layernorm_weight_bf16 = None @@ -2022,7 +2011,12 @@ class DeepseekV4DecoderLayer(nn.Module): use_fused = self.use_fused_mhc_post_pre if prev_residual is not None and use_fused: - residual, post, comb, hidden_states = _get_mhc_ops().mhc_fused_post_pre( + input_norm_weight = ( + self._input_layernorm_weight_bf16 + if self._input_layernorm_weight_bf16 is not None + else self.input_layernorm.weight.data + ) + fused = apply_mhc_post_pre_boundary( hidden_states, prev_residual, prev_post, @@ -2030,19 +2024,63 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, + self.hc_mult, self.rms_norm_eps, self.hc_eps, - self.hc_eps, _MHC_POST_MULT_VALUE, self.hc_sinkhorn_iters, - norm_weight=( - self._input_layernorm_weight_bf16 - if self._input_layernorm_weight_bf16 is not None - else self.input_layernorm.weight.data - ), - norm_eps=self.input_layernorm.variance_epsilon, + input_norm_weight, + self.input_layernorm.variance_epsilon, + fn_transpose=False, ) - x_quant = None + if fused is not None: + residual, hidden_states, post, comb, norm_fused = fused + if not norm_fused: + # The Triton fused post+pre returns the layer input WITHOUT + # the input layernorm applied (norm_fused=False). Apply it + # (fp8-quant on aiter gfx95) before attention, exactly as the + # unfused hc_pre path below does; otherwise unnormalized + # activations reach self_attn. + if _use_aiter and _is_gfx95_supported: + x_quant, hidden_states = _fused_rmsnorm_fp8_quant( + hidden_states, + self.input_layernorm.weight, + self.rms_norm_eps, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + x_quant = None + else: + x_quant = None + else: + # Fused dispatch declined: close the previous layer's deferred + # mHC post (prev_residual/prev_post/prev_comb) before opening this + # layer's pre. Skipping hc_post here would drop the previous-layer + # post state and corrupt all subsequent layers. + hidden_states = self.hc_post( + hidden_states, prev_residual, prev_post, prev_comb + ) + residual = hidden_states + hidden_states, post, comb, norm_fused = self.hc_pre( + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm=self.input_layernorm, + forward_batch=forward_batch, + ) + if not norm_fused: + if _use_aiter and _is_gfx95_supported: + x_quant, hidden_states = _fused_rmsnorm_fp8_quant( + hidden_states, + self.input_layernorm.weight, + self.rms_norm_eps, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + x_quant = None + else: + x_quant = None else: residual = hidden_states hidden_states, post, comb, norm_fused = self.hc_pre( @@ -2075,12 +2113,17 @@ class DeepseekV4DecoderLayer(nn.Module): ) if use_fused: - fused_mhc = try_fused_hc_post_pre( + post_attn_norm_weight = ( + self._post_attention_layernorm_weight_bf16 + if self._post_attention_layernorm_weight_bf16 is not None + else self.post_attention_layernorm.weight.data + ) + fused = apply_mhc_post_pre_boundary( hidden_states, residual, post, comb, - self.hc_ffn_fn.T, + self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base, self.hc_mult, @@ -2088,32 +2131,30 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_eps, _MHC_POST_MULT_VALUE, self.hc_sinkhorn_iters, - _is_gfx95_supported, + post_attn_norm_weight, + self.post_attention_layernorm.variance_epsilon, + fn_transpose=True, ) - if fused_mhc is not None: - residual, hidden_states, post, comb, norm_fused = fused_mhc + if fused is not None: + residual, hidden_states, post, comb, norm_fused = fused + if not norm_fused: + # The Triton fused post+pre skips the post-attention + # layernorm (norm_fused=False); apply it before the MoE, + # matching the unfused hc_pre path below. + hidden_states = self.post_attention_layernorm(hidden_states) else: - residual, post, comb, hidden_states = _get_mhc_ops().mhc_fused_post_pre( + hidden_states = self.hc_post(hidden_states, residual, post, comb) + residual = hidden_states + hidden_states, post, comb, norm_fused = self.hc_pre( hidden_states, - residual, - post.unsqueeze(-1) if post.ndim == 2 else post, - comb, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - _MHC_POST_MULT_VALUE, - self.hc_sinkhorn_iters, - norm_weight=( - self._post_attention_layernorm_weight_bf16 - if self._post_attention_layernorm_weight_bf16 is not None - else self.post_attention_layernorm.weight.data - ), - norm_eps=self.post_attention_layernorm.variance_epsilon, + norm=self.post_attention_layernorm, + forward_batch=forward_batch, ) - norm_fused = True + if not norm_fused: + hidden_states = self.post_attention_layernorm(hidden_states) else: hidden_states = self.hc_post(hidden_states, residual, post, comb) residual = hidden_states @@ -2352,11 +2393,63 @@ class DeepseekV4DecoderLayer(nn.Module): def op_mhc_post_attn_pre_mlp(self, state): # Close the attention mHC (hc_post), then open the FFN-side mHC pre + # post-attention layernorm. Produces the 2D MoE input. + # + # Pop each boundary tensor from the state EXACTLY ONCE, up front, and + # reuse the locals for both the fused attempt and the non-fused + # fallback. apply_mhc_post_pre_boundary() returns None when it declines + # to fuse -- most importantly for the 0-token DP two-batch-overlap idle + # ubatch -- in which case control must fall through to the unfused + # hc_post. Popping in the fused call's arguments and again in the + # fallback would double-pop -> KeyError: 'hidden_states_after_attn' on + # every idle DP rank. use_fused_mhc_post_pre is on whenever the aiter + # gfx95 mHC path is available (is_cross_layer_mhc_fusion_enabled), so + # this fallback is reached under DP regardless of the TileLang env. + hidden_states_after_attn = state.pop("hidden_states_after_attn") + attn_residual = state.pop("attn_residual") + attn_post = state.pop("attn_post") + attn_comb = state.pop("attn_comb") + + if self.use_fused_mhc_post_pre: + post_attn_norm_weight = ( + self._post_attention_layernorm_weight_bf16 + if self._post_attention_layernorm_weight_bf16 is not None + else self.post_attention_layernorm.weight.data + ) + fused = apply_mhc_post_pre_boundary( + hidden_states_after_attn, + attn_residual, + attn_post, + attn_comb, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.hc_mult, + self.rms_norm_eps, + self.hc_eps, + _MHC_POST_MULT_VALUE, + self.hc_sinkhorn_iters, + post_attn_norm_weight, + self.post_attention_layernorm.variance_epsilon, + fn_transpose=True, + ) + if fused is not None: + ffn_residual, hidden_states, post, comb, norm_fused = fused + if not norm_fused: + # The Triton fused post+pre skips the post-attention + # layernorm (norm_fused=False); apply it before the MoE, + # matching the unfused hc_pre path below. + hidden_states = self.post_attention_layernorm(hidden_states) + state.ffn_residual = ffn_residual + state.ffn_post = post + state.ffn_comb = comb + state.hidden_states_mlp_input = hidden_states + return + hidden_states = self.hc_post( - state.pop("hidden_states_after_attn"), - state.pop("attn_residual"), - state.pop("attn_post"), - state.pop("attn_comb"), + hidden_states_after_attn, + attn_residual, + attn_post, + attn_comb, ) ffn_residual = hidden_states hidden_states, post, comb, norm_fused = self.hc_pre( @@ -2630,7 +2723,7 @@ 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_fused_mhc_post_pre_enabled() + self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled() if self.dsa_enable_prefill_cp: self.cp_size = get_parallel().attn_cp_size diff --git a/test/registered/unit/models/test_deepseek_v4_amd_fused_mhc.py b/test/registered/unit/models/test_deepseek_v4_amd_fused_mhc.py new file mode 100644 index 000000000..51357a0e7 --- /dev/null +++ b/test/registered/unit/models/test_deepseek_v4_amd_fused_mhc.py @@ -0,0 +1,400 @@ +import unittest +from unittest import mock + +from sglang.srt.environ import envs +from sglang.srt.models.deepseek_common.amd import deepseek_v4_fused_mhc +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + + +class TestAmdFusedMhcCrossLayerGating(unittest.TestCase): + """Gating and dispatch-preference tests (CPU, no kernels required).""" + + def test_tilelang_fuse_flag_enables_cross_layer_fusion(self): + with ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True), + envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(True), + envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True), + ): + self.assertTrue(deepseek_v4_fused_mhc.is_cross_layer_mhc_fusion_enabled()) + + @mock.patch.object(deepseek_v4_fused_mhc, "is_sm120_supported", return_value=True) + def test_sm120_enables_fusion_with_tilelang_pre_disabled(self, _mock_sm120): + # Regression (PR review): consolidating _is_fused_mhc_post_pre_enabled into + # this module must preserve the SM120 special case. SM120 disables the + # standalone TileLang pre path, but mhc_fused_post_pre dispatches + # independently, so fuse+post enabled with the pre flag OFF must still + # enable fusion when SM120 is supported. The pre-fix consolidation + # required the pre flag unconditionally and silently disabled fusion on + # SM120. + with ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True), + envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True), + envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False), + ): + self.assertTrue(deepseek_v4_fused_mhc._is_fused_mhc_post_pre_enabled()) + + @mock.patch.object(deepseek_v4_fused_mhc, "is_sm120_supported", return_value=False) + def test_no_sm120_still_requires_tilelang_pre(self, _mock_sm120): + # Negative branch: the (pre OR sm120) clause must not degrade to + # always-true. With SM120 unsupported and the pre flag off, fuse+post + # alone must not enable the standalone TileLang fused path. + with ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True), + envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True), + envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False), + ): + self.assertFalse(deepseek_v4_fused_mhc._is_fused_mhc_post_pre_enabled()) + + def test_is_fused_mhc_post_pre_enabled_policy(self): + # Full gating table for _is_fused_mhc_post_pre_enabled, migrated from the + # removed test_deepseek_v4_fused_mhc_policy.py now that the helper lives + # in this module (it used to patch deepseek_v4.is_sm120_supported / + # deepseek_v4._is_fused_mhc_post_pre_enabled, both gone after the + # consolidation -> the registered CPU test AttributeError'd). Fusion + # requires the opt-in flag AND TileLang post AND (TileLang pre OR SM120). + cases = [ + # (fuse, pre, post, sm120, expected) + (True, False, True, True, True), # SM120 waives the standalone pre flag + (True, False, True, False, False), # non-SM120 still needs the pre flag + (True, True, True, False, True), # non-SM120 with the pre flag on + (False, False, True, True, False), # fusion opt-in is required + (True, False, False, True, False), # TileLang post is required + ] + for fuse, pre, post, sm120, expected in cases: + with self.subTest(fuse=fuse, pre=pre, post=post, sm120=sm120): + with ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(fuse), + envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(pre), + envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(post), + mock.patch.object( + deepseek_v4_fused_mhc, + "is_sm120_supported", + return_value=sm120, + ), + ): + self.assertEqual( + deepseek_v4_fused_mhc._is_fused_mhc_post_pre_enabled(), + expected, + ) + + @mock.patch.object(deepseek_v4_fused_mhc, "is_gfx95_supported", return_value=True) + @mock.patch.object(deepseek_v4_fused_mhc, "get_bool_env_var", return_value=True) + @mock.patch.object(deepseek_v4_fused_mhc, "_is_hip", True) + def test_aiter_gfx95_enables_cross_layer_fusion(self, _mock_aiter, _mock_gfx95): + # TileLang flags off: fusion must still enable via the aiter gfx95 path. + with ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(False), + envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False), + envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(False), + ): + self.assertTrue(deepseek_v4_fused_mhc.is_cross_layer_mhc_fusion_enabled()) + + @mock.patch.object(deepseek_v4_fused_mhc, "is_gfx95_supported", return_value=False) + @mock.patch.object(deepseek_v4_fused_mhc, "get_bool_env_var", return_value=True) + @mock.patch.object(deepseek_v4_fused_mhc, "_is_hip", True) + def test_aiter_cross_layer_disabled_without_gfx95(self, _mock_aiter, _mock_gfx95): + with ( + envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(False), + envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False), + envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(False), + ): + self.assertFalse(deepseek_v4_fused_mhc.is_cross_layer_mhc_fusion_enabled()) + + @mock.patch.object(deepseek_v4_fused_mhc, "is_gfx95_supported", return_value=True) + @mock.patch.object(deepseek_v4_fused_mhc, "get_bool_env_var", return_value=False) + @mock.patch.object(deepseek_v4_fused_mhc, "_is_hip", True) + def test_aiter_path_skips_without_sglang_use_aiter(self, _mock_aiter, _mock_gfx95): + result = deepseek_v4_fused_mhc.try_aiter_fused_mhc_post_pre( + layer_input=mock.Mock(shape=(32, 7168), dim=2, device="cpu"), + residual=mock.Mock(dim=3), + post=mock.Mock(), + comb=mock.Mock(), + hc_fn=mock.Mock(), + hc_scale=mock.Mock(), + hc_base=mock.Mock(), + rms_eps=1e-6, + hc_eps=1e-6, + hc_post_mult=2.0, + sinkhorn_iters=20, + norm_weight=mock.Mock(), + norm_eps=1e-6, + ) + self.assertIsNone(result) + + @mock.patch.object( + deepseek_v4_fused_mhc, + "try_aiter_fused_mhc_post_pre", + return_value=("res", "hs", "post", "comb", True), + ) + @mock.patch.object(deepseek_v4_fused_mhc, "try_fused_hc_post_pre") + def test_boundary_prefers_aiter_over_triton(self, mock_triton, mock_aiter): + result = deepseek_v4_fused_mhc.try_mhc_fused_post_pre_boundary( + layer_input=mock.Mock(shape=(32, 7168), dim=2), + residual=mock.Mock(dim=3), + post=mock.Mock(), + comb=mock.Mock(), + hc_fn=mock.Mock(), + hc_scale=mock.Mock(), + hc_base=mock.Mock(), + hc_mult=4, + rms_eps=1e-6, + hc_eps=1e-6, + hc_post_mult=2.0, + sinkhorn_iters=20, + norm_weight=mock.Mock(), + norm_eps=1e-6, + fn_transpose=True, + is_gfx95_supported_flag=True, + ) + self.assertEqual(result, ("res", "hs", "post", "comb", True)) + mock_triton.assert_not_called() + mock_aiter.assert_called_once() + + @mock.patch.object( + deepseek_v4_fused_mhc, "try_aiter_fused_mhc_post_pre", return_value=None + ) + @mock.patch.object( + deepseek_v4_fused_mhc, + "try_fused_hc_post_pre", + return_value=("res", "hs", "post", "comb", False), + ) + def test_boundary_falls_back_to_triton(self, mock_triton, mock_aiter): + hc_fn = mock.Mock() + hc_fn.T = "transposed_fn" + result = deepseek_v4_fused_mhc.try_mhc_fused_post_pre_boundary( + layer_input=mock.Mock(shape=(32, 7168), dim=2), + residual=mock.Mock(dim=3), + post=mock.Mock(), + comb=mock.Mock(), + hc_fn=hc_fn, + hc_scale=mock.Mock(), + hc_base=mock.Mock(), + hc_mult=4, + rms_eps=1e-6, + hc_eps=1e-6, + hc_post_mult=2.0, + sinkhorn_iters=20, + norm_weight=mock.Mock(), + norm_eps=1e-6, + fn_transpose=True, + is_gfx95_supported_flag=True, + ) + self.assertEqual(result, ("res", "hs", "post", "comb", False)) + mock_aiter.assert_called_once() + # fn_transpose=True must hand the Triton kernel the transposed fn. + self.assertEqual(mock_triton.call_args.args[4], "transposed_fn") + + +class TestAmdFusedMhcAttnBoundaryFallback(unittest.TestCase): + """Regression: the attn-side boundary fallback must close the previous + layer's deferred mHC post before opening the current layer's pre. + + When ``apply_mhc_post_pre_boundary`` declines to fuse (returns ``None``) -- + reachable after an aiter import/kernel failure permanently disables the fused + path -- the fallback must call + ``hc_post(hidden_states, prev_residual, prev_post, prev_comb)`` before + ``hc_pre``. The pre-fix code ran ``hc_pre`` directly on the raw input and + dropped ``prev_residual``/``prev_post``/``prev_comb``, so the previous + layer's deferred post was never applied and every subsequent layer computed + on corrupted activations. + + Drives the real ``DeepseekV4DecoderLayer.forward`` on a mocked layer with the + fused dispatcher forced to ``None`` and halts at ``self_attn`` via a sentinel, + so only the boundary fallback executes. + """ + + def test_fallback_closes_previous_post_before_pre(self): + try: + from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer + except Exception as e: # pragma: no cover - env without full model deps + self.skipTest(f"deepseek_v4 import unavailable: {e}") + + class _StopForward(Exception): + pass + + layer = mock.Mock() + layer.use_fused_mhc_post_pre = True + layer._input_layernorm_weight_bf16 = None + closed_post = object() + layer.hc_post.return_value = closed_post + # norm_fused=True keeps the fallback off the fp8-quant / layernorm branch. + layer.hc_pre.return_value = (object(), object(), object(), True) + layer.self_attn.maybe_use_decode_attn_tp.side_effect = _StopForward + + hs_in = object() + prev_residual, prev_post, prev_comb = object(), object(), object() + + with ( + mock.patch( + "sglang.srt.models.deepseek_v4.apply_mhc_post_pre_boundary", + return_value=None, + ), + self.assertRaises(_StopForward), + ): + DeepseekV4DecoderLayer.forward( + layer, + positions=object(), + hidden_states=hs_in, + input_ids=object(), + forward_batch=object(), + input_ids_global=object(), + prev_residual=prev_residual, + prev_post=prev_post, + prev_comb=prev_comb, + ) + + # The deferred previous-layer post must be closed with exactly the + # prev_* tensors, and hc_pre must then run on the closed result. + layer.hc_post.assert_called_once_with( + hs_in, prev_residual, prev_post, prev_comb + ) + layer.hc_pre.assert_called_once() + self.assertIs(layer.hc_pre.call_args.args[0], closed_post) + + +class TestAmdFusedMhcNormFusedHandling(unittest.TestCase): + """Regression: a fused-success result with ``norm_fused=False`` must have its + layernorm applied at the call site before the activation reaches attention. + + ``try_fused_hc_post_pre`` (the Triton fused post+pre) always returns + ``norm_fused=False`` -- it does not apply the input/post-attention layernorm. + The boundary dispatcher reaches it whenever the aiter kernel declines + (notably after an aiter import/kernel failure permanently disables the aiter + path). The pre-fix fused-success branch unpacked the tuple and fed the raw + (unnormalized) hidden_states straight into ``self_attn`` with ``x_quant=None``, + silently corrupting every subsequent layer. The fix mirrors the unfused + ``hc_pre`` branch: apply the input layernorm when ``norm_fused`` is False. + + Drives the real ``DeepseekV4DecoderLayer.forward`` with the boundary forced to + return ``norm_fused=False`` and halts at ``self_attn`` via a sentinel. + """ + + def test_fused_success_applies_input_layernorm_when_not_norm_fused(self): + try: + import sglang.srt.models.deepseek_v4 as deepseek_v4 + from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer + except Exception as e: # pragma: no cover - env without full model deps + self.skipTest(f"deepseek_v4 import unavailable: {e}") + + class _StopForward(Exception): + pass + + layer = mock.Mock() + layer.use_fused_mhc_post_pre = True + layer._input_layernorm_weight_bf16 = None + + fused_hs = object() + residual, post, comb = object(), object(), object() + # Fused dispatch SUCCEEDS but reports the input layernorm was NOT applied + # (norm_fused=False) -- the Triton fused post+pre contract. + normed = object() + layer.input_layernorm.return_value = normed + layer.self_attn.maybe_use_decode_attn_tp.side_effect = _StopForward + + # Force the non-aiter (torch layernorm) branch deterministically so the + # test does not depend on the runner arch and needs no real tensors. + with ( + mock.patch.object(deepseek_v4, "_use_aiter", False), + mock.patch.object(deepseek_v4, "_is_gfx95_supported", False), + mock.patch( + "sglang.srt.models.deepseek_v4.apply_mhc_post_pre_boundary", + return_value=(residual, fused_hs, post, comb, False), + ), + self.assertRaises(_StopForward), + ): + DeepseekV4DecoderLayer.forward( + layer, + positions=object(), + hidden_states=object(), + input_ids=object(), + forward_batch=object(), + input_ids_global=object(), + prev_residual=object(), + prev_post=object(), + prev_comb=object(), + ) + + # The fused (unnormalized) layer input must be run through the input + # layernorm before attention. Pre-fix this was never called on the + # fused-success path. + layer.input_layernorm.assert_called_once_with(fused_hs) + + +def _hardware_available() -> bool: + try: + import torch + + if not ( + torch.cuda.is_available() and deepseek_v4_fused_mhc.is_gfx95_supported() + ): + return False + from aiter.ops.mhc import mhc_fused_post_pre # noqa: F401 + + return True + except Exception: + return False + + +@unittest.skipUnless( + _hardware_available(), "requires a gfx95 device with aiter mHC kernels" +) +class TestAmdFusedMhcNumerical(unittest.TestCase): + """On-device equivalence of the aiter fused kernel vs unfused mhc_post+mhc_pre. + + Asserts the proven invariants: ``next_residual`` is bit-exact and + ``layer_input``/``post_mix`` match within bf16 tolerance. ``comb_mix`` is + intentionally not asserted here -- its raw-tensor value differs between the + fused and unfused kernels at the production Sinkhorn setting, and correctness + is established end-to-end (fused-on vs fused-off token match). See the PR + description; extend this test once the end-to-end sign-off pins the expected + comb_mix convention. + """ + + def _run(self, m, hc_mult=4, hidden=7168, sinkhorn_iters=20): + import torch + from aiter.ops import mhc + + dev = "cuda:0" + torch.manual_seed(0) + hc_mult3 = hc_mult * 2 + hc_mult * hc_mult + li = (torch.randn(m, hidden, device=dev) * 0.02).bfloat16() + res = (torch.randn(m, hc_mult, hidden, device=dev) * 0.02).bfloat16() + post = torch.randn(m, hc_mult, device=dev) * 0.02 + comb = torch.randn(m, hc_mult, hc_mult, device=dev) * 0.02 + fn = (torch.randn(hc_mult3, hc_mult * hidden, device=dev) * 0.02).bfloat16() + scl = torch.ones(hc_mult3, device=dev) + base = torch.zeros(hc_mult3, device=dev) + nw = torch.ones(hidden, device=dev).bfloat16() + kw = dict( + rms_eps=1e-6, + hc_pre_eps=1e-6, + hc_sinkhorn_eps=1e-6, + hc_post_mult_value=2.0, + sinkhorn_repeat=sinkhorn_iters, + norm_weight=nw, + norm_eps=1e-6, + ) + post_mix, _comb_mix, li_out, next_res = mhc.mhc_fused_post_pre( + li, res, post, comb, fn, scl, base, force_fused=True, **kw + ) + ref_next = torch.empty_like(res) + mhc.mhc_post(ref_next, li, res, post, comb) + ref_post, _ref_comb, ref_li = mhc.mhc_pre(ref_next, fn, scl, base, **kw) + + torch.testing.assert_close(next_res, ref_next, rtol=0, atol=0) + torch.testing.assert_close(li_out.float(), ref_li.float(), rtol=3e-2, atol=3e-2) + torch.testing.assert_close( + post_mix.float(), ref_post.float(), rtol=3e-2, atol=3e-2 + ) + + def test_equivalence_decode(self): + self._run(m=32) + + def test_equivalence_prefill(self): + self._run(m=96) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_deepseek_v4_fused_mhc_policy.py b/test/registered/unit/models/test_deepseek_v4_fused_mhc_policy.py deleted file mode 100644 index d5e1c3518..000000000 --- a/test/registered/unit/models/test_deepseek_v4_fused_mhc_policy.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Unit tests for the DeepSeek-V4 fused-MHC enable policy.""" - -import unittest -from unittest.mock import patch - -import sglang.srt.models.deepseek_v4 as deepseek_v4 -from sglang.srt.environ import envs -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=2, suite="base-a-test-cpu") - - -class TestDeepseekV4FusedMHCPolicy(CustomTestCase): - def _is_enabled( - self, - *, - fuse: bool, - tilelang_pre: bool, - tilelang_post: bool, - sm120: bool, - ) -> bool: - with ( - envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(fuse), - envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(tilelang_pre), - envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(tilelang_post), - patch.object(deepseek_v4, "is_sm120_supported", return_value=sm120), - ): - return deepseek_v4._is_fused_mhc_post_pre_enabled() - - def test_sm120_allows_fused_opt_in_with_standalone_pre_disabled(self): - self.assertTrue( - self._is_enabled( - fuse=True, - tilelang_pre=False, - tilelang_post=True, - sm120=True, - ) - ) - - def test_other_platform_still_requires_tilelang_pre(self): - self.assertFalse( - self._is_enabled( - fuse=True, - tilelang_pre=False, - tilelang_post=True, - sm120=False, - ) - ) - self.assertTrue( - self._is_enabled( - fuse=True, - tilelang_pre=True, - tilelang_post=True, - sm120=False, - ) - ) - - def test_fusion_opt_in_and_tilelang_post_remain_required(self): - self.assertFalse( - self._is_enabled( - fuse=False, - tilelang_pre=False, - tilelang_post=True, - sm120=True, - ) - ) - self.assertFalse( - self._is_enabled( - fuse=True, - tilelang_pre=False, - tilelang_post=False, - sm120=True, - ) - ) - - -if __name__ == "__main__": - unittest.main()