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..c4979abc8a 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,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) 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..401ccdd31d --- /dev/null +++ b/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 diff --git a/test/start_scripts/glm53/glm53_pd_2p1d.sh b/test/start_scripts/glm53/glm53_pd_2p1d.sh new file mode 100755 index 0000000000..1c90d996ca --- /dev/null +++ b/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 [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 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..c92ae5b86a --- /dev/null +++ b/unit_tests/models/glm5_next/test_fused_mtp_norm.py @@ -0,0 +1,126 @@ +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", [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)) + expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + 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", [4096, 1024]) +def test_model_integration(hidden: int) -> 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) + + 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)