perf(glm): fuse MTP input normalization and concatenation

lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py

@@ -1,5 +1,3 @@
-import torch
-
 from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
 
 
@@ -12,11 +10,17 @@ 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]
-        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)
-        return layer_weight.eh_proj_weight_.mm(torch.cat((input_embeddings, previous_hidden), dim=-1))
+        from lightllm.models.glm5_next_mtp.triton_kernel.fused_eh_norm import fused_mtp_norm
+
+        normalized_embeddings = 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=normalized_embeddings)
 
     def context_forward(self, input_ids, infer_state, layer_weight):
         input_embeddings = super().context_forward(input_ids, infer_state, layer_weight)

lightllm/models/glm5_next_mtp/triton_kernel/fused_eh_norm.py

@@ -0,0 +1,97 @@
+# 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 _rms_norm(values: tl.tensor, weight: tl.tensor, eps: float, hidden_size: tl.constexpr) -> tl.tensor:
+    variance = tl.sum(values * values, axis=0) / hidden_size
+    reciprocal_std = 1 / tl.sqrt(variance + eps)
+    return (values * reciprocal_std) * weight
+
+
+@triton.jit
+def _fused_mtp_norm_kernel(
+    embeddings: tl.tensor,
+    previous_hidden: tl.tensor,
+    main_norm_weight: tl.tensor,
+    embedding_norm_weight: tl.tensor,
+    hidden_norm_weight: tl.tensor,
+    output: tl.tensor,
+    embeddings_stride_0: tl.constexpr,
+    embeddings_stride_1: tl.constexpr,
+    previous_hidden_stride_0: tl.constexpr,
+    previous_hidden_stride_1: tl.constexpr,
+    hidden_size: tl.constexpr,
+    eps: float,
+):
+    token_index = tl.program_id(axis=0).to(tl.int64)
+    hidden_index = tl.arange(0, hidden_size)
+    embedding_values = tl.load(embeddings + token_index * embeddings_stride_0 + hidden_index * embeddings_stride_1).to(
+        tl.float32
+    )
+    previous_hidden_values = tl.load(
+        previous_hidden + token_index * previous_hidden_stride_0 + hidden_index * previous_hidden_stride_1
+    ).to(tl.float32)
+    main_norm_scale = tl.load(main_norm_weight + hidden_index).to(tl.float32)
+    embedding_norm_scale = tl.load(embedding_norm_weight + hidden_index).to(tl.float32)
+    hidden_norm_scale = tl.load(hidden_norm_weight + hidden_index).to(tl.float32)
+    # The second hidden norm reads the rounded output of main_norm.
+    main_normalized_hidden = (
+        _rms_norm(values=previous_hidden_values, weight=main_norm_scale, eps=eps, hidden_size=hidden_size)
+        .to(previous_hidden.dtype.element_ty)
+        .to(tl.float32)
+    )
+    normalized_embeddings = _rms_norm(
+        values=embedding_values, weight=embedding_norm_scale, eps=eps, hidden_size=hidden_size
+    ).to(embeddings.dtype.element_ty)
+    normalized_previous_hidden = _rms_norm(
+        values=main_normalized_hidden, weight=hidden_norm_scale, eps=eps, hidden_size=hidden_size
+    ).to(previous_hidden.dtype.element_ty)
+    tl.store(
+        embeddings + token_index * embeddings_stride_0 + hidden_index * embeddings_stride_1,
+        normalized_embeddings,
+    )
+    tl.store(
+        previous_hidden + token_index * previous_hidden_stride_0 + hidden_index * previous_hidden_stride_1,
+        normalized_previous_hidden,
+    )
+    tl.store(output + token_index * 2 * hidden_size + hidden_index, normalized_embeddings)
+    tl.store(
+        output + token_index * 2 * hidden_size + hidden_size + hidden_index,
+        normalized_previous_hidden,
+    )
+
+
+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:
+    """Apply the three GLM MTP RMS norms in place and concatenate the results."""
+    num_tokens, hidden_size = embeddings.shape
+    output = torch.empty(size=(num_tokens, 2 * hidden_size), device=embeddings.device, dtype=embeddings.dtype)
+    _fused_mtp_norm_kernel[(num_tokens,)](
+        embeddings=embeddings,
+        previous_hidden=previous_hidden,
+        main_norm_weight=main_weight,
+        embedding_norm_weight=embedding_weight,
+        hidden_norm_weight=hidden_weight,
+        output=output,
+        embeddings_stride_0=embeddings.stride(0),
+        embeddings_stride_1=embeddings.stride(1),
+        previous_hidden_stride_0=previous_hidden.stride(0),
+        previous_hidden_stride_1=previous_hidden.stride(1),
+        hidden_size=hidden_size,
+        eps=eps,
+        num_warps=8,
+    )
+    return output

test/start_scripts/glm53/glm53_pd_2p1d.sh

