From a270ee1969e90ece7f5dc84d7caf652733358333 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Sat, 4 Jul 2026 08:06:02 +0000 Subject: [PATCH 01/10] feat: opt mega moe perf --- .../fused_moe/grouped_fused_moe_ep.py | 171 ++++++++++++++++-- lightllm/common/quantization/__init__.py | 26 +-- lightllm/common/quantization/deepgemm.py | 23 ++- 3 files changed, 185 insertions(+), 35 deletions(-) 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 cb2e370cb9..14ec00384b 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,6 +1,7 @@ """Fused MoE kernel.""" import torch import triton +import triton.language as tl from typing import Any, Callable, Dict, Optional, Tuple from lightllm.distributed import dist_group_manager from lightllm.utils.log_utils import init_logger @@ -42,6 +43,154 @@ def use_sm100_mega_moe(quant_method: Any) -> bool: return is_sm100_gpu() and quant_method.method_name == "deepgemm-fp4fp8-b32" +def _per_token_cast_to_fp8_packed_ue8m0(hidden_states: torch.Tensor, gran_k: int): + from deep_gemm.utils import per_token_cast_to_fp8 + + hidden_states, scale = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=gran_k, + use_packed_ue8m0=False, + ) + assert scale.size(-1) % 4 == 0, "packed UE8M0 scale requires scale groups divisible by 4" + scale = (scale.view(torch.int32) >> 23).to(torch.uint8).view(torch.int32) + return hidden_states, scale + + +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True), exp + + +@triton.jit +def _mega_moe_quant_topk_to_buffer_kernel( + x_ptr, + x_out_ptr, + x_sf_out_ptr, + topk_idx_ptr, + topk_idx_out_ptr, + topk_weights_ptr, + topk_weights_out_ptr, + stride_x_m: tl.constexpr, + stride_x_k: tl.constexpr, + stride_x_out_m: tl.constexpr, + stride_x_out_k: tl.constexpr, + stride_x_sf_out_m: tl.constexpr, + stride_x_sf_out_k: tl.constexpr, + stride_topk_idx_m: tl.constexpr, + stride_topk_idx_k: tl.constexpr, + stride_topk_idx_out_m: tl.constexpr, + stride_topk_idx_out_k: tl.constexpr, + stride_topk_weights_m: tl.constexpr, + stride_topk_weights_k: tl.constexpr, + stride_topk_weights_out_m: tl.constexpr, + stride_topk_weights_out_k: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + TOPK: tl.constexpr, + GROUP_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + token_id = tl.program_id(0) + pack_id = tl.program_id(1) + offsets = tl.arange(0, BLOCK) + cols = pack_id * BLOCK + offsets + + x = tl.load(x_ptr + token_id * stride_x_m + cols * stride_x_k).to(tl.float32) + abs_x = tl.abs(x) + group_id = offsets // GROUP_SIZE + + amax0 = tl.max(tl.where(group_id == 0, abs_x, 0.0)) + amax1 = tl.max(tl.where(group_id == 1, abs_x, 0.0)) + amax2 = tl.max(tl.where(group_id == 2, abs_x, 0.0)) + amax3 = tl.max(tl.where(group_id == 3, abs_x, 0.0)) + + scale0, exp0 = _ceil_to_ue8m0(tl.maximum(amax0, 1.0e-4) / FP8_MAX) + scale1, exp1 = _ceil_to_ue8m0(tl.maximum(amax1, 1.0e-4) / FP8_MAX) + scale2, exp2 = _ceil_to_ue8m0(tl.maximum(amax2, 1.0e-4) / FP8_MAX) + scale3, exp3 = _ceil_to_ue8m0(tl.maximum(amax3, 1.0e-4) / FP8_MAX) + + scale = tl.where( + group_id == 0, + scale0, + tl.where(group_id == 1, scale1, tl.where(group_id == 2, scale2, scale3)), + ) + x_q = tl.clamp(x / scale, FP8_MIN, FP8_MAX).to(x_out_ptr.dtype.element_ty) + tl.store(x_out_ptr + token_id * stride_x_out_m + cols * stride_x_out_k, x_q) + + packed_scale = exp0 | (exp1 << 8) | (exp2 << 16) | (exp3 << 24) + tl.store(x_sf_out_ptr + token_id * stride_x_sf_out_m + pack_id * stride_x_sf_out_k, packed_scale) + + if pack_id == 0: + topk_offsets = tl.arange(0, TOPK) + topk_idx = tl.load(topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k) + topk_weights = tl.load( + topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k + ) + tl.store( + topk_idx_out_ptr + token_id * stride_topk_idx_out_m + topk_offsets * stride_topk_idx_out_k, + topk_idx.to(topk_idx_out_ptr.dtype.element_ty), + ) + tl.store( + topk_weights_out_ptr + + token_id * stride_topk_weights_out_m + + topk_offsets * stride_topk_weights_out_k, + topk_weights.to(topk_weights_out_ptr.dtype.element_ty), + ) + + +def _prepare_mega_moe_buffer( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + buffer: Any, + group_size: int, +): + num_tokens, hidden_size = hidden_states.shape + if num_tokens == 0: + return + assert hidden_size % (group_size * 4) == 0, "packed UE8M0 scale requires four FP8 groups per int32" + assert hidden_states.is_contiguous(), "hidden_states must be contiguous" + assert buffer.x.shape[0] >= num_tokens and buffer.x.shape[1] == hidden_size + assert buffer.x_sf.shape[0] >= num_tokens and buffer.x_sf.shape[1] == hidden_size // group_size // 4 + + block = group_size * 4 + finfo = torch.finfo(buffer.x.dtype) + _mega_moe_quant_topk_to_buffer_kernel[(num_tokens, hidden_size // block)]( + hidden_states, + buffer.x, + buffer.x_sf, + topk_ids, + buffer.topk_idx, + topk_weights, + buffer.topk_weights, + hidden_states.stride(0), + hidden_states.stride(1), + buffer.x.stride(0), + buffer.x.stride(1), + buffer.x_sf.stride(0), + buffer.x_sf.stride(1), + topk_ids.stride(0), + topk_ids.stride(1), + buffer.topk_idx.stride(0), + buffer.topk_idx.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + buffer.topk_weights.stride(0), + buffer.topk_weights.stride(1), + FP8_MIN=finfo.min, + FP8_MAX=finfo.max, + TOPK=topk_ids.shape[1], + GROUP_SIZE=group_size, + BLOCK=block, + num_warps=4, + num_stages=4, + ) + + 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: @@ -123,8 +272,6 @@ def mega_moe_impl( if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") - from deep_gemm.utils import per_token_cast_to_fp8 - buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) if buffer is None: raise RuntimeError("SM100 Mega MoE requires dist_group_manager.ep_mega_moe_buffer to be initialized") @@ -135,19 +282,10 @@ def mega_moe_impl( f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" ) - qinput_tensor = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) state = _get_mega_moe_cache_state(w13, w2) l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) - buffer.x[:num_tokens].copy_(qinput_tensor[0]) - buffer.x_sf[:num_tokens].copy_(qinput_tensor[1]) - buffer.topk_idx[:num_tokens].copy_(topk_ids) - buffer.topk_weights[:num_tokens].copy_(topk_weights) + _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) deep_gemm.fp8_fp4_mega_moe( @@ -167,14 +305,7 @@ def quantize_fused_experts_input( ): check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): - from deep_gemm.utils import per_token_cast_to_fp8 - - return per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) + return _per_token_cast_to_fp8_packed_ue8m0(hidden_states, quant_method.block_size) block_size_k = 0 if w13.weight.ndim == 3: diff --git a/lightllm/common/quantization/__init__.py b/lightllm/common/quantization/__init__.py index cd534d53ec..f676e35d64 100644 --- a/lightllm/common/quantization/__init__.py +++ b/lightllm/common/quantization/__init__.py @@ -43,6 +43,7 @@ def _parse_network_config(self, network_config): self.quantized_weight = False self.static_activation = False self.hf_quantization_config = None + self._mapping_expert_quant_method() return self.quantized_weight = True activation_scheme = network_config.get("activation_scheme", "dynamic") @@ -50,6 +51,19 @@ def _parse_network_config(self, network_config): self.hf_quantization_config = hf_quantization_config self.hf_quantization_method = hf_quantization_config["quant_method"] self._mapping_quant_method() + self._mapping_expert_quant_method() + + def _mapping_expert_quant_method(self): + expert_dtype = self.expert_dtype or self.network_config_.get("expert_dtype", None) + if expert_dtype is None: + return + target = self._get_expert_quant_type(expert_dtype) + for layer_num in range(self.layer_num): + if self.expert_dtype is not None: + self.quant_cfg[layer_num]["fused_moe"] = target + else: + self.quant_cfg[layer_num].setdefault("fused_moe", target) + logger.info(f"select fused_moe quant way from expert_dtype=`{expert_dtype}`: {target}") def _mapping_quant_method(self): if self.hf_quantization_method == "fp8": @@ -63,18 +77,6 @@ def _mapping_quant_method(self): self.quant_type = "vllm-fp8w8a8-b128" logger.info(f"select fp8w8a8-b128 quant way: {self.quant_type}") - # fp8 量化下，部分 MoE 模型（如 DeepSeek-V4），可以单独声明 expert 权重精度， - # 按其值给 fused_moe 选用对应的 deepgemm 量化方法。 - expert_dtype = self.expert_dtype or self.network_config_.get("expert_dtype", None) - if expert_dtype is None: - return - target = self._get_expert_quant_type(expert_dtype) - for layer_num in range(self.layer_num): - if self.expert_dtype is not None: - self.quant_cfg[layer_num]["fused_moe"] = target - else: - self.quant_cfg[layer_num].setdefault("fused_moe", target) - logger.info(f"select fused_moe quant way from expert_dtype=`{expert_dtype}`: {target}") elif self.hf_quantization_method == "awq": self.quant_type = "awq" if is_awq_marlin_compatible(self.hf_quantization_config): diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index ec1ee90fd4..328e9a71c9 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -182,11 +182,28 @@ def _create_weight( out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims assert in_dim % 2 == 0, "FP4 packed weight requires even input dimension" assert in_dim % self.block_size == 0, "FP4 scale dimension must be divisible by block_size" + scales_per_int32 = 4 + scale_layout_k = self.block_size * scales_per_int32 + assert in_dim % scale_layout_k == 0, ( + f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" + ) expert_prefix = (num_experts,) if num_experts > 1 else () weight = torch.empty(expert_prefix + (out_dim, in_dim // 2), dtype=torch.int8).cuda(device_id) - weight_scale = torch.empty(expert_prefix + (out_dim, in_dim // self.block_size), dtype=torch.int32).cuda( - device_id - ) + scale_dim = in_dim // scale_layout_k + if num_experts > 1: + weight_scale = torch.empty_strided( + (num_experts, out_dim, scale_dim), + (out_dim * scale_dim, 1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) + else: + weight_scale = torch.empty_strided( + (out_dim, scale_dim), + (1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) mm_param_list = self._split_weight_pack( mm_param, From 272bfb24bc8cb1a39b6c408cc01d9bf52b2e8f11 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 19 Jun 2026 00:13:53 +0800 Subject: [PATCH 02/10] support glm52 --- docs/CN/source/models/supported_models.rst | 4 +- docs/EN/source/models/supported_models.rst | 3 +- .../attention/nsa/flashmla_sparse.py | 9 +- .../attention/nsa/fp8_flashmla_sparse.py | 39 +- lightllm/common/basemodel/basemodel.py | 7 +- .../layer_weights/meta_weights/norm_weight.py | 16 + .../fused_moe/grouped_fused_moe.py | 57 +- .../fused_moe/moe_silu_and_mul_group_quant.py | 118 ++++ .../triton_kernel/norm/fused_add_rmsnorm.py | 98 ++++ ...oken_group_quant_deepseek3_2mem_manager.py | 44 +- .../kv_cache_mem_manager/operator/deepseek.py | 106 ++++ ...=32,dtype=torch.bfloat16}_NVIDIA_H200.json | 98 ++++ ...=64,dtype=torch.bfloat16}_NVIDIA_H200.json | 98 ++++ ...out_dtype=torch.bfloat16}_NVIDIA_H200.json | 98 ++++ ...out_dtype=torch.bfloat16}_NVIDIA_H200.json | 518 ++++++++++++++++++ lightllm/distributed/communication_op.py | 28 + lightllm/distributed/flashinfer_all_reduce.py | 23 + lightllm/models/__init__.py | 1 + .../layer_infer/transformer_layer_infer.py | 86 ++- .../layer_infer/transformer_layer_infer.py | 82 ++- .../layer_weights/transformer_layer_weight.py | 5 +- .../triton_kernel/fp8_mqa_logits.py | 191 +++---- lightllm/models/glm5_2/__init__.py | 3 + lightllm/models/glm5_2/indexshare.py | 12 + .../models/glm5_2/layer_infer/__init__.py | 3 + .../layer_infer/transformer_layer_infer.py | 101 ++++ .../models/glm5_2/layer_weights/__init__.py | 3 + .../layer_weights/transformer_layer_weight.py | 46 ++ lightllm/models/glm5_2/model.py | 54 ++ lightllm/models/glm5_2_mtp/__init__.py | 3 + lightllm/models/glm5_2_mtp/model.py | 98 ++++ lightllm/server/build_prompt.py | 34 +- .../model_infer/mode_backend/base_backend.py | 6 +- lightllm/server/tokenizer.py | 41 ++ lightllm/utils/config_utils.py | 8 + lightllm/utils/kv_cache_utils.py | 27 + test/acc/test_gsmk.py | 6 +- .../test_fp8_dsa_cpu_cache.py | 134 +++++ .../triton_kernel/test_fused_add_rmsnorm.py | 80 +++ .../test_moe_silu_and_mul_group_quant.py | 56 ++ .../triton_kernel/test_fp8_mqa_logits.py | 203 +++++++ 41 files changed, 2446 insertions(+), 201 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_group_quant.py create mode 100644 lightllm/common/basemodel/triton_kernel/norm/fused_add_rmsnorm.py create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=12288,out_dtype=torch.bfloat16}_NVIDIA_H200.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json create mode 100644 lightllm/models/glm5_2/__init__.py create mode 100644 lightllm/models/glm5_2/indexshare.py create mode 100644 lightllm/models/glm5_2/layer_infer/__init__.py create mode 100644 lightllm/models/glm5_2/layer_infer/transformer_layer_infer.py create mode 100644 lightllm/models/glm5_2/layer_weights/__init__.py create mode 100644 lightllm/models/glm5_2/layer_weights/transformer_layer_weight.py create mode 100644 lightllm/models/glm5_2/model.py create mode 100644 lightllm/models/glm5_2_mtp/__init__.py create mode 100644 lightllm/models/glm5_2_mtp/model.py create mode 100644 test/cpu_cache_kernel/test_fp8_dsa_cpu_cache.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_fused_add_rmsnorm.py create mode 100644 unit_tests/common/fused_moe/test_moe_silu_and_mul_group_quant.py create mode 100644 unit_tests/models/deepseek3_2/triton_kernel/test_fp8_mqa_logits.py diff --git a/docs/CN/source/models/supported_models.rst b/docs/CN/source/models/supported_models.rst index 3d66d2e073..eba38bd791 100755 --- a/docs/CN/source/models/supported_models.rst +++ b/docs/CN/source/models/supported_models.rst @@ -56,6 +56,8 @@ lightllm 支持大多数的主流的开源大语言模型以及多模态模型 - * - `Qwen3-Moe `_ - + * - `GLM-5.2 `_ + - 支持 BF16/FP8 和 MTP。 多模态模型 @@ -93,4 +95,4 @@ Reward模型 * - `internLM-reward `_ - :code:`--use_reward_model` * - `Qwen2-Reward `_ - - :code:`--use_reward_model` \ No newline at end of file + - :code:`--use_reward_model` diff --git a/docs/EN/source/models/supported_models.rst b/docs/EN/source/models/supported_models.rst index 1b1d4fcd03..105ced14c9 100755 --- a/docs/EN/source/models/supported_models.rst +++ b/docs/EN/source/models/supported_models.rst @@ -56,6 +56,8 @@ Large Language Models - * - `DeepSeek-V3.2 `_ - + * - `GLM-5.2 `_ + - Supports BF16/FP8 and MTP. Multimodal Models ^^^^^^^^^^^^^^^^^ @@ -94,4 +96,3 @@ Reward Models - :code:`--use_reward_model` * - `Qwen2-Reward `_ - :code:`--use_reward_model` - diff --git a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py index c3456f4b7a..49fb1c2dde 100644 --- a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py @@ -3,6 +3,7 @@ import dataclasses import torch +import torch.nn.functional as F from typing import Tuple, TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl @@ -86,6 +87,12 @@ def _nsa_prefill_att( if topk_mem_indices.ndim == 2: topk_mem_indices = topk_mem_indices.unsqueeze(1) + real_head_num = q.shape[1] + head_block_size = 64 + pad_head_num = (-real_head_num) % head_block_size + if pad_head_num: + q = F.pad(q, (0, 0, 0, pad_head_num)) + mla_out, _, _ = flash_mla_sparse_fwd( q=q, kv=kv, @@ -93,7 +100,7 @@ def _nsa_prefill_att( sm_scale=softmax_scale, d_v=kv_lora_rank, ) - return mla_out + return mla_out[:, :real_head_num, :] @dataclasses.dataclass diff --git a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py index 539ade769e..5527c614c0 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py @@ -1,5 +1,6 @@ import dataclasses import torch +import torch.nn.functional as F from typing import TYPE_CHECKING, Tuple from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState @@ -70,10 +71,9 @@ def _nsa_prefill_att( packed_kv: torch.Tensor, att_control: AttControl, ) -> torch.Tensor: - import flash_mla + from flash_mla import flash_mla_sparse_fwd nsa_dict = att_control.nsa_prefill_dict - topk_indices = nsa_dict["topk_indices"] softmax_scale = nsa_dict["softmax_scale"] kv_lora_rank = nsa_dict["kv_lora_rank"] topk_mem_indices = nsa_dict["topk_mem_indices"] @@ -91,18 +91,25 @@ def _nsa_prefill_att( ) else: kv = prefill_cache_kv + topk_indices = nsa_dict["topk_indices"] if topk_indices.ndim == 2: topk_indices = topk_indices.unsqueeze(1) - mla_out, _, _ = flash_mla.flash_mla_sparse_fwd( + real_head_num = q.shape[1] + head_block_size = 64 + pad_head_num = (-real_head_num) % head_block_size + if pad_head_num: + q = F.pad(q, (0, 0, 0, pad_head_num)) + + mla_out, _, _ = flash_mla_sparse_fwd( q=q, kv=kv, indices=topk_indices, sm_scale=softmax_scale, d_v=kv_lora_rank, ) - return mla_out + return mla_out[:, :real_head_num, :] @dataclasses.dataclass @@ -141,9 +148,9 @@ def init_state(self): ragged_mem_index=self.ragged_mem_index, hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, ) - import flash_mla + from flash_mla import get_mla_metadata - self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() + self.flashmla_sched_meta, _ = get_mla_metadata() return def decode_att( @@ -164,7 +171,7 @@ def _nsa_decode_att( packed_kv: torch.Tensor, att_control: AttControl, ) -> torch.Tensor: - import flash_mla + from flash_mla import flash_mla_with_kvcache nsa_dict = att_control.nsa_decode_dict topk_mem_indices = nsa_dict["topk_mem_indices"] @@ -177,13 +184,26 @@ def _nsa_decode_att( q_nope, q_rope = q q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous() + + real_head_num = q_all.shape[2] + if real_head_num <= 64: + padded_head_num = 64 + elif real_head_num <= 128: + padded_head_num = 128 + else: + padded_head_num = real_head_num + if padded_head_num != real_head_num: + q_all = F.pad(q_all, (0, 0, 0, padded_head_num - real_head_num)) + + # Token-granular "pages": indices=topk_mem_indices are absolute KV-pool slots, + # so the kv cache is viewed as (num_tokens, 1, 1, bytes) and no block_table is needed. kv = torch.as_strided( packed_kv, size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), ) - o_tensor, _ = flash_mla.flash_mla_with_kvcache( + o_tensor, _ = flash_mla_with_kvcache( q=q_all, k_cache=kv, block_table=None, @@ -193,6 +213,7 @@ def _nsa_decode_att( softmax_scale=softmax_scale, causal=False, is_fp8_kvcache=True, - indices=topk_mem_indices, + indices=topk_mem_indices.to(dtype=torch.int32), ) + o_tensor = o_tensor[:, :, :real_head_num, :] return o_tensor[:, 0, :, :] # [b, 1, h, d] -> [b, h, d] diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e83de684a7..16cd19e2f4 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1177,12 +1177,7 @@ def _init_padded_req(self): def _gen_special_model_input(self, token_num: int): special_model_input = {} - is_mtp_draft_model = ( - "Deepseek3MTPModel" in str(self.__class__) - or "Qwen3MOEMTPModel" in str(self.__class__) - or "MistralMTPModel" in str(self.__class__) - or "Glm4MoeLiteMTPModel" in str(self.__class__) - ) + is_mtp_draft_model = getattr(self, "is_mtp_draft_model", False) if is_mtp_draft_model: special_model_input["mtp_draft_input_hiddens"] = torch.randn( token_num, self.config["hidden_size"], dtype=self.data_type, device="cuda" 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..e01287fac5 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py @@ -3,6 +3,7 @@ from .base_weight import BaseWeightTpl from lightllm.utils.dist_utils import get_current_device_id, get_current_rank_in_dp, get_dp_world_size from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.common.basemodel.triton_kernel.norm.fused_add_rmsnorm import fused_add_rmsnorm_forward from lightllm.common.basemodel.triton_kernel.norm.layernorm import layernorm_forward from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_fused_forward from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward @@ -71,6 +72,21 @@ def __call__( ) -> torch.Tensor: return self._forward(input=input, eps=eps, out=out, alloc_func=alloc_func) + def fused_add_forward( + self, + residual: torch.Tensor, + x: torch.Tensor, + eps: float, + out: Optional[torch.Tensor] = None, + alloc_func=torch.empty, + ) -> torch.Tensor: + """Fused residual-add + RMSNorm: ``residual <- residual + x`` (in place) and return + ``rmsnorm(residual) * weight``. Bit-identical to a plain ``residual.add_(x)`` followed + by ``__call__`` but in a single Triton launch. CUDA/MUSA (Triton) only.""" + if out is None: + out = alloc_func(residual.shape, dtype=residual.dtype, device=residual.device) + return fused_add_rmsnorm_forward(residual=residual, x=x, weight=self.weight, eps=eps, out=out) + class GatedRMSNormWeight(RMSNormWeight): def _triton_forward( 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 c6eeb781dc..15033c55ee 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 @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import torch import triton import triton.language as tl @@ -26,6 +27,7 @@ from lightllm.utils.vllm_utils import vllm_ops from lightllm.utils.device_utils import triton_support_tensor_descriptor from .moe_silu_and_mul import silu_and_mul_fwd +from .moe_silu_and_mul_group_quant import silu_and_mul_group_quant_fwd from .moe_sum_reduce import moe_sum_reduce from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 from lightllm.utils.torch_ops_utils import direct_register_custom_op @@ -35,6 +37,11 @@ logger = init_logger(__name__) +# Fuse the down-projection's per-token-group fp8 quant into the silu_and_mul that produces its +# input (eliminates a separate per_token_group_quant launch per MoE layer). Numerically matches +# the unfused path; gate exists for A/B and rollback. +_ENABLE_FUSED_SILU_QUANT = os.environ.get("LIGHTLLM_FUSED_SILU_QUANT", "1") == "1" + @triton.jit def moe_align_kernel( @@ -860,8 +867,12 @@ def grouped_matmul( assert BLOCK_SIZE_K == triton.next_power_of_2(BLOCK_SIZE_K) if use_fp8_w8a8: + if token_inputs.dtype == expert_weights.dtype: + # token_inputs were already fp8-quantized by a fused producer (e.g. the fused + # silu+mul+group-quant on the down projection); token_input_scale is its scale. + assert token_input_scale is not None, "pre-quantized fp8 input requires its scale" # 当权重使用 block wise 量化时，激活也使用 per token， group size 量化 - if block_size_k == 0: + elif block_size_k == 0: # input 使用 per token 量化 token_inputs, token_input_scale = vllm_ops.scaled_fp8_quant( token_inputs, token_input_scale, use_per_token_if_dynamic=True @@ -1082,18 +1093,44 @@ def fused_experts_impl( bias=w1_bias, ) - silu_and_mul_fwd( - intermediate_cache1.view(-1, N), - intermediate_cache2.view(-1, N // 2), - limit=limit, - alpha=alpha, - layout=layout, - ) + # Fuse the down-projection's per-token-group fp8 quant into silu_and_mul when the down + # weight is block-wise fp8 quantized: the silu output is emitted directly as fp8 + scales, + # so grouped_matmul below skips its internal per_token_group_quant launch. + use_fused_silu_quant = _ENABLE_FUSED_SILU_QUANT and use_fp8_w8a8 and w2_scale is not None and w2_scale.ndim == 3 + if use_fused_silu_quant: + down_token_num = curr_topk_ids.numel() + down_k = N // 2 + down_group_size = down_k // w2_scale.shape[2] + down_inputs = alloc_tensor_func( + (down_token_num, down_k), device=hidden_states.device, dtype=torch.float8_e4m3fn + ) + down_input_scale = alloc_tensor_func( + (down_token_num, down_k // down_group_size), device=hidden_states.device, dtype=torch.float32 + ) + silu_and_mul_group_quant_fwd( + intermediate_cache1.view(-1, N), + down_inputs, + down_input_scale, + down_group_size, + layout=layout, + limit=limit, + alpha=alpha, + ) + else: + silu_and_mul_fwd( + intermediate_cache1.view(-1, N), + intermediate_cache2.view(-1, N // 2), + limit=limit, + alpha=alpha, + layout=layout, + ) + down_inputs = intermediate_cache2.view(-1, N // 2) + down_input_scale = a2_scale grouped_matmul( curr_topk_ids.numel(), - intermediate_cache2.view(-1, N // 2), - a2_scale, + down_inputs, + down_input_scale, expert_to_token_num, expert_to_tokens, expert_to_weights=expert_to_weights, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_group_quant.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_group_quant.py new file mode 100644 index 0000000000..b28987dc68 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_group_quant.py @@ -0,0 +1,118 @@ +import torch + +import triton +import triton.language as tl +from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu + + +@triton.jit +def _silu_and_mul_group_quant_kernel( + input_ptr, + stride_input_m, + out_q_ptr, + stride_out_q_m, + out_s_ptr, + stride_out_s_m, + size_n, # width of the silu output (= down-proj K), a multiple of GROUP_SIZE + limit, + alpha, + fp8_min, + fp8_max, + eps, + GROUP_SIZE: tl.constexpr, + layout: tl.constexpr, # "blocked" or "interleaved" + USE_LIMIT_AND_ALPHA: tl.constexpr, + USE_TANH_APPROXIMATE_GELU: tl.constexpr, +): + # One program per (token row, quant group of GROUP_SIZE columns). Computes silu(gate)*up for + # the group, then a per-group fp8 quant — byte-identical layout to + # silu_and_mul_fwd followed by per_token_group_quant_fp8 (row-major scales), in one launch. + m_index = tl.program_id(0).to(tl.int64) + group_index = tl.program_id(1) + cols = group_index * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + if layout == "interleaved": + # [gate0, up0, gate1, up1, ...] + gate_off = m_index * stride_input_m + cols * 2 + up_off = gate_off + 1 + else: + # [gate0, gate1, ..., up0, up1, ...] + gate_off = m_index * stride_input_m + cols + up_off = m_index * stride_input_m + cols + size_n + gate = tl.load(input_ptr + gate_off).to(tl.float32) + up = tl.load(input_ptr + up_off) + + if USE_LIMIT_AND_ALPHA: + gate = tl.minimum(gate, limit) + up = tl.minimum(tl.maximum(up, -limit), limit) + gate = 1 / (1 + tl.exp(-gate * alpha)) * gate + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = (up + 1) * gate + else: + if 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 + gate = 0.5 * gate * (1.0 + tanh_val) + else: + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + + # quantize the (bf16-rounded) silu output per group, matching per_token_group_quant_fp8. + gate_up_f = gate_up.to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(gate_up_f)), eps) + out_s = _absmax / fp8_max + out_q = tl.clamp(gate_up_f / out_s, fp8_min, fp8_max).to(out_q_ptr.dtype.element_ty) + tl.store(out_q_ptr + m_index * stride_out_q_m + cols, out_q) + tl.store(out_s_ptr + m_index * stride_out_s_m + group_index, out_s) + + +def silu_and_mul_group_quant_fwd( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, + group_size: int, + layout: str = "blocked", + limit=None, + alpha=None, +): + """Fused silu_and_mul + per-token-group fp8 quant. + + ``input`` [M, 2*N] (gate|up) -> ``output_q`` [M, N] fp8 + ``output_s`` [M, N//group_size] + float32 row-major. Equivalent to ``silu_and_mul_fwd(input, tmp); per_token_group_quant_fp8(tmp)`` + but in one kernel, so the down-projection's activation quant disappears as a separate launch. + """ + assert input.is_contiguous() + assert output_q.is_contiguous() and output_q.dtype == torch.float8_e4m3fn + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + size_m = input.shape[0] + size_n = input.shape[-1] // 2 + assert size_n % group_size == 0, f"silu output width {size_n} not divisible by group_size {group_size}" + assert output_q.shape[0] == size_m and output_q.shape[1] == size_n + assert output_s.shape[0] == size_m and output_s.shape[1] == size_n // group_size + + finfo = torch.finfo(torch.float8_e4m3fn) + fp8_max = finfo.max + fp8_min = -fp8_max + # grid: token rows on dim-0 (up to 2**31, may be large for prefill), groups on dim-1 (small). + grid = (size_m, size_n // group_size) + _silu_and_mul_group_quant_kernel[grid]( + input, + input.stride(0), + output_q, + output_q.stride(0), + output_s, + output_s.stride(0), + size_n, + limit if limit is not None else 0.0, + alpha if alpha is not None else 0.0, + fp8_min, + fp8_max, + 1e-10, + GROUP_SIZE=group_size, + layout=layout, + USE_LIMIT_AND_ALPHA=limit is not None and alpha is not None, + USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), + num_warps=1, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/norm/fused_add_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/fused_add_rmsnorm.py new file mode 100644 index 0000000000..bf84f498ce --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/norm/fused_add_rmsnorm.py @@ -0,0 +1,98 @@ +import torch + +import triton +import triton.language as tl +import os + +fused_add_rmsnorm_num_warps = int(os.getenv("RMSNORM_WARPS", "8")) + + +@triton.jit +def _fused_add_rmsnorm_fwd( + X, # pointer to the addend (e.g. attention / ffn output) + RESIDUAL, # pointer to the residual, updated in place to (residual + x) + Y, # pointer to the normalized output + W, # pointer to the weights + x_stride0, + residual_stride0, + y_stride0, + N, # number of columns + eps, # epsilon to avoid division by zero + HAS_WEIGHT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + X += row * x_stride0 + RESIDUAL += row * residual_stride0 + Y += row * y_stride0 + # pass 1: residual = residual + x (in place), accumulate variance of the updated residual + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + r = tl.load(RESIDUAL + cols, mask=mask, other=0.0).to(tl.float32) + # round the updated residual to the storage dtype first, then accumulate variance + # from the rounded value so this matches the unfused (store; reload; rmsnorm) path + # bit-for-bit instead of using the higher-precision fp32 sum. + s = (r + x).to(RESIDUAL.dtype.element_ty) + tl.store(RESIDUAL + cols, s, mask=mask) + s = s.to(tl.float32) + _var += s * s + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # pass 2: normalize the (rounded) updated residual and optionally apply weight + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + s = tl.load(RESIDUAL + cols, mask=mask, other=0.0).to(tl.float32) + y = s * rstd + if HAS_WEIGHT: + y = y * w + tl.store(Y + cols * 1, y.to(Y.dtype.element_ty), mask=mask) + + +def fused_add_rmsnorm_forward( + residual: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor +) -> torch.Tensor: + """Fused residual-add + RMSNorm. + + Computes ``residual <- residual + x`` (in place) and ``out <- rmsnorm(residual) * weight`` + in a single kernel, eliminating the separate elementwise-add launch and the extra HBM + round-trip of the hidden state. ``residual`` and ``x`` share shape; ``residual`` is + updated in place (so the running residual stream is preserved for the next layer). + """ + assert residual.shape == x.shape + r_arg = residual.view(-1, residual.shape[-1]) + x_arg = x.view(-1, x.shape[-1]) + y_arg = out.view(-1, out.shape[-1]) + assert r_arg.shape == x_arg.shape == y_arg.shape + if weight is not None: + assert r_arg.shape[-1] == weight.shape[0] + # contiguous last dim (the kernel assumes unit stride within a row) + assert r_arg.stride(1) == 1 and x_arg.stride(1) == 1 and y_arg.stride(1) == 1 + assert out.data_ptr() == y_arg.data_ptr() + M, N = r_arg.shape + MAX_FUSED_SIZE = 65536 // residual.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_SIZE: + raise RuntimeError("fused_add_rmsnorm doesn't support feature dim >= 64KB.") + if BLOCK_SIZE > 16384: + BLOCK_SIZE = 16384 + _fused_add_rmsnorm_fwd[(M,)]( + x_arg, + r_arg, + y_arg, + weight, + x_arg.stride(0), + r_arg.stride(0), + y_arg.stride(0), + N, + eps, + HAS_WEIGHT=weight is not None, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=fused_add_rmsnorm_num_warps, + ) + return out diff --git a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py index b72587545f..c5c0b157a0 100644 --- a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py @@ -6,7 +6,6 @@ class FP8PerTokenGroupQuantDeepseek3_2MemoryManager(Deepseek2MemoryManager): - operator_class = FP8PerTokenGroupQuantDeepseek3_2MemOperator kv_nope_dim = 512 @@ -27,24 +26,45 @@ class FP8PerTokenGroupQuantDeepseek3_2MemoryManager(Deepseek2MemoryManager): # 132 bytes = 128 + 4 indexer_bytes_per_token = indexer_head_dim + 4 - # 16-byte 对齐，满足FlashMLA的对齐要求 - alignment = 16 - total_bytes_per_token = ( - (flashmla_bytes_per_token + indexer_bytes_per_token + alignment - 1) // alignment * alignment - ) - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): assert head_num == 1, "DeepSeek-V3.2 DSA FP8 path expects MQA-style head_num == 1" self.prefill_dtype = dtype - super().__init__(size, torch.uint8, head_num, self.total_bytes_per_token, layer_num, always_copy, mem_fraction) + super().__init__( + size, torch.uint8, head_num, self.flashmla_bytes_per_token, layer_num, always_copy, mem_fraction + ) + + def get_cell_size(self): + return self.layer_num * (self.flashmla_bytes_per_token + self.indexer_bytes_per_token) + + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + token_num = size + 1 + self.kv_buffer = torch.empty( + (layer_num, token_num, head_num, self.flashmla_bytes_per_token), + dtype=dtype, + device="cuda", + ) + self.indexer_k_buffer = torch.empty( + (layer_num, token_num, head_num, self.indexer_bytes_per_token), + dtype=dtype, + device="cuda", + ) def get_att_input_params(self, layer_index: int) -> Any: - return self.kv_buffer[layer_index][:, :, : self.flashmla_bytes_per_token] + return self.kv_buffer[layer_index] def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: - begin = self.flashmla_bytes_per_token - end = begin + self.indexer_bytes_per_token - return self.kv_buffer[layer_index][:, :, begin:end] + return self.indexer_k_buffer[layer_index] + + def _free_buffers(self): + self.kv_buffer = None + self.indexer_k_buffer = None + + def get_index_kv_buffer(self, index): + return {"kv_buffer": self.kv_buffer[:, index], "indexer_k_buffer": self.indexer_k_buffer[:, index]} + + def load_index_kv_buffer(self, index, load_tensor_dict): + self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) + self.indexer_k_buffer[:, index].copy_(load_tensor_dict["indexer_k_buffer"]) def get_prefill_kv_cache_and_remap_indices( self, diff --git a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py b/lightllm/common/kv_cache_mem_manager/operator/deepseek.py index 6e05b96e10..919f0016fc 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py +++ b/lightllm/common/kv_cache_mem_manager/operator/deepseek.py @@ -50,6 +50,22 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: class FP8PerTokenGroupQuantDeepseek3_2MemOperator(BaseMemManagerOperator): + def _get_cpu_cache_views(self, cpu_cache_client): + from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( + FP8PerTokenGroupQuantDeepseek3_2MemoryManager, + ) + + mem_manager: FP8PerTokenGroupQuantDeepseek3_2MemoryManager = self.mem_manager + cpu_cache_meta = cpu_cache_client.kv_cache_tensor_meta + split = mem_manager.flashmla_bytes_per_token + end = split + mem_manager.indexer_bytes_per_token + assert cpu_cache_meta.data_type is torch.uint8 + assert cpu_cache_meta.num_heads == 1 + assert cpu_cache_meta.head_dim >= end + + cpu_cache_tensor = cpu_cache_client.cpu_kv_cache_tensor + return cpu_cache_tensor[:, :, :, :, :split], cpu_cache_tensor[:, :, :, :, split:end] + def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( FP8PerTokenGroupQuantDeepseek3_2MemoryManager, @@ -78,3 +94,93 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: o_rope, ) return + + def load_cpu_cache_to_gpu(self, mem_indexes: torch.Tensor, page_indexes: torch.Tensor, cpu_cache_client, req): + assert mem_indexes.is_cuda and page_indexes.is_cuda + from lightllm.utils.envs_utils import get_env_start_args + from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size + from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( + FP8PerTokenGroupQuantDeepseek3_2MemoryManager, + ) + from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu + + args = get_env_start_args() + assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 + mem_manager: FP8PerTokenGroupQuantDeepseek3_2MemoryManager = self.mem_manager + cpu_kv_cache, cpu_indexer_k_cache = self._get_cpu_cache_views(cpu_cache_client) + + rank_in_dp = get_current_rank_in_dp() + dp_world_size = get_dp_world_size() + load_cpu_kv_to_gpu( + gpu_mem_indexes=mem_indexes, + gpu_kv_cache=mem_manager.kv_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=cpu_kv_cache, + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + tp_index=rank_in_dp, + tp_world_size=dp_world_size, + grid_num=16, + ) + load_cpu_kv_to_gpu( + gpu_mem_indexes=mem_indexes, + gpu_kv_cache=mem_manager.indexer_k_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=cpu_indexer_k_cache, + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + tp_index=rank_in_dp, + tp_world_size=dp_world_size, + grid_num=16, + ) + return + + def offload_gpu_kv_to_cpu_cache( + self, + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + page_readies: torch.Tensor, + cpu_cache_client, + req, + ): + assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda + from lightllm.utils.envs_utils import get_env_start_args + from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size + from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( + FP8PerTokenGroupQuantDeepseek3_2MemoryManager, + ) + from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu + + args = get_env_start_args() + assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 + assert len(mem_indexes) // args.cpu_cache_token_page_size == len(page_indexes) + mem_manager: FP8PerTokenGroupQuantDeepseek3_2MemoryManager = self.mem_manager + cpu_kv_cache, cpu_indexer_k_cache = self._get_cpu_cache_views(cpu_cache_client) + + rank_in_dp = get_current_rank_in_dp() + dp_world_size = get_dp_world_size() + offload_gpu_kv_to_cpu( + token_indexes=mem_indexes, + gpu_kv_cache=mem_manager.kv_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=cpu_kv_cache, + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + page_readies=page_readies, + tp_index=rank_in_dp, + tp_world_size=dp_world_size, + grid_num=16, + ) + offload_gpu_kv_to_cpu( + token_indexes=mem_indexes, + gpu_kv_cache=mem_manager.indexer_k_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=cpu_indexer_k_cache, + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + page_readies=page_readies, + tp_index=rank_in_dp, + tp_world_size=dp_world_size, + grid_num=16, + ) + return diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json new file mode 100644 index 0000000000..eaef18b7e2 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json @@ -0,0 +1,98 @@ +{ + "1": { + "BLOCK_SEQ": 16, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 4, + "num_warps": 2 + }, + "100": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 1, + "num_warps": 2 + }, + "1024": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 4, + "num_stages": 5, + "num_warps": 1 + }, + "128": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 8, + "num_stages": 2, + "num_warps": 2 + }, + "16": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 3, + "num_warps": 1 + }, + "2048": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 2, + "num_warps": 1 + }, + "256": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 8, + "num_stages": 2, + "num_warps": 1 + }, + "32": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 2, + "num_warps": 2 + }, + "4096": { + "BLOCK_SEQ": 4, + "HEAD_PARALLEL_NUM": 4, + "num_stages": 1, + "num_warps": 2 + }, + "5120": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 4, + "num_warps": 1 + }, + "6144": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 2, + "num_warps": 1 + }, + "64": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 4, + "num_warps": 1 + }, + "6400": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 2, + "num_warps": 1 + }, + "7168": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 4, + "num_warps": 1 + }, + "8": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SEQ": 2, + "HEAD_PARALLEL_NUM": 1, + "num_stages": 4, + "num_warps": 2 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json new file mode 100644 index 0000000000..4d1b132f7b --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=1,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json @@ -0,0 +1,98 @@ +{ + "1": { + "BLOCK_SEQ": 2, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 5, + "num_warps": 4 + }, + "100": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 1, + "num_warps": 2 + }, + "1024": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 4, + "num_stages": 5, + "num_warps": 1 + }, + "128": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 5, + "num_warps": 1 + }, + "256": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 8, + "num_stages": 2, + "num_warps": 2 + }, + "32": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 1, + "num_warps": 2 + }, + "4096": { + "BLOCK_SEQ": 2, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 1, + "num_warps": 2 + }, + "5120": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 4, + "num_stages": 2, + "num_warps": 1 + }, + "6144": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 5, + "num_warps": 1 + }, + "64": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 1, + "num_warps": 2 + }, + "6400": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 8, + "num_stages": 4, + "num_warps": 1 + }, + "7168": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 4, + "num_stages": 2, + "num_warps": 1 + }, + "8": { + "BLOCK_SEQ": 1, + "HEAD_PARALLEL_NUM": 16, + "num_stages": 1, + "num_warps": 8 + }, + "8192": { + "BLOCK_SEQ": 4, + "HEAD_PARALLEL_NUM": 2, + "num_stages": 5, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=12288,out_dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=12288,out_dtype=torch.bfloat16}_NVIDIA_H200.json new file mode 100644 index 0000000000..6c3a7ac834 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=12288,out_dtype=torch.bfloat16}_NVIDIA_H200.json @@ -0,0 +1,98 @@ +{ + "1": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 8 + }, + "100": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 4 + }, + "2048": { + "BLOCK_M": 64, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "4096": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "5120": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "6144": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "6400": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "7168": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "8": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "8192": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json new file mode 100644 index 0000000000..e1f78b6925 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json @@ -0,0 +1,518 @@ +{ + "1": { + "BLOCK_M": 32, + "BLOCK_N": 64, + "NUM_STAGES": 4, + "num_warps": 8 + }, + "100": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 4 + }, + "1024": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10240": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10368": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10496": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1152": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "128": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "1408": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "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": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "1664": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1792": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "17920": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18048": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "18176": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18304": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18688": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18816": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1920": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "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 + }, + "256": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 4 + }, + "2688": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "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 + }, + "3072": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 4 + }, + "33920": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3456": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "34560": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34816": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34944": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "35200": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "35328": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3584": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3712": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "384": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 4 + }, + "3840": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "3968": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "42112": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4224": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "42624": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "43008": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "43264": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "43776": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "49664": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "50944": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "5120": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "51200": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "51328": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "51456": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "52352": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "52480": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "52992": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "53120": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "53248": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "53632": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "53888": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "58368": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "58880": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "59008": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "59136": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "59520": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "59776": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "59904": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "60032": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "6144": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 4 + }, + "640": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "6400": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "66176": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "66432": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "67200": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "67328": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "67456": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "67840": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "68096": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "68480": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "7168": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "8": { + "BLOCK_M": 1, + "BLOCK_N": 64, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "8192": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "9472": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "9728": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "9984": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index f15badde25..54af0a891d 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -100,6 +100,14 @@ def all_reduce(self, input_: torch.Tensor) -> None: return return dist.all_reduce(input_, group=self.device_group) + def all_reduce_residual_rmsnorm(self, inp, residual, rms_weight, eps, alloc_func): + # Fused AR + residual-add + RMSNorm via flashinfer when the message is small enough + # for the oneshot-lamport fast path; otherwise return None so the caller falls back to + # a plain all_reduce + a separate (fused_add) rmsnorm. + if self.flashinfer_reduce is not None and self.flashinfer_reduce.should_use(inp): + return self.flashinfer_reduce.allreduce_residual_rmsnorm(inp, residual, rms_weight, eps, alloc_func) + return None + def all_gather_into_tensor(self, output_: torch.Tensor, input_: torch.Tensor, async_op: bool = False) -> None: return dist.all_gather_into_tensor(output_, input_, group=self.device_group, async_op=async_op) @@ -235,6 +243,26 @@ def all_reduce( return dist.all_reduce(input_, op, group, async_op) +def all_reduce_residual_rmsnorm( + inp: torch.Tensor, + residual: torch.Tensor, + rms_weight: torch.Tensor, + eps: float, + group: Optional[Union[ProcessGroup, CustomProcessGroup]], + alloc_func, +): + """Fused all-reduce + residual-add + RMSNorm (SGLang #22390). + + Returns ``(norm_out, residual_out)`` when a fused fast path (flashinfer) is available, + otherwise ``None`` so the caller can fall back to ``all_reduce`` + a separate + (fused-add) RMSNorm. ``inp`` is the un-reduced tensor; ``residual`` is added after the + reduction. + """ + if isinstance(group, CustomProcessGroup): + return group.all_reduce_residual_rmsnorm(inp, residual, rms_weight, eps, alloc_func) + return None + + def all_gather_into_tensor( output_: torch.Tensor, input_: torch.Tensor, diff --git a/lightllm/distributed/flashinfer_all_reduce.py b/lightllm/distributed/flashinfer_all_reduce.py index 35cb0aaedf..6afbba4e07 100644 --- a/lightllm/distributed/flashinfer_all_reduce.py +++ b/lightllm/distributed/flashinfer_all_reduce.py @@ -134,3 +134,26 @@ def all_reduce(self, inp: torch.Tensor) -> torch.Tensor: pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce, # launch_with_pdl=True, # TODO: learn pdl and ensure no other side effects. ) + + def allreduce_residual_rmsnorm(self, inp, residual, rms_weight, eps, alloc_func): + """Fused all-reduce + residual-add + RMSNorm (flashinfer kARResidualRMSNorm). + + Computes ``residual_out = residual + allreduce(inp)`` and + ``norm_out = rmsnorm(residual_out) * rms_weight`` in one kernel — the SGLang + #22390 fusion. Returns ``(norm_out, residual_out)``; both are freshly allocated + (the kernel is out-of-place). ``inp`` must already satisfy ``should_use``. + """ + norm_out = alloc_func(inp.shape, dtype=inp.dtype, device=inp.device) + residual_out = alloc_func(residual.shape, dtype=residual.dtype, device=residual.device) + flashinfer_comm.allreduce_fusion( + input=inp, + workspace=self._workspace, + pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + rms_gamma=rms_weight, + rms_eps=eps, + fp32_acc=True, + ) + return norm_out, residual_out diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py index f619b1d88f..56cd13bb3f 100644 --- a/lightllm/models/__init__.py +++ b/lightllm/models/__init__.py @@ -21,6 +21,7 @@ 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_2.model import Glm5_2TpPartModel 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 be819c94a0..c7ffb31cc4 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -10,6 +10,7 @@ from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe from functools import partial from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale +from lightllm.distributed.communication_op import all_reduce_residual_rmsnorm from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_global_world_size from lightllm.utils.log_utils import init_logger @@ -50,6 +51,13 @@ def __init__(self, layer_num, network_config): mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim) self.softmax_scale = self.softmax_scale * mscale * mscale self.enable_cc_method = not os.getenv("DISABLE_CC_METHOD", "False").upper() in ["ON", "TRUE", "1"] + # Fuse the post-attention residual add into the ffn RMSNorm (one Triton launch instead + # of a separate add_ + rmsnorm). Bit-identical; gate exists only for A/B measurement. + self.enable_fused_add_norm = os.environ.get("LIGHTLLM_FUSED_ADD_RMSNORM", "1") == "1" + # Additionally fold the attention-output all-reduce into that residual-add + RMSNorm via + # flashinfer kARResidualRMSNorm (SGLang #22390). Only fires when flashinfer AR is the + # active backend (small messages / low concurrency); falls back otherwise. + self.enable_fused_ar_norm = os.environ.get("LIGHTLLM_FUSED_AR_RMSNORM", "1") == "1" super().__init__(layer_num, network_config) self.num_heads = network_config["num_attention_heads"] self.num_kv_heads = network_config["num_key_value_heads"] @@ -200,7 +208,11 @@ def _get_qkv( return q, cache_kv def _get_o( - self, input: torch.Tensor, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight + self, + input: torch.Tensor, + infer_state: Deepseek2InferStateInfo, + layer_weight: Deepseek2TransformerLayerWeight, + reduce: bool = True, ) -> torch.Tensor: if infer_state.need_dp_prefill_balance: input = infer_state._all_to_all_balance_get(data=input) @@ -208,7 +220,10 @@ def _get_o( if input.shape[2] == self.kv_lora_rank: input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) o_tensor = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) - o_tensor = self._tpsp_reduce(input=o_tensor, infer_state=infer_state) + # reduce=False leaves o un-reduced so the caller can fuse the all-reduce into the + # following residual-add + RMSNorm (flashinfer kARResidualRMSNorm). + if reduce: + o_tensor = self._tpsp_reduce(input=o_tensor, infer_state=infer_state) return o_tensor def _moe_ffn_tp( @@ -288,6 +303,73 @@ def _ffn_ep_impl( return ffn2_out + def _fused_add_ffn_norm(self, input_embdings: torch.Tensor, o: torch.Tensor, infer_state, layer_weight): + # Fuse the post-attention residual add (input_embdings += o) into the following ffn + # RMSNorm in a single Triton launch — eliminates one tiny elementwise-add kernel per + # layer. Bit-identical to `input_embdings.add_(o); self._ffn_norm(input_embdings)`. + if self.enable_fused_add_norm: + return layer_weight.ffn_norm_weight_.fused_add_forward( + residual=input_embdings.view(-1, self.embed_dim_), + x=o.view(-1, self.embed_dim_), + eps=self.eps_, + alloc_func=self.alloc_tensor, + ) + input_embdings.add_(o.view(-1, self.embed_dim_)) + return self._ffn_norm(input_embdings, infer_state, layer_weight) + + def _attn_out_add_ffn_norm(self, o_attn, input_embdings, infer_state, layer_weight): + """Combine the attention-output all-reduce, the residual add, and the ffn RMSNorm. + + Fast path (flashinfer kARResidualRMSNorm): all three fold into one kernel; ``o`` is kept + un-reduced and the reduction happens inside the fused op. Returns the new (normed_input, + residual). Falls back to the standard all-reduce + (fused-add) RMSNorm when flashinfer + AR is not the active backend (large messages / SP mode / disabled). + """ + if self.enable_fused_ar_norm and self.tp_world_size_ > 1 and not get_env_start_args().enable_tpsp_mix_mode: + o = self._get_o(o_attn, infer_state, layer_weight, reduce=False).view(-1, self.embed_dim_) + fused = all_reduce_residual_rmsnorm( + inp=o, + residual=input_embdings.view(-1, self.embed_dim_), + rms_weight=layer_weight.ffn_norm_weight_.weight, + eps=self.eps_, + group=infer_state.dist_group, + alloc_func=self.alloc_tensor, + ) + if fused is not None: + norm_out, residual_out = fused + return norm_out, residual_out + # flashinfer not applicable for this message size: finish the all-reduce normally. + o = self._tpsp_reduce(input=o, infer_state=infer_state) + return self._fused_add_ffn_norm(input_embdings, o, infer_state, layer_weight), input_embdings + + o = self._get_o(o_attn, infer_state, layer_weight, reduce=True) + return self._fused_add_ffn_norm(input_embdings, o, infer_state, layer_weight), input_embdings + + def context_forward(self, input_embdings, infer_state: Deepseek2InferStateInfo, layer_weight): + input1 = self._att_norm(input_embdings, infer_state, layer_weight) + o = self.context_attention_forward(input1, infer_state, layer_weight) + input1 = self._fused_add_ffn_norm(input_embdings, o, infer_state, layer_weight) + o = None + ffn_out = self._ffn(input1, infer_state, layer_weight) + input1 = None + input_embdings.add_(ffn_out.view(-1, self.embed_dim_)) + return input_embdings + + def token_forward(self, input_embdings, infer_state: Deepseek2InferStateInfo, layer_weight): + input1 = self._att_norm(input_embdings, infer_state, layer_weight) + # Inline the decode attention so the output projection stays un-reduced and its + # all-reduce can fold into the residual-add + ffn RMSNorm (see _attn_out_add_ffn_norm). + q, cache_kv = self._get_qkv(input1, infer_state, layer_weight) + self._post_cache_kv(cache_kv, infer_state, layer_weight) + o = self._token_attention_kernel(q, infer_state, layer_weight) + input1 = None + input1, input_embdings = self._attn_out_add_ffn_norm(o, input_embdings, infer_state, layer_weight) + o = None + ffn_out = self._ffn(input1, infer_state, layer_weight) + input1 = None + input_embdings.add_(ffn_out.view(-1, self.embed_dim_)) + return input_embdings + def overlap_tpsp_token_forward( self, input_embdings: torch.Tensor, 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..b9d64ff703 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +import os import torch from typing import Any from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo @@ -12,6 +13,8 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import all_gather_into_tensor +_ENABLE_PAGED_INDEXER_LOGITS = os.environ.get("LIGHTLLM_PAGED_INDEXER_LOGITS", "1") == "1" + class Deepseek3_2TransformerLayerInfer(Deepseek2TransformerLayerInfer): def __init__(self, layer_num, network_config): @@ -120,7 +123,7 @@ def _token_attention_kernel( # 计算 topk mem indices att_state = infer_state.decode_att_state - topk_mem_indices, _ = self.indexer._get_indices( + topk_mem_indices, topk_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, @@ -135,6 +138,7 @@ def _token_attention_kernel( nsa_decode_dict={ "layer_index": self.layer_num_, "topk_mem_indices": topk_mem_indices, + "topk_indices": topk_indices, "softmax_scale": self.softmax_scale, "kv_lora_rank": self.kv_lora_rank, "qk_rope_head_dim": self.qk_rope_head_dim, @@ -160,13 +164,15 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.qk_rope_head_dim = network_config["qk_rope_head_dim"] self.index_head_dim = network_config["index_head_dim"] self.eps = network_config["rms_norm_eps"] - self.block_size = network_config["quantization_config"]["weight_block_size"][0] - self.scale_fmt = network_config["quantization_config"]["scale_fmt"] + quantization_config = network_config.get("quantization_config") or {} + self.block_size = quantization_config.get("weight_block_size", [128, 128])[0] + self.scale_fmt = quantization_config.get("scale_fmt", "ue8m0") 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.indexer_rope_interleave = network_config.get("indexer_rope_interleave", False) def _get_indices( self, @@ -176,12 +182,11 @@ 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) 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, ) @@ -214,28 +219,44 @@ def _get_indices( 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, - ) - import deep_gemm + if not infer_state.is_prefill and _ENABLE_PAGED_INDEXER_LOGITS: + # Decode: compute logits directly on the paged indexer-K pool through + # ragged_mem_index — skips materializing a dense [batch * max_kv_seq_len] K copy. + from ..triton_kernel.fp8_mqa_logits import fp8_paged_mqa_logits + + logits = fp8_paged_mqa_logits( + q_fp8=q_fp8, + indexer_k_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + weights=weights.squeeze(-1).float(), + ragged_mem_index=att_state.ragged_mem_index, + cu_seqlen_ks=ks, + cu_seqlen_ke=ke, + max_kv_seq_len=infer_state.max_kv_seq_len, + ) + else: + # 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, + ) + + 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, - ) + 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 @@ -262,7 +283,7 @@ 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) + return hadamard_transform(x, scale=hidden_size**-0.5) def _get_q_k_bf16( self, @@ -276,8 +297,11 @@ def _get_q_k_bf16( k = layer_weight.k_norm_(k, eps=self.eps) - # 为什么 indexer 和主模型用的q k 的 rotary的排布方式不一样，这不是脱裤子放屁麻。 - from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd + if self.indexer_rope_interleave: + from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd + else: + # DeepSeek-V3.2 indexer RoPE uses the non-interleaved layout. + from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd rotary_emb_fwd( q[:, :, : self.qk_rope_head_dim], diff --git a/lightllm/models/deepseek3_2/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek3_2/layer_weights/transformer_layer_weight.py index c7c71c8273..86f67f6aa4 100644 --- a/lightllm/models/deepseek3_2/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/deepseek3_2/layer_weights/transformer_layer_weight.py @@ -19,7 +19,6 @@ def _init_weight(self): self._init_indexer_weight() def _init_indexer_weight(self): - prefix = f"model.layers.{self.layer_num_}.self_attn.indexer" assert self.index_n_heads % self.tp_world_size_ == 0 @@ -28,7 +27,7 @@ def _init_indexer_weight(self): 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, + quant_method=self.get_quant_method("indexer_wq_b"), tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, ) @@ -37,7 +36,7 @@ def _init_indexer_weight(self): out_dims=[self.index_head_dim], weight_names=f"{prefix}.wk.weight", data_type=self.data_type_, - quant_method=None, + quant_method=self.get_quant_method("indexer_wk"), tp_rank=0, tp_world_size=1, ) diff --git a/lightllm/models/deepseek3_2/triton_kernel/fp8_mqa_logits.py b/lightllm/models/deepseek3_2/triton_kernel/fp8_mqa_logits.py index 1c1f72b7d7..2fb38212a9 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/fp8_mqa_logits.py +++ b/lightllm/models/deepseek3_2/triton_kernel/fp8_mqa_logits.py @@ -1,141 +1,114 @@ +import torch import triton import triton.language as tl -import torch @triton.jit def _fp8_paged_mqa_logits_kernel( Q_ptr, - KV_ptr, - KVScale_ptr, - Weights_ptr, - MemIndex_ptr, - CuSeqlenKs_ptr, - CuSeqlenKe_ptr, - Output_ptr, - seq_len, - seq_len_kv, - num_heads, - head_dim, stride_q_seq, stride_q_head, - stride_q_dim, - stride_kv_pool, - stride_kv_dim, + K_ptr, + stride_k_pool, + KScale_ptr, + stride_kscale_pool, + Weights_ptr, stride_w_seq, - stride_w_head, + RaggedMemIndex_ptr, + CuSeqlenKs_ptr, + CuSeqlenKe_ptr, + Out_ptr, stride_o_seq, - stride_o_kv, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, + HEAD_NUM: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, ): + cur_q_index = tl.program_id(0) + block_n_index = tl.program_id(1) - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - - start_m = pid_m * BLOCK_SIZE_M - start_n = pid_n * BLOCK_SIZE_N - - offs_m = start_m + tl.arange(0, BLOCK_SIZE_M) - offs_n = start_n + tl.arange(0, BLOCK_SIZE_N) + ks = tl.load(CuSeqlenKs_ptr + cur_q_index) + ke = tl.load(CuSeqlenKe_ptr + cur_q_index) + kv_len = ke - ks + start_n = block_n_index * BLOCK_N + if start_n >= kv_len: + return - logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + offs_n = start_n + tl.arange(0, BLOCK_N) + mask_n = offs_n < kv_len + offs_h = tl.arange(0, HEAD_NUM) + offs_d = tl.arange(0, HEAD_DIM) - mask_m = offs_m < seq_len - mask_n = offs_n < seq_len_kv + mem_index = tl.load(RaggedMemIndex_ptr + ks + offs_n, mask=mask_n, other=0) - mem_indices = tl.load(MemIndex_ptr + offs_n, mask=mask_n, other=0) + # q: [HEAD_NUM, HEAD_DIM] fp8, k: [BLOCK_N, HEAD_DIM] fp8 gathered from the kv pool. + q = tl.load(Q_ptr + cur_q_index * stride_q_seq + offs_h[:, None] * stride_q_head + offs_d[None, :]) + k = tl.load(K_ptr + mem_index[:, None] * stride_k_pool + offs_d[None, :], mask=mask_n[:, None], other=0.0) - scales = tl.load(KVScale_ptr + mem_indices, mask=mask_n, other=1.0) + # relu commutes with the (positive) fp8 scales: q_scale is folded into weights by the + # caller, k_scale is applied on the reduced column below. + logits = tl.dot(q, tl.trans(k)) # [HEAD_NUM, BLOCK_N] fp32 + logits = tl.maximum(logits, 0.0) - for h in range(num_heads): - weights = tl.load(Weights_ptr + offs_m * stride_w_seq + h * stride_w_head, mask=mask_m, other=0.0) - score = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + weights = tl.load(Weights_ptr + cur_q_index * stride_w_seq + offs_h) # [HEAD_NUM] + k_scale = tl.load(KScale_ptr + mem_index * stride_kscale_pool, mask=mask_n, other=0.0) # [BLOCK_N] + out = tl.sum(logits * weights[:, None], axis=0) * k_scale - for d_block in range(tl.cdiv(head_dim, BLOCK_SIZE_D)): - d_start = d_block * BLOCK_SIZE_D - offs_d = d_start + tl.arange(0, BLOCK_SIZE_D) - mask_d = offs_d < head_dim - - q_ptrs = Q_ptr + offs_m[:, None] * stride_q_seq + h * stride_q_head + offs_d[None, :] * stride_q_dim - mask_q = (offs_m[:, None] < seq_len) & mask_d[None, :] - q = tl.load(q_ptrs, mask=mask_q, other=0.0).to(tl.float32) - - k_ptrs = KV_ptr + mem_indices[:, None] * stride_kv_pool + offs_d[None, :] * stride_kv_dim - mask_k = mask_n[:, None] & mask_d[None, :] - k = tl.load(k_ptrs, mask=mask_k, other=0.0).to(tl.float32) - - k = k * scales[:, None] - - score += tl.dot(q, tl.trans(k)) - score = tl.maximum(score, 0.0) - logits += score * weights[:, None] - - mask_ks = tl.load(CuSeqlenKs_ptr + offs_m, mask=mask_m, other=0) - mask_ke = tl.load(CuSeqlenKe_ptr + offs_m, mask=mask_m, other=seq_len_kv) - - mask_lo = offs_n[None, :] >= mask_ks[:, None] - mask_hi = offs_n[None, :] < mask_ke[:, None] - mask_valid = mask_lo & mask_hi & mask_m[:, None] & mask_n[None, :] - - logits = tl.where(mask_valid, logits, float("-inf")) - - # Store output - out_ptrs = Output_ptr + offs_m[:, None] * stride_o_seq + offs_n[None, :] * stride_o_kv - mask_out = (offs_m[:, None] < seq_len) & (offs_n[None, :] < seq_len_kv) - tl.store(out_ptrs, logits, mask=mask_out) + # row-compacted layout (matches deep_gemm.fp8_mqa_logits): row m's keys land at columns [0, ke-ks) + tl.store(Out_ptr + cur_q_index * stride_o_seq + offs_n, out, mask=mask_n) +@torch.no_grad() def fp8_paged_mqa_logits( - q: torch.Tensor, - kv: torch.Tensor, - kv_scale: torch.Tensor, + q_fp8: torch.Tensor, + indexer_k_buffer: torch.Tensor, weights: torch.Tensor, - mem_index: torch.Tensor, + ragged_mem_index: torch.Tensor, cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, - out: torch.Tensor = None, + max_kv_seq_len: int, ) -> torch.Tensor: - seq_len, num_heads, head_dim = q.shape - seq_len_kv = mem_index.shape[0] + """Indexer logits over the paged (token-granular) indexer-K pool. - if out is None: - output = torch.empty((seq_len, seq_len_kv), device=q.device, dtype=torch.float32) - else: - output = out + Row m holds k_scale_j * sum_h(weights[m, h] * relu(q[m, h] @ k_j)) for the ragged keys + j in [cu_seqlen_ks[m], cu_seqlen_ke[m]), row-compacted at columns [0, ke-ks) — the same + output contract as deep_gemm.fp8_mqa_logits(clean_logits=False) on the densely extracted + K, but reading K directly from the kv pool through ragged_mem_index. - BLOCK_SIZE_M = 16 - BLOCK_SIZE_N = 64 - BLOCK_SIZE_D = 128 + q_fp8: [q_token_num, head_num, 128] float8_e4m3fn + indexer_k_buffer: [pool_size, 1, 132] uint8 (128B fp8 K + 4B fp32 scale per token) + weights: [q_token_num, head_num] float32 (q_scale folded in) + ragged_mem_index / cu_seqlen_ks / cu_seqlen_ke: from gen_nsa_ks_ke + """ + q_token_num, head_num, head_dim = q_fp8.shape + assert head_dim == 128 and head_num in (16, 32, 64, 128) + assert indexer_k_buffer.dtype == torch.uint8 and indexer_k_buffer.shape[2] == 132 - grid = (triton.cdiv(seq_len, BLOCK_SIZE_M), triton.cdiv(seq_len_kv, BLOCK_SIZE_N)) + k_fp8 = indexer_k_buffer[:, 0, 0:128].view(dtype=torch.float8_e4m3fn) + k_scale = indexer_k_buffer[:, 0, 128:132].view(dtype=torch.float32) + logits = torch.empty((q_token_num, max_kv_seq_len), dtype=torch.float32, device=q_fp8.device) + + BLOCK_N = 128 + grid = (q_token_num, triton.cdiv(max_kv_seq_len, BLOCK_N)) _fp8_paged_mqa_logits_kernel[grid]( - q, - kv, - kv_scale, - weights, - mem_index, - cu_seqlen_ks, - cu_seqlen_ke, - output, - seq_len, - seq_len_kv, - num_heads, - head_dim, - q.stride(0), - q.stride(1), - q.stride(2), - kv.stride(0), - kv.stride(1), - weights.stride(0), - weights.stride(1), - output.stride(0), - output.stride(1), - BLOCK_SIZE_M=BLOCK_SIZE_M, - BLOCK_SIZE_N=BLOCK_SIZE_N, - BLOCK_SIZE_D=BLOCK_SIZE_D, + q_fp8, + stride_q_seq=q_fp8.stride(0), + stride_q_head=q_fp8.stride(1), + K_ptr=k_fp8, + stride_k_pool=k_fp8.stride(0), + KScale_ptr=k_scale, + stride_kscale_pool=k_scale.stride(0), + Weights_ptr=weights, + stride_w_seq=weights.stride(0), + RaggedMemIndex_ptr=ragged_mem_index, + CuSeqlenKs_ptr=cu_seqlen_ks, + CuSeqlenKe_ptr=cu_seqlen_ke, + Out_ptr=logits, + stride_o_seq=logits.stride(0), + HEAD_NUM=head_num, + HEAD_DIM=head_dim, + BLOCK_N=BLOCK_N, + num_warps=4, + num_stages=3, ) - - return output + return logits diff --git a/lightllm/models/glm5_2/__init__.py b/lightllm/models/glm5_2/__init__.py new file mode 100644 index 0000000000..bf95461887 --- /dev/null +++ b/lightllm/models/glm5_2/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.glm5_2.model import Glm5_2TpPartModel + +__all__ = ["Glm5_2TpPartModel"] diff --git a/lightllm/models/glm5_2/indexshare.py b/lightllm/models/glm5_2/indexshare.py new file mode 100644 index 0000000000..617e1f650c --- /dev/null +++ b/lightllm/models/glm5_2/indexshare.py @@ -0,0 +1,12 @@ +def owns_indexer_layer(layer_id: int, config: dict) -> bool: + num_hidden_layers = config.get("num_hidden_layers") + if num_hidden_layers is not None and layer_id >= num_hidden_layers: + return True + + pattern = config.get("index_topk_pattern") + if pattern is not None and 0 <= layer_id < len(pattern): + return pattern[layer_id] != "S" + + freq = config.get("index_topk_freq", 1) + offset = config.get("index_skip_topk_offset", 2) + return max(layer_id - offset + 1, 0) % freq == 0 diff --git a/lightllm/models/glm5_2/layer_infer/__init__.py b/lightllm/models/glm5_2/layer_infer/__init__.py new file mode 100644 index 0000000000..1fa3327f6a --- /dev/null +++ b/lightllm/models/glm5_2/layer_infer/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.glm5_2.layer_infer.transformer_layer_infer import Glm5_2TransformerLayerInfer + +__all__ = ["Glm5_2TransformerLayerInfer"] diff --git a/lightllm/models/glm5_2/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_2/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..00cbad6c1f --- /dev/null +++ b/lightllm/models/glm5_2/layer_infer/transformer_layer_infer.py @@ -0,0 +1,101 @@ +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer +from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer, NsaInfer +from lightllm.models.glm5_2.indexshare import owns_indexer_layer + + +class Glm5_2TransformerLayerInfer(Deepseek3_2TransformerLayerInfer): + def __init__(self, layer_num, network_config): + self.has_indexer = owns_indexer_layer(layer_num, network_config) + Deepseek2TransformerLayerInfer.__init__(self, layer_num, network_config) + self.indexer = ( + NsaInfer(layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_) + if self.has_indexer + else None + ) + + def _get_or_reuse_topk_indices(self, infer_state, att_state, layer_weight): + if getattr(infer_state, "glm5_2_reuse_mtp_topk_indices", False): + model_input = getattr(infer_state, "glm5_2_model_input", None) + cached_topk = getattr(model_input, "glm5_2_mtp_topk_cache", None) + if cached_topk is not None: + infer_state.glm5_2_indexshare_topk_cache = cached_topk + return cached_topk + + if self.indexer is not None: + topk_mem_indices, topk_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=att_state, + layer_weight=layer_weight, + ) + infer_state.glm5_2_indexshare_topk_cache = (topk_mem_indices, topk_indices) + if getattr(infer_state, "glm5_2_reuse_mtp_topk_indices", False): + model_input = getattr(infer_state, "glm5_2_model_input", None) + if model_input is not None: + model_input.glm5_2_mtp_topk_cache = infer_state.glm5_2_indexshare_topk_cache + return topk_mem_indices, topk_indices + + if not hasattr(infer_state, "glm5_2_indexshare_topk_cache"): + raise RuntimeError( + f"GLM-5.2 layer {self.layer_num_} needs cached IndexShare top-k indices, " + "but no previous indexer layer has produced them." + ) + return infer_state.glm5_2_indexshare_topk_cache + + def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): + 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) + q_all = q_nope if self.qk_rope_head_dim == 0 else torch.cat([q_nope, q_rope], dim=-1) + + att_state = infer_state.prefill_att_state + topk_mem_indices, topk_indices = self._get_or_reuse_topk_indices(infer_state, att_state, layer_weight) + del infer_state.get_topk_indices_params + + att_control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_mem_indices": topk_mem_indices, + "topk_indices": topk_indices, + "prefill_cache_kv": kv, + "softmax_scale": self.softmax_scale, + "kv_lora_rank": self.kv_lora_rank, + }, + ) + + return infer_state.prefill_att_state.prefill_att( + q=q_all, + k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + v=None, + att_control=att_control, + ) + + def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): + 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) + + att_state = infer_state.decode_att_state + topk_mem_indices, topk_indices = self._get_or_reuse_topk_indices(infer_state, att_state, layer_weight) + del infer_state.get_topk_indices_params + + att_control = AttControl( + nsa_decode=True, + nsa_decode_dict={ + "layer_index": self.layer_num_, + "topk_mem_indices": topk_mem_indices, + "topk_indices": topk_indices, + "softmax_scale": self.softmax_scale, + "kv_lora_rank": self.kv_lora_rank, + "qk_rope_head_dim": self.qk_rope_head_dim, + }, + ) + + 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=att_control, + ) diff --git a/lightllm/models/glm5_2/layer_weights/__init__.py b/lightllm/models/glm5_2/layer_weights/__init__.py new file mode 100644 index 0000000000..01013d5657 --- /dev/null +++ b/lightllm/models/glm5_2/layer_weights/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.glm5_2.layer_weights.transformer_layer_weight import Glm5_2TransformerLayerWeight + +__all__ = ["Glm5_2TransformerLayerWeight"] diff --git a/lightllm/models/glm5_2/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_2/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..066c551b1b --- /dev/null +++ b/lightllm/models/glm5_2/layer_weights/transformer_layer_weight.py @@ -0,0 +1,46 @@ +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import FusedMoeWeight, ROWMMWeight +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 lightllm.models.glm5_2.indexshare import owns_indexer_layer + + +class Glm5_2TransformerLayerWeight(Deepseek3_2TransformerLayerWeight): + def _parse_config(self): + super()._parse_config() + self.has_indexer = owns_indexer_layer(self.layer_num_, self.network_config_) + + def _init_weight(self): + Deepseek2TransformerLayerWeight._init_weight(self) + if self.has_indexer: + self._init_indexer_weight() + + def _init_moe(self): + if self.num_fused_shared_experts == 0: + self._load_mlp(f"model.layers.{self.layer_num_}.mlp.shared_experts", is_shared_experts=True) + + self.moe_gate = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.n_routed_experts], + weight_names=f"model.layers.{self.layer_num_}.mlp.gate.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.experts = FusedMoeWeight( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + e_score_correction_bias_name=self.e_score_correction_bias_name, + weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts", + n_routed_experts=self.n_routed_experts, + hidden_size=self.n_embed, + moe_intermediate_size=self.network_config_["moe_intermediate_size"], + data_type=self.data_type_, + quant_method=self.quant_cfg.get_quant_method(self.layer_num_, "fused_moe"), + num_fused_shared_experts=self.num_fused_shared_experts, + layer_num=self.layer_num_, + network_config=self.network_config_, + ) diff --git a/lightllm/models/glm5_2/model.py b/lightllm/models/glm5_2/model.py new file mode 100644 index 0000000000..10c1bb3f73 --- /dev/null +++ b/lightllm/models/glm5_2/model.py @@ -0,0 +1,54 @@ +import torch + +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel +from lightllm.models.glm5_2.layer_infer.transformer_layer_infer import Glm5_2TransformerLayerInfer +from lightllm.models.glm5_2.layer_weights.transformer_layer_weight import Glm5_2TransformerLayerWeight +from lightllm.models.registry import ModelRegistry + + +@ModelRegistry("glm_moe_dsa") +class Glm5_2TpPartModel(Deepseek3_2TpPartModel): + transformer_weight_class = Glm5_2TransformerLayerWeight + transformer_layer_infer_class = Glm5_2TransformerLayerInfer + + def _init_config(self): + super()._init_config() + if "scoring_func" not in self.config: + self.config["scoring_func"] = "sigmoid" + if self.config.get("rope_theta") is None: + self.config["rope_theta"] = self.config.get("rope_parameters", {}).get("rope_theta", 1000000.0) + + def _init_custom(self): + self._init_glm5_2_rotary() + dist_group_manager.new_deepep_group( + self.config["n_routed_experts"], + self.config["hidden_size"], + self.config.get("num_experts_per_tok", 1), + self.config.get("moe_intermediate_size", self.config.get("intermediate_size")), + ) + + def _create_inferstate(self, model_input, microbatch_index: int = 0): + infer_state = super()._create_inferstate(model_input, microbatch_index=microbatch_index) + infer_state.glm5_2_model_input = model_input + infer_state.glm5_2_reuse_mtp_topk_indices = ( + getattr(self, "is_mtp_draft_model", False) + and self.config.get("index_share_for_mtp_iteration", False) + and not model_input.is_prefill + ) + return infer_state + + def _init_glm5_2_rotary(self): + rope_theta = self.config.get("rope_theta", 1000000.0) + qk_rope_head_dim = self.config.get("qk_rope_head_dim", 64) + max_position_embeddings = self.config.get("max_position_embeddings", 1048576) + + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, qk_rope_head_dim, 2, device="cpu", dtype=torch.float32) / qk_rope_head_dim) + ) + max_seq_len = max(max_position_embeddings, self.max_seq_length) + t = torch.arange(max_seq_len, device="cpu", dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + + self._cos_cached = torch.cos(freqs).to(self.data_type).cuda() + self._sin_cached = torch.sin(freqs).to(self.data_type).cuda() diff --git a/lightllm/models/glm5_2_mtp/__init__.py b/lightllm/models/glm5_2_mtp/__init__.py new file mode 100644 index 0000000000..8798b32f3a --- /dev/null +++ b/lightllm/models/glm5_2_mtp/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.glm5_2_mtp.model import Glm5_2MTPModel + +__all__ = ["Glm5_2MTPModel"] diff --git a/lightllm/models/glm5_2_mtp/model.py b/lightllm/models/glm5_2_mtp/model.py new file mode 100644 index 0000000000..e1ea4a383e --- /dev/null +++ b/lightllm/models/glm5_2_mtp/model.py @@ -0,0 +1,98 @@ +from typing import List + +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.common.basemodel.basemodel import load_hf_weights +from lightllm.models.deepseek_mtp.layer_infer.pre_layer_infer import Deepseek3MTPPreLayerInfer +from lightllm.models.glm4_moe_lite_mtp.layer_weights.pre_and_post_layer_weight import ( + Glm4MoeLiteMTPPreAndPostLayerWeight, +) +from lightllm.models.glm5_2.model import Glm5_2TpPartModel + + +class Glm5_2MTPModel(Glm5_2TpPartModel): + is_mtp_draft_model = True + + pre_and_post_weight_class = Glm4MoeLiteMTPPreAndPostLayerWeight + pre_layer_infer_class = Deepseek3MTPPreLayerInfer + + def __init__(self, kvargs: dict): + self._pre_init(kvargs) + super().__init__(kvargs) + + def _pre_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") + + 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_start = self.config["num_hidden_layers"] + num_mtp_layers = self.config.get("num_nextn_predict_layers", 1) + + 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( + i, + self.data_type, + network_config=self.config, + quant_cfg=self.quant_cfg, + ) + for i in range(mtp_layer_start, mtp_layer_start + num_mtp_layers) + ] + + load_hf_weights( + self.data_type, + weight_dir=self.weight_dir_, + pre_post_layer=self.pre_post_weight, + transformer_layer_list=self.trans_layers_weight, + weight_dict=self.weight_dict, + ) + + self.pre_post_weight.verify_load() + [weight.verify_load() for weight in self.trans_layers_weight] + + 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 _load_hf_weights(self): + # GLM-5.2 MTP only loads the nextn layer in _init_weights(); avoid the + # base-class second pass over the same tensors, which creates large + # temporary CUDA buffers for FP8 dequantization during startup. + return + + 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) + + total_pre_layers_num = len(self.main_model.layers_infer) + total_pre_layers_num += sum( + [len(previous_model.layers_infer) for previous_model in self.mtp_previous_draft_models] + ) + + num_mtp_layers = self.config.get("num_nextn_predict_layers", 1) + self.layers_infer = [ + self.transformer_layer_infer_class(i, network_config=self.config) + for i in range(total_pre_layers_num, total_pre_layers_num + num_mtp_layers) + ] + + def _init_some_value(self): + super()._init_some_value() + self.layers_num = self.config.get("num_nextn_predict_layers", 1) + + def autotune_layers(self): + return self.config.get("num_nextn_predict_layers", 1) diff --git a/lightllm/server/build_prompt.py b/lightllm/server/build_prompt.py index 54d22a0d0d..16304a5b74 100644 --- a/lightllm/server/build_prompt.py +++ b/lightllm/server/build_prompt.py @@ -10,6 +10,13 @@ tokenizer = None +def _set_chat_template(chat_template_str: str) -> None: + if hasattr(tokenizer, "tokenizer"): + tokenizer.tokenizer.chat_template = chat_template_str + else: + tokenizer.chat_template = chat_template_str + + def init_tokenizer(args): global tokenizer @@ -17,30 +24,29 @@ def init_tokenizer(args): chat_path = args.chat_template if chat_path is not None: with open(chat_path, "r", encoding="utf-8") as f: - chat_template_str = f.read() - if hasattr(tokenizer, "tokenizer"): - tokenizer.tokenizer.chat_template = chat_template_str - else: - tokenizer.chat_template = chat_template_str + _set_chat_template(f.read()) return + default_jinja_chat_template_path = os.path.join(args.model_dir, "chat_template.jinja") + if os.path.exists(default_jinja_chat_template_path): + try: + with open(default_jinja_chat_template_path, "r", encoding="utf-8") as f: + _set_chat_template(f.read()) + logger.info(f"Loaded chat_template.jinja from {default_jinja_chat_template_path}") + return + except Exception as e: + logger.warning(f"Failed to load chat_template.jinja from {default_jinja_chat_template_path}: {e}") + # 如果 tokenizer 目录下存在chat_template.json， 同时不存在 chat_template.jinja, # 则加载其并赋值给tokenizer 的 chat_template 对象。 - if not os.path.exists(os.path.join(args.model_dir, "chat_template.jinja")) and os.path.exists( - os.path.join(args.model_dir, "chat_template.json") - ): + if os.path.exists(os.path.join(args.model_dir, "chat_template.json")): default_chat_template_path = os.path.join(args.model_dir, "chat_template.json") try: with open(default_chat_template_path, "r", encoding="utf-8") as f: template_data = json.load(f) if "chat_template" in template_data: # Set it directly on the tokenizer object so apply_chat_template can use it - if hasattr(tokenizer, "tokenizer"): - # 多模态 tokenizer - tokenizer.tokenizer.chat_template = template_data["chat_template"] - else: - tokenizer.chat_template = template_data["chat_template"] - + _set_chat_template(template_data["chat_template"]) logger.info(f"Loaded chat_template.json from {default_chat_template_path}") except Exception as e: logger.warning(f"Failed to load chat_template.json from {default_chat_template_path}: {e}") 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 a65dfb1bbb..f931f87f19 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -45,6 +45,7 @@ from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel from lightllm.models.mistral_mtp.model import MistralMTPModel from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel +from lightllm.models.glm5_2_mtp.model import Glm5_2MTPModel from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token from lightllm.server.pd_io_struct import PDChunckedTransTaskRet @@ -342,6 +343,9 @@ def init_mtp_draft_model(self, main_kvargs: dict): elif mtp_model_cfg["model_type"] == "glm4_moe_lite": assert self.args.mtp_mode in ["vanilla_with_att", "eagle_with_att"] self.draft_models.append(Glm4MoeLiteMTPModel(mtp_model_kvargs)) + elif mtp_model_cfg["model_type"] == "glm_moe_dsa": + assert self.args.mtp_mode in ["vanilla_with_att", "eagle_with_att"] + self.draft_models.append(Glm5_2MTPModel(mtp_model_kvargs)) else: raise ValueError(f"Unsupported MTP model type: {model_type}") @@ -584,7 +588,6 @@ def _get_classed_reqs( can_alloc_token_num = g_infer_context.get_can_alloc_token_num() for req_obj in ready_reqs: - if req_obj.filter_mark: finished_reqs.append(req_obj) continue @@ -787,7 +790,6 @@ def _sample_and_scatter_token( b_prefill_has_output_cpu: torch.Tensor = None, mask_func: Optional[Callable] = None, ): - if mask_func is not None: assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index e1a4e421d1..7ddf985fb0 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -16,6 +16,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +import os from typing import List, Tuple, Union from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast @@ -39,6 +41,41 @@ _FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer" +def _load_tokenizers_backend_tokenizer( + tokenizer_name: str, + error: ValueError, + *args, + **kwargs, +) -> Union[PreTrainedTokenizerFast, None]: + if "Tokenizer class TokenizersBackend does not exist or is not currently imported" not in str(error): + return None + if not os.path.isdir(tokenizer_name): + return None + + tokenizer_file = os.path.join(tokenizer_name, "tokenizer.json") + tokenizer_config_file = os.path.join(tokenizer_name, "tokenizer_config.json") + if not os.path.exists(tokenizer_file) or not os.path.exists(tokenizer_config_file): + return None + + with open(tokenizer_config_file, "r", encoding="utf-8") as fp: + tokenizer_config = json.load(fp) + if tokenizer_config.get("tokenizer_class") != "TokenizersBackend": + return None + + special_token_kwargs = { + name: tokenizer_config[name] + for name in ("bos_token", "eos_token", "unk_token", "sep_token", "pad_token", "cls_token", "mask_token") + if tokenizer_config.get(name) is not None + } + if tokenizer_config.get("extra_special_tokens"): + special_token_kwargs["additional_special_tokens"] = tokenizer_config["extra_special_tokens"] + if tokenizer_config.get("model_max_length") is not None: + special_token_kwargs["model_max_length"] = tokenizer_config["model_max_length"] + + logger.info("Loading TokenizersBackend tokenizer through tokenizer.json fast-tokenizer fallback.") + return PreTrainedTokenizerFast(tokenizer_file=tokenizer_file, *args, **special_token_kwargs, **kwargs) + + def get_tokenizer( tokenizer_name: str, tokenizer_mode: str = "auto", @@ -71,6 +108,10 @@ def get_tokenizer( logger.warning(f"load fast tokenizer fail: {str(e)}") kwargs["use_fast"] = False tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=trust_remote_code, *args, **kwargs) + except ValueError as e: + tokenizer = _load_tokenizers_backend_tokenizer(tokenizer_name, e, *args, **kwargs) + if tokenizer is None: + raise if not isinstance(tokenizer, PreTrainedTokenizerFast): logger.info( diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 85d21477b0..5f5324d772 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -482,6 +482,10 @@ def get_tool_call_parser_for_model(model_path: str) -> Optional[str]: if model_type is None: return None + # GLM-5 / GLM-5.2 DSA models use the GLM-4.7 tool-call format. + if model_type == "glm_moe_dsa": + return "glm47" + # Qwen3.5 series if model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"]: return "qwen3_coder" @@ -511,6 +515,10 @@ def get_reasoning_parser_for_model(model_path: str) -> Optional[str]: if model_type is None: return None + # GLM-5 / GLM-5.2 DSA models share the GLM-4.5-style thinking tags. + if model_type == "glm_moe_dsa": + return "glm45" + # Qwen3.5 and Qwen3 series if model_type in [ "qwen3", diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 494908cb10..57a0474d65 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -22,6 +22,8 @@ PPLINT8KVMemoryManager, PPLINT4KVMemoryManager, Deepseek2MemoryManager, + Deepseek3_2MemoryManager, + FP8PerTokenGroupQuantDeepseek3_2MemoryManager, Qwen3NextMemManager, ) @@ -92,6 +94,31 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": scale_head_dim=0, scale_data_type=get_llm_data_type(), ) + elif mem_manager_class is Deepseek3_2MemoryManager: + cpu_cache_meta = CpuKVCacheMeta( + page_num=0, + token_page_size=args.cpu_cache_token_page_size, + layer_num=get_layer_num(args.model_dir), + num_heads=1, + head_dim=512 + 64 + (144 // 2), + data_type=get_llm_data_type(), + scale_head_dim=0, + scale_data_type=get_llm_data_type(), + ) + elif mem_manager_class is FP8PerTokenGroupQuantDeepseek3_2MemoryManager: + cpu_cache_meta = CpuKVCacheMeta( + page_num=0, + token_page_size=args.cpu_cache_token_page_size, + layer_num=get_layer_num(args.model_dir), + num_heads=1, + head_dim=( + FP8PerTokenGroupQuantDeepseek3_2MemoryManager.flashmla_bytes_per_token + + FP8PerTokenGroupQuantDeepseek3_2MemoryManager.indexer_bytes_per_token + ), + data_type=torch.uint8, + scale_head_dim=0, + scale_data_type=torch.uint8, + ) elif mem_manager_class is MemoryManager: cpu_cache_meta = CpuKVCacheMeta( page_num=0, diff --git a/test/acc/test_gsmk.py b/test/acc/test_gsmk.py index 16ebc6e095..ec3fce5c29 100644 --- a/test/acc/test_gsmk.py +++ b/test/acc/test_gsmk.py @@ -142,6 +142,7 @@ def parse_args(): parser.add_argument("--num-questions", type=int, default=1000) parser.add_argument("--result-file", type=str, default="result.jsonl") parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--max-new-tokens", type=int, default=1024) return parser.parse_args() @@ -151,7 +152,7 @@ def main(args): # Read data url_data = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" - filename = download_and_cache_file(url_data) + filename = args.data_path if os.path.exists(args.data_path) else download_and_cache_file(url_data) lines = list(read_jsonl(filename)) # Construct prompts @@ -183,7 +184,7 @@ def get_one_answer(i): answer = call_generate_lightllm( prompt=few_shot_examples + questions[i], temperature=0, - max_tokens=1024, + max_tokens=args.max_new_tokens, stop=["Question", "Assistant:", "<|separator|>", "Human:", "\n\nQuestion"], url=url, ) @@ -231,6 +232,7 @@ def get_one_answer(i): "other": { "num_questions": args.num_questions, "parallel": args.parallel, + "max_new_tokens": args.max_new_tokens, }, } fout.write(json.dumps(value) + "\n") diff --git a/test/cpu_cache_kernel/test_fp8_dsa_cpu_cache.py b/test/cpu_cache_kernel/test_fp8_dsa_cpu_cache.py new file mode 100644 index 0000000000..4aa9a522ba --- /dev/null +++ b/test/cpu_cache_kernel/test_fp8_dsa_cpu_cache.py @@ -0,0 +1,134 @@ +from types import SimpleNamespace + +import torch + + +def test_fp8_dsa_cpu_cache_meta_includes_flashmla_and_indexer(monkeypatch): + import lightllm.utils.kv_cache_utils as kv_cache_utils + from lightllm.common.kv_cache_mem_manager import FP8PerTokenGroupQuantDeepseek3_2MemoryManager + + kv_cache_utils.calcu_cpu_cache_meta.cache_clear() + + monkeypatch.setattr( + kv_cache_utils, + "get_env_start_args", + lambda: SimpleNamespace( + enable_cpu_cache=True, + model_dir="/fake/glm-5.2", + cpu_cache_token_page_size=64, + cpu_cache_storage_size=1, + mtp_mode=None, + ), + ) + monkeypatch.setattr(kv_cache_utils, "is_linear_att_mixed_model", lambda _: False) + monkeypatch.setattr(kv_cache_utils, "select_mem_manager_class", lambda: FP8PerTokenGroupQuantDeepseek3_2MemoryManager) + monkeypatch.setattr(kv_cache_utils, "get_layer_num", lambda _: 2) + + meta = kv_cache_utils.calcu_cpu_cache_meta() + + expected_head_dim = ( + FP8PerTokenGroupQuantDeepseek3_2MemoryManager.flashmla_bytes_per_token + + FP8PerTokenGroupQuantDeepseek3_2MemoryManager.indexer_bytes_per_token + ) + assert meta.data_type is torch.uint8 + assert meta.num_heads == 1 + assert meta.head_dim == expected_head_dim + assert meta.scale_head_dim == 0 + assert meta.calcu_one_page_size() == 64 * 2 * expected_head_dim + + +def test_bf16_dsa_cpu_cache_meta_keeps_padded_indexer_bytes(monkeypatch): + import lightllm.utils.kv_cache_utils as kv_cache_utils + from lightllm.common.kv_cache_mem_manager import Deepseek3_2MemoryManager + + kv_cache_utils.calcu_cpu_cache_meta.cache_clear() + + monkeypatch.setattr( + kv_cache_utils, + "get_env_start_args", + lambda: SimpleNamespace( + enable_cpu_cache=True, + model_dir="/fake/glm-5.2", + cpu_cache_token_page_size=64, + cpu_cache_storage_size=1, + mtp_mode=None, + ), + ) + monkeypatch.setattr(kv_cache_utils, "is_linear_att_mixed_model", lambda _: False) + monkeypatch.setattr(kv_cache_utils, "select_mem_manager_class", lambda: Deepseek3_2MemoryManager) + monkeypatch.setattr(kv_cache_utils, "get_layer_num", lambda _: 2) + monkeypatch.setattr(kv_cache_utils, "get_llm_data_type", lambda: torch.bfloat16) + + meta = kv_cache_utils.calcu_cpu_cache_meta() + + assert meta.data_type is torch.bfloat16 + assert meta.num_heads == 1 + assert meta.head_dim == 512 + 64 + (144 // 2) + assert meta.scale_head_dim == 0 + assert meta.calcu_one_page_size() == 64 * 2 * meta.head_dim * torch.bfloat16.itemsize + + +def test_fp8_dsa_operator_implements_cpu_cache_methods(): + from lightllm.common.kv_cache_mem_manager.operator.base import BaseMemManagerOperator + from lightllm.common.kv_cache_mem_manager.operator.deepseek import FP8PerTokenGroupQuantDeepseek3_2MemOperator + + assert ( + FP8PerTokenGroupQuantDeepseek3_2MemOperator.load_cpu_cache_to_gpu + is not BaseMemManagerOperator.load_cpu_cache_to_gpu + ) + assert ( + FP8PerTokenGroupQuantDeepseek3_2MemOperator.offload_gpu_kv_to_cpu_cache + is not BaseMemManagerOperator.offload_gpu_kv_to_cpu_cache + ) + + +def test_fp8_dsa_operator_moves_flashmla_and_indexer_views(monkeypatch): + import lightllm.common.basemodel.triton_kernel.kv_cache_offload as kv_cache_offload + import lightllm.utils.dist_utils as dist_utils + import lightllm.utils.envs_utils as envs_utils + from lightllm.common.kv_cache_mem_manager import FP8PerTokenGroupQuantDeepseek3_2MemoryManager + from lightllm.common.kv_cache_mem_manager.operator.deepseek import FP8PerTokenGroupQuantDeepseek3_2MemOperator + + class CudaSeq: + is_cuda = True + + def __init__(self, n): + self.n = n + + def __len__(self): + return self.n + + split = FP8PerTokenGroupQuantDeepseek3_2MemoryManager.flashmla_bytes_per_token + indexer_size = FP8PerTokenGroupQuantDeepseek3_2MemoryManager.indexer_bytes_per_token + cpu_tensor = torch.empty((4, 2, 64, 1, split + indexer_size), dtype=torch.uint8) + client = SimpleNamespace( + kv_cache_tensor_meta=SimpleNamespace(data_type=torch.uint8, num_heads=1, head_dim=split + indexer_size), + cpu_kv_cache_tensor=cpu_tensor, + ) + mem_manager = SimpleNamespace( + flashmla_bytes_per_token=split, + indexer_bytes_per_token=indexer_size, + kv_buffer=object(), + indexer_k_buffer=object(), + ) + op = FP8PerTokenGroupQuantDeepseek3_2MemOperator(mem_manager) + mem_indexes = CudaSeq(128) + page_indexes = CudaSeq(2) + page_readies = CudaSeq(2) + + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: SimpleNamespace(cpu_cache_token_page_size=64)) + monkeypatch.setattr(dist_utils, "get_current_rank_in_dp", lambda: 0) + monkeypatch.setattr(dist_utils, "get_dp_world_size", lambda: 1) + + load_calls = [] + offload_calls = [] + monkeypatch.setattr(kv_cache_offload, "load_cpu_kv_to_gpu", lambda **kwargs: load_calls.append(kwargs)) + monkeypatch.setattr(kv_cache_offload, "offload_gpu_kv_to_cpu", lambda **kwargs: offload_calls.append(kwargs)) + + op.load_cpu_cache_to_gpu(mem_indexes, page_indexes, client, req=None) + op.offload_gpu_kv_to_cpu_cache(mem_indexes, page_indexes, page_readies, client, req=None) + + assert [call["gpu_kv_cache"] for call in load_calls] == [mem_manager.kv_buffer, mem_manager.indexer_k_buffer] + assert [call["cpu_kv_cache"].shape[-1] for call in load_calls] == [split, indexer_size] + assert [call["gpu_kv_cache"] for call in offload_calls] == [mem_manager.kv_buffer, mem_manager.indexer_k_buffer] + assert [call["cpu_kv_cache"].shape[-1] for call in offload_calls] == [split, indexer_size] diff --git a/unit_tests/common/basemodel/triton_kernel/test_fused_add_rmsnorm.py b/unit_tests/common/basemodel/triton_kernel/test_fused_add_rmsnorm.py new file mode 100644 index 0000000000..78d3d70f0c --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_fused_add_rmsnorm.py @@ -0,0 +1,80 @@ +"""Correctness (+ __main__ microbench) for fused_add_rmsnorm vs the unfused (add_ ; rmsnorm) sequence.""" +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.common.basemodel.triton_kernel.norm.fused_add_rmsnorm import fused_add_rmsnorm_forward + + +def _ref(residual, x, weight, eps): + # exactly what token_forward does today: in-place residual add, then a separate rmsnorm + res = residual.clone() + res.add_(x) + out = rmsnorm_forward(res, weight, eps) + return out, res + + +@pytest.mark.parametrize( + "M, N, dtype", + [(M, 6144, torch.bfloat16) for M in [1, 2, 4, 8, 16, 32]] # GLM-5.2 hidden + + [ + (1, 7168, torch.bfloat16), # DeepSeek-V3 hidden + (1, 4096, torch.float16), + (13, 6144, torch.bfloat16), + ], +) +def test_fused_add_rmsnorm_matches_unfused(M, N, dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + + eps = 1e-5 + torch.manual_seed(0) + residual = (-2.3 + 0.5 * torch.randn(M, N, dtype=dtype, device="cuda")).contiguous() + x = (0.7 * torch.randn(M, N, dtype=dtype, device="cuda")).contiguous() + weight = torch.rand(N, dtype=dtype, device="cuda") + + out_ref, res_ref = _ref(residual, x, weight, eps) + + res_fused = residual.clone() + out_fused = torch.empty_like(res_fused) + fused_add_rmsnorm_forward(res_fused, x, weight, eps, out=out_fused) + + # residual update must be bit-identical to a plain add_ + assert torch.equal(res_fused, res_ref), "residual (residual+x) must bit-match a plain add_" + # variance is taken from the rounded sum, matching the unfused path bit-for-bit + assert torch.equal(out_fused, out_ref), "normalized output must bit-match the unfused path" + + +def bench(M, N, dtype=torch.bfloat16, eps=1e-5, iters=200): + residual = torch.randn(M, N, dtype=dtype, device="cuda") + x = torch.randn(M, N, dtype=dtype, device="cuda") + weight = torch.rand(N, dtype=dtype, device="cuda") + out = torch.empty_like(residual) + + def run_unfused(): + residual.add_(x) + rmsnorm_forward(residual, weight, eps, out=out) + + def run_fused(): + fused_add_rmsnorm_forward(residual, x, weight, eps, out=out) + + for fn, name in [(run_unfused, "add_+rmsnorm"), (run_fused, "fused_add_rmsnorm")]: + for _ in range(20): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + print(f" {name:<20} {start.elapsed_time(end) / iters * 1e3:.2f} us/call") + + +if __name__ == "__main__": + print("=== microbench (decode-shaped, bs in 1..32 @ N=6144) ===") + for M in [1, 4, 16, 32]: + print(f"M={M}:") + bench(M, 6144) + pytest.main([__file__, "-q"]) diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_group_quant.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_group_quant.py new file mode 100644 index 0000000000..469d14a290 --- /dev/null +++ b/unit_tests/common/fused_moe/test_moe_silu_and_mul_group_quant.py @@ -0,0 +1,56 @@ +"""Correctness: fused silu_and_mul + group fp8 quant vs unfused silu_and_mul_fwd then per_token_group_quant_fp8.""" +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_group_quant import silu_and_mul_group_quant_fwd + + +def _dequant(q, s, group_size): + # q: [M, N] fp8, s: [M, N//group] fp32 -> [M, N] fp32 + M, N = q.shape + qf = q.to(torch.float32).reshape(M, N // group_size, group_size) + return (qf * s.reshape(M, N // group_size, 1)).reshape(M, N) + + +@pytest.mark.parametrize( + "M, N, layout", + [(M, 1536, "blocked") for M in [1, 8, 16, 64, 256]] # glm/deepseek-ish moe intermediate + + [ + (8, 2048, "blocked"), + (13, 768, "blocked"), + (8, 1536, "interleaved"), + ], +) +def test_silu_and_mul_group_quant_matches_unfused(M, N, layout): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + + group_size = 128 + dtype = torch.bfloat16 + torch.manual_seed(0) + inp = torch.randn(M, 2 * N, dtype=dtype, device="cuda").contiguous() + + # unfused reference + tmp = torch.empty(M, N, dtype=dtype, device="cuda") + silu_and_mul_fwd(inp, tmp, layout=layout) + q_ref, s_ref = per_token_group_quant_fp8(tmp, group_size, dtype=torch.float8_e4m3fn) + + # fused + q = torch.empty(M, N, dtype=torch.float8_e4m3fn, device="cuda") + s = torch.empty(M, N // group_size, dtype=torch.float32, device="cuda") + silu_and_mul_group_quant_fwd(inp, q, s, group_size, layout=layout) + + dq_ref = _dequant(q_ref, s_ref.reshape(M, N // group_size), group_size) + dq = _dequant(q, s, group_size) + cos = torch.nn.functional.cosine_similarity(dq.flatten(), dq_ref.flatten(), dim=0).item() + s_match = (s.reshape(-1) - s_ref.reshape(-1)).abs().max().item() + q_exact = (q.to(torch.float32) == q_ref.to(torch.float32)).float().mean().item() + + assert cos > 0.9999, f"cosine too low: {cos}" + assert s_match < 1e-6 or q_exact > 0.99, "scale/quant diverged beyond fp8 rounding" + + +if __name__ == "__main__": + pytest.main([__file__, "-q"]) diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_fp8_mqa_logits.py b/unit_tests/models/deepseek3_2/triton_kernel/test_fp8_mqa_logits.py new file mode 100644 index 0000000000..74c707ea64 --- /dev/null +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_fp8_mqa_logits.py @@ -0,0 +1,203 @@ +import pytest +import torch +import triton + +from lightllm.models.deepseek3_2.triton_kernel.fp8_mqa_logits import fp8_paged_mqa_logits + + +def _require_deep_gemm(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + try: + import deep_gemm # noqa: F401 + except ImportError: + pytest.skip("deep_gemm is not available") + import deep_gemm + + return deep_gemm + + +def _build_case(batch_size, seq_len, mtp_step, head_num=32, head_dim=128, pool_size=None, seed=42): + """Build decode-shaped inputs mirroring NsaInfer._get_indices + gen_nsa_ks_ke.""" + torch.manual_seed(seed) + device = "cuda" + req_num = batch_size + q_num = req_num * (mtp_step + 1) + + b_seq_len_req = torch.randint(max(1, seq_len // 2), seq_len + 1, (req_num,), device=device, dtype=torch.int32) + b_seq_len_req[0] = seq_len # keep max_kv_seq_len == seq_len + total_tokens = int(b_seq_len_req.sum().item()) + pool_size = pool_size or (total_tokens + 1024) + + # token-granular pool with shuffled slot assignment + perm = torch.randperm(pool_size, device=device, dtype=torch.int32) + indexer_k_buffer = torch.empty((pool_size, 1, 132), dtype=torch.uint8, device=device) + k_bf16 = torch.randn(pool_size, head_dim, device=device, dtype=torch.float32) + k_scale = (torch.rand(pool_size, device=device, dtype=torch.float32) + 0.5) / 100.0 + indexer_k_buffer[:, 0, 0:128] = k_bf16.to(torch.float8_e4m3fn).view(torch.uint8) + indexer_k_buffer[:, 0, 128:132] = k_scale.view(-1, 1).view(torch.uint8) + + # ragged layout + per-row key windows (decode: q rows of one req share its K, ke grows by 1) + ragged_mem_index = torch.empty(total_tokens, dtype=torch.int32, device=device) + ks = torch.empty(q_num, dtype=torch.int32, device=device) + ke = torch.empty(q_num, dtype=torch.int32, device=device) + pre_sum = 0 + slot_cursor = 0 + req_to_slots = [] + for r in range(req_num): + L = int(b_seq_len_req[r].item()) + slots = perm[slot_cursor : slot_cursor + L] + slot_cursor += L + req_to_slots.append(slots) + ragged_mem_index[pre_sum : pre_sum + L] = slots + for j in range(mtp_step + 1): + row = r * (mtp_step + 1) + j + ks[row] = pre_sum + ke[row] = pre_sum + (L - mtp_step) + j + pre_sum += L + + q_fp8 = torch.randn(q_num, head_num, head_dim, device=device, dtype=torch.float32).to(torch.float8_e4m3fn) + weights = torch.randn(q_num, head_num, device=device, dtype=torch.float32) / head_num + + max_kv_seq_len = seq_len + out_col_num = q_num * max_kv_seq_len // (mtp_step + 1) # K rows for the extract path + return { + "q_fp8": q_fp8, + "indexer_k_buffer": indexer_k_buffer, + "weights": weights, + "ragged_mem_index": ragged_mem_index, + "ks": ks, + "ke": ke, + "out_col_num": out_col_num, + "max_kv_seq_len": max_kv_seq_len, + "b_seq_len_req": b_seq_len_req, + "mtp_step": mtp_step, + "req_num": req_num, + "total_tokens": total_tokens, + } + + +def _ref_logits_deep_gemm(deep_gemm, case): + """Reference: today's extract + deep_gemm.fp8_mqa_logits path (extraction done in torch).""" + ragged = case["ragged_mem_index"][: case["total_tokens"]].long() + buf = case["indexer_k_buffer"] + k_fp8 = buf[:, 0, 0:128].view(torch.float8_e4m3fn)[ragged].contiguous() + k_scale = buf[:, 0, 128:132].view(torch.float32)[ragged].reshape(-1).contiguous() + pad = case["out_col_num"] - k_fp8.shape[0] + assert pad >= 0 + if pad: + k_fp8 = torch.cat([k_fp8, torch.zeros(pad, 128, dtype=k_fp8.dtype, device=k_fp8.device)]) + k_scale = torch.cat([k_scale, torch.zeros(pad, dtype=k_scale.dtype, device=k_scale.device)]) + return deep_gemm.fp8_mqa_logits( + case["q_fp8"], + (k_fp8, k_scale), + case["weights"], + case["ks"], + case["ke"], + clean_logits=False, + max_seqlen_k=case["max_kv_seq_len"], + ) + + +@pytest.mark.parametrize( + "batch_size,seq_len,mtp_step", + [ + (1, 128, 0), + (8, 1024, 0), + (32, 4096, 0), + (8, 1024, 1), + (16, 2048, 3), + (64, 512, 0), + ], +) +def test_fp8_paged_mqa_logits_matches_deep_gemm(batch_size, seq_len, mtp_step): + deep_gemm = _require_deep_gemm() + case = _build_case(batch_size, seq_len, mtp_step) + + ref = _ref_logits_deep_gemm(deep_gemm, case) + out = fp8_paged_mqa_logits( + q_fp8=case["q_fp8"], + indexer_k_buffer=case["indexer_k_buffer"], + weights=case["weights"], + ragged_mem_index=case["ragged_mem_index"], + cu_seqlen_ks=case["ks"], + cu_seqlen_ke=case["ke"], + max_kv_seq_len=case["max_kv_seq_len"], + ) + + assert out.shape == ref.shape + # rows are compacted: row m valid at columns [0, ke-ks) + for row in range(out.shape[0]): + length = int(case["ke"][row]) - int(case["ks"][row]) + torch.testing.assert_close(out[row, :length], ref[row, :length], rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + # quick perf A/B vs today's real path (extract_indexer_ks + deep_gemm) + from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import extract_indexer_ks + + deep_gemm = _require_deep_gemm() + + def _bench(fn): + return triton.testing.do_bench_cudagraph(fn, return_mode="median") + + print(f"{'batch':>5} {'seq':>6} {'mtp':>3} | {'old (extract+dg)':>16} {'new (paged)':>12} {'speedup':>8}") + for batch, seq, mtp in [ + (1, 4096, 0), + (8, 4096, 0), + (32, 4096, 0), + (8, 16384, 0), + (32, 16384, 0), + (64, 8192, 0), + (128, 2048, 0), + (8, 60000, 0), + (32, 60000, 0), + (8, 16384, 1), + ]: + case = _build_case(batch, seq, mtp) + q_num = case["q_fp8"].shape[0] + + # emulate the old path's inputs (b_seq_len/b_req_idx expanded per mtp row) + b_seq_len = case["b_seq_len_req"].repeat_interleave(mtp + 1).to(torch.int32) + b_req_idx = torch.arange(case["req_num"], device="cuda", dtype=torch.int32).repeat_interleave(mtp + 1) + req_to_token = torch.zeros((case["req_num"], seq), dtype=torch.int32, device="cuda") + pre = 0 + for r in range(case["req_num"]): + L = int(case["b_seq_len_req"][r]) + req_to_token[r, :L] = case["ragged_mem_index"][pre : pre + L] + pre += L + + def run_old(): + k_fp8_, k_scale_ = extract_indexer_ks( + I_buffer=case["indexer_k_buffer"], + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + req_to_token_indexs=req_to_token, + out_token_num=q_num * seq, + max_kv_seq_len=seq, + mtp_step=mtp, + ) + return deep_gemm.fp8_mqa_logits( + case["q_fp8"], + (k_fp8_, k_scale_), + case["weights"], + case["ks"], + case["ke"], + clean_logits=False, + max_seqlen_k=seq, + ) + + def run_new(): + return fp8_paged_mqa_logits( + q_fp8=case["q_fp8"], + indexer_k_buffer=case["indexer_k_buffer"], + weights=case["weights"], + ragged_mem_index=case["ragged_mem_index"], + cu_seqlen_ks=case["ks"], + cu_seqlen_ke=case["ke"], + max_kv_seq_len=case["max_kv_seq_len"], + ) + + t_old = _bench(run_old) + t_new = _bench(run_new) + print(f"{batch:>5} {seq:>6} {mtp:>3} | {t_old*1000:>13.1f}us {t_new*1000:>9.1f}us {t_old/t_new:>7.2f}x") From b830fbf4a3edeebb84a44a37bd82fe7b8dcfaacb Mon Sep 17 00:00:00 2001 From: shihaobai <1798930569@qq.com> Date: Mon, 29 Jun 2026 14:34:20 +0000 Subject: [PATCH 03/10] static cost --- docs/CN/source/getting_started/benchmark.rst | 62 +- docs/EN/source/getting_started/benchmark.rst | 55 +- .../benchmark/static_inference/model_infer.py | 447 ------ .../static_inference/model_infer_mtp.py | 285 ---- .../static_inference/static_benchmark.py | 1425 +++++++++++++++++ test/benchmark/static_inference/test_model.py | 47 +- 6 files changed, 1509 insertions(+), 812 deletions(-) delete mode 100644 test/benchmark/static_inference/model_infer.py delete mode 100644 test/benchmark/static_inference/model_infer_mtp.py create mode 100644 test/benchmark/static_inference/static_benchmark.py diff --git a/docs/CN/source/getting_started/benchmark.rst b/docs/CN/source/getting_started/benchmark.rst index 9fdaf37ed7..85c8b475b6 100644 --- a/docs/CN/source/getting_started/benchmark.rst +++ b/docs/CN/source/getting_started/benchmark.rst @@ -133,13 +133,18 @@ Prompt Cache 测试 静态推理性能测试 (Static Inference Benchmark) ---------------------------------------------- -静态推理测试用于评估模型在固定输入条件下的推理性能, 主要评估算子的优劣 -模型推理测试 (model_infer.py) -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +静态推理测试用于评估模型在固定输入条件下的推理性能, 主要评估算子的优劣。 +统一入口为 ``test/benchmark/static_inference/test_model.py``，核心实现集中在 +``test/benchmark/static_inference/static_benchmark.py``。 + +模型推理测试 +~~~~~~~~~~~~ **主要特性：** - 支持 prefill 和 decode 阶段性能测试 +- 支持 prefill 静态 TPS 的多输入长度、多 batch size 和 chunked prefill +- 支持 decode 静态 TPS 的多 batch size、多上下文长度和多输出长度 - 支持 microbatch overlap 优化 - 支持多 GPU 并行推理 - 提供详细的吞吐量统计 @@ -150,23 +155,28 @@ Prompt Cache 测试 python test/benchmark/static_inference/test_model.py \ --model_dir /path/to/model \ - --batch_size 32 \ - --input_len 1024 \ - --output_len 128 \ + --benchmark all \ + --batch_sizes 8,16,32 \ + --input_lens 1024,2048 \ + --context_lens 1024,4096 \ + --output_lens 128 \ + --chunked_prefill_sizes 512 \ --tp 2 \ --data_type bf16 **主要参数：** - ``--model_dir``: 模型路径 -- ``--batch_size``: 批次大小 -- ``--input_len``: 输入序列长度 -- ``--output_len``: 输出序列长度 +- ``--benchmark``: 测试阶段，可选 ``all``、``prefill``、``decode`` +- ``--batch_size`` / ``--batch_sizes``: 单个或多个批次大小 +- ``--input_len`` / ``--input_lens``: prefill 输入序列长度 +- ``--context_lens``: decode 阶段上下文长度 +- ``--output_len`` / ``--output_lens``: decode 输出长度 +- ``--chunked_prefill_sizes``: prefill chunk 大小，默认 ``4096``；使用 ``full``、``none`` 或 ``0`` 表示不分块 - ``--tp``: Tensor Parallel 并行度 - ``--data_type``: 数据类型 (bf16/fp16/fp32) -- ``--enable_prefill_microbatch_overlap``: 启用 prefill microbatch overlap，仅适用于DeepSeek模型的EP模式 -- ``--enable_decode_microbatch_overlap``: 启用 decode microbatch overlap，仅适用于DeepSeek模型的EP模式 -- ``--torch_profile``: 启用 torch profiler 进行性能分析 +- ``--enable_prefill_microbatch_overlap``: 启用 prefill microbatch overlap，仅适用于 DeepSeek 模型的 EP 模式 +- ``--enable_decode_microbatch_overlap``: 启用 decode microbatch overlap，仅适用于 DeepSeek 模型的 EP 模式 .. note:: 这里没有列举完整的启动参数，静态测试脚本也共享lightllm的启动参数，更多启动配置可以参考 :ref:`tutorial/api_server_args_zh` 。 @@ -177,10 +187,13 @@ Prompt Cache 测试 - Decode 阶段吞吐量 (tokens/s) - 各阶段延迟统计 -多结果预测性能测试 (model_infer_mtp.py) -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +多结果预测性能测试 +~~~~~~~~~~~~~~~~~~ -多结果预测静态性能测试，默认百分百接受率，用来评估多结果预测的极限性能。目前只支持DeepSeek 系列模型 +多结果预测静态性能测试默认 ``--mtp_accept_rate 1.0``，即接受全部 draft token； +可调低该值模拟更低接受率下的 MTP decode 吞吐。 +DeepSeek R1 可以使用 ``/mtc/models/DeepSeek-R1`` 和 ``/mtc/models/DeepSeek-R1-NextN`` 这类 +主模型/草稿模型结构。 **使用方法：** @@ -188,19 +201,22 @@ Prompt Cache 测试 python test/benchmark/static_inference/test_model.py \ --model_dir /path/to/main_model \ - --mtp_mode deepseekv3 \ - --mtp_step 1 \ + --benchmark decode \ + --mtp_mode eagle_with_att \ + --mtp_step 2 \ --mtp_draft_model_dir /path/to/draft_model \ - --batch_size 32 \ - --input_len 1024 \ - --output_len 128 + --mtp_accept_rate 0.8 \ + --batch_sizes 8,16 \ + --context_lens 1024,4096 \ + --output_lens 128 参数说明： - ``--model_dir``: 主模型路径 -- ``--mtp_mode``: 指定多结果预测的模型，目前只支持deepseekv2/v3/r1 -- ``--mtp_step``: 每次forward step产生的token 数量，默认为1 +- ``--mtp_mode``: MTP 模式，如 ``eagle_with_att``、``vanilla_with_att``、``eagle_no_att``、``vanilla_no_att`` +- ``--mtp_step``: 每次 decode 额外预测的 draft token 数量 - ``--mtp_draft_model_dir``: 草稿模型路径 +- ``--mtp_accept_rate``: 每个 draft token 的模拟接受概率，采样过程不计入 decode 耗时 Vision Transformer 测试 (test_vit.py) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -215,4 +231,4 @@ Vision Transformer 测试 (test_vit.py) --model_dir ./InternVL2/InternVL2-8B/ \ --batch_size 1 \ --image_size 448 \ - --world_size 2 \ No newline at end of file + --world_size 2 diff --git a/docs/EN/source/getting_started/benchmark.rst b/docs/EN/source/getting_started/benchmark.rst index 60052fd8f6..35b6a1ecfd 100755 --- a/docs/EN/source/getting_started/benchmark.rst +++ b/docs/EN/source/getting_started/benchmark.rst @@ -132,13 +132,17 @@ Static Inference Performance Testing (Static Inference Benchmark) ------------------------------------------------------------------ Static inference testing is used to evaluate model inference performance under fixed input conditions, mainly evaluating operator quality. +The unified entry is ``test/benchmark/static_inference/test_model.py``. The +core implementation lives in ``test/benchmark/static_inference/static_benchmark.py``. -Model Inference Testing (model_infer.py) -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Model Inference Testing +~~~~~~~~~~~~~~~~~~~~~~~ **Main Features:** - Supports prefill and decode stage performance testing +- Supports prefill static TPS with multiple input lengths, batch sizes, and chunked prefill sizes +- Supports decode static TPS with multiple batch sizes, context lengths, and output lengths - Supports microbatch overlap optimization - Supports multi-GPU parallel inference - Provides detailed throughput statistics @@ -149,23 +153,28 @@ Model Inference Testing (model_infer.py) python test/benchmark/static_inference/test_model.py \ --model_dir /path/to/model \ - --batch_size 32 \ - --input_len 1024 \ - --output_len 128 \ + --benchmark all \ + --batch_sizes 8,16,32 \ + --input_lens 1024,2048 \ + --context_lens 1024,4096 \ + --output_lens 128 \ + --chunked_prefill_sizes 512 \ --tp 2 \ --data_type bf16 **Main Parameters:** - ``--model_dir``: Model path -- ``--batch_size``: Batch size -- ``--input_len``: Input sequence length -- ``--output_len``: Output sequence length +- ``--benchmark``: Benchmark stage, one of ``all``, ``prefill``, or ``decode`` +- ``--batch_size`` / ``--batch_sizes``: Single or multiple batch sizes +- ``--input_len`` / ``--input_lens``: Prefill input lengths +- ``--context_lens``: Decode context lengths +- ``--output_len`` / ``--output_lens``: Decode output lengths +- ``--chunked_prefill_sizes``: Prefill chunk sizes, default ``4096``; use ``full``, ``none``, or ``0`` for unchunked prefill - ``--tp``: Tensor Parallel degree - ``--data_type``: Data type (bf16/fp16/fp32) - ``--enable_prefill_microbatch_overlap``: Enable prefill microbatch overlap, only applicable to DeepSeek model EP mode - ``--enable_decode_microbatch_overlap``: Enable decode microbatch overlap, only applicable to DeepSeek model EP mode -- ``--torch_profile``: Enable torch profiler for performance analysis .. note:: Complete startup parameters are not listed here. Static testing scripts also share Lightllm's startup parameters. For more startup configurations, please refer to :ref:`tutorial/api_server_args_zh`. @@ -176,10 +185,14 @@ Model Inference Testing (model_infer.py) - Decode stage throughput (tokens/s) - Latency statistics for each stage -Multi-Token Prediction Performance Testing (model_infer_mtp.py) -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Multi-Token Prediction Performance Testing +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Multi-token prediction static performance testing with 100% acceptance rate by default, used to evaluate the ultimate performance of multi-token prediction. Currently only supports DeepSeek series models. +Multi-token prediction static performance testing defaults to +``--mtp_accept_rate 1.0``, which accepts all draft tokens. Lower values simulate +MTP decode throughput with lower acceptance. DeepSeek R1 can use a main/draft +model pair such as ``/mtc/models/DeepSeek-R1`` and +``/mtc/models/DeepSeek-R1-NextN``. **Usage:** @@ -187,13 +200,19 @@ Multi-token prediction static performance testing with 100% acceptance rate by d python test/benchmark/static_inference/test_model.py \ --model_dir /path/to/main_model \ - --mtp_mode deepseekv3 \ - --mtp_step 1 \ + --benchmark decode \ + --mtp_mode eagle_with_att \ + --mtp_step 2 \ --mtp_draft_model_dir /path/to/draft_model \ - --batch_size 32 \ - --input_len 1024 \ - --output_len 128 + --mtp_accept_rate 0.8 \ + --batch_sizes 8,16 \ + --context_lens 1024,4096 \ + --output_lens 128 Parameter Description: -- ``--model_dir``: Main model path \ No newline at end of file +- ``--model_dir``: Main model path +- ``--mtp_mode``: MTP mode, for example ``eagle_with_att``, ``vanilla_with_att``, ``eagle_no_att``, or ``vanilla_no_att`` +- ``--mtp_step``: Number of extra draft tokens predicted per decode step +- ``--mtp_draft_model_dir``: Draft model path +- ``--mtp_accept_rate``: Simulated per-draft-token accept probability; sampling is excluded from decode timing diff --git a/test/benchmark/static_inference/model_infer.py b/test/benchmark/static_inference/model_infer.py deleted file mode 100644 index f2c900af09..0000000000 --- a/test/benchmark/static_inference/model_infer.py +++ /dev/null @@ -1,447 +0,0 @@ -import os -import torch -import numpy as np -from multiprocessing import Queue -import multiprocessing -from transformers import PretrainedConfig -from lightllm.utils.dist_utils import init_distributed_env, get_current_rank_in_dp -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.models import get_model -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from torch.profiler import profile, record_function, ProfilerActivity -from lightllm.utils.log_utils import init_logger -import torch.cuda as cuda - -logger = init_logger(__name__) - - -def test_model_inference(args): - ans_queue = Queue() - workers = [] - dp_size = args.get("dp", 1) - - for rank_id in range(args.node_rank * args.tp, (args.node_rank + 1) * args.tp): - model_kvargs = { - "args": args, - "nccl_host": args.nccl_host, - "data_type": args.data_type, - "nccl_port": args.nccl_port, - "rank_id": rank_id, - "world_size": args.tp, - "dp_size": dp_size, - "weight_dir": args.model_dir, - "quant_type": args.quant_type, - "load_way": "HF", - "max_total_token_num": args.max_total_token_num, - "graph_max_len_in_batch": args.max_req_total_len, - "graph_max_batch_size": args.graph_max_batch_size, - "mem_fraction": args.mem_fraction, - "max_req_num": 2048, - "batch_max_tokens": 1024, - "run_mode": "normal", - "max_seq_length": args.max_req_total_len, - "disable_cudagraph": args.disable_cudagraph, - "llm_prefill_att_backend": args.llm_prefill_att_backend, - "llm_decode_att_backend": args.llm_decode_att_backend, - "vit_att_backend": args.vit_att_backend, - "llm_kv_type": args.llm_kv_type, - "llm_kv_quant_group_size": args.llm_kv_quant_group_size, - } - proc = multiprocessing.Process( - target=tppart_model_infer, - args=(args, model_kvargs, args.batch_size, args.input_len, args.output_len, ans_queue), - ) - proc.start() - workers.append(proc) - - for proc in workers: - proc.join() - - assert not ans_queue.empty() - while not ans_queue.empty(): - assert ans_queue.get() - return - - -def overlap_prefill( - model_part, - batch_size, - max_len_in_batch, - input_ids, - mem_indexes, - b_req_idx, - b_mtp_index, - b_seq_len, - total_token_num, - b_ready_cache_len, -): - _0_batch_size = batch_size // 2 - _0_total_token_num = total_token_num // 2 - _0_input_ids = input_ids[: total_token_num // 2] - _0_mem_indexes = mem_indexes[: total_token_num // 2] - _0_b_req_idx = b_req_idx[: batch_size // 2] - _0_b_mtp_index = b_mtp_index[: batch_size // 2] - _0_b_seq_len = b_seq_len[: batch_size // 2] - _o_b_ready_cache_len = b_ready_cache_len[: batch_size // 2] - micro_batch1 = ModelInput( - batch_size=_0_batch_size, - total_token_num=_0_total_token_num, - input_ids=_0_input_ids, - b_req_idx=_0_b_req_idx, - b_mtp_index=_0_b_mtp_index, - b_seq_len=_0_b_seq_len, - is_prefill=True, - b_ready_cache_len=_o_b_ready_cache_len, - multimodal_params=[{"images": [], "audios": []} for _ in range(_0_batch_size)], - mem_indexes_cpu=_0_mem_indexes, - ) - - _1_batch_size = batch_size - batch_size // 2 - _1_total_token_num = total_token_num - total_token_num // 2 - _1_input_ids = input_ids[total_token_num // 2 :] - _1_mem_indexes = mem_indexes[total_token_num // 2 :] - _1_b_req_idx = b_req_idx[batch_size // 2 :] - _1_b_mtp_index = b_mtp_index[batch_size // 2 :] - _1_b_seq_len = b_seq_len[batch_size // 2 :] - _1_b_ready_cache_len = b_ready_cache_len[batch_size // 2 :] - - micro_batch2 = ModelInput( - batch_size=_1_batch_size, - total_token_num=_1_total_token_num, - input_ids=_1_input_ids, - b_req_idx=_1_b_req_idx, - b_mtp_index=_1_b_mtp_index, - b_seq_len=_1_b_seq_len, - is_prefill=True, - b_ready_cache_len=_1_b_ready_cache_len, - multimodal_params=[{"images": [], "audios": []} for _ in range(_1_batch_size)], - mem_indexes_cpu=_1_mem_indexes, - ) - - output, output1 = model_part.microbatch_overlap_prefill(micro_batch1, micro_batch2) - logits = output.logits - logits1 = output1.logits - return torch.cat((logits, logits1), dim=0) - - -def overlap_decode( - model_part, batch_size, max_len_in_batch, input_ids, mem_indexes, b_req_idx, b_mtp_index, b_seq_len, total_token_num -): - _0_batch_size = batch_size // 2 - _0_total_token_num = total_token_num // 2 - _0_input_ids = input_ids[: batch_size // 2] - _0_mem_indexes = mem_indexes[: batch_size // 2] - _0_b_req_idx = b_req_idx[: batch_size // 2] - _0_b_mtp_index = b_mtp_index[: batch_size // 2] - _0_b_seq_len = b_seq_len[: batch_size // 2] - micro_batch1 = ModelInput( - batch_size=_0_batch_size, - total_token_num=_0_total_token_num, - input_ids=_0_input_ids, - b_req_idx=_0_b_req_idx, - b_mtp_index=_0_b_mtp_index, - b_seq_len=_0_b_seq_len, - mem_indexes_cpu=_0_mem_indexes, - multimodal_params=[{"images": [], "audios": []} for _ in range(_0_batch_size)], - ) - - _1_batch_size = batch_size - batch_size // 2 - _1_total_token_num = total_token_num - total_token_num // 2 - _1_input_ids = input_ids[batch_size // 2 :] - _1_mem_indexes = mem_indexes[batch_size // 2 :] - _1_b_req_idx = b_req_idx[batch_size // 2 :] - _1_b_mtp_index = b_mtp_index[batch_size // 2 :] - _1_b_seq_len = b_seq_len[batch_size // 2 :] - - micro_batch2 = ModelInput( - batch_size=_1_batch_size, - total_token_num=_1_total_token_num, - input_ids=_1_input_ids, - b_req_idx=_1_b_req_idx, - b_mtp_index=_1_b_mtp_index, - b_seq_len=_1_b_seq_len, - mem_indexes_cpu=_1_mem_indexes, - multimodal_params=[{"images": [], "audios": []} for _ in range(_1_batch_size)], - ) - - output, output1 = model_part.microbatch_overlap_decode(micro_batch1, micro_batch2) - logits = output.logits - logits1 = output1.logits - return torch.cat((logits, logits1), dim=0) - - -def prefill( - model_part, - batch_size, - max_len_in_batch, - input_ids, - mem_indexes, - b_req_idx, - b_mtp_index, - b_seq_len, - total_token_num, - b_ready_cache_len, -): - b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cpu") - b_prefill_start_loc = b_seq_len.cumsum(dim=0, dtype=torch.int32) - b_seq_len - model_input = ModelInput( - batch_size=batch_size, - total_token_num=total_token_num, - max_q_seq_len=max_len_in_batch, - max_kv_seq_len=max_len_in_batch, - max_cache_len=0, - input_ids=input_ids, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mtp_index=b_mtp_index, - mem_indexes_cpu=mem_indexes, - is_prefill=True, - b_ready_cache_len=b_ready_cache_len, # b_ready_cache_len - b_prefill_start_loc=b_prefill_start_loc, - prefix_total_token_num=0, # the default kvcache len is zero. - multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], - ) - - model_output = model_part.forward(model_input) - return model_output.logits - - -def decode( - model_part, batch_size, max_len_in_batch, input_ids, mem_indexes, b_req_idx, b_mtp_index, b_seq_len, total_token_num -): - model_input = ModelInput( - batch_size=batch_size, - total_token_num=total_token_num, - max_q_seq_len=1, - max_kv_seq_len=max_len_in_batch, - input_ids=input_ids, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mtp_index=b_mtp_index, - mem_indexes_cpu=mem_indexes, - is_prefill=False, - multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], - ) - model_output = model_part.forward(model_input) - return model_output.logits - - -def torch_profile(fn, log_dir=None): - torch.cuda.synchronize() - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=False, - profile_memory=False, - on_trace_ready=torch.profiler.tensorboard_trace_handler(log_dir), - ) as prof: - fn() - if get_current_rank_in_dp() == 0: - print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) - - -def run_forward_once( - model_kvargs, input_len, output_len, batch_size, model_part, enable_overlap, enable_torch_profile=False -): - test_data = np.vstack([np.random.randint(0, 50256, input_len) for _ in range(batch_size)]) - test_data = test_data.reshape(-1) - test_data = torch.from_numpy(test_data) - import torch.distributed as dist - - dist.barrier() - import time - - dp_size = model_kvargs["dp_size"] - - torch.cuda.synchronize() - prefill_start_time = time.time() - - b_req_idx = torch.tensor( - [model_part.req_manager.alloc() for _ in range(batch_size)], dtype=torch.int32, device="cpu" - ) - b_seq_len = torch.zeros(batch_size, dtype=torch.int32, device="cpu") - b_ready_cache_len = torch.zeros(batch_size, dtype=torch.int32, device="cpu") - for i in range(batch_size): - b_seq_len[i] = input_len - - total_token_num = batch_size * input_len - mem_indexes = model_part.req_manager.mem_manager.alloc(test_data.shape[0]) - b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cpu") - rank_id = model_kvargs["rank_id"] - - if enable_overlap: - prefill_fn = overlap_prefill - decode_fn = overlap_decode - else: - prefill_fn = prefill - decode_fn = decode - - logits = prefill_fn( - model_part, - batch_size, - input_len, - test_data, - mem_indexes, - b_req_idx, - b_mtp_index, - b_seq_len, - total_token_num, - b_ready_cache_len, # b_ready_cache_len - ) - - prob_out = torch.softmax(logits, dim=-1) - predict_ids = torch.argmax(prob_out, dim=1, keepdim=True) - _ = predict_ids.detach().cpu().numpy() - - torch.cuda.synchronize() - - if rank_id == 0: - print( - f"prefill time cost: {(time.time() - prefill_start_time) * 1000}, " - f"prefill throughput: {dp_size * batch_size * input_len / (time.time() - prefill_start_time)} tokens/s" - ) - - if enable_torch_profile: - print("Profile Prefill") - try: - torch_profile( - lambda: prefill_fn( - model_part, - batch_size, - input_len, - test_data, - mem_indexes, - b_req_idx, - b_mtp_index, - b_seq_len, - total_token_num, - b_ready_cache_len, # b_ready_cache_len - ), - log_dir=f"./logs/forward_prefill_{model_kvargs['rank_id']}", - ) - except Exception as e: - print(str(e)) - raise - - for i in range(output_len): - torch.cuda.synchronize() - step_start = time.time() - total_token_num += batch_size - b_seq_len += 1 - mem_indexes = model_part.req_manager.mem_manager.alloc(predict_ids.shape[0]) - max_len_in_batch = input_len + i + 1 - logits = decode_fn( - model_part, - batch_size, - max_len_in_batch, - predict_ids.view(-1), - mem_indexes, - b_req_idx, - b_mtp_index, - b_seq_len, - total_token_num, - ) - if enable_torch_profile and i == output_len - 1: - try: - torch_profile( - lambda: decode_fn( - model_part, - batch_size, - max_len_in_batch, - predict_ids.view(-1), - mem_indexes, - b_req_idx, - b_mtp_index, - b_seq_len, - total_token_num, - ), - log_dir=f"./logs/forward_decode_{model_kvargs['rank_id']}", - ) - except Exception as e: - print(str(e)) - raise - - prob_out = torch.softmax(logits, dim=-1) - predict_ids = torch.argmax(prob_out, dim=1, keepdim=True) - _ = predict_ids.detach().cpu().numpy() - torch.cuda.synchronize() - if i % 100 == 0 or i == output_len - 1: - if rank_id == 0: - print( - f"i: {i}, step cost time: {(time.time() - step_start) * 1000} ms, " - f"throughput: {dp_size * batch_size / (time.time() - step_start)} tokens/s" - ) - - model_part.mem_manager.free_all() - model_part.req_manager.free_all() - torch.cuda.synchronize() - torch.cuda.empty_cache() - - -def tppart_model_infer(args, model_kvargs, batch_size, input_len, output_len, ans_queue): - args = get_env_start_args() - import triton.profiler as proton - import torch - from lightllm.distributed import dist_group_manager - from lightllm.utils.dist_utils import set_current_device_id - - if isinstance(batch_size, int): - batch_size = [batch_size] - else: - batch_size = [2, 8, 16, 32, 64, 128] - print(batch_size) - - import torch.distributed as dist - - enable_decode_overlap = args.enable_decode_microbatch_overlap - group_size = 1 - if enable_decode_overlap or args.enable_prefill_microbatch_overlap: - for bs in batch_size: - assert bs % 2 == 0, "batch size must be even number" - group_size = 2 - init_distributed_env(model_kvargs) - dist_group_manager.create_groups(group_size=group_size) - model_cfg, _ = PretrainedConfig.get_config_dict(model_kvargs["weight_dir"]) - dist.barrier() - - torch.cuda.empty_cache() - enable_overlap = args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap - - model_part, _ = get_model(model_cfg, model_kvargs) - - rank_id = model_kvargs["rank_id"] - for b in batch_size: - if rank_id == 0: - print(f"Testing batch size {b}") - - # warm up - run_forward_once( - model_kvargs, - input_len, - output_len=10, - batch_size=b, - model_part=model_part, - enable_overlap=enable_overlap, - enable_torch_profile=False, - ) - - # test - run_forward_once( - model_kvargs, - input_len, - output_len, - batch_size=b, - model_part=model_part, - enable_overlap=enable_overlap, - enable_torch_profile=args.torch_profile, - ) - if rank_id == 0: - print("=" * 50) - - ans_queue.put(True) - - try: - ans_queue.close() - ans_queue.join_thread() - except Exception: - pass - os._exit(0) diff --git a/test/benchmark/static_inference/model_infer_mtp.py b/test/benchmark/static_inference/model_infer_mtp.py deleted file mode 100644 index 72f06a919c..0000000000 --- a/test/benchmark/static_inference/model_infer_mtp.py +++ /dev/null @@ -1,285 +0,0 @@ -import os -import torch -import numpy as np -from multiprocessing import Queue -import multiprocessing -from transformers import PretrainedConfig -from lightllm.utils.dist_utils import init_distributed_env, get_current_rank_in_dp -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.models import get_model -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from lightllm.server.core.objs.start_args_type import StartArgs -from torch.profiler import profile, record_function, ProfilerActivity -from lightllm.utils.log_utils import init_logger -from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel -import torch.cuda as cuda - -logger = init_logger(__name__) - - -def init_mtp_model(args: StartArgs, kvargs, main_model): - mtp_step = args.mtp_step - draft_models = [] - - os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1" - mtp_model_kvargs = kvargs - mtp_model_kvargs.update( - { - "weight_dir": args.mtp_draft_model_dir, - "max_total_token_num": main_model.mem_manager.size, - "disable_chunked_prefill": True, - "mtp_mode": args.mtp_mode, - "main_model": main_model, - } - ) - for i in range(mtp_step): - mtp_model_cfg, _ = PretrainedConfig.get_config_dict(args.mtp_draft_model_dir) - mtp_model_kvargs.update( - { - "weight_dir": args.spec_model_dir, - "max_total_token_num": main_model.mem_manager.size, - "disable_chunked_prefill": True, - "mtp_mode": args.mtp_mode, - "main_model": main_model, - "mem_layer_start": main_model.config["num_hidden_layers"] + i * mtp_model_cfg["num_hidden_layers"], - } - ) - draft_models.append(Deepseek3MTPModel(mtp_model_kvargs)) - return draft_models - - -def test_model_inference_mtp(args): - ans_queue = Queue() - workers = [] - dp_size = args.get("dp", 1) - - for rank_id in range(args.node_rank * args.tp, (args.node_rank + 1) * args.tp): - model_kvargs = { - "args": args, - "nccl_host": args.nccl_host, - "data_type": args.data_type, - "nccl_port": args.nccl_port, - "rank_id": rank_id, - "world_size": args.tp, - "dp_size": dp_size, - "weight_dir": args.model_dir, - "quant_type": args.quant_type, - "load_way": "HF", - "max_total_token_num": args.max_total_token_num, - "graph_max_len_in_batch": args.max_req_total_len, - "graph_max_batch_size": args.graph_max_batch_size, - "mem_faction": args.mem_fraction, - "max_req_num": 2000, - "batch_max_tokens": 2048, - "run_mode": "normal", - "max_seq_length": args.max_req_total_len, - "spec_algo": args.spec_algo, - "disable_cudagraph": args.disable_cudagraph, - } - proc = multiprocessing.Process( - target=tppart_model_infer, - args=(args, model_kvargs, args.batch_size, args.input_len, args.output_len, ans_queue), - ) - proc.start() - workers.append(proc) - - for proc in workers: - proc.join() - - assert not ans_queue.empty() - while not ans_queue.empty(): - assert ans_queue.get() - return - - -def torch_profile(fn, log_dir=None): - torch.cuda.synchronize() - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=False, - profile_memory=False, - on_trace_ready=torch.profiler.tensorboard_trace_handler(log_dir), - ) as prof: - fn() - if get_current_rank_in_dp() == 0: - print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10)) - - -def run_forward_once(args, input_len, output_len, batch_size, main_model, draft_models, warmup=False): - import time - - torch.cuda.synchronize() - prefill_start_time = time.time() - - test_data = np.vstack([np.random.randint(0, 50256, input_len) for _ in range(batch_size)]) - test_data = test_data.reshape(-1) - test_data = torch.from_numpy(test_data).cuda() - - b_req_idx = torch.tensor( - [main_model.req_manager.alloc() for _ in range(batch_size)], dtype=torch.int32, device="cuda" - ) - b_seq_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda") - b_ready_cache_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda") - for i in range(batch_size): - b_seq_len[i] = input_len - - total_token_num = input_len * batch_size - mem_indexes = main_model.req_manager.mem_manager.alloc(test_data.shape[0]).cuda() - # Main model Prefill - model_input = ModelInput( - batch_size=batch_size, - total_token_num=total_token_num, - input_ids=test_data, - mem_indexes=mem_indexes, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - is_prefill=True, - b_ready_cache_len=b_ready_cache_len, - multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], - ) - - model_output: ModelOutput = main_model.forward(model_input) - prob_out = torch.softmax(model_output.logits, dim=-1) - predict_ids = torch.argmax(prob_out, dim=1, keepdim=True) - predict_ids = predict_ids.detach().cpu().numpy() - - draft_ids = [predict_ids] - - # Draft model Prefill - # For simplicity, we'll just take the input of main_model to draft model. - model_input.mtp_draft_input_hiddens = model_output.mtp_main_output_hiddens - for draft_model_id in range(len(draft_models)): - draft_model = draft_models[draft_model_id] - model_output = draft_model.forward(model_input) - prob_out = torch.softmax(model_output.logits, dim=-1) - predict_ids = torch.argmax(prob_out, dim=1, keepdim=True) - predict_ids = predict_ids.detach().cpu().numpy() - draft_ids.append(predict_ids) - model_input.mtp_draft_input_hiddens = model_output.mtp_main_output_hiddens - - torch.cuda.synchronize() - prefill_end_time = time.time() - if get_current_rank_in_dp() == 0 and not warmup: - print("prefill time cost:", (prefill_end_time - prefill_start_time) * 1000) - print( - f"Prefill throughput: {batch_size * input_len * args.dp / (prefill_end_time - prefill_start_time)} tokens/s" - ) - - torch.cuda.synchronize() - - decode_input_ids = np.stack(draft_ids, axis=-1).reshape(-1) - decode_input_ids = torch.from_numpy(decode_input_ids).cuda() - - # build main decode input: - nopad_b_seq_idx = [] - nopad_b_seq_len = [] - nopad_total_token_num = 0 - nopad_max_len_in_batch = 0 - - for i in range(batch_size): - nopad_b_seq_idx.append(b_req_idx[i]) - seq_len = b_seq_len[i].item() - nopad_b_seq_len.append(seq_len + 1) - nopad_total_token_num += seq_len + 1 - nopad_max_len_in_batch = max(nopad_max_len_in_batch, b_seq_len[i] + 1) - - for step in range(len(draft_models)): - nopad_b_seq_idx.append(b_req_idx[i]) - nopad_b_seq_len.append(seq_len + step + 2) - nopad_total_token_num += seq_len + step + 2 - nopad_max_len_in_batch = max(nopad_max_len_in_batch, seq_len + step + 2) - - nopad_b_seq_idx = torch.tensor(nopad_b_seq_idx, dtype=torch.int32, device="cuda") - nopad_b_seq_len = torch.tensor(nopad_b_seq_len, dtype=torch.int32, device="cuda") - mem_indexes = main_model.req_manager.mem_manager.alloc(batch_size * (len(draft_models) + 1)).cuda() - - model_input = ModelInput( - batch_size=batch_size * (len(draft_models) + 1), - total_token_num=nopad_total_token_num, - input_ids=decode_input_ids, - mem_indexes=mem_indexes, - b_req_idx=nopad_b_seq_idx, - b_seq_len=nopad_b_seq_len, - is_prefill=False, - multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size * (len(draft_models) + 1))], - ) - - # Main decode - for i in range(0, output_len, len(draft_models) + 1): - torch.cuda.synchronize() - step_start_time = time.time() - model_output = main_model.forward( - model_input, - ) - prob_out = torch.softmax(model_output.logits, dim=-1) - predict_ids = torch.argmax(prob_out, dim=1, keepdim=True) - - # draft decode - model_input.input_ids = predict_ids.reshape(-1) - model_input.mtp_draft_input_hiddens = model_output.mtp_main_output_hiddens - - for draft_model_id in range(len(draft_models)): - draft_model = draft_models[draft_model_id] - model_output = draft_model.forward( - model_input, - ) - prob_out = torch.softmax(model_output.logits, dim=-1) - predict_ids = torch.argmax(prob_out, dim=1, keepdim=True) - model_input.input_ids = predict_ids.reshape(-1) - model_input.mtp_draft_input_hiddens = model_output.mtp_main_output_hiddens - - # accept all draft ids by default. - model_input.input_ids = predict_ids.reshape(-1) - model_input.mtp_draft_input_hiddens = model_output.mtp_main_output_hiddens - torch.cuda.synchronize() - if i % 100 == 0 or i == output_len - 1: - step_end_time = time.time() - if get_current_rank_in_dp() == 0 and not warmup: - step_time = step_end_time - step_start_time - print(i, " step cost time:", step_time * 1000) - print(f"Decode throughput: {batch_size * (len(draft_models) + 1) * args.dp / step_time} tokens/s") - - main_model.mem_manager.free_all() - main_model.req_manager.free_all() - - -def tppart_model_infer(args, model_kvargs, batch_sizes, input_len, output_len, ans_queue): - args = get_env_start_args() - import triton.profiler as proton - import torch - from lightllm.distributed import dist_group_manager - from lightllm.utils.dist_utils import set_current_device_id - - import torch.distributed as dist - - enable_decode_overlap = args.enable_decode_microbatch_overlap - group_size = 1 - if enable_decode_overlap or args.enable_prefill_microbatch_overlap: - group_size = 2 - init_distributed_env(model_kvargs) - dist_group_manager.create_groups(group_size=group_size) - model_cfg, _ = PretrainedConfig.get_config_dict(model_kvargs["weight_dir"]) - dist.barrier() - - torch.cuda.empty_cache() - - main_model, _ = get_model(model_cfg, model_kvargs) - draft_models = init_mtp_model(args, model_kvargs, main_model) - if isinstance(batch_sizes, int): - batch_sizes = [batch_sizes] - - for batch_size in batch_sizes: - # warm up - run_forward_once(args, input_len, output_len, batch_size, main_model, draft_models, warmup=True) - torch.cuda.synchronize() - run_forward_once(args, input_len, output_len, batch_size, main_model, draft_models, warmup=False) - dist.barrier() - - ans_queue.put(True) - - try: - ans_queue.close() - ans_queue.join_thread() - except Exception: - pass - os._exit(0) diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py new file mode 100644 index 0000000000..1ea309d46c --- /dev/null +++ b/test/benchmark/static_inference/static_benchmark.py @@ -0,0 +1,1425 @@ +"""Static forward benchmark for LightLLM model parts. + +The entry uses synthetic token ids and measures forward-only TPS for prefill, +chunked prefill, decode, and MTP decode cases. +""" + +import argparse +import math +import os +import sys +import time +import traceback +from dataclasses import asdict, dataclass, replace +from pathlib import Path +from types import SimpleNamespace +from typing import Dict, List, Optional, Sequence + +import numpy as np +import torch +import torch.multiprocessing as mp +from transformers import PretrainedConfig + + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.append(str(REPO_ROOT)) + +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.models import get_model +from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel +from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel +from lightllm.models.mistral_mtp.model import MistralMTPModel +from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.router.model_infer.mode_backend.mtp_pre_process import ( + prepare_mtp_prefill_inputs, +) +from lightllm.utils.config_utils import get_dtype, get_vocab_size +from lightllm.utils.dist_utils import init_distributed_env +from lightllm.utils.envs_utils import set_env_start_args + + +DEFAULT_BATCH_SIZES = [2, 8, 16, 32, 64, 128] +MTP_MODES = {"vanilla_with_att", "eagle_with_att", "vanilla_no_att", "eagle_no_att"} +PREFILL_TABLE_HEADERS = [ + "ctx", + "hit", + "bs", + "max_total_token_num", + "uncached", + "cached", + "tokens", + "ms", + "qps", + "tok/s", + "logical_tok/s", +] +DECODE_TABLE_HEADERS = [ + "ctx", + "bs", + "accept", + "max_total_token_num", + "ms", + "qps", + "tok/s", + "itl_ms", +] + + +@dataclass(frozen=True) +class BenchmarkCase: + name: str + stage: str + batch_size: int + context_len: int + output_len: int + chunked_prefill_size: Optional[int] = None + profiled_max_total_token_num: Optional[int] = None + profiled_batch_divisor: Optional[int] = None + cache_hit_rate: float = 0.0 + prefill_uncached_len: Optional[int] = None + prefill_step_tokens_per_req: Optional[int] = None + prefill_batch_size_by_batch_max_tokens: Optional[int] = None + + +@dataclass +class BenchmarkResult: + case: str + stage: str + batch_size: int + context_len: int + output_len: int + chunked_prefill_size: Optional[int] + elapsed_ms: float + measured_tokens: int + qps: float + tps: float + profiled_max_total_token_num: Optional[int] = None + profiled_batch_divisor: Optional[int] = None + ttft_ms: Optional[float] = None + inter_token_latency_ms: Optional[float] = None + cache_hit_rate: float = 0.0 + prefill_uncached_len: Optional[int] = None + prefill_cached_len: Optional[int] = None + prefill_step_tokens_per_req: Optional[int] = None + mtp_accept_rate: Optional[float] = None + logical_tps: Optional[float] = None + + +class TokenSource: + def __init__(self, args: SimpleNamespace): + self.vocab_size = max(2, int(get_vocab_size(args.model_dir) or 0)) + self.rng = np.random.default_rng(args.seed) + + def batch(self, batch_size: int, need_len: int) -> np.ndarray: + return self.rng.integers(low=0, high=self.vocab_size, size=(batch_size, need_len), dtype=np.int64) + + +def cpu_i32_full(shape, value) -> torch.Tensor: + return torch.full(shape, value, dtype=torch.int32, device="cpu") + + +def cpu_i32_zeros(size: int) -> torch.Tensor: + return torch.zeros(size, dtype=torch.int32, device="cpu") + + +def empty_multimodal_params(batch_size: int) -> List[Dict]: + return [{"images": [], "audios": []} for _ in range(batch_size)] + + +class StaticBenchmarkExecutor: + def __init__( + self, + args: SimpleNamespace, + model, + draft_models: List, + token_source: TokenSource, + ): + self.args = args + self.model = model + self.draft_models = draft_models + self.token_source = token_source + self.dp = int(args.dp or 1) + + def _case_iters(self, warmup: bool) -> int: + return self.args.warmup_iters if warmup else self.args.bench_iters + + def run_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult: + if case.stage == "prefill": + return self._run_prefill_case(case, warmup) + if case.stage == "decode": + return self._run_decode_case(case, warmup) + raise ValueError(f"unknown benchmark stage: {case.stage}") + + def _run_prefill_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult: + """Measure full uncached prefill, chunked by production admission size.""" + uncached_len = int(case.prefill_uncached_len or case.context_len) + cached_len = max(0, case.context_len - uncached_len) + tokens = self.token_source.batch(case.batch_size, uncached_len) + elapsed = 0.0 + measured_tokens = case.batch_size * uncached_len + + for _ in range(self._case_iters(warmup)): + self._reset_model_cache() + req_idx = self._alloc_req_indexes(case.batch_size) + if cached_len > 0: + self._materialize_cached_prefix(req_idx, cached_len) + inputs = self._build_prefill_inputs( + token_rows=tokens, + req_idx=req_idx, + prompt_len=uncached_len, + chunk_size=case.chunked_prefill_size, + initial_ready_cache_len=cached_len, + ) + torch.cuda.synchronize() + start = time.perf_counter() + output = None + for model_input in inputs: + output = self._forward_prefill_input(model_input, allow_overlap=True) + torch.cuda.synchronize() + elapsed += time.perf_counter() - start + self._touch_output(output) + + self._reset_model_cache() + return self._make_result(case, elapsed, measured_tokens, warmup) + + def _run_decode_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult: + mtp_enabled = self._mtp_enabled() + token_rows = self.token_source.batch(case.batch_size, case.context_len) if mtp_enabled else None + measured_tokens = case.batch_size * case.output_len + elapsed = 0.0 + ttft_elapsed = 0.0 + decode_step_count = 0 + iters = self._case_iters(warmup) + + for _ in range(iters): + self._reset_model_cache() + if mtp_enabled: + torch.cuda.synchronize() + ttft_start = time.perf_counter() + req_idx, seq_len, next_ids = self._prefill_for_decode(case, token_rows, mtp_enabled) + torch.cuda.synchronize() + ttft_elapsed += time.perf_counter() - ttft_start + step_elapsed, step_count = self._run_mtp_decode_steps( + case=case, + req_idx=req_idx, + seq_len=seq_len, + next_ids=next_ids, + ) + elapsed += step_elapsed + decode_step_count += step_count + else: + req_idx, seq_len, next_ids = self._materialize_context_for_decode(case) + elapsed += self._run_plain_decode_steps( + case=case, + req_idx=req_idx, + seq_len=seq_len, + next_ids=next_ids, + ) + decode_step_count += case.output_len + + self._reset_model_cache() + inter_token_latency_ms = elapsed * 1000.0 / max(1, decode_step_count) if iters > 0 else None + return self._make_result( + case, + elapsed, + measured_tokens, + warmup, + ttft_elapsed_s=ttft_elapsed if mtp_enabled else None, + inter_token_latency_ms=inter_token_latency_ms, + ) + + def _materialize_context_for_decode(self, case: BenchmarkCase): + """Allocate historical KV slots so decode can be measured without prefill.""" + req_idx = self._alloc_req_indexes(case.batch_size) + self._materialize_cached_prefix(req_idx, case.context_len) + seq_len = cpu_i32_full((case.batch_size,), case.context_len) + next_ids = torch.from_numpy(np.ascontiguousarray(self.token_source.batch(case.batch_size, 1).reshape(-1))).to( + torch.int64 + ) + return req_idx, seq_len, next_ids + + def _prefill_for_decode(self, case: BenchmarkCase, token_rows: np.ndarray, mtp_enabled: bool): + req_idx = self._alloc_req_indexes(case.batch_size) + inputs = self._build_prefill_inputs( + token_rows=token_rows, + req_idx=req_idx, + prompt_len=case.context_len, + chunk_size=None, + ) + output = None + for model_input in inputs: + output = self._forward_prefill_input(model_input, allow_overlap=not mtp_enabled) + assert output is not None + self._touch_output(output) + + next_ids = self._argmax_ids(output.logits) + + seq_len = cpu_i32_full((case.batch_size,), case.context_len) + if mtp_enabled: + next_ids = self._fill_mtp_prefill_kv(case, inputs[-1], output, next_ids) + return req_idx, seq_len, next_ids + + def _fill_mtp_prefill_kv( + self, + case: BenchmarkCase, + main_prefill_input: ModelInput, + main_output: ModelOutput, + first_next_ids: torch.Tensor, + ): + draft_input = main_prefill_input + draft_output = main_output + current_next_ids = first_next_ids.cuda(non_blocking=True) + mtp_candidates = [current_next_ids.detach().cpu()] + for draft_index in range(self._num_mtp_modules()): + draft_input = prepare_mtp_prefill_inputs( + model_input=draft_input, + b_next_token_ids=current_next_ids, + mtp_draft_input_hiddens=draft_output.mtp_main_output_hiddens, + ) + draft_output = self.draft_models[draft_index].forward(draft_input) + current_next_ids = self._argmax_ids(draft_output.logits).cuda(non_blocking=True) + mtp_candidates.append(current_next_ids.detach().cpu()) + + step_width = self._mtp_step_width() + while len(mtp_candidates) < step_width: + mtp_candidates.append(mtp_candidates[-1]) + next_ids = torch.stack(mtp_candidates[:step_width], dim=1) + return next_ids + + def _run_plain_decode_steps( + self, + case: BenchmarkCase, + req_idx: torch.Tensor, + seq_len: torch.Tensor, + next_ids: torch.Tensor, + ) -> float: + elapsed = 0.0 + for step in range(case.output_len): + seq_len += 1 + model_input = self._make_decode_input( + batch_size=case.batch_size, + req_idx=req_idx, + mtp_index=cpu_i32_zeros(case.batch_size), + seq_len=seq_len, + input_ids=next_ids.reshape(-1), + max_kv_seq_len=int(seq_len.max().item()), + mem_token_num=case.batch_size, + ) + torch.cuda.synchronize() + start = time.perf_counter() + output = self._forward_decode_input(model_input, allow_overlap=True) + torch.cuda.synchronize() + elapsed += time.perf_counter() - start + self._touch_output(output) + + next_ids = self._argmax_ids(output.logits) + return elapsed + + def _run_mtp_decode_steps( + self, + case: BenchmarkCase, + req_idx: torch.Tensor, + seq_len: torch.Tensor, + next_ids: torch.Tensor, + ) -> tuple: + elapsed = 0.0 + step_count = 0 + generated_len = 0 + step_width = self._mtp_step_width() + base_req_idx, b_mtp_index = self._build_mtp_decode_index_tensors(req_idx, step_width) + current_candidates = next_ids + + while generated_len < case.output_len: + accepted_width = self._sample_mtp_accept_width(step_width, case.output_len - generated_len) + if current_candidates.ndim == 1: + current_candidates = current_candidates[:, None].repeat(1, step_width) + + b_seq_len = self._build_mtp_seq_len(seq_len, step_width) + model_input = self._make_decode_input( + batch_size=case.batch_size * step_width, + req_idx=base_req_idx, + mtp_index=b_mtp_index, + seq_len=b_seq_len, + input_ids=current_candidates.reshape(-1), + max_kv_seq_len=int(b_seq_len.max().item()), + mem_token_num=case.batch_size * step_width, + ) + + torch.cuda.synchronize() + start = time.perf_counter() + output = self.model.forward(model_input) + candidate_rows, temporary_mem = self._run_mtp_draft_decode( + model_input=model_input, + model_output=output, + real_batch_size=case.batch_size, + step_width=step_width, + ) + torch.cuda.synchronize() + elapsed += time.perf_counter() - start + self._touch_output(output) + if temporary_mem is not None: + self.model.req_manager.mem_manager.free(temporary_mem) + + self._free_rejected_mtp_mem( + model_input=model_input, + real_batch_size=case.batch_size, + step_width=step_width, + accepted_width=accepted_width, + ) + current_candidates = ( + self._select_mtp_candidates( + candidate_rows=candidate_rows, + real_batch_size=case.batch_size, + step_width=step_width, + accepted_width=accepted_width, + ) + .detach() + .cpu() + ) + seq_len += accepted_width + generated_len += accepted_width + step_count += 1 + + return elapsed, step_count + + def _run_mtp_draft_decode( + self, + model_input: ModelInput, + model_output: ModelOutput, + real_batch_size: int, + step_width: int, + ): + draft_input = model_input + draft_output = model_output + draft_next_ids = self._argmax_ids(model_output.logits).cuda(non_blocking=True) + generated = [draft_next_ids.detach()] + + temporary_mem = None + if self.args.mtp_mode.startswith("eagle"): + extra_mem_cpu = self.model.req_manager.mem_manager.alloc(real_batch_size * self.args.mtp_step) + temporary_mem = extra_mem_cpu + extra_mem = extra_mem_cpu.cuda(non_blocking=True) + else: + extra_mem = None + + for step in range(self.args.mtp_step): + draft_input.input_ids = draft_next_ids + draft_input.mtp_draft_input_hiddens = draft_output.mtp_main_output_hiddens + draft_model = self.draft_models[step % self._num_mtp_modules()] + draft_output = draft_model.forward(draft_input) + draft_next_ids = self._argmax_ids(draft_output.logits).cuda(non_blocking=True) + generated.append(draft_next_ids.detach()) + + if self.args.mtp_mode.startswith("eagle") and step + 1 < self.args.mtp_step: + draft_input.b_seq_len += 1 + draft_input.max_kv_seq_len += 1 + mem_i = extra_mem[step * real_batch_size : (step + 1) * real_batch_size] + draft_input.mem_indexes = torch.cat( + [ + draft_input.mem_indexes.view(-1, step_width)[:, 1:], + mem_i.view(-1, 1), + ], + dim=1, + ).reshape(-1) + + return torch.stack(generated[:step_width], dim=1), temporary_mem + + def _sample_mtp_accept_width(self, step_width: int, remaining_tokens: int) -> int: + """Sample accepted MTP width outside the timed decode section.""" + accept_rate = float(self.args.mtp_accept_rate) + accepted_width = 1 + for _ in range(step_width - 1): + if self.token_source.rng.random() >= accept_rate: + break + accepted_width += 1 + return max(1, min(accepted_width, remaining_tokens)) + + def _select_mtp_candidates( + self, + candidate_rows: torch.Tensor, + real_batch_size: int, + step_width: int, + accepted_width: int, + ) -> torch.Tensor: + row_ids = torch.arange(real_batch_size, device=candidate_rows.device) * step_width + accepted_width - 1 + return candidate_rows.index_select(0, row_ids) + + def _free_rejected_mtp_mem( + self, + model_input: ModelInput, + real_batch_size: int, + step_width: int, + accepted_width: int, + ): + if accepted_width >= step_width: + return + rejected_mem = ( + model_input.mem_indexes_cpu.view(real_batch_size, step_width)[:, accepted_width:].contiguous().reshape(-1) + ) + if rejected_mem.numel() > 0: + self.model.req_manager.mem_manager.free(rejected_mem) + + def _build_prefill_inputs( + self, + token_rows: np.ndarray, + req_idx: torch.Tensor, + prompt_len: int, + chunk_size: Optional[int], + initial_ready_cache_len: int = 0, + ) -> List[ModelInput]: + if not chunk_size or chunk_size <= 0 or chunk_size >= prompt_len: + return [ + self._make_prefill_input( + token_rows[:, :prompt_len], + req_idx, + ready_cache_len=initial_ready_cache_len, + ) + ] + + inputs = [] + for start in range(0, prompt_len, chunk_size): + end = min(prompt_len, start + chunk_size) + inputs.append( + self._make_prefill_input( + token_rows[:, start:end], + req_idx, + ready_cache_len=initial_ready_cache_len + start, + ) + ) + return inputs + + def _materialize_cached_prefix(self, req_idx: torch.Tensor, cached_len: int): + """Allocate dummy prefix KV so cache-hit cases consume real capacity.""" + if cached_len <= 0: + return + batch_size = int(req_idx.shape[0]) + need_tokens = batch_size * cached_len + mem_indexes = self.model.req_manager.mem_manager.alloc(need_tokens) + if mem_indexes is None: + raise RuntimeError(f"failed to allocate cached prefix: bs={batch_size} cached_len={cached_len}") + req_idx_gpu = req_idx.cuda(non_blocking=True) + mem_indexes_gpu = mem_indexes.reshape(batch_size, cached_len).cuda(non_blocking=True) + self.model.req_manager.req_to_token_indexs[req_idx_gpu, :cached_len] = mem_indexes_gpu + + def _make_prefill_input(self, token_chunk: np.ndarray, req_idx: torch.Tensor, ready_cache_len: int) -> ModelInput: + batch_size, q_len = token_chunk.shape + seq_len_value = ready_cache_len + q_len + b_seq_len = cpu_i32_full((batch_size,), seq_len_value) + b_ready_cache_len = cpu_i32_full((batch_size,), ready_cache_len) + b_q_seq_len = b_seq_len - b_ready_cache_len + b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len + input_ids = torch.from_numpy(np.ascontiguousarray(token_chunk.reshape(-1))).to(torch.int64) + mem_indexes = self.model.req_manager.mem_manager.alloc(input_ids.shape[0]) + return ModelInput( + batch_size=batch_size, + total_token_num=int(b_seq_len.sum().item()), + max_q_seq_len=q_len, + max_kv_seq_len=seq_len_value, + max_cache_len=ready_cache_len, + prefix_total_token_num=ready_cache_len * batch_size, + input_ids=input_ids, + b_req_idx=req_idx, + b_mtp_index=cpu_i32_zeros(batch_size), + b_seq_len=b_seq_len, + mem_indexes_cpu=mem_indexes, + is_prefill=True, + b_ready_cache_len=b_ready_cache_len, + b_prefill_start_loc=b_prefill_start_loc, + b_prefill_has_output_cpu=[False] * batch_size, + multimodal_params=empty_multimodal_params(batch_size), + ) + + def _make_decode_input( + self, + batch_size: int, + req_idx: torch.Tensor, + mtp_index: torch.Tensor, + seq_len: torch.Tensor, + input_ids: torch.Tensor, + max_kv_seq_len: int, + mem_token_num: int, + ) -> ModelInput: + mem_indexes = self.model.req_manager.mem_manager.alloc(mem_token_num) + return ModelInput( + batch_size=batch_size, + total_token_num=int(seq_len.sum().item()), + max_q_seq_len=1, + max_kv_seq_len=max_kv_seq_len, + input_ids=input_ids.to(torch.int64).cpu(), + b_req_idx=req_idx, + b_mtp_index=mtp_index, + b_seq_len=seq_len, + mem_indexes_cpu=mem_indexes, + is_prefill=False, + multimodal_params=empty_multimodal_params(batch_size), + ) + + def _forward_prefill_input(self, model_input: ModelInput, allow_overlap: bool) -> ModelOutput: + if allow_overlap and self.args.enable_prefill_microbatch_overlap and model_input.batch_size > 1: + micro_input0, micro_input1 = self._split_prefill_input(model_input) + output0, output1 = self.model.microbatch_overlap_prefill(micro_input0, micro_input1) + return self._merge_model_outputs(output0, output1) + return self.model.forward(model_input) + + def _forward_decode_input(self, model_input: ModelInput, allow_overlap: bool) -> ModelOutput: + if allow_overlap and self.args.enable_decode_microbatch_overlap and model_input.batch_size > 1: + micro_input0, micro_input1 = self._split_decode_input(model_input) + output0, output1 = self.model.microbatch_overlap_decode(micro_input0, micro_input1) + return self._merge_model_outputs(output0, output1) + return self.model.forward(model_input) + + def _split_prefill_input(self, model_input: ModelInput): + split_batch = model_input.batch_size // 2 + q_lens = model_input.b_seq_len - model_input.b_ready_cache_len + split_tokens = int(q_lens[:split_batch].sum().item()) + return ( + self._slice_prefill_input(model_input, 0, split_batch, 0, split_tokens), + self._slice_prefill_input( + model_input, + split_batch, + model_input.batch_size, + split_tokens, + int(q_lens.sum().item()), + ), + ) + + def _slice_prefill_input( + self, + model_input: ModelInput, + batch_start: int, + batch_end: int, + token_start: int, + token_end: int, + ) -> ModelInput: + b_seq_len = model_input.b_seq_len[batch_start:batch_end].clone() + b_ready_cache_len = model_input.b_ready_cache_len[batch_start:batch_end].clone() + b_q_seq_len = b_seq_len - b_ready_cache_len + b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len + has_output = model_input.b_prefill_has_output_cpu + return ModelInput( + batch_size=batch_end - batch_start, + total_token_num=int(b_seq_len.sum().item()), + max_q_seq_len=int(b_q_seq_len.max().item()), + max_kv_seq_len=int(b_seq_len.max().item()), + max_cache_len=int(b_ready_cache_len.max().item()), + prefix_total_token_num=int(b_ready_cache_len.sum().item()), + input_ids=model_input.input_ids[token_start:token_end].contiguous(), + b_req_idx=model_input.b_req_idx[batch_start:batch_end].clone(), + b_mtp_index=model_input.b_mtp_index[batch_start:batch_end].clone(), + b_seq_len=b_seq_len, + mem_indexes_cpu=model_input.mem_indexes_cpu[token_start:token_end].contiguous(), + is_prefill=True, + b_ready_cache_len=b_ready_cache_len, + b_prefill_start_loc=b_prefill_start_loc, + b_prefill_has_output_cpu=(has_output[batch_start:batch_end] if has_output is not None else None), + multimodal_params=model_input.multimodal_params[batch_start:batch_end], + ) + + def _split_decode_input(self, model_input: ModelInput): + split_batch = model_input.batch_size // 2 + return ( + self._slice_decode_input(model_input, 0, split_batch), + self._slice_decode_input(model_input, split_batch, model_input.batch_size), + ) + + def _slice_decode_input(self, model_input: ModelInput, batch_start: int, batch_end: int) -> ModelInput: + b_seq_len = model_input.b_seq_len[batch_start:batch_end].clone() + input_ids = model_input.input_ids + if input_ids is not None: + input_ids = input_ids[batch_start:batch_end].contiguous() + return ModelInput( + batch_size=batch_end - batch_start, + total_token_num=int(b_seq_len.sum().item()), + max_q_seq_len=model_input.max_q_seq_len, + max_kv_seq_len=int(b_seq_len.max().item()), + input_ids=input_ids, + b_req_idx=model_input.b_req_idx[batch_start:batch_end].clone(), + b_mtp_index=model_input.b_mtp_index[batch_start:batch_end].clone(), + b_seq_len=b_seq_len, + mem_indexes_cpu=model_input.mem_indexes_cpu[batch_start:batch_end].contiguous(), + is_prefill=False, + multimodal_params=model_input.multimodal_params[batch_start:batch_end], + ) + + def _merge_model_outputs(self, output0: ModelOutput, output1: ModelOutput) -> ModelOutput: + mtp_hiddens = None + if output0.mtp_main_output_hiddens is not None and output1.mtp_main_output_hiddens is not None: + mtp_hiddens = torch.cat( + (output0.mtp_main_output_hiddens, output1.mtp_main_output_hiddens), + dim=0, + ) + return ModelOutput( + logits=torch.cat((output0.logits, output1.logits), dim=0), + prefill_mem_indexes_ready_event=output0.prefill_mem_indexes_ready_event, + mtp_main_output_hiddens=mtp_hiddens, + ) + + def _build_mtp_decode_index_tensors(self, req_idx: torch.Tensor, step_width: int): + batch_size = int(req_idx.shape[0]) + return ( + req_idx.repeat_interleave(step_width).to(torch.int32).cpu(), + torch.arange(step_width, dtype=torch.int32).repeat(batch_size), + ) + + def _build_mtp_seq_len(self, base_seq_len: torch.Tensor, step_width: int) -> torch.Tensor: + offsets = torch.arange(1, step_width + 1, dtype=torch.int32) + return (base_seq_len[:, None].to(torch.int32) + offsets[None, :]).reshape(-1) + + def _alloc_req_indexes(self, batch_size: int) -> torch.Tensor: + req_indexes = [self.model.req_manager.alloc() for _ in range(batch_size)] + if any(index is None for index in req_indexes): + raise RuntimeError(f"failed to allocate {batch_size} request indexes") + return torch.tensor(req_indexes, dtype=torch.int32, device="cpu") + + def _reset_model_cache(self): + self.model.mem_manager.free_all() + self.model.req_manager.free_all() + torch.cuda.synchronize() + torch.cuda.empty_cache() + + def _argmax_ids(self, logits: torch.Tensor) -> torch.Tensor: + return torch.argmax(logits, dim=-1).detach().cpu().to(torch.int64) + + def _touch_output(self, output: Optional[ModelOutput]): + if output is not None and output.logits is not None: + _ = output.logits.shape + + def _make_result( + self, + case: BenchmarkCase, + elapsed_s: float, + measured_tokens: int, + warmup: bool, + ttft_elapsed_s: Optional[float] = None, + inter_token_latency_ms: Optional[float] = None, + ) -> BenchmarkResult: + """Convert raw timings into reported TPS and latency metrics.""" + iters = self._case_iters(warmup) + scaled_tokens = measured_tokens * self.dp * iters + qps = case.batch_size * self.dp * iters / elapsed_s if elapsed_s > 0 else 0.0 + tps = scaled_tokens / elapsed_s if elapsed_s > 0 else 0.0 + ttft_ms = ttft_elapsed_s * 1000.0 / max(1, iters) if ttft_elapsed_s is not None else None + logical_tps = None + prefill_uncached_len = case.prefill_uncached_len + prefill_cached_len = None + if case.stage == "prefill": + uncached_len = int(case.prefill_uncached_len or case.context_len) + prefill_uncached_len = uncached_len + prefill_cached_len = max(0, case.context_len - uncached_len) + token_count = case.batch_size * case.context_len * self.dp * iters + logical_tps = token_count / elapsed_s if elapsed_s > 0 else 0.0 + return BenchmarkResult( + case=case.name, + stage=case.stage, + batch_size=case.batch_size, + context_len=case.context_len, + output_len=case.output_len, + chunked_prefill_size=case.chunked_prefill_size, + elapsed_ms=elapsed_s * 1000.0, + measured_tokens=scaled_tokens, + qps=qps, + tps=tps, + profiled_max_total_token_num=case.profiled_max_total_token_num, + profiled_batch_divisor=case.profiled_batch_divisor, + ttft_ms=ttft_ms, + inter_token_latency_ms=inter_token_latency_ms, + cache_hit_rate=case.cache_hit_rate, + prefill_uncached_len=prefill_uncached_len, + prefill_cached_len=prefill_cached_len, + prefill_step_tokens_per_req=case.prefill_step_tokens_per_req, + mtp_accept_rate=( + float(self.args.mtp_accept_rate) if case.stage == "decode" and self._mtp_enabled() else None + ), + logical_tps=logical_tps, + ) + + def _mtp_enabled(self) -> bool: + return self.args.mtp_mode in MTP_MODES and self.args.mtp_step > 0 + + def _mtp_step_width(self) -> int: + return int(self.args.mtp_step) + 1 + + def _num_mtp_modules(self) -> int: + if not self._mtp_enabled(): + return 0 + if self.args.mtp_mode.startswith("eagle"): + return 1 + return int(self.args.mtp_step) + + +def parse_typed_list(value: Optional[str], fallback: Sequence, cast) -> List: + if value is None or value == "": + return list(fallback) + if isinstance(value, cast): + return [value] + normalized = str(value).replace(",", " ") + return [cast(item) for item in normalized.split() if item.strip()] + + +def parse_int_list(value: Optional[str], fallback: Sequence[int]) -> List[int]: + return parse_typed_list(value, fallback, int) + + +def parse_float_list(value: Optional[str], fallback: Sequence[float]) -> List[float]: + return parse_typed_list(value, fallback, float) + + +def parse_chunk_sizes(value: Optional[str], fallback: Optional[int]) -> List[Optional[int]]: + if value is None: + return [fallback] if fallback else [None] + chunks: List[Optional[int]] = [] + for item in str(value).replace(",", " ").split(): + item = item.strip().lower() + if item in {"none", "full", "0", "-1"}: + chunks.append(None) + else: + chunks.append(int(item)) + return chunks or [None] + + +def prefill_uncached_len(context_len: int, cache_hit_rate: float) -> int: + """Return the uncached suffix length for a prompt-cache hit ratio.""" + if cache_hit_rate < 0.0 or cache_hit_rate >= 1.0: + raise ValueError(f"cache hit rate must satisfy 0 <= hit < 1, got {cache_hit_rate}") + uncached = int(math.ceil(context_len * (1.0 - cache_hit_rate))) + return max(1, min(context_len, uncached)) + + +def prefill_step_tokens_per_req(uncached_len: int, chunked_prefill_size: Optional[int]) -> int: + """Return tokens handled per request in one production prefill step.""" + if chunked_prefill_size and chunked_prefill_size > 0: + return max(1, min(uncached_len, int(chunked_prefill_size))) + return max(1, uncached_len) + + +def format_cache_hit_suffix(cache_hit_rate: float) -> str: + return f"{cache_hit_rate:.4f}".rstrip("0").rstrip(".").replace(".", "p") + + +def apply_max_batch_size(batch_size: int, max_batch_size: int) -> int: + """Apply the benchmark-wide auto batch-size upper bound.""" + if max_batch_size > 0: + batch_size = min(batch_size, int(max_batch_size)) + return max(1, batch_size) + + +def prefill_batch_size_from_batch_max_tokens( + batch_max_tokens: int, + step_tokens_per_req: int, + max_batch_size: int, +) -> int: + """Compute prefill BS from batch_max_tokens before KV-capacity capping.""" + batch_size = max(1, int(batch_max_tokens) // max(1, step_tokens_per_req)) + return apply_max_batch_size(batch_size, max_batch_size) + + +def build_prefill_cases( + args: SimpleNamespace, + input_lens: Sequence[int], + chunk_sizes: Sequence[Optional[int]], + cache_hit_rates: Sequence[float], +) -> List[BenchmarkCase]: + """Build full-prefill cases using batch_max_tokens per chunk step.""" + if args.batch_max_tokens is None: + raise ValueError("prefill benchmark requires --batch_max_tokens") + cases: List[BenchmarkCase] = [] + for input_len in input_lens: + for chunk_size in chunk_sizes: + for cache_hit_rate in cache_hit_rates: + uncached_len = prefill_uncached_len(input_len, cache_hit_rate) + step_tokens = prefill_step_tokens_per_req(uncached_len, chunk_size) + bs = prefill_batch_size_from_batch_max_tokens( + args.batch_max_tokens, + step_tokens, + args.max_batch_size, + ) + chunk_name = chunk_size if chunk_size else "none" + hit_name = format_cache_hit_suffix(cache_hit_rate) + cases.append( + BenchmarkCase( + name=( + f"prefill_bs{bs}_in{input_len}_hit{hit_name}" + f"_uncached{uncached_len}_chunk{chunk_name}" + f"_btok{args.batch_max_tokens}" + ), + stage="prefill", + batch_size=bs, + context_len=input_len, + output_len=0, + chunked_prefill_size=chunk_size, + cache_hit_rate=cache_hit_rate, + prefill_uncached_len=uncached_len, + prefill_step_tokens_per_req=step_tokens, + prefill_batch_size_by_batch_max_tokens=bs, + ) + ) + return cases + + +def build_decode_cases( + args: SimpleNamespace, + batch_sizes: Sequence[int], + context_lens: Sequence[int], + output_lens: Sequence[int], +) -> List[BenchmarkCase]: + """Build decode cases; profile mode resolves BS after model load.""" + decode_batch_sizes = [1] if args.decode_batch_size_mode == "profile" else batch_sizes + cases: List[BenchmarkCase] = [] + for bs in decode_batch_sizes: + for context_len in context_lens: + for output_len in output_lens: + profile_suffix = "_profilebs" if args.decode_batch_size_mode == "profile" else "" + cases.append( + BenchmarkCase( + name=(f"decode_bs{bs}_ctx{context_len}_out{output_len}" f"{profile_suffix}"), + stage="decode", + batch_size=bs, + context_len=context_len, + output_len=output_len, + ) + ) + return cases + + +def build_cases(args: SimpleNamespace) -> List[BenchmarkCase]: + """Expand CLI list options into concrete prefill/decode benchmark cases.""" + batch_sizes = parse_int_list(args.batch_sizes, [args.batch_size] if args.batch_size else DEFAULT_BATCH_SIZES) + input_lens = parse_int_list(args.input_lens, [args.input_len]) + context_lens = parse_int_list(args.context_lens, input_lens) + output_lens = parse_int_list(args.output_lens, [args.output_len]) + chunk_sizes = parse_chunk_sizes(args.chunked_prefill_sizes, args.chunked_prefill_size) + cache_hit_rates = parse_float_list(args.prefill_cache_hit_rates, [0.0]) + + cases: List[BenchmarkCase] = [] + if args.benchmark in {"all", "prefill"}: + cases.extend(build_prefill_cases(args, input_lens, chunk_sizes, cache_hit_rates)) + if args.benchmark in {"all", "decode"}: + cases.extend(build_decode_cases(args, batch_sizes, context_lens, output_lens)) + return cases + + +def decode_profile_batch_divisor(args: SimpleNamespace, case: BenchmarkCase) -> int: + """Reserve KV capacity for context, generated tokens, and MTP expansion.""" + mtp_width = int(args.mtp_step) + 1 if args.mtp_mode in MTP_MODES else 1 + return max(1, case.context_len + case.output_len + mtp_width + 8) + + +def resolve_profile_decode_cases( + args: SimpleNamespace, + cases: Sequence[BenchmarkCase], + profiled_max_total_token_num: int, +) -> List[BenchmarkCase]: + """Replace decode profile placeholders with capacity-derived max BS.""" + if args.decode_batch_size_mode != "profile": + return list(cases) + + resolved: List[BenchmarkCase] = [] + for case in cases: + if case.stage != "decode": + resolved.append(case) + continue + + divisor = decode_profile_batch_divisor(args, case) + batch_size = max(1, int(profiled_max_total_token_num) // divisor) + batch_size = apply_max_batch_size(batch_size, args.max_batch_size) + + resolved.append( + replace( + case, + name=( + f"decode_bs{batch_size}_ctx{case.context_len}" + f"_out{case.output_len}_profile{profiled_max_total_token_num}" + ), + batch_size=batch_size, + profiled_max_total_token_num=int(profiled_max_total_token_num), + profiled_batch_divisor=divisor, + ) + ) + + return resolved + + +def resolve_batch_max_prefill_cases( + args: SimpleNamespace, + cases: Sequence[BenchmarkCase], + profiled_max_total_token_num: int, +) -> List[BenchmarkCase]: + """Cap prefill BS by profiled KV capacity after the model is loaded.""" + resolved: List[BenchmarkCase] = [] + capacity_tokens = int(profiled_max_total_token_num) + for case in cases: + if case.stage != "prefill": + resolved.append(case) + continue + + if case.context_len <= 0: + raise ValueError(f"invalid prefill context_len={case.context_len}") + bs_by_capacity = capacity_tokens // case.context_len + if bs_by_capacity <= 0: + raise ValueError( + "single prefill request does not fit profiled token capacity: " + f"context_len={case.context_len} capacity={profiled_max_total_token_num}" + ) + + bs_by_batch = int(case.prefill_batch_size_by_batch_max_tokens or case.batch_size) + batch_size = min(bs_by_batch, bs_by_capacity) + batch_size = apply_max_batch_size(batch_size, args.max_batch_size) + + chunk_name = case.chunked_prefill_size if case.chunked_prefill_size else "none" + hit_name = format_cache_hit_suffix(case.cache_hit_rate) + resolved.append( + replace( + case, + name=( + f"prefill_bs{batch_size}_in{case.context_len}_hit{hit_name}" + f"_uncached{case.prefill_uncached_len}_chunk{chunk_name}" + f"_btok{args.batch_max_tokens}" + f"_cap{profiled_max_total_token_num}" + ), + batch_size=batch_size, + profiled_max_total_token_num=int(profiled_max_total_token_num), + profiled_batch_divisor=case.context_len, + ) + ) + + return resolved + + +def normalize_args(args: argparse.Namespace, cases: Sequence[BenchmarkCase]) -> SimpleNamespace: + """Fill LightLLM startup args needed before model construction.""" + if args.data_type is None: + args.data_type = get_dtype(args.model_dir) + + if args.quant_type is None: + args.quant_type = "none" + + if not 0.0 <= float(args.mtp_accept_rate) <= 1.0: + raise ValueError(f"--mtp_accept_rate must be in [0, 1], got {args.mtp_accept_rate}") + + max_batch = max(case.batch_size for case in cases) + max_context = max(case.context_len for case in cases) + max_output = max(case.output_len for case in cases) + mtp_width = (args.mtp_step + 1) if args.mtp_mode in MTP_MODES else 1 + max_runtime_len = max_context + max_output + mtp_width + 2 + + if args.max_req_total_len is None: + args.max_req_total_len = max_runtime_len + else: + args.max_req_total_len = max(args.max_req_total_len, max_runtime_len) + + if args.graph_max_len_in_batch == 0: + args.graph_max_len_in_batch = args.max_req_total_len + + max_prefill_chunk = ( + max( + min(case.context_len, case.chunked_prefill_size or case.context_len) + for case in cases + if case.stage == "prefill" + ) + if any(case.stage == "prefill" for case in cases) + else max_context + ) + if args.batch_max_tokens is None: + args.batch_max_tokens = max(max_batch * max_prefill_chunk, max_batch * mtp_width, 1) + + decode_batch_size_needs_profile = ( + args.benchmark in {"all", "decode"} + and args.decode_batch_size_mode == "profile" + and args.max_total_token_num is None + ) + prefill_batch_size_needs_profile = args.benchmark in {"all", "prefill"} and args.max_total_token_num is None + needs_profiled_batch_size = decode_batch_size_needs_profile or prefill_batch_size_needs_profile + + if args.max_total_token_num is None and not needs_profiled_batch_size: + args.max_total_token_num = max_batch * (args.max_req_total_len + mtp_width + 8) + if args.max_total_token_num is not None: + args.max_total_token_num = max(args.max_total_token_num, args.batch_max_tokens + 1, args.max_req_total_len) + + if decode_batch_size_needs_profile and args.max_batch_size > 0: + args.running_max_req_size = max(args.running_max_req_size, int(args.max_batch_size)) + # Profile decode BS is resolved after model load. Use the cap as the + # pre-load upper bound so request slots and optional decode graphs agree. + if not args.disable_cudagraph: + args.graph_max_batch_size = max(args.graph_max_batch_size, int(args.max_batch_size)) + if prefill_batch_size_needs_profile: + args.running_max_req_size = max(args.running_max_req_size, max_batch) + + if args.graph_max_batch_size < max_batch: + args.graph_max_batch_size = max_batch + + if args.nccl_port is None: + args.nccl_port = 28765 + + if args.mtp_mode in MTP_MODES: + if args.mtp_step <= 0: + raise ValueError("--mtp_mode requires --mtp_step > 0") + if not args.mtp_draft_model_dir: + raise ValueError("--mtp_mode requires --mtp_draft_model_dir") + args.mtp_draft_model_dir = normalize_mtp_draft_dirs(args.mtp_mode, args.mtp_step, args.mtp_draft_model_dir) + else: + args.mtp_mode = None + args.mtp_step = 0 + args.mtp_draft_model_dir = None + + return SimpleNamespace(**vars(args)) + + +def normalize_mtp_draft_dirs(mtp_mode: str, mtp_step: int, draft_dirs: Sequence[str]) -> List[str]: + expected = 1 if mtp_mode.startswith("eagle") else mtp_step + if isinstance(draft_dirs, str): + draft_dirs = [draft_dirs] + draft_dirs = list(draft_dirs) + if len(draft_dirs) == 1 and expected > 1: + return draft_dirs * expected + if len(draft_dirs) != expected: + raise ValueError(f"{mtp_mode} expects {expected} draft model dir(s), got {len(draft_dirs)}") + return draft_dirs + + +def build_model_kvargs(args: SimpleNamespace, rank_id: int) -> Dict: + return { + "args": args, + "nccl_host": args.nccl_host, + "nccl_port": args.nccl_port, + "rank_id": rank_id, + "world_size": args.tp, + "dp_size": args.dp, + "weight_dir": args.model_dir, + "data_type": args.data_type, + "quant_type": args.quant_type, + "quant_cfg": args.quant_cfg, + "expert_dtype": args.expert_dtype, + "load_way": "HF", + "max_total_token_num": args.max_total_token_num, + "graph_max_len_in_batch": args.graph_max_len_in_batch, + "graph_max_batch_size": args.graph_max_batch_size, + "mem_fraction": args.mem_fraction, + "max_req_num": max(args.running_max_req_size, args.graph_max_batch_size), + "batch_max_tokens": args.batch_max_tokens, + "run_mode": "normal", + "max_seq_length": args.max_req_total_len, + "disable_cudagraph": args.disable_cudagraph, + "llm_prefill_att_backend": args.llm_prefill_att_backend, + "llm_decode_att_backend": args.llm_decode_att_backend, + "vit_att_backend": args.vit_att_backend, + "llm_kv_type": args.llm_kv_type, + "llm_kv_quant_group_size": args.llm_kv_quant_group_size, + } + + +def init_mtp_draft_models(args: SimpleNamespace, main_kvargs: Dict, main_model) -> List: + if args.mtp_mode not in MTP_MODES: + return [] + + os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1" + draft_models = [] + for draft_dir in args.mtp_draft_model_dir: + mtp_cfg, _ = PretrainedConfig.get_config_dict(draft_dir) + model_type = mtp_cfg.get("model_type", "") + mtp_kvargs = { + "weight_dir": draft_dir, + "max_total_token_num": main_model.mem_manager.size, + "load_way": main_kvargs["load_way"], + "max_req_num": main_kvargs["max_req_num"], + "max_seq_length": main_kvargs["max_seq_length"], + "is_token_healing": False, + "return_all_prompt_logics": False, + "disable_chunked_prefill": args.disable_chunked_prefill, + "data_type": main_kvargs["data_type"], + "graph_max_batch_size": main_kvargs["graph_max_batch_size"], + "graph_max_len_in_batch": main_kvargs["graph_max_len_in_batch"], + "disable_cudagraph": main_kvargs["disable_cudagraph"], + "mem_fraction": main_kvargs["mem_fraction"], + "batch_max_tokens": main_kvargs["batch_max_tokens"], + "quant_type": main_kvargs["quant_type"], + "quant_cfg": main_kvargs["quant_cfg"], + "expert_dtype": main_kvargs["expert_dtype"], + "run_mode": "normal", + "main_model": main_model, + "mtp_previous_draft_models": draft_models.copy(), + } + if model_type == "deepseek_v3": + assert args.mtp_mode in { + "vanilla_with_att", + "eagle_with_att", + }, f"{model_type} MTP requires *_with_att mode" + draft_models.append(Deepseek3MTPModel(mtp_kvargs)) + elif model_type == "qwen3_moe": + assert args.mtp_mode in { + "vanilla_no_att", + "eagle_no_att", + }, f"{model_type} MTP requires *_no_att mode" + draft_models.append(Qwen3MOEMTPModel(mtp_kvargs)) + elif model_type == "mistral": + assert args.mtp_mode in { + "vanilla_no_att", + "eagle_no_att", + }, f"{model_type} MTP requires *_no_att mode" + draft_models.append(MistralMTPModel(mtp_kvargs)) + elif model_type == "glm4_moe_lite": + assert args.mtp_mode in { + "vanilla_with_att", + "eagle_with_att", + }, f"{model_type} MTP requires *_with_att mode" + draft_models.append(Glm4MoeLiteMTPModel(mtp_kvargs)) + else: + raise ValueError(f"unsupported MTP draft model_type={model_type} from {draft_dir}") + return draft_models + + +def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue): + try: + args = SimpleNamespace(**args_dict) + cases = [BenchmarkCase(**case) for case in case_dicts] + set_env_start_args(args) + + from lightllm.distributed import dist_group_manager + import torch.distributed as dist + + model_kvargs = build_model_kvargs(args, rank_id) + group_size = 2 if (args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap) else 1 + if group_size == 2: + for case in cases: + assert case.batch_size % 2 == 0, "microbatch overlap requires even batch_size" + + init_distributed_env(model_kvargs) + dist_group_manager.create_groups(group_size=group_size) + model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) + dist.barrier() + + torch.cuda.empty_cache() + model, _ = get_model(model_cfg, model_kvargs) + cases = resolve_batch_max_prefill_cases(args, cases, model.mem_manager.size) + cases = resolve_profile_decode_cases(args, cases, model.mem_manager.size) + draft_models = init_mtp_draft_models(args, model_kvargs, model) + token_source = TokenSource(args) + executor = StaticBenchmarkExecutor(args, model, draft_models, token_source) + + results = [] + for case in cases: + if args.warmup_iters > 0: + executor.run_case(case, warmup=True) + result = executor.run_case(case, warmup=False) + if rank_id == 0: + results.append(asdict(result)) + dist.barrier() + + ans_queue.put({"ok": True, "rank": rank_id, "results": results}) + except Exception: + ans_queue.put({"ok": False, "rank": rank_id, "traceback": traceback.format_exc()}) + finally: + try: + ans_queue.close() + ans_queue.join_thread() + except Exception: + pass + os._exit(0) + + +def fmt_optional(value, precision: int = 2) -> str: + if value is None: + return "-" + if isinstance(value, float): + return f"{value:.{precision}f}" + return str(value) + + +def print_aligned_table(headers: Sequence[str], rows: Sequence[Sequence[str]]): + """Print a compact right-aligned ASCII table.""" + if not rows: + return + widths = [len(str(header)) for header in headers] + for row in rows: + for index, value in enumerate(row): + widths[index] = max(widths[index], len(str(value))) + + def format_row(row: Sequence[str]) -> str: + return " ".join(str(value).rjust(widths[index]) for index, value in enumerate(row)) + + print(format_row(headers), flush=True) + print(" ".join("-" * width for width in widths), flush=True) + for row in rows: + print(format_row(row), flush=True) + + +def prefill_table_row(result: BenchmarkResult) -> List[str]: + """Format one prefill result row for stdout table output.""" + return [ + str(result.context_len), + f"{result.cache_hit_rate:.2f}", + str(result.batch_size), + fmt_optional(result.profiled_max_total_token_num, 0), + fmt_optional(result.prefill_uncached_len, 0), + fmt_optional(result.prefill_cached_len, 0), + str(result.measured_tokens), + f"{result.elapsed_ms:.3f}", + f"{result.qps:.2f}", + f"{result.tps:.2f}", + fmt_optional(result.logical_tps, 2), + ] + + +def decode_table_row(result: BenchmarkResult) -> List[str]: + """Format one decode result row for stdout table output.""" + return [ + str(result.context_len), + str(result.batch_size), + fmt_optional(result.mtp_accept_rate, 2), + fmt_optional(result.profiled_max_total_token_num, 0), + f"{result.elapsed_ms:.3f}", + f"{result.qps:.2f}", + f"{result.tps:.2f}", + fmt_optional(result.inter_token_latency_ms, 3), + ] + + +def print_results_table(results: Sequence[BenchmarkResult]): + """Print separate prefill/decode tables for measured results.""" + prefill_rows = [prefill_table_row(result) for result in results if result.stage == "prefill"] + decode_rows = [decode_table_row(result) for result in results if result.stage == "decode"] + + if prefill_rows: + print("\n[prefill]", flush=True) + print_aligned_table(PREFILL_TABLE_HEADERS, prefill_rows) + if decode_rows: + print("\n[decode]", flush=True) + print_aligned_table(DECODE_TABLE_HEADERS, decode_rows) + + +def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List[Dict]: + ctx = mp.get_context("spawn") + ans_queue = ctx.Queue() + workers = [] + rank_start = args.node_rank * args.tp + rank_end = (args.node_rank + 1) * args.tp + case_dicts = [asdict(case) for case in cases] + args_dict = vars(args) + + for rank_id in range(rank_start, rank_end): + proc = ctx.Process(target=run_worker, args=(args_dict, case_dicts, rank_id, ans_queue)) + proc.start() + workers.append(proc) + + for proc in workers: + proc.join() + + messages = [] + while not ans_queue.empty(): + messages.append(ans_queue.get()) + + failed = [message for message in messages if not message.get("ok")] + failed.extend( + { + "ok": False, + "rank": index, + "traceback": f"worker exited with code {proc.exitcode}", + } + for index, proc in enumerate(workers) + if proc.exitcode not in (0, None) + ) + if failed: + for item in failed: + print( + f"rank {item.get('rank')} failed:\n{item.get('traceback')}", + file=sys.stderr, + ) + raise RuntimeError(f"{len(failed)} worker(s) failed") + + results = [] + for message in messages: + results.extend(message.get("results") or []) + result_objs = [BenchmarkResult(**result) for result in results] + print_results_table(result_objs) + return results + + +def add_static_benchmark_args(parser: argparse.ArgumentParser): + parser.add_argument("--benchmark", choices=["all", "prefill", "decode"], default="all") + parser.add_argument("--batch_size", type=int, default=None, help="legacy single batch size") + parser.add_argument( + "--batch_sizes", + type=str, + default=None, + help="comma/space separated batch sizes", + ) + parser.add_argument("--input_len", type=int, default=64, help="legacy single prefill/context length") + parser.add_argument( + "--input_lens", + type=str, + default=None, + help="comma/space separated prefill lengths", + ) + parser.add_argument( + "--context_lens", + type=str, + default=None, + help="comma/space separated decode context lengths", + ) + parser.add_argument("--output_len", type=int, default=512, help="legacy single decode output length") + parser.add_argument( + "--output_lens", + type=str, + default=None, + help="comma/space separated decode output lengths", + ) + parser.add_argument( + "--chunked_prefill_sizes", + type=str, + default=4096, + help=("comma/space separated prefill chunk sizes; default is 4096 " "(full/none/0 select unchunked prefill)"), + ) + parser.add_argument( + "--prefill_cache_hit_rates", + type=str, + default=None, + help=( + "comma/space separated cache hit rates for prefill, e.g. " + "'0,0.5,0.8,0.9'; uncached tokens are ceil(input_len * (1-hit))" + ), + ) + parser.add_argument( + "--max_batch_size", + type=int, + default=2048, + help="upper bound for auto-computed prefill/decode batch size; <=0 disables it", + ) + parser.add_argument( + "--decode_batch_size_mode", + choices=["explicit", "profile"], + default="explicit", + help=( + "explicit uses --batch_size/--batch_sizes; profile computes decode BS " + "from profiled max_total_token_num per context" + ), + ) + parser.add_argument( + "--mtp_accept_rate", + type=float, + default=1.0, + help=("per-draft-token MTP acceptance probability; sampling is outside " "the timed decode section"), + ) + parser.add_argument("--warmup_iters", type=int, default=1) + parser.add_argument("--bench_iters", type=int, default=1) + parser.add_argument("--seed", type=int, default=1234) + + +def main(argv: Optional[Sequence[str]] = None): + parser = make_argument_parser() + add_static_benchmark_args(parser) + args = parser.parse_args(argv) + if args.benchmark in {"all", "prefill"} and args.batch_max_tokens is None: + args.batch_max_tokens = 8192 + cases = build_cases(args) + if not cases: + raise ValueError("no benchmark cases were generated") + args = normalize_args(args, cases) + set_env_start_args(args) + + run_benchmark(args, cases) + + +if __name__ == "__main__": + mp.set_start_method("spawn", force=True) + main() diff --git a/test/benchmark/static_inference/test_model.py b/test/benchmark/static_inference/test_model.py index 5b3751bcc3..ccedad05cf 100644 --- a/test/benchmark/static_inference/test_model.py +++ b/test/benchmark/static_inference/test_model.py @@ -1,46 +1,15 @@ -import os import sys +from pathlib import Path -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) -import unittest -from model_infer import test_model_inference -from model_infer_mtp import test_model_inference_mtp -from lightllm.server.api_cli import make_argument_parser -from lightllm.utils.envs_utils import set_env_start_args, get_env_start_args -from lightllm.utils.config_utils import get_config_json, get_dtype +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = Path(__file__).resolve().parents[3] +for path in (SCRIPT_DIR, REPO_ROOT): + if str(path) not in sys.path: + sys.path.append(str(path)) - -class TestModelInfer(unittest.TestCase): - def test_model_infer(self): - args = get_env_start_args() - if args.data_type is None: - args.data_type = get_dtype(args.model_dir) - if args.mtp_mode == "deepseekv3": - test_model_inference_mtp(args) - else: - test_model_inference(args) - return +from static_benchmark import main if __name__ == "__main__": - import torch - - parser = make_argument_parser() - parser.add_argument("--batch_size", type=int, default=None, help="batch size") - parser.add_argument("--input_len", type=int, default=64, help="input sequence length") - parser.add_argument("--output_len", type=int, default=128, help="output sequence length") - parser.add_argument( - "--profile", - action="store_true", - help="Whether or not to allow for custom models defined on the Hub in their own modeling files.", - ) - parser.add_argument( - "--torch_profile", - action="store_true", - help="Enable torch profiler to profile the model", - ) - args = parser.parse_args() - set_env_start_args(args) - torch.multiprocessing.set_start_method("spawn") - unittest.main(argv=["first-arg-is-ignored"]) + main() From 06b60c7c596e727f174d0f231f4a4fe746150bbc Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 6 Jul 2026 20:03:54 +0800 Subject: [PATCH 04/10] fix: fix bugs --- test/benchmark/static_inference/static_benchmark.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 1ea309d46c..2ddb1f3f02 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -35,7 +35,7 @@ from lightllm.server.router.model_infer.mode_backend.mtp_pre_process import ( prepare_mtp_prefill_inputs, ) -from lightllm.utils.config_utils import get_dtype, get_vocab_size +from lightllm.utils.config_utils import auto_set_fused_shared_experts, get_dtype, get_vocab_size from lightllm.utils.dist_utils import init_distributed_env from lightllm.utils.envs_utils import set_env_start_args @@ -551,6 +551,7 @@ def _make_decode_input( b_req_idx=req_idx, b_mtp_index=mtp_index, b_seq_len=seq_len, + b_position_delta=cpu_i32_zeros(batch_size), mem_indexes_cpu=mem_indexes, is_prefill=False, multimodal_params=empty_multimodal_params(batch_size), @@ -638,6 +639,7 @@ def _slice_decode_input(self, model_input: ModelInput, batch_start: int, batch_e b_req_idx=model_input.b_req_idx[batch_start:batch_end].clone(), b_mtp_index=model_input.b_mtp_index[batch_start:batch_end].clone(), b_seq_len=b_seq_len, + b_position_delta=model_input.b_position_delta[batch_start:batch_end].clone(), mem_indexes_cpu=model_input.mem_indexes_cpu[batch_start:batch_end].contiguous(), is_prefill=False, multimodal_params=model_input.multimodal_params[batch_start:batch_end], @@ -1409,6 +1411,7 @@ def main(argv: Optional[Sequence[str]] = None): parser = make_argument_parser() add_static_benchmark_args(parser) args = parser.parse_args(argv) + auto_set_fused_shared_experts(args) if args.benchmark in {"all", "prefill"} and args.batch_max_tokens is None: args.batch_max_tokens = 8192 cases = build_cases(args) From 0c138e27117673458792506f88da68be675c813b Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 6 Jul 2026 12:19:58 +0000 Subject: [PATCH 05/10] feat: support UE8M0 --- .../quantization/fp8act_quant_kernel.py | 77 +++++++++++++------ .../fp8w8a8_block_quant_kernel.py | 35 +++++++-- lightllm/common/quantization/deepgemm.py | 3 +- 3 files changed, 84 insertions(+), 31 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py index 0a68372887..af4aff5c88 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py @@ -14,6 +14,14 @@ pass +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py @triton.jit def _per_token_group_quant_fp8( @@ -25,21 +33,19 @@ def _per_token_group_quant_fp8( eps, fp8_min, fp8_max, - xs_m, xs_n, - xs_row_major: tl.constexpr, + xs_stride_m, + xs_stride_n, BLOCK: tl.constexpr, NEED_MASK: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, ): g_id = tl.program_id(0) y_ptr += g_id * y_stride y_q_ptr += g_id * y_stride - if xs_row_major: - y_s_ptr += g_id - else: - row_id = g_id // xs_n - col_id = g_id % xs_n - y_s_ptr += col_id * xs_m + row_id # col major + row_id = g_id // xs_n + col_id = g_id % xs_n + y_s_ptr += row_id * xs_stride_m + col_id * xs_stride_n cols = tl.arange(0, BLOCK) # N <= BLOCK @@ -52,8 +58,11 @@ def _per_token_group_quant_fp8( y = tl.load(y_ptr + cols, mask=mask, other=other).to(tl.float32) # Quant - _absmax = tl.maximum(tl.max(tl.abs(y)), eps) - y_s = _absmax / fp8_max + _absmax = tl.max(tl.abs(y)) + if USE_UE8M0_SCALE: + y_s = _ceil_to_ue8m0(tl.maximum(_absmax, 1.0e-4) / fp8_max) + else: + y_s = tl.maximum(_absmax, eps) / fp8_max y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) tl.store(y_q_ptr + cols, y_q, mask=mask) @@ -67,6 +76,7 @@ def lightllm_per_token_group_quant_fp8( x_s: torch.Tensor, eps: float = 1e-10, dtype: torch.dtype = torch.float8_e4m3fn, + use_ue8m0_scales: bool = False, ): """group-wise, per-token quantization on input tensor `x`. Args: @@ -80,8 +90,8 @@ def lightllm_per_token_group_quant_fp8( assert x.shape[-1] % group_size == 0, "the last dimension of `x` cannot be divisible by `group_size`" assert x.is_contiguous(), "`x` is not contiguous" - xs_row_major = x_s.is_contiguous() - xs_m, xs_n = x_s.shape + xs_n = x_s.shape[-1] + xs_stride_m, xs_stride_n = x_s.stride() finfo = torch.finfo(dtype) fp8_max = finfo.max @@ -102,11 +112,12 @@ def lightllm_per_token_group_quant_fp8( eps, fp8_min=fp8_min, fp8_max=fp8_max, - xs_m=xs_m, xs_n=xs_n, - xs_row_major=xs_row_major, + xs_stride_m=xs_stride_m, + xs_stride_n=xs_stride_n, BLOCK=BLOCK, NEED_MASK=BLOCK != group_size, + USE_UE8M0_SCALE=use_ue8m0_scales, num_warps=num_warps, num_stages=num_stages, ) @@ -121,12 +132,13 @@ def per_token_group_quant_fp8( column_major_scales: bool = False, scale_tma_aligned: bool = False, alloc_func: Callable = torch.empty, + use_ue8m0_scales: bool = False, ): x_q = alloc_func(x.shape, dtype=dtype, device=x.device) x_s = None # Adapted from # https://github.com/sgl-project/sglang/blob/7e257cd666c0d639626487987ea8e590da1e9395/python/sglang/srt/layers/quantization/fp8_kernel.py#L290 - if HAS_SGL_KERNEL: + if HAS_SGL_KERNEL and not use_ue8m0_scales: finfo = torch.finfo(dtype) fp8_max, fp8_min = finfo.max, finfo.min @@ -157,14 +169,35 @@ def per_token_group_quant_fp8( sgl_ops.sgl_per_token_group_quant_fp8(x, x_q, x_s, group_size, 1e-10, fp8_min, fp8_max, False, enable_v2=True) else: # 使用LightLLM kernel进行量化 - x_s = alloc_func( - x.shape[:-1] + (x.shape[-1] // group_size,), - device=x.device, - dtype=torch.float32, + if column_major_scales: + if scale_tma_aligned: + aligned_size = (x.shape[-2] + 3) // 4 * 4 + x_s = alloc_func( + x.shape[:-2] + (x.shape[-1] // group_size, aligned_size), + device=x.device, + dtype=torch.float32, + ).permute(-1, -2)[: x.shape[-2], :] + else: + x_s = alloc_func( + (x.shape[-1] // group_size,) + x.shape[:-1], + device=x.device, + dtype=torch.float32, + ).permute(-1, -2) + else: + x_s = alloc_func( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + lightllm_per_token_group_quant_fp8( + x, + group_size, + x_q, + x_s, + eps=eps, + dtype=dtype, + use_ue8m0_scales=use_ue8m0_scales, ) - lightllm_per_token_group_quant_fp8(x, group_size, x_q, x_s, eps=1e-10, dtype=torch.float8_e4m3fn) - if column_major_scales and scale_tma_aligned: - x_s = tma_align_input_scale(x_s) return x_q, x_s diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py index 3881cfe4b8..1c2caa967d 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py @@ -5,7 +5,15 @@ @triton.jit -def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + +@triton.jit +def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr, USE_UE8M0_SCALE: tl.constexpr): pid_m = tl.program_id(axis=0) pid_n = tl.program_id(axis=1) n_blocks = tl.cdiv(N, BLOCK_SIZE) @@ -20,14 +28,21 @@ def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): amax = tl.max(tl.abs(x)) max_fp8e4m3_val = 448.0 - scale = amax / max_fp8e4m3_val - y = (x / (scale + 1e-6)).to(y_ptr.dtype.element_ty) + if USE_UE8M0_SCALE: + scale = _ceil_to_ue8m0(tl.maximum(amax, 1.0e-4) / max_fp8e4m3_val) + denom = scale + else: + scale = amax / max_fp8e4m3_val + denom = scale + 1e-6 + y = (x / denom).to(y_ptr.dtype.element_ty) tl.store(y_ptr + offs, y, mask=mask) tl.store(s_ptr + pid_m * n_blocks + pid_n, scale) -def mm_weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: +def mm_weight_quant( + x: torch.Tensor, block_size: int = 128, use_ue8m0_scales: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous(), "Input tensor must be contiguous" M, N = x.size() @@ -38,11 +53,15 @@ def mm_weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tenso s_scales = torch.empty((num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) - weight_quant_kernel[grid](x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size) + weight_quant_kernel[grid]( + x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales + ) return y_quant, s_scales -def weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: +def weight_quant( + x: torch.Tensor, block_size: int = 128, use_ue8m0_scales: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous(), "Input tensor must be contiguous" x = x.cuda(get_current_device_id()) if x.dim() == 3: @@ -51,8 +70,8 @@ def weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, num_blocks_n = triton.cdiv(x.shape[2], block_size) s_scales = torch.empty((x.shape[0], num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) for i in range(x.shape[0]): - y_quant[i], s_scales[i] = mm_weight_quant(x[i], block_size) + y_quant[i], s_scales[i] = mm_weight_quant(x[i], block_size, use_ue8m0_scales=use_ue8m0_scales) return y_quant, s_scales else: - y_quant, s_scales = mm_weight_quant(x, block_size) + y_quant, s_scales = mm_weight_quant(x, block_size, use_ue8m0_scales=use_ue8m0_scales) return y_quant, s_scales diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index 328e9a71c9..daaceadc36 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -62,7 +62,7 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant device = output.weight.device - weight, scale = weight_quant(weight.cuda(device), self.block_size) + weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=True) output.weight.copy_(weight) output.weight_scale.copy_(scale) return @@ -90,6 +90,7 @@ def apply( column_major_scales=True, scale_tma_aligned=True, alloc_func=alloc_func, + use_ue8m0_scales=True, ) if out is None: From 9b65b431df1f4fbd57a0dcecdb2e8d7c85644c66 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 7 Jul 2026 11:07:55 +0800 Subject: [PATCH 06/10] feat: glm5.2 perf opt --- .../attention/nsa/fp8_flashmla_sparse.py | 12 ++++++--- .../transformer_layer_infer_template.py | 27 ++++++++----------- ...oken_group_quant_deepseek3_2mem_manager.py | 2 ++ .../layer_infer/transformer_layer_infer.py | 7 +++-- .../prefill_compact_kv_flashmla_fp8.py | 26 +++++++++++------- .../triton_kernel/topk_index_to_mem_index.py | 2 +- lightllm/utils/config_utils.py | 1 + 7 files changed, 44 insertions(+), 33 deletions(-) diff --git a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py index 5527c614c0..0930bbc074 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py @@ -76,22 +76,24 @@ def _nsa_prefill_att( nsa_dict = att_control.nsa_prefill_dict softmax_scale = nsa_dict["softmax_scale"] kv_lora_rank = nsa_dict["kv_lora_rank"] - topk_mem_indices = nsa_dict["topk_mem_indices"] + topk_mem_indices = nsa_dict.get("topk_mem_indices") + topk_indices = nsa_dict["topk_indices"] prefill_cache_kv = nsa_dict["prefill_cache_kv"] if self.infer_state.prefix_total_token_num > 0: # 当前推理生成的token kv部分从 prefill_cache_kv 中获取，历史 # 部分kv 从 packed_kv 中获取, 并进行反量化，这样可以避免 prefill_cache_kv # 部分的数据进行重复的反量化操作，提升整体的性能。 + use_full_ragged_kv = topk_mem_indices is None or self.ragged_mem_index.numel() <= topk_indices.numel() kv, topk_indices = self.infer_state.mem_manager.get_prefill_kv_cache_and_remap_indices( packed_kv=packed_kv, - topk_indices=topk_mem_indices, + topk_indices=topk_indices if use_full_ragged_kv else topk_mem_indices, prefill_mem_index=self.infer_state.mem_index, prefill_cache_kv=prefill_cache_kv, + ragged_mem_index=self.ragged_mem_index if use_full_ragged_kv else None, ) else: kv = prefill_cache_kv - topk_indices = nsa_dict["topk_indices"] if topk_indices.ndim == 2: topk_indices = topk_indices.unsqueeze(1) @@ -213,7 +215,9 @@ def _nsa_decode_att( softmax_scale=softmax_scale, causal=False, is_fp8_kvcache=True, - indices=topk_mem_indices.to(dtype=torch.int32), + indices=( + topk_mem_indices if topk_mem_indices.dtype == torch.int32 else topk_mem_indices.to(dtype=torch.int32) + ), ) o_tensor = o_tensor[:, :, :real_head_num, :] return o_tensor[:, 0, :, :] # [b, 1, h, d] -> [b, h, d] 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..70c5e1930a 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 @@ -113,28 +113,23 @@ def _context_attention_wrapper_run( pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() pre_capture_graph.__exit__(None, None, None) - def get_o_shape_dtype_device(): - # 在一个新的 graph 中尝试运行，并不是为了捕获图，是为了尝试得到 o 的形状等信息 - with torch.cuda.graph(cuda_graph=torch.cuda.CUDAGraph()): - __o = self._context_attention_kernel(_q, _cache_kv, infer_state, layer_weight) - o_shape = __o.shape - o_dtype = __o.dtype - o_device = __o.device - del __o - - import gc - - gc.collect() - torch.cuda.empty_cache() - return o_shape, o_dtype, o_device - - o_shape, o_dtype, o_device = get_o_shape_dtype_device() + out_dim = ( + self.kv_lora_rank + if infer_state.prefill_att_state.__class__.__name__.startswith("Nsa") and hasattr(self, "kv_lora_rank") + else getattr(self, "v_head_dim", q.shape[-1]) + ) + o_shape = (*q.shape[:-1], out_dim) + o_dtype = q.dtype + o_device = q.device infer_state.prefill_cuda_graph_create_graph_obj() infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() o = torch.empty(o_shape, dtype=o_dtype, device=o_device) _o = tensor_to_no_ref_tensor(o) + graph_get_topk_indices_params = getattr(infer_state, "get_topk_indices_params", None) def att_func(new_infer_state: InferStateInfo): + if graph_get_topk_indices_params is not None: + new_infer_state.get_topk_indices_params = graph_get_topk_indices_params 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/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py index c5c0b157a0..f60be858b4 100644 --- a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py @@ -72,6 +72,7 @@ def get_prefill_kv_cache_and_remap_indices( topk_indices: torch.Tensor, prefill_mem_index: torch.Tensor, prefill_cache_kv: torch.Tensor, + ragged_mem_index: torch.Tensor = None, ): from lightllm.models.deepseek3_2.triton_kernel.prefill_compact_kv_flashmla_fp8 import ( get_prefill_kv_cache_and_remap_indices_triton, @@ -83,4 +84,5 @@ def get_prefill_kv_cache_and_remap_indices( prefill_mem_index=prefill_mem_index, prefill_cache_kv=prefill_cache_kv, prefill_dtype=self.prefill_dtype, + ragged_mem_index=ragged_mem_index, ) 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 b9d64ff703..ef8362fcb3 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -169,7 +169,7 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.scale_fmt = quantization_config.get("scale_fmt", "ue8m0") 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.indexer_rope_interleave = network_config.get("indexer_rope_interleave", False) @@ -266,6 +266,9 @@ def _get_indices( topk=self.index_topk, ) b_topk_index = torch.where(b_topk_index != -1, b_topk_index + ks.view(-1, 1), -1) + if infer_state.is_prefill and att_state.ragged_mem_index.numel() <= b_topk_index.numel(): + return None, b_topk_index + # 将 topk index 转化为 mem index from ..triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index @@ -283,7 +286,7 @@ 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) + return hadamard_transform(x, scale=hidden_size ** -0.5) def _get_q_k_bf16( self, diff --git a/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py b/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py index 0d0447fec8..9b6c927f56 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py +++ b/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py @@ -111,15 +111,12 @@ def get_prefill_kv_cache_and_remap_indices_triton( prefill_mem_index: torch.Tensor, prefill_cache_kv: torch.Tensor, prefill_dtype: torch.dtype, + ragged_mem_index: torch.Tensor = None, ): squeeze_h_kv = topk_mem_indices.ndim == 2 if squeeze_h_kv: topk_mem_indices = topk_mem_indices.unsqueeze(1) - original_shape = topk_mem_indices.shape - flat_topk = topk_mem_indices.reshape(-1).contiguous().to(torch.int32) - valid_mask = flat_topk != -1 - valid_topk = flat_topk[valid_mask] table_size = packed_kv.shape[0] prefill_row_table = torch.full((table_size,), -1, dtype=torch.int32, device=packed_kv.device) @@ -130,11 +127,21 @@ def get_prefill_kv_cache_and_remap_indices_triton( num_warps=4, ) - unique_mem_index, inverse = torch.unique(valid_topk, sorted=False, return_inverse=True) - unique_mem_index = unique_mem_index.to(torch.int32) - unique_count = unique_mem_index.numel() - remapped_flat = torch.full_like(flat_topk, -1) - remapped_flat[valid_mask] = inverse.to(torch.int32) + if ragged_mem_index is None: + original_shape = topk_mem_indices.shape + flat_topk = topk_mem_indices.reshape(-1).contiguous().to(torch.int32) + valid_mask = flat_topk != -1 + valid_topk = flat_topk[valid_mask] + unique_mem_index, inverse = torch.unique(valid_topk, sorted=False, return_inverse=True) + unique_mem_index = unique_mem_index.to(torch.int32) + unique_count = unique_mem_index.numel() + remapped_flat = torch.full_like(flat_topk, -1) + remapped_flat[valid_mask] = inverse.to(torch.int32) + remapped = remapped_flat.view(original_shape) + else: + unique_mem_index = ragged_mem_index.contiguous().to(torch.int32) + unique_count = unique_mem_index.numel() + remapped = topk_mem_indices compact_kv = torch.empty((unique_count, 1, 576), dtype=prefill_dtype, device=packed_kv.device) packed_nope = packed_kv[:, :, :512].view(torch.float8_e4m3fn).view(-1, 512) @@ -169,7 +176,6 @@ def get_prefill_kv_cache_and_remap_indices_triton( num_warps=4, ) - remapped = remapped_flat.view(original_shape) if squeeze_h_kv: remapped = remapped.squeeze(1) return compact_kv, remapped 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..9a0879c896 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 @@ -31,7 +31,7 @@ def trans_topk_index_to_mem_index(topk_index: torch.Tensor, ragged_mem_index: to grid = (topk_index.shape[0],) - topk_mem_index = torch.empty_like(topk_index) + topk_mem_index = torch.empty(topk_index.shape, dtype=torch.int32, device=topk_index.device) _trans_topk_index_to_mem_index[grid]( topk_index=topk_index, diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 5f5324d772..cc5005a47a 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -170,6 +170,7 @@ def auto_set_fused_shared_experts(args) -> None: "deepseek_v3", "deepseek_v31", "deepseek_v32", + "glm_moe_dsa", "qwen3_next", "qwen3_5", "qwen3_5_text", From fdebbecbd1940c04cdeefd69bfe0f2f3d56509b9 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 7 Jul 2026 07:03:47 +0000 Subject: [PATCH 07/10] fix: fix dpep bug in static_benchmark.py --- .../static_inference/static_benchmark.py | 117 +++++++++++++++--- 1 file changed, 103 insertions(+), 14 deletions(-) diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 2ddb1f3f02..45359398fb 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -7,6 +7,7 @@ import argparse import math import os +import queue import sys import time import traceback @@ -140,7 +141,6 @@ def __init__( self.model = model self.draft_models = draft_models self.token_source = token_source - self.dp = int(args.dp or 1) def _case_iters(self, warmup: bool) -> int: return self.args.warmup_iters if warmup else self.args.bench_iters @@ -699,8 +699,8 @@ def _make_result( ) -> BenchmarkResult: """Convert raw timings into reported TPS and latency metrics.""" iters = self._case_iters(warmup) - scaled_tokens = measured_tokens * self.dp * iters - qps = case.batch_size * self.dp * iters / elapsed_s if elapsed_s > 0 else 0.0 + scaled_tokens = measured_tokens * iters + qps = case.batch_size * iters / elapsed_s if elapsed_s > 0 else 0.0 tps = scaled_tokens / elapsed_s if elapsed_s > 0 else 0.0 ttft_ms = ttft_elapsed_s * 1000.0 / max(1, iters) if ttft_elapsed_s is not None else None logical_tps = None @@ -710,7 +710,7 @@ def _make_result( uncached_len = int(case.prefill_uncached_len or case.context_len) prefill_uncached_len = uncached_len prefill_cached_len = max(0, case.context_len - uncached_len) - token_count = case.batch_size * case.context_len * self.dp * iters + token_count = case.batch_size * case.context_len * iters logical_tps = token_count / elapsed_s if elapsed_s > 0 else 0.0 return BenchmarkResult( case=case.name, @@ -1204,8 +1204,7 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) if args.warmup_iters > 0: executor.run_case(case, warmup=True) result = executor.run_case(case, warmup=False) - if rank_id == 0: - results.append(asdict(result)) + results.append(asdict(result)) dist.barrier() ans_queue.put({"ok": True, "rank": rank_id, "results": results}) @@ -1290,6 +1289,83 @@ def print_results_table(results: Sequence[BenchmarkResult]): print_aligned_table(DECODE_TABLE_HEADERS, decode_rows) +def _dp_size(args: SimpleNamespace) -> int: + return max(1, int(args.dp or 1)) + + +def _dp_world_size(args: SimpleNamespace) -> int: + return max(1, int(args.tp) // _dp_size(args)) + + +def _is_dp_group_leader(args: SimpleNamespace, rank_id: int) -> bool: + return rank_id % _dp_world_size(args) == 0 + + +def _raw_decode_step_count(result: Dict) -> int: + inter_token_latency_ms = result.get("inter_token_latency_ms") + if inter_token_latency_ms is None or inter_token_latency_ms <= 0: + return 0 + return max(1, int(round(float(result["elapsed_ms"]) / float(inter_token_latency_ms)))) + + +def aggregate_rank_results(args: SimpleNamespace, messages: Sequence[Dict]) -> List[Dict]: + """Aggregate rank-local measurements into one global result per case.""" + by_case: Dict[int, List[Dict]] = {} + for message in messages: + rank_id = int(message["rank"]) + for case_index, result in enumerate(message.get("results") or []): + by_case.setdefault(case_index, []).append({"rank": rank_id, "result": result}) + + aggregated_results: List[Dict] = [] + iters = int(args.bench_iters) + for case_index in sorted(by_case): + rank_items = sorted(by_case[case_index], key=lambda item: int(item["rank"])) + leader_items = [item for item in rank_items if _is_dp_group_leader(args, int(item["rank"]))] + if not leader_items: + leader_items = rank_items + + first = dict(leader_items[0]["result"]) + elapsed_ms = max(float(item["result"]["elapsed_ms"]) for item in rank_items) + elapsed_s = elapsed_ms / 1000.0 + batch_size = sum(int(item["result"]["batch_size"]) for item in leader_items) + measured_tokens = sum(int(item["result"]["measured_tokens"]) for item in leader_items) + + first["batch_size"] = batch_size + first["elapsed_ms"] = elapsed_ms + first["measured_tokens"] = measured_tokens + first["qps"] = batch_size * iters / elapsed_s if elapsed_s > 0 else 0.0 + first["tps"] = measured_tokens / elapsed_s if elapsed_s > 0 else 0.0 + + profiled_values = [ + int(item["result"]["profiled_max_total_token_num"]) + for item in leader_items + if item["result"].get("profiled_max_total_token_num") is not None + ] + if profiled_values: + first["profiled_max_total_token_num"] = sum(profiled_values) + + if first["stage"] == "prefill": + logical_tokens = batch_size * int(first["context_len"]) * iters + first["logical_tps"] = logical_tokens / elapsed_s if elapsed_s > 0 else 0.0 + + slowest_item = max(rank_items, key=lambda item: float(item["result"]["elapsed_ms"])) + decode_step_count = _raw_decode_step_count(slowest_item["result"]) + if decode_step_count > 0: + first["inter_token_latency_ms"] = elapsed_ms / decode_step_count + + ttft_values = [ + float(item["result"]["ttft_ms"]) + for item in rank_items + if item["result"].get("ttft_ms") is not None + ] + if ttft_values: + first["ttft_ms"] = max(ttft_values) + + aggregated_results.append(first) + + return aggregated_results + + def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List[Dict]: ctx = mp.get_context("spawn") ans_queue = ctx.Queue() @@ -1304,23 +1380,38 @@ def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List proc.start() workers.append(proc) + messages = [] + while len(messages) < len(workers): + try: + messages.append(ans_queue.get(timeout=5)) + continue + except queue.Empty: + if all(not proc.is_alive() for proc in workers): + break + for proc in workers: proc.join() - messages = [] - while not ans_queue.empty(): - messages.append(ans_queue.get()) - failed = [message for message in messages if not message.get("ok")] + reported_ranks = {int(message["rank"]) for message in messages if "rank" in message} failed.extend( { "ok": False, - "rank": index, + "rank": rank_start + index, "traceback": f"worker exited with code {proc.exitcode}", } for index, proc in enumerate(workers) if proc.exitcode not in (0, None) ) + failed.extend( + { + "ok": False, + "rank": rank_start + index, + "traceback": "worker did not report a result", + } + for index, proc in enumerate(workers) + if rank_start + index not in reported_ranks and proc.exitcode in (0, None) + ) if failed: for item in failed: print( @@ -1329,9 +1420,7 @@ def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List ) raise RuntimeError(f"{len(failed)} worker(s) failed") - results = [] - for message in messages: - results.extend(message.get("results") or []) + results = aggregate_rank_results(args, messages) result_objs = [BenchmarkResult(**result) for result in results] print_results_table(result_objs) return results From 66898718e70d7ef7cb78ead6b85451eb0f869a0f Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Fri, 10 Jul 2026 13:37:50 +0800 Subject: [PATCH 08/10] fix: fix multi nodes results in static_benchmark.py --- .../static_inference/static_benchmark.py | 40 +++++++++++++++---- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 45359398fb..a38115e2dd 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -30,6 +30,7 @@ from lightllm.models import get_model from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel +from lightllm.models.glm5_2_mtp.model import Glm5_2MTPModel from lightllm.models.mistral_mtp.model import MistralMTPModel from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel from lightllm.server.api_cli import make_argument_parser @@ -1166,6 +1167,12 @@ def init_mtp_draft_models(args: SimpleNamespace, main_kvargs: Dict, main_model) "eagle_with_att", }, f"{model_type} MTP requires *_with_att mode" draft_models.append(Glm4MoeLiteMTPModel(mtp_kvargs)) + elif model_type == "glm_moe_dsa": + assert args.mtp_mode in { + "vanilla_with_att", + "eagle_with_att", + }, f"{model_type} MTP requires *_with_att mode" + draft_models.append(Glm5_2MTPModel(mtp_kvargs)) else: raise ValueError(f"unsupported MTP draft model_type={model_type} from {draft_dir}") return draft_models @@ -1207,7 +1214,12 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) results.append(asdict(result)) dist.barrier() - ans_queue.put({"ok": True, "rank": rank_id, "results": results}) + message = {"ok": True, "rank": rank_id, "results": results} + all_rank_messages = [None] * int(args.tp) if rank_id == 0 else None + dist.gather_object(message, all_rank_messages, dst=0) + if rank_id == 0: + message["all_rank_messages"] = all_rank_messages + ans_queue.put(message) except Exception: ans_queue.put({"ok": False, "rank": rank_id, "traceback": traceback.format_exc()}) finally: @@ -1354,9 +1366,7 @@ def aggregate_rank_results(args: SimpleNamespace, messages: Sequence[Dict]) -> L first["inter_token_latency_ms"] = elapsed_ms / decode_step_count ttft_values = [ - float(item["result"]["ttft_ms"]) - for item in rank_items - if item["result"].get("ttft_ms") is not None + float(item["result"]["ttft_ms"]) for item in rank_items if item["result"].get("ttft_ms") is not None ] if ttft_values: first["ttft_ms"] = max(ttft_values) @@ -1367,11 +1377,17 @@ def aggregate_rank_results(args: SimpleNamespace, messages: Sequence[Dict]) -> L def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List[Dict]: + if args.nnodes <= 0 or args.tp % args.nnodes != 0: + raise ValueError(f"--tp must be divisible by --nnodes, got tp={args.tp} nnodes={args.nnodes}") + if args.node_rank < 0 or args.node_rank >= args.nnodes: + raise ValueError(f"--node_rank must be in [0, {args.nnodes}), got {args.node_rank}") + ctx = mp.get_context("spawn") ans_queue = ctx.Queue() workers = [] - rank_start = args.node_rank * args.tp - rank_end = (args.node_rank + 1) * args.tp + node_world_size = args.tp // args.nnodes + rank_start = args.node_rank * node_world_size + rank_end = rank_start + node_world_size case_dicts = [asdict(case) for case in cases] args_dict = vars(args) @@ -1420,7 +1436,17 @@ def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List ) raise RuntimeError(f"{len(failed)} worker(s) failed") - results = aggregate_rank_results(args, messages) + if args.node_rank != 0: + return [] + + all_rank_messages = next( + (message["all_rank_messages"] for message in messages if "all_rank_messages" in message), + None, + ) + if all_rank_messages is None: + raise RuntimeError("global rank 0 did not report aggregated rank results") + + results = aggregate_rank_results(args, all_rank_messages) result_objs = [BenchmarkResult(**result) for result in results] print_results_table(result_objs) return results From 55d9822b9d3aee13c831b65de2d2281312512b53 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Fri, 10 Jul 2026 16:52:41 +0800 Subject: [PATCH 09/10] fix: fix a bug in mtp --- lightllm/models/glm5_2_mtp/model.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/lightllm/models/glm5_2_mtp/model.py b/lightllm/models/glm5_2_mtp/model.py index e1ea4a383e..88f5447d36 100644 --- a/lightllm/models/glm5_2_mtp/model.py +++ b/lightllm/models/glm5_2_mtp/model.py @@ -27,6 +27,21 @@ def _init_custom(self): self._cos_cached = self.main_model._cos_cached self._sin_cached = self.main_model._sin_cached + def _init_quant(self): + super()._init_quant() + expert_dtype = self.expert_dtype or self.config.get("expert_dtype", None) + if expert_dtype is None: + return + + target = self.quant_cfg._get_expert_quant_type(expert_dtype) + mtp_layer_start = self.config["num_hidden_layers"] + num_mtp_layers = self.config.get("num_nextn_predict_layers", 1) + for layer_num in range(mtp_layer_start, mtp_layer_start + num_mtp_layers): + if self.expert_dtype is not None: + self.quant_cfg.quant_cfg[layer_num]["fused_moe"] = target + else: + self.quant_cfg.quant_cfg[layer_num].setdefault("fused_moe", target) + def _init_req_manager(self): self.req_manager = self.main_model.req_manager From 0bdf16b6fb96b86b80555c927b1e0900d8803baf Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Fri, 10 Jul 2026 19:13:05 +0800 Subject: [PATCH 10/10] fix: fix OOM bug for mtp --- .../fused_moe/grouped_fused_moe_ep.py | 21 +++-- .../static_inference/static_benchmark.py | 80 ++++++++++++++++--- 2 files changed, 82 insertions(+), 19 deletions(-) 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 14ec00384b..5a43785f6b 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 @@ -135,9 +135,7 @@ def _mega_moe_quant_topk_to_buffer_kernel( topk_idx.to(topk_idx_out_ptr.dtype.element_ty), ) tl.store( - topk_weights_out_ptr - + token_id * stride_topk_weights_out_m - + topk_offsets * stride_topk_weights_out_k, + topk_weights_out_ptr + token_id * stride_topk_weights_out_m + topk_offsets * stride_topk_weights_out_k, topk_weights.to(topk_weights_out_ptr.dtype.element_ty), ) @@ -244,12 +242,21 @@ def _get_mega_moe_cache_state(w13: Any, w2: Any): return _MEGA_MOE_STATES.setdefault(state_key, {}) +def _transform_mega_moe_weights_in_place(w13: Any, w2: Any): + """Convert to Mega MoE layout without retaining a second weight copy.""" + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( + (w13.weight, w13.weight_scale), + (w2.weight, w2.weight_scale), + ) + w13.weight.copy_(transformed_l1[0]) + w13.weight_scale.copy_(transformed_l1[1]) + w2.weight_scale.copy_(transformed_l2[1]) + return (w13.weight, w13.weight_scale), (w2.weight, w2.weight_scale) + + def _get_mega_moe_weights(w13: Any, w2: Any, state: Dict[str, Any]): if "weight_cache" not in state: - state["weight_cache"] = deep_gemm.transform_weights_for_mega_moe( - (w13.weight, w13.weight_scale), - (w2.weight, w2.weight_scale), - ) + state["weight_cache"] = _transform_mega_moe_weights_in_place(w13, w2) return state["weight_cache"] diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index a38115e2dd..ab7db1a51f 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -130,6 +130,18 @@ def empty_multimodal_params(batch_size: int) -> List[Dict]: return [{"images": [], "audios": []} for _ in range(batch_size)] +def mtp_prefill_chunk_size(batch_size: int, prompt_len: int, batch_max_tokens: int) -> int: + """Bound each MTP setup prefill step by the production token budget.""" + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if batch_max_tokens < batch_size: + raise ValueError( + "MTP prefill requires at least one token per request in a step: " + f"batch_size={batch_size}, batch_max_tokens={batch_max_tokens}" + ) + return max(1, min(int(prompt_len), int(batch_max_tokens) // int(batch_size))) + + class StaticBenchmarkExecutor: def __init__( self, @@ -243,23 +255,30 @@ def _materialize_context_for_decode(self, case: BenchmarkCase): def _prefill_for_decode(self, case: BenchmarkCase, token_rows: np.ndarray, mtp_enabled: bool): req_idx = self._alloc_req_indexes(case.batch_size) + chunk_size = ( + mtp_prefill_chunk_size(case.batch_size, case.context_len, self.args.batch_max_tokens) + if mtp_enabled + else None + ) inputs = self._build_prefill_inputs( token_rows=token_rows, req_idx=req_idx, prompt_len=case.context_len, - chunk_size=None, + chunk_size=chunk_size, ) output = None + next_ids = None for model_input in inputs: output = self._forward_prefill_input(model_input, allow_overlap=not mtp_enabled) + self._touch_output(output) + next_ids = self._argmax_ids(output.logits) + if mtp_enabled: + # Main and draft KV must advance together for every chunk. + next_ids = self._fill_mtp_prefill_kv(case, model_input, output, next_ids) assert output is not None - self._touch_output(output) - - next_ids = self._argmax_ids(output.logits) + assert next_ids is not None seq_len = cpu_i32_full((case.batch_size,), case.context_len) - if mtp_enabled: - next_ids = self._fill_mtp_prefill_kv(case, inputs[-1], output, next_ids) return req_idx, seq_len, next_ids def _fill_mtp_prefill_kv( @@ -808,6 +827,18 @@ def apply_max_batch_size(batch_size: int, max_batch_size: int) -> int: return max(1, batch_size) +def max_decode_batch_size(cases: Sequence[BenchmarkCase]) -> int: + return max((int(case.batch_size) for case in cases if case.stage == "decode"), default=0) + + +def cap_graph_batch_size(configured_size: int, cases: Sequence[BenchmarkCase]) -> int: + """Avoid capturing decode graphs larger than every benchmark case.""" + decode_batch_size = max_decode_batch_size(cases) + if decode_batch_size <= 0: + return max(1, int(configured_size)) + return max(1, min(int(configured_size), decode_batch_size)) + + def prefill_batch_size_from_batch_max_tokens( batch_max_tokens: int, step_tokens_per_req: int, @@ -1042,15 +1073,11 @@ def normalize_args(args: argparse.Namespace, cases: Sequence[BenchmarkCase]) -> if decode_batch_size_needs_profile and args.max_batch_size > 0: args.running_max_req_size = max(args.running_max_req_size, int(args.max_batch_size)) - # Profile decode BS is resolved after model load. Use the cap as the - # pre-load upper bound so request slots and optional decode graphs agree. - if not args.disable_cudagraph: - args.graph_max_batch_size = max(args.graph_max_batch_size, int(args.max_batch_size)) if prefill_batch_size_needs_profile: args.running_max_req_size = max(args.running_max_req_size, max_batch) - if args.graph_max_batch_size < max_batch: - args.graph_max_batch_size = max_batch + if not args.disable_cudagraph and args.decode_batch_size_mode != "profile": + args.graph_max_batch_size = cap_graph_batch_size(args.graph_max_batch_size, cases) if args.nccl_port is None: args.nccl_port = 28765 @@ -1178,6 +1205,25 @@ def init_mtp_draft_models(args: SimpleNamespace, main_kvargs: Dict, main_model) return draft_models +def init_deferred_cudagraph(args: SimpleNamespace, cases: Sequence[BenchmarkCase], model_kvargs: Dict, model) -> None: + """Capture profile-mode graphs after the real decode batch is known.""" + graph_batch_size = cap_graph_batch_size(args.graph_max_batch_size, cases) + args.graph_max_batch_size = graph_batch_size + model_kvargs["graph_max_batch_size"] = graph_batch_size + model_kvargs["disable_cudagraph"] = False + + if args.enable_decode_microbatch_overlap: + graph_batch_size //= 2 + model.graph_max_batch_size = graph_batch_size * (int(args.mtp_step) + 1) + model.disable_cudagraph = False + # Attention backends may size persistent graph buffers during construction. + # Rebuild them after replacing the temporary profile-time batch limit. + model._init_att_backend() + model._init_att_backend1() + torch.cuda.empty_cache() + model._init_cudagraph() + + def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue): try: args = SimpleNamespace(**args_dict) @@ -1188,6 +1234,14 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) import torch.distributed as dist model_kvargs = build_model_kvargs(args, rank_id) + defer_cudagraph = ( + not args.disable_cudagraph + and args.decode_batch_size_mode == "profile" + and any(case.stage == "decode" for case in cases) + ) + if defer_cudagraph: + model_kvargs["disable_cudagraph"] = True + model_kvargs["graph_max_batch_size"] = 2 if args.enable_decode_microbatch_overlap else 1 group_size = 2 if (args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap) else 1 if group_size == 2: for case in cases: @@ -1202,6 +1256,8 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) model, _ = get_model(model_cfg, model_kvargs) cases = resolve_batch_max_prefill_cases(args, cases, model.mem_manager.size) cases = resolve_profile_decode_cases(args, cases, model.mem_manager.size) + if defer_cudagraph: + init_deferred_cudagraph(args, cases, model_kvargs, model) draft_models = init_mtp_draft_models(args, model_kvargs, model) token_source = TokenSource(args) executor = StaticBenchmarkExecutor(args, model, draft_models, token_source)