minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Perf] Fuse the glm5_next mHC attn->MLP boundary (#39200)

Browse Source 
Co-authored-by: mmangkad <mohammad.angkad@radixark.ai>
This commit is contained in:
Mohammad Miadh Angkad
2026-09-20 16:20:37 -07:00
committed by GitHub
co-authored by mmangkad
parent 983e643854
commit 2fa6b94e34
3 changed files with 162 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/communicator_mhc.py
+21 -1
View File
@@ -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__(
python/sglang/srt/models/glm5_next.py
+61 -1
View File
@@ -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(
test/registered/kernels/ops/layernorm/test_mhc_kernels.py
+80 -1
View File
@@ -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