@@ -0,0 +1,157 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ "$#" -lt 3 || "$#" -gt 4 ]]; then
+  echo "Usage: $0 <port> <model_dir> <cpu_cache_size_gb> [chat_template]" >&2
+  exit 2
+fi
+
+PORT="$1"
+MODEL_DIR="$2"
+CPU_CACHE_SIZE="$3"
+CHAT_TEMPLATE_ARGS=()
+if [[ -n "${4:-}" ]]; then
+  CHAT_TEMPLATE_ARGS=(--chat_template "$4")
+fi
+
+# GLM-5.3 Flash requires BF16 for dynamic FP8 KV.
+LLM_KV_TYPE="${LLM_KV_TYPE:-fp8kv_dsa}"
+DATA_TYPE="${DATA_TYPE:-bfloat16}"
+MTP_STEP="${MTP_STEP:-3}"
+PD_KV_PAGE_SIZE=16384
+if [[ "$LLM_KV_TYPE" == "fp8kv_dsa" ]]; then
+  PD_KV_PAGE_SIZE=32768
+fi
+
+export LOADWORKER="${LOADWORKER:-8}"
+export LIGHTLLM_TRITON_AUTOTUNE_LEVEL="${LIGHTLLM_TRITON_AUTOTUNE_LEVEL:-1}"
+export LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING=1
+export LIGHTLLM_LOG_LEVEL="${LIGHTLLM_LOG_LEVEL:-info}"
+
+# Keep same-host PD WebSocket traffic away from HTTP/WebSocket proxies.
+export NO_PROXY="${NO_PROXY:+${NO_PROXY},}127.0.0.1,localhost"
+export no_proxy="${NO_PROXY}"
+
+P_COMMON_ARGS=(
+  --model_dir "${MODEL_DIR}"
+  --model_name glm53
+  --data_type "${DATA_TYPE}"
+  --llm_kv_type "${LLM_KV_TYPE}"
+  --graph_max_batch_size 8
+  --running_max_req_size 64
+  --mem_fraction 0.80
+  --max_image_token_count 4096
+  --max_image_pixels 3686400
+  --batch_max_tokens 8192
+  --linear_att_cache_size 500
+  --linear_att_hash_page_size 2048
+  --linear_att_page_block_num 8
+  --linear_att_ssm_data_type float32
+  --enable_fused_shared_experts
+  --mtp_mode eagle_with_att
+  --mtp_draft_model_dir "${MODEL_DIR}"
+  --mtp_step "${MTP_STEP}"
+  --tool_call_parser glm47
+  --reasoning_parser glm45
+  "${CHAT_TEMPLATE_ARGS[@]}"
+  --pd_trans_mode nccl
+  --pd_kv_page_size "${PD_KV_PAGE_SIZE}"
+  --pd_kv_page_num 2
+  --pd_master_ip 127.0.0.1
+  --pd_master_port "${PORT}"
+)
+
+D_COMMON_ARGS=(
+  --model_dir "${MODEL_DIR}"
+  --model_name glm53
+  --data_type "${DATA_TYPE}"
+  --llm_kv_type "${LLM_KV_TYPE}"
+  --graph_max_batch_size 64
+  --graph_max_len_in_batch 65536
+  --running_max_req_size 64
+  --mem_fraction 0.80
+  --batch_max_tokens 256
+  --linear_att_cache_size 500
+  --linear_att_hash_page_size 2048
+  --linear_att_ssm_data_type float32
+  --enable_fused_shared_experts
+  --mtp_mode eagle_with_att
+  --mtp_draft_model_dir "${MODEL_DIR}"
+  --mtp_step "${MTP_STEP}"
+  --tool_call_parser glm47
+  --reasoning_parser glm45
+  "${CHAT_TEMPLATE_ARGS[@]}"
+  --pd_trans_mode nccl
+  --pd_kv_page_size "${PD_KV_PAGE_SIZE}"
+  --pd_kv_page_num 2
+  --pd_master_ip 127.0.0.1
+  --pd_master_port "${PORT}"
+)
+
+PIDS=()
+
+cleanup() {
+  kill -TERM "${PIDS[@]}" 2>/dev/null || true
+  wait "${PIDS[@]}" 2>/dev/null || true
+}
+
+trap cleanup EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+# Prefill 1: TP4 on GPUs 0-3.
+CUDA_VISIBLE_DEVICES=0,1,2,3 python -m lightllm.server.api_server \
+  "${P_COMMON_ARGS[@]}" \
+  --run_mode prefill \
+  --disable_cudagraph \
+  --enable_cpu_cache \
+  --cpu_cache_storage_size "${CPU_CACHE_SIZE}" \
+  --tp 4 \
+  --host 0.0.0.0 \
+  --port 28761 \
+  --nccl_port 29761 &
+PIDS+=("$!")
+
+# Prefill 2: TP4 on GPUs 4-7.
+CUDA_VISIBLE_DEVICES=4,5,6,7 python -m lightllm.server.api_server \
+  "${P_COMMON_ARGS[@]}" \
+  --run_mode prefill \
+  --disable_cudagraph \
+  --enable_cpu_cache \
+  --cpu_cache_storage_size "${CPU_CACHE_SIZE}" \
+  --tp 4 \
+  --host 0.0.0.0 \
+  --port 28762 \
+  --nccl_port 29762 &
+PIDS+=("$!")
+
+# Decode: TP4 on GPUs 8-11.
+CUDA_VISIBLE_DEVICES=8,9,10,11 python -m lightllm.server.api_server \
+  "${D_COMMON_ARGS[@]}" \
+  --run_mode decode \
+  --tp 4 \
+  --host 0.0.0.0 \
+  --port 28763 \
+  --nccl_port 29763 &
+PIDS+=("$!")
+
+# Start PD Master last, after all Prefill and Decode processes have started.
+CUDA_VISIBLE_DEVICES= python -m lightllm.server.api_server \
+  --model_dir "${MODEL_DIR}" \
+  --model_name glm53 \
+  --run_mode pd_master \
+  --pd_master_mode 2p1d \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --max_image_token_count 4096 \
+  --max_image_pixels 3686400 \
+  --tool_call_parser glm47 \
+  --reasoning_parser glm45 \
+  "${CHAT_TEMPLATE_ARGS[@]}" &
+PIDS+=("$!")
+
+echo "GLM-5.3 Flash 2P1D is starting at http://127.0.0.1:${PORT} (P0/P1/D: TP4)"
+
+# A fixed 2P1D deployment is incomplete as soon as any component exits.
+wait -n "${PIDS[@]}"
+exit 1