From 8a8f330e88f159d73588126ed999be8b55b9c239 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 19:10:24 +0800 Subject: [PATCH 01/28] feat: support GLM-5.3 Flash inference --- .../common/basemodel/attention/__init__.py | 1 + .../common/basemodel/attention/base_att.py | 12 + .../basemodel/attention/create_utils.py | 4 +- .../basemodel/attention/linear/__init__.py | 2 + .../common/basemodel/attention/linear/kda.py | 313 +++++ .../basemodel/attention/nsa/__init__.py | 1 + .../attention/nsa/flashmla_sparse.py | 126 +- .../attention/nsa/tilelang_sparse.py | 84 ++ lightllm/common/basemodel/basemodel.py | 130 +- lightllm/common/basemodel/batch_objs.py | 67 +- lightllm/common/basemodel/cuda_graph.py | 26 +- lightllm/common/basemodel/hidden_collector.py | 6 + lightllm/common/basemodel/infer_struct.py | 7 + .../transformer_layer_infer_template.py | 29 + .../fused_moe/fused_moe_weight.py | 39 +- .../fused_moe/impl/deepgemm_impl.py | 96 +- .../fused_moe/impl/triton_impl.py | 218 +++ lightllm/common/basemodel/mtp_manager.py | 37 +- .../common/basemodel/prefill_cuda_graph.py | 140 +- .../fused_moe/deepep_legacy_layout.py | 261 ++++ .../fused_moe/grouped_fused_moe.py | 11 + .../fused_moe/grouped_fused_moe_ep.py | 349 ++++- .../triton_kernel/fused_moe/grouped_topk.py | 166 +++ .../fused_moe/moe_silu_and_mul.py | 11 +- .../moe_silu_and_mul_mix_quant_ep.py | 22 +- .../triton_kernel/linear_att/causal_conv1d.py | 29 +- .../linear_att/fla/ops/__init__.py | 4 + .../linear_att/fla/ops/chunk_delta_h.py | 17 +- .../linear_att/fla/ops/fused_recurrent.py | 71 +- .../triton_kernel/linear_att/fla/ops/index.py | 5 +- .../triton_kernel/linear_att/fla/ops/kda.py | 1170 +++++++++++++++++ .../triton_kernel/linear_att/fla/ops/op.py | 1 + .../linear_att/fla/ops/solve_tril.py | 4 +- .../post_process/greedy_sample.py | 109 ++ .../post_process/vocab_parallel_greedy.py | 126 ++ .../triton_kernel/transpose_convert.py | 65 + .../linear_att_cache_manager/config_objs.py | 33 +- lightllm/distributed/communication_op.py | 84 +- lightllm/distributed/symm_mem_all_reduce.py | 47 +- lightllm/models/__init__.py | 2 + .../layer_infer/transformer_layer_infer.py | 45 +- .../layer_infer/transformer_layer_infer.py | 63 +- .../triton_kernel/extract_indexer_ks.py | 149 ++- .../triton_kernel/topk_index_to_mem_index.py | 12 +- .../gemma4/layer_infer/post_layer_infer.py | 12 +- lightllm/models/glm5_next/__init__.py | 1 + lightllm/models/glm5_next/infer_struct.py | 5 + .../models/glm5_next/layer_infer/__init__.py | 3 + .../layer_infer/transformer_layer_infer.py | 325 +++++ .../glm5_next/layer_weights/__init__.py | 4 + .../pre_and_post_layer_weight.py | 20 + .../layer_weights/transformer_layer_weight.py | 286 ++++ lightllm/models/glm5_next/mem_manager.py | 55 + lightllm/models/glm5_next/model.py | 200 +++ .../glm5_next/triton_kernel/__init__.py | 3 + .../models/glm5_next/triton_kernel/mhc.py | 604 +++++++++ lightllm/models/glm5_next_mtp/__init__.py | 3 + .../glm5_next_mtp/layer_infer/__init__.py | 5 + .../layer_infer/post_layer_infer.py | 97 ++ .../layer_infer/pre_layer_infer.py | 19 + .../glm5_next_mtp/layer_weights/__init__.py | 5 + .../pre_and_post_layer_weight.py | 49 + lightllm/models/glm5_next_mtp/model.py | 92 ++ .../glm5_next_mtp/triton_kernel/__init__.py | 5 + .../triton_kernel/zero_position_embedding.py | 47 + .../llama/layer_infer/post_layer_infer.py | 25 +- .../layer_infer/post_layer_infer.py | 6 +- .../layer_infer/transformer_layer_infer.py | 14 +- lightllm/server/api_cli.py | 50 +- lightllm/server/api_start.py | 22 +- lightllm/server/core/objs/start_args_type.py | 7 +- lightllm/server/router/manager.py | 61 +- .../model_infer/mode_backend/base_backend.py | 48 +- .../mode_backend/chunked_prefill/impl.py | 73 +- .../chunked_prefill/impl_for_reward_model.py | 1 + .../mode_backend/diverse_backend/impl.py | 9 +- .../mode_backend/dp_backend/impl.py | 51 +- .../mode_backend/generic_post_process.py | 43 +- .../mode_backend/generic_pre_process.py | 11 + .../dp_overlap_proposers/eagle_with_att.py | 2 +- .../mtp_speculative/planner/lightspec.py | 10 +- .../proposers/eagle_with_att.py | 25 +- lightllm/utils/device_utils.py | 5 + lightllm/utils/envs_utils.py | 30 +- .../kernel/test_extract_indexer_ks_dynamic.py | 109 ++ test/kernel/test_glm5_grouped_topk.py | 229 ++++ test/kernel/test_glm5_mhc.py | 290 ++++ test/kernel/test_glm5_sglang_moe_compat.py | 251 ++++ test/kernel/test_glm5_short_decode_helpers.py | 23 + test/kernel/test_glm5_strided_causal_conv.py | 134 ++ test/kernel/test_glm5_vocab_parallel_top1.py | 112 ++ test/test_moe_prefill_dispatch.py | 47 + tools/analyze_torch_trace.py | 67 + tools/bench_glm53_allreduce.py | 169 +++ tools/bench_glm53_kda_chunk_h.py | 125 ++ tools/bench_glm53_sglang_moe.py | 264 ++++ tools/bench_glm53_sparse_prefill.py | 164 +++ tools/bench_glm53_sparse_prefill_tp.py | 264 ++++ tools/check_glm53_symm_out_of_place.py | 60 + .../attention/test_flashmla_sparse_tp.py | 86 ++ .../attention/test_tilelang_sparse.py | 27 + .../basemodel/test_cuda_graph_layout.py | 28 + .../common/basemodel/test_hidden_collector.py | 4 + .../common/basemodel/test_model_output.py | 9 +- .../common/basemodel/test_mtp_manager.py | 38 +- .../test_prefill_cuda_graph_selection.py | 121 ++ .../test_sglang_triton_moe_config.py | 100 ++ .../linear_att/test_kda_fused_gate.py | 80 ++ .../test_vocab_parallel_greedy.py | 65 + .../test_topk_index_to_mem_index.py | 22 +- .../models/gemma4/test_post_layer_infer.py | 99 ++ .../test_vocab_parallel_greedy_output.py | 74 ++ .../test_chunked_prefill_mega_moe_overlap.py | 64 + .../test_vocab_parallel_greedy_sampling.py | 91 ++ .../mtp_speculative/test_planner.py | 4 +- .../server/router/test_dp_model_capacity.py | 39 + .../server/router/test_prefill_coalescing.py | 69 + unit_tests/server/test_mtp_start_args.py | 6 +- unit_tests/utils/test_envs_utils.py | 74 ++ 119 files changed, 9666 insertions(+), 269 deletions(-) create mode 100644 lightllm/common/basemodel/attention/linear/kda.py create mode 100644 lightllm/common/basemodel/attention/nsa/tilelang_sparse.py create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py create mode 100644 lightllm/common/basemodel/triton_kernel/transpose_convert.py create mode 100644 lightllm/models/glm5_next/__init__.py create mode 100644 lightllm/models/glm5_next/infer_struct.py create mode 100644 lightllm/models/glm5_next/layer_infer/__init__.py create mode 100644 lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py create mode 100644 lightllm/models/glm5_next/layer_weights/__init__.py create mode 100644 lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py create mode 100644 lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py create mode 100644 lightllm/models/glm5_next/mem_manager.py create mode 100644 lightllm/models/glm5_next/model.py create mode 100644 lightllm/models/glm5_next/triton_kernel/__init__.py create mode 100644 lightllm/models/glm5_next/triton_kernel/mhc.py create mode 100644 lightllm/models/glm5_next_mtp/__init__.py create mode 100644 lightllm/models/glm5_next_mtp/layer_infer/__init__.py create mode 100644 lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py create mode 100644 lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py create mode 100644 lightllm/models/glm5_next_mtp/layer_weights/__init__.py create mode 100644 lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py create mode 100644 lightllm/models/glm5_next_mtp/model.py create mode 100644 lightllm/models/glm5_next_mtp/triton_kernel/__init__.py create mode 100644 lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py create mode 100644 test/kernel/test_extract_indexer_ks_dynamic.py create mode 100644 test/kernel/test_glm5_grouped_topk.py create mode 100644 test/kernel/test_glm5_mhc.py create mode 100644 test/kernel/test_glm5_sglang_moe_compat.py create mode 100644 test/kernel/test_glm5_short_decode_helpers.py create mode 100644 test/kernel/test_glm5_strided_causal_conv.py create mode 100644 test/kernel/test_glm5_vocab_parallel_top1.py create mode 100644 test/test_moe_prefill_dispatch.py create mode 100644 tools/analyze_torch_trace.py create mode 100644 tools/bench_glm53_allreduce.py create mode 100644 tools/bench_glm53_kda_chunk_h.py create mode 100644 tools/bench_glm53_sglang_moe.py create mode 100644 tools/bench_glm53_sparse_prefill.py create mode 100644 tools/bench_glm53_sparse_prefill_tp.py create mode 100644 tools/check_glm53_symm_out_of_place.py create mode 100644 unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py create mode 100644 unit_tests/common/basemodel/attention/test_tilelang_sparse.py create mode 100644 unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py create mode 100644 unit_tests/common/basemodel/test_sglang_triton_moe_config.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py create mode 100644 unit_tests/models/gemma4/test_post_layer_infer.py create mode 100644 unit_tests/models/test_vocab_parallel_greedy_output.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py create mode 100644 unit_tests/server/router/test_dp_model_capacity.py create mode 100644 unit_tests/server/router/test_prefill_coalescing.py create mode 100644 unit_tests/utils/test_envs_utils.py diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py index 10cd3b0864..1cd6d0aff7 100644 --- a/lightllm/common/basemodel/attention/__init__.py +++ b/lightllm/common/basemodel/attention/__init__.py @@ -13,6 +13,7 @@ # NSA backend from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend +from .nsa.tilelang_sparse import NsaTilelangSparseAttBackend from .create_utils import ( get_prefill_att_backend_class, diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index a7e2d8122a..6cd87e9f2d 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -120,6 +120,18 @@ class BasePrefillAttState(ABC): backend: BaseAttBackend = None infer_state: "InferStateInfo" = None + def copy_for_prefill_cuda_graph(self, new_state: "BasePrefillAttState"): + """Refresh fixed-address attention metadata before graph replay.""" + for attr_name, attr_value in vars(new_state).items(): + if isinstance(attr_value, torch.Tensor): + graph_attr = getattr(self, attr_name, None) + if ( + graph_attr is not None + and graph_attr.data_ptr() != attr_value.data_ptr() + and graph_attr.shape == attr_value.shape + ): + graph_attr.copy_(attr_value, non_blocking=True) + @abstractmethod def init_state(self): pass diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py index f708357765..f7ac37943e 100644 --- a/lightllm/common/basemodel/attention/create_utils.py +++ b/lightllm/common/basemodel/attention/create_utils.py @@ -17,6 +17,7 @@ from .flashinfer.mla import MlaFlashInferAttBackend from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend +from .nsa.tilelang_sparse import NsaTilelangSparseAttBackend logger = init_logger(__name__) @@ -56,7 +57,8 @@ nsa_data_type_to_backend = { "None": { "flashmla_sparse": NsaFlashMlaSparseAttBackend, - # Future backends: "fa3", "tilelang", "aiter" + "tilelang": NsaTilelangSparseAttBackend, + # Future backends: "fa3", "aiter" }, "fp8kv_dsa": { "flashmla_sparse": NsaFlashMlaFp8SparseAttBackend, diff --git a/lightllm/common/basemodel/attention/linear/__init__.py b/lightllm/common/basemodel/attention/linear/__init__.py index bfc8445390..b1a3d1ca34 100644 --- a/lightllm/common/basemodel/attention/linear/__init__.py +++ b/lightllm/common/basemodel/attention/linear/__init__.py @@ -5,6 +5,7 @@ ) from .flashqla import FlashQlaLinearAttBackend from .triton import TritonLinearAttBackend +from .kda import KDALinearAttBackend __all__ = [ "LinearAttBackend", @@ -12,4 +13,5 @@ "LinearAttDecodeAttState", "FlashQlaLinearAttBackend", "TritonLinearAttBackend", + "KDALinearAttBackend", ] diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py new file mode 100644 index 0000000000..b596d18e3f --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""KDA attention backend for GLM-5-Next.""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class KDALinearAttBackend(BaseAttBackend): + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_projection_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = get_env_start_args().mtp_step + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return KDAPrefillAttState(backend=self, infer_state=infer_state) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return KDADecodeAttState(backend=self, infer_state=infer_state) + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_projection_size, dim=-1) + + def reshape_qkv(self, value: torch.Tensor, *, decode: bool): + if decode: + return value.view(-1, 1, self.tp_num_heads, self.head_dim) + return value.view(1, -1, self.tp_num_heads, self.head_dim) + + +@dataclasses.dataclass +class KDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + chunk_indices: torch.Tensor = None + seq_lens_cpu: list[int] = None + + def init_state(self): + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) + self.seq_lens_cpu = ( + self.infer_state.b1_cu_q_seq_len[1:] + - self.infer_state.b1_cu_q_seq_len[:-1] + ).tolist() + # prepare_chunk_indices performs a GPU-to-CPU shape sync. Build it + # before entering CUDA Graph capture and copy its fixed-size contents + # through BasePrefillAttState on replay. + self.chunk_indices = prepare_chunk_indices(self.infer_state.b1_cu_q_seq_len, 64) + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + # MTP widens each request's convolution cache so the verify path can + # retain all candidate states. Prefill only populates the canonical + # history, whose width remains kernel_size - 1. + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer_weight.get_merged_kda_conv_weight(), + bias=None, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + seq_lens_cpu=self.seq_lens_cpu, + ).transpose(0, 1) + + q, k, v = [backend.reshape_qkv(x, decode=False) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(1, -1, backend.tp_projection_size) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + + initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view( + 1, -1, backend.tp_num_heads, backend.head_dim + ), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + chunk_indices=self.chunk_indices, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to( + ssm_states.dtype, copy=False + ) + return output + + +@dataclasses.dataclass +class KDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + draft_step = self.backend.model.mtp_manager.get_decode_draft_step( + self.backend.model.is_mtp_draft_model + ) + if draft_step == 0: + self._init_normal_decode_state() + elif self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(draft_step + 1) + else: + self._init_fixed_mtp_decode_state(draft_step) + + def _init_normal_decode_state(self): + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx + + def _init_dynamic_mtp_decode_state(self, mtp_size: int): + ( + self.b1_mtp_cu_q_seq_len, + self.b_conv_buffer_idx, + self.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=self.infer_state.b_req_idx, + b_mtp_index=self.infer_state.b_mtp_index, + req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, + ) + self._init_mtp_ssm_buffer_idx(mtp_size) + + def _init_fixed_mtp_decode_state(self, draft_step: int): + mtp_size = draft_step + 1 + batch_size = self.infer_state.batch_size + assert batch_size % mtp_size == 0, ( + "KDA fixed-layout decode requires batch_size to be divisible by draft_step + 1, " + f"got batch_size={batch_size}, draft_step={draft_step}." + ) + + att_batch_size = batch_size // mtp_size + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, + batch_size + 1, + mtp_size, + dtype=torch.int32, + device=self.infer_state.b_req_idx.device, + ) + self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[ + self.b_conv_buffer_idx + ] + 1 + self._init_mtp_ssm_buffer_idx(mtp_size) + + def _init_mtp_ssm_buffer_idx(self, mtp_size: int): + att_batch_size = self.b_conv_buffer_idx.shape[0] + b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) + state_offsets = torch.arange( + mtp_size, + device=self.infer_state.b_req_idx.device, + dtype=self.infer_state.b_req_idx.dtype, + ).view(1, mtp_size) + self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode + params = att_control.linear_att_decode_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + if draft_step > 0: + return self._kda_mtp_kernel( + mixed_qkv=mixed_qkv, + raw_gate=raw_gate, + raw_beta=raw_beta, + conv_states=conv_states, + ssm_states=ssm_states, + layer_weight=layer_weight, + draft_step=draft_step, + ) + + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states, + layer_weight.get_merged_kda_conv_weight(), + bias=None, + activation="silu", + conv_state_indices=self.b_conv_buffer_idx, + ) + q, k, v = [backend.reshape_qkv(x, decode=True) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(-1, 1, backend.tp_projection_size) + raw_beta = raw_beta.view(-1, 1, backend.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer_weight.linear_A_log.weight, + gate_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + ssm_state_indices=self.b_ssm_buffer_idx, + ) + return output + + def _kda_mtp_kernel( + self, + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + layer_weight, + draft_step: int, + ): + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + backend: KDALinearAttBackend = self.backend + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer_weight.get_merged_kda_conv_weight(), + mtp_step=draft_step, + bias=None, + activation="silu", + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + + q, k, v = [backend.reshape_qkv(x, decode=False) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(1, -1, backend.tp_projection_size) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + assert self.b_ssm_buffer_idx.dim() == 2, "KDA MTP SSM buffer idx must be 2D [N, S+1]" + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer_weight.linear_A_log.weight, + gate_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + cu_seqlens=self.b1_mtp_cu_q_seq_len.to(torch.long), + ssm_state_indices=self.b_ssm_buffer_idx, + ssm_state_write_indices=self.b_ssm_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + ) + return output diff --git a/lightllm/common/basemodel/attention/nsa/__init__.py b/lightllm/common/basemodel/attention/nsa/__init__.py index f9db52dc2b..b183461504 100644 --- a/lightllm/common/basemodel/attention/nsa/__init__.py +++ b/lightllm/common/basemodel/attention/nsa/__init__.py @@ -10,6 +10,7 @@ NsaFlashMlaFp8SparsePrefillAttState, NsaFlashMlaFp8SparseDecodeAttState, ) +from .tilelang_sparse import NsaTilelangSparseAttBackend, NsaTilelangSparsePrefillAttState __all__ = [ "NsaFlashMlaSparseAttBackend", diff --git a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py index c25f1d31e7..6bf4ca4334 100644 --- a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py @@ -3,15 +3,70 @@ import dataclasses import torch +import torch.distributed as dist from typing import Tuple, TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from lightllm.utils.dist_utils import get_current_device_id +from lightllm.utils.dist_utils import get_current_rank_in_dp +from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo +_TP_HEAD_TOKEN_TRANSPOSE_MIN_TOKENS = 4096 + + +def _copy_received_head_shards(received: torch.Tensor, output: torch.Tensor, world_size: int) -> None: + """Transpose all-to-all receive order from rank-major to token-major heads.""" + + tokens, local_heads, head_dim = received.shape + tokens_per_rank = tokens // world_size + output.view(tokens_per_rank, world_size, local_heads, head_dim).copy_( + received.view(world_size, tokens_per_rank, local_heads, head_dim).permute(1, 0, 2, 3) + ) + + +def _copy_token_shard_for_head_scatter(output: torch.Tensor, send: torch.Tensor, world_size: int) -> None: + """Transpose token-major global heads into all-to-all destination order.""" + + tokens_per_rank, global_heads, head_dim = output.shape + local_heads = global_heads // world_size + send.view(world_size, tokens_per_rank, local_heads, head_dim).copy_( + output.view(tokens_per_rank, world_size, local_heads, head_dim).permute(1, 0, 2, 3) + ) + + +def _should_use_tp_head_token_transpose( + q: torch.Tensor, + infer_state: "InferStateInfo", + required_heads: int, +) -> bool: + """Select the exact TP transpose only for its validated serving layout.""" + + args = get_env_start_args() + world_size = infer_state.dist_group.dp_world_size + return ( + world_size > 1 + and q.is_contiguous() + and q.shape[0] >= _TP_HEAD_TOKEN_TRANSPOSE_MIN_TOKENS + and q.shape[0] % world_size == 0 + and q.shape[1] * world_size == required_heads + and infer_state.max_cache_len == 0 + and not infer_state.need_dp_prefill_balance + and not infer_state.use_replicated_attention_ep + and not args.enable_tpsp_mix_mode + and not args.enable_prefill_cudagraph + and not args.enable_prefill_microbatch_overlap + and not args.enable_prefill_decode_mixed + ) + + +def _alloc_like(input_: torch.Tensor, shape: Tuple[int, ...]) -> torch.Tensor: + return torch.empty(shape, dtype=input_.dtype, device=input_.device) + + class NsaFlashMlaSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) @@ -86,13 +141,70 @@ def _nsa_prefill_att( if topk_mem_indices.ndim == 2: topk_mem_indices = topk_mem_indices.unsqueeze(1) + # The FlashMLA sparse kernels require 64 query heads on Hopper and + # 128 on Blackwell. Tensor parallelism can leave fewer local heads + # (GLM-5.3 has 64 / TP8 = 8), so pad the inactive heads and trim the + # result just as SGLang's DSA backend does. + num_tokens, num_heads, head_dim = q.shape + device_sm_major = torch.cuda.get_device_capability(q.device)[0] + required_heads = 128 if device_sm_major >= 10 else 64 + need_padding = num_heads % required_heads != 0 + if need_padding and _should_use_tp_head_token_transpose(q, self.infer_state, required_heads): + world_size = self.infer_state.dist_group.dp_world_size + rank = get_current_rank_in_dp() + tokens_per_rank = num_tokens // world_size + + received_q = _alloc_like(q, q.shape) + dist.all_to_all_single( + received_q, + q, + group=self.infer_state.dist_group.device_group, + ) + transposed_q = _alloc_like(q, (tokens_per_rank, world_size * num_heads, head_dim)) + _copy_received_head_shards(received_q, transposed_q, world_size) + del received_q + + token_start = rank * tokens_per_rank + local_indices = topk_mem_indices[token_start : token_start + tokens_per_rank] + transposed_out, _, _ = flash_mla_sparse_fwd( + q=transposed_q, + kv=kv, + indices=local_indices, + sm_scale=softmax_scale, + d_v=kv_lora_rank, + ) + + send_out = _alloc_like(q, q.shape) + _copy_token_shard_for_head_scatter(transposed_out, send_out, world_size) + del transposed_q, transposed_out + output = _alloc_like(q, q.shape) + dist.all_to_all_single( + output, + send_out, + group=self.infer_state.dist_group.device_group, + ) + del send_out + return output + + if need_padding: + assert required_heads % num_heads == 0, ( + f"num_heads {num_heads} cannot be padded to {required_heads}; " + "the tensor-parallel size is unsupported" + ) + q_input = q.new_zeros((num_tokens, required_heads, head_dim)) + q_input[:, :num_heads, :] = q + else: + q_input = q + mla_out, _, _ = flash_mla_sparse_fwd( - q=q, + q=q_input, kv=kv, indices=topk_mem_indices, sm_scale=softmax_scale, d_v=kv_lora_rank, ) + if need_padding: + mla_out = mla_out[:, :num_heads, :] return mla_out @@ -173,14 +285,20 @@ def _nsa_decode_att( q_nope, q_rope = q # Extract k_rope and kv_nope from the KV buffer - k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) - kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) + only_qv = qk_rope_head_dim == 0 + if only_qv: + k_rope = None + kv_nope = kv[:, :, :kv_lora_rank].view(-1, 1, 1, kv_lora_rank) + else: + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) o_tensor = flash_attn_with_kvcache( - q=q_rope, + q=None if only_qv else q_rope, k_cache=k_rope, v_cache=kv_nope, qv=q_nope, + only_qv=only_qv, page_table=topk_mem_indices, cache_seqlens=self.nsa_cache_seqlens, cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, diff --git a/lightllm/common/basemodel/attention/nsa/tilelang_sparse.py b/lightllm/common/basemodel/attention/nsa/tilelang_sparse.py new file mode 100644 index 0000000000..33578d6335 --- /dev/null +++ b/lightllm/common/basemodel/attention/nsa/tilelang_sparse.py @@ -0,0 +1,84 @@ +"""TileLang sparse prefill attention for NSA/DSA models. + +The kernel is provided by SGLang's kernel package. Decode intentionally keeps +using the FlashMLA implementation: TileLang is selected only for the much +larger prefill workload where it is advantageous on Hopper. +""" + +import dataclasses + +import torch + +from ..base_att import AttControl +from .flashmla_sparse import ( + NsaFlashMlaSparseAttBackend, + NsaFlashMlaSparsePrefillAttState, +) + + +def pad_sparse_indices(indices: torch.Tensor, block_size: int = 64) -> torch.Tensor: + """Mask-pad the sparse index table to the TileLang block width.""" + + if indices.ndim == 2: + indices = indices.unsqueeze(1) + if indices.ndim != 3: + raise ValueError(f"Expected a 2D or 3D sparse index tensor, got shape {tuple(indices.shape)}") + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + + padding = (-indices.shape[-1]) % block_size + if padding: + indices = torch.cat( + (indices, indices.new_full((*indices.shape[:-1], padding), -1)), + dim=-1, + ) + return indices + + +class NsaTilelangSparseAttBackend(NsaFlashMlaSparseAttBackend): + """Use TileLang for prefill and FlashMLA for decode.""" + + def create_att_prefill_state(self, infer_state): + return NsaTilelangSparsePrefillAttState(backend=self, infer_state=infer_state) + + +@dataclasses.dataclass +class NsaTilelangSparsePrefillAttState(NsaFlashMlaSparsePrefillAttState): + def _nsa_prefill_att( + self, + q: torch.Tensor, + kv: torch.Tensor, + att_control: AttControl, + ) -> torch.Tensor: + from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + + nsa_dict = att_control.nsa_prefill_dict + # GLM packs its persistent MLA KV and FP8 indexer key in one allocation, + # so the MLA slice has a padded token stride. TileLang requires packed + # KV. For an uncached prefill, the freshly projected batch KV is the + # same ragged ordering addressed by topk_indices; materialize that + # compact view instead of copying the whole persistent cache. Prefix + # cache requests retain the FlashMLA path, which supports the padded + # persistent layout and global memory indices. + if self.infer_state.max_cache_len != 0: + return super()._nsa_prefill_att(q=q, kv=kv, att_control=att_control) + + compact_kv = nsa_dict["prefill_cache_kv"].contiguous() + topk_indices = pad_sparse_indices(nsa_dict["topk_indices"]) + if topk_indices.dtype != torch.int32: + topk_indices = topk_indices.to(torch.int32) + + output = tilelang_sparse_fwd( + q=q, + kv=compact_kv, + indices=topk_indices, + sm_scale=nsa_dict["softmax_scale"], + d_v=nsa_dict["kv_lora_rank"], + ) + # TileLang's generated kernel retains the synthetic batch dimension + # inserted by its Python wrapper. LightLLM uses token-major tensors. + if output.ndim == 4: + if output.shape[0] != 1: + raise RuntimeError(f"Unexpected TileLang sparse output shape {tuple(output.shape)}") + output = output.squeeze(0) + return output diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f2b6bae085..f9278362bb 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1,4 +1,5 @@ import os +import math # os.environ["CUDA_LAUNCH_BLOCKING"] = "1" import gc @@ -23,6 +24,9 @@ from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from lightllm.utils.log_utils import init_logger from lightllm.utils.dist_utils import get_dp_world_size from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory @@ -58,6 +62,7 @@ class TpPartBaseModel: is_mtp_draft_model = False + replicated_attention_ep = False # weight class pre_and_post_weight_class = None @@ -94,10 +99,17 @@ def __init__(self, kvargs): if get_env_start_args().enable_decode_microbatch_overlap else self.graph_max_batch_size ) + self.logical_graph_max_batch_size = self.graph_max_batch_size self.mtp_manager = MtpManager.get_instance() - self.graph_max_batch_size = self.graph_max_batch_size * self.mtp_manager.get_decode_batch_multiplier( + self.decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier( + self.is_mtp_draft_model + ) + cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier( self.is_mtp_draft_model ) + self.graph_max_batch_size = ( + self.graph_max_batch_size * cuda_graph_batch_multiplier + ) self.graph_max_len_in_batch = kvargs.get("graph_max_len_in_batch", 8192) self.disable_cudagraph = kvargs.get("disable_cudagraph", False) @@ -107,6 +119,12 @@ def __init__(self, kvargs): self.mem_fraction = kvargs.get("mem_fraction", 0.9) self.tp_world_size_ = get_dp_world_size() self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode + self.use_replicated_attention_ep = ( + self.replicated_attention_ep + and self.enable_tpsp_mix_mode + and self.args.enable_ep_moe + and not self.is_mtp_draft_model + ) self.torch_memory_saver = TorchMemorySaverWrapper(self.args.enable_torch_memory_saver) self.prefill_graph: PrefillCudaGraph = None @@ -279,19 +297,35 @@ def _init_att_backend1(self): return def _init_cudagraph(self): - decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier(self.is_mtp_draft_model) + cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier( + self.is_mtp_draft_model + ) cuda_graph_grow_step_size = self.mtp_manager.get_decode_cuda_graph_grow_step_size(self.is_mtp_draft_model) + extra_batch_sizes = None + if self.mtp_manager.draft_model_needs_logical_batch_graphs(self.is_mtp_draft_model): + # Recurrent EAGLE alternates between a full-width verify-layout + # extend and one-row-per-request recursive draft forwards. Keep + # the ordinary logical schedule in addition to the widened one so + # neither phase falls back to eager execution or excessive padding. + extra_batch_sizes = CudaGraph.gen_cuda_graph_batch_sizes( + batch_step_size_before_split=1, + split_batch_size=self.args.graph_split_batch_size, + batch_step_size_after_split=self.args.graph_grow_step_size, + max_batch_size=self.logical_graph_max_batch_size, + tp_world_size=self.tp_world_size_, + ) self.graph = ( None if self.disable_cudagraph else CudaGraph( batch_step_size_before_split=cuda_graph_grow_step_size, - split_batch_size=self.args.graph_split_batch_size * decode_batch_multiplier, + split_batch_size=self.args.graph_split_batch_size * cuda_graph_batch_multiplier, batch_step_size_after_split=self.args.graph_grow_step_size * cuda_graph_grow_step_size, max_batch_size=self.graph_max_batch_size, max_len_in_batch=self.graph_max_len_in_batch, tp_world_size=self.tp_world_size_, capture_infer_cost=self.args.mtp_dynamic_verify, + extra_batch_sizes=extra_batch_sizes, ) ) if self.graph is not None: @@ -382,13 +416,24 @@ def forward(self, model_input: ModelInput): else: return self._decode(model_input) + def _is_cuda_graph_output_compatible(self, *model_inputs: ModelInput) -> bool: + """Whether inputs match the dense/sparse output captured at startup.""" + + return ( + self.is_mtp_draft_model + or not is_vocab_parallel_greedy_enabled() + or all(model_input.use_vocab_parallel_greedy for model_input in model_inputs) + ) + def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0): infer_state = self.infer_state_class() + infer_state.use_replicated_attention_ep = self.use_replicated_attention_ep infer_state.hidden_collector = self.hidden_collector_prototype.new_instance() infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.is_token_healing = self.is_token_healing infer_state.return_all_prompt_logics = self.return_all_prompt_logics + infer_state.use_vocab_parallel_greedy = self.is_mtp_draft_model or model_input.use_vocab_parallel_greedy infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len @@ -539,6 +584,9 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba return model_output new_model_output = copy.copy(model_output) new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -551,6 +599,9 @@ def _create_unpad_prefill_model_output( new_model_output = copy.copy(padded_model_output) # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。 new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -584,7 +635,13 @@ def _prefill( if self.args.enable_tpsp_mix_mode: infer_handle_token_num = triton.cdiv(infer_handle_token_num, self.tp_world_size_) * self.tp_world_size_ - if self.prefill_graph is not None and self.prefill_graph.can_run(handle_token_num=infer_handle_token_num): + if self.prefill_graph is not None and self.prefill_graph.can_run( + handle_token_num=infer_handle_token_num, + batch_size=model_input.batch_size, + max_q_seq_len=model_input.max_q_seq_len, + max_kv_seq_len=model_input.max_kv_seq_len, + max_cache_len=model_input.max_cache_len, + ): infer_handle_token_num = self.prefill_graph.find_closest_graph_handle_token_num( handle_token_num=infer_handle_token_num ) @@ -643,14 +700,24 @@ def _decode( # 向上对齐到 TP world size 的整数倍，保证后续切分得到合法 shape。 infer_batch_size = max(1, origin_batch_size) if self.args.enable_tpsp_mix_mode: - infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_ + decode_alignment = math.lcm( + self.tp_world_size_, self.decode_batch_multiplier + ) + infer_batch_size = ( + triton.cdiv(infer_batch_size, decode_alignment) + * decode_alignment + ) # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = self.graph is not None and self.graph.can_run( - batch_size=infer_batch_size, - max_len_in_batch=infer_max_kv_seq_len, + use_cuda_graph = ( + self._is_cuda_graph_output_compatible(model_input) + and self.graph is not None + and self.graph.can_run( + batch_size=infer_batch_size, + max_len_in_batch=infer_max_kv_seq_len, + ) ) need_capture = False if use_cuda_graph: @@ -690,7 +757,8 @@ def _context_forward(self, infer_state: InferStateInfo): infer_state.prepare_prefill_dp_balance() input_embs = infer_state._all_to_all_balance_get(data=input_embs) - input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) + if not self.use_replicated_attention_ep: + input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) input_tensors = [input_embs] if Autotuner.is_autotune_warmup(): infer_state.hidden_collector = NoopHiddenCollector() @@ -710,7 +778,13 @@ def prefill_func(input_tensors, _infer_state): handle_token_num = infer_state.input_ids.shape[0] - if self.prefill_graph is not None and self.prefill_graph.can_run(handle_token_num=handle_token_num): + if self.prefill_graph is not None and self.prefill_graph.can_run( + handle_token_num=handle_token_num, + batch_size=infer_state.batch_size, + max_q_seq_len=infer_state.max_q_seq_len, + max_kv_seq_len=infer_state.max_kv_seq_len, + max_cache_len=infer_state.max_cache_len, + ): finded_handle_token_num = self.prefill_graph.find_closest_graph_handle_token_num( handle_token_num=handle_token_num ) @@ -733,7 +807,9 @@ def prefill_func(input_tensors, _infer_state): input_embs = output_tensors[0] - last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) + last_input_embs = input_embs + if not self.use_replicated_attention_ep: + last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) if infer_state.need_dp_prefill_balance: last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs) @@ -742,6 +818,8 @@ def prefill_func(input_tensors, _infer_state): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -756,14 +834,17 @@ def _token_forward(self, infer_state: InferStateInfo): input_ids = infer_state.input_ids cuda_input_ids = input_ids input_embs = self.pre_infer.token_forward(cuda_input_ids, infer_state, self.pre_post_weight) - input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) + if not self.use_replicated_attention_ep: + input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) for i in range(self.layers_num): layer = self.layers_infer[i] input_embs: torch.Tensor = layer.token_forward(input_embs, infer_state, self.trans_layers_weight[i]) hidden_collector.add(layer_index=i, hidden=input_embs) - last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) + last_input_embs = input_embs + if not self.use_replicated_attention_ep: + last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) predict_logits: torch.Tensor = self.post_infer.token_forward( last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight ) @@ -771,6 +852,8 @@ def _token_forward(self, infer_state: InferStateInfo): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -900,9 +983,18 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 origin_batch_size1 = model_input1.batch_size max_len_in_batch = max(2, model_input0.max_kv_seq_len, model_input1.max_kv_seq_len) infer_batch_size = max(1, origin_batch_size0, origin_batch_size1) - infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_ + decode_alignment = math.lcm( + self.tp_world_size_, self.decode_batch_multiplier + ) + infer_batch_size = ( + triton.cdiv(infer_batch_size, decode_alignment) * decode_alignment + ) - if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): + if ( + self._is_cuda_graph_output_compatible(model_input0, model_input1) + and self.graph is not None + and self.graph.can_run(infer_batch_size, max_len_in_batch) + ): infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) need_capture = self.graph.need_capture(infer_batch_size) padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) @@ -1025,11 +1117,15 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1074,10 +1170,14 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), ) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index ae645d4b7b..f12b407f24 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -55,6 +55,10 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None + # The router enables sparse vocabulary output only when target sampling is + # exact, unmodified greedy. Draft models always consume greedy proposals. + use_vocab_parallel_greedy: bool = False + def to_cuda(self): self.check_input() @@ -140,13 +144,15 @@ class ModelMtpOutputCollector: # - 未启用 MTP 时不收集投机特征，该字段同样为 None。 spec_hidden: Optional[torch.Tensor] = None - # DSpark block draft 模型直接生成的 token id，形状通常为 - # [request_count * block_size]。 - # - 仅 DSpark 启用 Markov head（markov_rank > 0）时由 head 直接生成并返回。 - # - DSpark 未启用 Markov head 时为 None，调用方从普通 logits 执行 argmax。 - # - Vanilla MTP、EAGLE、EAGLE3、DFlash 以及未启用 MTP 的模型均不使用该字段。 + # Draft head 直接生成的 token id。Block drafter 通常返回 + # [request_count * block_size]，autoregressive drafter 通常返回 + # [request_count]。未提供时，调用方从普通 logits 执行 argmax。 draft_token_ids: Optional[torch.Tensor] = None + # 与 draft_token_ids 一一对应的精确最大 token 概率。Vocab-parallel + # draft head 可直接返回该值，避免为了动态 verify 聚合完整词表。 + draft_token_probs: Optional[torch.Tensor] = None + # DSpark confidence head 输出的原始置信度 logits，形状通常为 # [request_count, block_size]，供动态 MTP verify 计算各 draft 位置的调度分数。 # - 仅 DSpark checkpoint 启用 confidence head 时返回；动态 verify 模式要求该字段存在。 @@ -159,6 +165,8 @@ def to_no_ref_tensor(self) -> None: self.spec_hidden = tensor_to_no_ref_tensor(self.spec_hidden) if self.draft_token_ids is not None: self.draft_token_ids = tensor_to_no_ref_tensor(self.draft_token_ids) + if self.draft_token_probs is not None: + self.draft_token_probs = tensor_to_no_ref_tensor(self.draft_token_probs) if self.confidence_logits is not None: self.confidence_logits = tensor_to_no_ref_tensor(self.confidence_logits) @@ -168,6 +176,8 @@ def unpad_decode(self, padded_batch_size: int, origin_batch_size: int) -> "Model collector.spec_hidden = collector.spec_hidden[:origin_batch_size] if collector.draft_token_ids is not None: collector.draft_token_ids = collector.draft_token_ids[:origin_batch_size] + if collector.draft_token_probs is not None: + collector.draft_token_probs = collector.draft_token_probs[:origin_batch_size] if collector.confidence_logits is not None: confidence_row_count = collector.confidence_logits.shape[0] assert confidence_row_count > 0 and padded_batch_size % confidence_row_count == 0 @@ -200,10 +210,57 @@ class ModelOutput: # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None + # Sparse vocabulary output. Each logit column maps to the corresponding + # global token id; logsumexp still covers the complete vocabulary. + logits_token_ids: Optional[torch.Tensor] = None + logits_logsumexp: Optional[torch.Tensor] = None + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() + assert (self.logits_token_ids is None) == (self.logits_logsumexp is None) + if self.logits_token_ids is not None: + assert self.logits.ndim == 2 + assert self.logits_token_ids.shape == self.logits.shape + assert self.logits_token_ids.dtype in (torch.int32, torch.int64) + assert self.logits_token_ids.device == self.logits.device + assert self.logits_logsumexp.shape == (self.logits.shape[0],) + assert self.logits_logsumexp.dtype == torch.float32 + assert self.logits_logsumexp.device == self.logits.device def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) + if self.logits_token_ids is not None: + self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids) + self.logits_logsumexp = tensor_to_no_ref_tensor(self.logits_logsumexp) self.mtp_collector.to_no_ref_tensor() + + @property + def has_vocab_parallel_logits(self) -> bool: + return self.logits_token_ids is not None + + def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": + """Select logit rows without dropping sparse-vocabulary metadata.""" + + return ModelOutput( + logits=self.logits.index_select(0, index), + logits_token_ids=( + self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None + ), + logits_logsumexp=( + self.logits_logsumexp.index_select(0, index) if self.logits_logsumexp is not None else None + ), + ) + + @classmethod + def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput": + """Concatenate compatible dense or sparse-vocabulary logit rows.""" + + assert outputs + sparse = outputs[0].has_vocab_parallel_logits + assert all(output.has_vocab_parallel_logits == sparse for output in outputs) + return cls( + logits=torch.cat([output.logits for output in outputs], dim=0), + logits_token_ids=(torch.cat([output.logits_token_ids for output in outputs], dim=0) if sparse else None), + logits_logsumexp=(torch.cat([output.logits_logsumexp for output in outputs], dim=0) if sparse else None), + ) diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 5849cccf54..acaab19f1a 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,14 +1,18 @@ import os +import math import torch import torch.distributed as dist import copy import bisect import triton -from typing import Optional +from typing import Iterable, Optional from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from lightllm.utils.torch_memory_saver_utils import ( TorchMemorySaverWrapper, MemoryTag, @@ -46,7 +50,13 @@ def gen_cuda_graph_batch_sizes( batch_sizes = sorted({size for size in batch_sizes if size < max_batch_size} | {max_batch_size}) if args.enable_tpsp_mix_mode: - batch_sizes = sorted({triton.cdiv(size, tp_world_size) * tp_world_size for size in batch_sizes}) + # Static speculative verification stores one request in a fixed + # block of ``batch_step_size_before_split`` rows. TP/SP padding + # must preserve both that block and an even split across TP ranks. + alignment = math.lcm(tp_world_size, batch_step_size_before_split) + batch_sizes = sorted( + {triton.cdiv(size, alignment) * alignment for size in batch_sizes} + ) assert batch_sizes[-1] == max_batch_size return batch_sizes @@ -59,6 +69,7 @@ def __init__( max_len_in_batch=8192, tp_world_size: int = 1, capture_infer_cost: bool = False, + extra_batch_sizes: Optional[Iterable[int]] = None, ): self.graph = {} self.tp_world_size = tp_world_size @@ -78,6 +89,15 @@ def __init__( max_batch_size=self.max_batch_size, tp_world_size=self.tp_world_size, ) + if extra_batch_sizes is not None: + self.cuda_graph_batch_sizes = sorted( + set(self.cuda_graph_batch_sizes) + | { + int(batch_size) + for batch_size in extra_batch_sizes + if 0 < int(batch_size) <= self.max_batch_size + } + ) logger.info(f"cuda graph batch_sizes: {self.cuda_graph_batch_sizes}") def can_run(self, batch_size, max_len_in_batch): @@ -279,6 +299,7 @@ def warmup(self, model): b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), is_prefill=False, multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(batch_size), ) model_output: ModelOutput = model.forward(model_input) @@ -340,6 +361,7 @@ def warmup_overlap(self, model): b_shared_radix_node_id=b_shared_radix_node_id, b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(batch_size), ) decode_batches.append(micro_batch) diff --git a/lightllm/common/basemodel/hidden_collector.py b/lightllm/common/basemodel/hidden_collector.py index 3eb946fe82..605ed68b59 100644 --- a/lightllm/common/basemodel/hidden_collector.py +++ b/lightllm/common/basemodel/hidden_collector.py @@ -91,6 +91,7 @@ def add_mtp_outputs( self, draft_token_ids: Optional[torch.Tensor], confidence_logits: Optional[torch.Tensor], + draft_token_probs: Optional[torch.Tensor] = None, ) -> None: """Collect optional token/confidence outputs produced by an MTP head. @@ -133,6 +134,7 @@ class MtpHeadOutputCollector(NoopHiddenCollector): def __init__(self) -> None: self.draft_token_ids: Optional[torch.Tensor] = None + self.draft_token_probs: Optional[torch.Tensor] = None self.confidence_logits: Optional[torch.Tensor] = None def new_instance(self) -> HiddenCollector: @@ -142,16 +144,20 @@ def add_mtp_outputs( self, draft_token_ids: Optional[torch.Tensor], confidence_logits: Optional[torch.Tensor], + draft_token_probs: Optional[torch.Tensor] = None, ) -> None: self.draft_token_ids = draft_token_ids + self.draft_token_probs = draft_token_probs self.confidence_logits = confidence_logits def finish_output(self, infer_state) -> ModelMtpOutputCollector: output = ModelMtpOutputCollector( draft_token_ids=self.draft_token_ids, + draft_token_probs=self.draft_token_probs, confidence_logits=self.confidence_logits, ) self.draft_token_ids = None + self.draft_token_probs = None self.confidence_logits = None return output diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 91c6e99699..bd1809ad5e 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -53,6 +53,9 @@ def __init__(self): self.is_token_healing: bool = False self.return_all_prompt_logics: bool = False + self.use_vocab_parallel_greedy: bool = False + self.logits_token_ids: Optional[torch.Tensor] = None + self.logits_logsumexp: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。 @@ -395,4 +398,8 @@ def copy_for_prefill_cuda_graph(self, new_infer_state: "InferStateInfo"): attr_ = getattr(self, attr_name, None) if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr() and attr_.shape == attr_value.shape: attr_.copy_(attr_value, non_blocking=True) + + self.prefill_att_state.copy_for_prefill_cuda_graph(new_infer_state.prefill_att_state) + if self.prefill_att_state1 is not None: + self.prefill_att_state1.copy_for_prefill_cuda_graph(new_infer_state.prefill_att_state1) return diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index f0cc129c09..5bf32bca2e 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -6,6 +6,7 @@ from lightllm.distributed import all_reduce from typing import Tuple from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor +from lightllm.utils.envs_utils import get_env_start_args class TransformerLayerInferTpl(TransformerLayerInfer): @@ -102,6 +103,13 @@ def _context_attention_wrapper_run( self, q: torch.Tensor, cache_kv: torch.Tensor, infer_state: InferStateInfo, layer_weight ) -> torch.Tensor: if torch.cuda.is_current_stream_capturing(): + # Exact-layout graphs can opt into capturing attention as part of + # the main graph. This avoids one CPU callback and one graph split + # per attention layer. The fixed layout is validated by + # PrefillCudaGraph before capture starts. + if bool(getattr(get_env_start_args(), "prefill_cudagraph_capture_attention", False)): + return self._context_attention_kernel(q, cache_kv, infer_state, layer_weight) + q = q.contiguous() # cache_kv is None for layers that own no K/V slot (e.g. gemma4 # KV-shared layers, which read K/V from a prior layer's cache and @@ -110,6 +118,25 @@ def _context_attention_wrapper_run( cache_kv = cache_kv.contiguous() if cache_kv is not None else None _q = tensor_to_no_ref_tensor(q) _cache_kv = tensor_to_no_ref_tensor(cache_kv) if cache_kv is not None else None + + # Some attention implementations stash graph-produced tensors in + # a short-lived infer-state dict between QKV projection and the + # CPU attention callback (for example NSA indexer inputs). Python + # attribute assignment is not replayed by CUDA Graph, and the + # shape-probing call below may consume/delete the dict. Preserve + # fixed-address, non-owning views and restore them for every CPU + # callback invocation. + callback_tensor_dicts = {} + for attr_name, attr_value in vars(infer_state).items(): + if ( + isinstance(attr_value, dict) + and attr_value + and all(isinstance(value, torch.Tensor) for value in attr_value.values()) + ): + callback_tensor_dicts[attr_name] = { + key: tensor_to_no_ref_tensor(value.contiguous()) + for key, value in attr_value.items() + } pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() pre_capture_graph.__exit__(None, None, None) @@ -135,6 +162,8 @@ def get_o_shape_dtype_device(): _o = tensor_to_no_ref_tensor(o) def att_func(new_infer_state: InferStateInfo): + for attr_name, attr_value in callback_tensor_dicts.items(): + setattr(new_infer_state, attr_name, attr_value) tmp_o = self._context_attention_kernel(_q, _cache_kv, new_infer_state, layer_weight) assert tmp_o.shape == _o.shape _o.copy_(tmp_o) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 7f369c4fd8..5a7a9d53d8 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -68,6 +68,9 @@ def __init__( routed_expert_counter_tensor=self.routed_expert_counter_tensor, auto_update_redundancy_expert=self.auto_update_redundancy_expert, ) + self.fuse_moe_impl.swiglu_limit = self.swiglu_limit + self.fuse_moe_impl.swiglu_alpha = self.swiglu_alpha + self.fuse_moe_impl.swiglu_clamp_up_add_one = self.swiglu_clamp_up_add_one self.lock = threading.Lock() self._create_weight() @@ -79,6 +82,11 @@ def _init_config(self, network_config: Dict[str, Any]): self.num_experts_per_tok = network_config["num_experts_per_tok"] self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0) self.scoring_func = network_config.get("scoring_func", "softmax") + self.swiglu_limit = network_config.get("swiglu_limit") + self.swiglu_alpha = network_config.get("swiglu_alpha", 1.0) + self.swiglu_clamp_up_add_one = network_config.get( + "swiglu_clamp_up_add_one", True + ) def _init_redundancy_expert_params(self): self.redundancy_expert_num = get_redundancy_expert_num() @@ -285,14 +293,41 @@ def load_hf_weights(self, weights): self._load_weight(self.redundancy_expert_idx_to_local_idx, weights) def verify_load(self): - weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list) + if getattr(self, "_sm90_mega_moe_weights_prepared", False): + weight_load_ok = all( + all(_weight_pack.load_ok) for _weight_pack in self.w2_list + ) + else: + weight_load_ok = all( + all(_weight_pack.load_ok) + for _weight_pack in self.w1_list + self.w2_list + self.w3_list + ) per_expert_scale_load_ok = ( True if self.per_expert_scale is None else getattr(self.per_expert_scale, "load_ok", False) ) e_score_correction_bias_load_ok = ( True if self.e_score_correction_bias is None else getattr(self.e_score_correction_bias, "load_ok", False) ) - return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok + load_ok = weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok + if load_ok and self.enable_ep_moe and not getattr( + self, "_sm90_mega_moe_weights_prepared", False + ): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( + prepare_sm90_mega_moe_weights, + use_sm90_mega_moe, + ) + + if use_sm90_mega_moe(self.quant_method): + prepare_sm90_mega_moe_weights(self.w13) + # The loader-only gate/up views retain the original storage. + # Drop them after replacement so all layers do not keep a + # second full copy of their routed-expert L1 weights. + self.w1 = None + self.w3 = None + self.w1_list = [] + self.w3_list = [] + self._sm90_mega_moe_weights_prepared = True + return load_ok def _create_weight(self): intermediate_size = self.split_inter_size diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index a5ba656c9c..955067942f 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -12,7 +12,9 @@ get_ep_num_sms, masked_group_gemm, chunked_expanded_moe_forward, + legacy_normal_moe_forward, quantize_fused_experts_input, + use_sm90_mega_moe, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.triton_utils.autotuner import Autotuner @@ -77,16 +79,24 @@ def _fused_experts( router_logits: Optional[torch.Tensor] = None, is_prefill: Optional[bool] = None, ): + fused_topk_ids = ( + topk_ids + if use_sm90_mega_moe(self.quant_method) + else topk_ids.to(torch.long) + ) output = fused_experts( hidden_states=input_tensor, w13=w13, w2=w2, topk_weights=topk_weights, - topk_idx=topk_ids.to(torch.long), + topk_idx=fused_topk_ids, num_experts=self.total_expert_num_contain_redundancy, # number of all experts contain redundancy quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_clamp_up_add_one=self.swiglu_clamp_up_add_one, ) return output @@ -163,6 +173,49 @@ def dispatch( overlap_event: Optional[Any] = None, ): buffer = dist_group_manager.ep_buffer + if dist_group_manager.ep_prefill_uses_legacy_buffer: + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + layout_event, + ) = buffer.get_dispatch_layout( + topk_idx, + self.total_expert_num_contain_redundancy, + previous_event=overlap_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + ( + recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + handle, + _, + ) = buffer.dispatch( + qinput_tensor, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=layout_event, + async_finish=False, + allocate_on_comm_stream=False, + expert_alignment=128, + ) + return ( + recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + handle, + lambda: None, + ) + num_max_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_prefill() recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch( qinput_tensor, @@ -206,6 +259,9 @@ def masked_group_gemm( w2_weight, w2_scale, expected_m=expected_m, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_clamp_up_add_one=self.swiglu_clamp_up_add_one, ) def prefilled_group_gemm( @@ -223,6 +279,22 @@ def prefilled_group_gemm( ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale + if dist_group_manager.ep_prefill_uses_legacy_buffer: + return legacy_normal_moe_forward( + num_recv_tokens_per_expert_list, + recv_x, + recv_topk_idx, + recv_topk_weights, + w13_weight, + w13_scale, + w2_weight, + w2_scale, + hidden_dtype, + self.swiglu_limit, + self.swiglu_alpha, + self.swiglu_clamp_up_add_one, + ) + assert recv_topk_idx is None all_tokens = sum(num_recv_tokens_per_expert_list) if all_tokens > 0: @@ -239,6 +311,9 @@ def prefilled_group_gemm( block_size_k=self.quant_method.block_size, workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(microbatch_index), hidden_dtype=hidden_dtype, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_clamp_up_add_one=self.swiglu_clamp_up_add_one, ) else: gather_out = torch.empty( @@ -254,7 +329,13 @@ def prefilled_group_gemm( N = w13_weight.shape[1] _gemm_out_a = torch.zeros((1, N), device=recv_x[0].device, dtype=hidden_dtype) _silu_out = torch.zeros((1, N // 2), device=recv_x[0].device, dtype=hidden_dtype) - silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out) + silu_and_mul_fwd( + _gemm_out_a.view(-1, N), + _silu_out, + limit=self.swiglu_limit, + alpha=self.swiglu_alpha, + clamp_up_add_one=self.swiglu_clamp_up_add_one, + ) _gemm_out_a, _silu_out = None, None del recv_x return gather_out @@ -278,6 +359,17 @@ def combine( overlap_event: Optional[Any] = None, ): # normal combine + if dist_group_manager.ep_prefill_uses_legacy_buffer: + combined_x, _, _ = dist_group_manager.ep_buffer.combine( + gemm_out_b, + handle, + topk_weights=None, + previous_event=overlap_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + return combined_x, lambda: None + combined_x, _, event = dist_group_manager.ep_buffer.combine( gemm_out_b, handle, diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index 1d6a38c069..d13f633fb2 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -1,3 +1,8 @@ +import os +from contextlib import contextmanager, nullcontext +from functools import lru_cache +from types import SimpleNamespace + import torch from typing import Callable, Optional from lightllm.common.quantization.no_quant import WeightPack @@ -5,6 +10,161 @@ from .base_impl import FuseMoeBaseImpl +def _use_sglang_triton_moe() -> bool: + return os.getenv("LIGHTLLM_USE_SGLANG_TRITON_MOE", "0").upper() in { + "ON", + "TRUE", + "1", + } + + +@lru_cache(maxsize=1) +def _get_sglang_fused_experts_impl(): + """Load SGLang's tuned Triton MoE without its server RuntimeContext. + + LightLLM and SGLang share the same standard block-FP8 expert layout, but + SGLang's standalone kernel helpers consult two unrelated process-global + server flags. Supply their disabled defaults locally so the kernel can be + used as an optional backend without initializing an SGLang server. + """ + + from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe as sglang_fused_moe, + ) + from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe_triton_config as sglang_fused_moe_config, + ) + + # Despite the upstream flag name, for this in-place standalone invocation + # the optimization only folds the top-k expert sum into the down GEMM. It + # does not perform the TP all-reduce, which remains owned by LightLLM. + enable_fused_sum = os.getenv("LIGHTLLM_SGLANG_FUSED_MOE_SUM", "1").upper() in { + "ON", + "TRUE", + "1", + } + standalone_exec = SimpleNamespace( + moe=SimpleNamespace(enable_fused_moe_sum_all_reduce=enable_fused_sum), + deterministic=SimpleNamespace(enable_deterministic_inference=False), + ) + sglang_fused_moe.get_exec = lambda: standalone_exec + sglang_fused_moe_config.get_exec = lambda: standalone_exec + from sglang.srt.layers.moe.moe_runner.triton_utils import override_config + + return sglang_fused_moe, override_config + + +@contextmanager +def _override_sglang_moe_configs( + sglang_fused_moe, + override_config, + up_config: dict, + down_config: Optional[dict], +): + """Temporarily select independently measured up/down MoE configs.""" + + if down_config is None: + with override_config(up_config): + yield + return + + original = sglang_fused_moe.try_get_optimal_moe_config + + def resolve_config(*args, return_down_config=False, **kwargs): + if return_down_config: + return up_config, (down_config, None) + return up_config + + sglang_fused_moe.try_get_optimal_moe_config = resolve_config + try: + yield + finally: + sglang_fused_moe.try_get_optimal_moe_config = original + + +@lru_cache(maxsize=None) +def _get_sglang_triton_moe_configs( + w13_shape: tuple[int, ...], + w2_shape: tuple[int, ...], + topk: int, + is_prefill: bool, + token_count: int, +): + """Return measured H100 (up, down) configs for GLM-5's TP8 shape.""" + + if "H100" not in torch.cuda.get_device_name(torch.cuda.current_device()): + return None + if is_prefill: + if token_count < 8192: + up_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + return up_config, None + + # The two GLM expert projections have different output widths and + # diverge at large prefill batches. Measurements on H100 TP8 show + # that N=128 remains best for the 4096->512 up projection, while the + # 256->4096 down projection crosses over to N=64 at roughly 24K + # tokens. Selecting them independently avoids making the faster up + # projection pay for the down projection's narrower tile. + up_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3, + } + down_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64 if token_count >= 24576 else 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32 if token_count >= 24576 else 8, + "num_warps": 4, + "num_stages": 3, + } + return up_config, down_config + if w13_shape != (289, 512, 4096): + return None + if w2_shape != (289, 4096, 256) or topk != 9: + return None + up_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + } + # The draft graph has eight physical rows; the MTP2 verify graph has 24. + # Separate down-projection sweeps found different winners for those two + # hot shapes while retaining BLOCK_SIZE_M=16 for shared route alignment. + if token_count <= 8: + down_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3, + } + else: + down_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + } + return up_config, down_config + + class FuseMoeTriton(FuseMoeBaseImpl): def __init__( self, @@ -92,6 +252,57 @@ def _fused_experts( w2_weight, w2_scale = w2.weight, w2.weight_scale use_fp8_w8a8 = w13_weight.dtype == torch.float8_e4m3fn + if _use_sglang_triton_moe(): + block_size = getattr(self.quant_method, "block_size", None) + if not use_fp8_w8a8 or block_size != 128: + raise RuntimeError( + "LIGHTLLM_USE_SGLANG_TRITON_MOE currently requires " + "block-wise FP8 expert weights with block size 128" + ) + if ( + getattr(self, "swiglu_limit", None) is not None + and getattr(self, "swiglu_clamp_up_add_one", True) + ): + raise RuntimeError( + "SGLang Triton MoE does not support clamp_up_add_one=True" + ) + + sglang_fused_moe, override_config = _get_sglang_fused_experts_impl() + tuned_configs = _get_sglang_triton_moe_configs( + tuple(w13_weight.shape), + tuple(w2_weight.shape), + topk_ids.shape[1], + bool(is_prefill), + input_tensor.shape[0], + ) + config_context = ( + _override_sglang_moe_configs( + sglang_fused_moe, + override_config, + *tuned_configs, + ) + if tuned_configs is not None + else nullcontext() + ) + with config_context: + sglang_fused_moe.fused_experts_impl( + hidden_states=input_tensor, + w1=w13_weight, + w2=w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w13_scale, + w2_scale=w2_scale, + block_shape=[block_size, block_size], + routed_scaling_factor=1.0, + filter_expert=False, + swiglu_limit=getattr(self, "swiglu_limit", None), + gate_up_interleaved=False, + ) + return input_tensor + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import fused_experts fused_experts( @@ -104,6 +315,13 @@ def _fused_experts( use_fp8_w8a8=use_fp8_w8a8, w1_scale=w13_scale, w2_scale=w2_scale, + limit=getattr(self, "swiglu_limit", None), + alpha=( + getattr(self, "swiglu_alpha", 1.0) + if getattr(self, "swiglu_limit", None) is not None + else None + ), + clamp_up_add_one=getattr(self, "swiglu_clamp_up_add_one", True), ) return input_tensor diff --git a/lightllm/common/basemodel/mtp_manager.py b/lightllm/common/basemodel/mtp_manager.py index be6c477b99..16bcdbebb5 100644 --- a/lightllm/common/basemodel/mtp_manager.py +++ b/lightllm/common/basemodel/mtp_manager.py @@ -16,6 +16,7 @@ class MtpManager: _instance: ClassVar[Optional["MtpManager"]] = None _CHAINED_DRAFT_MODES = ("vanilla_with_att", "vanilla_no_att") _RECURRENT_DRAFT_MODES = ("eagle_with_att", "eagle_no_att", "eagle3") + _RECURRENT_ATTN_DRAFT_MODES = ("eagle_with_att", "eagle3") _BLOCK_DRAFT_MODES = ("dspark", "dflash") @classmethod @@ -41,7 +42,10 @@ def get_decode_batch_multiplier(self, is_draft_model: bool) -> int: if not is_draft_model: return verify_width - # Chained MTP runs every draft module over the expanded verify layout. + # Chained MTP runs every draft module over the expanded verify layout, + # but each physical row is still an independent one-token draft decode + # from the attention backend's point of view. CUDA Graph sizing must + # account for the wider physical batch separately (see below). if spec_mode in self._CHAINED_DRAFT_MODES: return 1 @@ -55,12 +59,39 @@ def get_decode_batch_multiplier(self, is_draft_model: bool) -> int: return 1 + def get_decode_cuda_graph_batch_multiplier(self, is_draft_model: bool) -> int: + """Return physical decode rows per logical request for graph sizing. + + Chained draft models keep normal one-token attention semantics, while + their proposer forwards the complete target verification layout through + every draft depth. Consequently a logical batch of ``N`` requests has + ``N * (mtp_step + 1)`` physical rows and needs graphs captured at that + width even though :meth:`get_decode_batch_multiplier` returns one. + """ + + # Attention-backed recurrent EAGLE normally decodes one row per + # request, but its first forward after every target verification runs + # over the complete expanded verify layout to commit draft KV. That + # extend forward must be graph-covered as well. + if is_draft_model and self.args.mtp_mode in ( + *self._CHAINED_DRAFT_MODES, + *self._RECURRENT_ATTN_DRAFT_MODES, + ): + return self.args.mtp_step + 1 + return self.get_decode_batch_multiplier(is_draft_model=is_draft_model) + + def draft_model_needs_logical_batch_graphs(self, is_draft_model: bool) -> bool: + """Whether a widened draft graph also needs one-row-per-request sizes.""" + + return is_draft_model and self.args.mtp_mode in self._RECURRENT_ATTN_DRAFT_MODES + def get_decode_cuda_graph_grow_step_size(self, is_draft_model: bool) -> int: """Return the batch-size stride used to capture decode CUDA Graphs.""" - # Draft model CUDA Graphs follow the drafter's physical decode layout. + # Draft model CUDA Graphs follow the drafter's physical forward layout, + # which is wider than its attention semantics for chained MTP. if is_draft_model: - return self.get_decode_batch_multiplier(is_draft_model=True) + return self.get_decode_cuda_graph_batch_multiplier(is_draft_model=True) # Main model CUDA Graphs use unit growth for dynamically compacted verify rows. else: if self.args.mtp_dynamic_verify: diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index bf6039a48f..aad9a97802 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -10,6 +10,9 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from .infer_struct import InferStateInfo from .cuda_graph import CudaGraph @@ -33,16 +36,69 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): if self.args.batch_max_tokens is not None: self.max_handle_token_num = min(self.max_handle_token_num, self.args.batch_max_tokens) - graph_handle_token_nums = ( - list(range(4, 33, 4)) - + list(range(48, 257, 16)) - + list(range(288, 513, 32)) - + list(range(576, 1024 + 1, 64)) - + list(range(1280, 4096 + 1, 256)) - + list(range(4608, self.max_handle_token_num + 1, 512)) - ) - graph_handle_token_nums = [e for e in graph_handle_token_nums if e <= self.max_handle_token_num] - graph_handle_token_nums.append(self.max_handle_token_num) + configured_token_nums = getattr(self.args, "prefill_cudagraph_token_nums", None) + configured_batch_sizes = getattr(self.args, "prefill_cudagraph_batch_sizes", None) + self.use_exact_token_nums = configured_token_nums is not None + self.capture_attention = bool(getattr(self.args, "prefill_cudagraph_capture_attention", False)) + if self.capture_attention and not self.use_exact_token_nums: + raise ValueError( + "--prefill_cudagraph_capture_attention requires " + "--prefill_cudagraph_token_nums and --prefill_cudagraph_batch_sizes" + ) + self.exact_batch_size_by_token_num = {} + if self.use_exact_token_nums: + if configured_batch_sizes is None: + raise ValueError( + "--prefill_cudagraph_batch_sizes is required with " + "--prefill_cudagraph_token_nums" + ) + if len(configured_token_nums) != len(configured_batch_sizes): + raise ValueError( + "--prefill_cudagraph_token_nums and --prefill_cudagraph_batch_sizes " + "must contain the same number of entries" + ) + + graph_handle_token_nums = [] + for raw_token_num, raw_batch_size in zip(configured_token_nums, configured_batch_sizes): + token_num = int(raw_token_num) + batch_size = int(raw_batch_size) + if not 0 < token_num <= self.max_handle_token_num: + continue + if batch_size <= 0: + raise ValueError("--prefill_cudagraph_batch_sizes entries must be positive") + if token_num % batch_size != 0: + raise ValueError( + "Each --prefill_cudagraph_token_nums entry must be divisible by its paired " + "--prefill_cudagraph_batch_sizes entry" + ) + old_batch_size = self.exact_batch_size_by_token_num.get(token_num) + if old_batch_size is not None and old_batch_size != batch_size: + raise ValueError( + f"Conflicting prefill CUDA Graph batch sizes for token count {token_num}: " + f"{old_batch_size} and {batch_size}" + ) + self.exact_batch_size_by_token_num[token_num] = batch_size + graph_handle_token_nums.append(token_num) + if not graph_handle_token_nums: + raise ValueError( + "--prefill_cudagraph_token_nums must contain a positive value no larger than " + "--prefill_cudagraph_max_handle_token and --batch_max_tokens" + ) + else: + if configured_batch_sizes is not None: + raise ValueError( + "--prefill_cudagraph_batch_sizes requires --prefill_cudagraph_token_nums" + ) + graph_handle_token_nums = ( + list(range(4, 33, 4)) + + list(range(48, 257, 16)) + + list(range(288, 513, 32)) + + list(range(576, 1024 + 1, 64)) + + list(range(1280, 4096 + 1, 256)) + + list(range(4608, self.max_handle_token_num + 1, 512)) + ) + graph_handle_token_nums = [e for e in graph_handle_token_nums if e <= self.max_handle_token_num] + graph_handle_token_nums.append(self.max_handle_token_num) graph_handle_token_nums = list(set[int](graph_handle_token_nums)) graph_handle_token_nums.sort() @@ -55,8 +111,30 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): self.graph_handle_token_nums = graph_handle_token_nums logger.info(f"prefill cuda graph graph_handle_token_nums: {self.graph_handle_token_nums}") + if self.exact_batch_size_by_token_num: + logger.info( + "prefill cuda graph exact layouts (token_num -> batch_size): " + f"{self.exact_batch_size_by_token_num}" + ) - def can_run(self, handle_token_num: int): + def can_run( + self, + handle_token_num: int, + batch_size: Optional[int] = None, + max_q_seq_len: Optional[int] = None, + max_kv_seq_len: Optional[int] = None, + max_cache_len: Optional[int] = None, + ): + if self.use_exact_token_nums: + configured_batch_size = self.exact_batch_size_by_token_num.get(handle_token_num) + if configured_batch_size is None or batch_size != configured_batch_size: + return False + uniform_seq_len = handle_token_num // configured_batch_size + return ( + max_q_seq_len == uniform_seq_len + and max_kv_seq_len == uniform_seq_len + and max_cache_len == 0 + ) return handle_token_num <= self.max_handle_token_num def need_capture(self, handle_token_num: int): @@ -191,23 +269,35 @@ def warmup(self, model): # prefill cuda graph init for handle_token_num in self.graph_handle_token_nums[::-1]: - logger.info(f"Capture prefill cudagraph, handle_token_num: {handle_token_num}") + batch_size = self.exact_batch_size_by_token_num.get(handle_token_num, 1) + if handle_token_num % batch_size != 0: + raise ValueError( + f"Prefill CUDA Graph token count {handle_token_num} is not divisible by batch size {batch_size}" + ) + seq_len = handle_token_num // batch_size + logger.info( + "Capture prefill cudagraph, " + f"handle_token_num: {handle_token_num}, batch_size: {batch_size}, seq_len: {seq_len}" + ) total_token_num = handle_token_num input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int64, device="cuda") mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda() - b_req_idx = torch.tensor([model.req_manager.HOLD_REQUEST_ID], dtype=torch.int32, device="cuda") - b_seq_len = torch.empty(1, dtype=torch.int32, device="cuda") - b_seq_len.fill_(total_token_num) - b_mtp_index = torch.zeros(1, dtype=torch.int32, device="cuda") - b_is_decode_req = torch.zeros(1, dtype=torch.bool, device="cuda") - b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda") - b_prefill_start_loc = torch.zeros(1, dtype=torch.int32, device="cuda") + b_req_idx = torch.full( + (batch_size,), model.req_manager.HOLD_REQUEST_ID, dtype=torch.int32, device="cuda" + ) + b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + b_is_decode_req = torch.zeros(batch_size, dtype=torch.bool, device="cuda") + b_ready_cache_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + b_prefill_start_loc = torch.arange( + 0, total_token_num, seq_len, dtype=torch.int32, device="cuda" + ) model_input = ModelInput( - batch_size=1, + batch_size=batch_size, total_token_num=total_token_num, - max_q_seq_len=total_token_num, - max_kv_seq_len=total_token_num, + max_q_seq_len=seq_len, + max_kv_seq_len=seq_len, max_cache_len=0, input_ids=input_ids, mem_indexes=mem_indexes, @@ -218,8 +308,9 @@ def warmup(self, model): b_ready_cache_len=b_ready_cache_len, b_prefill_start_loc=b_prefill_start_loc, is_prefill=True, - b_prefill_has_output_cpu=[False], - multimodal_params=[{"images": [], "audios": []}], + b_prefill_has_output_cpu=[False for _ in range(batch_size)], + multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) model_output: ModelOutput = model.forward(model_input) @@ -281,6 +372,7 @@ def warmup_overlap(self, model): is_prefill=True, b_prefill_has_output_cpu=[False], multimodal_params=[{"images": [], "audios": []}], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py new file mode 100644 index 0000000000..ac53255d4f --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py @@ -0,0 +1,261 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ep_scatter_offsets( + padded_tokens_per_expert, + valid_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts: tl.constexpr, + block_experts: tl.constexpr, + block_tokens: tl.constexpr, +): + expert_id = tl.program_id(0) + expert_offsets = tl.arange(0, block_experts) + counts = tl.load( + padded_tokens_per_expert + expert_offsets, + mask=expert_offsets < num_experts, + other=0, + ) + starts = tl.cumsum(counts) - counts + tl.store(expert_start_loc + expert_offsets, starts, mask=expert_offsets < num_experts) + + expert_start = tl.load(expert_start_loc + expert_id) + padded_count = tl.load(padded_tokens_per_expert + expert_id) + valid_count = tl.load(valid_tokens_per_expert + expert_id) + offsets = tl.arange(0, block_tokens) + for token_start in tl.range(0, padded_count, block_tokens, num_stages=4): + token_offsets = token_start + offsets + tl.store( + m_indices + expert_start + token_offsets, + tl.where(token_offsets < valid_count, expert_id, -1), + ) + + +@triton.jit +def _ep_scatter_tokens( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + hidden_size: tl.constexpr, + hidden_size_pad: tl.constexpr, + scale_hidden_size: tl.constexpr, + scale_hidden_size_pad: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_size = tl.num_programs(0) + hidden_offsets = tl.arange(0, hidden_size_pad) + hidden_mask = hidden_offsets < hidden_size + scale_offsets = tl.arange(0, scale_hidden_size_pad) + scale_mask = scale_offsets < scale_hidden_size + + for token_id_int32 in range(start_token_id, total_token_num, grid_size): + token_id = token_id_int32.to(tl.int64) + token = tl.load( + recv_x + token_id * recv_x_stride0 + hidden_offsets * recv_x_stride1, + mask=hidden_mask, + ) + token_scale = tl.load( + recv_x_scale + token_id * recv_x_scale_stride0 + scale_offsets * recv_x_scale_stride1, + mask=scale_mask, + ) + for topk_offset_int32 in tl.range(0, topk_num, 1, num_stages=4): + topk_offset = topk_offset_int32.to(tl.int64) + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_offset * recv_topk_stride1) + if expert_id >= 0: + destination_int32 = tl.atomic_add(expert_start_loc + expert_id, 1) + destination = destination_int32.to(tl.int64) + tl.store( + output_index + token_id * output_index_stride0 + topk_offset * output_index_stride1, + destination_int32, + ) + tl.store( + output_tensor + + destination * output_tensor_stride0 + + hidden_offsets * output_tensor_stride1, + token, + mask=hidden_mask, + ) + tl.store( + output_tensor_scale + + destination * output_tensor_scale_stride0 + + scale_offsets * output_tensor_scale_stride1, + token_scale, + mask=scale_mask, + ) + + +@torch.no_grad() +def ep_scatter( + recv_x: torch.Tensor, + recv_x_scale: torch.Tensor, + recv_topk: torch.Tensor, + padded_tokens_per_expert: torch.Tensor, + valid_tokens_per_expert: torch.Tensor, + expert_start_loc: torch.Tensor, + output_tensor: torch.Tensor, + output_tensor_scale: torch.Tensor, + m_indices: torch.Tensor, + output_index: torch.Tensor, +): + block_tokens = 128 + num_experts = padded_tokens_per_expert.shape[0] + hidden_size = recv_x.shape[1] + scale_hidden_size = recv_x_scale.shape[1] + assert m_indices.shape[0] % block_tokens == 0 + + _ep_scatter_offsets[(num_experts,)]( + padded_tokens_per_expert, + valid_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts=num_experts, + block_experts=triton.next_power_of_2(num_experts), + block_tokens=block_tokens, + num_warps=8, + ) + _ep_scatter_tokens[(min(recv_topk.shape[0], 8192),)]( + recv_topk.shape[0], + expert_start_loc, + recv_x, + recv_x.stride(0), + recv_x.stride(1), + recv_x_scale, + recv_x_scale.stride(0), + recv_x_scale.stride(1), + recv_topk, + recv_topk.stride(0), + recv_topk.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + output_tensor_scale, + output_tensor_scale.stride(0), + output_tensor_scale.stride(1), + output_index, + output_index.stride(0), + output_index.stride(1), + topk_num=recv_topk.shape[1], + hidden_size=hidden_size, + hidden_size_pad=triton.next_power_of_2(hidden_size), + scale_hidden_size=scale_hidden_size, + scale_hidden_size_pad=triton.next_power_of_2(scale_hidden_size), + num_warps=8, + ) + + +@triton.jit +def _ep_gather_kernel( + total_token_num, + input_tensor, + input_tensor_stride0, + input_tensor_stride1, + recv_topk_ids, + recv_topk_ids_stride0, + recv_topk_ids_stride1, + recv_topk_weights, + recv_topk_weights_stride0, + recv_topk_weights_stride1, + input_index, + input_index_stride0, + input_index_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + topk_num: tl.constexpr, + block_hidden: tl.constexpr, +): + hidden_block_int32 = tl.program_id(0) + hidden_block = hidden_block_int32.to(tl.int64) + start_token_int32 = tl.program_id(1) + grid_size = tl.num_programs(1) + hidden_offsets = tl.arange(0, block_hidden) + + for token_int32 in range(start_token_int32, total_token_num, grid_size): + token = token_int32.to(tl.int64) + accumulator = tl.zeros([block_hidden], dtype=tl.float32) + for topk_offset_int32 in range(0, topk_num): + topk_offset = topk_offset_int32.to(tl.int64) + expert_id = tl.load( + recv_topk_ids + token * recv_topk_ids_stride0 + topk_offset * recv_topk_ids_stride1 + ) + if expert_id >= 0: + source_int32 = tl.load( + input_index + token * input_index_stride0 + topk_offset * input_index_stride1 + ) + source = source_int32.to(tl.int64) + weight = tl.load( + recv_topk_weights + + token * recv_topk_weights_stride0 + + topk_offset * recv_topk_weights_stride1 + ) + value = tl.load( + input_tensor + + source * input_tensor_stride0 + + hidden_block * block_hidden + + hidden_offsets * input_tensor_stride1 + ) + accumulator += value.to(tl.float32) * weight + tl.store( + output_tensor + + token * output_tensor_stride0 + + hidden_block * block_hidden + + hidden_offsets * output_tensor_stride1, + accumulator.to(output_tensor.dtype.element_ty), + ) + + +@torch.no_grad() +def ep_gather( + input_tensor: torch.Tensor, + recv_topk_ids: torch.Tensor, + recv_topk_weights: torch.Tensor, + input_index: torch.Tensor, + output_tensor: torch.Tensor, +): + hidden_size = input_tensor.shape[1] + block_hidden = 1024 if hidden_size % 1024 == 0 else 128 + assert hidden_size % block_hidden == 0 + grid = (triton.cdiv(hidden_size, block_hidden), min(output_tensor.shape[0], 1024)) + _ep_gather_kernel[grid]( + output_tensor.shape[0], + input_tensor, + input_tensor.stride(0), + input_tensor.stride(1), + recv_topk_ids, + recv_topk_ids.stride(0), + recv_topk_ids.stride(1), + recv_topk_weights, + recv_topk_weights.stride(0), + recv_topk_weights.stride(1), + input_index, + input_index.stride(0), + input_index.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + topk_num=recv_topk_ids.shape[1], + block_hidden=block_hidden, + num_warps=2, + ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index e10adf7758..ed3bf33cfe 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -1009,6 +1009,7 @@ def fused_experts_impl( layout="blocked", limit=None, alpha=None, + clamp_up_add_one=True, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -1087,6 +1088,7 @@ def fused_experts_impl( intermediate_cache2.view(-1, N // 2), limit=limit, alpha=alpha, + clamp_up_add_one=clamp_up_add_one, layout=layout, ) @@ -1133,6 +1135,7 @@ def inplace_fused_experts_impl( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: fused_experts_impl( hidden_states, @@ -1152,6 +1155,7 @@ def inplace_fused_experts_impl( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -1173,6 +1177,7 @@ def inplace_fused_experts_impl_fake( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: pass @@ -1203,6 +1208,7 @@ def outplace_fused_experts_impl( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: return fused_experts_impl( hidden_states, @@ -1222,6 +1228,7 @@ def outplace_fused_experts_impl( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -1243,6 +1250,7 @@ def outplace_fused_experts_impl_fake( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: return torch.empty_like(hidden_states) @@ -1274,6 +1282,7 @@ def fused_experts( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): if inplace: torch.ops.lightllm.inplace_fused_experts_impl( @@ -1293,6 +1302,7 @@ def fused_experts( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return hidden_states else: @@ -1313,4 +1323,5 @@ def fused_experts( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 58d4d45514..a12da6488f 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -1,5 +1,7 @@ """Fused MoE kernel.""" +import os + import torch import triton import triton.language as tl @@ -12,7 +14,9 @@ ) from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import ( per_token_group_quant_fp8, + tma_align_input_scale, ) +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_legacy_layout import ep_gather, ep_scatter from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ( ep_build_m_indices, ep_compact_metadata, @@ -24,7 +28,7 @@ get_deepep_num_max_dispatch_tokens_per_rank_decode, ) from lightllm.common.triton_utils.autotuner import Autotuner -from lightllm.utils.device_utils import is_sm100_gpu +from lightllm.utils.device_utils import is_sm90_gpu, is_sm100_gpu from lightllm.utils.sgl_utils import HAS_SGL_KERNEL from lightllm.utils.tensor_buffer_manager import TensorBufferManager @@ -51,6 +55,22 @@ def use_sm100_mega_moe(quant_method: Any) -> bool: return is_sm100_gpu() and quant_method.method_name == "fp4fp8-b32-deepgemm" +def use_sm90_mega_moe(quant_method: Any) -> bool: + return ( + is_sm90_gpu() + and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() + in {"1", "ON", "TRUE"} + and quant_method.method_name == "fp8w8a8-b128-deepgemm" + and HAS_DEEPGEMM + and hasattr(deep_gemm, "fp8_mega_moe") + and hasattr(deep_gemm, "mega_moe_pre_dispatch_sm90") + ) + + +def use_mega_moe(quant_method: Any) -> bool: + return use_sm90_mega_moe(quant_method) or use_sm100_mega_moe(quant_method) + + def check_ep_expert_dtype(quant_method: Any): expert_dtype = getattr(quant_method, "method_name", None) if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES: @@ -75,6 +95,9 @@ def masked_group_gemm( w2: torch.Tensor, w2_scale: torch.Tensor, expected_m: int, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): padded_m = recv_x[0].shape[1] E, N, _ = w1.shape @@ -86,7 +109,16 @@ def masked_group_gemm( qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) - silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + gemm_out_a, + qsilu_out, + qsilu_out_scale, + block_size, + masked_m, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) @@ -120,6 +152,82 @@ def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, return stats +def prepare_sm90_mega_moe_weights(w13: Any) -> None: + """Replace standard gate/up FP8 weights with SM90 Mega-MoE layout.""" + + if getattr(w13, "sm90_mega_moe_prepared", False): + return + weight = w13.weight + num_groups, n, *rest = weight.shape + granularity = 8 + half = n // 2 + assert half % granularity == 0 + gate = weight[:, :half].reshape( + num_groups, half // granularity, granularity, *rest + ) + up = weight[:, half:].reshape( + num_groups, half // granularity, granularity, *rest + ) + w13.weight = torch.stack((gate, up), dim=2).reshape( + num_groups, n, *rest + ) + w13.sm90_mega_moe_prepared = True + + +def _sm90_mega_moe_impl( + hidden_states: torch.Tensor, + w13: Any, + w2: Any, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation_clamp: Optional[float], +): + buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) + if buffer is None: + raise RuntimeError("SM90 Mega MoE buffer is not initialized") + if hidden_states.shape[0] > buffer.num_max_tokens_per_rank: + raise RuntimeError( + f"SM90 Mega MoE got {hidden_states.shape[0]} tokens, exceeding " + f"num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" + ) + if not getattr(w13, "sm90_mega_moe_prepared", False): + raise RuntimeError("SM90 Mega MoE weights were not prepared after loading") + + num_tokens = hidden_states.shape[0] + deep_gemm.mega_moe_pre_dispatch_sm90( + hidden_states, + topk_ids.to(torch.int32), + topk_weights.to(torch.float32), + buffer.x, + buffer.x_sf, + buffer.topk_idx, + buffer.topk_weights, + num_tokens=num_tokens, + group_size=128, + routed_scaling_factor=1.0, + ) + # Match DeepGEMM's SM90 integration exactly. In particular, do not pass + # cumulative expert statistics here: the production SGLang path leaves the + # optional pointer null, and the SM90 persistent collective must not gain an + # extra per-layer side effect across repeated decode launches. + output = torch.empty( + (max(num_tokens, 1), hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + deep_gemm.fp8_mega_moe( + output, + (w13.weight, w13.weight_scale), + (w2.weight, w2.weight_scale), + buffer, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=activation_clamp, + fast_math=True, + ) + return output[:num_tokens] + + def mega_moe_impl( hidden_states: torch.Tensor, w13: Any, @@ -127,7 +235,17 @@ def mega_moe_impl( topk_weights: torch.Tensor, topk_ids: torch.Tensor, quant_method: Any, + swiglu_limit: Optional[float] = None, ): + if use_sm90_mega_moe(quant_method): + return _sm90_mega_moe_impl( + hidden_states, + w13, + w2, + topk_weights, + topk_ids, + activation_clamp=swiglu_limit, + ) if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") @@ -174,7 +292,7 @@ def quantize_fused_experts_input( quant_method: Any, ): check_ep_expert_dtype(quant_method) - if use_sm100_mega_moe(quant_method): + if use_mega_moe(quant_method): from deep_gemm.utils import per_token_cast_to_fp8 return per_token_cast_to_fp8( @@ -191,6 +309,114 @@ def quantize_fused_experts_input( return per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w13.weight.dtype) +def legacy_normal_moe_forward( + num_recv_tokens_per_expert_list: List[int], + recv_x: Tuple[torch.Tensor, torch.Tensor], + recv_topk_idx: torch.Tensor, + recv_topk_weights: torch.Tensor, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + hidden_dtype: torch.dtype, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, +) -> torch.Tensor: + """Run MoE on the token layout returned by DeepEP's legacy normal buffer.""" + hidden_size = recv_x[0].shape[1] + intermediate_twice = w1.shape[1] + intermediate_size = intermediate_twice // 2 + block_size_k = w1.shape[2] // w1_scale.shape[2] + all_tokens = sum(num_recv_tokens_per_expert_list) + gather_out = torch.empty( + (recv_x[0].shape[0], hidden_size), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + if all_tokens == 0: + if Autotuner.is_autotune_warmup(): + gemm_out_a = torch.zeros((1, intermediate_twice), device=recv_x[0].device, dtype=hidden_dtype) + silu_out = torch.zeros((1, intermediate_size), device=recv_x[0].device, dtype=hidden_dtype) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) + return gather_out + + input_tensor = torch.empty( + (all_tokens, hidden_size), + device=recv_x[0].device, + dtype=recv_x[0].dtype, + ) + input_scale = torch.empty( + (all_tokens, hidden_size // block_size_k), + device=recv_x[0].device, + dtype=recv_x[1].dtype, + ) + m_indices = torch.empty(all_tokens, device=recv_x[0].device, dtype=torch.int32) + output_index = torch.empty_like(recv_topk_idx, dtype=torch.int32) + padded_counts = torch.tensor( + num_recv_tokens_per_expert_list, + dtype=torch.int32, + pin_memory=True, + device="cpu", + ).cuda(non_blocking=True) + valid_expert_ids = recv_topk_idx[recv_topk_idx >= 0].to(torch.int64) + valid_counts = torch.bincount(valid_expert_ids, minlength=w1.shape[0]).to(torch.int32) + expert_start_loc = torch.empty_like(padded_counts) + ep_scatter( + recv_x[0], + recv_x[1], + recv_topk_idx, + padded_counts, + valid_counts, + expert_start_loc, + input_tensor, + input_scale, + m_indices, + output_index, + ) + + gemm_out_a = torch.empty( + (all_tokens, intermediate_twice), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + input_scale = tma_align_input_scale(input_scale) + deepgemm_grouped_fp8_nt_contiguous((input_tensor, input_scale), (w1, w1_scale), gemm_out_a, m_indices) + silu_out = torch.empty( + (all_tokens, intermediate_size), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) + quant_silu = per_token_group_quant_fp8( + silu_out, + block_size_k, + dtype=w2.dtype, + column_major_scales=True, + scale_tma_aligned=True, + ) + gemm_out_b = torch.empty( + (all_tokens, hidden_size), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + deepgemm_grouped_fp8_nt_contiguous(quant_silu, (w2, w2_scale), gemm_out_b, m_indices) + ep_gather(gemm_out_b, recv_topk_idx, recv_topk_weights, output_index, gather_out) + return gather_out + + def fused_experts( hidden_states: torch.Tensor, w13: Any, @@ -201,10 +427,21 @@ def fused_experts( quant_method: Any, is_prefill: Optional[bool], previous_event: Optional[Any] = None, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): check_ep_expert_dtype(quant_method) - if use_sm100_mega_moe(quant_method): - return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) + if use_mega_moe(quant_method): + return mega_moe_impl( + hidden_states, + w13, + w2, + topk_weights, + topk_idx, + quant_method, + swiglu_limit=swiglu_limit, + ) buffer = dist_group_manager.ep_buffer if is_prefill else dist_group_manager.ep_low_latency_buffer return fused_experts_impl( @@ -222,6 +459,9 @@ def fused_experts( w1_scale=w13.weight_scale, w2_scale=w2.weight_scale, previous_event=previous_event, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_clamp_up_add_one=swiglu_clamp_up_add_one, ) @@ -240,6 +480,9 @@ def fused_experts_impl( w1_scale: Optional[torch.Tensor] = None, w2_scale: Optional[torch.Tensor] = None, previous_event: Optional[Any] = None, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -263,6 +506,64 @@ def fused_experts_impl( if is_prefill: qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w1.dtype) allocate_on_comm_stream = previous_event is not None + if dist_group_manager.ep_prefill_uses_legacy_buffer: + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + layout_event, + ) = buffer.get_dispatch_layout( + topk_idx, + num_experts, + previous_event=previous_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + ( + recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + handle, + _, + ) = buffer.dispatch( + (qinput_tensor, input_scale), + topk_idx=topk_idx, + topk_weights=topk_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=layout_event, + async_finish=False, + allocate_on_comm_stream=False, + expert_alignment=128, + ) + gather_out = legacy_normal_moe_forward( + num_recv_tokens_per_expert_list, + recv_x, + recv_topk_idx, + recv_topk_weights, + w1, + w1_scale, + w2, + w2_scale, + hidden_states.dtype, + swiglu_limit, + swiglu_alpha, + swiglu_clamp_up_add_one, + ) + combined_x, _, _ = buffer.combine( + gather_out, + handle, + topk_weights=None, + previous_event=previous_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + return combined_x + # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: # recv_x[0]: [num_expanded_tokens, hidden] # recv_x[1]: [num_expanded_tokens, hidden // block_size_k], with a @@ -309,6 +610,9 @@ def fused_experts_impl( block_size_k=block_size_k, workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(), hidden_dtype=hidden_states.dtype, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_clamp_up_add_one=swiglu_clamp_up_add_one, ) else: gather_out = torch.empty( @@ -324,7 +628,13 @@ def fused_experts_impl( N = w1.shape[1] _gemm_out_a = torch.zeros((1, N), device=hidden_states.device, dtype=hidden_states.dtype) _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) - silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out) + silu_and_mul_fwd( + _gemm_out_a.view(-1, N), + _silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) _gemm_out_a, _silu_out = None, None del recv_x @@ -350,7 +660,19 @@ def fused_experts_impl( return_recv_hook=False, ) # deepgemm - gemm_out_b = masked_group_gemm(recv_x, masked_m, hidden_states.dtype, w1, w1_scale, w2, w2_scale, expected_m) + gemm_out_b = masked_group_gemm( + recv_x, + masked_m, + hidden_states.dtype, + w1, + w1_scale, + w2, + w2_scale, + expected_m, + swiglu_limit, + swiglu_alpha, + swiglu_clamp_up_add_one, + ) # low latency combine combined_x, event_overlap, hook = buffer.low_latency_combine( gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False @@ -468,6 +790,9 @@ def chunked_expanded_moe_forward( block_size_k: int, workspace: torch.Tensor, # [workspace_bytes], uint8 hidden_dtype: torch.dtype, # scalar dtype descriptor + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" alignment = 128 @@ -535,7 +860,13 @@ def chunked_expanded_moe_forward( gemm_out_a, m_indices[chunk_start:chunk_end], ) - silu_and_mul_fwd(gemm_out_a, silu_out) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) workspace_manager.free(gemm_out_a) del gemm_out_a @@ -588,6 +919,8 @@ def _deepgemm_grouped_fp8_nt_masked( expected_m: int, ): if HAS_DEEPGEMM: + if hasattr(deep_gemm, "fp8_m_grouped_gemm_nt_masked"): + return deep_gemm.fp8_m_grouped_gemm_nt_masked(input_tuple, w_tuple, out, masked_m, expected_m) if hasattr(deep_gemm, "m_grouped_fp8_gemm_nt_masked"): return deep_gemm.m_grouped_fp8_gemm_nt_masked(input_tuple, w_tuple, out, masked_m, expected_m) if hasattr(deep_gemm, "m_grouped_gemm_fp8_fp8_bf16_nt_masked"): diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py index fb0323cd4b..2f30054793 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py @@ -89,6 +89,139 @@ def argsort(x, x_1, ids, dim: tl.core.constexpr = None, descending: tl.core.cons return x, x_1, ids +@triton.jit +def single_group_sigmoid_topk_kernel( + gating_output_ptr, + gating_output_stride_m, + correction_bias_ptr, + out_topk_weights, + out_topk_weights_stride_m, + out_topk_ids, + out_topk_ids_stride_m, + total_expert_num, + HAS_CORRECTION_BIAS: tl.constexpr, + EXPERT_BLOCK_SIZE: tl.constexpr, + TOPK_BLOCK_SIZE: tl.constexpr, + TOPK_NUM: tl.constexpr, + RENORMALIZE: tl.constexpr, +): + """Select biased sigmoid top-k when every expert belongs to one group. + + GLM-5.3-Flash declares one expert group. The generic grouped kernel still + materializes a scratch buffer, synchronizes twice, and bitonic-sorts the + full power-of-two expert block. Repeated reductions are substantially + cheaper for the small fixed top-k used by MoE decode. + """ + + token_index = tl.program_id(axis=0) + offs_n = tl.arange(0, EXPERT_BLOCK_SIZE) + valid = offs_n < total_expert_num + hidden_states = tl.load( + gating_output_ptr + token_index * gating_output_stride_m + offs_n, + mask=valid, + other=-float("inf"), + ).to(tl.float32) + old_scores = tl.sigmoid(hidden_states) + if HAS_CORRECTION_BIAS: + scores = old_scores + tl.load( + correction_bias_ptr + offs_n, mask=valid, other=0.0 + ) + else: + scores = old_scores + scores = tl.where(valid, scores, -float("inf")) + + for topk_index in tl.static_range(0, TOPK_NUM): + selected_index = tl.argmax(scores, axis=0) + selected_weight = tl.sum( + tl.where(offs_n == selected_index, old_scores, 0.0), axis=0 + ) + tl.store( + out_topk_weights + + token_index * out_topk_weights_stride_m + + topk_index, + selected_weight, + ) + tl.store( + out_topk_ids + token_index * out_topk_ids_stride_m + topk_index, + selected_index, + ) + scores = tl.where(offs_n == selected_index, -float("inf"), scores) + + if RENORMALIZE: + topk_offs = tl.arange(0, TOPK_BLOCK_SIZE) + topk_mask = topk_offs < TOPK_NUM + weights = tl.load( + out_topk_weights + + token_index * out_topk_weights_stride_m + + topk_offs, + mask=topk_mask, + other=0.0, + ) + weight_sum = tl.sum(weights, axis=0) + tl.store( + out_topk_weights + + token_index * out_topk_weights_stride_m + + topk_offs, + weights / weight_sum, + mask=topk_mask, + ) + + +@triton.jit +def single_group_sigmoid_topk_bitonic_kernel( + gating_output_ptr, + gating_output_stride_m, + correction_bias_ptr, + out_topk_weights, + out_topk_weights_stride_m, + out_topk_ids, + out_topk_ids_stride_m, + total_expert_num, + HAS_CORRECTION_BIAS: tl.constexpr, + EXPERT_BLOCK_SIZE: tl.constexpr, + TOPK_NUM: tl.constexpr, + RENORMALIZE: tl.constexpr, +): + """Scratch-free single-group path retaining the legacy bitonic order.""" + + token_index = tl.program_id(axis=0) + offs_n = tl.arange(0, EXPERT_BLOCK_SIZE) + hidden_states = tl.load( + gating_output_ptr + token_index * gating_output_stride_m + offs_n, + mask=offs_n < total_expert_num, + other=-10000000.0, + ).to(tl.float32) + old_scores = tl.sigmoid(hidden_states) + if HAS_CORRECTION_BIAS: + scores = old_scores + tl.load( + correction_bias_ptr + offs_n, + mask=offs_n < total_expert_num, + other=-10000000.0, + ) + else: + scores = old_scores + + _, sorted_scores, sorted_indexes = argsort( + scores, old_scores, offs_n, descending=True + ) + if RENORMALIZE: + sum_scores = tl.sum( + tl.where(offs_n < TOPK_NUM, sorted_scores, 0.0) + ) + sorted_scores = sorted_scores / sum_scores + + tl.store( + out_topk_weights + token_index * out_topk_weights_stride_m + offs_n, + sorted_scores, + mask=offs_n < TOPK_NUM, + ) + tl.store( + out_topk_ids + token_index * out_topk_ids_stride_m + offs_n, + sorted_indexes, + mask=offs_n < TOPK_NUM, + ) + + @triton.jit def grouped_topk_kernel( gating_output_ptr, @@ -212,6 +345,7 @@ def triton_grouped_topk( topk_group: int = 0, scoring_func: str = "softmax", group_score_used_topk_num=2, + use_single_group_fast_path: bool = True, ): if correction_bias is not None: @@ -220,6 +354,38 @@ def triton_grouped_topk( has_correction_bias = False token_num, total_expert_num = gating_output.shape + + if ( + use_single_group_fast_path + and num_expert_group == 1 + and topk_group == 1 + and scoring_func == "sigmoid" + ): + out_topk_weights = torch.empty( + (token_num, topk), dtype=torch.float32, device="cuda" + ) + out_topk_ids = torch.empty( + (token_num, topk), dtype=torch.long, device="cuda" + ) + single_group_sigmoid_topk_kernel[(token_num,)]( + gating_output, + gating_output.stride(0), + correction_bias, + out_topk_weights, + out_topk_weights.stride(0), + out_topk_ids, + out_topk_ids.stride(0), + total_expert_num, + HAS_CORRECTION_BIAS=has_correction_bias, + EXPERT_BLOCK_SIZE=triton.next_power_of_2(total_expert_num), + TOPK_BLOCK_SIZE=triton.next_power_of_2(topk), + TOPK_NUM=topk, + RENORMALIZE=renormalize, + num_warps=1, + num_stages=1, + ) + return out_topk_weights, out_topk_ids + if gating_output.dtype == torch.float64: dtype = torch.float64 else: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py index a63d92692e..13d7540cc3 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py @@ -24,6 +24,7 @@ def _silu_and_mul_kernel_fast( NEED_MASK: tl.constexpr, layout: tl.constexpr = "blocked", # "blocked" or "interleaved" USE_LIMIT_AND_ALPHA: tl.constexpr = False, + CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, ): stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) @@ -70,11 +71,9 @@ def _silu_and_mul_kernel_fast( up = tl.minimum(tl.maximum(up, -limit), limit) gate = 1 / (1 + tl.exp(-gate * alpha)) * gate gate = gate.to(input_ptr.dtype.element_ty) - tl.store( - output_ptr + out_offsets, - (up + 1) * gate, - mask=mask, - ) + if CLAMP_UP_ADD_ONE: + up += 1 + tl.store(output_ptr + out_offsets, up * gate, mask=mask) else: if USE_TANH_APPROXIMATE_GELU: # tanh-approx GELU, matching Gemma's gelu_pytorch_tanh MLP. @@ -120,6 +119,7 @@ def silu_and_mul_fwd( layout="blocked", limit=None, alpha=None, + clamp_up_add_one=True, run_config=None, ): assert input.stride(-1) == 1 @@ -171,6 +171,7 @@ def silu_and_mul_fwd( num_warps=num_warps, layout=layout, USE_LIMIT_AND_ALPHA=USE_LIMIT_AND_ALPHA, + CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), ) return diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index aa91f15ed9..827caf0f95 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -24,9 +24,13 @@ def _silu_and_mul_post_quant_kernel( size_n, fp8_max, fp8_min, + limit: tl.constexpr, + alpha: tl.constexpr, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, + USE_LIMIT_AND_ALPHA: tl.constexpr = False, + CLAMP_UP_ADD_ONE: tl.constexpr = True, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -51,7 +55,13 @@ def _silu_and_mul_post_quant_kernel( for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) - if USE_TANH_APPROXIMATE_GELU: + if USE_LIMIT_AND_ALPHA: + gate = tl.minimum(gate, limit) + up = tl.minimum(tl.maximum(up, -limit), limit) + gate = gate / (1 + tl.exp(-gate * alpha)) + if CLAMP_UP_ADD_ONE: + up += 1 + elif USE_TANH_APPROXIMATE_GELU: gate_cubed = gate * gate * gate tanh_arg = 0.7978845608028654 * (gate + 0.044715 * gate_cubed) tanh_val = 2.0 / (1.0 + tl.exp(-2.0 * tanh_arg)) - 1.0 @@ -80,6 +90,9 @@ def silu_and_mul_masked_post_quant_fwd( output_scale: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, + limit=None, + alpha=None, + clamp_up_add_one=True, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -121,6 +134,9 @@ def silu_and_mul_masked_post_quant_fwd( finfo = torch.finfo(torch.float8_e4m3fn) fp8_max = finfo.max fp8_min = -fp8_max + assert (limit is None and alpha is None) or ( + limit is not None and alpha is not None + ) _silu_and_mul_post_quant_kernel[grid]( input, @@ -133,9 +149,13 @@ def silu_and_mul_masked_post_quant_fwd( size_n, fp8_max, fp8_min, + limit=limit, + alpha=alpha, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), + USE_LIMIT_AND_ALPHA=limit is not None, + CLAMP_UP_ADD_ONE=clamp_up_add_one, num_warps=num_warps, ) return diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index 2bf325340f..2afeec8f60 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -1,6 +1,6 @@ # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/causal_conv1d.py -from typing import Optional +from typing import List, Optional import torch @@ -15,6 +15,7 @@ def causal_conv1d_fn( conv_states: Optional[torch.Tensor] = None, activation: Optional[str] = "silu", pad_slot_id: int = -1, + seq_lens_cpu: Optional[List[int]] = None, **kwargs, ): """ @@ -48,6 +49,32 @@ def causal_conv1d_fn( """ if activation not in [None, "silu", "swish"]: raise NotImplementedError("activation must be None, silu, or swish") + # The compiled CUDA kernel requires sequence-contiguous input. KDA's + # projection is token-major, so its dim-major transpose would otherwise + # materialize a large copy at every linear-attention layer. SGLang's + # Triton fallback accepts arbitrary input/output strides and is already a + # runtime dependency in the optimized GLM-5 deployment image. + if x.stride(-1) != 1 and seq_lens_cpu is not None: + try: + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn as causal_conv1d_strided, + ) + except ImportError: + pass + else: + return causal_conv1d_strided( + x, + weight, + bias, + conv_states=conv_states, + query_start_loc=query_start_loc, + seq_lens_cpu=seq_lens_cpu, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + pad_slot_id=pad_slot_id, + ) + from sgl_kernel import causal_conv1d_fwd if x.stride(-1) != 1: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..aa91129088 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -8,8 +8,12 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from .chunk import chunk_gated_delta_rule from .fused_recurrent import fused_recurrent_gated_delta_rule +from .kda import chunk_kda_with_fused_gate, fused_recurrent_kda, kda_safe_gate __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", + "fused_recurrent_kda", + "chunk_kda_with_fused_gate", + "kda_safe_gate", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py index 97933b2ac2..2ab864eaf4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py @@ -24,6 +24,7 @@ { "USE_G": lambda args: args["g"] is not None, "USE_GK": lambda args: args["gk"] is not None, + "USE_EXP2": lambda args: args["use_exp2"], "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "STORE_FINAL_STATE": lambda args: args["ht"] is not None, "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, @@ -43,6 +44,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht, cu_seqlens, chunk_offsets, + use_exp2, T, H: tl.constexpr, Hg: tl.constexpr, @@ -52,6 +54,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( BV: tl.constexpr, USE_G: tl.constexpr, USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, @@ -169,7 +172,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[:, None] + b_h1 *= (tl.exp2(b_gk_last1) if USE_EXP2 else exp(b_gk_last1))[:, None] if K > 64: o_k2 = 64 + o_k1 b_gk_last2 = tl.load( @@ -177,7 +180,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k2 < K), other=0.0, ) - b_h2 *= exp(b_gk_last2)[:, None] + b_h2 *= (tl.exp2(b_gk_last2) if USE_EXP2 else exp(b_gk_last2))[:, None] if K > 128: o_k3 = 128 + o_k1 b_gk_last3 = tl.load( @@ -185,7 +188,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k3 < K), other=0.0, ) - b_h3 *= exp(b_gk_last3)[:, None] + b_h3 *= (tl.exp2(b_gk_last3) if USE_EXP2 else exp(b_gk_last3))[:, None] if K > 192: o_k4 = 192 + o_k1 b_gk_last4 = tl.load( @@ -193,7 +196,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k4 < K), other=0.0, ) - b_h4 *= exp(b_gk_last4)[:, None] + b_h4 *= (tl.exp2(b_gk_last4) if USE_EXP2 else exp(b_gk_last4))[:, None] b_v = b_v.to(k.dtype.element_ty) p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) @@ -264,6 +267,8 @@ def chunk_gated_delta_rule_fwd_h( chunk_size: int = 64, # SY: remove this argument and force chunk size 64? save_new_value: bool = True, cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.Tensor | None = None, + use_exp2: bool = False, run_config=None, ) -> tuple[torch.Tensor, torch.Tensor]: # This kernel is slightly different from fla to support Q/K with different head numbers. @@ -272,7 +277,8 @@ def chunk_gated_delta_rule_fwd_h( H = u.shape[-2] BT = chunk_size - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) # N: the actual number of sequences in the batch with either equal or variable lengths if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None @@ -311,6 +317,7 @@ def chunk_gated_delta_rule_fwd_h( ht=final_state, cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, + use_exp2=use_exp2, T=T, H=H, Hg=Hg, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 5dfbd6e4ab..4af05e31bc 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -8,6 +8,8 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 +import os + import torch import triton @@ -16,6 +18,9 @@ from .op import exp +ENABLE_FAST_MTP_KDA = os.getenv("LIGHTLLM_ENABLE_FAST_MTP_KDA", "0") == "1" + + @triton.heuristics( { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, @@ -45,6 +50,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( a_raw, # [B*T, HV] raw alpha values (before softplus) b_raw, # [B*T, HV] raw beta values (before sigmoid) scale, + kda_lower_bound, N: tl.int64, # num of sequences T: tl.int64, # num of tokens B: tl.constexpr, @@ -105,8 +111,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if FUSE_GATING: # Fused gating: load per-head constants once, compute g/beta inline per token b_A_log = tl.load(A_log + i_hv).to(tl.float32) - b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) - p_a_raw = a_raw + bos * stride_a_tok + i_hv + if IS_KDA: + p_dt_bias = dt_bias + i_hv * K + o_k + p_a_raw = a_raw + bos * stride_a_tok + i_hv * K + o_k + else: + b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) + p_a_raw = a_raw + bos * stride_a_tok + i_hv p_b_raw = b_raw + bos * stride_b_tok + i_hv else: if IS_BETA_HEADWISE: @@ -151,16 +161,23 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_q = b_q * scale # [BK, BV] if FUSE_GATING: - # Compute g = -exp(A_log) * softplus(a_raw + dt_bias) inline - b_a = tl.load(p_a_raw).to(tl.float32) - x = b_a + b_dt_bias - softplus_x = tl.where( - SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, - (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), - x, - ) - b_g = -tl.exp(b_A_log) * softplus_x - b_h *= exp(b_g) + if IS_KDA: + # GLM-5 safe gate is a K-wide vector per value head. + b_a = tl.load(p_a_raw, mask=mask_k, other=0).to(tl.float32) + b_dt_bias = tl.load(p_dt_bias, mask=mask_k, other=0).to(tl.float32) + b_gk = kda_lower_bound * tl.sigmoid(tl.exp(b_A_log) * (b_a + b_dt_bias)) + b_h *= exp(b_gk[:, None]) + else: + # GDN scalar gate: g = -exp(A_log) * softplus(a_raw + dt_bias). + b_a = tl.load(p_a_raw).to(tl.float32) + x = b_a + b_dt_bias + softplus_x = tl.where( + SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, + (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + b_h *= exp(b_g) # Compute beta = sigmoid(b_raw) inline b_b = tl.load(p_b_raw).to(tl.float32) b_beta = tl.sigmoid(b_b) @@ -267,6 +284,8 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + is_kda: bool = False, + kda_lower_bound: float = -5.0, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -277,15 +296,30 @@ def fused_recurrent_gated_delta_rule_fwd( q, stride_q_tok = _ensure_qkv_token_strided(q, H * K) k, stride_k_tok = _ensure_qkv_token_strided(k, H * K) v, stride_v_tok = _ensure_qkv_token_strided(v, HV * V) - a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV) + a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV * K if is_kda else HV) b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw, HV) BK = triton.next_power_of_2(K) + is_spec_verify = ( + ENABLE_FAST_MTP_KDA and cu_seqlens is not None and num_accepted_tokens is not None + ) if T == 1: # Decode path: use larger BV to reduce kernel instances (4 blocks instead of 16) # and more warps for better SM utilization at T=1 where there's no pipelining benefit BV = min(triton.next_power_of_2(V), 32) num_warps = 4 num_stages = 1 + elif is_spec_verify: + # MTP verification is represented as one packed tensor, so ``T`` is the + # total physical token count even though every sequence is only a few + # tokens long. Treating this as prefill creates 16 single-warp value + # tiles per head for GLM-5 (V=128). The short recurrent loop benefits + # from the same BV=32 shape used by SGLang's KDA target-verify kernel, + # cutting the launch grid to four value tiles. Keep this opt-in because + # the wider tile can alter bf16 rounding at the 1e-5 level even though + # the recurrence is mathematically identical. + BV = min(triton.next_power_of_2(V), 32) + num_warps = 1 + num_stages = 3 else: # Prefill path: small BV for better pipelining across sequence length BV = min(triton.next_power_of_2(V), 8) @@ -346,6 +380,7 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw=a_raw, b_raw=b_raw, scale=scale, + kda_lower_bound=kda_lower_bound, N=N, T=T, B=B, @@ -371,7 +406,7 @@ def fused_recurrent_gated_delta_rule_fwd( IS_BETA_HEADWISE=False if fuse_gating else (beta.ndim == v.ndim), USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, - IS_KDA=False, + IS_KDA=is_kda, FUSE_GATING=fuse_gating, num_warps=num_warps, num_stages=num_stages, @@ -402,6 +437,8 @@ def forward( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + is_kda: bool = False, + kda_lower_bound: float = -5.0, ): # q/k/v/a_raw/b_raw may be non-contiguous column views of one projection # output; the kernel handles them via per-token strides (no copies). @@ -424,6 +461,8 @@ def forward( a_raw=a_raw, b_raw=b_raw, out=out, + is_kda=is_kda, + kda_lower_bound=kda_lower_bound, ) return o, final_state @@ -449,6 +488,8 @@ def fused_recurrent_gated_delta_rule( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + is_kda: bool = False, + kda_lower_bound: float = -5.0, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -531,5 +572,7 @@ def fused_recurrent_gated_delta_rule( a_raw, b_raw, out, + is_kda, + kda_lower_bound, ) return o, final_state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py index 8b1d59fc63..d09f3e82a2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py @@ -27,4 +27,7 @@ def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torc @tensor_cache def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: - return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) + # new_tensor([0]) stages a regular CPU tensor before copying it to CUDA, + # which is illegal inside CUDA Graph capture. new_zeros allocates and + # initializes directly on the source tensor's device. + return torch.cat([cu_seqlens.new_zeros(1), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py new file mode 100644 index 0000000000..d2d1e4b877 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -0,0 +1,1170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang + +"""KDA helpers built on LightLLM's continuous-batching recurrent kernel.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .chunk_delta_h import chunk_gated_delta_rule_fwd_h +from .cumsum import chunk_local_cumsum +from .fused_recurrent import ENABLE_FAST_MTP_KDA, fused_recurrent_gated_delta_rule +from .index import prepare_chunk_indices +from .l2norm import l2norm_fwd +from .op import exp2, log +from .solve_tril import solve_tril + + +FLA_CHUNK_SIZE = 64 +RCP_LN2 = 1.4426950216293335 + + +def cdiv(a: int, b: int) -> int: + return -(a // -b) + + +def next_power_of_2(n: int) -> int: + return 1 if n < 1 else 1 << (n - 1).bit_length() + + +def kda_safe_gate( + raw_gate: torch.Tensor, + a_log: torch.Tensor, + gate_bias: torch.Tensor, + lower_bound: float = -5.0, +) -> torch.Tensor: + """GLM-5 bounded KDA decay in fp32. + + ``raw_gate`` is ``[..., heads, key_dim]``; ``a_log`` is per-head and + ``gate_bias`` is per head/key coordinate. + """ + + head_count = a_log.numel() + key_dim = gate_bias.numel() // head_count + gate = raw_gate.float().view(*raw_gate.shape[:-1], head_count, key_dim) + amplitude = a_log.float().reshape( + *((1,) * (gate.ndim - 2)), head_count, 1 + ).exp() + bias = gate_bias.float().reshape( + *((1,) * (gate.ndim - 2)), head_count, key_dim + ) + return lower_bound * torch.sigmoid(amplitude * (gate + bias)) + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + a_log: torch.Tensor, + gate_bias: torch.Tensor, + initial_state: torch.Tensor, + *, + lower_bound: float = -5.0, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + ssm_state_write_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run GLM-5 KDA with bounded decay and sigmoid beta.""" + + if ENABLE_FAST_MTP_KDA: + value_heads = v.shape[2] + key_dim = q.shape[-1] + assert a_log.numel() == value_heads + assert gate_bias.numel() == value_heads * key_dim + # Keep the raw projections in their packed token-major views and form + # the GLM-5 safe gate plus beta inside the recurrent kernel. This + # removes the fp32 cast/add/exp/sigmoid/multiply launches from every KDA + # layer, which is especially important for the model's 34 KDA layers. + return fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=initial_state, + inplace_final_state=inplace_final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=gate_bias, + a_raw=raw_gate.reshape(-1, value_heads * key_dim), + b_raw=raw_beta.reshape(-1, value_heads), + out=out, + is_kda=True, + kda_lower_bound=lower_bound, + ) + + gate = kda_safe_gate(raw_gate, a_log, gate_bias, lower_bound) + beta = raw_beta.float().sigmoid() + return fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + initial_state=initial_state, + inplace_final_state=inplace_final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + out=out, + is_kda=True, + ) +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) + ) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_g = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + b_kt = tl.make_block_ptr( + k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + p_gk = tl.make_block_ptr( + g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=["BK", "BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K + + +def chunk_kda_scaled_dot_kkt_fwd( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + assert K <= 256 + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BC = min(16, BT) + NC = cdiv(BT, BC) + BK = max(next_power_of_2(K), 16) + A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + grid = (NT, NC * NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + return A, Aqk + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, i_v * BV), + (BK, BV), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BD": BD}, num_warps=num_warps) + for BD in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "D", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_cumsum_fwd_kernel( + g, + A, + y, + g_bias, + cu_seqlens, + chunk_indices, + cumsum_scale, + beta, + threshold, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + p_y = tl.make_block_ptr( + y + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + o_d = i_d * BD + tl.arange(0, BD) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + b_a = tl.load(A + i_h).to(tl.float32) + b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) + if SAFE_GATE: + # y = lower_bound * sigmoid(exp(A) * (g + g_bias)); bounded to + # (lower_bound, 0). Mirrors the SGlang safe_gate branch used by GLM5-Next + # checkpoints whose linear_attn_config["safe_gate"] is True. + b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g))) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = b_a * b_softplus + + # Out-of-bounds rows (load returns 0, but softplus/bias can still make + # b_gate non-zero) participate in the dot product. They only contribute to + # out-of-bounds output rows, which are masked away by `boundary_check` on + # the store, so visible output matches unfused gate + chunk-local cumsum. + o_t = tl.arange(0, BT) + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, + safe_gate: bool = False, + lower_bound: float = -5.0, +) -> torch.Tensor: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + + def grid(meta): + return (cdiv(meta["D"], meta["BD"]), NT, B * H) + + kda_gate_cumsum_fwd_kernel[grid]( + g=raw_g, + A=A_log, + y=y, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + beta=beta, + threshold=threshold, + SAFE_GATE=safe_gate, + LOWER_BOUND=lower_bound, + T=T, + H=H, + D=D, + BT=chunk_size, + ) + return y + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + # the intra Aqk is kept in fp32 + # the computation has very marginal effect on the entire throughput + A, Aqk = chunk_kda_scaled_dot_kkt_fwd( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, +): + chunk_size = FLA_CHUNK_SIZE + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + g = fused_kda_gate_chunk_cumsum( + raw_g, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, + **kwargs, +): + """Run chunk KDA from raw gate projection using fused gate+cumsum.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + beta=beta.contiguous(), + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return o, final_state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py index 2f69aa981d..6bfb1db11c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py @@ -13,6 +13,7 @@ from .utils import is_gather_supported exp = tl.exp +exp2 = tl.exp2 log = tl.log log2 = tl.log2 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py index b5b6cfc369..8b0293bfd5 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py @@ -411,6 +411,7 @@ def merge_16x16_to_64x64_inverse_kernel( def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float, ) -> torch.Tensor: """ @@ -433,7 +434,8 @@ def solve_tril( output_dtype = A.dtype if output_dtype is None else output_dtype B, T, H, BT = A.shape - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) Ai = torch.zeros_like(A, dtype=output_dtype) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py new file mode 100644 index 0000000000..8e7e99e1ab --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py @@ -0,0 +1,109 @@ +"""Local greedy statistics for distributed vocabulary shards.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _greedy_sample_stage1_kernel( + logits, + partial_max, + partial_sum, + partial_argmax, + stride_row, + stride_col, + vocab_size: tl.constexpr, + num_blocks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + block = tl.program_id(1) + offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + values = tl.load( + logits + row * stride_row + offsets * stride_col, + mask=offsets < vocab_size, + other=-float("inf"), + ) + values = values.to(tl.float32) + + block_max = tl.max(values, axis=0) + block_sum = tl.sum(tl.exp(values - block_max), axis=0) + block_argmax = tl.argmax(values, axis=0) + block * BLOCK_SIZE + output_offset = row * num_blocks + block + tl.store(partial_max + output_offset, block_max) + tl.store(partial_sum + output_offset, block_sum) + tl.store(partial_argmax + output_offset, block_argmax) + + +@triton.jit +def _greedy_sample_stage2_stats_kernel( + partial_max, + partial_sum, + partial_argmax, + output_stats, + output_argmax, + num_blocks: tl.constexpr, + batch_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < num_blocks + input_offset = row * num_blocks + offsets + block_max = tl.load(partial_max + input_offset, mask=mask, other=-float("inf")) + block_sum = tl.load(partial_sum + input_offset, mask=mask, other=0.0) + block_argmax = tl.load(partial_argmax + input_offset, mask=mask, other=0x7FFFFFFF) + + global_max = tl.max(block_max, axis=0) + global_sum = tl.sum(block_sum * tl.exp(block_max - global_max), axis=0) + candidate_ids = tl.where(block_max == global_max, block_argmax, 0x7FFFFFFF) + global_argmax = tl.min(candidate_ids, axis=0) + tl.store(output_stats + row, global_max) + tl.store(output_stats + batch_size + row, global_max + tl.log(global_sum)) + tl.store(output_argmax + row, global_argmax) + + +def _launch_stage1(logits: torch.Tensor, scratch: torch.Tensor, block_size: int, num_blocks: int) -> None: + batch_size, vocab_size = logits.shape + _greedy_sample_stage1_kernel[(batch_size, num_blocks)]( + logits, + scratch[0], + scratch[1], + scratch[2], + logits.stride(0), + logits.stride(1), + vocab_size=vocab_size, + num_blocks=num_blocks, + BLOCK_SIZE=block_size, + num_warps=8, + ) + + +@torch.no_grad() +def greedy_sample_local_stats(logits: torch.Tensor, alloc_func=torch.empty) -> torch.Tensor: + """Return local max, logsumexp and argmax rows for distributed greedy sampling.""" + + assert logits.ndim == 2 and logits.is_cuda and logits.is_contiguous() + batch_size, vocab_size = logits.shape + block_size = 4096 + num_blocks = triton.cdiv(vocab_size, block_size) + scratch = alloc_func((3, batch_size, num_blocks), dtype=torch.float32, device=logits.device) + # The third FP32 row carries INT32 argmax bits. Keeping one fixed-size + # payload gives the distributed reducer a single collective without losing + # token-id precision through a numeric int-to-float conversion. + output_stats = alloc_func((3, batch_size), dtype=torch.float32, device=logits.device) + + _launch_stage1(logits, scratch, block_size, num_blocks) + _greedy_sample_stage2_stats_kernel[(batch_size,)]( + scratch[0], + scratch[1], + scratch[2], + output_stats, + output_stats[2].view(torch.int32), + num_blocks=num_blocks, + batch_size=batch_size, + BLOCK_SIZE=triton.next_power_of_2(num_blocks), + num_warps=4, + ) + return output_stats diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py new file mode 100644 index 0000000000..d7dea3e861 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py @@ -0,0 +1,126 @@ +"""Greedy sampling directly from tensor-parallel vocabulary shards.""" + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) +from lightllm.common.basemodel.triton_kernel.transpose_convert import ( + transpose_convert_2d, +) +from lightllm.distributed.communication_op import all_gather_into_tensor +from lightllm.utils.envs_utils import enable_env_vars + + +VOCAB_PARALLEL_GREEDY_ENV = "LIGHTLLM_VOCAB_PARALLEL_GREEDY" + + +def is_vocab_parallel_greedy_enabled() -> bool: + return enable_env_vars(VOCAB_PARALLEL_GREEDY_ENV) + + +@triton.jit +def _combine_vocab_parallel_stats_kernel( + gathered_stats, + gathered_argmax, + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size: tl.constexpr, + tp_world_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < token_num + rank_stride = 3 * token_num + + global_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) + global_id = tl.full((BLOCK_SIZE,), 0x7FFFFFFF, tl.int32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_max = tl.load( + gathered_stats + rank_base + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + local_id = tl.load( + gathered_argmax + rank_base + 2 * token_num + token_offsets, + mask=token_mask, + other=0x7FFFFFFF, + ) + local_id += (rank * vocab_size) // tp_world_size + wins = (local_max > global_max) | ((local_max == global_max) & (local_id < global_id)) + global_max = tl.where(wins, local_max, global_max) + global_id = tl.where(wins, local_id, global_id) + + global_sum = tl.zeros((BLOCK_SIZE,), tl.float32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_lse = tl.load( + gathered_stats + rank_base + token_num + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + global_sum += tl.exp(local_lse - global_max) + + tl.store(output_logits + token_offsets, global_max, mask=token_mask) + tl.store(output_token_ids + token_offsets, global_id, mask=token_mask) + tl.store(output_logsumexp + token_offsets, global_max + tl.log(global_sum), mask=token_mask) + + +@torch.no_grad() +def vocab_parallel_greedy( + local_logits: torch.Tensor, + *, + vocab_size: int, + tp_world_size: int, + group, + alloc_func, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return exact sparse logits, global token ids and full-vocab logsumexp.""" + + assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() + local_vocab_size, token_num = local_logits.shape + assert local_vocab_size in { + vocab_size // tp_world_size, + (vocab_size + tp_world_size - 1) // tp_world_size, + } + + transposed_logits = alloc_func( + (token_num, local_vocab_size), + dtype=local_logits.dtype, + device=local_logits.device, + ) + transpose_convert_2d(local_logits, transposed_logits) + local_stats = greedy_sample_local_stats(transposed_logits, alloc_func=alloc_func) + + if tp_world_size == 1: + gathered_stats = local_stats.view(1, 3, token_num) + else: + gathered_stats = alloc_func((tp_world_size, 3, token_num), dtype=torch.float32, device=local_logits.device) + all_gather_into_tensor( + output_=gathered_stats, + input_=local_stats, + group=group, + async_op=False, + ) + + output_logits = alloc_func((token_num, 1), dtype=torch.float32, device=local_logits.device) + output_token_ids = alloc_func((token_num, 1), dtype=torch.int64, device=local_logits.device) + output_logsumexp = alloc_func((token_num,), dtype=torch.float32, device=local_logits.device) + _combine_vocab_parallel_stats_kernel[(triton.cdiv(token_num, 256),)]( + gathered_stats, + gathered_stats.view(torch.int32), + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size=vocab_size, + tp_world_size=tp_world_size, + BLOCK_SIZE=256, + num_warps=4, + ) + return output_logits, output_token_ids, output_logsumexp diff --git a/lightllm/common/basemodel/triton_kernel/transpose_convert.py b/lightllm/common/basemodel/triton_kernel/transpose_convert.py new file mode 100644 index 0000000000..b618d69526 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/transpose_convert.py @@ -0,0 +1,65 @@ +"""Tiled transpose kernels used by the post-layer logits path.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _transpose_convert_2d_kernel( + input_ptr, + output_ptr, + rows, + cols, + input_stride_0, + input_stride_1, + output_stride_0, + output_stride_1, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + row_offsets = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + col_offsets = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + input_offsets = row_offsets[:, None] * input_stride_0 + col_offsets[None, :] * input_stride_1 + mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) + values = tl.load(input_ptr + input_offsets, mask=mask) + + output_offsets = col_offsets[:, None] * output_stride_0 + row_offsets[None, :] * output_stride_1 + tl.store(output_ptr + output_offsets, tl.trans(values), mask=tl.trans(mask)) + + +@torch.no_grad() +def transpose_convert_2d( + input_tensor: torch.Tensor, + output_tensor: torch.Tensor, + *, + block_rows: int = 64, + block_cols: int = 64, + num_warps: int = 8, + num_stages: int = 1, +) -> torch.Tensor: + """Transpose a contiguous 2-D CUDA tensor while converting its dtype.""" + + assert input_tensor.is_cuda and output_tensor.is_cuda + assert input_tensor.device == output_tensor.device + assert input_tensor.ndim == 2 and output_tensor.ndim == 2 + assert output_tensor.shape == (input_tensor.shape[1], input_tensor.shape[0]) + assert input_tensor.is_contiguous() and output_tensor.is_contiguous() + + rows, cols = input_tensor.shape + grid = (triton.cdiv(rows, block_rows), triton.cdiv(cols, block_cols)) + _transpose_convert_2d_kernel[grid]( + input_tensor, + output_tensor, + rows, + cols, + input_tensor.stride(0), + input_tensor.stride(1), + output_tensor.stride(0), + output_tensor.stride(1), + BLOCK_ROWS=block_rows, + BLOCK_COLS=block_cols, + num_warps=num_warps, + num_stages=num_stages, + ) + return output_tensor diff --git a/lightllm/common/linear_att_cache_manager/config_objs.py b/lightllm/common/linear_att_cache_manager/config_objs.py index f588ec7d5c..adc90f9486 100644 --- a/lightllm/common/linear_att_cache_manager/config_objs.py +++ b/lightllm/common/linear_att_cache_manager/config_objs.py @@ -113,7 +113,14 @@ def load_from_args() -> "LinearAttCacheConfig": model_cfg, _ = PretrainedConfig.get_config_dict(model_path) model_type = model_cfg["model_type"] - assert model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"] + assert model_type in [ + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "glm5_next", + "glm5_next_text", + ] llm_config = model_cfg try: llm_config = llm_config["text_config"] @@ -122,6 +129,30 @@ def load_from_args() -> "LinearAttCacheConfig": n_layer = llm_config["num_hidden_layers"] + if model_type in ["glm5_next", "glm5_next_text"]: + linear = llm_config["linear_attn_config"] + tp_world_size = get_env_start_args().tp // get_env_start_args().dp + return LinearAttCacheConfig( + tp_world_size=tp_world_size, + full_att_all_num_kv_heads=1, + full_att_dtype=get_torch_dtype(args.data_type), + full_att_num_kv_heads=1, + full_att_head_dim=llm_config["kv_lora_rank"], + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + num_linear_k_heads=linear["num_heads"] // tp_world_size, + num_linear_v_heads=linear["num_heads"] // tp_world_size, + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + linear_layer_num=len(linear["kda_layers"]), + conv_state_dtype=get_torch_dtype(args.data_type), + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=4, + all_layer_num=n_layer, + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), + ) + tp_world_size = get_env_start_args().tp // get_env_start_args().dp return LinearAttCacheConfig( tp_world_size=tp_world_size, diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 93c603212d..4d95e68e97 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -38,7 +38,7 @@ create_new_group_for_current_dp, create_dp_special_inter_group, ) -from lightllm.utils.device_utils import get_device_sm_count, is_sm100_gpu +from lightllm.utils.device_utils import get_device_sm_count, is_sm90_gpu, is_sm100_gpu from lightllm.utils.torch_dtype_utils import get_torch_dtype logger = init_logger(__name__) @@ -57,6 +57,9 @@ class CustomProcessGroup: def __init__(self): self.symm_mem_reduce = None self.flashinfer_reduce = None + self.symm_mem_out_of_place = os.getenv( + "LIGHTLLM_SYMM_MEM_OUT_OF_PLACE", "0" + ).upper() in {"1", "ON", "TRUE"} self.dp_world_size = get_dp_world_size() self.device_group = create_new_group_for_current_dp("nccl") if get_env_start_args().enable_dp_prefill_balance: @@ -97,7 +100,10 @@ def all_reduce(self, input_: torch.Tensor) -> None: input_.data = self.flashinfer_reduce.all_reduce(input_) return if self.symm_mem_reduce is not None and self.symm_mem_reduce.should_use(input_): - self.symm_mem_reduce.all_reduce(input_) + if self.symm_mem_out_of_place: + input_.data = self.symm_mem_reduce.all_reduce_out_of_place(input_) + else: + self.symm_mem_reduce.all_reduce(input_) return return dist.all_reduce(input_, group=self.device_group) @@ -109,6 +115,7 @@ class DistributeGroupManager: def __init__(self): self.groups = [] self.ep_buffer = None + self.ep_prefill_uses_legacy_buffer = False self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None self.ep_num_sms = None @@ -175,6 +182,7 @@ def new_deepep_group( decode_num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode() if not enable_ep_moe: self.ep_buffer = None + self.ep_prefill_uses_legacy_buffer = False self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None self.ep_num_sms = None @@ -194,14 +202,38 @@ def new_deepep_group( self.ll_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank self.ll_hidden = hidden_size self.ll_num_experts = n_routed_experts + get_redundancy_expert_num() * global_world_size - self.ep_buffer = deep_ep.ElasticBuffer( - deepep_group, - num_max_tokens_per_rank=self.ll_num_tokens, - hidden=self.ll_hidden, - num_topk=num_experts_per_tok, - use_fp8_dispatch=True, - allow_multiple_reduction=True, - ) + # ElasticBuffer uses NCCL GIN even for an intra-node group. GIN is not + # available on many otherwise fully-connected NVLink hosts. The legacy + # normal DeepEP buffer is the native intra-node transport and avoids that + # unnecessary network dependency; keep ElasticBuffer for multi-node EP. + self.ep_prefill_uses_legacy_buffer = get_env_start_args().nnodes == 1 + if self.ep_prefill_uses_legacy_buffer: + hidden_bytes = self.ll_hidden * get_torch_dtype(get_env_start_args().data_type).itemsize + dispatch_config = deep_ep.Buffer.get_dispatch_config(global_world_size) + combine_config = deep_ep.Buffer.get_combine_config(global_world_size) + num_nvl_bytes = max( + dispatch_config.get_nvl_buffer_size_hint(hidden_bytes, global_world_size), + combine_config.get_nvl_buffer_size_hint(hidden_bytes, global_world_size), + ) + num_rdma_bytes = max( + dispatch_config.get_rdma_buffer_size_hint(hidden_bytes, global_world_size), + combine_config.get_rdma_buffer_size_hint(hidden_bytes, global_world_size), + ) + self.ep_buffer = deep_ep.Buffer( + deepep_group, + num_nvl_bytes=num_nvl_bytes, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=False, + ) + else: + self.ep_buffer = deep_ep.ElasticBuffer( + deepep_group, + num_max_tokens_per_rank=self.ll_num_tokens, + hidden=self.ll_hidden, + num_topk=num_experts_per_tok, + use_fp8_dispatch=True, + allow_multiple_reduction=True, + ) self.ep_mega_moe_buffer = None self.ep_low_latency_buffer = None @@ -209,7 +241,13 @@ def new_deepep_group( raise ValueError("No valid MoE quant method was found while initializing DeepEP buffers") mega_moe_quant_method = "fp4fp8-b32-deepgemm" + sm90_mega_moe_quant_method = "fp8w8a8-b128-deepgemm" is_sm100 = is_sm100_gpu() + enable_sm90_mega_moe = ( + is_sm90_gpu() + and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() + in {"1", "ON", "TRUE"} + ) # Buffer 选择规则： # 1. 非 SM100 不支持 Mega MoE，只初始化 legacy low-latency buffer； @@ -225,6 +263,14 @@ def new_deepep_group( ) enable_mega_moe_buffer = has_mega_moe_layer enable_low_latency_buffer = has_legacy_moe_layer + elif enable_sm90_mega_moe: + has_mega_moe_layer = sm90_mega_moe_quant_method in expert_quant_method_names + has_legacy_moe_layer = any( + method_name != sm90_mega_moe_quant_method + for method_name in expert_quant_method_names + ) + enable_mega_moe_buffer = has_mega_moe_layer + enable_low_latency_buffer = has_legacy_moe_layer else: enable_mega_moe_buffer = False enable_low_latency_buffer = True @@ -265,6 +311,9 @@ def new_deepep_group( import deep_gemm + mega_moe_kwargs = {} + if enable_sm90_mega_moe: + mega_moe_kwargs.update(use_fp8_dispatch=True, activation="swiglu") self.ep_mega_moe_buffer = deep_gemm.get_symm_buffer_for_mega_moe( deepep_group, self.ll_num_experts, @@ -272,14 +321,25 @@ def new_deepep_group( num_experts_per_tok, self.ll_hidden, moe_intermediate_size, + **mega_moe_kwargs, ) logger.info( - "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, expert_quant_method_names=%s", + "Initialize DeepEP MoE buffers: legacy_prefill=%s, low_latency=%s, mega_moe=%s, " + "expert_quant_method_names=%s", + self.ep_prefill_uses_legacy_buffer, enable_low_latency_buffer, enable_mega_moe_buffer, sorted(expert_quant_method_names), ) - theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) + theoretical_sms = ( + 0 + if enable_mega_moe_buffer and not enable_low_latency_buffer + else ( + deep_ep.Buffer.num_sms + if self.ep_prefill_uses_legacy_buffer + else self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) + ) + ) self._set_num_sms_for_deep_gemm(theoretical_sms) def _set_num_sms_for_deep_gemm(self, deepep_sms: int): diff --git a/lightllm/distributed/symm_mem_all_reduce.py b/lightllm/distributed/symm_mem_all_reduce.py index 2256a5093f..48f4204a87 100644 --- a/lightllm/distributed/symm_mem_all_reduce.py +++ b/lightllm/distributed/symm_mem_all_reduce.py @@ -55,7 +55,16 @@ def __init__(self, group: ProcessGroup, device, dtype: torch.dtype = torch.bfloa self.use_multimem = self.world_size in _WORLD_SIZES_MULTIMEM.get(cap_str, []) try: - self.buffer = torch_symm_mem.empty(self.max_size // dtype.itemsize, device=device, dtype=dtype) + # The optional out-of-place path adopts this storage through + # ``Tensor.data``. PyTorch cannot safely mix a normal tensor's + # version-counter state into an inference tensor, so make the + # reusable communication workspace explicitly inference-only. + with torch.inference_mode(): + self.buffer = torch_symm_mem.empty( + self.max_size // dtype.itemsize, + device=device, + dtype=dtype, + ) handle = torch_symm_mem.rendezvous(self.buffer, group.group_name) except RuntimeError as e: logger.warning("SymmMemAllreduce: rendezvous failed (%s). Disabling.", e) @@ -82,11 +91,33 @@ def should_use(self, inp: torch.Tensor) -> bool: # CustomProcessGroup.all_reduce: FlashInfer claims small messages first. return nbytes < self.max_size + def _reduce_to_workspace(self, inp: torch.Tensor) -> torch.Tensor: + # Mutating an inference tensor is only legal inside inference mode. + # Enter it explicitly as capability probes also call this helper from + # ordinary no-grad code. + with torch.inference_mode(): + n = inp.numel() + output = self.buffer[:n] + output.copy_(inp.view(-1)) + if self.use_multimem: + torch.ops.symm_mem.multimem_all_reduce_( + output, "sum", self.group.group_name + ) + else: + torch.ops.symm_mem.two_shot_all_reduce_( + output, "sum", self.group.group_name + ) + return output.view_as(inp) + def all_reduce(self, inp: torch.Tensor) -> None: - n = inp.numel() - self.buffer[:n].copy_(inp.view(-1)) - if self.use_multimem: - torch.ops.symm_mem.multimem_all_reduce_(self.buffer[:n], "sum", self.group.group_name) - else: - torch.ops.symm_mem.two_shot_all_reduce_(self.buffer[:n], "sum", self.group.group_name) - inp.view(-1).copy_(self.buffer[:n]) + inp.copy_(self._reduce_to_workspace(inp)) + + def all_reduce_out_of_place(self, inp: torch.Tensor) -> torch.Tensor: + """Reduce into symmetric memory and return its shaped workspace view. + + The result aliases the single reusable communication workspace. This + is suitable for the model's serialized current-stream all-reduces: + each result is consumed before the next reduction overwrites it. + """ + + return self._reduce_to_workspace(inp) diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py index c7e9a59aad..b464f65340 100644 --- a/lightllm/models/__init__.py +++ b/lightllm/models/__init__.py @@ -21,6 +21,8 @@ from lightllm.models.deepseek2.model import Deepseek2TpPartModel from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel from lightllm.models.glm4_moe_lite.model import Glm4MoeLiteTpPartModel +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel from lightllm.models.internvl.model import ( InternVLLlamaTpPartModel, InternVLPhi3TpPartModel, diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 3254031056..962e72d1c7 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -8,7 +8,7 @@ from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - use_sm100_mega_moe, + use_mega_moe, ) from functools import partial from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale @@ -213,6 +213,12 @@ def _get_o( o_tensor = self._tpsp_reduce(input=o_tensor, infer_state=infer_state) return o_tensor + def _shared_ffn_tp(self, input, infer_state, layer_weight): + """Shared-expert FFN hook for model-specific activation semantics.""" + return LlamaTransformerLayerInfer._ffn_tp( + self, input, infer_state, layer_weight + ) + def _moe_ffn_tp( self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight ) -> torch.Tensor: @@ -222,7 +228,9 @@ def _moe_ffn_tp( # if fused_shared_experts is not enabled, compute shared_output if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: - shared_output = LlamaTransformerLayerInfer._ffn_tp(self, hidden_states, infer_state, layer_weight) + shared_output = self._shared_ffn_tp( + hidden_states, infer_state, layer_weight + ) moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) @@ -234,6 +242,7 @@ def _moe_ffn_tp( use_grouped_topk=self.n_group, topk_group=self.topk_group, num_expert_group=self.n_group, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) @@ -249,7 +258,9 @@ def _moe_ffn_edp( hidden_states = input token_num, hidden_dim = hidden_states.shape if self.n_shared_experts is not None: - shared_output = LlamaTransformerLayerInfer._ffn_tp(self, hidden_states, infer_state, layer_weight) + shared_output = self._shared_ffn_tp( + hidden_states, infer_state, layer_weight + ) moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) @@ -300,7 +311,7 @@ def overlap_tpsp_token_forward( infer_state1: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_token_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -324,7 +335,9 @@ def overlap_tpsp_token_forward( # 0 shared expert if self.n_shared_experts is not None: - _0_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _0_input1, infer_state, layer_weight) + _0_shared_output = self._shared_ffn_tp( + _0_input1, infer_state, layer_weight + ) # 0 dispatch ( @@ -359,7 +372,9 @@ def overlap_tpsp_token_forward( # 1 shared expert if self.n_shared_experts is not None: - _1_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _1_input1, infer_state1, layer_weight) + _1_shared_output = self._shared_ffn_tp( + _1_input1, infer_state1, layer_weight + ) # 1 dispatch ( @@ -426,7 +441,7 @@ def overlap_tpsp_context_forward( infer_state1: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_context_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -495,17 +510,21 @@ def overlap_tpsp_context_forward( # 0 shared expert if self.n_shared_experts is not None: - _0_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _0_input1, infer_state, layer_weight) + _0_shared_output = self._shared_ffn_tp( + _0_input1, infer_state, layer_weight + ) # 1 shared expert if self.n_shared_experts is not None: - _1_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _1_input1, infer_state1, layer_weight) + _1_shared_output = self._shared_ffn_tp( + _1_input1, infer_state1, layer_weight + ) # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( _0_num_recv_tokens_per_expert_list, - _0_handle.num_unaligned_recv_tokens_per_expert, - _0_handle.recv_src_metadata, + getattr(_0_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_0_handle, "recv_src_metadata", None), _0_recv_x, _0_recv_topk_idx, _0_recv_topk_weight, @@ -536,8 +555,8 @@ def overlap_tpsp_context_forward( # 1 moe calc _1_moe_out = layer_weight.experts.prefilled_group_gemm( _1_num_recv_tokens_per_expert_list, - _1_handle.num_unaligned_recv_tokens_per_expert, - _1_handle.recv_src_metadata, + getattr(_1_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_1_handle, "recv_src_metadata", None), _1_recv_x, _1_recv_topk_idx, _1_recv_topk_weight, diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index d6eaebe2fd..bfdc789914 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -8,7 +8,10 @@ from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import destindex_copy_indexer_ks -from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import extract_indexer_ks +from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import ( + extract_indexer_ks, + extract_indexer_ks_dynamic, +) from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import all_gather_into_tensor @@ -167,6 +170,11 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.index_n_heads_scale = (self.index_n_heads ** -0.5) * self.softmax_scale self.tp_world_size_ = tp_world_size self.tp_index_n_heads = self.index_n_heads // self.tp_world_size_ + # Most NSA models only instantiate the target model, so their decode + # layout follows the process-wide MTP setting. A model that reuses an + # NSA layer as a recurrent drafter can override this with its own + # physical decode width (normally zero extra rows). + self.decode_mtp_step = None def _get_indices( self, @@ -181,7 +189,7 @@ def _get_indices( if self.tp_world_size_ > 1: q_merge = torch.empty( - size=(self.tp_world_size_ * q.numel()), + size=(self.tp_world_size_ * q.numel(),), dtype=q.dtype, device=q.device, ) @@ -196,11 +204,12 @@ def _get_indices( q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) + indexer_k_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_) destindex_copy_indexer_ks( K_fp8=k_fp8, K_scale=k_scale, DestLoc=infer_state.mem_index, - O_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + O_buffer=indexer_k_buffer, ) weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale @@ -213,17 +222,39 @@ def _get_indices( if infer_state.is_prefill: mtp_step = 0 else: - mtp_step = get_env_start_args().mtp_step - # Use efficient Triton kernel to extract FP8 keys and scales from buffer - k_fp8_, k_scale_ = extract_indexer_ks( - I_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), - b_seq_len=infer_state.b_seq_len, - b_req_idx=infer_state.b_req_idx, - req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, - out_token_num=infer_state.b_seq_len.shape[0] * infer_state.max_kv_seq_len, - max_kv_seq_len=infer_state.max_kv_seq_len, - mtp_step=mtp_step, + mtp_step = ( + get_env_start_args().mtp_step + if self.decode_mtp_step is None + else self.decode_mtp_step + ) + # LightSpec compacts each request to a variable number of contiguous + # verify rows. Its sparse-index K packing must follow request boundaries + # instead of assuming the fixed process-wide MTP width. + use_dynamic_layout = ( + not infer_state.is_prefill + and mtp_step > 0 + and get_env_start_args().mtp_dynamic_verify ) + if use_dynamic_layout: + k_fp8_, k_scale_ = extract_indexer_ks_dynamic( + I_buffer=indexer_k_buffer, + b_seq_len=infer_state.b_seq_len, + b_req_idx=infer_state.b_req_idx, + b_mtp_index=infer_state.b_mtp_index, + req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, + max_kv_seq_len=infer_state.max_kv_seq_len, + max_request_num=infer_state.req_manager.max_request_num, + ) + else: + k_fp8_, k_scale_ = extract_indexer_ks( + I_buffer=indexer_k_buffer, + b_seq_len=infer_state.b_seq_len, + b_req_idx=infer_state.b_req_idx, + req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, + out_token_num=infer_state.b_seq_len.shape[0] * infer_state.max_kv_seq_len, + max_kv_seq_len=infer_state.max_kv_seq_len, + mtp_step=mtp_step, + ) import deep_gemm @@ -244,12 +275,16 @@ def _get_indices( lengths=lengths, topk=self.index_topk, ) - b_topk_index = torch.where(b_topk_index != -1, b_topk_index + ks.view(-1, 1), -1) + # The long-prefill score matrix can be tens of GiB. fast_topk_v2 has + # already consumed it, so release its storage before materializing the + # 2048-wide global and memory-index tables below. + del logits # 将 topk index 转化为 mem index from ..triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index b_topk_mem_index = trans_topk_index_to_mem_index( topk_index=b_topk_index, + ragged_start_index=ks, ragged_mem_index=att_state.ragged_mem_index, ) diff --git a/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py b/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py index f02fc30942..4b2dc0e067 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py +++ b/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py @@ -43,7 +43,12 @@ def _fwd_kernel_extract_indexer_ks( store_start_index = tl.sum(b_seq_len) for i in range(token_start_index, cur_seq_len, tl.num_programs(1)): - mem_index = tl.load(req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n) + # Token slots can exceed INT32_MAX / stride_in_fp8_bs when the KV + # cache packs multiple states into one row. Promote before pointer + # arithmetic so large cache offsets do not wrap around. + mem_index = tl.load( + req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n + ).to(tl.int64) in_fp8_ptrs = in_fp8 + mem_index * stride_in_fp8_bs + 0 * stride_in_fp8_h + stride_in_fp8_d * offs_d kv_fp8 = tl.load(in_fp8_ptrs) @@ -60,6 +65,86 @@ def _fwd_kernel_extract_indexer_ks( return +@triton.jit +def _fwd_kernel_extract_indexer_ks_dynamic( + in_fp8, + stride_in_fp8_bs, + stride_in_fp8_h, + stride_in_fp8_d, + in_fp8_scale, + stride_in_scale_bs, + stride_in_scale_h, + stride_in_scale_d, + req_to_token_indexs, + stride_req_to_token_m, + stride_req_to_token_n, + b_seq_len, + b_req_idx, + b_mtp_index, + batch_size, + O_fp8, + stride_o_fp8_bs, + stride_o_fp8_d, + O_scale, + stride_o_scale_bs, + stride_o_scale_d, + BLOCK_DMODEL: tl.constexpr, + BLOCK_BATCH: tl.constexpr, +): + """Extract one packed K sequence per variable-width MTP request. + + Dynamic verification keeps each request's selected rows contiguous and + starts every request (and every CUDA-graph padding row) at MTP index zero. + A request's last row therefore has either no successor or a successor with + MTP index zero. Using that boundary instead of a process-wide fixed width + lets the sparse indexer consume LightSpec's compacted verify layout. + """ + + cur_row = tl.program_id(0) + token_start_index = tl.program_id(1) + next_mtp_index = tl.load( + b_mtp_index + cur_row + 1, + mask=cur_row + 1 < batch_size, + other=0, + ) + is_group_end = (cur_row == batch_size - 1) | (next_mtp_index == 0) + if not is_group_end: + return + + cur_req_idx = tl.load(b_req_idx + cur_row) + cur_seq_len = tl.load(b_seq_len + cur_row) + + rows = tl.arange(0, BLOCK_BATCH) + next_rows = rows + 1 + next_mtp_indices = tl.load( + b_mtp_index + next_rows, + mask=next_rows < batch_size, + other=0, + ) + prior_group_ends = (rows < cur_row) & (next_mtp_indices == 0) + prior_group_seq_lens = tl.load( + b_seq_len + rows, + mask=prior_group_ends, + other=0, + ) + store_start_index = tl.sum(prior_group_seq_lens) + + offs_d = tl.arange(0, BLOCK_DMODEL) + for i in range(token_start_index, cur_seq_len, tl.num_programs(1)): + mem_index = tl.load( + req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n + ).to(tl.int64) + + in_fp8_ptrs = in_fp8 + mem_index * stride_in_fp8_bs + stride_in_fp8_d * offs_d + kv_fp8 = tl.load(in_fp8_ptrs) + in_scale_ptr = in_fp8_scale + mem_index * stride_in_scale_bs + kv_scale = tl.load(in_scale_ptr) + + o_fp8_ptrs = O_fp8 + (store_start_index + i) * stride_o_fp8_bs + stride_o_fp8_d * offs_d + tl.store(o_fp8_ptrs, kv_fp8) + tl.store(O_scale + (store_start_index + i) * stride_o_scale_bs, kv_scale) + + @torch.no_grad() def extract_indexer_ks( I_buffer: torch.Tensor, @@ -113,3 +198,65 @@ def extract_indexer_ks( ) return O_fp8, O_scale.squeeze(-1) + + +@torch.no_grad() +def extract_indexer_ks_dynamic( + I_buffer: torch.Tensor, + b_seq_len: torch.Tensor, + b_req_idx: torch.Tensor, + b_mtp_index: torch.Tensor, + req_to_token_indexs: torch.Tensor, + max_kv_seq_len: int, + max_request_num: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Extract packed indexer keys for a variable-width MTP verify batch.""" + + head_dim = 128 + batch_size = b_seq_len.shape[0] + assert I_buffer.dtype == torch.uint8, f"Expected I_buffer dtype=uint8, got {I_buffer.dtype}" + assert I_buffer.shape[2] == 132, f"Expected I_buffer last dim=132, got {I_buffer.shape[2]}" + assert b_req_idx.shape == b_seq_len.shape == b_mtp_index.shape + + in_fp8 = I_buffer[:, :, 0:128].view(dtype=torch.float8_e4m3fn) + in_fp8_scale = I_buffer[:, :, 128:132].view(dtype=torch.float32) + + # At most max_request_num rows belong to distinct real requests. Any + # remaining CUDA-graph padding rows have a fixed sequence length of two. + # This static upper bound avoids allocating batch_size * max_kv_seq_len for + # every captured graph while remaining safe for every compacted layout. + max_real_groups = min(batch_size, max_request_num) + output_capacity = max_real_groups * max_kv_seq_len + (batch_size - max_real_groups) * 2 + O_fp8 = torch.empty((output_capacity, head_dim), dtype=torch.float8_e4m3fn, device=I_buffer.device) + O_scale = torch.empty((output_capacity, 1), dtype=torch.float32, device=I_buffer.device) + + grid = (batch_size, min(256, max_kv_seq_len)) + _fwd_kernel_extract_indexer_ks_dynamic[grid]( + in_fp8, + stride_in_fp8_bs=in_fp8.stride(0), + stride_in_fp8_h=in_fp8.stride(1), + stride_in_fp8_d=in_fp8.stride(2), + in_fp8_scale=in_fp8_scale, + stride_in_scale_bs=in_fp8_scale.stride(0), + stride_in_scale_h=in_fp8_scale.stride(1), + stride_in_scale_d=in_fp8_scale.stride(2), + req_to_token_indexs=req_to_token_indexs, + stride_req_to_token_m=req_to_token_indexs.stride(0), + stride_req_to_token_n=req_to_token_indexs.stride(1), + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + batch_size=batch_size, + O_fp8=O_fp8, + stride_o_fp8_bs=O_fp8.stride(0), + stride_o_fp8_d=O_fp8.stride(1), + O_scale=O_scale, + stride_o_scale_bs=O_scale.stride(0), + stride_o_scale_d=O_scale.stride(1), + BLOCK_DMODEL=head_dim, + BLOCK_BATCH=triton.next_power_of_2(batch_size), + num_warps=1, + num_stages=1, + ) + + return O_fp8, O_scale.squeeze(-1) diff --git a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py index 12786c6619..4b45bdb5b7 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py +++ b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py @@ -9,6 +9,7 @@ def _trans_topk_index_to_mem_index( topk_index, topk_index_stride_b, topk_index_stride_k, + ragged_start_index, ragged_mem_index, topk_mem_index, topk_mem_index_stride_b, @@ -19,6 +20,9 @@ def _trans_topk_index_to_mem_index( offs_d = tl.arange(0, BLOCK_DMODEL) topk_index_ptrs = topk_index + cur_index * topk_index_stride_b + offs_d * topk_index_stride_k topk_indices = tl.load(topk_index_ptrs) + ragged_start = tl.load(ragged_start_index + cur_index) + topk_indices = tl.where(topk_indices != -1, topk_indices + ragged_start, -1) + tl.store(topk_index_ptrs, topk_indices) dest_mem_index = ragged_mem_index + topk_indices mem_index = tl.load(dest_mem_index, mask=topk_indices != -1, other=-1) @@ -26,8 +30,13 @@ def _trans_topk_index_to_mem_index( @torch.no_grad() -def trans_topk_index_to_mem_index(topk_index: torch.Tensor, ragged_mem_index: torch.Tensor): +def trans_topk_index_to_mem_index( + topk_index: torch.Tensor, + ragged_start_index: torch.Tensor, + ragged_mem_index: torch.Tensor, +): assert topk_index.shape[1] == 2048, f"Expected topk_index shape[1]=2048, got {topk_index.shape[1]}" + assert ragged_start_index.shape == (topk_index.shape[0],) grid = (topk_index.shape[0],) @@ -37,6 +46,7 @@ def trans_topk_index_to_mem_index(topk_index: torch.Tensor, ragged_mem_index: to topk_index=topk_index, topk_index_stride_b=topk_index.stride(0), topk_index_stride_k=topk_index.stride(1), + ragged_start_index=ragged_start_index, ragged_mem_index=ragged_mem_index, topk_mem_index=topk_mem_index, topk_mem_index_stride_b=topk_mem_index.stride(0), diff --git a/lightllm/models/gemma4/layer_infer/post_layer_infer.py b/lightllm/models/gemma4/layer_infer/post_layer_infer.py index b736a2d6c1..98158bd79c 100644 --- a/lightllm/models/gemma4/layer_infer/post_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/post_layer_infer.py @@ -5,18 +5,18 @@ class Gemma4PostLayerInfer(LlamaPostLayerInfer): """ Same final RMSNorm + tied lm_head path as Llama, with an extra tanh-based - logit softcap at the end: logits = softcap * tanh(logits / softcap). + transform before sampling: logits = softcap * tanh(logits / softcap). """ def __init__(self, network_config): super().__init__(network_config) self.final_logit_softcapping = float(network_config.get("final_logit_softcapping")) - def token_forward(self, input_embdings, infer_state, layer_weight): - logits = super().token_forward(input_embdings, infer_state, layer_weight) + def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor: if self.final_logit_softcapping is not None and self.final_logit_softcapping > 0: cap = self.final_logit_softcapping - logits = torch.tanh(logits / cap) * cap - if infer_state.prompt_logics is not None: - infer_state.prompt_logics = torch.tanh(infer_state.prompt_logics / cap) * cap + # The historical dense path first materializes FP32 logits and then + # applies the softcap. Preserve that numerical contract when the + # transform runs on a local BF16/FP16 vocabulary shard. + return torch.tanh(logits.float() / cap) * cap return logits diff --git a/lightllm/models/glm5_next/__init__.py b/lightllm/models/glm5_next/__init__.py new file mode 100644 index 0000000000..728282aa05 --- /dev/null +++ b/lightllm/models/glm5_next/__init__.py @@ -0,0 +1 @@ +"""GLM-5-Next model support.""" diff --git a/lightllm/models/glm5_next/infer_struct.py b/lightllm/models/glm5_next/infer_struct.py new file mode 100644 index 0000000000..74d7ae824c --- /dev/null +++ b/lightllm/models/glm5_next/infer_struct.py @@ -0,0 +1,5 @@ +from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo + + +class Glm5NextInferStateInfo(Deepseek2InferStateInfo): + pass diff --git a/lightllm/models/glm5_next/layer_infer/__init__.py b/lightllm/models/glm5_next/layer_infer/__init__.py new file mode 100644 index 0000000000..1235114fbb --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/__init__.py @@ -0,0 +1,3 @@ +from .transformer_layer_infer import Glm5NextTransformerLayerInfer + +__all__ = ["Glm5NextTransformerLayerInfer"] diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..4f6c6bfbd5 --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.distributed.communication_op import all_reduce +from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import ( + Deepseek3_2TransformerLayerInfer, + NsaInfer, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( + silu_and_mul_fwd, +) +from lightllm.models.glm5_next.triton_kernel.mhc import ( + hc_contract, + hc_expand, + hc_post, + hc_pre_norm, +) +from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.utils.envs_utils import get_env_start_args + + +class Glm5NextNsaInfer(NsaInfer): + """GLM indexer projection without rotary dimensions.""" + + def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): + q = layer_weight.wq_b_proj_.mm(q_lora).view( + -1, self.tp_index_n_heads, self.index_head_dim + ) + k = layer_weight.wk_proj_.mm(hidden_states.to(q_lora.dtype)) + k = layer_weight.k_norm_(k, eps=self.eps) + return self._rotate_activation(q), self._rotate_activation(k) + + def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): + # GLM stores weights_proj in FP32, so its activation must match before + # delegating to the shared NSA scoring and top-k implementation. + return super()._get_indices( + hidden_states.float(), q_lora, infer_state, att_state, layer_weight + ) + +class Glm5NextTransformerLayerInfer(Deepseek3_2TransformerLayerInfer): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + self.num_hidden_layers = network_config["num_hidden_layers"] + self.autotune_layer_num = network_config.get( + "autotune_layer_num", self.num_hidden_layers + ) + self.is_mtp_layer = layer_num >= self.num_hidden_layers + self.is_linear_attention_layer = ( + not self.is_mtp_layer + and network_config["layer_types"][layer_num] == "linear_attention" + ) + self.mhc_streams = network_config.get("hc_mult", 4) + self.hc_eps = network_config.get("hc_eps", 1e-6) + self.hc_sinkhorn_iters = network_config.get("hc_sinkhorn_iters", 20) + self.swiglu_limit = network_config["swiglu_limit"] + linear = network_config["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.tp_linear_num_heads = self.linear_num_heads // self.tp_world_size_ + self.tp_linear_projection_size = self.tp_linear_num_heads * self.linear_head_dim + if not self.is_linear_attention_layer: + self.indexer = Glm5NextNsaInfer( + layer_idx=self.layer_num_, + network_config=self.network_config_, + tp_world_size=self.tp_world_size_, + ) + # GLM's recurrent EAGLE drafter processes one row per logical + # request. Only target-model decode uses the widened verification + # layout of mtp_step + 1 rows. + self.indexer.decode_mtp_step = ( + 0 if self.is_mtp_layer else get_env_start_args().mtp_step + ) + + def _ffn_tp(self, input, infer_state, layer_weight): + """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" + + input = input.view(-1, self.embed_dim_) + up_gate_out = layer_weight.gate_up_proj.mm(input) + ffn1_out = self.alloc_tensor( + (input.size(0), up_gate_out.size(1) // 2), input.dtype + ) + silu_and_mul_fwd( + up_gate_out, + ffn1_out, + limit=self.swiglu_limit, + alpha=1.0, + clamp_up_add_one=False, + ) + return layer_weight.down_proj.mm(ffn1_out) + + def _shared_ffn_tp(self, input, infer_state, layer_weight): + return self._ffn_tp(input, infer_state, layer_weight) + + def _get_qkv(self, input, infer_state, layer_weight): + if self.is_linear_attention_layer: + raise AssertionError("KDA projections use _kda_projections") + + input = input.view(-1, self.embed_dim_) + if not infer_state.use_replicated_attention_ep: + input = self._tpsp_allgather(input=input, infer_state=infer_state) + if infer_state.need_dp_prefill_balance: + input = infer_state._all_to_all_unbalance_get(data=input) + + q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split( + [self.q_lora_rank, self.kv_lora_rank], dim=-1 + ) + q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_) + infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q} + q = layer_weight.q_b_proj_.mm(q).view( + -1, self.tp_q_head_num_, self.qk_nope_head_dim + ) + cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank) + rmsnorm_forward( + cache_kv[:, :, : self.kv_lora_rank], + weight=layer_weight.kv_a_layernorm_.weight, + eps=self.eps_, + out=cache_kv[:, :, : self.kv_lora_rank], + ) + return q, cache_kv + + def _get_o(self, input, infer_state, layer_weight): + if not infer_state.use_replicated_attention_ep: + return super()._get_o(input, infer_state, layer_weight) + + if infer_state.need_dp_prefill_balance: + input = infer_state._all_to_all_balance_get(data=input) + if input.shape[2] == self.kv_lora_rank: + input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) + output = layer_weight.o_weight_.mm( + input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim) + ) + all_reduce(output, group=infer_state.dist_group) + return output + + def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): + if self.is_linear_attention_layer: + raise AssertionError("KDA uses its dedicated backend") + q_nope = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1) + topk_mem_indices, _ = self.indexer._get_indices( + hidden_states=infer_state.get_topk_indices_params["hidden_states"], + q_lora=infer_state.get_topk_indices_params["q_lora"], + infer_state=infer_state, + att_state=infer_state.decode_att_state, + layer_weight=layer_weight, + ) + del infer_state.get_topk_indices_params + q_rope = q_nope[..., :0] + return infer_state.decode_att_state.decode_att( + q=(q_nope, q_rope), + k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + v=None, + att_control=AttControl( + nsa_decode=True, + nsa_decode_dict={ + "layer_index": self.layer_num_, + "topk_mem_indices": topk_mem_indices, + "softmax_scale": self.softmax_scale, + "kv_lora_rank": self.kv_lora_rank, + "qk_rope_head_dim": 0, + }, + ), + ) + + def _kda_projections(self, input, infer_state, layer_weight): + # KDA shards heads across TP ranks, so every rank still needs every + # token before updating its recurrent head state. In TP/SP mode the + # layer input is sequence-sharded; gather it here just like the MLA + # projection path and reduce-scatter the output in _kda_post. + input = input.view(-1, self.embed_dim_) + if not infer_state.use_replicated_attention_ep: + input = self._tpsp_allgather(input=input, infer_state=infer_state) + projected = layer_weight.linear_qkvb_proj.mm(input) + qkv_size = 3 * self.tp_linear_projection_size + mixed_qkv, raw_beta = projected.split( + [qkv_size, self.tp_linear_num_heads], dim=-1 + ) + fg_a = layer_weight.linear_fg_a_proj.mm(input) + f_a, g_a = fg_a.split(self.linear_head_dim, dim=-1) + raw_gate, norm_gate = layer_weight.project_kda_fg_b(f_a, g_a) + return mixed_qkv, raw_gate, raw_beta, norm_gate + + def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): + tokens = norm_gate.shape[0] + core_output = core_output.view(-1, self.linear_head_dim) + norm_gate = norm_gate.contiguous().view(-1, self.linear_head_dim) + output = layer_weight.linear_o_norm( + input=core_output, + eps=self.eps_, + alloc_func=self.alloc_tensor, + ) + output.mul_(norm_gate.float().sigmoid().to(output.dtype)) + output = layer_weight.linear_o_proj.mm(output.view(tokens, -1)) + if infer_state.use_replicated_attention_ep: + all_reduce(output, group=infer_state.dist_group) + return output + return self._tpsp_reduce(input=output, infer_state=infer_state) + + def context_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().context_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( + input_embeddings, infer_state, layer_weight + ) + core_output = infer_state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_prefill=True, + linear_att_prefill_dict={ + "mixed_qkv": mixed_qkv, + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, norm_gate, infer_state, layer_weight) + + def token_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().token_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( + input_embeddings, infer_state, layer_weight + ) + core_output = infer_state.decode_att_state1.decode_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_decode=True, + linear_att_decode_dict={ + "mixed_qkv": mixed_qkv, + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, norm_gate, infer_state, layer_weight) + + def _hc_pre(self, streams, layer_weight, prefix, norm_weight): + return hc_pre_norm( + x=streams, + fn=getattr(layer_weight, f"hc_{prefix}_fn").weight, + scale=getattr(layer_weight, f"hc_{prefix}_scale").weight, + base=getattr(layer_weight, f"hc_{prefix}_base").weight, + norm_weight=norm_weight.weight, + streams=self.mhc_streams, + rms_eps=self.eps_, + norm_eps=self.eps_, + hc_eps=self.hc_eps, + sinkhorn_iters=self.hc_sinkhorn_iters, + ) + + def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): + streams = input_embeddings + if self.layer_num_ == 0: + streams = hc_expand(streams.view(-1, self.embed_dim_), self.mhc_streams) + + layer_input, residual_mix, post_mix = self._hc_pre( + streams, layer_weight, "attn", layer_weight.att_norm_weight_ + ) + if prefill: + layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) + else: + layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) + streams = hc_post( + layer_output, streams, residual_mix, post_mix, self.mhc_streams + ) + + layer_input, residual_mix, post_mix = self._hc_pre( + streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_ + ) + if infer_state.use_replicated_attention_ep: + if self.is_moe: + local_input = self._tpsp_sp_split( + input=layer_input, infer_state=infer_state + ) + local_output = self._ffn(local_input, infer_state, layer_weight) + layer_output = self._tpsp_allgather( + input=local_output, infer_state=infer_state + ) + else: + layer_output = self._ffn_tp(layer_input, infer_state, layer_weight) + all_reduce(layer_output, group=infer_state.dist_group) + else: + layer_output = self._ffn(layer_input, infer_state, layer_weight) + streams = hc_post( + layer_output, streams, residual_mix, post_mix, self.mhc_streams + ) + is_autotune_last_layer = ( + Autotuner.is_autotune_warmup() + and self.layer_num_ == self.autotune_layer_num - 1 + ) + if self.layer_num_ == self.num_hidden_layers - 1 or is_autotune_last_layer: + return hc_contract(streams, self.mhc_streams) + return streams + + def context_forward(self, input_embeddings, infer_state, layer_weight): + if self.is_mtp_layer: + return super().context_forward( + input_embeddings, infer_state, layer_weight + ) + return self._forward_mhc( + input_embeddings, infer_state, layer_weight, prefill=True + ) + + def token_forward(self, input_embeddings, infer_state, layer_weight): + if self.is_mtp_layer: + return super().token_forward( + input_embeddings, infer_state, layer_weight + ) + return self._forward_mhc( + input_embeddings, infer_state, layer_weight, prefill=False + ) diff --git a/lightllm/models/glm5_next/layer_weights/__init__.py b/lightllm/models/glm5_next/layer_weights/__init__.py new file mode 100644 index 0000000000..4136e1883d --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/__init__.py @@ -0,0 +1,4 @@ +from .pre_and_post_layer_weight import Glm5NextPreAndPostLayerWeight +from .transformer_layer_weight import Glm5NextTransformerLayerWeight + +__all__ = ["Glm5NextPreAndPostLayerWeight", "Glm5NextTransformerLayerWeight"] diff --git a/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..b336f00632 --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 + +from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import ( + LlamaPreAndPostLayerWeight, +) + + +def add_language_model_aliases(weights: dict) -> None: + """Expose GLM's nested language-model keys under LightLLM names.""" + + prefix = "model.language_model." + for name in list(weights): + if name.startswith(prefix): + weights.setdefault("model." + name[len(prefix) :], weights[name]) + + +class Glm5NextPreAndPostLayerWeight(LlamaPreAndPostLayerWeight): + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..616cb16c0f --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.layer_weights.transformer_layer_weight import ( + TransformerLayerWeight, +) +from lightllm.common.basemodel.layer_weights.meta_weights import ( + COLMMWeight, + LayerNormWeight, + ParameterWeight, + RMSNormWeight, + ROWMMWeight, + TpParameterWeight, +) +from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import ( + Deepseek2TransformerLayerWeight, +) +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import ( + Deepseek3_2TransformerLayerWeight, +) +from .pre_and_post_layer_weight import add_language_model_aliases + + +class Glm5NextTransformerLayerWeight(Deepseek3_2TransformerLayerWeight): + def _parse_config(self): + super()._parse_config() + # The released sparse MLA keeps kv_b_proj in BF16 even though the + # surrounding projections are native FP8. Its compressed-context + # shortcut assumes a quantized kv_b matrix, so GLM uses the BMM split. + self.enable_cc_method = False + self.is_mtp_layer = self.layer_num_ >= self.network_config_["num_hidden_layers"] + self.is_linear_attention_layer = ( + not self.is_mtp_layer + and self.network_config_["layer_types"][self.layer_num_] + == "linear_attention" + ) + linear = self.network_config_["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.linear_projection_size = self.linear_num_heads * self.linear_head_dim + self.linear_conv_kernel_size = linear["short_conv_kernel_size"] + self.mhc_streams = self.network_config_.get("hc_mult", 4) + + def _init_weight(self): + if self.is_linear_attention_layer: + self._init_kda() + else: + Deepseek2TransformerLayerWeight._init_qkvo(self) + self._init_indexer_weight() + + if self.is_moe: + self._init_moe() + else: + self._init_ffn() + self._init_glm_norms() + if not self.is_mtp_layer: + self._init_mhc() + + def _init_kda(self): + prefix = f"model.layers.{self.layer_num_}.self_attn" + projection = self.linear_projection_size + head_count = self.linear_num_heads + head_dim = self.linear_head_dim + + self.linear_qkvb_proj = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[projection, projection, projection, head_count], + weight_names=[ + f"{prefix}.q_proj.weight", + f"{prefix}.k_proj.weight", + f"{prefix}.v_proj.weight", + f"{prefix}.b_proj.weight", + ], + data_type=self.data_type_, + quant_method=None, + ) + # f_a and g_a are replicated across TP ranks. + self.linear_fg_a_proj = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[head_dim, head_dim], + weight_names=[f"{prefix}.f_a_proj.weight", f"{prefix}.g_a_proj.weight"], + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.linear_fg_b_proj = ROWMMWeight( + in_dim=head_dim, + out_dims=[projection, projection], + weight_names=[f"{prefix}.f_b_proj.weight", f"{prefix}.g_b_proj.weight"], + data_type=self.data_type_, + quant_method=None, + ) + self.linear_qkv_conv1d = ROWMMWeight( + in_dim=self.linear_conv_kernel_size, + out_dims=[projection, projection, projection], + weight_names=[ + f"{prefix}.q_conv1d.weight", + f"{prefix}.k_conv1d.weight", + f"{prefix}.v_conv1d.weight", + ], + data_type=self.data_type_, + quant_method=None, + ) + self.linear_A_log = TpParameterWeight( + weight_name=f"{prefix}.A_log", + data_type=torch.float32, + weight_shape=(head_count,), + ) + self.linear_dt_bias = TpParameterWeight( + weight_name=f"{prefix}.dt_bias", + data_type=torch.float32, + weight_shape=(projection,), + ) + self.linear_o_norm = RMSNormWeight( + dim=head_dim, + weight_name=f"{prefix}.o_norm.weight", + data_type=self.data_type_, + ) + self.linear_o_proj = COLMMWeight( + in_dim=projection, + out_dims=[self.n_embed], + weight_names=f"{prefix}.o_proj.weight", + data_type=self.data_type_, + quant_method=None, + ) + + def _init_indexer_weight(self): + """Initialize GLM's NoPE, K-pool indexer parameters. + + The head-weight projection intentionally accumulates in fp32. Both + reference engines do this because bf16 head weights can change close + K-pool rankings on difficult long-context prompts. + """ + + prefix = f"model.layers.{self.layer_num_}.self_attn.indexer" + self.wq_b_proj_ = ROWMMWeight( + in_dim=self.q_lora_rank, + out_dims=[self.index_n_heads * self.index_head_dim], + weight_names=f"{prefix}.wq_b.weight", + data_type=self.data_type_, + quant_method=None, + ) + self.wk_proj_ = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.wk.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.k_norm_ = LayerNormWeight( + dim=self.index_head_dim, + weight_name=f"{prefix}.k_norm.weight", + data_type=self.data_type_, + bias_name=f"{prefix}.k_norm.bias", + ) + self.weights_proj_ = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_n_heads], + weight_names=f"{prefix}.weights_proj.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_gate = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.index_kpool_compress_gate", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_ape = ParameterWeight( + weight_name=f"{prefix}.index_kpool_compress_ape", + data_type=torch.float32, + weight_shape=(self.network_config_["index_kpool"], self.index_head_dim), + ) + + def _init_glm_norms(self): + prefix = f"model.layers.{self.layer_num_}" + self.att_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.input_layernorm.weight", + data_type=self.data_type_, + ) + self.ffn_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.post_attention_layernorm.weight", + data_type=self.data_type_, + ) + if not self.is_linear_attention_layer: + self.kv_a_layernorm_ = RMSNormWeight( + dim=self.kv_lora_rank, + weight_name=f"{prefix}.self_attn.kv_a_layernorm.weight", + data_type=self.data_type_, + ) + self.q_a_layernorm_ = RMSNormWeight( + dim=self.q_lora_rank, + weight_name=f"{prefix}.self_attn.q_a_layernorm.weight", + data_type=self.data_type_, + ) + + def _init_mhc(self): + prefix = f"model.layers.{self.layer_num_}" + streams = self.mhc_streams + mix_size = (2 + streams) * streams + flattened_hidden = streams * self.n_embed + self.hc_attn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_attn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_attn_base = ParameterWeight( + weight_name=f"{prefix}.hc_attn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_attn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_attn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + self.hc_ffn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_ffn_base = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_ffn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + + def get_merged_kda_conv_weight(self): + return self.linear_qkv_conv1d.mm_param.weight + + def project_kda_fg_b(self, f_a: torch.Tensor, g_a: torch.Tensor): + method = self.linear_fg_b_proj.quant_method + f = method.apply(f_a, self.linear_fg_b_proj.mm_param_list[0]) + g = method.apply(g_a, self.linear_fg_b_proj.mm_param_list[1]) + return f, g + + def _preprocess_kda_weights(self, weights): + prefix = f"model.layers.{self.layer_num_}.self_attn" + for projection in ("q", "k", "v"): + name = f"{prefix}.{projection}_conv1d.weight" + if name in weights and weights[name].ndim == 3: + weights[name] = weights[name].squeeze(1) + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + + # GLM checkpoints nest the shared expert under + # ``mlp.shared_experts``. This class deliberately bypasses + # Deepseek2TransformerLayerWeight.load_hf_weights below, so perform + # the fused-shared remap here before the generic loader consumes the + # expert tensors. + if self.num_fused_shared_experts > 0: + self._rename_shared_experts( + weights, + self.experts.quant_method.weight_scale_suffix, + ) + + if self.is_linear_attention_layer: + self._preprocess_kda_weights(weights) + return TransformerLayerWeight.load_hf_weights(self, weights) + + kv_b_name = f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight" + if kv_b_name in weights: + k_b_proj, v_b_proj = self._split_kv_b_proj(weights[kv_b_name]) + weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = k_b_proj + weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = v_b_proj + + return TransformerLayerWeight.load_hf_weights(self, weights) diff --git a/lightllm/models/glm5_next/mem_manager.py b/lightllm/models/glm5_next/mem_manager.py new file mode 100644 index 0000000000..f813701d99 --- /dev/null +++ b/lightllm/models/glm5_next/mem_manager.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from lightllm.common.kv_cache_mem_manager.operator import LinearAttMemOperator +from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( + Qwen3NextMemManager, +) + + +class Glm5NextMemOperator(LinearAttMemOperator): + def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): + layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import ( + destindex_copy_kv, + ) + + output = self.mem_manager.kv_buffer[layer_index][ + :, :, : self.mem_manager.mla_head_dim + ] + destindex_copy_kv(kv, mem_index, output) + + +class Glm5NextMemManager(Qwen3NextMemManager): + """Packed sparse-MLA KV, DSA index keys, and KDA recurrent states.""" + + operator_class = Glm5NextMemOperator + indexer_padding_bytes = 144 + indexer_payload_bytes = 132 + + def __init__(self, *args, **kwargs): + self.mla_head_dim = kwargs.get("head_dim", args[3] if len(args) > 3 else None) + super().__init__(*args, **kwargs) + + def get_cell_size(self): + bytes_per_token = self.mla_head_dim * self.dtype.itemsize + self.indexer_padding_bytes + return bytes_per_token * self.layer_num + + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + assert head_num == 1 and dtype in (torch.bfloat16, torch.float16) + padding_elements = self.indexer_padding_bytes // dtype.itemsize + self.kv_buffer = torch.empty( + (layer_num, size + 1, head_num, head_dim + padding_elements), + dtype=dtype, + device="cuda", + ) + self._init_linear_att_buffers() + + def get_att_input_params(self, layer_index): + packed_index = self.linear_config.get_full_att_kv_layer_index(layer_index) + return self.kv_buffer[packed_index][:, :, : self.mla_head_dim] + + def get_indexer_k_buffer(self, layer_index): + packed_index = self.linear_config.get_full_att_kv_layer_index(layer_index) + return self.kv_buffer[packed_index].view(torch.uint8)[:, :, -self.indexer_payload_bytes :] diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py new file mode 100644 index 0000000000..d56e8f4b53 --- /dev/null +++ b/lightllm/models/glm5_next/model.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os + +import torch + +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector +from lightllm.common.basemodel.attention.linear import KDALinearAttBackend +from lightllm.common.build_utils import repair_config +from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.req_manager import ReqManagerForMamba +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel +from lightllm.models.glm5_next.infer_struct import Glm5NextInferStateInfo +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import ( + Glm5NextTransformerLayerInfer, +) +from lightllm.models.glm5_next.layer_weights.pre_and_post_layer_weight import ( + Glm5NextPreAndPostLayerWeight, +) +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import ( + Glm5NextTransformerLayerWeight, +) +from lightllm.models.glm5_next.mem_manager import Glm5NextMemManager +from lightllm.models.registry import ModelRegistry +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args + + +class Glm5NextPostNormHiddenCollector(FinalHiddenCollector): + """Expose the post-final-norm hidden state expected by GLM NextN. + + The generic EAGLE collector returns the decoder output before the model's + final RMSNorm. GLM's NextN block was trained against the normalized target + hidden and also recycles its own normalized hidden between recurrent draft + steps (matching the vLLM and SGLang implementations). + """ + + def __init__(self, norm_weight, eps: float): + super().__init__() + self.norm_weight = norm_weight + self.eps = eps + self.draft_token_ids = None + self.draft_token_probs = None + + def new_instance(self): + return Glm5NextPostNormHiddenCollector(self.norm_weight, self.eps) + + def add_final_hidden(self, final_hidden: torch.Tensor) -> None: + self.final_hidden = self.norm_weight(input=final_hidden, eps=self.eps) + + def add_mtp_outputs( + self, + draft_token_ids: torch.Tensor | None, + confidence_logits: torch.Tensor | None, + draft_token_probs: torch.Tensor | None = None, + ) -> None: + assert confidence_logits is None + self.draft_token_ids = draft_token_ids + self.draft_token_probs = draft_token_probs + + def finish_output(self, infer_state): + output = super().finish_output(infer_state) + output.draft_token_ids = self.draft_token_ids + output.draft_token_probs = self.draft_token_probs + self.draft_token_ids = None + self.draft_token_probs = None + return output + + +@ModelRegistry(["glm5_next", "glm5_next_text"]) +class Glm5NextTpPartModel(Deepseek3_2TpPartModel): + # Keep attention/mHC rows replicated so their row-parallel projections can + # use the fast custom all-reduce. Only shard rows around EP MoE, then + # all-gather once. The recurrent draft model opts out below. + replicated_attention_ep = True + pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight + transformer_weight_class = Glm5NextTransformerLayerWeight + transformer_layer_infer_class = Glm5NextTransformerLayerInfer + infer_state_class = Glm5NextInferStateInfo + + def _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json"), "r") as config_file: + outer_config = json.load(config_file) + self.outer_config = outer_config + self.config = dict(outer_config.get("text_config", outer_config)) + self.vision_config = outer_config.get("vision_config") + if "quantization_config" in outer_config: + self.config["quantization_config"] = outer_config["quantization_config"] + # GLM-5.3's index cache follows the official ue8m0-scale path; + # the released HF config does not spell this implementation detail + # out in quantization_config. + self.config["quantization_config"].setdefault("scale_fmt", "ue8m0") + # The checkpoint uses the standard clamped SwiGLU, not GPT-OSS's + # clamped (up + 1) variant supported by the shared kernel. + self.config["swiglu_clamp_up_add_one"] = False + # The generic autotune warmup executes only a representative prefix. + # Let the last layer in that prefix contract mHC's residual streams + # before the shared LM head consumes its output. + self.config["autotune_layer_num"] = 4 + repair_config(self.config, same_names=["num_attention_heads", "n_head"]) + repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + if self.finetune_config: + self.config["vocab_size"] = self.finetune_config.vocab_size + + def autotune_layers(self): + return 4 + + def _init_hidden_collector(self): + collector = self.mtp_manager.create_hidden_collector(model=self) + if isinstance(collector, FinalHiddenCollector): + collector = Glm5NextPostNormHiddenCollector( + norm_weight=self.pre_post_weight.final_norm_weight_, + eps=self.config["rms_norm_eps"], + ) + self.hidden_collector_prototype = collector + + def _make_linear_config(self): + linear = self.config["linear_attn_config"] + start_args: StartArgs = get_env_start_args() + state_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} + return LinearAttCacheConfig( + tp_world_size=self.tp_world_size_, + full_att_all_num_kv_heads=1, + full_att_dtype=self.data_type, + full_att_num_kv_heads=1, + full_att_head_dim=self.config["kv_lora_rank"], + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + num_linear_k_heads=linear["num_heads"] // self.tp_world_size_, + num_linear_v_heads=linear["num_heads"] // self.tp_world_size_, + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + linear_layer_num=len(linear["kda_layers"]), + conv_state_dtype=self.data_type, + ssm_state_dtype=state_dtypes[start_args.linear_att_ssm_data_type], + full_attention_interval=4, + all_layer_num=self.config["n_layer"], + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), + ) + + def _init_req_manager(self): + max_sequence_length = 0 + if self.batch_max_tokens is not None: + max_sequence_length = max(max_sequence_length, self.batch_max_tokens) + if self.max_seq_length is not None: + max_sequence_length = max(max_sequence_length, self.max_seq_length) + self.linear_config = self._make_linear_config() + self.req_manager = ReqManagerForMamba( + self.max_req_num, + max_sequence_length, + None, + linear_config=self.linear_config, + ) + + def _init_mem_manager(self): + self.linear_config = getattr(self, "linear_config", self._make_linear_config()) + self.mem_manager = Glm5NextMemManager( + size=self.max_total_token_num, + dtype=self.data_type, + num_kv_heads=1, + head_dim=self.config["kv_lora_rank"], + full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), + linear_config=self.linear_config, + mem_fraction=self.mem_fraction, + ) + + def _init_att_backend1(self): + if getattr(self, "is_mtp_draft_model", False): + self.prefill_att_backend1 = None + self.decode_att_backend1 = None + return + self.prefill_att_backend1 = KDALinearAttBackend(model=self) + self.decode_att_backend1 = self.prefill_att_backend1 + + def _init_custom(self): + # GLM-5 sparse MLA is entirely NoPE. Keep zero-width tables so the + # generic infer-state position setup remains valid without allocating + # a million-token rotary cache. + max_length = max( + self.config["max_position_embeddings"], self.max_seq_length or 0 + ) + self._cos_cached = torch.empty( + (max_length, 0), dtype=self.data_type, device="cuda" + ) + self._sin_cached = torch.empty_like(self._cos_cached) + dist_group_manager.new_deepep_group( + n_routed_experts=self.config["n_routed_experts"], + hidden_size=self.config["hidden_size"], + expert_quant_method_names=dist_group_manager.get_moe_quant_methods( + self.trans_layers_weight + ), + num_experts_per_tok=self.config["num_experts_per_tok"], + moe_intermediate_size=self.config["moe_intermediate_size"], + ) diff --git a/lightllm/models/glm5_next/triton_kernel/__init__.py b/lightllm/models/glm5_next/triton_kernel/__init__.py new file mode 100644 index 0000000000..7db51ba063 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/__init__.py @@ -0,0 +1,3 @@ +from .mhc import hc_contract, hc_expand, hc_post, hc_pre, hc_pre_norm + +__all__ = ["hc_contract", "hc_expand", "hc_post", "hc_pre", "hc_pre_norm"] diff --git a/lightllm/models/glm5_next/triton_kernel/mhc.py b/lightllm/models/glm5_next/triton_kernel/mhc.py new file mode 100644 index 0000000000..77869a51f0 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/mhc.py @@ -0,0 +1,604 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC operators used by GLM-5-Next. + +The public entry points use fused Triton kernels for the small, launch-bound +mixing operations. The explicit PyTorch implementations remain available as +correctness oracles: all mixing math is accumulated in fp32 and only the +collapsed layer input / expanded residual streams are cast back to the +activation dtype. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + + +@triton.jit +def _hc_prepare_kernel( + mixes, + scale, + base, + pre, + post, + residual_mix, + mix_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, +): + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + + pre_raw = tl.load(mixes + token * mix_stride_m + stream_offsets) + post_raw = tl.load(mixes + token * mix_stride_m + STREAMS + stream_offsets) + pre_values = tl.sigmoid( + pre_raw * tl.load(scale) + tl.load(base + stream_offsets) + ) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid( + post_raw * tl.load(scale + 1) + + tl.load(base + STREAMS + stream_offsets) + ) + + logits = tl.load( + mixes + token * mix_stride_m + 2 * STREAMS + matrix_offsets + ) + logits = logits * tl.load(scale + 2) + tl.load( + base + 2 * STREAMS + matrix_offsets + ) + logits = tl.reshape(logits, (STREAMS, STREAMS)) + logits = logits - tl.max(logits, axis=1)[:, None] + matrix = tl.exp(logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + + # The checkpoint definition starts with a column normalization, then + # alternates row and column normalizations for the remaining iterations. + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_prepare_prenorm_kernel( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_stride_s: tl.constexpr, + gemm_stride_m: tl.constexpr, + sqrsum_stride_s: tl.constexpr, + sqrsum_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + FLATTENED_HIDDEN: tl.constexpr, + RMS_EPS: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, + N_SPLITS: tl.constexpr, +): + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + pre_raw = tl.zeros((STREAMS,), dtype=tl.float32) + post_raw = tl.zeros((STREAMS,), dtype=tl.float32) + matrix_raw = tl.zeros((STREAMS * STREAMS,), dtype=tl.float32) + sqrsum = 0.0 + for split in tl.static_range(N_SPLITS): + partial_base = ( + gemm_partial + split * gemm_stride_s + token * gemm_stride_m + ) + pre_raw += tl.load(partial_base + stream_offsets) + post_raw += tl.load(partial_base + STREAMS + stream_offsets) + matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) + sqrsum += tl.load( + sqrsum_partial + + split * sqrsum_stride_s + + token * sqrsum_stride_m + ) + inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) + pre_raw *= inv_rms + post_raw *= inv_rms + matrix_raw *= inv_rms + + pre_values = tl.sigmoid( + pre_raw * tl.load(scale) + tl.load(base + stream_offsets) + ) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid( + post_raw * tl.load(scale + 1) + + tl.load(base + STREAMS + stream_offsets) + ) + + logits = ( + matrix_raw * tl.load(scale + 2) + + tl.load(base + 2 * STREAMS + matrix_offsets) + ) + logits = tl.reshape(logits, (STREAMS, STREAMS)) + logits = logits - tl.max(logits, axis=1)[:, None] + matrix = tl.exp(logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_pre_combine_kernel( + x, + pre, + output, + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + tl.store( + output + token * out_stride_m + hidden_offsets, + accumulator, + mask=hidden_mask, + ) + + +@triton.jit +def _hc_pre_combine_norm_kernel( + x, + pre, + norm_weight, + output, + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + NORM_EPS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_offsets = tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + + # hc_pre returns bf16 before the following RMSNorm in the checkpoint + # definition. Preserve that rounding point while keeping both operations + # in one kernel. + rounded = accumulator.to(tl.bfloat16).to(tl.float32) + variance = tl.sum(rounded * rounded, axis=0) / hidden + inv_rms = tl.rsqrt(variance + NORM_EPS) + weight = tl.load(norm_weight + hidden_offsets, mask=hidden_mask, other=0.0) + tl.store( + output + token * out_stride_m + hidden_offsets, + rounded * inv_rms * weight, + mask=hidden_mask, + ) + + +@triton.jit +def _hc_post_4stream_kernel( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden: tl.constexpr, + layer_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + mix_stride_m: tl.constexpr, + post_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + + # GLM-5 always uses four mHC streams. Compute all four outputs in one + # program so the layer output and residual streams are read only once. + # The previous output-stream grid reread each residual stream four times; + # that becomes bandwidth-bound for large prefills. + layer_value = tl.load( + layer_output + token * layer_stride_m + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + residual_base = residual + token * residual_stride_m + hidden_offsets + residual_0 = tl.load( + residual_base, mask=hidden_mask, other=0.0 + ).to(tl.float32) + residual_1 = tl.load( + residual_base + hidden, mask=hidden_mask, other=0.0 + ).to(tl.float32) + residual_2 = tl.load( + residual_base + 2 * hidden, mask=hidden_mask, other=0.0 + ).to(tl.float32) + residual_3 = tl.load( + residual_base + 3 * hidden, mask=hidden_mask, other=0.0 + ).to(tl.float32) + + post_base = post_mix + token * post_stride_m + mix_base = residual_mix + token * mix_stride_m + accumulator_0 = layer_value * tl.load(post_base) + accumulator_1 = layer_value * tl.load(post_base + 1) + accumulator_2 = layer_value * tl.load(post_base + 2) + accumulator_3 = layer_value * tl.load(post_base + 3) + + accumulator_0 += residual_0 * tl.load(mix_base) + accumulator_0 += residual_1 * tl.load(mix_base + 4) + accumulator_0 += residual_2 * tl.load(mix_base + 8) + accumulator_0 += residual_3 * tl.load(mix_base + 12) + accumulator_1 += residual_0 * tl.load(mix_base + 1) + accumulator_1 += residual_1 * tl.load(mix_base + 5) + accumulator_1 += residual_2 * tl.load(mix_base + 9) + accumulator_1 += residual_3 * tl.load(mix_base + 13) + accumulator_2 += residual_0 * tl.load(mix_base + 2) + accumulator_2 += residual_1 * tl.load(mix_base + 6) + accumulator_2 += residual_2 * tl.load(mix_base + 10) + accumulator_2 += residual_3 * tl.load(mix_base + 14) + accumulator_3 += residual_0 * tl.load(mix_base + 3) + accumulator_3 += residual_1 * tl.load(mix_base + 7) + accumulator_3 += residual_2 * tl.load(mix_base + 11) + accumulator_3 += residual_3 * tl.load(mix_base + 15) + + output_base = output + token * out_stride_m + hidden_offsets + tl.store(output_base, accumulator_0, mask=hidden_mask) + tl.store(output_base + hidden, accumulator_1, mask=hidden_mask) + tl.store(output_base + 2 * hidden, accumulator_2, mask=hidden_mask) + tl.store(output_base + 3 * hidden, accumulator_3, mask=hidden_mask) + + +def hc_expand(x: torch.Tensor, streams: int) -> torch.Tensor: + """Expand ``[tokens, hidden]`` into flattened residual streams.""" + + assert x.ndim == 2 + return x.unsqueeze(1).expand(-1, streams, -1).reshape(x.shape[0], -1) + + +def hc_contract(x: torch.Tensor, streams: int) -> torch.Tensor: + """Contract flattened residual streams by taking their mean.""" + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + return x.view(x.shape[0], streams, -1).mean(dim=1) + + +def hc_pre_reference( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + streams: int, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute mHC pre-mixes. + + Returns ``(layer_input, residual_mix, post_mix)``. ``x`` and + ``layer_input`` use the activation dtype; both mix tensors are fp32. + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + residual = x.view(tokens, streams, hidden) + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre_raw = mixes[:, :streams] + post_raw = mixes[:, streams : 2 * streams] + residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) + + pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps + post = post_multiplier * torch.sigmoid( + post_raw * scale[1] + base[streams : 2 * streams] + ) + residual_mix = ( + residual_raw * scale[2] + base[2 * streams :].view(streams, streams) + ).softmax(dim=-1) + residual_mix = residual_mix + hc_eps + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + for _ in range(sinkhorn_iters - 1): + residual_mix = residual_mix / ( + residual_mix.sum(dim=-1, keepdim=True) + hc_eps + ) + residual_mix = residual_mix / ( + residual_mix.sum(dim=-2, keepdim=True) + hc_eps + ) + + layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) + return layer_input, residual_mix, post + + +def hc_post_reference( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """Mix a sublayer output back into the flattened residual streams.""" + + tokens, hidden = layer_output.shape + residual_3d = residual.view(tokens, streams, hidden) + mixed_residual = ( + residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2) + ).sum(dim=1) + out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual + return out.to(layer_output.dtype).reshape(tokens, streams * hidden) + + +def hc_pre( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + streams: int, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute mHC pre-mixes with fused Sinkhorn and residual combining.""" + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert x.is_contiguous() and fn.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty( + (tokens, streams, streams), dtype=torch.float32, device=x.device + ) + _hc_prepare_kernel[(tokens,)]( + mixes, + scale, + base, + pre, + post, + residual_mix, + mixes.stride(0), + pre.stride(0), + residual_mix.stride(0), + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + num_warps=1, + ) + + layer_input = torch.empty( + (tokens, hidden), dtype=x.dtype, device=x.device + ) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_pre_combine_kernel[(tokens, triton.cdiv(hidden, block_h))]( + x, + pre, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post + + +def _compute_prenorm_splits( + tokens: int, flattened_hidden: int, device: torch.device +) -> int: + grid_size = triton.cdiv(tokens, 64) + k_blocks = triton.cdiv(flattened_hidden, 64) + sms = torch.cuda.get_device_properties(device).multi_processor_count + return max(1, min(sms // max(grid_size, 1), k_blocks // 4)) + + +def hc_pre_norm( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + streams: int, + rms_eps: float, + norm_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fuse mHC pre-mixing with its immediately following RMSNorm. + + DeepGEMM computes the fp32 projection and residual square sum in one + split-K kernel. Triton then reduces those partials, runs Sinkhorn, forms + the bf16 residual collapse, and applies RMSNorm. + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert x.dtype == torch.bfloat16 and fn.dtype == torch.float32 + assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + try: + import deep_gemm + + prenorm_gemm = deep_gemm.tf32_hc_prenorm_gemm + except (AttributeError, ImportError): + from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import ( + rmsnorm_forward, + ) + + layer_input, residual_mix, post = hc_pre( + x, + fn, + scale, + base, + streams, + rms_eps, + hc_eps, + sinkhorn_iters, + post_multiplier, + ) + layer_input = rmsnorm_forward( + layer_input, weight=norm_weight, eps=norm_eps + ) + return layer_input, residual_mix, post + + mix_size = (2 + streams) * streams + n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) + gemm_partial = torch.empty( + (n_splits, tokens, mix_size), dtype=torch.float32, device=x.device + ) + sqrsum_partial = torch.empty( + (n_splits, tokens), dtype=torch.float32, device=x.device + ) + prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty( + (tokens, streams, streams), dtype=torch.float32, device=x.device + ) + _hc_prepare_prenorm_kernel[(tokens,)]( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_partial.stride(0), + gemm_partial.stride(1), + sqrsum_partial.stride(0), + sqrsum_partial.stride(1), + pre.stride(0), + residual_mix.stride(0), + FLATTENED_HIDDEN=flattened_hidden, + RMS_EPS=rms_eps, + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + N_SPLITS=n_splits, + num_warps=1, + ) + + layer_input = torch.empty( + (tokens, hidden), dtype=x.dtype, device=x.device + ) + block_h = triton.next_power_of_2(hidden) + _hc_pre_combine_norm_kernel[(tokens,)]( + x, + pre, + norm_weight, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + NORM_EPS=norm_eps, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post + + +def hc_post( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """Mix a sublayer output into residual streams with one Triton launch.""" + + tokens, hidden = layer_output.shape + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert layer_output.is_contiguous() and residual.is_contiguous() + output = torch.empty( + (tokens, streams * hidden), + dtype=layer_output.dtype, + device=layer_output.device, + ) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_post_4stream_kernel[(tokens, triton.cdiv(hidden, block_h))]( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden=hidden, + layer_stride_m=layer_output.stride(0), + residual_stride_m=residual.stride(0), + mix_stride_m=residual_mix.stride(0), + post_stride_m=post_mix.stride(0), + out_stride_m=output.stride(0), + BLOCK_H=block_h, + num_warps=8, + ) + return output diff --git a/lightllm/models/glm5_next_mtp/__init__.py b/lightllm/models/glm5_next_mtp/__init__.py new file mode 100644 index 0000000000..d55bc8b398 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + +__all__ = ["Glm5NextMTPModel"] diff --git a/lightllm/models/glm5_next_mtp/layer_infer/__init__.py b/lightllm/models/glm5_next_mtp/layer_infer/__init__.py new file mode 100644 index 0000000000..0ec4b5a01b --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/__init__.py @@ -0,0 +1,5 @@ +from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import ( + Glm5NextMTPPreLayerInfer, +) + +__all__ = ["Glm5NextMTPPreLayerInfer"] diff --git a/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py new file mode 100644 index 0000000000..5708ce30ce --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Callable + +import torch + +from lightllm.distributed.communication_op import all_gather_into_tensor +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + + +def vocab_parallel_top1( + local_logits: torch.Tensor, + local_vocab_start_id: int, + tp_world_size: int, + dist_group, + alloc_func: Callable, +) -> torch.Tensor: + """Return global greedy token ids without materializing global logits.""" + + assert local_logits.ndim == 2 + local_max_values, local_max_indexes = torch.max(local_logits, dim=0) + local_token_ids = local_max_indexes + local_vocab_start_id + if tp_world_size == 1: + return local_token_ids + + local_winners = torch.stack( + [local_max_values.float(), local_token_ids.float()], + dim=-1, + ).contiguous() + token_num = local_winners.shape[0] + gathered_winners = alloc_func( + (tp_world_size * token_num, 2), + dtype=torch.float32, + device=local_logits.device, + ) + all_gather_into_tensor( + gathered_winners, + local_winners, + group=dist_group, + async_op=False, + ) + gathered_winners = gathered_winners.view(tp_world_size, token_num, 2) + winning_ranks = torch.argmax(gathered_winners[:, :, 0], dim=0) + token_rows = torch.arange(token_num, dtype=torch.long, device=local_logits.device) + return gathered_winners[winning_ranks, token_rows, 1].long() + + +def vocab_parallel_top1_and_prob( + local_logits: torch.Tensor, + local_vocab_start_id: int, + tp_world_size: int, + dist_group, + alloc_func: Callable, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return global Top-1 ids and exact softmax probabilities. + + Only three FP32 values per token and rank are gathered: the local maximum, + its global token id, and the local log-sum-exp. This replaces the normal + all-gather of every vocabulary logit while preserving the result of a + global FP32 argmax/softmax reduction. + """ + + assert local_logits.ndim == 2 + local_logits_fp32 = local_logits.float() + local_max_values, local_max_indexes = torch.max(local_logits_fp32, dim=0) + local_logsumexp = torch.logsumexp(local_logits_fp32, dim=0) + local_token_ids = local_max_indexes + local_vocab_start_id + + if tp_world_size == 1: + return local_token_ids, torch.exp(local_max_values - local_logsumexp) + + local_stats = torch.stack( + [local_max_values, local_token_ids.float(), local_logsumexp], + dim=-1, + ).contiguous() + token_num = local_stats.shape[0] + gathered_stats = alloc_func( + (tp_world_size * token_num, 3), + dtype=torch.float32, + device=local_logits.device, + ) + all_gather_into_tensor( + gathered_stats, + local_stats, + group=dist_group, + async_op=False, + ) + gathered_stats = gathered_stats.view(tp_world_size, token_num, 3) + winning_ranks = torch.argmax(gathered_stats[:, :, 0], dim=0) + token_rows = torch.arange(token_num, dtype=torch.long, device=local_logits.device) + winning_stats = gathered_stats[winning_ranks, token_rows] + global_logsumexp = torch.logsumexp(gathered_stats[:, :, 2], dim=0) + return winning_stats[:, 1].long(), torch.exp(winning_stats[:, 0] - global_logsumexp) + + +class Glm5NextMTPPostLayerInfer(LlamaPostLayerInfer): + """GLM NextN uses the common exact vocab-parallel greedy output path.""" 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 new file mode 100644 index 0000000000..6fc7826a75 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 + +from lightllm.models.deepseek_mtp.layer_infer.pre_layer_infer import ( + Deepseek3MTPPreLayerInfer, +) +from lightllm.models.glm5_next_mtp.triton_kernel.zero_position_embedding import ( + zero_position_embedding_, +) + + +class Glm5NextMTPPreLayerInfer(Deepseek3MTPPreLayerInfer): + """GLM NextN input fusion with the trained position-zero convention.""" + + def _mtp_context_forward(self, input_embdings, infer_state, layer_weight): + # GLM's fused_eh_norm reference zeros the token embedding at absolute + # position zero before applying enorm. The first target hidden remains + # intact; only the missing previous-token embedding is suppressed. + zero_position_embedding_(input_embdings, infer_state.position_ids) + return super()._mtp_context_forward(input_embdings, infer_state, layer_weight) diff --git a/lightllm/models/glm5_next_mtp/layer_weights/__init__.py b/lightllm/models/glm5_next_mtp/layer_weights/__init__.py new file mode 100644 index 0000000000..8507c67bd8 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_weights/__init__.py @@ -0,0 +1,5 @@ +from lightllm.models.glm5_next_mtp.layer_weights.pre_and_post_layer_weight import ( + Glm5NextMTPPreAndPostLayerWeight, +) + +__all__ = ["Glm5NextMTPPreAndPostLayerWeight"] diff --git a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..f697114829 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,49 @@ +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import ( + EmbeddingWeight, + LMHeadWeight, + RMSNormWeight, + ROWMMWeight, +) +from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.layer_weights.pre_and_post_layer_weight import ( + add_language_model_aliases, +) + + +class Glm5NextMTPPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config, quant_cfg: Quantcfg): + super().__init__(data_type, network_config) + layer_idx = network_config["num_hidden_layers"] + hidden_size = network_config["hidden_size"] + prefix = f"model.layers.{layer_idx}" + self.eh_proj_weight_ = ROWMMWeight( + in_dim=hidden_size * 2, + out_dims=[hidden_size], + weight_names=f"{prefix}.eh_proj.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.enorm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name=f"{prefix}.enorm.weight", + data_type=self.data_type_, + ) + self.hnorm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name=f"{prefix}.hnorm.weight", + data_type=self.data_type_, + ) + self.final_norm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name=f"{prefix}.shared_head.norm.weight", + data_type=self.data_type_, + ) + self.wte_weight_: EmbeddingWeight = None + self.lm_head_weight_: LMHeadWeight = None + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py new file mode 100644 index 0000000000..c609c5b597 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/model.py @@ -0,0 +1,92 @@ +from typing import List + +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import ( + Glm5NextMTPPreLayerInfer, +) +from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import ( + Glm5NextMTPPostLayerInfer, +) +from lightllm.models.draft_registry import DraftModelRegistry +from lightllm.models.glm5_next_mtp.layer_weights.pre_and_post_layer_weight import ( + Glm5NextMTPPreAndPostLayerWeight, +) +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import ( + Glm5NextTransformerLayerInfer, +) +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import ( + Glm5NextTransformerLayerWeight, +) +from lightllm.models.glm5_next.model import Glm5NextTpPartModel + + +@DraftModelRegistry( + model_type=("glm5_next", "glm5_next_text"), + spec_modes=("vanilla_with_att", "eagle_with_att"), +) +class Glm5NextMTPModel(Glm5NextTpPartModel): + """GLM-5.3's shared one-layer NextN draft model.""" + + is_mtp_draft_model = True + replicated_attention_ep = False + pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight + pre_layer_infer_class = Glm5NextMTPPreLayerInfer + post_layer_infer_class = Glm5NextMTPPostLayerInfer + transformer_weight_class = Glm5NextTransformerLayerWeight + transformer_layer_infer_class = Glm5NextTransformerLayerInfer + + def __init__(self, kvargs: dict): + self.main_model: TpPartBaseModel = kvargs.pop("main_model") + self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop( + "mtp_previous_draft_models" + ) + super().__init__(kvargs) + + def _init_custom(self): + self._cos_cached = self.main_model._cos_cached + self._sin_cached = self.main_model._sin_cached + + def _init_req_manager(self): + self.req_manager = self.main_model.req_manager + + def _init_mem_manager(self): + self.mem_manager = self.main_model.mem_manager + + def _init_weights(self, start_layer_index=None): + assert start_layer_index is None + mtp_layer = self.config["num_hidden_layers"] + self.pre_post_weight = self.pre_and_post_weight_class( + self.data_type, network_config=self.config, quant_cfg=self.quant_cfg + ) + self.trans_layers_weight = [ + self.transformer_weight_class( + mtp_layer, + self.data_type, + network_config=self.config, + quant_cfg=self.quant_cfg, + ) + ] + self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ + self.pre_post_weight.lm_head_weight_ = ( + self.main_model.pre_post_weight.lm_head_weight_ + ) + + def _init_infer_layer(self, start_layer_index=None): + assert start_layer_index is None + self.pre_infer = self.pre_layer_infer_class(network_config=self.config) + self.post_infer = self.post_layer_infer_class(network_config=self.config) + logical_layer = len(self.main_model.layers_infer) + sum( + len(model.layers_infer) for model in self.mtp_previous_draft_models + ) + self.layers_infer = [ + self.transformer_layer_infer_class( + logical_layer, network_config=self.config + ) + ] + + def _init_some_value(self): + super()._init_some_value() + self.layers_num = 1 + + def autotune_layers(self): + return 1 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..daeefbafca --- /dev/null +++ b/lightllm/models/glm5_next_mtp/triton_kernel/__init__.py @@ -0,0 +1,5 @@ +from lightllm.models.glm5_next_mtp.triton_kernel.zero_position_embedding import ( + zero_position_embedding_, +) + +__all__ = ["zero_position_embedding_"] diff --git a/lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py b/lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py new file mode 100644 index 0000000000..ea0e48dbe7 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +import triton +import triton.language as tl + + +@triton.jit +def _zero_position_embedding_kernel( + embeddings, + stride_embeddings_m, + stride_embeddings_n, + position_ids, + hidden_size: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + if tl.load(position_ids + row) != 0: + return + + offsets = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + tl.store( + embeddings + row * stride_embeddings_m + offsets * stride_embeddings_n, + 0.0, + mask=offsets < hidden_size, + ) + + +@torch.no_grad() +def zero_position_embedding_(embeddings: torch.Tensor, position_ids: torch.Tensor) -> None: + """Zero MTP token embeddings at absolute position zero in place.""" + + assert embeddings.is_cuda and position_ids.is_cuda + assert embeddings.ndim == 2 and position_ids.shape == embeddings.shape[:1] + block_n = 256 + grid = (embeddings.shape[0], triton.cdiv(embeddings.shape[1], block_n)) + _zero_position_embedding_kernel[grid]( + embeddings=embeddings, + stride_embeddings_m=embeddings.stride(0), + stride_embeddings_n=embeddings.stride(1), + position_ids=position_ids, + hidden_size=embeddings.shape[1], + BLOCK_N=block_n, + num_warps=4, + num_stages=1, + ) diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index bb6e4f3735..06698205d6 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -7,6 +7,9 @@ from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import PostLayerInferTpl +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + vocab_parallel_greedy, +) from lightllm.distributed.communication_op import all_gather @@ -21,6 +24,10 @@ def __init__(self, network_config): def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor: return layer_weight.final_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor) + def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor: + """Apply model-specific elementwise transforms before greedy reduction.""" + return logits + def _slice_get_last_input(self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo): embed_dim_ = input_embdings.shape[1] if infer_state.is_prefill and infer_state.is_token_healing: @@ -80,7 +87,7 @@ def _token_forward( if prompt_logics_hiddens is not None: prompt_token_num = prompt_logics_hiddens.shape[0] infer_state.prompt_logics = self._lm_head_and_gather( - prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state + prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state, force_full_logits=True ) return ans_logics @@ -91,6 +98,7 @@ def _lm_head_and_gather( token_num: int, layer_weight: LlamaPreAndPostLayerWeight, infer_state: LlamaInferStateInfo, + force_full_logits: bool = False, ) -> torch.Tensor: normed = self._norm(hidden, infer_state, layer_weight) normed = normed.permute(1, 0).view(-1, token_num) @@ -98,6 +106,19 @@ def _lm_head_and_gather( normed = None vocab_size = layer_weight.lm_head_weight_.vocab_size + if infer_state.use_vocab_parallel_greedy and not force_full_logits: + logic_batch = self._apply_logit_postprocessing(logic_batch) + logits, token_ids, logsumexp = vocab_parallel_greedy( + logic_batch, + vocab_size=vocab_size, + tp_world_size=self.tp_world_size_, + group=infer_state.dist_group, + alloc_func=self.alloc_tensor, + ) + infer_state.logits_token_ids = token_ids + infer_state.logits_logsumexp = logsumexp + return logits + if self.tp_world_size_ == 1: gather_data = logic_batch else: @@ -114,7 +135,7 @@ def _lm_head_and_gather( ans_logics = self.alloc_tensor((token_num, vocab_size), dtype=torch.float32) ans_logics[:, :] = gather_data.permute(1, 0) gather_data = None - return ans_logics + return self._apply_logit_postprocessing(ans_logics) def token_forward( self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight diff --git a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py index 5a74cd988e..eb4481fd6b 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -181,7 +181,11 @@ def token_forward( logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) block_logits = logits.reshape(num_reqs, self.block_size_, -1) - sampled_tokens = torch.argmax(block_logits, dim=-1) + if infer_state.logits_token_ids is None: + sampled_tokens = torch.argmax(block_logits, dim=-1) + else: + assert block_logits.shape[-1] == 1 + sampled_tokens = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_) confidence_logits = self.predict_confidence_logits( block_hidden, anchor_token_ids=anchor_token_ids, diff --git a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py index 7311c4d141..e317f08f2d 100644 --- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py @@ -7,7 +7,7 @@ from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - use_sm100_mega_moe, + use_mega_moe, ) from lightllm.utils.dist_utils import get_global_world_size from lightllm.utils.envs_utils import get_env_start_args @@ -138,7 +138,7 @@ def overlap_tpsp_token_forward( infer_state1: LlamaInferStateInfo, layer_weight: Qwen3MOETransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_token_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -250,7 +250,7 @@ def overlap_tpsp_context_forward( infer_state1: LlamaInferStateInfo, layer_weight: Qwen3MOETransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_context_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -318,8 +318,8 @@ def overlap_tpsp_context_forward( # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( _0_num_recv_tokens_per_expert_list, - _0_handle.num_unaligned_recv_tokens_per_expert, - _0_handle.recv_src_metadata, + getattr(_0_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_0_handle, "recv_src_metadata", None), _0_recv_x, _0_recv_topk_idx, _0_recv_topk_weight, @@ -350,8 +350,8 @@ def overlap_tpsp_context_forward( # 1 moe calc _1_moe_out = layer_weight.experts.prefilled_group_gemm( _1_num_recv_tokens_per_expert_list, - _1_handle.num_unaligned_recv_tokens_per_expert, - _1_handle.recv_src_metadata, + getattr(_1_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_1_handle, "recv_src_metadata", None), _1_recv_x, _1_recv_topk_idx, _1_recv_topk_weight, diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 72499b5b1f..c2322adc6f 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -233,6 +233,16 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--running_max_req_size", type=int, default=256, help="the max size for forward requests in the same time" ) + parser.add_argument( + "--per_dp_max_req_size", + type=int, + default=None, + help=( + "Optional request-state capacity allocated by each local DP replica. " + "Defaults to running_max_req_size for backward compatibility; lowering it can " + "substantially reduce hybrid linear-attention state memory in balanced DP workloads." + ), + ) parser.add_argument("--nnodes", type=int, default=1, help="the number of nodes") parser.add_argument("--node_rank", type=int, default=0, help="the rank of the current node") parser.add_argument( @@ -424,11 +434,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--llm_prefill_att_backend", type=str, nargs="+", - choices=["auto", "triton", "fa3", "flashinfer", "flashqla"], + choices=["auto", "triton", "fa3", "flashinfer", "flashqla", "tilelang"], default=["auto"], help="""prefill attention kernel used in llm. auto: automatically select best backend based on GPU and available packages (priority: fa3 > flashinfer > triton) + for NSA/DSA models, tilelang selects SGLang's sparse prefill kernel for hybrid linear-attention models, the second value selects the linear-attention backend (priority: flashqla > triton); when omitted, it defaults to auto""", ) @@ -601,6 +612,36 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=8192, help="max handle token num for prefill cudagraph", ) + parser.add_argument( + "--prefill_cudagraph_token_nums", + nargs="+", + type=int, + default=None, + help=( + "Optional exact prefill token counts to capture. Requests with other token counts " + "run eagerly instead of being padded to one of these graphs." + ), + ) + parser.add_argument( + "--prefill_cudagraph_batch_sizes", + nargs="+", + type=int, + default=None, + help=( + "Batch size paired with each --prefill_cudagraph_token_nums entry. " + "Exact-layout graphs require uniform uncached sequences, so each token count " + "must be divisible by its paired batch size." + ), + ) + parser.add_argument( + "--prefill_cudagraph_capture_attention", + action="store_true", + help=( + "Capture prefill attention inside the main CUDA Graph instead of running it " + "between graph segments. This experimental mode requires exact token-count and " + "batch-size layouts." + ), + ) parser.add_argument( "--graph_max_batch_size", @@ -792,6 +833,13 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=0.03, help="""The interval of the schedule time, default is 30ms.""", ) + parser.add_argument( + "--prefill_coalesce_interval", + type=float, + default=0.0, + help="""Maximum time in seconds to collect a burst of waiting requests before scheduling + a prefill batch. Disabled by default. A full runnable batch is scheduled immediately.""", + ) parser.add_argument( "--afs_image_embed_dir", type=str, diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 4152bb107b..0f89969103 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -1,8 +1,9 @@ import multiprocessing as mp import os -import uuid -import subprocess import math +import subprocess +import sys +import uuid from lightllm.utils.start_utils import process_manager from .metrics.manager import start_metric_manager from .embed_cache.manager import start_cache_manager @@ -36,10 +37,10 @@ def _set_envs_and_config(args: StartArgs): def _launch_subprocesses(args: StartArgs): _set_envs_and_config(args) - if args.mtp_mode is not None: + if args.mtp_mode is not None and args.mtp_dynamic_verify: assert ( not args.disable_cudagraph or args.run_mode == "prefill" - ), "--disable_cudagraph is only supported on Prefill nodes when --mtp_mode is enabled" + ), "--disable_cudagraph is only supported on Prefill nodes when --mtp_dynamic_verify is enabled" auto_set_max_req_total_len(args) auto_set_fused_shared_experts(args) @@ -416,12 +417,19 @@ def _hypercorn_config_args(args: StartArgs): return ["--keep-alive", "10"] +def _hypercorn_entrypoint(): + # Launch from the active Python environment instead of relying on a + # console-script wrapper being present on PATH (for example, packages + # installed with pip --target do not install that wrapper). + return [sys.executable, "-m", "hypercorn"] + + def normal_or_p_d_start(args: StartArgs): process_manager = _launch_subprocesses(args) # 启动 Hypercorn command = [ - "hypercorn", + *_hypercorn_entrypoint(), *_hypercorn_config_args(args), "--workers", f"{args.httpserver_workers}", @@ -484,7 +492,7 @@ def pd_master_start(args: StartArgs): ) command = [ - "hypercorn", + *_hypercorn_entrypoint(), *_hypercorn_config_args(args), "--workers", "1", @@ -572,7 +580,7 @@ def config_server_start(args): start_redis_service(args) command = [ - "hypercorn", + *_hypercorn_entrypoint(), *_hypercorn_config_args(args), "--workers", "1", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0d0b9c014b..23f81c1df9 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -81,6 +81,7 @@ class StartArgs: ) chat_template: Optional[str] = field(default=None) running_max_req_size: int = field(default=256) + per_dp_max_req_size: Optional[int] = field(default=None) tp: int = field(default=1) dp: int = field(default=1) nnodes: int = field(default=1) @@ -151,6 +152,9 @@ class StartArgs: disable_cudagraph: bool = field(default=False) enable_prefill_cudagraph: bool = field(default=False) prefill_cudagraph_max_handle_token: int = field(default=8192) + prefill_cudagraph_token_nums: Optional[List[int]] = field(default=None) + prefill_cudagraph_batch_sizes: Optional[List[int]] = field(default=None) + prefill_cudagraph_capture_attention: bool = field(default=False) graph_max_batch_size: int = field(default=256) graph_split_batch_size: int = field(default=32) graph_grow_step_size: int = field(default=16) @@ -162,7 +166,7 @@ class StartArgs: expert_dtype: Optional[str] = field(default=None, metadata={"choices": ["fp8", "fp4"]}) llm_prefill_att_backend: List[str] = field( default_factory=lambda: ["auto"], - metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashqla"]}, + metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashqla", "tilelang"]}, ) llm_decode_att_backend: List[str] = field( default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "flashinfer"]} @@ -220,6 +224,7 @@ class StartArgs: enable_mps: bool = field(default=False) multinode_router_gloo_port: int = field(default=20001) schedule_time_interval: float = field(default=0.03) + prefill_coalesce_interval: float = field(default=0.0) use_dynamic_prompt_cache: bool = field(default=False) enable_rl: bool = field(default=False) enable_torch_memory_saver: bool = field(default=False) diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index 6359f99191..fd13ab04d1 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -1,4 +1,5 @@ import time +import math import uvloop import asyncio import pickle @@ -42,6 +43,23 @@ logger = init_logger(__name__) +def resolve_model_max_req_num(args: StartArgs) -> int: + """Resolve the request-state capacity allocated by each model process.""" + + configured = getattr(args, "per_dp_max_req_size", None) + if configured is None: + return args.running_max_req_size + + local_dp_size = max(1, args.dp // args.nnodes) + min_balanced_capacity = math.ceil(args.running_max_req_size / local_dp_size) + if configured < min_balanced_capacity: + raise ValueError( + "per_dp_max_req_size must be at least ceil(running_max_req_size / local_dp_size): " + f"got {configured}, minimum {min_balanced_capacity}" + ) + return configured + + class RouterManager(RouterMultiNodeTpHelper, RouterRlOpHelper, object): def __init__(self, args: StartArgs): self.args = args @@ -52,6 +70,8 @@ def __init__(self, args: StartArgs): self.node_rank = args.node_rank self.dp_size = args.dp self.schedule_time_interval = args.schedule_time_interval # 默认30ms 的调度周期 + self.prefill_coalesce_interval = max(0.0, args.prefill_coalesce_interval) + self._prefill_coalesce_deadline = None # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容 self.dp_size_in_node = max(1, args.dp // self.nnodes) self.dp_world_size = self.world_size // self.dp_size @@ -148,7 +168,7 @@ async def wait_to_model_ready(self): "weight_dir": self.model_weightdir, "load_way": self.load_way, "max_total_token_num": self.max_total_token_num, - "max_req_num": self.args.running_max_req_size, + "max_req_num": resolve_model_max_req_num(self.args), # MTP length stopping is asynchronous, so up to mtp_step accepted # positions may already be committed when FINISHED_LENGTH is observed. # The overlapped iteration then needs mtp_step positions for target @@ -445,6 +465,43 @@ def _generate_new_batch(self): self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch) return + def _should_defer_prefill_batch(self): + """Hold a partial request burst briefly so prefill can run as one batch.""" + if self.prefill_coalesce_interval <= 0.0: + return False + + waiting_req_num = self.req_queue.get_wait_req_num() + if waiting_req_num <= 0: + self._prefill_coalesce_deadline = None + return False + + scheduled_req_num = 0 + if self.running_batch is not None: + scheduled_req_num += len(self.running_batch.reqs) + if self.schedule_new_batch is not None: + scheduled_req_num += len(self.schedule_new_batch.reqs) + runnable_slots = max(0, self.args.running_max_req_size - scheduled_req_num) + + now = time.monotonic() + if self._prefill_coalesce_deadline is None: + self._prefill_coalesce_deadline = now + self.prefill_coalesce_interval + + # There is no reason to wait once all currently runnable slots can be filled. + if runnable_slots > 0 and waiting_req_num >= runnable_slots: + self._prefill_coalesce_deadline = None + return False + + # Preserve the original deadline while the running batch is full. This lets a + # queued burst launch immediately when capacity becomes available. + if runnable_slots == 0: + return True + + if now < self._prefill_coalesce_deadline: + return True + + self._prefill_coalesce_deadline = None + return False + async def _recv_new_reqs_and_schedule(self): if not hasattr(self, "recv_max_count"): self.recv_max_count = 64 @@ -471,7 +528,7 @@ async def _recv_new_reqs_and_schedule(self): if self.is_multinode_tp: self.multinode_tp_generate_new_batch() else: - if self._get_paused_req_num() == 0: + if self._get_paused_req_num() == 0 and not self._should_defer_prefill_batch(): self._generate_new_batch() return diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6eb6413a02..af7eab0639 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -377,12 +377,20 @@ def _async_copy_next_token_infos_to_pin_mem( ) return next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu - def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tensor) -> torch.Tensor: + def _get_next_token_ranks(self, model_output: ModelOutput, next_token_ids: torch.Tensor) -> torch.Tensor: """计算（或占位）每个 next token 在 vocab 上的 1-based rank（GPU tensor）。 仅 ``--enable_rl`` 时做真实 rank；否则返回 GPU 常量 ``-1``，避免 O(batch * vocab) 比较。 下游 async_copy 在同样条件下会忽略该返回值。 """ + if model_output.has_vocab_parallel_logits: + assert model_output.logits.shape[1] == 1 + return g_pin_mem_manager.get_const_gpu_tensor( + key="next_token_ranks", + shape=next_token_ids.shape, + fill_value=1 if self.args.enable_rl else -1, + dtype=torch.int32, + ) if not self.args.enable_rl: return g_pin_mem_manager.get_const_gpu_tensor( key="next_token_ranks", @@ -390,8 +398,8 @@ def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tens fill_value=-1, dtype=torch.int32, ) - selected_logits = logits.gather(1, next_token_ids.long().view(-1, 1)) - return (logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 + selected_logits = model_output.logits.gather(1, next_token_ids.long().view(-1, 1)) + return (model_output.logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 def _capture_prompt_logprobs_if_needed( self, @@ -862,18 +870,35 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: return [g_infer_context.requests_mapping[req_id] for req_id in req_ids] def _gen_argmax_token_ids(self, model_output: ModelOutput): + if model_output.mtp_collector.draft_token_ids is not None: + return model_output.mtp_collector.draft_token_ids logits = model_output.logits - return torch.argmax(logits, dim=-1) + candidate_indexes = torch.argmax(logits, dim=-1) + return self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): + if model_output.mtp_collector.draft_token_ids is not None: + draft_token_probs = model_output.mtp_collector.draft_token_probs + if draft_token_probs is None: + raise RuntimeError("draft head returned token ids without token probabilities") + return model_output.mtp_collector.draft_token_ids, draft_token_probs logits = model_output.logits - probs = torch.softmax(logits, dim=-1) - max_probs, draft_next_token_ids_gpu = torch.max(probs, dim=-1) - return draft_next_token_ids_gpu, max_probs + if model_output.has_vocab_parallel_logits: + max_logits, candidate_indexes = torch.max(logits, dim=-1) + token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) + return token_ids, torch.exp(max_logits - model_output.logits_logsumexp) + max_probs, token_ids = torch.max(torch.softmax(logits, dim=-1), dim=-1) + return token_ids, max_probs + + @staticmethod + def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor): + if not model_output.has_vocab_parallel_logits: + return candidate_indexes + return model_output.logits_token_ids.gather(1, candidate_indexes.long().view(-1, 1)).view(-1).long() def _sample_and_scatter_token( self, - logits: torch.Tensor, + model_output: ModelOutput, b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, run_reqs: List[InferReq], @@ -882,12 +907,15 @@ def _sample_and_scatter_token( mask_func: Optional[Callable] = None, ): + logits = model_output.logits + if mask_func is not None: + assert not model_output.has_vocab_parallel_logits, "constrained sampling requires dense logits" assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + next_token_ids, next_token_logprobs = sample(model_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_has_out = None if is_prefill: b_has_out = g_pin_mem_manager.gen_from_list( diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..c9bcc18971 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,3 +1,5 @@ +import os + import torch import time from typing import List @@ -27,6 +29,17 @@ class ChunkedPrefillBackend(ModeBackend): def __init__(self) -> None: super().__init__() + # Mega-MoE owns one symmetric communication workspace per rank. Keep + # the CPU pre/post pipeline enabled, but do not let its two host threads + # enqueue a second model forward while that workspace is still live. + self._serialize_sm90_mega_moe_forwards = os.getenv( + "LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0" + ).upper() in { + "1", + "ON", + "TRUE", + } + # 用于控制每一步是执行prefill 和 decode 还是跳过 self.control_state_machine = ControlState() @@ -41,6 +54,24 @@ def __init__(self) -> None: self.classed_req_strict_prefill = False return + def _record_forward_completion(self) -> torch.cuda.Event: + sync_event = torch.cuda.Event() + sync_event.record() + return sync_event + + def _notify_next_forward_when_safe(self, event_pack: OverlapEventPack, sync_event: torch.cuda.Event): + if self._serialize_sm90_mega_moe_forwards: + # Mega-MoE owns one symmetric communication workspace per rank. Let + # this iteration's CPU bookkeeping overlap its GPU work, but wait for + # the workspace to be released before waking the next forward thread. + sync_event.synchronize() + event_pack.notify_forward_and_wait_post_handle() + else: + # Preserve the regular two-forward overlap path. + event_pack.notify_forward_and_wait_post_handle() + sync_event.synchronize() + return + def init_spec_engine(self): self.spec_engine = SpecEngine( backend=self, @@ -107,11 +138,13 @@ def prefill_normal( ): # 第一阶段: 模型推理 model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + if self.prefill_mask_func is not None: + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -123,16 +156,14 @@ def prefill_normal( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) self._post_handle( run_reqs=run_reqs, next_token_ids=next_token_ids_cpu, @@ -152,26 +183,26 @@ def decode_normal( decode_reqs: List[InferReq], ): model_input, run_reqs = prepare_decode_inputs(decode_reqs) + if self.decode_mask_func is not None: + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, is_prefill=False, mask_func=self.decode_mask_func, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=False) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) self._post_handle( run_reqs=run_reqs, next_token_ids=next_token_ids_cpu, @@ -191,6 +222,8 @@ def prefill_mtp( prefill_reqs: List[InferReq], ): model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + if self.prefill_mask_func is not None: + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) @@ -200,7 +233,7 @@ def prefill_mtp( next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -219,16 +252,14 @@ def prefill_mtp( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) self._post_handle( run_reqs=run_reqs, @@ -251,6 +282,8 @@ def decode_mtp( ): """Run the speculative draft-and-verify decode flow.""" model_input, run_reqs = prepare_decode_inputs(decode_reqs) + if self.decode_mask_func is not None: + model_input.use_vocab_parallel_greedy = False spec_engine = self.spec_engine req_num = len(decode_reqs) @@ -272,11 +305,11 @@ def decode_mtp( selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] next_token_ids, next_token_logprobs = sample( - model_output.logits, + model_output, run_reqs, self.eos_id, ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( @@ -331,8 +364,7 @@ def decode_mtp( next_token_ranks=next_token_ranks, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() @@ -350,8 +382,7 @@ def decode_mtp( update_packs = self._pre_post_handle(verify_ok_reqs, is_chuncked_mode=False) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) spec_engine.update_planner_statics( plan=spec_plan, diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py index dfb1020820..65161ca8fa 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py @@ -17,6 +17,7 @@ def reward_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferR assert self.disable_chunked_prefill is True model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + model_input.use_vocab_parallel_greedy = False model_output = self.model.forward(model_input) scores: torch.Tensor = model_output.logits diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index 21979cbef0..2b6eb0b9d3 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -38,12 +38,11 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq model_input, group_run_reqs = prepare_prefill_inputs( group_reqs, is_chuncked_mode=not self.disable_chunked_prefill ) + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) - logits = model_output.logits - batch_idx, run_reqs = self._diverse_copy( master_reqs=group_reqs, b_prefill_has_out=model_input.b_prefill_has_output_cpu ) @@ -60,11 +59,11 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq non_blocking=True ) - logits = logits[batch_idx] + sampled_output = model_output.index_select_logits_rows(batch_idx) b_mtp_index = model_input.b_mtp_index[batch_idx] - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + next_token_ids, next_token_logprobs = sample(sampled_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(sampled_output, next_token_ids) scatter_token( next_token_ids=next_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..a8c1bd8a40 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -187,7 +187,7 @@ def prefill_normal( next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -240,7 +240,7 @@ def decode_normal(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -287,13 +287,8 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) self._capture_prompt_logprobs_if_needed(model_input1, run_reqs1, model_output1.prompt_logics) - logits0 = model_output0.logits - logits1 = model_output1.logits - req_num0, req_num1 = len(run_reqs0), len(run_reqs1) - logits = torch.empty((req_num0 + req_num1, logits0.shape[1]), dtype=logits0.dtype, device=logits0.device) - logits[0:req_num0, :].copy_(logits0, non_blocking=True) - logits[req_num0 : req_num0 + req_num1, :].copy_(logits1, non_blocking=True) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) run_reqs = run_reqs0 + run_reqs1 b_has_out_cpu = model_input0.b_prefill_has_output_cpu + model_input1.b_prefill_has_output_cpu @@ -307,7 +302,7 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -356,7 +351,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output0, model_output1 = self.model.microbatch_overlap_decode(model_input0, model_input1) if req_num0 + req_num1 > 0: - logits = torch.cat((model_output0.logits, model_output1.logits), dim=0) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) b_req_idx = torch.cat((model_input0.b_req_idx, model_input1.b_req_idx), dim=0) b_mtp_index = torch.cat((model_input0.b_mtp_index, model_input1.b_mtp_index), dim=0) ( @@ -365,7 +360,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -421,7 +416,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -499,11 +494,11 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): if req_num > 0: next_token_ids, next_token_logprobs = sample( - model_output.logits, + model_output, run_reqs, self.eos_id, ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_req_mtp_start_loc = gen_b_req_mtp_start_loc( b_mtp_index=model_input.b_mtp_index, @@ -649,17 +644,9 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) self._capture_prompt_logprobs_if_needed(model_input1, run_reqs1, model_output1.prompt_logics) - logits0 = model_output0.logits - logits1 = model_output1.logits req_num0, req_num1 = len(run_reqs0), len(run_reqs1) req_num = req_num0 + req_num1 - logits = torch.empty( - (req_num0 + req_num1, logits0.shape[1]), - dtype=logits0.dtype, - device=logits0.device, - ) - logits[0:req_num0, :].copy_(logits0, non_blocking=True) - logits[req_num0 : (req_num0 + req_num1), :].copy_(logits1, non_blocking=True) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) run_reqs = run_reqs0 + run_reqs1 b_has_out_cpu = model_input0.b_prefill_has_output_cpu + model_input1.b_prefill_has_output_cpu @@ -673,7 +660,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, run_reqs=run_reqs, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, @@ -681,7 +668,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I b_prefill_has_output_cpu=b_has_out_cpu, ) else: - next_token_ids = torch.empty((0,), dtype=torch.int64, device=logits.device) + next_token_ids = torch.empty((0,), dtype=torch.int64, device=sampled_output.logits.device) target_next_token_ids_gpu0 = next_token_ids[:req_num0] target_next_token_ids_gpu1 = next_token_ids[req_num0:] @@ -769,20 +756,12 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf verify_row_num0 = model_input0.batch_size verify_row_num1 = model_input1.batch_size verify_row_num = verify_row_num0 + verify_row_num1 - logits0 = model_output0.logits - logits1 = model_output1.logits run_reqs = run_reqs0 + run_reqs1 if req_num > 0: assert len(run_reqs) == verify_row_num - logits = torch.empty( - (verify_row_num, logits0.shape[1]), - dtype=logits0.dtype, - device=logits0.device, - ) - logits[:verify_row_num0, :].copy_(logits0, non_blocking=True) - logits[verify_row_num0:, :].copy_(logits1, non_blocking=True) - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) + next_token_ids, next_token_logprobs = sample(sampled_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(sampled_output, next_token_ids) ( next_token_ids_cpu, next_token_logprobs_cpu, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 5b29ea0510..ad428d32af 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,5 +1,9 @@ import torch from typing import List, Tuple +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids @@ -8,7 +12,44 @@ from lightllm.utils.envs_utils import get_env_start_args -def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): +def _can_use_unmodified_greedy_logits(reqs: List[InferReq]) -> bool: + """Whether sampling is exactly argmax over the incoming logits.""" + + for req_obj in reqs: + sample_param = req_obj.sampling_param + shm_param = sample_param.shm_param + if shm_param.top_k != 1 or shm_param.temperature != 1.0: + return False + if ( + shm_param.presence_penalty != 0.0 + or shm_param.frequency_penalty != 0.0 + or shm_param.repetition_penalty != 1.0 + ): + return False + if shm_param.exponential_decay_length_penalty.to_tuple()[1] != 1.0: + return False + out_token_len = req_obj.get_cur_total_len() - req_obj.shm_req.input_len + if out_token_len < shm_param.min_new_tokens - 1: + return False + if sample_param.invalid_token_ids: + return False + return True + + +def can_use_vocab_parallel_greedy(reqs: List[InferReq]) -> bool: + return is_vocab_parallel_greedy_enabled() and _can_use_unmodified_greedy_logits(reqs) + + +def sample(model_output: ModelOutput, reqs: List[InferReq], eos_id: List[int] = [2]): + logits = model_output.logits + if model_output.has_vocab_parallel_logits: + if not _can_use_unmodified_greedy_logits(reqs): + raise RuntimeError("vocab-parallel logits require unmodified greedy requests") + candidate_indexes = torch.argmax(logits, dim=-1, keepdim=True) + token_ids = model_output.logits_token_ids.gather(1, candidate_indexes).view(-1).long() + selected_logits = logits.gather(1, candidate_indexes).view(-1) + return token_ids, selected_logits - model_output.logits_logsumexp + ( b_req_idx, b_temperatures, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 22731439c4..94633dfb55 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -3,6 +3,9 @@ from typing import List, Tuple from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.common.basemodel.batch_objs import ModelInput +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + can_use_vocab_parallel_greedy, +) INT64_MAX = torch.iinfo(torch.int64).max @@ -87,6 +90,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> is_prefill=True, b_prefill_has_output_cpu=b_prefill_has_output, multimodal_params=batch_multimodal_params, + use_vocab_parallel_greedy=can_use_vocab_parallel_greedy(run_reqs), ) return model_input, run_reqs @@ -160,6 +164,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_shared_radix_node_id=b_shared_radix_node_id, is_prefill=False, multimodal_params=multimodal_params, + use_vocab_parallel_greedy=can_use_vocab_parallel_greedy(run_reqs), ) return model_input, run_reqs @@ -176,6 +181,9 @@ def overlap_prepare_decode_inputs(req_objs: List[InferReq]): model_input1, run_reqs1 = prepare_decode_inputs( req_objs=decode_reqs1, ) + use_vocab_parallel_greedy = can_use_vocab_parallel_greedy(run_reqs0 + run_reqs1) + model_input0.use_vocab_parallel_greedy = use_vocab_parallel_greedy + model_input1.use_vocab_parallel_greedy = use_vocab_parallel_greedy return model_input0, run_reqs0, decode_reqs0, model_input1, run_reqs1, decode_reqs1 @@ -211,6 +219,9 @@ def overlap_prepare_prefill_inputs(req_objs: List[InferReq]): req_objs=right_reqs, is_chuncked_mode=True, ) + use_vocab_parallel_greedy = can_use_vocab_parallel_greedy(run_reqs0 + run_reqs1) + model_input0.use_vocab_parallel_greedy = use_vocab_parallel_greedy + model_input1.use_vocab_parallel_greedy = use_vocab_parallel_greedy return model_input0, run_reqs0, model_input1, run_reqs1 diff --git a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py index 6b8c23e8fd..137a580ae6 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py @@ -141,7 +141,7 @@ def propose_next_overlap( req_num_by_batch, ) ): - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) draft_token_probs = draft_token_probs.float() diff --git a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py index b25c62b9b7..a021021c86 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py +++ b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py @@ -211,7 +211,15 @@ def _get_draft_cost_ms(self, req_num: int, verify_batch_size: int, draft_step: i if self.spec_mode in ("vanilla_no_att", "eagle_no_att"): return self.draft_infer_costs.estimate(req_num) * draft_step - if self.spec_mode in ("vanilla_with_att", "eagle_with_att", "eagle3"): + if self.spec_mode == "vanilla_with_att": + assert draft_step > 0, f"{self.spec_mode} requires draft_step to be greater than 0" + # Every level in the chained Vanilla proposer forwards the full + # (possibly compacted) verify layout so that its fixed-depth KV + # state covers every committed position. Pricing only the first + # level at B and the remaining levels at N systematically makes a + # wide verify budget look cheaper than the work actually run. + return self.draft_infer_costs.estimate(verify_batch_size) * draft_step + if self.spec_mode in ("eagle_with_att", "eagle3"): assert draft_step > 0, f"{self.spec_mode} requires draft_step to be greater than 0" draft_cost_ms = self.draft_infer_costs.estimate(verify_batch_size) if draft_step > 1: diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 3d2c0a0e86..fc3155b87f 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -79,14 +79,25 @@ def propose_next( verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden extend_output = draft_model.forward(verify_draft_input) - # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 - # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) - if self.enable_dynmaic_mtp: - draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) - schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) + # 第一列 proposal 来自每个请求的 accepted tail。Specialized draft + # heads can return vocab-parallel Top-1 directly; generic heads retain + # the full-logits fallback. + head_token_ids = extend_output.mtp_collector.draft_token_ids + if head_token_ids is not None: + draft_token_ids = head_token_ids.index_select(0, accepted_tail_rows) + if self.enable_dynmaic_mtp: + head_token_probs = extend_output.mtp_collector.draft_token_probs + if head_token_probs is None: + raise RuntimeError("draft head returned token ids without token probabilities") + draft_token_probs = head_token_probs.index_select(0, accepted_tail_rows) + schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) else: - draft_token_ids = self._gen_argmax_token_ids(accepted_tail_output) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) + if self.enable_dynmaic_mtp: + draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) + schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) + else: + draft_token_ids = self._gen_argmax_token_ids(accepted_tail_output) proposal_token_ids_by_step.append(draft_token_ids.unsqueeze(1)) if draft_step == 1: diff --git a/lightllm/utils/device_utils.py b/lightllm/utils/device_utils.py index 58bff90560..1bf4ff7fee 100644 --- a/lightllm/utils/device_utils.py +++ b/lightllm/utils/device_utils.py @@ -45,6 +45,11 @@ def is_sm100_gpu(): return torch.cuda.get_device_capability()[0] == 10 +@lru_cache(maxsize=None) +def is_sm90_gpu(): + return torch.cuda.get_device_capability() == (9, 0) + + @lru_cache(maxsize=None) def get_device_sm_regs_num(): import triton diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 4fed9509a9..93266684a8 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -83,8 +83,34 @@ def get_deepep_num_max_dispatch_tokens_per_rank_prefill(): @lru_cache(maxsize=None) def get_deepep_num_max_dispatch_tokens_per_rank_decode(): - # 该参数需要大于单卡最大batch size，且是8的倍数。该参数与显存占用直接相关，值越大，显存占用越大，如果出现显存不足，可以尝试调小该值 - return int(os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", 256)) + # DeepEP requires this limit to cover the physical rows passed to one + # decode dispatch and to be a multiple of eight. Speculative target + # decode widens every logical request to ``mtp_step + 1`` rows, so the old + # fixed default of 256 failed as soon as (for example) a 64-request MTP5 + # CUDA Graph reached 264 rows. + configured = os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", None) + if configured is not None: + return int(configured) + + args = get_env_start_args() + logical_batch_size = max( + int(args.get("graph_max_batch_size", 0) or 0), + int(args.get("running_max_req_size", 0) or 0), + 1, + ) + verify_width = ( + int(args.get("mtp_step", 0) or 0) + 1 + if args.get("mtp_mode", None) is not None + else 1 + ) + required_tokens = logical_batch_size * verify_width + # In TP/SP + EP mode each rank dispatches only its sequence-parallel slice + # to DeepEP. CUDA Graph batch sizes are TP-aligned, but use ceil here as a + # defensive bound for non-graph decode as well. + if args.get("enable_tpsp_mix_mode", False) and args.get("enable_ep_moe", False): + tp = max(int(args.get("tp", 1) or 1), 1) + required_tokens = (required_tokens + tp - 1) // tp + return ((max(required_tokens, 256) + 7) // 8) * 8 @lru_cache(maxsize=None) diff --git a/test/kernel/test_extract_indexer_ks_dynamic.py b/test/kernel/test_extract_indexer_ks_dynamic.py new file mode 100644 index 0000000000..f948d598e6 --- /dev/null +++ b/test/kernel/test_extract_indexer_ks_dynamic.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import ( + extract_indexer_ks, + extract_indexer_ks_dynamic, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _indexer_buffer(num_slots: int): + keys = ( + torch.arange(num_slots * 128, device="cuda", dtype=torch.float32) + .remainder_(31) + .sub_(15) + .to(torch.float8_e4m3fn) + .view(num_slots, 1, 128) + ) + scales = torch.arange(1, num_slots + 1, device="cuda", dtype=torch.float32).view(num_slots, 1, 1) + buffer = torch.empty((num_slots, 1, 132), device="cuda", dtype=torch.uint8) + buffer[:, :, :128].copy_(keys.view(torch.uint8)) + buffer[:, :, 128:132].copy_(scales.view(torch.uint8).view(num_slots, 1, 4)) + return buffer, keys.view(num_slots, 128), scales.view(num_slots) + + +def _assert_packed(output_keys, output_scales, source_keys, source_scales, expected_slots): + expected_slots = torch.tensor(expected_slots, device="cuda", dtype=torch.int64) + expected_keys = source_keys.index_select(0, expected_slots) + expected_scales = source_scales.index_select(0, expected_slots) + torch.testing.assert_close( + output_keys[: len(expected_slots)].float(), + expected_keys.float(), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + output_scales[: len(expected_slots)], + expected_scales, + rtol=0, + atol=0, + ) + + +def test_dynamic_layout_matches_fixed_full_width_layout(): + buffer, source_keys, source_scales = _indexer_buffer(8) + req_to_token = torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7]], + device="cuda", + dtype=torch.int32, + ) + b_req_idx = torch.tensor([0, 0, 0, 1, 1, 1], device="cuda", dtype=torch.int32) + b_mtp_index = torch.tensor([0, 1, 2, 0, 1, 2], device="cuda", dtype=torch.int32) + b_seq_len = torch.tensor([2, 3, 4, 1, 2, 3], device="cuda", dtype=torch.int32) + + fixed_keys, fixed_scales = extract_indexer_ks( + I_buffer=buffer, + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + req_to_token_indexs=req_to_token, + out_token_num=24, + max_kv_seq_len=4, + mtp_step=2, + ) + dynamic_keys, dynamic_scales = extract_indexer_ks_dynamic( + I_buffer=buffer, + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_token_indexs=req_to_token, + max_kv_seq_len=4, + max_request_num=2, + ) + + torch.testing.assert_close(dynamic_keys[:7].float(), fixed_keys[:7].float(), rtol=0, atol=0) + torch.testing.assert_close(dynamic_scales[:7], fixed_scales[:7], rtol=0, atol=0) + + +def test_dynamic_layout_packs_variable_request_widths(): + buffer, source_keys, source_scales = _indexer_buffer(12) + req_to_token = torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], + device="cuda", + dtype=torch.int32, + ) + b_req_idx = torch.tensor([0, 0, 0, 1, 2], device="cuda", dtype=torch.int32) + b_mtp_index = torch.tensor([0, 1, 2, 0, 0], device="cuda", dtype=torch.int32) + b_seq_len = torch.tensor([2, 3, 4, 3, 2], device="cuda", dtype=torch.int32) + + output_keys, output_scales = extract_indexer_ks_dynamic( + I_buffer=buffer, + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_token_indexs=req_to_token, + max_kv_seq_len=4, + max_request_num=3, + ) + + _assert_packed( + output_keys, + output_scales, + source_keys, + source_scales, + expected_slots=[0, 1, 2, 3, 4, 5, 6, 8, 9], + ) diff --git a/test/kernel/test_glm5_grouped_topk.py b/test/kernel/test_glm5_grouped_topk.py new file mode 100644 index 0000000000..dde6413988 --- /dev/null +++ b/test/kernel/test_glm5_grouped_topk.py @@ -0,0 +1,229 @@ +#!/usr/bin/env python3 + +import argparse + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import ( + single_group_sigmoid_topk_kernel, + single_group_sigmoid_topk_bitonic_kernel, + triton_grouped_topk, +) + + +def torch_reference(gating_output, correction_bias, topk): + scores = gating_output.float().sigmoid() + choice_scores = scores + correction_bias + topk_ids = torch.topk( + choice_scores, k=topk, dim=-1, largest=True, sorted=True + ).indices + topk_weights = torch.gather(scores, 1, topk_ids) + topk_weights /= topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + + +def run_topk(gating_output, correction_bias, *, fast): + return triton_grouped_topk( + hidden_states=None, + gating_output=gating_output, + correction_bias=correction_bias, + topk=8, + renormalize=True, + num_expert_group=1, + topk_group=1, + scoring_func="sigmoid", + group_score_used_topk_num=1, + use_single_group_fast_path=fast, + ) + + +def run_fast_with_warps(gating_output, correction_bias, num_warps): + tokens, experts = gating_output.shape + topk_weights = torch.empty((tokens, 8), dtype=torch.float32, device="cuda") + topk_ids = torch.empty((tokens, 8), dtype=torch.long, device="cuda") + single_group_sigmoid_topk_kernel[(tokens,)]( + gating_output, + gating_output.stride(0), + correction_bias, + topk_weights, + topk_weights.stride(0), + topk_ids, + topk_ids.stride(0), + experts, + HAS_CORRECTION_BIAS=True, + EXPERT_BLOCK_SIZE=triton.next_power_of_2(experts), + TOPK_BLOCK_SIZE=8, + TOPK_NUM=8, + RENORMALIZE=True, + num_warps=num_warps, + num_stages=1, + ) + return topk_weights, topk_ids + + +def run_scratch_free_bitonic(gating_output, correction_bias): + tokens, experts = gating_output.shape + topk_weights = torch.empty((tokens, 8), dtype=torch.float32, device="cuda") + topk_ids = torch.empty((tokens, 8), dtype=torch.long, device="cuda") + single_group_sigmoid_topk_bitonic_kernel[(tokens,)]( + gating_output, + gating_output.stride(0), + correction_bias, + topk_weights, + topk_weights.stride(0), + topk_ids, + topk_ids.stride(0), + experts, + HAS_CORRECTION_BIAS=True, + EXPERT_BLOCK_SIZE=triton.next_power_of_2(experts), + TOPK_NUM=8, + RENORMALIZE=True, + num_warps=4, + num_stages=1, + ) + return topk_weights, topk_ids + + +def assert_correct(tokens): + generator = torch.Generator(device="cuda").manual_seed(20260828 + tokens) + gating_output = torch.randn( + (tokens, 288), generator=generator, dtype=torch.float32, device="cuda" + ) + correction_bias = torch.randn( + (288,), generator=generator, dtype=torch.float32, device="cuda" + ) + + ref_weights, ref_ids = torch_reference(gating_output, correction_bias, 8) + fast_weights, fast_ids = run_topk( + gating_output, correction_bias, fast=True + ) + generic_weights, generic_ids = run_topk( + gating_output, correction_bias, fast=False + ) + + torch.testing.assert_close(fast_ids, ref_ids, rtol=0, atol=0) + torch.testing.assert_close(generic_ids, ref_ids, rtol=0, atol=0) + torch.testing.assert_close(fast_weights, ref_weights, rtol=1e-5, atol=1e-6) + torch.testing.assert_close( + generic_weights, ref_weights, rtol=1e-5, atol=1e-6 + ) + weight_delta = (fast_weights - generic_weights).abs() + print( + f"tokens={tokens}: exact expert ids, weights match reference; " + f"fast/generic bitwise={torch.equal(fast_weights, generic_weights)} " + f"max_delta={weight_delta.max().item():.9g}" + ) + return gating_output, correction_bias + + +def capture_graph(fn): + for _ in range(3): + outputs = fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = fn() + return graph, outputs + + +def graph_ms(graph, iterations): + for _ in range(20): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def benchmark(tokens, iterations): + gating_output, correction_bias = assert_correct(tokens) + fast_graph, fast_outputs = capture_graph( + lambda: run_topk(gating_output, correction_bias, fast=True) + ) + generic_graph, generic_outputs = capture_graph( + lambda: run_topk(gating_output, correction_bias, fast=False) + ) + bitonic_graph, bitonic_outputs = capture_graph( + lambda: run_scratch_free_bitonic(gating_output, correction_bias) + ) + fast_ms = graph_ms(fast_graph, iterations) + generic_ms = graph_ms(generic_graph, iterations) + bitonic_ms = graph_ms(bitonic_graph, iterations) + torch.testing.assert_close(fast_outputs[1], generic_outputs[1], rtol=0, atol=0) + torch.testing.assert_close(bitonic_outputs[0], generic_outputs[0], rtol=0, atol=0) + torch.testing.assert_close(bitonic_outputs[1], generic_outputs[1], rtol=0, atol=0) + speedup = generic_ms / fast_ms + print( + f"tokens={tokens}: graph fast={fast_ms:.6f} ms " + f"bitonic={bitonic_ms:.6f} ms generic={generic_ms:.6f} ms " + f"speedup={speedup:.2f}x bitonic_speedup={generic_ms / bitonic_ms:.2f}x" + ) + return speedup + + +def tune_warps(tokens, iterations): + gating_output, correction_bias = assert_correct(tokens) + ref_weights, ref_ids = torch_reference(gating_output, correction_bias, 8) + generic_weights, _ = run_topk(gating_output, correction_bias, fast=False) + results = [] + for num_warps in (1, 2, 4, 8, 16): + graph, outputs = capture_graph( + lambda num_warps=num_warps: run_fast_with_warps( + gating_output, correction_bias, num_warps + ) + ) + graph.replay() + torch.cuda.synchronize() + if not torch.equal(outputs[1], ref_ids) or not torch.allclose( + outputs[0], ref_weights, rtol=1e-5, atol=1e-6 + ): + print(f"tokens={tokens}: num_warps={num_warps} INVALID") + continue + elapsed_ms = graph_ms(graph, iterations) + results.append((elapsed_ms, num_warps)) + max_delta = (outputs[0] - generic_weights).abs().max().item() + print( + f"tokens={tokens}: num_warps={num_warps} " + f"graph={elapsed_ms:.6f} ms " + f"generic_bitwise={torch.equal(outputs[0], generic_weights)} " + f"max_delta={max_delta:.9g}" + ) + best_ms, best_warps = min(results) + print( + f"tokens={tokens}: best num_warps={best_warps} graph={best_ms:.6f} ms" + ) + + +def test_glm5_single_group_topk(): + if not torch.cuda.is_available(): + return + for tokens in (1, 8, 48, 256): + assert_correct(tokens) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--tune-warps", action="store_true") + parser.add_argument("--iterations", type=int, default=5000) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required") + if args.tune_warps: + for tokens in (1, 8, 48, 256): + tune_warps(tokens, args.iterations) + elif args.benchmark: + for tokens in (1, 8, 48, 256): + benchmark(tokens, args.iterations) + else: + test_glm5_single_group_topk() + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_mhc.py b/test/kernel/test_glm5_mhc.py new file mode 100644 index 0000000000..bc9456bacc --- /dev/null +++ b/test/kernel/test_glm5_mhc.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Numerical and launch-overhead check for the fused GLM-5 mHC kernels.""" + +import argparse + +import torch + +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next.triton_kernel.mhc import ( + hc_post, + hc_post_reference, + hc_pre_norm, + hc_pre_reference, +) + + +def _max_error(actual: torch.Tensor, expected: torch.Tensor) -> float: + return (actual.float() - expected.float()).abs().max().item() + + +def _time_ms(function, iterations: int = 100) -> float: + for _ in range(5): + function() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--large-post", action="store_true") + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + torch.manual_seed(1234) + device = torch.device("cuda") + streams = 4 + hidden = 4096 + for tokens in (1, 8, 48): + x = torch.randn( + (tokens, streams * hidden), device=device, dtype=torch.bfloat16 + ) + fn = 0.005 * torch.randn( + ((2 + streams) * streams, streams * hidden), + device=device, + dtype=torch.float32, + ) + scale = torch.randn((3,), device=device, dtype=torch.float32) + base = torch.randn( + ((2 + streams) * streams,), device=device, dtype=torch.float32 + ) + layer_output = torch.randn( + (tokens, hidden), device=device, dtype=torch.bfloat16 + ) + norm_weight = torch.randn( + (hidden,), device=device, dtype=torch.bfloat16 + ) + arguments = (x, fn, scale, base, streams, 1e-6, 1e-6, 20) + + expected_pre = hc_pre_reference(*arguments) + expected_pre = ( + rmsnorm_forward(expected_pre[0], weight=norm_weight, eps=1e-6), + expected_pre[1], + expected_pre[2], + ) + actual_pre = hc_pre_norm( + x, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + torch.cuda.synchronize() + pre_errors = tuple( + _max_error(actual, expected) + for actual, expected in zip(actual_pre, expected_pre) + ) + assert pre_errors[0] <= 0.03125, pre_errors + # DeepGEMM intentionally uses TF32 and a split-K reduction, matching + # the optimized serving path rather than torch.mm's accumulation order. + assert pre_errors[1] <= 5e-4, pre_errors + assert pre_errors[2] <= 5e-4, pre_errors + + expected_post = hc_post_reference( + layer_output, x, expected_pre[1], expected_pre[2], streams + ) + actual_post = hc_post( + layer_output, x, actual_pre[1], actual_pre[2], streams + ) + torch.cuda.synchronize() + post_error = _max_error(actual_post, expected_post) + assert post_error <= 0.03125, post_error + + def reference_path(): + layer_input, residual_mix, post_mix = hc_pre_reference(*arguments) + rmsnorm_forward(layer_input, weight=norm_weight, eps=1e-6) + return hc_post_reference( + layer_output, x, residual_mix, post_mix, streams + ) + + def fused_path(): + _, residual_mix, post_mix = hc_pre_norm( + x, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + return hc_post( + layer_output, x, residual_mix, post_mix, streams + ) + + def fused_pre_path(): + return hc_pre_norm( + x, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + + def fused_post_path(): + return hc_post( + layer_output, x, actual_pre[1], actual_pre[2], streams + ) + + reference_ms = _time_ms(reference_path) + fused_ms = _time_ms(fused_path) + fused_pre_ms = _time_ms(fused_pre_path) + fused_post_ms = _time_ms(fused_post_path) + print( + f"PASS tokens={tokens} pre_errors={pre_errors} " + f"post_error={post_error:.8f} reference_ms={reference_ms:.4f} " + f"fused_ms={fused_ms:.4f} pre_ms={fused_pre_ms:.4f} " + f"post_ms={fused_post_ms:.4f} " + f"speedup={reference_ms / fused_ms:.2f}x" + ) + + if args.large_post: + from sglang.kernels.ops.layernorm.mhc import mhc_post_tilelang + + tokens = 17152 + residual = torch.randn( + (tokens, streams * hidden), device=device, dtype=torch.bfloat16 + ) + layer_output = torch.randn( + (tokens, hidden), device=device, dtype=torch.bfloat16 + ) + residual_mix = torch.randn( + (tokens, streams, streams), device=device, dtype=torch.float32 + ) + post_mix = torch.randn( + (tokens, streams), device=device, dtype=torch.float32 + ) + actual = hc_post( + layer_output, residual, residual_mix, post_mix, streams + ) + def sgl_mhc_post(): + output = torch.empty_like( + residual.view(tokens, streams, hidden) + ) + mhc_post_tilelang( + residual_mix, + residual.view(tokens, streams, hidden), + post_mix, + layer_output, + output, + streams, + hidden, + ) + return output + + sglang_output = sgl_mhc_post().view(tokens, -1) + torch.cuda.synchronize() + cross_error = _max_error(actual, sglang_output) + assert cross_error <= 0.0625, cross_error + triton_ms = _time_ms( + lambda: hc_post( + layer_output, residual, residual_mix, post_mix, streams + ), + args.iterations, + ) + tilelang_ms = _time_ms( + sgl_mhc_post, + args.iterations, + ) + print( + f"PASS large_post tokens={tokens} cross_error={cross_error:.8f} " + f"triton_ms={triton_ms:.4f} tilelang_ms={tilelang_ms:.4f} " + f"triton_over_tilelang={tilelang_ms / triton_ms:.2f}x" + ) + + # Compare the full DeepGEMM prenorm + mHC-pre fusion used by both + # runtimes. The standalone kernel test has no SGLang process group, + # so disable only its symmetric-allocation context. + import contextlib + import sglang.kernels.ops.layernorm.mhc as sgl_mhc + + sgl_mhc.get_tp_group = lambda: None + sgl_mhc.is_allocation_symmetric = lambda: False + sgl_mhc.use_symmetric_memory = ( + lambda *_args, **_kwargs: contextlib.nullcontext() + ) + fn = 0.005 * torch.randn( + ((2 + streams) * streams, streams * hidden), + device=device, + dtype=torch.float32, + ) + scale = torch.randn((3,), device=device, dtype=torch.float32) + base = torch.randn( + ((2 + streams) * streams,), device=device, dtype=torch.float32 + ) + norm_weight = torch.randn( + (hidden,), device=device, dtype=torch.bfloat16 + ) + + def lightllm_pre(): + return hc_pre_norm( + residual, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + + def sglang_pre(): + return sgl_mhc.mhc_pre( + residual.view(tokens, streams, hidden), + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 2.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + lightllm_result = lightllm_pre() + sglang_result = sglang_pre() + torch.cuda.synchronize() + pre_cross_errors = ( + _max_error(lightllm_result[0], sglang_result[2]), + _max_error(lightllm_result[1], sglang_result[1]), + _max_error(lightllm_result[2], sglang_result[0].squeeze(-1)), + ) + assert pre_cross_errors[0] <= 0.0625, pre_cross_errors + assert pre_cross_errors[1] <= 5e-4, pre_cross_errors + assert pre_cross_errors[2] <= 5e-4, pre_cross_errors + lightllm_pre_ms = _time_ms(lightllm_pre, args.iterations) + sglang_pre_ms = _time_ms(sglang_pre, args.iterations) + print( + f"PASS large_pre tokens={tokens} errors={pre_cross_errors} " + f"lightllm_ms={lightllm_pre_ms:.4f} " + f"sglang_ms={sglang_pre_ms:.4f} " + f"sglang_speedup={lightllm_pre_ms / sglang_pre_ms:.2f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_sglang_moe_compat.py b/test/kernel/test_glm5_sglang_moe_compat.py new file mode 100644 index 0000000000..6edcd3b262 --- /dev/null +++ b/test/kernel/test_glm5_sglang_moe_compat.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Numerically compare LightLLM and SGLang's GLM-5 FP8 MoE paths. + +This is an integration probe for development images that provide both +packages. It intentionally uses GLM-5's TP8 decode shapes, including the +fused shared expert, block-wise FP8 scales, and clamped SwiGLU. +""" + +import argparse +import itertools +import json +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import ( + fused_experts as lightllm_fused_experts, +) +from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe as sglang_fused_moe +from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe_triton_config as sglang_fused_moe_config, +) +from sglang.srt.layers.moe.moe_runner.triton_utils import override_config + + +def _fp8_randn(shape, *, scale=0.02): + return (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * scale).to( + torch.float8_e4m3fn + ) + + +def _graph_ms(fn, source, iterations): + static_input = source.clone() + for _ in range(3): + static_input.copy_(source) + fn(static_input) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_input.copy_(source) + fn(static_input) + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + elapsed_ms = start.elapsed_time(end) / iterations + graph.reset() + return elapsed_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", action="store_true") + parser.add_argument( + "--fuse-sum", + action="store_true", + help="Fuse the top-k sum into SGLang's down-projection kernel.", + ) + parser.add_argument( + "--tune-configs", + action="store_true", + help="Search a compact H100 config set for GLM-5's TP8 decode shape.", + ) + parser.add_argument( + "--tune-tma-configs", + action="store_true", + help="Search the compact small-M set with SGLang's up-projection TMA path.", + ) + parser.add_argument( + "--num-tokens", + type=int, + default=48, + help="Physical token count used by the MoE probe (48 main, 8 MTP draft).", + ) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + torch.manual_seed(123) + + num_tokens = args.num_tokens + num_experts = 289 + topk = 9 + hidden_size = 4096 + tp_intermediate_size = 256 + + hidden_states = torch.randn( + (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16 + ) + w13 = _fp8_randn((num_experts, 2 * tp_intermediate_size, hidden_size)) + w2 = _fp8_randn((num_experts, hidden_size, tp_intermediate_size)) + w13_scale = torch.ones( + (num_experts, 2 * tp_intermediate_size // 128, hidden_size // 128), + device="cuda", + dtype=torch.float32, + ) + w2_scale = torch.ones( + (num_experts, hidden_size // 128, tp_intermediate_size // 128), + device="cuda", + dtype=torch.float32, + ) + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), device="cuda", dtype=torch.int64 + ) + topk_weights = torch.rand( + (num_tokens, topk), device="cuda", dtype=torch.float32 + ) + topk_weights.mul_(2.5 / topk_weights.sum(dim=-1, keepdim=True)) + + def run_lightllm(output): + lightllm_fused_experts( + hidden_states=output, + w1=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w13_scale, + w2_scale=w2_scale, + limit=10.0, + alpha=1.0, + clamp_up_add_one=False, + ) + + lightllm_output = hidden_states.clone() + run_lightllm(lightllm_output) + + # SGLang's fused sequence consults this runtime flag. A standalone + # LightLLM process has no SGLang RuntimeContext, so provide the selected + # fused top-k sum behavior explicitly for this comparison. + standalone_exec = SimpleNamespace( + moe=SimpleNamespace(enable_fused_moe_sum_all_reduce=args.fuse_sum), + deterministic=SimpleNamespace(enable_deterministic_inference=False), + ) + sglang_fused_moe.get_exec = lambda: standalone_exec + sglang_fused_moe_config.get_exec = lambda: standalone_exec + def run_sglang(output): + sglang_fused_moe.fused_experts_impl( + hidden_states=output, + w1=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w13_scale, + w2_scale=w2_scale, + block_shape=[128, 128], + routed_scaling_factor=1.0, + filter_expert=False, + swiglu_limit=10.0, + gate_up_interleaved=False, + ) + + sglang_output = hidden_states.clone() + run_sglang(sglang_output) + torch.cuda.synchronize() + + diff = (lightllm_output.float() - sglang_output.float()).abs() + reference = lightllm_output.float().abs() + print( + "max_abs=%.6f mean_abs=%.6f max_ref=%.6f mean_ref=%.6f" + % ( + diff.max().item(), + diff.mean().item(), + reference.max().item(), + reference.mean().item(), + ) + ) + torch.testing.assert_close( + sglang_output.float(), + lightllm_output.float(), + rtol=0.08, + atol=0.08, + ) + print("GLM-5 SGLang MoE compatibility: PASS") + + if args.benchmark: + lightllm_ms = _graph_ms(run_lightllm, hidden_states, args.iterations) + sglang_ms = _graph_ms(run_sglang, hidden_states, args.iterations) + print( + "graph_ms lightllm=%.6f sglang=%.6f speedup=%.3fx" + % (lightllm_ms, sglang_ms, lightllm_ms / sglang_ms) + ) + + if args.tune_configs or args.tune_tma_configs: + # GLM-5 decode has few physical tokens (48 for the main model and 8 + # for each draft) spread over 289 experts. SGLang's generic block-FP8 + # fallback uses BLOCK_SIZE_M=64, which can pad this sparse workload + # excessively. Search the small-M region instead of the generic + # 1920-config tuning space. + search_space = itertools.product( + (16, 32), + (64, 128), + (128,) if args.tune_tma_configs else (64, 128), + (1, 16), + (4, 8), + (2, 3), + ) + candidates = [ + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_m, + "num_warps": num_warps, + "num_stages": num_stages, + **({"USE_TMA": True} if args.tune_tma_configs else {}), + } + for block_m, block_n, block_k, group_m, num_warps, num_stages in search_space + ] + default_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + candidates.append(default_config) + + results = [] + for config in candidates: + try: + with override_config(config): + graph_ms = _graph_ms(run_sglang, hidden_states, args.iterations) + except Exception as exc: + print("config_failed=%s error=%r" % (json.dumps(config, sort_keys=True), exc)) + continue + results.append((graph_ms, config)) + print( + "config_ms=%.6f config=%s" + % (graph_ms, json.dumps(config, sort_keys=True)) + ) + + results.sort(key=lambda item: item[0]) + if not results: + raise RuntimeError("all SGLang MoE tuning candidates failed") + print("top_configs=%s" % json.dumps(results[:10], sort_keys=True)) + print("best_config=%s" % json.dumps(results[0][1], sort_keys=True)) + print("best_graph_ms=%.6f" % results[0][0]) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_short_decode_helpers.py b/test/kernel/test_glm5_short_decode_helpers.py new file mode 100644 index 0000000000..4b276a771b --- /dev/null +++ b/test/kernel/test_glm5_short_decode_helpers.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from lightllm.models.glm5_next_mtp.triton_kernel.zero_position_embedding import ( + zero_position_embedding_, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_zero_position_embedding_only_changes_position_zero_rows(): + embeddings = torch.arange(15, device="cuda", dtype=torch.bfloat16).view(3, 5) + original = embeddings.clone() + position_ids = torch.tensor([0, 1, 0], device="cuda", dtype=torch.int64) + + zero_position_embedding_(embeddings, position_ids) + + torch.testing.assert_close(embeddings[0], torch.zeros_like(embeddings[0]), rtol=0, atol=0) + torch.testing.assert_close(embeddings[1], original[1], rtol=0, atol=0) + torch.testing.assert_close(embeddings[2], torch.zeros_like(embeddings[2]), rtol=0, atol=0) diff --git a/test/kernel/test_glm5_strided_causal_conv.py b/test/kernel/test_glm5_strided_causal_conv.py new file mode 100644 index 0000000000..f4b03d37a5 --- /dev/null +++ b/test/kernel/test_glm5_strided_causal_conv.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Check GLM-5 KDA's copy-free strided causal convolution path.""" + +import argparse + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, +) + + +def _inputs(seq_lens: list[int], dim: int, cache_lines: int): + total_tokens = sum(seq_lens) + token_major = torch.randn( + (total_tokens, dim), device="cuda", dtype=torch.bfloat16 + ) + weight = torch.randn((dim, 4), device="cuda", dtype=torch.bfloat16) + conv_states = torch.randn( + (cache_lines, dim, 3), device="cuda", dtype=torch.bfloat16 + ) + cache_indices = torch.arange( + len(seq_lens), device="cuda", dtype=torch.int32 + ) + has_initial_state = torch.tensor( + [(index % 2) == 1 for index in range(len(seq_lens))], + device="cuda", + dtype=torch.bool, + ) + query_start_loc = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], + device="cuda", + dtype=torch.int32, + ) + return ( + token_major, + weight, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + ) + + +def _run(arguments, seq_lens: list[int], *, copy_free: bool): + ( + token_major, + weight, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + ) = arguments + return causal_conv1d_fn( + token_major.transpose(0, 1), + weight, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_states, + activation="silu", + seq_lens_cpu=seq_lens if copy_free else None, + ) + + +def _time_ms(function, iterations: int) -> float: + for _ in range(3): + function() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--iterations", type=int, default=20) + args = parser.parse_args() + + torch.manual_seed(1234) + seq_lens = [7, 1, 19, 5] + inputs = _inputs(seq_lens, dim=64, cache_lines=8) + reference_inputs = tuple( + value.clone() if isinstance(value, torch.Tensor) else value + for value in inputs + ) + actual = _run(inputs, seq_lens, copy_free=True) + expected = _run(reference_inputs, seq_lens, copy_free=False) + torch.cuda.synchronize() + output_error = (actual.float() - expected.float()).abs().max().item() + state_error = ( + inputs[2].float() - reference_inputs[2].float() + ).abs().max().item() + # The two kernels accumulate the four taps in a different order; one BF16 + # ULP at this random input scale is expected. + assert output_error <= 0.0625, output_error + assert state_error == 0.0, state_error + assert actual.transpose(0, 1).is_contiguous() + print( + f"PASS correctness output_error={output_error:.8f} " + f"state_error={state_error:.8f}" + ) + + if args.benchmark: + seq_lens = [268] * 64 + strided_inputs = _inputs(seq_lens, dim=3072, cache_lines=64) + copied_inputs = tuple( + value.clone() if isinstance(value, torch.Tensor) else value + for value in strided_inputs + ) + strided_ms = _time_ms( + lambda: _run(strided_inputs, seq_lens, copy_free=True), + args.iterations, + ) + copied_ms = _time_ms( + lambda: _run(copied_inputs, seq_lens, copy_free=False), + args.iterations, + ) + print( + f"PASS benchmark tokens={sum(seq_lens)} dim=3072 " + f"strided_ms={strided_ms:.4f} copied_ms={copied_ms:.4f} " + f"speedup={copied_ms / strided_ms:.2f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_vocab_parallel_top1.py b/test/kernel/test_glm5_vocab_parallel_top1.py new file mode 100644 index 0000000000..811ce8398c --- /dev/null +++ b/test/kernel/test_glm5_vocab_parallel_top1.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest +import torch +import torch.distributed as dist + +from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import ( + vocab_parallel_top1, + vocab_parallel_top1_and_prob, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_vocab_parallel_top1_matches_full_vocab_softmax(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + initialized_here = not dist.is_initialized() + if initialized_here: + dist.init_process_group(backend="nccl", device_id=torch.device("cuda", local_rank)) + + try: + rank = dist.get_rank() + world_size = dist.get_world_size() + assert world_size >= 2 + + torch.manual_seed(20260828) + token_num = 8 + vocab_size = 154880 + assert vocab_size % world_size == 0 + full_logits = torch.randn( + (vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + ) + # Exercise rank-boundary winners and deterministic ties. Global + # argmax must choose the lower vocabulary id just like torch.argmax. + shard_size = vocab_size // world_size + full_logits[shard_size - 1, 0] = 20 + full_logits[shard_size, 0] = 20 + full_logits[-1, 1] = 21 + + local_start = rank * shard_size + local_logits = full_logits[local_start : local_start + shard_size].contiguous() + def alloc_func(shape, dtype, device): + return torch.empty(shape, dtype=dtype, device=device) + token_ids, token_probs = vocab_parallel_top1_and_prob( + local_logits=local_logits, + local_vocab_start_id=local_start, + tp_world_size=world_size, + dist_group=dist.group.WORLD, + alloc_func=alloc_func, + ) + + expected_probs, expected_ids = torch.softmax(full_logits.float().t(), dim=-1).max(dim=-1) + torch.testing.assert_close(token_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(token_probs, expected_probs, rtol=2e-5, atol=2e-7) + torch.testing.assert_close( + vocab_parallel_top1( + local_logits, + local_start, + world_size, + dist.group.WORLD, + alloc_func, + ), + expected_ids, + rtol=0, + atol=0, + ) + + if os.environ.get("LIGHTLLM_TEST_SKIP_CUDA_GRAPH") == "1": + return + + # The draft head normally runs inside a decode CUDA Graph. Replay + # with changed logits to ensure the tiny collective and returned + # tensors are captured as live outputs rather than stale warmup data. + graph_logits = local_logits.clone() + for _ in range(2): + vocab_parallel_top1_and_prob( + graph_logits, + local_start, + world_size, + dist.group.WORLD, + alloc_func, + ) + dist.barrier() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_ids, graph_probs = vocab_parallel_top1_and_prob( + graph_logits, + local_start, + world_size, + dist.group.WORLD, + alloc_func, + ) + + torch.manual_seed(20260829) + replay_full_logits = torch.randn_like(full_logits) + replay_full_logits[shard_size * 2 - 1, 2] = 22 + graph_logits.copy_(replay_full_logits[local_start : local_start + shard_size]) + graph.replay() + torch.cuda.synchronize() + expected_probs, expected_ids = torch.softmax(replay_full_logits.float().t(), dim=-1).max(dim=-1) + torch.testing.assert_close(graph_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(graph_probs, expected_probs, rtol=2e-5, atol=2e-7) + finally: + if initialized_here: + dist.destroy_process_group() diff --git a/test/test_moe_prefill_dispatch.py b/test/test_moe_prefill_dispatch.py new file mode 100644 index 0000000000..f7f70d45b4 --- /dev/null +++ b/test/test_moe_prefill_dispatch.py @@ -0,0 +1,47 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import ( + Deepseek2TransformerLayerInfer, +) + + +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_tp_moe_propagates_prefill_stage(is_prefill): + captured = {} + + class Gate: + data_type_ = torch.float32 + + @staticmethod + def mm(hidden_states): + return torch.zeros((hidden_states.shape[0], 4)) + + class Experts: + @staticmethod + def experts(*args, **kwargs): + captured.update(kwargs) + + layer = Deepseek2TransformerLayerInfer.__new__( + Deepseek2TransformerLayerInfer + ) + layer.embed_dim_ = 4 + layer.n_shared_experts = None + layer.num_experts_per_tok = 2 + layer.norm_topk_prob = True + layer.n_group = 1 + layer.topk_group = 1 + layer_weight = SimpleNamespace( + moe_gate=Gate(), + experts=Experts(), + num_fused_shared_experts=0, + ) + infer_state = SimpleNamespace(is_prefill=is_prefill) + + output = layer._moe_ffn_tp(torch.zeros((3, 4)), infer_state, layer_weight) + + assert output.shape == (3, 4) + assert captured["is_prefill"] is is_prefill + assert captured["infer_state"] is infer_state diff --git a/tools/analyze_torch_trace.py b/tools/analyze_torch_trace.py new file mode 100644 index 0000000000..c19c330920 --- /dev/null +++ b/tools/analyze_torch_trace.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 + +"""Aggregate CUDA kernel time from one or more Torch profiler traces.""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +import gzip +import json +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("traces", nargs="+", type=Path) + parser.add_argument("--graph-only", action="store_true") + parser.add_argument("--graph-id", type=int) + parser.add_argument("--top", type=int, default=40) + args = parser.parse_args() + + totals: dict[str, list[float]] = defaultdict(lambda: [0.0, 0.0, 0.0]) + graph_totals: dict[int, list[float]] = defaultdict(lambda: [0.0, 0.0]) + for path in args.traces: + opener = gzip.open if path.suffix == ".gz" else open + with opener(path, "rt", encoding="utf-8") as trace_file: + events = json.load(trace_file)["traceEvents"] + for event in events: + if event.get("cat") != "kernel" or "dur" not in event: + continue + graph_id = int(event.get("args", {}).get("graph id", 0)) + if args.graph_only and graph_id == 0: + continue + if args.graph_id is not None and graph_id != args.graph_id: + continue + duration = float(event["dur"]) + name = event.get("name", "") + totals[name][0] += duration + totals[name][1] += 1 + totals[name][2] = max(totals[name][2], duration) + graph_totals[graph_id][0] += duration + graph_totals[graph_id][1] += 1 + + total_us = sum(values[0] for values in totals.values()) + total_calls = int(sum(values[1] for values in totals.values())) + print(f"CUDA total: {total_us / 1000:.3f} ms across {total_calls} kernels") + print("Graph totals:") + for graph_id, (duration, count) in sorted( + graph_totals.items(), key=lambda item: item[1][0], reverse=True + ): + print( + f" graph={graph_id:<5} total_ms={duration / 1000:10.3f} " + f"share={duration / total_us:7.2%} calls={int(count)}" + ) + print("Kernel totals:") + for name, (duration, count, maximum) in sorted( + totals.items(), key=lambda item: item[1][0], reverse=True + )[: args.top]: + print( + f" total_ms={duration / 1000:10.3f} share={duration / total_us:7.2%} " + f"calls={int(count):7d} avg_us={duration / count:9.3f} " + f"max_us={maximum:9.3f} {name}" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_allreduce.py b/tools/bench_glm53_allreduce.py new file mode 100644 index 0000000000..eb56a169f7 --- /dev/null +++ b/tools/bench_glm53_allreduce.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +"""Compare NCCL and symmetric-memory all-reduce for GLM-5.3 TP8 tensors.""" + +import argparse +import os + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +import flashinfer.comm as flashinfer_comm +from flashinfer.comm.mnnvl import TorchDistBackend + + +def elapsed_ms(fn, warmup: int, iterations: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def graph_elapsed_ms(fn, warmup: int, iterations: int) -> float: + """Measure replay cost, matching LightLLM's decode CUDA-graph path.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + for _ in range(warmup): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def max_rank(value: float) -> float: + tensor = torch.tensor(value, device="cuda", dtype=torch.float64) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return tensor.item() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=17152) + parser.add_argument("--hidden-size", type=int, default=4096) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iterations", type=int, default=10) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + group_name = dist.group.WORLD.group_name + rank = dist.get_rank() + + shape = (args.tokens, args.hidden_size) + # Keep the repeated in-place NCCL input finite without adding a reset copy to + # the timed region. Correctness of the transport setup is checked separately. + source = torch.zeros(shape, device="cuda", dtype=torch.bfloat16) + + nccl_input = source.clone() + dist.all_reduce(nccl_input) + torch.cuda.synchronize() + torch.testing.assert_close(nccl_input, torch.zeros_like(nccl_input), rtol=0, atol=0) + nccl_input.copy_(source) + torch.cuda.synchronize() + nccl_ms = elapsed_ms( + lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations + ) + nccl_graph_ms = graph_elapsed_ms( + lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations + ) + + buffer = symm_mem.empty(source.numel(), device="cuda", dtype=source.dtype) + handle = symm_mem.rendezvous(buffer, group_name) + if getattr(handle, "multicast_ptr", 0) == 0: + raise RuntimeError("symmetric-memory multicast pointer is unavailable") + symm_output = torch.empty_like(source) + + def symm_all_reduce() -> None: + buffer.copy_(source.view(-1)) + torch.ops.symm_mem.multimem_all_reduce_(buffer, "sum", group_name) + symm_output.view(-1).copy_(buffer) + + def symm_all_reduce_out_of_place() -> torch.Tensor: + buffer.copy_(source.view(-1)) + torch.ops.symm_mem.multimem_all_reduce_(buffer, "sum", group_name) + return buffer.view_as(source) + + symm_all_reduce() + torch.cuda.synchronize() + torch.testing.assert_close(symm_output, torch.zeros_like(symm_output), rtol=0, atol=0) + torch.cuda.synchronize() + symm_ms = elapsed_ms(symm_all_reduce, args.warmup, args.iterations) + symm_graph_ms = graph_elapsed_ms(symm_all_reduce, args.warmup, args.iterations) + symm_out_ms = elapsed_ms( + symm_all_reduce_out_of_place, args.warmup, args.iterations + ) + symm_out_graph_ms = graph_elapsed_ms( + symm_all_reduce_out_of_place, args.warmup, args.iterations + ) + + nccl_max_ms = max_rank(nccl_ms) + symm_max_ms = max_rank(symm_ms) + symm_out_max_ms = max_rank(symm_out_ms) + nccl_graph_max_ms = max_rank(nccl_graph_ms) + symm_graph_max_ms = max_rank(symm_graph_ms) + symm_out_graph_max_ms = max_rank(symm_out_graph_ms) + + cpu_group = dist.new_group( + list(range(dist.get_world_size())), backend="gloo" + ) + workspace = flashinfer_comm.create_allreduce_fusion_workspace( + backend="trtllm", + world_size=dist.get_world_size(), + rank=rank, + max_token_num=args.tokens, + hidden_dim=args.hidden_size, + dtype=source.dtype, + comm_backend=TorchDistBackend(group=cpu_group), + ) + + def flashinfer_all_reduce() -> torch.Tensor: + return flashinfer_comm.allreduce_fusion( + input=source, + workspace=workspace, + pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce, + ) + + fi_output = flashinfer_all_reduce() + torch.cuda.synchronize() + torch.testing.assert_close(fi_output, torch.zeros_like(fi_output), rtol=0, atol=0) + fi_ms = elapsed_ms(flashinfer_all_reduce, args.warmup, args.iterations) + fi_graph_ms = graph_elapsed_ms(flashinfer_all_reduce, args.warmup, args.iterations) + fi_max_ms = max_rank(fi_ms) + fi_graph_max_ms = max_rank(fi_graph_ms) + + if rank == 0: + nbytes = source.numel() * source.element_size() + print( + f"shape={shape} bytes={nbytes} nccl_ms={nccl_max_ms:.6f} " + f"symm_multimem_ms={symm_max_ms:.6f} symm_out_ms={symm_out_max_ms:.6f} " + f"flashinfer_ms={fi_max_ms:.6f} " + f"best={min((nccl_max_ms, 'nccl'), (symm_max_ms, 'symm'), (symm_out_max_ms, 'symm_out'), (fi_max_ms, 'flashinfer'))[1]} " + f"graph_nccl_ms={nccl_graph_max_ms:.6f} " + f"graph_symm_ms={symm_graph_max_ms:.6f} " + f"graph_symm_out_ms={symm_out_graph_max_ms:.6f} " + f"graph_flashinfer_ms={fi_graph_max_ms:.6f} " + f"graph_best={min((nccl_graph_max_ms, 'nccl'), (symm_graph_max_ms, 'symm'), (symm_out_graph_max_ms, 'symm_out'), (fi_graph_max_ms, 'flashinfer'))[1]}" + ) + workspace.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_kda_chunk_h.py b/tools/bench_glm53_kda_chunk_h.py new file mode 100644 index 0000000000..028a40fc1d --- /dev/null +++ b/tools/bench_glm53_kda_chunk_h.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Benchmark GLM-5.3 KDA chunk-state kernel configs on a real packed shape.""" + +from __future__ import annotations + +import argparse +import itertools +import json +import statistics + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--sequences", type=int, default=64) + parser.add_argument("--sequence-length", type=int, default=268) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--repeats", type=int, default=7) + return parser.parse_args() + + +@torch.inference_mode() +def main() -> None: + args = parse_args() + device = torch.device("cuda") + dtype = torch.bfloat16 + batch = 1 + heads = 8 + key_dim = 128 + value_dim = 128 + total_tokens = args.sequences * args.sequence_length + + shape = (batch, total_tokens, heads, key_dim) + k = torch.zeros(shape, device=device, dtype=dtype) + w = torch.zeros(shape, device=device, dtype=dtype) + u = torch.zeros((batch, total_tokens, heads, value_dim), device=device, dtype=dtype) + # The fused safe-gate+cumsum stage keeps cumulative decay in fp32; the + # exp2 path in the state kernel expects that exact dtype. + gk = torch.zeros(shape, device=device, dtype=torch.float32) + initial_state = torch.zeros( + (args.sequences, heads, key_dim, value_dim), device=device, dtype=dtype + ) + cu_seqlens = torch.arange( + 0, + total_tokens + 1, + args.sequence_length, + device=device, + dtype=torch.int32, + ) + chunk_indices = prepare_chunk_indices(cu_seqlens, 64) + + def run(config: dict[str, int]) -> None: + output = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=initial_state, + output_final_state=True, + save_new_value=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=64, + use_exp2=True, + run_config=config, + ) + del output + + results = [] + for value_tile, num_warps, num_stages in itertools.product( + (32, 64), (2, 4), (2, 3, 4) + ): + config = { + "BV": value_tile, + "num_warps": num_warps, + "num_stages": num_stages, + } + for _ in range(args.warmup): + run(config) + torch.cuda.synchronize() + + samples_ms = [] + for _ in range(args.repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run(config) + end.record() + end.synchronize() + samples_ms.append(start.elapsed_time(end)) + results.append( + { + **config, + "median_ms": statistics.median(samples_ms), + "min_ms": min(samples_ms), + } + ) + + results.sort(key=lambda item: item["median_ms"]) + print( + json.dumps( + { + "shape": { + "sequences": args.sequences, + "sequence_length": args.sequence_length, + "total_tokens": total_tokens, + "heads": heads, + "key_dim": key_dim, + "value_dim": value_dim, + }, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_sglang_moe.py b/tools/bench_glm53_sglang_moe.py new file mode 100644 index 0000000000..10cfc0f03b --- /dev/null +++ b/tools/bench_glm53_sglang_moe.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""Tune SGLang TritonMoE tiles for GLM-5.3-Flash TP8 on H100.""" + +import argparse +import json +import traceback +from contextlib import contextmanager +from types import SimpleNamespace + +import torch + + +SMALL_M_CONFIGS = [ + # Decode runs with 24 physical rows for c8 + MTP2. The original sweep + # barely sampled this regime, so cover the useful tile/scheduling axes. + (8, 32, 8, 4, 3), + (8, 64, 8, 4, 3), + (8, 128, 8, 4, 3), + (8, 64, 4, 4, 3), + (8, 64, 16, 4, 3), + (8, 64, 32, 4, 3), + (8, 128, 4, 4, 3), + (8, 128, 16, 4, 3), + (8, 128, 32, 4, 3), + (8, 64, 8, 4, 2), + (8, 64, 8, 4, 4), + (8, 128, 8, 4, 2), + (8, 128, 8, 4, 4), + (8, 64, 8, 8, 3), + (8, 128, 8, 8, 3), + (16, 32, 16, 4, 3), + (16, 64, 4, 4, 3), + (16, 64, 8, 4, 3), + (16, 64, 16, 4, 3), + (16, 64, 32, 4, 3), + (16, 128, 4, 4, 3), + (16, 128, 8, 4, 3), + (16, 128, 16, 4, 3), + (16, 128, 32, 4, 3), + (16, 64, 16, 4, 2), + (16, 64, 16, 4, 4), + (16, 128, 16, 4, 2), + (16, 128, 16, 4, 4), + (16, 64, 16, 8, 3), + (16, 128, 16, 8, 3), + (32, 64, 16, 4, 3), + (32, 128, 16, 4, 3), +] + + +LARGE_M_CONFIGS = [ + (64, 128, 32, 4, 3), + (64, 128, 8, 4, 3), + (64, 128, 16, 4, 3), + (64, 128, 64, 4, 3), + (32, 128, 16, 4, 3), + (128, 128, 32, 4, 3), + (128, 128, 32, 8, 3), + (64, 64, 32, 4, 3), + (64, 256, 32, 8, 3), + (128, 256, 32, 8, 3), + (64, 128, 32, 4, 2), + (64, 128, 32, 4, 4), +] + + +CONFIGS = SMALL_M_CONFIGS + LARGE_M_CONFIGS + + +def make_config(values): + block_m, block_n, group_m, num_warps, num_stages = values + return { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": group_m, + "num_warps": num_warps, + "num_stages": num_stages, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=17152) + parser.add_argument( + "--tp-size", + type=int, + default=8, + choices=(4, 8), + help="Tensor-parallel size; GLM-5 has 2048 total expert intermediate rows.", + ) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--max-configs", type=int, default=len(CONFIGS)) + parser.add_argument( + "--config-set", + choices=("all", "small", "large"), + default="all", + help="Restrict the sweep to decode-sized or prefill-sized tile configs.", + ) + parser.add_argument( + "--fuse-sum", + action="store_true", + help="Fuse the top-k expert sum into SGLang's down-projection kernel.", + ) + parser.add_argument( + "--tune-down", + action="store_true", + help="Tune a separate down-projection config with the best measured up config.", + ) + parser.add_argument( + "--fixed-up-config", + choices=("decode", "prefill"), + default="decode", + help="Fixed up-projection config used by --tune-down.", + ) + args = parser.parse_args() + + from sglang.srt.layers.moe.moe_runner.triton_utils import override_config + from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe + from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe_triton_config, + ) + + standalone_exec = SimpleNamespace( + moe=SimpleNamespace(enable_fused_moe_sum_all_reduce=args.fuse_sum), + deterministic=SimpleNamespace(enable_deterministic_inference=False), + ) + fused_moe.get_exec = lambda: standalone_exec + fused_moe_triton_config.get_exec = lambda: standalone_exec + + device = torch.device("cuda:0") + experts, hidden, intermediate, topk = 289, 4096, 2048 // args.tp_size, 9 + x = torch.zeros((args.tokens, hidden), dtype=torch.bfloat16, device=device) + w1 = torch.zeros( + (experts, intermediate * 2, hidden), dtype=torch.float8_e4m3fn, device=device + ) + w2 = torch.zeros( + (experts, hidden, intermediate), dtype=torch.float8_e4m3fn, device=device + ) + w1_scale = torch.ones((experts, 4, 32), dtype=torch.float32, device=device) + w2_scale = torch.ones((experts, 32, 2), dtype=torch.float32, device=device) + rows = torch.arange(args.tokens, dtype=torch.int64, device=device)[:, None] + cols = torch.arange(topk, dtype=torch.int64, device=device)[None, :] + topk_ids = (rows * topk + cols) % experts + topk_weights = torch.full( + (args.tokens, topk), 1.0 / topk, dtype=torch.float32, device=device + ) + + fixed_up_config = make_config( + { + "decode": (16, 64, 16, 4, 3), + "prefill": (64, 128, 64, 4, 3), + }[args.fixed_up_config] + ) + + @contextmanager + def config_context(config): + if not args.tune_down: + with override_config(config): + yield + return + original = fused_moe.try_get_optimal_moe_config + + def resolve_config(*resolve_args, return_down_config=False, **resolve_kwargs): + if return_down_config: + return fixed_up_config, (config, None) + return fixed_up_config + + fused_moe.try_get_optimal_moe_config = resolve_config + try: + yield + finally: + fused_moe.try_get_optimal_moe_config = original + + def run(config): + with config_context(config): + fused_moe.fused_experts_impl( + hidden_states=x, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_shape=[128, 128], + routed_scaling_factor=1.0, + filter_expert=False, + swiglu_limit=10.0, + gate_up_interleaved=False, + ) + + def component_times(config): + original = fused_moe.invoke_fused_moe_kernel + events = [] + + def timed_kernel(*kernel_args, **kernel_kwargs): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = original(*kernel_args, **kernel_kwargs) + end.record() + events.append((start, end)) + return result + + fused_moe.invoke_fused_moe_kernel = timed_kernel + try: + run(config) + torch.cuda.synchronize() + finally: + fused_moe.invoke_fused_moe_kernel = original + if len(events) != 2: + raise RuntimeError(f"expected two MoE GEMM calls, got {len(events)}") + return events[0][0].elapsed_time(events[0][1]), events[1][0].elapsed_time(events[1][1]) + + results = [] + selected_configs = { + "all": CONFIGS, + "small": SMALL_M_CONFIGS, + "large": LARGE_M_CONFIGS, + }[args.config_set] + if args.tune_down: + # Both GEMMs share the same token alignment, so the down projection's + # BLOCK_SIZE_M must match the fixed up projection. + selected_configs = [ + values + for values in CONFIGS + if values[0] == fixed_up_config["BLOCK_SIZE_M"] + ] + for values in selected_configs[: args.max_configs]: + config = make_config(values) + try: + for _ in range(args.warmup): + run(config) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(args.iters): + run(config) + end.record() + end.synchronize() + elapsed_ms = start.elapsed_time(end) / args.iters + up_ms, down_ms = component_times(config) + result = {"ms": elapsed_ms, "up_ms": up_ms, "down_ms": down_ms, **config} + if args.tune_down: + result["fixed_up_config"] = fixed_up_config + except Exception as exc: + result = {"error": f"{type(exc).__name__}: {exc}", **config} + if not results: + traceback.print_exc() + torch.cuda.synchronize() + results.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + + valid = [result for result in results if "ms" in result] + if valid: + print("BEST", json.dumps(min(valid, key=lambda result: result["ms"]), sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_sparse_prefill.py b/tools/bench_glm53_sparse_prefill.py new file mode 100644 index 0000000000..3e512abdb1 --- /dev/null +++ b/tools/bench_glm53_sparse_prefill.py @@ -0,0 +1,164 @@ +"""Compare GLM-5.3 sparse prefill kernels at serving-sized shapes. + +This benchmark deliberately includes FlashMLA's TP head-padding allocation and +copy, because that work is part of every LightLLM prefill attention layer. +TileLang receives the native eight TP-local query heads. +""" + +import argparse +import gc +import json +import statistics +from collections.abc import Callable + +import torch + + +def make_causal_indices(tokens: int, sequence_length: int, topk: int) -> torch.Tensor: + """Build the packed-request causal index layout used by the serving test.""" + + token_ids = torch.arange(tokens, dtype=torch.int32, device="cuda") + positions = token_ids.remainder(sequence_length) + sequence_starts = token_ids - positions + columns = torch.arange(topk, dtype=torch.int32, device="cuda").view(1, -1) + indices = sequence_starts.view(-1, 1) + columns + indices.masked_fill_(columns > positions.view(-1, 1), -1) + return indices.unsqueeze(1) + + +def measure_ms( + name: str, + operation: Callable[[], torch.Tensor], + warmup: int, + iterations: int, +) -> dict[str, float | str]: + result = None + for _ in range(warmup): + result = operation() + torch.cuda.synchronize() + del result + + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + baseline_bytes = torch.cuda.memory_allocated() + samples = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = operation() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + del result + + peak_delta_bytes = torch.cuda.max_memory_allocated() - baseline_bytes + return { + "name": name, + "median_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "peak_delta_gib": peak_delta_bytes / 2**30, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, required=True) + parser.add_argument("--sequence-length", type=int, default=1024) + parser.add_argument("--topk", type=int, default=2048) + parser.add_argument("--local-heads", type=int, default=8) + parser.add_argument("--required-heads", type=int, default=64) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=5) + args = parser.parse_args() + + if args.topk % 64: + raise ValueError("TileLang requires --topk to be divisible by 64") + if args.required_heads % args.local_heads: + raise ValueError("--required-heads must be divisible by --local-heads") + + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + + torch.manual_seed(0) + q = torch.randn( + (args.tokens, args.local_heads, args.head_dim), + dtype=torch.bfloat16, + device="cuda", + ) + kv = torch.randn( + (args.tokens, 1, args.head_dim), + dtype=torch.bfloat16, + device="cuda", + ) + indices = make_causal_indices(args.tokens, args.sequence_length, args.topk) + scale = args.head_dim**-0.5 + + padded_q = q.new_zeros((args.tokens, args.required_heads, args.head_dim)) + padded_q[:, : args.local_heads].copy_(q) + + def flashmla_kernel() -> torch.Tensor: + return flash_mla_sparse_fwd( + padded_q, + kv, + indices, + scale, + d_v=args.head_dim, + )[0] + + def flashmla_lightllm_path() -> torch.Tensor: + q_input = q.new_zeros((args.tokens, args.required_heads, args.head_dim)) + q_input[:, : args.local_heads].copy_(q) + return flash_mla_sparse_fwd( + q_input, + kv, + indices, + scale, + d_v=args.head_dim, + )[0][:, : args.local_heads] + + def tilelang_lightllm_path() -> torch.Tensor: + output = tilelang_sparse_fwd( + q, + kv, + indices, + scale, + d_v=args.head_dim, + ) + return output.squeeze(0) if output.ndim == 4 else output + + results = [] + for name, operation in ( + ("flashmla_padded_kernel", flashmla_kernel), + ("flashmla_lightllm_path", flashmla_lightllm_path), + ("tilelang_lightllm_path", tilelang_lightllm_path), + ): + results.append(measure_ms(name, operation, args.warmup, args.iterations)) + gc.collect() + torch.cuda.empty_cache() + + flash_ms = next(r["median_ms"] for r in results if r["name"] == "flashmla_lightllm_path") + tile_ms = next(r["median_ms"] for r in results if r["name"] == "tilelang_lightllm_path") + print( + json.dumps( + { + "shape": { + "tokens": args.tokens, + "sequence_length": args.sequence_length, + "topk": args.topk, + "local_heads": args.local_heads, + "required_heads": args.required_heads, + "head_dim": args.head_dim, + }, + "results": results, + "tilelang_speedup_over_lightllm_flashmla": flash_ms / tile_ms, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_sparse_prefill_tp.py b/tools/bench_glm53_sparse_prefill_tp.py new file mode 100644 index 0000000000..b9628b6564 --- /dev/null +++ b/tools/bench_glm53_sparse_prefill_tp.py @@ -0,0 +1,264 @@ +"""Prototype TP head/token transpose for GLM-5.3 sparse prefill attention.""" + +import argparse +import gc +import json +import statistics +from collections.abc import Callable + +import torch +import torch.distributed as dist + + +def make_causal_indices(tokens: int, sequence_length: int, topk: int) -> torch.Tensor: + token_ids = torch.arange(tokens, dtype=torch.int32, device="cuda") + positions = token_ids.remainder(sequence_length) + sequence_starts = token_ids - positions + columns = torch.arange(topk, dtype=torch.int32, device="cuda").view(1, -1) + indices = sequence_starts.view(-1, 1) + columns + indices.masked_fill_(columns > positions.view(-1, 1), -1) + return indices.unsqueeze(1) + + +def head_shards_to_token_shards(q: torch.Tensor, world_size: int) -> torch.Tensor: + tokens, local_heads, head_dim = q.shape + if tokens % world_size: + raise ValueError(f"tokens={tokens} must be divisible by world_size={world_size}") + tokens_per_rank = tokens // world_size + received = torch.empty_like(q) + dist.all_to_all_single(received, q) + return ( + received.view(world_size, tokens_per_rank, local_heads, head_dim) + .permute(1, 0, 2, 3) + .contiguous() + .view(tokens_per_rank, world_size * local_heads, head_dim) + ) + + +def token_shards_to_head_shards(output: torch.Tensor, world_size: int) -> torch.Tensor: + tokens_per_rank, global_heads, head_dim = output.shape + if global_heads % world_size: + raise ValueError(f"global_heads={global_heads} must be divisible by world_size={world_size}") + local_heads = global_heads // world_size + send = ( + output.view(tokens_per_rank, world_size, local_heads, head_dim) + .permute(1, 0, 2, 3) + .contiguous() + .view(world_size * tokens_per_rank, local_heads, head_dim) + ) + received = torch.empty_like(send) + dist.all_to_all_single(received, send) + return received + + +def global_max(value: float, device: torch.device) -> float: + tensor = torch.tensor(value, dtype=torch.float64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return tensor.item() + + +def measure_ms( + name: str, + operation: Callable[[], torch.Tensor], + warmup: int, + iterations: int, + device: torch.device, +) -> dict[str, float | str]: + result = None + for _ in range(warmup): + result = operation() + torch.cuda.synchronize() + del result + + gc.collect() + torch.cuda.empty_cache() + dist.barrier() + torch.cuda.reset_peak_memory_stats() + baseline_bytes = torch.cuda.memory_allocated() + samples = [] + for _ in range(iterations): + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = operation() + end.record() + end.synchronize() + samples.append(global_max(start.elapsed_time(end), device)) + del result + + peak_delta_gib = global_max( + (torch.cuda.max_memory_allocated() - baseline_bytes) / 2**30, + device, + ) + return { + "name": name, + "median_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "peak_delta_gib": peak_delta_gib, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, required=True) + parser.add_argument("--sequence-length", type=int, default=1024) + parser.add_argument("--topk", type=int, default=2048) + parser.add_argument("--local-heads", type=int, default=8) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--validate", action="store_true") + args = parser.parse_args() + + local_rank = int(__import__("os").environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.device("cuda", local_rank) + global_heads = args.local_heads * world_size + if args.tokens % world_size: + raise ValueError("--tokens must be divisible by the TP world size") + if global_heads != 64: + raise ValueError(f"FlashMLA on Hopper requires 64 global heads, got {global_heads}") + + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + torch.manual_seed(1000 + rank) + q = torch.randn( + (args.tokens, args.local_heads, args.head_dim), + dtype=torch.bfloat16, + device=device, + ) + torch.manual_seed(0) + kv = torch.randn( + (args.tokens, 1, args.head_dim), + dtype=torch.bfloat16, + device=device, + ) + indices = make_causal_indices(args.tokens, args.sequence_length, args.topk) + scale = args.head_dim**-0.5 + token_start = rank * (args.tokens // world_size) + token_end = token_start + args.tokens // world_size + local_indices = indices[token_start:token_end] + comm_workspace = torch.empty_like(q) + transpose_workspace = torch.empty( + (args.tokens // world_size, global_heads, args.head_dim), + dtype=q.dtype, + device=q.device, + ) + + def padded_flashmla() -> torch.Tensor: + q_input = q.new_zeros((args.tokens, global_heads, args.head_dim)) + q_input[:, : args.local_heads].copy_(q) + return flash_mla_sparse_fwd( + q_input, + kv, + indices, + scale, + d_v=args.head_dim, + )[0][:, : args.local_heads] + + def transposed_flashmla() -> torch.Tensor: + transposed_q = head_shards_to_token_shards(q, world_size) + transposed_output = flash_mla_sparse_fwd( + transposed_q, + kv, + local_indices, + scale, + d_v=args.head_dim, + )[0] + return token_shards_to_head_shards(transposed_output, world_size) + + def transposed_flashmla_workspace() -> torch.Tensor: + dist.all_to_all_single(comm_workspace, q) + comm_rank_major = comm_workspace.view( + world_size, + args.tokens // world_size, + args.local_heads, + args.head_dim, + ) + transpose_workspace.view( + args.tokens // world_size, + world_size, + args.local_heads, + args.head_dim, + ).copy_(comm_rank_major.permute(1, 0, 2, 3)) + transposed_output = flash_mla_sparse_fwd( + transpose_workspace, + kv, + local_indices, + scale, + d_v=args.head_dim, + )[0] + comm_workspace.view( + world_size, + args.tokens // world_size, + args.local_heads, + args.head_dim, + ).copy_( + transposed_output.view( + args.tokens // world_size, + world_size, + args.local_heads, + args.head_dim, + ).permute(1, 0, 2, 3) + ) + dist.all_to_all_single(transpose_workspace.view_as(q), comm_workspace) + return transpose_workspace.view_as(q) + + validation = None + if args.validate: + baseline = padded_flashmla().contiguous() + candidate = transposed_flashmla_workspace().clone() + torch.cuda.synchronize() + difference = (baseline - candidate).abs().float() + validation = { + "max_abs_diff": global_max(difference.max().item(), device), + "mean_abs_diff_max_rank": global_max(difference.mean().item(), device), + "allclose": bool(torch.allclose(baseline, candidate, rtol=0.05, atol=0.05)), + } + validation_tensor = torch.tensor(int(validation["allclose"]), device=device) + dist.all_reduce(validation_tensor, op=dist.ReduceOp.MIN) + validation["allclose"] = bool(validation_tensor.item()) + del baseline, candidate, difference + torch.cuda.empty_cache() + + results = [] + for name, operation in ( + ("flashmla_head_padding", padded_flashmla), + ("flashmla_tp_head_token_transpose", transposed_flashmla), + ("flashmla_tp_transpose_workspace", transposed_flashmla_workspace), + ): + results.append(measure_ms(name, operation, args.warmup, args.iterations, device)) + + if rank == 0: + baseline_ms = results[0]["median_ms"] + candidate_ms = results[-1]["median_ms"] + print( + json.dumps( + { + "shape": { + "tokens_per_rank_before_transpose": args.tokens, + "tokens_per_rank_after_transpose": args.tokens // world_size, + "sequence_length": args.sequence_length, + "topk": args.topk, + "local_heads": args.local_heads, + "global_heads": global_heads, + "head_dim": args.head_dim, + "world_size": world_size, + }, + "validation": validation, + "results": results, + "transpose_speedup_over_padding": baseline_ms / candidate_ms, + }, + indent=2, + ) + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/check_glm53_symm_out_of_place.py b/tools/check_glm53_symm_out_of_place.py new file mode 100644 index 0000000000..32443b5a12 --- /dev/null +++ b/tools/check_glm53_symm_out_of_place.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Exercise the production SymmMem alias path in normal and inference modes.""" + +import os +from contextlib import nullcontext + +import torch +import torch.distributed as dist + +from lightllm.distributed.symm_mem_all_reduce import SymmMemAllreduce + + +def _check(reducer: SymmMemAllreduce, *, inference: bool) -> None: + rank = dist.get_rank() + expected = sum(range(1, dist.get_world_size() + 1)) + context = torch.inference_mode() if inference else nullcontext() + with context: + for rows in (64, 32, 64): + value = torch.full( + (rows, 4096), + rank + 1, + dtype=torch.bfloat16, + device="cuda", + ) + value.data = reducer.all_reduce_out_of_place(value) + # Match the model's immediate post-reduction view operation. + viewed = value.view(rows, 4, 1024) + torch.testing.assert_close( + viewed, + torch.full_like(viewed, expected), + rtol=0, + atol=0, + ) + + +def main() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group( + "nccl", device_id=torch.device("cuda", local_rank) + ) + reducer = SymmMemAllreduce( + dist.group.WORLD, + torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + if reducer.disabled: + raise RuntimeError("SymmMemAllreduce unexpectedly disabled") + if not reducer.buffer.is_inference(): + raise RuntimeError("SymmMem workspace is not an inference tensor") + _check(reducer, inference=False) + _check(reducer, inference=True) + dist.barrier() + if dist.get_rank() == 0: + print("SymmMem out-of-place normal+inference alias checks passed") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py b/unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py new file mode 100644 index 0000000000..cfb5eabf92 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py @@ -0,0 +1,86 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.nsa import flashmla_sparse + + +def _disabled_modes(): + return SimpleNamespace( + enable_tpsp_mix_mode=False, + enable_prefill_cudagraph=False, + enable_prefill_microbatch_overlap=False, + enable_prefill_decode_mixed=False, + ) + + +def _infer_state(world_size=8): + return SimpleNamespace( + dist_group=SimpleNamespace(dp_world_size=world_size), + max_cache_len=0, + need_dp_prefill_balance=False, + use_replicated_attention_ep=False, + ) + + +def test_copy_received_head_shards_makes_token_major_global_heads(): + # Two source ranks sent the same destination's two-token chunk. + received = torch.tensor([[[10]], [[11]], [[20]], [[21]]]) + output = torch.empty((2, 2, 1), dtype=received.dtype) + + flashmla_sparse._copy_received_head_shards(received, output, world_size=2) + + torch.testing.assert_close(output, torch.tensor([[[10], [20]], [[11], [21]]])) + + +def test_copy_token_shard_for_head_scatter_makes_destination_blocks(): + output = torch.tensor([[[10], [20]], [[11], [21]]]) + send = torch.empty((4, 1, 1), dtype=output.dtype) + + flashmla_sparse._copy_token_shard_for_head_scatter(output, send, world_size=2) + + torch.testing.assert_close(send, torch.tensor([[[10]], [[11]], [[20]], [[21]]])) + + +def test_tp_head_token_transpose_accepts_validated_glm_tp8_layout(monkeypatch): + monkeypatch.setattr(flashmla_sparse, "get_env_start_args", _disabled_modes) + q = torch.empty((4096, 8, 2)) + + assert flashmla_sparse._should_use_tp_head_token_transpose(q, _infer_state(), required_heads=64) + + +@pytest.mark.parametrize( + ("tokens", "heads", "world_size", "state_change", "mode_change"), + [ + (2048, 8, 8, {}, {}), + (4097, 8, 8, {}, {}), + (4096, 16, 8, {}, {}), + (4096, 8, 1, {}, {}), + (4096, 8, 8, {"max_cache_len": 1}, {}), + (4096, 8, 8, {"need_dp_prefill_balance": True}, {}), + (4096, 8, 8, {"use_replicated_attention_ep": True}, {}), + (4096, 8, 8, {}, {"enable_tpsp_mix_mode": True}), + (4096, 8, 8, {}, {"enable_prefill_cudagraph": True}), + (4096, 8, 8, {}, {"enable_prefill_microbatch_overlap": True}), + (4096, 8, 8, {}, {"enable_prefill_decode_mixed": True}), + ], +) +def test_tp_head_token_transpose_rejects_unvalidated_layouts( + monkeypatch, + tokens, + heads, + world_size, + state_change, + mode_change, +): + modes = _disabled_modes() + for name, value in mode_change.items(): + setattr(modes, name, value) + monkeypatch.setattr(flashmla_sparse, "get_env_start_args", lambda: modes) + state = _infer_state(world_size) + for name, value in state_change.items(): + setattr(state, name, value) + q = torch.empty((tokens, heads, 2)) + + assert not flashmla_sparse._should_use_tp_head_token_transpose(q, state, required_heads=64) diff --git a/unit_tests/common/basemodel/attention/test_tilelang_sparse.py b/unit_tests/common/basemodel/attention/test_tilelang_sparse.py new file mode 100644 index 0000000000..448276703c --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_tilelang_sparse.py @@ -0,0 +1,27 @@ +import pytest +import torch + +from lightllm.common.basemodel.attention.nsa.tilelang_sparse import pad_sparse_indices + + +def test_pad_sparse_indices_adds_masked_block_tail(): + indices = torch.arange(65, dtype=torch.int32).view(1, 65) + + padded = pad_sparse_indices(indices) + + assert padded.shape == (1, 1, 128) + torch.testing.assert_close(padded[0, 0, :65], indices[0]) + assert torch.all(padded[0, 0, 65:] == -1) + + +def test_pad_sparse_indices_preserves_aligned_storage(): + indices = torch.zeros((4, 1, 128), dtype=torch.int32) + + assert pad_sparse_indices(indices) is indices + + +def test_pad_sparse_indices_rejects_invalid_shape_or_block(): + with pytest.raises(ValueError, match="2D or 3D"): + pad_sparse_indices(torch.zeros((2, 3, 4, 5), dtype=torch.int32)) + with pytest.raises(ValueError, match="positive"): + pad_sparse_indices(torch.zeros((2, 64), dtype=torch.int32), block_size=0) diff --git a/unit_tests/common/basemodel/test_cuda_graph_layout.py b/unit_tests/common/basemodel/test_cuda_graph_layout.py index 17ddcd98da..ce84bf5585 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_layout.py +++ b/unit_tests/common/basemodel/test_cuda_graph_layout.py @@ -76,3 +76,31 @@ def test_batch_step_size_after_split_controls_capture_range(_graph_args): 42, 56, ] + + +def test_extra_batch_sizes_are_merged_and_bounded(_graph_args): + graph = CudaGraph( + batch_step_size_before_split=6, + split_batch_size=24, + batch_step_size_after_split=96, + max_batch_size=384, + extra_batch_sizes=[1, 2, 4, 20, 36, 52, 64, 999], + ) + + assert graph.cuda_graph_batch_sizes == [ + 1, + 2, + 4, + 6, + 12, + 18, + 20, + 24, + 36, + 52, + 64, + 120, + 216, + 312, + 384, + ] diff --git a/unit_tests/common/basemodel/test_hidden_collector.py b/unit_tests/common/basemodel/test_hidden_collector.py index 22748b7294..6d9937d941 100644 --- a/unit_tests/common/basemodel/test_hidden_collector.py +++ b/unit_tests/common/basemodel/test_hidden_collector.py @@ -103,18 +103,22 @@ def test_noop_collector_keeps_normal_forward_output_minimal(): def test_mtp_head_output_collector_returns_and_clears_outputs(): collector = MtpHeadOutputCollector() draft_token_ids = torch.arange(6) + draft_token_probs = torch.linspace(0.1, 0.6, 6) confidence_logits = torch.arange(6).view(2, 3) collector.add_mtp_outputs( draft_token_ids=draft_token_ids, + draft_token_probs=draft_token_probs, confidence_logits=confidence_logits, ) output = collector.finish_output(infer_state=None) assert output.spec_hidden is None assert output.draft_token_ids is draft_token_ids + assert output.draft_token_probs is draft_token_probs assert output.confidence_logits is confidence_logits assert collector.draft_token_ids is None + assert collector.draft_token_probs is None assert collector.confidence_logits is None diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 6f9477e294..ca169f9396 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -11,13 +11,19 @@ def test_decode_unpad_slices_spec_output_with_logits(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(24).view(6, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), + mtp_collector=ModelMtpOutputCollector( + spec_hidden=torch.arange(18).view(6, 3), + draft_token_ids=torch.arange(6), + draft_token_probs=torch.linspace(0.1, 0.6, 6), + ), ) unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4) assert unpadded.logits.shape == (4, 4) assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) + assert unpadded.mtp_collector.draft_token_ids.shape == (4,) + assert unpadded.mtp_collector.draft_token_probs.shape == (4,) # Unpadding returns a shallow output copy and leaves the graph-owned # tensors on the original ModelOutput intact. assert output.logits.shape == (6, 4) @@ -118,6 +124,7 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): model = TpPartBaseModel.__new__(TpPartBaseModel) model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode) model.tp_world_size_ = tp_world_size + model.decode_batch_multiplier = 1 model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEX=99) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) diff --git a/unit_tests/common/basemodel/test_mtp_manager.py b/unit_tests/common/basemodel/test_mtp_manager.py index c7ad04f56b..c1aa7a8870 100644 --- a/unit_tests/common/basemodel/test_mtp_manager.py +++ b/unit_tests/common/basemodel/test_mtp_manager.py @@ -53,8 +53,8 @@ def test_decode_batch_multiplier(monkeypatch, spec_mode, is_draft_model, expecte [ (False, False, 8), (True, False, 1), - (False, True, 1), - (True, True, 1), + (False, True, 8), + (True, True, 8), ], ) def test_decode_cuda_graph_grow_step_size(monkeypatch, dynamic_verify, is_draft_model, expected): @@ -68,6 +68,40 @@ def test_decode_cuda_graph_grow_step_size(monkeypatch, dynamic_verify, is_draft_ assert MtpManager.get_instance().get_decode_cuda_graph_grow_step_size(is_draft_model) == expected +@pytest.mark.parametrize( + "spec_mode,is_draft_model,expected", + [ + ("eagle_with_att", True, 8), + ("eagle3", True, 8), + ("eagle_no_att", True, 1), + ("vanilla_with_att", True, 8), + ("eagle_with_att", False, 8), + ], +) +def test_decode_cuda_graph_batch_multiplier(monkeypatch, spec_mode, is_draft_model, expected): + args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + manager = MtpManager.get_instance() + assert manager.get_decode_cuda_graph_batch_multiplier(is_draft_model) == expected + + +@pytest.mark.parametrize( + "spec_mode,expected", + [ + ("eagle_with_att", True), + ("eagle3", True), + ("eagle_no_att", False), + ("vanilla_with_att", False), + ], +) +def test_recurrent_attention_draft_keeps_logical_graph_schedule(monkeypatch, spec_mode, expected): + args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + assert MtpManager.get_instance().draft_model_needs_logical_batch_graphs(True) is expected + + @pytest.mark.parametrize( "spec_mode,is_draft_model,expected", [ diff --git a/unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py b/unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py new file mode 100644 index 0000000000..efd1f384fe --- /dev/null +++ b/unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace + +import pytest + +import lightllm.common.basemodel.prefill_cuda_graph as prefill_cuda_graph + + +class _DecodeGraph: + mempool = object() + + +def _args(**overrides): + values = { + "enable_prefill_microbatch_overlap": False, + "prefill_cudagraph_max_handle_token": 32768, + "prefill_cudagraph_token_nums": None, + "prefill_cudagraph_batch_sizes": None, + "prefill_cudagraph_capture_attention": False, + "batch_max_tokens": 32768, + "enable_tpsp_mix_mode": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_explicit_token_nums_only_run_exact_shapes(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[17152, 352, 17152], + prefill_cudagraph_batch_sizes=[64, 1, 64], + ), + ) + + graph = prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + assert graph.graph_handle_token_nums == [352, 17152] + assert graph.can_run(352, batch_size=1, max_q_seq_len=352, max_kv_seq_len=352, max_cache_len=0) + assert graph.can_run(17152, batch_size=64, max_q_seq_len=268, max_kv_seq_len=268, max_cache_len=0) + assert not graph.can_run(17152, batch_size=1, max_q_seq_len=17152, max_kv_seq_len=17152, max_cache_len=0) + assert not graph.can_run(17152, batch_size=64, max_q_seq_len=269, max_kv_seq_len=269, max_cache_len=0) + assert not graph.can_run(17152, batch_size=64, max_q_seq_len=268, max_kv_seq_len=300, max_cache_len=32) + assert not graph.can_run(351, batch_size=1, max_q_seq_len=351, max_kv_seq_len=351, max_cache_len=0) + assert not graph.can_run(16000, batch_size=64, max_q_seq_len=250, max_kv_seq_len=250, max_cache_len=0) + + +def test_explicit_token_nums_respect_limits(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_max_handle_token=20000, + batch_max_tokens=18000, + prefill_cudagraph_token_nums=[0, 17152, 20000], + prefill_cudagraph_batch_sizes=[1, 64, 64], + ), + ) + + graph = prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + assert graph.graph_handle_token_nums == [17152] + + +def test_explicit_token_nums_reject_empty_valid_set(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[0, 40000], + prefill_cudagraph_batch_sizes=[1, 64], + ), + ) + + with pytest.raises(ValueError, match="prefill_cudagraph_token_nums"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + +def test_explicit_layout_rejects_missing_or_mismatched_batch_sizes(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args(prefill_cudagraph_token_nums=[17152]), + ) + with pytest.raises(ValueError, match="prefill_cudagraph_batch_sizes is required"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[17152, 352], + prefill_cudagraph_batch_sizes=[64], + ), + ) + with pytest.raises(ValueError, match="same number of entries"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + +def test_explicit_layout_rejects_nonuniform_sequence_shape(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[17153], + prefill_cudagraph_batch_sizes=[64], + ), + ) + with pytest.raises(ValueError, match="must be divisible"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + +def test_attention_capture_requires_exact_layout(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args(prefill_cudagraph_capture_attention=True), + ) + + with pytest.raises(ValueError, match="requires --prefill_cudagraph_token_nums"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) diff --git a/unit_tests/common/basemodel/test_sglang_triton_moe_config.py b/unit_tests/common/basemodel/test_sglang_triton_moe_config.py new file mode 100644 index 0000000000..64898b4156 --- /dev/null +++ b/unit_tests/common/basemodel/test_sglang_triton_moe_config.py @@ -0,0 +1,100 @@ +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import ( + triton_impl, +) + + +def test_glm5_decode_uses_measured_h100_moe_config(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + up_config, down_config = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + False, + 24, + ) + + assert up_config["BLOCK_SIZE_M"] == 16 + assert up_config["BLOCK_SIZE_N"] == 64 + assert up_config["num_stages"] == 3 + assert down_config["BLOCK_SIZE_M"] == up_config["BLOCK_SIZE_M"] + assert down_config["BLOCK_SIZE_N"] == 64 + assert down_config["num_stages"] == 2 + + +def test_glm5_draft_uses_measured_h100_down_config(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + up_config, down_config = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + False, + 8, + ) + + assert up_config["BLOCK_SIZE_M"] == 16 + assert down_config["GROUP_SIZE_M"] == 8 + assert down_config["num_stages"] == 3 + + +def test_glm5_prefill_uses_measured_h100_moe_config(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + (up_config, down_config) = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + True, + 8192, + ) + + assert up_config["BLOCK_SIZE_M"] == 64 + assert up_config["BLOCK_SIZE_N"] == 128 + assert up_config["GROUP_SIZE_M"] == 64 + assert down_config["BLOCK_SIZE_M"] == up_config["BLOCK_SIZE_M"] + assert down_config["BLOCK_SIZE_N"] == 128 + assert down_config["GROUP_SIZE_M"] == 8 + + +def test_glm5_large_prefill_uses_independent_down_tile(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + up_config, down_config = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + True, + 65536, + ) + + assert up_config["BLOCK_SIZE_N"] == 128 + assert up_config["GROUP_SIZE_M"] == 64 + assert down_config["BLOCK_SIZE_M"] == up_config["BLOCK_SIZE_M"] + assert down_config["BLOCK_SIZE_N"] == 64 + assert down_config["GROUP_SIZE_M"] == 32 + + +def test_non_h100_has_no_sglang_override(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H20") + + assert ( + triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + False, + 24, + ) + is None + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py new file mode 100644 index 0000000000..087dd2d7a4 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py @@ -0,0 +1,80 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import kda_safe_gate + + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + + +@pytest.mark.parametrize("sequence_length", [3, 6]) +def test_fused_kda_gate_matches_materialized_gate(sequence_length): + torch.manual_seed(2026 + sequence_length) + request_count, query_heads, value_heads, key_dim, value_dim = 3, 2, 4, 64, 64 + token_count = request_count * sequence_length + slot_count = token_count + 8 + + q = torch.randn(1, token_count, query_heads, key_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, token_count, value_heads, value_dim, device="cuda", dtype=torch.bfloat16) + raw_gate = torch.randn(1, token_count, value_heads * key_dim, device="cuda", dtype=torch.bfloat16) + raw_beta = torch.randn(1, token_count, value_heads, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(value_heads, device="cuda", dtype=torch.float32) * 0.1 + gate_bias = torch.randn(value_heads * key_dim, device="cuda", dtype=torch.float32) * 0.1 + state = torch.randn( + slot_count, + value_heads, + key_dim, + value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + state_indices = torch.arange(token_count, device="cuda", dtype=torch.int32).view( + request_count, sequence_length + ) + cu_seqlens = torch.arange(request_count + 1, device="cuda", dtype=torch.int64) * sequence_length + accepted = torch.full( + (request_count,), sequence_length, device="cuda", dtype=torch.int32 + ) + + reference_state = state.clone() + reference, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=kda_safe_gate(raw_gate, a_log, gate_bias), + beta=raw_beta.float().sigmoid(), + initial_state=reference_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + ssm_state_write_indices=state_indices, + num_accepted_tokens=accepted, + use_qk_l2norm_in_kernel=True, + is_kda=True, + ) + + fused_state = state.clone() + fused, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=fused_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + ssm_state_write_indices=state_indices, + num_accepted_tokens=accepted, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=gate_bias, + a_raw=raw_gate.reshape(token_count, value_heads * key_dim), + b_raw=raw_beta.reshape(token_count, value_heads), + is_kda=True, + kda_lower_bound=-5.0, + ) + + torch.testing.assert_close(fused, reference, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(fused_state, reference_state, rtol=2e-3, atol=2e-3) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py new file mode 100644 index 0000000000..279268a85c --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py @@ -0,0 +1,65 @@ +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") + + +@pytest.mark.parametrize("token_num", [1, 7, 64]) +def test_vocab_parallel_greedy_matches_full_logits(monkeypatch, token_num): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy") + tp_world_size = 4 + local_vocab_size = 8192 + vocab_size = tp_world_size * local_vocab_size + generator = torch.Generator(device="cuda").manual_seed(20260826 + token_num) + local_logits_by_rank = [ + torch.randn( + (local_vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + for _ in range(tp_world_size) + ] + + # Exercise deterministic tie-breaking both between local reduction blocks + # and across tensor-parallel ranks. The smallest global token id must win. + local_logits_by_rank[0][4097, 0] = 20.0 + local_logits_by_rank[0][3, 0] = 20.0 + local_logits_by_rank[3][2, 0] = 20.0 + + local_stats_by_rank = [ + greedy_sample_local_stats(local_logits.transpose(0, 1).contiguous()) for local_logits in local_logits_by_rank + ] + + def fake_all_gather_into_tensor(output_, input_, **_kwargs): + for output, local_stats in zip(output_, local_stats_by_rank): + output.copy_(local_stats) + + monkeypatch.setattr(module, "all_gather_into_tensor", fake_all_gather_into_tensor) + actual_logits, actual_ids, actual_logsumexp = module.vocab_parallel_greedy( + local_logits_by_rank[0], + vocab_size=vocab_size, + tp_world_size=tp_world_size, + group=None, + alloc_func=torch.empty, + ) + actual_ids = actual_ids.view(-1) + actual_logprobs = actual_logits.view(-1) - actual_logsumexp + + full_logits = torch.cat(local_logits_by_rank, dim=0).transpose(0, 1).float() + expected_ids = full_logits.argmax(dim=1) + expected_logits = full_logits.gather(1, expected_ids[:, None]).view(-1) + expected_logsumexp = torch.logsumexp(full_logits, dim=1) + expected_logprobs = torch.log_softmax(full_logits, dim=1).gather(1, expected_ids[:, None]).squeeze(1) + + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual_logits.view(-1), expected_logits, rtol=0, atol=0) + torch.testing.assert_close(actual_logsumexp, expected_logsumexp, rtol=2e-4, atol=2e-4) + torch.testing.assert_close(actual_logprobs, expected_logprobs, rtol=2e-4, atol=2e-4) diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py index 5e3aaf38d0..981c1dbb66 100644 --- a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py @@ -10,14 +10,24 @@ def test_trans_topk_index_to_mem_index(): # Create topk_index tensor with some valid indices and some -1 (padding) topk_index = torch.zeros((batch_size, topk), dtype=torch.int32, device="cuda") - topk_index[:, 0:2048] = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + topk_index[:, 0:2047] = torch.arange(0, 2047, dtype=torch.int32, device="cuda") + topk_index[:, -1] = -1 + ragged_start_index = torch.tensor([2], dtype=torch.int32, device="cuda") # Create ragged_mem_index lookup table - ragged_mem_index = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10 - - topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_mem_index) - - assert torch.equal(topk_mem_index, (torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10).view(1, -1)) + ragged_mem_index = torch.arange(0, 2050, dtype=torch.int32, device="cuda") + 10 + + topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_start_index, ragged_mem_index) + + expected_index = torch.cat( + ( + torch.arange(2, 2049, dtype=torch.int32, device="cuda"), + torch.tensor([-1], dtype=torch.int32, device="cuda"), + ) + ).view(1, -1) + expected_mem_index = torch.where(expected_index != -1, expected_index + 10, -1) + assert torch.equal(topk_index, expected_index) + assert torch.equal(topk_mem_index, expected_mem_index) if __name__ == "__main__": diff --git a/unit_tests/models/gemma4/test_post_layer_infer.py b/unit_tests/models/gemma4/test_post_layer_infer.py new file mode 100644 index 0000000000..1d6db329d3 --- /dev/null +++ b/unit_tests/models/gemma4/test_post_layer_infer.py @@ -0,0 +1,99 @@ +from types import SimpleNamespace + +import torch + +import lightllm.models.llama.layer_infer.post_layer_infer as llama_post_layer +from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer + + +def test_vocab_parallel_greedy_softcaps_local_logits_before_reduction(monkeypatch): + post_layer = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) + post_layer.final_logit_softcapping = 2.0 + post_layer.tp_world_size_ = 2 + post_layer.alloc_tensor = torch.empty + post_layer._norm = lambda hidden, infer_state, layer_weight: hidden + + local_logits = torch.tensor( + [ + [4.0, -4.0], + [2.0, -2.0], + [1.0, -1.0], + ], + dtype=torch.bfloat16, + ) + + class LMHead: + vocab_size = 6 + + def __call__(self, input, alloc_func): + return local_logits + + sparse_logits = torch.tensor([[5.0], [1.0]]) + token_ids = torch.tensor([[5], [1]]) + logsumexp = torch.tensor([5.25, 1.5]) + captured = {} + + def fake_vocab_parallel_greedy(logits, **kwargs): + captured["logits"] = logits + return sparse_logits, token_ids, logsumexp + + monkeypatch.setattr(llama_post_layer, "vocab_parallel_greedy", fake_vocab_parallel_greedy) + + infer_state = SimpleNamespace( + dist_group=None, + use_vocab_parallel_greedy=True, + logits_token_ids=None, + logits_logsumexp=None, + ) + + result = post_layer._lm_head_and_gather( + hidden=torch.empty((2, 3)), + token_num=2, + layer_weight=SimpleNamespace(lm_head_weight_=LMHead()), + infer_state=infer_state, + ) + + expected_logits = ( + torch.tanh(local_logits.float() / post_layer.final_logit_softcapping) * post_layer.final_logit_softcapping + ) + assert captured["logits"].dtype == torch.float32 + torch.testing.assert_close(captured["logits"], expected_logits) + assert result is sparse_logits + assert infer_state.logits_token_ids is token_ids + assert infer_state.logits_logsumexp is logsumexp + + +def test_full_logits_softcap_after_float32_conversion(monkeypatch): + post_layer = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) + post_layer.final_logit_softcapping = 2.0 + post_layer.tp_world_size_ = 1 + post_layer.alloc_tensor = torch.empty + post_layer._norm = lambda hidden, infer_state, layer_weight: hidden + + local_logits = torch.tensor( + [ + [1.234375, -3.140625], + [2.71875, -0.333984375], + [0.10009765625, -1.609375], + ], + dtype=torch.bfloat16, + ) + + class LMHead: + vocab_size = 3 + + def __call__(self, input, alloc_func): + return local_logits + + result = post_layer._lm_head_and_gather( + hidden=torch.empty((2, 3), dtype=torch.bfloat16), + token_num=2, + layer_weight=SimpleNamespace(lm_head_weight_=LMHead()), + infer_state=SimpleNamespace(dist_group=None, use_vocab_parallel_greedy=True), + force_full_logits=True, + ) + + expected = torch.tanh(local_logits.T.float() / post_layer.final_logit_softcapping) + expected *= post_layer.final_logit_softcapping + assert result.dtype == torch.float32 + torch.testing.assert_close(result, expected) diff --git a/unit_tests/models/test_vocab_parallel_greedy_output.py b/unit_tests/models/test_vocab_parallel_greedy_output.py new file mode 100644 index 0000000000..88e5a7c282 --- /dev/null +++ b/unit_tests/models/test_vocab_parallel_greedy_output.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + + +def test_argmax_restores_global_token_ids_and_exact_probabilities(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput( + logits=torch.tensor([[3.0, 1.0], [0.0, 5.0]]), + logits_token_ids=torch.tensor([[30, 10], [100, 500]]), + logits_logsumexp=torch.tensor([4.0, 5.25]), + ) + + token_ids = backend._gen_argmax_token_ids(output) + token_ids_with_prob, probs = backend._gen_argmax_token_ids_and_prob(output) + + torch.testing.assert_close(token_ids, torch.tensor([30, 500])) + torch.testing.assert_close(token_ids_with_prob, token_ids) + torch.testing.assert_close(probs, torch.exp(torch.tensor([-1.0, -0.25]))) + + +def test_dense_argmax_keeps_column_index_semantics(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput(logits=torch.tensor([[1.0, 4.0, 2.0]])) + + torch.testing.assert_close(backend._gen_argmax_token_ids(output), torch.tensor([1])) + + +def test_dspark_confidence_path_receives_global_token_ids(): + post = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post.block_size_ = 2 + post.markov_rank_ = 0 + post._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) + sparse_logits = torch.tensor([[4.0], [5.0], [7.0], [9.0]]) + sparse_token_ids = torch.tensor([[40], [50], [70], [90]]) + + def gather_vocab_parallel(*args, **kwargs): + infer_state = args[3] + infer_state.logits_token_ids = sparse_token_ids + return sparse_logits + + post._lm_head_and_gather = gather_vocab_parallel + observed = {} + + def predict_confidence(block_hidden, anchor_token_ids, sampled_tokens, layer_weight): + observed["sampled_tokens"] = sampled_tokens + return None + + post.predict_confidence_logits = predict_confidence + + class Collector: + def add_mtp_outputs(self, **kwargs): + self.outputs = kwargs + + collector = Collector() + infer_state = SimpleNamespace( + is_prefill=False, + input_ids=torch.tensor([1, 0, 2, 0]), + logits_token_ids=None, + hidden_collector=collector, + ) + + returned_logits = post.token_forward( + input_embdings=torch.ones((4, 3)), + infer_state=infer_state, + layer_weight=object(), + ) + + torch.testing.assert_close(returned_logits, sparse_logits) + torch.testing.assert_close(observed["sampled_tokens"], torch.tensor([[40, 50], [70, 90]])) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py b/unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py new file mode 100644 index 0000000000..7d85ad21c5 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py @@ -0,0 +1,64 @@ +from lightllm.server.router.model_infer.mode_backend.chunked_prefill import ( + impl as chunked_prefill_impl, +) +from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ( + ChunkedPrefillBackend, +) + + +class _FakeEvent: + def __init__(self, calls=None): + self.recorded = False + self.synchronized = False + self.calls = calls + + def record(self): + self.recorded = True + + def synchronize(self): + self.synchronized = True + if self.calls is not None: + self.calls.append("synchronize") + + +class _FakeEventPack: + def __init__(self, calls): + self.calls = calls + + def notify_forward_and_wait_post_handle(self): + self.calls.append("notify_forward") + + +def test_record_forward_completion_keeps_cpu_bookkeeping_async(monkeypatch): + event = _FakeEvent() + monkeypatch.setattr(chunked_prefill_impl.torch.cuda, "Event", lambda: event) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend._serialize_sm90_mega_moe_forwards = True + + assert backend._record_forward_completion() is event + assert event.recorded + assert not event.synchronized + + +def test_mega_moe_waits_before_notifying_next_forward(): + calls = [] + event = _FakeEvent(calls) + event_pack = _FakeEventPack(calls) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend._serialize_sm90_mega_moe_forwards = True + + backend._notify_next_forward_when_safe(event_pack, event) + + assert calls == ["synchronize", "notify_forward"] + + +def test_regular_path_notifies_next_forward_before_waiting(): + calls = [] + event = _FakeEvent(calls) + event_pack = _FakeEventPack(calls) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend._serialize_sm90_mega_moe_forwards = False + + backend._notify_next_forward_when_safe(event_pack, event) + + assert calls == ["notify_forward", "synchronize"] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py new file mode 100644 index 0000000000..fddcda9380 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py @@ -0,0 +1,91 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.utils.envs_utils import enable_env_vars +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + _can_use_unmodified_greedy_logits, + can_use_vocab_parallel_greedy, + sample, +) + + +def make_req(**overrides): + values = { + "top_k": 1, + "temperature": 1.0, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "repetition_penalty": 1.0, + "min_new_tokens": 1, + "decay_factor": 1.0, + "invalid_token_ids": [], + "output_len": 0, + } + values.update(overrides) + shm_param = SimpleNamespace( + top_k=values["top_k"], + temperature=values["temperature"], + presence_penalty=values["presence_penalty"], + frequency_penalty=values["frequency_penalty"], + repetition_penalty=values["repetition_penalty"], + min_new_tokens=values["min_new_tokens"], + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (1, values["decay_factor"])), + ) + input_len = 10 + return SimpleNamespace( + sampling_param=SimpleNamespace( + shm_param=shm_param, + invalid_token_ids=values["invalid_token_ids"], + ), + shm_req=SimpleNamespace(input_len=input_len), + get_cur_total_len=lambda: input_len + values["output_len"], + ) + + +def test_accepts_unmodified_greedy_requests(): + assert _can_use_unmodified_greedy_logits([make_req(), make_req(output_len=5)]) + + +def test_feature_gate_requires_environment_and_eligible_batch(monkeypatch): + monkeypatch.delenv("LIGHTLLM_VOCAB_PARALLEL_GREEDY", raising=False) + enable_env_vars.cache_clear() + assert not can_use_vocab_parallel_greedy([make_req()]) + + monkeypatch.setenv("LIGHTLLM_VOCAB_PARALLEL_GREEDY", "1") + enable_env_vars.cache_clear() + assert can_use_vocab_parallel_greedy([make_req()]) + assert not can_use_vocab_parallel_greedy([make_req(top_k=2)]) + enable_env_vars.cache_clear() + + +def test_samples_explicit_vocab_parallel_output_exactly(): + model_output = ModelOutput( + logits=torch.tensor([[9.0], [4.0]], dtype=torch.float32), + logits_token_ids=torch.tensor([[17], [3]], dtype=torch.int64), + logits_logsumexp=torch.tensor([9.25, 5.5], dtype=torch.float32), + ) + + token_ids, token_logprobs = sample(model_output, [make_req(), make_req()]) + + torch.testing.assert_close(token_ids, torch.tensor([17, 3])) + torch.testing.assert_close(token_logprobs, torch.tensor([-0.25, -1.5])) + + +@pytest.mark.parametrize( + "override", + [ + {"top_k": 2}, + {"temperature": 0.5}, + {"presence_penalty": 0.1}, + {"frequency_penalty": 0.1}, + {"repetition_penalty": 1.1}, + {"decay_factor": 1.1}, + {"invalid_token_ids": [7]}, + {"min_new_tokens": 2}, + ], +) +def test_rejects_logits_modifiers(override): + assert not _can_use_unmodified_greedy_logits([make_req(**override)]) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py index 7baf34061b..9b6995a6b6 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -739,7 +739,7 @@ def test_lightspec_eagle_draft_always_keeps_the_extend_candidate(): assert plan.draft_step >= 1 -def test_vanilla_with_attention_planner_prices_extend_then_normal_batches(): +def test_vanilla_with_attention_planner_prices_every_chained_level_at_verify_batch(): planner = build_lightspec_planner(spec_mode="vanilla_with_att") planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.1) @@ -749,7 +749,7 @@ def test_vanilla_with_attention_planner_prices_extend_then_normal_batches(): draft_step=3, ) - assert np.isclose(draft_cost_ms, 0.8) + assert np.isclose(draft_cost_ms, 1.2) with pytest.raises(AssertionError, match="requires draft_step to be greater than 0"): planner._get_draft_cost_ms(req_num=4, verify_batch_size=8, draft_step=0) diff --git a/unit_tests/server/router/test_dp_model_capacity.py b/unit_tests/server/router/test_dp_model_capacity.py new file mode 100644 index 0000000000..236e368372 --- /dev/null +++ b/unit_tests/server/router/test_dp_model_capacity.py @@ -0,0 +1,39 @@ +from types import SimpleNamespace + +import pytest + +from lightllm.server.router.manager import resolve_model_max_req_num + + +def test_model_request_capacity_defaults_to_global_running_limit(): + args = SimpleNamespace( + running_max_req_size=256, + per_dp_max_req_size=None, + dp=8, + nnodes=1, + ) + + assert resolve_model_max_req_num(args) == 256 + + +def test_model_request_capacity_accepts_balanced_per_dp_limit(): + args = SimpleNamespace( + running_max_req_size=256, + per_dp_max_req_size=40, + dp=8, + nnodes=1, + ) + + assert resolve_model_max_req_num(args) == 40 + + +def test_model_request_capacity_rejects_less_than_balanced_share(): + args = SimpleNamespace( + running_max_req_size=256, + per_dp_max_req_size=31, + dp=8, + nnodes=1, + ) + + with pytest.raises(ValueError, match="minimum 32"): + resolve_model_max_req_num(args) diff --git a/unit_tests/server/router/test_prefill_coalescing.py b/unit_tests/server/router/test_prefill_coalescing.py new file mode 100644 index 0000000000..0c0c1bec4a --- /dev/null +++ b/unit_tests/server/router/test_prefill_coalescing.py @@ -0,0 +1,69 @@ +from types import SimpleNamespace + +from lightllm.server.router import manager as router_manager +from lightllm.server.router.manager import RouterManager + + +class _ReqQueue: + def __init__(self, waiting_req_num): + self.waiting_req_num = waiting_req_num + + def get_wait_req_num(self): + return self.waiting_req_num + + +def _make_router(waiting_req_num, interval=0.5, running_req_num=0, pending_req_num=0): + router = RouterManager.__new__(RouterManager) + router.args = SimpleNamespace(running_max_req_size=64) + router.prefill_coalesce_interval = interval + router._prefill_coalesce_deadline = None + router.req_queue = _ReqQueue(waiting_req_num) + router.running_batch = None if running_req_num == 0 else SimpleNamespace(reqs=[None] * running_req_num) + router.schedule_new_batch = None if pending_req_num == 0 else SimpleNamespace(reqs=[None] * pending_req_num) + return router + + +def test_prefill_coalescing_is_disabled_by_default(): + router = _make_router(waiting_req_num=1, interval=0.0) + + assert not router._should_defer_prefill_batch() + + +def test_partial_burst_waits_only_until_original_deadline(monkeypatch): + now = [10.0] + monkeypatch.setattr(router_manager.time, "monotonic", lambda: now[0]) + router = _make_router(waiting_req_num=12) + + assert router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline == 10.5 + + now[0] = 10.49 + router.req_queue.waiting_req_num = 48 + assert router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline == 10.5 + + now[0] = 10.5 + assert not router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline is None + + +def test_full_runnable_batch_skips_coalescing(monkeypatch): + monkeypatch.setattr(router_manager.time, "monotonic", lambda: 10.0) + router = _make_router(waiting_req_num=48, running_req_num=16) + + assert not router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline is None + + +def test_expired_burst_launches_as_soon_as_running_batch_clears(monkeypatch): + now = [10.0] + monkeypatch.setattr(router_manager.time, "monotonic", lambda: now[0]) + router = _make_router(waiting_req_num=40, running_req_num=64) + + assert router._should_defer_prefill_batch() + now[0] = 11.0 + assert router._should_defer_prefill_batch() + + router.running_batch = None + assert not router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline is None diff --git a/unit_tests/server/test_mtp_start_args.py b/unit_tests/server/test_mtp_start_args.py index 87e12479d4..0e5504bc98 100644 --- a/unit_tests/server/test_mtp_start_args.py +++ b/unit_tests/server/test_mtp_start_args.py @@ -4,11 +4,11 @@ from lightllm.server.core.objs.start_args_type import StartArgs -def test_mtp_requires_cuda_graph(monkeypatch): +def test_dynamic_mtp_requires_cuda_graph(monkeypatch): monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - args = StartArgs(mtp_mode="vanilla_no_att", disable_cudagraph=True) + args = StartArgs(mtp_mode="vanilla_no_att", mtp_dynamic_verify=True, disable_cudagraph=True) - with pytest.raises(AssertionError, match="only supported on Prefill nodes"): + with pytest.raises(AssertionError, match="--mtp_dynamic_verify"): _launch_subprocesses(args) diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py new file mode 100644 index 0000000000..55e95bc3d5 --- /dev/null +++ b/unit_tests/utils/test_envs_utils.py @@ -0,0 +1,74 @@ +from easydict import EasyDict + +from lightllm.utils import envs_utils + + +def _set_start_args(monkeypatch, **kwargs): + args = EasyDict( + graph_max_batch_size=kwargs.get("graph_max_batch_size", 256), + running_max_req_size=kwargs.get("running_max_req_size", 256), + mtp_mode=kwargs.get("mtp_mode"), + mtp_step=kwargs.get("mtp_step", 0), + enable_tpsp_mix_mode=kwargs.get("enable_tpsp_mix_mode", False), + enable_ep_moe=kwargs.get("enable_ep_moe", False), + tp=kwargs.get("tp", 1), + ) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + +def test_deepep_decode_limit_covers_speculative_physical_batch(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + _set_start_args( + monkeypatch, + graph_max_batch_size=64, + running_max_req_size=64, + mtp_mode="eagle_with_att", + mtp_step=5, + ) + + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 384 + + +def test_deepep_decode_limit_uses_sequence_parallel_local_batch(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + _set_start_args( + monkeypatch, + graph_max_batch_size=64, + running_max_req_size=64, + mtp_mode="eagle_with_att", + mtp_step=5, + enable_tpsp_mix_mode=True, + enable_ep_moe=True, + tp=8, + ) + + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 256 + + +def test_deepep_decode_limit_is_aligned_and_keeps_minimum(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + _set_start_args( + monkeypatch, + graph_max_batch_size=63, + running_max_req_size=32, + mtp_mode="eagle_with_att", + mtp_step=4, + ) + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 320 + + _set_start_args(monkeypatch, graph_max_batch_size=8, running_max_req_size=8) + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 256 + + +def test_deepep_decode_limit_honors_explicit_override(monkeypatch): + monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", "512") + _set_start_args( + monkeypatch, + graph_max_batch_size=64, + running_max_req_size=64, + mtp_mode="eagle_with_att", + mtp_step=5, + ) + + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 512 From 2b1fbcb2337c70ea486e693d1784b143a5643a00 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 19:10:46 +0800 Subject: [PATCH 02/28] build: package GLM-5.3 H100 runtime image --- .dockerignore | 46 +++++++++++++++ GLM53_H100_DEPLOY.md | 59 ++++++++++++++++++++ docker/Dockerfile.glm53-h100 | 80 +++++++++++++++++++++++++++ docker/requirements-glm53-runtime.txt | 2 + tools/run_glm53_h100_container.sh | 26 +++++++++ 5 files changed, 213 insertions(+) create mode 100644 .dockerignore create mode 100644 GLM53_H100_DEPLOY.md create mode 100644 docker/Dockerfile.glm53-h100 create mode 100644 docker/requirements-glm53-runtime.txt create mode 100755 tools/run_glm53_h100_container.sh diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000..40e0ac5d5c --- /dev/null +++ b/.dockerignore @@ -0,0 +1,46 @@ +.git +.git/** +.github +.idea +.vscode +.pytest_cache +.mypy_cache +.ruff_cache +**/__pycache__ +**/*.pyc +**/*.pyo + +# Local environments, credentials, and host configuration. +.env +.env.* +.aws +.sco +.netrc +*.key +*.pem +id_rsa* +venv +.venv + +# Models, caches, traces, and experiment output never belong in an image. +models +checkpoints +cache* +.cache +traces +trace-* +*.log +*.jsonl +*.pt +*.pth +*.safetensors + +# Runtime images only need the package and release Docker inputs. +assets +demos +docs +skills +test +unit_tests +tools +format_out diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md new file mode 100644 index 0000000000..4af5d1e784 --- /dev/null +++ b/GLM53_H100_DEPLOY.md @@ -0,0 +1,59 @@ +# GLM-5.3-Flash on H100 TP8 + +This branch includes a self-contained LightLLM image variant for one eight-GPU +H100 node. The image contains the LightLLM source and runtime dependency; only +the model and compiler-cache directories are mounted from the host. + +## Build + +Use an immutable tag containing the full source revision: + +```bash +revision="$(git rev-parse HEAD)" +created="$(date -u +%Y-%m-%dT%H:%M:%SZ)" +version="v1.2.0-h100-tp8-${revision:0:12}" + +docker buildx build --load --platform linux/amd64 \ + -f docker/Dockerfile.glm53-h100 \ + --build-arg "OCI_CREATED=${created}" \ + --build-arg "OCI_REVISION=${revision}" \ + --build-arg "OCI_VERSION=${version}" \ + -t "lightllm-glm53:${version}" \ + . +``` + +After verification, the host-local convenience alias may point to the same +image ID: + +```bash +docker tag "lightllm-glm53:${version}" lightllm-glm53:h100-tp8 +``` + +## Run on h100 + +The default image command is the measured no-speculation, concurrency-256 +profile on port 8002. Run it in the foreground with: + +```bash +LIGHTLLM_GLM53_IMAGE="lightllm-glm53:${version}" \ + tools/run_glm53_h100_container.sh +``` + +Override `LIGHTLLM_GLM53_MODEL_DIR`, `LIGHTLLM_GLM53_CACHE_DIR`, or +`LIGHTLLM_GLM53_TRITON_CACHE_DIR` when the host paths differ. In another shell, +wait for the model list endpoint: + +```bash +curl --fail --show-error http://127.0.0.1:8002/v1/models +``` + +Stop the foreground process with `Ctrl-C`. If it was detached externally, use +`sudo docker stop --timeout 30 glm53-lightllm`. + +## Measured profile + +The final pre-release candidate reached 4169.40 output tokens/s at concurrency +256 with random 1024-token inputs and 256-token outputs. This is 5.74% below +the measured vLLM result, so the earlier three-percent stretch goal remains +unmet at concurrency 256. Concurrency 16 and 64 exceeded the corresponding +vLLM and SGLang measurements when run with MTP2. diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 new file mode 100644 index 0000000000..072e20f1e2 --- /dev/null +++ b/docker/Dockerfile.glm53-h100 @@ -0,0 +1,80 @@ +# syntax=docker/dockerfile:1.7 + +# This is the exact SGLang GLM-5.3 runtime used for the H100 measurements. +ARG BASE_IMAGE=lmsysorg/sglang@sha256:e6f5482505e7502f791fe4615ad1fbec118cbbd6b44e98f2479b16b98b985ad6 +FROM ${BASE_IMAGE} + +ARG OCI_CREATED +ARG OCI_REVISION +ARG OCI_SOURCE=https://github.com/sufubao/LightLLM +ARG OCI_VERSION +ARG BASE_NAME=lmsysorg/sglang:glm-5.3-flash +ARG BASE_DIGEST=sha256:e6f5482505e7502f791fe4615ad1fbec118cbbd6b44e98f2479b16b98b985ad6 + +LABEL org.opencontainers.image.created="${OCI_CREATED}" \ + org.opencontainers.image.revision="${OCI_REVISION}" \ + org.opencontainers.image.source="${OCI_SOURCE}" \ + org.opencontainers.image.version="${OCI_VERSION}" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100 TP8" \ + org.opencontainers.image.base.name="${BASE_NAME}" \ + org.opencontainers.image.base.digest="${BASE_DIGEST}" \ + ai.lightllm.model="GLM-5.3-Flash" \ + ai.lightllm.accelerator="NVIDIA H100 80GB" \ + ai.lightllm.tensor-parallel-size="8" \ + ai.lightllm.profile="throughput-c256" + +ENV PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + LOADWORKER=18 \ + LIGHTLLM_USE_SGLANG_TRITON_MOE=1 \ + LIGHTLLM_SGLANG_FUSED_MOE_SUM=1 \ + LIGHTLLM_VOCAB_PARALLEL_GREEDY=1 \ + LIGHTLLM_SYMM_MEM_OUT_OF_PLACE=1 \ + LIGHTLLM_LOG_LEVEL=warning \ + LIGHTLLM_ENABLE_FAST_MTP_KDA=1 \ + NCCL_CUMEM_ENABLE=1 \ + NCCL_NVLS_ENABLE=1 \ + CUDA_DEVICE_MAX_CONNECTIONS=8 \ + TRTLLM_ENABLE_PDL=1 \ + NCCL_GRAPH_MIXING_SUPPORT=0 + +WORKDIR /opt/lightllm + +COPY docker/requirements-glm53-runtime.txt /tmp/requirements-glm53-runtime.txt +RUN python -m pip install --no-cache-dir --no-deps --require-hashes \ + -r /tmp/requirements-glm53-runtime.txt && \ + rm /tmp/requirements-glm53-runtime.txt + +COPY setup.py LICENSE README.md ./ +COPY lightllm ./lightllm +RUN python -m pip install --no-cache-dir --no-deps . && \ + python -c "import lightllm.server.api_server; print('LightLLM GLM-5.3 runtime import OK')" + +EXPOSE 8002 +STOPSIGNAL SIGTERM + +HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ + CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 + +CMD ["python", "-m", "lightllm.server.api_server", \ + "--model_dir", "/model", \ + "--model_name", "glm-5.3-flash", \ + "--tp", "8", \ + "--host", "0.0.0.0", \ + "--port", "8002", \ + "--httpserver_workers", "16", \ + "--mem_fraction", ".90", \ + "--max_total_token_num", "335000", \ + "--running_max_req_size", "256", \ + "--max_req_total_len", "65500", \ + "--batch_max_tokens", "65536", \ + "--chunked_prefill_size", "65536", \ + "--linear_att_ssm_data_type", "bfloat16", \ + "--graph_max_batch_size", "256", \ + "--graph_split_batch_size", "4", \ + "--graph_grow_step_size", "16", \ + "--enable_fused_shared_experts", \ + "--schedule_time_interval", "0.001", \ + "--prefill_coalesce_interval", "0.5", \ + "--reasoning_parser", "glm45", \ + "--tool_call_parser", "glm47"] diff --git a/docker/requirements-glm53-runtime.txt b/docker/requirements-glm53-runtime.txt new file mode 100644 index 0000000000..c8318222e8 --- /dev/null +++ b/docker/requirements-glm53-runtime.txt @@ -0,0 +1,2 @@ +atomics==1.0.3 \ + --hash=sha256:19ed27f1ae3fe3353e9103e6d6f54af83e56a929f78ba061ae2e8e9435101daa diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh new file mode 100755 index 0000000000..e29ddefb05 --- /dev/null +++ b/tools/run_glm53_h100_container.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +set -euo pipefail + +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:h100-tp8}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm}" +model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" +cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" +triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" + +if [[ ! -d "${model_dir}" ]]; then + echo "model directory does not exist: ${model_dir}" >&2 + exit 1 +fi + +mkdir -p "${cache_dir}" "${triton_cache_dir}" + +exec sudo docker run --rm --name "${name}" \ + --gpus all \ + --ipc host \ + --network host \ + --ulimit memlock=-1:-1 \ + --ulimit nofile=1048576:1048576 \ + -v "${model_dir}:/model:ro" \ + -v "${cache_dir}:/root/.cache" \ + -v "${triton_cache_dir}:/root/.triton" \ + "${image}" From 3fbcbd507f0c2942e0737a7ae7555b7ac3269c6b Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 19:17:08 +0800 Subject: [PATCH 03/28] ci: publish signed GLM-5.3 H100 image --- .github/workflows/glm53-h100-image.yml | 98 ++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 .github/workflows/glm53-h100-image.yml diff --git a/.github/workflows/glm53-h100-image.yml b/.github/workflows/glm53-h100-image.yml new file mode 100644 index 0000000000..38c6561844 --- /dev/null +++ b/.github/workflows/glm53-h100-image.yml @@ -0,0 +1,98 @@ +name: GLM-5.3 H100 image + +on: + push: + branches: + - support-glm-5-3-flash + paths: + - .github/workflows/glm53-h100-image.yml + - .dockerignore + - docker/Dockerfile.glm53-h100 + - docker/requirements-glm53-runtime.txt + - lightllm/** + - setup.py + workflow_dispatch: + +permissions: + contents: read + packages: write + id-token: write + +env: + REGISTRY: ghcr.io + IMAGE_NAME: ${{ github.repository }} + +jobs: + build-publish-sign: + runs-on: ubuntu-latest + + steps: + - name: Free disk space + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be + with: + tool-cache: true + android: true + dotnet: true + haskell: true + large-packages: true + swap-storage: false + docker-images: false + + - name: Check out source + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 + + - name: Resolve immutable image identity + id: identity + shell: bash + run: | + version="v1.2.0-h100-tp8-${GITHUB_SHA}" + created="$(git show -s --format=%cI HEAD)" + echo "version=${version}" >> "${GITHUB_OUTPUT}" + echo "created=${created}" >> "${GITHUB_OUTPUT}" + echo "image=${REGISTRY}/${IMAGE_NAME}:${version}" >> "${GITHUB_OUTPUT}" + + - name: Set up Buildx + uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 + + - name: Log in to GHCR + uses: docker/login-action@184bdaa0721073962dff0199f1fb9940f07167d1 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and publish image + id: build + uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 + with: + context: . + file: docker/Dockerfile.glm53-h100 + platforms: linux/amd64 + push: true + provenance: mode=max + sbom: true + tags: ${{ steps.identity.outputs.image }} + build-args: | + OCI_CREATED=${{ steps.identity.outputs.created }} + OCI_REVISION=${{ github.sha }} + OCI_VERSION=${{ steps.identity.outputs.version }} + cache-from: type=gha,scope=glm53-h100 + cache-to: type=gha,mode=max,scope=glm53-h100 + + - name: Install cosign + uses: sigstore/cosign-installer@d58896d6a1865668819e1d91763c7751a165e159 + + - name: Sign published digest + env: + IMAGE: ${{ steps.identity.outputs.image }} + DIGEST: ${{ steps.build.outputs.digest }} + run: cosign sign --yes "${IMAGE%@*}@${DIGEST}" + + - name: Verify signature + env: + IMAGE: ${{ steps.identity.outputs.image }} + DIGEST: ${{ steps.build.outputs.digest }} + run: | + cosign verify "${IMAGE%@*}@${DIGEST}" \ + --certificate-identity-regexp "^https://github.com/${GITHUB_REPOSITORY}/.github/workflows/glm53-h100-image.yml@refs/heads/support-glm-5-3-flash$" \ + --certificate-oidc-issuer "https://token.actions.githubusercontent.com" From ec79db8c2630f72866fa4298f7e2b268355937d0 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 19:48:04 +0800 Subject: [PATCH 04/28] build: harden GLM-5.3 runtime image --- .github/workflows/glm53-h100-image.yml | 15 ++++++++ docker/Dockerfile.glm53-h100 | 47 ++++++++++++++++++------- docker/glm53-h100.openvex.json | 48 ++++++++++++++++++++++++++ docker/requirements-glm53-runtime.txt | 6 ++++ 4 files changed, 104 insertions(+), 12 deletions(-) create mode 100644 docker/glm53-h100.openvex.json diff --git a/.github/workflows/glm53-h100-image.yml b/.github/workflows/glm53-h100-image.yml index 38c6561844..0891d1d319 100644 --- a/.github/workflows/glm53-h100-image.yml +++ b/.github/workflows/glm53-h100-image.yml @@ -8,6 +8,7 @@ on: - .github/workflows/glm53-h100-image.yml - .dockerignore - docker/Dockerfile.glm53-h100 + - docker/glm53-h100.openvex.json - docker/requirements-glm53-runtime.txt - lightllm/** - setup.py @@ -88,6 +89,16 @@ jobs: DIGEST: ${{ steps.build.outputs.digest }} run: cosign sign --yes "${IMAGE%@*}@${DIGEST}" + - name: Attach security VEX + env: + IMAGE: ${{ steps.identity.outputs.image }} + DIGEST: ${{ steps.build.outputs.digest }} + run: | + cosign attest --yes \ + --predicate docker/glm53-h100.openvex.json \ + --type openvex \ + "${IMAGE%@*}@${DIGEST}" + - name: Verify signature env: IMAGE: ${{ steps.identity.outputs.image }} @@ -96,3 +107,7 @@ jobs: cosign verify "${IMAGE%@*}@${DIGEST}" \ --certificate-identity-regexp "^https://github.com/${GITHUB_REPOSITORY}/.github/workflows/glm53-h100-image.yml@refs/heads/support-glm-5-3-flash$" \ --certificate-oidc-issuer "https://token.actions.githubusercontent.com" + cosign verify-attestation "${IMAGE%@*}@${DIGEST}" \ + --type openvex \ + --certificate-identity-regexp "^https://github.com/${GITHUB_REPOSITORY}/.github/workflows/glm53-h100-image.yml@refs/heads/support-glm-5-3-flash$" \ + --certificate-oidc-issuer "https://token.actions.githubusercontent.com" diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index 072e20f1e2..a0add96fad 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -2,7 +2,36 @@ # This is the exact SGLang GLM-5.3 runtime used for the H100 measurements. ARG BASE_IMAGE=lmsysorg/sglang@sha256:e6f5482505e7502f791fe4615ad1fbec118cbbd6b44e98f2479b16b98b985ad6 -FROM ${BASE_IMAGE} +FROM ${BASE_IMAGE} AS prepared + +WORKDIR /opt/lightllm + +COPY docker/requirements-glm53-runtime.txt /tmp/requirements-glm53-runtime.txt +RUN python -m pip install --no-cache-dir --no-deps --require-hashes \ + -r /tmp/requirements-glm53-runtime.txt && \ + rm /tmp/requirements-glm53-runtime.txt + +COPY setup.py LICENSE README.md ./ +COPY lightllm ./lightllm +COPY docker/glm53-h100.openvex.json /usr/share/doc/lightllm/glm53-h100.openvex.json +RUN python -m pip install --no-cache-dir --no-deps . && \ + rm -rf \ + /etc/ssh/ssh_host_*_key* \ + /opt/nvidia/nsight-compute/2025.3.1/host/target-linux-x64/plugins/efa_metrics \ + /opt/nvidia/nsight-systems-cli/2026.4.1/target-linux-x64/plugins/efa_metrics \ + /root/.cargo/registry \ + /sgl-workspace/sglang/python/sglang/multimodal_gen \ + /sgl-workspace/sglang/python/sglang/srt/disaggregation \ + /sgl-workspace/sglang/scripts/playground/replay_request_dump.py && \ + python -c "import lightllm.server.api_server; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ + test ! -e /etc/ssh/ssh_host_rsa_key && \ + test ! -e /sgl-workspace/sglang/python/sglang/multimodal_gen && \ + test ! -e /sgl-workspace/sglang/python/sglang/srt/disaggregation + +# Flatten the prepared root filesystem so deleted base-image host keys and +# developer-only vulnerable components are absent from the published layers, +# rather than merely hidden by an OCI whiteout in a later layer. +FROM scratch AS runtime ARG OCI_CREATED ARG OCI_REVISION @@ -11,6 +40,8 @@ ARG OCI_VERSION ARG BASE_NAME=lmsysorg/sglang:glm-5.3-flash ARG BASE_DIGEST=sha256:e6f5482505e7502f791fe4615ad1fbec118cbbd6b44e98f2479b16b98b985ad6 +COPY --from=prepared / / + LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ org.opencontainers.image.source="${OCI_SOURCE}" \ @@ -21,7 +52,8 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100 80GB" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="throughput-c256" + ai.lightllm.profile="throughput-c256" \ + ai.lightllm.security-profile="flattened-no-sglang-server-components" ENV PYTHONUNBUFFERED=1 \ PYTHONDONTWRITEBYTECODE=1 \ @@ -40,18 +72,9 @@ ENV PYTHONUNBUFFERED=1 \ WORKDIR /opt/lightllm -COPY docker/requirements-glm53-runtime.txt /tmp/requirements-glm53-runtime.txt -RUN python -m pip install --no-cache-dir --no-deps --require-hashes \ - -r /tmp/requirements-glm53-runtime.txt && \ - rm /tmp/requirements-glm53-runtime.txt - -COPY setup.py LICENSE README.md ./ -COPY lightllm ./lightllm -RUN python -m pip install --no-cache-dir --no-deps . && \ - python -c "import lightllm.server.api_server; print('LightLLM GLM-5.3 runtime import OK')" - EXPOSE 8002 STOPSIGNAL SIGTERM +ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 diff --git a/docker/glm53-h100.openvex.json b/docker/glm53-h100.openvex.json new file mode 100644 index 0000000000..90a83eebd1 --- /dev/null +++ b/docker/glm53-h100.openvex.json @@ -0,0 +1,48 @@ +{ + "@context": "https://openvex.dev/ns/v0.2.0", + "@id": "https://github.com/sufubao/LightLLM/security/vex/glm53-h100-sglang-unused-components", + "author": "https://github.com/sufubao/LightLLM", + "timestamp": "2026-08-28T11:45:00Z", + "version": 1, + "statements": [ + { + "vulnerability": { + "name": "CVE-2026-3059" + }, + "products": [ + { + "@id": "pkg:pypi/sglang@0.0.0.dev1%2Bgf609d677b" + } + ], + "status": "not_affected", + "justification": "vulnerable_code_not_present", + "impact_statement": "The image removes sglang/multimodal_gen; LightLLM only imports SGLang's Triton MoE kernels." + }, + { + "vulnerability": { + "name": "CVE-2026-3060" + }, + "products": [ + { + "@id": "pkg:pypi/sglang@0.0.0.dev1%2Bgf609d677b" + } + ], + "status": "not_affected", + "justification": "vulnerable_code_not_present", + "impact_statement": "The image removes sglang/srt/disaggregation; LightLLM does not expose SGLang's disaggregation service." + }, + { + "vulnerability": { + "name": "CVE-2026-3989" + }, + "products": [ + { + "@id": "pkg:pypi/sglang@0.0.0.dev1%2Bgf609d677b" + } + ], + "status": "not_affected", + "justification": "vulnerable_code_not_present", + "impact_statement": "The image removes SGLang's replay_request_dump.py developer utility." + } + ] +} diff --git a/docker/requirements-glm53-runtime.txt b/docker/requirements-glm53-runtime.txt index c8318222e8..b65a9e33d5 100644 --- a/docker/requirements-glm53-runtime.txt +++ b/docker/requirements-glm53-runtime.txt @@ -1,2 +1,8 @@ atomics==1.0.3 \ --hash=sha256:19ed27f1ae3fe3353e9103e6d6f54af83e56a929f78ba061ae2e8e9435101daa +diffusers==0.38.0 \ + --hash=sha256:18e53f9e539096320470f62c6360a6fd5727ff28cffda566265316e13fcdb612 +msgpack==1.2.1 \ + --hash=sha256:020e881a764b20d8d7ca1a54fc01b8175519d108e3c3f194fddc200bda95951a +setuptools==78.1.1 \ + --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 From 2b688850e0074c675425b827062a0d88583578fd Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 20:29:51 +0800 Subject: [PATCH 05/28] build: remove vulnerable packaging tools --- docker/Dockerfile.glm53-h100 | 13 +++++++++++-- docker/requirements-glm53-runtime.txt | 8 ++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index a0add96fad..6d88204e32 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -22,8 +22,17 @@ RUN python -m pip install --no-cache-dir --no-deps . && \ /root/.cargo/registry \ /sgl-workspace/sglang/python/sglang/multimodal_gen \ /sgl-workspace/sglang/python/sglang/srt/disaggregation \ - /sgl-workspace/sglang/scripts/playground/replay_request_dump.py && \ - python -c "import lightllm.server.api_server; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ + /sgl-workspace/sglang/scripts/playground/replay_request_dump.py \ + /opt/sglang/bin/pip \ + /opt/sglang/bin/pip3 \ + /opt/sglang/bin/pip3.12 \ + /opt/sglang/lib/python3.12/site-packages/pip \ + /opt/sglang/lib/python3.12/site-packages/pip-*.dist-info \ + /opt/sglang/lib/python3.12/site-packages/setuptools \ + /opt/sglang/lib/python3.12/site-packages/setuptools-*.dist-info \ + /opt/sglang/lib/python3.12/site-packages/wheel \ + /opt/sglang/lib/python3.12/site-packages/wheel-*.dist-info && \ + python -c "import importlib.util, lightllm.server.api_server, msgpack; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; assert msgpack.__version__ == '1.2.2'; assert importlib.util.find_spec('pip') is None; assert importlib.util.find_spec('setuptools') is None; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ test ! -e /etc/ssh/ssh_host_rsa_key && \ test ! -e /sgl-workspace/sglang/python/sglang/multimodal_gen && \ test ! -e /sgl-workspace/sglang/python/sglang/srt/disaggregation diff --git a/docker/requirements-glm53-runtime.txt b/docker/requirements-glm53-runtime.txt index b65a9e33d5..682612f92d 100644 --- a/docker/requirements-glm53-runtime.txt +++ b/docker/requirements-glm53-runtime.txt @@ -2,7 +2,7 @@ atomics==1.0.3 \ --hash=sha256:19ed27f1ae3fe3353e9103e6d6f54af83e56a929f78ba061ae2e8e9435101daa diffusers==0.38.0 \ --hash=sha256:18e53f9e539096320470f62c6360a6fd5727ff28cffda566265316e13fcdb612 -msgpack==1.2.1 \ - --hash=sha256:020e881a764b20d8d7ca1a54fc01b8175519d108e3c3f194fddc200bda95951a -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 +msgpack==1.2.2 \ + --hash=sha256:77c2e018417dc1d66f235e383877ee885b60ade9d29e494dd581e08af2cb1923 +setuptools==84.0.0 \ + --hash=sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670 From 624d3813931260abcd978ff177c9ba1cdf36c89f Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 21:09:12 +0800 Subject: [PATCH 06/28] build: restore flattened runtime environment --- docker/Dockerfile.glm53-h100 | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index 6d88204e32..f717cea838 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -64,7 +64,19 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ ai.lightllm.profile="throughput-c256" \ ai.lightllm.security-profile="flattened-no-sglang-server-components" -ENV PYTHONUNBUFFERED=1 \ +ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ + LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/local/cuda/lib64 \ + LIBRARY_PATH=/usr/local/cuda/lib64/stubs \ + CUDA_HOME=/usr/local/cuda \ + CUDA_VERSION=13.0.3 \ + NVARCH=x86_64 \ + NVIDIA_REQUIRE_CUDA="cuda>=13.0" \ + NVIDIA_VISIBLE_DEVICES=all \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility \ + LANG=en_US.UTF-8 \ + LANGUAGE=en_US:en \ + LC_ALL=en_US.UTF-8 \ + PYTHONUNBUFFERED=1 \ PYTHONDONTWRITEBYTECODE=1 \ LOADWORKER=18 \ LIGHTLLM_USE_SGLANG_TRITON_MOE=1 \ @@ -81,14 +93,17 @@ ENV PYTHONUNBUFFERED=1 \ WORKDIR /opt/lightllm +RUN test "$(command -v python)" = /opt/sglang/bin/python && \ + python -c "import lightllm.server.api_server, torch; print(torch.__version__)" + EXPOSE 8002 STOPSIGNAL SIGTERM ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ - CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 + CMD /opt/sglang/bin/python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 -CMD ["python", "-m", "lightllm.server.api_server", \ +CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ "--tp", "8", \ From 8e3dc81374f28435d95bd8d64ae849bb5d8df937 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 21:35:12 +0800 Subject: [PATCH 07/28] build: close GLM-5.3 runtime dependencies --- docker/Dockerfile.glm53-h100 | 4 ++-- docker/requirements-glm53-runtime.txt | 20 +++++++++++++++++++ .../models/qwen3_omni_moe_thinker/model.py | 1 - lightllm/server/multimodal_params.py | 3 ++- 4 files changed, 24 insertions(+), 4 deletions(-) diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index f717cea838..f5d2a744e8 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -32,7 +32,7 @@ RUN python -m pip install --no-cache-dir --no-deps . && \ /opt/sglang/lib/python3.12/site-packages/setuptools-*.dist-info \ /opt/sglang/lib/python3.12/site-packages/wheel \ /opt/sglang/lib/python3.12/site-packages/wheel-*.dist-info && \ - python -c "import importlib.util, lightllm.server.api_server, msgpack; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; assert msgpack.__version__ == '1.2.2'; assert importlib.util.find_spec('pip') is None; assert importlib.util.find_spec('setuptools') is None; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ + python -c "import frozendict, hypercorn, importlib.metadata as metadata, importlib.util, lightllm.server.api_start, msgpack, rpyc, ujson; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; assert metadata.version('hypercorn') == '0.18.0'; assert msgpack.__version__ == '1.2.2'; assert rpyc.__version__ == '6.0.2'; assert importlib.util.find_spec('pip') is None; assert importlib.util.find_spec('setuptools') is None; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ test ! -e /etc/ssh/ssh_host_rsa_key && \ test ! -e /sgl-workspace/sglang/python/sglang/multimodal_gen && \ test ! -e /sgl-workspace/sglang/python/sglang/srt/disaggregation @@ -94,7 +94,7 @@ ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sb WORKDIR /opt/lightllm RUN test "$(command -v python)" = /opt/sglang/bin/python && \ - python -c "import lightllm.server.api_server, torch; print(torch.__version__)" + python -c "import hypercorn, importlib.metadata as metadata, lightllm.server.api_start, rpyc, torch; assert metadata.version('hypercorn') == '0.18.0'; assert rpyc.__version__ == '6.0.2'; print(torch.__version__)" EXPOSE 8002 STOPSIGNAL SIGTERM diff --git a/docker/requirements-glm53-runtime.txt b/docker/requirements-glm53-runtime.txt index 682612f92d..8a0e41751e 100644 --- a/docker/requirements-glm53-runtime.txt +++ b/docker/requirements-glm53-runtime.txt @@ -2,7 +2,27 @@ atomics==1.0.3 \ --hash=sha256:19ed27f1ae3fe3353e9103e6d6f54af83e56a929f78ba061ae2e8e9435101daa diffusers==0.38.0 \ --hash=sha256:18e53f9e539096320470f62c6360a6fd5727ff28cffda566265316e13fcdb612 +frozendict==2.4.7 \ + --hash=sha256:972af65924ea25cf5b4d9326d549e69a9a4918d8a76a9d3a7cd174d98b237550 +h2==4.4.1 \ + --hash=sha256:0e25f1462b23c9cb82d9eb02e28bc706dac2a68cb457c6a0d74d63c8a2a5d0e6 +hpack==4.2.0 \ + --hash=sha256:858ac0b02280fa582b5080d68db0899c62a80375e0e5413a74970c5e518b6986 +hypercorn==0.18.0 \ + --hash=sha256:225e268f2c1c2f28f6d8f6db8f40cb8c992963610c5725e13ccfcddccb24b1cd +hyperframe==6.1.0 \ + --hash=sha256:b03380493a519fce58ea5af42e4a42317bf9bd425596f7a0835ffce80f1a42e5 msgpack==1.2.2 \ --hash=sha256:77c2e018417dc1d66f235e383877ee885b60ade9d29e494dd581e08af2cb1923 +plumbum==2.0.2 \ + --hash=sha256:a865771826c1d2cce0ca8b2a2c48aa0447ed7a518859e389a95f158e9c369547 +priority==2.0.0 \ + --hash=sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa +rpyc==6.0.2 \ + --hash=sha256:8072308ad30725bc281c42c011fc8c922be15f3eeda6eafb2917cafe1b6f00ec setuptools==84.0.0 \ --hash=sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670 +ujson==5.13.0 \ + --hash=sha256:ea7204e9fa7538bfbb1396e1ee8c2bbcd3818b3633ef5bb14d4fdea52994d14d +wsproto==1.3.2 \ + --hash=sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584 diff --git a/lightllm/models/qwen3_omni_moe_thinker/model.py b/lightllm/models/qwen3_omni_moe_thinker/model.py index e4ea3ccc34..32b959675e 100644 --- a/lightllm/models/qwen3_omni_moe_thinker/model.py +++ b/lightllm/models/qwen3_omni_moe_thinker/model.py @@ -1,6 +1,5 @@ import os import json -import librosa import copy from functools import lru_cache from io import BytesIO diff --git a/lightllm/server/multimodal_params.py b/lightllm/server/multimodal_params.py index ed3535a69f..6a417fd1d1 100644 --- a/lightllm/server/multimodal_params.py +++ b/lightllm/server/multimodal_params.py @@ -1,7 +1,6 @@ """Multimodal parameters for text generation.""" import asyncio import os -import librosa import base64 import numpy as np from typing import List, Tuple, Optional @@ -54,6 +53,8 @@ async def preload(self, request: Request): raise ValueError(f"cannot read audio which type is {self._type}!") # check if valid audio bytes + import librosa + audio_values, _ = await asyncio.to_thread(librosa.load, BytesIO(audio_data), sr=16000) audio_values = np.asarray(audio_values, dtype=np.float32) From 013da1d186d382d9fc6f6d59f9318f1e11c820f1 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 22:25:30 +0800 Subject: [PATCH 08/28] ci: normalize GHCR image name --- .github/workflows/glm53-h100-image.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/glm53-h100-image.yml b/.github/workflows/glm53-h100-image.yml index 0891d1d319..4c89d4cf08 100644 --- a/.github/workflows/glm53-h100-image.yml +++ b/.github/workflows/glm53-h100-image.yml @@ -21,7 +21,6 @@ permissions: env: REGISTRY: ghcr.io - IMAGE_NAME: ${{ github.repository }} jobs: build-publish-sign: @@ -46,11 +45,12 @@ jobs: id: identity shell: bash run: | + repository="${GITHUB_REPOSITORY,,}" version="v1.2.0-h100-tp8-${GITHUB_SHA}" created="$(git show -s --format=%cI HEAD)" echo "version=${version}" >> "${GITHUB_OUTPUT}" echo "created=${created}" >> "${GITHUB_OUTPUT}" - echo "image=${REGISTRY}/${IMAGE_NAME}:${version}" >> "${GITHUB_OUTPUT}" + echo "image=${REGISTRY}/${repository}:${version}" >> "${GITHUB_OUTPUT}" - name: Set up Buildx uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 From 9151a7c31c48752cff02a86d6ed3dd1c89f26c87 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 28 Aug 2026 22:44:39 +0800 Subject: [PATCH 09/28] ci: remove GHCR publishing --- .github/workflows/glm53-h100-image.yml | 113 ------------------------- 1 file changed, 113 deletions(-) delete mode 100644 .github/workflows/glm53-h100-image.yml diff --git a/.github/workflows/glm53-h100-image.yml b/.github/workflows/glm53-h100-image.yml deleted file mode 100644 index 4c89d4cf08..0000000000 --- a/.github/workflows/glm53-h100-image.yml +++ /dev/null @@ -1,113 +0,0 @@ -name: GLM-5.3 H100 image - -on: - push: - branches: - - support-glm-5-3-flash - paths: - - .github/workflows/glm53-h100-image.yml - - .dockerignore - - docker/Dockerfile.glm53-h100 - - docker/glm53-h100.openvex.json - - docker/requirements-glm53-runtime.txt - - lightllm/** - - setup.py - workflow_dispatch: - -permissions: - contents: read - packages: write - id-token: write - -env: - REGISTRY: ghcr.io - -jobs: - build-publish-sign: - runs-on: ubuntu-latest - - steps: - - name: Free disk space - uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be - with: - tool-cache: true - android: true - dotnet: true - haskell: true - large-packages: true - swap-storage: false - docker-images: false - - - name: Check out source - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 - - - name: Resolve immutable image identity - id: identity - shell: bash - run: | - repository="${GITHUB_REPOSITORY,,}" - version="v1.2.0-h100-tp8-${GITHUB_SHA}" - created="$(git show -s --format=%cI HEAD)" - echo "version=${version}" >> "${GITHUB_OUTPUT}" - echo "created=${created}" >> "${GITHUB_OUTPUT}" - echo "image=${REGISTRY}/${repository}:${version}" >> "${GITHUB_OUTPUT}" - - - name: Set up Buildx - uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 - - - name: Log in to GHCR - uses: docker/login-action@184bdaa0721073962dff0199f1fb9940f07167d1 - with: - registry: ${{ env.REGISTRY }} - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Build and publish image - id: build - uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 - with: - context: . - file: docker/Dockerfile.glm53-h100 - platforms: linux/amd64 - push: true - provenance: mode=max - sbom: true - tags: ${{ steps.identity.outputs.image }} - build-args: | - OCI_CREATED=${{ steps.identity.outputs.created }} - OCI_REVISION=${{ github.sha }} - OCI_VERSION=${{ steps.identity.outputs.version }} - cache-from: type=gha,scope=glm53-h100 - cache-to: type=gha,mode=max,scope=glm53-h100 - - - name: Install cosign - uses: sigstore/cosign-installer@d58896d6a1865668819e1d91763c7751a165e159 - - - name: Sign published digest - env: - IMAGE: ${{ steps.identity.outputs.image }} - DIGEST: ${{ steps.build.outputs.digest }} - run: cosign sign --yes "${IMAGE%@*}@${DIGEST}" - - - name: Attach security VEX - env: - IMAGE: ${{ steps.identity.outputs.image }} - DIGEST: ${{ steps.build.outputs.digest }} - run: | - cosign attest --yes \ - --predicate docker/glm53-h100.openvex.json \ - --type openvex \ - "${IMAGE%@*}@${DIGEST}" - - - name: Verify signature - env: - IMAGE: ${{ steps.identity.outputs.image }} - DIGEST: ${{ steps.build.outputs.digest }} - run: | - cosign verify "${IMAGE%@*}@${DIGEST}" \ - --certificate-identity-regexp "^https://github.com/${GITHUB_REPOSITORY}/.github/workflows/glm53-h100-image.yml@refs/heads/support-glm-5-3-flash$" \ - --certificate-oidc-issuer "https://token.actions.githubusercontent.com" - cosign verify-attestation "${IMAGE%@*}@${DIGEST}" \ - --type openvex \ - --certificate-identity-regexp "^https://github.com/${GITHUB_REPOSITORY}/.github/workflows/glm53-h100-image.yml@refs/heads/support-glm-5-3-flash$" \ - --certificate-oidc-issuer "https://token.actions.githubusercontent.com" From aa7cab9ad12452fb08bfc316461242640b1b976c Mon Sep 17 00:00:00 2001 From: sufubao <47234901+sufubao@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:54:56 +0800 Subject: [PATCH 10/28] feat(pd): expose stage and node load metrics (#1509) --- .../httpserver_for_pd_master/manager.py | 60 +++++++-- lightllm/server/metrics/manager.py | 5 +- lightllm/server/metrics/metrics.py | 16 ++- .../test_pd_master_metrics.py | 124 ++++++++++++++++++ 4 files changed, 188 insertions(+), 17 deletions(-) create mode 100644 test/test_pd_selector/test_pd_master_metrics.py diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index b422bf7703..db1d923fe3 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -41,13 +41,16 @@ def __init__( self.metric_client = MetricClient(get_shm_port_args().metric_port) self.id_gen = ReqIDGenerator() - self.pd_manager = PDManager(args) + self.pd_manager = PDManager(args, self.metric_client) self.req_id_to_out_inf: Dict[int, ReqStatus] = {} self.infos_queues = None # 这个需要延迟初始化，否则使用的loop不对 self.health_timeout = int(os.getenv("HEALTH_TIMEOUT", "200")) self.latest_success_infer_time = time.time() self.running_request_count = 0 + self.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} + for stage in self.pd_stage_waiting_request_counts: + self.metric_client.gauge_set("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}) self.tokenizer = get_tokenizer(args.model_dir, args.tokenizer_mode, trust_remote_code=args.trust_remote_code) @@ -359,6 +362,28 @@ async def raise_if_disconnected() -> None: except asyncio.TimeoutError: continue + def _change_pd_stage_waiting_requests(self, stage: str, delta: int) -> None: + self.pd_stage_waiting_request_counts[stage] += delta + self.metric_client.gauge_set( + "lightllm_pd_master_stage_waiting_requests", + self.pd_stage_waiting_request_counts[stage], + labels={"stage": stage}, + ) + + async def _wait_for_pd_stage( + self, + event: asyncio.Event, + request: Request, + timeout: float, + group_request_id: int, + stage: str, + ) -> None: + self._change_pd_stage_waiting_requests(stage, 1) + try: + await self._wait_for_event_or_disconnect(event, request, timeout, group_request_id, stage) + finally: + self._change_pd_stage_waiting_requests(stage, -1) + async def _log_req_header(self, request: Request, group_request_id: int): x_request_id = request.headers.get("X-Request-Id", "") x_session_id = request.headers.get("X-Session-Id", "") @@ -393,7 +418,7 @@ async def fetch_pd_stream( await p_node.websocket.send_bytes(pickle.dumps((ObjType.REQ, (prompt, sampling_params, multimodal_params)))) try: - await self._wait_for_event_or_disconnect( + await self._wait_for_pd_stage( prefill_prompt_ids_event, request, timeout=60, @@ -414,7 +439,7 @@ async def fetch_pd_stream( ) try: - await self._wait_for_event_or_disconnect( + await self._wait_for_pd_stage( up_status_event, request, timeout=180, @@ -436,13 +461,19 @@ async def fetch_pd_stream( first_token_gen = False needs_prefill_first_token = decode_node_info.ready_kv_len != len(prompt_ids) - 1 - prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( - req_status=req_status, - request=request, - group_request_id=group_request_id, - needs_prefill_first_token=needs_prefill_first_token, - ready_kv_len=decode_node_info.ready_kv_len, - ) + if needs_prefill_first_token: + self._change_pd_stage_waiting_requests("prefill", 1) + try: + prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( + req_status=req_status, + request=request, + group_request_id=group_request_id, + needs_prefill_first_token=needs_prefill_first_token, + ready_kv_len=decode_node_info.ready_kv_len, + ) + finally: + if needs_prefill_first_token: + self._change_pd_stage_waiting_requests("prefill", -1) while True: await req_status.wait_to_ready() @@ -754,8 +785,9 @@ async def put_tokens_to_front(self, token_list: List[Tuple[int, str, dict, Finis class PDManager: - def __init__(self, args: StartArgs): + def __init__(self, args: StartArgs, metric_client=None): self.args: StartArgs = args + self.metric_client = metric_client self.prefill_nodes: List[PD_Client_Obj] = [] self.decode_nodes: List[PD_Client_Obj] = [] self.url_to_pd_nodes: Dict[str, PD_Client_Obj] = {} @@ -879,6 +911,12 @@ def update_node_load_info(self, load_info: Optional[dict]): total_token_usage_rate = load_info["total_token_usage_rate"] pd_client = self.url_to_pd_nodes.get(client_ip_port) pd_client.run_status.total_token_usage_rate = total_token_usage_rate + if self.metric_client is not None: + self.metric_client.gauge_set( + "lightllm_pd_node_token_usage_ratio", + total_token_usage_rate, + labels={"role": pd_client.mode, "endpoint": client_ip_port}, + ) except BaseException as e: logger.warning(f"udpate node load info failed, load_info: {load_info} error: {str(e)}") return diff --git a/lightllm/server/metrics/manager.py b/lightllm/server/metrics/manager.py index 22f6426a77..f8f56dc8d3 100644 --- a/lightllm/server/metrics/manager.py +++ b/lightllm/server/metrics/manager.py @@ -55,8 +55,9 @@ def exposed_counter_inc_by(self, name: str, amount: float) -> None: def exposed_histogram_observe(self, name: str, value: float, label: str = None) -> None: return self.monitor.histogram_observe(name, value, label) - def exposed_gauge_set(self, name: str, value: float) -> None: - return self.monitor.gauge_set(name, value) + def exposed_gauge_set(self, name: str, value: float, labels: dict = None) -> None: + local_labels = None if labels is None else {key: labels[key] for key in labels} + return self.monitor.gauge_set(name, value, local_labels) def exposed_generate_latest(self) -> bytes: data = generate_latest(self.monitor.registry) diff --git a/lightllm/server/metrics/metrics.py b/lightllm/server/metrics/metrics.py index 0d42462c3f..b58b587b26 100644 --- a/lightllm/server/metrics/metrics.py +++ b/lightllm/server/metrics/metrics.py @@ -32,6 +32,8 @@ "lightllm_cache_hit_rate": "Prefix cache hit rate of latest completed request", "lightllm_gen_throughput": "Generation throughput of latest completed request (tokens/s)", "lightllm_num_running_reqs": "Number of running requests", + "lightllm_pd_node_token_usage_ratio": "Token capacity usage ratio reported by a PD node", + "lightllm_pd_master_stage_waiting_requests": "Requests waiting for a PD stage to become ready", } @@ -111,6 +113,8 @@ def init_metrics(self, args): self.create_gauge("lightllm_cache_hit_rate") self.create_gauge("lightllm_gen_throughput") self.create_gauge("lightllm_num_running_reqs") + self.create_gauge("lightllm_pd_node_token_usage_ratio", labelnames=["role", "endpoint"]) + self.create_gauge("lightllm_pd_master_stage_waiting_requests", labelnames=["stage"]) def create_histogram(self, name, buckets, labelnames=None): all_labels = ["model_name"] + (labelnames or []) @@ -122,8 +126,9 @@ def create_counter(self, name, labelnames=None): counter = Counter(name, MONITOR_INFO[name], labelnames=all_labels, registry=self.registry) self.monitor_registry[name] = counter - def create_gauge(self, name): - gauge = Gauge(name, MONITOR_INFO[name], labelnames=["model_name"], registry=self.registry) + def create_gauge(self, name, labelnames=None): + all_labels = ["model_name"] + (labelnames or []) + gauge = Gauge(name, MONITOR_INFO[name], labelnames=all_labels, registry=self.registry) self.monitor_registry[name] = gauge def counter_inc(self, name, label=None): @@ -141,8 +146,11 @@ def histogram_observe(self, name, value, label=None): else: self.monitor_registry[name].labels(model_name=self.model_name, method=label).observe(value) - def gauge_set(self, name, value): - self.monitor_registry[name].labels(model_name=self.model_name).set(value) + def gauge_set(self, name, value, labels=None): + metric_labels = {"model_name": self.model_name} + if labels: + metric_labels.update(labels) + self.monitor_registry[name].labels(**metric_labels).set(value) def push_metrices(self): if self.gateway_url is not None: diff --git a/test/test_pd_selector/test_pd_master_metrics.py b/test/test_pd_selector/test_pd_master_metrics.py new file mode 100644 index 0000000000..0243565cd5 --- /dev/null +++ b/test/test_pd_selector/test_pd_master_metrics.py @@ -0,0 +1,124 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, call + +import pytest +from prometheus_client import generate_latest + +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager +from lightllm.server.metrics.manager import MetricServer +from lightllm.server.metrics.metrics import Monitor +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError + + +class RPyCDictProxyLike: + def __init__(self, values): + self.values = values + + def __iter__(self): + return iter(self.values) + + def __getitem__(self, key): + return self.values[key] + + +def test_pd_master_exports_node_load_with_role_and_endpoint_labels(): + args = SimpleNamespace( + metric_gateway=None, + job_name="lightllm", + grouping_key=None, + enable_monitor_auth=False, + model_name="test-model", + max_req_total_len=1024, + mtp_step=0, + select_p_d_node_strategy="random", + ) + monitor = Monitor(args) + metric_client = MagicMock() + metric_client.gauge_set.side_effect = monitor.gauge_set + manager = PDManager(args, metric_client) + manager.url_to_pd_nodes["10.0.0.1:28761"] = PD_Client_Obj( + node_id=1, + client_ip_port="10.0.0.1:28761", + mode="prefill", + start_args={}, + ) + manager.url_to_pd_nodes["10.0.0.2:28764"] = PD_Client_Obj( + node_id=2, + client_ip_port="10.0.0.2:28764", + mode="decode", + start_args={}, + ) + + manager.update_node_load_info({"client_ip_port": "10.0.0.1:28761", "total_token_usage_rate": 0.25}) + manager.update_node_load_info({"client_ip_port": "10.0.0.2:28764", "total_token_usage_rate": 0.75}) + monitor.gauge_set("lightllm_pd_master_stage_waiting_requests", 2, labels={"stage": "decode"}) + + metrics = generate_latest(monitor.registry).decode() + assert ( + 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.1:28761",model_name="test-model",role="prefill"} 0.25' + in metrics + ) + assert ( + 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.2:28764",model_name="test-model",role="decode"} 0.75' + in metrics + ) + assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 2.0' in metrics + + +def test_metric_server_copies_rpyc_label_proxy_before_updating_gauge(): + args = SimpleNamespace( + metric_gateway=None, + job_name="lightllm", + grouping_key=None, + enable_monitor_auth=False, + model_name="test-model", + max_req_total_len=1024, + mtp_step=0, + push_interval=10, + ) + server = MetricServer(args) + + server.exposed_gauge_set( + "lightllm_pd_master_stage_waiting_requests", + 3, + RPyCDictProxyLike({"stage": "decode"}), + ) + + metrics = generate_latest(server.monitor.registry).decode() + assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 3.0' in metrics + + +@pytest.mark.parametrize("stage", ["prefill", "decode"]) +@pytest.mark.parametrize( + "outcome", + [ + None, + ServerBusyError(), + asyncio.CancelledError(), + ClientDisconnected(group_request_id=123), + RuntimeError("worker failed"), + ], + ids=["success", "timeout", "cancellation", "disconnect", "exception"], +) +def test_pd_master_stage_waiting_gauge_is_balanced_for_every_exit(stage, outcome): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.metric_client = MagicMock() + manager.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} + manager._wait_for_event_or_disconnect = AsyncMock(side_effect=outcome) + + if outcome is None: + await manager._wait_for_pd_stage(AsyncMock(), AsyncMock(), 1, 123, stage) + else: + with pytest.raises(type(outcome)): + await manager._wait_for_pd_stage(AsyncMock(), AsyncMock(), 1, 123, stage) + + assert manager.pd_stage_waiting_request_counts[stage] == 0 + assert manager.metric_client.gauge_set.call_args_list == [ + call("lightllm_pd_master_stage_waiting_requests", 1, labels={"stage": stage}), + call("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}), + ] + + asyncio.run(run()) From 3e0d9f70ce3e1070504ac812606aff7be6744b5a Mon Sep 17 00:00:00 2001 From: sufubao <47234901+sufubao@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:10:21 +0800 Subject: [PATCH 11/28] perf(mtp): optimize Qwen3.5 GDN verification kernel (#1513) --- .../common/basemodel/attention/linear/gdn.py | 51 +++-- .../linear_att/mtp_fused_recurrent.py | 130 +++++++++-- .../triton_kernel/norm/gated_rmsnorm.py | 28 ++- .../layer_infer/transformer_layer_infer.py | 1 - .../test_mtp_fused_recurrent_autotune.py | 201 ++++++++++++++++++ .../test_mtp_fused_recurrent_equiv.py | 1 + .../triton_kernel/test_gated_rmsnorm.py | 43 ++++ .../qwen3next/test_transformer_layer_infer.py | 35 +++ 8 files changed, 455 insertions(+), 35 deletions(-) create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py create mode 100644 unit_tests/models/qwen3next/test_transformer_layer_infer.py diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index ca6ceaec43..85a19107f8 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -3,7 +3,8 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args, get_triton_autotune_level +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs @@ -247,7 +248,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): dtype=torch.int32, device=self.infer_state.b_req_idx.device, ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0] self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) @@ -388,18 +389,38 @@ def _gdn_mtp_kernel( # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=self.b_ssm_buffer_idx, - ssm_state_write_indices=self.b_ssm_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, + fixed_seq_len = ( + 0 + if backend.uses_dynamic_spec_verify_layout() + else backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + 1 + ) + tune_now = ( + get_triton_autotune_level() in [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE] + and getattr(infer_state, "is_cuda_graph", False) + and infer_state.microbatch_index == 0 + and layer_weight.layer_num_ == 0 + and not torch.cuda.is_current_stream_capturing() + and not Autotuner.is_autotune_warmup() ) + if tune_now: + Autotuner.start_autotune_warmup() + try: + core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + cu_seqlens=cu_seqlens_q, + ssm_state_indices=self.b_ssm_buffer_idx, + ssm_state_write_indices=self.b_ssm_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + a_raw=a, + b_raw=b, + fixed_seq_len=fixed_seq_len, + ) + finally: + if tune_now: + Autotuner.end_autotune_warmup() return core_attn_out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py index 2eb5ef5333..37642369fc 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py @@ -15,6 +15,46 @@ import triton import triton.language as tl +from lightllm.common.triton_utils.autotuner import autotune + + +_MTP_RECURRENT_BV_SIZES = (4, 8, 16, 32, 64) + + +def _get_mtp_fused_recurrent_configs(): + return [ + {"BV": bv, "num_warps": num_warps, "num_stages": num_stages} + for bv in _MTP_RECURRENT_BV_SIZES + for num_warps in (1, 2, 4, 8) + for num_stages in (1, 2, 3) + ] + + +def _get_mtp_fused_recurrent_static_key(q, v, initial_state, fixed_seq_len): + return { + "H": q.shape[2], + "HV": v.shape[2], + "K": q.shape[3], + "V": v.shape[3], + "dtype": str(q.dtype), + "state_dtype": str(initial_state.dtype), + "fixed_seq_len": int(fixed_seq_len), + } + + +def _get_mtp_fused_recurrent_run_key(q, cu_seqlens): + sequence_count = int(cu_seqlens.shape[0] - 1) + total_tokens = int(q.shape[1]) + return sequence_count * 10 ** 9 + total_tokens + + +def _default_mtp_fused_recurrent_config(V): + return { + "BV": min(triton.next_power_of_2(V), 8), + "num_warps": 1, + "num_stages": 3, + } + # --------------------------------------------------------------------------- # Triton kernel @@ -59,14 +99,17 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( stride_write_indices_tok: tl.constexpr, SOFTPLUS_BETA: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, + FIXED_SEQ_LEN: tl.constexpr, ): i_v, i_n, i_hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_h = i_hv // (HV // H) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = eos - bos + if FIXED_SEQ_LEN > 0: + bos = i_n * FIXED_SEQ_LEN + T: tl.constexpr = FIXED_SEQ_LEN + else: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos if T == 0: return @@ -133,11 +176,18 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( # --------------------------------------------------------------------------- -# Public API — directly launches the triton kernel (no autograd.Function) +# Autotuned implementation — directly launches the Triton kernel # --------------------------------------------------------------------------- -def mtp_fused_recurrent_gated_delta_rule( +@autotune( + kernel_name="mtp_fused_recurrent_gated_delta_rule:v1", + configs_gen_func=_get_mtp_fused_recurrent_configs, + static_key_func=_get_mtp_fused_recurrent_static_key, + run_key_func=_get_mtp_fused_recurrent_run_key, + mutates_args=["initial_state"], +) +def _mtp_fused_recurrent_gated_delta_rule_autotuned( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -150,6 +200,8 @@ def mtp_fused_recurrent_gated_delta_rule( dt_bias: torch.Tensor, a_raw: torch.Tensor, b_raw: torch.Tensor, + fixed_seq_len: int, + run_config: dict = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused recurrent gated delta rule with fused gating (GDN layer). @@ -170,6 +222,8 @@ def mtp_fused_recurrent_gated_delta_rule( dt_bias: ``[HV]`` per-head dt bias. a_raw: ``[T, HV]`` raw alpha. b_raw: ``[T, HV]`` raw beta. + fixed_seq_len: Compile-time sequence length for dense fixed-width MTP + verification. Zero keeps the variable-length ``cu_seqlens`` path. Returns: ``(o, final_state)`` where ``o`` is ``[1, T, HV, V]`` and @@ -184,6 +238,12 @@ def mtp_fused_recurrent_gated_delta_rule( V = v.shape[-1] HV = v.shape[2] N = len(cu_seqlens) - 1 + fixed_seq_len = int(fixed_seq_len) + assert fixed_seq_len >= 0 + if fixed_seq_len: + assert q.shape[1] == N * fixed_seq_len, ( + f"fixed_seq_len={fixed_seq_len} requires {N * fixed_seq_len} tokens, " f"got {q.shape[1]}" + ) q, stride_q_tok = _ensure_qkv_token_strided(q) k, stride_k_tok = _ensure_qkv_token_strided(k) v, stride_v_tok = _ensure_qkv_token_strided(v) @@ -191,17 +251,18 @@ def mtp_fused_recurrent_gated_delta_rule( b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw) BK = triton.next_power_of_2(K) assert K == BK, f"K={K} must be a power of 2" - BV = min(triton.next_power_of_2(V), 8) - num_warps = 1 - num_stages = 3 + if run_config is None: + run_config = _default_mtp_fused_recurrent_config(V) + BV = run_config["BV"] + num_warps = run_config.get("num_warps", 1) + num_stages = run_config["num_stages"] NV = triton.cdiv(V, BV) - o = q.new_empty(v.shape) + output = q.new_empty(v.shape) final_state = initial_state - stride_init_state_token = initial_state.stride(0) stride_final_state_token = final_state.stride(0) - stride_o_tok = o.stride(1) + stride_o_tok = output.stride(1) assert stride_o_tok == HV * V, f"stride_o_tok={stride_o_tok} must be HV*V" stride_state_hv = K * V @@ -216,7 +277,7 @@ def mtp_fused_recurrent_gated_delta_rule( q=q, k=k, v=v, - o=o, + o=output, h0=initial_state, ht=final_state, cu_seqlens=cu_seqlens, @@ -249,10 +310,49 @@ def mtp_fused_recurrent_gated_delta_rule( stride_write_indices_tok=stride_write_indices_tok, SOFTPLUS_BETA=1.0, SOFTPLUS_THRESHOLD=20.0, + FIXED_SEQ_LEN=fixed_seq_len, num_warps=num_warps, num_stages=num_stages, ) - return o, final_state + return output, final_state + + +def mtp_fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + ssm_state_write_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + a_raw: torch.Tensor, + b_raw: torch.Tensor, + fixed_seq_len: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the autotuned fused recurrent GDN kernel. + + ``fixed_seq_len=0`` preserves the variable-length behavior of the original + public API. The private autotuned implementation always receives the value + explicitly so its cache key does not depend on generic default handling. + """ + return _mtp_fused_recurrent_gated_delta_rule_autotuned( + q=q, + k=k, + v=v, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + A_log=A_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + fixed_seq_len=fixed_seq_len, + ) # --------------------------------------------------------------------------- diff --git a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py index c62c5eb5d2..b42d1eeaa8 100644 --- a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py +++ b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py @@ -18,20 +18,26 @@ def gated_rmsnorm_forward_kernel( Z, # pointer to the other branch (required, not optional) stride_x_row, # how much to increase the pointer when moving by 1 row stride_y_row, - stride_z_row, + stride_z_token, + stride_z_head, M, # number of rows in X N, # number of columns in X eps, # epsilon to avoid division by zero BLOCK_N: tl.constexpr, HAS_BIAS: tl.constexpr, NORM_BEFORE_GATE: tl.constexpr, + Z_HEADS: tl.constexpr, ): # Map the program id to the row of X and Y it should compute. row = tl.program_id(0) group = tl.program_id(1) X += row * stride_x_row + group * N Y += row * stride_y_row + group * N - Z += row * stride_z_row + group * N + # X is flattened to [tokens * heads, N], while GDN can keep Z as a + # zero-copy [tokens, heads, N] slice of its packed projection output. + z_token = row // Z_HEADS + z_head = row % Z_HEADS + Z += z_token * stride_z_token + z_head * stride_z_head + group * N W += group * N if HAS_BIAS: B += group * N @@ -112,8 +118,20 @@ def gated_rmsnorm_forward( assert x.stride(-1) == 1 # z is required for gated_rmsnorm assert z is not None, "z cannot be None for gated_rmsnorm_forward" + # Accept GDN's strided 3D gate without materializing a flattened copy. + assert z.ndim in (2, 3), f"z must be [M, N] or [tokens, heads, N], got shape={z.shape}" assert z.stride(-1) == 1 - assert z.shape == (M, N) + if z.ndim == 2: + assert z.shape == (M, N) + z_heads = 1 + stride_z_token = z.stride(0) + stride_z_head = 0 + else: + assert z.shape[-1] == N, f"z.shape[-1]={z.shape[-1]} must match N={N}" + assert z.shape[0] * z.shape[1] == M, f"z token/head rows={z.shape[0] * z.shape[1]} must match M={M}" + z_heads = z.shape[1] + stride_z_token = z.stride(0) + stride_z_head = z.stride(1) assert weight.shape == (N,) assert weight.stride(-1) == 1 if bias is not None: @@ -156,12 +174,14 @@ def gated_rmsnorm_forward( z, x.stride(0), out.stride(0), - z.stride(0), + stride_z_token, + stride_z_head, M, group_size, eps, BLOCK_N=BLOCK_N, NORM_BEFORE_GATE=norm_before_gate, + Z_HEADS=z_heads, num_warps=num_warps, ) return out diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 92d68c9fd2..2b7b6a3ae4 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -319,7 +319,6 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py new file mode 100644 index 0000000000..c70452aa08 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -0,0 +1,201 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.linear import gdn as gdn_module +from lightllm.common.basemodel.triton_kernel.linear_att import causal_conv1d_mtp +from lightllm.common.basemodel.triton_kernel.linear_att import mtp_fused_recurrent as mtp_recurrent +from lightllm.common.triton_utils.autotuner import Autotuner + + +def _layout_tensors(total_tokens=8, K=64, V=128): + H, HV = 1, 2 + q = torch.empty((1, total_tokens, H, K), dtype=torch.bfloat16) + k = torch.empty_like(q) + v = torch.empty((1, total_tokens, HV, V), dtype=torch.bfloat16) + return q, k, v + + +def test_v1_config_space_and_fallback(): + configs = mtp_recurrent._get_mtp_fused_recurrent_configs() + assert len(configs) == 60 + assert len({tuple(sorted(config.items())) for config in configs}) == len(configs) + assert {config["BV"] for config in configs} == {4, 8, 16, 32, 64} + assert {config["num_warps"] for config in configs} == {1, 2, 4, 8} + assert {config["num_stages"] for config in configs} == {1, 2, 3} + assert mtp_recurrent._default_mtp_fused_recurrent_config(128) == { + "BV": 8, + "num_warps": 1, + "num_stages": 3, + } + + +def test_key_separates_fixed_and_dynamic_specializations(): + q, _, v = _layout_tensors() + initial_state = torch.empty((16, 2, 64, 128), dtype=torch.bfloat16) + cu_seqlens = torch.empty(3, dtype=torch.int64) + + dynamic_key = mtp_recurrent._get_mtp_fused_recurrent_static_key(q, v, initial_state, fixed_seq_len=0) + fixed_key = mtp_recurrent._get_mtp_fused_recurrent_static_key(q, v, initial_state, fixed_seq_len=4) + + assert dynamic_key["fixed_seq_len"] == 0 + assert fixed_key["fixed_seq_len"] == 4 + assert dynamic_key != fixed_key + assert mtp_recurrent._get_mtp_fused_recurrent_run_key(q, cu_seqlens) == 2 * 10 ** 9 + 8 + + +def test_tuner_declares_mutated_state(): + assert mtp_recurrent._mtp_fused_recurrent_gated_delta_rule_autotuned.mutates_args == ["initial_state"] + + +def test_gdn_cuda_graph_warmup_scopes_v1_autotune_to_first_layer_microbatch(monkeypatch): + Autotuner.end_autotune_warmup() + calls = [] + q, k, v = _layout_tensors() + backend = SimpleNamespace( + activation="silu", + conv_kernel_dim=4, + model=SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 3), + ), + uses_dynamic_spec_verify_layout=lambda: False, + _rearrange_mixed_qkv=lambda _, decode: (q, k, v), + ) + state = gdn_module.LinearAttDecodeAttState(backend=backend) + state.b1_mtp_cu_q_seq_len = torch.tensor([0, 4, 8], dtype=torch.int32) + state.b_conv_buffer_idx = torch.tensor([0, 1], dtype=torch.int32) + state.b_ssm_buffer_idx = torch.zeros((2, 4), dtype=torch.int32) + state.b_num_accepted_tokens = torch.ones(2, dtype=torch.int32) + infer_state = SimpleNamespace(is_cuda_graph=True, microbatch_index=0) + layer_weight = SimpleNamespace( + layer_num_=0, + linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=torch.empty(1)), bias=None), + linear_A_log=SimpleNamespace(weight=torch.empty(1)), + linear_dt_bias=SimpleNamespace(weight=torch.empty(1)), + ) + + monkeypatch.setattr(gdn_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr(causal_conv1d_mtp, "causal_conv1d_update", lambda mixed_qkv, *args, **kwargs: mixed_qkv) + + def fake_recurrent(**kwargs): + calls.append(Autotuner.is_autotune_warmup()) + return kwargs["q"], kwargs["initial_state"] + + monkeypatch.setattr(gdn_module, "mtp_fused_recurrent_gated_delta_rule", fake_recurrent) + + args = ( + torch.empty((8, 1)), + torch.empty(1), + torch.empty(1), + torch.empty((8, 2)), + torch.empty((8, 2)), + infer_state, + layer_weight, + ) + state._gdn_mtp_kernel(*args) + layer_weight.layer_num_ = 1 + state._gdn_mtp_kernel(*args) + + assert calls == [True, False] + assert not Autotuner.is_autotune_warmup() + + +def _make_fixed_case(): + torch.manual_seed(11) + batch, seq_len, H, HV, K, V = 2, 4, 1, 2, 64, 128 + total_tokens = batch * seq_len + q = torch.randn((1, total_tokens, H, K), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn((1, total_tokens, HV, V), device="cuda", dtype=torch.bfloat16) + initial_state = torch.randn((10, HV, K, V), device="cuda", dtype=torch.bfloat16) + cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int64) * seq_len + read_indices = torch.tensor([[0, 0, 0, 0], [1, 1, 1, 1]], device="cuda", dtype=torch.int32) + write_indices = torch.tensor([[2, 3, 4, 5], [6, 7, 8, 9]], device="cuda", dtype=torch.int32) + accepted = torch.full((batch,), seq_len, device="cuda", dtype=torch.int32) + A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + a_raw = torch.randn((total_tokens, HV), device="cuda", dtype=torch.bfloat16) + b_raw = torch.randn_like(a_raw) + return ( + q, + k, + v, + initial_state, + cu_seqlens, + read_indices, + write_indices, + accepted, + A_log, + dt_bias, + a_raw, + b_raw, + ) + + +def _launch(case, run_config, *, fixed_seq_len): + ( + q, + k, + v, + initial_state, + cu_seqlens, + read_indices, + write_indices, + accepted, + A_log, + dt_bias, + a_raw, + b_raw, + ) = case + state = initial_state.clone() + kwargs = { + "q": q, + "k": k, + "v": v, + "initial_state": state, + "cu_seqlens": cu_seqlens, + "ssm_state_indices": read_indices, + "ssm_state_write_indices": write_indices, + "num_accepted_tokens": accepted, + "A_log": A_log, + "dt_bias": dt_bias, + "a_raw": a_raw, + "b_raw": b_raw, + "fixed_seq_len": fixed_seq_len, + "run_config": run_config, + } + output, _ = mtp_recurrent._mtp_fused_recurrent_gated_delta_rule_autotuned(**kwargs) + return output, state, kwargs + + +_RETAINED_CONFIGS = mtp_recurrent._get_mtp_fused_recurrent_configs() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("fixed_seq_len", [0, 4]) +@pytest.mark.parametrize("run_config", _RETAINED_CONFIGS) +def test_retained_configs_match_fallback(run_config, fixed_seq_len): + case = _make_fixed_case() + fallback = mtp_recurrent._default_mtp_fused_recurrent_config(128) + expected_output, expected_state, _ = _launch(case, fallback, fixed_seq_len=fixed_seq_len) + output, state, _ = _launch(case, run_config, fixed_seq_len=fixed_seq_len) + + torch.testing.assert_close(output.float(), expected_output.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state.float(), expected_state.float(), atol=5.0, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_tuner_warmup_does_not_mutate_caller_state(): + case = _make_fixed_case() + fallback = mtp_recurrent._default_mtp_fused_recurrent_config(128) + _, state, kwargs = _launch(case, fallback, fixed_seq_len=4) + state.copy_(case[3]) + state_before = state.clone() + tuner = mtp_recurrent._mtp_fused_recurrent_gated_delta_rule_autotuned + + tuner.kernel_warmup(tuner._static_key(**kwargs), **kwargs) + + torch.testing.assert_close(state, state_before, atol=0, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..a54ab21ecb 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -60,6 +60,7 @@ def _run_both( dt_bias=dt_bias, a_raw=a_raw, b_raw=b_raw, + fixed_seq_len=0, ) return o_old, o_new, fs_old, fs_new diff --git a/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py new file mode 100644 index 0000000000..f03fef4629 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py @@ -0,0 +1,43 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import ( + gated_rmsnorm_forward, +) + + +@pytest.mark.parametrize("norm_before_gate", [True, False]) +def test_gated_rmsnorm_accepts_strided_3d_gate(norm_before_gate): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for gated RMSNorm test") + + torch.manual_seed(123) + tokens, heads, head_dim = 5, 4, 128 + packed_dim = heads * head_dim + 256 + x = torch.randn((tokens * heads, head_dim), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16) + packed = torch.randn((tokens, packed_dim), device="cuda", dtype=torch.bfloat16) + z = packed[:, 128 : 128 + heads * head_dim].view(tokens, heads, head_dim) + assert not z.is_contiguous() + + run_config = {"BLOCK_N": head_dim, "num_warps": 1} + expected = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z.contiguous().view(-1, head_dim), + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z, + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/models/qwen3next/test_transformer_layer_infer.py b/unit_tests/models/qwen3next/test_transformer_layer_infer.py new file mode 100644 index 0000000000..08334539d1 --- /dev/null +++ b/unit_tests/models/qwen3next/test_transformer_layer_infer.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import ( + Qwen3NextTransformerLayerInfer, +) + + +def test_linear_post_passes_strided_gate_without_materializing(): + tokens, heads, head_dim = 3, 4, 8 + packed_dim = heads * head_dim + 16 + packed = torch.randn((tokens, packed_dim)) + z = packed[:, 8 : 8 + heads * head_dim].view(tokens, heads, head_dim) + core_attn_out = torch.randn((1, tokens, heads, head_dim)) + calls = [] + + def linear_norm(input, gate_value, eps): + calls.append((input, gate_value, eps)) + return input + + layer_weight = SimpleNamespace( + linear_norm=linear_norm, + linear_out_proj=SimpleNamespace(mm=lambda input: input), + ) + layer_infer = object.__new__(Qwen3NextTransformerLayerInfer) + layer_infer.eps_ = 1e-6 + + output = layer_infer._linear_post(core_attn_out, z, layer_weight) + + assert len(calls) == 1 + assert calls[0][0].shape == (tokens * heads, head_dim) + assert calls[0][1] is z + assert not calls[0][1].is_contiguous() + assert output.shape == (tokens, heads * head_dim) From cc25c63d4317489a1ac1a005d43682512f4769bd Mon Sep 17 00:00:00 2001 From: sufubao <47234901+sufubao@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:30:57 +0800 Subject: [PATCH 12/28] perf: avoid full-vocab all-gather for draft greedy sampling (#1517) --- lightllm/common/basemodel/basemodel.py | 19 +++ lightllm/common/basemodel/batch_objs.py | 36 ++++++ lightllm/common/basemodel/infer_struct.py | 3 + .../post_process/greedy_sample.py | 109 ++++++++++++++++ .../post_process/vocab_parallel_greedy.py | 118 ++++++++++++++++++ .../triton_kernel/transpose_convert.py | 65 ++++++++++ .../llama/layer_infer/post_layer_infer.py | 18 ++- .../layer_infer/post_layer_infer.py | 6 +- .../model_infer/mode_backend/base_backend.py | 18 ++- .../dp_overlap_proposers/eagle_with_att.py | 2 +- .../proposers/eagle_with_att.py | 2 +- .../common/basemodel/test_model_output.py | 24 ++++ .../test_vocab_parallel_greedy.py | 65 ++++++++++ .../test_vocab_parallel_greedy_output.py | 74 +++++++++++ 14 files changed, 551 insertions(+), 8 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py create mode 100644 lightllm/common/basemodel/triton_kernel/transpose_convert.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py create mode 100644 unit_tests/models/test_vocab_parallel_greedy_output.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e80f2b552f..b561595b98 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -384,6 +384,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.return_all_prompt_logics = self.return_all_prompt_logics + infer_state.use_vocab_parallel_greedy = self.is_mtp_draft_model infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len @@ -534,6 +535,9 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba return model_output new_model_output = copy.copy(model_output) new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -546,6 +550,9 @@ def _create_unpad_prefill_model_output( new_model_output = copy.copy(padded_model_output) # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。 new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -737,6 +744,8 @@ def prefill_func(input_tensors, _infer_state): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -766,6 +775,8 @@ def _token_forward(self, infer_state: InferStateInfo): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -1020,11 +1031,15 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1069,10 +1084,14 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), ) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index ae645d4b7b..edddb749e5 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -200,10 +200,46 @@ class ModelOutput: # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None + # Vocab-parallel outputs keep logits as logits while mapping each sparse + # column back to its global token id. logits_logsumexp is computed over the + # complete vocabulary, so sparse argmax probabilities remain exact. + # Both fields are None for historical dense logits. + logits_token_ids: Optional[torch.Tensor] = None + logits_logsumexp: Optional[torch.Tensor] = None + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() + assert (self.logits_token_ids is None) == (self.logits_logsumexp is None) + if self.logits_token_ids is not None: + assert self.logits.ndim == 2 + assert self.logits_token_ids.shape == self.logits.shape + assert self.logits_token_ids.dtype in (torch.int32, torch.int64) + assert self.logits_token_ids.device == self.logits.device + assert self.logits_logsumexp.shape == (self.logits.shape[0],) + assert self.logits_logsumexp.dtype == torch.float32 + assert self.logits_logsumexp.device == self.logits.device def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) + if self.logits_token_ids is not None: + self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids) + self.logits_logsumexp = tensor_to_no_ref_tensor(self.logits_logsumexp) self.mtp_collector.to_no_ref_tensor() + + @property + def has_vocab_parallel_logits(self) -> bool: + return self.logits_token_ids is not None + + def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": + """Select logit rows without dropping their vocabulary metadata.""" + + return ModelOutput( + logits=self.logits.index_select(0, index), + logits_token_ids=( + self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None + ), + logits_logsumexp=( + self.logits_logsumexp.index_select(0, index) if self.logits_logsumexp is not None else None + ), + ) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 29648aa78e..5c1e361729 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -52,6 +52,9 @@ def __init__(self): self.mem_index: torch.Tensor = None self.return_all_prompt_logics: bool = False + self.use_vocab_parallel_greedy: bool = False + self.logits_token_ids: Optional[torch.Tensor] = None + self.logits_logsumexp: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。 diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py new file mode 100644 index 0000000000..8e7e99e1ab --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py @@ -0,0 +1,109 @@ +"""Local greedy statistics for distributed vocabulary shards.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _greedy_sample_stage1_kernel( + logits, + partial_max, + partial_sum, + partial_argmax, + stride_row, + stride_col, + vocab_size: tl.constexpr, + num_blocks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + block = tl.program_id(1) + offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + values = tl.load( + logits + row * stride_row + offsets * stride_col, + mask=offsets < vocab_size, + other=-float("inf"), + ) + values = values.to(tl.float32) + + block_max = tl.max(values, axis=0) + block_sum = tl.sum(tl.exp(values - block_max), axis=0) + block_argmax = tl.argmax(values, axis=0) + block * BLOCK_SIZE + output_offset = row * num_blocks + block + tl.store(partial_max + output_offset, block_max) + tl.store(partial_sum + output_offset, block_sum) + tl.store(partial_argmax + output_offset, block_argmax) + + +@triton.jit +def _greedy_sample_stage2_stats_kernel( + partial_max, + partial_sum, + partial_argmax, + output_stats, + output_argmax, + num_blocks: tl.constexpr, + batch_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < num_blocks + input_offset = row * num_blocks + offsets + block_max = tl.load(partial_max + input_offset, mask=mask, other=-float("inf")) + block_sum = tl.load(partial_sum + input_offset, mask=mask, other=0.0) + block_argmax = tl.load(partial_argmax + input_offset, mask=mask, other=0x7FFFFFFF) + + global_max = tl.max(block_max, axis=0) + global_sum = tl.sum(block_sum * tl.exp(block_max - global_max), axis=0) + candidate_ids = tl.where(block_max == global_max, block_argmax, 0x7FFFFFFF) + global_argmax = tl.min(candidate_ids, axis=0) + tl.store(output_stats + row, global_max) + tl.store(output_stats + batch_size + row, global_max + tl.log(global_sum)) + tl.store(output_argmax + row, global_argmax) + + +def _launch_stage1(logits: torch.Tensor, scratch: torch.Tensor, block_size: int, num_blocks: int) -> None: + batch_size, vocab_size = logits.shape + _greedy_sample_stage1_kernel[(batch_size, num_blocks)]( + logits, + scratch[0], + scratch[1], + scratch[2], + logits.stride(0), + logits.stride(1), + vocab_size=vocab_size, + num_blocks=num_blocks, + BLOCK_SIZE=block_size, + num_warps=8, + ) + + +@torch.no_grad() +def greedy_sample_local_stats(logits: torch.Tensor, alloc_func=torch.empty) -> torch.Tensor: + """Return local max, logsumexp and argmax rows for distributed greedy sampling.""" + + assert logits.ndim == 2 and logits.is_cuda and logits.is_contiguous() + batch_size, vocab_size = logits.shape + block_size = 4096 + num_blocks = triton.cdiv(vocab_size, block_size) + scratch = alloc_func((3, batch_size, num_blocks), dtype=torch.float32, device=logits.device) + # The third FP32 row carries INT32 argmax bits. Keeping one fixed-size + # payload gives the distributed reducer a single collective without losing + # token-id precision through a numeric int-to-float conversion. + output_stats = alloc_func((3, batch_size), dtype=torch.float32, device=logits.device) + + _launch_stage1(logits, scratch, block_size, num_blocks) + _greedy_sample_stage2_stats_kernel[(batch_size,)]( + scratch[0], + scratch[1], + scratch[2], + output_stats, + output_stats[2].view(torch.int32), + num_blocks=num_blocks, + batch_size=batch_size, + BLOCK_SIZE=triton.next_power_of_2(num_blocks), + num_warps=4, + ) + return output_stats diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py new file mode 100644 index 0000000000..fbfa4191e6 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py @@ -0,0 +1,118 @@ +"""Greedy sampling directly from tensor-parallel vocabulary shards.""" + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) +from lightllm.common.basemodel.triton_kernel.transpose_convert import ( + transpose_convert_2d, +) +from lightllm.distributed.communication_op import all_gather_into_tensor + + +@triton.jit +def _combine_vocab_parallel_stats_kernel( + gathered_stats, + gathered_argmax, + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size: tl.constexpr, + tp_world_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < token_num + rank_stride = 3 * token_num + + global_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) + global_id = tl.full((BLOCK_SIZE,), 0x7FFFFFFF, tl.int32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_max = tl.load( + gathered_stats + rank_base + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + local_id = tl.load( + gathered_argmax + rank_base + 2 * token_num + token_offsets, + mask=token_mask, + other=0x7FFFFFFF, + ) + local_id += (rank * vocab_size) // tp_world_size + wins = (local_max > global_max) | ((local_max == global_max) & (local_id < global_id)) + global_max = tl.where(wins, local_max, global_max) + global_id = tl.where(wins, local_id, global_id) + + global_sum = tl.zeros((BLOCK_SIZE,), tl.float32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_lse = tl.load( + gathered_stats + rank_base + token_num + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + global_sum += tl.exp(local_lse - global_max) + + tl.store(output_logits + token_offsets, global_max, mask=token_mask) + tl.store(output_token_ids + token_offsets, global_id, mask=token_mask) + tl.store(output_logsumexp + token_offsets, global_max + tl.log(global_sum), mask=token_mask) + + +@torch.no_grad() +def vocab_parallel_greedy( + local_logits: torch.Tensor, + *, + vocab_size: int, + tp_world_size: int, + group, + alloc_func, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return exact sparse logits, global token ids and full-vocab logsumexp.""" + + assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() + local_vocab_size, token_num = local_logits.shape + assert local_vocab_size in { + vocab_size // tp_world_size, + (vocab_size + tp_world_size - 1) // tp_world_size, + } + + transposed_logits = alloc_func( + (token_num, local_vocab_size), + dtype=local_logits.dtype, + device=local_logits.device, + ) + transpose_convert_2d(local_logits, transposed_logits) + local_stats = greedy_sample_local_stats(transposed_logits, alloc_func=alloc_func) + + if tp_world_size == 1: + gathered_stats = local_stats.view(1, 3, token_num) + else: + gathered_stats = alloc_func((tp_world_size, 3, token_num), dtype=torch.float32, device=local_logits.device) + all_gather_into_tensor( + output_=gathered_stats, + input_=local_stats, + group=group, + async_op=False, + ) + + output_logits = alloc_func((token_num, 1), dtype=torch.float32, device=local_logits.device) + output_token_ids = alloc_func((token_num, 1), dtype=torch.int64, device=local_logits.device) + output_logsumexp = alloc_func((token_num,), dtype=torch.float32, device=local_logits.device) + _combine_vocab_parallel_stats_kernel[(triton.cdiv(token_num, 256),)]( + gathered_stats, + gathered_stats.view(torch.int32), + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size=vocab_size, + tp_world_size=tp_world_size, + BLOCK_SIZE=256, + num_warps=4, + ) + return output_logits, output_token_ids, output_logsumexp diff --git a/lightllm/common/basemodel/triton_kernel/transpose_convert.py b/lightllm/common/basemodel/triton_kernel/transpose_convert.py new file mode 100644 index 0000000000..b618d69526 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/transpose_convert.py @@ -0,0 +1,65 @@ +"""Tiled transpose kernels used by the post-layer logits path.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _transpose_convert_2d_kernel( + input_ptr, + output_ptr, + rows, + cols, + input_stride_0, + input_stride_1, + output_stride_0, + output_stride_1, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + row_offsets = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + col_offsets = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + input_offsets = row_offsets[:, None] * input_stride_0 + col_offsets[None, :] * input_stride_1 + mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) + values = tl.load(input_ptr + input_offsets, mask=mask) + + output_offsets = col_offsets[:, None] * output_stride_0 + row_offsets[None, :] * output_stride_1 + tl.store(output_ptr + output_offsets, tl.trans(values), mask=tl.trans(mask)) + + +@torch.no_grad() +def transpose_convert_2d( + input_tensor: torch.Tensor, + output_tensor: torch.Tensor, + *, + block_rows: int = 64, + block_cols: int = 64, + num_warps: int = 8, + num_stages: int = 1, +) -> torch.Tensor: + """Transpose a contiguous 2-D CUDA tensor while converting its dtype.""" + + assert input_tensor.is_cuda and output_tensor.is_cuda + assert input_tensor.device == output_tensor.device + assert input_tensor.ndim == 2 and output_tensor.ndim == 2 + assert output_tensor.shape == (input_tensor.shape[1], input_tensor.shape[0]) + assert input_tensor.is_contiguous() and output_tensor.is_contiguous() + + rows, cols = input_tensor.shape + grid = (triton.cdiv(rows, block_rows), triton.cdiv(cols, block_cols)) + _transpose_convert_2d_kernel[grid]( + input_tensor, + output_tensor, + rows, + cols, + input_tensor.stride(0), + input_tensor.stride(1), + output_tensor.stride(0), + output_tensor.stride(1), + BLOCK_ROWS=block_rows, + BLOCK_COLS=block_cols, + num_warps=num_warps, + num_stages=num_stages, + ) + return output_tensor diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 6e4b15a55d..11d916b9e7 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -7,6 +7,9 @@ from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import PostLayerInferTpl +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + vocab_parallel_greedy, +) from lightllm.distributed.communication_op import all_gather @@ -64,7 +67,7 @@ def _token_forward( if prompt_logics_hiddens is not None: prompt_token_num = prompt_logics_hiddens.shape[0] infer_state.prompt_logics = self._lm_head_and_gather( - prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state + prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state, force_full_logits=True ) return ans_logics @@ -75,6 +78,7 @@ def _lm_head_and_gather( token_num: int, layer_weight: LlamaPreAndPostLayerWeight, infer_state: LlamaInferStateInfo, + force_full_logits: bool = False, ) -> torch.Tensor: normed = self._norm(hidden, infer_state, layer_weight) normed = normed.permute(1, 0).view(-1, token_num) @@ -82,6 +86,18 @@ def _lm_head_and_gather( normed = None vocab_size = layer_weight.lm_head_weight_.vocab_size + if infer_state.use_vocab_parallel_greedy and not force_full_logits: + logits, token_ids, logsumexp = vocab_parallel_greedy( + logic_batch, + vocab_size=vocab_size, + tp_world_size=self.tp_world_size_, + group=infer_state.dist_group, + alloc_func=self.alloc_tensor, + ) + infer_state.logits_token_ids = token_ids + infer_state.logits_logsumexp = logsumexp + return logits + if self.tp_world_size_ == 1: gather_data = logic_batch else: diff --git a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py index 5a74cd988e..eb4481fd6b 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -181,7 +181,11 @@ def token_forward( logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) block_logits = logits.reshape(num_reqs, self.block_size_, -1) - sampled_tokens = torch.argmax(block_logits, dim=-1) + if infer_state.logits_token_ids is None: + sampled_tokens = torch.argmax(block_logits, dim=-1) + else: + assert block_logits.shape[-1] == 1 + sampled_tokens = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_) confidence_logits = self.predict_confidence_logits( block_hidden, anchor_token_ids=anchor_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 28f2abf74b..3848a6bbcb 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -861,13 +861,23 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: def _gen_argmax_token_ids(self, model_output: ModelOutput): logits = model_output.logits - return torch.argmax(logits, dim=-1) + candidate_indexes = torch.argmax(logits, dim=-1) + return self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): logits = model_output.logits - probs = torch.softmax(logits, dim=-1) - max_probs, draft_next_token_ids_gpu = torch.max(probs, dim=-1) - return draft_next_token_ids_gpu, max_probs + if model_output.has_vocab_parallel_logits: + max_logits, candidate_indexes = torch.max(logits, dim=-1) + token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) + return token_ids, torch.exp(max_logits - model_output.logits_logsumexp) + max_probs, token_ids = torch.max(torch.softmax(logits, dim=-1), dim=-1) + return token_ids, max_probs + + @staticmethod + def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor): + if not model_output.has_vocab_parallel_logits: + return candidate_indexes + return model_output.logits_token_ids.gather(1, candidate_indexes.long().view(-1, 1)).view(-1).long() def _sample_and_scatter_token( self, diff --git a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py index 6b8c23e8fd..137a580ae6 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py @@ -141,7 +141,7 @@ def propose_next_overlap( req_num_by_batch, ) ): - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) draft_token_probs = draft_token_probs.float() diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 3d2c0a0e86..f2c730ccd5 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -81,7 +81,7 @@ def propose_next( # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 6f9477e294..8ec94ae5dc 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -7,20 +7,40 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput +def test_vocab_parallel_metadata_follows_row_selection(): + output = ModelOutput( + logits=torch.tensor([[8.0], [7.0], [6.0]]), + logits_token_ids=torch.tensor([[4], [9], [2]]), + logits_logsumexp=torch.tensor([8.5, 7.25, 6.75]), + ) + + selected = output.index_select_logits_rows(torch.tensor([2, 0])) + + torch.testing.assert_close(selected.logits.view(-1), torch.tensor([6.0, 8.0])) + torch.testing.assert_close(selected.logits_token_ids.view(-1), torch.tensor([2, 4])) + torch.testing.assert_close(selected.logits_logsumexp, torch.tensor([6.75, 8.5])) + + def test_decode_unpad_slices_spec_output_with_logits(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(24).view(6, 4), + logits_token_ids=torch.arange(100, 124).view(6, 4), + logits_logsumexp=torch.arange(6, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), ) unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4) assert unpadded.logits.shape == (4, 4) + assert unpadded.logits_token_ids.shape == (4, 4) + assert unpadded.logits_logsumexp.shape == (4,) assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) # Unpadding returns a shallow output copy and leaves the graph-owned # tensors on the original ModelOutput intact. assert output.logits.shape == (6, 4) + assert output.logits_token_ids.shape == (6, 4) + assert output.logits_logsumexp.shape == (6,) assert output.mtp_collector.spec_hidden.shape == (6, 3) @@ -28,6 +48,8 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(20).view(5, 4), + logits_token_ids=torch.arange(100, 120).view(5, 4), + logits_logsumexp=torch.arange(5, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), prompt_logics=torch.arange(32).view(8, 4), ) @@ -39,6 +61,8 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): ) assert unpadded.logits.shape == (3, 4) + assert unpadded.logits_token_ids.shape == (3, 4) + assert unpadded.logits_logsumexp.shape == (3,) assert unpadded.mtp_collector.spec_hidden.shape == (6, 3) assert unpadded.prompt_logics.shape == (6, 4) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py new file mode 100644 index 0000000000..279268a85c --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py @@ -0,0 +1,65 @@ +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") + + +@pytest.mark.parametrize("token_num", [1, 7, 64]) +def test_vocab_parallel_greedy_matches_full_logits(monkeypatch, token_num): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy") + tp_world_size = 4 + local_vocab_size = 8192 + vocab_size = tp_world_size * local_vocab_size + generator = torch.Generator(device="cuda").manual_seed(20260826 + token_num) + local_logits_by_rank = [ + torch.randn( + (local_vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + for _ in range(tp_world_size) + ] + + # Exercise deterministic tie-breaking both between local reduction blocks + # and across tensor-parallel ranks. The smallest global token id must win. + local_logits_by_rank[0][4097, 0] = 20.0 + local_logits_by_rank[0][3, 0] = 20.0 + local_logits_by_rank[3][2, 0] = 20.0 + + local_stats_by_rank = [ + greedy_sample_local_stats(local_logits.transpose(0, 1).contiguous()) for local_logits in local_logits_by_rank + ] + + def fake_all_gather_into_tensor(output_, input_, **_kwargs): + for output, local_stats in zip(output_, local_stats_by_rank): + output.copy_(local_stats) + + monkeypatch.setattr(module, "all_gather_into_tensor", fake_all_gather_into_tensor) + actual_logits, actual_ids, actual_logsumexp = module.vocab_parallel_greedy( + local_logits_by_rank[0], + vocab_size=vocab_size, + tp_world_size=tp_world_size, + group=None, + alloc_func=torch.empty, + ) + actual_ids = actual_ids.view(-1) + actual_logprobs = actual_logits.view(-1) - actual_logsumexp + + full_logits = torch.cat(local_logits_by_rank, dim=0).transpose(0, 1).float() + expected_ids = full_logits.argmax(dim=1) + expected_logits = full_logits.gather(1, expected_ids[:, None]).view(-1) + expected_logsumexp = torch.logsumexp(full_logits, dim=1) + expected_logprobs = torch.log_softmax(full_logits, dim=1).gather(1, expected_ids[:, None]).squeeze(1) + + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual_logits.view(-1), expected_logits, rtol=0, atol=0) + torch.testing.assert_close(actual_logsumexp, expected_logsumexp, rtol=2e-4, atol=2e-4) + torch.testing.assert_close(actual_logprobs, expected_logprobs, rtol=2e-4, atol=2e-4) diff --git a/unit_tests/models/test_vocab_parallel_greedy_output.py b/unit_tests/models/test_vocab_parallel_greedy_output.py new file mode 100644 index 0000000000..88e5a7c282 --- /dev/null +++ b/unit_tests/models/test_vocab_parallel_greedy_output.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + + +def test_argmax_restores_global_token_ids_and_exact_probabilities(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput( + logits=torch.tensor([[3.0, 1.0], [0.0, 5.0]]), + logits_token_ids=torch.tensor([[30, 10], [100, 500]]), + logits_logsumexp=torch.tensor([4.0, 5.25]), + ) + + token_ids = backend._gen_argmax_token_ids(output) + token_ids_with_prob, probs = backend._gen_argmax_token_ids_and_prob(output) + + torch.testing.assert_close(token_ids, torch.tensor([30, 500])) + torch.testing.assert_close(token_ids_with_prob, token_ids) + torch.testing.assert_close(probs, torch.exp(torch.tensor([-1.0, -0.25]))) + + +def test_dense_argmax_keeps_column_index_semantics(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput(logits=torch.tensor([[1.0, 4.0, 2.0]])) + + torch.testing.assert_close(backend._gen_argmax_token_ids(output), torch.tensor([1])) + + +def test_dspark_confidence_path_receives_global_token_ids(): + post = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post.block_size_ = 2 + post.markov_rank_ = 0 + post._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) + sparse_logits = torch.tensor([[4.0], [5.0], [7.0], [9.0]]) + sparse_token_ids = torch.tensor([[40], [50], [70], [90]]) + + def gather_vocab_parallel(*args, **kwargs): + infer_state = args[3] + infer_state.logits_token_ids = sparse_token_ids + return sparse_logits + + post._lm_head_and_gather = gather_vocab_parallel + observed = {} + + def predict_confidence(block_hidden, anchor_token_ids, sampled_tokens, layer_weight): + observed["sampled_tokens"] = sampled_tokens + return None + + post.predict_confidence_logits = predict_confidence + + class Collector: + def add_mtp_outputs(self, **kwargs): + self.outputs = kwargs + + collector = Collector() + infer_state = SimpleNamespace( + is_prefill=False, + input_ids=torch.tensor([1, 0, 2, 0]), + logits_token_ids=None, + hidden_collector=collector, + ) + + returned_logits = post.token_forward( + input_embdings=torch.ones((4, 3)), + infer_state=infer_state, + layer_weight=object(), + ) + + torch.testing.assert_close(returned_logits, sparse_logits) + torch.testing.assert_close(observed["sampled_tokens"], torch.tensor([[40, 50], [70, 90]])) From 79f67134fe03088eb53cddff1a1ade3cad3f7c39 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 29 Aug 2026 02:48:13 +0800 Subject: [PATCH 13/28] feat(glm5): add multimodal and 1m deployment support --- GLM53_H100_DEPLOY.md | 107 ++++++++++++---- docker/Dockerfile.glm53-h100 | 41 ++++--- docker/requirements-glm53-runtime.txt | 2 + .../common/basemodel/attention/linear/kda.py | 21 +--- lightllm/common/basemodel/basemodel.py | 33 ++--- lightllm/common/basemodel/cuda_graph.py | 10 +- .../transformer_layer_infer_template.py | 3 +- .../fused_moe/fused_moe_weight.py | 15 +-- .../fused_moe/impl/deepgemm_impl.py | 15 +-- .../fused_moe/impl/triton_impl.py | 15 +-- .../common/basemodel/prefill_cuda_graph.py | 26 +--- .../fused_moe/deepep_legacy_layout.py | 16 +-- .../fused_moe/grouped_fused_moe_ep.py | 24 +--- .../triton_kernel/fused_moe/grouped_topk.py | 43 ++----- .../moe_silu_and_mul_mix_quant_ep.py | 4 +- .../linear_att/fla/ops/fused_recurrent.py | 4 +- .../triton_kernel/linear_att/fla/ops/kda.py | 70 +++-------- lightllm/distributed/communication_op.py | 17 ++- lightllm/distributed/symm_mem_all_reduce.py | 8 +- .../layer_infer/transformer_layer_infer.py | 28 ++--- .../layer_infer/transformer_layer_infer.py | 12 +- .../triton_kernel/extract_indexer_ks.py | 12 +- lightllm/models/glm5_next/glm5_next_visual.py | 115 ++++++++++++++++++ .../glm5_next/layer_infer/pre_layer_infer.py | 7 ++ .../layer_infer/transformer_layer_infer.py | 93 ++++---------- .../layer_weights/transformer_layer_weight.py | 4 +- lightllm/models/glm5_next/mem_manager.py | 4 +- lightllm/models/glm5_next/model.py | 24 ++-- lightllm/models/glm5_next/tokenizer.py | 73 +++++++++++ .../models/glm5_next/triton_kernel/mhc.py | 109 ++++------------- lightllm/models/glm5_next_mtp/model.py | 14 +-- .../mode_backend/chunked_prefill/impl.py | 4 +- lightllm/server/tokenizer.py | 8 ++ .../visualserver/model_infer/model_rpc.py | 5 + lightllm/utils/config_utils.py | 2 + lightllm/utils/envs_utils.py | 6 +- test/kernel/test_glm5_grouped_topk.py | 48 ++------ test/kernel/test_glm5_mhc.py | 82 ++++--------- test/kernel/test_glm5_sglang_moe_compat.py | 27 ++-- test/kernel/test_glm5_strided_causal_conv.py | 31 ++--- test/kernel/test_glm5_vocab_parallel_top1.py | 2 + test/test_glm5_next_multimodal.py | 96 +++++++++++++++ test/test_moe_prefill_dispatch.py | 4 +- tools/analyze_torch_trace.py | 10 +- tools/bench_glm53_allreduce.py | 36 +++--- tools/bench_glm53_kda_chunk_h.py | 8 +- tools/bench_glm53_sglang_moe.py | 18 +-- tools/bench_glm53_sparse_prefill.py | 14 +-- tools/bench_glm53_sparse_prefill_tp.py | 21 +--- tools/check_glm53_symm_out_of_place.py | 4 +- tools/run_glm53_h100_container.sh | 8 +- .../linear_att/test_kda_fused_gate.py | 8 +- 52 files changed, 677 insertions(+), 734 deletions(-) create mode 100644 lightllm/models/glm5_next/glm5_next_visual.py create mode 100644 lightllm/models/glm5_next/layer_infer/pre_layer_infer.py create mode 100644 lightllm/models/glm5_next/tokenizer.py create mode 100644 test/test_glm5_next_multimodal.py diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 4af5d1e784..6d2863fd58 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -1,17 +1,27 @@ -# GLM-5.3-Flash on H100 TP8 +# GLM-5.3-Flash multimodal TP8 deployment -This branch includes a self-contained LightLLM image variant for one eight-GPU -H100 node. The image contains the LightLLM source and runtime dependency; only -the model and compiler-cache directories are mounted from the host. +This branch packages LightLLM text and vision inference for GLM-5.3-Flash on +one eight-GPU H100 or H200 node. The image contains the LightLLM source and +runtime dependencies. Mount the model and compiler caches from the host. -## Build +The default command serves the OpenAI-compatible API on port 8002 with: + +- tensor parallel size 8; +- image encoding as eight data-parallel workers; +- a 1,048,576-token request limit; +- up to 256 active requests; +- 8,192-token chunked prefill, which bounds the DSA score matrix during + million-token requests, with CUDA graphs disabled for this profile; +- `glm45` reasoning and `glm47` tool-call parsers. + +## Build a local image Use an immutable tag containing the full source revision: ```bash revision="$(git rev-parse HEAD)" created="$(date -u +%Y-%m-%dT%H:%M:%SZ)" -version="v1.2.0-h100-tp8-${revision:0:12}" +version="v1.3.0-glm53-vl-1m-tp8-${revision}" docker buildx build --load --platform linux/amd64 \ -f docker/Dockerfile.glm53-h100 \ @@ -20,40 +30,83 @@ docker buildx build --load --platform linux/amd64 \ --build-arg "OCI_VERSION=${version}" \ -t "lightllm-glm53:${version}" \ . + +docker tag "lightllm-glm53:${version}" lightllm-glm53:vl-1m-tp8 ``` -After verification, the host-local convenience alias may point to the same -image ID: +## Run on the local H200 node ```bash -docker tag "lightllm-glm53:${version}" lightllm-glm53:h100-tp8 +LIGHTLLM_GLM53_IMAGE=lightllm-glm53:vl-1m-tp8 \ +LIGHTLLM_GLM53_MODEL_DIR=/nvme/sufubao/models/GLM-5.3-Flash \ +LIGHTLLM_GLM53_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-lightllm-h200 \ +LIGHTLLM_GLM53_TRITON_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-triton-h200 \ +LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-deep-gemm-h200 \ +tools/run_glm53_h100_container.sh ``` -## Run on h100 +## Run the published image on H100 -The default image command is the measured no-speculation, concurrency-256 -profile on port 8002. Run it in the foreground with: +Set `IMAGE` to the immutable registry tag or digest listed in the pull request. +The fixed container name is `glm53-lightllm-vl-1m`. The image default uses the +8,192-token prefill chunk validated on H200. The command below overrides the +default with a conservative 1,024-token chunk for an 80 GB H100, where the DSA +score matrix has much less temporary-memory headroom. ```bash -LIGHTLLM_GLM53_IMAGE="lightllm-glm53:${version}" \ - tools/run_glm53_h100_container.sh +IMAGE=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm: + +sudo docker pull "$IMAGE" +sudo docker run -d \ + --name glm53-lightllm-vl-1m \ + --restart unless-stopped \ + --gpus all \ + --ipc=host \ + --network=host \ + --ulimit memlock=-1:-1 \ + --ulimit nofile=1048576:1048576 \ + -v /home/devsft/models/GLM-5.3-Flash:/model:ro \ + -v /home/devsft/cache-glm53-lightllm:/root/.cache \ + -v /home/devsft/cache-glm53-triton:/root/.triton \ + -v /home/devsft/cache-glm53-deep-gemm:/root/.deep_gemm \ + "$IMAGE" \ + /opt/sglang/bin/python -m lightllm.server.api_server \ + --model_dir /model \ + --model_name glm-5.3-flash \ + --tp 8 \ + --host 0.0.0.0 \ + --port 8002 \ + --httpserver_workers 16 \ + --mem_fraction .90 \ + --max_total_token_num 1048612 \ + --running_max_req_size 256 \ + --max_req_total_len 1048576 \ + --batch_max_tokens 65536 \ + --chunked_prefill_size 1024 \ + --linear_att_ssm_data_type bfloat16 \ + --linear_att_cache_size 256 \ + --disable_cudagraph \ + --enable_fused_shared_experts \ + --max_image_pixels 6272000 \ + --max_image_token_count 8000 \ + --visual_tp 1 \ + --visual_dp 8 \ + --visual_infer_batch_size 8 \ + --cache_capacity 64 \ + --schedule_time_interval 0.001 \ + --prefill_coalesce_interval 0.5 \ + --reasoning_parser glm45 \ + --tool_call_parser glm47 ``` -Override `LIGHTLLM_GLM53_MODEL_DIR`, `LIGHTLLM_GLM53_CACHE_DIR`, or -`LIGHTLLM_GLM53_TRITON_CACHE_DIR` when the host paths differ. In another shell, -wait for the model list endpoint: +Wait for the model endpoint, then stop the deployment when required: ```bash curl --fail --show-error http://127.0.0.1:8002/v1/models +sudo docker stop --timeout 30 glm53-lightllm-vl-1m ``` -Stop the foreground process with `Ctrl-C`. If it was detached externally, use -`sudo docker stop --timeout 30 glm53-lightllm`. - -## Measured profile - -The final pre-release candidate reached 4169.40 output tokens/s at concurrency -256 with random 1024-token inputs and 256-token outputs. This is 5.74% below -the measured vLLM result, so the earlier three-percent stretch goal remains -unmet at concurrency 256. Concurrency 16 and 64 exceeded the corresponding -vLLM and SGLang measurements when run with MTP2. +The H200 accuracy, long-context, and throughput results in the pull request are +measured through the OpenAI-compatible endpoint. H100 performance is not +inferred from H200 data, and the conservative H100 override above must be +validated independently on the target host before production traffic. diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index f5d2a744e8..9d25686bed 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -1,9 +1,9 @@ -# syntax=docker/dockerfile:1.7 - -# This is the exact SGLang GLM-5.3 runtime used for the H100 measurements. -ARG BASE_IMAGE=lmsysorg/sglang@sha256:e6f5482505e7502f791fe4615ad1fbec118cbbd6b44e98f2479b16b98b985ad6 +# This is the exact SGLang GLM-5.3 runtime used for the H200 validation. +ARG BASE_IMAGE=lmsysorg/sglang@sha256:92afb4c878eef9cbb17ca9a2c1d15d5cda58585f90bbf5915a79f0f6284aad10 FROM ${BASE_IMAGE} AS prepared +ARG PIP_INDEX_URL + WORKDIR /opt/lightllm COPY docker/requirements-glm53-runtime.txt /tmp/requirements-glm53-runtime.txt @@ -32,7 +32,7 @@ RUN python -m pip install --no-cache-dir --no-deps . && \ /opt/sglang/lib/python3.12/site-packages/setuptools-*.dist-info \ /opt/sglang/lib/python3.12/site-packages/wheel \ /opt/sglang/lib/python3.12/site-packages/wheel-*.dist-info && \ - python -c "import frozendict, hypercorn, importlib.metadata as metadata, importlib.util, lightllm.server.api_start, msgpack, rpyc, ujson; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; assert metadata.version('hypercorn') == '0.18.0'; assert msgpack.__version__ == '1.2.2'; assert rpyc.__version__ == '6.0.2'; assert importlib.util.find_spec('pip') is None; assert importlib.util.find_spec('setuptools') is None; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ + python -c "import frozendict, hypercorn, importlib.metadata as metadata, importlib.util, lightllm.server.api_start, msgpack, redis, rpyc, ujson; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; assert metadata.version('hypercorn') == '0.18.0'; assert msgpack.__version__ == '1.2.2'; assert redis.__version__ == '7.3.0'; assert rpyc.__version__ == '6.0.2'; assert importlib.util.find_spec('pip') is None; assert importlib.util.find_spec('setuptools') is None; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ test ! -e /etc/ssh/ssh_host_rsa_key && \ test ! -e /sgl-workspace/sglang/python/sglang/multimodal_gen && \ test ! -e /sgl-workspace/sglang/python/sglang/srt/disaggregation @@ -47,7 +47,7 @@ ARG OCI_REVISION ARG OCI_SOURCE=https://github.com/sufubao/LightLLM ARG OCI_VERSION ARG BASE_NAME=lmsysorg/sglang:glm-5.3-flash -ARG BASE_DIGEST=sha256:e6f5482505e7502f791fe4615ad1fbec118cbbd6b44e98f2479b16b98b985ad6 +ARG BASE_DIGEST=sha256:92afb4c878eef9cbb17ca9a2c1d15d5cda58585f90bbf5915a79f0f6284aad10 COPY --from=prepared / / @@ -55,13 +55,13 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ org.opencontainers.image.source="${OCI_SOURCE}" \ org.opencontainers.image.version="${OCI_VERSION}" \ - org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100 TP8" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 multimodal" \ org.opencontainers.image.base.name="${BASE_NAME}" \ org.opencontainers.image.base.digest="${BASE_DIGEST}" \ ai.lightllm.model="GLM-5.3-Flash" \ - ai.lightllm.accelerator="NVIDIA H100 80GB" \ + ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="throughput-c256" \ + ai.lightllm.profile="multimodal-1m-c256" \ ai.lightllm.security-profile="flattened-no-sglang-server-components" ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ @@ -94,7 +94,7 @@ ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sb WORKDIR /opt/lightllm RUN test "$(command -v python)" = /opt/sglang/bin/python && \ - python -c "import hypercorn, importlib.metadata as metadata, lightllm.server.api_start, rpyc, torch; assert metadata.version('hypercorn') == '0.18.0'; assert rpyc.__version__ == '6.0.2'; print(torch.__version__)" + python -c "import hypercorn, importlib.metadata as metadata, lightllm.server.api_start, redis, rpyc, torch; assert metadata.version('hypercorn') == '0.18.0'; assert redis.__version__ == '7.3.0'; assert rpyc.__version__ == '6.0.2'; print(torch.__version__)" EXPOSE 8002 STOPSIGNAL SIGTERM @@ -103,6 +103,10 @@ ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ CMD /opt/sglang/bin/python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 +# LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is +# therefore request_limit + 36 so /v1/models advertises the full 1,048,576. +# The 8,192-token DSA prefill chunk below is the H200-validated default; use +# the conservative 1,024-token override documented for an 80 GB H100. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -111,16 +115,21 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--port", "8002", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ - "--max_total_token_num", "335000", \ + "--max_total_token_num", "1048612", \ "--running_max_req_size", "256", \ - "--max_req_total_len", "65500", \ + "--max_req_total_len", "1048576", \ "--batch_max_tokens", "65536", \ - "--chunked_prefill_size", "65536", \ + "--chunked_prefill_size", "8192", \ "--linear_att_ssm_data_type", "bfloat16", \ - "--graph_max_batch_size", "256", \ - "--graph_split_batch_size", "4", \ - "--graph_grow_step_size", "16", \ + "--linear_att_cache_size", "256", \ + "--disable_cudagraph", \ "--enable_fused_shared_experts", \ + "--max_image_pixels", "6272000", \ + "--max_image_token_count", "8000", \ + "--visual_tp", "1", \ + "--visual_dp", "8", \ + "--visual_infer_batch_size", "8", \ + "--cache_capacity", "64", \ "--schedule_time_interval", "0.001", \ "--prefill_coalesce_interval", "0.5", \ "--reasoning_parser", "glm45", \ diff --git a/docker/requirements-glm53-runtime.txt b/docker/requirements-glm53-runtime.txt index 8a0e41751e..08740b3440 100644 --- a/docker/requirements-glm53-runtime.txt +++ b/docker/requirements-glm53-runtime.txt @@ -20,6 +20,8 @@ priority==2.0.0 \ --hash=sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa rpyc==6.0.2 \ --hash=sha256:8072308ad30725bc281c42c011fc8c922be15f3eeda6eafb2917cafe1b6f00ec +redis==7.3.0 \ + --hash=sha256:9d4fcb002a12a5e3c3fbe005d59c48a2cc231f87fbb2f6b70c2d89bb64fec364 setuptools==84.0.0 \ --hash=sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670 ujson==5.13.0 \ diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index b596d18e3f..ac31c02c73 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -67,10 +67,7 @@ class KDAPrefillAttState(BasePrefillAttState): def init_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) - self.seq_lens_cpu = ( - self.infer_state.b1_cu_q_seq_len[1:] - - self.infer_state.b1_cu_q_seq_len[:-1] - ).tolist() + self.seq_lens_cpu = (self.infer_state.b1_cu_q_seq_len[1:] - self.infer_state.b1_cu_q_seq_len[:-1]).tolist() # prepare_chunk_indices performs a GPU-to-CPU shape sync. Build it # before entering CUDA Graph capture and copy its fixed-size contents # through BasePrefillAttState on replay. @@ -120,9 +117,7 @@ def prefill_att( q=q, k=k, v=v, - raw_g=raw_gate.view( - 1, -1, backend.tp_num_heads, backend.head_dim - ), + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), beta=raw_beta.float().sigmoid(), A_log=layer_weight.linear_A_log.weight, g_bias=layer_weight.linear_dt_bias.weight, @@ -134,9 +129,7 @@ def prefill_att( safe_gate=True, lower_bound=backend.lower_bound, ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to( - ssm_states.dtype, copy=False - ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) return output @@ -148,9 +141,7 @@ class KDADecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def init_state(self): - draft_step = self.backend.model.mtp_manager.get_decode_draft_step( - self.backend.model.is_mtp_draft_model - ) + draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step == 0: self._init_normal_decode_state() elif self.backend.uses_dynamic_spec_verify_layout(): @@ -192,9 +183,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[ - self.b_conv_buffer_idx - ] + 1 + self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f9278362bb..4aafcc8663 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -101,15 +101,9 @@ def __init__(self, kvargs): ) self.logical_graph_max_batch_size = self.graph_max_batch_size self.mtp_manager = MtpManager.get_instance() - self.decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier( - self.is_mtp_draft_model - ) - cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier( - self.is_mtp_draft_model - ) - self.graph_max_batch_size = ( - self.graph_max_batch_size * cuda_graph_batch_multiplier - ) + self.decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier(self.is_mtp_draft_model) + cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier(self.is_mtp_draft_model) + self.graph_max_batch_size = self.graph_max_batch_size * cuda_graph_batch_multiplier self.graph_max_len_in_batch = kvargs.get("graph_max_len_in_batch", 8192) self.disable_cudagraph = kvargs.get("disable_cudagraph", False) @@ -297,9 +291,7 @@ def _init_att_backend1(self): return def _init_cudagraph(self): - cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier( - self.is_mtp_draft_model - ) + cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier(self.is_mtp_draft_model) cuda_graph_grow_step_size = self.mtp_manager.get_decode_cuda_graph_grow_step_size(self.is_mtp_draft_model) extra_batch_sizes = None if self.mtp_manager.draft_model_needs_logical_batch_graphs(self.is_mtp_draft_model): @@ -700,13 +692,8 @@ def _decode( # 向上对齐到 TP world size 的整数倍，保证后续切分得到合法 shape。 infer_batch_size = max(1, origin_batch_size) if self.args.enable_tpsp_mix_mode: - decode_alignment = math.lcm( - self.tp_world_size_, self.decode_batch_multiplier - ) - infer_batch_size = ( - triton.cdiv(infer_batch_size, decode_alignment) - * decode_alignment - ) + decode_alignment = math.lcm(self.tp_world_size_, self.decode_batch_multiplier) + infer_batch_size = triton.cdiv(infer_batch_size, decode_alignment) * decode_alignment # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 @@ -983,12 +970,8 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 origin_batch_size1 = model_input1.batch_size max_len_in_batch = max(2, model_input0.max_kv_seq_len, model_input1.max_kv_seq_len) infer_batch_size = max(1, origin_batch_size0, origin_batch_size1) - decode_alignment = math.lcm( - self.tp_world_size_, self.decode_batch_multiplier - ) - infer_batch_size = ( - triton.cdiv(infer_batch_size, decode_alignment) * decode_alignment - ) + decode_alignment = math.lcm(self.tp_world_size_, self.decode_batch_multiplier) + infer_batch_size = triton.cdiv(infer_batch_size, decode_alignment) * decode_alignment if ( self._is_cuda_graph_output_compatible(model_input0, model_input1) diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index acaab19f1a..b2467f7dfc 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -54,9 +54,7 @@ def gen_cuda_graph_batch_sizes( # block of ``batch_step_size_before_split`` rows. TP/SP padding # must preserve both that block and an even split across TP ranks. alignment = math.lcm(tp_world_size, batch_step_size_before_split) - batch_sizes = sorted( - {triton.cdiv(size, alignment) * alignment for size in batch_sizes} - ) + batch_sizes = sorted({triton.cdiv(size, alignment) * alignment for size in batch_sizes}) assert batch_sizes[-1] == max_batch_size return batch_sizes @@ -92,11 +90,7 @@ def __init__( if extra_batch_sizes is not None: self.cuda_graph_batch_sizes = sorted( set(self.cuda_graph_batch_sizes) - | { - int(batch_size) - for batch_size in extra_batch_sizes - if 0 < int(batch_size) <= self.max_batch_size - } + | {int(batch_size) for batch_size in extra_batch_sizes if 0 < int(batch_size) <= self.max_batch_size} ) logger.info(f"cuda graph batch_sizes: {self.cuda_graph_batch_sizes}") diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index 5bf32bca2e..da57c8f81b 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -134,8 +134,7 @@ def _context_attention_wrapper_run( and all(isinstance(value, torch.Tensor) for value in attr_value.values()) ): callback_tensor_dicts[attr_name] = { - key: tensor_to_no_ref_tensor(value.contiguous()) - for key, value in attr_value.items() + key: tensor_to_no_ref_tensor(value.contiguous()) for key, value in attr_value.items() } pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() pre_capture_graph.__exit__(None, None, None) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 5a7a9d53d8..367e424a22 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -84,9 +84,7 @@ def _init_config(self, network_config: Dict[str, Any]): self.scoring_func = network_config.get("scoring_func", "softmax") self.swiglu_limit = network_config.get("swiglu_limit") self.swiglu_alpha = network_config.get("swiglu_alpha", 1.0) - self.swiglu_clamp_up_add_one = network_config.get( - "swiglu_clamp_up_add_one", True - ) + self.swiglu_clamp_up_add_one = network_config.get("swiglu_clamp_up_add_one", True) def _init_redundancy_expert_params(self): self.redundancy_expert_num = get_redundancy_expert_num() @@ -294,13 +292,10 @@ def load_hf_weights(self, weights): def verify_load(self): if getattr(self, "_sm90_mega_moe_weights_prepared", False): - weight_load_ok = all( - all(_weight_pack.load_ok) for _weight_pack in self.w2_list - ) + weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w2_list) else: weight_load_ok = all( - all(_weight_pack.load_ok) - for _weight_pack in self.w1_list + self.w2_list + self.w3_list + all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list ) per_expert_scale_load_ok = ( True if self.per_expert_scale is None else getattr(self.per_expert_scale, "load_ok", False) @@ -309,9 +304,7 @@ def verify_load(self): True if self.e_score_correction_bias is None else getattr(self.e_score_correction_bias, "load_ok", False) ) load_ok = weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok - if load_ok and self.enable_ep_moe and not getattr( - self, "_sm90_mega_moe_weights_prepared", False - ): + if load_ok and self.enable_ep_moe and not getattr(self, "_sm90_mega_moe_weights_prepared", False): from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( prepare_sm90_mega_moe_weights, use_sm90_mega_moe, diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 955067942f..1878b4071e 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -79,11 +79,7 @@ def _fused_experts( router_logits: Optional[torch.Tensor] = None, is_prefill: Optional[bool] = None, ): - fused_topk_ids = ( - topk_ids - if use_sm90_mega_moe(self.quant_method) - else topk_ids.to(torch.long) - ) + fused_topk_ids = topk_ids if use_sm90_mega_moe(self.quant_method) else topk_ids.to(torch.long) output = fused_experts( hidden_states=input_tensor, w13=w13, @@ -187,14 +183,7 @@ def dispatch( async_finish=False, allocate_on_comm_stream=False, ) - ( - recv_x, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - handle, - _, - ) = buffer.dispatch( + (recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, _,) = buffer.dispatch( qinput_tensor, topk_idx=topk_idx, topk_weights=topk_weights, diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index d13f633fb2..cbeea19baa 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -259,13 +259,8 @@ def _fused_experts( "LIGHTLLM_USE_SGLANG_TRITON_MOE currently requires " "block-wise FP8 expert weights with block size 128" ) - if ( - getattr(self, "swiglu_limit", None) is not None - and getattr(self, "swiglu_clamp_up_add_one", True) - ): - raise RuntimeError( - "SGLang Triton MoE does not support clamp_up_add_one=True" - ) + if getattr(self, "swiglu_limit", None) is not None and getattr(self, "swiglu_clamp_up_add_one", True): + raise RuntimeError("SGLang Triton MoE does not support clamp_up_add_one=True") sglang_fused_moe, override_config = _get_sglang_fused_experts_impl() tuned_configs = _get_sglang_triton_moe_configs( @@ -316,11 +311,7 @@ def _fused_experts( w1_scale=w13_scale, w2_scale=w2_scale, limit=getattr(self, "swiglu_limit", None), - alpha=( - getattr(self, "swiglu_alpha", 1.0) - if getattr(self, "swiglu_limit", None) is not None - else None - ), + alpha=(getattr(self, "swiglu_alpha", 1.0) if getattr(self, "swiglu_limit", None) is not None else None), clamp_up_add_one=getattr(self, "swiglu_clamp_up_add_one", True), ) return input_tensor diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index aad9a97802..244bfd0114 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -48,10 +48,7 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): self.exact_batch_size_by_token_num = {} if self.use_exact_token_nums: if configured_batch_sizes is None: - raise ValueError( - "--prefill_cudagraph_batch_sizes is required with " - "--prefill_cudagraph_token_nums" - ) + raise ValueError("--prefill_cudagraph_batch_sizes is required with " "--prefill_cudagraph_token_nums") if len(configured_token_nums) != len(configured_batch_sizes): raise ValueError( "--prefill_cudagraph_token_nums and --prefill_cudagraph_batch_sizes " @@ -86,9 +83,7 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): ) else: if configured_batch_sizes is not None: - raise ValueError( - "--prefill_cudagraph_batch_sizes requires --prefill_cudagraph_token_nums" - ) + raise ValueError("--prefill_cudagraph_batch_sizes requires --prefill_cudagraph_token_nums") graph_handle_token_nums = ( list(range(4, 33, 4)) + list(range(48, 257, 16)) @@ -113,8 +108,7 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): logger.info(f"prefill cuda graph graph_handle_token_nums: {self.graph_handle_token_nums}") if self.exact_batch_size_by_token_num: logger.info( - "prefill cuda graph exact layouts (token_num -> batch_size): " - f"{self.exact_batch_size_by_token_num}" + "prefill cuda graph exact layouts (token_num -> batch_size): " f"{self.exact_batch_size_by_token_num}" ) def can_run( @@ -130,11 +124,7 @@ def can_run( if configured_batch_size is None or batch_size != configured_batch_size: return False uniform_seq_len = handle_token_num // configured_batch_size - return ( - max_q_seq_len == uniform_seq_len - and max_kv_seq_len == uniform_seq_len - and max_cache_len == 0 - ) + return max_q_seq_len == uniform_seq_len and max_kv_seq_len == uniform_seq_len and max_cache_len == 0 return handle_token_num <= self.max_handle_token_num def need_capture(self, handle_token_num: int): @@ -282,16 +272,12 @@ def warmup(self, model): total_token_num = handle_token_num input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int64, device="cuda") mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda() - b_req_idx = torch.full( - (batch_size,), model.req_manager.HOLD_REQUEST_ID, dtype=torch.int32, device="cuda" - ) + b_req_idx = torch.full((batch_size,), model.req_manager.HOLD_REQUEST_ID, dtype=torch.int32, device="cuda") b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda") b_is_decode_req = torch.zeros(batch_size, dtype=torch.bool, device="cuda") b_ready_cache_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda") - b_prefill_start_loc = torch.arange( - 0, total_token_num, seq_len, dtype=torch.int32, device="cuda" - ) + b_prefill_start_loc = torch.arange(0, total_token_num, seq_len, dtype=torch.int32, device="cuda") model_input = ModelInput( batch_size=batch_size, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py index ac53255d4f..0a40488422 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py @@ -91,9 +91,7 @@ def _ep_scatter_tokens( destination_int32, ) tl.store( - output_tensor - + destination * output_tensor_stride0 - + hidden_offsets * output_tensor_stride1, + output_tensor + destination * output_tensor_stride0 + hidden_offsets * output_tensor_stride1, token, mask=hidden_mask, ) @@ -197,18 +195,12 @@ def _ep_gather_kernel( accumulator = tl.zeros([block_hidden], dtype=tl.float32) for topk_offset_int32 in range(0, topk_num): topk_offset = topk_offset_int32.to(tl.int64) - expert_id = tl.load( - recv_topk_ids + token * recv_topk_ids_stride0 + topk_offset * recv_topk_ids_stride1 - ) + expert_id = tl.load(recv_topk_ids + token * recv_topk_ids_stride0 + topk_offset * recv_topk_ids_stride1) if expert_id >= 0: - source_int32 = tl.load( - input_index + token * input_index_stride0 + topk_offset * input_index_stride1 - ) + source_int32 = tl.load(input_index + token * input_index_stride0 + topk_offset * input_index_stride1) source = source_int32.to(tl.int64) weight = tl.load( - recv_topk_weights - + token * recv_topk_weights_stride0 - + topk_offset * recv_topk_weights_stride1 + recv_topk_weights + token * recv_topk_weights_stride0 + topk_offset * recv_topk_weights_stride1 ) value = tl.load( input_tensor diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index a12da6488f..1eef8114f0 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -58,8 +58,7 @@ def use_sm100_mega_moe(quant_method: Any) -> bool: def use_sm90_mega_moe(quant_method: Any) -> bool: return ( is_sm90_gpu() - and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() - in {"1", "ON", "TRUE"} + and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() in {"1", "ON", "TRUE"} and quant_method.method_name == "fp8w8a8-b128-deepgemm" and HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_mega_moe") @@ -162,15 +161,9 @@ def prepare_sm90_mega_moe_weights(w13: Any) -> None: granularity = 8 half = n // 2 assert half % granularity == 0 - gate = weight[:, :half].reshape( - num_groups, half // granularity, granularity, *rest - ) - up = weight[:, half:].reshape( - num_groups, half // granularity, granularity, *rest - ) - w13.weight = torch.stack((gate, up), dim=2).reshape( - num_groups, n, *rest - ) + gate = weight[:, :half].reshape(num_groups, half // granularity, granularity, *rest) + up = weight[:, half:].reshape(num_groups, half // granularity, granularity, *rest) + w13.weight = torch.stack((gate, up), dim=2).reshape(num_groups, n, *rest) w13.sm90_mega_moe_prepared = True @@ -520,14 +513,7 @@ def fused_experts_impl( async_finish=False, allocate_on_comm_stream=False, ) - ( - recv_x, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - handle, - _, - ) = buffer.dispatch( + (recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, _,) = buffer.dispatch( (qinput_tensor, input_scale), topk_idx=topk_idx, topk_weights=topk_weights, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py index 2f30054793..a8bc8b9e58 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py @@ -123,22 +123,16 @@ def single_group_sigmoid_topk_kernel( ).to(tl.float32) old_scores = tl.sigmoid(hidden_states) if HAS_CORRECTION_BIAS: - scores = old_scores + tl.load( - correction_bias_ptr + offs_n, mask=valid, other=0.0 - ) + scores = old_scores + tl.load(correction_bias_ptr + offs_n, mask=valid, other=0.0) else: scores = old_scores scores = tl.where(valid, scores, -float("inf")) for topk_index in tl.static_range(0, TOPK_NUM): selected_index = tl.argmax(scores, axis=0) - selected_weight = tl.sum( - tl.where(offs_n == selected_index, old_scores, 0.0), axis=0 - ) + selected_weight = tl.sum(tl.where(offs_n == selected_index, old_scores, 0.0), axis=0) tl.store( - out_topk_weights - + token_index * out_topk_weights_stride_m - + topk_index, + out_topk_weights + token_index * out_topk_weights_stride_m + topk_index, selected_weight, ) tl.store( @@ -151,17 +145,13 @@ def single_group_sigmoid_topk_kernel( topk_offs = tl.arange(0, TOPK_BLOCK_SIZE) topk_mask = topk_offs < TOPK_NUM weights = tl.load( - out_topk_weights - + token_index * out_topk_weights_stride_m - + topk_offs, + out_topk_weights + token_index * out_topk_weights_stride_m + topk_offs, mask=topk_mask, other=0.0, ) weight_sum = tl.sum(weights, axis=0) tl.store( - out_topk_weights - + token_index * out_topk_weights_stride_m - + topk_offs, + out_topk_weights + token_index * out_topk_weights_stride_m + topk_offs, weights / weight_sum, mask=topk_mask, ) @@ -201,13 +191,9 @@ def single_group_sigmoid_topk_bitonic_kernel( else: scores = old_scores - _, sorted_scores, sorted_indexes = argsort( - scores, old_scores, offs_n, descending=True - ) + _, sorted_scores, sorted_indexes = argsort(scores, old_scores, offs_n, descending=True) if RENORMALIZE: - sum_scores = tl.sum( - tl.where(offs_n < TOPK_NUM, sorted_scores, 0.0) - ) + sum_scores = tl.sum(tl.where(offs_n < TOPK_NUM, sorted_scores, 0.0)) sorted_scores = sorted_scores / sum_scores tl.store( @@ -355,18 +341,9 @@ def triton_grouped_topk( token_num, total_expert_num = gating_output.shape - if ( - use_single_group_fast_path - and num_expert_group == 1 - and topk_group == 1 - and scoring_func == "sigmoid" - ): - out_topk_weights = torch.empty( - (token_num, topk), dtype=torch.float32, device="cuda" - ) - out_topk_ids = torch.empty( - (token_num, topk), dtype=torch.long, device="cuda" - ) + if use_single_group_fast_path and num_expert_group == 1 and topk_group == 1 and scoring_func == "sigmoid": + out_topk_weights = torch.empty((token_num, topk), dtype=torch.float32, device="cuda") + out_topk_ids = torch.empty((token_num, topk), dtype=torch.long, device="cuda") single_group_sigmoid_topk_kernel[(token_num,)]( gating_output, gating_output.stride(0), diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index 827caf0f95..f3b3426acf 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -134,9 +134,7 @@ def silu_and_mul_masked_post_quant_fwd( finfo = torch.finfo(torch.float8_e4m3fn) fp8_max = finfo.max fp8_min = -fp8_max - assert (limit is None and alpha is None) or ( - limit is not None and alpha is not None - ) + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) _silu_and_mul_post_quant_kernel[grid]( input, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 4af05e31bc..3e782b69b6 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -299,9 +299,7 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV * K if is_kda else HV) b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw, HV) BK = triton.next_power_of_2(K) - is_spec_verify = ( - ENABLE_FAST_MTP_KDA and cu_seqlens is not None and num_accepted_tokens is not None - ) + is_spec_verify = ENABLE_FAST_MTP_KDA and cu_seqlens is not None and num_accepted_tokens is not None if T == 1: # Decode path: use larger BV to reduce kernel instances (4 blocks instead of 16) # and more warps for better SM utilization at T=1 where there's no pipelining benefit diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py index d2d1e4b877..0cf7ce8fd8 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -46,12 +46,8 @@ def kda_safe_gate( head_count = a_log.numel() key_dim = gate_bias.numel() // head_count gate = raw_gate.float().view(*raw_gate.shape[:-1], head_count, key_dim) - amplitude = a_log.float().reshape( - *((1,) * (gate.ndim - 2)), head_count, 1 - ).exp() - bias = gate_bias.float().reshape( - *((1,) * (gate.ndim - 2)), head_count, key_dim - ) + amplitude = a_log.float().reshape(*((1,) * (gate.ndim - 2)), head_count, 1).exp() + bias = gate_bias.float().reshape(*((1,) * (gate.ndim - 2)), head_count, key_dim) return lower_bound * torch.sigmoid(amplitude * (gate + bias)) @@ -122,6 +118,8 @@ def fused_recurrent_kda( out=out, is_kda=True, ) + + @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @triton.autotune( configs=[ @@ -179,29 +177,17 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_kt = tl.make_block_ptr(k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K @@ -223,13 +209,9 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) + p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) @@ -471,9 +453,7 @@ def recompute_w_u_fwd_kernel( p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) b_b = tl.load(p_b, boundary_check=(0,)) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) b_A = tl.load(p_A, boundary_check=(0, 1)) for i_v in range(tl.cdiv(V, BV)): @@ -553,9 +533,7 @@ def recompute_w_u_fwd_kernel( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - b_gn = tl.load( - gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 - ) + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0) b_kg = b_k * exp2(b_gn - b_gk) p_kg = tl.make_block_ptr( @@ -726,9 +704,7 @@ def chunk_gla_fwd_kernel_o( (BT, BV), (1, 0), ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) # [BT, BV] b_v = tl.load(p_v, boundary_check=(0, 1)) # [BT, BT] @@ -786,11 +762,7 @@ def grid(meta): } ) @triton.autotune( - configs=[ - triton.Config({"BD": BD}, num_warps=num_warps) - for BD in [32, 64] - for num_warps in [2, 4, 8] - ], + configs=[triton.Config({"BD": BD}, num_warps=num_warps) for BD in [32, 64] for num_warps in [2, 4, 8]], key=["H", "D", "BT", "IS_VARLEN"], ) @triton.jit(do_not_specialize=["T"]) @@ -892,9 +864,7 @@ def fused_kda_gate_chunk_cumsum( lower_bound: float = -5.0, ) -> torch.Tensor: if cu_seqlens is not None: - assert raw_g.shape[0] == 1, ( - "Only batch size 1 is supported when cu_seqlens are provided" - ) + assert raw_g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" B, T, H, D = raw_g.shape if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) @@ -1018,11 +988,7 @@ def chunk_kda_fwd( cu_seqlens: torch.Tensor | None = None, ): chunk_size = FLA_CHUNK_SIZE - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is not None - else None - ) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None g = chunk_local_cumsum( g, chunk_size=chunk_size, diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 4d95e68e97..5b67ffa424 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -57,9 +57,7 @@ class CustomProcessGroup: def __init__(self): self.symm_mem_reduce = None self.flashinfer_reduce = None - self.symm_mem_out_of_place = os.getenv( - "LIGHTLLM_SYMM_MEM_OUT_OF_PLACE", "0" - ).upper() in {"1", "ON", "TRUE"} + self.symm_mem_out_of_place = os.getenv("LIGHTLLM_SYMM_MEM_OUT_OF_PLACE", "0").upper() in {"1", "ON", "TRUE"} self.dp_world_size = get_dp_world_size() self.device_group = create_new_group_for_current_dp("nccl") if get_env_start_args().enable_dp_prefill_balance: @@ -243,11 +241,11 @@ def new_deepep_group( mega_moe_quant_method = "fp4fp8-b32-deepgemm" sm90_mega_moe_quant_method = "fp8w8a8-b128-deepgemm" is_sm100 = is_sm100_gpu() - enable_sm90_mega_moe = ( - is_sm90_gpu() - and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() - in {"1", "ON", "TRUE"} - ) + enable_sm90_mega_moe = is_sm90_gpu() and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() in { + "1", + "ON", + "TRUE", + } # Buffer 选择规则： # 1. 非 SM100 不支持 Mega MoE，只初始化 legacy low-latency buffer； @@ -266,8 +264,7 @@ def new_deepep_group( elif enable_sm90_mega_moe: has_mega_moe_layer = sm90_mega_moe_quant_method in expert_quant_method_names has_legacy_moe_layer = any( - method_name != sm90_mega_moe_quant_method - for method_name in expert_quant_method_names + method_name != sm90_mega_moe_quant_method for method_name in expert_quant_method_names ) enable_mega_moe_buffer = has_mega_moe_layer enable_low_latency_buffer = has_legacy_moe_layer diff --git a/lightllm/distributed/symm_mem_all_reduce.py b/lightllm/distributed/symm_mem_all_reduce.py index 48f4204a87..e23c0a478e 100644 --- a/lightllm/distributed/symm_mem_all_reduce.py +++ b/lightllm/distributed/symm_mem_all_reduce.py @@ -100,13 +100,9 @@ def _reduce_to_workspace(self, inp: torch.Tensor) -> torch.Tensor: output = self.buffer[:n] output.copy_(inp.view(-1)) if self.use_multimem: - torch.ops.symm_mem.multimem_all_reduce_( - output, "sum", self.group.group_name - ) + torch.ops.symm_mem.multimem_all_reduce_(output, "sum", self.group.group_name) else: - torch.ops.symm_mem.two_shot_all_reduce_( - output, "sum", self.group.group_name - ) + torch.ops.symm_mem.two_shot_all_reduce_(output, "sum", self.group.group_name) return output.view_as(inp) def all_reduce(self, inp: torch.Tensor) -> None: diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 962e72d1c7..a1343e4165 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -215,9 +215,7 @@ def _get_o( def _shared_ffn_tp(self, input, infer_state, layer_weight): """Shared-expert FFN hook for model-specific activation semantics.""" - return LlamaTransformerLayerInfer._ffn_tp( - self, input, infer_state, layer_weight - ) + return LlamaTransformerLayerInfer._ffn_tp(self, input, infer_state, layer_weight) def _moe_ffn_tp( self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight @@ -228,9 +226,7 @@ def _moe_ffn_tp( # if fused_shared_experts is not enabled, compute shared_output if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: - shared_output = self._shared_ffn_tp( - hidden_states, infer_state, layer_weight - ) + shared_output = self._shared_ffn_tp(hidden_states, infer_state, layer_weight) moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) @@ -258,9 +254,7 @@ def _moe_ffn_edp( hidden_states = input token_num, hidden_dim = hidden_states.shape if self.n_shared_experts is not None: - shared_output = self._shared_ffn_tp( - hidden_states, infer_state, layer_weight - ) + shared_output = self._shared_ffn_tp(hidden_states, infer_state, layer_weight) moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) @@ -335,9 +329,7 @@ def overlap_tpsp_token_forward( # 0 shared expert if self.n_shared_experts is not None: - _0_shared_output = self._shared_ffn_tp( - _0_input1, infer_state, layer_weight - ) + _0_shared_output = self._shared_ffn_tp(_0_input1, infer_state, layer_weight) # 0 dispatch ( @@ -372,9 +364,7 @@ def overlap_tpsp_token_forward( # 1 shared expert if self.n_shared_experts is not None: - _1_shared_output = self._shared_ffn_tp( - _1_input1, infer_state1, layer_weight - ) + _1_shared_output = self._shared_ffn_tp(_1_input1, infer_state1, layer_weight) # 1 dispatch ( @@ -510,15 +500,11 @@ def overlap_tpsp_context_forward( # 0 shared expert if self.n_shared_experts is not None: - _0_shared_output = self._shared_ffn_tp( - _0_input1, infer_state, layer_weight - ) + _0_shared_output = self._shared_ffn_tp(_0_input1, infer_state, layer_weight) # 1 shared expert if self.n_shared_experts is not None: - _1_shared_output = self._shared_ffn_tp( - _1_input1, infer_state1, layer_weight - ) + _1_shared_output = self._shared_ffn_tp(_1_input1, infer_state1, layer_weight) # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index bfdc789914..20dcbaed38 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -222,19 +222,11 @@ def _get_indices( if infer_state.is_prefill: mtp_step = 0 else: - mtp_step = ( - get_env_start_args().mtp_step - if self.decode_mtp_step is None - else self.decode_mtp_step - ) + mtp_step = get_env_start_args().mtp_step if self.decode_mtp_step is None else self.decode_mtp_step # LightSpec compacts each request to a variable number of contiguous # verify rows. Its sparse-index K packing must follow request boundaries # instead of assuming the fixed process-wide MTP width. - use_dynamic_layout = ( - not infer_state.is_prefill - and mtp_step > 0 - and get_env_start_args().mtp_dynamic_verify - ) + use_dynamic_layout = not infer_state.is_prefill and mtp_step > 0 and get_env_start_args().mtp_dynamic_verify if use_dynamic_layout: k_fp8_, k_scale_ = extract_indexer_ks_dynamic( I_buffer=indexer_k_buffer, diff --git a/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py b/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py index 4b2dc0e067..06caa7dae6 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py +++ b/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py @@ -46,9 +46,9 @@ def _fwd_kernel_extract_indexer_ks( # Token slots can exceed INT32_MAX / stride_in_fp8_bs when the KV # cache packs multiple states into one row. Promote before pointer # arithmetic so large cache offsets do not wrap around. - mem_index = tl.load( - req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n - ).to(tl.int64) + mem_index = tl.load(req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n).to( + tl.int64 + ) in_fp8_ptrs = in_fp8 + mem_index * stride_in_fp8_bs + 0 * stride_in_fp8_h + stride_in_fp8_d * offs_d kv_fp8 = tl.load(in_fp8_ptrs) @@ -131,9 +131,9 @@ def _fwd_kernel_extract_indexer_ks_dynamic( offs_d = tl.arange(0, BLOCK_DMODEL) for i in range(token_start_index, cur_seq_len, tl.num_programs(1)): - mem_index = tl.load( - req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n - ).to(tl.int64) + mem_index = tl.load(req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n).to( + tl.int64 + ) in_fp8_ptrs = in_fp8 + mem_index * stride_in_fp8_bs + stride_in_fp8_d * offs_d kv_fp8 = tl.load(in_fp8_ptrs) diff --git a/lightllm/models/glm5_next/glm5_next_visual.py b/lightllm/models/glm5_next/glm5_next_visual.py new file mode 100644 index 0000000000..6825d49697 --- /dev/null +++ b/lightllm/models/glm5_next/glm5_next_visual.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import json +import os +from io import BytesIO +from typing import List + +import torch +from PIL import Image +from safetensors import safe_open +from transformers import AutoConfig, AutoProcessor + +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.utils.torch_dtype_utils import get_torch_dtype + + +class Glm5NextVisionModel: + """LightLLM visual-server adapter around the official GLM-5 vision tower.""" + + def __init__(self, data_type="bfloat16"): + self.data_type = data_type if isinstance(data_type, torch.dtype) else get_torch_dtype(data_type) + self.device = torch.device("cpu") + self.vision_tower = None + self.image_processor = None + + @staticmethod + def _weight_files(weight_dir): + index_path = os.path.join(weight_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path, "r", encoding="utf-8") as stream: + weight_map = json.load(stream)["weight_map"] + return sorted(set(weight_map.values())) + return sorted(name for name in os.listdir(weight_dir) if name.endswith(".safetensors")) + + @classmethod + def _load_prefix_state_dict(cls, weight_dir, prefix): + state_dict = {} + for file_name in cls._weight_files(weight_dir): + with safe_open(os.path.join(weight_dir, file_name), framework="pt", device="cpu") as stream: + for key in stream.keys(): + if key.startswith(prefix): + state_dict[key[len(prefix) :]] = stream.get_tensor(key) + return state_dict + + def load_model(self, weight_dir): + try: + from transformers.models.glm5_next.modeling_glm5_next import ( + Glm5NextVisionModel as HFGlm5NextVisionModel, + ) + except ImportError as exc: + raise ImportError("GLM-5 vision requires a Transformers build with glm5_next support") from exc + + config = AutoConfig.from_pretrained(weight_dir, trust_remote_code=True) + if config.vision_config is None: + raise ValueError("GLM-5 checkpoint does not contain vision_config") + # Direct submodel construction bypasses AutoModel's normal attention + # implementation selection. SDPA keeps large-image attention + # memory-efficient and is available in the pinned PyTorch runtime. + config.vision_config._attn_implementation = "sdpa" + self.vision_tower = HFGlm5NextVisionModel(config.vision_config).eval() + self.image_processor = AutoProcessor.from_pretrained(weight_dir).image_processor + + state_dict = self._load_prefix_state_dict(weight_dir, "model.visual.") + missing, unexpected = self.vision_tower.load_state_dict(state_dict, strict=False) + if missing or unexpected: + raise RuntimeError(f"GLM-5 vision weight mismatch: missing={missing}, unexpected={unexpected}") + return self + + def cuda(self): + self.device = torch.device("cuda") + self.vision_tower = self.vision_tower.to(device=self.device, dtype=self.data_type) + return self + + @torch.inference_mode() + def forward(self, pixel_values, image_grid_thw): + output = self.vision_tower( + hidden_states=pixel_values.to(self.device, dtype=self.data_type, non_blocking=True), + grid_thw=image_grid_thw.to(self.device, non_blocking=True), + ) + return output.pooler_output.to(self.data_type) + + @torch.inference_mode() + def encode(self, images: List[ImageItem]): + pil_images = [] + uuids = [] + for image_item in images: + if not isinstance(image_item, ImageItem): + raise TypeError(f"Unsupported GLM-5 image input type: {type(image_item)}") + uuids.append(image_item.uuid) + image_data = read_shm(get_shm_name_data(image_item.uuid)) + with Image.open(BytesIO(image_data)) as image: + pil_images.append(image.convert("RGB")) + + if not pil_images: + return None + + image_inputs = self.image_processor(pil_images, return_tensors="pt") + pixel_values = image_inputs["pixel_values"] + image_grid_thw = image_inputs["image_grid_thw"] + + merge_area = self.image_processor.merge_size ** 2 + token_nums = [int(grid.prod().item() // merge_area) for grid in image_grid_thw] + valid_ids = [] + valid_start = 0 + for image_item, token_num in zip(images, token_nums): + if image_item.token_num is not None and image_item.token_num != token_num: + raise ValueError(f"GLM-5 image token mismatch: allocated={image_item.token_num}, encoded={token_num}") + valid_ids.append([valid_start, valid_start + token_num]) + valid_start += token_num + + image_embeds = self.forward(pixel_values, image_grid_thw) + if image_embeds.shape[0] != valid_start: + raise ValueError(f"GLM-5 image embed length mismatch: embeds={image_embeds.shape[0]}, tokens={valid_start}") + return image_embeds, uuids, valid_ids diff --git a/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py b/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..ff0e48c71a --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py @@ -0,0 +1,7 @@ +from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer + + +class Glm5NextMultimodalPreLayerInfer(LlamaMultimodalPreLayerInfer): + """Merge cached GLM-5 vision embeddings into the text-token stream.""" + + pass diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 4f6c6bfbd5..0d03f9c60c 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -28,9 +28,7 @@ class Glm5NextNsaInfer(NsaInfer): """GLM indexer projection without rotary dimensions.""" def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): - q = layer_weight.wq_b_proj_.mm(q_lora).view( - -1, self.tp_index_n_heads, self.index_head_dim - ) + q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.tp_index_n_heads, self.index_head_dim) k = layer_weight.wk_proj_.mm(hidden_states.to(q_lora.dtype)) k = layer_weight.k_norm_(k, eps=self.eps) return self._rotate_activation(q), self._rotate_activation(k) @@ -38,21 +36,17 @@ def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): # GLM stores weights_proj in FP32, so its activation must match before # delegating to the shared NSA scoring and top-k implementation. - return super()._get_indices( - hidden_states.float(), q_lora, infer_state, att_state, layer_weight - ) + return super()._get_indices(hidden_states.float(), q_lora, infer_state, att_state, layer_weight) + class Glm5NextTransformerLayerInfer(Deepseek3_2TransformerLayerInfer): def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.num_hidden_layers = network_config["num_hidden_layers"] - self.autotune_layer_num = network_config.get( - "autotune_layer_num", self.num_hidden_layers - ) + self.autotune_layer_num = network_config.get("autotune_layer_num", self.num_hidden_layers) self.is_mtp_layer = layer_num >= self.num_hidden_layers self.is_linear_attention_layer = ( - not self.is_mtp_layer - and network_config["layer_types"][layer_num] == "linear_attention" + not self.is_mtp_layer and network_config["layer_types"][layer_num] == "linear_attention" ) self.mhc_streams = network_config.get("hc_mult", 4) self.hc_eps = network_config.get("hc_eps", 1e-6) @@ -72,18 +66,14 @@ def __init__(self, layer_num, network_config): # GLM's recurrent EAGLE drafter processes one row per logical # request. Only target-model decode uses the widened verification # layout of mtp_step + 1 rows. - self.indexer.decode_mtp_step = ( - 0 if self.is_mtp_layer else get_env_start_args().mtp_step - ) + self.indexer.decode_mtp_step = 0 if self.is_mtp_layer else get_env_start_args().mtp_step def _ffn_tp(self, input, infer_state, layer_weight): """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" input = input.view(-1, self.embed_dim_) up_gate_out = layer_weight.gate_up_proj.mm(input) - ffn1_out = self.alloc_tensor( - (input.size(0), up_gate_out.size(1) // 2), input.dtype - ) + ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) silu_and_mul_fwd( up_gate_out, ffn1_out, @@ -106,14 +96,10 @@ def _get_qkv(self, input, infer_state, layer_weight): if infer_state.need_dp_prefill_balance: input = infer_state._all_to_all_unbalance_get(data=input) - q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split( - [self.q_lora_rank, self.kv_lora_rank], dim=-1 - ) + q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split([self.q_lora_rank, self.kv_lora_rank], dim=-1) q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_) infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q} - q = layer_weight.q_b_proj_.mm(q).view( - -1, self.tp_q_head_num_, self.qk_nope_head_dim - ) + q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim) cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank) rmsnorm_forward( cache_kv[:, :, : self.kv_lora_rank], @@ -131,9 +117,7 @@ def _get_o(self, input, infer_state, layer_weight): input = infer_state._all_to_all_balance_get(data=input) if input.shape[2] == self.kv_lora_rank: input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) - output = layer_weight.o_weight_.mm( - input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim) - ) + output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) all_reduce(output, group=infer_state.dist_group) return output @@ -176,9 +160,7 @@ def _kda_projections(self, input, infer_state, layer_weight): input = self._tpsp_allgather(input=input, infer_state=infer_state) projected = layer_weight.linear_qkvb_proj.mm(input) qkv_size = 3 * self.tp_linear_projection_size - mixed_qkv, raw_beta = projected.split( - [qkv_size, self.tp_linear_num_heads], dim=-1 - ) + mixed_qkv, raw_beta = projected.split([qkv_size, self.tp_linear_num_heads], dim=-1) fg_a = layer_weight.linear_fg_a_proj.mm(input) f_a, g_a = fg_a.split(self.linear_head_dim, dim=-1) raw_gate, norm_gate = layer_weight.project_kda_fg_b(f_a, g_a) @@ -203,9 +185,7 @@ def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): def context_attention_forward(self, input_embeddings, infer_state, layer_weight): if not self.is_linear_attention_layer: return super().context_attention_forward(input_embeddings, infer_state, layer_weight) - mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( - input_embeddings, infer_state, layer_weight - ) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) core_output = infer_state.prefill_att_state1.prefill_att( q=None, k=None, @@ -227,9 +207,7 @@ def context_attention_forward(self, input_embeddings, infer_state, layer_weight) def token_attention_forward(self, input_embeddings, infer_state, layer_weight): if not self.is_linear_attention_layer: return super().token_attention_forward(input_embeddings, infer_state, layer_weight) - mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( - input_embeddings, infer_state, layer_weight - ) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) core_output = infer_state.decode_att_state1.decode_att( q=None, k=None, @@ -267,59 +245,36 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): if self.layer_num_ == 0: streams = hc_expand(streams.view(-1, self.embed_dim_), self.mhc_streams) - layer_input, residual_mix, post_mix = self._hc_pre( - streams, layer_weight, "attn", layer_weight.att_norm_weight_ - ) + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) if prefill: layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) else: layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) - streams = hc_post( - layer_output, streams, residual_mix, post_mix, self.mhc_streams - ) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) - layer_input, residual_mix, post_mix = self._hc_pre( - streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_ - ) + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) if infer_state.use_replicated_attention_ep: if self.is_moe: - local_input = self._tpsp_sp_split( - input=layer_input, infer_state=infer_state - ) + local_input = self._tpsp_sp_split(input=layer_input, infer_state=infer_state) local_output = self._ffn(local_input, infer_state, layer_weight) - layer_output = self._tpsp_allgather( - input=local_output, infer_state=infer_state - ) + layer_output = self._tpsp_allgather(input=local_output, infer_state=infer_state) else: layer_output = self._ffn_tp(layer_input, infer_state, layer_weight) all_reduce(layer_output, group=infer_state.dist_group) else: layer_output = self._ffn(layer_input, infer_state, layer_weight) - streams = hc_post( - layer_output, streams, residual_mix, post_mix, self.mhc_streams - ) - is_autotune_last_layer = ( - Autotuner.is_autotune_warmup() - and self.layer_num_ == self.autotune_layer_num - 1 - ) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + is_autotune_last_layer = Autotuner.is_autotune_warmup() and self.layer_num_ == self.autotune_layer_num - 1 if self.layer_num_ == self.num_hidden_layers - 1 or is_autotune_last_layer: return hc_contract(streams, self.mhc_streams) return streams def context_forward(self, input_embeddings, infer_state, layer_weight): if self.is_mtp_layer: - return super().context_forward( - input_embeddings, infer_state, layer_weight - ) - return self._forward_mhc( - input_embeddings, infer_state, layer_weight, prefill=True - ) + return super().context_forward(input_embeddings, infer_state, layer_weight) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=True) def token_forward(self, input_embeddings, infer_state, layer_weight): if self.is_mtp_layer: - return super().token_forward( - input_embeddings, infer_state, layer_weight - ) - return self._forward_mhc( - input_embeddings, infer_state, layer_weight, prefill=False - ) + return super().token_forward(input_embeddings, infer_state, layer_weight) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=False) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index 616cb16c0f..4f1a6725dc 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -33,9 +33,7 @@ def _parse_config(self): self.enable_cc_method = False self.is_mtp_layer = self.layer_num_ >= self.network_config_["num_hidden_layers"] self.is_linear_attention_layer = ( - not self.is_mtp_layer - and self.network_config_["layer_types"][self.layer_num_] - == "linear_attention" + not self.is_mtp_layer and self.network_config_["layer_types"][self.layer_num_] == "linear_attention" ) linear = self.network_config_["linear_attn_config"] self.linear_num_heads = linear["num_heads"] diff --git a/lightllm/models/glm5_next/mem_manager.py b/lightllm/models/glm5_next/mem_manager.py index f813701d99..20ab8e445f 100644 --- a/lightllm/models/glm5_next/mem_manager.py +++ b/lightllm/models/glm5_next/mem_manager.py @@ -15,9 +15,7 @@ def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): destindex_copy_kv, ) - output = self.mem_manager.kv_buffer[layer_index][ - :, :, : self.mem_manager.mla_head_dim - ] + output = self.mem_manager.kv_buffer[layer_index][:, :, : self.mem_manager.mla_head_dim] destindex_copy_kv(kv, mem_index, output) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index d56e8f4b53..f74ced0f78 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -18,6 +18,9 @@ from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import ( Glm5NextTransformerLayerInfer, ) +from lightllm.models.glm5_next.layer_infer.pre_layer_infer import ( + Glm5NextMultimodalPreLayerInfer, +) from lightllm.models.glm5_next.layer_weights.pre_and_post_layer_weight import ( Glm5NextPreAndPostLayerWeight, ) @@ -182,19 +185,22 @@ def _init_custom(self): # GLM-5 sparse MLA is entirely NoPE. Keep zero-width tables so the # generic infer-state position setup remains valid without allocating # a million-token rotary cache. - max_length = max( - self.config["max_position_embeddings"], self.max_seq_length or 0 - ) - self._cos_cached = torch.empty( - (max_length, 0), dtype=self.data_type, device="cuda" - ) + max_length = max(self.config["max_position_embeddings"], self.max_seq_length or 0) + self._cos_cached = torch.empty((max_length, 0), dtype=self.data_type, device="cuda") self._sin_cached = torch.empty_like(self._cos_cached) dist_group_manager.new_deepep_group( n_routed_experts=self.config["n_routed_experts"], hidden_size=self.config["hidden_size"], - expert_quant_method_names=dist_group_manager.get_moe_quant_methods( - self.trans_layers_weight - ), + expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), num_experts_per_tok=self.config["num_experts_per_tok"], moe_intermediate_size=self.config["moe_intermediate_size"], ) + + +@ModelRegistry( + "glm5_next", + is_multimodal=True, + condition=lambda model_cfg: model_cfg.get("vision_config") is not None, +) +class Glm5NextMultimodalTpPartModel(Glm5NextTpPartModel): + pre_layer_infer_class = Glm5NextMultimodalPreLayerInfer diff --git a/lightllm/models/glm5_next/tokenizer.py b/lightllm/models/glm5_next/tokenizer.py new file mode 100644 index 0000000000..fbd2843ba7 --- /dev/null +++ b/lightllm/models/glm5_next/tokenizer.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import copy +from typing import List, Union + +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from lightllm.server.multimodal_params import ImageItem, MultimodalParams + + +class Glm5NextTokenizer(QWen2VLTokenizer): + """Multimodal tokenizer adapter for GLM-5 Next checkpoints. + + The released GLM-5.3-Flash chat template deliberately renders OpenAI + image parts as a text-only reminder. LightLLM carries image bytes through + ``MultimodalParams``, so template-facing image parts must instead become + one GLM image placeholder. ``encode`` (inherited from QWen2VLTokenizer) + replaces that placeholder with the virtual token range allocated by the + embedding cache. + """ + + image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + + def __init__(self, tokenizer=None, image_processor=None, **kwargs): + BaseMultiModalTokenizer.__init__(self, tokenizer) + self.image_processor = image_processor + model_cfg = kwargs["model_cfg"] + self.image_start_id = model_cfg["image_start_token_id"] + self.image_end_id = model_cfg["image_end_token_id"] + self.image_token_id = model_cfg["image_token_id"] + self.patch_size = image_processor.patch_size + self.merge_size = image_processor.merge_size + self.min_image_tokens = image_processor.min_image_tokens + self.max_image_tokens = image_processor.max_image_tokens + + def get_image_token_length(self, img: ImageItem): + if img.image_w <= 0 or img.image_h <= 0: + raise ValueError(f"invalid GLM-5 image size: {img.image_w}x{img.image_h}") + patch_num = self.image_processor.get_number_of_image_patches(img.image_h, img.image_w) + token_num = patch_num // (self.merge_size ** 2) + if token_num <= 0: + raise ValueError(f"GLM-5 image produced no visual tokens: {img.image_w}x{img.image_h}") + return token_num + + def apply_chat_template(self, conversation=None, messages=None, **kwargs): + source = conversation if conversation is not None else messages + if source is None: + return self.tokenizer.apply_chat_template(conversation=conversation, messages=messages, **kwargs) + + normalized = copy.deepcopy(source) + for message in normalized: + content = message.get("content") + if not isinstance(content, list): + continue + rendered_parts = [] + for part in content: + if isinstance(part, dict) and part.get("type") in ("image", "image_url"): + rendered_parts.append({"type": "text", "text": self.image_placeholder}) + else: + rendered_parts.append(part) + message["content"] = rendered_parts + + if conversation is not None: + return self.tokenizer.apply_chat_template(conversation=normalized, **kwargs) + return self.tokenizer.apply_chat_template(messages=normalized, **kwargs) + + def encode( + self, + prompt: Union[str, List[int]], + multimodal_params: MultimodalParams = None, + **kwargs, + ): + return super().encode(prompt, multimodal_params=multimodal_params, **kwargs) diff --git a/lightllm/models/glm5_next/triton_kernel/mhc.py b/lightllm/models/glm5_next/triton_kernel/mhc.py index 77869a51f0..87eadf7b50 100644 --- a/lightllm/models/glm5_next/triton_kernel/mhc.py +++ b/lightllm/models/glm5_next/triton_kernel/mhc.py @@ -41,20 +41,11 @@ def _hc_prepare_kernel( pre_raw = tl.load(mixes + token * mix_stride_m + stream_offsets) post_raw = tl.load(mixes + token * mix_stride_m + STREAMS + stream_offsets) - pre_values = tl.sigmoid( - pre_raw * tl.load(scale) + tl.load(base + stream_offsets) - ) + HC_EPS - post_values = POST_MULTIPLIER * tl.sigmoid( - post_raw * tl.load(scale + 1) - + tl.load(base + STREAMS + stream_offsets) - ) + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) - logits = tl.load( - mixes + token * mix_stride_m + 2 * STREAMS + matrix_offsets - ) - logits = logits * tl.load(scale + 2) + tl.load( - base + 2 * STREAMS + matrix_offsets - ) + logits = tl.load(mixes + token * mix_stride_m + 2 * STREAMS + matrix_offsets) + logits = logits * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) logits = tl.reshape(logits, (STREAMS, STREAMS)) logits = logits - tl.max(logits, axis=1)[:, None] matrix = tl.exp(logits) @@ -107,34 +98,20 @@ def _hc_prepare_prenorm_kernel( matrix_raw = tl.zeros((STREAMS * STREAMS,), dtype=tl.float32) sqrsum = 0.0 for split in tl.static_range(N_SPLITS): - partial_base = ( - gemm_partial + split * gemm_stride_s + token * gemm_stride_m - ) + partial_base = gemm_partial + split * gemm_stride_s + token * gemm_stride_m pre_raw += tl.load(partial_base + stream_offsets) post_raw += tl.load(partial_base + STREAMS + stream_offsets) matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) - sqrsum += tl.load( - sqrsum_partial - + split * sqrsum_stride_s - + token * sqrsum_stride_m - ) + sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) pre_raw *= inv_rms post_raw *= inv_rms matrix_raw *= inv_rms - pre_values = tl.sigmoid( - pre_raw * tl.load(scale) + tl.load(base + stream_offsets) - ) + HC_EPS - post_values = POST_MULTIPLIER * tl.sigmoid( - post_raw * tl.load(scale + 1) - + tl.load(base + STREAMS + stream_offsets) - ) + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) - logits = ( - matrix_raw * tl.load(scale + 2) - + tl.load(base + 2 * STREAMS + matrix_offsets) - ) + logits = matrix_raw * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) logits = tl.reshape(logits, (STREAMS, STREAMS)) logits = logits - tl.max(logits, axis=1)[:, None] matrix = tl.exp(logits) @@ -256,18 +233,10 @@ def _hc_post_4stream_kernel( other=0.0, ).to(tl.float32) residual_base = residual + token * residual_stride_m + hidden_offsets - residual_0 = tl.load( - residual_base, mask=hidden_mask, other=0.0 - ).to(tl.float32) - residual_1 = tl.load( - residual_base + hidden, mask=hidden_mask, other=0.0 - ).to(tl.float32) - residual_2 = tl.load( - residual_base + 2 * hidden, mask=hidden_mask, other=0.0 - ).to(tl.float32) - residual_3 = tl.load( - residual_base + 3 * hidden, mask=hidden_mask, other=0.0 - ).to(tl.float32) + residual_0 = tl.load(residual_base, mask=hidden_mask, other=0.0).to(tl.float32) + residual_1 = tl.load(residual_base + hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_2 = tl.load(residual_base + 2 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_3 = tl.load(residual_base + 3 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) post_base = post_mix + token * post_stride_m mix_base = residual_mix + token * mix_stride_m @@ -345,21 +314,13 @@ def hc_pre_reference( residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps - post = post_multiplier * torch.sigmoid( - post_raw * scale[1] + base[streams : 2 * streams] - ) - residual_mix = ( - residual_raw * scale[2] + base[2 * streams :].view(streams, streams) - ).softmax(dim=-1) + post = post_multiplier * torch.sigmoid(post_raw * scale[1] + base[streams : 2 * streams]) + residual_mix = (residual_raw * scale[2] + base[2 * streams :].view(streams, streams)).softmax(dim=-1) residual_mix = residual_mix + hc_eps residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) for _ in range(sinkhorn_iters - 1): - residual_mix = residual_mix / ( - residual_mix.sum(dim=-1, keepdim=True) + hc_eps - ) - residual_mix = residual_mix / ( - residual_mix.sum(dim=-2, keepdim=True) + hc_eps - ) + residual_mix = residual_mix / (residual_mix.sum(dim=-1, keepdim=True) + hc_eps) + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) return layer_input, residual_mix, post @@ -376,9 +337,7 @@ def hc_post_reference( tokens, hidden = layer_output.shape residual_3d = residual.view(tokens, streams, hidden) - mixed_residual = ( - residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2) - ).sum(dim=1) + mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual return out.to(layer_output.dtype).reshape(tokens, streams * hidden) @@ -408,9 +367,7 @@ def hc_pre( pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) post = torch.empty_like(pre) - residual_mix = torch.empty( - (tokens, streams, streams), dtype=torch.float32, device=x.device - ) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) _hc_prepare_kernel[(tokens,)]( mixes, scale, @@ -428,9 +385,7 @@ def hc_pre( num_warps=1, ) - layer_input = torch.empty( - (tokens, hidden), dtype=x.dtype, device=x.device - ) + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) block_h = min(triton.next_power_of_2(hidden), 1024) _hc_pre_combine_kernel[(tokens, triton.cdiv(hidden, block_h))]( x, @@ -447,9 +402,7 @@ def hc_pre( return layer_input, residual_mix, post -def _compute_prenorm_splits( - tokens: int, flattened_hidden: int, device: torch.device -) -> int: +def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.device) -> int: grid_size = triton.cdiv(tokens, 64) k_blocks = triton.cdiv(flattened_hidden, 64) sms = torch.cuda.get_device_properties(device).multi_processor_count @@ -503,26 +456,18 @@ def hc_pre_norm( sinkhorn_iters, post_multiplier, ) - layer_input = rmsnorm_forward( - layer_input, weight=norm_weight, eps=norm_eps - ) + layer_input = rmsnorm_forward(layer_input, weight=norm_weight, eps=norm_eps) return layer_input, residual_mix, post mix_size = (2 + streams) * streams n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) - gemm_partial = torch.empty( - (n_splits, tokens, mix_size), dtype=torch.float32, device=x.device - ) - sqrsum_partial = torch.empty( - (n_splits, tokens), dtype=torch.float32, device=x.device - ) + gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device) + sqrsum_partial = torch.empty((n_splits, tokens), dtype=torch.float32, device=x.device) prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) post = torch.empty_like(pre) - residual_mix = torch.empty( - (tokens, streams, streams), dtype=torch.float32, device=x.device - ) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) _hc_prepare_prenorm_kernel[(tokens,)]( gemm_partial, sqrsum_partial, @@ -547,9 +492,7 @@ def hc_pre_norm( num_warps=1, ) - layer_input = torch.empty( - (tokens, hidden), dtype=x.dtype, device=x.device - ) + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) block_h = triton.next_power_of_2(hidden) _hc_pre_combine_norm_kernel[(tokens,)]( x, diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index c609c5b597..7c10ade411 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -37,9 +37,7 @@ class Glm5NextMTPModel(Glm5NextTpPartModel): def __init__(self, kvargs: dict): self.main_model: TpPartBaseModel = kvargs.pop("main_model") - self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop( - "mtp_previous_draft_models" - ) + self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models") super().__init__(kvargs) def _init_custom(self): @@ -67,9 +65,7 @@ def _init_weights(self, start_layer_index=None): ) ] self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ - self.pre_post_weight.lm_head_weight_ = ( - self.main_model.pre_post_weight.lm_head_weight_ - ) + self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ def _init_infer_layer(self, start_layer_index=None): assert start_layer_index is None @@ -78,11 +74,7 @@ def _init_infer_layer(self, start_layer_index=None): logical_layer = len(self.main_model.layers_infer) + sum( len(model.layers_infer) for model in self.mtp_previous_draft_models ) - self.layers_infer = [ - self.transformer_layer_infer_class( - logical_layer, network_config=self.config - ) - ] + self.layers_infer = [self.transformer_layer_infer_class(logical_layer, network_config=self.config)] def _init_some_value(self): super()._init_some_value() diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index c9bcc18971..c639c8c5c9 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -32,9 +32,7 @@ def __init__(self) -> None: # Mega-MoE owns one symmetric communication workspace per rank. Keep # the CPU pre/post pipeline enabled, but do not let its two host threads # enqueue a second model forward while that workspace is still live. - self._serialize_sm90_mega_moe_forwards = os.getenv( - "LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0" - ).upper() in { + self._serialize_sm90_mega_moe_forwards = os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() in { "1", "ON", "TRUE", diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index e1a4e421d1..bf2acfbbb2 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -122,6 +122,14 @@ def get_tokenizer( tokenizer = QWen3_5Tokenizer( tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg ) + elif model_type == "glm5_next" and model_cfg.get("vision_config") is not None: + from transformers import AutoProcessor + from ..models.glm5_next.tokenizer import Glm5NextTokenizer + + processor = AutoProcessor.from_pretrained(tokenizer_name) + tokenizer = Glm5NextTokenizer( + tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg + ) elif model_cfg.get("thinker_config") is not None: from transformers import AutoProcessor diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 68e0a97ca1..e67e381ac8 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -21,6 +21,7 @@ from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel +from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionModel from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.dist_utils import init_vision_distributed_env from lightllm.utils.envs_utils import get_env_start_args @@ -91,6 +92,10 @@ def exposed_init_model(self, kvargs): self.model = ( Qwen3VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) + elif self.model_type == "glm5_next" and model_cfg.get("vision_config") is not None: + if self.vit_tp != 1: + raise ValueError("GLM-5 vision supports --visual_tp 1; use --visual_dp for parallelism") + self.model = Glm5NextVisionModel(data_type=self.data_type) elif model_cfg["architectures"][0] == "TarsierForConditionalGeneration": self.model = TarsierVisionTransformerPretrainedModel(**model_cfg).eval().bfloat16() elif self.model_type == "llava": diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index bf247d090e..0e24e5f12e 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -392,6 +392,8 @@ def has_vision_module(model_path: str) -> bool: # Qwen3VisionTransformerPretrainedModel model_cfg["vision_config"] return True + elif model_type == "glm5_next": + return model_cfg.get("vision_config") is not None elif model_cfg["architectures"][0] == "TarsierForConditionalGeneration": # TarsierVisionTransformerPretrainedModel return True diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 93266684a8..2cd48338e9 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -98,11 +98,7 @@ def get_deepep_num_max_dispatch_tokens_per_rank_decode(): int(args.get("running_max_req_size", 0) or 0), 1, ) - verify_width = ( - int(args.get("mtp_step", 0) or 0) + 1 - if args.get("mtp_mode", None) is not None - else 1 - ) + verify_width = int(args.get("mtp_step", 0) or 0) + 1 if args.get("mtp_mode", None) is not None else 1 required_tokens = logical_batch_size * verify_width # In TP/SP + EP mode each rank dispatches only its sequence-parallel slice # to DeepEP. CUDA Graph batch sizes are TP-aligned, but use ceil here as a diff --git a/test/kernel/test_glm5_grouped_topk.py b/test/kernel/test_glm5_grouped_topk.py index dde6413988..7efc318c41 100644 --- a/test/kernel/test_glm5_grouped_topk.py +++ b/test/kernel/test_glm5_grouped_topk.py @@ -15,9 +15,7 @@ def torch_reference(gating_output, correction_bias, topk): scores = gating_output.float().sigmoid() choice_scores = scores + correction_bias - topk_ids = torch.topk( - choice_scores, k=topk, dim=-1, largest=True, sorted=True - ).indices + topk_ids = torch.topk(choice_scores, k=topk, dim=-1, largest=True, sorted=True).indices topk_weights = torch.gather(scores, 1, topk_ids) topk_weights /= topk_weights.sum(dim=-1, keepdim=True) return topk_weights, topk_ids @@ -87,27 +85,17 @@ def run_scratch_free_bitonic(gating_output, correction_bias): def assert_correct(tokens): generator = torch.Generator(device="cuda").manual_seed(20260828 + tokens) - gating_output = torch.randn( - (tokens, 288), generator=generator, dtype=torch.float32, device="cuda" - ) - correction_bias = torch.randn( - (288,), generator=generator, dtype=torch.float32, device="cuda" - ) + gating_output = torch.randn((tokens, 288), generator=generator, dtype=torch.float32, device="cuda") + correction_bias = torch.randn((288,), generator=generator, dtype=torch.float32, device="cuda") ref_weights, ref_ids = torch_reference(gating_output, correction_bias, 8) - fast_weights, fast_ids = run_topk( - gating_output, correction_bias, fast=True - ) - generic_weights, generic_ids = run_topk( - gating_output, correction_bias, fast=False - ) + fast_weights, fast_ids = run_topk(gating_output, correction_bias, fast=True) + generic_weights, generic_ids = run_topk(gating_output, correction_bias, fast=False) torch.testing.assert_close(fast_ids, ref_ids, rtol=0, atol=0) torch.testing.assert_close(generic_ids, ref_ids, rtol=0, atol=0) torch.testing.assert_close(fast_weights, ref_weights, rtol=1e-5, atol=1e-6) - torch.testing.assert_close( - generic_weights, ref_weights, rtol=1e-5, atol=1e-6 - ) + torch.testing.assert_close(generic_weights, ref_weights, rtol=1e-5, atol=1e-6) weight_delta = (fast_weights - generic_weights).abs() print( f"tokens={tokens}: exact expert ids, weights match reference; " @@ -143,15 +131,9 @@ def graph_ms(graph, iterations): def benchmark(tokens, iterations): gating_output, correction_bias = assert_correct(tokens) - fast_graph, fast_outputs = capture_graph( - lambda: run_topk(gating_output, correction_bias, fast=True) - ) - generic_graph, generic_outputs = capture_graph( - lambda: run_topk(gating_output, correction_bias, fast=False) - ) - bitonic_graph, bitonic_outputs = capture_graph( - lambda: run_scratch_free_bitonic(gating_output, correction_bias) - ) + fast_graph, fast_outputs = capture_graph(lambda: run_topk(gating_output, correction_bias, fast=True)) + generic_graph, generic_outputs = capture_graph(lambda: run_topk(gating_output, correction_bias, fast=False)) + bitonic_graph, bitonic_outputs = capture_graph(lambda: run_scratch_free_bitonic(gating_output, correction_bias)) fast_ms = graph_ms(fast_graph, iterations) generic_ms = graph_ms(generic_graph, iterations) bitonic_ms = graph_ms(bitonic_graph, iterations) @@ -174,15 +156,11 @@ def tune_warps(tokens, iterations): results = [] for num_warps in (1, 2, 4, 8, 16): graph, outputs = capture_graph( - lambda num_warps=num_warps: run_fast_with_warps( - gating_output, correction_bias, num_warps - ) + lambda num_warps=num_warps: run_fast_with_warps(gating_output, correction_bias, num_warps) ) graph.replay() torch.cuda.synchronize() - if not torch.equal(outputs[1], ref_ids) or not torch.allclose( - outputs[0], ref_weights, rtol=1e-5, atol=1e-6 - ): + if not torch.equal(outputs[1], ref_ids) or not torch.allclose(outputs[0], ref_weights, rtol=1e-5, atol=1e-6): print(f"tokens={tokens}: num_warps={num_warps} INVALID") continue elapsed_ms = graph_ms(graph, iterations) @@ -195,9 +173,7 @@ def tune_warps(tokens, iterations): f"max_delta={max_delta:.9g}" ) best_ms, best_warps = min(results) - print( - f"tokens={tokens}: best num_warps={best_warps} graph={best_ms:.6f} ms" - ) + print(f"tokens={tokens}: best num_warps={best_warps} graph={best_ms:.6f} ms") def test_glm5_single_group_topk(): diff --git a/test/kernel/test_glm5_mhc.py b/test/kernel/test_glm5_mhc.py index bc9456bacc..4cb7a1f504 100644 --- a/test/kernel/test_glm5_mhc.py +++ b/test/kernel/test_glm5_mhc.py @@ -44,24 +44,16 @@ def main() -> None: streams = 4 hidden = 4096 for tokens in (1, 8, 48): - x = torch.randn( - (tokens, streams * hidden), device=device, dtype=torch.bfloat16 - ) + x = torch.randn((tokens, streams * hidden), device=device, dtype=torch.bfloat16) fn = 0.005 * torch.randn( ((2 + streams) * streams, streams * hidden), device=device, dtype=torch.float32, ) scale = torch.randn((3,), device=device, dtype=torch.float32) - base = torch.randn( - ((2 + streams) * streams,), device=device, dtype=torch.float32 - ) - layer_output = torch.randn( - (tokens, hidden), device=device, dtype=torch.bfloat16 - ) - norm_weight = torch.randn( - (hidden,), device=device, dtype=torch.bfloat16 - ) + base = torch.randn(((2 + streams) * streams,), device=device, dtype=torch.float32) + layer_output = torch.randn((tokens, hidden), device=device, dtype=torch.bfloat16) + norm_weight = torch.randn((hidden,), device=device, dtype=torch.bfloat16) arguments = (x, fn, scale, base, streams, 1e-6, 1e-6, 20) expected_pre = hc_pre_reference(*arguments) @@ -83,22 +75,15 @@ def main() -> None: 20, ) torch.cuda.synchronize() - pre_errors = tuple( - _max_error(actual, expected) - for actual, expected in zip(actual_pre, expected_pre) - ) + pre_errors = tuple(_max_error(actual, expected) for actual, expected in zip(actual_pre, expected_pre)) assert pre_errors[0] <= 0.03125, pre_errors # DeepGEMM intentionally uses TF32 and a split-K reduction, matching # the optimized serving path rather than torch.mm's accumulation order. assert pre_errors[1] <= 5e-4, pre_errors assert pre_errors[2] <= 5e-4, pre_errors - expected_post = hc_post_reference( - layer_output, x, expected_pre[1], expected_pre[2], streams - ) - actual_post = hc_post( - layer_output, x, actual_pre[1], actual_pre[2], streams - ) + expected_post = hc_post_reference(layer_output, x, expected_pre[1], expected_pre[2], streams) + actual_post = hc_post(layer_output, x, actual_pre[1], actual_pre[2], streams) torch.cuda.synchronize() post_error = _max_error(actual_post, expected_post) assert post_error <= 0.03125, post_error @@ -106,9 +91,7 @@ def main() -> None: def reference_path(): layer_input, residual_mix, post_mix = hc_pre_reference(*arguments) rmsnorm_forward(layer_input, weight=norm_weight, eps=1e-6) - return hc_post_reference( - layer_output, x, residual_mix, post_mix, streams - ) + return hc_post_reference(layer_output, x, residual_mix, post_mix, streams) def fused_path(): _, residual_mix, post_mix = hc_pre_norm( @@ -123,9 +106,7 @@ def fused_path(): 1e-6, 20, ) - return hc_post( - layer_output, x, residual_mix, post_mix, streams - ) + return hc_post(layer_output, x, residual_mix, post_mix, streams) def fused_pre_path(): return hc_pre_norm( @@ -142,9 +123,7 @@ def fused_pre_path(): ) def fused_post_path(): - return hc_post( - layer_output, x, actual_pre[1], actual_pre[2], streams - ) + return hc_post(layer_output, x, actual_pre[1], actual_pre[2], streams) reference_ms = _time_ms(reference_path) fused_ms = _time_ms(fused_path) @@ -162,25 +141,14 @@ def fused_post_path(): from sglang.kernels.ops.layernorm.mhc import mhc_post_tilelang tokens = 17152 - residual = torch.randn( - (tokens, streams * hidden), device=device, dtype=torch.bfloat16 - ) - layer_output = torch.randn( - (tokens, hidden), device=device, dtype=torch.bfloat16 - ) - residual_mix = torch.randn( - (tokens, streams, streams), device=device, dtype=torch.float32 - ) - post_mix = torch.randn( - (tokens, streams), device=device, dtype=torch.float32 - ) - actual = hc_post( - layer_output, residual, residual_mix, post_mix, streams - ) + residual = torch.randn((tokens, streams * hidden), device=device, dtype=torch.bfloat16) + layer_output = torch.randn((tokens, hidden), device=device, dtype=torch.bfloat16) + residual_mix = torch.randn((tokens, streams, streams), device=device, dtype=torch.float32) + post_mix = torch.randn((tokens, streams), device=device, dtype=torch.float32) + actual = hc_post(layer_output, residual, residual_mix, post_mix, streams) + def sgl_mhc_post(): - output = torch.empty_like( - residual.view(tokens, streams, hidden) - ) + output = torch.empty_like(residual.view(tokens, streams, hidden)) mhc_post_tilelang( residual_mix, residual.view(tokens, streams, hidden), @@ -197,9 +165,7 @@ def sgl_mhc_post(): cross_error = _max_error(actual, sglang_output) assert cross_error <= 0.0625, cross_error triton_ms = _time_ms( - lambda: hc_post( - layer_output, residual, residual_mix, post_mix, streams - ), + lambda: hc_post(layer_output, residual, residual_mix, post_mix, streams), args.iterations, ) tilelang_ms = _time_ms( @@ -220,21 +186,15 @@ def sgl_mhc_post(): sgl_mhc.get_tp_group = lambda: None sgl_mhc.is_allocation_symmetric = lambda: False - sgl_mhc.use_symmetric_memory = ( - lambda *_args, **_kwargs: contextlib.nullcontext() - ) + sgl_mhc.use_symmetric_memory = lambda *_args, **_kwargs: contextlib.nullcontext() fn = 0.005 * torch.randn( ((2 + streams) * streams, streams * hidden), device=device, dtype=torch.float32, ) scale = torch.randn((3,), device=device, dtype=torch.float32) - base = torch.randn( - ((2 + streams) * streams,), device=device, dtype=torch.float32 - ) - norm_weight = torch.randn( - (hidden,), device=device, dtype=torch.bfloat16 - ) + base = torch.randn(((2 + streams) * streams,), device=device, dtype=torch.float32) + norm_weight = torch.randn((hidden,), device=device, dtype=torch.bfloat16) def lightllm_pre(): return hc_pre_norm( diff --git a/test/kernel/test_glm5_sglang_moe_compat.py b/test/kernel/test_glm5_sglang_moe_compat.py index 6edcd3b262..7f9a7ce217 100644 --- a/test/kernel/test_glm5_sglang_moe_compat.py +++ b/test/kernel/test_glm5_sglang_moe_compat.py @@ -25,9 +25,7 @@ def _fp8_randn(shape, *, scale=0.02): - return (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * scale).to( - torch.float8_e4m3fn - ) + return (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * scale).to(torch.float8_e4m3fn) def _graph_ms(fn, source, iterations): @@ -89,9 +87,7 @@ def main(): hidden_size = 4096 tp_intermediate_size = 256 - hidden_states = torch.randn( - (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16 - ) + hidden_states = torch.randn((num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16) w13 = _fp8_randn((num_experts, 2 * tp_intermediate_size, hidden_size)) w2 = _fp8_randn((num_experts, hidden_size, tp_intermediate_size)) w13_scale = torch.ones( @@ -104,12 +100,8 @@ def main(): device="cuda", dtype=torch.float32, ) - topk_ids = torch.randint( - 0, num_experts, (num_tokens, topk), device="cuda", dtype=torch.int64 - ) - topk_weights = torch.rand( - (num_tokens, topk), device="cuda", dtype=torch.float32 - ) + topk_ids = torch.randint(0, num_experts, (num_tokens, topk), device="cuda", dtype=torch.int64) + topk_weights = torch.rand((num_tokens, topk), device="cuda", dtype=torch.float32) topk_weights.mul_(2.5 / topk_weights.sum(dim=-1, keepdim=True)) def run_lightllm(output): @@ -140,6 +132,7 @@ def run_lightllm(output): ) sglang_fused_moe.get_exec = lambda: standalone_exec sglang_fused_moe_config.get_exec = lambda: standalone_exec + def run_sglang(output): sglang_fused_moe.fused_experts_impl( hidden_states=output, @@ -184,10 +177,7 @@ def run_sglang(output): if args.benchmark: lightllm_ms = _graph_ms(run_lightllm, hidden_states, args.iterations) sglang_ms = _graph_ms(run_sglang, hidden_states, args.iterations) - print( - "graph_ms lightllm=%.6f sglang=%.6f speedup=%.3fx" - % (lightllm_ms, sglang_ms, lightllm_ms / sglang_ms) - ) + print("graph_ms lightllm=%.6f sglang=%.6f speedup=%.3fx" % (lightllm_ms, sglang_ms, lightllm_ms / sglang_ms)) if args.tune_configs or args.tune_tma_configs: # GLM-5 decode has few physical tokens (48 for the main model and 8 @@ -234,10 +224,7 @@ def run_sglang(output): print("config_failed=%s error=%r" % (json.dumps(config, sort_keys=True), exc)) continue results.append((graph_ms, config)) - print( - "config_ms=%.6f config=%s" - % (graph_ms, json.dumps(config, sort_keys=True)) - ) + print("config_ms=%.6f config=%s" % (graph_ms, json.dumps(config, sort_keys=True))) results.sort(key=lambda item: item[0]) if not results: diff --git a/test/kernel/test_glm5_strided_causal_conv.py b/test/kernel/test_glm5_strided_causal_conv.py index f4b03d37a5..0f1428a47d 100644 --- a/test/kernel/test_glm5_strided_causal_conv.py +++ b/test/kernel/test_glm5_strided_causal_conv.py @@ -13,16 +13,10 @@ def _inputs(seq_lens: list[int], dim: int, cache_lines: int): total_tokens = sum(seq_lens) - token_major = torch.randn( - (total_tokens, dim), device="cuda", dtype=torch.bfloat16 - ) + token_major = torch.randn((total_tokens, dim), device="cuda", dtype=torch.bfloat16) weight = torch.randn((dim, 4), device="cuda", dtype=torch.bfloat16) - conv_states = torch.randn( - (cache_lines, dim, 3), device="cuda", dtype=torch.bfloat16 - ) - cache_indices = torch.arange( - len(seq_lens), device="cuda", dtype=torch.int32 - ) + conv_states = torch.randn((cache_lines, dim, 3), device="cuda", dtype=torch.bfloat16) + cache_indices = torch.arange(len(seq_lens), device="cuda", dtype=torch.int32) has_initial_state = torch.tensor( [(index % 2) == 1 for index in range(len(seq_lens))], device="cuda", @@ -87,34 +81,23 @@ def main() -> None: torch.manual_seed(1234) seq_lens = [7, 1, 19, 5] inputs = _inputs(seq_lens, dim=64, cache_lines=8) - reference_inputs = tuple( - value.clone() if isinstance(value, torch.Tensor) else value - for value in inputs - ) + reference_inputs = tuple(value.clone() if isinstance(value, torch.Tensor) else value for value in inputs) actual = _run(inputs, seq_lens, copy_free=True) expected = _run(reference_inputs, seq_lens, copy_free=False) torch.cuda.synchronize() output_error = (actual.float() - expected.float()).abs().max().item() - state_error = ( - inputs[2].float() - reference_inputs[2].float() - ).abs().max().item() + state_error = (inputs[2].float() - reference_inputs[2].float()).abs().max().item() # The two kernels accumulate the four taps in a different order; one BF16 # ULP at this random input scale is expected. assert output_error <= 0.0625, output_error assert state_error == 0.0, state_error assert actual.transpose(0, 1).is_contiguous() - print( - f"PASS correctness output_error={output_error:.8f} " - f"state_error={state_error:.8f}" - ) + print(f"PASS correctness output_error={output_error:.8f} " f"state_error={state_error:.8f}") if args.benchmark: seq_lens = [268] * 64 strided_inputs = _inputs(seq_lens, dim=3072, cache_lines=64) - copied_inputs = tuple( - value.clone() if isinstance(value, torch.Tensor) else value - for value in strided_inputs - ) + copied_inputs = tuple(value.clone() if isinstance(value, torch.Tensor) else value for value in strided_inputs) strided_ms = _time_ms( lambda: _run(strided_inputs, seq_lens, copy_free=True), args.iterations, diff --git a/test/kernel/test_glm5_vocab_parallel_top1.py b/test/kernel/test_glm5_vocab_parallel_top1.py index 811ce8398c..b4522f5ebc 100644 --- a/test/kernel/test_glm5_vocab_parallel_top1.py +++ b/test/kernel/test_glm5_vocab_parallel_top1.py @@ -45,8 +45,10 @@ def test_vocab_parallel_top1_matches_full_vocab_softmax(): local_start = rank * shard_size local_logits = full_logits[local_start : local_start + shard_size].contiguous() + def alloc_func(shape, dtype, device): return torch.empty(shape, dtype=dtype, device=device) + token_ids, token_probs = vocab_parallel_top1_and_prob( local_logits=local_logits, local_vocab_start_id=local_start, diff --git a/test/test_glm5_next_multimodal.py b/test/test_glm5_next_multimodal.py new file mode 100644 index 0000000000..7521b2d9d0 --- /dev/null +++ b/test/test_glm5_next_multimodal.py @@ -0,0 +1,96 @@ +from transformers.configuration_utils import PretrainedConfig + +from lightllm.models.glm5_next.model import ( + Glm5NextMultimodalTpPartModel, + Glm5NextTpPartModel, +) +from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer +from lightllm.models.registry import get_model_class +from lightllm.server.multimodal_params import MultimodalParams +from lightllm.utils.config_utils import has_vision_module + + +class _FakeImageProcessor: + patch_size = 14 + merge_size = 2 + min_image_tokens = 16 + max_image_tokens = 8000 + + @staticmethod + def get_number_of_image_patches(height, width): + assert (height, width) == (448, 448) + return 1024 + + +class _FakeTokenizer: + def __init__(self): + self.last_conversation = None + + def apply_chat_template(self, conversation, **kwargs): + self.last_conversation = conversation + return conversation[0]["content"][0]["text"] + + @staticmethod + def encode(prompt): + assert prompt == Glm5NextTokenizer.image_placeholder + return [7, 10, 12, 11, 8] + + +def _make_tokenizer(): + model_cfg = { + "image_start_token_id": 10, + "image_end_token_id": 11, + "image_token_id": 12, + } + return Glm5NextTokenizer( + tokenizer=_FakeTokenizer(), + image_processor=_FakeImageProcessor(), + model_cfg=model_cfg, + ) + + +def test_glm5_image_prompt_and_virtual_tokens(): + tokenizer = _make_tokenizer() + conversation = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}], + } + ] + + prompt = tokenizer.apply_chat_template(conversation=conversation) + assert prompt == tokenizer.image_placeholder + assert conversation[0]["content"][0]["type"] == "image_url" + + multimodal_params = MultimodalParams(images=[{"type": "base64", "data": ""}]) + image = multimodal_params.images[0] + image.image_w = image.image_h = 448 + image.token_num = tokenizer.get_image_token_length(image) + image.token_id = 1000 + + input_ids = tokenizer.encode(prompt, multimodal_params=multimodal_params) + assert image.token_num == 256 + assert image.start_idx == 2 + assert input_ids == [7, 10, *range(1000, 1256), 11, 8] + + +def test_glm5_vision_config_selects_multimodal_model(monkeypatch): + assert get_model_class({"model_type": "glm5_next"}) is Glm5NextTpPartModel + assert ( + get_model_class({"model_type": "glm5_next", "vision_config": {"hidden_size": 1024}}) + is Glm5NextMultimodalTpPartModel + ) + + monkeypatch.setattr( + PretrainedConfig, + "get_config_dict", + staticmethod( + lambda _: ( + {"model_type": "glm5_next", "vision_config": {"hidden_size": 1024}}, + {}, + ) + ), + ) + has_vision_module.cache_clear() + assert has_vision_module("unused") is True + has_vision_module.cache_clear() diff --git a/test/test_moe_prefill_dispatch.py b/test/test_moe_prefill_dispatch.py index f7f70d45b4..eabac2d6c0 100644 --- a/test/test_moe_prefill_dispatch.py +++ b/test/test_moe_prefill_dispatch.py @@ -24,9 +24,7 @@ class Experts: def experts(*args, **kwargs): captured.update(kwargs) - layer = Deepseek2TransformerLayerInfer.__new__( - Deepseek2TransformerLayerInfer - ) + layer = Deepseek2TransformerLayerInfer.__new__(Deepseek2TransformerLayerInfer) layer.embed_dim_ = 4 layer.n_shared_experts = None layer.num_experts_per_tok = 2 diff --git a/tools/analyze_torch_trace.py b/tools/analyze_torch_trace.py index c19c330920..0da113e23b 100644 --- a/tools/analyze_torch_trace.py +++ b/tools/analyze_torch_trace.py @@ -45,17 +45,15 @@ def main() -> None: total_calls = int(sum(values[1] for values in totals.values())) print(f"CUDA total: {total_us / 1000:.3f} ms across {total_calls} kernels") print("Graph totals:") - for graph_id, (duration, count) in sorted( - graph_totals.items(), key=lambda item: item[1][0], reverse=True - ): + for graph_id, (duration, count) in sorted(graph_totals.items(), key=lambda item: item[1][0], reverse=True): print( f" graph={graph_id:<5} total_ms={duration / 1000:10.3f} " f"share={duration / total_us:7.2%} calls={int(count)}" ) print("Kernel totals:") - for name, (duration, count, maximum) in sorted( - totals.items(), key=lambda item: item[1][0], reverse=True - )[: args.top]: + for name, (duration, count, maximum) in sorted(totals.items(), key=lambda item: item[1][0], reverse=True)[ + : args.top + ]: print( f" total_ms={duration / 1000:10.3f} share={duration / total_us:7.2%} " f"calls={int(count):7d} avg_us={duration / count:9.3f} " diff --git a/tools/bench_glm53_allreduce.py b/tools/bench_glm53_allreduce.py index eb56a169f7..de377f2319 100644 --- a/tools/bench_glm53_allreduce.py +++ b/tools/bench_glm53_allreduce.py @@ -77,12 +77,8 @@ def main() -> None: torch.testing.assert_close(nccl_input, torch.zeros_like(nccl_input), rtol=0, atol=0) nccl_input.copy_(source) torch.cuda.synchronize() - nccl_ms = elapsed_ms( - lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations - ) - nccl_graph_ms = graph_elapsed_ms( - lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations - ) + nccl_ms = elapsed_ms(lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations) + nccl_graph_ms = graph_elapsed_ms(lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations) buffer = symm_mem.empty(source.numel(), device="cuda", dtype=source.dtype) handle = symm_mem.rendezvous(buffer, group_name) @@ -106,12 +102,8 @@ def symm_all_reduce_out_of_place() -> torch.Tensor: torch.cuda.synchronize() symm_ms = elapsed_ms(symm_all_reduce, args.warmup, args.iterations) symm_graph_ms = graph_elapsed_ms(symm_all_reduce, args.warmup, args.iterations) - symm_out_ms = elapsed_ms( - symm_all_reduce_out_of_place, args.warmup, args.iterations - ) - symm_out_graph_ms = graph_elapsed_ms( - symm_all_reduce_out_of_place, args.warmup, args.iterations - ) + symm_out_ms = elapsed_ms(symm_all_reduce_out_of_place, args.warmup, args.iterations) + symm_out_graph_ms = graph_elapsed_ms(symm_all_reduce_out_of_place, args.warmup, args.iterations) nccl_max_ms = max_rank(nccl_ms) symm_max_ms = max_rank(symm_ms) @@ -120,9 +112,7 @@ def symm_all_reduce_out_of_place() -> torch.Tensor: symm_graph_max_ms = max_rank(symm_graph_ms) symm_out_graph_max_ms = max_rank(symm_out_graph_ms) - cpu_group = dist.new_group( - list(range(dist.get_world_size())), backend="gloo" - ) + cpu_group = dist.new_group(list(range(dist.get_world_size())), backend="gloo") workspace = flashinfer_comm.create_allreduce_fusion_workspace( backend="trtllm", world_size=dist.get_world_size(), @@ -150,16 +140,28 @@ def flashinfer_all_reduce() -> torch.Tensor: if rank == 0: nbytes = source.numel() * source.element_size() + eager_best = min( + (nccl_max_ms, "nccl"), + (symm_max_ms, "symm"), + (symm_out_max_ms, "symm_out"), + (fi_max_ms, "flashinfer"), + )[1] + graph_best = min( + (nccl_graph_max_ms, "nccl"), + (symm_graph_max_ms, "symm"), + (symm_out_graph_max_ms, "symm_out"), + (fi_graph_max_ms, "flashinfer"), + )[1] print( f"shape={shape} bytes={nbytes} nccl_ms={nccl_max_ms:.6f} " f"symm_multimem_ms={symm_max_ms:.6f} symm_out_ms={symm_out_max_ms:.6f} " f"flashinfer_ms={fi_max_ms:.6f} " - f"best={min((nccl_max_ms, 'nccl'), (symm_max_ms, 'symm'), (symm_out_max_ms, 'symm_out'), (fi_max_ms, 'flashinfer'))[1]} " + f"best={eager_best} " f"graph_nccl_ms={nccl_graph_max_ms:.6f} " f"graph_symm_ms={symm_graph_max_ms:.6f} " f"graph_symm_out_ms={symm_out_graph_max_ms:.6f} " f"graph_flashinfer_ms={fi_graph_max_ms:.6f} " - f"graph_best={min((nccl_graph_max_ms, 'nccl'), (symm_graph_max_ms, 'symm'), (symm_out_graph_max_ms, 'symm_out'), (fi_graph_max_ms, 'flashinfer'))[1]}" + f"graph_best={graph_best}" ) workspace.destroy() dist.destroy_process_group() diff --git a/tools/bench_glm53_kda_chunk_h.py b/tools/bench_glm53_kda_chunk_h.py index 028a40fc1d..0c5fed8578 100644 --- a/tools/bench_glm53_kda_chunk_h.py +++ b/tools/bench_glm53_kda_chunk_h.py @@ -43,9 +43,7 @@ def main() -> None: # The fused safe-gate+cumsum stage keeps cumulative decay in fp32; the # exp2 path in the state kernel expects that exact dtype. gk = torch.zeros(shape, device=device, dtype=torch.float32) - initial_state = torch.zeros( - (args.sequences, heads, key_dim, value_dim), device=device, dtype=dtype - ) + initial_state = torch.zeros((args.sequences, heads, key_dim, value_dim), device=device, dtype=dtype) cu_seqlens = torch.arange( 0, total_tokens + 1, @@ -73,9 +71,7 @@ def run(config: dict[str, int]) -> None: del output results = [] - for value_tile, num_warps, num_stages in itertools.product( - (32, 64), (2, 4), (2, 3, 4) - ): + for value_tile, num_warps, num_stages in itertools.product((32, 64), (2, 4), (2, 3, 4)): config = { "BV": value_tile, "num_warps": num_warps, diff --git a/tools/bench_glm53_sglang_moe.py b/tools/bench_glm53_sglang_moe.py index 10cfc0f03b..6b2a378b30 100644 --- a/tools/bench_glm53_sglang_moe.py +++ b/tools/bench_glm53_sglang_moe.py @@ -132,20 +132,14 @@ def main(): device = torch.device("cuda:0") experts, hidden, intermediate, topk = 289, 4096, 2048 // args.tp_size, 9 x = torch.zeros((args.tokens, hidden), dtype=torch.bfloat16, device=device) - w1 = torch.zeros( - (experts, intermediate * 2, hidden), dtype=torch.float8_e4m3fn, device=device - ) - w2 = torch.zeros( - (experts, hidden, intermediate), dtype=torch.float8_e4m3fn, device=device - ) + w1 = torch.zeros((experts, intermediate * 2, hidden), dtype=torch.float8_e4m3fn, device=device) + w2 = torch.zeros((experts, hidden, intermediate), dtype=torch.float8_e4m3fn, device=device) w1_scale = torch.ones((experts, 4, 32), dtype=torch.float32, device=device) w2_scale = torch.ones((experts, 32, 2), dtype=torch.float32, device=device) rows = torch.arange(args.tokens, dtype=torch.int64, device=device)[:, None] cols = torch.arange(topk, dtype=torch.int64, device=device)[None, :] topk_ids = (rows * topk + cols) % experts - topk_weights = torch.full( - (args.tokens, topk), 1.0 / topk, dtype=torch.float32, device=device - ) + topk_weights = torch.full((args.tokens, topk), 1.0 / topk, dtype=torch.float32, device=device) fixed_up_config = make_config( { @@ -224,11 +218,7 @@ def timed_kernel(*kernel_args, **kernel_kwargs): if args.tune_down: # Both GEMMs share the same token alignment, so the down projection's # BLOCK_SIZE_M must match the fixed up projection. - selected_configs = [ - values - for values in CONFIGS - if values[0] == fixed_up_config["BLOCK_SIZE_M"] - ] + selected_configs = [values for values in CONFIGS if values[0] == fixed_up_config["BLOCK_SIZE_M"]] for values in selected_configs[: args.max_configs]: config = make_config(values) try: diff --git a/tools/bench_glm53_sparse_prefill.py b/tools/bench_glm53_sparse_prefill.py index 3e512abdb1..df39ebbc3b 100644 --- a/tools/bench_glm53_sparse_prefill.py +++ b/tools/bench_glm53_sparse_prefill.py @@ -58,7 +58,7 @@ def measure_ms( "median_ms": statistics.median(samples), "min_ms": min(samples), "max_ms": max(samples), - "peak_delta_gib": peak_delta_bytes / 2**30, + "peak_delta_gib": peak_delta_bytes / 2 ** 30, } @@ -94,7 +94,7 @@ def main() -> None: device="cuda", ) indices = make_causal_indices(args.tokens, args.sequence_length, args.topk) - scale = args.head_dim**-0.5 + scale = args.head_dim ** -0.5 padded_q = q.new_zeros((args.tokens, args.required_heads, args.head_dim)) padded_q[:, : args.local_heads].copy_(q) @@ -111,13 +111,9 @@ def flashmla_kernel() -> torch.Tensor: def flashmla_lightllm_path() -> torch.Tensor: q_input = q.new_zeros((args.tokens, args.required_heads, args.head_dim)) q_input[:, : args.local_heads].copy_(q) - return flash_mla_sparse_fwd( - q_input, - kv, - indices, - scale, - d_v=args.head_dim, - )[0][:, : args.local_heads] + return flash_mla_sparse_fwd(q_input, kv, indices, scale, d_v=args.head_dim,)[ + 0 + ][:, : args.local_heads] def tilelang_lightllm_path() -> torch.Tensor: output = tilelang_sparse_fwd( diff --git a/tools/bench_glm53_sparse_prefill_tp.py b/tools/bench_glm53_sparse_prefill_tp.py index b9628b6564..9f886f85da 100644 --- a/tools/bench_glm53_sparse_prefill_tp.py +++ b/tools/bench_glm53_sparse_prefill_tp.py @@ -88,7 +88,7 @@ def measure_ms( del result peak_delta_gib = global_max( - (torch.cuda.max_memory_allocated() - baseline_bytes) / 2**30, + (torch.cuda.max_memory_allocated() - baseline_bytes) / 2 ** 30, device, ) return { @@ -139,7 +139,7 @@ def main() -> None: device=device, ) indices = make_causal_indices(args.tokens, args.sequence_length, args.topk) - scale = args.head_dim**-0.5 + scale = args.head_dim ** -0.5 token_start = rank * (args.tokens // world_size) token_end = token_start + args.tokens // world_size local_indices = indices[token_start:token_end] @@ -153,13 +153,9 @@ def main() -> None: def padded_flashmla() -> torch.Tensor: q_input = q.new_zeros((args.tokens, global_heads, args.head_dim)) q_input[:, : args.local_heads].copy_(q) - return flash_mla_sparse_fwd( - q_input, - kv, - indices, - scale, - d_v=args.head_dim, - )[0][:, : args.local_heads] + return flash_mla_sparse_fwd(q_input, kv, indices, scale, d_v=args.head_dim,)[ + 0 + ][:, : args.local_heads] def transposed_flashmla() -> torch.Tensor: transposed_q = head_shards_to_token_shards(q, world_size) @@ -193,12 +189,7 @@ def transposed_flashmla_workspace() -> torch.Tensor: scale, d_v=args.head_dim, )[0] - comm_workspace.view( - world_size, - args.tokens // world_size, - args.local_heads, - args.head_dim, - ).copy_( + comm_workspace.view(world_size, args.tokens // world_size, args.local_heads, args.head_dim,).copy_( transposed_output.view( args.tokens // world_size, world_size, diff --git a/tools/check_glm53_symm_out_of_place.py b/tools/check_glm53_symm_out_of_place.py index 32443b5a12..be8f928160 100644 --- a/tools/check_glm53_symm_out_of_place.py +++ b/tools/check_glm53_symm_out_of_place.py @@ -36,9 +36,7 @@ def _check(reducer: SymmMemAllreduce, *, inference: bool) -> None: def main() -> None: local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) - dist.init_process_group( - "nccl", device_id=torch.device("cuda", local_rank) - ) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) reducer = SymmMemAllreduce( dist.group.WORLD, torch.cuda.current_device(), diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh index e29ddefb05..a2ad31d7c2 100755 --- a/tools/run_glm53_h100_container.sh +++ b/tools/run_glm53_h100_container.sh @@ -1,18 +1,19 @@ #!/usr/bin/env bash set -euo pipefail -image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:h100-tp8}" -name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm}" +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:vl-1m-tp8}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-vl-1m}" model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" +deep_gemm_cache_dir="${LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR:-/home/devsft/cache-glm53-deep-gemm}" if [[ ! -d "${model_dir}" ]]; then echo "model directory does not exist: ${model_dir}" >&2 exit 1 fi -mkdir -p "${cache_dir}" "${triton_cache_dir}" +mkdir -p "${cache_dir}" "${triton_cache_dir}" "${deep_gemm_cache_dir}" exec sudo docker run --rm --name "${name}" \ --gpus all \ @@ -23,4 +24,5 @@ exec sudo docker run --rm --name "${name}" \ -v "${model_dir}:/model:ro" \ -v "${cache_dir}:/root/.cache" \ -v "${triton_cache_dir}:/root/.triton" \ + -v "${deep_gemm_cache_dir}:/root/.deep_gemm" \ "${image}" diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py index 087dd2d7a4..44a5ac8c44 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py @@ -33,13 +33,9 @@ def test_fused_kda_gate_matches_materialized_gate(sequence_length): device="cuda", dtype=torch.bfloat16, ) - state_indices = torch.arange(token_count, device="cuda", dtype=torch.int32).view( - request_count, sequence_length - ) + state_indices = torch.arange(token_count, device="cuda", dtype=torch.int32).view(request_count, sequence_length) cu_seqlens = torch.arange(request_count + 1, device="cuda", dtype=torch.int64) * sequence_length - accepted = torch.full( - (request_count,), sequence_length, device="cuda", dtype=torch.int32 - ) + accepted = torch.full((request_count,), sequence_length, device="cuda", dtype=torch.int32) reference_state = state.clone() reference, _ = fused_recurrent_gated_delta_rule( From f4b94b916d7518a9b75daed08629898233791c27 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 29 Aug 2026 04:35:17 +0800 Subject: [PATCH 14/28] build: split flattened runtime for registry upload --- docker/Dockerfile.glm53-h100 | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index 9d25686bed..3dc3e3753a 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -39,7 +39,14 @@ RUN python -m pip install --no-cache-dir --no-deps . && \ # Flatten the prepared root filesystem so deleted base-image host keys and # developer-only vulnerable components are absent from the published layers, -# rather than merely hidden by an OCI whiteout in a later layer. +# rather than merely hidden by an OCI whiteout in a later layer. Keep the +# Python NVIDIA runtime in a separate clean layer so private registries do not +# have to accept a single compressed blob larger than their upload gateway +# limit; the combined filesystem is identical to the prepared stage. +FROM prepared AS prepared_without_python_nvidia + +RUN rm -rf /opt/sglang/lib/python3.12/site-packages/nvidia + FROM scratch AS runtime ARG OCI_CREATED @@ -49,7 +56,9 @@ ARG OCI_VERSION ARG BASE_NAME=lmsysorg/sglang:glm-5.3-flash ARG BASE_DIGEST=sha256:92afb4c878eef9cbb17ca9a2c1d15d5cda58585f90bbf5915a79f0f6284aad10 -COPY --from=prepared / / +COPY --from=prepared_without_python_nvidia / / +COPY --from=prepared /opt/sglang/lib/python3.12/site-packages/nvidia \ + /opt/sglang/lib/python3.12/site-packages/nvidia LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ From 14798e45a07256b7f78aeaf8d1b0827e7b523060 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 29 Aug 2026 13:10:53 +0800 Subject: [PATCH 15/28] docker: shard flattened GLM runtime layers --- docker/Dockerfile.glm53-h100 | 70 +++++++++++++++---- docker/split_glm53_rootfs.py | 130 +++++++++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+), 13 deletions(-) create mode 100644 docker/split_glm53_rootfs.py diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index 3dc3e3753a..d3bffca40a 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -7,14 +7,16 @@ ARG PIP_INDEX_URL WORKDIR /opt/lightllm COPY docker/requirements-glm53-runtime.txt /tmp/requirements-glm53-runtime.txt -RUN python -m pip install --no-cache-dir --no-deps --require-hashes \ +RUN --mount=type=cache,id=glm53-pip,target=/root/.cache/pip \ + python -m pip install --no-deps --require-hashes \ -r /tmp/requirements-glm53-runtime.txt && \ rm /tmp/requirements-glm53-runtime.txt COPY setup.py LICENSE README.md ./ COPY lightllm ./lightllm COPY docker/glm53-h100.openvex.json /usr/share/doc/lightllm/glm53-h100.openvex.json -RUN python -m pip install --no-cache-dir --no-deps . && \ +RUN --mount=type=cache,id=glm53-pip,target=/root/.cache/pip \ + python -m pip install --no-deps . && \ rm -rf \ /etc/ssh/ssh_host_*_key* \ /opt/nvidia/nsight-compute/2025.3.1/host/target-linux-x64/plugins/efa_metrics \ @@ -38,14 +40,26 @@ RUN python -m pip install --no-cache-dir --no-deps . && \ test ! -e /sgl-workspace/sglang/python/sglang/srt/disaggregation # Flatten the prepared root filesystem so deleted base-image host keys and -# developer-only vulnerable components are absent from the published layers, -# rather than merely hidden by an OCI whiteout in a later layer. Keep the -# Python NVIDIA runtime in a separate clean layer so private registries do not -# have to accept a single compressed blob larger than their upload gateway -# limit; the combined filesystem is identical to the prepared stage. -FROM prepared AS prepared_without_python_nvidia - -RUN rm -rf /opt/sglang/lib/python3.12/site-packages/nvidia +# developer-only vulnerable components are absent from published layers, +# rather than merely hidden by an OCI whiteout. Size-bound the clean content +# layers because the private registry's cross-region gateway cannot reliably +# accept one multi-gigabyte upload stream; the merged filesystem is unchanged. +FROM prepared AS partitioned + +COPY docker/split_glm53_rootfs.py /tmp/split_glm53_rootfs.py +RUN /usr/bin/python3 /tmp/split_glm53_rootfs.py + +FROM prepared AS prepared_remainder + +RUN rm -rf \ + /opt \ + /root \ + /sgl-workspace \ + /usr/include \ + /usr/lib/x86_64-linux-gnu \ + /usr/libexec \ + /usr/local \ + /usr/share FROM scratch AS runtime @@ -56,9 +70,39 @@ ARG OCI_VERSION ARG BASE_NAME=lmsysorg/sglang:glm-5.3-flash ARG BASE_DIGEST=sha256:92afb4c878eef9cbb17ca9a2c1d15d5cda58585f90bbf5915a79f0f6284aad10 -COPY --from=prepared_without_python_nvidia / / -COPY --from=prepared /opt/sglang/lib/python3.12/site-packages/nvidia \ - /opt/sglang/lib/python3.12/site-packages/nvidia +COPY --from=prepared_remainder / / +COPY --from=partitioned /__image_layers/00/ / +COPY --from=partitioned /__image_layers/01/ / +COPY --from=partitioned /__image_layers/02/ / +COPY --from=partitioned /__image_layers/03/ / +COPY --from=partitioned /__image_layers/04/ / +COPY --from=partitioned /__image_layers/05/ / +COPY --from=partitioned /__image_layers/06/ / +COPY --from=partitioned /__image_layers/07/ / +COPY --from=partitioned /__image_layers/08/ / +COPY --from=partitioned /__image_layers/09/ / +COPY --from=partitioned /__image_layers/10/ / +COPY --from=partitioned /__image_layers/11/ / +COPY --from=partitioned /__image_layers/12/ / +COPY --from=partitioned /__image_layers/13/ / +COPY --from=partitioned /__image_layers/14/ / +COPY --from=partitioned /__image_layers/15/ / +COPY --from=partitioned /__image_layers/16/ / +COPY --from=partitioned /__image_layers/17/ / +COPY --from=partitioned /__image_layers/18/ / +COPY --from=partitioned /__image_layers/19/ / +COPY --from=partitioned /__image_layers/20/ / +COPY --from=partitioned /__image_layers/21/ / +COPY --from=partitioned /__image_layers/22/ / +COPY --from=partitioned /__image_layers/23/ / +COPY --from=partitioned /__image_layers/24/ / +COPY --from=partitioned /__image_layers/25/ / +COPY --from=partitioned /__image_layers/26/ / +COPY --from=partitioned /__image_layers/27/ / +COPY --from=partitioned /__image_layers/28/ / +COPY --from=partitioned /__image_layers/29/ / +COPY --from=partitioned /__image_layers/30/ / +COPY --from=partitioned /__image_layers/31/ / LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ diff --git a/docker/split_glm53_rootfs.py b/docker/split_glm53_rootfs.py new file mode 100644 index 0000000000..5dde20a1c8 --- /dev/null +++ b/docker/split_glm53_rootfs.py @@ -0,0 +1,130 @@ +#!/usr/bin/python3 +"""Move large runtime trees into deterministic, size-bounded image layers.""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path +from typing import NamedTuple + + +ATOM_LIMIT = 512 * 1024 * 1024 +BUCKET_COUNT = 32 +BUCKET_LIMIT = 1_250_000_000 +LAYER_ROOT = Path("/__image_layers") +SOURCES = ( + Path("/opt"), + Path("/root"), + Path("/sgl-workspace"), + Path("/usr/include"), + Path("/usr/lib/x86_64-linux-gnu"), + Path("/usr/libexec"), + Path("/usr/local"), + Path("/usr/share"), +) + + +class Node(NamedTuple): + path: Path + size: int + children: tuple["Node", ...] + + +def scan_tree(path: Path, directory_stats: dict[Path, os.stat_result]) -> Node: + metadata = path.lstat() + if not stat.S_ISDIR(metadata.st_mode): + return Node(path, metadata.st_size, ()) + + directory_stats[path] = metadata + children = tuple( + scan_tree(Path(entry.path), directory_stats) for entry in sorted(os.scandir(path), key=lambda item: item.name) + ) + return Node(path, metadata.st_size + sum(child.size for child in children), children) + + +def iter_atoms(node: Node): + if not node.children or node.size <= ATOM_LIMIT: + yield node.path, node.size + return + for child in node.children: + yield from iter_atoms(child) + + +def make_parent_directories( + source: Path, + destination: Path, + directory_stats: dict[Path, os.stat_result], + created_directories: dict[Path, Path], +) -> None: + source_parent = Path("/") + destination_parent = destination.parent + relative_parents = source.relative_to("/").parts[:-1] + destination_cursor = LAYER_ROOT / destination.relative_to(LAYER_ROOT).parts[0] + + for component in relative_parents: + source_parent /= component + destination_cursor /= component + if not destination_cursor.exists(): + destination_cursor.mkdir() + created_directories[destination_cursor] = source_parent + + assert destination_cursor == destination_parent + + +def restore_directory_metadata(destination: Path, source: Path, directory_stats: dict[Path, os.stat_result]) -> None: + metadata = directory_stats[source] + os.chown(destination, metadata.st_uid, metadata.st_gid, follow_symlinks=False) + os.chmod(destination, stat.S_IMODE(metadata.st_mode), follow_symlinks=False) + os.utime( + destination, + ns=(metadata.st_atime_ns, metadata.st_mtime_ns), + follow_symlinks=False, + ) + + +def main() -> None: + Path(__file__).unlink() + LAYER_ROOT.mkdir(mode=0o755) + + directory_stats: dict[Path, os.stat_result] = {} + for source in SOURCES: + parent = source.parent + while parent != Path("/"): + directory_stats.setdefault(parent, parent.lstat()) + parent = parent.parent + nodes = tuple(scan_tree(source, directory_stats) for source in SOURCES) + atoms = sorted( + (atom for node in nodes for atom in iter_atoms(node)), + key=lambda item: (-item[1], str(item[0])), + ) + + bucket_sizes = [0] * BUCKET_COUNT + assignments: list[tuple[Path, int]] = [] + for source, size in atoms: + bucket = min(range(BUCKET_COUNT), key=lambda index: (bucket_sizes[index], index)) + bucket_sizes[bucket] += size + assignments.append((source, bucket)) + + largest_bucket = max(bucket_sizes) + if largest_bucket > BUCKET_LIMIT: + raise RuntimeError(f"largest rootfs bucket is {largest_bucket} bytes, over {BUCKET_LIMIT}") + + created_directories: dict[Path, Path] = {} + for index in range(BUCKET_COUNT): + (LAYER_ROOT / f"{index:02d}").mkdir(mode=0o755) + + for source, bucket in sorted(assignments, key=lambda item: str(item[0])): + destination = LAYER_ROOT / f"{bucket:02d}" / source.relative_to("/") + make_parent_directories(source, destination, directory_stats, created_directories) + source.rename(destination) + + for destination, source in sorted(created_directories.items(), key=lambda item: len(item[0].parts), reverse=True): + restore_directory_metadata(destination, source, directory_stats) + + for index, size in enumerate(bucket_sizes): + print(f"rootfs layer {index:02d}: {size} uncompressed bytes") + + +if __name__ == "__main__": + main() From efff73a280db119be82fc7a94b133199516c0ac9 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 29 Aug 2026 13:15:32 +0800 Subject: [PATCH 16/28] docker: preserve metadata during layer copy-up --- docker/split_glm53_rootfs.py | 67 +++++++++++++++++++++++++++++++++++- 1 file changed, 66 insertions(+), 1 deletion(-) diff --git a/docker/split_glm53_rootfs.py b/docker/split_glm53_rootfs.py index 5dde20a1c8..e5af15989b 100644 --- a/docker/split_glm53_rootfs.py +++ b/docker/split_glm53_rootfs.py @@ -3,7 +3,9 @@ from __future__ import annotations +import errno import os +import shutil import stat from pathlib import Path from typing import NamedTuple @@ -83,6 +85,68 @@ def restore_directory_metadata(destination: Path, source: Path, directory_stats: ) +def copy_path( + source: Path, + destination: Path, + copied_inodes: dict[tuple[int, int], Path], +) -> None: + metadata = source.lstat() + inode = (metadata.st_dev, metadata.st_ino) + + if stat.S_ISLNK(metadata.st_mode): + destination.symlink_to(os.readlink(source)) + os.lchown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + return + + if stat.S_ISDIR(metadata.st_mode): + destination.mkdir(mode=stat.S_IMODE(metadata.st_mode)) + for entry in sorted(os.scandir(source), key=lambda item: item.name): + copy_path(Path(entry.path), destination / entry.name, copied_inodes) + os.chown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + return + + if stat.S_ISREG(metadata.st_mode): + previous = copied_inodes.get(inode) + if metadata.st_nlink > 1 and previous is not None: + destination.hardlink_to(previous) + else: + shutil.copyfile(source, destination, follow_symlinks=False) + copied_inodes[inode] = destination + os.chown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + return + + if stat.S_ISFIFO(metadata.st_mode): + os.mkfifo(destination, stat.S_IMODE(metadata.st_mode)) + elif stat.S_ISCHR(metadata.st_mode) or stat.S_ISBLK(metadata.st_mode): + os.mknod(destination, metadata.st_mode, metadata.st_rdev) + else: + raise RuntimeError(f"unsupported file type while copying {source}") + os.chown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + + +def move_path( + source: Path, + destination: Path, + copied_inodes: dict[tuple[int, int], Path], +) -> None: + try: + source.rename(destination) + return + except OSError as error: + if error.errno != errno.EXDEV: + raise + + copy_path(source, destination, copied_inodes) + if source.is_dir() and not source.is_symlink(): + shutil.rmtree(source) + else: + source.unlink() + + def main() -> None: Path(__file__).unlink() LAYER_ROOT.mkdir(mode=0o755) @@ -114,10 +178,11 @@ def main() -> None: for index in range(BUCKET_COUNT): (LAYER_ROOT / f"{index:02d}").mkdir(mode=0o755) + copied_inodes: dict[tuple[int, int], Path] = {} for source, bucket in sorted(assignments, key=lambda item: str(item[0])): destination = LAYER_ROOT / f"{bucket:02d}" / source.relative_to("/") make_parent_directories(source, destination, directory_stats, created_directories) - source.rename(destination) + move_path(source, destination, copied_inodes) for destination, source in sorted(created_directories.items(), key=lambda item: len(item[0].parts), reverse=True): restore_directory_metadata(destination, source, directory_stats) From 1b408dea79848835af05574832688a6b642329e9 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 29 Aug 2026 22:17:50 +0800 Subject: [PATCH 17/28] docs: record validated H100 deployment results --- GLM53_H100_DEPLOY.md | 57 +++++++++++++++++++++++++++++++++++--------- 1 file changed, 46 insertions(+), 11 deletions(-) diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 6d2863fd58..6990d5c3f3 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -47,14 +47,16 @@ tools/run_glm53_h100_container.sh ## Run the published image on H100 -Set `IMAGE` to the immutable registry tag or digest listed in the pull request. -The fixed container name is `glm53-lightllm-vl-1m`. The image default uses the -8,192-token prefill chunk validated on H200. The command below overrides the -default with a conservative 1,024-token chunk for an 80 GB H100, where the DSA -score matrix has much less temporary-memory headroom. +This command was validated on one eight-GPU H100 80 GB node. The fixed +container name is `glm53-lightllm-vl-1m`. Keep `batch_max_tokens` at 8,192 and +use a conservative 1,024-token prefill chunk: a 65,536-token batch maximum +OOMed during the server's startup length check because the DSA score matrix +exceeded the H100's temporary-memory headroom. Keep FlashInfer all-reduce +disabled for this profile; without that override, the first inference request +stalled across ranks on the tested host. ```bash -IMAGE=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm: +IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:2a664580a495215a5bfb48d96bf118a8321d7accde589e505de283d6ea5753b2" sudo docker pull "$IMAGE" sudo docker run -d \ @@ -81,11 +83,12 @@ sudo docker run -d \ --max_total_token_num 1048612 \ --running_max_req_size 256 \ --max_req_total_len 1048576 \ - --batch_max_tokens 65536 \ + --batch_max_tokens 8192 \ --chunked_prefill_size 1024 \ --linear_att_ssm_data_type bfloat16 \ --linear_att_cache_size 256 \ --disable_cudagraph \ + --disable_flashinfer_allreduce \ --enable_fused_shared_experts \ --max_image_pixels 6272000 \ --max_image_token_count 8000 \ @@ -106,7 +109,39 @@ curl --fail --show-error http://127.0.0.1:8002/v1/models sudo docker stop --timeout 30 glm53-lightllm-vl-1m ``` -The H200 accuracy, long-context, and throughput results in the pull request are -measured through the OpenAI-compatible endpoint. H100 performance is not -inferred from H200 data, and the conservative H100 override above must be -validated independently on the target host before production traffic. +## H100 validation + +The published digest and command above were exercised through the +OpenAI-compatible endpoint on one eight-GPU H100 80 GB node on 2026-08-29. +Every inference, evaluation, and benchmark command was recorded with `exp`. + +| Check | Result | +| --- | --- | +| Text smoke | `123 + 456` returned `579` in 15.34 s | +| Synthetic vision smoke | Identified the red square in 17.22 s | +| Exact 1M context needle | Both tokenizer and API counted exactly 1,000,000 prompt tokens; recovered `ZEBRA-4821` in 171.49 s | +| Sampled peak during 1M request | Approximately 80,063 MiB of 81,559 MiB per GPU; no OOM | +| SGLang-style latency workload | 10/10 requests; 3,309 input and 3,700 output tokens in 319.65 s; 21.93 total tok/s and 11.58 output tok/s | +| SGLang-style throughput workload | 1,000/1,000 requests; 504,929 input and 494,908 output tokens in 1,962.10 s; 509.58 total tok/s and 252.23 output tok/s | +| GSM8K | 99/100; all completed in 118.17 s and none reached the 2,048-token cap | +| MMMU vision | 64/100; all completed in 752.75 s and 39 reached the 2,048-token cap | + +The comprehensive black-box checker reported 14 passes, 6 failures, and 5 +skips. Core discovery, native generation and streaming, OpenAI chat streaming, +Responses API, multi-output, recovery, vision, tool parsing, and reasoning +parsing passed. Strict parity/determinism checks for completions text, seeded +token IDs, blocked-token filtering, prompt-cache hit reporting, and concurrent +output differed; Anthropic request translation returned HTTP 400. Treat the +deployment as operational for its validated native/OpenAI paths, not as a +claim that every optional compatibility path is green. + +The SGLang-style result files are: + +- `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c1.jsonl`; +- `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c100.jsonl`. + +Relevant experiment run prefixes are `260829-210446` (health checker), +`260829-210839` (text), `260829-210908` (vision), `260829-211036` +(exact 1M), `260829-211521` (concurrency 1), `260829-212146` +(concurrency 100), `260829-215835` (GSM8K), and `260829-220106` +(MMMU). From 3dc38408129fe6eca7d11db463afd0aa0ca27bda Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 30 Aug 2026 01:07:13 +0800 Subject: [PATCH 18/28] docs: add same-host GLM-5.3 engine benchmarks --- GLM53_H100_DEPLOY.md | 59 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 3 deletions(-) diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 6990d5c3f3..5c9067876a 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -126,6 +126,53 @@ Every inference, evaluation, and benchmark command was recorded with `exp`. | GSM8K | 99/100; all completed in 118.17 s and none reached the 2,048-token cap | | MMMU vision | 64/100; all completed in 752.75 s and 39 reached the 2,048-token cap | +### Same-host LightLLM, SGLang, and vLLM comparison + +The three engines were measured sequentially on the same otherwise-idle +8xH100 80 GB host on 2026-08-30. All used the same local FP8 checkpoint, TP8, +BF16 KV cache, a declared 1,048,576-token context, no speculative decoding, +and an 8,192-token prefill budget. The common SGLang `bench_serving` client +used seed 42, temperature 0, streaming, ignored EOS, an infinite request rate, +and one excluded warmup request. `random-range-ratio=0` sampled input and +output lengths from 1 through 1,000 tokens. The exact c1 samples contained +3,309 input and 3,700 output tokens; the exact c100 samples contained 504,929 +input and 494,908 output tokens. Every engine completed every request. + +| Engine | Pinned build and material server differences | +| --- | --- | +| LightLLM | Published image above; TP8, 1,024-token chunks, CUDA graph disabled, FlashInfer all-reduce disabled | +| SGLang | `lmsysorg/sglang@sha256:0836f0160fa785e424e68d13ef88ddd548f87e6e11ad9f0e4de982e4f9188aaf` (`0.0.0.dev1+gf609d677b`); TP8/EP8, DeepGEMM, `mem-fraction-static=0.80`, 1,024-token chunks | +| vLLM | `vllm/vllm-openai:glm53-flash-x86_64-cu130@sha256:2e771fa615452282cc331eb418b3ef21636fce355bea0491fca89e6d362ab703` (`0.1.dev20051+g487ecf187`); TP8, 256 sequences, chunked prefill and prefix caching, `gpu-memory-utilization=0.90` | + +| Engine | c1 duration | c1 output / total tok/s | c100 duration | c100 output / total tok/s | c100 total vs LightLLM | +| --- | ---: | ---: | ---: | ---: | ---: | +| LightLLM | 319.65 s | 11.58 / 21.93 | 1,962.10 s | 252.23 / 509.58 | 1.00x | +| SGLang | 91.23 s | 40.56 / 76.83 | 328.59 s | 1,506.16 / 3,042.82 | 5.97x | +| vLLM | 36.30 s | 101.92 / 193.07 | 270.74 s | 1,827.97 / 3,692.96 | 7.25x | + +At c1, SGLang and vLLM delivered 3.50x and 8.81x LightLLM's total +throughput, respectively; vLLM was 2.51x SGLang. At c100, vLLM was 1.214x +SGLang. SGLang and vLLM also passed the same arithmetic smoke check by +returning `579` for `37 * 16 - 13`. + +The SGLang server first failed during CUDA-graph capture at +`mem-fraction-static=0.90`; the reported run used its suggested `0.80` while +retaining the 1M context declaration and enough KV capacity for this workload. +vLLM used the dedicated pre-merge GLM-5.3 image because the ordinary public +image on the host did not register `Glm5Next`; its official recipe requires +BF16 KV on Hopper. Its startup also warned that the H100/288-expert combination +lacked a model-specific MoE tuning table. SGLang and vLLM expose reasoning +tokens under different streaming fields, and LightLLM's chunk shape is not +recognized consistently by this client, so cross-engine TTFT/ITL and +retokenized-text counts are not compared. Successful requests, API usage token +counts, wall time, and the aggregate throughput figures above are directly +comparable. + +References: [SGLang GLM-5 benchmark recipe](https://github.com/sgl-project/sglang/blob/main/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx), +[SGLang serving benchmark](https://github.com/sgl-project/sglang/blob/main/docs/cookbook/base/benchmarks/autoregressive_model_benchmark.mdx), +[vLLM GLM-5.3-Flash recipe](https://recipes.vllm.ai/zai-org/GLM-5.3-Flash), +and [vLLM GLM-5.3 support PR](https://github.com/vllm-project/vllm/pull/53906). + The comprehensive black-box checker reported 14 passes, 6 failures, and 5 skips. Core discovery, native generation and streaming, OpenAI chat streaming, Responses API, multi-output, recovery, vision, tool parsing, and reasoning @@ -135,13 +182,19 @@ output differed; Anthropic request translation returned HTTP 400. Treat the deployment as operational for its validated native/OpenAI paths, not as a claim that every optional compatibility path is green. -The SGLang-style result files are: +The benchmark result files are: - `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c1.jsonl`; -- `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c100.jsonl`. +- `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c100.jsonl`; +- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/sglang_h100_random_1k_c1.jsonl`; +- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/sglang_h100_random_1k_c100.jsonl`; +- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/vllm_h100_random_1k_c1.jsonl`; +- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/vllm_h100_random_1k_c100.jsonl`. Relevant experiment run prefixes are `260829-210446` (health checker), `260829-210839` (text), `260829-210908` (vision), `260829-211036` (exact 1M), `260829-211521` (concurrency 1), `260829-212146` (concurrency 100), `260829-215835` (GSM8K), and `260829-220106` -(MMMU). +(MMMU). Comparison runs are `260830-002938`, `260830-003031`, and +`260830-003246` for SGLang smoke/c1/c100, and `260830-005212`, +`260830-005226`, and `260830-005412` for vLLM smoke/c1/c100. From 8efce02695dd0873a8b2e83b09d180e08908c80f Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 30 Aug 2026 03:13:00 +0800 Subject: [PATCH 19/28] perf(glm5): tune H100 concurrency profile --- docker/Dockerfile.glm53-h100 | 28 +- ...=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json | 7 + ...torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json | 68 +++++ ...torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json | 278 ++++++++++++++++++ 4 files changed, 371 insertions(+), 10 deletions(-) create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index d3bffca40a..0ffdf42f7a 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -114,7 +114,7 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="multimodal-1m-c256" \ + ai.lightllm.profile="multimodal-1m-c100-ep8-tpsp-prefill-overlap" \ ai.lightllm.security-profile="flattened-no-sglang-server-components" ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ @@ -156,10 +156,10 @@ ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ CMD /opt/sglang/bin/python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 -# LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is +# LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is # therefore request_limit + 36 so /v1/models advertises the full 1,048,576. -# The 8,192-token DSA prefill chunk below is the H200-validated default; use -# the conservative 1,024-token override documented for an 80 GB H100. +# This EP8 + TP/SP profile is the exact full-multimodal configuration validated +# at 100-way concurrency on an eight-H100 80 GB node. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -169,14 +169,22 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ "--max_total_token_num", "1048612", \ - "--running_max_req_size", "256", \ + "--running_max_req_size", "104", \ "--max_req_total_len", "1048576", \ - "--batch_max_tokens", "65536", \ - "--chunked_prefill_size", "8192", \ + "--batch_max_tokens", "4096", \ + "--chunked_prefill_size", "1024", \ "--linear_att_ssm_data_type", "bfloat16", \ - "--linear_att_cache_size", "256", \ - "--disable_cudagraph", \ - "--enable_fused_shared_experts", \ + "--linear_att_cache_size", "104", \ + "--graph_max_batch_size", "104", \ + "--graph_split_batch_size", "8", \ + "--graph_grow_step_size", "16", \ + "--graph_max_len_in_batch", "2048", \ + "--disable_flashinfer_allreduce", \ + "--enable_ep_moe", \ + "--enable_tpsp_mix_mode", \ + "--enable_prefill_microbatch_overlap", \ + "--disable_aggressive_schedule", \ + "--router_max_wait_tokens", "64", \ "--max_image_pixels", "6272000", \ "--max_image_token_count", "8000", \ "--visual_tp", "1", \ diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..dfbc9746cc --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,7 @@ +{ + "8": { + "BV": 32, + "num_stages": 4, + "num_warps": 4 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..89c7d19be9 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,68 @@ +{ + "1": { + "BLOCK_M": 256, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 8 + }, + "1024": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "128": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "16": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 4 + }, + "2048": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "32": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "4096": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "8": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..9c0dd1629a --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,278 @@ +{ + "1": { + "BLOCK_M": 128, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 8 + }, + "10112": { + "BLOCK_M": 64, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10240": { + "BLOCK_M": 64, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10496": { + "BLOCK_M": 64, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1152": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "128": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1280": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1536": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 8 + }, + "1792": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18304": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18432": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18560": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1920": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "19200": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2176": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2304": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "2816": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2944": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "3328": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34048": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34688": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34944": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "35072": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "35584": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3584": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "3712": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "384": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3968": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4096": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4224": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4352": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4480": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4608": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4736": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "640": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "768": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "8": { + "BLOCK_M": 1, + "BLOCK_N": 64, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "9856": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + } +} From cc6319a4350c70e062218d85a1cbf1e8d6890cd5 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 30 Aug 2026 03:19:08 +0800 Subject: [PATCH 20/28] build: add offline GLM-5.3 release overlay --- docker/Dockerfile.glm53-h100-overlay | 72 ++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 docker/Dockerfile.glm53-h100-overlay diff --git a/docker/Dockerfile.glm53-h100-overlay b/docker/Dockerfile.glm53-h100-overlay new file mode 100644 index 0000000000..cabc2c1ef0 --- /dev/null +++ b/docker/Dockerfile.glm53-h100-overlay @@ -0,0 +1,72 @@ +# Fast release path for environments where the original Docker Hub base is +# unavailable. The pinned base is the already-flattened GLM-5.3 runtime; this +# layer replaces only the public LightLLM source and the validated default CMD. +ARG BASE_IMAGE=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:2a664580a495215a5bfb48d96bf118a8321d7accde589e505de283d6ea5753b2 +FROM ${BASE_IMAGE} + +ARG OCI_CREATED +ARG OCI_REVISION +ARG OCI_SOURCE=https://github.com/sufubao/LightLLM +ARG OCI_VERSION +ARG BASE_NAME=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm +ARG BASE_DIGEST=sha256:2a664580a495215a5bfb48d96bf118a8321d7accde589e505de283d6ea5753b2 + +WORKDIR /opt/lightllm + +COPY lightllm ./lightllm + +RUN test -f 'lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json' && \ + test -f 'lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json' && \ + test -f 'lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json' && \ + python -c "import lightllm.server.api_start; print('Updated LightLLM GLM-5.3 runtime import OK')" + +LABEL org.opencontainers.image.created="${OCI_CREATED}" \ + org.opencontainers.image.revision="${OCI_REVISION}" \ + org.opencontainers.image.source="${OCI_SOURCE}" \ + org.opencontainers.image.version="${OCI_VERSION}" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 multimodal" \ + org.opencontainers.image.base.name="${BASE_NAME}" \ + org.opencontainers.image.base.digest="${BASE_DIGEST}" \ + ai.lightllm.model="GLM-5.3-Flash" \ + ai.lightllm.accelerator="NVIDIA H100/H200" \ + ai.lightllm.tensor-parallel-size="8" \ + ai.lightllm.profile="multimodal-1m-c100-ep8-tpsp-prefill-overlap" \ + ai.lightllm.security-profile="flattened-base-plus-source-overlay" + +# LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is +# therefore request_limit + 36 so /v1/models advertises the full 1,048,576. +CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ + "--model_dir", "/model", \ + "--model_name", "glm-5.3-flash", \ + "--tp", "8", \ + "--host", "0.0.0.0", \ + "--port", "8002", \ + "--httpserver_workers", "16", \ + "--mem_fraction", ".90", \ + "--max_total_token_num", "1048612", \ + "--running_max_req_size", "104", \ + "--max_req_total_len", "1048576", \ + "--batch_max_tokens", "4096", \ + "--chunked_prefill_size", "1024", \ + "--linear_att_ssm_data_type", "bfloat16", \ + "--linear_att_cache_size", "104", \ + "--graph_max_batch_size", "104", \ + "--graph_split_batch_size", "8", \ + "--graph_grow_step_size", "16", \ + "--graph_max_len_in_batch", "2048", \ + "--disable_flashinfer_allreduce", \ + "--enable_ep_moe", \ + "--enable_tpsp_mix_mode", \ + "--enable_prefill_microbatch_overlap", \ + "--disable_aggressive_schedule", \ + "--router_max_wait_tokens", "64", \ + "--max_image_pixels", "6272000", \ + "--max_image_token_count", "8000", \ + "--visual_tp", "1", \ + "--visual_dp", "8", \ + "--visual_infer_batch_size", "8", \ + "--cache_capacity", "64", \ + "--schedule_time_interval", "0.001", \ + "--prefill_coalesce_interval", "0.5", \ + "--reasoning_parser", "glm45", \ + "--tool_call_parser", "glm47"] From c639f6b4954d615124a74b50b40afb253bad5654 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 30 Aug 2026 03:25:38 +0800 Subject: [PATCH 21/28] docs: publish final GLM-5.3 H100 profile --- GLM53_H100_DEPLOY.md | 268 ++++++++++++------------------ tools/run_glm53_h100_container.sh | 4 +- 2 files changed, 108 insertions(+), 164 deletions(-) diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 5c9067876a..60fab17761 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -1,200 +1,144 @@ -# GLM-5.3-Flash multimodal TP8 deployment +# GLM-5.3-Flash H100/H200 deployment and validation This branch packages LightLLM text and vision inference for GLM-5.3-Flash on -one eight-GPU H100 or H200 node. The image contains the LightLLM source and -runtime dependencies. Mount the model and compiler caches from the host. +one eight-GPU H100 or H200 node. The release profile keeps the 1,048,576-token +request limit and multimodal workers while optimizing 100-way concurrency with +EP8, TP/SP mixing, prefill microbatch overlap, and CUDA graphs. -The default command serves the OpenAI-compatible API on port 8002 with: +## Release image -- tensor parallel size 8; -- image encoding as eight data-parallel workers; -- a 1,048,576-token request limit; -- up to 256 active requests; -- 8,192-token chunked prefill, which bounds the DSA score matrix during - million-token requests, with CUDA graphs disabled for this profile; -- `glm45` reasoning and `glm47` tool-call parsers. +The image is published only to the requested private registry: -## Build a local image - -Use an immutable tag containing the full source revision: - -```bash -revision="$(git rev-parse HEAD)" -created="$(date -u +%Y-%m-%dT%H:%M:%SZ)" -version="v1.3.0-glm53-vl-1m-tp8-${revision}" - -docker buildx build --load --platform linux/amd64 \ - -f docker/Dockerfile.glm53-h100 \ - --build-arg "OCI_CREATED=${created}" \ - --build-arg "OCI_REVISION=${revision}" \ - --build-arg "OCI_VERSION=${version}" \ - -t "lightllm-glm53:${version}" \ - . - -docker tag "lightllm-glm53:${version}" lightllm-glm53:vl-1m-tp8 +```text +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.3.0-glm53-vl-1m-c100-ep8-cc6319a4350c70e062218d85a1cbf1e8d6890cd5 ``` -## Run on the local H200 node +Immutable manifest: -```bash -LIGHTLLM_GLM53_IMAGE=lightllm-glm53:vl-1m-tp8 \ -LIGHTLLM_GLM53_MODEL_DIR=/nvme/sufubao/models/GLM-5.3-Flash \ -LIGHTLLM_GLM53_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-lightllm-h200 \ -LIGHTLLM_GLM53_TRITON_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-triton-h200 \ -LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-deep-gemm-h200 \ -tools/run_glm53_h100_container.sh +```text +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:f456207d3869996c9cfaa17df73058296e8de3720ef3d7eda4d56abb4719ff14 ``` -## Run the published image on H100 +The corresponding local image is +`lightllm-glm53:vl-1m-c100-ep8` (image ID +`sha256:f5b1c4282af30f656275d9726637e8d86e7a2dd7b5a103069fb81c6bea92831d`). +The image embeds source revision `cc6319a4350c70e062218d85a1cbf1e8d6890cd5`, +the validated H100 autotune records, and the complete server command. -This command was validated on one eight-GPU H100 80 GB node. The fixed -container name is `glm53-lightllm-vl-1m`. Keep `batch_max_tokens` at 8,192 and -use a conservative 1,024-token prefill chunk: a 65,536-token batch maximum -OOMed during the server's startup length check because the DSA score matrix -exceeded the H100's temporary-memory headroom. Keep FlashInfer all-reduce -disabled for this profile; without that override, the first inference request -stalled across ranks on the tested host. +## Deploy on the H100 node + +No command override or autotune-config mount is required: ```bash -IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:2a664580a495215a5bfb48d96bf118a8321d7accde589e505de283d6ea5753b2" +IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:f456207d3869996c9cfaa17df73058296e8de3720ef3d7eda4d56abb4719ff14" sudo docker pull "$IMAGE" sudo docker run -d \ - --name glm53-lightllm-vl-1m \ + --name glm53-lightllm-vl-1m-c100-ep8 \ --restart unless-stopped \ + --network host \ + --ipc host \ + --security-opt label=disable \ --gpus all \ - --ipc=host \ - --network=host \ - --ulimit memlock=-1:-1 \ + --ulimit memlock=-1 \ --ulimit nofile=1048576:1048576 \ - -v /home/devsft/models/GLM-5.3-Flash:/model:ro \ - -v /home/devsft/cache-glm53-lightllm:/root/.cache \ -v /home/devsft/cache-glm53-triton:/root/.triton \ -v /home/devsft/cache-glm53-deep-gemm:/root/.deep_gemm \ - "$IMAGE" \ - /opt/sglang/bin/python -m lightllm.server.api_server \ - --model_dir /model \ - --model_name glm-5.3-flash \ - --tp 8 \ - --host 0.0.0.0 \ - --port 8002 \ - --httpserver_workers 16 \ - --mem_fraction .90 \ - --max_total_token_num 1048612 \ - --running_max_req_size 256 \ - --max_req_total_len 1048576 \ - --batch_max_tokens 8192 \ - --chunked_prefill_size 1024 \ - --linear_att_ssm_data_type bfloat16 \ - --linear_att_cache_size 256 \ - --disable_cudagraph \ - --disable_flashinfer_allreduce \ - --enable_fused_shared_experts \ - --max_image_pixels 6272000 \ - --max_image_token_count 8000 \ - --visual_tp 1 \ - --visual_dp 8 \ - --visual_infer_batch_size 8 \ - --cache_capacity 64 \ - --schedule_time_interval 0.001 \ - --prefill_coalesce_interval 0.5 \ - --reasoning_parser glm45 \ - --tool_call_parser glm47 + -v /home/devsft/models/GLM-5.3-Flash:/model:ro \ + -v /home/devsft/cache-glm53-lightllm:/root/.cache \ + "$IMAGE" ``` -Wait for the model endpoint, then stop the deployment when required: +The endpoint is `http://127.0.0.1:8002/v1`. Startup takes several minutes on +the tested host. Check it with: ```bash +sudo docker ps --filter name=glm53-lightllm-vl-1m-c100-ep8 curl --fail --show-error http://127.0.0.1:8002/v1/models -sudo docker stop --timeout 30 glm53-lightllm-vl-1m ``` -## H100 validation - -The published digest and command above were exercised through the -OpenAI-compatible endpoint on one eight-GPU H100 80 GB node on 2026-08-29. -Every inference, evaluation, and benchmark command was recorded with `exp`. +## Run the local image on H200 -| Check | Result | -| --- | --- | -| Text smoke | `123 + 456` returned `579` in 15.34 s | -| Synthetic vision smoke | Identified the red square in 17.22 s | -| Exact 1M context needle | Both tokenizer and API counted exactly 1,000,000 prompt tokens; recovered `ZEBRA-4821` in 171.49 s | -| Sampled peak during 1M request | Approximately 80,063 MiB of 81,559 MiB per GPU; no OOM | -| SGLang-style latency workload | 10/10 requests; 3,309 input and 3,700 output tokens in 319.65 s; 21.93 total tok/s and 11.58 output tok/s | -| SGLang-style throughput workload | 1,000/1,000 requests; 504,929 input and 494,908 output tokens in 1,962.10 s; 509.58 total tok/s and 252.23 output tok/s | -| GSM8K | 99/100; all completed in 118.17 s and none reached the 2,048-token cap | -| MMMU vision | 64/100; all completed in 752.75 s and 39 reached the 2,048-token cap | +```bash +LIGHTLLM_GLM53_IMAGE=lightllm-glm53:vl-1m-c100-ep8 \ +LIGHTLLM_GLM53_MODEL_DIR=/nvme/sufubao/models/GLM-5.3-Flash \ +LIGHTLLM_GLM53_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-lightllm-h200 \ +LIGHTLLM_GLM53_TRITON_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-triton-h200 \ +LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-deep-gemm-h200 \ +tools/run_glm53_h100_container.sh +``` -### Same-host LightLLM, SGLang, and vLLM comparison +## Same-host c100 comparison -The three engines were measured sequentially on the same otherwise-idle +LightLLM, vLLM, and SGLang were run sequentially on the same otherwise-idle 8xH100 80 GB host on 2026-08-30. All used the same local FP8 checkpoint, TP8, -BF16 KV cache, a declared 1,048,576-token context, no speculative decoding, -and an 8,192-token prefill budget. The common SGLang `bench_serving` client -used seed 42, temperature 0, streaming, ignored EOS, an infinite request rate, -and one excluded warmup request. `random-range-ratio=0` sampled input and -output lengths from 1 through 1,000 tokens. The exact c1 samples contained -3,309 input and 3,700 output tokens; the exact c100 samples contained 504,929 -input and 494,908 output tokens. Every engine completed every request. - -| Engine | Pinned build and material server differences | -| --- | --- | -| LightLLM | Published image above; TP8, 1,024-token chunks, CUDA graph disabled, FlashInfer all-reduce disabled | -| SGLang | `lmsysorg/sglang@sha256:0836f0160fa785e424e68d13ef88ddd548f87e6e11ad9f0e4de982e4f9188aaf` (`0.0.0.dev1+gf609d677b`); TP8/EP8, DeepGEMM, `mem-fraction-static=0.80`, 1,024-token chunks | -| vLLM | `vllm/vllm-openai:glm53-flash-x86_64-cu130@sha256:2e771fa615452282cc331eb418b3ef21636fce355bea0491fca89e6d362ab703` (`0.1.dev20051+g487ecf187`); TP8, 256 sequences, chunked prefill and prefix caching, `gpu-memory-utilization=0.90` | - -| Engine | c1 duration | c1 output / total tok/s | c100 duration | c100 output / total tok/s | c100 total vs LightLLM | -| --- | ---: | ---: | ---: | ---: | ---: | -| LightLLM | 319.65 s | 11.58 / 21.93 | 1,962.10 s | 252.23 / 509.58 | 1.00x | -| SGLang | 91.23 s | 40.56 / 76.83 | 328.59 s | 1,506.16 / 3,042.82 | 5.97x | -| vLLM | 36.30 s | 101.92 / 193.07 | 270.74 s | 1,827.97 / 3,692.96 | 7.25x | - -At c1, SGLang and vLLM delivered 3.50x and 8.81x LightLLM's total -throughput, respectively; vLLM was 2.51x SGLang. At c100, vLLM was 1.214x -SGLang. SGLang and vLLM also passed the same arithmetic smoke check by -returning `579` for `37 * 16 - 13`. - -The SGLang server first failed during CUDA-graph capture at -`mem-fraction-static=0.90`; the reported run used its suggested `0.80` while -retaining the 1M context declaration and enough KV capacity for this workload. -vLLM used the dedicated pre-merge GLM-5.3 image because the ordinary public -image on the host did not register `Glm5Next`; its official recipe requires -BF16 KV on Hopper. Its startup also warned that the H100/288-expert combination -lacked a model-specific MoE tuning table. SGLang and vLLM expose reasoning -tokens under different streaming fields, and LightLLM's chunk shape is not -recognized consistently by this client, so cross-engine TTFT/ITL and -retokenized-text counts are not compared. Successful requests, API usage token -counts, wall time, and the aggregate throughput figures above are directly -comparable. +BF16 KV cache, declared 1,048,576-token context, and no speculative decoding. +The common SGLang `bench_serving` client used seed 42, temperature 0, +streaming, ignored EOS, infinite request rate, one excluded warmup request, +and random input/output lengths of 1--1,000 tokens. The fixed 1,000-request +sample contained 504,929 input and 494,908 generated tokens. Every engine +completed 1,000/1,000 requests. + +| Engine | Duration | Request/s | Input tok/s | Output tok/s | Total tok/s | LightLLM lead | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| **LightLLM** | **233.37 s** | **4.285** | **2,163.67** | **2,120.73** | **4,284.40** | -- | +| vLLM | 270.74 s | 3.694 | 1,864.99 | 1,827.97 | 3,692.96 | **+16.02%** | +| SGLang | 328.59 s | 3.043 | 1,536.66 | 1,506.16 | 3,042.82 | **+40.80%** | + +The compared server builds and material settings were: + +- LightLLM: this release profile; EP8, TP/SP mixing, prefill microbatch + overlap, 1,024-token chunks, graph batches through 104, and FlashInfer + all-reduce disabled. +- vLLM: `vllm/vllm-openai:glm53-flash-x86_64-cu130@sha256:2e771fa615452282cc331eb418b3ef21636fce355bea0491fca89e6d362ab703` + (`0.1.dev20051+g487ecf187`); TP8, 256 sequences, chunked prefill, prefix + caching, and `gpu-memory-utilization=0.90`. +- SGLang: `lmsysorg/sglang@sha256:0836f0160fa785e424e68d13ef88ddd548f87e6e11ad9f0e4de982e4f9188aaf` + (`0.0.0.dev1+gf609d677b`); TP8/EP8, DeepGEMM, + `mem-fraction-static=0.80`, and 1,024-token chunks. + +Cross-engine TTFT/ITL and client-retokenized text are intentionally omitted: +the engines expose reasoning text under different streaming fields. The API +usage token totals, successful-request count, wall time, and aggregate +throughput above are directly comparable. + +Result files: + +- LightLLM: `/nvme/sufubao/m39-home/results/glm53_h100_optimization/lightllm_ep8_tpsp_prefill_overlap_graph_c104_b4096_wait64_tuned_random_1k_c100_full1000.jsonl` +- vLLM: `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/vllm_h100_random_1k_c100.jsonl` +- SGLang: `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/sglang_h100_random_1k_c100.jsonl` + +The LightLLM run is archived by `exp` under +`~/experiments/runs/260830-025933-usr-bin-timeout-2400-docker-run-rm-network-host-`. + +## Accuracy and capability checks + +The final optimized execution path was evaluated before publication. Every +evaluation was recorded with `exp`. + +| Check | Result | Experiment | +| --- | --- | --- | +| GSM8K, fixed first 100, 5-shot, greedy | **99/100**, 100/100 completed, no 2,048-token truncation | `260830-030435-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-PYTHONPA` | +| MMMU, fixed 100, full multimodal, greedy | **63/100**, 100/100 completed; 34 answers reached the 2,048-token cap | `260830-030533-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-PYTHONPA` | +| Exact 1M-context needle | API and tokenizer both counted 1,000,000 prompt tokens; recovered `ZEBRA-4821` | `260829-211036` | +| Text smoke | Arithmetic answer `579` | `260829-210839` | +| Synthetic vision smoke | Correctly identified the red square | `260829-210908` | + +The previous conservative profile scored 64/100 on the same MMMU slice; the +optimized profile's one-answer difference is within this 100-example sample, +while it reduced capped answers from 39 to 34. The multimodal and 1M-context +features remain enabled in the published command. References: [SGLang GLM-5 benchmark recipe](https://github.com/sgl-project/sglang/blob/main/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx), [SGLang serving benchmark](https://github.com/sgl-project/sglang/blob/main/docs/cookbook/base/benchmarks/autoregressive_model_benchmark.mdx), [vLLM GLM-5.3-Flash recipe](https://recipes.vllm.ai/zai-org/GLM-5.3-Flash), and [vLLM GLM-5.3 support PR](https://github.com/vllm-project/vllm/pull/53906). -The comprehensive black-box checker reported 14 passes, 6 failures, and 5 -skips. Core discovery, native generation and streaming, OpenAI chat streaming, -Responses API, multi-output, recovery, vision, tool parsing, and reasoning -parsing passed. Strict parity/determinism checks for completions text, seeded -token IDs, blocked-token filtering, prompt-cache hit reporting, and concurrent -output differed; Anthropic request translation returned HTTP 400. Treat the -deployment as operational for its validated native/OpenAI paths, not as a -claim that every optional compatibility path is green. - -The benchmark result files are: - -- `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c1.jsonl`; -- `/nvme/sufubao/m39-home/results/glm53_h100_bench/lightllm_h100_sglang_style_c100.jsonl`; -- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/sglang_h100_random_1k_c1.jsonl`; -- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/sglang_h100_random_1k_c100.jsonl`; -- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/vllm_h100_random_1k_c1.jsonl`; -- `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/vllm_h100_random_1k_c100.jsonl`. - -Relevant experiment run prefixes are `260829-210446` (health checker), -`260829-210839` (text), `260829-210908` (vision), `260829-211036` -(exact 1M), `260829-211521` (concurrency 1), `260829-212146` -(concurrency 100), `260829-215835` (GSM8K), and `260829-220106` -(MMMU). Comparison runs are `260830-002938`, `260830-003031`, and -`260830-003246` for SGLang smoke/c1/c100, and `260830-005212`, -`260830-005226`, and `260830-005412` for vLLM smoke/c1/c100. +## Rebuild note + +`docker/Dockerfile.glm53-h100` is the complete reproducible build. The release +was produced with `docker/Dockerfile.glm53-h100-overlay` from the previous +immutable, flattened private image because Docker Hub base-metadata requests +timed out during release. The overlay replaces the complete LightLLM source, +checks the three new autotune records and import path, and pins both its source +revision and base digest in OCI labels. diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh index a2ad31d7c2..2e0d7d32e0 100755 --- a/tools/run_glm53_h100_container.sh +++ b/tools/run_glm53_h100_container.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash set -euo pipefail -image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:vl-1m-tp8}" -name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-vl-1m}" +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:vl-1m-c100-ep8}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-vl-1m-c100-ep8}" model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" From 62b8a9da25f7519822657c8126017fbc2793a08a Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 30 Aug 2026 16:32:23 +0800 Subject: [PATCH 22/28] perf(glm5): publish all-concurrency TP8 profile --- docker/Dockerfile.glm53-h100 | 19 +++++++++---------- docker/Dockerfile.glm53-h100-overlay | 15 +++++++-------- tools/run_glm53_h100_container.sh | 4 ++-- 3 files changed, 18 insertions(+), 20 deletions(-) diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index 0ffdf42f7a..7f38d4ab58 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -114,7 +114,7 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="multimodal-1m-c100-ep8-tpsp-prefill-overlap" \ + ai.lightllm.profile="multimodal-1m-tp8-c256-no-prompt-cache" \ ai.lightllm.security-profile="flattened-no-sglang-server-components" ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ @@ -158,8 +158,8 @@ HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ # LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is # therefore request_limit + 36 so /v1/models advertises the full 1,048,576. -# This EP8 + TP/SP profile is the exact full-multimodal configuration validated -# at 100-way concurrency on an eight-H100 80 GB node. +# This pure-TP8 profile is the exact full-multimodal configuration validated +# from 1 through 256 concurrent requests on an eight-H100 80 GB node. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -169,20 +169,19 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ "--max_total_token_num", "1048612", \ - "--running_max_req_size", "104", \ + "--running_max_req_size", "256", \ "--max_req_total_len", "1048576", \ - "--batch_max_tokens", "4096", \ + "--batch_max_tokens", "16384", \ "--chunked_prefill_size", "1024", \ "--linear_att_ssm_data_type", "bfloat16", \ - "--linear_att_cache_size", "104", \ - "--graph_max_batch_size", "104", \ + "--linear_att_cache_size", "256", \ + "--graph_max_batch_size", "256", \ "--graph_split_batch_size", "8", \ "--graph_grow_step_size", "16", \ "--graph_max_len_in_batch", "2048", \ "--disable_flashinfer_allreduce", \ - "--enable_ep_moe", \ - "--enable_tpsp_mix_mode", \ - "--enable_prefill_microbatch_overlap", \ + "--enable_fused_shared_experts", \ + "--disable_dynamic_prompt_cache", \ "--disable_aggressive_schedule", \ "--router_max_wait_tokens", "64", \ "--max_image_pixels", "6272000", \ diff --git a/docker/Dockerfile.glm53-h100-overlay b/docker/Dockerfile.glm53-h100-overlay index cabc2c1ef0..0055926840 100644 --- a/docker/Dockerfile.glm53-h100-overlay +++ b/docker/Dockerfile.glm53-h100-overlay @@ -30,7 +30,7 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="multimodal-1m-c100-ep8-tpsp-prefill-overlap" \ + ai.lightllm.profile="multimodal-1m-tp8-c256-no-prompt-cache" \ ai.lightllm.security-profile="flattened-base-plus-source-overlay" # LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is @@ -44,20 +44,19 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ "--max_total_token_num", "1048612", \ - "--running_max_req_size", "104", \ + "--running_max_req_size", "256", \ "--max_req_total_len", "1048576", \ - "--batch_max_tokens", "4096", \ + "--batch_max_tokens", "16384", \ "--chunked_prefill_size", "1024", \ "--linear_att_ssm_data_type", "bfloat16", \ - "--linear_att_cache_size", "104", \ - "--graph_max_batch_size", "104", \ + "--linear_att_cache_size", "256", \ + "--graph_max_batch_size", "256", \ "--graph_split_batch_size", "8", \ "--graph_grow_step_size", "16", \ "--graph_max_len_in_batch", "2048", \ "--disable_flashinfer_allreduce", \ - "--enable_ep_moe", \ - "--enable_tpsp_mix_mode", \ - "--enable_prefill_microbatch_overlap", \ + "--enable_fused_shared_experts", \ + "--disable_dynamic_prompt_cache", \ "--disable_aggressive_schedule", \ "--router_max_wait_tokens", "64", \ "--max_image_pixels", "6272000", \ diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh index 2e0d7d32e0..5448832e06 100755 --- a/tools/run_glm53_h100_container.sh +++ b/tools/run_glm53_h100_container.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash set -euo pipefail -image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:vl-1m-c100-ep8}" -name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-vl-1m-c100-ep8}" +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:vl-1m-tp8-c256}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-vl-1m-tp8-c256}" model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" From 65256aab59e755630c0999f3de24ca137d26a29e Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 30 Aug 2026 16:42:02 +0800 Subject: [PATCH 23/28] docs: record final GLM-5.3 release results --- GLM53_H100_DEPLOY.md | 120 +++++++++++++++++-------------------------- 1 file changed, 46 insertions(+), 74 deletions(-) diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 60fab17761..44d428b080 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -2,39 +2,38 @@ This branch packages LightLLM text and vision inference for GLM-5.3-Flash on one eight-GPU H100 or H200 node. The release profile keeps the 1,048,576-token -request limit and multimodal workers while optimizing 100-way concurrency with -EP8, TP/SP mixing, prefill microbatch overlap, and CUDA graphs. +request limit and uses pure TP8, CUDA graphs through batch 256, a 16,384-token +batch budget, fused shared experts, and no dynamic prompt cache. ## Release image The image is published only to the requested private registry: ```text -registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.3.0-glm53-vl-1m-c100-ep8-cc6319a4350c70e062218d85a1cbf1e8d6890cd5 +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.3.0-glm53-vl-1m-tp8-c256-62b8a9da25f7519822657c8126017fbc2793a08a ``` Immutable manifest: ```text -registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:f456207d3869996c9cfaa17df73058296e8de3720ef3d7eda4d56abb4719ff14 +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:3d07b3e9964cae15001e8136f1d19bd0b655df58488546e32f1dd15ffc9dbab7 ``` -The corresponding local image is -`lightllm-glm53:vl-1m-c100-ep8` (image ID -`sha256:f5b1c4282af30f656275d9726637e8d86e7a2dd7b5a103069fb81c6bea92831d`). -The image embeds source revision `cc6319a4350c70e062218d85a1cbf1e8d6890cd5`, -the validated H100 autotune records, and the complete server command. +The corresponding local image is `lightllm-glm53:vl-1m-tp8-c256` (image ID +`sha256:34939d00288e2e52ccecd3a241185648e2929a2a6f48bec1173592d337969dff`). +The image embeds source revision +`62b8a9da25f7519822657c8126017fbc2793a08a` and the complete server command. ## Deploy on the H100 node No command override or autotune-config mount is required: ```bash -IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:f456207d3869996c9cfaa17df73058296e8de3720ef3d7eda4d56abb4719ff14" +IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:3d07b3e9964cae15001e8136f1d19bd0b655df58488546e32f1dd15ffc9dbab7" sudo docker pull "$IMAGE" sudo docker run -d \ - --name glm53-lightllm-vl-1m-c100-ep8 \ + --name glm53-lightllm-vl-1m-tp8-c256 \ --restart unless-stopped \ --network host \ --ipc host \ @@ -53,14 +52,14 @@ The endpoint is `http://127.0.0.1:8002/v1`. Startup takes several minutes on the tested host. Check it with: ```bash -sudo docker ps --filter name=glm53-lightllm-vl-1m-c100-ep8 +sudo docker ps --filter name=glm53-lightllm-vl-1m-tp8-c256 curl --fail --show-error http://127.0.0.1:8002/v1/models ``` ## Run the local image on H200 ```bash -LIGHTLLM_GLM53_IMAGE=lightllm-glm53:vl-1m-c100-ep8 \ +LIGHTLLM_GLM53_IMAGE=lightllm-glm53:vl-1m-tp8-c256 \ LIGHTLLM_GLM53_MODEL_DIR=/nvme/sufubao/models/GLM-5.3-Flash \ LIGHTLLM_GLM53_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-lightllm-h200 \ LIGHTLLM_GLM53_TRITON_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-triton-h200 \ @@ -68,77 +67,50 @@ LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-deep-gemm- tools/run_glm53_h100_container.sh ``` -## Same-host c100 comparison - -LightLLM, vLLM, and SGLang were run sequentially on the same otherwise-idle -8xH100 80 GB host on 2026-08-30. All used the same local FP8 checkpoint, TP8, -BF16 KV cache, declared 1,048,576-token context, and no speculative decoding. -The common SGLang `bench_serving` client used seed 42, temperature 0, -streaming, ignored EOS, infinite request rate, one excluded warmup request, -and random input/output lengths of 1--1,000 tokens. The fixed 1,000-request -sample contained 504,929 input and 494,908 generated tokens. Every engine -completed 1,000/1,000 requests. - -| Engine | Duration | Request/s | Input tok/s | Output tok/s | Total tok/s | LightLLM lead | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | -| **LightLLM** | **233.37 s** | **4.285** | **2,163.67** | **2,120.73** | **4,284.40** | -- | -| vLLM | 270.74 s | 3.694 | 1,864.99 | 1,827.97 | 3,692.96 | **+16.02%** | -| SGLang | 328.59 s | 3.043 | 1,536.66 | 1,506.16 | 3,042.82 | **+40.80%** | - -The compared server builds and material settings were: - -- LightLLM: this release profile; EP8, TP/SP mixing, prefill microbatch - overlap, 1,024-token chunks, graph batches through 104, and FlashInfer - all-reduce disabled. -- vLLM: `vllm/vllm-openai:glm53-flash-x86_64-cu130@sha256:2e771fa615452282cc331eb418b3ef21636fce355bea0491fca89e6d362ab703` - (`0.1.dev20051+g487ecf187`); TP8, 256 sequences, chunked prefill, prefix - caching, and `gpu-memory-utilization=0.90`. -- SGLang: `lmsysorg/sglang@sha256:0836f0160fa785e424e68d13ef88ddd548f87e6e11ad9f0e4de982e4f9188aaf` - (`0.0.0.dev1+gf609d677b`); TP8/EP8, DeepGEMM, - `mem-fraction-static=0.80`, and 1,024-token chunks. - -Cross-engine TTFT/ITL and client-retokenized text are intentionally omitted: -the engines expose reasoning text under different streaming fields. The API -usage token totals, successful-request count, wall time, and aggregate -throughput above are directly comparable. - -Result files: - -- LightLLM: `/nvme/sufubao/m39-home/results/glm53_h100_optimization/lightllm_ep8_tpsp_prefill_overlap_graph_c104_b4096_wait64_tuned_random_1k_c100_full1000.jsonl` -- vLLM: `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/vllm_h100_random_1k_c100.jsonl` -- SGLang: `/nvme/sufubao/m39-home/results/glm53_h100_engine_compare/sglang_h100_random_1k_c100.jsonl` - -The LightLLM run is archived by `exp` under -`~/experiments/runs/260830-025933-usr-bin-timeout-2400-docker-run-rm-network-host-`. +## Same-host concurrency comparison + +LightLLM, vLLM, and SGLang ran sequentially on the same otherwise-idle 8xH100 +80 GB node with the same local FP8 checkpoint, TP8, BF16 KV cache, declared +1,048,576-token context, no speculative decoding, and prompt/prefix caching +disabled. The unchanged SGLang `bench_serving` client used seed 42, +temperature 0, streaming, ignored EOS, infinite request rate, one excluded +warmup request, and random 1--1,000-token inputs and outputs. + +The request counts for c1/c8/c16/c64/c128/c256 were +10/80/160/640/1,000/1,000. Every request completed. Each engine cell is +`output tok/s / total tok/s`; the percentage is LightLLM's lead over the +faster competing engine. + +| Concurrency | LightLLM | SGLang | vLLM | Lead | +| ---: | ---: | ---: | ---: | ---: | +| 1 | **106.03 / 200.86** | 45.75 / 86.66 | 72.91 / 138.11 | **+45.44%** | +| 8 | **556.56 / 1,097.02** | 300.77 / 592.83 | 483.78 / 953.56 | **+15.04%** | +| 16 | **917.99 / 1,734.29** | 850.45 / 1,606.69 | 812.67 / 1,535.31 | **+7.94%** | +| 64 | **2,296.01 / 4,696.03** | 1,680.11 / 3,436.32 | 1,943.30 / 3,974.64 | **+18.15%** | +| 128 | **3,615.77 / 7,304.76** | 2,298.01 / 4,642.55 | 2,980.61 / 6,021.57 | **+21.31%** | +| 256 | **5,032.38 / 10,166.66** | 2,939.85 / 5,939.23 | 4,109.91 / 8,303.03 | **+22.45%** | + +The lowest-margin c16 point was repeated at 1,757.52 total tok/s. Result files +are under +`/nvme/sufubao/m39-home/results/glm53_h100_optimization_round3/`; the original +SGLang and vLLM comparison is in +`/nvme/sufubao/m39-home/results/glm53_h100_engine_concurrency_sweep/`. ## Accuracy and capability checks -The final optimized execution path was evaluated before publication. Every -evaluation was recorded with `exp`. +Every evaluation was recorded with `exp`. | Check | Result | Experiment | | --- | --- | --- | -| GSM8K, fixed first 100, 5-shot, greedy | **99/100**, 100/100 completed, no 2,048-token truncation | `260830-030435-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-PYTHONPA` | -| MMMU, fixed 100, full multimodal, greedy | **63/100**, 100/100 completed; 34 answers reached the 2,048-token cap | `260830-030533-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-PYTHONPA` | +| GSM8K, fixed first 100, 5-shot, greedy | **99/100**, 100/100 completed, no 2,048-token truncation | `260830-162659-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-PYTHONPA` | +| MMMU, fixed 100, full multimodal, greedy | **64/100**, 100/100 completed; 36 answers reached the 2,048-token cap | `260830-163016-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-HF-DATAS` | | Exact 1M-context needle | API and tokenizer both counted 1,000,000 prompt tokens; recovered `ZEBRA-4821` | `260829-211036` | -| Text smoke | Arithmetic answer `579` | `260829-210839` | | Synthetic vision smoke | Correctly identified the red square | `260829-210908` | -The previous conservative profile scored 64/100 on the same MMMU slice; the -optimized profile's one-answer difference is within this 100-example sample, -while it reduced capped answers from 39 to 34. The multimodal and 1M-context -features remain enabled in the published command. - -References: [SGLang GLM-5 benchmark recipe](https://github.com/sgl-project/sglang/blob/main/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx), -[SGLang serving benchmark](https://github.com/sgl-project/sglang/blob/main/docs/cookbook/base/benchmarks/autoregressive_model_benchmark.mdx), -[vLLM GLM-5.3-Flash recipe](https://recipes.vllm.ai/zai-org/GLM-5.3-Flash), -and [vLLM GLM-5.3 support PR](https://github.com/vllm-project/vllm/pull/53906). - ## Rebuild note `docker/Dockerfile.glm53-h100` is the complete reproducible build. The release was produced with `docker/Dockerfile.glm53-h100-overlay` from the previous -immutable, flattened private image because Docker Hub base-metadata requests -timed out during release. The overlay replaces the complete LightLLM source, -checks the three new autotune records and import path, and pins both its source -revision and base digest in OCI labels. +immutable private release because its runtime layers were already available. +The overlay replaces the complete LightLLM source and pins the source revision +and base digest in OCI labels. From 1a0085f07651a07b662f735b4cb8d8d0405bc573 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 31 Aug 2026 08:39:51 +0800 Subject: [PATCH 24/28] perf: accelerate GLM-5.3 long-context index prefill --- docker/Dockerfile.glm53-h100 | 34 +- docker/Dockerfile.glm53-h100-overlay | 32 +- .../layer_infer/transformer_layer_infer.py | 357 +++++++++++++++--- .../triton_kernel/topk_index_to_mem_index.py | 24 +- .../layer_infer/transformer_layer_infer.py | 102 +++-- tools/run_glm53_h100_container.sh | 4 +- tools/run_glm53_long_prompt_bench.py | 67 ++++ .../test_topk_index_to_mem_index.py | 20 +- 8 files changed, 520 insertions(+), 120 deletions(-) create mode 100644 tools/run_glm53_long_prompt_bench.py diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index 7f38d4ab58..a127fcc38f 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -108,13 +108,13 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ org.opencontainers.image.source="${OCI_SOURCE}" \ org.opencontainers.image.version="${OCI_VERSION}" \ - org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 multimodal" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 512K text" \ org.opencontainers.image.base.name="${BASE_NAME}" \ org.opencontainers.image.base.digest="${BASE_DIGEST}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="multimodal-1m-tp8-c256-no-prompt-cache" \ + ai.lightllm.profile="text-512k-tp8-mtp4-c256-no-prompt-cache" \ ai.lightllm.security-profile="flattened-no-sglang-server-components" ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ @@ -156,10 +156,9 @@ ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ CMD /opt/sglang/bin/python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 -# LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is -# therefore request_limit + 36 so /v1/models advertises the full 1,048,576. -# This pure-TP8 profile is the exact full-multimodal configuration validated -# from 1 through 256 concurrent requests on an eight-H100 80 GB node. +# Exact 512K-input/1K-output text profile validated on an eight-H100 80 GB +# node. The request queue accepts 256 clients; one 512K request is GPU-resident +# at a time. MTP accelerates the 1K-token decode phase. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -168,29 +167,26 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--port", "8002", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ - "--max_total_token_num", "1048612", \ + "--max_total_token_num", "560000", \ "--running_max_req_size", "256", \ - "--max_req_total_len", "1048576", \ - "--batch_max_tokens", "16384", \ - "--chunked_prefill_size", "1024", \ + "--max_req_total_len", "532480", \ + "--batch_max_tokens", "32768", \ + "--chunked_prefill_size", "32768", \ "--linear_att_ssm_data_type", "bfloat16", \ "--linear_att_cache_size", "256", \ - "--graph_max_batch_size", "256", \ - "--graph_split_batch_size", "8", \ - "--graph_grow_step_size", "16", \ - "--graph_max_len_in_batch", "2048", \ + "--graph_max_batch_size", "1", \ + "--disable_cudagraph", \ "--disable_flashinfer_allreduce", \ "--enable_fused_shared_experts", \ "--disable_dynamic_prompt_cache", \ "--disable_aggressive_schedule", \ "--router_max_wait_tokens", "64", \ - "--max_image_pixels", "6272000", \ - "--max_image_token_count", "8000", \ - "--visual_tp", "1", \ - "--visual_dp", "8", \ - "--visual_infer_batch_size", "8", \ + "--disable_vision", \ "--cache_capacity", "64", \ "--schedule_time_interval", "0.001", \ "--prefill_coalesce_interval", "0.5", \ + "--mtp_mode", "eagle_with_att", \ + "--mtp_step", "4", \ + "--mtp_draft_model_dir", "/model", \ "--reasoning_parser", "glm45", \ "--tool_call_parser", "glm47"] diff --git a/docker/Dockerfile.glm53-h100-overlay b/docker/Dockerfile.glm53-h100-overlay index 0055926840..83528ee6fc 100644 --- a/docker/Dockerfile.glm53-h100-overlay +++ b/docker/Dockerfile.glm53-h100-overlay @@ -24,17 +24,18 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ org.opencontainers.image.source="${OCI_SOURCE}" \ org.opencontainers.image.version="${OCI_VERSION}" \ - org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 multimodal" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 512K text" \ org.opencontainers.image.base.name="${BASE_NAME}" \ org.opencontainers.image.base.digest="${BASE_DIGEST}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="multimodal-1m-tp8-c256-no-prompt-cache" \ + ai.lightllm.profile="text-512k-tp8-mtp4-c256-no-prompt-cache" \ ai.lightllm.security-profile="flattened-base-plus-source-overlay" -# LightLLM reserves 36 tokens in its HTTP admission guard. The token pool is -# therefore request_limit + 36 so /v1/models advertises the full 1,048,576. +# Exact 512K-input/1K-output text profile validated on an eight-H100 80 GB +# node. The request queue accepts 256 clients; one 512K request is GPU-resident +# at a time. MTP accelerates the 1K-token decode phase. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -43,29 +44,26 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--port", "8002", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ - "--max_total_token_num", "1048612", \ + "--max_total_token_num", "560000", \ "--running_max_req_size", "256", \ - "--max_req_total_len", "1048576", \ - "--batch_max_tokens", "16384", \ - "--chunked_prefill_size", "1024", \ + "--max_req_total_len", "532480", \ + "--batch_max_tokens", "32768", \ + "--chunked_prefill_size", "32768", \ "--linear_att_ssm_data_type", "bfloat16", \ "--linear_att_cache_size", "256", \ - "--graph_max_batch_size", "256", \ - "--graph_split_batch_size", "8", \ - "--graph_grow_step_size", "16", \ - "--graph_max_len_in_batch", "2048", \ + "--graph_max_batch_size", "1", \ + "--disable_cudagraph", \ "--disable_flashinfer_allreduce", \ "--enable_fused_shared_experts", \ "--disable_dynamic_prompt_cache", \ "--disable_aggressive_schedule", \ "--router_max_wait_tokens", "64", \ - "--max_image_pixels", "6272000", \ - "--max_image_token_count", "8000", \ - "--visual_tp", "1", \ - "--visual_dp", "8", \ - "--visual_infer_batch_size", "8", \ + "--disable_vision", \ "--cache_capacity", "64", \ "--schedule_time_interval", "0.001", \ "--prefill_coalesce_interval", "0.5", \ + "--mtp_mode", "eagle_with_att", \ + "--mtp_step", "4", \ + "--mtp_draft_model_dir", "/model", \ "--reasoning_parser", "glm45", \ "--tool_call_parser", "glm47"] diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index 20dcbaed38..b24e7d41bd 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,13 +1,21 @@ +import os + import torch from typing import Any from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo -from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer -from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight +from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import ( + Deepseek2TransformerLayerInfer, +) +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import ( + Deepseek3_2TransformerLayerWeight, +) from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant -from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import destindex_copy_indexer_ks +from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import ( + destindex_copy_indexer_ks, +) from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import ( extract_indexer_ks, extract_indexer_ks_dynamic, @@ -22,7 +30,9 @@ def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.indexer = NsaInfer( - layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_ + layer_idx=self.layer_num_, + network_config=self.network_config_, + tp_world_size=self.tp_world_size_, ) return @@ -49,8 +59,12 @@ def _get_qkv( q = layer_weight.q_b_proj_.mm(q) cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim) - q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim) - q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q = q.view( + -1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim + ) + q_nope, q_rope = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) rmsnorm_forward( cache_kv[:, :, : self.kv_lora_rank], weight=layer_weight.kv_a_layernorm_.weight, @@ -75,7 +89,9 @@ def _context_attention_kernel( out=None, ) -> torch.Tensor: # Model-specific q projection (uses layer weights) - q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_nope, q_rope = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) q_all = torch.cat([q_nope, q_rope], dim=-1) @@ -118,7 +134,10 @@ def _token_attention_kernel( out=None, ): # Model-specific q projection (uses layer weights) - q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] + q_nope, q_rope = ( + q[:, :, : -self.qk_rope_head_dim], + q[:, :, -self.qk_rope_head_dim :], + ) q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) # 计算 topk mem indices @@ -154,6 +173,12 @@ def _token_attention_kernel( class NsaInfer: + _MQA_LOGITS_BYTES_PER_ELEM = 4 + _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 + _MQA_LOGITS_FREE_MEM_FRACTION = 0.2 + _FLASHMLA_SPARSE_TOPK_ALIGNMENT = 128 + _mqa_logits_budget_bytes = {} + def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): super().__init__() self.layer_idx_ = layer_idx @@ -167,9 +192,12 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.scale_fmt = network_config["quantization_config"]["scale_fmt"] self.softmax_scale = (self.index_head_dim) ** (-0.5) self.index_n_heads = network_config["index_n_heads"] - self.index_n_heads_scale = (self.index_n_heads ** -0.5) * self.softmax_scale + self.index_n_heads_scale = (self.index_n_heads**-0.5) * self.softmax_scale self.tp_world_size_ = tp_world_size self.tp_index_n_heads = self.index_n_heads // self.tp_world_size_ + self.index_kpool = network_config.get("index_kpool", 1) + self.index_kpool_compress = network_config.get("index_kpool_compress", False) + self._kpool_indexer_k_buffer = None # Most NSA models only instantiate the target model, so their decode # layout follows the process-wide MTP setting. A model that reuses an # NSA layer as a recurrent drafter can override this with its own @@ -184,8 +212,12 @@ def _get_indices( att_state: Any, layer_weight: Deepseek3_2TransformerLayerWeight, ): - - q, k = self._get_q_k_bf16(hidden_states, q_lora, infer_state, layer_weight) + q_k = self._get_q_k_bf16(hidden_states, q_lora, infer_state, layer_weight) + if len(q_k) == 2: + q, k = q_k + raw_k = None + else: + q, k, raw_k = q_k if self.tp_world_size_ > 1: q_merge = torch.empty( @@ -193,9 +225,16 @@ def _get_indices( dtype=q.dtype, device=q.device, ) - all_gather_into_tensor(output_=q_merge, input_=q.view(-1), group=infer_state.dist_group, async_op=False) + all_gather_into_tensor( + output_=q_merge, + input_=q.view(-1), + group=infer_state.dist_group, + async_op=False, + ) q = ( - q_merge.view(self.tp_world_size_, q.shape[0], self.tp_index_n_heads, q.shape[2]) + q_merge.view( + self.tp_world_size_, q.shape[0], self.tp_index_n_heads, q.shape[2] + ) .transpose(0, 1) .contiguous() .view(q.shape[0], self.index_n_heads, q.shape[2]) @@ -212,22 +251,62 @@ def _get_indices( O_buffer=indexer_k_buffer, ) - weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale + weights = ( + layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale + ) weights = weights.unsqueeze(-1) * q_scale ks = att_state.ks ke = att_state.ke lengths = att_state.lengths - if infer_state.is_prefill: + use_kpool = ( + infer_state.is_prefill + and self.index_kpool > 1 + and self.index_kpool_compress + and raw_k is not None + and infer_state.b_seq_len.shape[0] == 1 + and infer_state.mem_index.shape[0] == q_fp8.shape[0] + ) + if use_kpool: + ( + k_fp8_, + k_scale_, + score_ks, + score_ke, + score_lengths, + ) = self._prepare_kpool_scoring( + raw_k=raw_k, + hidden_states=hidden_states, + q_lora=q_lora, + infer_state=infer_state, + layer_weight=layer_weight, + indexer_k_buffer=indexer_k_buffer, + lengths=lengths, + ) + use_kpool = k_fp8_ is not None + + if use_kpool: + mtp_step = 0 + elif infer_state.is_prefill: mtp_step = 0 else: - mtp_step = get_env_start_args().mtp_step if self.decode_mtp_step is None else self.decode_mtp_step + mtp_step = ( + get_env_start_args().mtp_step + if self.decode_mtp_step is None + else self.decode_mtp_step + ) # LightSpec compacts each request to a variable number of contiguous # verify rows. Its sparse-index K packing must follow request boundaries # instead of assuming the fixed process-wide MTP width. - use_dynamic_layout = not infer_state.is_prefill and mtp_step > 0 and get_env_start_args().mtp_dynamic_verify - if use_dynamic_layout: + use_dynamic_layout = ( + not infer_state.is_prefill + and mtp_step > 0 + and get_env_start_args().mtp_dynamic_verify + ) + if use_kpool: + pass + elif use_dynamic_layout: k_fp8_, k_scale_ = extract_indexer_ks_dynamic( I_buffer=indexer_k_buffer, b_seq_len=infer_state.b_seq_len, @@ -243,36 +322,93 @@ def _get_indices( b_seq_len=infer_state.b_seq_len, b_req_idx=infer_state.b_req_idx, req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, - out_token_num=infer_state.b_seq_len.shape[0] * infer_state.max_kv_seq_len, + out_token_num=infer_state.b_seq_len.shape[0] + * infer_state.max_kv_seq_len, max_kv_seq_len=infer_state.max_kv_seq_len, mtp_step=mtp_step, ) + if not use_kpool: + score_ks, score_ke, score_lengths = ks, ke, lengths import deep_gemm - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - (k_fp8_, k_scale_), - weights.squeeze(-1), - ks, - ke, - clean_logits=False, - max_seqlen_k=infer_state.max_kv_seq_len, - ) - from sgl_kernel import fast_topk_v2 - b_topk_index = fast_topk_v2( - score=logits, - lengths=lengths, - topk=self.index_topk, + weights = weights.squeeze(-1) + query_token_num = q_fp8.shape[0] + kv_token_num = k_fp8_.shape[0] + query_chunk_size = self._get_mqa_logits_chunk_size( + query_token_num=query_token_num, + kv_token_num=kv_token_num, + device=q_fp8.device, ) - # The long-prefill score matrix can be tens of GiB. fast_topk_v2 has - # already consumed it, so release its storage before materializing the - # 2048-wide global and memory-index tables below. - del logits + + if use_kpool: + output_topk = ( + ( + self.index_topk + + self.index_kpool + - 1 + + self._FLASHMLA_SPARSE_TOPK_ALIGNMENT + - 1 + ) + // self._FLASHMLA_SPARSE_TOPK_ALIGNMENT + * self._FLASHMLA_SPARSE_TOPK_ALIGNMENT + ) + # K-pool appends at most pool_size - 1 always-selected tail + # tokens. Allocate the FlashMLA-aligned result once and leave the + # alignment gap masked, avoiding two full-size top-k buffers. + b_topk_index = torch.full( + (query_token_num, output_topk), + -1, + dtype=torch.int32, + device=q_fp8.device, + ) + else: + b_topk_index = torch.empty( + (query_token_num, self.index_topk), + dtype=torch.int32, + device=q_fp8.device, + ) + for start in range(0, query_token_num, query_chunk_size): + end = min(start + query_chunk_size, query_token_num) + logits = deep_gemm.fp8_mqa_logits( + q_fp8[start:end], + (k_fp8_, k_scale_), + weights[start:end], + score_ks[start:end], + score_ke[start:end], + # fast top-k already masks columns past each row's valid + # length, while DeepGEMM's small-M path rejects clean_logits. + clean_logits=False, + max_seqlen_k=kv_token_num, + ) + if use_kpool: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + topk_from_pooled_history_logits, + ) + + topk_index = topk_from_pooled_history_logits( + logits=logits, + group_lengths=score_lengths[start:end], + pool_size=self.index_kpool, + topk=self.index_topk, + seq_lens=lengths[start:end], + ) + else: + topk_index = fast_topk_v2( + score=logits, + lengths=score_lengths[start:end], + topk=self.index_topk, + ) + b_topk_index[start:end, : topk_index.shape[1]].copy_(topk_index) + # The long-prefill score matrix can be tens of GiB. Release each + # chunk before computing the next one. + del logits, topk_index # 将 topk index 转化为 mem index - from ..triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index + from ..triton_kernel.topk_index_to_mem_index import ( + trans_topk_index_to_mem_index, + ) b_topk_mem_index = trans_topk_index_to_mem_index( topk_index=b_topk_index, @@ -282,14 +418,149 @@ def _get_indices( return b_topk_mem_index, b_topk_index + def _prepare_kpool_scoring( + self, + raw_k, + hidden_states, + q_lora, + infer_state, + layer_weight, + indexer_k_buffer, + lengths, + ): + pool_size = self.index_kpool + query_token_num = raw_k.shape[0] + total_seq_len = infer_state.max_kv_seq_len + prefix_len = total_seq_len - query_token_num + if prefix_len < 0 or prefix_len % pool_size != 0: + return None, None, None, None, None + + if self._kpool_indexer_k_buffer is None: + self._kpool_indexer_k_buffer = torch.empty_like(indexer_k_buffer) + + gate_score = layer_weight.index_kpool_compress_gate.mm( + hidden_states.to(q_lora.dtype) + ) + closed_pool_num = query_token_num // pool_size + if closed_pool_num: + closed_token_num = closed_pool_num * pool_size + slot_k = raw_k[:closed_token_num].view( + closed_pool_num, pool_size, self.index_head_dim + ) + slot_score = gate_score[:closed_token_num].view( + closed_pool_num, pool_size, self.index_head_dim + ) + write_locs = infer_state.mem_index[:closed_token_num].view( + closed_pool_num, pool_size + )[:, -1] + self._compress_kpool_keys( + slot_k=slot_k, + slot_score=slot_score, + write_locs=write_locs, + layer_weight=layer_weight, + ) + + pool_seq_len = total_seq_len // pool_size + if pool_seq_len == 0: + return None, None, None, None, None + req_idx = infer_state.b_req_idx[-1:].to(torch.int64) + token_positions = torch.arange( + pool_size - 1, + pool_seq_len * pool_size, + pool_size, + dtype=torch.int64, + device=raw_k.device, + ) + mem_indices = infer_state.req_manager.req_to_token_indexs[ + req_idx, token_positions + ].view(-1) + packed_k = self._kpool_indexer_k_buffer[mem_indices, 0] + k_fp8 = ( + packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) + ) + k_scale = ( + packed_k[:, self.index_head_dim : self.index_head_dim + 4] + .contiguous() + .view(torch.float32) + .view(-1) + ) + + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( + torch.int32 + ) + score_ks = torch.zeros_like(pool_lengths) + return k_fp8, k_scale, score_ks, pool_lengths, pool_lengths + + def _compress_kpool_keys( + self, + slot_k, + slot_score, + write_locs, + layer_weight, + ): + from types import SimpleNamespace + + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + kpool_softmax_rotate_write_cache, + ) + + compressed_k, compressed_scale = kpool_softmax_rotate_write_cache( + pool=SimpleNamespace(page_size=64, index_head_dim=self.index_head_dim), + buf=self._kpool_indexer_k_buffer, + slot_k=slot_k, + slot_score=slot_score, + ape=layer_weight.index_kpool_compress_ape.weight, + loc=write_locs.to(torch.int64), + round_scale=self.scale_fmt is not None, + return_compressed=True, + write_cache=False, + ) + destindex_copy_indexer_ks( + K_fp8=compressed_k, + K_scale=compressed_scale, + DestLoc=write_locs, + O_buffer=self._kpool_indexer_k_buffer, + ) + + @classmethod + def _get_mqa_logits_chunk_size( + cls, query_token_num: int, kv_token_num: int, device: torch.device + ) -> int: + score_element_num = query_token_num * kv_token_num + if score_element_num < cls._MQA_LOGITS_STATIC_SKIP_ELEMS: + return query_token_num + + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + logits_budget_bytes = cls._mqa_logits_budget_bytes.get(device_index) + if logits_budget_bytes is None: + free_memory_bytes, _ = torch.cuda.mem_get_info(device_index) + free_mem_fraction = float( + os.getenv( + "LIGHTLLM_MQA_LOGITS_FREE_MEM_FRACTION", + cls._MQA_LOGITS_FREE_MEM_FRACTION, + ) + ) + free_mem_fraction = min(max(free_mem_fraction, 0.01), 0.9) + logits_budget_bytes = max(1, int(free_memory_bytes * free_mem_fraction)) + cls._mqa_logits_budget_bytes[device_index] = logits_budget_bytes + + bytes_per_query = kv_token_num * cls._MQA_LOGITS_BYTES_PER_ELEM + return max(1, min(query_token_num, logits_budget_bytes // bytes_per_query)) + @staticmethod def _rotate_activation(x: torch.Tensor) -> torch.Tensor: assert x.dtype == torch.bfloat16 - from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform + from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import ( + hadamard_transform, + ) hidden_size = x.size(-1) - assert (hidden_size & (hidden_size - 1)) == 0, "Hidden size must be a power of 2 for Hadamard transform." - return hadamard_transform(x, scale=hidden_size ** -0.5) + assert ( + hidden_size & (hidden_size - 1) + ) == 0, "Hidden size must be a power of 2 for Hadamard transform." + return hadamard_transform(x, scale=hidden_size**-0.5) def _get_q_k_bf16( self, @@ -298,7 +569,9 @@ def _get_q_k_bf16( infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek3_2TransformerLayerWeight, ): - q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.tp_index_n_heads, self.index_head_dim) + q = layer_weight.wq_b_proj_.mm(q_lora).view( + -1, self.tp_index_n_heads, self.index_head_dim + ) k = layer_weight.wk_proj_.mm(hidden_states) k = layer_weight.k_norm_(k, eps=self.eps) diff --git a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py index 4b45bdb5b7..57054023b8 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py +++ b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py @@ -14,19 +14,29 @@ def _trans_topk_index_to_mem_index( topk_mem_index, topk_mem_index_stride_b, topk_mem_index_stride_k, + topk_width, BLOCK_DMODEL: tl.constexpr, ): cur_index = tl.program_id(0) offs_d = tl.arange(0, BLOCK_DMODEL) - topk_index_ptrs = topk_index + cur_index * topk_index_stride_b + offs_d * topk_index_stride_k - topk_indices = tl.load(topk_index_ptrs) + mask = offs_d < topk_width + topk_index_ptrs = ( + topk_index + cur_index * topk_index_stride_b + offs_d * topk_index_stride_k + ) + topk_indices = tl.load(topk_index_ptrs, mask=mask, other=-1) ragged_start = tl.load(ragged_start_index + cur_index) topk_indices = tl.where(topk_indices != -1, topk_indices + ragged_start, -1) - tl.store(topk_index_ptrs, topk_indices) + tl.store(topk_index_ptrs, topk_indices, mask=mask) dest_mem_index = ragged_mem_index + topk_indices - mem_index = tl.load(dest_mem_index, mask=topk_indices != -1, other=-1) - tl.store(topk_mem_index + cur_index * topk_mem_index_stride_b + offs_d * topk_mem_index_stride_k, mem_index) + mem_index = tl.load(dest_mem_index, mask=mask & (topk_indices != -1), other=-1) + tl.store( + topk_mem_index + + cur_index * topk_mem_index_stride_b + + offs_d * topk_mem_index_stride_k, + mem_index, + mask=mask, + ) @torch.no_grad() @@ -35,7 +45,6 @@ def trans_topk_index_to_mem_index( ragged_start_index: torch.Tensor, ragged_mem_index: torch.Tensor, ): - assert topk_index.shape[1] == 2048, f"Expected topk_index shape[1]=2048, got {topk_index.shape[1]}" assert ragged_start_index.shape == (topk_index.shape[0],) grid = (topk_index.shape[0],) @@ -51,7 +60,8 @@ def trans_topk_index_to_mem_index( topk_mem_index=topk_mem_index, topk_mem_index_stride_b=topk_mem_index.stride(0), topk_mem_index_stride_k=topk_mem_index.stride(1), - BLOCK_DMODEL=2048, + topk_width=topk_index.shape[1], + BLOCK_DMODEL=triton.next_power_of_2(topk_index.shape[1]), num_warps=8, ) return topk_mem_index diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 0d03f9c60c..71e3eef1a6 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -28,25 +28,32 @@ class Glm5NextNsaInfer(NsaInfer): """GLM indexer projection without rotary dimensions.""" def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): - q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.tp_index_n_heads, self.index_head_dim) + q = layer_weight.wq_b_proj_.mm(q_lora).view( + -1, self.tp_index_n_heads, self.index_head_dim + ) k = layer_weight.wk_proj_.mm(hidden_states.to(q_lora.dtype)) k = layer_weight.k_norm_(k, eps=self.eps) - return self._rotate_activation(q), self._rotate_activation(k) + return self._rotate_activation(q), self._rotate_activation(k), k def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): # GLM stores weights_proj in FP32, so its activation must match before # delegating to the shared NSA scoring and top-k implementation. - return super()._get_indices(hidden_states.float(), q_lora, infer_state, att_state, layer_weight) + return super()._get_indices( + hidden_states.float(), q_lora, infer_state, att_state, layer_weight + ) class Glm5NextTransformerLayerInfer(Deepseek3_2TransformerLayerInfer): def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.num_hidden_layers = network_config["num_hidden_layers"] - self.autotune_layer_num = network_config.get("autotune_layer_num", self.num_hidden_layers) + self.autotune_layer_num = network_config.get( + "autotune_layer_num", self.num_hidden_layers + ) self.is_mtp_layer = layer_num >= self.num_hidden_layers self.is_linear_attention_layer = ( - not self.is_mtp_layer and network_config["layer_types"][layer_num] == "linear_attention" + not self.is_mtp_layer + and network_config["layer_types"][layer_num] == "linear_attention" ) self.mhc_streams = network_config.get("hc_mult", 4) self.hc_eps = network_config.get("hc_eps", 1e-6) @@ -66,14 +73,18 @@ def __init__(self, layer_num, network_config): # GLM's recurrent EAGLE drafter processes one row per logical # request. Only target-model decode uses the widened verification # layout of mtp_step + 1 rows. - self.indexer.decode_mtp_step = 0 if self.is_mtp_layer else get_env_start_args().mtp_step + self.indexer.decode_mtp_step = ( + 0 if self.is_mtp_layer else get_env_start_args().mtp_step + ) def _ffn_tp(self, input, infer_state, layer_weight): """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" input = input.view(-1, self.embed_dim_) up_gate_out = layer_weight.gate_up_proj.mm(input) - ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) + ffn1_out = self.alloc_tensor( + (input.size(0), up_gate_out.size(1) // 2), input.dtype + ) silu_and_mul_fwd( up_gate_out, ffn1_out, @@ -96,10 +107,14 @@ def _get_qkv(self, input, infer_state, layer_weight): if infer_state.need_dp_prefill_balance: input = infer_state._all_to_all_unbalance_get(data=input) - q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split([self.q_lora_rank, self.kv_lora_rank], dim=-1) + q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split( + [self.q_lora_rank, self.kv_lora_rank], dim=-1 + ) q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_) infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q} - q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim) + q = layer_weight.q_b_proj_.mm(q).view( + -1, self.tp_q_head_num_, self.qk_nope_head_dim + ) cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank) rmsnorm_forward( cache_kv[:, :, : self.kv_lora_rank], @@ -117,7 +132,9 @@ def _get_o(self, input, infer_state, layer_weight): input = infer_state._all_to_all_balance_get(data=input) if input.shape[2] == self.kv_lora_rank: input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) - output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) + output = layer_weight.o_weight_.mm( + input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim) + ) all_reduce(output, group=infer_state.dist_group) return output @@ -160,7 +177,9 @@ def _kda_projections(self, input, infer_state, layer_weight): input = self._tpsp_allgather(input=input, infer_state=infer_state) projected = layer_weight.linear_qkvb_proj.mm(input) qkv_size = 3 * self.tp_linear_projection_size - mixed_qkv, raw_beta = projected.split([qkv_size, self.tp_linear_num_heads], dim=-1) + mixed_qkv, raw_beta = projected.split( + [qkv_size, self.tp_linear_num_heads], dim=-1 + ) fg_a = layer_weight.linear_fg_a_proj.mm(input) f_a, g_a = fg_a.split(self.linear_head_dim, dim=-1) raw_gate, norm_gate = layer_weight.project_kda_fg_b(f_a, g_a) @@ -184,8 +203,12 @@ def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): def context_attention_forward(self, input_embeddings, infer_state, layer_weight): if not self.is_linear_attention_layer: - return super().context_attention_forward(input_embeddings, infer_state, layer_weight) - mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) + return super().context_attention_forward( + input_embeddings, infer_state, layer_weight + ) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( + input_embeddings, infer_state, layer_weight + ) core_output = infer_state.prefill_att_state1.prefill_att( q=None, k=None, @@ -206,8 +229,12 @@ def context_attention_forward(self, input_embeddings, infer_state, layer_weight) def token_attention_forward(self, input_embeddings, infer_state, layer_weight): if not self.is_linear_attention_layer: - return super().token_attention_forward(input_embeddings, infer_state, layer_weight) - mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) + return super().token_attention_forward( + input_embeddings, infer_state, layer_weight + ) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( + input_embeddings, infer_state, layer_weight + ) core_output = infer_state.decode_att_state1.decode_att( q=None, k=None, @@ -245,26 +272,45 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): if self.layer_num_ == 0: streams = hc_expand(streams.view(-1, self.embed_dim_), self.mhc_streams) - layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) + layer_input, residual_mix, post_mix = self._hc_pre( + streams, layer_weight, "attn", layer_weight.att_norm_weight_ + ) if prefill: - layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) + layer_output = self.context_attention_forward( + layer_input, infer_state, layer_weight + ) else: - layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) - streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + layer_output = self.token_attention_forward( + layer_input, infer_state, layer_weight + ) + streams = hc_post( + layer_output, streams, residual_mix, post_mix, self.mhc_streams + ) - layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) + layer_input, residual_mix, post_mix = self._hc_pre( + streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_ + ) if infer_state.use_replicated_attention_ep: if self.is_moe: - local_input = self._tpsp_sp_split(input=layer_input, infer_state=infer_state) + local_input = self._tpsp_sp_split( + input=layer_input, infer_state=infer_state + ) local_output = self._ffn(local_input, infer_state, layer_weight) - layer_output = self._tpsp_allgather(input=local_output, infer_state=infer_state) + layer_output = self._tpsp_allgather( + input=local_output, infer_state=infer_state + ) else: layer_output = self._ffn_tp(layer_input, infer_state, layer_weight) all_reduce(layer_output, group=infer_state.dist_group) else: layer_output = self._ffn(layer_input, infer_state, layer_weight) - streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) - is_autotune_last_layer = Autotuner.is_autotune_warmup() and self.layer_num_ == self.autotune_layer_num - 1 + streams = hc_post( + layer_output, streams, residual_mix, post_mix, self.mhc_streams + ) + is_autotune_last_layer = ( + Autotuner.is_autotune_warmup() + and self.layer_num_ == self.autotune_layer_num - 1 + ) if self.layer_num_ == self.num_hidden_layers - 1 or is_autotune_last_layer: return hc_contract(streams, self.mhc_streams) return streams @@ -272,9 +318,13 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): def context_forward(self, input_embeddings, infer_state, layer_weight): if self.is_mtp_layer: return super().context_forward(input_embeddings, infer_state, layer_weight) - return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=True) + return self._forward_mhc( + input_embeddings, infer_state, layer_weight, prefill=True + ) def token_forward(self, input_embeddings, infer_state, layer_weight): if self.is_mtp_layer: return super().token_forward(input_embeddings, infer_state, layer_weight) - return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=False) + return self._forward_mhc( + input_embeddings, infer_state, layer_weight, prefill=False + ) diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh index 5448832e06..2455b41f5c 100755 --- a/tools/run_glm53_h100_container.sh +++ b/tools/run_glm53_h100_container.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash set -euo pipefail -image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:vl-1m-tp8-c256}" -name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-vl-1m-tp8-c256}" +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:text-512k-mtp4}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-text-512k-mtp4}" model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" diff --git a/tools/run_glm53_long_prompt_bench.py b/tools/run_glm53_long_prompt_bench.py new file mode 100644 index 0000000000..1408a47df8 --- /dev/null +++ b/tools/run_glm53_long_prompt_bench.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Run SGLang bench_serving with one cached long prompt reused per request.""" + +import copy +import os +import pickle +import runpy +import sys +import types +from pathlib import Path + +from sglang.benchmark.datasets.random import RandomDataset + + +def _install_prompt_cache() -> None: + cache_path = Path(os.environ["GLM53_LONG_PROMPT_CACHE"]) + original_load = RandomDataset.load + + def cached_load(self, tokenizer, model_id=None): + expected = { + "input_len": self.input_len, + "output_len": self.output_len, + "range_ratio": self.range_ratio, + } + if cache_path.exists(): + with cache_path.open("rb") as cache_file: + cached = pickle.load(cache_file) + if cached["metadata"] != expected: + raise ValueError( + f"Long-prompt cache metadata mismatch: {cached['metadata']} != {expected}" + ) + row = cached["row"] + print(f"Loaded long prompt from {cache_path}", flush=True) + else: + num_requests = self.num_requests + self.num_requests = 1 + try: + row = original_load(self, tokenizer, model_id)[0] + finally: + self.num_requests = num_requests + cache_path.parent.mkdir(parents=True, exist_ok=True) + with cache_path.open("wb") as cache_file: + pickle.dump({"metadata": expected, "row": row}, cache_file) + print(f"Cached long prompt at {cache_path}", flush=True) + + return [copy.copy(row) for _ in range(self.num_requests)] + + RandomDataset.load = cached_load + + +def main() -> None: + _install_prompt_cache() + + disaggregation_utils = types.ModuleType("sglang.srt.disaggregation.utils") + disaggregation_utils.FAKE_BOOTSTRAP_HOST = "fake" + sys.modules[disaggregation_utils.__name__] = disaggregation_utils + + network_utils = types.ModuleType("sglang.srt.utils.network") + network_utils.NetworkAddress = object + sys.modules[network_utils.__name__] = network_utils + + sys.argv = ["sglang.bench_serving", *sys.argv[1:]] + runpy.run_module("sglang.bench_serving", run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py index 981c1dbb66..60b9995d37 100644 --- a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py @@ -1,27 +1,33 @@ import torch import pytest -from lightllm.models.deepseek3_2.triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index +from lightllm.models.deepseek3_2.triton_kernel.topk_index_to_mem_index import ( + trans_topk_index_to_mem_index, +) -def test_trans_topk_index_to_mem_index(): +@pytest.mark.parametrize("topk", [2048, 2176]) +def test_trans_topk_index_to_mem_index(topk): """Test trans_topk_index_to_mem_index converts topk indices to memory indices correctly.""" batch_size = 1 - topk = 2048 # Create topk_index tensor with some valid indices and some -1 (padding) topk_index = torch.zeros((batch_size, topk), dtype=torch.int32, device="cuda") - topk_index[:, 0:2047] = torch.arange(0, 2047, dtype=torch.int32, device="cuda") + topk_index[:, 0 : topk - 1] = torch.arange( + 0, topk - 1, dtype=torch.int32, device="cuda" + ) topk_index[:, -1] = -1 ragged_start_index = torch.tensor([2], dtype=torch.int32, device="cuda") # Create ragged_mem_index lookup table - ragged_mem_index = torch.arange(0, 2050, dtype=torch.int32, device="cuda") + 10 + ragged_mem_index = torch.arange(0, topk + 2, dtype=torch.int32, device="cuda") + 10 - topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_start_index, ragged_mem_index) + topk_mem_index = trans_topk_index_to_mem_index( + topk_index, ragged_start_index, ragged_mem_index + ) expected_index = torch.cat( ( - torch.arange(2, 2049, dtype=torch.int32, device="cuda"), + torch.arange(2, topk + 1, dtype=torch.int32, device="cuda"), torch.tensor([-1], dtype=torch.int32, device="cuda"), ) ).view(1, -1) From c3f39a8276e9caa6d55f862e564b4034c390a7af Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 1 Sep 2026 02:17:26 +0800 Subject: [PATCH 25/28] perf(glm5): accelerate K-pool serving --- docker/Dockerfile.glm53-h100 | 24 +- docker/Dockerfile.glm53-h100-overlay | 27 +-- lightllm/common/basemodel/basemodel.py | 31 +++ lightllm/common/basemodel/batch_objs.py | 17 ++ lightllm/common/basemodel/infer_struct.py | 2 + .../layer_weights/meta_weights/norm_weight.py | 21 +- .../triton_kernel/norm/gated_rmsnorm.py | 8 +- .../layer_infer/transformer_layer_infer.py | 227 +++++++++++++++--- .../triton_kernel/hadamard_transform.py | 69 ++++++ .../triton_kernel/indexer_weight_scale.py | 39 +++ .../layer_infer/transformer_layer_infer.py | 41 +++- .../layer_weights/transformer_layer_weight.py | 74 ++++-- .../mode_backend/generic_pre_process.py | 9 + test/kernel/test_glm5_kda_fusions.py | 118 +++++++++ tools/run_glm53_h100_container.sh | 6 +- 15 files changed, 622 insertions(+), 91 deletions(-) create mode 100644 lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py create mode 100644 test/kernel/test_glm5_kda_fusions.py diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 index a127fcc38f..902897ed3b 100644 --- a/docker/Dockerfile.glm53-h100 +++ b/docker/Dockerfile.glm53-h100 @@ -108,13 +108,13 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ org.opencontainers.image.source="${OCI_SOURCE}" \ org.opencontainers.image.version="${OCI_VERSION}" \ - org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 512K text" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 optimized text" \ org.opencontainers.image.base.name="${BASE_NAME}" \ org.opencontainers.image.base.digest="${BASE_DIGEST}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="text-512k-tp8-mtp4-c256-no-prompt-cache" \ + ai.lightllm.profile="text-16k256-tp8-kpool-c256-no-prompt-cache" \ ai.lightllm.security-profile="flattened-no-sglang-server-components" ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ @@ -137,7 +137,6 @@ ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sb LIGHTLLM_VOCAB_PARALLEL_GREEDY=1 \ LIGHTLLM_SYMM_MEM_OUT_OF_PLACE=1 \ LIGHTLLM_LOG_LEVEL=warning \ - LIGHTLLM_ENABLE_FAST_MTP_KDA=1 \ NCCL_CUMEM_ENABLE=1 \ NCCL_NVLS_ENABLE=1 \ CUDA_DEVICE_MAX_CONNECTIONS=8 \ @@ -156,9 +155,9 @@ ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ CMD /opt/sglang/bin/python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 -# Exact 512K-input/1K-output text profile validated on an eight-H100 80 GB -# node. The request queue accepts 256 clients; one 512K request is GPU-resident -# at a time. MTP accelerates the 1K-token decode phase. +# Exact 16K-input/256-output profile validated on an eight-H100 80 GB node. +# The larger KV capacity keeps a 64-request prompt wave resident, and decode +# CUDA graphs cover batches through 128. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -167,26 +166,23 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--port", "8002", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ - "--max_total_token_num", "560000", \ + "--max_total_token_num", "1120000", \ "--running_max_req_size", "256", \ "--max_req_total_len", "532480", \ "--batch_max_tokens", "32768", \ "--chunked_prefill_size", "32768", \ "--linear_att_ssm_data_type", "bfloat16", \ "--linear_att_cache_size", "256", \ - "--graph_max_batch_size", "1", \ - "--disable_cudagraph", \ + "--graph_max_batch_size", "128", \ + "--graph_split_batch_size", "8", \ + "--graph_grow_step_size", "8", \ + "--graph_max_len_in_batch", "32768", \ "--disable_flashinfer_allreduce", \ "--enable_fused_shared_experts", \ "--disable_dynamic_prompt_cache", \ - "--disable_aggressive_schedule", \ "--router_max_wait_tokens", "64", \ "--disable_vision", \ "--cache_capacity", "64", \ "--schedule_time_interval", "0.001", \ - "--prefill_coalesce_interval", "0.5", \ - "--mtp_mode", "eagle_with_att", \ - "--mtp_step", "4", \ - "--mtp_draft_model_dir", "/model", \ "--reasoning_parser", "glm45", \ "--tool_call_parser", "glm47"] diff --git a/docker/Dockerfile.glm53-h100-overlay b/docker/Dockerfile.glm53-h100-overlay index 83528ee6fc..72b46a925c 100644 --- a/docker/Dockerfile.glm53-h100-overlay +++ b/docker/Dockerfile.glm53-h100-overlay @@ -1,7 +1,7 @@ # Fast release path for environments where the original Docker Hub base is # unavailable. The pinned base is the already-flattened GLM-5.3 runtime; this # layer replaces only the public LightLLM source and the validated default CMD. -ARG BASE_IMAGE=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:2a664580a495215a5bfb48d96bf118a8321d7accde589e505de283d6ea5753b2 +ARG BASE_IMAGE=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:3801951c8697edcab68ac5d38efc1807fcda0cba3d7946125a04103c302a3845 FROM ${BASE_IMAGE} ARG OCI_CREATED @@ -9,7 +9,7 @@ ARG OCI_REVISION ARG OCI_SOURCE=https://github.com/sufubao/LightLLM ARG OCI_VERSION ARG BASE_NAME=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm -ARG BASE_DIGEST=sha256:2a664580a495215a5bfb48d96bf118a8321d7accde589e505de283d6ea5753b2 +ARG BASE_DIGEST=sha256:3801951c8697edcab68ac5d38efc1807fcda0cba3d7946125a04103c302a3845 WORKDIR /opt/lightllm @@ -24,18 +24,18 @@ LABEL org.opencontainers.image.created="${OCI_CREATED}" \ org.opencontainers.image.revision="${OCI_REVISION}" \ org.opencontainers.image.source="${OCI_SOURCE}" \ org.opencontainers.image.version="${OCI_VERSION}" \ - org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 512K text" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 optimized text" \ org.opencontainers.image.base.name="${BASE_NAME}" \ org.opencontainers.image.base.digest="${BASE_DIGEST}" \ ai.lightllm.model="GLM-5.3-Flash" \ ai.lightllm.accelerator="NVIDIA H100/H200" \ ai.lightllm.tensor-parallel-size="8" \ - ai.lightllm.profile="text-512k-tp8-mtp4-c256-no-prompt-cache" \ + ai.lightllm.profile="text-16k256-tp8-kpool-c256-no-prompt-cache" \ ai.lightllm.security-profile="flattened-base-plus-source-overlay" -# Exact 512K-input/1K-output text profile validated on an eight-H100 80 GB -# node. The request queue accepts 256 clients; one 512K request is GPU-resident -# at a time. MTP accelerates the 1K-token decode phase. +# Exact 16K-input/256-output profile validated on an eight-H100 80 GB node. +# The larger KV capacity keeps a 64-request prompt wave resident, and decode +# CUDA graphs cover batches through 128. CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--model_dir", "/model", \ "--model_name", "glm-5.3-flash", \ @@ -44,26 +44,23 @@ CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ "--port", "8002", \ "--httpserver_workers", "16", \ "--mem_fraction", ".90", \ - "--max_total_token_num", "560000", \ + "--max_total_token_num", "1120000", \ "--running_max_req_size", "256", \ "--max_req_total_len", "532480", \ "--batch_max_tokens", "32768", \ "--chunked_prefill_size", "32768", \ "--linear_att_ssm_data_type", "bfloat16", \ "--linear_att_cache_size", "256", \ - "--graph_max_batch_size", "1", \ - "--disable_cudagraph", \ + "--graph_max_batch_size", "128", \ + "--graph_split_batch_size", "8", \ + "--graph_grow_step_size", "8", \ + "--graph_max_len_in_batch", "32768", \ "--disable_flashinfer_allreduce", \ "--enable_fused_shared_experts", \ "--disable_dynamic_prompt_cache", \ - "--disable_aggressive_schedule", \ "--router_max_wait_tokens", "64", \ "--disable_vision", \ "--cache_capacity", "64", \ "--schedule_time_interval", "0.001", \ - "--prefill_coalesce_interval", "0.5", \ - "--mtp_mode", "eagle_with_att", \ - "--mtp_step", "4", \ - "--mtp_draft_model_dir", "/model", \ "--reasoning_parser", "glm45", \ "--tool_call_parser", "glm47"] diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 1d6376d4d9..f28362f66c 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -396,6 +396,32 @@ def _init_hidden_collector(self): @torch.no_grad() def forward(self, model_input: ModelInput): + if model_input.is_prefill: + pool_size = int(self.config.get("index_kpool", 1) or 1) + if pool_size > 1: + # These tensors are normally still on CPU here. Computing the + # batch-wide predicate once avoids synchronizing every NSA + # layer merely to decide whether its batched K-pool path is + # safe. Requiring both boundaries to align also makes every + # query chunk an integer number of pools. + model_input.kpool_prefill_aligned = bool( + torch.all( + (model_input.b_ready_cache_len.remainder(pool_size) == 0) + & (model_input.b_seq_len.remainder(pool_size) == 0) + & (model_input.b_input_len.remainder(pool_size) == 0) + ).item() + ) + else: + pool_size = int(self.config.get("index_kpool", 1) or 1) + model_input.kpool_decode_aligned = bool( + pool_size > 1 + and os.getenv("LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0").upper() + in {"1", "ON", "TRUE"} + and self.args.mtp_mode is None + and self.args.disable_dynamic_prompt_cache + and self.args.chunked_prefill_size % pool_size == 0 + and torch.all(model_input.b_input_len.remainder(pool_size) == 0).item() + ) model_input.to_cuda() assert model_input.mem_indexes.is_cuda @@ -426,6 +452,8 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.max_q_seq_len = model_input.max_q_seq_len infer_state.max_kv_seq_len = model_input.max_kv_seq_len infer_state.max_cache_len = model_input.max_cache_len + infer_state.kpool_prefill_aligned = model_input.kpool_prefill_aligned + infer_state.kpool_decode_aligned = model_input.kpool_decode_aligned assert model_input.b_req_idx.shape[0] == model_input.b_seq_len.shape[0] infer_state.b_req_idx = model_input.b_req_idx infer_state.b_seq_len = model_input.b_seq_len @@ -486,6 +514,9 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s new_model_input.b_mtp_index, (0, padded_batch_size), mode="constant", value=0 ) new_model_input.b_seq_len = F.pad(new_model_input.b_seq_len, (0, padded_batch_size), mode="constant", value=2) + new_model_input.b_input_len = F.pad( + new_model_input.b_input_len, (0, padded_batch_size), mode="constant", value=2 + ) if new_model_input.b_position_delta is not None: new_model_input.b_position_delta = F.pad( new_model_input.b_position_delta, (0, padded_batch_size), mode="constant", value=0 diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index 0d06514f5d..cd0b4e75dd 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -20,6 +20,9 @@ class ModelInput: b_req_idx: torch.Tensor = None b_mtp_index: torch.Tensor = None b_seq_len: torch.Tensor = None + # Original prompt length per request/row. This remains on CPU and is used + # only for one-time model-specific fast-path eligibility decisions. + b_input_len: torch.Tensor = None # 在 prefill 阶段，用于在 enable_prefill_decode_mixed 开启下， # 用于标识请求是否为 decode 请求混合在 prefill 请求中。 # 其对应的 input_ids 需要特殊处理, 从 req_to_next_token_ids 中获取。 @@ -55,6 +58,13 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None + # Whether every request in this prefill batch starts and ends on an + # index-kpool boundary. BaseModel computes this once while the length + # tensors are still on CPU; NSA layers reuse it without a per-layer GPU + # synchronization. + kpool_prefill_aligned: bool = False + kpool_decode_aligned: bool = False + # The router enables sparse vocabulary output only when target sampling is # exact, unmodified greedy. Draft models always consume greedy proposals. use_vocab_parallel_greedy: bool = False @@ -89,6 +99,11 @@ def __post_init__(self): self.check_input() def check_input(self): + if self.b_input_len is None: + # Internal warmup/padding inputs predate this metadata. Their + # synthetic prompt occupies the whole sequence, so b_seq_len is + # the correct conservative default. + self.b_input_len = self.b_seq_len if self.input_ids is not None: assert ( self.input_ids.dtype == torch.int64 @@ -97,12 +112,14 @@ def check_input(self): assert self.b_req_idx is not None assert self.b_mtp_index is not None assert self.b_seq_len is not None + assert self.b_input_len is not None assert self.multimodal_params is not None assert self.mem_indexes is not None or self.mem_indexes_cpu is not None assert self.b_req_idx.shape == (self.batch_size,) assert self.b_mtp_index.shape == self.b_req_idx.shape assert self.b_seq_len.shape == self.b_req_idx.shape + assert self.b_input_len.shape == self.b_req_idx.shape assert len(self.multimodal_params) == self.batch_size if self.is_prefill: diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 3df658b071..f81275279a 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -45,6 +45,8 @@ def __init__(self): # max_cache_len 用于 prefill 阶段标识请求中最大 cache的kv 的长度 self.max_cache_len: int = None self.is_prefill: bool = None + self.kpool_prefill_aligned: bool = False + self.kpool_decode_aligned: bool = False self.mem_manager: MemoryManager = None self.req_manager: ReqManager = None diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py index ee9d1923c3..fe80974569 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py @@ -73,6 +73,17 @@ def __call__( class GatedRMSNormWeight(RMSNormWeight): + def __init__( + self, + dim: int, + weight_name: str, + data_type: torch.dtype, + activation: str = "silu", + ): + super().__init__(dim=dim, weight_name=weight_name, data_type=data_type) + assert activation in ("silu", "sigmoid") + self.activation = activation + def _triton_forward( self, input: torch.Tensor, @@ -86,7 +97,15 @@ def _triton_forward( ), f"input.ndim: {input.ndim} != 2 or weight.ndim: {self.weight.ndim} != 1" if out is None: out = alloc_func(input.shape, dtype=input.dtype, device=input.device) - return gated_rmsnorm_forward(x=input, weight=self.weight, bias=None, eps=eps, z=gate_value, out=out) + return gated_rmsnorm_forward( + x=input, + weight=self.weight, + bias=None, + eps=eps, + z=gate_value, + out=out, + activation=self.activation, + ) def _cuda_forward( self, diff --git a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py index b42d1eeaa8..59f5ab2434 100644 --- a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py +++ b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py @@ -26,6 +26,7 @@ def gated_rmsnorm_forward_kernel( BLOCK_N: tl.constexpr, HAS_BIAS: tl.constexpr, NORM_BEFORE_GATE: tl.constexpr, + SIGMOID_GATE: tl.constexpr, Z_HEADS: tl.constexpr, ): # Map the program id to the row of X and Y it should compute. @@ -46,7 +47,7 @@ def gated_rmsnorm_forward_kernel( x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) if not NORM_BEFORE_GATE: z = tl.load(Z + cols, mask=cols < N).to(tl.float32) - x *= z * tl.sigmoid(z) + x *= tl.sigmoid(z) if SIGMOID_GATE else z * tl.sigmoid(z) # RMS norm: compute variance directly without mean subtraction xbar = tl.where(cols < N, x, 0.0) var = tl.sum(xbar * xbar, axis=0) / N @@ -61,7 +62,7 @@ def gated_rmsnorm_forward_kernel( y = x_hat * w + b if HAS_BIAS else x_hat * w if NORM_BEFORE_GATE: z = tl.load(Z + cols, mask=mask).to(tl.float32) - y *= z * tl.sigmoid(z) + y *= tl.sigmoid(z) if SIGMOID_GATE else z * tl.sigmoid(z) # Write output tl.store(Y + cols, y, mask=mask) @@ -108,6 +109,7 @@ def gated_rmsnorm_forward( out: torch.Tensor = None, group_size: int = None, norm_before_gate: bool = True, + activation: str = "silu", run_config: dict = None, ): M, N = x.shape @@ -118,6 +120,7 @@ def gated_rmsnorm_forward( assert x.stride(-1) == 1 # z is required for gated_rmsnorm assert z is not None, "z cannot be None for gated_rmsnorm_forward" + assert activation in ("silu", "sigmoid"), f"unsupported gate activation: {activation}" # Accept GDN's strided 3D gate without materializing a flattened copy. assert z.ndim in (2, 3), f"z must be [M, N] or [tokens, heads, N], got shape={z.shape}" assert z.stride(-1) == 1 @@ -181,6 +184,7 @@ def gated_rmsnorm_forward( eps, BLOCK_N=BLOCK_N, NORM_BEFORE_GATE=norm_before_gate, + SIGMOID_GATE=activation == "sigmoid", Z_HEADS=z_heads, num_warps=num_warps, ) diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index b24e7d41bd..c67d9de612 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -197,7 +197,11 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.tp_index_n_heads = self.index_n_heads // self.tp_world_size_ self.index_kpool = network_config.get("index_kpool", 1) self.index_kpool_compress = network_config.get("index_kpool_compress", False) - self._kpool_indexer_k_buffer = None + self.enable_kpool_decode_fastpath = os.getenv( + "LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0" + ).upper() in {"1", "ON", "TRUE"} + self._kpool_tail_k = None + self._kpool_tail_score = None # Most NSA models only instantiate the target model, so their decode # layout follows the process-wide MTP setting. A model that reuses an # NSA layer as a recurrent drafter can override this with its own @@ -240,8 +244,8 @@ def _get_indices( .view(q.shape[0], self.index_n_heads, q.shape[2]) ) - q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) - k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) + q_fp8, q_scale = self._quantize_indexer_activation(q) + k_fp8, k_scale = self._quantize_indexer_activation(k) indexer_k_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_) destindex_copy_indexer_ks( @@ -251,24 +255,31 @@ def _get_indices( O_buffer=indexer_k_buffer, ) - weights = ( - layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale + weights = self._scale_indexer_weights( + layer_weight.weights_proj_.mm(hidden_states), q_scale ) - weights = weights.unsqueeze(-1) * q_scale ks = att_state.ks ke = att_state.ke lengths = att_state.lengths - use_kpool = ( + use_kpool_prefill = ( infer_state.is_prefill and self.index_kpool > 1 and self.index_kpool_compress and raw_k is not None - and infer_state.b_seq_len.shape[0] == 1 + and infer_state.kpool_prefill_aligned and infer_state.mem_index.shape[0] == q_fp8.shape[0] ) - if use_kpool: + use_kpool_decode = ( + not infer_state.is_prefill + and self.enable_kpool_decode_fastpath + and infer_state.kpool_decode_aligned + and raw_k is not None + and get_env_start_args().mtp_mode is None + ) + use_kpool = use_kpool_prefill or use_kpool_decode + if use_kpool_prefill: ( k_fp8_, k_scale_, @@ -276,6 +287,25 @@ def _get_indices( score_ke, score_lengths, ) = self._prepare_kpool_scoring( + raw_k=raw_k, + hidden_states=hidden_states, + q_lora=q_lora, + infer_state=infer_state, + layer_weight=layer_weight, + indexer_k_buffer=indexer_k_buffer, + ragged_mem_index=att_state.ragged_mem_index, + ks=ks, + lengths=lengths, + ) + use_kpool = k_fp8_ is not None + elif use_kpool_decode: + ( + k_fp8_, + k_scale_, + score_ks, + score_ke, + score_lengths, + ) = self._prepare_kpool_decode_scoring( raw_k=raw_k, hidden_states=hidden_states, q_lora=q_lora, @@ -334,7 +364,6 @@ def _get_indices( from sgl_kernel import fast_topk_v2 - weights = weights.squeeze(-1) query_token_num = q_fp8.shape[0] kv_token_num = k_fp8_.shape[0] query_chunk_size = self._get_mqa_logits_chunk_size( @@ -418,6 +447,16 @@ def _get_indices( return b_topk_mem_index, b_topk_index + def _quantize_indexer_activation(self, value: torch.Tensor): + return act_quant(value, self.block_size, self.scale_fmt) + + def _scale_indexer_weights( + self, weights: torch.Tensor, q_scale: torch.Tensor + ) -> torch.Tensor: + return ( + weights.mul(self.index_n_heads_scale).unsqueeze(-1).mul(q_scale) + ).squeeze(-1) + def _prepare_kpool_scoring( self, raw_k, @@ -426,22 +465,25 @@ def _prepare_kpool_scoring( infer_state, layer_weight, indexer_k_buffer, + ragged_mem_index, + ks, lengths, ): pool_size = self.index_kpool query_token_num = raw_k.shape[0] - total_seq_len = infer_state.max_kv_seq_len - prefix_len = total_seq_len - query_token_num - if prefix_len < 0 or prefix_len % pool_size != 0: + if not infer_state.kpool_prefill_aligned: + return None, None, None, None, None + # Without decode K-pool enabled the zero-prefix path is intentionally + # transient, so no pooled history exists for a later chunk. + if infer_state.max_cache_len > 0 and not self.enable_kpool_decode_fastpath: return None, None, None, None, None - - if self._kpool_indexer_k_buffer is None: - self._kpool_indexer_k_buffer = torch.empty_like(indexer_k_buffer) gate_score = layer_weight.index_kpool_compress_gate.mm( hidden_states.to(q_lora.dtype) ) closed_pool_num = query_token_num // pool_size + compressed_k = None + compressed_scale = None if closed_pool_num: closed_token_num = closed_pool_num * pool_size slot_k = raw_k[:closed_token_num].view( @@ -453,28 +495,44 @@ def _prepare_kpool_scoring( write_locs = infer_state.mem_index[:closed_token_num].view( closed_pool_num, pool_size )[:, -1] - self._compress_kpool_keys( + compressed_k, compressed_scale = self._compress_kpool_keys( slot_k=slot_k, slot_score=slot_score, write_locs=write_locs, layer_weight=layer_weight, + output_buffer=indexer_k_buffer, + persist=self.enable_kpool_decode_fastpath, ) - pool_seq_len = total_seq_len // pool_size - if pool_seq_len == 0: + # The common serving path has no prefix and each request fits in this + # aligned chunk. DeepGEMM can consume the freshly compressed pools + # directly; allocating and gathering a second max-token-sized cache + # would waste several GiB for GLM-5.3 TP8 and can make c64 OOM. + if infer_state.max_cache_len == 0: + if compressed_k is None: + return None, None, None, None, None + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( + torch.int32 + ) + score_ks = torch.div(ks, pool_size, rounding_mode="floor").to(torch.int32) + score_ke = score_ks + pool_lengths + return compressed_k, compressed_scale, score_ks, score_ke, pool_lengths + + pooled_token_num = infer_state.total_token_num // pool_size + if pooled_token_num == 0: return None, None, None, None, None - req_idx = infer_state.b_req_idx[-1:].to(torch.int64) - token_positions = torch.arange( + # Every request boundary is pool-aligned, so selecting each pool's + # final token from the ragged request-major layout preserves exactly + # the packed-K order expected by DeepGEMM's per-row ks/ke ranges. + pooled_ragged_positions = torch.arange( pool_size - 1, - pool_seq_len * pool_size, + infer_state.total_token_num, pool_size, dtype=torch.int64, device=raw_k.device, ) - mem_indices = infer_state.req_manager.req_to_token_indexs[ - req_idx, token_positions - ].view(-1) - packed_k = self._kpool_indexer_k_buffer[mem_indices, 0] + mem_indices = ragged_mem_index[pooled_ragged_positions].to(torch.int64) + packed_k = indexer_k_buffer[mem_indices, 0] k_fp8 = ( packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) ) @@ -488,8 +546,103 @@ def _prepare_kpool_scoring( pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( torch.int32 ) - score_ks = torch.zeros_like(pool_lengths) - return k_fp8, k_scale, score_ks, pool_lengths, pool_lengths + score_ks = torch.div(ks, pool_size, rounding_mode="floor").to(torch.int32) + score_ke = score_ks + pool_lengths + return k_fp8, k_scale, score_ks, score_ke, pool_lengths + + def _prepare_kpool_decode_scoring( + self, + raw_k, + hidden_states, + q_lora, + infer_state, + layer_weight, + indexer_k_buffer, + lengths, + ): + """Update one-token K-pool tails and gather pooled decode history. + + This fast path deliberately targets plain decode. MTP verify has + multiple speculative rows per request and needs acceptance-aware tail + rollback, so it remains on the regular indexer path. + """ + + pool_size = self.index_kpool + batch_size = raw_k.shape[0] + if batch_size == 0: + return None, None, None, None, None + + if self._kpool_tail_k is None: + tail_shape = ( + infer_state.req_manager.max_request_num + 1, + pool_size, + self.index_head_dim, + ) + self._kpool_tail_k = torch.empty( + tail_shape, dtype=torch.bfloat16, device=raw_k.device + ) + self._kpool_tail_score = torch.empty_like(self._kpool_tail_k) + + req_idx = infer_state.b_req_idx.to(torch.int64) + positions = infer_state.b_seq_len.to(torch.int64) - 1 + tail_slots = torch.remainder(positions, pool_size) + gate_score = layer_weight.index_kpool_compress_gate.mm( + hidden_states.to(q_lora.dtype) + ) + self._kpool_tail_k[req_idx, tail_slots] = raw_k + self._kpool_tail_score[req_idx, tail_slots] = gate_score + + # Compress every row's current tail. Only pool-end token locations are + # gathered as history, so writes at intermediate token locations are + # harmless and avoid a dynamic nonzero/host synchronization. + self._compress_kpool_keys( + slot_k=self._kpool_tail_k[req_idx], + slot_score=self._kpool_tail_score[req_idx], + write_locs=infer_state.mem_index, + layer_weight=layer_weight, + output_buffer=indexer_k_buffer, + persist=True, + ) + + max_pool_len = infer_state.max_kv_seq_len // pool_size + if max_pool_len == 0: + return None, None, None, None, None + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( + torch.int32 + ) + endpoint_positions = ( + torch.arange(max_pool_len, dtype=torch.int64, device=raw_k.device) + * pool_size + + pool_size + - 1 + ) + last_endpoint = torch.clamp( + pool_lengths.to(torch.int64) * pool_size - 1, min=0 + ) + safe_positions = torch.minimum( + endpoint_positions.unsqueeze(0), last_endpoint.unsqueeze(1) + ) + mem_indices = infer_state.req_manager.req_to_token_indexs[ + req_idx.unsqueeze(1), safe_positions + ].to(torch.int64) + packed_k = indexer_k_buffer[mem_indices.reshape(-1), 0] + k_fp8 = ( + packed_k[:, : self.index_head_dim] + .contiguous() + .view(torch.float8_e4m3fn) + ) + k_scale = ( + packed_k[:, self.index_head_dim : self.index_head_dim + 4] + .contiguous() + .view(torch.float32) + .view(-1) + ) + score_ks = ( + torch.arange(batch_size, dtype=torch.int32, device=raw_k.device) + * max_pool_len + ) + score_ke = score_ks + pool_lengths + return k_fp8, k_scale, score_ks, score_ke, pool_lengths def _compress_kpool_keys( self, @@ -497,6 +650,8 @@ def _compress_kpool_keys( slot_score, write_locs, layer_weight, + output_buffer, + persist, ): from types import SimpleNamespace @@ -506,7 +661,7 @@ def _compress_kpool_keys( compressed_k, compressed_scale = kpool_softmax_rotate_write_cache( pool=SimpleNamespace(page_size=64, index_head_dim=self.index_head_dim), - buf=self._kpool_indexer_k_buffer, + buf=output_buffer, slot_k=slot_k, slot_score=slot_score, ape=layer_weight.index_kpool_compress_ape.weight, @@ -515,12 +670,14 @@ def _compress_kpool_keys( return_compressed=True, write_cache=False, ) - destindex_copy_indexer_ks( - K_fp8=compressed_k, - K_scale=compressed_scale, - DestLoc=write_locs, - O_buffer=self._kpool_indexer_k_buffer, - ) + if persist: + destindex_copy_indexer_ks( + K_fp8=compressed_k, + K_scale=compressed_scale, + DestLoc=write_locs, + O_buffer=output_buffer, + ) + return compressed_k, compressed_scale @classmethod def _get_mqa_logits_chunk_size( diff --git a/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py b/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py index eabf703f56..07f1f16a59 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py +++ b/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py @@ -41,6 +41,41 @@ def _hadamard_transform_kernel( tl.store(Y + offsets, x * scale, mask=mask) +@triton.jit +def _hadamard_transform_quant_fp8_kernel( + X, + Y, + S, + n_rows, + scale: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + row_mask = rows < n_rows + cols = tl.arange(0, BLOCK_N) + offsets = rows[:, None] * BLOCK_N + cols[None, :] + x = tl.load(X + offsets, mask=row_mask[:, None], other=0.0).to(tl.float32) + + x = _butterfly_stage(x, 64, 1, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 32, 2, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 16, 4, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 8, 8, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 4, 16, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 2, 32, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 1, 64, BLOCK_R, BLOCK_N) + + # Match the unfused path's bf16 Hadamard output before FP8 quantization. + x = (x * scale).to(tl.bfloat16).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) + quant_scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0)))) + y = tl.minimum(tl.maximum(x / quant_scale[:, None], -448.0), 448.0) + + tl.store(Y + offsets, y, mask=row_mask[:, None]) + tl.store(S + rows, quant_scale, mask=row_mask) + + @functools.lru_cache(maxsize=None) def _target_programs(device_index: int) -> int: return torch.cuda.get_device_properties(device_index).multi_processor_count * 2 @@ -78,3 +113,37 @@ def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: assert x.size(-1) == 128, "DeepSeek-V3.2 Hadamard transform expects hidden size 128" return _hadamard_transform_triton(x, scale) + + +def hadamard_transform_quant_fp8( + x: torch.Tensor, scale: float = 1.0 +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse Hadamard-128 with the following ue8m0 FP8 quantization.""" + + assert x.is_cuda, "hadamard_transform_quant_fp8 only supports CUDA tensors" + assert x.dtype == torch.bfloat16, "Hadamard transform expects bfloat16 input" + assert x.size(-1) == 128, "Hadamard transform expects hidden size 128" + if not x.is_contiguous(): + x = x.contiguous() + + original_shape = x.shape + rows = x.numel() // 128 + output = torch.empty_like(x, dtype=torch.float8_e4m3fn) + output_scale = torch.empty( + (*original_shape[:-1], 1), dtype=torch.float32, device=x.device + ) + if rows == 0: + return output, output_scale + + block_r = 32 + _hadamard_transform_quant_fp8_kernel[(triton.cdiv(rows, block_r),)]( + x, + output, + output_scale, + rows, + scale, + BLOCK_R=block_r, + BLOCK_N=128, + num_warps=2, + ) + return output.view(original_shape), output_scale diff --git a/lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py b/lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py new file mode 100644 index 0000000000..391d826f21 --- /dev/null +++ b/lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py @@ -0,0 +1,39 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _indexer_weight_scale_kernel( + weights, + q_scale, + size, + scale: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < size + value = tl.load(weights + offsets, mask=mask).to(tl.float32) + query_scale = tl.load(q_scale + offsets, mask=mask).to(tl.float32) + tl.store(weights + offsets, value * scale * query_scale, mask=mask) + + +def scale_indexer_weights_(weights: torch.Tensor, q_scale: torch.Tensor, scale: float) -> torch.Tensor: + """Apply the indexer query/head scales in one in-place kernel.""" + + assert weights.dtype == torch.float32 + assert weights.is_contiguous() and q_scale.is_contiguous() + assert weights.numel() == q_scale.numel() + size = weights.numel() + if size == 0: + return weights + block_size = 256 + _indexer_weight_scale_kernel[(triton.cdiv(size, block_size),)]( + weights, + q_scale, + size, + scale, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return weights diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 71e3eef1a6..709f29843f 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -33,7 +33,26 @@ def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): ) k = layer_weight.wk_proj_.mm(hidden_states.to(q_lora.dtype)) k = layer_weight.k_norm_(k, eps=self.eps) - return self._rotate_activation(q), self._rotate_activation(k), k + return q, k, k + + def _quantize_indexer_activation(self, value: torch.Tensor): + from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import ( + hadamard_transform_quant_fp8, + ) + + assert self.block_size == 128 and self.scale_fmt == "ue8m0" + return hadamard_transform_quant_fp8( + value, scale=self.index_head_dim**-0.5 + ) + + def _scale_indexer_weights( + self, weights: torch.Tensor, q_scale: torch.Tensor + ) -> torch.Tensor: + from lightllm.models.deepseek3_2.triton_kernel.indexer_weight_scale import ( + scale_indexer_weights_, + ) + + return scale_indexer_weights_(weights, q_scale, self.index_n_heads_scale) def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): # GLM stores weights_proj in FP32, so its activation must match before @@ -175,26 +194,32 @@ def _kda_projections(self, input, infer_state, layer_weight): input = input.view(-1, self.embed_dim_) if not infer_state.use_replicated_attention_ep: input = self._tpsp_allgather(input=input, infer_state=infer_state) - projected = layer_weight.linear_qkvb_proj.mm(input) + projected = layer_weight.linear_qkvbfg_a_proj.mm(input) qkv_size = 3 * self.tp_linear_projection_size - mixed_qkv, raw_beta = projected.split( - [qkv_size, self.tp_linear_num_heads], dim=-1 + mixed_qkv, raw_beta, f_a, g_a = projected.split( + [ + qkv_size, + self.tp_linear_num_heads, + self.linear_head_dim, + self.linear_head_dim, + ], + dim=-1, ) - fg_a = layer_weight.linear_fg_a_proj.mm(input) - f_a, g_a = fg_a.split(self.linear_head_dim, dim=-1) raw_gate, norm_gate = layer_weight.project_kda_fg_b(f_a, g_a) return mixed_qkv, raw_gate, raw_beta, norm_gate def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): tokens = norm_gate.shape[0] core_output = core_output.view(-1, self.linear_head_dim) - norm_gate = norm_gate.contiguous().view(-1, self.linear_head_dim) + norm_gate = norm_gate.view( + tokens, self.tp_linear_num_heads, self.linear_head_dim + ) output = layer_weight.linear_o_norm( input=core_output, + gate_value=norm_gate, eps=self.eps_, alloc_func=self.alloc_tensor, ) - output.mul_(norm_gate.float().sigmoid().to(output.dtype)) output = layer_weight.linear_o_proj.mm(output.view(tokens, -1)) if infer_state.use_replicated_attention_ep: all_reduce(output, group=infer_state.dist_group) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index 4f1a6725dc..0f82f41c02 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -9,12 +9,20 @@ ) from lightllm.common.basemodel.layer_weights.meta_weights import ( COLMMWeight, + GatedRMSNormWeight, LayerNormWeight, ParameterWeight, RMSNormWeight, ROWMMWeight, TpParameterWeight, ) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import ( + get_row_slice_mixin, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import ( + MMWeightTpl, +) +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import ( Deepseek2TransformerLayerWeight, ) @@ -24,6 +32,50 @@ from .pre_and_post_layer_weight import add_language_model_aliases +class Glm5NextMergedKdaProjection(MMWeightTpl): + """One KDA input GEMM with TP-sharded q/k/v/b and replicated f_a/g_a.""" + + def __init__( + self, + in_dim: int, + projection: int, + head_count: int, + head_dim: int, + weight_names: list[str], + data_type: torch.dtype, + tp_rank: int | None = None, + tp_world_size: int | None = None, + ): + tp_rank = get_current_rank_in_dp() if tp_rank is None else tp_rank + tp_world_size = get_dp_world_size() if tp_world_size is None else tp_world_size + assert projection % tp_world_size == 0 + assert head_count % tp_world_size == 0 + super().__init__( + in_dim=in_dim, + out_dims=[ + projection // tp_world_size, + projection // tp_world_size, + projection // tp_world_size, + head_count // tp_world_size, + head_dim, + head_dim, + ], + weight_names=weight_names, + bias_names=None, + data_type=data_type, + quant_method=None, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + self.sharded_slicer = get_row_slice_mixin( + "none", tp_rank=tp_rank, tp_world_size=tp_world_size + ) + self.replicated_slicer = get_row_slice_mixin("none", tp_rank=0, tp_world_size=1) + + def _get_param_slicer(self, sub_child_index: int): + return self.replicated_slicer if sub_child_index >= 4 else self.sharded_slicer + + class Glm5NextTransformerLayerWeight(Deepseek3_2TransformerLayerWeight): def _parse_config(self): super()._parse_config() @@ -63,27 +115,20 @@ def _init_kda(self): head_count = self.linear_num_heads head_dim = self.linear_head_dim - self.linear_qkvb_proj = ROWMMWeight( + self.linear_qkvbfg_a_proj = Glm5NextMergedKdaProjection( in_dim=self.n_embed, - out_dims=[projection, projection, projection, head_count], + projection=projection, + head_count=head_count, + head_dim=head_dim, weight_names=[ f"{prefix}.q_proj.weight", f"{prefix}.k_proj.weight", f"{prefix}.v_proj.weight", f"{prefix}.b_proj.weight", + f"{prefix}.f_a_proj.weight", + f"{prefix}.g_a_proj.weight", ], data_type=self.data_type_, - quant_method=None, - ) - # f_a and g_a are replicated across TP ranks. - self.linear_fg_a_proj = ROWMMWeight( - in_dim=self.n_embed, - out_dims=[head_dim, head_dim], - weight_names=[f"{prefix}.f_a_proj.weight", f"{prefix}.g_a_proj.weight"], - data_type=self.data_type_, - quant_method=None, - tp_rank=0, - tp_world_size=1, ) self.linear_fg_b_proj = ROWMMWeight( in_dim=head_dim, @@ -113,10 +158,11 @@ def _init_kda(self): data_type=torch.float32, weight_shape=(projection,), ) - self.linear_o_norm = RMSNormWeight( + self.linear_o_norm = GatedRMSNormWeight( dim=head_dim, weight_name=f"{prefix}.o_norm.weight", data_type=self.data_type_, + activation="sigmoid", ) self.linear_o_proj = COLMMWeight( in_dim=projection, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 94633dfb55..124e3d6d90 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -16,6 +16,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> input_ids = [] b_req_idx = [] b_seq_len = [] + b_input_len = [] b_q_seq_len = [] batch_multimodal_params = [] b_ready_cache_len = [] @@ -41,6 +42,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> input_id = input_token_ids[req.cur_kv_len :] b_seq_len.append(seq_len) + b_input_len.append(req.shm_req.input_len) b_q_seq_len.append(input_token_len) input_ids.append(input_id) total_token_num += seq_len @@ -62,6 +64,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cpu") b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu") b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu") + b_input_len = torch.tensor(b_input_len, dtype=torch.int32, device="cpu") b_is_decode_req = torch.tensor(b_is_decode_req, dtype=torch.bool, device="cpu") b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu") b_ready_cache_len = torch.tensor(b_ready_cache_len, dtype=torch.int32, device="cpu") @@ -84,6 +87,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, b_seq_len=b_seq_len, + b_input_len=b_input_len, b_is_decode_req=b_is_decode_req, b_ready_cache_len=b_ready_cache_len, b_prefill_start_loc=b_prefill_start_loc, @@ -102,6 +106,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx = [] b_mtp_index = [] b_seq_len = [] + b_input_len = [] b_q_seq_len = [] multimodal_params = [] for req in req_objs: @@ -110,6 +115,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In seq_len = req.get_cur_total_len() assert req.cur_kv_len == seq_len - 1, f"{req.cur_kv_len} {seq_len}" b_seq_len.append(seq_len) + b_input_len.append(req.shm_req.input_len) b_q_seq_len.append(1) total_token_num += seq_len b_mtp_index.append(0) @@ -120,6 +126,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx.append(req.req_idx) seq_len += 1 b_seq_len.append(seq_len) + b_input_len.append(req.shm_req.input_len) total_token_num += seq_len b_mtp_index.append(step + 1) multimodal_params.append(req.multimodal_params) @@ -132,6 +139,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu") b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu") + b_input_len = torch.tensor(b_input_len, dtype=torch.int32, device="cpu") b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu") b_position_delta = build_b_position_delta(multimodal_params) @@ -159,6 +167,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, b_seq_len=b_seq_len, + b_input_len=b_input_len, b_position_delta=b_position_delta, b_shared_seq_len=b_shared_seq_len, b_shared_radix_node_id=b_shared_radix_node_id, diff --git a/test/kernel/test_glm5_kda_fusions.py b/test/kernel/test_glm5_kda_fusions.py new file mode 100644 index 0000000000..9042a41b92 --- /dev/null +++ b/test/kernel/test_glm5_kda_fusions.py @@ -0,0 +1,118 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import ( + gated_rmsnorm_forward, +) +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import ( + Glm5NextMergedKdaProjection, +) +from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant +from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import ( + hadamard_transform, + hadamard_transform_quant_fp8, +) +from lightllm.models.deepseek3_2.triton_kernel.indexer_weight_scale import ( + scale_indexer_weights_, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("activation", ["silu", "sigmoid"]) +def test_gated_rmsnorm_activations(activation): + torch.manual_seed(0) + x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(8, 4, 128, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(128, device="cuda", dtype=torch.bfloat16) + + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=gate, + activation=activation, + run_config={"BLOCK_N": 128, "num_warps": 4}, + ) + + x_float = x.float() + expected = x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + 1e-6) + expected *= weight.float() + gate_float = gate.view_as(x).float() + if activation == "silu": + expected *= gate_float * gate_float.sigmoid() + else: + expected *= gate_float.sigmoid() + + torch.testing.assert_close(actual, expected.to(actual.dtype), rtol=0, atol=0) + + +def test_merged_kda_projection_mixes_sharded_and_replicated_weights(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + projection, head_count, head_dim, in_dim = 8, 4, 2, 3 + names = [f"weight_{index}" for index in range(6)] + merged = Glm5NextMergedKdaProjection( + in_dim=in_dim, + projection=projection, + head_count=head_count, + head_dim=head_dim, + weight_names=names, + data_type=torch.float32, + tp_rank=1, + tp_world_size=2, + ) + full_weights = { + names[0]: torch.arange(0, 24, dtype=torch.float32).view(8, 3), + names[1]: torch.arange(24, 48, dtype=torch.float32).view(8, 3), + names[2]: torch.arange(48, 72, dtype=torch.float32).view(8, 3), + names[3]: torch.arange(72, 84, dtype=torch.float32).view(4, 3), + names[4]: torch.arange(84, 90, dtype=torch.float32).view(2, 3), + names[5]: torch.arange(90, 96, dtype=torch.float32).view(2, 3), + } + merged.load_hf_weights(full_weights) + + expected_weight = torch.cat( + [ + full_weights[names[0]][4:], + full_weights[names[1]][4:], + full_weights[names[2]][4:], + full_weights[names[3]][2:], + full_weights[names[4]], + full_weights[names[5]], + ] + ).cuda() + torch.testing.assert_close(merged.mm_param.weight, expected_weight) + + x = torch.arange(6, dtype=torch.float32, device="cuda").view(2, 3) + torch.testing.assert_close(merged.mm(x, use_custom_tensor_mananger=False), x @ expected_weight.T) + assert merged.verify_load() + + +@pytest.mark.parametrize("shape", [(17, 128), (3, 32, 128)]) +def test_fused_hadamard_fp8_quant_matches_two_kernel_chain(shape): + torch.manual_seed(1) + value = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + expected_value, expected_scale = act_quant( + hadamard_transform(value, scale=128**-0.5), + block_size=128, + scale_fmt="ue8m0", + ) + + actual_value, actual_scale = hadamard_transform_quant_fp8(value, scale=128**-0.5) + + assert torch.equal(actual_value, expected_value) + torch.testing.assert_close(actual_scale, expected_scale, rtol=0, atol=0) + + +def test_fused_indexer_weight_scale_matches_torch_chain(): + torch.manual_seed(2) + weights = torch.randn(257, 32, device="cuda", dtype=torch.float32) + q_scale = torch.rand(257, 32, 1, device="cuda", dtype=torch.float32) + scale = 128**-0.5 * 32**-0.5 + expected = (weights * scale).unsqueeze(-1).mul(q_scale).squeeze(-1) + + actual = scale_indexer_weights_(weights.clone(), q_scale, scale) + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh index 2455b41f5c..b563085fb6 100755 --- a/tools/run_glm53_h100_container.sh +++ b/tools/run_glm53_h100_container.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash set -euo pipefail -image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:text-512k-mtp4}" -name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-text-512k-mtp4}" +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:16k256-kpool}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-16k256-kpool}" model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" @@ -19,6 +19,8 @@ exec sudo docker run --rm --name "${name}" \ --gpus all \ --ipc host \ --network host \ + -e LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH=1 \ + -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ --ulimit memlock=-1:-1 \ --ulimit nofile=1048576:1048576 \ -v "${model_dir}:/model:ro" \ From 33132607ec76175d3a5b56b305bb03f8c971f875 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 1 Sep 2026 02:22:11 +0800 Subject: [PATCH 26/28] docs: publish GLM-5.3 optimized H100 profile --- GLM53_H100_DEPLOY.md | 129 +++++++++++++++++-------------------------- 1 file changed, 52 insertions(+), 77 deletions(-) diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 44d428b080..0ecc81e4df 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -1,65 +1,58 @@ -# GLM-5.3-Flash H100/H200 deployment and validation +# GLM-5.3-Flash H100/H200 部署与验证 -This branch packages LightLLM text and vision inference for GLM-5.3-Flash on -one eight-GPU H100 or H200 node. The release profile keeps the 1,048,576-token -request limit and uses pure TP8, CUDA graphs through batch 256, a 16,384-token -batch budget, fused shared experts, and no dynamic prompt cache. +该配置使用 TP8、BF16 KV cache、32K batch token、1.12M token cache 和 +CUDA Graph 128，面向 16K 输入、256 输出的高并发文本服务。 -## Release image +## 镜像 -The image is published only to the requested private registry: +私有仓库标签： ```text -registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.3.0-glm53-vl-1m-tp8-c256-62b8a9da25f7519822657c8126017fbc2793a08a +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.5.0-glm53-16k256-kpool-c3f39a82 ``` -Immutable manifest: +不可变镜像： ```text -registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:3d07b3e9964cae15001e8136f1d19bd0b655df58488546e32f1dd15ffc9dbab7 +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:76a968231ffb0a4c8996bc9a0fb981b88b3afbf25b7ad77a07fa07fccc04901d ``` -The corresponding local image is `lightllm-glm53:vl-1m-tp8-c256` (image ID -`sha256:34939d00288e2e52ccecd3a241185648e2929a2a6f48bec1173592d337969dff`). -The image embeds source revision -`62b8a9da25f7519822657c8126017fbc2793a08a` and the complete server command. +本机镜像为 `lightllm-glm53:16k256-kpool`，镜像 ID 为 +`sha256:d388173fbbc1fc2e8961785248331cecd08683af3c5f28ff67bc530d770d1409`。 +同一不可变镜像已拉取到 H100 节点。 -## Deploy on the H100 node +## H100 部署命令 -No command override or autotune-config mount is required: +镜像已内置服务参数；启动时显式开启 K-pool decode 快路径： ```bash -IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:3d07b3e9964cae15001e8136f1d19bd0b655df58488546e32f1dd15ffc9dbab7" +IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:76a968231ffb0a4c8996bc9a0fb981b88b3afbf25b7ad77a07fa07fccc04901d" sudo docker pull "$IMAGE" sudo docker run -d \ - --name glm53-lightllm-vl-1m-tp8-c256 \ + --name glm53-lightllm-16k256-kpool \ --restart unless-stopped \ --network host \ --ipc host \ - --security-opt label=disable \ --gpus all \ + -e LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH=1 \ + -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ --ulimit memlock=-1 \ --ulimit nofile=1048576:1048576 \ - -v /home/devsft/cache-glm53-triton:/root/.triton \ - -v /home/devsft/cache-glm53-deep-gemm:/root/.deep_gemm \ -v /home/devsft/models/GLM-5.3-Flash:/model:ro \ -v /home/devsft/cache-glm53-lightllm:/root/.cache \ + -v /home/devsft/cache-glm53-triton:/root/.triton \ + -v /home/devsft/cache-glm53-deep-gemm:/root/.deep_gemm \ "$IMAGE" ``` -The endpoint is `http://127.0.0.1:8002/v1`. Startup takes several minutes on -the tested host. Check it with: - -```bash -sudo docker ps --filter name=glm53-lightllm-vl-1m-tp8-c256 -curl --fail --show-error http://127.0.0.1:8002/v1/models -``` +接口为 `http://127.0.0.1:8002/v1`。该快路径要求请求 prompt 长度按 +16 token 对齐，并使用镜像内置的关闭 prompt cache、关闭 MTP 配置；其他负载可移除 +`LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH`，使用保守路径。 -## Run the local image on H200 +本机 H200 可直接运行： ```bash -LIGHTLLM_GLM53_IMAGE=lightllm-glm53:vl-1m-tp8-c256 \ LIGHTLLM_GLM53_MODEL_DIR=/nvme/sufubao/models/GLM-5.3-Flash \ LIGHTLLM_GLM53_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-lightllm-h200 \ LIGHTLLM_GLM53_TRITON_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-triton-h200 \ @@ -67,50 +60,32 @@ LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-deep-gemm- tools/run_glm53_h100_container.sh ``` -## Same-host concurrency comparison - -LightLLM, vLLM, and SGLang ran sequentially on the same otherwise-idle 8xH100 -80 GB node with the same local FP8 checkpoint, TP8, BF16 KV cache, declared -1,048,576-token context, no speculative decoding, and prompt/prefix caching -disabled. The unchanged SGLang `bench_serving` client used seed 42, -temperature 0, streaming, ignored EOS, infinite request rate, one excluded -warmup request, and random 1--1,000-token inputs and outputs. - -The request counts for c1/c8/c16/c64/c128/c256 were -10/80/160/640/1,000/1,000. Every request completed. Each engine cell is -`output tok/s / total tok/s`; the percentage is LightLLM's lead over the -faster competing engine. - -| Concurrency | LightLLM | SGLang | vLLM | Lead | -| ---: | ---: | ---: | ---: | ---: | -| 1 | **106.03 / 200.86** | 45.75 / 86.66 | 72.91 / 138.11 | **+45.44%** | -| 8 | **556.56 / 1,097.02** | 300.77 / 592.83 | 483.78 / 953.56 | **+15.04%** | -| 16 | **917.99 / 1,734.29** | 850.45 / 1,606.69 | 812.67 / 1,535.31 | **+7.94%** | -| 64 | **2,296.01 / 4,696.03** | 1,680.11 / 3,436.32 | 1,943.30 / 3,974.64 | **+18.15%** | -| 128 | **3,615.77 / 7,304.76** | 2,298.01 / 4,642.55 | 2,980.61 / 6,021.57 | **+21.31%** | -| 256 | **5,032.38 / 10,166.66** | 2,939.85 / 5,939.23 | 4,109.91 / 8,303.03 | **+22.45%** | - -The lowest-margin c16 point was repeated at 1,757.52 total tok/s. Result files -are under -`/nvme/sufubao/m39-home/results/glm53_h100_optimization_round3/`; the original -SGLang and vLLM comparison is in -`/nvme/sufubao/m39-home/results/glm53_h100_engine_concurrency_sweep/`. - -## Accuracy and capability checks - -Every evaluation was recorded with `exp`. - -| Check | Result | Experiment | -| --- | --- | --- | -| GSM8K, fixed first 100, 5-shot, greedy | **99/100**, 100/100 completed, no 2,048-token truncation | `260830-162659-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-PYTHONPA` | -| MMMU, fixed 100, full multimodal, greedy | **64/100**, 100/100 completed; 36 answers reached the 2,048-token cap | `260830-163016-bin-bash-lc-export-OPENAI-API-KEY-EMPTY-HF-DATAS` | -| Exact 1M-context needle | API and tokenizer both counted 1,000,000 prompt tokens; recovered `ZEBRA-4821` | `260829-211036` | -| Synthetic vision smoke | Correctly identified the red square | `260829-210908` | - -## Rebuild note - -`docker/Dockerfile.glm53-h100` is the complete reproducible build. The release -was produced with `docker/Dockerfile.glm53-h100-overlay` from the previous -immutable private release because its runtime layers were already available. -The overlay replaces the complete LightLLM source and pins the source revision -and base digest in OCI labels. +## 16K/256 性能 + +LightLLM 与 vLLM 在同一台空闲的 8×H100 80GB 节点顺序测试，使用同一 FP8 +checkpoint、TP8、BF16 KV cache、32K batch token、关闭 prompt/prefix cache、 +seed 42、temperature 0、无限请求速率和一次不计入结果的 warmup。每档请求数等于 +并发数，每个请求输入 16,384 token、输出 256 token。 + +| 并发 | LightLLM 总吞吐 tok/s | vLLM 总吞吐 tok/s | LightLLM / vLLM | +| ---: | ---: | ---: | ---: | +| 1 | 5,373.86 | 6,542.89 | 82.13% | +| 8 | 16,444.11 | 17,635.33 | 93.25% | +| 16 | 20,669.61 | 25,246.99 | 81.87% | +| 64 | 25,614.05 | 31,485.12 | 81.35% | +| 128 | 25,588.22 | 27,988.09 | 91.43% | +| 256 | 25,623.93 | 31,206.94 | 82.11% | + +结果文件位于 +`/nvme/sufubao/m39-home/results/glm53_h100_16k_256_kpool_decode_ab/` 和 +`/nvme/sufubao/m39-home/results/glm53_h100_16k_256_vllm_opt/`；所有启动、性能与 +精度实验均由 `exp` 归档到 `~/experiments/runs/`。 + +## 精度 + +| 验证 | 结果 | +| --- | --- | +| GSM8K 固定前 100 题、5-shot、greedy | **99/100**；100/100 完成；无截断 | +| 精确 16,384-token needle | **PASS**；API 与本地 tokenizer 均计数 16,384；找回 `ZEBRA-4821` | + +相关实验记录为 `260901-020939` 和 `260901-021013`。 From 996cef93143f43562e45cf376680dc9bd9163931 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 1 Sep 2026 02:25:33 +0800 Subject: [PATCH 27/28] style: format GLM-5.3 changes --- lightllm/common/basemodel/basemodel.py | 3 +- .../layer_infer/transformer_layer_infer.py | 173 +++++------------- .../triton_kernel/hadamard_transform.py | 8 +- .../triton_kernel/topk_index_to_mem_index.py | 8 +- .../layer_infer/transformer_layer_infer.py | 108 +++-------- .../layer_weights/transformer_layer_weight.py | 4 +- test/kernel/test_glm5_kda_fusions.py | 6 +- tools/run_glm53_long_prompt_bench.py | 4 +- .../test_topk_index_to_mem_index.py | 8 +- 9 files changed, 80 insertions(+), 242 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f28362f66c..46f49ebab5 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -415,8 +415,7 @@ def forward(self, model_input: ModelInput): pool_size = int(self.config.get("index_kpool", 1) or 1) model_input.kpool_decode_aligned = bool( pool_size > 1 - and os.getenv("LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0").upper() - in {"1", "ON", "TRUE"} + and os.getenv("LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0").upper() in {"1", "ON", "TRUE"} and self.args.mtp_mode is None and self.args.disable_dynamic_prompt_cache and self.args.chunked_prefill_size % pool_size == 0 diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index c67d9de612..a31eb3e407 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -59,12 +59,8 @@ def _get_qkv( q = layer_weight.q_b_proj_.mm(q) cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim) - q = q.view( - -1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim - ) - q_nope, q_rope = torch.split( - q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) + q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim) + q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) rmsnorm_forward( cache_kv[:, :, : self.kv_lora_rank], weight=layer_weight.kv_a_layernorm_.weight, @@ -89,9 +85,7 @@ def _context_attention_kernel( out=None, ) -> torch.Tensor: # Model-specific q projection (uses layer weights) - q_nope, q_rope = torch.split( - q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) + q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) q_all = torch.cat([q_nope, q_rope], dim=-1) @@ -192,14 +186,16 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.scale_fmt = network_config["quantization_config"]["scale_fmt"] self.softmax_scale = (self.index_head_dim) ** (-0.5) self.index_n_heads = network_config["index_n_heads"] - self.index_n_heads_scale = (self.index_n_heads**-0.5) * self.softmax_scale + self.index_n_heads_scale = (self.index_n_heads ** -0.5) * self.softmax_scale self.tp_world_size_ = tp_world_size self.tp_index_n_heads = self.index_n_heads // self.tp_world_size_ self.index_kpool = network_config.get("index_kpool", 1) self.index_kpool_compress = network_config.get("index_kpool_compress", False) - self.enable_kpool_decode_fastpath = os.getenv( - "LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0" - ).upper() in {"1", "ON", "TRUE"} + self.enable_kpool_decode_fastpath = os.getenv("LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0").upper() in { + "1", + "ON", + "TRUE", + } self._kpool_tail_k = None self._kpool_tail_score = None # Most NSA models only instantiate the target model, so their decode @@ -236,9 +232,7 @@ def _get_indices( async_op=False, ) q = ( - q_merge.view( - self.tp_world_size_, q.shape[0], self.tp_index_n_heads, q.shape[2] - ) + q_merge.view(self.tp_world_size_, q.shape[0], self.tp_index_n_heads, q.shape[2]) .transpose(0, 1) .contiguous() .view(q.shape[0], self.index_n_heads, q.shape[2]) @@ -255,9 +249,7 @@ def _get_indices( O_buffer=indexer_k_buffer, ) - weights = self._scale_indexer_weights( - layer_weight.weights_proj_.mm(hidden_states), q_scale - ) + weights = self._scale_indexer_weights(layer_weight.weights_proj_.mm(hidden_states), q_scale) ks = att_state.ks ke = att_state.ke @@ -280,13 +272,7 @@ def _get_indices( ) use_kpool = use_kpool_prefill or use_kpool_decode if use_kpool_prefill: - ( - k_fp8_, - k_scale_, - score_ks, - score_ke, - score_lengths, - ) = self._prepare_kpool_scoring( + (k_fp8_, k_scale_, score_ks, score_ke, score_lengths,) = self._prepare_kpool_scoring( raw_k=raw_k, hidden_states=hidden_states, q_lora=q_lora, @@ -299,13 +285,7 @@ def _get_indices( ) use_kpool = k_fp8_ is not None elif use_kpool_decode: - ( - k_fp8_, - k_scale_, - score_ks, - score_ke, - score_lengths, - ) = self._prepare_kpool_decode_scoring( + (k_fp8_, k_scale_, score_ks, score_ke, score_lengths,) = self._prepare_kpool_decode_scoring( raw_k=raw_k, hidden_states=hidden_states, q_lora=q_lora, @@ -321,19 +301,11 @@ def _get_indices( elif infer_state.is_prefill: mtp_step = 0 else: - mtp_step = ( - get_env_start_args().mtp_step - if self.decode_mtp_step is None - else self.decode_mtp_step - ) + mtp_step = get_env_start_args().mtp_step if self.decode_mtp_step is None else self.decode_mtp_step # LightSpec compacts each request to a variable number of contiguous # verify rows. Its sparse-index K packing must follow request boundaries # instead of assuming the fixed process-wide MTP width. - use_dynamic_layout = ( - not infer_state.is_prefill - and mtp_step > 0 - and get_env_start_args().mtp_dynamic_verify - ) + use_dynamic_layout = not infer_state.is_prefill and mtp_step > 0 and get_env_start_args().mtp_dynamic_verify if use_kpool: pass elif use_dynamic_layout: @@ -352,8 +324,7 @@ def _get_indices( b_seq_len=infer_state.b_seq_len, b_req_idx=infer_state.b_req_idx, req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, - out_token_num=infer_state.b_seq_len.shape[0] - * infer_state.max_kv_seq_len, + out_token_num=infer_state.b_seq_len.shape[0] * infer_state.max_kv_seq_len, max_kv_seq_len=infer_state.max_kv_seq_len, mtp_step=mtp_step, ) @@ -374,13 +345,7 @@ def _get_indices( if use_kpool: output_topk = ( - ( - self.index_topk - + self.index_kpool - - 1 - + self._FLASHMLA_SPARSE_TOPK_ALIGNMENT - - 1 - ) + (self.index_topk + self.index_kpool - 1 + self._FLASHMLA_SPARSE_TOPK_ALIGNMENT - 1) // self._FLASHMLA_SPARSE_TOPK_ALIGNMENT * self._FLASHMLA_SPARSE_TOPK_ALIGNMENT ) @@ -450,12 +415,8 @@ def _get_indices( def _quantize_indexer_activation(self, value: torch.Tensor): return act_quant(value, self.block_size, self.scale_fmt) - def _scale_indexer_weights( - self, weights: torch.Tensor, q_scale: torch.Tensor - ) -> torch.Tensor: - return ( - weights.mul(self.index_n_heads_scale).unsqueeze(-1).mul(q_scale) - ).squeeze(-1) + def _scale_indexer_weights(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: + return (weights.mul(self.index_n_heads_scale).unsqueeze(-1).mul(q_scale)).squeeze(-1) def _prepare_kpool_scoring( self, @@ -478,23 +439,15 @@ def _prepare_kpool_scoring( if infer_state.max_cache_len > 0 and not self.enable_kpool_decode_fastpath: return None, None, None, None, None - gate_score = layer_weight.index_kpool_compress_gate.mm( - hidden_states.to(q_lora.dtype) - ) + gate_score = layer_weight.index_kpool_compress_gate.mm(hidden_states.to(q_lora.dtype)) closed_pool_num = query_token_num // pool_size compressed_k = None compressed_scale = None if closed_pool_num: closed_token_num = closed_pool_num * pool_size - slot_k = raw_k[:closed_token_num].view( - closed_pool_num, pool_size, self.index_head_dim - ) - slot_score = gate_score[:closed_token_num].view( - closed_pool_num, pool_size, self.index_head_dim - ) - write_locs = infer_state.mem_index[:closed_token_num].view( - closed_pool_num, pool_size - )[:, -1] + slot_k = raw_k[:closed_token_num].view(closed_pool_num, pool_size, self.index_head_dim) + slot_score = gate_score[:closed_token_num].view(closed_pool_num, pool_size, self.index_head_dim) + write_locs = infer_state.mem_index[:closed_token_num].view(closed_pool_num, pool_size)[:, -1] compressed_k, compressed_scale = self._compress_kpool_keys( slot_k=slot_k, slot_score=slot_score, @@ -511,9 +464,7 @@ def _prepare_kpool_scoring( if infer_state.max_cache_len == 0: if compressed_k is None: return None, None, None, None, None - pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( - torch.int32 - ) + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to(torch.int32) score_ks = torch.div(ks, pool_size, rounding_mode="floor").to(torch.int32) score_ke = score_ks + pool_lengths return compressed_k, compressed_scale, score_ks, score_ke, pool_lengths @@ -533,19 +484,10 @@ def _prepare_kpool_scoring( ) mem_indices = ragged_mem_index[pooled_ragged_positions].to(torch.int64) packed_k = indexer_k_buffer[mem_indices, 0] - k_fp8 = ( - packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) - ) - k_scale = ( - packed_k[:, self.index_head_dim : self.index_head_dim + 4] - .contiguous() - .view(torch.float32) - .view(-1) - ) + k_fp8 = packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) + k_scale = packed_k[:, self.index_head_dim : self.index_head_dim + 4].contiguous().view(torch.float32).view(-1) - pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( - torch.int32 - ) + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to(torch.int32) score_ks = torch.div(ks, pool_size, rounding_mode="floor").to(torch.int32) score_ke = score_ks + pool_lengths return k_fp8, k_scale, score_ks, score_ke, pool_lengths @@ -578,17 +520,13 @@ def _prepare_kpool_decode_scoring( pool_size, self.index_head_dim, ) - self._kpool_tail_k = torch.empty( - tail_shape, dtype=torch.bfloat16, device=raw_k.device - ) + self._kpool_tail_k = torch.empty(tail_shape, dtype=torch.bfloat16, device=raw_k.device) self._kpool_tail_score = torch.empty_like(self._kpool_tail_k) req_idx = infer_state.b_req_idx.to(torch.int64) positions = infer_state.b_seq_len.to(torch.int64) - 1 tail_slots = torch.remainder(positions, pool_size) - gate_score = layer_weight.index_kpool_compress_gate.mm( - hidden_states.to(q_lora.dtype) - ) + gate_score = layer_weight.index_kpool_compress_gate.mm(hidden_states.to(q_lora.dtype)) self._kpool_tail_k[req_idx, tail_slots] = raw_k self._kpool_tail_score[req_idx, tail_slots] = gate_score @@ -607,40 +545,17 @@ def _prepare_kpool_decode_scoring( max_pool_len = infer_state.max_kv_seq_len // pool_size if max_pool_len == 0: return None, None, None, None, None - pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to( - torch.int32 - ) + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to(torch.int32) endpoint_positions = ( - torch.arange(max_pool_len, dtype=torch.int64, device=raw_k.device) - * pool_size - + pool_size - - 1 + torch.arange(max_pool_len, dtype=torch.int64, device=raw_k.device) * pool_size + pool_size - 1 ) - last_endpoint = torch.clamp( - pool_lengths.to(torch.int64) * pool_size - 1, min=0 - ) - safe_positions = torch.minimum( - endpoint_positions.unsqueeze(0), last_endpoint.unsqueeze(1) - ) - mem_indices = infer_state.req_manager.req_to_token_indexs[ - req_idx.unsqueeze(1), safe_positions - ].to(torch.int64) + last_endpoint = torch.clamp(pool_lengths.to(torch.int64) * pool_size - 1, min=0) + safe_positions = torch.minimum(endpoint_positions.unsqueeze(0), last_endpoint.unsqueeze(1)) + mem_indices = infer_state.req_manager.req_to_token_indexs[req_idx.unsqueeze(1), safe_positions].to(torch.int64) packed_k = indexer_k_buffer[mem_indices.reshape(-1), 0] - k_fp8 = ( - packed_k[:, : self.index_head_dim] - .contiguous() - .view(torch.float8_e4m3fn) - ) - k_scale = ( - packed_k[:, self.index_head_dim : self.index_head_dim + 4] - .contiguous() - .view(torch.float32) - .view(-1) - ) - score_ks = ( - torch.arange(batch_size, dtype=torch.int32, device=raw_k.device) - * max_pool_len - ) + k_fp8 = packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) + k_scale = packed_k[:, self.index_head_dim : self.index_head_dim + 4].contiguous().view(torch.float32).view(-1) + score_ks = torch.arange(batch_size, dtype=torch.int32, device=raw_k.device) * max_pool_len score_ke = score_ks + pool_lengths return k_fp8, k_scale, score_ks, score_ke, pool_lengths @@ -680,9 +595,7 @@ def _compress_kpool_keys( return compressed_k, compressed_scale @classmethod - def _get_mqa_logits_chunk_size( - cls, query_token_num: int, kv_token_num: int, device: torch.device - ) -> int: + def _get_mqa_logits_chunk_size(cls, query_token_num: int, kv_token_num: int, device: torch.device) -> int: score_element_num = query_token_num * kv_token_num if score_element_num < cls._MQA_LOGITS_STATIC_SKIP_ELEMS: return query_token_num @@ -714,10 +627,8 @@ def _rotate_activation(x: torch.Tensor) -> torch.Tensor: ) hidden_size = x.size(-1) - assert ( - hidden_size & (hidden_size - 1) - ) == 0, "Hidden size must be a power of 2 for Hadamard transform." - return hadamard_transform(x, scale=hidden_size**-0.5) + assert (hidden_size & (hidden_size - 1)) == 0, "Hidden size must be a power of 2 for Hadamard transform." + return hadamard_transform(x, scale=hidden_size ** -0.5) def _get_q_k_bf16( self, @@ -726,9 +637,7 @@ def _get_q_k_bf16( infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek3_2TransformerLayerWeight, ): - q = layer_weight.wq_b_proj_.mm(q_lora).view( - -1, self.tp_index_n_heads, self.index_head_dim - ) + q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.tp_index_n_heads, self.index_head_dim) k = layer_weight.wk_proj_.mm(hidden_states) k = layer_weight.k_norm_(k, eps=self.eps) diff --git a/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py b/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py index 07f1f16a59..5e6e1941f2 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py +++ b/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py @@ -115,9 +115,7 @@ def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: return _hadamard_transform_triton(x, scale) -def hadamard_transform_quant_fp8( - x: torch.Tensor, scale: float = 1.0 -) -> tuple[torch.Tensor, torch.Tensor]: +def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[torch.Tensor, torch.Tensor]: """Fuse Hadamard-128 with the following ue8m0 FP8 quantization.""" assert x.is_cuda, "hadamard_transform_quant_fp8 only supports CUDA tensors" @@ -129,9 +127,7 @@ def hadamard_transform_quant_fp8( original_shape = x.shape rows = x.numel() // 128 output = torch.empty_like(x, dtype=torch.float8_e4m3fn) - output_scale = torch.empty( - (*original_shape[:-1], 1), dtype=torch.float32, device=x.device - ) + output_scale = torch.empty((*original_shape[:-1], 1), dtype=torch.float32, device=x.device) if rows == 0: return output, output_scale diff --git a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py index 57054023b8..6c0996aaeb 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py +++ b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py @@ -20,9 +20,7 @@ def _trans_topk_index_to_mem_index( cur_index = tl.program_id(0) offs_d = tl.arange(0, BLOCK_DMODEL) mask = offs_d < topk_width - topk_index_ptrs = ( - topk_index + cur_index * topk_index_stride_b + offs_d * topk_index_stride_k - ) + topk_index_ptrs = topk_index + cur_index * topk_index_stride_b + offs_d * topk_index_stride_k topk_indices = tl.load(topk_index_ptrs, mask=mask, other=-1) ragged_start = tl.load(ragged_start_index + cur_index) topk_indices = tl.where(topk_indices != -1, topk_indices + ragged_start, -1) @@ -31,9 +29,7 @@ def _trans_topk_index_to_mem_index( dest_mem_index = ragged_mem_index + topk_indices mem_index = tl.load(dest_mem_index, mask=mask & (topk_indices != -1), other=-1) tl.store( - topk_mem_index - + cur_index * topk_mem_index_stride_b - + offs_d * topk_mem_index_stride_k, + topk_mem_index + cur_index * topk_mem_index_stride_b + offs_d * topk_mem_index_stride_k, mem_index, mask=mask, ) diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 709f29843f..b71ad99947 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -28,9 +28,7 @@ class Glm5NextNsaInfer(NsaInfer): """GLM indexer projection without rotary dimensions.""" def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): - q = layer_weight.wq_b_proj_.mm(q_lora).view( - -1, self.tp_index_n_heads, self.index_head_dim - ) + q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.tp_index_n_heads, self.index_head_dim) k = layer_weight.wk_proj_.mm(hidden_states.to(q_lora.dtype)) k = layer_weight.k_norm_(k, eps=self.eps) return q, k, k @@ -41,13 +39,9 @@ def _quantize_indexer_activation(self, value: torch.Tensor): ) assert self.block_size == 128 and self.scale_fmt == "ue8m0" - return hadamard_transform_quant_fp8( - value, scale=self.index_head_dim**-0.5 - ) + return hadamard_transform_quant_fp8(value, scale=self.index_head_dim ** -0.5) - def _scale_indexer_weights( - self, weights: torch.Tensor, q_scale: torch.Tensor - ) -> torch.Tensor: + def _scale_indexer_weights(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: from lightllm.models.deepseek3_2.triton_kernel.indexer_weight_scale import ( scale_indexer_weights_, ) @@ -57,22 +51,17 @@ def _scale_indexer_weights( def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): # GLM stores weights_proj in FP32, so its activation must match before # delegating to the shared NSA scoring and top-k implementation. - return super()._get_indices( - hidden_states.float(), q_lora, infer_state, att_state, layer_weight - ) + return super()._get_indices(hidden_states.float(), q_lora, infer_state, att_state, layer_weight) class Glm5NextTransformerLayerInfer(Deepseek3_2TransformerLayerInfer): def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.num_hidden_layers = network_config["num_hidden_layers"] - self.autotune_layer_num = network_config.get( - "autotune_layer_num", self.num_hidden_layers - ) + self.autotune_layer_num = network_config.get("autotune_layer_num", self.num_hidden_layers) self.is_mtp_layer = layer_num >= self.num_hidden_layers self.is_linear_attention_layer = ( - not self.is_mtp_layer - and network_config["layer_types"][layer_num] == "linear_attention" + not self.is_mtp_layer and network_config["layer_types"][layer_num] == "linear_attention" ) self.mhc_streams = network_config.get("hc_mult", 4) self.hc_eps = network_config.get("hc_eps", 1e-6) @@ -92,18 +81,14 @@ def __init__(self, layer_num, network_config): # GLM's recurrent EAGLE drafter processes one row per logical # request. Only target-model decode uses the widened verification # layout of mtp_step + 1 rows. - self.indexer.decode_mtp_step = ( - 0 if self.is_mtp_layer else get_env_start_args().mtp_step - ) + self.indexer.decode_mtp_step = 0 if self.is_mtp_layer else get_env_start_args().mtp_step def _ffn_tp(self, input, infer_state, layer_weight): """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" input = input.view(-1, self.embed_dim_) up_gate_out = layer_weight.gate_up_proj.mm(input) - ffn1_out = self.alloc_tensor( - (input.size(0), up_gate_out.size(1) // 2), input.dtype - ) + ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) silu_and_mul_fwd( up_gate_out, ffn1_out, @@ -126,14 +111,10 @@ def _get_qkv(self, input, infer_state, layer_weight): if infer_state.need_dp_prefill_balance: input = infer_state._all_to_all_unbalance_get(data=input) - q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split( - [self.q_lora_rank, self.kv_lora_rank], dim=-1 - ) + q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split([self.q_lora_rank, self.kv_lora_rank], dim=-1) q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_) infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q} - q = layer_weight.q_b_proj_.mm(q).view( - -1, self.tp_q_head_num_, self.qk_nope_head_dim - ) + q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim) cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank) rmsnorm_forward( cache_kv[:, :, : self.kv_lora_rank], @@ -151,9 +132,7 @@ def _get_o(self, input, infer_state, layer_weight): input = infer_state._all_to_all_balance_get(data=input) if input.shape[2] == self.kv_lora_rank: input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) - output = layer_weight.o_weight_.mm( - input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim) - ) + output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) all_reduce(output, group=infer_state.dist_group) return output @@ -211,9 +190,7 @@ def _kda_projections(self, input, infer_state, layer_weight): def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): tokens = norm_gate.shape[0] core_output = core_output.view(-1, self.linear_head_dim) - norm_gate = norm_gate.view( - tokens, self.tp_linear_num_heads, self.linear_head_dim - ) + norm_gate = norm_gate.view(tokens, self.tp_linear_num_heads, self.linear_head_dim) output = layer_weight.linear_o_norm( input=core_output, gate_value=norm_gate, @@ -228,12 +205,8 @@ def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): def context_attention_forward(self, input_embeddings, infer_state, layer_weight): if not self.is_linear_attention_layer: - return super().context_attention_forward( - input_embeddings, infer_state, layer_weight - ) - mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( - input_embeddings, infer_state, layer_weight - ) + return super().context_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) core_output = infer_state.prefill_att_state1.prefill_att( q=None, k=None, @@ -254,12 +227,8 @@ def context_attention_forward(self, input_embeddings, infer_state, layer_weight) def token_attention_forward(self, input_embeddings, infer_state, layer_weight): if not self.is_linear_attention_layer: - return super().token_attention_forward( - input_embeddings, infer_state, layer_weight - ) - mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections( - input_embeddings, infer_state, layer_weight - ) + return super().token_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) core_output = infer_state.decode_att_state1.decode_att( q=None, k=None, @@ -297,45 +266,26 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): if self.layer_num_ == 0: streams = hc_expand(streams.view(-1, self.embed_dim_), self.mhc_streams) - layer_input, residual_mix, post_mix = self._hc_pre( - streams, layer_weight, "attn", layer_weight.att_norm_weight_ - ) + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) if prefill: - layer_output = self.context_attention_forward( - layer_input, infer_state, layer_weight - ) + layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) else: - layer_output = self.token_attention_forward( - layer_input, infer_state, layer_weight - ) - streams = hc_post( - layer_output, streams, residual_mix, post_mix, self.mhc_streams - ) + layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) - layer_input, residual_mix, post_mix = self._hc_pre( - streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_ - ) + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) if infer_state.use_replicated_attention_ep: if self.is_moe: - local_input = self._tpsp_sp_split( - input=layer_input, infer_state=infer_state - ) + local_input = self._tpsp_sp_split(input=layer_input, infer_state=infer_state) local_output = self._ffn(local_input, infer_state, layer_weight) - layer_output = self._tpsp_allgather( - input=local_output, infer_state=infer_state - ) + layer_output = self._tpsp_allgather(input=local_output, infer_state=infer_state) else: layer_output = self._ffn_tp(layer_input, infer_state, layer_weight) all_reduce(layer_output, group=infer_state.dist_group) else: layer_output = self._ffn(layer_input, infer_state, layer_weight) - streams = hc_post( - layer_output, streams, residual_mix, post_mix, self.mhc_streams - ) - is_autotune_last_layer = ( - Autotuner.is_autotune_warmup() - and self.layer_num_ == self.autotune_layer_num - 1 - ) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + is_autotune_last_layer = Autotuner.is_autotune_warmup() and self.layer_num_ == self.autotune_layer_num - 1 if self.layer_num_ == self.num_hidden_layers - 1 or is_autotune_last_layer: return hc_contract(streams, self.mhc_streams) return streams @@ -343,13 +293,9 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): def context_forward(self, input_embeddings, infer_state, layer_weight): if self.is_mtp_layer: return super().context_forward(input_embeddings, infer_state, layer_weight) - return self._forward_mhc( - input_embeddings, infer_state, layer_weight, prefill=True - ) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=True) def token_forward(self, input_embeddings, infer_state, layer_weight): if self.is_mtp_layer: return super().token_forward(input_embeddings, infer_state, layer_weight) - return self._forward_mhc( - input_embeddings, infer_state, layer_weight, prefill=False - ) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=False) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index 0f82f41c02..32e712f010 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -67,9 +67,7 @@ def __init__( tp_rank=tp_rank, tp_world_size=tp_world_size, ) - self.sharded_slicer = get_row_slice_mixin( - "none", tp_rank=tp_rank, tp_world_size=tp_world_size - ) + self.sharded_slicer = get_row_slice_mixin("none", tp_rank=tp_rank, tp_world_size=tp_world_size) self.replicated_slicer = get_row_slice_mixin("none", tp_rank=0, tp_world_size=1) def _get_param_slicer(self, sub_child_index: int): diff --git a/test/kernel/test_glm5_kda_fusions.py b/test/kernel/test_glm5_kda_fusions.py index 9042a41b92..83489ecb8c 100644 --- a/test/kernel/test_glm5_kda_fusions.py +++ b/test/kernel/test_glm5_kda_fusions.py @@ -95,12 +95,12 @@ def test_fused_hadamard_fp8_quant_matches_two_kernel_chain(shape): torch.manual_seed(1) value = torch.randn(shape, device="cuda", dtype=torch.bfloat16) expected_value, expected_scale = act_quant( - hadamard_transform(value, scale=128**-0.5), + hadamard_transform(value, scale=128 ** -0.5), block_size=128, scale_fmt="ue8m0", ) - actual_value, actual_scale = hadamard_transform_quant_fp8(value, scale=128**-0.5) + actual_value, actual_scale = hadamard_transform_quant_fp8(value, scale=128 ** -0.5) assert torch.equal(actual_value, expected_value) torch.testing.assert_close(actual_scale, expected_scale, rtol=0, atol=0) @@ -110,7 +110,7 @@ def test_fused_indexer_weight_scale_matches_torch_chain(): torch.manual_seed(2) weights = torch.randn(257, 32, device="cuda", dtype=torch.float32) q_scale = torch.rand(257, 32, 1, device="cuda", dtype=torch.float32) - scale = 128**-0.5 * 32**-0.5 + scale = 128 ** -0.5 * 32 ** -0.5 expected = (weights * scale).unsqueeze(-1).mul(q_scale).squeeze(-1) actual = scale_indexer_weights_(weights.clone(), q_scale, scale) diff --git a/tools/run_glm53_long_prompt_bench.py b/tools/run_glm53_long_prompt_bench.py index 1408a47df8..a31795d498 100644 --- a/tools/run_glm53_long_prompt_bench.py +++ b/tools/run_glm53_long_prompt_bench.py @@ -26,9 +26,7 @@ def cached_load(self, tokenizer, model_id=None): with cache_path.open("rb") as cache_file: cached = pickle.load(cache_file) if cached["metadata"] != expected: - raise ValueError( - f"Long-prompt cache metadata mismatch: {cached['metadata']} != {expected}" - ) + raise ValueError(f"Long-prompt cache metadata mismatch: {cached['metadata']} != {expected}") row = cached["row"] print(f"Loaded long prompt from {cache_path}", flush=True) else: diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py index 60b9995d37..a666e48646 100644 --- a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py @@ -12,18 +12,14 @@ def test_trans_topk_index_to_mem_index(topk): # Create topk_index tensor with some valid indices and some -1 (padding) topk_index = torch.zeros((batch_size, topk), dtype=torch.int32, device="cuda") - topk_index[:, 0 : topk - 1] = torch.arange( - 0, topk - 1, dtype=torch.int32, device="cuda" - ) + topk_index[:, 0 : topk - 1] = torch.arange(0, topk - 1, dtype=torch.int32, device="cuda") topk_index[:, -1] = -1 ragged_start_index = torch.tensor([2], dtype=torch.int32, device="cuda") # Create ragged_mem_index lookup table ragged_mem_index = torch.arange(0, topk + 2, dtype=torch.int32, device="cuda") + 10 - topk_mem_index = trans_topk_index_to_mem_index( - topk_index, ragged_start_index, ragged_mem_index - ) + topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_start_index, ragged_mem_index) expected_index = torch.cat( ( From 1496b5506d2652ecf3307791724e4d6106f8dadb Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 1 Sep 2026 02:27:50 +0800 Subject: [PATCH 28/28] docs: point to formatted GLM-5.3 image --- GLM53_H100_DEPLOY.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md index 0ecc81e4df..f4c9f089f6 100644 --- a/GLM53_H100_DEPLOY.md +++ b/GLM53_H100_DEPLOY.md @@ -8,17 +8,17 @@ CUDA Graph 128，面向 16K 输入、256 输出的高并发文本服务。 私有仓库标签： ```text -registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.5.0-glm53-16k256-kpool-c3f39a82 +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.5.0-glm53-16k256-kpool-996cef93 ``` 不可变镜像： ```text -registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:76a968231ffb0a4c8996bc9a0fb981b88b3afbf25b7ad77a07fa07fccc04901d +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:0260e9884e46de899f4b845aa3796d0603b7b6ee7593c1fded35b7cee4462169 ``` 本机镜像为 `lightllm-glm53:16k256-kpool`，镜像 ID 为 -`sha256:d388173fbbc1fc2e8961785248331cecd08683af3c5f28ff67bc530d770d1409`。 +`sha256:8fbb91ee50dde5af4289d9e3cc75dddf93a8dd6fc2a201222a28612a5679cf0c`。 同一不可变镜像已拉取到 H100 节点。 ## H100 部署命令 @@ -26,7 +26,7 @@ registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:76a968231ffb0a4c899 镜像已内置服务参数；启动时显式开启 K-pool decode 快路径： ```bash -IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:76a968231ffb0a4c8996bc9a0fb981b88b3afbf25b7ad77a07fa07fccc04901d" +IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:0260e9884e46de899f4b845aa3796d0603b7b6ee7593c1fded35b7cee4462169" sudo docker pull "$IMAGE" sudo docker run -d \