diff --git a/python/sglang/srt/layers/communicator_mhc.py b/python/sglang/srt/layers/communicator_mhc.py index 57f128253..d48f43458 100644 --- a/python/sglang/srt/layers/communicator_mhc.py +++ b/python/sglang/srt/layers/communicator_mhc.py @@ -72,6 +72,7 @@ class MHCState: hc_attn_pre: Callable hc_ffn_pre: Callable hc_post: Callable + hc_ffn_post_pre: Optional[Callable] = None h_res: Optional[torch.Tensor] = None h_post: Optional[torch.Tensor] = None @@ -94,9 +95,26 @@ class MHCState: def attn_to_mlp( self, hidden_states, residual, out_norm: Optional[torch.nn.Module] = None ): + out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm) + if self.hc_ffn_post_pre is not None and hidden_states.shape[0] != 0: + # Returns None when it declines -- no fused kernel for this platform + # or shape, or a shape the fusion is slower at -- and the chain runs. + fused = self.hc_ffn_post_pre( + hidden_states=hidden_states, + residual=residual, + h_res=self.h_res, + h_post=self.h_post, + out_norm_weight=out_norm_weight, + out_norm_eps=out_norm_eps, + ) + if fused is not None: + hidden_states, residual, self.h_res, self.h_post, norm_fused = fused + if out_norm is not None and not norm_fused: + hidden_states = out_norm(hidden_states) + return hidden_states, residual + hidden_states = self.hc_post(hidden_states, residual, self.h_res, self.h_post) residual = hidden_states - out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm) hidden_states, self.h_res, self.h_post, norm_fused = self.hc_ffn_pre( hidden_states, out_norm_weight, out_norm_eps ) @@ -406,6 +424,7 @@ class MHCLayerCommunicator(LayerCommunicator): hc_attn_pre: Callable, hc_ffn_pre: Callable, hc_post: Callable, + hc_ffn_post_pre: Optional[Callable] = None, ): self.is_first_layer = is_first_layer self.mhc = MHCState( @@ -413,6 +432,7 @@ class MHCLayerCommunicator(LayerCommunicator): hc_attn_pre=hc_attn_pre, hc_ffn_pre=hc_ffn_pre, hc_post=hc_post, + hc_ffn_post_pre=hc_ffn_post_pre, ) super().__init__( diff --git a/python/sglang/srt/models/glm5_next.py b/python/sglang/srt/models/glm5_next.py index 50aafeab7..3538914fd 100644 --- a/python/sglang/srt/models/glm5_next.py +++ b/python/sglang/srt/models/glm5_next.py @@ -76,6 +76,10 @@ from sglang.srt.model_loader.weight_utils import ( default_weight_loader, sharded_weight_loader, ) +from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import ( + apply_mhc_post_pre_boundary, + is_cross_layer_mhc_fusion_enabled, +) from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( DeepseekV2WeightLoaderMixin, ) @@ -124,6 +128,15 @@ if _use_aiter_gfx95: logger = logging.getLogger(__name__) +# Matches DeepSeek-V4's _MHC_POST_MULT_VALUE; the fused and unfused boundaries +# must agree on it. +_MHC_POST_MULT_VALUE = 2.0 + +# Conservative cap, not the crossover: at GLM-5.3-Flash's hc_mult=4 and +# hidden_size=4096 the fusion wins to 16 tokens and reaches parity at 24, with +# 17-23 unmeasured. Past it the pre-norm GEMM drops mhc_pre's split-K kernel. +_MHC_FUSED_BOUNDARY_MAX_TOKENS = 16 + @torch.compile def swiglu_clamped(y: torch.Tensor, limit: float): @@ -788,6 +801,13 @@ class Glm5NextDecoderLayer(nn.Module): hc_attn_pre=self.hc_attn_pre, hc_ffn_pre=self.hc_ffn_pre, hc_post=self.hc_post, + # Resolved once: env and platform are frozen after startup, + # and None keeps the dispatch off the per-boundary path. + hc_ffn_post_pre=( + self.hc_ffn_post_pre + if is_cross_layer_mhc_fusion_enabled() + else None + ), ) self.layer_communicator = MHCLayerCommunicator( **shared_kwargs, @@ -808,7 +828,7 @@ class Glm5NextDecoderLayer(nn.Module): rms_eps=self.config.rms_norm_eps, hc_eps=self.config.hc_eps, sinkhorn_iters=self.config.hc_sinkhorn_iters, - post_mult_value=2.0, + post_mult_value=_MHC_POST_MULT_VALUE, hc_norm_weight=None, out_norm_weight=out_norm_weight, out_norm_eps=out_norm_eps, @@ -834,6 +854,46 @@ class Glm5NextDecoderLayer(nn.Module): out_norm_eps, ) + def hc_ffn_post_pre( + self, hidden_states, residual, h_res, h_post, out_norm_weight, out_norm_eps + ): + # Fuses hc_post into the pre-norm GEMM; the mhc_pre big-fuse stage + # still launches separately, so this is two launches instead of three. + assert self.config.mhc, "hc_ffn_post_pre is only valid when config.mhc=True" + num_tokens, hidden_size = hidden_states.shape + if num_tokens > _MHC_FUSED_BOUNDARY_MAX_TOKENS: + return None + hc_mult = self.config.hc_mult + fused = apply_mhc_post_pre_boundary( + layer_input=hidden_states, + residual=residual.view(num_tokens, hc_mult, hidden_size), + post=h_post.view(num_tokens, hc_mult), + comb=h_res.view(num_tokens, hc_mult, hc_mult), + hc_fn=self.hc_ffn_fn, + hc_scale=self.hc_ffn_scale, + hc_base=self.hc_ffn_base, + hc_mult=hc_mult, + rms_eps=self.config.rms_norm_eps, + hc_eps=self.config.hc_eps, + hc_post_mult=_MHC_POST_MULT_VALUE, + sinkhorn_iters=self.config.hc_sinkhorn_iters, + norm_weight=out_norm_weight, + norm_eps=out_norm_eps, + # Matches DeepSeek-V4's two hc_ffn_fn boundaries; the Triton tier's + # parameter is hc_fn_t and this fn has the same [mix_hc, hc_dim] layout. + fn_transpose=True, + ) + if fused is None: + return None + next_residual, layer_input, post, comb, norm_fused = fused + return ( + layer_input, + next_residual.reshape(num_tokens, -1), + comb.reshape(num_tokens, hc_mult * hc_mult), + post.reshape(num_tokens, hc_mult), + norm_fused, + ) + def hc_post(self, hidden_states, residual, h_res, h_post): assert self.config.mhc, "hc_post is only valid when config.mhc=True" return _hc_post_fn( diff --git a/test/registered/kernels/ops/layernorm/test_mhc_kernels.py b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py index 099d2b4cc..47e0c181b 100644 --- a/test/registered/kernels/ops/layernorm/test_mhc_kernels.py +++ b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py @@ -1,4 +1,5 @@ from contextlib import nullcontext +from types import SimpleNamespace import pytest import torch @@ -25,7 +26,7 @@ def stated_tp_group(): @pytest.mark.parametrize("hidden_size", [4096, 7168]) -@pytest.mark.parametrize("num_tokens", [0, 1, 8, 17, 32, 64]) +@pytest.mark.parametrize("num_tokens", [0, 1, 6, 8, 17, 32, 64]) @pytest.mark.parametrize("use_norm", [False, True]) def test_mhc_fused_post_pre_matches_unfused( monkeypatch, hidden_size, num_tokens, use_norm, stated_tp_group @@ -107,6 +108,18 @@ def test_mhc_fused_post_pre_matches_unfused( norm_eps=norm_eps, ) + if hidden_size == 4096 and num_tokens in (0, 1, 6, 17): + _check_glm_boundary( + x, + residual, + post_prev, + comb_prev, + fn, + hc_scale, + hc_base, + use_norm=use_norm, + ) + torch.cuda.synchronize() if num_tokens == 0: assert residual_out.shape == residual.shape @@ -136,6 +149,72 @@ def test_mhc_fused_post_pre_matches_unfused( torch.testing.assert_close(layer_out, layer_ref, atol=layer_atol, rtol=layer_rtol) +def _check_glm_boundary(x, residual, post, comb, fn, scale, base, *, use_norm): + from sglang.srt.environ import envs + from sglang.srt.layers.communicator_mhc import MHCState + from sglang.srt.layers.layernorm import RMSNorm + from sglang.srt.models.glm5_next import Glm5NextDecoderLayer + + layer = Glm5NextDecoderLayer.__new__(Glm5NextDecoderLayer) + torch.nn.Module.__init__(layer) + layer.config = SimpleNamespace( + mhc=True, + hc_mult=4, + rms_norm_eps=1e-6, + hc_eps=1e-6, + hc_sinkhorn_iters=20, + ) + layer.hc_ffn_fn = torch.nn.Parameter(fn) + layer.hc_ffn_scale = torch.nn.Parameter(scale) + layer.hc_ffn_base = torch.nn.Parameter(base) + norm = RMSNorm(x.shape[-1], eps=1e-6).to(x) if use_norm else None + states = [ + MHCState( + hc_mult=4, + hc_attn_pre=layer.hc_attn_pre, + hc_ffn_pre=layer.hc_ffn_pre, + hc_post=layer.hc_post, + hc_ffn_post_pre=callback, + h_res=comb.flatten(1), + h_post=post.flatten(1), + ) + for callback in (None, layer.hc_ffn_post_pre) + ] + # Literal, not derived from the cutoff constant: deriving it makes this a + # mirror that stays green when the cutoff moves. None is the empty batch, + # which attn_to_mlp short-circuits before reaching the callback. + fused_expected = {1: True, 6: True, 17: False}[x.shape[0]] if x.shape[0] else None + with envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True): + if x.shape[0] > 0: + declined = ( + layer.hc_ffn_post_pre( + hidden_states=x, + residual=residual.flatten(1), + h_res=comb.flatten(1), + h_post=post.flatten(1), + out_norm_weight=None, + out_norm_eps=None, + ) + is None + ) + assert declined is not fused_expected, ( + f"num_tokens={x.shape[0]} fused={not declined}, " + f"expected fused={fused_expected}" + ) + outputs = [s.attn_to_mlp(x, residual.flatten(1), norm) for s in states] + torch.testing.assert_close(outputs[0][0], outputs[1][0], atol=2e-2, rtol=2e-2) + torch.testing.assert_close(outputs[0][1], outputs[1][1], atol=0, rtol=0) + torch.testing.assert_close(states[0].h_res, states[1].h_res, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(states[0].h_post, states[1].h_post, atol=1e-3, rtol=1e-3) + # The next combine must consume the FFN mixing matrices, not attention's. + torch.testing.assert_close( + states[0].mlp_combine(x, outputs[0][1]), + states[1].mlp_combine(x, outputs[1][1]), + atol=2e-3, + rtol=2e-2, + ) + + if __name__ == "__main__": import sys