diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py new file mode 100644 index 0000000000..367d9b6c9e --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""KDA attention backend for GLM-5-Next.""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseAttBackend, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class KDALinearAttBackend(BaseAttBackend): + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return KDAPrefillAttState(backend=self, infer_state=infer_state) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return KDADecodeAttState(backend=self, infer_state=infer_state) + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + +@dataclasses.dataclass +class KDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states = conv_states[..., : backend.conv_kernel_size - 1] + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer_weight.get_merged_kda_conv_weight(), + bias=None, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) + k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) + v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) + raw_gate = raw_gate.view(1, -1, backend.tp_hidden_size) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + + initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +@dataclasses.dataclass +class KDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + mtp_step = self.backend.mtp_step + if mtp_step == 0: + self._init_normal_decode_state() + elif self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(mtp_step + 1) + else: + self._init_fixed_mtp_decode_state(mtp_step) + + def _init_normal_decode_state(self): + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx + + def _init_dynamic_mtp_decode_state(self, mtp_size: int): + ( + self.b1_mtp_cu_q_seq_len, + self.b_conv_buffer_idx, + self.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=self.infer_state.b_req_idx, + b_mtp_index=self.infer_state.b_mtp_index, + req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, + ) + self._init_mtp_ssm_buffer_idx(mtp_size) + + def _init_fixed_mtp_decode_state(self, mtp_step: int): + mtp_size = mtp_step + 1 + batch_size = self.infer_state.batch_size + assert batch_size % mtp_size == 0, ( + "KDA fixed-layout decode requires batch_size to be divisible by mtp_step + 1, " + f"got batch_size={batch_size}, mtp_step={mtp_step}." + ) + + att_batch_size = batch_size // mtp_size + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, + batch_size + 1, + mtp_size, + dtype=torch.int32, + device=self.infer_state.b_req_idx.device, + ) + self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self._init_mtp_ssm_buffer_idx(mtp_size) + + def _init_mtp_ssm_buffer_idx(self, mtp_size: int): + att_batch_size = self.b_conv_buffer_idx.shape[0] + # Each request owns mtp_size consecutive recurrent-state slots. + b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) + state_offsets = torch.arange( + mtp_size, + device=self.infer_state.b_req_idx.device, + dtype=self.infer_state.b_req_idx.dtype, + ).view(1, mtp_size) + self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode + params = att_control.linear_att_decode_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if backend.mtp_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.mtp_step, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update(mixed_qkv, conv_states, layer_weight.get_merged_kda_conv_weight(), **conv_kwargs) + q, k, v = backend.split_qkv(mixed_qkv) + shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) + q = q.view(*shape, backend.tp_num_heads, backend.head_dim) + k = k.view(*shape, backend.tp_num_heads, backend.head_dim) + v = v.view(*shape, backend.tp_num_heads, backend.head_dim) + raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) + raw_beta = raw_beta.view(*shape, backend.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer_weight.linear_A_log.weight, + gate_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + ssm_state_indices=self.b_ssm_buffer_idx, + cu_seqlens=self.b1_mtp_cu_q_seq_len, + num_accepted_tokens=self.b_num_accepted_tokens, + ) + return output diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py new file mode 100644 index 0000000000..c46116d413 --- /dev/null +++ b/lightllm/common/basemodel/attention/nsa/glm5_next.py @@ -0,0 +1,74 @@ +import dataclasses + +import torch + +from .flashmla_sparse import ( + NsaFlashMlaSparseAttBackend, + NsaFlashMlaSparsePrefillAttState, + NsaFlashMlaSparseDecodeAttState, +) + + +class Glm5NextSparseAttBackend(NsaFlashMlaSparseAttBackend): + def create_att_prefill_state(self, infer_state): + return Glm5NextSparsePrefillState(backend=self, infer_state=infer_state) + + def create_att_decode_state(self, infer_state): + return Glm5NextSparseDecodeState(backend=self, infer_state=infer_state) + + +@dataclasses.dataclass +class Glm5NextSparsePrefillState(NsaFlashMlaSparsePrefillAttState): + def _nsa_prefill_att(self, q, kv, att_control): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + tokens, heads, dim = q.shape + # FlashMLA uses 64-head tiles. + padded_heads = ((heads + 63) // 64) * 64 + padded_q = q + if padded_heads != heads: + padded_q = q.new_zeros((tokens, padded_heads, dim)) + padded_q[:, :heads] = q + params = att_control.nsa_prefill_dict + out, _, _ = flash_mla_sparse_fwd( + q=padded_q, + kv=kv, + indices=params["topk_mem_indices"].unsqueeze(1), + sm_scale=params["softmax_scale"], + d_v=params["kv_lora_rank"], + ) + return out[:, :heads] + + +@dataclasses.dataclass +class Glm5NextSparseDecodeState(NsaFlashMlaSparseDecodeAttState): + def init_state(self): + super().init_state() + pool = self.backend.model.config["index_kpool"] + topk = self.backend.model.config["index_topk"] + self.nsa_cache_seqlens = ( + torch.minimum(self.lengths // pool * pool, torch.full_like(self.lengths, topk)) + self.lengths % pool + ) + self.nsa_cu_seqlens_k_new = torch.nn.functional.pad(self.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0)) + + def _nsa_decode_att(self, q, kv, att_control): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + params = att_control.nsa_decode_dict + q_nope, _ = q + kv_nope = kv.view(-1, 1, 1, params["kv_lora_rank"]) + # NoPE has no RoPE Q/K tensors; only_qv uses q_nope and kv_nope directly. + return flash_attn_with_kvcache( + q=None, + qv=q_nope, + k_cache=None, + v_cache=kv_nope, + page_table=params["topk_mem_indices"], + cache_seqlens=self.nsa_cache_seqlens, + cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, + cu_seqlens_k_new=self.nsa_cu_seqlens_k_new, + max_seqlen_q=self.infer_state.max_q_seq_len, + softmax_scale=params["softmax_scale"], + causal=False, + only_qv=True, + ) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 7f369c4fd8..fd3de5c9b4 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -137,6 +137,9 @@ def experts( is_prefill: Optional[bool] = None, infer_state=None, shared_expert_gate: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> torch.Tensor: # Captures MoE topk expert ids for routed-experts metadata when enabled. moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) @@ -156,6 +159,9 @@ def experts( moe_capture_callback=moe_capture_callback, per_expert_scale=self.per_expert_scale, shared_expert_gate=shared_expert_gate, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) def low_latency_dispatch( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py index 240bc726ca..7ce350a762 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py @@ -179,6 +179,7 @@ def experts( layout="interleaved", alpha=self.alpha, limit=self.limit, + clamp_up_add_one=True, ) return output_tensor diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index 1e3ad4b196..e242870d69 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -67,5 +67,8 @@ def __call__( per_expert_scale: Optional[torch.Tensor] = None, # Qwen3.5 uses this gate to control fused shared expert aggregation weights. shared_expert_gate: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> torch.Tensor: pass diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 024be9f55c..09b538a33d 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -76,6 +76,9 @@ def _fused_experts( topk_ids: torch.Tensor, router_logits: Optional[torch.Tensor] = None, is_prefill: Optional[bool] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): output = fused_experts( hidden_states=input_tensor, @@ -87,6 +90,9 @@ def _fused_experts( quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return output diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py index 0094b09b1c..5ccdbb4e9a 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py @@ -30,7 +30,12 @@ def _fused_experts( topk_ids: torch.Tensor, router_logits: Optional[torch.Tensor] = None, is_prefill: Optional[bool] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): + if alpha is not None or limit is not None: + raise NotImplementedError("FuseMoeMarlin does not support clamped SwiGLU") w1_weight, w1_scale, w1_zero_point = w13.weight, w13.weight_scale, w13.weight_zero_point w2_weight, w2_scale, w2_zero_point = w2.weight, w2.weight_scale, w2.weight_zero_point diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index 1d6a38c069..f748aea467 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -1,33 +1,10 @@ import torch from typing import Callable, Optional from lightllm.common.quantization.no_quant import WeightPack -from lightllm.common.quantization.quantize_method import QuantizationMethod from .base_impl import FuseMoeBaseImpl class FuseMoeTriton(FuseMoeBaseImpl): - def __init__( - self, - n_routed_experts: int, - num_fused_shared_experts: int, - routed_scaling_factor: float, - quant_method: QuantizationMethod, - redundancy_expert_num: int, - redundancy_expert_ids_tensor: torch.Tensor, - routed_expert_counter_tensor: torch.Tensor, - auto_update_redundancy_expert: bool, - ): - super().__init__( - n_routed_experts=n_routed_experts, - num_fused_shared_experts=num_fused_shared_experts, - routed_scaling_factor=routed_scaling_factor, - quant_method=quant_method, - redundancy_expert_num=redundancy_expert_num, - redundancy_expert_ids_tensor=redundancy_expert_ids_tensor, - routed_expert_counter_tensor=routed_expert_counter_tensor, - auto_update_redundancy_expert=auto_update_redundancy_expert, - ) - def create_workspace(self): return None @@ -87,6 +64,9 @@ def _fused_experts( topk_ids: torch.Tensor, router_logits: Optional[torch.Tensor] = None, is_prefill: bool = False, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale @@ -104,6 +84,9 @@ def _fused_experts( use_fp8_w8a8=use_fp8_w8a8, w1_scale=w13_scale, w2_scale=w2_scale, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return input_tensor @@ -125,6 +108,9 @@ def __call__( moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None, per_expert_scale: Optional[torch.Tensor] = None, shared_expert_gate: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): topk_weights, topk_ids, origin_topk_ids = self._select_experts( input_tensor=input_tensor, @@ -151,5 +137,8 @@ def __call__( topk_ids=topk_ids, router_logits=router_logits, is_prefill=is_prefill, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return output 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..03ff393811 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py @@ -73,6 +73,17 @@ def __call__( class GatedRMSNormWeight(RMSNormWeight): + def __init__( + self, + dim: int, + weight_name: str, + data_type: torch.dtype, + gate_type: str = "silu", + ): + super().__init__(dim=dim, weight_name=weight_name, data_type=data_type) + assert gate_type in ("silu", "sigmoid") + self.gate_type = gate_type + def _triton_forward( self, input: torch.Tensor, @@ -86,7 +97,15 @@ def _triton_forward( ), f"input.ndim: {input.ndim} != 2 or weight.ndim: {self.weight.ndim} != 1" if out is None: out = alloc_func(input.shape, dtype=input.dtype, device=input.device) - return gated_rmsnorm_forward(x=input, weight=self.weight, bias=None, eps=eps, z=gate_value, out=out) + return gated_rmsnorm_forward( + x=input, + weight=self.weight, + bias=None, + eps=eps, + z=gate_value, + out=out, + gate_type=self.gate_type, + ) def _cuda_forward( self, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index e10adf7758..d7bfe7341a 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 @@ -114,7 +114,6 @@ def moe_align1_kernel( TOKEN_BLOCK_SIZE: tl.constexpr, NUM_STAGE: tl.constexpr, ): - expert_id = tl.program_id(axis=0) off_n = tl.arange(0, TOKEN_BLOCK_SIZE) @@ -406,7 +405,6 @@ def moe_align2_kernel( BLOCK_M: tl.constexpr, BLOCK_EXPERT: tl.constexpr, ): - expert_id = tl.program_id(axis=0) off_expert = tl.arange(0, BLOCK_EXPERT) expert_to_token_num = tl.load(experts_token_num_ptr + off_expert, mask=off_expert < expert_num, other=0) @@ -1009,6 +1007,7 @@ def fused_experts_impl( layout="blocked", limit=None, alpha=None, + clamp_up_add_one=True, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -1087,6 +1086,7 @@ def fused_experts_impl( intermediate_cache2.view(-1, N // 2), limit=limit, alpha=alpha, + clamp_up_add_one=clamp_up_add_one, layout=layout, ) @@ -1133,6 +1133,7 @@ def inplace_fused_experts_impl( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: fused_experts_impl( hidden_states, @@ -1152,6 +1153,7 @@ def inplace_fused_experts_impl( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -1173,6 +1175,7 @@ def inplace_fused_experts_impl_fake( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: pass @@ -1203,6 +1206,7 @@ def outplace_fused_experts_impl( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: return fused_experts_impl( hidden_states, @@ -1222,6 +1226,7 @@ def outplace_fused_experts_impl( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -1243,6 +1248,7 @@ def outplace_fused_experts_impl_fake( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: return torch.empty_like(hidden_states) @@ -1274,6 +1280,7 @@ def fused_experts( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): if inplace: torch.ops.lightllm.inplace_fused_experts_impl( @@ -1293,6 +1300,7 @@ def fused_experts( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return hidden_states else: @@ -1313,4 +1321,5 @@ def fused_experts( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index ca39376bab..fc868a85d2 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 @@ -75,6 +75,9 @@ def masked_group_gemm( w2: torch.Tensor, w2_scale: torch.Tensor, expected_m: int, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): padded_m = recv_x[0].shape[1] E, N, _ = w1.shape @@ -86,7 +89,16 @@ def masked_group_gemm( qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) - silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + gemm_out_a, + qsilu_out, + qsilu_out_scale, + block_size, + masked_m, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) @@ -201,9 +213,15 @@ def fused_experts( quant_method: Any, is_prefill: Optional[bool], previous_event: Optional[Any] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): + if limit is not None: + raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU") return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) buffer = dist_group_manager.ep_buffer if is_prefill else dist_group_manager.ep_low_latency_buffer @@ -222,6 +240,9 @@ def fused_experts( w1_scale=w13.weight_scale, w2_scale=w2.weight_scale, previous_event=previous_event, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -240,6 +261,9 @@ def fused_experts_impl( w1_scale: Optional[torch.Tensor] = None, w2_scale: Optional[torch.Tensor] = None, previous_event: Optional[Any] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -309,6 +333,9 @@ def fused_experts_impl( block_size_k=block_size_k, workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(), hidden_dtype=hidden_states.dtype, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) else: gather_out = torch.empty( @@ -324,7 +351,13 @@ def fused_experts_impl( N = w1.shape[1] _gemm_out_a = torch.zeros((1, N), device=hidden_states.device, dtype=hidden_states.dtype) _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) - silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out) + silu_and_mul_fwd( + _gemm_out_a.view(-1, N), + _silu_out, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) _gemm_out_a, _silu_out = None, None del recv_x @@ -350,7 +383,19 @@ def fused_experts_impl( return_recv_hook=False, ) # deepgemm - gemm_out_b = masked_group_gemm(recv_x, masked_m, hidden_states.dtype, w1, w1_scale, w2, w2_scale, expected_m) + gemm_out_b = masked_group_gemm( + recv_x, + masked_m, + hidden_states.dtype, + w1, + w1_scale, + w2, + w2_scale, + expected_m, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) # low latency combine combined_x, event_overlap, hook = buffer.low_latency_combine( gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False @@ -468,6 +513,9 @@ def chunked_expanded_moe_forward( block_size_k: int, workspace: torch.Tensor, # [workspace_bytes], uint8 hidden_dtype: torch.dtype, # scalar dtype descriptor + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" alignment = 128 @@ -535,7 +583,7 @@ def chunked_expanded_moe_forward( gemm_out_a, m_indices[chunk_start:chunk_end], ) - silu_and_mul_fwd(gemm_out_a, silu_out) + silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) workspace_manager.free(gemm_out_a) del gemm_out_a diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py index a63d92692e..8169cefc3c 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py @@ -24,6 +24,7 @@ def _silu_and_mul_kernel_fast( NEED_MASK: tl.constexpr, layout: tl.constexpr = "blocked", # "blocked" or "interleaved" USE_LIMIT_AND_ALPHA: tl.constexpr = False, + CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, ): stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) @@ -70,11 +71,9 @@ def _silu_and_mul_kernel_fast( up = tl.minimum(tl.maximum(up, -limit), limit) gate = 1 / (1 + tl.exp(-gate * alpha)) * gate gate = gate.to(input_ptr.dtype.element_ty) - tl.store( - output_ptr + out_offsets, - (up + 1) * gate, - mask=mask, - ) + if CLAMP_UP_ADD_ONE: + up += 1 + tl.store(output_ptr + out_offsets, up * gate, mask=mask) else: if USE_TANH_APPROXIMATE_GELU: # tanh-approx GELU, matching Gemma's gelu_pytorch_tanh MLP. @@ -121,6 +120,7 @@ def silu_and_mul_fwd( limit=None, alpha=None, run_config=None, + clamp_up_add_one=True, ): assert input.stride(-1) == 1 assert output.is_contiguous() @@ -171,6 +171,7 @@ def silu_and_mul_fwd( num_warps=num_warps, layout=layout, USE_LIMIT_AND_ALPHA=USE_LIMIT_AND_ALPHA, + CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), ) return diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index aa91f15ed9..661b267088 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -27,6 +27,10 @@ def _silu_and_mul_post_quant_kernel( BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, + USE_LIMIT_AND_ALPHA: tl.constexpr = False, + alpha: tl.constexpr = None, + limit: tl.constexpr = None, + CLAMP_UP_ADD_ONE: tl.constexpr = True, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -51,7 +55,13 @@ def _silu_and_mul_post_quant_kernel( for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) - if USE_TANH_APPROXIMATE_GELU: + if USE_LIMIT_AND_ALPHA: + gate = tl.minimum(gate, limit) + up = tl.minimum(tl.maximum(up, -limit), limit) + gate = gate / (1 + tl.exp(-gate * alpha)) + if CLAMP_UP_ADD_ONE: + up += 1 + elif USE_TANH_APPROXIMATE_GELU: gate_cubed = gate * gate * gate tanh_arg = 0.7978845608028654 * (gate + 0.044715 * gate_cubed) tanh_val = 2.0 / (1.0 + tl.exp(-2.0 * tanh_arg)) - 1.0 @@ -60,6 +70,9 @@ def _silu_and_mul_post_quant_kernel( gate = gate / (1 + tl.exp(-gate)) gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate + if USE_LIMIT_AND_ALPHA: + # Match the BF16/FP16 activation stored before prefill's FP8 quantization. + gate_up = gate_up.to(input_ptr.dtype.element_ty).to(tl.float32) _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) output_s = _absmax / fp8_max output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to(output_ptr.dtype.element_ty) @@ -80,6 +93,9 @@ def silu_and_mul_masked_post_quant_fwd( output_scale: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, + alpha=None, + limit=None, + clamp_up_add_one=True, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -89,6 +105,7 @@ def silu_and_mul_masked_post_quant_fwd( masked_m shape [expert_num], """ + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) assert input.is_contiguous() assert output.dtype == torch.float8_e4m3fn assert output.is_contiguous() @@ -136,6 +153,10 @@ def silu_and_mul_masked_post_quant_fwd( BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), + USE_LIMIT_AND_ALPHA=limit is not None, + alpha=alpha, + limit=limit, + CLAMP_UP_ADD_ONE=clamp_up_add_one, num_warps=num_warps, ) return diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py index 97933b2ac2..849bccb9e9 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py @@ -24,6 +24,7 @@ { "USE_G": lambda args: args["g"] is not None, "USE_GK": lambda args: args["gk"] is not None, + "USE_EXP2": lambda args: args["use_exp2"], "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "STORE_FINAL_STATE": lambda args: args["ht"] is not None, "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, @@ -43,6 +44,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht, cu_seqlens, chunk_offsets, + use_exp2, T, H: tl.constexpr, Hg: tl.constexpr, @@ -52,11 +54,22 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( BV: tl.constexpr, USE_G: tl.constexpr, USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, IS_VARLEN: tl.constexpr, ): + """Scan chunks to compute their residuals E and recurrent state boundaries. + + In the KDA path, k=Kg, v=U, w=W, g=None, gk=G, and use_exp2=True: + E = U - W @ S_in # [BT, BV] + S_out = exp2(G_last)[:, None] * S_in + Kg.T @ E # [K, BV] + Each program owns one (sequence, head, V tile), keeps [K, BV] state in fp32, + and visits that sequence's chunks in order. K is split into up to four + 64-row tiles within the program; V tiles are independent. h stores S_in + before each update, v_new stores E, and ht optionally stores the final state. + """ i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: @@ -111,7 +124,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) - # main recurrence + # Recur across chunks inside this program; preserve each entry state for output. for i_t in range(NT): p_h1 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) @@ -125,6 +138,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( p_h4 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + # W @ S_in: [BT, K] @ [K, BV], reducing all K tiles before forming E. p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) @@ -141,7 +155,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_w = tl.load(p_w, boundary_check=(0, 1)) b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v # E = U - W @ S_in. if SAVE_NEW_VALUE: p_v = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) @@ -163,13 +177,14 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 = b_h4 * b_g_last if USE_GK: + # KDA entry-state contribution at the chunk end: exp2(G_last) * S_in. o_k1 = tl.arange(0, 64) b_gk_last1 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[:, None] + b_h1 *= (tl.exp2(b_gk_last1) if USE_EXP2 else exp(b_gk_last1))[:, None] if K > 64: o_k2 = 64 + o_k1 b_gk_last2 = tl.load( @@ -177,7 +192,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k2 < K), other=0.0, ) - b_h2 *= exp(b_gk_last2)[:, None] + b_h2 *= (tl.exp2(b_gk_last2) if USE_EXP2 else exp(b_gk_last2))[:, None] if K > 128: o_k3 = 128 + o_k1 b_gk_last3 = tl.load( @@ -185,7 +200,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k3 < K), other=0.0, ) - b_h3 *= exp(b_gk_last3)[:, None] + b_h3 *= (tl.exp2(b_gk_last3) if USE_EXP2 else exp(b_gk_last3))[:, None] if K > 192: o_k4 = 192 + o_k1 b_gk_last4 = tl.load( @@ -193,9 +208,11 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k4 < K), other=0.0, ) - b_h4 *= exp(b_gk_last4)[:, None] + b_h4 *= (tl.exp2(b_gk_last4) if USE_EXP2 else exp(b_gk_last4))[:, None] b_v = b_v.to(k.dtype.element_ty) + # Add the chunk's writes. In KDA, k already holds the end-decayed Kg: + # [64, BT] @ [BT, BV] supplies one row tile of Kg.T @ E. p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) b_k = tl.load(p_k, boundary_check=(0, 1)) b_h1 += tl.dot(b_k, b_v) @@ -265,7 +282,16 @@ def chunk_gated_delta_rule_fwd_h( save_new_value: bool = True, cu_seqlens: torch.LongTensor | None = None, run_config=None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: + """Return chunk-entry states h, token residuals v_new, and optional final state. + + For KDA, pass k=Kg: [B, T, H, K], w=W with the same shape, u=U: [B, T, H, V], + and gk=G: [B, T, H, K]. h has shape [B, NT, H, K, V], v_new has u.shape, + and final_state is fp32 [N, H, K, V]. Packed inputs have B=1, N requests, + and NT total chunks. LightLLM autotune chooses BV/warps/stages for grid + (ceil(V/BV), N*H); the chunk dimension is the sequential loop inside each program. + """ # This kernel is slightly different from fla to support Q/K with different head numbers. # In fla, Q/K always have the same head number, so Hg is always equal to H. B, T, Hg, K, V = *k.shape, u.shape[-1] @@ -311,6 +337,7 @@ def chunk_gated_delta_rule_fwd_h( ht=final_state, cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, + use_exp2=use_exp2, T=T, H=H, Hg=Hg, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py new file mode 100644 index 0000000000..6d627078dd --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -0,0 +1,1420 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang + +"""Packed chunkwise KDA prefill with per-channel decay gates. + +q/k/g: [1, total_tokens, head_num, head_dim]; v: [1, total_tokens, head_num, value_head_dim]. +For one head/chunk, chunk_size=64; Q/K: [chunk_size, head_dim]; V/E: [chunk_size, value_head_dim]. +S: [head_dim, value_head_dim]. +G is the chunk-local log2 prefix sum of the decay gate; P = exp2(G). +The forward pass consists of: + 1. Activate the gate and compute G independently for each chunk. + 2. Build the strictly lower-triangular update coupling Akk and causal Aqk. + 3. Compute R = (I + Akk)^-1, U = R @ (beta * V), W = R @ (beta * K * P), + and Kg = K * exp2(G_last - G), independently for each chunk. + 4. Recur across chunks: E = U - W @ S_in; + S_out = P_last[:, None] * S_in + Kg.T @ E. Save S_in for each chunk. + 5. Compute O = scale * (Q * P) @ S_in + Aqk @ E in parallel across chunks. +Here beta is broadcast over key/value channels. E is stored as v_new. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from lightllm.common.triton_utils.autotuner import autotune + +from .chunk_delta_h import chunk_gated_delta_rule_fwd_h +from .cumsum import chunk_local_cumsum +from .index import prepare_chunk_indices +from .l2norm import l2norm_fwd +from triton.language import exp2, log +from .solve_tril import solve_tril + + +FLA_CHUNK_SIZE = 64 +RCP_LN2 = 1.4426950216293335 + + +def cdiv(a: int, b: int) -> int: + return -(a // -b) + + +def next_power_of_2(n: int) -> int: + return 1 if n < 1 else 1 << (n - 1).bit_length() + + +def kda_safe_gate( + raw_gate: torch.Tensor, + a_log: torch.Tensor, + gate_bias: torch.Tensor, + lower_bound: float = -5.0, +) -> torch.Tensor: + """GLM-5 bounded KDA decay in fp32. + + ``raw_gate`` is ``[..., heads, key_dim]``; ``a_log`` is per-head and + ``gate_bias`` is per head/key coordinate. + """ + + head_num = a_log.numel() + head_dim = gate_bias.numel() // head_num + gate = raw_gate.float().view(*raw_gate.shape[:-1], head_num, head_dim) + amplitude = a_log.float().reshape(*((1,) * (gate.ndim - 2)), head_num, 1).exp() + bias = gate_bias.float().reshape(*((1,) * (gate.ndim - 2)), head_num, head_dim) + return lower_bound * torch.sigmoid(amplitude * (gate + bias)) + + +@triton.jit +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + Akk, + Aqk, + scale, + cu_seqlens, + chunk_indices, + head_num: tl.constexpr, + head_dim: tl.constexpr, + chunk_size: tl.constexpr, + subchunk_size: tl.constexpr, + head_block_size: tl.constexpr, + subchunk_num: tl.constexpr, +): + """Build Akk/Aqk on off-diagonal [subchunk_size, subchunk_size] tiles below the chunk diagonal. + + One program handles (global chunk, row/column subtile, head). With g=G: + Akk[i,j] = beta_i * sum_d(k_i[d] * k_j[d] * exp2(G_i[d] - G_j[d])) + Aqk[i,j] = scale * sum_d(q_i[d] * k_j[d] * exp2(G_i[d] - G_j[d])) + The row subtile is strictly after the column subtile, so every pair has i>j. + Accumulate [subchunk_size, head_block_size] @ [head_block_size, subchunk_size] over head_dim. + All chunks/heads/tiles are independent. + """ + global_chunk_id, subtile_id, head_id = tl.program_id(0), tl.program_id(1), tl.program_id(2) + row_subtile_id, col_subtile_id = subtile_id // subchunk_num, subtile_id % subchunk_num + if row_subtile_id <= col_subtile_id: + return + + seq_id = tl.load(chunk_indices + global_chunk_id * 2).to(tl.int32) + chunk_id_in_seq = tl.load(chunk_indices + global_chunk_id * 2 + 1).to(tl.int32) + seq_start = tl.load(cu_seqlens + seq_id).to(tl.int32) + seq_end = tl.load(cu_seqlens + seq_id + 1).to(tl.int32) + seq_len = seq_end - seq_start + chunk_start = chunk_id_in_seq * chunk_size + row_start = chunk_start + row_subtile_id * subchunk_size + col_start = chunk_start + col_subtile_id * subchunk_size + if row_start >= seq_len: + return + + q += (seq_start * head_num + head_id) * head_dim + k += (seq_start * head_num + head_id) * head_dim + g += (seq_start * head_num + head_id) * head_dim + Akk += (seq_start * head_num + head_id) * chunk_size + Aqk += (seq_start * head_num + head_id) * chunk_size + + p_beta = tl.make_block_ptr( + base=beta + seq_start * head_num + head_id, + shape=(seq_len,), + strides=(head_num,), + offsets=(row_start,), + block_shape=(subchunk_size,), + order=(0,), + ) + beta_tile = tl.load(p_beta, boundary_check=(0,)) + + kk_tile = tl.zeros([subchunk_size, subchunk_size], dtype=tl.float32) + qk_tile = tl.zeros([subchunk_size, subchunk_size], dtype=tl.float32) + for key_block_id in range(tl.cdiv(head_dim, head_block_size)): + p_q = tl.make_block_ptr( + base=q, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(row_start, key_block_id * head_block_size), + block_shape=(subchunk_size, head_block_size), + order=(1, 0), + ) + p_k = tl.make_block_ptr( + base=k, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(row_start, key_block_id * head_block_size), + block_shape=(subchunk_size, head_block_size), + order=(1, 0), + ) + p_g = tl.make_block_ptr( + base=g, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(row_start, key_block_id * head_block_size), + block_shape=(subchunk_size, head_block_size), + order=(1, 0), + ) + p_k_col = tl.make_block_ptr( + base=k, + shape=(head_dim, seq_len), + strides=(1, head_num * head_dim), + offsets=(key_block_id * head_block_size, col_start), + block_shape=(head_block_size, subchunk_size), + order=(0, 1), + ) + p_g_cols = tl.make_block_ptr( + base=g, + shape=(head_dim, seq_len), + strides=(1, head_num * head_dim), + offsets=(key_block_id * head_block_size, col_start), + block_shape=(head_block_size, subchunk_size), + order=(0, 1), + ) + + channel_offsets = key_block_id * head_block_size + tl.arange(0, head_block_size) + valid_channels = channel_offsets < head_dim + # Use the first row's G as a shared anchor: the two decay factors + # multiply to exp2(G_i - G_j) without separately forming exp2(-G_j). + # [head_block_size,] + g_anchor = tl.load(g + row_start * head_num * head_dim + channel_offsets, mask=valid_channels, other=0) + # [subchunk_size, head_block_size] + g_rows = tl.load(p_g, boundary_check=(0, 1)) + gated_k_rows = tl.load(p_k, boundary_check=(0, 1)) * exp2(g_rows - g_anchor[None, :]) + # [head_block_size, subchunk_size] + g_cols = tl.load(p_g_cols, boundary_check=(0, 1)) + k_cols = tl.load(p_k_col, boundary_check=(0, 1)) + # [head_block_size, subchunk_size] + gated_k_cols = k_cols * exp2(g_anchor[:, None] - g_cols) + kk_tile += tl.dot(gated_k_rows, gated_k_cols) + + q_rows = tl.load(p_q, boundary_check=(0, 1)) + gated_q_rows = q_rows * exp2(g_rows - g_anchor[None, :]) * scale + qk_tile += tl.dot(gated_q_rows, gated_k_cols) + + kk_tile *= beta_tile[:, None] + + p_Akk = tl.make_block_ptr( + base=Akk, + shape=(seq_len, chunk_size), + strides=(head_num * chunk_size, 1), + offsets=(row_start, col_subtile_id * subchunk_size), + block_shape=(subchunk_size, subchunk_size), + order=(1, 0), + ) + tl.store(p_Akk, kk_tile.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr( + base=Aqk, + shape=(seq_len, chunk_size), + strides=(head_num * chunk_size, 1), + offsets=(row_start, col_subtile_id * subchunk_size), + block_shape=(subchunk_size, subchunk_size), + order=(1, 0), + ) + tl.store(p_Aqk, qk_tile.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + Akk, + Aqk, + scale, + cu_seqlens, + chunk_indices, + head_num: tl.constexpr, + head_dim: tl.constexpr, + chunk_size: tl.constexpr, + subchunk_size: tl.constexpr, + head_block_size: tl.constexpr, +): + """Complete Akk/Aqk inside each diagonal [subchunk_size, subchunk_size] subtile of a chunk. + + One program handles (global chunk, diagonal subtile, head), keeping Q/K/G + tiles of shape [subchunk_size, head_block_size]. Each loop iteration fixes column j and reduces over head_dim + for all subchunk_size rows. Akk uses i>j because E_i reads the state before its own write; + Aqk uses i>=j because o_i reads the state after that write. Only Akk has beta_i. + """ + global_chunk_id, subtile_id, head_id = tl.program_id(0), tl.program_id(1), tl.program_id(2) + seq_id = tl.load(chunk_indices + global_chunk_id * 2).to(tl.int32) + chunk_id_in_seq = tl.load(chunk_indices + global_chunk_id * 2 + 1).to(tl.int32) + seq_start = tl.load(cu_seqlens + seq_id).to(tl.int32) + seq_end = tl.load(cu_seqlens + seq_id + 1).to(tl.int32) + seq_len = seq_end - seq_start + chunk_start = chunk_id_in_seq * chunk_size + row_start = chunk_start + subtile_id * subchunk_size + if row_start >= seq_len: + return + + row_offsets = tl.arange(0, subchunk_size) + channel_offsets = tl.arange(0, head_block_size) + valid_channels = channel_offsets < head_dim + valid_rows = (row_start + row_offsets) < seq_len + output_offsets = ( + (seq_start + row_start + row_offsets) * head_num * chunk_size + + head_id * chunk_size + + subtile_id * subchunk_size + ) + + p_q = tl.make_block_ptr( + base=q + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(row_start, 0), + block_shape=(subchunk_size, head_block_size), + order=(1, 0), + ) + p_k = tl.make_block_ptr( + base=k + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(row_start, 0), + block_shape=(subchunk_size, head_block_size), + order=(1, 0), + ) + p_g = tl.make_block_ptr( + base=g + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(row_start, 0), + block_shape=(subchunk_size, head_block_size), + order=(1, 0), + ) + q_rows = tl.load(p_q, boundary_check=(0, 1)) + beta_k_rows = tl.load(p_k, boundary_check=(0, 1)) + g_rows = tl.load(p_g, boundary_check=(0, 1)) + + p_beta = beta + (seq_start + row_start + row_offsets) * head_num + head_id + beta_k_rows = beta_k_rows * tl.load(p_beta, mask=valid_rows, other=0)[:, None] + + p_k_col = k + (seq_start + row_start) * head_num * head_dim + head_id * head_dim + channel_offsets + p_g_col = g + (seq_start + row_start) * head_num * head_dim + head_id * head_dim + channel_offsets + + for j in range(0, min(subchunk_size, seq_len - row_start)): + k_col = tl.load(p_k_col, mask=valid_channels, other=0).to(tl.float32) + g_col = tl.load(p_g_col, mask=valid_channels, other=0).to(tl.float32) + decayed_k_col = k_col[None, :] * exp2(g_rows - g_col[None, :]) + kk_col = tl.sum(beta_k_rows * decayed_k_col, 1) + kk_col = tl.where(row_offsets > j, kk_col, 0.0) + qk_col = tl.sum(q_rows * decayed_k_col, 1) + qk_col = tl.where(row_offsets >= j, qk_col * scale, 0.0) + tl.store(Akk + output_offsets + j, kk_col, mask=valid_rows) + tl.store(Aqk + output_offsets + j, qk_col, mask=valid_rows) + p_k_col += head_num * head_dim + p_g_col += head_num * head_dim + + +def _get_kda_kkt_sub_inter_configs(): + return [ + {"head_block_size": head_block_size, "num_warps": num_warps, "num_stages": num_stages} + for head_block_size in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ] + + +def _get_kda_kkt_sub_intra_configs(): + return [{"num_warps": num_warps} for num_warps in [1, 2, 4, 8]] + + +def _get_kda_kkt_static_key(k, gk, beta, Akk): + return { + "head_num": k.shape[2], + "head_dim": k.shape[3], + "chunk_size": Akk.shape[-1], + "dtype": str(k.dtype).removeprefix("torch."), + "gate_dtype": str(gk.dtype).removeprefix("torch."), + "beta_dtype": str(beta.dtype).removeprefix("torch."), + "out_dtype": str(Akk.dtype).removeprefix("torch."), + } + + +@autotune( + kernel_name="chunk_kda_scaled_dot_kkt_sub_inter:v1", + configs_gen_func=_get_kda_kkt_sub_inter_configs, + static_key_func=_get_kda_kkt_static_key, + run_key_func=lambda k: k.shape[1], +) +def _chunk_kda_scaled_dot_kkt_sub_inter(q, k, gk, beta, Akk, Aqk, scale, cu_seqlens, chunk_indices, run_config=None): + """Tune and launch off-diagonal Akk/Aqk tiles, overwriting only those output tiles.""" + head_num, head_dim = k.shape[-2:] + chunk_size = Akk.shape[-1] + subchunk_size = min(16, chunk_size) + subchunk_num = cdiv(chunk_size, subchunk_size) + chunk_num = len(chunk_indices) + if run_config is None: + run_config = {"head_block_size": 64, "num_warps": 4, "num_stages": 2} + + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[(chunk_num, subchunk_num * subchunk_num, head_num)]( + q=q, + k=k, + g=gk, + beta=beta, + Akk=Akk, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + head_num=head_num, + head_dim=head_dim, + chunk_size=chunk_size, + subchunk_size=subchunk_size, + head_block_size=run_config.get("head_block_size", 64), + subchunk_num=subchunk_num, + num_warps=run_config.get("num_warps", 4), + num_stages=run_config.get("num_stages", 2), + ) + + +@autotune( + kernel_name="chunk_kda_scaled_dot_kkt_sub_intra:v1", + configs_gen_func=_get_kda_kkt_sub_intra_configs, + static_key_func=_get_kda_kkt_static_key, + run_key_func=lambda k: k.shape[1], +) +def _chunk_kda_scaled_dot_kkt_sub_intra(q, k, gk, beta, Akk, Aqk, scale, cu_seqlens, chunk_indices, run_config=None): + """Tune and launch diagonal Akk/Aqk subtiles independently of off-diagonal tiles.""" + head_num, head_dim = k.shape[-2:] + chunk_size = Akk.shape[-1] + subchunk_size = min(16, chunk_size) + chunk_num = len(chunk_indices) + if run_config is None: + run_config = {"num_warps": 4} + + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[(chunk_num, cdiv(chunk_size, subchunk_size), head_num)]( + q=q, + k=k, + g=gk, + beta=beta, + Akk=Akk, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + head_num=head_num, + head_dim=head_dim, + chunk_size=chunk_size, + subchunk_size=subchunk_size, + head_block_size=max(next_power_of_2(head_dim), 16), + num_warps=run_config.get("num_warps", 4), + ) + + +def chunk_kda_scaled_dot_kkt_fwd( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build the update coupling Akk and the output weights Aqk for every chunk. + + q/k/gk: [1, total_tokens, head_num, head_dim], where gk is the chunk-local log2 prefix sum G. + beta: [1, total_tokens, head_num]. Both returned tensors have shape [1, total_tokens, head_num, chunk_size]; each + token row stores weights for the chunk_size positions in its own chunk: + Akk[i,j] = beta_i * sum_d(k_i[d] * k_j[d] * exp2(G_i[d] - G_j[d])), i>j. + Aqk[i,j] = scale * sum_d(q_i[d] * k_j[d] * exp2(G_i[d] - G_j[d])), i>=j. + Entries outside the respective causal masks are zero. solve_tril computes R=(I+Akk)^-1; + Aqk is retained for O = scale * (Q * exp2(G)) @ S_in + Aqk @ E. + The two subtile kernels have separate LightLLM autotune configurations. + """ + _, total_tokens, head_num, head_dim = k.shape + assert head_dim <= 256 + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + Akk = torch.zeros(1, total_tokens, head_num, chunk_size, device=k.device, dtype=output_dtype) + Aqk = torch.zeros(1, total_tokens, head_num, chunk_size, device=k.device, dtype=output_dtype) + _chunk_kda_scaled_dot_kkt_sub_inter( + q=q, + k=k, + gk=gk, + beta=beta, + Akk=Akk, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + _chunk_kda_scaled_dot_kkt_sub_intra( + q=q, + k=k, + gk=gk, + beta=beta, + Akk=Akk, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return Akk, Aqk + + +@triton.jit +def recompute_w_u_fwd_kernel( + k, + kg, + v, + beta, + w, + u, + R, + gk, + cu_seqlens, + chunk_indices, + head_num: tl.constexpr, + head_dim: tl.constexpr, + value_head_dim: tl.constexpr, + chunk_size: tl.constexpr, + head_block_size: tl.constexpr, + value_block_size: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + """Precompute U/W/Kg from R=(I+Akk)^-1 without reading recurrent state. + + For one chunk/head, P=exp2(gk); R: [chunk_size, chunk_size]; + K/P: [chunk_size, head_dim]; V: [chunk_size, value_head_dim]: + U = R @ (beta[:, None] * V) # [chunk_size, value_head_dim] + W = R @ (beta[:, None] * K * P) # [chunk_size, head_dim] + Kg = K * exp2(G_last - G) # [chunk_size, head_dim], writes decayed to chunk end + Later E = U - W @ S_in and S_out = P_last[:, None] * S_in + Kg.T @ E. + One program handles (global chunk, head), looping over blocks of key/value channels. + Q * P is computed directly in the output kernel. + """ + global_chunk_id, head_id = tl.program_id(0), tl.program_id(1) + seq_id = tl.load(chunk_indices + global_chunk_id * 2).to(tl.int32) + chunk_id_in_seq = tl.load(chunk_indices + global_chunk_id * 2 + 1).to(tl.int32) + seq_start = tl.load(cu_seqlens + seq_id).to(tl.int32) + seq_end = tl.load(cu_seqlens + seq_id + 1).to(tl.int32) + seq_len = seq_end - seq_start + chunk_start = chunk_id_in_seq * chunk_size + last_idx = min(chunk_start + chunk_size, seq_len) - 1 + p_beta = tl.make_block_ptr( + base=beta + seq_start * head_num + head_id, + shape=(seq_len,), + strides=(head_num,), + offsets=(chunk_start,), + block_shape=(chunk_size,), + order=(0,), + ) + beta_tile = tl.load(p_beta, boundary_check=(0,)) + + p_r = tl.make_block_ptr( + base=R + (seq_start * head_num + head_id) * chunk_size, + shape=(seq_len, chunk_size), + strides=(head_num * chunk_size, 1), + offsets=(chunk_start, 0), + block_shape=(chunk_size, chunk_size), + order=(1, 0), + ) + r_tile = tl.load(p_r, boundary_check=(0, 1)) + + for value_block_id in range(tl.cdiv(value_head_dim, value_block_size)): + p_v = tl.make_block_ptr( + base=v + (seq_start * head_num + head_id) * value_head_dim, + shape=(seq_len, value_head_dim), + strides=(head_num * value_head_dim, 1), + offsets=(chunk_start, value_block_id * value_block_size), + block_shape=(chunk_size, value_block_size), + order=(1, 0), + ) + p_u = tl.make_block_ptr( + base=u + (seq_start * head_num + head_id) * value_head_dim, + shape=(seq_len, value_head_dim), + strides=(head_num * value_head_dim, 1), + offsets=(chunk_start, value_block_id * value_block_size), + block_shape=(chunk_size, value_block_size), + order=(1, 0), + ) + v_tile = tl.load(p_v, boundary_check=(0, 1)) + weighted_v = (v_tile * beta_tile[:, None]).to(v_tile.dtype) + # U: [chunk_size, chunk_size] @ [chunk_size, value_block_size] -> [chunk_size, value_block_size]. + u_tile = tl.dot(r_tile, weighted_v, input_precision=DOT_PRECISION) + tl.store(p_u, u_tile.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for key_block_id in range(tl.cdiv(head_dim, head_block_size)): + p_w = tl.make_block_ptr( + base=w + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(chunk_start, key_block_id * head_block_size), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + p_k = tl.make_block_ptr( + base=k + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(chunk_start, key_block_id * head_block_size), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + k_tile = tl.load(p_k, boundary_check=(0, 1)) + weighted_k = k_tile * beta_tile[:, None] + + p_gk = tl.make_block_ptr( + base=gk + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(chunk_start, key_block_id * head_block_size), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + g_tile = tl.load(p_gk, boundary_check=(0, 1)) + weighted_k *= exp2(g_tile) + + # Kg carries each write forward to the last valid token of this chunk. + channel_offsets = key_block_id * head_block_size + tl.arange(0, head_block_size) + valid_channels = channel_offsets < head_dim + g_last = tl.load( + gk + ((seq_start + last_idx) * head_num + head_id) * head_dim + channel_offsets, + mask=valid_channels, + other=0.0, + ) + kg_tile = k_tile * exp2(g_last - g_tile) + p_kg = tl.make_block_ptr( + base=kg + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(chunk_start, key_block_id * head_block_size), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + tl.store(p_kg, kg_tile.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + # W maps the chunk-entry state to corrections: [chunk_size, chunk_size] @ [chunk_size, head_block_size]. + w_tile = tl.dot(r_tile, weighted_k.to(k_tile.dtype)) + tl.store(p_w, w_tile.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def _get_kda_w_u_configs(): + return [{"num_warps": num_warps, "num_stages": num_stages} for num_warps in [2, 4, 8] for num_stages in [2, 3, 4]] + + +def _get_kda_w_u_static_key(k, v, beta, R, gk): + return { + "head_num": k.shape[2], + "head_dim": k.shape[3], + "value_head_dim": v.shape[-1], + "chunk_size": R.shape[-1], + "dtype": str(k.dtype).removeprefix("torch."), + "v_dtype": str(v.dtype).removeprefix("torch."), + "r_dtype": str(R.dtype).removeprefix("torch."), + "gate_dtype": str(gk.dtype).removeprefix("torch."), + "beta_dtype": str(beta.dtype).removeprefix("torch."), + } + + +@autotune( + kernel_name="kda_recompute_w_u_fwd:v1", + configs_gen_func=_get_kda_w_u_configs, + static_key_func=_get_kda_w_u_static_key, + run_key_func=lambda k: k.shape[1], +) +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + R: torch.Tensor, + gk: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor | None = None, + run_config: dict | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return W/U/Kg for the packed chunk recurrence. + + R: [1, total_tokens, head_num, chunk_size] stores (I+Akk)^-1; gk: [1, total_tokens, head_num, head_dim] stores G. + W/Kg: [1, total_tokens, head_num, head_dim]; U: [1, total_tokens, head_num, value_head_dim]. + These factors depend only on the current chunk, so grid (chunk_num, head_num) computes all chunks + independently before the state scan. LightLLM tunes num_warps/num_stages; + the channel tiles head_block_size=value_block_size=64 stay fixed. + """ + head_num, head_dim = k.shape[-2:] + value_head_dim = v.shape[-1] + chunk_size = R.shape[-1] + head_block_size = 64 + value_block_size = 64 + if run_config is None: + run_config = {"num_warps": 4, "num_stages": 2} + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + chunk_num = len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) + recompute_w_u_fwd_kernel[(chunk_num, head_num)]( + k=k, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + R=R, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + head_num=head_num, + head_dim=head_dim, + value_head_dim=value_head_dim, + chunk_size=chunk_size, + head_block_size=head_block_size, + value_block_size=value_block_size, + DOT_PRECISION="ieee", + num_warps=run_config.get("num_warps", 4), + num_stages=run_config.get("num_stages", 2), + ) + return w, u, kg + + +@triton.jit +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + Aqk, + cu_seqlens, + chunk_indices, + scale, + head_num: tl.constexpr, + head_dim: tl.constexpr, + value_head_dim: tl.constexpr, + chunk_size: tl.constexpr, + head_block_size: tl.constexpr, + value_block_size: tl.constexpr, +): + """Compute O = scale * (Q * exp2(G)) @ S_in + Aqk @ E for each chunk. + + h holds the state before each chunk; v holds E (v_new); Aqk holds the output weights. + One program handles (value channel block, global chunk, head), producing [chunk_size, value_block_size]. + The history term reduces [chunk_size, head_block_size] @ [head_block_size, value_block_size] over head_dim. + The local term is [chunk_size, chunk_size] @ [chunk_size, value_block_size]. + All chunks can run independently after h/E are saved. + """ + value_block_id, global_chunk_id, head_id = tl.program_id(0), tl.program_id(1), tl.program_id(2) + seq_id = tl.load(chunk_indices + global_chunk_id * 2).to(tl.int32) + chunk_id_in_seq = tl.load(chunk_indices + global_chunk_id * 2 + 1).to(tl.int32) + seq_start = tl.load(cu_seqlens + seq_id).to(tl.int32) + seq_end = tl.load(cu_seqlens + seq_id + 1).to(tl.int32) + seq_len = seq_end - seq_start + chunk_start = chunk_id_in_seq * chunk_size + + causal_mask = tl.arange(0, chunk_size)[:, None] >= tl.arange(0, chunk_size)[None, :] + + output_tile = tl.zeros([chunk_size, value_block_size], dtype=tl.float32) + for key_block_id in range(tl.cdiv(head_dim, head_block_size)): + p_q = tl.make_block_ptr( + base=q + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(chunk_start, key_block_id * head_block_size), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + p_g = tl.make_block_ptr( + base=g + (seq_start * head_num + head_id) * head_dim, + shape=(seq_len, head_dim), + strides=(head_num * head_dim, 1), + offsets=(chunk_start, key_block_id * head_block_size), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + p_h = tl.make_block_ptr( + base=h + (global_chunk_id * head_num + head_id) * head_dim * value_head_dim, + shape=(head_dim, value_head_dim), + strides=(value_head_dim, 1), + offsets=(key_block_id * head_block_size, value_block_id * value_block_size), + block_shape=(head_block_size, value_block_size), + order=(1, 0), + ) + + # [chunk_size, head_block_size] + q_tile = tl.load(p_q, boundary_check=(0, 1)) + q_tile = (q_tile * scale).to(q_tile.dtype) + # [chunk_size, head_block_size] + g_tile = tl.load(p_g, boundary_check=(0, 1)) + # [chunk_size, head_block_size] + gated_q_tile = (q_tile * exp2(g_tile)).to(q_tile.dtype) + # Chunk-entry state tile S_in: [head_block_size, value_block_size]. + state_tile = tl.load(p_h, boundary_check=(0, 1)) + # Historical contribution: scale * (Q * exp2(G)) @ S_in, [chunk_size, value_block_size]. + output_tile += tl.dot(gated_q_tile, state_tile.to(gated_q_tile.dtype)) + p_e = tl.make_block_ptr( + base=v + (seq_start * head_num + head_id) * value_head_dim, + shape=(seq_len, value_head_dim), + strides=(head_num * value_head_dim, 1), + offsets=(chunk_start, value_block_id * value_block_size), + block_shape=(chunk_size, value_block_size), + order=(1, 0), + ) + p_o = tl.make_block_ptr( + base=o + (seq_start * head_num + head_id) * value_head_dim, + shape=(seq_len, value_head_dim), + strides=(head_num * value_head_dim, 1), + offsets=(chunk_start, value_block_id * value_block_size), + block_shape=(chunk_size, value_block_size), + order=(1, 0), + ) + p_Aqk = tl.make_block_ptr( + base=Aqk + (seq_start * head_num + head_id) * chunk_size, + shape=(seq_len, chunk_size), + strides=(head_num * chunk_size, 1), + offsets=(chunk_start, 0), + block_shape=(chunk_size, chunk_size), + order=(1, 0), + ) + # E: [chunk_size, value_block_size], already includes beta and earlier-token corrections. + residual_tile = tl.load(p_e, boundary_check=(0, 1)) + # Aqk: [chunk_size, chunk_size], causal including the diagonal; scale is already included. + qk_tile = tl.load(p_Aqk, boundary_check=(0, 1)) + qk_tile = tl.where(causal_mask, qk_tile, 0.0).to(residual_tile.dtype) + output_tile += tl.dot(qk_tile, residual_tile, allow_tf32=False) # Current-chunk contribution Aqk @ E. + tl.store(p_o, output_tile.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def _get_kda_output_configs(): + return [ + { + "head_block_size": head_block_size, + "value_block_size": value_block_size, + "num_warps": num_warps, + "num_stages": num_stages, + } + for head_block_size in [32, 64] + for value_block_size in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ] + + +def _get_kda_output_static_key(q, v, g, Aqk, h, o, chunk_size): + return { + "head_num": q.shape[2], + "head_dim": q.shape[3], + "value_head_dim": v.shape[-1], + "chunk_size": chunk_size, + "dtype": str(q.dtype).removeprefix("torch."), + "v_dtype": str(v.dtype).removeprefix("torch."), + "gate_dtype": str(g.dtype).removeprefix("torch."), + "aqk_dtype": str(Aqk.dtype).removeprefix("torch."), + "state_dtype": str(h.dtype).removeprefix("torch."), + "out_dtype": str(o.dtype).removeprefix("torch."), + } + + +@autotune( + kernel_name="kda_chunk_gla_fwd_o_gk:v1", + configs_gen_func=_get_kda_output_configs, + static_key_func=_get_kda_output_static_key, + run_key_func=lambda q: q.shape[1], +) +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + Aqk: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + run_config: dict | None = None, +): + """Write all chunk outputs from saved entry states h and corrections v=E. + + q/g: [1, total_tokens, head_num, head_dim]; v/o: [1, total_tokens, head_num, value_head_dim]. + Aqk: [1, total_tokens, head_num, chunk_size]. + h: [1, chunk_num, head_num, head_dim, value_head_dim], where chunk_num is the total chunk count across all requests. + The caller supplies the output buffer o. Each launch fully overwrites it, + allowing repeated tuning runs without changing h or E. LightLLM selects + head_block_size/value_block_size and launch parameters. + Grid: (ceil(value_head_dim/value_block_size), chunk_num, head_num). + """ + head_num, head_dim = q.shape[-2:] + value_head_dim = v.shape[-1] + + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + chunk_num = len(chunk_indices) + + if run_config is None: + run_config = {"head_block_size": 64, "value_block_size": 64, "num_warps": 4, "num_stages": 2} + head_block_size = run_config.get("head_block_size", 64) + value_block_size = run_config.get("value_block_size", 64) + + grid = (cdiv(value_head_dim, value_block_size), chunk_num, head_num) + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + Aqk=Aqk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + head_num=head_num, + head_dim=head_dim, + value_head_dim=value_head_dim, + chunk_size=chunk_size, + head_block_size=head_block_size, + value_block_size=value_block_size, + num_warps=run_config.get("num_warps", 4), + num_stages=run_config.get("num_stages", 2), + ) + return o + + +@triton.heuristics({"HAS_BIAS": lambda args: args["g_bias"] is not None}) +@triton.jit +def kda_gate_cumsum_fwd_kernel( + g, + A_log, + y, + g_bias, + cu_seqlens, + chunk_indices, + # Element strides for input/output [total_tokens, head_num, head_dim]: token, head, channel. + stride_g_token: tl.constexpr, + stride_g_head: tl.constexpr, + stride_g_dim: tl.constexpr, + stride_y_token: tl.constexpr, + stride_y_head: tl.constexpr, + stride_y_dim: tl.constexpr, + cumsum_scale, + beta, + threshold, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, + head_dim: tl.constexpr, + chunk_size: tl.constexpr, + head_block_size: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + """Fuse raw-gate activation with G_i = sum_{r<=i} ell_r / ln(2) in each chunk. + + g/y: [total_tokens, head_num, head_dim]; A_log: [head_num]. + Activate ell with the bounded sigmoid or negative softplus, then multiply + a [chunk_size, chunk_size] lower-triangular matrix of ones by ell: [chunk_size, head_block_size]. The resulting + log2 prefixes reset at every sequence/chunk boundary. beta here is the + scalar softplus parameter, not the per-token KDA update strength. + """ + # One program handles one [chunk_size, head_block_size] tile for one request/head. + head_block_id, global_chunk_id, head_id = tl.program_id(0), tl.program_id(1), tl.program_id(2) + # chunk_indices[global_chunk_id] = (seq_id, chunk_id_in_seq). + seq_id = tl.load(chunk_indices + global_chunk_id * 2).to(tl.int32) + chunk_id_in_seq = tl.load(chunk_indices + global_chunk_id * 2 + 1).to(tl.int32) + seq_start = tl.load(cu_seqlens + seq_id).to(tl.int32) + seq_end = tl.load(cu_seqlens + seq_id + 1).to(tl.int32) + seq_len = seq_end - seq_start + chunk_start = chunk_id_in_seq * chunk_size + head_dim_start = head_block_id * head_block_size + + # Fix the request/head, then view [total_tokens, head_num, head_dim] as a [seq_len, head_dim] matrix. + # Moving one token/channel advances by stride_*_token/stride_*_dim elements. + g_seq_head = g + seq_start * stride_g_token + head_id * stride_g_head + y_seq_head = y + seq_start * stride_y_token + head_id * stride_y_head + p_g = tl.make_block_ptr( + base=g_seq_head, + shape=(seq_len, head_dim), + strides=(stride_g_token, stride_g_dim), + offsets=(chunk_start, head_dim_start), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + p_y = tl.make_block_ptr( + base=y_seq_head, + shape=(seq_len, head_dim), + strides=(stride_y_token, stride_y_dim), + offsets=(chunk_start, head_dim_start), + block_shape=(chunk_size, head_block_size), + order=(1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + if HAS_BIAS: + head_dim_indices = head_dim_start + tl.arange(0, head_block_size) + b_bias = tl.load( + g_bias + head_id * head_dim + head_dim_indices, mask=head_dim_indices < head_dim, other=0.0 + ).to(tl.float32) + b_g = b_g + b_bias[None, :] + + b_a = tl.load(A_log + head_id).to(tl.float32) + b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) + if SAFE_GATE: + # log_gate = lower_bound * sigmoid(exp(A_log) * (raw_g + bias)). + # For lower_bound < 0, log_gate is in [lower_bound, 0]; decay = exp(log_gate). + b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g))) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = b_a * b_softplus + + # Out-of-bounds rows (load returns 0, but softplus/bias can still make + # b_gate non-zero) participate in the dot product. They only contribute to + # out-of-bounds output rows, which are masked away by `boundary_check` on + # the store, so visible output matches unfused gate + chunk-local cumsum. + o_t = tl.arange(0, chunk_size) + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +def _get_kda_gate_cumsum_configs(): + return [ + {"head_block_size": head_block_size, "num_warps": num_warps} + for head_block_size in [32, 64] + for num_warps in [2, 4, 8] + ] + + +def _get_kda_gate_cumsum_static_key(raw_g, g_bias, chunk_size, output_dtype, safe_gate): + return { + "head_num": raw_g.shape[1], + "head_dim": raw_g.shape[2], + "chunk_size": chunk_size, + "SAFE_GATE": safe_gate, + "HAS_BIAS": g_bias is not None, + "dtype": str(raw_g.dtype).removeprefix("torch."), + "out_dtype": str(output_dtype or raw_g.dtype).removeprefix("torch."), + } + + +@autotune( + kernel_name="fused_kda_gate_chunk_cumsum:v1", + configs_gen_func=_get_kda_gate_cumsum_configs, + static_key_func=_get_kda_gate_cumsum_static_key, + run_key_func=lambda raw_g: raw_g.shape[0], # Total packed token count. +) +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + A_log: torch.Tensor, + cu_seqlens: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, + safe_gate: bool = False, + lower_bound: float = -5.0, + run_config: dict | None = None, +) -> torch.Tensor: + """Activate packed decay gates and return chunk-local log2 prefix sums in [total_tokens, head_num, head_dim]. + + raw_g: [total_tokens, head_num, head_dim], packed tokens, local heads, and key channels. + Input/output addressing uses each tensor's strides, measured in elements. + A_log: [head_num]; g_bias: [head_num * head_dim] or [head_num, head_dim], or None to skip the bias. + cu_seqlens: [request_num + 1], required token boundaries for request_num packed requests. + run_config: optional LightLLM autotune config with head_block_size and num_warps. + + Returns G_i = sum_{r=chunk_start..i} ell_r / ln(2), where ell is the + activated log decay. Subsequent kernels use P_i=exp2(G_i) for entry-state + decay and exp2(G_i-G_j) for writes propagated from token j to token i. + """ + assert raw_g.ndim == 3, "raw_g must have packed shape [total_tokens, head_num, head_dim]" + assert cu_seqlens is not None, "cu_seqlens is required for packed KDA prefill" + head_num, head_dim = raw_g.shape[1:] + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + chunk_num = len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + + if run_config is None: + run_config = {"head_block_size": 32, "num_warps": 4} + head_block_size = run_config.get("head_block_size", 32) + num_warps = run_config.get("num_warps", 4) + + grid = (cdiv(head_dim, head_block_size), chunk_num, head_num) + kda_gate_cumsum_fwd_kernel[grid]( + g=raw_g, + A_log=A_log, + y=y, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + stride_g_token=raw_g.stride(0), + stride_g_head=raw_g.stride(1), + stride_g_dim=raw_g.stride(2), + stride_y_token=y.stride(0), + stride_y_head=y.stride(1), + stride_y_dim=y.stride(2), + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + beta=beta, + threshold=threshold, + SAFE_GATE=safe_gate, + LOWER_BOUND=lower_bound, + head_dim=head_dim, + chunk_size=chunk_size, + head_block_size=head_block_size, + num_warps=num_warps, + ) + return y + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + """Compute KDA state updates and outputs from gate prefix sums. + + q/k/g: [1, total_tokens, head_num, head_dim]; v: [1, total_tokens, head_num, value_head_dim]. + beta: [1, total_tokens, head_num]. + + Execution order and parallelism: + 1. Prepare Akk, Aqk, R, U, W, and Kg using only each chunk's inputs. + Each of these stages can process different chunks in parallel. + 2. Update the state in chunk order within each request: one chunk's output state + becomes the next chunk's input state. Save the state before processing each + chunk, S_in, in h. Compute E = U - W @ S_in and save it in v_new; + E contains the value correction used in each token's state update. + Different requests, heads, and blocks of value channels can run in parallel. + 3. Once h and v_new are saved, compute all chunk outputs in parallel using these tensors. + + h: [1, chunk_num, head_num, head_dim, value_head_dim]; v_new: [1, total_tokens, head_num, value_head_dim]. + The output kernel writes into v, overwriting its original values. The returned o shares v's memory. + """ + # 1. Build Akk and Aqk for all chunks: [1, total_tokens, head_num, chunk_size], where chunk_size is the chunk size. + # Akk weights earlier writes in each token's value correction; Aqk weights writes in the chunk output. + # Keep Aqk in fp32 until the output matrix multiplication. Only Akk is passed to solve_tril. + Akk, Aqk = chunk_kda_scaled_dot_kkt_fwd( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + output_dtype=torch.float32, + ) + # 2. Compute R=(I+Akk)^-1, then use matrix multiplication to obtain U/W for all tokens in each chunk. + R = solve_tril( + A=Akk, + cu_seqlens=cu_seqlens, + output_dtype=k.dtype, + ) + del Akk + # 3. Prepare U/W/Kg for the state update without reading the input state S_in. + # U=R@(beta*V), W=R@(beta*K*exp2(G)), Kg=K*exp2(G_last-G). + w, u, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + R=R, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del R + # 4. Update state in chunk order within each request. Requests, heads, and blocks of value channels run in parallel. + # Save the state before each chunk, S_in, in h; compute E=U-W@S_in and write it to v_new. + # S_out=exp2(G_last)[:, None]*S_in + Kg.T@E becomes the next chunk's S_in. + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + use_exp2=True, + ) + del w, u, kg + # 5. Read each chunk's saved h=S_in and v_new=E to compute outputs in parallel: + # O=scale*(Q*exp2(G))@S_in + Aqk@E combines the input state and the writes inside the chunk. + # Passing o=v overwrites the original values, which are no longer needed at this stage. + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + Aqk=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor, +): + """Convert activated log decays g=ell to chunk-local log2 prefixes, then run KDA. + + g: [1, total_tokens, head_num, head_dim] contains per-token natural-log decays, not raw gate logits. + chunk_local_cumsum followed by RCP_LN2 supplies G to the shared chunk path. + """ + chunk_size = FLA_CHUNK_SIZE + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, +): + """Activate packed raw gates and build G in one kernel before the shared KDA path. + + raw_g: [1, total_tokens, head_num, head_dim]; cu_seqlens separates requests. The gate kernel uses + the [total_tokens, head_num, head_dim] view and returns fp32 chunk-local log2 prefixes of the same shape. + """ + assert ( + raw_g.ndim == 4 and raw_g.shape[0] == 1 + ), "KDA prefill expects packed gates shaped [1, total_tokens, head_num, head_dim]" + chunk_size = FLA_CHUNK_SIZE + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # The gate kernel uses [total_tokens, head_num, head_dim]; downstream FLA ops add a leading batch dimension. + # Removing/restoring the leading dimension only creates tensor views. + g = fused_kda_gate_chunk_cumsum( + raw_g.squeeze(0), + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + ).unsqueeze(0) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + **kwargs, +): + assert q.ndim == 4 and q.shape[0] == 1, "KDA prefill expects packed q shaped [1, total_tokens, head_num, head_dim]" + assert cu_seqlens is not None, "cu_seqlens is required for packed KDA prefill" + if scale is None: + head_dim = k.shape[-1] + scale = head_dim ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + cu_seqlens: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, + **kwargs, +): + """KDA prefill with 64-token chunks and fused gate activation. + + Shapes: + q/k/raw_g: [1, total_tokens, head_num, head_dim]; v: [1, total_tokens, head_num, value_head_dim]. + beta: [1, total_tokens, head_num]. + A_log: [head_num]; g_bias: [head_num * head_dim] or [head_num, head_dim]. + initial_state: [request_num, head_num, head_dim, value_head_dim]. + cu_seqlens: [request_num + 1]; chunk_indices: [chunk_num, 2]. + + Per-token definition (one head; q/k/v/delta are column vectors; S: [head_dim, value_head_dim]): + ell_t is the activated log decay from the gate kernel. + S_decay = diag(exp(ell_t)) @ S_prev + delta_t = beta_t * (v_t - S_decay.T @ k_t) # Compute the value correction. + S_t = S_decay + outer(k_t, delta_t) # Write the correction into the state. + o_t = scale * (S_t.T @ q_t) # Read the updated state with q_t. + + beta is supplied after sigmoid. If use_qk_l2norm_in_kernel=True, normalize q/k first + with x / sqrt(sum(x * x) + 1e-6). The default scale is head_dim ** -0.5. + + Chunk equations (one head, chunk_size=64; Q/K/V hold token vectors in rows): + Q/K/P: [chunk_size, head_dim]; V/E/O: [chunk_size, value_head_dim]. + Akk/Aqk: [chunk_size, chunk_size]; S_in: [head_dim, value_head_dim]. + S_in is the state before the chunk; P=exp2(G); E_i=delta_i.T. + + Expanding the decayed state before token i's write gives: + S_decay_i = diag(P_i) @ S_in + + sum_{j 1: + start = tl.load(CuSeqLens + row) + end = tl.load(CuSeqLens + row + 1) + if start == end: + return + accepted = tl.load(Accepted + row) - 1 + state_idx = tl.load(Idx + row * MTP_SIZE + accepted) + else: + start, end = row, row + 1 + state_idx = tl.load(Idx + row) + ki = tl.arange(0, D) + vi = tl.program_id(0) * BV + tl.arange(0, BV) + bias = tl.load(Bias + head * D + ki) + amplitude = tl.exp(tl.load(A + head)) + state_offset = head * D * D + ki[:, None] * D + vi[None, :] + state = tl.load(State + state_idx * H * D * D + state_offset).to(tl.float32) + for token in range(start, end): + q = tl.load(Q + token * SQ + head * D + ki).to(tl.float32) + k = tl.load(K + token * SK + head * D + ki).to(tl.float32) + v = tl.load(V + token * SV + head * D + vi).to(tl.float32) + q *= tl.rsqrt(tl.sum(q * q, 0) + 1e-6) * (D ** -0.5) + k *= tl.rsqrt(tl.sum(k * k, 0) + 1e-6) + gate = tl.load(G + token * SG + head * D + ki).to(tl.float32) + decay = tl.exp(LOWER * tl.sigmoid(amplitude * (gate + bias))) + beta = tl.sigmoid(tl.load(B + token * SB + head).to(tl.float32)) + state *= decay[:, None] + delta = (v - tl.sum(state * k[:, None], 0)) * beta + state += k[:, None] * delta[None, :] + if MTP_SIZE > 1: + state_idx = tl.load(Idx + row * MTP_SIZE + token - start) + tl.store(State + state_idx * H * D * D + state_offset, state) + out = tl.sum(state * q[:, None], 0) + tl.store(O + (token * H + head) * D + vi, out) + # Match successive single-token calls when the state cache is BF16. + state = state.to(State.dtype.element_ty).to(tl.float32) + + +def fused_recurrent_kda( + q, + k, + v, + raw_gate, + raw_beta, + a_log, + gate_bias, + initial_state, + ssm_state_indices, + lower_bound=-5.0, + inplace_final_state=True, + cu_seqlens=None, + num_accepted_tokens=None, +): + assert inplace_final_state + batch, _, heads, dim = q.shape + assert dim == 128 + mtp_size = 1 + token_axis = 0 + if cu_seqlens is not None: + assert q.shape[0] == 1 and ssm_state_indices.ndim == 2 + batch, mtp_size = ssm_state_indices.shape + assert mtp_size > 1 and cu_seqlens.numel() == batch + 1 + assert num_accepted_tokens is not None and num_accepted_tokens.numel() == batch + token_axis = 1 + else: + assert q.shape[1] == 1 and ssm_state_indices.ndim == 1 + out = torch.empty_like(v, memory_format=torch.contiguous_format) + _kda_decode[(triton.cdiv(dim, 32), batch * heads)]( + q, + k, + v, + raw_gate, + raw_beta, + a_log, + gate_bias, + initial_state, + ssm_state_indices, + cu_seqlens, + num_accepted_tokens, + out, + q.stride(token_axis), + k.stride(token_axis), + v.stride(token_axis), + raw_gate.stride(token_axis), + raw_beta.stride(token_axis), + heads, + dim, + lower_bound, + 32, + mtp_size, + num_warps=4, + ) + return out, initial_state diff --git a/lightllm/common/basemodel/triton_kernel/mhc/__init__.py b/lightllm/common/basemodel/triton_kernel/mhc/__init__.py new file mode 100644 index 0000000000..af5a6bb969 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/mhc/__init__.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC：将残差扩展为 S 条 stream，在每个 attention / FFN 子层前合并、子层后更新。 + +T = tokens，S = streams，H = 每条 stream 的 hidden size，K = S * H，M = 2 * S + S * S。 +x: [T, K] 为当前子层的残差状态，R = x.view(T, S, H)。 +fn: [M, K]，scale: [3]，base: [M] 为可学习参数。 + +1. hc_expand：把 embedding [T, H] 复制 S 份，得到初始 x [T, K]。 + +2. hc_pre_norm：生成混合权重、合并 stream，并完成子层 RMSNorm： + inv_rms = rsqrt(mean(x.float() ** 2, dim=-1, keepdim=True) + rms_eps) + mixes = (x.float() @ fn.T) * inv_rms + mixes 按 S、S、S*S 拆成 pre_raw、post_raw、residual_raw，最后一组 view 为 [T, S, S]： + pre = sigmoid(pre_raw * scale[0] + base[:S]) + hc_eps + post_mix = post_multiplier * sigmoid(post_raw * scale[1] + base[S:2*S]) + residual_logits = residual_raw * scale[2] + base[2*S:].view(S, S) + residual_logits 经 softmax 和 Sinkhorn 得到 residual_mix [T, S, S]，行列和接近 1。 + 用 pre [T, S] 合并残差，得到子层输入： + layer_input[t, h] = sum_i pre[t, i] * R[t, i, h] # [T, H] + 合并结果先转 bf16，再在 H 维做带 norm_weight 的 RMSNorm。 + 投影和平方和默认由 DeepGEMM 计算；LIGHTLLM_DISABLE_DEEPGEMM_MHC=1 改用 PyTorch。 + 两种后端共用后续的 Triton 融合 kernel。 + +3. hc_post：用子层输出 y [T, H] 更新残差： + out[t, j, h] = post_mix[t, j] * y[t, h] + sum_i residual_mix[t, i, j] * R[t, i, h] + out 展平为 [T, K]，作为下一子层的 x；i、j 分别是输入、输出 stream。 + +4. hc_contract：模型末尾沿 S 维取均值，[T, K] -> [T, H]。 + +hc_pre_norm / hc_post 仅支持 S = 4；hc_pre_norm 要求 bf16 激活和 fp32 fn。 +""" + +from .post import hc_post +from .pre_norm import hc_pre_norm +from .streams import hc_contract, hc_expand + +__all__ = [ + "hc_contract", + "hc_expand", + "hc_post", + "hc_pre_norm", +] diff --git a/lightllm/common/basemodel/triton_kernel/mhc/post.py b/lightllm/common/basemodel/triton_kernel/mhc/post.py new file mode 100644 index 0000000000..b4050e0765 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/mhc/post.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC post-mixing of sublayer outputs into residual streams.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_post_4stream_kernel( + layer_output, # 输入 [T, H]，激活 dtype。 + residual, # 输入 [T, 4 * H]，激活 dtype；当前子层 pre-mixing 前的残差。 + residual_mix, # 输入 [T, 4, 4]，fp32；最后两维为输入、输出 stream。 + post_mix, # 输入 [T, 4]，fp32。 + output, # 输出 [T, 4 * H]，与 layer_output 相同 dtype。 + hidden: tl.constexpr, + layer_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + mix_stride_m: tl.constexpr, + post_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """将子层输出和原始残差混合为四条新 stream：[T, H] + [T, 4, H] -> [T, 4, H]。 + + T = tokens，H = hidden。每个 program 同时计算四条 stream 的 BLOCK_H 个 + hidden 元素，沿 residual_mix 的输入 stream 维归约；fp32 累加后写入 output。 + 输出按 [T, 4 * H] 展平存储，不返回 Tensor。 + """ + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + + # Compute all four outputs in one program so the layer output and + # residual streams are read only once, reducing prefill memory traffic. + layer_value = tl.load( + layer_output + token * layer_stride_m + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + residual_base = residual + token * residual_stride_m + hidden_offsets + residual_0 = tl.load(residual_base, mask=hidden_mask, other=0.0).to(tl.float32) + residual_1 = tl.load(residual_base + hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_2 = tl.load(residual_base + 2 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_3 = tl.load(residual_base + 3 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + + post_base = post_mix + token * post_stride_m + mix_base = residual_mix + token * mix_stride_m + accumulator_0 = layer_value * tl.load(post_base) + accumulator_1 = layer_value * tl.load(post_base + 1) + accumulator_2 = layer_value * tl.load(post_base + 2) + accumulator_3 = layer_value * tl.load(post_base + 3) + + accumulator_0 += residual_0 * tl.load(mix_base) + accumulator_0 += residual_1 * tl.load(mix_base + 4) + accumulator_0 += residual_2 * tl.load(mix_base + 8) + accumulator_0 += residual_3 * tl.load(mix_base + 12) + accumulator_1 += residual_0 * tl.load(mix_base + 1) + accumulator_1 += residual_1 * tl.load(mix_base + 5) + accumulator_1 += residual_2 * tl.load(mix_base + 9) + accumulator_1 += residual_3 * tl.load(mix_base + 13) + accumulator_2 += residual_0 * tl.load(mix_base + 2) + accumulator_2 += residual_1 * tl.load(mix_base + 6) + accumulator_2 += residual_2 * tl.load(mix_base + 10) + accumulator_2 += residual_3 * tl.load(mix_base + 14) + accumulator_3 += residual_0 * tl.load(mix_base + 3) + accumulator_3 += residual_1 * tl.load(mix_base + 7) + accumulator_3 += residual_2 * tl.load(mix_base + 11) + accumulator_3 += residual_3 * tl.load(mix_base + 15) + + output_base = output + token * out_stride_m + hidden_offsets + tl.store(output_base, accumulator_0, mask=hidden_mask) + tl.store(output_base + hidden, accumulator_1, mask=hidden_mask) + tl.store(output_base + 2 * hidden, accumulator_2, mask=hidden_mask) + tl.store(output_base + 3 * hidden, accumulator_3, mask=hidden_mask) + + +def hc_post( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """用一次 Triton launch 将子层输出混回残差 streams。 + + T = tokens，S = streams，H = 单条 stream 的 hidden size。 + + Args: + layer_output: [T, H]，连续的 attention / FFN 子层输出。 + residual: [T, S * H]，连续的、该子层 pre-mixing 前保存的原始残差。 + residual_mix: [T, S, S]，fp32；来自 hc_pre_norm。 + post_mix: [T, S]，fp32；来自同一次 pre 调用。 + streams: 残差 stream 数 S；当前 Triton 入口要求 S = 4。 + + Returns: + 新残差 [T, S * H]，dtype 与 layer_output 相同；输入张量保持不变。 + 逻辑输出 [T, S, H] 以 fp32 计算： + out[t, j, h] = post_mix[t, j] * layer_output[t, h] + + sum_i residual_mix[t, i, j] * residual.view(T, S, H)[t, i, h]。 + 矩阵第一个 stream 维 i 是输入，第二个 stream 维 j 是输出。 + """ + + tokens, hidden = layer_output.shape + assert streams == 4, "the fused mHC kernel is specialized for four streams" + assert layer_output.is_contiguous() and residual.is_contiguous() + output = torch.empty( + (tokens, streams * hidden), + dtype=layer_output.dtype, + device=layer_output.device, + ) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_post_4stream_kernel[(tokens, triton.cdiv(hidden, block_h))]( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden=hidden, + layer_stride_m=layer_output.stride(0), + residual_stride_m=residual.stride(0), + mix_stride_m=residual_mix.stride(0), + post_stride_m=post_mix.stride(0), + out_stride_m=output.stride(0), + BLOCK_H=block_h, + num_warps=8, + ) + return output diff --git a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py new file mode 100644 index 0000000000..2710040673 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC pre-mixing fused with the following RMSNorm.""" + +from __future__ import annotations + +import os +from typing import Tuple + +import torch +import triton +import triton.language as tl + +LIGHTLLM_DISABLE_DEEPGEMM_MHC = os.getenv("LIGHTLLM_DISABLE_DEEPGEMM_MHC", "False").upper() in ["ON", "TRUE", "1"] + + +@triton.jit +def _hc_prepare_prenorm_kernel( + gemm_partial, # 输入 [P, T, M]，fp32；投影的 split-K 部分和。 + sqrsum_partial, # 输入 [P, T]，fp32；残差平方和的 split-K 部分和。 + scale, # 输入 [3]，fp32；pre、post、residual 三组缩放系数。 + base, # 输入 [M]，fp32；M = 2 * S + S * S。 + pre, # 输出 [T, S]，fp32。 + post, # 输出 [T, S]，fp32。 + residual_mix, # 输出 [T, S, S]，fp32；最后两维为输入、输出 stream。 + gemm_stride_s, + gemm_stride_m: tl.constexpr, + sqrsum_stride_s, + sqrsum_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + FLATTENED_HIDDEN: tl.constexpr, + RMS_EPS: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, + N_SPLITS: tl.constexpr, +): + """归并分片并准备混合权重：[P, T, M] / [P, T] -> [T, S] / [T, S, S]。 + + T = tokens，S = STREAMS，M = 2 * S + S * S，P = N_SPLITS。 + FLATTENED_HIDDEN = S * H，用于从平方和计算每个 token 的 RMS 倒数。 + 每个 program 处理一个 token，写入 pre/post/residual_mix，不返回 Tensor。 + """ + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + pre_raw = tl.zeros((STREAMS,), dtype=tl.float32) + post_raw = tl.zeros((STREAMS,), dtype=tl.float32) + matrix_raw = tl.zeros((STREAMS * STREAMS,), dtype=tl.float32) + sqrsum = 0.0 + for split in tl.static_range(N_SPLITS): + partial_base = gemm_partial + split * gemm_stride_s + token * gemm_stride_m + pre_raw += tl.load(partial_base + stream_offsets) + post_raw += tl.load(partial_base + STREAMS + stream_offsets) + matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) + sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) + inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) + pre_raw *= inv_rms + post_raw *= inv_rms + matrix_raw *= inv_rms + + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) + + residual_logits = matrix_raw * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) + residual_logits = tl.reshape(residual_logits, (STREAMS, STREAMS)) + residual_logits = residual_logits - tl.max(residual_logits, axis=1)[:, None] + matrix = tl.exp(residual_logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_pre_combine_norm_kernel( + x, # 输入 [T, S * H]，bf16；逻辑形状 [T, S, H]。 + pre, # 输入 [T, S]，fp32。 + norm_weight, # 输入 [H]，子层 RMSNorm 权重，通常为 bf16。 + output, # 输出 [T, H]，bf16；加权合并后再做 RMSNorm。 + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + NORM_EPS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """融合 stream 合并与子层 RMSNorm：[T, S, H] -> [T, H]。 + + T = tokens，S = STREAMS，H = hidden。每个 program 处理一个 token 的 + 整行 hidden；先把合并结果舍入到 bf16，再以 fp32 计算均方值、归一化并乘 norm_weight。 + 结果写入 output，不返回 Tensor。 + """ + token = tl.program_id(0) + hidden_offsets = tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + + # Preserve the checkpoint's bf16 rounding between stream mixing and RMSNorm. + rounded = accumulator.to(tl.bfloat16).to(tl.float32) + variance = tl.sum(rounded * rounded, axis=0) / hidden + inv_rms = tl.rsqrt(variance + NORM_EPS) + weight = tl.load(norm_weight + hidden_offsets, mask=hidden_mask, other=0.0) + tl.store( + output + token * out_stride_m + hidden_offsets, + rounded * inv_rms * weight, + mask=hidden_mask, + ) + + +def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.device) -> int: + """由 T = tokens、K = flattened_hidden 和设备 SM 数确定 split-K 分片数 P。 + + tokens 和 flattened_hidden 为整数，device 指定 CUDA 设备；返回整数 P。 + P 决定 gemm_partial [P, T, M] 和 sqrsum_partial [P, T] 的首维， + 其中 M = 2 * S + S * S，S 为 stream 数。 + """ + grid_size = triton.cdiv(tokens, 64) + k_blocks = triton.cdiv(flattened_hidden, 64) + sms = torch.cuda.get_device_properties(device).multi_processor_count + return max(1, min(sms // max(grid_size, 1), k_blocks // 4)) + + +def hc_pre_norm( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + streams: int, + rms_eps: float, + norm_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """将 mHC pre-mixing 与紧随其后的子层 RMSNorm 融合。 + + T = tokens，S = streams，H = 单条 stream 的 hidden size，M = 2 * S + S * S。 + + Args: + x: [T, S * H]，连续的 bf16 残差激活；逻辑形状 [T, S, H]。 + fn: [M, S * H]，连续的 fp32 投影权重。 + scale: [3]，fp32；分别缩放 pre、post、residual logits。 + base: [M]，fp32；按 [S, S, S * S] 分组的偏置。 + norm_weight: [H]，连续的子层 RMSNorm 权重，通常为 bf16。 + streams: 残差 stream 数 S；当前要求 S = 4。 + rms_eps: 在 S * H 维计算投影前 RMS 倒数时使用的 epsilon。 + norm_eps: 合并 stream 后，在 H 维做子层 RMSNorm 时使用的 epsilon。 + hc_eps: pre sigmoid 后的偏移量及 Sinkhorn 的稳定项。 + sinkhorn_iters: Sinkhorn 迭代次数。 + post_multiplier: post sigmoid 的乘数，默认 2.0。 + + Returns: + (layer_input, residual_mix, post_mix)：分别为 bf16 [T, H]、fp32 [T, S, S]、fp32 [T, S]。 + layer_input 已完成子层 RMSNorm；residual_mix[t, i, j] 对应输入 i 到输出 j。 + x 保持不变，留给对应 hc_post 使用。 + + 默认用 DeepGEMM 融合投影与平方和计算，Triton 完成 Sinkhorn、stream 合并和 RMSNorm。 + 合并结果先舍入到 bf16，再做 RMSNorm。 + LIGHTLLM_DISABLE_DEEPGEMM_MHC=1 禁用 DeepGEMM，改用 PyTorch 计算投影与平方和。 + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused mHC kernel is specialized for four streams" + assert x.dtype == torch.bfloat16 and fn.dtype == torch.float32 + assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + if LIGHTLLM_DISABLE_DEEPGEMM_MHC: + # PyTorch produces one complete partition: [1, T, M] and [1, T]. + x_fp32 = x.float() + gemm_partial = (x_fp32 @ fn.T).unsqueeze(0) + sqrsum_partial = x_fp32.square().sum(dim=-1).unsqueeze(0) + n_splits = 1 + else: + from deep_gemm import tf32_hc_prenorm_gemm + + mix_size = (2 + streams) * streams + n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) + gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device) + sqrsum_partial = torch.empty((n_splits, tokens), dtype=torch.float32, device=x.device) + tf32_hc_prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + _hc_prepare_prenorm_kernel[(tokens,)]( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_partial.stride(0), + gemm_partial.stride(1), + sqrsum_partial.stride(0), + sqrsum_partial.stride(1), + pre.stride(0), + residual_mix.stride(0), + FLATTENED_HIDDEN=flattened_hidden, + RMS_EPS=rms_eps, + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + N_SPLITS=n_splits, + num_warps=1, + ) + + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) + block_h = triton.next_power_of_2(hidden) + _hc_pre_combine_norm_kernel[(tokens,)]( + x, + pre, + norm_weight, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + NORM_EPS=norm_eps, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post diff --git a/lightllm/common/basemodel/triton_kernel/mhc/streams.py b/lightllm/common/basemodel/triton_kernel/mhc/streams.py new file mode 100644 index 0000000000..4c2c56757c --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/mhc/streams.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Expansion and mean contraction of flattened residual streams.""" + +from __future__ import annotations + +import torch + + +def hc_expand(x: torch.Tensor, streams: int) -> torch.Tensor: + """将每个 token 的 embedding 复制到 S 条 residual stream。 + + 输入 x: [T, H]，streams: 标量 S；T = tokens，H = hidden size。 + 输出: [T, S * H]，dtype/device 与 x 相同，逻辑形状为 [T, S, H]。 + output.view(T, S, H)[t, s, h] = x[t, h]，沿 stream 维复制完整的 hidden 向量。 + """ + + assert x.ndim == 2 + return x.unsqueeze(1).expand(-1, streams, -1).reshape(x.shape[0], -1) + + +def hc_contract(x: torch.Tensor, streams: int) -> torch.Tensor: + """沿 stream 维取均值，将模型末尾的多条残差收缩成一个 hidden 向量。 + + 输入 x: [T, S * H]，streams: 标量 S；T = tokens，H = 单条 stream 的 hidden size。 + 输出: [T, H]，dtype/device 与 x 相同，即 x.view(T, S, H).mean(dim=1)。 + 此操作不含可学习权重；模型若使用 hc_head，应调用其专用的加权收缩算子。 + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + return x.view(x.shape[0], streams, -1).mean(dim=1) diff --git a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py index c62c5eb5d2..36203999d1 100644 --- a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py +++ b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py @@ -18,20 +18,27 @@ def gated_rmsnorm_forward_kernel( Z, # pointer to the other branch (required, not optional) stride_x_row, # how much to increase the pointer when moving by 1 row stride_y_row, - stride_z_row, + stride_z_token, + stride_z_head, M, # number of rows in X N, # number of columns in X eps, # epsilon to avoid division by zero BLOCK_N: tl.constexpr, HAS_BIAS: tl.constexpr, NORM_BEFORE_GATE: tl.constexpr, + SIGMOID_GATE: tl.constexpr, + Z_HEADS: tl.constexpr, ): # Map the program id to the row of X and Y it should compute. row = tl.program_id(0) group = tl.program_id(1) X += row * stride_x_row + group * N Y += row * stride_y_row + group * N - Z += row * stride_z_row + group * N + # X is flattened to [tokens * heads, N], while GDN can keep Z as a + # zero-copy [tokens, heads, N] slice of its packed projection output. + z_token = row // Z_HEADS + z_head = row % Z_HEADS + Z += z_token * stride_z_token + z_head * stride_z_head + group * N W += group * N if HAS_BIAS: B += group * N @@ -40,7 +47,7 @@ def gated_rmsnorm_forward_kernel( x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) if not NORM_BEFORE_GATE: z = tl.load(Z + cols, mask=cols < N).to(tl.float32) - x *= z * tl.sigmoid(z) + x *= tl.sigmoid(z) if SIGMOID_GATE else z * tl.sigmoid(z) # RMS norm: compute variance directly without mean subtraction xbar = tl.where(cols < N, x, 0.0) var = tl.sum(xbar * xbar, axis=0) / N @@ -55,7 +62,7 @@ def gated_rmsnorm_forward_kernel( y = x_hat * w + b if HAS_BIAS else x_hat * w if NORM_BEFORE_GATE: z = tl.load(Z + cols, mask=mask).to(tl.float32) - y *= z * tl.sigmoid(z) + y *= tl.sigmoid(z) if SIGMOID_GATE else z * tl.sigmoid(z) # Write output tl.store(Y + cols, y, mask=mask) @@ -103,6 +110,7 @@ def gated_rmsnorm_forward( group_size: int = None, norm_before_gate: bool = True, run_config: dict = None, + gate_type: str = "silu", ): M, N = x.shape if group_size is None: @@ -112,8 +120,21 @@ def gated_rmsnorm_forward( assert x.stride(-1) == 1 # z is required for gated_rmsnorm assert z is not None, "z cannot be None for gated_rmsnorm_forward" + assert gate_type in ("silu", "sigmoid"), f"unsupported gate type: {gate_type}" + # Accept GDN's strided 3D gate without materializing a flattened copy. + assert z.ndim in (2, 3), f"z must be [M, N] or [tokens, heads, N], got shape={z.shape}" assert z.stride(-1) == 1 - assert z.shape == (M, N) + if z.ndim == 2: + assert z.shape == (M, N) + z_heads = 1 + stride_z_token = z.stride(0) + stride_z_head = 0 + else: + assert z.shape[-1] == N, f"z.shape[-1]={z.shape[-1]} must match N={N}" + assert z.shape[0] * z.shape[1] == M, f"z token/head rows={z.shape[0] * z.shape[1]} must match M={M}" + z_heads = z.shape[1] + stride_z_token = z.stride(0) + stride_z_head = z.stride(1) assert weight.shape == (N,) assert weight.stride(-1) == 1 if bias is not None: @@ -156,12 +177,15 @@ def gated_rmsnorm_forward( z, x.stride(0), out.stride(0), - z.stride(0), + stride_z_token, + stride_z_head, M, group_size, eps, BLOCK_N=BLOCK_N, NORM_BEFORE_GATE=norm_before_gate, + SIGMOID_GATE=gate_type == "sigmoid", + Z_HEADS=z_heads, num_warps=num_warps, ) return out diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py index 05544e149a..b3597b4bc8 100644 --- a/lightllm/common/kv_cache_mem_manager/__init__.py +++ b/lightllm/common/kv_cache_mem_manager/__init__.py @@ -8,6 +8,7 @@ from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager from .qwen3next_mem_manager import Qwen3NextMemManager +from .glm5_next_mem_manager import Glm5NextMemManager __all__ = [ "KvCacheAllocator", @@ -21,4 +22,5 @@ "FP8StaticPerHeadQuantMemManager", "FP8StaticPerTensorQuantMemManager", "Qwen3NextMemManager", + "Glm5NextMemManager", ] diff --git a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py new file mode 100644 index 0000000000..45a3be667c --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py @@ -0,0 +1,122 @@ +import torch + +from lightllm.common.kv_cache_mem_manager.operator import LinearAttMemOperator +from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( + Qwen3NextMemManager, + Qwen3NextLinearAttPageHelper, +) +from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv +from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io + + +class Glm5NextMemOperator(LinearAttMemOperator): + def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): + output = self.mem_manager.get_att_input_params(layer_index)[:, :, : kv.shape[-1]] + destindex_copy_kv(kv, mem_index, output) + + +class Glm5NextMemManager(Qwen3NextMemManager): + """One packed token buffer; KDA uses the standard big/small page pools.""" + + operator_class = Glm5NextMemOperator + + def __init__(self, *args, mla_head_dim=512, **kwargs): + self.mla_head_dim = mla_head_dim + super().__init__(*args, **kwargs) + + def get_cell_size(self): + return self.head_dim * self.dtype.itemsize * self.layer_num + + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + assert head_num == 1 + self.kv_buffer = torch.empty((layer_num, size + 1, 1, head_dim), dtype=dtype, device="cuda") + self._init_linear_att_buffers() + + def _layer_buffer(self, layer_index): + return self.kv_buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] + + def get_att_input_params(self, layer_index): + return self._layer_buffer(layer_index)[:, :, : self.mla_head_dim] + + def get_indexer_k_buffer(self, layer_index): + return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] + + def write_to_shm(self, req_manager): + self.req_to_indexer_tail = req_manager.req_to_indexer_tail + return super().write_to_shm(req_manager) + + def _create_att_state_page_helper(self): + return Glm5NextAttStatePageHelper(self) + + def get_paged_kv_move_buffer_shape(self, page_num, page_size): + # Packed MLA/index KV is replicated across TP ranks. + return (page_num, page_size, self.layer_num, self.head_num, self.head_dim) + + def write_mem_to_page_kv_move_buffer( + self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None + ): + if page_kind != "kv": + return super().write_mem_to_page_kv_move_buffer( + mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind, req_idx + ) + pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] + pin_indexes.numpy()[:] = mem_indexes + mla_page_io( + mem_indexes=pin_indexes.cuda(non_blocking=True), + page_tensor=self.kv_move_buffer[page_index], + kv_buffer=mem_managers[dp_index * dp_world_size].kv_buffer, + mode="write", + ) + + def read_page_kv_move_buffer_to_mem( + self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None + ): + if page_kind != "kv": + return super().read_page_kv_move_buffer_to_mem( + mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind, req_idx + ) + pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] + pin_indexes.numpy()[:] = mem_indexes + indexes = pin_indexes.cuda(non_blocking=True) + for mem in mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size]: + mla_page_io( + mem_indexes=indexes, + page_tensor=self.kv_move_buffer[page_index], + kv_buffer=mem.kv_buffer, + mode="read", + ) + + +class Glm5NextAttStatePageHelper(Qwen3NextLinearAttPageHelper): + """Append the replicated K-pool tail to the global Conv/SSM state page.""" + + def __init__(self, mem_manager): + super().__init__(mem_manager) + self.tail_dtype = mem_manager.req_to_indexer_tail.buffer.dtype + self.tail_shape = ( + self.linear_config.get_full_att_kv_layer_num_with_draft_model(), + *mem_manager.req_to_indexer_tail.buffer.shape[2:], + ) + self.tail_offset = ((self.state_nbytes + 15) // 16) * 16 + self.tail_nbytes = self.tail_shape[0] * self.tail_shape[1] * self.tail_shape[2] * self.tail_dtype.itemsize + self.state_nbytes = self.tail_offset + self.tail_nbytes + + def _view_page_to_tail(self, page_index): + page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) + return ( + page_bytes[self.tail_offset : self.tail_offset + self.tail_nbytes] + .view(self.tail_dtype) + .view(self.tail_shape) + ) + + def write_req_to_page(self, page_index, req_idx, dp_mems): + super().write_req_to_page(page_index, req_idx, dp_mems) + # The ring uses absolute token positions; sequence length already travels + # with the request, so restoring it needs no additional PD metadata. + self._view_page_to_tail(page_index).copy_(dp_mems[0].req_to_indexer_tail.buffer[:, req_idx], non_blocking=True) + + def read_page_to_req(self, page_index, req_idx, dp_mems): + super().read_page_to_req(page_index, req_idx, dp_mems) + tail = self._view_page_to_tail(page_index) + for mem in dp_mems: + mem.req_to_indexer_tail.buffer[:, req_idx].copy_(tail, non_blocking=True) diff --git a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py index 71158ac97a..e935539e37 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py @@ -6,7 +6,6 @@ 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.utils.log_utils import init_logger -from lightllm.common.state_cache_manager import LinearAttCacheConfig if TYPE_CHECKING: from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient @@ -22,7 +21,7 @@ class LinearAttMemOperator(BaseMemManagerOperator): def __init__(self, mem_manager): super().__init__(mem_manager) - self.linear_config = LinearAttCacheConfig.load_from_args() + self.linear_config = mem_manager.linear_config def load_cpu_cache_to_gpu( self, diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index ce22808e16..1962fa71b0 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -87,9 +87,13 @@ def write_to_shm(self, req_manager): def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - Qwen3NextLinearAttPageHelper(self).assert_page_size() + self.att_state_page_helper = self._create_att_state_page_helper() + self.att_state_page_helper.assert_page_size() return kv_move_buffer + def _create_att_state_page_helper(self): + return Qwen3NextLinearAttPageHelper(self) + def write_mem_to_page_kv_move_buffer( self, mem_indexes, @@ -112,7 +116,7 @@ def write_mem_to_page_kv_move_buffer( ) assert page_kind == "att_state", f"unknown page_kind={page_kind}" assert req_idx is not None - helper = Qwen3NextLinearAttPageHelper(self) + helper = self.att_state_page_helper dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) helper.write_req_to_page(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) return @@ -139,7 +143,7 @@ def read_page_kv_move_buffer_to_mem( ) assert page_kind == "att_state", f"unknown page_kind={page_kind}" assert req_idx is not None - helper = Qwen3NextLinearAttPageHelper(self) + helper = self.att_state_page_helper dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) helper.read_page_to_req(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) return diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index abcd4f491e..ca1f64d5b8 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,6 +1,13 @@ from .base import ReqManager from .linear_att import ReqManagerForMamba +from .glm5_next import Glm5NextReqManager from .hybrid_base import HybridAttentionReqManager from .req_sampling_params import ReqSamplingParamsManager -__all__ = ["ReqManager", "HybridAttentionReqManager", "ReqManagerForMamba", "ReqSamplingParamsManager"] +__all__ = [ + "ReqManager", + "HybridAttentionReqManager", + "ReqManagerForMamba", + "Glm5NextReqManager", + "ReqSamplingParamsManager", +] diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py new file mode 100644 index 0000000000..b4814afbf7 --- /dev/null +++ b/lightllm/common/req_manager/glm5_next.py @@ -0,0 +1,34 @@ +from lightllm.common.state_cache_manager import LayerCache +from lightllm.utils.envs_utils import get_env_start_args + +from .linear_att import ReqManagerForMamba + + +class Glm5NextReqManager(ReqManagerForMamba): + """KDA runtime state and the NSA indexer's incomplete four-token pool.""" + + def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config): + # Both checkpoint sizes are multiples of this hash page. A restored + # prefix therefore has no incomplete K-pool to serialize or replay. + assert ( + get_env_start_args().linear_att_hash_page_size % linear_config.index_kpool == 0 + ), "GLM K-pool requires cache pages aligned to index_kpool" + super().__init__(max_request_num, max_sequence_length, mem_manager, linear_config) + self.req_to_indexer_tail = LayerCache( + size=max_request_num + 1, + dtype=linear_config.full_att_dtype, + shape=(linear_config.index_kpool + self.mtp_step, 2 * linear_config.index_head_dim), + layer_num=linear_config.get_full_att_kv_layer_num_with_draft_model(), + device="cuda", + ) + + def get_indexer_tail_buffer(self, layer_index): + return self.req_to_indexer_tail.buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] + + def init_hybrid_attention_state(self, req): + super().init_hybrid_attention_state(req) + self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() + + def restore_state(self, req, state_cache_manager, buffer_idx): + super().restore_state(req, state_cache_manager, buffer_idx) + self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index af0479a158..306a20db80 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -1,13 +1,19 @@ from .base import StateCacheManager from .layer_cache import LayerCache from .linear_att import LinearAttCacheConfig, LinearAttCacheManager +from .glm5_next import Glm5NextCacheConfig def get_hybrid_cache_config(): """Return the model-specific layout used by hybrid CPU/disk cache pages.""" - from lightllm.utils.config_utils import is_linear_att_mixed_model + from transformers.configuration_utils import PretrainedConfig from lightllm.utils.envs_utils import get_env_start_args - if is_linear_att_mixed_model(get_env_start_args().model_dir): + args = get_env_start_args() + model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) + model_type = model_cfg["model_type"] + if model_type in ("glm5_next", "glm5_next_text"): + return Glm5NextCacheConfig.from_model_config(model_cfg, args) + if model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): return LinearAttCacheConfig.load_from_args() raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py new file mode 100644 index 0000000000..13195a74ac --- /dev/null +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -0,0 +1,63 @@ +import dataclasses + +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.torch_dtype_utils import get_torch_dtype + +from .linear_att import LinearAttCacheConfig + + +@dataclasses.dataclass +class Glm5NextCacheConfig(LinearAttCacheConfig): + """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + + index_kpool: int = 4 + index_head_dim: int = 128 + + INDEX_PADDING_BYTES = 144 + + @classmethod + def from_model_config(cls, config, args): + config = config.get("text_config", config) + linear = config["linear_attn_config"] + tp = args.tp // args.dp + assert linear["num_heads"] % tp == 0 + layers = config["num_hidden_layers"] + assert config["layer_types"] == [ + "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) + ] + dtype = get_torch_dtype(args.data_type) + # Raw index keys and compression scores only live in request tails. + packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize + return cls( + tp_world_size=tp, + full_att_all_num_kv_heads=1, + full_att_dtype=dtype, + full_att_num_kv_heads=1, + full_att_head_dim=packed_dim, + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + num_linear_k_heads=linear["num_heads"] // tp, + num_linear_v_heads=linear["num_heads"] // tp, + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + linear_layer_num=len(linear["kda_layers"]), + conv_state_dtype=dtype, + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=4, + all_layer_num=layers, + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, + index_kpool=config["index_kpool"], + index_head_dim=config["index_head_dim"], + ) + + def get_cpu_cache_full_att_bytes(self): + args = get_env_start_args() + page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + assert page_tokens == args.cpu_cache_token_page_size + return ( + self.full_att_head_dim + * self.full_att_dtype.itemsize + * self.get_full_att_kv_layer_num_with_draft_model() + * page_tokens + ) diff --git a/lightllm/models/builtin.py b/lightllm/models/builtin.py index 77a0d920ce..455b192d03 100644 --- a/lightllm/models/builtin.py +++ b/lightllm/models/builtin.py @@ -21,6 +21,8 @@ def _tarsier_text_model_is(model_type): ModelRegistry.register("gemma4", "lightllm.models.gemma4.model:Gemma4TpPartModel", is_multimodal=True) ModelRegistry.register("gemma", "lightllm.models.gemma_2b.model:Gemma_2bTpPartModel") ModelRegistry.register("glm4_moe_lite", "lightllm.models.glm4_moe_lite.model:Glm4MoeLiteTpPartModel") +ModelRegistry.register("glm5_next", "lightllm.models.glm5_next.model:Glm5NextTpPartModel", is_multimodal=True) +ModelRegistry.register("glm5_next_text", "lightllm.models.glm5_next.model:Glm5NextTpPartModel") ModelRegistry.register("gpt_oss", "lightllm.models.gpt_oss.model:GptOssTpPartModel") ModelRegistry.register("internlm", "lightllm.models.internlm.model:InternlmTpPartModel") ModelRegistry.register("internlm2", "lightllm.models.internlm2.model:Internlm2TpPartModel") @@ -155,6 +157,11 @@ def _tarsier_text_model_is(model_type): ("vanilla_with_att", "eagle_with_att"), "lightllm.models.glm4_moe_lite_mtp.model:Glm4MoeLiteMTPModel", ) +DraftModelRegistry.register( + ("glm5_next", "glm5_next_text"), + ("vanilla_with_att", "eagle_with_att"), + "lightllm.models.glm5_next_mtp.model:Glm5NextMTPModel", +) DraftModelRegistry.register( "mistral", ("vanilla_no_att", "eagle_no_att"), "lightllm.models.mistral_mtp.model:MistralMTPModel" ) diff --git a/lightllm/models/glm5_next/__init__.py b/lightllm/models/glm5_next/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/glm5_next_visual.py b/lightllm/models/glm5_next/glm5_next_visual.py new file mode 100644 index 0000000000..43637c2cbb --- /dev/null +++ b/lightllm/models/glm5_next/glm5_next_visual.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# Adapted from Hugging Face Transformers' GLM-5-Next vision encoder. + +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open + +from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm +from lightllm.server.visualserver import get_vit_attn_backend +from .vision_process import Glm5NextImageProcessor + + +class Glm5NextVisionRMSNorm(nn.Module): + def __init__(self, hidden_size, eps): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, hidden_states): + return rms_norm(hidden_states, self.weight, eps=self.eps, round_norm_before_weight=True) + + +class Glm5NextVisionMLP(nn.Module): + def __init__(self, hidden_size, intermediate_size, limit, bias=False): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=bias) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=bias) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=bias) + self.limit = limit + + def forward(self, x): + gate = self.gate_proj(x).clamp(max=self.limit) + up = self.up_proj(x).clamp(min=-self.limit, max=self.limit) + return self.down_proj(F.silu(gate) * up) + + +class Glm5NextVisionAttention(nn.Module): + def __init__(self, hidden_size, num_heads, eps, bias): + super().__init__() + self.num_heads = num_heads + self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=bias) + self.proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.q_norm = Glm5NextVisionRMSNorm(hidden_size // num_heads, eps) + self.k_norm = Glm5NextVisionRMSNorm(hidden_size // num_heads, eps) + + def forward(self, x, cu_seqlens, max_seqlen, rotary_cos, rotary_sin): + qkv = self.qkv(x).reshape(x.shape[0], 3, self.num_heads, -1) + q, k = qk_rms_norm(qkv, self.q_norm.weight, self.k_norm.weight, self.q_norm.eps) + v = qkv[:, 2] + q = apply_rotary_pos_emb_triton(q, rotary_cos, rotary_sin) + k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) + out = torch.empty_like(q) + get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + return self.proj(out.reshape(x.shape[0], -1)) + + +class Glm5NextVisionBlock(nn.Module): + def __init__(self, hidden_size, intermediate_size, num_heads, rms_norm_eps, swiglu_limit, attention_bias): + super().__init__() + self.norm1 = Glm5NextVisionRMSNorm(hidden_size, rms_norm_eps) + self.norm2 = Glm5NextVisionRMSNorm(hidden_size, rms_norm_eps) + self.attn = Glm5NextVisionAttention(hidden_size, num_heads, rms_norm_eps, attention_bias) + self.mlp = Glm5NextVisionMLP(hidden_size, intermediate_size, swiglu_limit, attention_bias) + + def forward(self, x, cu_seqlens, max_seqlen, rotary_cos, rotary_sin): + x = x + self.attn(self.norm1(x), cu_seqlens, max_seqlen, rotary_cos, rotary_sin) + return x + self.mlp(self.norm2(x)) + + +class Glm5NextVisionPatchEmbed(nn.Module): + def __init__(self, in_channels, hidden_size, patch_size, temporal_patch_size): + super().__init__() + self.in_channels = in_channels + self.kernel = (temporal_patch_size, patch_size, patch_size) + self.proj = nn.Conv3d(in_channels, hidden_size, kernel_size=self.kernel, stride=self.kernel) + + def forward(self, x): + x = x.reshape(-1, self.in_channels, *self.kernel) + return self.proj(x).flatten(1) + + +class Glm5NextVisionPatchMerger(Glm5NextVisionMLP): + def __init__(self, hidden_size, intermediate_size, limit): + super().__init__(hidden_size, intermediate_size, limit) + self.proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.post_projection_norm = nn.LayerNorm(hidden_size) + + def forward(self, x): + return super().forward(F.gelu(self.post_projection_norm(self.proj(x)))) + + +class Glm5NextVisionTransformer(Qwen2VisionTransformerPretrainedModel): + """GLM encoder using the existing variable-resolution image batching and cache interface.""" + + def __init__( + self, + kvargs, + hidden_size=1024, + out_hidden_size=4096, + depth=24, + intermediate_size=4096, + projection_intermediate_size=10240, + num_heads=16, + in_channels=3, + patch_size=14, + temporal_patch_size=2, + spatial_merge_size=2, + rms_norm_eps=1e-5, + swiglu_limit=10.0, + attention_bias=True, + hidden_act="silu", + rope_parameters=None, + **kwargs, + ): + nn.Module.__init__(self) + assert hidden_act == "silu" + self.data_type = kvargs.get("data_type", "bfloat16") + self._init_datatype() + self.hidden_size = out_hidden_size + self.spatial_merge_size = spatial_merge_size + self.patch_size = patch_size + self.patch_embed = Glm5NextVisionPatchEmbed(in_channels, hidden_size, patch_size, temporal_patch_size) + rope_parameters = rope_parameters or {"rope_type": "axial", "rope_theta": 10000.0} + assert rope_parameters["rope_type"] == "axial" + rotary_dim = hidden_size // num_heads // 2 + self.rotary_inv_freq = 1.0 / ( + rope_parameters["rope_theta"] ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim) + ) + self.blocks = nn.ModuleList( + [ + Glm5NextVisionBlock( + hidden_size, intermediate_size, num_heads, rms_norm_eps, swiglu_limit, attention_bias + ) + for _ in range(depth) + ] + ) + self.post_layernorm = Glm5NextVisionRMSNorm(hidden_size, rms_norm_eps) + self.downsample = nn.Conv2d( + hidden_size, out_hidden_size, kernel_size=spatial_merge_size, stride=spatial_merge_size + ) + self.merger = Glm5NextVisionPatchMerger(out_hidden_size, projection_intermediate_size, swiglu_limit) + + def load_model(self, weight_dir): + self.processor = Glm5NextImageProcessor.from_pretrained(weight_dir) + weights = {} + prefix = "model.visual." + for filename in os.listdir(weight_dir): + if filename.endswith(".safetensors"): + with safe_open(os.path.join(weight_dir, filename), framework="pt", device="cpu") as f: + for name in f.keys(): + if name.startswith(prefix): + weights[name[len(prefix) :]] = f.get_tensor(name) + self.load_state_dict(weights, strict=True) + + def rot_pos_emb(self, grid_thw, device): + positions = [] + size = self.spatial_merge_size + for t, h, w in grid_thw.tolist(): + rows, cols = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") + shape = (h // size, size, w // size, size) + rows = rows.reshape(shape).transpose(1, 2).flatten() + cols = cols.reshape(shape).transpose(1, 2).flatten() + positions.append(torch.stack((rows, cols), -1).repeat(t, 1)) + positions = torch.cat(positions).to(device) + # Keep frequencies in FP32 and evaluate trig on the same device as attention. + self.rotary_inv_freq = self.rotary_inv_freq.to(device) + angles = positions[..., None].float() * self.rotary_inv_freq + return angles.cos().flatten(1), angles.sin().flatten(1) + + def forward(self, hidden_states, grid_thw): + hidden_states = self.patch_embed(hidden_states) + # Grid metadata stays on the CPU; attention and rotary tensors are copied once per batch. + grid_thw = grid_thw.cpu() + rotary_cos, rotary_sin = self.rot_pos_emb(grid_thw, hidden_states.device) + lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + max_seqlen = lengths.max().item() + cu_seqlens = F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)).to(hidden_states.device) + for block in self.blocks: + hidden_states = block(hidden_states, cu_seqlens, max_seqlen, rotary_cos, rotary_sin) + hidden_states = self.post_layernorm(hidden_states) + size = self.spatial_merge_size + hidden_states = hidden_states.reshape(-1, size, size, hidden_states.shape[-1]).permute(0, 3, 1, 2) + hidden_states = self.downsample(hidden_states).reshape(-1, self.hidden_size) + return self.merger(hidden_states) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py new file mode 100644 index 0000000000..cf13ab83de --- /dev/null +++ b/lightllm/models/glm5_next/indexer.py @@ -0,0 +1,125 @@ +import torch +import triton + +from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops + +from .triton_kernel.index_quant import hadamard_transform_quant_fp8 +from .triton_kernel.kpool import compress_pools, gather_pools, gather_paged_pools, get_pool_ranges, expand_topk + + +class Glm5NextNsaInfer: + """K-pool indexing with pooled token KV and a small per-request raw tail.""" + + def __init__(self, layer_idx, network_config, tp_world_size): + self.layer_idx = layer_idx + self.topk = network_config["index_topk"] + self.heads = network_config["index_n_heads"] + self.dim = network_config["index_head_dim"] + self.eps = network_config["rms_norm_eps"] + + def select_topk_indices(self, logits, lengths, indices): + """Select row-relative indices, with valid entries before -1 padding.""" + if HAS_VLLM: + # next_n=1 treats each query as an independent row, including prefill. + # The decode entry splits long rows before merging their candidates. + # Tested with vLLM 0.22.1; persistent_topk can drop candidates (#51782). + vllm_ops.top_k_per_row_decode( + logits, 1, lengths, indices, logits.shape[0], logits.stride(0), logits.stride(1), indices.shape[1] + ) + else: + positions = torch.arange(logits.shape[1], device=logits.device) + logits.masked_fill_(positions[None, :] >= lengths[:, None], -float("inf")) + selected = torch.topk(logits, indices.shape[1], dim=-1, sorted=True).indices + indices.copy_(selected.masked_fill(selected >= lengths[:, None], -1)) + + def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): + k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.eps) + gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) + raw = torch.cat((k, gate), -1) + tail = infer_state.req_manager.get_indexer_tail_buffer(self.layer_idx) + packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) + compress_pools( + raw=raw, + tail=tail, + packed_buffer=packed_buffer, + ape=layer_weight.index_kpool_compress_ape.weight, + lengths=att_state.lengths, + starts=att_state.ks, + ragged=att_state.ragged_mem_index, + req_idx=infer_state.b_req_idx, + cu_q_lens=infer_state.b1_cu_q_seq_len, + seq_lens=infer_state.b_seq_len, + max_q_len=infer_state.max_q_seq_len, + mtp_index=None if infer_state.is_prefill else infer_state.b_mtp_index, + ) + + if infer_state.max_kv_seq_len <= self.topk: + return expand_topk(None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True) + + # The small indexer is replicated: no all-gather of query heads and + # identical pool selection on every TP rank. + q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.heads, self.dim) + q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim ** -0.5) + weights = layer_weight.weights_proj_.mm(hidden_states.float()) + weights = weights * (self.heads ** -0.5 * self.dim ** -0.5) * q_scale.squeeze(-1) + max_pools = triton.cdiv(infer_state.max_kv_seq_len, 4 * 128) * 128 + # Prefill usually has many queries: parallelize over Q and reuse K within each query tile. + # Decode has few queries, so paged MQA also splits K across SMs for parallelism. + if infer_state.is_prefill: + groups = self._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + else: + groups = self._get_decode_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) + + def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): + keys = gather_pools( + packed_buffer, + infer_state.req_manager.req_to_token_indexs, + infer_state.b_req_idx, + infer_state.b_seq_len, + max_pools, + ) + starts, ends, lengths = get_pool_ranges( + att_state.lengths, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools + ) + groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) + # Budget 64 MiB for FP32 logits. At 1M tokens this permits 64 queries per chunk, + # limiting Q parallelism even when the prefill batch contains many tokens. + chunk_size = max(1, min(q_fp8.shape[0], 16 * 1024 * 1024 // max_pools)) + import deep_gemm + + for start in range(0, q_fp8.shape[0], chunk_size): + end = min(start + chunk_size, q_fp8.shape[0]) + logits = deep_gemm.fp8_mqa_logits( + q_fp8[start:end], + keys, + weights[start:end], + starts[start:end], + ends[start:end], + clean_logits=False, + max_seqlen_k=max_pools, + ) + self.select_topk_indices(logits, lengths[start:end], groups[start:end]) + return groups + + def _get_decode_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): + import deep_gemm + + lengths = (att_state.lengths // 4).view(-1, 1) + pages, block_table = gather_paged_pools( + packed_buffer, + infer_state.req_manager.req_to_token_indexs, + infer_state.b_req_idx, + lengths, + max_pools, + ) + metadata = deep_gemm.get_paged_mqa_logits_metadata(lengths, 64, deep_gemm.get_num_sms()) + # Each MTP position has its own pool length; HOLD rows have no valid pools. + # Fixed-width mtp_step=2 gives Q=[3 * num_requests, 1, heads, dim], so next_n stays 1. + # All three verify positions are preserved despite the SM90 kernel's native next_n limit of 2. + logits = deep_gemm.fp8_paged_mqa_logits( + q_fp8.unsqueeze(1), pages, weights, lengths, block_table, metadata, max_pools, clean_logits=False + ) + groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) + self.select_topk_indices(logits, lengths.view(-1), groups) + return groups diff --git a/lightllm/models/glm5_next/layer_infer/__init__.py b/lightllm/models/glm5_next/layer_infer/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py b/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..1cca971ea6 --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 + +from lightllm.common.basemodel.triton_kernel.mhc import hc_expand +from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer + + +class Glm5NextPreLayerInfer(LlamaMultimodalPreLayerInfer): + """Initialize mHC residual streams after token embedding and TP reduction.""" + + def __init__(self, network_config): + super().__init__(network_config) + self.mhc_streams = network_config.get("hc_mult", 4) + + def context_forward(self, input_ids, infer_state, layer_weight): + input_embeddings = super().context_forward(input_ids, infer_state, layer_weight) + return hc_expand(input_embeddings, self.mhc_streams) + + def token_forward(self, input_ids, infer_state, layer_weight): + input_embeddings = super().token_forward(input_ids, infer_state, layer_weight) + return hc_expand(input_embeddings, self.mhc_streams) diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..620554bc06 --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel import TransformerLayerInferTpl +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( + silu_and_mul_fwd, +) +from lightllm.common.basemodel.triton_kernel.mhc import ( + hc_contract, + hc_post, + hc_pre_norm, +) +from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor + + +class Glm5NextTransformerLayerInfer(TransformerLayerInferTpl): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + self.eps_ = network_config["rms_norm_eps"] + self.embed_dim_ = network_config["hidden_size"] + self.num_hidden_layers = network_config["num_hidden_layers"] + self.autotune_layer_num = network_config.get("autotune_layer_num", self.num_hidden_layers) + self.is_linear_attention_layer = ( + layer_num < self.num_hidden_layers and network_config["layer_types"][layer_num] == "linear_attention" + ) + self.use_mhc = network_config.get("mhc", True) + self.mhc_streams = network_config.get("hc_mult", 4) + self.hc_eps = network_config.get("hc_eps", 1e-6) + self.hc_sinkhorn_iters = network_config.get("hc_sinkhorn_iters", 20) + self.swiglu_limit = network_config["swiglu_limit"] + self.enable_ep_moe = get_env_start_args().enable_ep_moe + self.is_moe = ( + network_config["n_routed_experts"] is not None + and layer_num >= network_config["first_k_dense_replace"] + and layer_num % network_config.get("moe_layer_freq", 1) == 0 + ) + self.n_shared_experts = network_config["n_shared_experts"] + self.num_experts_per_tok = network_config["num_experts_per_tok"] + self.norm_topk_prob = network_config["norm_topk_prob"] + self.n_group = network_config["n_group"] + self.topk_group = network_config["topk_group"] + linear = network_config["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.tp_linear_num_heads = self.linear_num_heads // self.tp_world_size_ + self.tp_linear_projection_size = self.tp_linear_num_heads * self.linear_head_dim + if not self.is_linear_attention_layer: + self.tp_q_head_num_ = network_config["num_attention_heads"] // self.tp_world_size_ + self.qk_nope_head_dim = network_config["qk_nope_head_dim"] + self.q_lora_rank = network_config["q_lora_rank"] + self.kv_lora_rank = network_config["kv_lora_rank"] + self.v_head_dim = network_config["v_head_dim"] + self.softmax_scale = self.qk_nope_head_dim ** -0.5 + self.indexer = Glm5NextNsaInfer( + layer_idx=self.layer_num_, + network_config=self.network_config_, + tp_world_size=self.tp_world_size_, + ) + + def _att_norm(self, input, infer_state, layer_weight): + return layer_weight.att_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor) + + def _ffn_norm(self, input, infer_state, layer_weight): + return layer_weight.ffn_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor) + + def _ffn(self, input, infer_state, layer_weight): + input = self._tpsp_allgather(input=input.view(-1, self.embed_dim_), infer_state=infer_state) + if self.is_moe: + output = self._moe_ffn(input, infer_state, layer_weight) + if self.enable_ep_moe: + # DeepEP combine already includes all routed experts. + return output + else: + output = self._ffn_tp(input, infer_state, layer_weight) + return self._tpsp_reduce(input=output, infer_state=infer_state) + + def _ffn_tp(self, input, infer_state, layer_weight): + """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" + + input = input.view(-1, self.embed_dim_) + up_gate_out = layer_weight.gate_up_proj.mm(input) + ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) + silu_and_mul_fwd( + up_gate_out, + ffn1_out, + limit=self.swiglu_limit, + alpha=1.0, + clamp_up_add_one=False, + ) + return layer_weight.down_proj.mm(ffn1_out) + + def _moe_ffn(self, input, infer_state, layer_weight) -> torch.Tensor: + hidden_states = input.view(-1, self.embed_dim_) + num_tokens, hidden_dim = hidden_states.shape + + # if fused_shared_experts is not enabled, compute shared_output + if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: + shared_output = self._ffn_tp(hidden_states, infer_state, layer_weight) + + moe_gate_dtype = layer_weight.moe_gate.data_type_ + router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) + output = layer_weight.experts.experts( + hidden_states, + router_logits=router_logits, + top_k=self.num_experts_per_tok, + renormalize=self.norm_topk_prob, + use_grouped_topk=self.n_group, + topk_group=self.topk_group, + num_expert_group=self.n_group, + is_prefill=infer_state.is_prefill, + infer_state=infer_state, + alpha=1.0, + limit=self.swiglu_limit, + clamp_up_add_one=False, + ) + + if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: + output.add_(shared_output) + + return output.view(num_tokens, hidden_dim) + + def _get_qkv(self, input, infer_state, layer_weight): + if self.is_linear_attention_layer: + raise AssertionError("KDA projections use _kda_projections") + + input = input.view(-1, self.embed_dim_) + input = self._tpsp_allgather(input=input, infer_state=infer_state) + if infer_state.need_dp_prefill_balance: + input = infer_state._all_to_all_unbalance_get(data=input) + + q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split([self.q_lora_rank, self.kv_lora_rank], dim=-1) + q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_) + infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q} + q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim) + cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank) + rmsnorm_forward( + cache_kv[:, :, : self.kv_lora_rank], + weight=layer_weight.kv_a_layernorm_.weight, + eps=self.eps_, + out=cache_kv[:, :, : self.kv_lora_rank], + ) + return q, cache_kv + + def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): + q = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1).contiguous() + 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=infer_state.prefill_att_state, + layer_weight=layer_weight, + ) + del infer_state.get_topk_indices_params + return infer_state.prefill_att_state.prefill_att( + q=q, + k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + v=None, + 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, + }, + ), + ) + + def _context_attention_wrapper_run(self, q, cache_kv, infer_state, layer_weight): + """Capture GLM sparse attention with the indexer's graph inputs. + + Sparse attention is replayed as a CPU-side step because the indexer + performs runtime request-state updates. Besides Q and KV, it needs + the hidden state and Q-LoRA projection saved by ``_get_qkv``. Python + assignments made during capture do not run during graph replay, so + preserve those tensors explicitly and restore them for every replay. + """ + if not torch.cuda.is_current_stream_capturing(): + return self._context_attention_kernel(q, cache_kv, infer_state, layer_weight) + + q = q.contiguous() + cache_kv = cache_kv.contiguous() + indexer_inputs = infer_state.get_topk_indices_params + hidden_states = indexer_inputs["hidden_states"].contiguous() + q_lora = indexer_inputs["q_lora"].contiguous() + _q = tensor_to_no_ref_tensor(q) + _cache_kv = tensor_to_no_ref_tensor(cache_kv) + _hidden_states = tensor_to_no_ref_tensor(hidden_states) + _q_lora = tensor_to_no_ref_tensor(q_lora) + pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() + pre_capture_graph.__exit__(None, None, None) + + def restore_indexer_inputs(state): + state.get_topk_indices_params = {"hidden_states": _hidden_states, "q_lora": _q_lora} + + def get_o_shape_dtype_device(): + restore_indexer_inputs(infer_state) + with torch.cuda.graph(cuda_graph=torch.cuda.CUDAGraph()): + output = self._context_attention_kernel(_q, _cache_kv, infer_state, layer_weight) + output_shape, output_dtype, output_device = output.shape, output.dtype, output.device + return output_shape, output_dtype, output_device + + output_shape, output_dtype, output_device = get_o_shape_dtype_device() + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + output = torch.empty(output_shape, dtype=output_dtype, device=output_device) + _output = tensor_to_no_ref_tensor(output) + + def sparse_att_func(new_infer_state): + restore_indexer_inputs(new_infer_state) + tmp_output = self._context_attention_kernel(_q, _cache_kv, new_infer_state, layer_weight) + assert tmp_output.shape == _output.shape + _output.copy_(tmp_output) + + infer_state.prefill_cuda_graph_add_cpu_runnning_func( + func=sparse_att_func, + after_graph=pre_capture_graph, + ) + return output + + def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): + if self.is_linear_attention_layer: + raise AssertionError("KDA uses its dedicated backend") + q_nope = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1) + topk_mem_indices, _ = self.indexer._get_indices( + hidden_states=infer_state.get_topk_indices_params["hidden_states"], + q_lora=infer_state.get_topk_indices_params["q_lora"], + infer_state=infer_state, + att_state=infer_state.decode_att_state, + layer_weight=layer_weight, + ) + del infer_state.get_topk_indices_params + q_rope = q_nope[..., :0] + return infer_state.decode_att_state.decode_att( + q=(q_nope, q_rope), + k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + v=None, + att_control=AttControl( + nsa_decode=True, + nsa_decode_dict={ + "layer_index": self.layer_num_, + "topk_mem_indices": topk_mem_indices, + "softmax_scale": self.softmax_scale, + "kv_lora_rank": self.kv_lora_rank, + "qk_rope_head_dim": 0, + }, + ), + ) + + def _get_o(self, input, infer_state, layer_weight): + if infer_state.need_dp_prefill_balance: + input = infer_state._all_to_all_balance_get(data=input) + # Both sparse prefill and decode return attention in MLA latent space. + input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) + output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) + return self._tpsp_reduce(input=output, infer_state=infer_state) + + def _kda_projections(self, hidden_states, infer_state, layer_weight): + # Gather sequence-sharded tokens for each rank's KDA heads; + # _kda_post reduces the output back to the sequence shard. + hidden_states = hidden_states.view(-1, self.embed_dim_) + hidden_states = self._tpsp_allgather(input=hidden_states, infer_state=infer_state) + qkv_gate_proj = layer_weight.linear_qkvbfg_a_proj.mm(hidden_states) + qkv_dim = 3 * self.tp_linear_projection_size + qkv, beta_logits, decay_gate_hidden, output_gate_hidden = qkv_gate_proj.split( + [ + qkv_dim, + self.tp_linear_num_heads, + self.linear_head_dim, + self.linear_head_dim, + ], + dim=-1, + ) + raw_decay_gate, raw_output_gate = layer_weight.project_kda_fg_b(decay_gate_hidden, output_gate_hidden) + return qkv, raw_decay_gate, beta_logits, raw_output_gate + + def _kda_post(self, core_output, raw_output_gate, infer_state, layer_weight): + tokens = raw_output_gate.shape[0] + core_output = core_output.view(-1, self.linear_head_dim) + raw_output_gate = raw_output_gate.view(tokens, self.tp_linear_num_heads, self.linear_head_dim) + output = layer_weight.linear_o_norm( + input=core_output, + gate_value=raw_output_gate, + eps=self.eps_, + alloc_func=self.alloc_tensor, + ) + output = layer_weight.linear_o_proj.mm(output.view(tokens, -1)) + return self._tpsp_reduce(input=output, infer_state=infer_state) + + def _kda_prefill_cuda_graph_wrapper( + self, + mixed_qkv: torch.Tensor, + raw_decay_gate: torch.Tensor, + beta_logits: torch.Tensor, + infer_state, + layer_weight, + ) -> torch.Tensor: + """Run KDA prefill between CUDA-graph segments. + + KDA updates request-owned convolution and SSM state, so its prefill + kernel must use the attention state constructed for the request being + replayed rather than the one used while capturing. Keep its inputs and + output at stable addresses, then register it as a CPU-side replay step + between the surrounding CUDA-graph segments. + """ + backend = infer_state.prefill_att_state1.backend + mixed_qkv = mixed_qkv.contiguous() + raw_decay_gate = raw_decay_gate.contiguous() + beta_logits = beta_logits.contiguous() + _mixed_qkv = tensor_to_no_ref_tensor(mixed_qkv) + _raw_decay_gate = tensor_to_no_ref_tensor(raw_decay_gate) + _beta_logits = tensor_to_no_ref_tensor(beta_logits) + + pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() + pre_capture_graph.__exit__(None, None, None) + + # chunk_kda_with_fused_gate returns [1, tokens, heads, head_dim]. + # Construct this shape directly instead of dry-running the kernel: + # kernel setup may synchronize with the host, which is illegal while a + # CUDA graph capture is active. + output_shape = (1, mixed_qkv.shape[0], backend.tp_num_heads, backend.head_dim) + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + output = torch.empty(output_shape, dtype=mixed_qkv.dtype, device=mixed_qkv.device) + _output = tensor_to_no_ref_tensor(output) + + def kda_prefill_func(new_infer_state): + tmp_output = new_infer_state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_prefill=True, + linear_att_prefill_dict={ + "mixed_qkv": _mixed_qkv, + "raw_gate": _raw_decay_gate, + "raw_beta": _beta_logits, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + assert tmp_output.shape == _output.shape + _output.copy_(tmp_output) + + infer_state.prefill_cuda_graph_add_cpu_runnning_func( + func=kda_prefill_func, + after_graph=pre_capture_graph, + ) + return output + + def context_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().context_attention_forward(input_embeddings, infer_state, layer_weight) + qkv, raw_decay_gate, beta_logits, raw_output_gate = self._kda_projections( + input_embeddings, infer_state, layer_weight + ) + if torch.cuda.is_current_stream_capturing(): + core_output = self._kda_prefill_cuda_graph_wrapper( + qkv, + raw_decay_gate, + beta_logits, + infer_state, + layer_weight, + ) + else: + core_output = infer_state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_prefill=True, + linear_att_prefill_dict={ + "mixed_qkv": qkv, + "raw_gate": raw_decay_gate, + "raw_beta": beta_logits, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) + + def token_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().token_attention_forward(input_embeddings, infer_state, layer_weight) + qkv, raw_decay_gate, beta_logits, raw_output_gate = self._kda_projections( + input_embeddings, infer_state, layer_weight + ) + core_output = infer_state.decode_att_state1.decode_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_decode=True, + linear_att_decode_dict={ + "mixed_qkv": qkv, + "raw_gate": raw_decay_gate, + "raw_beta": beta_logits, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) + + def _hc_pre(self, streams, layer_weight, prefix, norm_weight): + return hc_pre_norm( + x=streams, + fn=getattr(layer_weight, f"hc_{prefix}_fn").weight, + scale=getattr(layer_weight, f"hc_{prefix}_scale").weight, + base=getattr(layer_weight, f"hc_{prefix}_base").weight, + norm_weight=norm_weight.weight, + streams=self.mhc_streams, + rms_eps=self.eps_, + norm_eps=self.eps_, + hc_eps=self.hc_eps, + sinkhorn_iters=self.hc_sinkhorn_iters, + ) + + def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): + streams = input_embeddings + + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) + if prefill: + layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) + else: + layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) + layer_output = self._ffn(layer_input, infer_state, layer_weight) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + # Only prefill autotuning truncates the model. Decode autotuning still + # executes every layer and must keep all residual streams until the end. + is_autotune_last_layer = ( + prefill and Autotuner.is_autotune_warmup() and self.layer_num_ == self.autotune_layer_num - 1 + ) + if self.layer_num_ == self.num_hidden_layers - 1 or is_autotune_last_layer: + return hc_contract(streams, self.mhc_streams) + return streams + + def context_forward(self, input_embeddings, infer_state, layer_weight): + if not self.use_mhc: + return super().context_forward(input_embeddings, infer_state, layer_weight) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=True) + + def token_forward(self, input_embeddings, infer_state, layer_weight): + if not self.use_mhc: + return super().token_forward(input_embeddings, infer_state, layer_weight) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=False) diff --git a/lightllm/models/glm5_next/layer_weights/__init__.py b/lightllm/models/glm5_next/layer_weights/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..b336f00632 --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 + +from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import ( + LlamaPreAndPostLayerWeight, +) + + +def add_language_model_aliases(weights: dict) -> None: + """Expose GLM's nested language-model keys under LightLLM names.""" + + prefix = "model.language_model." + for name in list(weights): + if name.startswith(prefix): + weights.setdefault("model." + name[len(prefix) :], weights[name]) + + +class Glm5NextPreAndPostLayerWeight(LlamaPreAndPostLayerWeight): + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..32191809fe --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -0,0 +1,449 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.layer_weights.transformer_layer_weight import ( + TransformerLayerWeight, +) +from lightllm.common.basemodel.layer_weights.meta_weights import ( + COLMMWeight, + FusedMoeWeight, + GatedRMSNormWeight, + LayerNormWeight, + ParameterWeight, + RMSNormWeight, + ROWBMMWeight, + ROWMMWeight, + TpParameterWeight, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import ( + get_row_slice_mixin, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import ( + MMWeightTpl, +) +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.utils.envs_utils import get_env_start_args +from .pre_and_post_layer_weight import add_language_model_aliases + + +class Glm5NextMergedKdaProjection(MMWeightTpl): + """One KDA input GEMM with TP-sharded q/k/v/b and replicated f_a/g_a.""" + + def __init__( + self, + in_dim: int, + projection: int, + head_count: int, + head_dim: int, + weight_names: list[str], + data_type: torch.dtype, + tp_rank: int | None = None, + tp_world_size: int | None = None, + ): + tp_rank = get_current_rank_in_dp() if tp_rank is None else tp_rank + tp_world_size = get_dp_world_size() if tp_world_size is None else tp_world_size + assert projection % tp_world_size == 0 + assert head_count % tp_world_size == 0 + super().__init__( + in_dim=in_dim, + out_dims=[ + projection // tp_world_size, + projection // tp_world_size, + projection // tp_world_size, + head_count // tp_world_size, + head_dim, + head_dim, + ], + weight_names=weight_names, + bias_names=None, + data_type=data_type, + quant_method=None, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + self.sharded_slicer = get_row_slice_mixin("none", tp_rank=tp_rank, tp_world_size=tp_world_size) + self.replicated_slicer = get_row_slice_mixin("none", tp_rank=0, tp_world_size=1) + + def _get_param_slicer(self, sub_child_index: int): + return self.replicated_slicer if sub_child_index >= 4 else self.sharded_slicer + + +class Glm5NextTransformerLayerWeight(TransformerLayerWeight): + def _parse_config(self): + self.n_embed = self.network_config_["hidden_size"] + self.n_inter = self.network_config_["intermediate_size"] + self.moe_inter = self.network_config_.get("moe_intermediate_size", self.n_inter) + self.n_routed_experts = self.network_config_["n_routed_experts"] + self.is_moe = ( + self.n_routed_experts is not None + and self.layer_num_ >= self.network_config_["first_k_dense_replace"] + and self.layer_num_ % self.network_config_.get("moe_layer_freq", 1) == 0 + ) + self.num_fused_shared_experts = 0 + args = get_env_start_args() + if self.is_moe and args.enable_fused_shared_experts and not args.enable_ep_moe: + self.num_fused_shared_experts = self.network_config_.get("n_shared_experts", 0) + self.is_linear_attention_layer = ( + self.layer_num_ < self.network_config_["num_hidden_layers"] + and self.network_config_["layer_types"][self.layer_num_] == "linear_attention" + ) + linear = self.network_config_["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.linear_projection_size = self.linear_num_heads * self.linear_head_dim + self.linear_conv_kernel_size = linear["short_conv_kernel_size"] + self.mhc_streams = self.network_config_.get("hc_mult", 4) + if not self.is_linear_attention_layer: + self.num_attention_heads = self.network_config_["num_attention_heads"] + self.q_lora_rank = self.network_config_["q_lora_rank"] + self.kv_lora_rank = self.network_config_["kv_lora_rank"] + self.qk_nope_head_dim = self.network_config_["qk_nope_head_dim"] + self.v_head_dim = self.network_config_["v_head_dim"] + self.index_n_heads = self.network_config_["index_n_heads"] + self.index_head_dim = self.network_config_["index_head_dim"] + + def _init_weight(self): + if self.is_linear_attention_layer: + self._init_kda() + else: + self._init_mla() + self._init_indexer_weight() + + if self.is_moe: + self._init_moe() + else: + self._init_mlp(f"model.layers.{self.layer_num_}.mlp", self.n_inter) + self._init_glm_norms() + if self.network_config_.get("mhc", True): + self._init_mhc() + + def _init_mla(self): + prefix = f"model.layers.{self.layer_num_}.self_attn" + self.qkv_a_proj_with_mqa_ = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.q_lora_rank, self.kv_lora_rank], + weight_names=[f"{prefix}.q_a_proj.weight", f"{prefix}.kv_a_proj_with_mqa.weight"], + data_type=self.data_type_, + quant_method=self.get_quant_method("qkv_a_proj_with_mqa"), + tp_rank=0, + tp_world_size=1, + ) + self.q_b_proj_ = ROWMMWeight( + in_dim=self.q_lora_rank, + out_dims=[self.num_attention_heads * self.qk_nope_head_dim], + weight_names=f"{prefix}.q_b_proj.weight", + data_type=self.data_type_, + quant_method=self.get_quant_method("q_b_proj"), + ) + # The checkpoint keeps kv_b_proj in BF16 while the surrounding + # projections use FP8. Split it into the two unquantized BMMs. + self.k_b_proj_ = ROWBMMWeight( + dim0=self.num_attention_heads, + dim1=self.qk_nope_head_dim, + dim2=self.kv_lora_rank, + weight_names=f"{prefix}.k_b_proj.weight", + data_type=self.data_type_, + quant_method=None, + ) + self.v_b_proj_ = ROWBMMWeight( + dim0=self.num_attention_heads, + dim1=self.kv_lora_rank, + dim2=self.v_head_dim, + weight_names=f"{prefix}.v_b_proj.weight", + data_type=self.data_type_, + quant_method=None, + ) + self.o_weight_ = COLMMWeight( + in_dim=self.num_attention_heads * self.v_head_dim, + out_dims=[self.n_embed], + weight_names=f"{prefix}.o_proj.weight", + data_type=self.data_type_, + quant_method=self.get_quant_method("o_weight"), + ) + + def _init_mlp(self, prefix, intermediate_size): + # EP returns complete routed outputs on each rank, so shared experts + # must also be complete. Dense layers retain their TP weight shards. + tp_kwargs = {"tp_rank": 0, "tp_world_size": 1} if self.is_moe and get_env_start_args().enable_ep_moe else {} + self.gate_up_proj = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[intermediate_size, intermediate_size], + weight_names=[f"{prefix}.gate_proj.weight", f"{prefix}.up_proj.weight"], + data_type=self.data_type_, + quant_method=self.get_quant_method("gate_up_proj"), + **tp_kwargs, + ) + self.down_proj = COLMMWeight( + in_dim=intermediate_size, + out_dims=[self.n_embed], + weight_names=f"{prefix}.down_proj.weight", + data_type=self.data_type_, + quant_method=self.get_quant_method("down_proj"), + **tp_kwargs, + ) + + def _init_moe(self): + prefix = f"model.layers.{self.layer_num_}.mlp" + self.moe_gate = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.n_routed_experts], + weight_names=f"{prefix}.gate.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + if self.num_fused_shared_experts == 0: + self._init_mlp(f"{prefix}.shared_experts", self.moe_inter) + self.experts = FusedMoeWeight( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + e_score_correction_bias_name=f"{prefix}.gate.e_score_correction_bias", + weight_prefix=f"{prefix}.experts", + n_routed_experts=self.n_routed_experts, + hidden_size=self.n_embed, + moe_intermediate_size=self.moe_inter, + data_type=self.data_type_, + quant_method=self.get_quant_method("fused_moe"), + num_fused_shared_experts=self.num_fused_shared_experts, + layer_num=self.layer_num_, + network_config=self.network_config_, + ) + + def _init_kda(self): + prefix = f"model.layers.{self.layer_num_}.self_attn" + projection = self.linear_projection_size + head_count = self.linear_num_heads + head_dim = self.linear_head_dim + + self.linear_qkvbfg_a_proj = Glm5NextMergedKdaProjection( + in_dim=self.n_embed, + projection=projection, + head_count=head_count, + head_dim=head_dim, + weight_names=[ + f"{prefix}.q_proj.weight", + f"{prefix}.k_proj.weight", + f"{prefix}.v_proj.weight", + f"{prefix}.b_proj.weight", + f"{prefix}.f_a_proj.weight", + f"{prefix}.g_a_proj.weight", + ], + data_type=self.data_type_, + ) + self.linear_fg_b_proj = ROWMMWeight( + in_dim=head_dim, + out_dims=[projection, projection], + weight_names=[f"{prefix}.f_b_proj.weight", f"{prefix}.g_b_proj.weight"], + data_type=self.data_type_, + quant_method=None, + ) + self.linear_qkv_conv1d = ROWMMWeight( + in_dim=self.linear_conv_kernel_size, + out_dims=[projection, projection, projection], + weight_names=[ + f"{prefix}.q_conv1d.weight", + f"{prefix}.k_conv1d.weight", + f"{prefix}.v_conv1d.weight", + ], + data_type=self.data_type_, + quant_method=None, + ) + self.linear_A_log = TpParameterWeight( + weight_name=f"{prefix}.A_log", + data_type=torch.float32, + weight_shape=(head_count,), + ) + self.linear_dt_bias = TpParameterWeight( + weight_name=f"{prefix}.dt_bias", + data_type=torch.float32, + weight_shape=(projection,), + ) + self.linear_o_norm = GatedRMSNormWeight( + dim=head_dim, + weight_name=f"{prefix}.o_norm.weight", + data_type=self.data_type_, + gate_type="sigmoid", + ) + self.linear_o_proj = COLMMWeight( + in_dim=projection, + out_dims=[self.n_embed], + weight_names=f"{prefix}.o_proj.weight", + data_type=self.data_type_, + quant_method=None, + ) + + def _init_indexer_weight(self): + """Initialize GLM's NoPE, K-pool indexer parameters. + + The head-weight projection intentionally accumulates in fp32. Both + reference engines do this because bf16 head weights can change close + K-pool rankings on difficult long-context prompts. + """ + + prefix = f"model.layers.{self.layer_num_}.self_attn.indexer" + self.wq_b_proj_ = ROWMMWeight( + in_dim=self.q_lora_rank, + out_dims=[self.index_n_heads * self.index_head_dim], + weight_names=f"{prefix}.wq_b.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.wk_proj_ = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.wk.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.k_norm_ = LayerNormWeight( + dim=self.index_head_dim, + weight_name=f"{prefix}.k_norm.weight", + data_type=self.data_type_, + bias_name=f"{prefix}.k_norm.bias", + ) + self.weights_proj_ = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.index_n_heads], + weight_names=f"{prefix}.weights_proj.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_gate = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.index_kpool_compress_gate", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_ape = ParameterWeight( + weight_name=f"{prefix}.index_kpool_compress_ape", + data_type=torch.float32, + weight_shape=(self.network_config_["index_kpool"], self.index_head_dim), + ) + + def _init_glm_norms(self): + prefix = f"model.layers.{self.layer_num_}" + self.att_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.input_layernorm.weight", + data_type=self.data_type_, + ) + self.ffn_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.post_attention_layernorm.weight", + data_type=self.data_type_, + ) + if not self.is_linear_attention_layer: + self.kv_a_layernorm_ = RMSNormWeight( + dim=self.kv_lora_rank, + weight_name=f"{prefix}.self_attn.kv_a_layernorm.weight", + data_type=self.data_type_, + ) + self.q_a_layernorm_ = RMSNormWeight( + dim=self.q_lora_rank, + weight_name=f"{prefix}.self_attn.q_a_layernorm.weight", + data_type=self.data_type_, + ) + + def _init_mhc(self): + prefix = f"model.layers.{self.layer_num_}" + streams = self.mhc_streams + mix_size = (2 + streams) * streams + flattened_hidden = streams * self.n_embed + self.hc_attn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_attn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_attn_base = ParameterWeight( + weight_name=f"{prefix}.hc_attn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_attn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_attn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + self.hc_ffn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_ffn_base = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_ffn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + + def get_merged_kda_conv_weight(self): + return self.linear_qkv_conv1d.mm_param.weight + + def project_kda_fg_b(self, decay_gate_hidden: torch.Tensor, output_gate_hidden: torch.Tensor): + quant_method = self.linear_fg_b_proj.quant_method + raw_decay_gate = quant_method.apply(decay_gate_hidden, self.linear_fg_b_proj.mm_param_list[0]) + raw_output_gate = quant_method.apply(output_gate_hidden, self.linear_fg_b_proj.mm_param_list[1]) + return raw_decay_gate, raw_output_gate + + def _preprocess_kda_weights(self, weights): + prefix = f"model.layers.{self.layer_num_}.self_attn" + for projection in ("q", "k", "v"): + name = f"{prefix}.{projection}_conv1d.weight" + if name in weights and weights[name].ndim == 3: + weights[name] = weights[name].squeeze(1) + + def _split_kv_b_proj(self, weight): + weight = weight.view(self.num_attention_heads, self.qk_nope_head_dim + self.v_head_dim, self.kv_lora_rank) + k_weight, v_weight = weight.split([self.qk_nope_head_dim, self.v_head_dim], dim=1) + return k_weight.contiguous(), v_weight.transpose(1, 2).contiguous() + + def _rename_shared_experts(self, weights): + prefix = f"model.layers.{self.layer_num_}.mlp" + suffixes = ["weight"] + if self.quant_cfg.quantized_weight: + scale_suffix = self.experts.quant_method.weight_scale_suffix + assert scale_suffix is not None + suffixes.append(scale_suffix) + for index in range(self.num_fused_shared_experts): + expert_id = self.n_routed_experts + index + for projection in ("gate_proj", "down_proj", "up_proj"): + for suffix in suffixes: + source = f"{prefix}.shared_experts.{projection}.{suffix}" + if source in weights: + weights[f"{prefix}.experts.{expert_id}.{projection}.{suffix}"] = weights[source] + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + + # Fused shared experts use the same tensor layout as routed experts. + if self.num_fused_shared_experts > 0: + self._rename_shared_experts(weights) + + if self.is_linear_attention_layer: + self._preprocess_kda_weights(weights) + else: + kv_b_name = f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight" + if kv_b_name in weights: + k_b_proj, v_b_proj = self._split_kv_b_proj(weights[kv_b_name]) + weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = k_b_proj + weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = v_b_proj + + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py new file mode 100644 index 0000000000..cc8e86613e --- /dev/null +++ b/lightllm/models/glm5_next/model.py @@ -0,0 +1,98 @@ +import json +import os + +import torch +import triton + +from lightllm.common.build_utils import repair_config +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend +from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend +from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager +from lightllm.common.req_manager import Glm5NextReqManager +from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer +from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer +from .layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer +from .layer_weights.pre_and_post_layer_weight import Glm5NextPreAndPostLayerWeight +from .layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight + + +class Glm5NextTpPartModel(TpPartBaseModel): + pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight + transformer_weight_class = Glm5NextTransformerLayerWeight + pre_layer_infer_class = Glm5NextPreLayerInfer + post_layer_infer_class = LlamaPostLayerInfer + transformer_layer_infer_class = Glm5NextTransformerLayerInfer + + def _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json")) as f: + outer_config = json.load(f) + self.config = dict(outer_config.get("text_config", outer_config)) + if "quantization_config" in outer_config: + self.config["quantization_config"] = dict(outer_config["quantization_config"]) + self.config["autotune_layer_num"] = 4 + repair_config(self.config, same_names=["num_attention_heads", "n_head"]) + repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + + def _verify_params(self): + super()._verify_params() + assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" + assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" + + def autotune_layers(self): + return 4 + + def _init_some_value(self): + self.layers_num = self.config["n_layer"] + self.vocab_size = self.config["vocab_size"] + # MLA stores one replicated latent KV vector per token on every TP rank. + self.tp_k_head_num_ = 1 + self.tp_v_head_num_ = 0 + self.qk_nope_head_dim = self.config["qk_nope_head_dim"] + self.qk_rope_head_dim = self.config["qk_rope_head_dim"] + self.q_lora_rank = self.config["q_lora_rank"] + self.kv_lora_rank = self.config["kv_lora_rank"] + self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) + self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim + + def _init_req_manager(self): + self.linear_config = Glm5NextCacheConfig.from_model_config(self.config, self.args) + self.req_manager = Glm5NextReqManager( + self.max_req_num, + max(self.batch_max_tokens or 0, self.max_seq_length or 0), + None, + linear_config=self.linear_config, + ) + + def _init_mem_manager(self): + self.mem_manager = Glm5NextMemManager( + size=self.max_total_token_num, + dtype=self.data_type, + num_kv_heads=1, + head_dim=self.linear_config.full_att_head_dim, + full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), + linear_config=self.linear_config, + mem_fraction=self.mem_fraction, + ) + + def _init_att_backend(self): + self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) + self.decode_att_backend = self.prefill_att_backend + + def _init_att_backend1(self): + self.prefill_att_backend1 = KDALinearAttBackend(model=self) + self.decode_att_backend1 = self.prefill_att_backend1 + + def _init_custom(self): + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + if self.args.enable_ep_moe: + dist_group_manager.new_deepep_group( + n_routed_experts=self.config["n_routed_experts"], + hidden_size=self.config["hidden_size"], + expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), + num_experts_per_tok=self.config["num_experts_per_tok"], + moe_intermediate_size=self.config["moe_intermediate_size"], + ) diff --git a/lightllm/models/glm5_next/tokenizer.py b/lightllm/models/glm5_next/tokenizer.py new file mode 100644 index 0000000000..be416e8b52 --- /dev/null +++ b/lightllm/models/glm5_next/tokenizer.py @@ -0,0 +1,19 @@ +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from .vision_process import Glm5NextImageProcessor + + +class Glm5NextTokenizer(QWen2VLTokenizer): + def __init__(self, tokenizer, model_cfg, weight_dir): + BaseMultiModalTokenizer.__init__(self, tokenizer) + self.image_processor = Glm5NextImageProcessor.from_pretrained(weight_dir) + self.image_start_id = model_cfg["image_start_token_id"] + self.image_end_id = model_cfg["image_end_token_id"] + self.image_token_id = model_cfg["image_token_id"] + + def get_image_token_length(self, img): + height, width = self.image_processor.get_image_size(img.image_h, img.image_w) + factor = self.image_processor.patch_size * self.image_processor.merge_size + grid_h, grid_w = height // factor, width // factor + img.grid_thwd = (1, grid_h, grid_w, 0) + return grid_h * grid_w diff --git a/lightllm/models/glm5_next/triton_kernel/__init__.py b/lightllm/models/glm5_next/triton_kernel/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/triton_kernel/index_quant.py b/lightllm/models/glm5_next/triton_kernel/index_quant.py new file mode 100644 index 0000000000..67485d9b52 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/index_quant.py @@ -0,0 +1,69 @@ +import torch +import triton +import triton.language as tl +from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage + + +@triton.jit +def _hadamard_transform_quant_fp8_kernel( + X, + Y, + S, + n_rows, + scale: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + row_mask = rows < n_rows + cols = tl.arange(0, BLOCK_N) + offsets = rows[:, None] * BLOCK_N + cols[None, :] + x = tl.load(X + offsets, mask=row_mask[:, None], other=0.0).to(tl.float32) + + x = _butterfly_stage(x, 64, 1, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 32, 2, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 16, 4, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 8, 8, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 4, 16, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 2, 32, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 1, 64, BLOCK_R, BLOCK_N) + + # Match the unfused path's bf16 Hadamard output before FP8 quantization. + x = (x * scale).to(tl.bfloat16).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) + quant_scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0)))) + y = tl.minimum(tl.maximum(x / quant_scale[:, None], -448.0), 448.0) + + tl.store(Y + offsets, y, mask=row_mask[:, None]) + tl.store(S + rows, quant_scale, mask=row_mask) + + +def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse Hadamard-128 with the following ue8m0 FP8 quantization.""" + + assert x.is_cuda, "hadamard_transform_quant_fp8 only supports CUDA tensors" + assert x.dtype == torch.bfloat16, "Hadamard transform expects bfloat16 input" + assert x.size(-1) == 128, "Hadamard transform expects hidden size 128" + if not x.is_contiguous(): + x = x.contiguous() + + original_shape = x.shape + rows = x.numel() // 128 + output = torch.empty_like(x, dtype=torch.float8_e4m3fn) + output_scale = torch.empty((*original_shape[:-1], 1), dtype=torch.float32, device=x.device) + if rows == 0: + return output, output_scale + + block_r = 32 + _hadamard_transform_quant_fp8_kernel[(triton.cdiv(rows, block_r),)]( + x, + output, + output_scale, + rows, + scale, + BLOCK_R=block_r, + BLOCK_N=128, + num_warps=2, + ) + return output.view(original_shape), output_scale diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py new file mode 100644 index 0000000000..0e126aea6a --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -0,0 +1,359 @@ +import torch +import triton +import triton.language as tl + +from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage +from lightllm.utils.device_utils import get_device_sm_count + + +@triton.jit +def _get_query_block(CuQLens, BATCH: tl.constexpr, BLOCK: tl.constexpr): + block = tl.program_id(0) + # Request b owns [cu[b] // BLOCK + b, cu[b + 1] // BLOCK + b + 1). + # The extra block covers unaligned boundaries without padding to max_q_len. + if BATCH == 1: + batch = 0 + else: + batches = tl.arange(0, triton.next_power_of_2(BATCH)) + q_ends = tl.load(CuQLens + batches + 1, batches < BATCH, 0) + block_ends = q_ends // BLOCK + batches + 1 + batch = tl.sum(((block >= block_ends) & (batches < BATCH)).to(tl.int32), 0) + q_start = tl.load(CuQLens + batch) + q_end = tl.load(CuQLens + batch + 1) + local_block = block - (q_start // BLOCK + batch) + return batch, q_start, q_end, local_block + + +@triton.jit +def _compress_pools( + Raw, + Tail, + Packed, + Ape, + Lengths, + Starts, + Ragged, + ReqIdx, + CuQLens, + MtpIndex, + RAW_STRIDE: tl.constexpr, + TAIL_REQ_STRIDE: tl.constexpr, + TAIL_SLOT_STRIDE: tl.constexpr, + PACKED_STRIDE: tl.constexpr, + HOLD_REQ: tl.constexpr, + BATCH: tl.constexpr, + SINGLE_QUERY: tl.constexpr, + TAIL_SIZE: tl.constexpr, + HAS_MTP_INDEX: tl.constexpr, +): + if SINGLE_QUERY: + batch = tl.program_id(0) + row = batch + length = tl.load(Lengths + row) + valid = length > 0 and length % 4 == 0 + q_start = row + if HAS_MTP_INDEX: + q_start -= tl.load(MtpIndex + row) + else: + batch, q_start, q_end, pool_index = _get_query_block(CuQLens, BATCH, 4) + prefix = tl.load(Lengths + q_start, q_start < q_end, 1) - 1 + # Only process pool-closing rows, including a pool spanning the old tail. + length = (prefix // 4 + pool_index + 1) * 4 + row = q_start + length - prefix - 1 + valid = row < q_end + req = tl.load(ReqIdx + batch) + if valid and req != HOLD_REQ: + start = tl.load(Starts + row) + pool = tl.arange(0, 4) + cols = tl.arange(0, 128) + raw_rows = row - 3 + pool + from_chunk = raw_rows >= q_start + chunk_ptr = Raw + raw_rows[:, None] * RAW_STRIDE + cols[None, :] + tail_slots = (length - 4 + pool) % TAIL_SIZE + tail_ptr = Tail + req * TAIL_REQ_STRIDE + tail_slots[:, None] * TAIL_SLOT_STRIDE + cols[None, :] + raw = tl.where( + from_chunk[:, None], + tl.load(chunk_ptr, from_chunk[:, None], 0), + tl.load(tail_ptr, ~from_chunk[:, None], 0), + ).to(tl.float32) + score = tl.where( + from_chunk[:, None], + tl.load(chunk_ptr + 128, from_chunk[:, None], 0), + tl.load(tail_ptr + 128, ~from_chunk[:, None], 0), + ).to(tl.float32) + score += tl.load(Ape + pool[:, None] * 128 + cols[None, :]) + score = tl.exp(score - tl.max(score, 0)[None, :]) + weights = score / tl.sum(score, 0)[None, :] + key = tl.sum(raw * weights, 0).to(tl.bfloat16).to(tl.float32).reshape(1, 128) + for step in tl.static_range(7): + key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128) + key = (key * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32) + scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0))) + key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) + loc = tl.load(Ragged + start + length - 1).to(tl.int64) + dest = Packed + loc * PACKED_STRIDE + tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) + tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) + + +@triton.jit +def _save_pool_tails( + Raw, + Tail, + ReqIdx, + CuQLens, + SeqLens, + RAW_STRIDE: tl.constexpr, + TAIL_REQ_STRIDE: tl.constexpr, + TAIL_SLOT_STRIDE: tl.constexpr, + HOLD_REQ: tl.constexpr, + TAIL_SIZE: tl.constexpr, +): + batch = tl.program_id(0) + req = tl.load(ReqIdx + batch) + if req != HOLD_REQ: + length = tl.load(SeqLens + batch) + q_start = tl.load(CuQLens + batch) + q_end = tl.load(CuQLens + batch + 1) + slots = tl.arange(0, triton.next_power_of_2(TAIL_SIZE)) + cols = tl.arange(0, 256) + positions = length - TAIL_SIZE + slots + raw_rows = q_end - length + positions + # Absolute positions address a short ring, preserving the raw history + # needed after rejecting candidates or revisiting draft positions. + mask = ((slots < TAIL_SIZE) & (positions >= 0) & (raw_rows >= q_start))[:, None] + raw = tl.load(Raw + raw_rows[:, None] * RAW_STRIDE + cols[None, :], mask, 0) + tl.store( + Tail + req * TAIL_REQ_STRIDE + (positions % TAIL_SIZE)[:, None] * TAIL_SLOT_STRIDE + cols[None, :], + raw, + mask, + ) + + +def compress_pools( + raw, tail, packed_buffer, ape, lengths, starts, ragged, req_idx, cu_q_lens, seq_lens, max_q_len, mtp_index=None +): + batch = req_idx.numel() + single_query = max_q_len == 1 and lengths.numel() == batch + blocks = batch if single_query else lengths.numel() // 4 + batch + _compress_pools[(blocks,)]( + raw, + tail, + packed_buffer, + ape, + lengths, + starts, + ragged, + req_idx, + cu_q_lens, + mtp_index, + raw.stride(0), + tail.stride(0), + tail.stride(1), + packed_buffer.stride(0), + tail.shape[0] - 1, + batch, + single_query, + tail.shape[1], + mtp_index is not None, + num_warps=4, + ) + # Complete all boundary pools before replacing the previous chunk's tail. + _save_pool_tails[(req_idx.numel(),)]( + raw, + tail, + req_idx, + cu_q_lens, + seq_lens, + raw.stride(0), + tail.stride(0), + tail.stride(1), + tail.shape[0] - 1, + tail.shape[1], + num_warps=4, + ) + + +@triton.jit +def _get_pool_ranges( + Lengths, + CuQLens, + Starts, + Ends, + PoolLengths, + BATCH: tl.constexpr, + TOKENS: tl.constexpr, + MAX_POOLS: tl.constexpr, + SINGLE_QUERY: tl.constexpr, + BLOCK: tl.constexpr, +): + if SINGLE_QUERY: + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + batch = row + valid = row < TOKENS + else: + batch, q_start, q_end, local_block = _get_query_block(CuQLens, BATCH, BLOCK) + row = q_start + local_block * BLOCK + tl.arange(0, BLOCK) + valid = row < q_end + length = tl.load(Lengths + row, valid, 0) // 4 + start = batch * MAX_POOLS + tl.store(Starts + row, start, valid) + tl.store(Ends + row, start + length, valid) + tl.store(PoolLengths + row, length, valid) + + +def get_pool_ranges(lengths, cu_q_lens, max_q_len, max_pools): + """Build per-query pool bounds directly from the packed request boundaries.""" + starts, ends, pool_lengths = [torch.empty_like(lengths) for _ in range(3)] + batch, tokens = cu_q_lens.numel() - 1, lengths.numel() + single_query = max_q_len == 1 and tokens == batch + block = 256 + blocks = triton.cdiv(tokens, block) if single_query else tokens // block + batch + _get_pool_ranges[(blocks,)]( + lengths, cu_q_lens, starts, ends, pool_lengths, batch, tokens, max_pools, single_query, block, num_warps=4 + ) + return starts, ends, pool_lengths + + +@triton.jit +def _gather_pools( + Packed, + ReqTable, + ReqIdx, + SeqLen, + K, + Scale, + PACKED_STRIDE: tl.constexpr, + REQ_STRIDE: tl.constexpr, + POOLS: tl.constexpr, +): + pool, batch = tl.program_id(0), tl.program_id(1) + length = tl.load(SeqLen + batch) + req = tl.load(ReqIdx + batch) + valid = pool < length // 4 + loc = tl.load(ReqTable + req * REQ_STRIDE + pool * 4 + 3, valid, 0).to(tl.int64) + cols = tl.arange(0, 128) + packed = tl.load(Packed + loc * PACKED_STRIDE + cols, valid, 0) + scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) + row = batch.to(tl.int64) * POOLS + pool + tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) + tl.store(Scale + row, scale) + + +def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): + keys = torch.empty((req_idx.numel() * max_pools, 128), device=packed_buffer.device, dtype=torch.float8_e4m3fn) + scales = torch.empty((keys.shape[0],), device=keys.device, dtype=torch.float32) + # Long contexts can exceed grid Y's 65535-block limit; put pools on X. + _gather_pools[(max_pools, req_idx.numel())]( + packed_buffer, + req_table, + req_idx, + seq_len, + keys, + scales, + packed_buffer.stride(0), + req_table.stride(0), + max_pools, + num_warps=4, + ) + return keys, scales + + +@triton.jit +def _gather_paged_pools( + Packed, + ReqTable, + ReqIdx, + PoolLengths, + Pages, + PACKED_STRIDE: tl.constexpr, + REQ_STRIDE: tl.constexpr, + MAX_PAGES: tl.constexpr, +): + batch = tl.program_id(1) + req = tl.load(ReqIdx + batch) + length = tl.load(PoolLengths + batch) + rows = tl.arange(0, 64) + cols = tl.arange(0, 128) + # A fixed grid is replayable at 1M; the GPU length bounds the work. + for page in range(tl.program_id(0), tl.cdiv(length, 64), tl.num_programs(0)): + pools = page * 64 + rows + valid = pools < length + locs = tl.load(ReqTable + req * REQ_STRIDE + pools * 4 + 3, valid, 0).to(tl.int64) + keys = tl.load(Packed + locs[:, None] * PACKED_STRIDE + cols[None, :], valid[:, None], 0) + scales = tl.load((Packed + locs * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) + dest = Pages + (batch.to(tl.int64) * MAX_PAGES + page) * (64 * 132) + # DeepGEMM stores 64 FP8 keys followed by their 64 FP32 scales. + tl.store(dest + rows[:, None] * 128 + cols[None, :], keys) + tl.store((dest + 64 * 128).to(tl.pointer_type(tl.float32)) + rows, scales) + + +def gather_paged_pools(packed_buffer, req_table, req_idx, pool_lengths, max_pools): + """Pack valid pools into DeepGEMM pages; unused pages remain unread.""" + batch = req_idx.numel() + max_pages = triton.cdiv(max_pools, 64) + pages = torch.empty((batch * max_pages, 64, 1, 132), device=packed_buffer.device, dtype=torch.uint8) + block_table = torch.arange(batch * max_pages, device=pages.device, dtype=torch.int32).view(batch, max_pages) + blocks = min(max_pages, triton.cdiv(get_device_sm_count() * 4, batch)) + _gather_paged_pools[(blocks, batch)]( + packed_buffer, + req_table, + req_idx, + pool_lengths, + pages, + packed_buffer.stride(0), + req_table.stride(0), + max_pages, + num_warps=4, + ) + return pages, block_table + + +@triton.jit +def _expand_topk( + Groups, + Lengths, + Starts, + Ragged, + Out, + Relative, + TOPK: tl.constexpr, + WIDTH: tl.constexpr, + DENSE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + lane = tl.arange(0, BLOCK) + length = tl.load(Lengths + row) + start = tl.load(Starts + row) + if DENSE: + token = lane + valid = lane < length + else: + closed_tokens = tl.minimum(length // 4 * 4, TOPK) + group = tl.load(Groups + row * (TOPK // 4) + lane // 4, lane < closed_tokens, -1) + token = tl.where(lane < closed_tokens, group * 4 + lane % 4, length // 4 * 4 + lane - closed_tokens) + valid = (lane < closed_tokens + length % 4) & (token >= 0) + mem = tl.load(Ragged + start + token, valid & (lane < WIDTH), -1) + tl.store(Out + row * WIDTH + lane, mem, lane < WIDTH) + tl.store(Relative + row * WIDTH + lane, tl.where(valid, token, -1), lane < WIDTH) + + +def expand_topk(groups, lengths, starts, ragged, topk, dense=False): + width = triton.cdiv(topk + 3, 128) * 128 + out = torch.empty((lengths.numel(), width), dtype=torch.int32, device=lengths.device) + relative = torch.empty_like(out) + _expand_topk[(lengths.numel(),)]( + groups, + lengths, + starts, + ragged, + out, + relative, + topk, + width, + dense, + triton.next_power_of_2(width), + num_warps=4, + ) + return out, relative diff --git a/lightllm/models/glm5_next/vision_process.py b/lightllm/models/glm5_next/vision_process.py new file mode 100644 index 0000000000..1fd6843029 --- /dev/null +++ b/lightllm/models/glm5_next/vision_process.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 the HuggingFace Team. All rights reserved. +# Adapted from Hugging Face Transformers' GLM-5-Next image processor. + +import json +import math +import os + +import numpy as np +import torch +from torchvision.transforms.v2 import functional as F + +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor + + +def smart_resize(height, width, factor=28, min_image_tokens=16, max_image_tokens=8000): + """Choose an aligned canvas; padding, rather than stretching, preserves the image.""" + min_pixels, max_pixels = min_image_tokens * factor ** 2, max_image_tokens * factor ** 2 + + def align(value): + return math.ceil(value / factor) * factor + + target_h, target_w = align(height), align(width) + if target_h * target_w < min_pixels: + scale = math.sqrt(min_pixels / (height * width)) + target_h, target_w = align(max(1, math.ceil(height * scale))), align(max(1, math.ceil(width * scale))) + if target_h * target_w > max_pixels: + if max_pixels < factor ** 2: + raise ValueError("max_image_tokens must allow at least one aligned patch") + low, high = 1, height + target_h = target_w = factor + while low <= high: + content_h = (low + high) // 2 + content_w = max(1, math.floor(width * content_h / height)) + candidate_h, candidate_w = align(content_h), align(content_w) + if candidate_h * candidate_w <= max_pixels: + target_h, target_w = candidate_h, candidate_w + low = content_h + 1 + else: + high = content_h - 1 + return target_h, target_w + + +class Glm5NextImageProcessor(Qwen2VLImageProcessor): + def __init__(self, min_image_tokens=16, max_image_tokens=8000, patch_expand_factor=1, **kwargs): + super().__init__(**kwargs) + self.min_image_tokens = min_image_tokens + self.max_image_tokens = max_image_tokens + self.patch_expand_factor = patch_expand_factor + + @classmethod + def from_pretrained(cls, weight_dir): + with open(os.path.join(weight_dir, "processor_config.json")) as f: + return cls(**json.load(f)["image_processor"]) + + def get_image_size(self, height, width): + return smart_resize( + height, + width, + factor=self.patch_size * self.merge_size * self.patch_expand_factor, + min_image_tokens=self.min_image_tokens, + max_image_tokens=self.max_image_tokens, + ) + + def _preprocess_bydevice(self, image, device="cuda"): + pixels = torch.from_numpy(np.array(image.convert("RGB"))).permute(2, 0, 1).contiguous().to(device) + height, width = pixels.shape[-2:] + target_h, target_w = self.get_image_size(height, width) + factor = self.patch_size * self.merge_size * self.patch_expand_factor + scale = min(target_h / height, target_w / width) + if height * width >= factor ** 2 * self.min_image_tokens: + scale = min(1.0, scale) + content_h = max(1, min(target_h, math.floor(height * scale))) + content_w = max(1, min(target_w, math.floor(width * scale))) + if (content_h, content_w) != (height, width): + pixels = F.resize(pixels, [content_h, content_w], interpolation=self.interpolation, antialias=True) + pixels = F.pad(pixels, [0, 0, target_w - content_w, target_h - content_h], fill=0) + pixels = self.rescale_and_normalize( + pixels, self.do_rescale, self.rescale_factor, self.do_normalize, self.image_mean, self.image_std + ) + patch, merge, temporal = self.patch_size, self.merge_size, self.temporal_patch_size + grid_h, grid_w = target_h // patch, target_w // patch + pixels = pixels.reshape(3, grid_h // merge, merge, patch, grid_w // merge, merge, patch) + pixels = pixels.permute(1, 4, 2, 5, 0, 3, 6) + pixels = pixels.unsqueeze(5).expand(-1, -1, -1, -1, -1, temporal, -1, -1) + return pixels.reshape(grid_h * grid_w, 3 * temporal * patch ** 2), torch.tensor([[1, grid_h, grid_w]]) diff --git a/lightllm/models/glm5_next_mtp/__init__.py b/lightllm/models/glm5_next_mtp/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next_mtp/layer_infer/__init__.py b/lightllm/models/glm5_next_mtp/layer_infer/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..50d0461822 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py @@ -0,0 +1,27 @@ +import torch + +from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer + + +class Glm5NextMTPPreLayerInfer(LlamaMultimodalPreLayerInfer): + """Resolve shifted image tokens before the standard NextN embedding/hidden fusion.""" + + def __init__(self, network_config): + super().__init__(network_config) + self.eps_ = network_config["rms_norm_eps"] + + def _fuse_hidden(self, input_embeddings, infer_state, layer_weight): + previous_hidden = infer_state.mtp_draft_input_hiddens + assert input_embeddings.shape[0] == previous_hidden.shape[0] + layer_weight.main_norm_weight_(input=previous_hidden, eps=self.eps_, out=previous_hidden) + layer_weight.enorm_weight_(input=input_embeddings, eps=self.eps_, out=input_embeddings) + layer_weight.hnorm_weight_(input=previous_hidden, eps=self.eps_, out=previous_hidden) + return layer_weight.eh_proj_weight_.mm(torch.cat((input_embeddings, previous_hidden), dim=-1)) + + def context_forward(self, input_ids, infer_state, layer_weight): + input_embeddings = super().context_forward(input_ids, infer_state, layer_weight) + return self._fuse_hidden(input_embeddings, infer_state, layer_weight) + + def token_forward(self, input_ids, infer_state, layer_weight): + input_embeddings = super().token_forward(input_ids, infer_state, layer_weight) + return self._fuse_hidden(input_embeddings, infer_state, layer_weight) diff --git a/lightllm/models/glm5_next_mtp/layer_weights/__init__.py b/lightllm/models/glm5_next_mtp/layer_weights/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..06e5576f16 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,32 @@ +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import RMSNormWeight, ROWMMWeight +from lightllm.models.glm5_next.layer_weights.pre_and_post_layer_weight import add_language_model_aliases + + +class Glm5NextMTPPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config, quant_cfg): + super().__init__(data_type, network_config) + hidden_size = network_config["hidden_size"] + prefix = f"model.layers.{network_config['num_hidden_layers']}" + self.eh_proj_weight_ = ROWMMWeight( + in_dim=2 * hidden_size, + out_dims=[hidden_size], + weight_names=f"{prefix}.eh_proj.weight", + data_type=data_type, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.enorm_weight_ = RMSNormWeight(dim=hidden_size, weight_name=f"{prefix}.enorm.weight", data_type=data_type) + self.hnorm_weight_ = RMSNormWeight(dim=hidden_size, weight_name=f"{prefix}.hnorm.weight", data_type=data_type) + self.final_norm_weight_ = RMSNormWeight( + dim=hidden_size, weight_name=f"{prefix}.shared_head.norm.weight", data_type=data_type + ) + self.wte_weight_ = None + self.lm_head_weight_ = None + # Shared with the target model and injected by the draft model. + self.main_norm_weight_: RMSNormWeight = None + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py new file mode 100644 index 0000000000..19cd9857df --- /dev/null +++ b/lightllm/models/glm5_next_mtp/model.py @@ -0,0 +1,73 @@ +from copy import deepcopy + +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer +from .layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer +from .layer_weights.pre_and_post_layer_weight import Glm5NextMTPPreAndPostLayerWeight + + +class Glm5NextMTPModel(TpPartBaseModel): + is_mtp_draft_model = True + pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight + transformer_weight_class = Glm5NextTransformerLayerWeight + pre_layer_infer_class = Glm5NextMTPPreLayerInfer + post_layer_infer_class = LlamaPostLayerInfer + transformer_layer_infer_class = Glm5NextTransformerLayerInfer + + def __init__(self, kvargs): + self.main_model = kvargs.pop("main_model") + self.mtp_previous_draft_models = kvargs.pop("mtp_previous_draft_models") + super().__init__(kvargs) + + def _init_config(self): + # Reuse the validated target configuration without changing its mHC layout. + self.config = deepcopy(self.main_model.config) + assert self.config.get("num_nextn_predict_layers") == 1, "GLM NextN requires one native MTP block" + self.config["mhc"] = False + + def _init_weights(self, start_layer_index=None): + assert start_layer_index is None + self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, self.config, self.quant_cfg) + self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ + self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ + self.pre_post_weight.main_norm_weight_ = self.main_model.pre_post_weight.final_norm_weight_ + self.trans_layers_weight = [ + self.transformer_weight_class(self.config["num_hidden_layers"], self.data_type, self.config, self.quant_cfg) + ] + + def _init_infer_layer(self, start_layer_index=None): + assert start_layer_index is None + self.pre_infer = self.pre_layer_infer_class(self.config) + self.post_infer = self.post_layer_infer_class(self.config) + # Chained modules reuse the native weights but own distinct cache layers. + layer_index = len(self.main_model.layers_infer) + len(self.mtp_previous_draft_models) + self.layers_infer = [self.transformer_layer_infer_class(layer_index, self.config)] + + def _init_some_value(self): + self.layers_num = 1 + self.vocab_size = self.config["vocab_size"] + self.tp_k_head_num_ = 1 + self.tp_v_head_num_ = 0 + self.qk_nope_head_dim = self.config["qk_nope_head_dim"] + self.qk_rope_head_dim = self.config["qk_rope_head_dim"] + self.q_lora_rank = self.config["q_lora_rank"] + self.kv_lora_rank = self.config["kv_lora_rank"] + self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) + self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim + + def _init_req_manager(self): + self.req_manager = self.main_model.req_manager + self.linear_config = self.main_model.linear_config + + def _init_mem_manager(self): + self.mem_manager = self.main_model.mem_manager + + def _init_att_backend(self): + self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) + self.decode_att_backend = self.prefill_att_backend + + def autotune_layers(self): + return 1 diff --git a/lightllm/models/vit/triton_kernel/rms_norm_vit.py b/lightllm/models/vit/triton_kernel/rms_norm_vit.py index 387bedfcfe..1e44a0ae2e 100644 --- a/lightllm/models/vit/triton_kernel/rms_norm_vit.py +++ b/lightllm/models/vit/triton_kernel/rms_norm_vit.py @@ -17,6 +17,7 @@ def rms_norm_kernel( eps: tl.constexpr, N_COLS: tl.constexpr, BLOCK_N: tl.constexpr, + ROUND_NORM_BEFORE_WEIGHT: tl.constexpr, ): """Rms norm kernel.""" prog_id = tl.program_id(0) @@ -30,13 +31,21 @@ def rms_norm_kernel( var = tl.sum(xf * xf, 0) * float(1.0 / N_COLS) out = xf / tl.sqrt(var + eps) + if ROUND_NORM_BEFORE_WEIGHT: + out = out.to(x.dtype) out = (w * out).to(x.dtype) out_ptr = output + prog_id * out_row_stride tl.store(out_ptr + offsets * out_col_stride, out, mask=offsets < N_COLS) -def rms_norm(hidden_states: Tensor, weight: Tensor, eps: float = 1e-5, use_custom_tensor_mananger: bool = False): +def rms_norm( + hidden_states: Tensor, + weight: Tensor, + eps: float = 1e-5, + use_custom_tensor_mananger: bool = False, + round_norm_before_weight: bool = False, +): """Rms norm.""" assert hidden_states.is_contiguous(), "hidden_states must be contiguous" @@ -73,12 +82,88 @@ def rms_norm(hidden_states: Tensor, weight: Tensor, eps: float = 1e-5, use_custo eps=eps, N_COLS=hidden_dim, BLOCK_N=BLOCK_N, + ROUND_NORM_BEFORE_WEIGHT=round_norm_before_weight, num_warps=4, num_stages=3, ) return output.reshape(origin_shape) +@triton.jit +def qk_rms_norm_kernel( + input, + q_weight, + k_weight, + q_output, + k_output, + input_stride_token, + input_stride_qkv, + input_stride_head, + input_stride_dim, + output_stride_token, + output_stride_head, + output_stride_dim, + eps: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + offsets = tl.arange(0, BLOCK_N) + mask = offsets < HEAD_DIM + + input_offset = token * input_stride_token + head * input_stride_head + offsets * input_stride_dim + q = tl.load(input + input_offset, mask=mask, other=0.0).to(tl.float32) + k = tl.load(input + input_offset + input_stride_qkv, mask=mask, other=0.0).to(tl.float32) + q_rstd = tl.rsqrt(tl.sum(q * q, axis=0) / HEAD_DIM + eps) + k_rstd = tl.rsqrt(tl.sum(k * k, axis=0) / HEAD_DIM + eps) + + # GLM rounds the normalized value to the input dtype before applying the weight. + q = (q * q_rstd).to(input.dtype.element_ty) + k = (k * k_rstd).to(input.dtype.element_ty) + q = q * tl.load(q_weight + offsets, mask=mask, other=0.0) + k = k * tl.load(k_weight + offsets, mask=mask, other=0.0) + + output_offset = token * output_stride_token + head * output_stride_head + offsets * output_stride_dim + tl.store(q_output + output_offset, q, mask=mask) + tl.store(k_output + output_offset, k, mask=mask) + + +def qk_rms_norm(input: Tensor, q_weight: Tensor, k_weight: Tensor, eps: float) -> tuple[Tensor, Tensor]: + """Normalize Q and K from a packed ``[tokens, 3, heads, head_dim]`` QKV tensor. + + This avoids materializing the strided Q/K views before their per-head RMSNorm. + """ + + assert input.ndim == 4 and input.shape[1] == 3 and input.is_contiguous() + tokens, _, heads, head_dim = input.shape + assert q_weight.shape == k_weight.shape == (head_dim,) + q_output = torch.empty((tokens, heads, head_dim), dtype=input.dtype, device=input.device) + k_output = torch.empty_like(q_output) + input_stride_token, input_stride_qkv, input_stride_head, input_stride_dim = input.stride() + output_stride_token, output_stride_head, output_stride_dim = q_output.stride() + qk_rms_norm_kernel[(tokens, heads)]( + input, + q_weight, + k_weight, + q_output, + k_output, + input_stride_token, + input_stride_qkv, + input_stride_head, + input_stride_dim, + output_stride_token, + output_stride_head, + output_stride_dim, + eps=eps, + HEAD_DIM=head_dim, + BLOCK_N=triton.next_power_of_2(head_dim), + num_warps=4, + num_stages=3, + ) + return q_output, k_output + + def test(): def _rms_norm_ref(x: torch.Tensor, weight: torch.Tensor, eps: float): var = (x.float() ** 2).mean(dim=-1, keepdim=True) diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index 99a6c5833d..f4515b0f12 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -1234,17 +1234,16 @@ def __init__(self): self.func_detail_regex = re.compile( r"([^<\n]+?)(?:\n|(?=)|(?=))(.*?)", re.DOTALL ) + self.func_name_regex = re.compile(r"([^<\n]+?)(?:\n|(?=)|(?=))") + self._streaming_tool_name: Optional[str] = None # Extract arg_key/arg_value pairs self.func_arg_regex = re.compile(r"(.*?)\s*(.*?)", re.DOTALL) - self._last_arguments = "" - self._normal_text_buffer = "" - def has_tool_call(self, text: str) -> bool: """Check if the text contains a GLM-4.7 format tool call.""" return self.bot_token in text - def _parse_xml_arguments(self, arg_text: str) -> dict: + def _parse_xml_arguments(self, arg_text: str, tool: Tool) -> dict: """ Parse XML-style arguments into a dictionary. @@ -1258,10 +1257,15 @@ def _parse_xml_arguments(self, arg_text: str) -> dict: return {} args = {} + properties = (tool.function.parameters or {}).get("properties", {}) matches = self.func_arg_regex.findall(arg_text) for key, value in matches: key = key.strip() - value = value.strip() + # XML string values are literal: file content and edit targets may + # contain meaningful whitespace or text that happens to be JSON. + if properties.get(key, {}).get("type") == "string": + args[key] = value + continue # Try to parse value as JSON for complex types (arrays, objects, numbers, booleans) try: parsed_value = json.loads(value) @@ -1308,11 +1312,11 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult continue # Parse XML arguments to JSON - func_args = self._parse_xml_arguments(arg_text) + func_args = self._parse_xml_arguments(arg_text, tools[tool_indices[func_name]]) calls.append( ToolCallItem( - tool_index=tool_indices[func_name], + tool_index=len(calls), name=func_name, parameters=json.dumps(func_args, ensure_ascii=False), ) @@ -1324,134 +1328,52 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult return StreamingParseResult(normal_text=normal_text, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: - """ - Streaming incremental parsing for GLM-4.7 tool calls. - - This handles the streaming case where tool calls arrive incrementally. - """ + """Buffer XML arguments while emitting tool deltas to keep the stream alive.""" self._buffer += new_text - current_text = self._buffer - - # Check if we have a tool call starting - if not self.has_tool_call(current_text): - # Check for partial bot_token at the end - partial_len = self._ends_with_partial_token(current_text, self.bot_token) - if partial_len: - # Might be partial bot_token, keep buffering - return StreamingParseResult() - - # No tool call, emit as normal text - self._buffer = "" - # Clean up any stray end tokens - if self.eot_token in new_text: - new_text = new_text.replace(self.eot_token, "") - return StreamingParseResult(normal_text=new_text) - - # Build tool indices if not already built - if not hasattr(self, "_tool_indices"): - self._tool_indices = self._get_tool_indices(tools) - + normal_text = "" calls: List[ToolCallItem] = [] - - try: - # Check if we have a complete tool call - if self.eot_token in current_text: - # We have at least one complete tool call - # Parse all complete tool calls - result = self.detect_and_parse(current_text, tools) - - # Find the end of the last complete tool call - last_end = current_text.rfind(self.eot_token) - if last_end != -1: - remaining = current_text[last_end + len(self.eot_token) :] - self._buffer = remaining.lstrip() - else: - self._buffer = "" - - # Reset state for next tool call - self.current_tool_id = -1 - self.current_tool_name_sent = False - self._last_arguments = "" - - return result - - # We have a partial tool call - try to stream it - # Extract what we can from the partial tool call - tool_call_start = current_text.find(self.bot_token) - if tool_call_start == -1: - return StreamingParseResult() - - # Get content after - content_after_start = current_text[tool_call_start + len(self.bot_token) :] - - # Try to extract function name (first line after ) - newline_pos = content_after_start.find("\n") - if newline_pos == -1: - # Still waiting for function name to complete - return StreamingParseResult() - - func_name = content_after_start[:newline_pos].strip() - - # Initialize state if this is the first tool call - if self.current_tool_id == -1: - self.current_tool_id = 0 - self.prev_tool_call_arr = [] - self.streamed_args_for_tool = [""] - - # Ensure we have enough entries - while len(self.prev_tool_call_arr) <= self.current_tool_id: - self.prev_tool_call_arr.append({}) - while len(self.streamed_args_for_tool) <= self.current_tool_id: - self.streamed_args_for_tool.append("") - - # Check if function name is valid - if func_name and func_name in self._tool_indices: - if not self.current_tool_name_sent: - # Send function name first - calls.append( - ToolCallItem( - tool_index=self.current_tool_id, - name=func_name, - parameters="", - ) - ) - self.current_tool_name_sent = True - self.prev_tool_call_arr[self.current_tool_id] = { - "name": func_name, - "arguments": {}, - } + while self._buffer: + start = self._buffer.find(self.bot_token) + if start == -1: + keep = self._ends_with_partial_token(self._buffer, self.bot_token) + end = len(self._buffer) - keep + normal_text += self._buffer[:end] + self._buffer = self._buffer[end:] + break + + normal_text += self._buffer[:start] + self._buffer = self._buffer[start:] + end = self._buffer.find(self.eot_token) + if end == -1: + # Match the buffered Qwen3-Coder flow: announce the name once, + # then send empty argument deltas until the XML call is complete. + # Long Write/Bash arguments must not leave the HTTP stream idle. + name = None + if self._streaming_tool_name is None: + match = self.func_name_regex.match(self._buffer) + if match and match.group(1).strip() in self._get_tool_indices(tools): + name = match.group(1).strip() + self._streaming_tool_name = name + self.current_tool_id += 1 + if self._streaming_tool_name is not None: + calls.append(ToolCallItem(tool_index=self.current_tool_id, name=name, parameters="")) + break + + end += len(self.eot_token) + # Waiting for the closing tag avoids re-emitting partial JSON or + # rewriting already sent arguments when another XML key arrives. + result = self.detect_and_parse(self._buffer[:end], tools) + for call in result.calls: + if self._streaming_tool_name is None: + self.current_tool_id += 1 else: - # Stream arguments incrementally - arg_text = content_after_start[newline_pos + 1 :] - current_args = self._parse_xml_arguments(arg_text) - - if current_args: - current_args_json = json.dumps(current_args, ensure_ascii=False) - prev_args = self.prev_tool_call_arr[self.current_tool_id].get("arguments", {}) - prev_args_json = json.dumps(prev_args, ensure_ascii=False) if prev_args else "" - - if current_args_json != prev_args_json: - # Calculate the diff - sent = len(self.streamed_args_for_tool[self.current_tool_id]) - argument_diff = current_args_json[sent:] - - if argument_diff: - calls.append( - ToolCallItem( - tool_index=self.current_tool_id, - name=None, - parameters=argument_diff, - ) - ) - self.streamed_args_for_tool[self.current_tool_id] += argument_diff - - self.prev_tool_call_arr[self.current_tool_id]["arguments"] = current_args - - return StreamingParseResult(normal_text="", calls=calls) + call.name = None + call.tool_index = self.current_tool_id + calls.append(call) + self._streaming_tool_name = None + self._buffer = self._buffer[end:] - except Exception as e: - logger.error(f"Error in GLM-4.7 parse_streaming_increment: {e}") - return StreamingParseResult(normal_text="", calls=calls) + return StreamingParseResult(normal_text=normal_text, calls=calls) class DeepSeekV32Detector(BaseFormatDetector): diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index f17cc5aa66..5dadaf0b13 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -90,6 +90,10 @@ def get_tokenizer( image_processor = Qwen2VLImageProcessor.from_pretrained(tokenizer_name) tokenizer = Tarsier2Tokenizer(tokenizer=tokenizer, image_processor=image_processor, model_cfg=model_cfg) + elif model_type == "glm5_next" and model_cfg.get("vision_config") is not None: + from ..models.glm5_next.tokenizer import Glm5NextTokenizer + + tokenizer = Glm5NextTokenizer(tokenizer, model_cfg, tokenizer_name) elif model_type == "llava" or model_type == "internlmxcomposer2": from ..models.llava.model import LlavaTokenizer diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index b3dad0ef71..22d5a43e73 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -22,6 +22,7 @@ from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel +from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.dist_utils import init_vision_distributed_env from lightllm.utils.envs_utils import get_env_start_args @@ -80,6 +81,8 @@ def exposed_init_model(self, kvargs): self.model_type = model_cfg["model_type"] if self.model_type == "qwen": self.model = QWenVisionTransformer(**model_cfg["visual"]).eval().bfloat16() + elif self.model_type == "glm5_next": + self.model = Glm5NextVisionTransformer(kvargs, **model_cfg["vision_config"]).eval().bfloat16() elif self.model_type == "qwen2_vl": self.model = ( Qwen2VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 784359a4d4..8f76b4a53c 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -413,7 +413,7 @@ def has_vision_module(model_path: str) -> bool: # Qwen2_5_VisionTransformerPretrainedModel model_cfg["vision_config"] return True - elif model_type in ["qwen3_vl", "qwen3_vl_moe"]: + elif model_type in ["qwen3_vl", "qwen3_vl_moe", "glm5_next"]: # Qwen3VisionTransformerPretrainedModel model_cfg["vision_config"] return True @@ -477,7 +477,7 @@ def is_linear_att_mixed_model(model_path: str) -> bool: model_cfg, _ = PretrainedConfig.get_config_dict(model_path) model_type = model_cfg["model_type"] - if model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"]: + if model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", "glm5_next", "glm5_next_text"]: return True else: return False diff --git a/requirements.txt b/requirements.txt index 7e30557f12..41c9c26cb4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,6 @@ +# FlashInfer cubin wheels are published on the official FlashInfer index. +--extra-index-url https://flashinfer.ai/whl + anyio==3.7.1 black==23.12.0 boltons==23.0.0 @@ -38,7 +41,8 @@ packaging==24.2 pip==23.0.1 pluggy==1.2.0 plumbum==1.8.2 -protobuf==4.22.3 +# CUTLASS DSL 4.6 requires protobuf >= 6.30.2, < 7. +protobuf==6.33.5 pycparser==2.21 pydantic==2.10.3 pyOpenSSL==23.2.0 @@ -80,9 +84,15 @@ frozendict==2.4.6 atomics==1.0.3 easydict==1.13 hypercorn==0.18.0 -flashinfer-python==0.6.12 -flashinfer-cubin==0.6.12 -sglang-kernel==0.4.2.post1 +# SM90 NoPE MLA (512 latent dimensions, zero RoPE dimensions). +flashinfer-python[cu13]==0.6.18 +flashinfer-cubin==0.6.18 +cuda-tile==1.4.0 +nvidia-cutlass-dsl==4.6.2 +nvidia-cudnn-frontend==1.25.0 +nccl4py==0.3.1 +# FA3 only_qv supports NoPE MLA on Hopper; 0.4.5 keeps Torch 2.11 compatibility. +sglang-kernel==0.4.5 httpx==0.28.1 librosa==0.11.0 cuda_bindings==13.2.0 @@ -98,4 +108,4 @@ nixl==1.2.0 xformers==0.0.35 redis==7.3.0 litellm>=1.84.8,<1.85 -torch_memory_saver==0.0.9.post1 \ No newline at end of file +torch_memory_saver==0.0.9.post1 diff --git a/test/benchmark/service/benchmark_glm53_flash.py b/test/benchmark/service/benchmark_glm53_flash.py new file mode 100644 index 0000000000..5f71f69d8f --- /dev/null +++ b/test/benchmark/service/benchmark_glm53_flash.py @@ -0,0 +1,115 @@ +"""Fixed-token GLM-5.3 Flash serving benchmark (TTFT and decode measured separately).""" + +import argparse +import concurrent.futures +import json +import random +import statistics +import time +from pathlib import Path + +import requests +from lightllm.server.tokenizer import get_tokenizer + + +def generate(url, prompt, output_tokens): + start = time.perf_counter() + arrivals = [] + ids = [] + with requests.post( + url.rstrip("/") + "/generate_stream", + json={ + "inputs": prompt, + "parameters": { + "do_sample": False, + "ignore_eos": True, + "max_new_tokens": output_tokens, + }, + }, + stream=True, + timeout=600, + ) as response: + response.raise_for_status() + for line in response.iter_lines(chunk_size=1): + if not line.startswith(b"data:"): + continue + event = json.loads(line[5:]) + if "token" not in event: + raise RuntimeError(event) + arrivals.append(time.perf_counter()) + ids.append(event["token"]["id"]) + if len(ids) != output_tokens: + raise RuntimeError(f"Expected {output_tokens} output tokens, got {len(ids)}") + decode_seconds = arrivals[-1] - arrivals[0] + return { + "input_tokens": len(prompt), + "output_tokens": len(ids), + "ttft_ms": (arrivals[0] - start) * 1000, + "tpot_ms": decode_seconds / (len(ids) - 1) * 1000, + "decode_tokens_per_second": (len(ids) - 1) / decode_seconds, + "latency_seconds": arrivals[-1] - start, + "token_ids": ids, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", default="http://127.0.0.1:18153") + parser.add_argument("--model-dir", default="/nvme/models/GLM-5.3-Flash") + parser.add_argument("--input-tokens", type=int, nargs="+", default=[1024, 4096]) + parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 4, 8]) + parser.add_argument("--output-tokens", type=int, default=128) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--warmup-rounds", type=int, default=2) + parser.add_argument("--seed", type=int) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.output_tokens < 2: + parser.error("--output-tokens must be at least 2 to measure decode") + if args.seed is None: + args.seed = time.time_ns() + tokenizer = get_tokenizer(args.model_dir) + content = tokenizer.encode( + "The following document describes a language model inference service. " + "Explain its performance clearly and continue the discussion. ", + add_special_tokens=False, + ) + rng = random.Random(args.seed) + + def prompt(length): + # Vary the FIRST tokens, including warmups, so radix hits cannot + # inflate prefill throughput while prompt caching remains enabled. + return [rng.randrange(1000, 100000) for _ in range(min(length, 16))] + (content * (length // len(content) + 1))[ + : max(0, length - 16) + ] + + report = {"settings": {**vars(args), "output": str(args.output)}, "results": []} + for length in args.input_tokens: + for concurrency in args.concurrency: + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + for _ in range(args.warmup_rounds): + list(pool.map(lambda p: generate(args.url, p, 32), [prompt(length) for _ in range(concurrency)])) + for repeat in range(args.repeats): + prompts = [prompt(length) for _ in range(concurrency)] + start = time.perf_counter() + rows = list(pool.map(lambda p: generate(args.url, p, args.output_tokens), prompts)) + elapsed = time.perf_counter() - start + result = { + "input_tokens": length, + "concurrency": concurrency, + "repeat": repeat, + "ttft_ms_median": statistics.median(r["ttft_ms"] for r in rows), + "tpot_ms_median": statistics.median(r["tpot_ms"] for r in rows), + "per_request_decode_tps_median": statistics.median(r["decode_tokens_per_second"] for r in rows), + "output_tps_including_prefill": concurrency * args.output_tokens / elapsed, + "elapsed_seconds": elapsed, + "requests": rows, + } + report["results"].append(result) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2)) + print(json.dumps({k: v for k, v in result.items() if k != "requests"}), flush=True) + + +if __name__ == "__main__": + main() diff --git a/test/start_scripts/glm53/glm53_pd_1p1d.sh b/test/start_scripts/glm53/glm53_pd_1p1d.sh new file mode 100644 index 0000000000..4ad924dd74 --- /dev/null +++ b/test/start_scripts/glm53/glm53_pd_1p1d.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "$#" -lt 2 || "$#" -gt 3 ]]; then + echo "Usage: $0 [model_dir]" >&2 + exit 2 +fi + +PORT="$1" +NODE_IP="$2" +MODEL_DIR="${3:-/nvme/models/GLM-5.3-Flash}" + +# P/D must advertise an address reachable by the master and the other node. +export NO_PROXY="${NO_PROXY:+${NO_PROXY},}127.0.0.1,localhost,${NODE_IP}" +export no_proxy="${NO_PROXY}" +export LOADWORKER="${LOADWORKER:-8}" + +COMMON_ARGS=( + --model_dir "${MODEL_DIR}" + --model_name glm53 + --tp 4 + --batch_max_tokens 8192 + --running_max_req_size 64 + --mem_fraction 0.8 + --enable_fused_shared_experts + --tool_call_parser glm47 + --reasoning_parser glm45 + --linear_att_ssm_data_type float32 + --pd_trans_mode nccl + # One transfer page must also fit the global Conv/SSM/indexer-tail state. + --pd_kv_page_size 16384 + --pd_kv_page_num 2 + --pd_master_ip 127.0.0.1 + --pd_master_port "${PORT}" + --host "${NODE_IP}" +) + +PIDS=() +cleanup() { + kill -TERM "${PIDS[@]}" 2>/dev/null || true + wait "${PIDS[@]}" 2>/dev/null || true +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m lightllm.server.api_server \ + "${COMMON_ARGS[@]}" \ + --run_mode prefill \ + --disable_cudagraph \ + --port "$((PORT + 1))" \ + --nccl_port "$((PORT + 101))" & +PIDS+=("$!") + +CUDA_VISIBLE_DEVICES=4,5,6,7 python -m lightllm.server.api_server \ + "${COMMON_ARGS[@]}" \ + --run_mode decode \ + --graph_max_batch_size 64 \ + --graph_max_len_in_batch 65536 \ + --port "$((PORT + 2))" \ + --nccl_port "$((PORT + 102))" & +PIDS+=("$!") + +CUDA_VISIBLE_DEVICES= python -m lightllm.server.api_server \ + --model_dir "${MODEL_DIR}" \ + --model_name glm53 \ + --run_mode pd_master \ + --pd_master_mode 1p1d \ + --tool_call_parser glm47 \ + --reasoning_parser glm45 \ + --host 127.0.0.1 \ + --port "${PORT}" & +PIDS+=("$!") + +echo "GLM-5.3 Flash 1P1D is starting at http://127.0.0.1:${PORT} (P TP4 + D TP4)" +wait -n "${PIDS[@]}" +exit 1 diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py new file mode 100644 index 0000000000..4c33d8cf8a --- /dev/null +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py @@ -0,0 +1,106 @@ +"""Check the FlashInfer SM90 path used for unpadded GLM NoPE MLA.""" + +import inspect + +import pytest +import torch + +flashinfer = pytest.importorskip("flashinfer") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="FlashInfer NoPE FA3 MLA requires Hopper", +) + + +@pytest.mark.parametrize( + "num_tokens,num_heads,kv_dtype", + [ + (1, 16, torch.bfloat16), + (8, 16, torch.bfloat16), + (128, 16, torch.bfloat16), + (8, 64, torch.bfloat16), + (8, 16, torch.float8_e4m3fn), + ], +) +def test_nope_sparse_mla_eager_and_graph_replay(num_tokens, num_heads, kv_dtype): + wrapper_cls = flashinfer.mla.BatchMLAPagedAttentionWrapper + assert "ckv_scale_arr" in inspect.signature(wrapper_cls.run).parameters, "FlashInfer >= 0.6.18 is required" + + torch.manual_seed(17) + dim, num_slots, width = 512, 4096, 2176 + sm_scale = 0.0625 + q = torch.randn(num_tokens, num_heads, dim, dtype=torch.bfloat16, device="cuda") + q_pe = q[..., :0] + # BF16 KV shares its token row with the 144-byte index region. The + # attention kernel must respect the resulting non-contiguous token stride. + tail = 72 if kv_dtype == torch.bfloat16 else 0 + packed = torch.full((num_slots, 1, dim + tail), 7.0, dtype=kv_dtype, device="cuda") + ckv = packed[..., :dim] + raw_kv = torch.randn(num_slots, 1, dim, dtype=torch.bfloat16, device="cuda") + scale = 0.125 if kv_dtype == torch.float8_e4m3fn else 1.0 + ckv.copy_((raw_kv.float() / scale).to(kv_dtype)) + kpe = packed[..., dim:dim] + reference_kv = (ckv.float() * scale).to(torch.bfloat16).float() + scale_kwargs = {"ckv_scale": scale, "kpe_scale": 1.0} if scale != 1.0 else {} + + qo_indptr = torch.arange(num_tokens + 1, dtype=torch.int32) + kv_indptr = qo_indptr * width + slots = torch.zeros(num_tokens, width, dtype=torch.int32, device="cuda") + wrapper = wrapper_cls( + torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda"), + use_cuda_graph=True, + qo_indptr=qo_indptr.to("cuda"), + kv_indptr=kv_indptr.to("cuda"), + kv_indices=slots.flatten(), + kv_len_arr=torch.empty(num_tokens, dtype=torch.int32, device="cuda"), + backend="fa3", + ) + + def plan(step): + choices = [1, 3, 127, 128, 511, 2048, 2049, 2051] + lengths = [choices[(row + step) % len(choices)] for row in range(num_tokens)] + slots.zero_() + for row, length in enumerate(lengths): + slots[row, :length] = (torch.arange(length, device="cuda") * 5 + row * 31 + step * 7) % num_slots + # Lengths are compiled into the host-side plan. Replan before graph + # replay when they change, while keeping all captured buffers stable. + wrapper.plan( + qo_indptr, + kv_indptr, + slots.flatten(), + torch.tensor(lengths, dtype=torch.int32), + num_heads, + dim, + 0, + 1, + False, + sm_scale, + torch.bfloat16, + kv_dtype, + ) + return lengths + + def check(output, lengths): + expected = [] + for row, length in enumerate(lengths): + keys = reference_kv[slots[row, :length].long(), 0] + scores = q[row].float() @ keys.T * sm_scale + expected.append(scores.softmax(-1) @ keys) + torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) + + lengths = plan(0) + output = torch.empty_like(q) + wrapper.run(q, q_pe, ckv, kpe, out=output, **scale_kwargs) + check(output, lengths) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + wrapper.run(q, q_pe, ckv, kpe, out=output, **scale_kwargs) + graph.replay() + check(output, lengths) + + lengths = plan(1) + q.normal_() + graph.replay() + check(output, lengths) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py new file mode 100644 index 0000000000..525dfae756 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py @@ -0,0 +1,196 @@ +import collections +import functools +import inspect +import json +from itertools import accumulate + +import pytest +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import kda +from lightllm.common.kernel_config import KernelConfigs +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel + + +KERNELS = [ + ("_chunk_kda_scaled_dot_kkt_sub_inter", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter"), + ("_chunk_kda_scaled_dot_kkt_sub_intra", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra"), + ("recompute_w_u_fwd", "recompute_w_u_fwd_kernel"), + ("chunk_gla_fwd_o_gk", "chunk_gla_fwd_kernel_o"), +] + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.CLOSE_AUTOTUNE) + + +def wrapper_inputs(kernel, tokens=65): + q = torch.empty(1, tokens, 2, 128, dtype=torch.bfloat16) + v = torch.empty(1, tokens, 2, 80, dtype=q.dtype) + g = torch.empty_like(q, dtype=torch.float32) + chunk_indices = torch.tensor([[0, 0], *[[1, i] for i in range(triton.cdiv(tokens - 3, 64))]]) + values = dict( + q=q, + k=q, + v=v, + g=g, + gk=g, + beta=torch.empty(1, tokens, 2), + Akk=torch.empty(1, tokens, 2, 64), + Aqk=torch.empty(1, tokens, 2, 64), + R=torch.empty(1, tokens, 2, 64, dtype=q.dtype), + h=torch.empty(1, len(chunk_indices), 2, 128, 80, dtype=q.dtype), + o=torch.empty_like(v), + scale=128 ** -0.5, + cu_seqlens=torch.tensor([0, 3, tokens]), + chunk_indices=chunk_indices, + ) + return {name: values[name] for name in inspect.signature(kernel.fn).parameters if name in values} + + +@pytest.mark.parametrize("wrapper_name,jit_name", KERNELS) +def test_lightllm_tuning_selects_persists_and_reuses_launch_config(monkeypatch, tmp_path, wrapper_name, jit_name): + kernel = getattr(kda, wrapper_name) + assert isinstance(kernel, Autotuner) + inputs = wrapper_inputs(kernel) + candidates = [kernel.configs_gen_func()[0], kernel.configs_gen_func()[-1]] + launches, benchmarks = [], [] + + class LaunchRecorder: + def __getitem__(self, grid): + def launch(**kwargs): + launches.append((grid, kwargs)) + + return launch + + def bench(*args, run_config, **kwargs): + benchmarks.append(run_config) + return 1.0 if run_config == candidates[-1] else 2.0 + + monkeypatch.setattr(kda, jit_name, LaunchRecorder()) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr("lightllm.common.kernel_config.get_current_device_name", lambda: "test-device") + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + monkeypatch.setattr(kernel, "configs_gen_func", lambda: candidates) + monkeypatch.setattr(kernel, "_bench", bench) + + with Autotuner.autotune_warmup(): + kernel(**inputs) + assert benchmarks == candidates + for name, value in candidates[-1].items(): + assert launches[-1][1][name] == value + + cache_file = tmp_path / KernelConfigs.get_config_file_name(kernel._static_key(**inputs)) + assert json.loads(cache_file.read_text()) == {"65": candidates[-1]} + + # A new wrapper invocation must reload the same selected config from disk. + kernel.cached_configs.clear() + kernel.fast_match_configs.clear() + kernel.warmuped_configs_set.clear() + benchmarks.clear() + kernel(**inputs) + assert benchmarks == [] + for name, value in candidates[-1].items(): + assert launches[-1][1][name] == value + + # Packed B=1 still needs different run keys as the token count grows. + assert kernel._run_key(**inputs) == 65 + assert kernel._run_key(**wrapper_inputs(kernel, tokens=129)) == 129 + assert kernel._static_key(**inputs) == kernel._static_key(**wrapper_inputs(kernel, tokens=129)) + # Request boundaries are runtime data; all calls use the packed path. + single_request = dict(inputs, cu_seqlens=torch.tensor([0, 65])) + assert kernel._static_key(**inputs) == kernel._static_key(**single_request) + + +def recurrent_reference(q, k, v, g, beta, initial_state, sequences): + q, k, v, g, beta = [x.float().cpu() for x in (q, k, v, g, beta)] + q = q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() / q.shape[-1] ** 0.5 + k = k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() + output = torch.empty_like(v) + final = initial_state.float().cpu().clone() + for n, (batch, start, end) in enumerate(sequences): + state = final[n] + for t in range(start, end): + kt = k[batch, t] + state = state * g[batch, t].exp()[..., None] + delta = beta[batch, t, :, None] * (v[batch, t] - torch.einsum("hk,hkv->hv", kt, state)) + state = state + kt[..., None] * delta[:, None, :] + output[batch, t] = torch.einsum("hk,hkv->hv", q[batch, t], state) + final[n] = state + return output, final + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("config_index", [None, 0, -1], ids=["default", "first-candidate", "last-candidate"]) +@pytest.mark.parametrize( + "seq_lens,key_dim,value_dim", + [((65, 65), 64, 96), ((3, 65, 129), 128, 80)], + ids=["equal-length-packed", "ragged-packed"], +) +def test_kda_configs_match_token_recurrence(monkeypatch, config_index, seq_lens, key_dim, value_dim): + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + if config_index is not None: + for wrapper_name, _ in KERNELS: + kernel = getattr(kda, wrapper_name) + config = kernel.configs_gen_func()[config_index] + monkeypatch.setattr(kda, wrapper_name, functools.partial(kernel, run_config=config)) + + # Include partial final chunks, distinct K/V sizes, and V tiles with masked columns. + batch, tokens, heads = 1, sum(seq_lens), 2 + boundaries = [0, *accumulate(seq_lens)] + sequences = [(0, start, end) for start, end in zip(boundaries, boundaries[1:])] + cu_seqlens = torch.tensor(boundaries, device="cuda", dtype=torch.int32) + q, k = [torch.randn(batch, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16) for _ in range(2)] + v = torch.randn(batch, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16) + g = -torch.rand(batch, tokens, heads, key_dim, device="cuda") * 0.1 + beta = torch.rand(batch, tokens, heads, device="cuda") + initial = torch.randn(len(sequences), heads, key_dim, value_dim, device="cuda") * 0.1 + initial_before = initial.clone() + expected, expected_final = recurrent_reference(q, k, v, g, beta, initial, sequences) + + actual, final = kda.chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + torch.testing.assert_close(actual.float().cpu(), expected, atol=4e-3, rtol=3e-2) + torch.testing.assert_close(final.cpu(), expected_final, atol=8e-3, rtol=3e-2) + torch.testing.assert_close(initial, initial_before, atol=0, rtol=0) + + +@pytest.mark.parametrize("fused_gate", [False, True]) +@pytest.mark.parametrize("invalid_input", ["batch-dimension", "missing-boundaries"]) +def test_kda_requires_packed_inputs(fused_gate, invalid_input): + batch = 2 if invalid_input == "batch-dimension" else 1 + q = torch.empty(batch, 3, 2, 128) + inputs = dict( + q=q, + k=q, + v=q, + beta=torch.empty(batch, 3, 2), + cu_seqlens=None if invalid_input == "missing-boundaries" else torch.tensor([0, 3, 6]), + ) + if fused_gate: + kernel = kda.chunk_kda_with_fused_gate + inputs.update(raw_g=q, A_log=torch.empty(2), g_bias=None) + else: + kernel = kda.chunk_kda + inputs.update(g=q) + expected_message = "packed q" if invalid_input == "batch-dimension" else "cu_seqlens is required" + with pytest.raises(AssertionError, match=expected_message): + kernel(**inputs) diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py new file mode 100644 index 0000000000..03ae1f3c33 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 + +import sys +from types import ModuleType + +import pytest +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.basemodel.triton_kernel.mhc import hc_post, hc_pre_norm +from lightllm.common.basemodel.triton_kernel.mhc import pre_norm +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_iters, post_multiplier=2.0): + """Unfused FP32 reference; returns the merged input before sublayer RMSNorm.""" + assert x.ndim == 2 and x.shape[-1] % streams == 0 + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + residual = x.view(tokens, streams, hidden) + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre_raw = mixes[:, :streams] + post_raw = mixes[:, streams : 2 * streams] + residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) + + pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps + post = post_multiplier * torch.sigmoid(post_raw * scale[1] + base[streams : 2 * streams]) + residual_mix = (residual_raw * scale[2] + base[2 * streams :].view(streams, streams)).softmax(dim=-1) + residual_mix = residual_mix + hc_eps + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + for _ in range(sinkhorn_iters - 1): + residual_mix = residual_mix / (residual_mix.sum(dim=-1, keepdim=True) + hc_eps) + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + + layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) + return layer_input, residual_mix, post + + +def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams): + """Unfused FP32 reference; residual_mix maps input streams to output streams.""" + tokens, hidden = layer_output.shape + residual_3d = residual.view(tokens, streams, hidden) + mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) + out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual + return out.to(layer_output.dtype).reshape(tokens, streams * hidden) + + +@pytest.fixture(autouse=True) +def setup(): + torch.manual_seed(1525) + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + + +@pytest.mark.parametrize("tokens", [3, 19, 3073]) +@pytest.mark.parametrize("disable_deepgemm", [False, True], ids=["deepgemm", "torch"]) +def test_mhc_matches_reference(tokens, disable_deepgemm, monkeypatch): + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", disable_deepgemm) + if disable_deepgemm: + monkeypatch.setitem(sys.modules, "deep_gemm", None) + streams, hidden = 4, 4096 + x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 + scale = torch.randn(3, device="cuda") + base = torch.randn(24, device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + expected = _hc_pre_reference(x, fn, scale, base, streams, 1e-5, 1e-6, 20) + actual = hc_pre_norm(x, fn, scale, base, norm, streams, 1e-5, 1e-5, 1e-6, 20) + torch.testing.assert_close(actual[0], rmsnorm_forward(expected[0], norm, 1e-5), atol=0.04, rtol=0.03) + for a, b in zip(actual[1:], expected[1:]): + # DeepGEMM's mHC projection uses TF32 inputs and FP32 accumulation. + torch.testing.assert_close(a, b, atol=1e-4, rtol=1e-3) + layer_out = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close( + hc_post(layer_out, x, *actual[1:], streams), + _hc_post_reference(layer_out, x, *actual[1:], streams), + atol=0.04, + rtol=0.02, + ) + + +@pytest.mark.parametrize("missing", ["module", "kernel"]) +def test_mhc_reports_missing_deepgemm(monkeypatch, missing): + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) + monkeypatch.setitem(sys.modules, "deep_gemm", None if missing == "module" else ModuleType("deep_gemm")) + x = torch.empty(1, 16, dtype=torch.bfloat16) + fn = torch.empty(24, 16) + scale = torch.empty(3) + base = torch.empty(24) + norm = torch.empty(4, dtype=torch.bfloat16) + + with pytest.raises(ImportError, match="deep_gemm"): + hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20) diff --git a/unit_tests/common/fused_moe/test_activation_config.py b/unit_tests/common/fused_moe/test_activation_config.py new file mode 100644 index 0000000000..311b1fc652 --- /dev/null +++ b/unit_tests/common/fused_moe/test_activation_config.py @@ -0,0 +1,127 @@ +import dataclasses +import json + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.marlin_impl import FuseMoeMarlin +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for name, value in { + "GLOBAL_RANK": 0, + "GLOBAL_WORLD_SIZE": 1, + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_RANK_IN_NODE": 0, + "CURRENT_DEVICE_ID": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + name, str(value)) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("activation", ["silu", "clamped_silu", "clamped_silu_add_one", "gelu"]) +def test_call_parameters_reach_expert_activation(monkeypatch, activation): + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul.ffn_use_tanh_approximate_gelu", + lambda: activation == "gelu", + ) + dim, count = 128, 4 + config = {"norm_topk_prob": True, "num_experts_per_tok": 2, "scoring_func": "softmax"} + kwargs = { + "clamped_silu": {"alpha": 1.0, "limit": 10.0, "clamp_up_add_one": False}, + "clamped_silu_add_one": {"alpha": 1.702, "limit": 7.0}, + }.get(activation, {}) + weight = FusedMoeWeight( + gate_proj_name="gate", + up_proj_name="up", + down_proj_name="down", + e_score_correction_bias_name="", + weight_prefix="experts", + n_routed_experts=count, + hidden_size=dim, + moe_intermediate_size=dim, + data_type=torch.bfloat16, + quant_method=NoQuantization(), + network_config=config, + ) + eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) + weights = {} + for i in range(count): + weights[f"experts.{i}.gate.weight"] = eye * (1 + i / 4) + weights[f"experts.{i}.up.weight"] = eye * (2 + i / 8) + weights[f"experts.{i}.down.weight"] = eye + weight.load_hf_weights(weights) + assert weight.verify_load() + x = torch.linspace(-25, 25, 3 * dim, device="cuda", dtype=torch.bfloat16).view(3, dim) + router = torch.tensor([[1, 3, 2, 0], [4, 3, 1, 2], [1, 2, 3, 4]], device="cuda", dtype=torch.float32) + top = router.topk(2, dim=-1) + probs = top.values.softmax(-1) + expected = torch.zeros_like(x, dtype=torch.float32) + for row in range(x.shape[0]): + for choice in range(2): + i = int(top.indices[row, choice]) + gate = (x[row] * (1 + i / 4)).float() + up = (x[row] * (2 + i / 8)).float() + if kwargs: + gate = gate.clamp(max=kwargs["limit"]) + up = up.clamp(-kwargs["limit"], kwargs["limit"]) + gate = gate * torch.sigmoid(kwargs["alpha"] * gate) + if kwargs.get("clamp_up_add_one", True): + up += 1 + elif activation == "gelu": + gate = torch.nn.functional.gelu(gate, approximate="tanh") + else: + gate = torch.nn.functional.silu(gate) + expert_out = (gate.bfloat16() * up.bfloat16()).bfloat16() + expected[row] += (expert_out.float() * probs[row, choice]).bfloat16().float() + if kwargs: + default_output = weight.experts(x.clone(), router, 2, True, False, 0, 0) + actual = weight.experts(x.clone(), router, 2, True, False, 0, 0, **kwargs) + torch.testing.assert_close(actual, expected.bfloat16(), atol=0.125, rtol=0.01) + if kwargs: + # A clamped call must not change subsequent calls on the same weight. + actual_default = weight.experts(x.clone(), router, 2, True, False, 0, 0) + torch.testing.assert_close(actual_default, default_output, atol=0, rtol=0) + + +def test_marlin_rejects_clamp_at_call(monkeypatch): + backend = FuseMoeMarlin + monkeypatch.setattr(backend, "create_workspace", lambda self: None) + monkeypatch.setattr(backend, "_select_experts", lambda *args, **kwargs: (None, None, None)) + impl = backend( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=None, + redundancy_expert_num=0, + redundancy_expert_ids_tensor=None, + routed_expert_counter_tensor=None, + auto_update_redundancy_expert=False, + ) + with pytest.raises(NotImplementedError, match="does not support clamped SwiGLU"): + impl( + input_tensor=None, + router_logits=None, + w13=None, + w2=None, + correction_bias=None, + scoring_func="softmax", + top_k=2, + renormalize=True, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + alpha=1.0, + limit=10.0, + clamp_up_add_one=False, + ) diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py new file mode 100644 index 0000000000..fb4aa62b7d --- /dev/null +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -0,0 +1,73 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not ep.HAS_DEEPGEMM or torch.cuda.get_device_capability()[0] < 9, + reason="SM90+, DeepEP and DeepGEMM required", +) + + +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) +def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, alpha, limit, add_one): + # Identity projections isolate activation/FP8 behavior from GEMM rounding. + dim, padded = 128, 128 + x = torch.linspace(-24, 24, 5 * dim, device="cuda").view(5, dim).to(torch.float8_e4m3fn) + eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) + w1 = torch.cat([2 * eye, 3 * eye]).unsqueeze(0).repeat(2, 1, 1).to(torch.float8_e4m3fn) + w2 = eye.unsqueeze(0).repeat(2, 1, 1).to(torch.float8_e4m3fn) + w1_scale = torch.ones(2, 2, 1, device="cuda") + w2_scale = torch.ones(2, 1, 1, device="cuda") + counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) + activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) + recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) + recv[0, :3] = x[:3] + recv[1, :2] = x[3:] + + if is_prefill: + metadata = torch.zeros(5, 3, device="cuda", dtype=torch.int32) + metadata[:, 2] = torch.tensor([0, 1, 2, padded, padded + 1], device="cuda") + # Force two chunks so both expert groups execute the requested activation. + monkeypatch.setattr(ep, "_get_max_chunk_rows", lambda **kwargs: padded) + actual = ep.chunked_expanded_moe_forward( + num_recv_tokens_per_expert_list=[padded, padded], + num_unaligned_recv_tokens_per_expert=counts, + recv_x=(recv.view(2 * padded, dim), torch.ones(1, 2 * padded, device="cuda").T), + recv_topk_weights=torch.ones(2 * padded, device="cuda"), + recv_src_metadata=metadata, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + block_size_k=128, + workspace=torch.empty(4 * 1024 * 1024, device="cuda", dtype=torch.uint8), + hidden_dtype=torch.bfloat16, + **activation_args, + ) + else: + output = ep.masked_group_gemm( + (recv, torch.ones(2, padded, 1, device="cuda")), + counts, + torch.bfloat16, + w1, + w1_scale, + w2, + w2_scale, + expected_m=3, + **activation_args, + ) + actual = torch.cat([output[0, :3], output[1, :2]]) + + gate = (x.float() * 2).bfloat16().float() + up = (x.float() * 3).bfloat16().float() + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(-limit, limit) + int(add_one) + gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float() + activation = (gate * up).bfloat16().float() + scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448 + expected = ((activation / scales).to(torch.float8_e4m3fn).float() * scales).bfloat16() + torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.02) diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py index 8783f35a42..d1e7f21543 100644 --- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py +++ b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py @@ -108,5 +108,34 @@ def test_silu_and_mul_masked_skips_padded_tokens(): ) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("alpha,limit,add_one", [(1.0, 10.0, False), (1.702, 7.0, True)]) +def test_clamped_swiglu_quantization_matches_reference(dtype, alpha, limit, add_one): + torch.manual_seed(53) + x = (torch.randn(3, 5, 4096, device="cuda") * 16).to(dtype) + masked_m = torch.tensor([0, 2, 5], dtype=torch.int32, device="cuda") + out = torch.full((3, 5, 2048), 1.0, dtype=torch.float8_e4m3fn, device="cuda") + scales = torch.full((3, 5, 16), 7.0, dtype=torch.float32, device="cuda") + + silu_and_mul_masked_post_quant_fwd( + x, out, scales, 128, masked_m, alpha=alpha, limit=limit, clamp_up_add_one=add_one + ) + + gate, up = x.float().chunk(2, dim=-1) + gate = gate.clamp(max=limit) + gate = (gate * torch.sigmoid(alpha * gate)).to(dtype).float() + activation = (gate * (up.clamp(-limit, limit) + int(add_one))).to(dtype).float() + groups = activation.reshape(3, 5, 16, 128) + expected_scales = groups.abs().amax(-1).clamp(min=1e-10) / 448 + expected_q = (groups / expected_scales.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) + expected = (expected_q.float() * expected_scales.unsqueeze(-1)).reshape(3, 5, 2048) + actual = out.float() * scales.repeat_interleave(128, dim=-1) + for expert, count in enumerate(masked_m.tolist()): + torch.testing.assert_close(scales[expert, :count], expected_scales[expert, :count], atol=1e-6, rtol=1e-5) + torch.testing.assert_close(actual[expert, :count], expected[expert, :count], atol=0.05, rtol=0.02) + assert torch.all(out[expert, count:].float() == 1) + assert torch.all(scales[expert, count:] == 7) + + if __name__ == "__main__": pytest.main() diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py new file mode 100644 index 0000000000..ae40702d75 --- /dev/null +++ b/unit_tests/models/glm5_next/test_cache.py @@ -0,0 +1,213 @@ +import dataclasses +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager +from lightllm.common.req_manager import Glm5NextReqManager +from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def test_nope_cache_config_uses_native_mla_width(): + config = Glm5NextCacheConfig.from_model_config( + { + "kv_lora_rank": 512, + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "index_kpool": 4, + "index_head_dim": 128, + "linear_attn_config": { + "num_heads": 8, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [0, 1, 2], + }, + }, + StartArgs(tp=4, data_type="bfloat16"), + ) + assert config.full_att_head_dim == 584 + assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 + + +@pytest.mark.parametrize("small_page", [False, True]) +@pytest.mark.parametrize("tp_world_size", [1, 4]) +@pytest.mark.parametrize("mtp_step", [0, 2]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) + args = StartArgs( + tp=tp_world_size, + data_type="bfloat16", + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + cpu_cache_token_page_size=8, + mtp_step=mtp_step, + mtp_mode="eagle_with_att" if mtp_step else None, + ) + set_unique_server_name(args) + get_env_start_args.cache_clear() + set_env_start_args(dataclasses.asdict(args)) + config = Glm5NextCacheConfig( + tp_world_size=tp_world_size, + full_att_all_num_kv_heads=1, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=584, + global_linear_k_heads=2 * tp_world_size, + global_linear_v_heads=2 * tp_world_size, + num_linear_k_heads=2, + num_linear_v_heads=2, + head_linear_k_dim=128, + head_linear_v_dim=128, + conv_kernel_size=4, + linear_layer_num=3, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4, + all_layer_num=4, + draft_full_att_kv_layer_num=int(mtp_step > 0), + ) + full_layers = config.get_full_att_kv_layer_num_with_draft_model() + mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers, config) + req = Glm5NextReqManager(3, 16, mem, config) + att_kv = mem.get_att_input_params(3) + assert att_kv.shape == (17, 1, 512) + assert att_kv.stride(0) == 584 + index_bytes = mem.get_indexer_k_buffer(3) + index_bytes.random_(0, 256) + expected_index_bytes = index_bytes.clone() + new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") + destinations = torch.tensor([5, 9], dtype=torch.int32, device="cuda") + mem.operator.copy_kv_to_mem_manager(3, destinations, new_kv) + assert torch.equal(att_kv[destinations], new_kv) + assert torch.equal(index_bytes, expected_index_bytes) + assert req.get_indexer_tail_buffer(3).shape == (4, 4 + mtp_step, 256) + if mtp_step: + assert req.get_indexer_tail_buffer(4).shape == (4, 4 + mtp_step, 256) + assert mem.get_att_input_params(4).shape == att_kv.shape + cache = req.create_small_page_cache_manager(2) if small_page else mem.big_page_buffers + slot = cache.alloc_one_state_cache() + source_req = SimpleNamespace(req_idx=0) + req.init_hybrid_attention_state(source_req) + req.req_to_conv_state.buffer[:, 0].normal_() + req.req_to_ssm_state.buffer[:, 0].normal_() + conv = req.req_to_conv_state.buffer[:, 0, ..., :3].clone() + ssm = req.req_to_ssm_state.buffer[:, 0].clone() + if small_page: + req.save_state(0, slot, cache) + else: + req.save_big_page_states(torch.tensor([0], dtype=torch.int32, device="cuda"), [0], [slot]) + torch.cuda.synchronize() + req.req_to_conv_state.buffer[:, 0].zero_() + req.req_to_ssm_state.buffer[:, 0].zero_() + req.req_to_indexer_tail.buffer.normal_() + unchanged = req.req_to_indexer_tail.buffer[:, [0, 1, 3]].clone() + dest_req = SimpleNamespace(req_idx=2, shared_kv_node=SimpleNamespace(small_page_buffer_idx=slot)) + if small_page: + req.restore_small_page_state(dest_req) + else: + req.restore_big_page_state(slot, dest_req) + torch.cuda.synchronize() + assert torch.equal(req.req_to_conv_state.buffer[:, 2, ..., :3], conv) + assert torch.equal(req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1)], ssm) + # Both page sizes are aligned to complete pools. Restoring a prefix must + # discard stale tail values from a previously allocated request slot. + assert not req.req_to_indexer_tail.buffer[:, 2].any() + assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) + req.req_to_indexer_tail.buffer[:, 2].normal_() + req.init_hybrid_attention_state(dest_req) + assert not req.req_to_conv_state.buffer[:, 2].any() + assert not req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + assert not req.req_to_indexer_tail.buffer[:, 2].any() + assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) + # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only + # belong to live requests. Bytewise equality checks FP8 scale integrity. + packed_bytes = mem.kv_buffer.view(torch.uint8) + packed_bytes[:, 0].random_(0, 256) + mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) + assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) + assert mem.get_cell_size() == 584 * 2 * full_layers + assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 + + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_kv_buffer_to_cpu_cache, + copy_cpu_cache_to_kv_buffer, + ) + + cpu_pages = torch.zeros((1, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8, pin_memory=True) + indexes = torch.arange(8, dtype=torch.int32, device="cuda") + zero = torch.zeros(1, dtype=torch.int64, device="cuda") + ready = torch.zeros(1, dtype=torch.int32, pin_memory=True) + big = mem.big_page_buffers + packed_bytes[:, :8].random_(0, 256) + big.conv_state_cache.buffer[0].normal_() + big.ssm_state_cache.buffer[0].normal_() + expected_kv = packed_bytes[:, :8].clone() + expected_conv = big.conv_state_cache.buffer[0].clone() + expected_ssm = big.ssm_state_cache.buffer[0].clone() + common = dict( + mem_indexes=indexes, + page_indexes=zero, + big_page_buffer_ids=zero, + cpu_kv_conv_state=big.conv_state_cache.buffer, + cpu_kv_ssm_state=big.ssm_state_cache.buffer, + cpu_cache_tensor=cpu_pages, + tp_world_size=tp_world_size, + big_page_token_num=8, + linear_config=config, + ) + # Simulate TP writers to a shared CPU page: MLA/index KV is replicated, + # while each rank must retain its own KDA checkpoint region. + expected_states = [] + for rank in range(tp_world_size): + big.conv_state_cache.buffer[0].copy_(expected_conv + rank) + big.ssm_state_cache.buffer[0].copy_(expected_ssm + rank) + expected_states.append((big.conv_state_cache.buffer[0].clone(), big.ssm_state_cache.buffer[0].clone())) + copy_kv_buffer_to_cpu_cache(page_readies=ready, gpu_kv_full_att_state=mem.kv_buffer, tp_rank=rank, **common) + torch.cuda.synchronize() + for rank, (conv, ssm) in enumerate(expected_states): + packed_bytes[:, :8].zero_() + big.conv_state_cache.buffer[0].zero_() + big.ssm_state_cache.buffer[0].zero_() + copy_cpu_cache_to_kv_buffer(gpu_full_att_kv_state=mem.kv_buffer, tp_rank=rank, **common) + torch.cuda.synchronize() + assert torch.equal(packed_bytes[:, :8], expected_kv) + assert torch.equal(big.conv_state_cache.buffer[0], conv) + assert torch.equal(big.ssm_state_cache.buffer[0], ssm) + + +@pytest.mark.parametrize("gate_type", ["silu", "sigmoid"]) +def test_gated_norm_gate_type_and_strided_gate(gate_type): + from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward + + x = torch.randn(12, 128, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(3, 8, 128, device="cuda", dtype=torch.bfloat16)[:, :4] + weight = torch.randn(128, device="cuda", dtype=torch.bfloat16) + actual = gated_rmsnorm_forward(x, weight, None, 1e-5, gate, gate_type=gate_type) + expected = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-5) * weight.float() + z = gate.reshape(12, 128).float() + expected *= z.sigmoid() if gate_type == "sigmoid" else torch.nn.functional.silu(z) + torch.testing.assert_close(actual, expected.bfloat16(), atol=0.008, rtol=0.008) + + +@pytest.mark.parametrize("add_one", [False, True]) +def test_clamped_swiglu_preserves_gpt_oss_default(add_one): + from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd + + x = torch.linspace(-25, 25, 2048, device="cuda", dtype=torch.bfloat16).view(4, 512) + out = torch.empty(4, 256, device="cuda", dtype=torch.bfloat16) + kwargs = {} if add_one else {"clamp_up_add_one": False} + alpha, limit = (1.702, 7.0) if add_one else (1.0, 10.0) + silu_and_mul_fwd(x, out, limit=limit, alpha=alpha, **kwargs) + gate, up = x.float().chunk(2, -1) + gate = gate.clamp(max=limit) + gate = (gate * torch.sigmoid(alpha * gate)).bfloat16().float() + expected = gate * (up.clamp(-limit, limit) + int(add_one)) + torch.testing.assert_close(out, expected.bfloat16(), atol=0.008, rtol=0.008) diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py new file mode 100644 index 0000000000..fce877448c --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.indexer import HAS_VLLM, Glm5NextNsaInfer + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture +def indexer(): + return Glm5NextNsaInfer( + 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 + ) + + +def _require_vllm_topk(): + if not HAS_VLLM: + pytest.skip("vLLM top_k_per_row_decode required") + + +def _assert_topk(logits, lengths, indices): + topk = indices.shape[1] + valid = torch.arange(topk, device=logits.device)[None, :] < lengths[:, None] + assert (indices[~valid] == -1).all() + assert (indices[valid] >= 0).all() + assert (indices < lengths[:, None]).all() + for row, length in zip(indices, lengths.tolist()): + count = min(length, topk) + assert row[:count].unique().numel() == count + + got = logits.gather(1, indices.long().clamp_min(0)).masked_fill(~valid, -float("inf")) + positions = torch.arange(logits.shape[1], device=logits.device) + masked = logits.masked_fill(positions[None, :] >= lengths[:, None], -float("inf")) + # Compare values, since different equal-valued candidates are valid top-k. + torch.testing.assert_close(got.sort(descending=True).values, masked.topk(topk).values, rtol=0, atol=0) + + +@pytest.mark.parametrize("pools", [640, 8192, 65536, 262144]) +@pytest.mark.parametrize("distribution", ["normal", "concentrated", "ties"]) +def test_topk_variable_lengths(indexer, pools, distribution): + _require_vllm_topk() + torch.manual_seed(42) + # DeepGEMM logits can have a padded row stride, and output is a query chunk. + logits = torch.randn(8, pools + 128, device="cuda")[:, :pools] + if distribution == "concentrated": + logits.mul_(0.001).add_(1) + elif distribution == "ties": + logits.copy_(torch.randint(0, 4, logits.shape, device="cuda")) + lengths = torch.tensor([0, 1, 511, 512, 513, pools // 3, pools - 3, pools], device="cuda", dtype=torch.int32) + positions = torch.arange(pools, device="cuda") + logits.masked_fill_(positions[None, :] >= lengths[:, None], float("nan")) + output = torch.full((10, 512), -2, device="cuda", dtype=torch.int32) + indices = output[1:-1] + indexer.select_topk_indices(logits, lengths, indices) + _assert_topk(logits, lengths, indices) + assert (output[[0, -1]] == -2).all() + + +@pytest.mark.parametrize("pools", [8192, 262144]) +def test_topk_cuda_graph_with_changing_lengths(indexer, pools): + _require_vllm_topk() + logits = torch.zeros(8, pools, device="cuda") + lengths = torch.zeros(8, device="cuda", dtype=torch.int32) + indices = torch.empty(8, 512, device="cuda", dtype=torch.int32) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + indexer.select_topk_indices(logits, lengths, indices) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + indexer.select_topk_indices(logits, lengths, indices) + + for shift in range(3): + logits.normal_().mul_(0.001).add_(1) + lengths.copy_( + torch.tensor([0, 1, 511, 512, 513, pools // 3, pools - 3, pools], device="cuda", dtype=torch.int32).roll( + shift + ) + ) + logits.masked_fill_(torch.arange(pools, device="cuda")[None, :] >= lengths[:, None], float("nan")) + graph.replay() + _assert_topk(logits, lengths, indices) + + +def test_topk_without_vllm(indexer, monkeypatch): + from lightllm.models.glm5_next import indexer as indexer_module + + monkeypatch.setattr(indexer_module, "HAS_VLLM", False) + logits = torch.randn(5, 640, device="cuda") + lengths = torch.tensor([0, 1, 511, 512, 639], device="cuda", dtype=torch.int32) + indices = torch.empty(5, 512, device="cuda", dtype=torch.int32) + indexer.select_topk_indices(logits, lengths, indices) + _assert_topk(logits, lengths, indices) + + +@pytest.mark.parametrize("mtp_size", [1, 3, 6]) +@pytest.mark.parametrize("max_pools", [1024, 262144]) +def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, max_pools): + pytest.importorskip("deep_gemm") + _require_vllm_topk() + torch.manual_seed(27) + rows = 5 * mtp_size + storage = torch.empty(1024, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + packed[:, 0, :128] = torch.randn(1024, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + scales = torch.pow(2.0, torch.randint(-3, 2, (1024,), device="cuda").float()) + packed[:, 0, 128:] = scales.view(torch.uint8).view(-1, 4) + table = torch.full((5, max_pools * 4), -1, device="cuda", dtype=torch.int32) + table[:, 3::4] = torch.randint(0, 1024, (5, max_pools), device="cuda", dtype=torch.int32) + req_idx = torch.tensor([1, 3, 0, 2, 4], device="cuda", dtype=torch.int32).repeat_interleave(mtp_size) + lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) + q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn) + weights = torch.randn(rows, 32, device="cuda") * 0.05 + infer_state = SimpleNamespace( + b_req_idx=req_idx, + b_seq_len=lengths, + b1_cu_q_seq_len=torch.arange(rows + 1, device="cuda", dtype=torch.int32), + max_q_seq_len=1, + req_manager=SimpleNamespace(req_to_token_indexs=table), + ) + att_state = SimpleNamespace(lengths=lengths) + logits_outputs = [] + select_topk = indexer.select_topk_indices + + def capture_logits(logits, pool_lengths, indices): + logits_outputs.append(logits) + select_topk(logits, pool_lengths, indices) + + monkeypatch.setattr(indexer, "select_topk_indices", capture_logits) + + def run(): + return indexer._get_decode_indices(q, weights, packed, infer_state, att_state, max_pools) + + # Match startup capture with only empty HOLD rows, then replay real work. + run() + logits_outputs.clear() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + graph_logits = logits_outputs.pop() + for bases in ([0, 2047, 2051, max_pools * 4 - mtp_size + 1, 0], [7, 65, 511, 2053, 0]): + seqs = torch.tensor(bases, device="cuda", dtype=torch.int32)[:, None] + seqs = seqs + torch.arange(mtp_size, device="cuda", dtype=torch.int32)[None, :] + seqs[-1].zero_() + lengths.copy_(seqs.flatten()) + graph.replay() + indexer._get_prefill_indices(q, weights, packed, infer_state, att_state, max_pools) + reference_logits = logits_outputs.pop() + pool_lengths = lengths // 4 + valid = torch.arange(max_pools, device="cuda")[None, :] < pool_lengths[:, None] + torch.testing.assert_close(graph_logits[valid], reference_logits[valid], rtol=0, atol=0) + # Fragmented tables deliberately repeat keys; equal-score top-k ties + # may choose different indices. Check exact selected values and uniqueness. + _assert_topk(reference_logits, pool_lengths, actual) + assert (actual[-mtp_size:] == -1).all() + # Exercise request-slot reuse and speculative rollback on the same graph. + req_idx.copy_(req_idx.roll(mtp_size)) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py new file mode 100644 index 0000000000..3270ba972e --- /dev/null +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -0,0 +1,643 @@ +import dataclasses +import math +from types import SimpleNamespace + +import pytest +import torch +import triton + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( + chunk_kda_with_fused_gate, + fused_kda_gate_chunk_cumsum, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.models.glm5_next.triton_kernel.kpool import ( + compress_pools, + gather_pools, + gather_paged_pools, + get_pool_ranges, + expand_topk, +) +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def setup(): + torch.manual_seed(1525) + set_env_start_args(dataclasses.asdict(StartArgs())) + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + + +def _reference_kda(q, k, v, gate, beta, a, bias, state): + q = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) / q.shape[-1] ** 0.5 + k = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) + decay = (-5 * torch.sigmoid(a.exp()[:, None] * (gate.float() + bias))).exp() + state = state * decay[..., None] + delta = (v.float() - torch.einsum("hkv,hk->hv", state, k)) * beta.float().sigmoid()[:, None] + state = state + k[..., None] * delta[:, None, :] + return torch.einsum("hkv,hk->hv", state, q), state + + +@pytest.mark.parametrize("tokens", [1, 3, 65, 129]) +def test_kda_chunk_and_decode_match_recurrence(tokens): + heads, dim = 2, 128 + rand = lambda *shape: torch.randn(shape, device="cuda", dtype=torch.bfloat16) + q, k, v, gate = [rand(1, tokens, heads, dim) for _ in range(4)] + beta = rand(1, tokens, heads) + a = torch.randn(heads, device="cuda") + bias = torch.randn(heads, dim, device="cuda") + initial = torch.randn(1, heads, dim, dim, device="cuda") * 0.1 + expected = [] + state = initial[0].clone() + for i in range(tokens): + out, state = _reference_kda(q[0, i], k[0, i], v[0, i], gate[0, i], beta[0, i], a, bias, state) + expected.append(out) + expected = torch.stack(expected).unsqueeze(0) + actual, final = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v.clone(), + raw_g=gate, + beta=beta.float().sigmoid(), + A_log=a, + g_bias=bias.flatten(), + initial_state=initial, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=torch.tensor([0, tokens], dtype=torch.int32, device="cuda"), + safe_gate=True, + ) + torch.testing.assert_close(actual.float(), expected, atol=4e-3, rtol=3e-2) + torch.testing.assert_close(final[0], state, atol=8e-3, rtol=3e-2) + # Reuse a nonzero request slot; the neighboring requests must stay intact. + states = torch.randn(4, heads, dim, dim, device="cuda") + states[2] = initial[0] + unchanged = states[[0, 1, 3]].clone() + for i in range(tokens): + out, _ = fused_recurrent_kda( + q[:, i : i + 1], + k[:, i : i + 1], + v[:, i : i + 1], + gate[:, i : i + 1].reshape(1, 1, -1), + beta[:, i : i + 1], + a, + bias.flatten(), + states, + torch.tensor([2], device="cuda", dtype=torch.int32), + ) + torch.testing.assert_close(out[0, 0].float(), expected[0, i], atol=2e-3, rtol=1e-2) + torch.testing.assert_close(states[2], state, atol=2e-5, rtol=2e-4) + assert torch.equal(states[[0, 1, 3]], unchanged) + + +@pytest.mark.parametrize("seq_lens", [(65,), (3, 65, 129)]) +@pytest.mark.parametrize("safe_gate", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +def test_kda_gate_cumsum_packed_shape(seq_lens, safe_gate, strided): + heads, dim, chunk_size = 2, 128, 64 + if strided: + # All three input strides differ from contiguous [T, H, D]; empty_like + # allocates a contiguous output for this sliced, non-dense input view. + storage = torch.randn(sum(seq_lens) * 2, heads * 2, dim * 2, device="cuda", dtype=torch.bfloat16) + raw_g = storage[::2, ::2, ::2] + else: + raw_g = torch.randn(sum(seq_lens), heads, dim, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(heads, device="cuda") + bias = torch.randn(heads, dim, device="cuda") if len(seq_lens) > 1 else None + cu_seqlens = torch.tensor([0, *seq_lens], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + lower_bound = -3.0 + + actual = fused_kda_gate_chunk_cumsum( + raw_g, + A_log=a_log, + g_bias=bias.flatten() if bias is not None else None, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + + gate_input = raw_g.float() + (bias if bias is not None else 0) + amplitude = a_log.exp()[None, :, None] + if safe_gate: + log_gate = lower_bound * torch.sigmoid(amplitude * gate_input) + else: + log_gate = -amplitude * torch.nn.functional.softplus(gate_input) + expected = torch.empty_like(log_gate) + seq_start = 0 + for seq_len in seq_lens: + for offset in range(0, seq_len, chunk_size): + start = seq_start + offset + end = seq_start + min(offset + chunk_size, seq_len) + expected[start:end] = log_gate[start:end].cumsum(0) / math.log(2) + seq_start += seq_len + + assert actual.shape == raw_g.shape + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) + + +@pytest.mark.parametrize( + "q_lens", + [(3, 2, 4), (0, 1, 0, 1, 0), (1,) * 64, (257, 1, 513, 0), (8129,) + (1,) * 63], +) +def test_kpool_ranges_cuda_graph_with_changing_query_boundaries(q_lens): + # Move the longest query across the batch while replaying the same graph. + # Empty queries must not use the one-query-per-request decode shortcut. + tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + lengths = torch.empty(sum(q_lens), device="cuda", dtype=torch.int32) + cu_q_lens = torch.empty(len(q_lens) + 1, device="cuda", dtype=torch.int32) + max_pools = 4096 + + def set_inputs(counts, phase): + cu, visible_lengths, pool_starts = [0], [], [] + for batch, count in enumerate(counts): + prefix = batch * 7 + phase + visible_lengths.extend(range(prefix + 1, prefix + count + 1)) + pool_starts.extend([batch * max_pools] * count) + cu.append(cu[-1] + count) + lengths.copy_(tensor(visible_lengths)) + cu_q_lens.copy_(tensor(cu)) + expected_lengths = tensor([length // 4 for length in visible_lengths]) + expected_starts = tensor(pool_starts) + return expected_starts, expected_starts + expected_lengths, expected_lengths + + expected = set_inputs(q_lens, 0) + actual = get_pool_ranges(lengths, cu_q_lens, max(q_lens), max_pools) + for output, reference in zip(actual, expected): + assert torch.equal(output, reference) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = get_pool_ranges(lengths, cu_q_lens, max(q_lens), max_pools) + expected = set_inputs(q_lens[::-1], 3) + graph.replay() + for output, reference in zip(actual, expected): + assert torch.equal(output, reference) + + +@pytest.mark.parametrize("cuda_graph", [False, True]) +def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): + # Unaligned chunks, reordered requests, and pools spanning old/new tails. + seqs = [9, 7] + ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)] + req_idx = torch.tensor([1, 3], device="cuda", dtype=torch.int32) + table = torch.zeros(5, 12, device="cuda", dtype=torch.int32) + table[1, :9], table[3, :7] = ragged[:9], ragged[9:] + source = torch.randn(sum(seqs), 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda") + packed_storage = torch.zeros(40, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = packed_storage.view(torch.uint8)[:, :, -132:] + tail = torch.randn(5, 4, 256, device="cuda", dtype=torch.bfloat16) + unchanged = tail[[0, 2, 4]].clone() + tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) + for chunks in [[(1, 0, 3), (3, 0, 2)], [(3, 2, 3), (1, 3, 9)], [(3, 3, 7)]]: + raw, lengths, starts, locations = [], [], [], [] + cu_q_lens = [0] + for req, first, end in chunks: + offset = 0 if req == 1 else 9 + raw.append(source[offset + first : offset + end]) + lengths.extend(range(first + 1, end + 1)) + starts.extend([len(locations)] * (end - first)) + locations.extend(table[req, :end].tolist()) + cu_q_lens.append(cu_q_lens[-1] + end - first) + args = ( + torch.cat(raw), + tail, + packed, + ape, + tensor(lengths), + tensor(starts), + tensor(locations), + tensor([r for r, _, _ in chunks]), + tensor(cu_q_lens), + tensor([end for _, _, end in chunks]), + ) + max_q_len = max(end - first for _, first, end in chunks) + if cuda_graph: + # Compile before capture, then restore this chunk's input cache state. + saved_tail, saved_packed = tail.clone(), packed.clone() + compress_pools(*args, max_q_len=max_q_len) + tail.copy_(saved_tail) + packed.copy_(saved_packed) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + compress_pools(*args, max_q_len=max_q_len) + graph.replay() + else: + compress_pools(*args, max_q_len=max_q_len) + for req, _, end in chunks: + offset = 0 if req == 1 else 9 + assert torch.equal(tail[req, : end % 4], source[offset + end // 4 * 4 : offset + end]) + assert torch.equal(tail[[0, 2, 4]], unchanged) + # Model the full KV copy performed by cache offload/load or request move. + packed_storage = packed_storage.clone() + packed = packed_storage.view(torch.uint8)[:, :, -132:] + keys, scales = gather_pools(packed, table, req_idx, torch.tensor(seqs, device="cuda", dtype=torch.int32), 3) + for batch, start in enumerate([0, 9]): + for group in range(seqs[batch] // 4): + values = source[start + group * 4 : start + group * 4 + 4] + expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) + torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) + torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) + lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) + starts = torch.tensor([0, 9], device="cuda", dtype=torch.int32) + groups = torch.tensor([[1, 0], [0, -1]], device="cuda", dtype=torch.int32) + out, relative = expand_topk(groups, lengths, starts, ragged, topk=8) + for batch, expected in enumerate([[4, 5, 6, 7, 0, 1, 2, 3, 8], list(range(7))]): + assert relative[batch, : len(expected)].tolist() == expected + assert ( + out[batch, : len(expected)].tolist() + == ragged[starts[batch].item() + torch.tensor(expected, device="cuda")].tolist() + ) + assert (out[batch, len(expected) :] == -1).all() + + +def test_kpool_decode_cuda_graph_and_padding(): + # Replay one graph across all four tail phases, with reordered requests + # and repeated padding slots. Padding must not race on the hold state. + source = torch.randn(2, 8, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda") + tail = torch.zeros(5, 4, 256, device="cuda", dtype=torch.bfloat16) + storage = torch.zeros(32, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + table = torch.randperm(32, device="cuda", dtype=torch.int32).view(4, 8) + raw = torch.zeros(4, 256, device="cuda", dtype=torch.bfloat16) + req_idx = torch.tensor([1, 3, 4, 4], device="cuda", dtype=torch.int32) + lengths = torch.ones(4, device="cuda", dtype=torch.int32) + starts = torch.arange(4, device="cuda", dtype=torch.int32) * 8 + ragged = table.flatten().clone() + cu_q_lens = torch.arange(5, device="cuda", dtype=torch.int32) + + def run(): + compress_pools(raw, tail, packed, ape, lengths, starts, ragged, req_idx, cu_q_lens, lengths, max_q_len=1) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + tail.zero_() + for pos in range(8): + order = [pos % 2, 1 - pos % 2] + req_idx[:2] = torch.tensor([1 if i == 0 else 3 for i in order], device="cuda") + raw[:2] = source[order, pos] + lengths.fill_(pos + 1) + ragged[:16] = table[order].flatten() + graph.replay() + for batch, req in enumerate([1, 3]): + pending = (pos + 1) % 4 + assert torch.equal(tail[req, :pending], source[batch, (pos + 1) // 4 * 4 : pos + 1]) + if pending == 0: + values = source[batch, pos - 3 : pos + 1] + pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() + key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) + actual = packed[table[batch, pos].long(), 0] + assert torch.equal(actual[:128], key.view(torch.uint8)[0]) + assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) + assert not tail[[0, 2, 4]].any() + + +@pytest.mark.parametrize("max_pools", [640, 262144]) +def test_gather_paged_pools_valid_pages_and_graph_replay(max_pools): + storage = torch.zeros(64, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + source_keys = torch.randn(64, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + source_scales = torch.rand(64, device="cuda") + 0.1 + packed[:, 0, :128] = source_keys + packed[:, 0, 128:] = source_scales.view(torch.uint8).view(64, 4) + table = torch.full((4, max_pools * 4), -1, device="cuda", dtype=torch.int32) + locations = torch.randint(0, 64, (4, max_pools), device="cuda", dtype=torch.int32) + table[:, 3::4] = locations + req_idx = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + lengths = torch.tensor([max_pools, 3, 0], device="cuda", dtype=torch.int32) + + def check(pages, block_table): + page_bytes = pages.view(pages.shape[0], -1) + for row, (req, length) in enumerate(zip(req_idx.tolist(), lengths.tolist())): + page_ids = block_table[row, : triton.cdiv(length, 64)].long() + keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128) + scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten() + locs = locations[req, :length].long() + assert torch.equal(keys[:length], source_keys[locs]) + assert torch.equal(scales[:length], source_scales[locs]) + assert not keys[length:].any() + assert (scales[length:] == 1).all() + + check(*gather_paged_pools(packed, table, req_idx, lengths, max_pools)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pages, block_table = gather_paged_pools(packed, table, req_idx, lengths, max_pools) + # Empty HOLD rows must do no page writes, even with a 1M graph capacity. + pages.fill_(127) + lengths.zero_() + graph.replay() + assert (pages == 127).all() + for counts in ([513, 65, max_pools], [1, 0, 3]): + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + req_idx.copy_(torch.tensor([1, 3, 0], device="cuda", dtype=torch.int32)) + graph.replay() + check(pages, block_table) + + +@pytest.mark.parametrize("max_pools", [65535, 65536, 262144]) +def test_gather_pools_long_context_cuda_graph(max_pools): + # 1M-token graph warmup pads to 262144 pools even for empty hold requests. + # Also gather real pools at the end of that range from fragmented token KV. + storage = torch.zeros(64, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + source_keys = torch.randn(64, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + source_scales = torch.rand(64, device="cuda") + 0.1 + packed[:, 0, :128] = source_keys + packed[:, 0, 128:] = source_scales.view(torch.uint8).view(64, 4) + table = torch.full((4, max_pools * 4), -1, device="cuda", dtype=torch.int32) + locations = torch.randint(0, 64, (4, max_pools), device="cuda", dtype=torch.int32) + table[:, 3::4] = locations + req_idx = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + seq_len = torch.tensor([max_pools * 4, 14, 2], device="cuda", dtype=torch.int32) + + def check(keys, scales): + keys = keys.view(torch.uint8).view(3, max_pools, 128) + scales = scales.view(3, max_pools) + for batch, (req, count) in enumerate([(2, max_pools), (0, 3), (3, 0)]): + locs = locations[req, :count].long() + assert torch.equal(keys[batch, :count], source_keys[locs]) + assert torch.equal(scales[batch, :count], source_scales[locs]) + assert not keys[batch, count:].any() + assert (scales[batch, count:] == 1).all() + + check(*gather_pools(packed, table, req_idx, seq_len, max_pools)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + keys, scales = gather_pools(packed, table, req_idx, seq_len, max_pools) + # Replay with the same short hold lengths used by model startup. + original_lengths = seq_len.clone() + seq_len.fill_(2) + graph.replay() + assert not keys.view(torch.uint8).any() + assert (scales == 1).all() + seq_len.copy_(original_lengths) + graph.replay() + check(keys, scales) + + +def test_kpool_packed_kv_above_two_gib(): + # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. + loc = 2 ** 31 // (584 * 2) + 1 + storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) + tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda") + lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32) + zero = torch.zeros(4, device="cuda", dtype=torch.int32) + ragged = torch.tensor([0, 1, 2, loc], device="cuda", dtype=torch.int32) + cu_q_lens = torch.tensor([0, 4], device="cuda", dtype=torch.int32) + compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) + keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) + pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) + assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) + assert torch.equal(scales, expected_scale.flatten()) + + +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): + from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + + hidden = 4096 + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_world_size)) + pre_infer = Glm5NextTpPartModel.pre_layer_infer_class({"hc_mult": 4}) + embeddings = torch.randn(1, hidden, device="cuda", dtype=torch.bfloat16) + input_ids = torch.zeros(1, device="cuda", dtype=torch.long) + pre_weight = SimpleNamespace(wte_weight_=lambda input_ids, alloc_func: embeddings[input_ids]) + pre_weight.wte_weight_.weight = embeddings + pre_weight.wte_weight_.tp_vocab_start_id = 0 + pre_weight.wte_weight_.tp_vocab_end_id = 1 + infer_state = SimpleNamespace(dist_group=None, multimodal_params=[]) + + def all_reduce(input_embeddings, **kwargs): + # TP communication must operate on the original embedding width. + assert input_embeddings.shape == (1, hidden) + input_embeddings.mul_(tp_world_size) + + monkeypatch.setattr("lightllm.models.llama.layer_infer.pre_layer_infer.all_reduce", all_reduce) + monkeypatch.setattr("lightllm.models.qwen_vl.layer_infer.pre_layer_infer.all_reduce", all_reduce) + expected_streams = (embeddings * tp_world_size).unsqueeze(1).expand(-1, 4, -1) + weight = SimpleNamespace( + att_norm_weight_=SimpleNamespace(weight=torch.ones(hidden, device="cuda", dtype=torch.bfloat16)), + ffn_norm_weight_=SimpleNamespace(weight=torch.ones(hidden, device="cuda", dtype=torch.bfloat16)), + ) + for prefix in ("attn", "ffn"): + setattr(weight, f"hc_{prefix}_fn", SimpleNamespace(weight=torch.randn(24, 4 * hidden, device="cuda") * 0.005)) + setattr(weight, f"hc_{prefix}_base", SimpleNamespace(weight=torch.zeros(24, device="cuda"))) + setattr(weight, f"hc_{prefix}_scale", SimpleNamespace(weight=torch.ones(3, device="cuda"))) + layer = object.__new__(Glm5NextTransformerLayerInfer) + layer.use_mhc = True + layer.embed_dim_, layer.mhc_streams = hidden, 4 + layer.num_hidden_layers, layer.autotune_layer_num = 5, 4 + layer.eps_, layer.hc_eps, layer.hc_sinkhorn_iters = 1e-5, 1e-6, 20 + layer.token_attention_forward = layer.context_attention_forward = layer._ffn = lambda x, *_: x * 0.1 + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + x = pre_infer.token_forward(input_ids, infer_state, pre_weight) + torch.testing.assert_close(x.view(1, 4, hidden), expected_streams, atol=0, rtol=0) + for i in range(5): + layer.layer_num_ = i + x = layer.token_forward(x, None, weight) + assert x.shape == (1, hidden if i == 4 else 4 * hidden) + with Autotuner.autotune_warmup(): + x = pre_infer.context_forward(input_ids, infer_state, pre_weight) + torch.testing.assert_close(x.view(1, 4, hidden), expected_streams, atol=0, rtol=0) + for i in range(4): + layer.layer_num_ = i + x = layer.context_forward(x, None, weight) + assert x.shape == (1, hidden if i == 3 else 4 * hidden) + + +def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + + tokens, heads, dim = 67, 2, 128 + q, k, v = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + q = torch.nn.functional.normalize(q.float(), dim=-1).bfloat16() + k = torch.nn.functional.normalize(k.float(), dim=-1).bfloat16() + gate = -torch.rand(1, tokens, heads, device="cuda") + beta = torch.rand_like(gate) + state = torch.randn(1, heads, dim, dim, device="cuda") * 0.1 + initial = state.clone() + expected = [] + for i in range(tokens): + state *= gate[:, i, :, None, None].exp() + delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] + state += k[:, i, :, :, None].float() * delta[:, :, None, :] + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) + out, final = chunk_gated_delta_rule( + q, + k, + v, + gate, + beta, + initial_state=initial, + output_final_state=True, + cu_seqlens=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), + ) + torch.testing.assert_close(out.float(), torch.stack(expected, 1), atol=4e-3, rtol=3e-2) + torch.testing.assert_close(final, state, atol=8e-3, rtol=3e-2) + + +@pytest.mark.parametrize("heads", [16, 64, 128]) +def test_nope_attention_native_512_and_cuda_graph(heads): + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("FA3 NoPE requires Hopper") + # The 72-element tail stores indexer data, not a zero RoPE embedding. + packed = torch.randn(4096, 1, 584, dtype=torch.bfloat16, device="cuda") + kv = packed[:, :, :512] + # Match the head-major projection's non-contiguous query layout. + q = torch.randn(heads, 9, 512, dtype=torch.bfloat16, device="cuda").transpose(0, 1) + lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], dtype=torch.int32, device="cuda") + indexes = torch.randint(0, 4096, (9, 2176), dtype=torch.int32, device="cuda") + offsets = torch.arange(indexes.shape[1], device="cuda") + indexes.masked_fill_(offsets[None, :] >= lengths[:, None], -1) + + def reference(): + expected = [] + for i, length in enumerate(lengths.tolist()): + keys = kv[indexes[i, :length].long(), 0].float() + expected.append((q[i].float() @ keys.T * 0.0625).softmax(-1) @ keys) + return torch.stack(expected) + + expected = reference() + control = AttControl( + nsa_prefill_dict={"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} + ) + prefill_state = Glm5NextSparsePrefillState() + prefill = prefill_state._nsa_prefill_att(q, kv, control) + torch.testing.assert_close(prefill.float(), expected, atol=0.012, rtol=0.015) + decode = Glm5NextSparseDecodeState( + infer_state=SimpleNamespace( + b1_cu_q_seq_len=torch.arange(10, dtype=torch.int32, device="cuda"), max_q_seq_len=1 + ), + nsa_cache_seqlens=lengths, + nsa_cu_seqlens_k_new=torch.nn.functional.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), + ) + control.nsa_decode_dict = control.nsa_prefill_dict + out = decode._nsa_decode_att((q, q[..., :0]), kv, control) + torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + prefill_state._nsa_prefill_att(q, kv, control) + decode._nsa_decode_att((q, q[..., :0]), kv, control) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_prefill = prefill_state._nsa_prefill_att(q, kv, control) + graph_decode = decode._nsa_decode_att((q, q[..., :0]), kv, control) + graph.replay() + torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) + torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) + + q.normal_() + packed.normal_() + lengths.copy_(lengths.roll(1)) + decode.nsa_cu_seqlens_k_new[1:].copy_(lengths.cumsum(0, dtype=torch.int32)) + indexes.random_(0, 4096) + indexes.masked_fill_(offsets[None, :] >= lengths[:, None], -1) + graph.replay() + expected = reference() + torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) + torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) + + +@pytest.mark.parametrize("max_kv_seq_len", [2065, 1 << 20]) +def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): + from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer + + tokens, heads, dim = 2065, 32, 128 + hidden = torch.randn(tokens, dim, device="cuda", dtype=torch.bfloat16) + q_weight = torch.randn(dim, heads * dim, device="cuda", dtype=torch.bfloat16) * 0.1 + ape = torch.randn(4, dim, device="cuda") + weights = SimpleNamespace( + wk_proj_=SimpleNamespace(mm=lambda x: x), + k_norm_=lambda x, eps: x, + index_kpool_compress_gate=SimpleNamespace(mm=lambda x: x * 0.25), + index_kpool_compress_ape=SimpleNamespace(weight=ape), + wq_b_proj_=SimpleNamespace(mm=lambda x: x @ q_weight), + weights_proj_=SimpleNamespace(mm=lambda x: torch.ones(x.shape[0], heads, device=x.device)), + ) + storage = torch.zeros(tokens + 9, 1, 584, device="cuda", dtype=torch.bfloat16) + tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) + ragged = torch.randperm(tokens + 9, device="cuda", dtype=torch.int32)[:tokens] + manager = SimpleNamespace( + get_indexer_k_buffer=lambda _: storage.view(torch.uint8)[:, :, -132:], + ) + infer = SimpleNamespace( + mem_manager=manager, + is_prefill=True, + b_mtp_index=torch.zeros(1, device="cuda", dtype=torch.int32), + mem_index=ragged, + # A large capacity also exercises query chunking and uninitialized + # logits beyond each query's actual pool range. + max_kv_seq_len=max_kv_seq_len, + req_manager=SimpleNamespace(req_to_token_indexs=ragged[None], get_indexer_tail_buffer=lambda _: tail), + b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([tokens], device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), + max_q_seq_len=tokens, + ) + state = SimpleNamespace( + lengths=torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32), + ks=torch.zeros(tokens, device="cuda", dtype=torch.int32), + ragged_mem_index=ragged, + ) + indexer = Glm5NextNsaInfer( + 0, {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5}, 1 + ) + _, full = indexer._get_indices(hidden, hidden, infer, state, weights) + assert full[-1, 2048].item() == tokens - 1 # Always-selected incomplete tail. + assert full[0, 0].item() == 0 and (full[0, 1:] == -1).all() + assert full[-1, :2049].unique().numel() == 2049 + assert (full[-1, :2049] < tokens).all() + # Restore an aligned prefix using only token KV, then decode across the + # next pool boundary. The runtime tail starts empty, as on a page restore. + prefix = 2060 + storage = storage.clone() + storage[ragged[prefix:].long()] = 0 + tail.zero_() + infer.is_prefill = False + infer.b1_cu_q_seq_len[1] = 1 + infer.max_q_seq_len = 1 + for pos in range(prefix, tokens): + infer.mem_index = ragged[pos : pos + 1] + infer.max_kv_seq_len = pos + 1 + infer.b_seq_len.fill_(pos + 1) + state.lengths = torch.tensor([pos + 1], device="cuda", dtype=torch.int32) + state.ks = torch.zeros_like(state.lengths) + _, decoded = indexer._get_indices(hidden[pos : pos + 1], hidden[pos : pos + 1], infer, state, weights) + valid = 2048 + (pos + 1) % 4 + assert set(decoded[0, :valid].tolist()) == set(full[pos, :valid].tolist()) + + pool_count = tokens // 4 + pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) + pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) + query = (hidden[-1:] @ q_weight).view(heads, dim) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) + logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) + scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) + expected_groups = set(scores.topk(512).indices.tolist()) + assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py new file mode 100644 index 0000000000..63c2f9e886 --- /dev/null +++ b/unit_tests/models/glm5_next/test_layers.py @@ -0,0 +1,406 @@ +import dataclasses +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState +from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for name, value in { + "GLOBAL_RANK": 0, + "GLOBAL_WORLD_SIZE": 1, + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_RANK_IN_NODE": 0, + "CURRENT_DEVICE_ID": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + name, str(value)) + get_env_start_args.cache_clear() + torch.manual_seed(53) + yield + get_env_start_args.cache_clear() + + +@pytest.fixture +def config(): + return { + "hidden_size": 128, + "n_embed": 128, + "intermediate_size": 128, + "moe_intermediate_size": 128, + "num_hidden_layers": 4, + "n_layer": 4, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "vocab_size": 256, + "rms_norm_eps": 1e-5, + "q_lora_rank": 128, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 0, + "v_head_dim": 128, + "index_n_heads": 4, + "index_head_dim": 128, + "index_topk": 2048, + "index_kpool": 4, + "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4}, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "mhc": False, + "n_routed_experts": 4, + "n_shared_experts": 1, + "first_k_dense_replace": 1, + "num_experts_per_tok": 2, + "norm_topk_prob": True, + "scoring_func": "sigmoid", + "n_group": 1, + "topk_group": 1, + "swiglu_limit": 10.0, + "num_nextn_predict_layers": 1, + } + + +def random_weight(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.05 + + +def rmsnorm(x, eps): + return (x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps)).to(x.dtype) + + +def test_kda_projections_need_no_sparse_attention_config(config): + for key in ("q_lora_rank", "kv_lora_rank", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"): + del config[key] + for key in list(config): + if key.startswith("index_"): + del config[key] + layer = Glm5NextTransformerLayerInfer(0, config) + weight = Glm5NextTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + projections = { + name: random_weight(size, 128) + for name, size in {"q": 512, "k": 512, "v": 512, "b": 4, "f_a": 128, "g_a": 128, "f_b": 512, "g_b": 512}.items() + } + weight.load_hf_weights( + {f"model.language_model.layers.0.self_attn.{name}_proj.weight": value for name, value in projections.items()} + ) + hidden = random_weight(3, 128) + actual = layer._kda_projections(hidden, SimpleNamespace(), weight) + expected = ( + torch.cat([F.linear(hidden, projections[name]) for name in ("q", "k", "v")], dim=-1), + F.linear(F.linear(hidden, projections["f_a"]), projections["f_b"]), + F.linear(hidden, projections["b"]), + F.linear(F.linear(hidden, projections["g_a"]), projections["g_b"]), + ) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, atol=1e-3, rtol=0.01) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_mla_loading_preserves_replicated_qkv_and_sharded_bmm(config, monkeypatch, rank): + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + layer = Glm5NextTransformerLayerInfer(3, config) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + q_a, kv_a = random_weight(128, 128), random_weight(512, 128) + q_b, kv_b, o = random_weight(512, 128), random_weight(1024, 512), random_weight(128, 512) + tensors = { + "q_a_proj.weight": q_a, + "kv_a_proj_with_mqa.weight": kv_a, + "q_b_proj.weight": q_b, + "kv_b_proj.weight": kv_b, + "o_proj.weight": o, + "q_a_layernorm.weight": torch.ones(128, device="cuda", dtype=torch.bfloat16), + "kv_a_layernorm.weight": torch.ones(512, device="cuda", dtype=torch.bfloat16), + } + weight.load_hf_weights( + {f"model.language_model.layers.3.self_attn.{name}": value for name, value in tensors.items()} + ) + hidden = random_weight(3, 128) + state = SimpleNamespace(need_dp_prefill_balance=False) + q, kv = layer._get_qkv(hidden, state, weight) + q_lora = rmsnorm(F.linear(hidden, q_a), config["rms_norm_eps"]) + expected_q = F.linear(q_lora, q_b).view(3, 4, 128)[:, rank * 2 : (rank + 1) * 2] + expected_kv = rmsnorm(F.linear(hidden, kv_a), config["rms_norm_eps"]).unsqueeze(1) + torch.testing.assert_close(q, expected_q, atol=0.015, rtol=0.015) + torch.testing.assert_close(kv, expected_kv, atol=0.015, rtol=0.015) + head_weights = kv_b.view(4, 256, 512)[rank * 2 : (rank + 1) * 2] + projected = weight.k_b_proj_.bmm(q.transpose(0, 1)) + torch.testing.assert_close(projected, torch.bmm(q.transpose(0, 1), head_weights[:, :128])) + latent_output = random_weight(3, 2, 512) + monkeypatch.setattr(layer, "_tpsp_reduce", lambda input, infer_state: input) + actual = layer._get_o(latent_output, state, weight) + values = torch.bmm(latent_output.transpose(0, 1), head_weights[:, 128:].transpose(1, 2)) + expected = F.linear(values.transpose(0, 1).reshape(3, 256), o[:, rank * 256 : (rank + 1) * 256]) + torch.testing.assert_close(actual, expected) + + +def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("GLM sparse attention kernels require Hopper") + layer = Glm5NextTransformerLayerInfer(3, config) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + kv_b = random_weight(1024, 512) + weight.load_hf_weights({"model.layers.3.self_attn.kv_b_proj.weight": kv_b}) + q = random_weight(3, 4, 128) + # Use noncontiguous token slots and a padded sparse index table. + packed = random_weight(8, 1, 584) + kv = packed[:, :, :512] + slots = torch.tensor([7, 1, 5, 0, 3], device="cuda", dtype=torch.int32) + lengths = torch.tensor([2, 3, 5], device="cuda", dtype=torch.int32) + indices = torch.full((3, 128), -1, device="cuda", dtype=torch.int32) + for row, length in enumerate(lengths.tolist()): + indices[row, :length] = slots[:length] + layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) + state = SimpleNamespace( + get_topk_indices_params={"hidden_states": None, "q_lora": None}, + mem_manager=SimpleNamespace(get_att_input_params=lambda layer_index: kv), + prefill_att_state=Glm5NextSparsePrefillState(), + b1_cu_q_seq_len=torch.arange(4, device="cuda", dtype=torch.int32), + max_q_seq_len=1, + ) + state.decode_att_state = Glm5NextSparseDecodeState( + infer_state=state, + nsa_cache_seqlens=lengths, + nsa_cu_seqlens_k_new=F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), + ) + projected = torch.bmm(q.transpose(0, 1), kv_b.view(4, 256, 512)[:, :128]).transpose(0, 1) + expected = [] + for row, length in enumerate(lengths.tolist()): + keys = kv[slots[:length].long(), 0].float() + scores = projected[row].float() @ keys.T * config["qk_nope_head_dim"] ** -0.5 + expected.append(scores.softmax(-1) @ keys) + expected = torch.stack(expected) + prefill = layer._context_attention_kernel(q, None, state, weight) + state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + decode = layer._token_attention_kernel(q, state, weight) + for output in (prefill, decode): + torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=0.015) + + +@pytest.mark.parametrize("mode", ["dense", "moe", "fused_moe"]) +def test_ffn_dispatch_preserves_clamp_and_shared_expert(config, monkeypatch, mode): + args = StartArgs(enable_fused_shared_experts=mode == "fused_moe") + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + get_env_start_args.cache_clear() + index = 0 if mode == "dense" else 3 + layer = Glm5NextTransformerLayerInfer(index, config) + weight = Glm5NextTransformerLayerWeight(index, torch.bfloat16, config, Quantcfg(config)) + prefix = f"model.language_model.layers.{index}.mlp" + eye = torch.eye(128, device="cuda", dtype=torch.bfloat16) + tensors = {} + + def mlp_weights(name, gate_scale, up_scale): + tensors[f"{name}.gate_proj.weight"] = eye * gate_scale + tensors[f"{name}.up_proj.weight"] = eye * up_scale + tensors[f"{name}.down_proj.weight"] = eye + + if mode == "dense": + mlp_weights(prefix, 1.5, 0.75) + else: + mlp_weights(f"{prefix}.shared_experts", 1.5, 0.75) + for expert in range(4): + mlp_weights(f"{prefix}.experts.{expert}", 1 + expert / 4, 2 + expert / 8) + router = torch.zeros(4, 128, device="cuda", dtype=torch.float32) + router[:, :4] = torch.eye(4, device="cuda") * 0.2 + tensors[f"{prefix}.gate.weight"] = router + tensors[f"{prefix}.gate.e_score_correction_bias"] = torch.zeros(4, device="cuda") + weight.load_hf_weights(tensors) + x = torch.linspace(-25, 25, 3 * 128, device="cuda", dtype=torch.bfloat16).view(3, 128) + + def reference_mlp(gate_scale, up_scale): + gate = (x * gate_scale).float().clamp(max=10) + up = (x * up_scale).float().clamp(-10, 10) + return (F.silu(gate).bfloat16().float() * up).bfloat16() + + expected = reference_mlp(1.5, 0.75) + if mode != "dense": + scores = F.linear(x.float(), router).sigmoid() + probabilities, ids = scores.topk(2, dim=-1) + probabilities /= probabilities.sum(-1, keepdim=True) + routed = torch.zeros_like(x) + for expert in range(4): + coefficient = (probabilities * (ids == expert)).sum(-1, keepdim=True) + routed += (reference_mlp(1 + expert / 4, 2 + expert / 8).float() * coefficient).bfloat16() + expected += routed + actual = layer._ffn(x.clone(), SimpleNamespace(is_prefill=True), weight) + torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.015) + + +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_ep_ffn_uses_combined_output_and_replicated_shared_expert(config, monkeypatch, is_prefill): + args = StartArgs(tp=2, enable_ep_moe=True, enable_fused_shared_experts=True) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "1") + monkeypatch.setenv("LIGHTLLM_GLOBAL_WORLD_SIZE", "2") + monkeypatch.setenv("LIGHTLLM_GLOBAL_RANK", "1") + get_env_start_args.cache_clear() + layer = Glm5NextTransformerLayerInfer(3, config) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + assert weight.num_fused_shared_experts == 0 + prefix = "model.language_model.layers.3.mlp.shared_experts" + eye = torch.eye(128, device="cuda", dtype=torch.bfloat16) + weight.load_hf_weights({f"{prefix}.{name}.weight": eye for name in ("gate_proj", "up_proj", "down_proj")}) + x = torch.linspace(-25, 25, 3 * 128, device="cuda", dtype=torch.bfloat16).view(3, 128) + original = x.clone() + monkeypatch.setattr(weight.moe_gate, "mm", lambda x: torch.zeros(3, 4, device="cuda")) + monkeypatch.setattr(layer, "_tpsp_reduce", lambda **kwargs: pytest.fail("EP output must not be TP-reduced again")) + + def combined_output(hidden_states, **kwargs): + assert kwargs["is_prefill"] is is_prefill + assert kwargs["alpha"] == 1.0 and kwargs["limit"] == 10.0 + assert kwargs["clamp_up_add_one"] is False + return torch.full_like(hidden_states, 17) + + monkeypatch.setattr(weight.experts, "experts", combined_output) + actual = layer._ffn(x, SimpleNamespace(is_prefill=is_prefill), weight) + shared = (F.silu(x.float().clamp(max=10)).bfloat16().float() * x.float().clamp(-10, 10)).bfloat16() + torch.testing.assert_close(actual, shared + 17, atol=0.25, rtol=0.01) + torch.testing.assert_close(x, original, atol=0, rtol=0) + + +def test_fp8_shared_expert_scales_and_bf16_kv_b_load_together(config, monkeypatch): + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs(enable_fused_shared_experts=True))) + ) + get_env_start_args.cache_clear() + config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128]} + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + prefix = "model.language_model.layers.3" + tensors = {} + for expert in range(5): + name = f"experts.{expert}" if expert < 4 else "shared_experts" + for projection in ("gate_proj", "up_proj", "down_proj"): + tensors[f"{prefix}.mlp.{name}.{projection}.weight"] = random_weight(128, 128).to(torch.float8_e4m3fn) + tensors[f"{prefix}.mlp.{name}.{projection}.weight_scale_inv"] = torch.ones(1, 1, device="cuda") + tensors[f"{prefix}.mlp.gate.e_score_correction_bias"] = torch.zeros(4, device="cuda") + kv_b = random_weight(1024, 512) + tensors[f"{prefix}.self_attn.kv_b_proj.weight"] = kv_b + # There is deliberately no kv_b scale: this matrix stays BF16 in the FP8 checkpoint. + weight.load_hf_weights(tensors) + assert weight.experts.verify_load() + assert weight.k_b_proj_.verify_load() and weight.v_b_proj_.verify_load() + q = random_weight(4, 2, 128) + actual = weight.k_b_proj_.bmm(q) + expected = torch.bmm(q, kv_b.view(4, 256, 512)[:, :128]) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("method", ["context_forward", "token_forward"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_mtp_pre_layer_applies_main_norm_before_hidden_fusion(config, monkeypatch, method, dtype): + from lightllm.common.basemodel.layer_weights.meta_weights import RMSNormWeight + from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer + from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer + + hidden_size, eps = config["hidden_size"], config["rms_norm_eps"] + norms = [RMSNormWeight(hidden_size, name, dtype) for name in ("main_norm", "enorm", "hnorm")] + for norm in norms: + norm.weight.copy_(torch.randn_like(norm.weight)) + embeddings = torch.randn(5, hidden_size, device="cuda", dtype=dtype) + hidden = torch.randn_like(embeddings) + projection = torch.randn(2 * hidden_size, hidden_size, device="cuda", dtype=dtype) * 0.1 + weight = SimpleNamespace( + main_norm_weight_=norms[0], + enorm_weight_=norms[1], + hnorm_weight_=norms[2], + eh_proj_weight_=SimpleNamespace(mm=lambda x: x @ projection), + ) + state = SimpleNamespace(mtp_draft_input_hiddens=hidden.clone()) + monkeypatch.setattr(LlamaMultimodalPreLayerInfer, method, lambda self, ids, state, weight: embeddings.clone()) + actual = getattr(Glm5NextMTPPreLayerInfer(config), method)(None, state, weight) + + def reference_norm(value, norm): + return F.rms_norm(value.float(), (hidden_size,), norm.weight.float(), eps).to(dtype) + + normalized_hidden = reference_norm(reference_norm(hidden, norms[0]), norms[2]) + expected = torch.cat((reference_norm(embeddings, norms[1]), normalized_hidden), dim=-1) @ projection + torch.testing.assert_close(actual, expected, atol=0.015 if dtype == torch.bfloat16 else 1e-5, rtol=0.015) + + +@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) +def test_glm_weight_loading_requires_hf(config, model_class): + model = object.__new__(model_class) + model.config, model.tp_world_size_ = config, 1 + model.args, model.load_way = StartArgs(), "HF" + model._verify_params() + model.load_way = "DS" + with pytest.raises(AssertionError, match="only support HF format weights"): + model._verify_params() + + +@pytest.mark.parametrize("tp_world_size", [1, 2]) +def test_glm_rejects_tpsp(config, tp_world_size): + model = object.__new__(Glm5NextTpPartModel) + model.config, model.tp_world_size_ = config, tp_world_size + model.args, model.load_way = StartArgs(), "HF" + model._verify_params() + + model.args.enable_tpsp_mix_mode = True + with pytest.raises(AssertionError, match="GLM-5.3 Flash does not support TP/SP mixed mode"): + model._verify_params() + + +def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): + main = object.__new__(Glm5NextTpPartModel) + main.config = dict(config, mhc=True) + main.tp_world_size_ = 1 + main.args, main.load_way = StartArgs(), "HF" + main._verify_params() + main._init_some_value() + main.layers_infer = [object() for _ in range(config["num_hidden_layers"])] + main.pre_post_weight = SimpleNamespace(wte_weight_=object(), lm_head_weight_=object(), final_norm_weight_=object()) + main.req_manager, main.mem_manager, main.linear_config = object(), object(), object() + drafts = [] + for step in range(2): + draft = object.__new__(Glm5NextMTPModel) + draft.main_model, draft.mtp_previous_draft_models = main, list(drafts) + draft.tp_world_size_, draft.data_type = 1, torch.bfloat16 + draft.load_way = "HF" + draft._init_config() + draft._verify_params() + draft.quant_cfg = Quantcfg(draft.config) + draft._init_weights() + draft._init_infer_layer() + draft._init_some_value() + draft._init_req_manager() + draft._init_mem_manager() + draft._init_att_backend1() + assert main.config["mhc"] and not draft.config["mhc"] + assert draft.config["layer_types"] is not main.config["layer_types"] + assert draft.layers_num == 1 and draft.head_dim_ == 512 + assert draft.layers_infer[0].layer_num_ == 4 + step + assert draft.trans_layers_weight[0].layer_num_ == 4 + assert draft.pre_post_weight.wte_weight_ is main.pre_post_weight.wte_weight_ + assert draft.pre_post_weight.lm_head_weight_ is main.pre_post_weight.lm_head_weight_ + assert draft.pre_post_weight.main_norm_weight_ is main.pre_post_weight.final_norm_weight_ + assert draft.req_manager is main.req_manager and draft.mem_manager is main.mem_manager + assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + drafts.append(draft) + + # NoPE state initialization must not require a model-owned RoPE cache. + for model in (main, *drafts): + state = model.infer_state_class() + state.is_prefill = True + state.input_ids = torch.zeros(3, device="cuda", dtype=torch.int64) + state.b_ready_cache_len = torch.tensor([2], device="cuda", dtype=torch.int32) + state.b_seq_len = torch.tensor([5], device="cuda", dtype=torch.int32) + state.init_some_extra_state(model) + torch.testing.assert_close( + state.position_ids, torch.arange(2, 5, device="cuda", dtype=state.position_ids.dtype) + ) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py new file mode 100644 index 0000000000..dac19a44b5 --- /dev/null +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -0,0 +1,222 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_step, state_dtype): + torch.manual_seed(53) + width, heads, dim = mtp_step + 1, 2, 128 + # Nonconsecutive requests plus a zero-length graph-padding sequence. + req_ids = torch.tensor([3, 1, 5], device="cuda", dtype=torch.int32) + slots = req_ids[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) + states = torch.randn(6 * width, heads, dim, dim, device="cuda", dtype=state_dtype) * 0.1 + accepted = torch.tensor([width, 1, 1], device="cuda", dtype=torch.int32) + for lengths in ([width, 1, 0], [2, width, 0], [1, 1, 0]): + cu = torch.tensor([0, lengths[0], sum(lengths), sum(lengths)], device="cuda", dtype=torch.int32) + tokens = sum(lengths) + q, k, v, gate = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(4)] + beta = torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16) + a = torch.randn(heads, device="cuda") + bias = torch.randn(heads * dim, device="cuda") + reference = states.clone() + before = states.clone() + expected = torch.empty_like(v) + for request, length in enumerate(lengths): + if not length: + continue + start = sum(lengths[:request]) + current = reference[slots[request, accepted[request] - 1]].clone().unsqueeze(0) + for offset in range(length): + token = start + offset + out, _ = fused_recurrent_kda( + q[:, token : token + 1], + k[:, token : token + 1], + v[:, token : token + 1], + gate[:, token : token + 1].reshape(1, 1, -1), + beta[:, token : token + 1], + a, + bias, + current, + torch.zeros(1, device="cuda", dtype=torch.int32), + ) + expected[:, token : token + 1] = out + reference[slots[request, offset]] = current[0] + actual, _ = fused_recurrent_kda( + q, + k, + v, + gate.reshape(1, tokens, -1), + beta, + a, + bias, + states, + slots, + cu_seqlens=cu, + num_accepted_tokens=accepted, + ) + torch.testing.assert_close(actual, expected, atol=2e-3, rtol=1e-2) + torch.testing.assert_close(states, reference, atol=2e-3, rtol=1e-2) + assert torch.equal(states[slots[2]], before[slots[2]]) + # The next verify must begin at the accepted candidate, not the last written slot. + accepted = torch.tensor([min(2, lengths[0]), 1, 1], device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) +@pytest.mark.parametrize("prefix", [1, 2, 3, 4, 7, 15]) +def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix): + torch.manual_seed(53) + width, capacity = mtp_step + 1, 128 + raw = torch.randn(capacity, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) + ring = torch.zeros(2, 4 + mtp_step, 256, device="cuda", dtype=torch.bfloat16) + packed = torch.zeros(capacity, 132, device="cuda", dtype=torch.uint8) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + + def prefill(values, destination, tail): + count = values.shape[0] + compress_pools( + values, + tail, + destination, + ape, + torch.arange(1, count + 1, device="cuda", dtype=torch.int32), + torch.zeros(count, device="cuda", dtype=torch.int32), + ragged, + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([0, count], device="cuda", dtype=torch.int32), + torch.tensor([count], device="cuda", dtype=torch.int32), + count, + ) + + prefill(raw[:prefix], packed, ring) + # Alternate expanded verification, rewound single-token drafting, and re-verification. + for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: + raw[start : start + count] = torch.randn(count, 256, device="cuda", dtype=torch.bfloat16) + lengths = torch.arange(start + 1, start + count + 1, device="cuda", dtype=torch.int32) + compress_pools( + raw[start : start + count], + ring, + packed, + ape, + lengths, + torch.zeros(count, device="cuda", dtype=torch.int32), + ragged, + torch.zeros(count, device="cuda", dtype=torch.int32), + torch.arange(count + 1, device="cuda", dtype=torch.int32), + lengths, + 1, + mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), + ) + expected = torch.zeros_like(packed) + prefill(raw[: start + count], expected, torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16)) + closing = torch.arange(start, start + count, device="cuda") + closing = closing[(closing + 1) % 4 == 0] + assert torch.equal(packed[closing], expected[closing]) + + +@pytest.mark.parametrize("dynamic", [False, True]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): + from types import SimpleNamespace + + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + + torch.manual_seed(53) + heads, dim, width = 2, 128, 3 + hidden = heads * dim + lengths = [2, 1, 3] if dynamic else [3, 3, 3] + requests = [2, 0, 1] + request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] + offsets = [offset for length in lengths for offset in range(length)] + if dynamic: + request_rows += [3, 3] + offsets += [0, 0] + tokens = len(request_rows) + conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 + old_conv, expected_ssm = conv.clone(), ssm.clone() + accepted_offsets = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) + conv_weight = torch.randn(3 * hidden, 4, device="cuda", dtype=torch.bfloat16) * 0.1 + a, bias = torch.randn(heads, device="cuda"), torch.randn(hidden, device="cuda") + expected = [] + row = 0 + zero = torch.zeros(1, dtype=torch.int32, device="cuda") + for req, length in zip(requests, lengths): + offset = int(accepted_offsets[req]) + history = old_conv[req : req + 1, :, offset : offset + 3].contiguous() + state = expected_ssm[req * width + offset].clone().unsqueeze(0) + for token in range(length): + convolved = causal_conv1d_update( + mixed[row : row + 1].clone(), + history, + conv_weight, + bias=None, + activation="silu", + conv_state_indices=zero, + ) + q, k, v = [part.reshape(1, 1, heads, dim) for part in convolved.split(hidden, -1)] + output, _ = fused_recurrent_kda( + q, + k, + v, + gate[row : row + 1].reshape(1, 1, hidden), + beta[row : row + 1].reshape(1, 1, heads), + a, + bias, + state, + zero, + ) + expected.append(output.reshape(1, heads, dim)) + expected_ssm[req * width + token] = state[0] + row += 1 + backend = SimpleNamespace( + mtp_step=2, + tp_num_heads=heads, + head_dim=dim, + tp_hidden_size=hidden, + lower_bound=-5.0, + uses_dynamic_spec_verify_layout=lambda: dynamic, + split_qkv=lambda x: x.split(hidden, -1), + ) + manager = SimpleNamespace( + req_to_mtp_state_index=accepted_offsets, HOLD_REQUEST_ID=3, get_mamba_cache=lambda layer: (conv, ssm) + ) + infer = SimpleNamespace( + batch_size=tokens, + b_req_idx=torch.tensor(request_rows, device="cuda", dtype=torch.int32), + b_mtp_index=torch.tensor(offsets, device="cuda", dtype=torch.int32), + req_manager=manager, + ) + weight = SimpleNamespace( + get_merged_kda_conv_weight=lambda: conv_weight, + linear_A_log=SimpleNamespace(weight=a), + linear_dt_bias=SimpleNamespace(weight=bias), + ) + state = KDADecodeAttState(backend=backend, infer_state=infer) + state.init_state() + actual = state.decode_att( + None, + None, + None, + AttControl( + linear_att_decode=True, + linear_att_decode_dict=dict( + layer_weight=weight, layer_num=0, mixed_qkv=mixed, raw_gate=gate, raw_beta=beta + ), + ), + ) + torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) + torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) + if dynamic: + assert torch.equal(conv[3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py new file mode 100644 index 0000000000..57665cbbe3 --- /dev/null +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -0,0 +1,173 @@ +import dataclasses + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, MemoryManager +from lightllm.common.req_manager import Glm5NextReqManager +from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(params=[0, 2]) +def make_mems(monkeypatch, request): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) + mtp_step = request.param + args = StartArgs( + data_type="bfloat16", + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + mtp_step=mtp_step, + mtp_mode="eagle_with_att" if mtp_step else None, + ) + set_unique_server_name(args) + get_env_start_args.cache_clear() + set_env_start_args(dataclasses.asdict(args)) + + def check_shm_refs(mem, req): + assert mem.req_to_indexer_tail is req.req_to_indexer_tail + assert mem.req_to_conv_state is req.req_to_conv_state + assert mem.req_to_ssm_state is req.req_to_ssm_state + assert mem.big_page_buffers is None + + monkeypatch.setattr(MemoryManager, "write_to_shm", check_shm_refs) + + def create(tp): + config = Glm5NextCacheConfig( + tp_world_size=tp, + full_att_all_num_kv_heads=1, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=584, + global_linear_k_heads=8, + global_linear_v_heads=8, + num_linear_k_heads=8 // tp, + num_linear_v_heads=8 // tp, + head_linear_k_dim=128, + head_linear_v_dim=128, + conv_kernel_size=4, + linear_layer_num=6, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4, + all_layer_num=8, + draft_full_att_kv_layer_num=int(mtp_step > 0), + ) + mems = [] + for _ in range(tp): + mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) + req = Glm5NextReqManager(3, 32, mem, config) + mem.write_to_shm(req) + assert mem.big_page_buffers is not None + mems.append(mem) + mems[0].alloc_paged_kv_move_buffer(1, 2048) + return mems + + return create + + +@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) +@pytest.mark.parametrize("remainder", range(4)) +def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder): + source, dest = make_mems(prefill_tp), make_mems(decode_tp) + mtp_size = get_env_start_args().mtp_step + 1 + layers = source[0].layer_num + src_req, dst_req = 0, 2 + length = 8 + remainder + src_indexes = torch.randperm(32)[:length].tolist() + dst_indexes = torch.randperm(32)[:length].tolist() + packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") + conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") + ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") + # Non-live slots deliberately contain stale data. The sequence length, not + # the bytes in these slots, controls subsequent K-pool completion. + tail = torch.randn(layers, 3 + mtp_size, 256, dtype=torch.bfloat16, device="cuda") + for rank, mem in enumerate(source): + heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) + mem.kv_buffer.view(torch.uint8)[:, src_indexes] = packed + mem.req_to_conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) + mem.req_to_ssm_state.buffer[:, src_req * mtp_size].copy_(ssm[:, heads]) + mem.req_to_indexer_tail.buffer[:, src_req].copy_(tail) + + untouched = [] + for mem in dest: + mem.kv_buffer.zero_() + mem.req_to_conv_state.buffer.normal_() + mem.req_to_ssm_state.buffer.normal_() + mem.req_to_indexer_tail.buffer.normal_() + untouched.append( + [ + x.buffer[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() + for x, stride in ( + (mem.req_to_conv_state, 1), + (mem.req_to_ssm_state, mtp_size), + (mem.req_to_indexer_tail, 1), + ) + ] + ) + + source[0].write_mem_to_page_kv_move_buffer(src_indexes, 0, 0, source, prefill_tp) + # Model a byte-preserving transport between separate P/D page buffers. + dest[0].kv_move_buffer.view(torch.uint8).copy_(source[0].kv_move_buffer.view(torch.uint8)) + dest[0].read_page_kv_move_buffer_to_mem(dst_indexes, 0, 0, dest, decode_tp) + source[0].write_mem_to_page_kv_move_buffer([], 0, 0, source, prefill_tp, "att_state", src_req) + dest[0].kv_move_buffer.view(torch.uint8).copy_(source[0].kv_move_buffer.view(torch.uint8)) + dest[0].read_page_kv_move_buffer_to_mem([], 0, 0, dest, decode_tp, "att_state", dst_req) + torch.cuda.synchronize() + + for rank, mem in enumerate(dest): + heads = slice(rank * 8 // decode_tp, (rank + 1) * 8 // decode_tp) + assert torch.equal(mem.kv_buffer.view(torch.uint8)[:, dst_indexes], packed) + assert torch.equal(mem.req_to_conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) + assert torch.equal(mem.req_to_ssm_state.buffer[:, dst_req * mtp_size], ssm[:, heads]) + assert torch.equal(mem.req_to_indexer_tail.buffer[:, dst_req], tail) + for (state, stride), expected in zip( + ((mem.req_to_conv_state, 1), (mem.req_to_ssm_state, mtp_size), (mem.req_to_indexer_tail, 1)), + untouched[rank], + ): + indexes = torch.tensor([0, 1, 3], device="cuda") * stride + assert torch.equal(state.buffer[:, indexes], expected) + + # Continue through the next pool boundary on P and on the restored D state. + # Different physical KV/request slots must yield identical pooled FP8 bytes. + count = 4 - remainder + raw = torch.randn(count, 256, dtype=torch.bfloat16, device="cuda") + ape = torch.randn(4, 128, device="cuda") + tensor = lambda values: torch.tensor(values, dtype=torch.int32, device="cuda") + outputs = [] + for mem, req_idx, indexes in ((source[0], src_req, src_indexes), (dest[0], dst_req, dst_indexes)): + new_indexes = [i for i in range(32) if i not in indexes][:count] + all_indexes = indexes + new_indexes + compress_pools( + raw=raw, + tail=mem.req_to_indexer_tail.buffer[0], + packed_buffer=mem.get_indexer_k_buffer(3), + ape=ape, + lengths=tensor(range(length + 1, length + count + 1)), + starts=tensor([0] * count), + ragged=tensor(all_indexes), + req_idx=tensor([req_idx]), + cu_q_lens=tensor([0, count]), + seq_lens=tensor([length + count]), + max_q_len=count, + ) + outputs.append(mem.get_indexer_k_buffer(3)[new_indexes[-1]].clone()) + assert torch.equal(*outputs) + + +def test_pd_page_capacity_includes_tail_without_resizing(make_mems): + mem = make_mems(1)[0] + helper = mem.att_state_page_helper + # A page that fits Conv+SSM but not the appended tail must be rejected. + mem.kv_move_buffer = torch.empty((1, helper.tail_offset), dtype=torch.uint8, device="cuda") + with pytest.raises(AssertionError, match="smaller than global linear att state"): + helper.assert_page_size() + shape = mem.get_paged_kv_move_buffer_shape(2, 2048) + assert shape == (2, 2048, mem.layer_num, 1, 584) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py new file mode 100644 index 0000000000..6a25a8d4ea --- /dev/null +++ b/unit_tests/models/glm5_next/test_visual.py @@ -0,0 +1,216 @@ +import dataclasses +import json +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer +from lightllm.models.glm5_next.vision_process import Glm5NextImageProcessor, smart_resize +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import set_env_start_args + + +@pytest.mark.parametrize("width,height,canvas", [(7, 11, (168, 112)), (173, 89, (112, 196)), (448, 448, (448, 448))]) +def test_image_canvas_and_token_expansion(tmp_path, width, height, canvas): + assert smart_resize(height, width) == canvas + (tmp_path / "processor_config.json").write_text(json.dumps({"image_processor": {}})) + tokenizer = Glm5NextTokenizer( + None, {"image_start_token_id": 100, "image_end_token_id": 101, "image_token_id": 102}, str(tmp_path) + ) + image = ImageItem(type="base64", data="") + image.image_w, image.image_h = width, height + image.token_id = 1000 + image.token_num = tokenizer.get_image_token_length(image) + assert image.token_num == canvas[0] * canvas[1] // 28 ** 2 + multi = MultimodalParams() + multi.images = [image] + ids = tokenizer.encode([1, 100, 102, 101, 2], multi) + assert ids == [1, 100, *range(1000, 1000 + image.token_num), 101, 2] + assert image.start_idx == 2 + assert image.grid_thwd[-1] == 0 + assert tokenizer.encode([1, 2, 3]) == [1, 2, 3] + assert tokenizer.encode([100, 102, 101]) == [100, 101] + + +def test_image_patch_order_temporal_repeat_and_black_padding(): + pixels = np.zeros((29, 31, 3), dtype=np.uint8) + pixels[:, :, 0] = np.arange(31)[None, :] + pixels[:, :, 1] = np.arange(29)[:, None] + pixels[:, :, 2] = 127 + processor = Glm5NextImageProcessor(min_image_tokens=1, do_rescale=False, do_normalize=False) + patches, grid = processor._preprocess_bydevice(Image.fromarray(pixels), device="cpu") + assert grid.tolist() == [[1, 4, 4]] + patches = patches.reshape(16, 3, 2, 14, 14) + assert torch.equal(patches[:, :, 0], patches[:, :, 1]) + assert torch.equal(patches[0, :, 0], torch.from_numpy(pixels[:14, :14]).permute(2, 0, 1)) + assert torch.equal(patches[1, :, 0], torch.from_numpy(pixels[:14, 14:28]).permute(2, 0, 1)) + assert torch.equal(patches[2, :, 0], torch.from_numpy(pixels[14:28, :14]).permute(2, 0, 1)) + assert not patches[-1].any() + + +def test_vision_attention_uses_triton_rope_with_fa3_backend(monkeypatch): + import lightllm.models.glm5_next.glm5_next_visual as glm_visual + import lightllm.server.visualserver as visualserver + from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend + from lightllm.server.visualserver import set_vit_att_backend + + rope_inputs = [] + norm_inputs = [] + attn_inputs = [] + + def fake_qk_norm(qkv, q_weight, k_weight, eps): + norm_inputs.append((qkv, q_weight, k_weight, eps)) + return qkv[:, 0].contiguous(), qkv[:, 1].contiguous() + + def fake_rope(x, cos, sin): + rope_inputs.append((x, cos, sin)) + return x + + def fake_fa3(q, k, v, o, cu_seqlens, max_seqlen): + attn_inputs.append((q, k, v, cu_seqlens, max_seqlen)) + o.copy_(q) + return o + + monkeypatch.setattr(glm_visual, "qk_rms_norm", fake_qk_norm) + monkeypatch.setattr(glm_visual, "apply_rotary_pos_emb_triton", fake_rope) + monkeypatch.setattr(Fa3VitAttBackend, "_vit_att_fwd", staticmethod(fake_fa3)) + monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", visualserver.VIT_ATTN_BACKEND) + set_vit_att_backend("fa3") + + attention = glm_visual.Glm5NextVisionAttention(hidden_size=16, num_heads=2, eps=1e-5, bias=True) + x = torch.randn(3, 16) + cos = torch.randn(3, 4) + sin = torch.randn(3, 4) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + result = attention(x, cu_seqlens, 3, cos, sin) + + assert result.shape == x.shape + assert len(norm_inputs) == 1 + assert norm_inputs[0][0].shape == (3, 3, 2, 8) + assert [(q.shape, cos_.shape, sin_.shape) for q, cos_, sin_ in rope_inputs] == [ + ((3, 2, 8), (3, 4), (3, 4)), + ((3, 2, 8), (3, 4), (3, 4)), + ] + assert len(attn_inputs) == 1 + assert attn_inputs[0][0].shape == attn_inputs[0][1].shape == attn_inputs[0][2].shape == (3, 2, 8) + + +def test_vision_rms_norm_preserves_glm_rounding(monkeypatch): + import lightllm.models.glm5_next.glm5_next_visual as glm_visual + + norm_calls = [] + + def fake_rms_norm(x, weight, eps, round_norm_before_weight): + norm_calls.append((x, weight, eps, round_norm_before_weight)) + return x + + monkeypatch.setattr(glm_visual, "rms_norm", fake_rms_norm) + norm = glm_visual.Glm5NextVisionRMSNorm(hidden_size=8, eps=1e-5) + x = torch.randn(2, 8) + assert norm(x) is x + assert len(norm_calls) == 1 + assert norm_calls[0][0] is x + assert norm_calls[0][1] is norm.weight + assert norm_calls[0][2:] == (1e-5, True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_triton_glm_vision_norms_match_reference(): + from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm + + def glm_rms_norm(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + return normalized.to(x.dtype) * weight + + torch.manual_seed(53) + eps = 1e-5 + qkv = torch.randn(7, 3, 16, 64, dtype=torch.bfloat16, device="cuda") + q_weight = torch.randn(64, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(64, dtype=torch.bfloat16, device="cuda") + q, k = qk_rms_norm(qkv, q_weight, k_weight, eps) + + torch.testing.assert_close(q, glm_rms_norm(qkv[:, 0], q_weight, eps), atol=1e-2, rtol=1e-2) + torch.testing.assert_close(k, glm_rms_norm(qkv[:, 1], k_weight, eps), atol=1e-2, rtol=1e-2) + + x = torch.randn(7, 1024, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(1024, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close( + rms_norm(x, weight, eps, round_norm_before_weight=True), + glm_rms_norm(x, weight, eps), + atol=1e-2, + rtol=1e-2, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_vision_batch_preserves_independent_image_attention(): + from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer + from lightllm.server.visualserver import set_vit_att_backend + + torch.manual_seed(53) + set_vit_att_backend("sdpa") + model = ( + Glm5NextVisionTransformer( + {"data_type": "float32"}, + hidden_size=128, + out_hidden_size=256, + depth=2, + intermediate_size=256, + projection_intermediate_size=512, + num_heads=2, + ) + .cuda() + .eval() + ) + grid = torch.tensor([[1, 4, 6], [1, 6, 4]]) + x = torch.randn(48, 3 * 2 * 14 ** 2, device="cuda") + with torch.inference_mode(): + batch = model(x, grid) + individual = torch.cat([model(x[:24], grid[:1]), model(x[24:], grid[1:])]) + assert batch.shape == (12, 256) + torch.testing.assert_close(batch, individual, atol=2e-5, rtol=2e-5) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_target_and_mtp_prefill_embed_images_before_hidden_fusion(monkeypatch): + from lightllm.models.glm5_next.layer_infer.pre_layer_infer import Glm5NextPreLayerInfer + from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + set_env_start_args(dataclasses.asdict(StartArgs(mtp_mode="eagle_with_att", mtp_step=2))) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + hidden = 128 + config = {"hidden_size": hidden, "rms_norm_eps": 1e-5, "hc_mult": 4} + image_cache = torch.randn(3, 1, hidden, pin_memory=True) + monkeypatch.setattr(g_infer_context, "cpu_embed_cache_client", SimpleNamespace(cpu_embed_cache_tensor=image_cache)) + text = torch.randn(8, hidden, device="cuda") + ids = torch.tensor([1, 1000, 1001, 2], device="cuda") + metadata = {"token_id": 1000, "token_num": 2, "start_index_in_embed_cache": 1} + state = SimpleNamespace(multimodal_params=[{"images": [metadata], "audios": []}]) + weight = SimpleNamespace(wte_weight_=SimpleNamespace(weight=text, tp_vocab_start_id=0, tp_vocab_end_id=8)) + expected = torch.cat((text[1:2], image_cache[1:, 0].cuda(), text[2:3])) + expanded = Glm5NextPreLayerInfer(config).context_forward(ids, state, weight) + torch.testing.assert_close(expanded.reshape(4, 4, hidden), expected[:, None].expand(-1, 4, -1)) + + def norm(input, eps, out): + out.copy_(input * torch.rsqrt(input.square().mean(-1, keepdim=True) + eps)) + return out + + state.mtp_draft_input_hiddens = torch.randn(4, hidden, device="cuda") + main_norm_weight = torch.randn(hidden, device="cuda") + old_hidden = state.mtp_draft_input_hiddens.clone() + projection = torch.randn(2 * hidden, hidden, device="cuda") * 0.1 + weight.enorm_weight_ = weight.hnorm_weight_ = norm + weight.main_norm_weight_ = lambda input, eps, out: out.copy_( + torch.nn.functional.rms_norm(input, (hidden,), main_norm_weight, eps) + ) + weight.eh_proj_weight_ = SimpleNamespace(mm=lambda x: x @ projection) + actual = Glm5NextMTPPreLayerInfer(config).context_forward(ids, state, weight) + old_hidden = torch.nn.functional.rms_norm(old_hidden, (hidden,), main_norm_weight, 1e-5) + normalized = [x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) for x in (expected, old_hidden)] + torch.testing.assert_close(actual, torch.cat(normalized, -1) @ projection) diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py new file mode 100644 index 0000000000..f16001f8c7 --- /dev/null +++ b/unit_tests/server/test_dp_request_capacity.py @@ -0,0 +1,110 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from lightllm.server.core.objs.shm_req_manager import ShmReqManager +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.batch import Batch +from lightllm.server.router.req_queue.chunked_prefill.impl import ChunkedPrefillQueue +from lightllm.server.router.req_queue.dp_base_queue import DpQueue +from lightllm.utils.config_utils import get_running_max_req_size_per_dp + + +@pytest.mark.parametrize( + "global_capacity,dp,nnodes,expected", + [(32, 8, 1, 4), (33, 8, 1, 5), (4, 8, 1, 1), (32, 8, 2, 8), (32, 1, 2, 32), (32, 1, 1, 32)], +) +def test_hybrid_capacity_uses_local_dp_count(monkeypatch, global_capacity, dp, nnodes, expected): + monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) + args = StartArgs(running_max_req_size=global_capacity, dp=dp, nnodes=nnodes) + assert get_running_max_req_size_per_dp(args) == expected + # HTTP/SHM indexes remain global even when GPU request indexes are local. + monkeypatch.setattr("lightllm.server.core.objs.shm_req_manager.get_env_start_args", lambda: args) + assert ShmReqManager.get_max_req_num(None) == global_capacity + + +@pytest.mark.parametrize("mode", ["full_attention", "diverse", "cache_fetch"]) +def test_modes_requiring_existing_capacity_keep_global_default(monkeypatch, mode): + monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: mode != "full_attention") + args = StartArgs( + running_max_req_size=32, + dp=8, + diverse_mode=mode == "diverse", + enable_dp_prompt_cache_fetch=mode == "cache_fetch", + ) + assert get_running_max_req_size_per_dp(args) == 32 + + +@pytest.mark.parametrize("capacity", [0, -1]) +def test_reject_invalid_global_capacity(capacity): + args = StartArgs(running_max_req_size=capacity, dp=8) + with pytest.raises(ValueError, match="running_max_req_size"): + get_running_max_req_size_per_dp(args) + + +def _make_dp_queue(monkeypatch, balancer): + monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) + monkeypatch.setattr("lightllm.server.router.req_queue.base_queue.get_fixed_kv_len", lambda: 0) + args = StartArgs( + running_max_req_size=32, + dp=8, + max_total_token_num=1048576, + batch_max_tokens=4096, + router_token_ratio=0.85, + dp_balancer=balancer, + ) + router = SimpleNamespace( + router_statics=SimpleNamespace(ema_req_out_len=16), + get_used_tokens=lambda dp: 0, + shared_token_load=MagicMock(), + ) + return DpQueue(args, router, ChunkedPrefillQueue, dp_size_in_node=8) + + +def _req(request_id, dp=-1): + return SimpleNamespace( + request_id=request_id, + sample_params=SimpleNamespace(suggested_dp_index=dp, pd_high_priority_request=False), + is_aborted=False, + get_tuple_tokens=lambda busy, ema: (64, 16), + get_first_router_need_tokens=lambda: 64, + get_decode_need_tokens=lambda: 3, + ) + + +@pytest.mark.parametrize("balancer", ["bs_balancer", "round_robin"]) +def test_global_32_requests_run_as_four_per_rank(monkeypatch, balancer): + queue = _make_dp_queue(monkeypatch, balancer) + for i in range(32): + queue.extend([_req(i)]) + batch = queue.generate_new_batch(None) + assert batch.get_all_dp_req_num() == [4] * 8 + assert queue.get_wait_req_num() == 0 + + +def test_pinned_overflow_waits_until_a_local_slot_is_released(monkeypatch): + queue = _make_dp_queue(monkeypatch, "bs_balancer") + # A skewed client can use any global SHM slot, but only four may enter + # inference on this rank. The fifth must not exhaust its GPU request pool. + pinned = [_req(i, dp=2) for i in range(6)] + for req in pinned: + queue.extend([req]) + queue.extend([_req(6, dp=3)]) + batch = queue.generate_new_batch(None) + assert batch.get_all_dp_req_num() == [0, 0, 4, 1, 0, 0, 0, 0] + assert queue.get_wait_req_num() == 2 + assert queue.generate_new_batch(batch) is None + # Releasing a request on a different rank must not open a slot on rank 2. + batch.pop_req(6) + assert queue.generate_new_batch(batch) is None + batch.pop_req(0) + resumed = queue.generate_new_batch(batch) + assert [req.request_id for req in resumed.reqs] == [4] + batch.merge(resumed) + assert batch.get_all_dp_req_num()[2] == 4 + assert queue.get_wait_req_num() == 1 + assert queue.generate_new_batch(batch) is None + rest = queue.generate_new_batch(Batch(-1, [], dp_size_in_node=8)) + assert [req.request_id for req in rest.reqs] == [5] + assert queue.get_wait_req_num() == 0 diff --git a/unit_tests/server/test_glm47_tool_parser.py b/unit_tests/server/test_glm47_tool_parser.py new file mode 100644 index 0000000000..8f626d6343 --- /dev/null +++ b/unit_tests/server/test_glm47_tool_parser.py @@ -0,0 +1,127 @@ +import json + +import pytest + +from lightllm.server.api_models import Tool +from lightllm.server.function_call_parser import FunctionCallParser + + +TOOLS = [ + Tool.model_validate( + { + "type": "function", + "function": { + "name": "Read", + "parameters": {"type": "object", "properties": {"file_path": {"type": "string"}}}, + }, + } + ), + Tool.model_validate( + { + "type": "function", + "function": { + "name": "Write", + "parameters": { + "type": "object", + "properties": {"content": {"type": "string"}, "overwrite": {"type": "boolean"}}, + }, + }, + } + ), +] + + +def collect_stream(text, chunk_size): + parser = FunctionCallParser(TOOLS, "glm47") + content = "" + calls = {} + for start in range(0, len(text), chunk_size): + normal, deltas = parser.parse_stream_chunk(text[start : start + chunk_size]) + content += normal + for delta in deltas: + if delta.name is not None: + assert delta.tool_index not in calls, "a tool call must have exactly one name event" + calls[delta.tool_index] = {"name": delta.name, "arguments": ""} + calls[delta.tool_index]["arguments"] += delta.parameters + return content, calls + + +@pytest.mark.parametrize("separator", ["", "\n"]) +@pytest.mark.parametrize("chunk_size", [1, 7, 39, 4096]) +def test_repeated_tool_calls_have_separate_indices(separator, chunk_size): + text = ( + f'Write{separator}content{{"ready": true}}' + "overwritefalse" + f"Read{separator}file_path/tmp/a" + f"Read{separator}file_path/tmp/b" + ) + _, calls = collect_stream(text, chunk_size) + assert list(calls) == [0, 1, 2] + assert [call["name"] for call in calls.values()] == ["Write", "Read", "Read"] + assert [json.loads(call["arguments"]) for call in calls.values()] == [ + {"content": '{"ready": true}', "overwrite": False}, + {"file_path": "/tmp/a"}, + {"file_path": "/tmp/b"}, + ] + _, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text) + assert [call.tool_index for call in nonstream] == [0, 1, 2] + assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()] + + +@pytest.mark.parametrize("value", ['{"count": 1}', "123", "false", " return 1\n", '"quoted"']) +def test_string_arguments_preserve_type_and_whitespace(value): + text = f"Writecontent{value}" + _, calls = collect_stream(text, 1) + assert json.loads(calls[0]["arguments"])["content"] == value + + +@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) +def test_text_around_tool_calls_survives_chunk_boundaries(chunk_size): + text = ( + "Before Readfile_path/tmp/a" + " between Readfile_path/tmp/b after" + ) + content, calls = collect_stream(text, chunk_size) + assert content == "Before between after" + assert list(calls) == [0, 1] + + +def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name(): + parser = FunctionCallParser(TOOLS, "glm47") + normal, calls = parser.parse_stream_chunk("Writecontent") + assert normal == "" + assert len(calls) == 1 + assert calls[0].tool_index == 0 + assert calls[0].name == "Write" + assert calls[0].parameters == "" + + value = ' line 1\n{"literal": true}\n' + "x" * 100_000 + for chunk in (value[:13], value[13:27], value[27:]): + normal, calls = parser.parse_stream_chunk(chunk) + assert normal == "" + assert len(calls) == 1 + assert calls[0].tool_index == 0 + assert calls[0].name is None + assert calls[0].parameters == "" + + normal, calls = parser.parse_stream_chunk("") + assert normal == "" + assert len(calls) == 1 + assert calls[0].tool_index == 0 + assert calls[0].name is None + assert json.loads(calls[0].parameters) == {"content": value} + + _, calls = parser.parse_stream_chunk("Write\n") + assert len(calls) == 1 + assert calls[0].tool_index == 1 + assert calls[0].name == "Write" + + +def test_undefined_buffered_tool_does_not_consume_a_call_index(): + parser = FunctionCallParser(TOOLS, "glm47") + _, calls = parser.parse_stream_chunk("unknown\ncontentx") + assert calls == [] + _, calls = parser.parse_stream_chunk("Read\n") + assert len(calls) == 1 + assert calls[0].name == "Read" + assert calls[0].tool_index == 0