From c2e7e21a32ab7792ab92c4fb41ba78a30d539b9f Mon Sep 17 00:00:00 2001 From: baishihao Date: Thu, 12 Nov 2026 19:10:33 +0800 Subject: [PATCH] perf(glm): fuse MTP input normalization and concatenation --- .../layer_infer/pre_layer_infer.py | 19 +++ .../glm5_next_mtp/triton_kernel/__init__.py | 0 .../triton_kernel/fused_eh_norm.py | 81 +++++++++++ .../models/glm5_next/test_fused_mtp_norm.py | 130 ++++++++++++++++++ 4 files changed, 230 insertions(+) create mode 100644 lightllm/models/glm5_next_mtp/triton_kernel/__init__.py create mode 100644 lightllm/models/glm5_next_mtp/triton_kernel/fused_eh_norm.py create mode 100644 unit_tests/models/glm5_next/test_fused_mtp_norm.py diff --git a/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py index 50d0461822..3ec1967512 100644 --- a/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py +++ b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py @@ -1,5 +1,6 @@ import torch +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_num_warps from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer @@ -13,6 +14,24 @@ def __init__(self, network_config): def _fuse_hidden(self, input_embeddings, infer_state, layer_weight): previous_hidden = infer_state.mtp_draft_input_hiddens assert input_embeddings.shape[0] == previous_hidden.shape[0] + if ( + rmsnorm_num_warps == 8 + and input_embeddings.is_cuda + and input_embeddings.shape[-1] == 4096 + and input_embeddings.dtype in (torch.bfloat16, torch.float16) + and previous_hidden.dtype == input_embeddings.dtype + ): + from lightllm.models.glm5_next_mtp.triton_kernel.fused_eh_norm import fused_mtp_norm + + normed = fused_mtp_norm( + embeddings=input_embeddings, + previous_hidden=previous_hidden, + main_weight=layer_weight.main_norm_weight_.weight, + embedding_weight=layer_weight.enorm_weight_.weight, + hidden_weight=layer_weight.hnorm_weight_.weight, + eps=self.eps_, + ) + return layer_weight.eh_proj_weight_.mm(input_tensor=normed) layer_weight.main_norm_weight_(input=previous_hidden, eps=self.eps_, out=previous_hidden) layer_weight.enorm_weight_(input=input_embeddings, eps=self.eps_, out=input_embeddings) layer_weight.hnorm_weight_(input=previous_hidden, eps=self.eps_, out=previous_hidden) diff --git a/lightllm/models/glm5_next_mtp/triton_kernel/__init__.py b/lightllm/models/glm5_next_mtp/triton_kernel/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next_mtp/triton_kernel/fused_eh_norm.py b/lightllm/models/glm5_next_mtp/triton_kernel/fused_eh_norm.py new file mode 100644 index 0000000000..1838428c41 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/triton_kernel/fused_eh_norm.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# Layout adapted from vLLM's fused_eh_norm.py (5082c45853ca93ae809c58255e6126cfe20c415f). +# Copyright contributors to the vLLM project. +# Retain LightLLM's main norm, intermediate rounding and in-place updates. + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _norm(values: tl.tensor, weight: tl.tensor, eps: float, H: tl.constexpr) -> tl.tensor: + variance = tl.sum(values * values, axis=0) / H + reciprocal_std = 1 / tl.sqrt(variance + eps) + return (values * reciprocal_std) * weight + + +@triton.jit +def _fused_mtp_norm( + E: tl.tensor, + P: tl.tensor, + WM: tl.tensor, + WE: tl.tensor, + WH: tl.tensor, + O: tl.tensor, + es0: tl.constexpr, + es1: tl.constexpr, + ps0: tl.constexpr, + ps1: tl.constexpr, + H: tl.constexpr, + eps: float, +): + row = tl.program_id(axis=0).to(tl.int64) + col = tl.arange(0, H) + e = tl.load(E + row * es0 + col * es1).to(tl.float32) + p = tl.load(P + row * ps0 + col * ps1).to(tl.float32) + wm = tl.load(WM + col).to(tl.float32) + we = tl.load(WE + col).to(tl.float32) + wh = tl.load(WH + col).to(tl.float32) + # The second hidden norm reads the rounded output of main_norm. + normalized_main = _norm(values=p, weight=wm, eps=eps, H=H).to(P.dtype.element_ty).to(tl.float32) + normalized_e = _norm(values=e, weight=we, eps=eps, H=H).to(E.dtype.element_ty) + normalized_p = _norm(values=normalized_main, weight=wh, eps=eps, H=H).to(P.dtype.element_ty) + tl.store(E + row * es0 + col * es1, normalized_e) + tl.store(P + row * ps0 + col * ps1, normalized_p) + tl.store(O + row * 2 * H + col, normalized_e) + tl.store(O + row * 2 * H + H + col, normalized_p) + + +def fused_mtp_norm( + embeddings: torch.Tensor, + previous_hidden: torch.Tensor, + main_weight: torch.Tensor, + embedding_weight: torch.Tensor, + hidden_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Normalize GLM draft inputs in place and concatenate them for eh_proj.""" + assert embeddings.ndim == previous_hidden.ndim == 2 + assert embeddings.shape == previous_hidden.shape and embeddings.shape[1] == 4096 + assert embeddings.dtype == previous_hidden.dtype + assert all(w.shape == (4096,) and w.is_contiguous() for w in (main_weight, embedding_weight, hidden_weight)) + rows, hidden = embeddings.shape + out = torch.empty(size=(rows, 2 * hidden), device=embeddings.device, dtype=embeddings.dtype) + if rows: + _fused_mtp_norm[(rows,)]( + E=embeddings, + P=previous_hidden, + WM=main_weight, + WE=embedding_weight, + WH=hidden_weight, + O=out, + es0=embeddings.stride(0), + es1=embeddings.stride(1), + ps0=previous_hidden.stride(0), + ps1=previous_hidden.stride(1), + H=hidden, + eps=eps, + num_warps=8, + ) + return out diff --git a/unit_tests/models/glm5_next/test_fused_mtp_norm.py b/unit_tests/models/glm5_next/test_fused_mtp_norm.py new file mode 100644 index 0000000000..8449843f42 --- /dev/null +++ b/unit_tests/models/glm5_next/test_fused_mtp_norm.py @@ -0,0 +1,130 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next_mtp.layer_infer import pre_layer_infer as mtp +from lightllm.models.glm5_next_mtp.triton_kernel.fused_eh_norm import fused_mtp_norm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def reference( + embeddings: torch.Tensor, + previous_hidden: torch.Tensor, + weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + eps: float, +) -> torch.Tensor: + rmsnorm_forward(x=previous_hidden, weight=weights[0], eps=eps, out=previous_hidden) + rmsnorm_forward(x=embeddings, weight=weights[1], eps=eps, out=embeddings) + rmsnorm_forward(x=previous_hidden, weight=weights[2], eps=eps, out=previous_hidden) + return torch.cat(tensors=(embeddings, previous_hidden), dim=-1) + + +@pytest.mark.parametrize("rows", [0, 1, 4, 16, 64, 512, 8192]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("scale", [1e-4, 1.0, 100.0]) +def test_norm_rounding_and_inplace_updates(rows: int, dtype: torch.dtype, scale: float) -> None: + torch.manual_seed(seed=5514) + hidden = 4096 + e = (torch.randn(size=(rows, hidden + 16), dtype=dtype, device="cuda") * scale)[:, :hidden] + p = (torch.randn(size=(rows, hidden + 16), dtype=dtype, device="cuda") * scale)[:, :hidden] + er, pr = e.clone(), p.clone() + weights = tuple(torch.randn(size=(hidden,), dtype=dtype, device="cuda") * 0.2 + 1 for _ in range(3)) + if rows: + expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + else: + expected = torch.empty(size=(0, 2 * hidden), dtype=dtype, device="cuda") + actual = fused_mtp_norm( + embeddings=e, + previous_hidden=p, + main_weight=weights[0], + embedding_weight=weights[1], + hidden_weight=weights[2], + eps=1e-5, + ) + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) + torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) + + +def test_graph_replay_changes_both_inputs_and_preserves_strided_storage() -> None: + torch.manual_seed(seed=5515) + rows, hidden = 16, 4096 + e_storage = torch.randn(size=(rows, hidden * 2), device="cuda", dtype=torch.bfloat16) + p_storage = torch.randn(size=(rows, hidden * 2), device="cuda", dtype=torch.bfloat16) + e, p = e_storage[:, ::2], p_storage[:, ::2] + weights = tuple(torch.randn(size=(hidden,), device="cuda", dtype=e.dtype) * 0.2 + 1 for _ in range(3)) + fused_mtp_norm( + embeddings=e, + previous_hidden=p, + main_weight=weights[0], + embedding_weight=weights[1], + hidden_weight=weights[2], + eps=1e-5, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph=graph): + actual = fused_mtp_norm( + embeddings=e, + previous_hidden=p, + main_weight=weights[0], + embedding_weight=weights[1], + hidden_weight=weights[2], + eps=1e-5, + ) + for scale in (0.01, 10.0): + e_storage.normal_().mul_(scale) + p_storage.normal_().mul_(scale) + untouched_e, untouched_p = e_storage[:, 1::2].clone(), p_storage[:, 1::2].clone() + er, pr = e.clone(), p.clone() + expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + graph.replay() + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) + torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) + assert torch.equal(input=e_storage[:, 1::2], other=untouched_e) + assert torch.equal(input=p_storage[:, 1::2], other=untouched_p) + + +@pytest.mark.parametrize("hidden,warps", [(4096, 8), (4096, 4), (1024, 8)]) +def test_model_integration_and_fallback(hidden: int, warps: int, monkeypatch: pytest.MonkeyPatch) -> None: + torch.manual_seed(seed=5516) + e = torch.randn(size=(7, hidden), dtype=torch.bfloat16, device="cuda") + p = torch.randn_like(input=e) + er, pr = e.clone(), p.clone() + weights = tuple(torch.randn(size=(hidden,), dtype=e.dtype, device="cuda") * 0.2 + 1 for _ in range(3)) + + class Norm: + def __init__(self, weight: torch.Tensor) -> None: + self.weight = weight + + def __call__(self, input: torch.Tensor, eps: float, out: torch.Tensor) -> torch.Tensor: + return rmsnorm_forward(x=input, weight=self.weight, eps=eps, out=out) + + class Projection: + def __init__(self) -> None: + self.weight = torch.randn(size=(2 * hidden, 128), dtype=e.dtype, device=e.device) * 0.1 + + def mm(self, input_tensor: torch.Tensor) -> torch.Tensor: + return torch.mm(input=input_tensor, mat2=self.weight) + + monkeypatch.setattr(target=mtp, name="rmsnorm_num_warps", value=warps) + layer = mtp.Glm5NextMTPPreLayerInfer.__new__(mtp.Glm5NextMTPPreLayerInfer) + layer.eps_ = 1e-5 + layer_weights = SimpleNamespace( + main_norm_weight_=Norm(weight=weights[0]), + enorm_weight_=Norm(weight=weights[1]), + hnorm_weight_=Norm(weight=weights[2]), + eh_proj_weight_=Projection(), + ) + expected = layer_weights.eh_proj_weight_.mm( + input_tensor=reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + ) + actual = layer._fuse_hidden( + input_embeddings=e, infer_state=SimpleNamespace(mtp_draft_input_hiddens=p), layer_weight=layer_weights + ) + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) + torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0)