perf(glm): fuse MTP input normalization and concatenation

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)

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

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)