From 6650d1757a01a87b7d275c14db179d48c96fbcf5 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:49:21 +0000 Subject: [PATCH 01/24] feat: add minimal GLM-5.3 Flash TP inference with hybrid page cache --- .../fused_moe/impl/triton_impl.py | 5 + .../layer_weights/meta_weights/norm_weight.py | 21 +- .../fused_moe/grouped_fused_moe.py | 13 +- .../fused_moe/moe_silu_and_mul.py | 11 +- .../linear_att/fla/ops/chunk_delta_h.py | 17 +- .../linear_att/fla/ops/solve_tril.py | 4 +- .../triton_kernel/norm/gated_rmsnorm.py | 36 +- .../common/state_cache_manager/linear_att.py | 4 + lightllm/models/__init__.py | 2 + lightllm/models/glm5_next/README.md | 123 ++ lightllm/models/glm5_next/__init__.py | 0 lightllm/models/glm5_next/attention.py | 86 ++ lightllm/models/glm5_next/cache_config.py | 63 + lightllm/models/glm5_next/indexer.py | 77 ++ lightllm/models/glm5_next/kda_backend.py | 175 +++ .../models/glm5_next/layer_infer/__init__.py | 0 .../layer_infer/transformer_layer_infer.py | 259 ++++ .../glm5_next/layer_weights/__init__.py | 0 .../pre_and_post_layer_weight.py | 20 + .../layer_weights/transformer_layer_weight.py | 340 ++++++ lightllm/models/glm5_next/mem_manager.py | 44 + lightllm/models/glm5_next/model.py | 91 ++ lightllm/models/glm5_next/tokenizer.py | 17 + .../glm5_next/triton_kernel/__init__.py | 0 .../glm5_next/triton_kernel/index_quant.py | 69 ++ .../models/glm5_next/triton_kernel/kda.py | 1066 +++++++++++++++++ .../glm5_next/triton_kernel/kda_decode.py | 91 ++ .../models/glm5_next/triton_kernel/kpool.py | 137 +++ .../models/glm5_next/triton_kernel/mhc.py | 547 +++++++++ lightllm/server/tokenizer.py | 9 +- lightllm/utils/config_utils.py | 2 +- .../service/benchmark_glm53_flash.py | 115 ++ unit_tests/models/glm5_next/test_cache.py | 162 +++ unit_tests/models/glm5_next/test_kernels.py | 308 +++++ unit_tests/models/glm5_next/test_tokenizer.py | 27 + 35 files changed, 3918 insertions(+), 23 deletions(-) create mode 100644 lightllm/models/glm5_next/README.md create mode 100644 lightllm/models/glm5_next/__init__.py create mode 100644 lightllm/models/glm5_next/attention.py create mode 100644 lightllm/models/glm5_next/cache_config.py create mode 100644 lightllm/models/glm5_next/indexer.py create mode 100644 lightllm/models/glm5_next/kda_backend.py create mode 100644 lightllm/models/glm5_next/layer_infer/__init__.py create mode 100644 lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py create mode 100644 lightllm/models/glm5_next/layer_weights/__init__.py create mode 100644 lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py create mode 100644 lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py create mode 100644 lightllm/models/glm5_next/mem_manager.py create mode 100644 lightllm/models/glm5_next/model.py create mode 100644 lightllm/models/glm5_next/tokenizer.py create mode 100644 lightllm/models/glm5_next/triton_kernel/__init__.py create mode 100644 lightllm/models/glm5_next/triton_kernel/index_quant.py create mode 100644 lightllm/models/glm5_next/triton_kernel/kda.py create mode 100644 lightllm/models/glm5_next/triton_kernel/kda_decode.py create mode 100644 lightllm/models/glm5_next/triton_kernel/kpool.py create mode 100644 lightllm/models/glm5_next/triton_kernel/mhc.py create mode 100644 test/benchmark/service/benchmark_glm53_flash.py create mode 100644 unit_tests/models/glm5_next/test_cache.py create mode 100644 unit_tests/models/glm5_next/test_kernels.py create mode 100644 unit_tests/models/glm5_next/test_tokenizer.py 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..8596ab6f71 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 @@ -27,6 +27,8 @@ def __init__( routed_expert_counter_tensor=routed_expert_counter_tensor, auto_update_redundancy_expert=auto_update_redundancy_expert, ) + self.swiglu_limit = None + self.swiglu_clamp_up_add_one = True def create_workspace(self): return None @@ -104,6 +106,9 @@ def _fused_experts( use_fp8_w8a8=use_fp8_w8a8, w1_scale=w13_scale, w2_scale=w2_scale, + limit=self.swiglu_limit, + alpha=1.0 if self.swiglu_limit is not None else None, + clamp_up_add_one=self.swiglu_clamp_up_add_one, ) return input_tensor diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py index ee9d1923c3..fe80974569 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py @@ -73,6 +73,17 @@ def __call__( class GatedRMSNormWeight(RMSNormWeight): + def __init__( + self, + dim: int, + weight_name: str, + data_type: torch.dtype, + activation: str = "silu", + ): + super().__init__(dim=dim, weight_name=weight_name, data_type=data_type) + assert activation in ("silu", "sigmoid") + self.activation = activation + def _triton_forward( self, input: torch.Tensor, @@ -86,7 +97,15 @@ def _triton_forward( ), f"input.ndim: {input.ndim} != 2 or weight.ndim: {self.weight.ndim} != 1" if out is None: out = alloc_func(input.shape, dtype=input.dtype, device=input.device) - return gated_rmsnorm_forward(x=input, weight=self.weight, bias=None, eps=eps, z=gate_value, out=out) + return gated_rmsnorm_forward( + x=input, + weight=self.weight, + bias=None, + eps=eps, + z=gate_value, + out=out, + activation=self.activation, + ) def _cuda_forward( self, diff --git a/lightllm/common/basemodel/triton_kernel/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/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/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py index 97933b2ac2..3ca4d6e39b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py @@ -24,6 +24,7 @@ { "USE_G": lambda args: args["g"] is not None, "USE_GK": lambda args: args["gk"] is not None, + "USE_EXP2": lambda args: args["use_exp2"], "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "STORE_FINAL_STATE": lambda args: args["ht"] is not None, "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, @@ -43,6 +44,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht, cu_seqlens, chunk_offsets, + use_exp2, T, H: tl.constexpr, Hg: tl.constexpr, @@ -52,6 +54,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( BV: tl.constexpr, USE_G: tl.constexpr, USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, @@ -169,7 +172,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[:, None] + b_h1 *= (tl.exp2(b_gk_last1) if USE_EXP2 else exp(b_gk_last1))[:, None] if K > 64: o_k2 = 64 + o_k1 b_gk_last2 = tl.load( @@ -177,7 +180,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k2 < K), other=0.0, ) - b_h2 *= exp(b_gk_last2)[:, None] + b_h2 *= (tl.exp2(b_gk_last2) if USE_EXP2 else exp(b_gk_last2))[:, None] if K > 128: o_k3 = 128 + o_k1 b_gk_last3 = tl.load( @@ -185,7 +188,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k3 < K), other=0.0, ) - b_h3 *= exp(b_gk_last3)[:, None] + b_h3 *= (tl.exp2(b_gk_last3) if USE_EXP2 else exp(b_gk_last3))[:, None] if K > 192: o_k4 = 192 + o_k1 b_gk_last4 = tl.load( @@ -193,7 +196,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k4 < K), other=0.0, ) - b_h4 *= exp(b_gk_last4)[:, None] + b_h4 *= (tl.exp2(b_gk_last4) if USE_EXP2 else exp(b_gk_last4))[:, None] b_v = b_v.to(k.dtype.element_ty) p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) @@ -265,6 +268,8 @@ def chunk_gated_delta_rule_fwd_h( save_new_value: bool = True, cu_seqlens: torch.LongTensor | None = None, run_config=None, + chunk_indices: torch.Tensor | None = None, + use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # 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. @@ -272,7 +277,8 @@ def chunk_gated_delta_rule_fwd_h( H = u.shape[-2] BT = chunk_size - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) # N: the actual number of sequences in the batch with either equal or variable lengths if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None @@ -311,6 +317,7 @@ def chunk_gated_delta_rule_fwd_h( ht=final_state, cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, + use_exp2=use_exp2, T=T, H=H, Hg=Hg, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py index b5b6cfc369..1053d281a0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py @@ -412,6 +412,7 @@ def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float, + chunk_indices: torch.Tensor | None = None, ) -> torch.Tensor: """ Compute the inverse of the matrix I + A @@ -433,7 +434,8 @@ def solve_tril( output_dtype = A.dtype if output_dtype is None else output_dtype B, T, H, BT = A.shape - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) Ai = torch.zeros_like(A, dtype=output_dtype) diff --git a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py index c62c5eb5d2..962b502b5c 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, + activation: 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 activation in ("silu", "sigmoid"), f"unsupported gate activation: {activation}" + # Accept GDN's strided 3D gate without materializing a flattened copy. + assert z.ndim in (2, 3), f"z must be [M, N] or [tokens, heads, N], got shape={z.shape}" assert z.stride(-1) == 1 - 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=activation == "sigmoid", + Z_HEADS=z_heads, num_warps=num_warps, ) return out diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index aa7f5e439a..f2b54c9ffc 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -118,6 +118,10 @@ def load_from_args() -> "LinearAttCacheConfig": model_cfg, _ = PretrainedConfig.get_config_dict(model_path) model_type = model_cfg["model_type"] + if model_type in ("glm5_next", "glm5_next_text"): + from lightllm.models.glm5_next.cache_config import Glm5NextCacheConfig + + return Glm5NextCacheConfig.from_model_config(model_cfg, args) assert model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"] llm_config = model_cfg try: diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py index c7e9a59aad..0cbcf944ac 100644 --- a/lightllm/models/__init__.py +++ b/lightllm/models/__init__.py @@ -56,3 +56,5 @@ from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel from .draft_registry import get_draft_model_class from .registry import get_model, get_model_class + +from .glm5_next.model import Glm5NextTpPartModel diff --git a/lightllm/models/glm5_next/README.md b/lightllm/models/glm5_next/README.md new file mode 100644 index 0000000000..c503e6d412 --- /dev/null +++ b/lightllm/models/glm5_next/README.md @@ -0,0 +1,123 @@ +# GLM-5.3 Flash：首版推理支持 + +基于 main `cd2edb90`，参考 [PR #1525](https://github.com/ModelTC/LightLLM/pull/1525) +移植必要的模型逻辑。使用现有 `bsh_dsv4` 镜像，无需升级依赖或修改模型文件。 + +## 支持范围 + +- 文本生成、原生 FP8 权重、BF16 激活、普通 TP、chunked prefill、decode CUDA graph。 +- KDA、NoPE sparse MLA、4-token K-pool、mHC、GLM 的 sigmoid gated RMSNorm 和 clamped SwiGLU。 +- main 的大小页前缀缓存与请求状态管理；首版验证范围为单机 H200、TP4、8K 内上下文。 +- 暂不支持 MTP、EP、PD 分离、TP/SP 混合、微批重叠和 prefill CUDA graph；模型入口会显式拒绝这些组合。 + +## 缓存布局与流程 + +KDA 的卷积与 FP32 recurrent state 直接使用 `ReqManagerForMamba` 和 +`LinearAttCacheManager`。大小页保存、命中恢复、请求槽位回收、CPU 页搬运均沿用 main。 + +每个稀疏注意力层、每个 token 使用一个 904 元素的 BF16 容器： + +| 区域 | 内容 | +| --- | --- | +| 0–511 | MLA latent KV | +| 512–575 | 零填充，适配镜像已有的 576 维 FlashMLA/FA3 接口 | +| 576–703 | K-pool 原始 index key | +| 704–831 | K-pool 压缩 gate | +| 末尾 132 字节 | 128 维 FP8 压缩 key 与 FP32 scale，仅完整 pool 的末 token 有效 | + +一次 KV 搬运就能带走全部 K-pool 历史。跨 chunk、跨缓存命中的不足 4-token 尾部从原始 +token KV 重建，不引入独立的池化尾状态，也不改 scheduler/radix cache 的生命周期。 +代价是每个注意力层每 token 1808 字节的 KV；11 层合计 19,888 字节，TP 各 rank 复制。 +完整 pool 选择 512 组后展开为最多 2048 个 token，另保留当前未完成 pool 的尾部。 + +共享算子的扩展参数保持原默认值；GLM 显式启用 sigmoid gate、无 `up + 1` 的 clamp、 +KDA 的 exp2 gate。模型特有的权重、attention、索引和 tokenizer 适配放在本目录。 + +## 启动与测速 + +在工作区根目录，使用容器中现有环境（GPU 编号按实际空闲情况调整）： + +```bash +docker exec -d -w /mtc/baishihao/LightLLM \ + -e PYTHONPATH=/mtc/baishihao/LightLLM \ + -e CUDA_VISIBLE_DEVICES=2,3,4,6 -e LOADWORKER=4 \ + bsh_dsv4 bash -lc 'exec python -m lightllm.server.api_server \ + --model_dir /mtc/models/GLM-5.3-Flash --model_name glm53 \ + --tp 4 --host 127.0.0.1 --port 18153 --nccl_port 28153 \ + --max_total_token_num 32768 --max_req_total_len 8192 \ + --batch_max_tokens 2048 --chunked_prefill_size 1024 \ + --running_max_req_size 16 --graph_max_batch_size 8 --graph_max_len_in_batch 8192 \ + --linear_att_hash_page_size 128 --linear_att_page_block_num 4 \ + --linear_att_cache_size 32 --disable_vision --disable_audio \ + --enable_fused_shared_experts > /tmp/glm53.log 2>&1' +``` + +这组参数用 128-token 小页、512-token 大页，方便同时覆盖两种状态恢复路径。 +`linear_att_cache_size` 控制小页 checkpoint 的数量，运行中 KDA state 按请求槽位分配。 +CPU cache 开启时，`cpu_cache_token_page_size` 必须等于大页 token 数。 +本次服务测试未开启完整 CPU KV offload；CPU 页的数据布局与 TP 分片通过下述往返测试验证。 + +```bash +docker exec -w /mtc/baishihao/LightLLM \ + -e PYTHONPATH=/mtc/baishihao/LightLLM bsh_dsv4 \ + python test/benchmark/service/benchmark_glm53_flash.py \ + --input-tokens 1024 4096 --concurrency 1 4 8 \ + --output-tokens 128 --repeats 3 --output /tmp/glm53-benchmark.json +``` + +脚本每组预热，固定输入/输出 token 数，随机化首段 token 避免前缀命中影响 TTFT。 +分别记录首 token 延迟、decode TPOT、每请求 decode tok/s，以及包含 prefill 的输出吞吐。 +权重加载、编译和 graph 捕获不计入请求耗时。测速使用 `ignore_eos=true`，因此不是质量评估。 + +## 实测结果（2026-09-10） + +4 × H200，TP4，镜像 ID `c3f03de5e8dc`；PyTorch `2.11.0+cu130`、Triton `3.6.0`、 +transformers `4.57.1`。启动参数如上，保留 prefix cache，但各测速请求的首段 token 不同。 +先完成各形状的 JIT，再每组预热两轮、正式测三轮；输出固定 128 token。下表取轮次中位数。 + +| 输入 token | 并发 | TTFT（ms） | 每请求 decode（tok/s） | 总输出吞吐，含 prefill（tok/s） | +| ---: | ---: | ---: | ---: | ---: | +| 1024 | 1 | 431 | 99.8 | 75.1 | +| 1024 | 4 | 482 | 85.1 | 257.6 | +| 1024 | 8 | 747 | 68.5 | 391.6 | +| 4096 | 1 | 1160 | 99.6 | 52.6 | +| 4096 | 4 | 1461 | 85.7 | 173.6 | +| 4096 | 8* | 2154 | 64.1 | 150.5 | + +`*` 4096 × 8 再加输出超过本次 32768-token KV 容量，出现排队/分批；该行是容量压力测试， +不能当作完整八路同时 decode 的性能。提高 `max_total_token_num` 后应重新测量。 + +单请求约 10 ms/token 可以作为首版基线，尚未到硬件上限。32 次 decode graph 的独立 profile +中，每 rank 每 token 有 1723 个 GPU kernel，graph 跨度中位数约 9.9 ms。 +按 kernel duration 合计分类：MoE GEMM 约 33%，其他 GEMM/BMM 20%，mHC 13%， +MoE 路由/归并 10%，K-pool 索引 7%，TP 通信 5%，KDA recurrence/conv/gated norm 3%。 +这些百分比用于定位算子耗时，不包含 CPU 调度，也不把 prefill 的通信等待计入 decode。 +18B 激活参数之外，45 层的小矩阵运算、路由和 mHC 都会产生开销；后续应先优化 MoE/GEMM +与算子融合，而不是重写大小页管理流程。 + +首次遇到新形状可能触发十几秒的 JIT，表中是预热后速度。本版没有实现统一的 prefill 形状 +分桶，也没有跑完整精度基准或模型宣称的超长上下文测试。 + +## 验证 + +```bash +docker exec -w /mtc/baishihao/LightLLM \ + -e PYTHONPATH=/mtc/baishihao/LightLLM -e CUDA_VISIBLE_DEVICES=0 bsh_dsv4 \ + python -m pytest -q unit_tests/models/glm5_next +``` + +新增 19 项测试覆盖：KDA chunk/decode 对照 FP32 recurrence；跨 chunk、非连续 KV 和长序列 +K-pool；当前镜像的 NoPE attention；mHC 数值及解码调优；大小页 checkpoint、完整 CPU 页往返； +TP4 的 replicated MLA/index KV 与各 rank 独立 KDA state 区域；transformers 5 tokenizer 文件 +在现有 transformers 4 镜像中的兼容。 + +实际服务验证了中文自我介绍、`17 × 23 = 391`、Python 列表求和函数，并正常遇 EOS 结束。 +为了检查最终答案，使用 checkpoint 的 `reasoning_effort="low"` 模板并追加 ``； +原始默认 Max reasoning 模板也能生成连贯推理，但 256-token 输出上限可能在思考期间截断。 +2179-token 请求的冷/热输出 token ID 完全一致，热请求命中 2176-token 小页；扩展到 +3075 token 的请求命中 2048-token 大页前缀。四并发重复请求、流式中途断开后继续请求均通过， +日志确认请求槽位全部释放。 + +共享流程回归测试 `test_config_utils.py` 与 radix cache 的其余 14 项通过。 +`test_radix_cache.py::test_case10` 在原始 main 上同样失败：没有传入 mem_manager 的实例调用 +`flush_cache()` 触发断言。该既有问题不在本次模型支持中修改。 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/attention.py b/lightllm/models/glm5_next/attention.py new file mode 100644 index 0000000000..8c334d0fec --- /dev/null +++ b/lightllm/models/glm5_next/attention.py @@ -0,0 +1,86 @@ +import dataclasses + +import torch + +from lightllm.common.basemodel.attention.nsa.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): + query_batch: torch.Tensor = None + + def init_state(self): + super().init_state() + state = self.infer_state + self.query_batch = torch.repeat_interleave( + torch.arange(state.batch_size, device=state.b_req_idx.device, dtype=torch.int32), + state.b_q_seq_len, + output_size=state.input_ids.numel(), + ) + + def _nsa_prefill_att(self, q, kv, att_control): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + tokens, heads, dim = q.shape + # The installed Hopper kernel accepts 576-wide Q/K and 64 heads. + # Zero padding preserves NoPE attention and avoids a runtime fork. + padded_heads = ((heads + 63) // 64) * 64 + padded_q = q.new_zeros((tokens, padded_heads, dim + 64)) + padded_q[:, :heads, :dim] = 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=512, + ) + return out[:, :heads] + + +@dataclasses.dataclass +class Glm5NextSparseDecodeState(NsaFlashMlaSparseDecodeAttState): + query_batch: torch.Tensor = None + + def init_state(self): + super().init_state() + state = self.infer_state + self.query_batch = torch.arange(state.batch_size, device=state.b_req_idx.device, dtype=torch.int32) + 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 + + q_nope, _ = q + q_rope = q_nope.new_zeros((*q_nope.shape[:-1], 64)) + params = att_control.nsa_decode_dict + return flash_attn_with_kvcache( + q=q_rope, + qv=q_nope, + k_cache=kv[:, :, 512:].view(-1, 1, 1, 64), + v_cache=kv[:, :, :512].view(-1, 1, 1, 512), + 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, + ) diff --git a/lightllm/models/glm5_next/cache_config.py b/lightllm/models/glm5_next/cache_config.py new file mode 100644 index 0000000000..2cd9359e84 --- /dev/null +++ b/lightllm/models/glm5_next/cache_config.py @@ -0,0 +1,63 @@ +import dataclasses + +from lightllm.common.state_cache_manager import LinearAttCacheConfig +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.torch_dtype_utils import get_torch_dtype + + +@dataclasses.dataclass +class Glm5NextCacheConfig(LinearAttCacheConfig): + """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + + MLA_PADDING = 64 # Existing FlashMLA/FA3 kernels consume a 576-wide MLA key. + 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) + # Keep raw index keys and compression scores with each token, so a + # pool spanning a prefill chunk/cache hit never needs hidden tail state. + packed_dim = ( + config["kv_lora_rank"] + + cls.MLA_PADDING + + 2 * config["index_head_dim"] + + 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, + ) + + 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_main_model_full_att_layer_num() + * page_tokens + ) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py new file mode 100644 index 0000000000..f0ed56e792 --- /dev/null +++ b/lightllm/models/glm5_next/indexer.py @@ -0,0 +1,77 @@ +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv +from .triton_kernel.index_quant import hadamard_transform_quant_fp8 +from .triton_kernel.kpool import compress_pools, gather_pools, expand_topk + + +class Glm5NextNsaInfer: + """K-pool indexing with all persistent history stored in token KV.""" + + 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 _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).unsqueeze(1) + raw_buffer = infer_state.mem_manager.get_indexer_raw_buffer(self.layer_idx) + packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) + destindex_copy_kv(raw, infer_state.mem_index, raw_buffer) + compress_pools( + raw_buffer, + packed_buffer, + layer_weight.index_kpool_compress_ape.weight, + att_state.lengths, + att_state.ks, + att_state.ragged_mem_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 + 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, + ) + lengths = att_state.lengths // 4 + starts = att_state.query_batch * max_pools + ends = starts + lengths + groups = torch.empty((q.shape[0], self.topk // 4), dtype=torch.int32, device=q.device) + # Bound the transient score matrix independently of total batch length. + chunk_size = max(1, min(q.shape[0], 16 * 1024 * 1024 // max_pools)) + import deep_gemm + + pool_positions = torch.arange(max_pools, device=q.device) + + for start in range(0, q.shape[0], chunk_size): + end = min(start + chunk_size, q.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, + ) + # The current image's fast_topk_v2 only supports 2048 entries; + # K-pool selects 512 groups. Torch topk is CUDA-graph compatible. + logits.masked_fill_(pool_positions[None, :] >= lengths[start:end, None], -float("inf")) + groups[start:end] = torch.topk(logits, self.topk // 4, dim=-1, sorted=True).indices + return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) diff --git a/lightllm/models/glm5_next/kda_backend.py b/lightllm/models/glm5_next/kda_backend.py new file mode 100644 index 0000000000..e5690ee828 --- /dev/null +++ b/lightllm/models/glm5_next/kda_backend.py @@ -0,0 +1,175 @@ +# 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 .triton_kernel.kda import chunk_kda_with_fused_gate +from .triton_kernel.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class KDALinearAttBackend(BaseAttBackend): + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_projection_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return KDAPrefillAttState(backend=self, infer_state=infer_state) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return KDADecodeAttState(backend=self, infer_state=infer_state) + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_projection_size, dim=-1) + + def reshape_qkv(self, value: torch.Tensor, *, decode: bool): + if decode: + return value.view(-1, 1, self.tp_num_heads, self.head_dim) + return value.view(1, -1, self.tp_num_heads, self.head_dim) + + +@dataclasses.dataclass +class KDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + chunk_indices: torch.Tensor = None + + 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 + # Build variable-length chunk metadata once for all KDA layers. + self.chunk_indices = prepare_chunk_indices(self.infer_state.b1_cu_q_seq_len, 64) + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + 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.reshape_qkv(x, decode=False) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(1, -1, backend.tp_projection_size) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + + initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + chunk_indices=self.chunk_indices, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +@dataclasses.dataclass +class KDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + 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 + + 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) + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states, + layer_weight.get_merged_kda_conv_weight(), + bias=None, + activation="silu", + conv_state_indices=self.b_conv_buffer_idx, + ) + q, k, v = [backend.reshape_qkv(x, decode=True) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(-1, 1, backend.tp_projection_size) + raw_beta = raw_beta.view(-1, 1, backend.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer_weight.linear_A_log.weight, + gate_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + ssm_state_indices=self.b_ssm_buffer_idx, + ) + return output 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/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..6d6ae93cb1 --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import ( + Deepseek3_2TransformerLayerInfer, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( + silu_and_mul_fwd, +) +from lightllm.models.glm5_next.triton_kernel.mhc import ( + hc_contract, + hc_expand, + hc_post, + hc_pre_norm, +) +from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer + + +class Glm5NextTransformerLayerInfer(Deepseek3_2TransformerLayerInfer): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + self.num_hidden_layers = network_config["num_hidden_layers"] + self.autotune_layer_num = network_config.get("autotune_layer_num", self.num_hidden_layers) + self.is_linear_attention_layer = network_config["layer_types"][layer_num] == "linear_attention" + self.mhc_streams = network_config.get("hc_mult", 4) + self.hc_eps = network_config.get("hc_eps", 1e-6) + self.hc_sinkhorn_iters = network_config.get("hc_sinkhorn_iters", 20) + self.swiglu_limit = network_config["swiglu_limit"] + linear = network_config["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.tp_linear_num_heads = self.linear_num_heads // self.tp_world_size_ + self.tp_linear_projection_size = self.tp_linear_num_heads * self.linear_head_dim + if not self.is_linear_attention_layer: + self.indexer = Glm5NextNsaInfer( + layer_idx=self.layer_num_, + network_config=self.network_config_, + tp_world_size=self.tp_world_size_, + ) + + 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_tp(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)) + 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, + infer_state=infer_state, + ) + + if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: + hidden_states.add_(shared_output) + + return hidden_states.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 _token_attention_kernel(self, q, infer_state, layer_weight, out=None): + if self.is_linear_attention_layer: + raise AssertionError("KDA uses its dedicated backend") + q_nope = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1) + topk_mem_indices, _ = self.indexer._get_indices( + hidden_states=infer_state.get_topk_indices_params["hidden_states"], + q_lora=infer_state.get_topk_indices_params["q_lora"], + infer_state=infer_state, + att_state=infer_state.decode_att_state, + layer_weight=layer_weight, + ) + del infer_state.get_topk_indices_params + q_rope = q_nope[..., :0] + return infer_state.decode_att_state.decode_att( + q=(q_nope, q_rope), + k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + v=None, + att_control=AttControl( + nsa_decode=True, + nsa_decode_dict={ + "layer_index": self.layer_num_, + "topk_mem_indices": topk_mem_indices, + "softmax_scale": self.softmax_scale, + "kv_lora_rank": self.kv_lora_rank, + "qk_rope_head_dim": 0, + }, + ), + ) + + def _kda_projections(self, input, infer_state, layer_weight): + # KDA shards heads across TP ranks, so every rank still needs every + # token before updating its recurrent head state. In TP/SP mode the + # layer input is sequence-sharded; gather it here just like the MLA + # projection path and reduce-scatter the output in _kda_post. + input = input.view(-1, self.embed_dim_) + input = self._tpsp_allgather(input=input, infer_state=infer_state) + projected = layer_weight.linear_qkvbfg_a_proj.mm(input) + qkv_size = 3 * self.tp_linear_projection_size + mixed_qkv, raw_beta, f_a, g_a = projected.split( + [ + qkv_size, + self.tp_linear_num_heads, + self.linear_head_dim, + self.linear_head_dim, + ], + dim=-1, + ) + raw_gate, norm_gate = layer_weight.project_kda_fg_b(f_a, g_a) + return mixed_qkv, raw_gate, raw_beta, norm_gate + + def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): + tokens = norm_gate.shape[0] + core_output = core_output.view(-1, self.linear_head_dim) + norm_gate = norm_gate.view(tokens, self.tp_linear_num_heads, self.linear_head_dim) + output = layer_weight.linear_o_norm( + input=core_output, + gate_value=norm_gate, + eps=self.eps_, + alloc_func=self.alloc_tensor, + ) + output = layer_weight.linear_o_proj.mm(output.view(tokens, -1)) + return self._tpsp_reduce(input=output, infer_state=infer_state) + + def context_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().context_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) + core_output = infer_state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_prefill=True, + linear_att_prefill_dict={ + "mixed_qkv": mixed_qkv, + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, norm_gate, infer_state, layer_weight) + + def token_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().token_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) + core_output = infer_state.decode_att_state1.decode_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_decode=True, + linear_att_decode_dict={ + "mixed_qkv": mixed_qkv, + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, norm_gate, infer_state, layer_weight) + + def _hc_pre(self, streams, layer_weight, prefix, norm_weight): + return hc_pre_norm( + x=streams, + fn=getattr(layer_weight, f"hc_{prefix}_fn").weight, + scale=getattr(layer_weight, f"hc_{prefix}_scale").weight, + base=getattr(layer_weight, f"hc_{prefix}_base").weight, + norm_weight=norm_weight.weight, + streams=self.mhc_streams, + rms_eps=self.eps_, + norm_eps=self.eps_, + hc_eps=self.hc_eps, + sinkhorn_iters=self.hc_sinkhorn_iters, + ) + + def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): + streams = input_embeddings + if self.layer_num_ == 0: + streams = hc_expand(streams.view(-1, self.embed_dim_), self.mhc_streams) + + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) + if prefill: + layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) + else: + layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) + 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): + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=True) + + def token_forward(self, 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..3d39a3cda9 --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -0,0 +1,340 @@ +# 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, + GatedRMSNormWeight, + LayerNormWeight, + ParameterWeight, + RMSNormWeight, + ROWMMWeight, + TpParameterWeight, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import ( + get_row_slice_mixin, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import ( + MMWeightTpl, +) +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import ( + Deepseek2TransformerLayerWeight, +) +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import ( + Deepseek3_2TransformerLayerWeight, +) +from .pre_and_post_layer_weight import add_language_model_aliases + + +class Glm5NextMergedKdaProjection(MMWeightTpl): + """One KDA input GEMM with TP-sharded q/k/v/b and replicated f_a/g_a.""" + + def __init__( + self, + in_dim: int, + projection: int, + head_count: int, + head_dim: int, + weight_names: list[str], + data_type: torch.dtype, + tp_rank: int | None = None, + tp_world_size: int | None = None, + ): + tp_rank = get_current_rank_in_dp() if tp_rank is None else tp_rank + tp_world_size = get_dp_world_size() if tp_world_size is None else tp_world_size + assert projection % tp_world_size == 0 + assert head_count % tp_world_size == 0 + super().__init__( + in_dim=in_dim, + out_dims=[ + projection // tp_world_size, + projection // tp_world_size, + projection // tp_world_size, + head_count // tp_world_size, + head_dim, + head_dim, + ], + weight_names=weight_names, + bias_names=None, + data_type=data_type, + quant_method=None, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + self.sharded_slicer = get_row_slice_mixin("none", tp_rank=tp_rank, tp_world_size=tp_world_size) + self.replicated_slicer = get_row_slice_mixin("none", tp_rank=0, tp_world_size=1) + + def _get_param_slicer(self, sub_child_index: int): + return self.replicated_slicer if sub_child_index >= 4 else self.sharded_slicer + + +class Glm5NextTransformerLayerWeight(Deepseek3_2TransformerLayerWeight): + def _init_moe(self): + super()._init_moe() + self.experts.fuse_moe_impl.swiglu_limit = self.network_config_["swiglu_limit"] + self.experts.fuse_moe_impl.swiglu_clamp_up_add_one = False + self.moe_gate = ROWMMWeight( + in_dim=self.n_embed, + out_dims=[self.n_routed_experts], + weight_names=f"model.layers.{self.layer_num_}.mlp.gate.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + + def _parse_config(self): + super()._parse_config() + # The released sparse MLA keeps kv_b_proj in BF16 even though the + # surrounding projections are native FP8. Its compressed-context + # shortcut assumes a quantized kv_b matrix, so GLM uses the BMM split. + self.enable_cc_method = False + self.is_linear_attention_layer = self.network_config_["layer_types"][self.layer_num_] == "linear_attention" + linear = self.network_config_["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.linear_projection_size = self.linear_num_heads * self.linear_head_dim + self.linear_conv_kernel_size = linear["short_conv_kernel_size"] + self.mhc_streams = self.network_config_.get("hc_mult", 4) + + def _init_weight(self): + if self.is_linear_attention_layer: + self._init_kda() + else: + Deepseek2TransformerLayerWeight._init_qkvo(self) + self._init_indexer_weight() + + if self.is_moe: + self._init_moe() + else: + self._init_ffn() + self._init_glm_norms() + self._init_mhc() + + def _init_kda(self): + prefix = f"model.layers.{self.layer_num_}.self_attn" + projection = self.linear_projection_size + head_count = self.linear_num_heads + head_dim = self.linear_head_dim + + self.linear_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_, + activation="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.hidden_size, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.wk.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.k_norm_ = LayerNormWeight( + dim=self.index_head_dim, + weight_name=f"{prefix}.k_norm.weight", + data_type=self.data_type_, + bias_name=f"{prefix}.k_norm.bias", + ) + self.weights_proj_ = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_n_heads], + weight_names=f"{prefix}.weights_proj.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_gate = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.index_kpool_compress_gate", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_ape = ParameterWeight( + weight_name=f"{prefix}.index_kpool_compress_ape", + data_type=torch.float32, + weight_shape=(self.network_config_["index_kpool"], self.index_head_dim), + ) + + def _init_glm_norms(self): + prefix = f"model.layers.{self.layer_num_}" + self.att_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.input_layernorm.weight", + data_type=self.data_type_, + ) + self.ffn_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.post_attention_layernorm.weight", + data_type=self.data_type_, + ) + if not self.is_linear_attention_layer: + self.kv_a_layernorm_ = RMSNormWeight( + dim=self.kv_lora_rank, + weight_name=f"{prefix}.self_attn.kv_a_layernorm.weight", + data_type=self.data_type_, + ) + self.q_a_layernorm_ = RMSNormWeight( + dim=self.q_lora_rank, + weight_name=f"{prefix}.self_attn.q_a_layernorm.weight", + data_type=self.data_type_, + ) + + def _init_mhc(self): + prefix = f"model.layers.{self.layer_num_}" + streams = self.mhc_streams + mix_size = (2 + streams) * streams + flattened_hidden = streams * self.n_embed + self.hc_attn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_attn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_attn_base = ParameterWeight( + weight_name=f"{prefix}.hc_attn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_attn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_attn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + self.hc_ffn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_ffn_base = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_ffn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + + def get_merged_kda_conv_weight(self): + return self.linear_qkv_conv1d.mm_param.weight + + def project_kda_fg_b(self, f_a: torch.Tensor, g_a: torch.Tensor): + method = self.linear_fg_b_proj.quant_method + f = method.apply(f_a, self.linear_fg_b_proj.mm_param_list[0]) + g = method.apply(g_a, self.linear_fg_b_proj.mm_param_list[1]) + return f, g + + def _preprocess_kda_weights(self, weights): + prefix = f"model.layers.{self.layer_num_}.self_attn" + for projection in ("q", "k", "v"): + name = f"{prefix}.{projection}_conv1d.weight" + if name in weights and weights[name].ndim == 3: + weights[name] = weights[name].squeeze(1) + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + + # GLM checkpoints nest the shared expert under + # ``mlp.shared_experts``. This class deliberately bypasses + # Deepseek2TransformerLayerWeight.load_hf_weights below, so perform + # the fused-shared remap here before the generic loader consumes the + # expert tensors. + if self.num_fused_shared_experts > 0: + self._rename_shared_experts( + weights, + self.experts.quant_method.weight_scale_suffix, + ) + + if self.is_linear_attention_layer: + self._preprocess_kda_weights(weights) + return TransformerLayerWeight.load_hf_weights(self, weights) + + kv_b_name = f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight" + if kv_b_name in weights: + k_b_proj, v_b_proj = self._split_kv_b_proj(weights[kv_b_name]) + weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = k_b_proj + weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = v_b_proj + + return TransformerLayerWeight.load_hf_weights(self, weights) diff --git a/lightllm/models/glm5_next/mem_manager.py b/lightllm/models/glm5_next/mem_manager.py new file mode 100644 index 0000000000..cda714e2f1 --- /dev/null +++ b/lightllm/models/glm5_next/mem_manager.py @@ -0,0 +1,44 @@ +import torch + +from lightllm.common.kv_cache_mem_manager.operator import LinearAttMemOperator +from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager +from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + +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, index_head_dim=128, **kwargs): + self.mla_head_dim = mla_head_dim + self.index_head_dim = index_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.kv_buffer[:, :, :, self.mla_head_dim : self.mla_head_dim + 64].zero_() + 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 + 64] + + def get_indexer_raw_buffer(self, layer_index): + start = self.mla_head_dim + 64 + return self._layer_buffer(layer_index)[:, :, start : start + 2 * self.index_head_dim] + + def get_indexer_k_buffer(self, layer_index): + return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py new file mode 100644 index 0000000000..d9a7ff87d7 --- /dev/null +++ b/lightllm/models/glm5_next/model.py @@ -0,0 +1,91 @@ +import json +import os + +import torch +import triton + +from lightllm.common.build_utils import repair_config +from lightllm.common.req_manager import ReqManagerForMamba +from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel +from lightllm.models.registry import ModelRegistry +from .attention import Glm5NextSparseAttBackend +from .cache_config import Glm5NextCacheConfig +from .kda_backend import KDALinearAttBackend +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 +from .mem_manager import Glm5NextMemManager + + +@ModelRegistry(["glm5_next", "glm5_next_text"]) +class Glm5NextTpPartModel(Deepseek3_2TpPartModel): + pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight + transformer_weight_class = Glm5NextTransformerLayerWeight + 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["quantization_config"].setdefault("scale_fmt", "ue8m0") + self.config["autotune_layer_num"] = 4 + for names in ( + ["num_attention_heads", "n_head"], + ["hidden_size", "n_embd", "n_embed"], + ["num_hidden_layers", "n_layer"], + ): + repair_config(self.config, same_names=names) + + def _verify_params(self): + super()._verify_params() + args = self.args + assert self.data_type == torch.bfloat16, "GLM-5.3 Flash currently requires bfloat16 activations" + assert self.run_mode == "normal", "GLM-5.3 Flash v1 supports normal TP serving" + assert args.dp == 1 and not args.enable_tpsp_mix_mode, "GLM-5.3 Flash v1 uses plain tensor parallelism" + assert args.mtp_mode is None and args.mtp_step == 0, "GLM-5.3 Flash MTP is not implemented yet" + assert not args.enable_ep_moe, "GLM-5.3 Flash v1 uses tensor-parallel MoE" + assert not ( + args.enable_prefill_microbatch_overlap or args.enable_decode_microbatch_overlap + ), "GLM-5.3 Flash mHC does not support microbatch overlap yet" + assert not args.enable_prefill_cudagraph, "GLM-5.3 Flash v1 supports decode CUDA graphs" + assert args.llm_kv_type in (None, "None"), "GLM-5.3 Flash v1 uses BF16 MLA KV with FP8 index keys" + assert self.config["qk_rope_head_dim"] == 0 and self.config["kv_lora_rank"] == 512 + assert self.config["index_head_dim"] == 128 and self.config["index_kpool"] == 4 + + def autotune_layers(self): + return 4 + + def _init_req_manager(self): + self.linear_config = Glm5NextCacheConfig.from_model_config(self.config, self.args) + self.req_manager = ReqManagerForMamba( + 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_main_model_full_att_layer_num(), + 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)) + self._cos_cached = torch.empty((self.max_seq_length, 0), dtype=self.data_type, device="cuda") + self._sin_cached = torch.empty_like(self._cos_cached) diff --git a/lightllm/models/glm5_next/tokenizer.py b/lightllm/models/glm5_next/tokenizer.py new file mode 100644 index 0000000000..0866ff47f7 --- /dev/null +++ b/lightllm/models/glm5_next/tokenizer.py @@ -0,0 +1,17 @@ +from transformers import PreTrainedTokenizerFast +from transformers.models.auto.tokenization_auto import get_tokenizer_config + + +def get_glm5_next_tokenizer(tokenizer_name, *args, **kwargs): + """Load the checkpoint's Rust tokenizer with the existing transformers 4 image. + + GLM-5.3 exports the transformers 5 TokenizersBackend name and a list of + extra_special_tokens. Transformers 4 calls that list additional_special_tokens. + The tokenizer.json and checkpoint chat template remain authoritative. + """ + config = get_tokenizer_config(tokenizer_name, **kwargs) + extra_tokens = config.get("extra_special_tokens", {}) + if isinstance(extra_tokens, list): + kwargs.setdefault("additional_special_tokens", extra_tokens) + kwargs["extra_special_tokens"] = {} + return PreTrainedTokenizerFast.from_pretrained(tokenizer_name, *args, **kwargs) 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/kda.py b/lightllm/models/glm5_next/triton_kernel/kda.py new file mode 100644 index 0000000000..f2d42919a4 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/kda.py @@ -0,0 +1,1066 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang + +"""KDA helpers built on LightLLM's continuous-batching recurrent kernel.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.cumsum import chunk_local_cumsum +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd +from triton.language import exp2, log +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril import solve_tril + + +FLA_CHUNK_SIZE = 64 +RCP_LN2 = 1.4426950216293335 + + +def cdiv(a: int, b: int) -> int: + return -(a // -b) + + +def next_power_of_2(n: int) -> int: + return 1 if n < 1 else 1 << (n - 1).bit_length() + + +def kda_safe_gate( + raw_gate: torch.Tensor, + a_log: torch.Tensor, + gate_bias: torch.Tensor, + lower_bound: float = -5.0, +) -> torch.Tensor: + """GLM-5 bounded KDA decay in fp32. + + ``raw_gate`` is ``[..., heads, key_dim]``; ``a_log`` is per-head and + ``gate_bias`` is per head/key coordinate. + """ + + head_count = a_log.numel() + key_dim = gate_bias.numel() // head_count + gate = raw_gate.float().view(*raw_gate.shape[:-1], head_count, key_dim) + amplitude = a_log.float().reshape(*((1,) * (gate.ndim - 2)), head_count, 1).exp() + bias = gate_bias.float().reshape(*((1,) * (gate.ndim - 2)), head_count, key_dim) + return lower_bound * torch.sigmoid(amplitude * (gate + bias)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_kt = tl.make_block_ptr(k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=["BK", "BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K + + +def chunk_kda_scaled_dot_kkt_fwd( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + assert K <= 256 + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BC = min(16, BT) + NC = cdiv(BT, BC) + BK = max(next_power_of_2(K), 16) + A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + grid = (NT, NC * NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + return A, Aqk + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, i_v * BV), + (BK, BV), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[triton.Config({"BD": BD}, num_warps=num_warps) for BD in [32, 64] for num_warps in [2, 4, 8]], + key=["H", "D", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_cumsum_fwd_kernel( + g, + A, + y, + g_bias, + cu_seqlens, + chunk_indices, + cumsum_scale, + beta, + threshold, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + p_y = tl.make_block_ptr( + y + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + o_d = i_d * BD + tl.arange(0, BD) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + b_a = tl.load(A + i_h).to(tl.float32) + b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) + if SAFE_GATE: + # y = lower_bound * sigmoid(exp(A) * (g + g_bias)); bounded to + # (lower_bound, 0). Mirrors the SGlang safe_gate branch used by GLM5-Next + # checkpoints whose linear_attn_config["safe_gate"] is True. + b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g))) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = b_a * b_softplus + + # Out-of-bounds rows (load returns 0, but softplus/bias can still make + # b_gate non-zero) participate in the dot product. They only contribute to + # out-of-bounds output rows, which are masked away by `boundary_check` on + # the store, so visible output matches unfused gate + chunk-local cumsum. + o_t = tl.arange(0, BT) + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, + safe_gate: bool = False, + lower_bound: float = -5.0, +) -> torch.Tensor: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" + B, T, H, D = raw_g.shape + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + + def grid(meta): + return (cdiv(meta["D"], meta["BD"]), NT, B * H) + + kda_gate_cumsum_fwd_kernel[grid]( + g=raw_g, + A=A_log, + y=y, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + beta=beta, + threshold=threshold, + SAFE_GATE=safe_gate, + LOWER_BOUND=lower_bound, + T=T, + H=H, + D=D, + BT=chunk_size, + ) + return y + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + # the intra Aqk is kept in fp32 + # the computation has very marginal effect on the entire throughput + A, Aqk = chunk_kda_scaled_dot_kkt_fwd( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, +): + chunk_size = FLA_CHUNK_SIZE + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + g = fused_kda_gate_chunk_cumsum( + raw_g, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, + **kwargs, +): + """Run chunk KDA from raw gate projection using fused gate+cumsum.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + beta=beta.contiguous(), + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return o, final_state diff --git a/lightllm/models/glm5_next/triton_kernel/kda_decode.py b/lightllm/models/glm5_next/triton_kernel/kda_decode.py new file mode 100644 index 0000000000..6482114f28 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/kda_decode.py @@ -0,0 +1,91 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _kda_decode( + Q, + K, + V, + G, + B, + A, + Bias, + State, + Idx, + O, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SG: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + LOWER: tl.constexpr, + BV: tl.constexpr, +): + row_head = tl.program_id(1) + row, head = row_head // H, row_head % H + ki = tl.arange(0, D) + vi = tl.program_id(0) * BV + tl.arange(0, BV) + q = tl.load(Q + row * SQ + head * D + ki).to(tl.float32) + k = tl.load(K + row * SK + head * D + ki).to(tl.float32) + v = tl.load(V + row * 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 + row * SG + head * D + ki).to(tl.float32) + bias = tl.load(Bias + head * D + ki) + amplitude = tl.exp(tl.load(A + head)) + decay = tl.exp(LOWER * tl.sigmoid(amplitude * (gate + bias))) + beta = tl.sigmoid(tl.load(B + row * SB + head).to(tl.float32)) + req = tl.load(Idx + row) + ptr = State + (req * H + head) * D * D + ki[:, None] * D + vi[None, :] + state = tl.load(ptr).to(tl.float32) * decay[:, None] + delta = (v - tl.sum(state * k[:, None], 0)) * beta + state += k[:, None] * delta[None, :] + tl.store(ptr, state) + out = tl.sum(state * q[:, None], 0) + tl.store(O + row_head * D + vi, out) + + +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, +): + assert inplace_final_state and q.shape[1] == 1 + batch, _, heads, dim = q.shape + assert dim == 128 + 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, + out, + q.stride(0), + k.stride(0), + v.stride(0), + raw_gate.stride(0), + raw_beta.stride(0), + heads, + dim, + lower_bound, + 32, + num_warps=4, + ) + return out, initial_state 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..9dd267f565 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -0,0 +1,137 @@ +import torch +import triton +import triton.language as tl + +from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage + + +@triton.jit +def _compress_pools(Raw, Packed, Ape, Lengths, Starts, Ragged, RAW_STRIDE: tl.constexpr, PACKED_STRIDE: tl.constexpr): + row = tl.program_id(0) + length = tl.load(Lengths + row) + if length % 4 == 0 and length > 0: + start = tl.load(Starts + row) + pool = tl.arange(0, 4) + cols = tl.arange(0, 128) + locs = tl.load(Ragged + start + length - 4 + pool) + raw = tl.load(Raw + locs[:, None] * RAW_STRIDE + cols[None, :]).to(tl.float32) + score = tl.load(Raw + locs[:, None] * RAW_STRIDE + 128 + cols[None, :]).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) + 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)) + + +def compress_pools(raw_buffer, packed_buffer, ape, lengths, starts, ragged): + _compress_pools[(lengths.numel(),)]( + raw_buffer, + packed_buffer, + ape, + lengths, + starts, + ragged, + raw_buffer.stride(0), + packed_buffer.stride(0), + num_warps=4, + ) + + +@triton.jit +def _gather_pools( + Packed, + ReqTable, + ReqIdx, + SeqLen, + K, + Scale, + PACKED_STRIDE: tl.constexpr, + REQ_STRIDE: tl.constexpr, + POOLS: tl.constexpr, +): + batch, pool = 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) + 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) + tl.store(K + (batch * POOLS + pool) * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) + tl.store(Scale + batch * POOLS + pool, 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) + _gather_pools[(req_idx.numel(), max_pools)]( + 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 _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/triton_kernel/mhc.py b/lightllm/models/glm5_next/triton_kernel/mhc.py new file mode 100644 index 0000000000..87eadf7b50 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/mhc.py @@ -0,0 +1,547 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC operators used by GLM-5-Next. + +The public entry points use fused Triton kernels for the small, launch-bound +mixing operations. The explicit PyTorch implementations remain available as +correctness oracles: all mixing math is accumulated in fp32 and only the +collapsed layer input / expanded residual streams are cast back to the +activation dtype. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + + +@triton.jit +def _hc_prepare_kernel( + mixes, + scale, + base, + pre, + post, + residual_mix, + mix_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, +): + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + + pre_raw = tl.load(mixes + token * mix_stride_m + stream_offsets) + post_raw = tl.load(mixes + token * mix_stride_m + STREAMS + stream_offsets) + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) + + logits = tl.load(mixes + token * mix_stride_m + 2 * STREAMS + matrix_offsets) + logits = logits * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) + logits = tl.reshape(logits, (STREAMS, STREAMS)) + logits = logits - tl.max(logits, axis=1)[:, None] + matrix = tl.exp(logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + + # The checkpoint definition starts with a column normalization, then + # alternates row and column normalizations for the remaining iterations. + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_prepare_prenorm_kernel( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_stride_s: tl.constexpr, + gemm_stride_m: tl.constexpr, + sqrsum_stride_s: tl.constexpr, + sqrsum_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + FLATTENED_HIDDEN: tl.constexpr, + RMS_EPS: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, + N_SPLITS: tl.constexpr, +): + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + pre_raw = tl.zeros((STREAMS,), dtype=tl.float32) + post_raw = tl.zeros((STREAMS,), dtype=tl.float32) + matrix_raw = tl.zeros((STREAMS * STREAMS,), dtype=tl.float32) + sqrsum = 0.0 + for split in tl.static_range(N_SPLITS): + partial_base = gemm_partial + split * gemm_stride_s + token * gemm_stride_m + pre_raw += tl.load(partial_base + stream_offsets) + post_raw += tl.load(partial_base + STREAMS + stream_offsets) + matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) + sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) + inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) + pre_raw *= inv_rms + post_raw *= inv_rms + matrix_raw *= inv_rms + + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) + + logits = matrix_raw * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) + logits = tl.reshape(logits, (STREAMS, STREAMS)) + logits = logits - tl.max(logits, axis=1)[:, None] + matrix = tl.exp(logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_pre_combine_kernel( + x, + pre, + output, + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + tl.store( + output + token * out_stride_m + hidden_offsets, + accumulator, + mask=hidden_mask, + ) + + +@triton.jit +def _hc_pre_combine_norm_kernel( + x, + pre, + norm_weight, + output, + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + NORM_EPS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_offsets = tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + + # hc_pre returns bf16 before the following RMSNorm in the checkpoint + # definition. Preserve that rounding point while keeping both operations + # in one kernel. + rounded = accumulator.to(tl.bfloat16).to(tl.float32) + variance = tl.sum(rounded * rounded, axis=0) / hidden + inv_rms = tl.rsqrt(variance + NORM_EPS) + weight = tl.load(norm_weight + hidden_offsets, mask=hidden_mask, other=0.0) + tl.store( + output + token * out_stride_m + hidden_offsets, + rounded * inv_rms * weight, + mask=hidden_mask, + ) + + +@triton.jit +def _hc_post_4stream_kernel( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden: tl.constexpr, + layer_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + mix_stride_m: tl.constexpr, + post_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + + # GLM-5 always uses four mHC streams. Compute all four outputs in one + # program so the layer output and residual streams are read only once. + # The previous output-stream grid reread each residual stream four times; + # that becomes bandwidth-bound for large prefills. + layer_value = tl.load( + layer_output + token * layer_stride_m + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + residual_base = residual + token * residual_stride_m + hidden_offsets + residual_0 = tl.load(residual_base, mask=hidden_mask, other=0.0).to(tl.float32) + residual_1 = tl.load(residual_base + hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_2 = tl.load(residual_base + 2 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_3 = tl.load(residual_base + 3 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + + post_base = post_mix + token * post_stride_m + mix_base = residual_mix + token * mix_stride_m + accumulator_0 = layer_value * tl.load(post_base) + accumulator_1 = layer_value * tl.load(post_base + 1) + accumulator_2 = layer_value * tl.load(post_base + 2) + accumulator_3 = layer_value * tl.load(post_base + 3) + + accumulator_0 += residual_0 * tl.load(mix_base) + accumulator_0 += residual_1 * tl.load(mix_base + 4) + accumulator_0 += residual_2 * tl.load(mix_base + 8) + accumulator_0 += residual_3 * tl.load(mix_base + 12) + accumulator_1 += residual_0 * tl.load(mix_base + 1) + accumulator_1 += residual_1 * tl.load(mix_base + 5) + accumulator_1 += residual_2 * tl.load(mix_base + 9) + accumulator_1 += residual_3 * tl.load(mix_base + 13) + accumulator_2 += residual_0 * tl.load(mix_base + 2) + accumulator_2 += residual_1 * tl.load(mix_base + 6) + accumulator_2 += residual_2 * tl.load(mix_base + 10) + accumulator_2 += residual_3 * tl.load(mix_base + 14) + accumulator_3 += residual_0 * tl.load(mix_base + 3) + accumulator_3 += residual_1 * tl.load(mix_base + 7) + accumulator_3 += residual_2 * tl.load(mix_base + 11) + accumulator_3 += residual_3 * tl.load(mix_base + 15) + + output_base = output + token * out_stride_m + hidden_offsets + tl.store(output_base, accumulator_0, mask=hidden_mask) + tl.store(output_base + hidden, accumulator_1, mask=hidden_mask) + tl.store(output_base + 2 * hidden, accumulator_2, mask=hidden_mask) + tl.store(output_base + 3 * hidden, accumulator_3, mask=hidden_mask) + + +def hc_expand(x: torch.Tensor, streams: int) -> torch.Tensor: + """Expand ``[tokens, hidden]`` into flattened residual streams.""" + + assert x.ndim == 2 + return x.unsqueeze(1).expand(-1, streams, -1).reshape(x.shape[0], -1) + + +def hc_contract(x: torch.Tensor, streams: int) -> torch.Tensor: + """Contract flattened residual streams by taking their mean.""" + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + return x.view(x.shape[0], streams, -1).mean(dim=1) + + +def hc_pre_reference( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + streams: int, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute mHC pre-mixes. + + Returns ``(layer_input, residual_mix, post_mix)``. ``x`` and + ``layer_input`` use the activation dtype; both mix tensors are fp32. + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + residual = x.view(tokens, streams, hidden) + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre_raw = mixes[:, :streams] + post_raw = mixes[:, streams : 2 * streams] + residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) + + pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps + post = post_multiplier * torch.sigmoid(post_raw * scale[1] + base[streams : 2 * streams]) + residual_mix = (residual_raw * scale[2] + base[2 * streams :].view(streams, streams)).softmax(dim=-1) + residual_mix = residual_mix + hc_eps + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + for _ in range(sinkhorn_iters - 1): + residual_mix = residual_mix / (residual_mix.sum(dim=-1, keepdim=True) + hc_eps) + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + + layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) + return layer_input, residual_mix, post + + +def hc_post_reference( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """Mix a sublayer output back into the flattened residual streams.""" + + tokens, hidden = layer_output.shape + residual_3d = residual.view(tokens, streams, hidden) + mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) + out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual + return out.to(layer_output.dtype).reshape(tokens, streams * hidden) + + +def hc_pre( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + streams: int, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute mHC pre-mixes with fused Sinkhorn and residual combining.""" + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert x.is_contiguous() and fn.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + _hc_prepare_kernel[(tokens,)]( + mixes, + scale, + base, + pre, + post, + residual_mix, + mixes.stride(0), + pre.stride(0), + residual_mix.stride(0), + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + num_warps=1, + ) + + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_pre_combine_kernel[(tokens, triton.cdiv(hidden, block_h))]( + x, + pre, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post + + +def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.device) -> int: + grid_size = triton.cdiv(tokens, 64) + k_blocks = triton.cdiv(flattened_hidden, 64) + sms = torch.cuda.get_device_properties(device).multi_processor_count + return max(1, min(sms // max(grid_size, 1), k_blocks // 4)) + + +def hc_pre_norm( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + streams: int, + rms_eps: float, + norm_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fuse mHC pre-mixing with its immediately following RMSNorm. + + DeepGEMM computes the fp32 projection and residual square sum in one + split-K kernel. Triton then reduces those partials, runs Sinkhorn, forms + the bf16 residual collapse, and applies RMSNorm. + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert x.dtype == torch.bfloat16 and fn.dtype == torch.float32 + assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + try: + import deep_gemm + + prenorm_gemm = deep_gemm.tf32_hc_prenorm_gemm + except (AttributeError, ImportError): + from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import ( + rmsnorm_forward, + ) + + layer_input, residual_mix, post = hc_pre( + x, + fn, + scale, + base, + streams, + rms_eps, + hc_eps, + sinkhorn_iters, + post_multiplier, + ) + layer_input = rmsnorm_forward(layer_input, weight=norm_weight, eps=norm_eps) + return layer_input, residual_mix, post + + mix_size = (2 + streams) * streams + n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) + gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device) + sqrsum_partial = torch.empty((n_splits, tokens), dtype=torch.float32, device=x.device) + prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + _hc_prepare_prenorm_kernel[(tokens,)]( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_partial.stride(0), + gemm_partial.stride(1), + sqrsum_partial.stride(0), + sqrsum_partial.stride(1), + pre.stride(0), + residual_mix.stride(0), + FLATTENED_HIDDEN=flattened_hidden, + RMS_EPS=rms_eps, + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + N_SPLITS=n_splits, + num_warps=1, + ) + + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) + block_h = triton.next_power_of_2(hidden) + _hc_pre_combine_norm_kernel[(tokens,)]( + x, + pre, + norm_weight, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + NORM_EPS=norm_eps, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post + + +def hc_post( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """Mix a sublayer output into residual streams with one Triton launch.""" + + tokens, hidden = layer_output.shape + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert layer_output.is_contiguous() and residual.is_contiguous() + output = torch.empty( + (tokens, streams * hidden), + dtype=layer_output.dtype, + device=layer_output.device, + ) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_post_4stream_kernel[(tokens, triton.cdiv(hidden, block_h))]( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden=hidden, + layer_stride_m=layer_output.stride(0), + residual_stride_m=residual.stride(0), + mix_stride_m=residual_mix.stride(0), + post_stride_m=post_mix.stride(0), + out_stride_m=output.stride(0), + BLOCK_H=block_h, + num_warps=8, + ) + return output diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index e1a4e421d1..53404d3aff 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -63,6 +63,13 @@ def get_tokenizer( # tokenizer = convert_slow_tokenizer(tokenizer) # return tokenizer + model_cfg, _ = PretrainedConfig.get_config_dict(tokenizer_name) + model_type = model_cfg.get("model_type", "") + if model_type in ("glm5_next", "glm5_next_text"): + from ..models.glm5_next.tokenizer import get_glm5_next_tokenizer + + return get_glm5_next_tokenizer(tokenizer_name, *args, **kwargs) + try: tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=trust_remote_code, *args, **kwargs) except TypeError as e: @@ -78,8 +85,6 @@ def get_tokenizer( "slowdown. Consider using a fast tokenizer instead." ) - model_cfg, _ = PretrainedConfig.get_config_dict(tokenizer_name) - model_type = model_cfg.get("model_type", "") # DeepSeek-V3.2 custom tokenizer mode: wraps the HF tokenizer with # a Python-based apply_chat_template that uses encoding_dsv32.py. if model_type == "deepseek_v32": diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 4bd78d887f..bc7fb1437f 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -459,7 +459,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/test/benchmark/service/benchmark_glm53_flash.py b/test/benchmark/service/benchmark_glm53_flash.py new file mode 100644 index 0000000000..de3a451898 --- /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="/mtc/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/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py new file mode 100644 index 0000000000..a6caad8ec3 --- /dev/null +++ b/unit_tests/models/glm5_next/test_cache.py @@ -0,0 +1,162 @@ +import dataclasses +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.req_manager import ReqManagerForMamba +from lightllm.models.glm5_next.cache_config import Glm5NextCacheConfig +from lightllm.models.glm5_next.mem_manager import Glm5NextMemManager +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.mark.parametrize("small_page", [False, True]) +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size): + 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, + ) + 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=904, + 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, + ) + monkeypatch.setattr("lightllm.common.state_cache_manager.LinearAttCacheConfig.load_from_args", lambda: config) + mem = Glm5NextMemManager(16, torch.bfloat16, 1, 904, 1, config) + req = ReqManagerForMamba(3, 16, mem, config) + 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].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_() + 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], conv) + assert torch.equal(req.req_to_ssm_state.buffer[:, 2], ssm) + 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].any() + # KV moves must carry raw index keys, compression gates and pooled FP8 + # bytes together; bytewise equality catches omissions and scale corruption. + 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() == 904 * 2 + 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("activation", ["silu", "sigmoid"]) +def test_gated_norm_activation_and_strided_gate(activation): + 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, activation=activation) + 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 activation == "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} + silu_and_mul_fwd(x, out, limit=10.0, alpha=1.0, **kwargs) + gate, up = x.float().chunk(2, -1) + gate = torch.nn.functional.silu(gate.clamp(max=10)).bfloat16().float() + expected = gate * (up.clamp(-10, 10) + 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_kernels.py b/unit_tests/models/glm5_next/test_kernels.py new file mode 100644 index 0000000000..f719c090a0 --- /dev/null +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -0,0 +1,308 @@ +import dataclasses +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.models.glm5_next.triton_kernel.kda import chunk_kda_with_fused_gate +from lightllm.models.glm5_next.triton_kernel.kda_decode import fused_recurrent_kda +from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools, gather_pools, expand_topk +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 +from lightllm.models.glm5_next.triton_kernel.mhc import hc_pre_norm, hc_pre_reference, hc_post, hc_post_reference +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward + + +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) + + +def test_kpool_chunk_boundaries_and_fragmented_token_kv(): + # Two unaligned requests, with a pool completed after restoring token KV. + 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(4, 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, 904, device="cuda", dtype=torch.bfloat16) + raw = packed_storage[:, :, 576:832] + packed = packed_storage.view(torch.uint8)[:, :, -132:] + for first, end in [(0, 3), (3, 9), (9, 11), (11, 16)]: + raw[ragged[first:end].long(), 0] = source[first:end] + start = 0 if first < 9 else 9 + lengths = torch.arange(first - start + 1, end - start + 1, device="cuda", dtype=torch.int32) + starts = torch.full_like(lengths, start) + compress_pools(raw, packed, ape, lengths, starts, ragged) + # Model the full KV copy performed by cache offload/load or request move. + packed_storage = packed_storage.clone() + raw = packed_storage[:, :, 576:832] + 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_mhc_matches_reference(): + streams, hidden = 4, 4096 + x = torch.randn(3, 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(3, 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, + ) + + +def test_mhc_keeps_streams_through_decode_autotuning(): + from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer + + hidden = 4096 + 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.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 + x = torch.randn(1, hidden, device="cuda", dtype=torch.bfloat16) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + 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(): + 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) + + +def test_nope_attention_with_existing_image_kernels(): + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.models.glm5_next.attention import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + + packed = torch.randn(32, 1, 904, dtype=torch.bfloat16, device="cuda") + packed[:, :, 512:576] = 0 + kv = packed[:, :, :576] + q = torch.randn(3, 16, 512, dtype=torch.bfloat16, device="cuda") + indexes = torch.full((3, 128), -1, dtype=torch.int32, device="cuda") + selected = [[7, 2, 9], [8, 19, 3, 5, 1], [12, 24]] + expected = [] + for i, locs in enumerate(selected): + indexes[i, : len(locs)] = torch.tensor(locs, device="cuda") + keys = kv[locs, 0, :512].float() + expected.append((q[i].float() @ keys.T * 0.0625).softmax(-1) @ keys) + expected = torch.stack(expected) + control = AttControl(nsa_prefill_dict={"topk_mem_indices": indexes, "softmax_scale": 0.0625}) + prefill = Glm5NextSparsePrefillState()._nsa_prefill_att(q, kv, control) + torch.testing.assert_close(prefill.float(), expected, atol=0.012, rtol=0.015) + lengths = torch.tensor([len(x) for x in selected], dtype=torch.int32, device="cuda") + decode = Glm5NextSparseDecodeState( + infer_state=SimpleNamespace(b1_cu_q_seq_len=torch.arange(4, 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) + + +def test_kpool_indexer_long_prefill_and_cached_decode(): + from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer + + tokens, heads, dim = 2061, 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, 904, device="cuda", dtype=torch.bfloat16) + ragged = torch.randperm(tokens + 9, device="cuda", dtype=torch.int32)[:tokens] + manager = SimpleNamespace( + get_indexer_raw_buffer=lambda _: storage[:, :, 576:832], + get_indexer_k_buffer=lambda _: storage.view(torch.uint8)[:, :, -132:], + ) + infer = SimpleNamespace( + mem_manager=manager, + mem_index=ragged, + max_kv_seq_len=tokens, + req_manager=SimpleNamespace(req_to_token_indexs=ragged[None]), + b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([tokens], device="cuda", dtype=torch.int32), + ) + 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, + query_batch=torch.zeros(tokens, device="cuda", dtype=torch.int32), + ) + 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() + # A later decode reads pooled history from restored token KV. + storage = storage.clone() + infer.mem_index = ragged[-1:] + state.lengths = state.lengths[-1:] + state.ks = state.ks[-1:] + state.query_batch = state.query_batch[-1:] + _, decoded = indexer._get_indices(hidden[-1:], hidden[-1:], infer, state, weights) + assert set(decoded[0, :2049].tolist()) == set(full[-1, :2049].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_tokenizer.py b/unit_tests/models/glm5_next/test_tokenizer.py new file mode 100644 index 0000000000..1392a7f9ce --- /dev/null +++ b/unit_tests/models/glm5_next/test_tokenizer.py @@ -0,0 +1,27 @@ +import json + +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace + +from lightllm.server.tokenizer import get_tokenizer + + +def test_transformers5_checkpoint_on_existing_image(tmp_path): + backend = Tokenizer(WordLevel({"[UNK]": 0, "hello": 1, "<|user|>": 2, "<|assistant|>": 3}, unk_token="[UNK]")) + backend.pre_tokenizer = Whitespace() + backend.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "config.json").write_text(json.dumps({"model_type": "glm5_next"})) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "TokenizersBackend", + "extra_special_tokens": ["<|user|>", "<|assistant|>"], + "unk_token": "[UNK]", + } + ) + ) + (tmp_path / "chat_template.jinja").write_text("<|user|>{{ messages[0]['content'] }}<|assistant|>") + tokenizer = get_tokenizer(str(tmp_path)) + assert tokenizer.apply_chat_template([{"role": "user", "content": "hello"}], tokenize=True) == [2, 1, 3] + assert tokenizer.decode([2, 1, 3], skip_special_tokens=True) == "hello" From 9f9eb4d351a868d22b3ca768b7505d06461a052e Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 05:36:55 +0000 Subject: [PATCH 02/24] refactor: pass MoE SwiGLU configuration through constructors --- .../fused_moe/fused_moe_weight.py | 5 + .../fused_moe/gpt_oss_fused_moe_weight_tp.py | 1 + .../fused_moe/impl/deepgemm_impl.py | 2 + .../fused_moe/impl/marlin_impl.py | 2 + .../fused_moe/impl/triton_impl.py | 18 +++- .../layer_weights/transformer_layer_weight.py | 4 +- lightllm/models/glm5_next/README.md | 3 + .../layer_weights/transformer_layer_weight.py | 4 +- .../fused_moe/test_activation_config.py | 102 ++++++++++++++++++ unit_tests/models/glm5_next/test_cache.py | 8 +- 10 files changed, 138 insertions(+), 11 deletions(-) create mode 100644 unit_tests/common/fused_moe/test_activation_config.py 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..4d2b1b2baf 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 @@ -35,6 +35,9 @@ def __init__( layer_num: int = 0, network_config: Dict[str, Any] = None, per_expert_scale_name: str = "", + *, + swiglu_alpha: Optional[float] = None, + swiglu_limit: Optional[float] = None, ) -> None: super().__init__(data_type=data_type) self.w1_weight_name = gate_proj_name @@ -67,6 +70,8 @@ def __init__( redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, routed_expert_counter_tensor=self.routed_expert_counter_tensor, auto_update_redundancy_expert=self.auto_update_redundancy_expert, + swiglu_alpha=swiglu_alpha, + swiglu_limit=swiglu_limit, ) self.lock = threading.Lock() self._create_weight() 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/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 024be9f55c..5e3d8de674 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 @@ -20,6 +20,8 @@ class FuseMoeDeepGEMM(FuseMoeTriton): + supports_swiglu_clamp = False + def _select_experts( self, input_tensor: torch.Tensor, 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..5b50b53ebe 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 @@ -11,6 +11,8 @@ class FuseMoeMarlin(FuseMoeTriton): + supports_swiglu_clamp = False + def create_workspace(self): from lightllm.utils.vllm_utils import HAS_VLLM 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 8596ab6f71..007c434004 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 @@ -6,6 +6,8 @@ class FuseMoeTriton(FuseMoeBaseImpl): + supports_swiglu_clamp = True + def __init__( self, n_routed_experts: int, @@ -16,7 +18,16 @@ def __init__( redundancy_expert_ids_tensor: torch.Tensor, routed_expert_counter_tensor: torch.Tensor, auto_update_redundancy_expert: bool, + *, + swiglu_alpha: Optional[float] = None, + swiglu_limit: Optional[float] = None, ): + if (swiglu_alpha is None) != (swiglu_limit is None): + raise ValueError("swiglu_alpha and swiglu_limit must be specified together") + if swiglu_limit is not None and not self.supports_swiglu_clamp: + raise NotImplementedError(f"{type(self).__name__} does not support clamped SwiGLU") + self.swiglu_alpha = swiglu_alpha + self.swiglu_limit = swiglu_limit super().__init__( n_routed_experts=n_routed_experts, num_fused_shared_experts=num_fused_shared_experts, @@ -27,8 +38,6 @@ def __init__( routed_expert_counter_tensor=routed_expert_counter_tensor, auto_update_redundancy_expert=auto_update_redundancy_expert, ) - self.swiglu_limit = None - self.swiglu_clamp_up_add_one = True def create_workspace(self): return None @@ -107,8 +116,9 @@ def _fused_experts( w1_scale=w13_scale, w2_scale=w2_scale, limit=self.swiglu_limit, - alpha=1.0 if self.swiglu_limit is not None else None, - clamp_up_add_one=self.swiglu_clamp_up_add_one, + alpha=self.swiglu_alpha, + # GPT-OSS owns its up + 1 variant in its dedicated experts path. + clamp_up_add_one=False, ) return input_tensor diff --git a/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py index cff020ea40..7b5c915237 100644 --- a/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py @@ -223,7 +223,7 @@ def _load_mlp(self, mlp_prefix, is_shared_experts=False): quant_method=self.get_quant_method("down_proj"), ) - def _init_moe(self): + def _init_moe(self, *, swiglu_alpha=None, swiglu_limit=None): moe_intermediate_size = self.network_config_["moe_intermediate_size"] self.moe_gate = ROWMMWeight( in_dim=self.n_embed, @@ -256,6 +256,8 @@ def _init_moe(self): num_fused_shared_experts=self.num_fused_shared_experts, layer_num=self.layer_num_, network_config=self.network_config_, + swiglu_alpha=swiglu_alpha, + swiglu_limit=swiglu_limit, ) def _init_ffn(self): diff --git a/lightllm/models/glm5_next/README.md b/lightllm/models/glm5_next/README.md index c503e6d412..b20f947a32 100644 --- a/lightllm/models/glm5_next/README.md +++ b/lightllm/models/glm5_next/README.md @@ -32,6 +32,9 @@ token KV 重建，不引入独立的池化尾状态，也不改 scheduler/radix 共享算子的扩展参数保持原默认值；GLM 显式启用 sigmoid gate、无 `up + 1` 的 clamp、 KDA 的 exp2 gate。模型特有的权重、attention、索引和 tokenizer 适配放在本目录。 +MoE 的 `swiglu_alpha=1.0`、`swiglu_limit=10.0` 沿构造链传入 `FusedMoeWeight` 和 Triton +实现，初始化后不再修改实现对象。`up + 1` 由 GPT-OSS 的专用 experts 调用显式选择； +EP/Marlin 尚未实现的 clamped SwiGLU 组合会在初始化时拒绝。 ## 启动与测速 diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index 3d39a3cda9..ad2ea5b7b3 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -76,9 +76,7 @@ def _get_param_slicer(self, sub_child_index: int): class Glm5NextTransformerLayerWeight(Deepseek3_2TransformerLayerWeight): def _init_moe(self): - super()._init_moe() - self.experts.fuse_moe_impl.swiglu_limit = self.network_config_["swiglu_limit"] - self.experts.fuse_moe_impl.swiglu_clamp_up_add_one = False + super()._init_moe(swiglu_alpha=1.0, swiglu_limit=self.network_config_["swiglu_limit"]) self.moe_gate = ROWMMWeight( in_dim=self.n_embed, out_dims=[self.n_routed_experts], 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..2c09752e2d --- /dev/null +++ b/unit_tests/common/fused_moe/test_activation_config.py @@ -0,0 +1,102 @@ +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.deepgemm_impl import FuseMoeDeepGEMM +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", "gelu"]) +def test_constructor_config_reaches_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 = {"swiglu_alpha": 1.0, "swiglu_limit": 10.0} if activation == "clamped_silu" else {} + 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, + **kwargs, + ) + 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 activation == "clamped_silu": + gate = gate.clamp(max=10) + up = up.clamp(-10, 10) + gate = ( + torch.nn.functional.gelu(gate, approximate="tanh") + if activation == "gelu" + else torch.nn.functional.silu(gate) + ) + expert_out = (gate.bfloat16() * up.bfloat16()).bfloat16() + expected[row] += (expert_out.float() * probs[row, choice]).bfloat16().float() + actual = weight.experts(x.clone(), router, 2, True, False, 0, 0) + torch.testing.assert_close(actual, expected.bfloat16(), atol=0.125, rtol=0.01) + + +@pytest.mark.parametrize("backend", [FuseMoeDeepGEMM, FuseMoeMarlin]) +def test_unsupported_backend_rejects_clamp_at_construction(backend): + with pytest.raises(NotImplementedError, match="does not support clamped SwiGLU"): + 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, + swiglu_alpha=1.0, + swiglu_limit=10.0, + ) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index a6caad8ec3..2277980c94 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -155,8 +155,10 @@ def test_clamped_swiglu_preserves_gpt_oss_default(add_one): 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} - silu_and_mul_fwd(x, out, limit=10.0, alpha=1.0, **kwargs) + 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 = torch.nn.functional.silu(gate.clamp(max=10)).bfloat16().float() - expected = gate * (up.clamp(-10, 10) + int(add_one)) + 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) From 4f032e1dff9c46485607cb51d16c98f1c8d7f775 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 06:13:44 +0000 Subject: [PATCH 03/24] refactor: validate activation parameters in MoE backend constructors --- .../meta_weights/fused_moe/impl/deepgemm_impl.py | 7 ++++++- .../meta_weights/fused_moe/impl/marlin_impl.py | 5 ++++- .../meta_weights/fused_moe/impl/triton_impl.py | 4 ---- 3 files changed, 10 insertions(+), 6 deletions(-) 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 5e3d8de674..08e3caf582 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 @@ -20,7 +20,12 @@ class FuseMoeDeepGEMM(FuseMoeTriton): - supports_swiglu_clamp = False + def __init__(self, *args, swiglu_alpha: Optional[float] = None, swiglu_limit: Optional[float] = None, **kwargs): + if swiglu_alpha is not None or swiglu_limit is not None: + raise NotImplementedError( + "FuseMoeDeepGEMM does not support clamped SwiGLU: EP activation kernels need alpha/limit support" + ) + super().__init__(*args, **kwargs) def _select_experts( self, 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 5b50b53ebe..b587a0776f 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 @@ -11,7 +11,10 @@ class FuseMoeMarlin(FuseMoeTriton): - supports_swiglu_clamp = False + def __init__(self, *args, swiglu_alpha: Optional[float] = None, swiglu_limit: Optional[float] = None, **kwargs): + if swiglu_alpha is not None or swiglu_limit is not None: + raise NotImplementedError("FuseMoeMarlin does not support clamped SwiGLU") + super().__init__(*args, **kwargs) def create_workspace(self): from lightllm.utils.vllm_utils import HAS_VLLM 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 007c434004..7fa2c7ac93 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 @@ -6,8 +6,6 @@ class FuseMoeTriton(FuseMoeBaseImpl): - supports_swiglu_clamp = True - def __init__( self, n_routed_experts: int, @@ -24,8 +22,6 @@ def __init__( ): if (swiglu_alpha is None) != (swiglu_limit is None): raise ValueError("swiglu_alpha and swiglu_limit must be specified together") - if swiglu_limit is not None and not self.supports_swiglu_clamp: - raise NotImplementedError(f"{type(self).__name__} does not support clamped SwiGLU") self.swiglu_alpha = swiglu_alpha self.swiglu_limit = swiglu_limit super().__init__( From 6f5ec3312396a48cff21767b5c579c0d864f380f Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 06:24:09 +0000 Subject: [PATCH 04/24] refactor: pass MoE activation parameters at execution time --- .../fused_moe/fused_moe_weight.py | 11 +-- .../meta_weights/fused_moe/impl/base_impl.py | 3 + .../fused_moe/impl/deepgemm_impl.py | 14 ++-- .../fused_moe/impl/marlin_impl.py | 10 +-- .../fused_moe/impl/triton_impl.py | 48 ++++-------- .../layer_weights/transformer_layer_weight.py | 4 +- lightllm/models/glm5_next/README.md | 7 +- .../layer_infer/transformer_layer_infer.py | 3 + .../layer_weights/transformer_layer_weight.py | 2 +- .../fused_moe/test_activation_config.py | 76 +++++++++++++------ 10 files changed, 95 insertions(+), 83 deletions(-) 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 4d2b1b2baf..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 @@ -35,9 +35,6 @@ def __init__( layer_num: int = 0, network_config: Dict[str, Any] = None, per_expert_scale_name: str = "", - *, - swiglu_alpha: Optional[float] = None, - swiglu_limit: Optional[float] = None, ) -> None: super().__init__(data_type=data_type) self.w1_weight_name = gate_proj_name @@ -70,8 +67,6 @@ def __init__( redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, routed_expert_counter_tensor=self.routed_expert_counter_tensor, auto_update_redundancy_expert=self.auto_update_redundancy_expert, - swiglu_alpha=swiglu_alpha, - swiglu_limit=swiglu_limit, ) self.lock = threading.Lock() self._create_weight() @@ -142,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_) @@ -161,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/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 08e3caf582..d2f2df7c55 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 @@ -20,13 +20,6 @@ class FuseMoeDeepGEMM(FuseMoeTriton): - def __init__(self, *args, swiglu_alpha: Optional[float] = None, swiglu_limit: Optional[float] = None, **kwargs): - if swiglu_alpha is not None or swiglu_limit is not None: - raise NotImplementedError( - "FuseMoeDeepGEMM does not support clamped SwiGLU: EP activation kernels need alpha/limit support" - ) - super().__init__(*args, **kwargs) - def _select_experts( self, input_tensor: torch.Tensor, @@ -83,7 +76,14 @@ 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( + "FuseMoeDeepGEMM does not support clamped SwiGLU: EP activation kernels need alpha/limit support" + ) output = fused_experts( hidden_states=input_tensor, w13=w13, 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 b587a0776f..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 @@ -11,11 +11,6 @@ class FuseMoeMarlin(FuseMoeTriton): - def __init__(self, *args, swiglu_alpha: Optional[float] = None, swiglu_limit: Optional[float] = None, **kwargs): - if swiglu_alpha is not None or swiglu_limit is not None: - raise NotImplementedError("FuseMoeMarlin does not support clamped SwiGLU") - super().__init__(*args, **kwargs) - def create_workspace(self): from lightllm.utils.vllm_utils import HAS_VLLM @@ -35,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 7fa2c7ac93..cfbb4faff4 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,40 +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, - *, - swiglu_alpha: Optional[float] = None, - swiglu_limit: Optional[float] = None, - ): - if (swiglu_alpha is None) != (swiglu_limit is None): - raise ValueError("swiglu_alpha and swiglu_limit must be specified together") - self.swiglu_alpha = swiglu_alpha - self.swiglu_limit = swiglu_limit - 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 @@ -94,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 @@ -111,10 +84,9 @@ def _fused_experts( use_fp8_w8a8=use_fp8_w8a8, w1_scale=w13_scale, w2_scale=w2_scale, - limit=self.swiglu_limit, - alpha=self.swiglu_alpha, - # GPT-OSS owns its up + 1 variant in its dedicated experts path. - clamp_up_add_one=False, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return input_tensor @@ -136,7 +108,12 @@ 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, ): + if (alpha is None) != (limit is None): + raise ValueError("alpha and limit must be specified together") topk_weights, topk_ids, origin_topk_ids = self._select_experts( input_tensor=input_tensor, router_logits=router_logits, @@ -162,5 +139,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/models/deepseek2/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py index 7b5c915237..cff020ea40 100644 --- a/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py @@ -223,7 +223,7 @@ def _load_mlp(self, mlp_prefix, is_shared_experts=False): quant_method=self.get_quant_method("down_proj"), ) - def _init_moe(self, *, swiglu_alpha=None, swiglu_limit=None): + def _init_moe(self): moe_intermediate_size = self.network_config_["moe_intermediate_size"] self.moe_gate = ROWMMWeight( in_dim=self.n_embed, @@ -256,8 +256,6 @@ def _init_moe(self, *, swiglu_alpha=None, swiglu_limit=None): num_fused_shared_experts=self.num_fused_shared_experts, layer_num=self.layer_num_, network_config=self.network_config_, - swiglu_alpha=swiglu_alpha, - swiglu_limit=swiglu_limit, ) def _init_ffn(self): diff --git a/lightllm/models/glm5_next/README.md b/lightllm/models/glm5_next/README.md index b20f947a32..828dc97ba9 100644 --- a/lightllm/models/glm5_next/README.md +++ b/lightllm/models/glm5_next/README.md @@ -32,9 +32,10 @@ token KV 重建，不引入独立的池化尾状态，也不改 scheduler/radix 共享算子的扩展参数保持原默认值；GLM 显式启用 sigmoid gate、无 `up + 1` 的 clamp、 KDA 的 exp2 gate。模型特有的权重、attention、索引和 tokenizer 适配放在本目录。 -MoE 的 `swiglu_alpha=1.0`、`swiglu_limit=10.0` 沿构造链传入 `FusedMoeWeight` 和 Triton -实现，初始化后不再修改实现对象。`up + 1` 由 GPT-OSS 的专用 experts 调用显式选择； -EP/Marlin 尚未实现的 clamped SwiGLU 组合会在初始化时拒绝。 +MoE 在推理调用处显式传入 `alpha=1.0`、`limit=10.0`、`clamp_up_add_one=False`， +沿 `experts → __call__ → _fused_experts` 传给激活算子，不在通用 MoE 对象上保存激活配置。 +GPT-OSS 的专用 experts 调用仍显式选择 `up + 1`；EP/Marlin 尚未实现的 clamped SwiGLU +组合会在执行时拒绝。 ## 启动与测速 diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 6d6ae93cb1..c57a45f34f 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -78,6 +78,9 @@ def _moe_ffn_tp(self, input, infer_state, layer_weight) -> torch.Tensor: topk_group=self.topk_group, num_expert_group=self.n_group, 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: diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index ad2ea5b7b3..1690cebc46 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -76,7 +76,7 @@ def _get_param_slicer(self, sub_child_index: int): class Glm5NextTransformerLayerWeight(Deepseek3_2TransformerLayerWeight): def _init_moe(self): - super()._init_moe(swiglu_alpha=1.0, swiglu_limit=self.network_config_["swiglu_limit"]) + super()._init_moe() self.moe_gate = ROWMMWeight( in_dim=self.n_embed, out_dims=[self.n_routed_experts], diff --git a/unit_tests/common/fused_moe/test_activation_config.py b/unit_tests/common/fused_moe/test_activation_config.py index 2c09752e2d..ca86f27d86 100644 --- a/unit_tests/common/fused_moe/test_activation_config.py +++ b/unit_tests/common/fused_moe/test_activation_config.py @@ -30,15 +30,18 @@ def runtime(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("activation", ["silu", "clamped_silu", "gelu"]) -def test_constructor_config_reaches_expert_activation(monkeypatch, activation): +@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 = {"swiglu_alpha": 1.0, "swiglu_limit": 10.0} if activation == "clamped_silu" else {} + 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", @@ -51,7 +54,6 @@ def test_constructor_config_reaches_expert_activation(monkeypatch, activation): data_type=torch.bfloat16, quant_method=NoQuantization(), network_config=config, - **kwargs, ) eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) weights = {} @@ -71,32 +73,56 @@ def test_constructor_config_reaches_expert_activation(monkeypatch, activation): i = int(top.indices[row, choice]) gate = (x[row] * (1 + i / 4)).float() up = (x[row] * (2 + i / 8)).float() - if activation == "clamped_silu": - gate = gate.clamp(max=10) - up = up.clamp(-10, 10) - gate = ( - torch.nn.functional.gelu(gate, approximate="tanh") - if activation == "gelu" - else torch.nn.functional.silu(gate) - ) + 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() - actual = weight.experts(x.clone(), router, 2, True, False, 0, 0) + 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) @pytest.mark.parametrize("backend", [FuseMoeDeepGEMM, FuseMoeMarlin]) -def test_unsupported_backend_rejects_clamp_at_construction(backend): +def test_unsupported_backend_rejects_clamp_at_call(monkeypatch, backend): + 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"): - 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, - swiglu_alpha=1.0, - swiglu_limit=10.0, + 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, ) From 6e8dac64847f675b2fcbba490c08fa01c9d09b37 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 06:26:20 +0000 Subject: [PATCH 05/24] refactor: remove redundant MoE activation argument check --- .../layer_weights/meta_weights/fused_moe/impl/triton_impl.py | 2 -- 1 file changed, 2 deletions(-) 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 cfbb4faff4..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 @@ -112,8 +112,6 @@ def __call__( limit: Optional[float] = None, clamp_up_add_one: bool = True, ): - if (alpha is None) != (limit is None): - raise ValueError("alpha and limit must be specified together") topk_weights, topk_ids, origin_topk_ids = self._select_experts( input_tensor=input_tensor, router_logits=router_logits, From 8aa07beefe97e013b2da3bd897ee8ca937248282 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 06:37:27 +0000 Subject: [PATCH 06/24] refactor: name gated RMSNorm option gate_type --- .../basemodel/layer_weights/meta_weights/norm_weight.py | 8 ++++---- .../common/basemodel/triton_kernel/norm/gated_rmsnorm.py | 6 +++--- .../glm5_next/layer_weights/transformer_layer_weight.py | 2 +- unit_tests/models/glm5_next/test_cache.py | 8 ++++---- 4 files changed, 12 insertions(+), 12 deletions(-) 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 fe80974569..03ff393811 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py @@ -78,11 +78,11 @@ def __init__( dim: int, weight_name: str, data_type: torch.dtype, - activation: str = "silu", + gate_type: str = "silu", ): super().__init__(dim=dim, weight_name=weight_name, data_type=data_type) - assert activation in ("silu", "sigmoid") - self.activation = activation + assert gate_type in ("silu", "sigmoid") + self.gate_type = gate_type def _triton_forward( self, @@ -104,7 +104,7 @@ def _triton_forward( eps=eps, z=gate_value, out=out, - activation=self.activation, + gate_type=self.gate_type, ) def _cuda_forward( diff --git a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py index 962b502b5c..36203999d1 100644 --- a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py +++ b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py @@ -110,7 +110,7 @@ def gated_rmsnorm_forward( group_size: int = None, norm_before_gate: bool = True, run_config: dict = None, - activation: str = "silu", + gate_type: str = "silu", ): M, N = x.shape if group_size is None: @@ -120,7 +120,7 @@ def gated_rmsnorm_forward( assert x.stride(-1) == 1 # z is required for gated_rmsnorm assert z is not None, "z cannot be None for gated_rmsnorm_forward" - assert activation in ("silu", "sigmoid"), f"unsupported gate activation: {activation}" + 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 @@ -184,7 +184,7 @@ def gated_rmsnorm_forward( eps, BLOCK_N=BLOCK_N, NORM_BEFORE_GATE=norm_before_gate, - SIGMOID_GATE=activation == "sigmoid", + SIGMOID_GATE=gate_type == "sigmoid", Z_HEADS=z_heads, num_warps=num_warps, ) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index 1690cebc46..85c3dfed0e 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -168,7 +168,7 @@ def _init_kda(self): dim=head_dim, weight_name=f"{prefix}.o_norm.weight", data_type=self.data_type_, - activation="sigmoid", + gate_type="sigmoid", ) self.linear_o_proj = COLMMWeight( in_dim=projection, diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 2277980c94..10cc802a31 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -134,17 +134,17 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t assert torch.equal(big.ssm_state_cache.buffer[0], ssm) -@pytest.mark.parametrize("activation", ["silu", "sigmoid"]) -def test_gated_norm_activation_and_strided_gate(activation): +@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, activation=activation) + 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 activation == "sigmoid" else torch.nn.functional.silu(z) + 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) From 7ecea66db8d1f3a988ca6d5e36e84fe636debe05 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 07:26:16 +0000 Subject: [PATCH 07/24] refactor: share KDA decode and centralize attention backends --- .../basemodel/attention/linear/kda.py} | 19 ++-- .../basemodel/attention/nsa/glm5_next.py} | 2 +- .../linear_att/fla/ops/fused_recurrent.py | 81 ++++++++++++----- .../triton_kernel/linear_att/fla/ops}/kda.py | 12 +-- lightllm/models/glm5_next/model.py | 4 +- .../glm5_next/triton_kernel/kda_decode.py | 91 ------------------- .../test_fused_recurrent_strided.py | 59 ++++++++++++ unit_tests/models/glm5_next/test_kernels.py | 24 +++-- 8 files changed, 148 insertions(+), 144 deletions(-) rename lightllm/{models/glm5_next/kda_backend.py => common/basemodel/attention/linear/kda.py} (91%) rename lightllm/{models/glm5_next/attention.py => common/basemodel/attention/nsa/glm5_next.py} (97%) rename lightllm/{models/glm5_next/triton_kernel => common/basemodel/triton_kernel/linear_att/fla/ops}/kda.py (98%) delete mode 100644 lightllm/models/glm5_next/triton_kernel/kda_decode.py diff --git a/lightllm/models/glm5_next/kda_backend.py b/lightllm/common/basemodel/attention/linear/kda.py similarity index 91% rename from lightllm/models/glm5_next/kda_backend.py rename to lightllm/common/basemodel/attention/linear/kda.py index e5690ee828..0b084ad941 100644 --- a/lightllm/models/glm5_next/kda_backend.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -19,8 +19,10 @@ causal_conv1d_fn, causal_conv1d_update, ) -from .triton_kernel.kda import chunk_kda_with_fused_gate -from .triton_kernel.kda_decode import fused_recurrent_kda +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.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices if TYPE_CHECKING: @@ -157,19 +159,18 @@ def decode_att( conv_state_indices=self.b_conv_buffer_idx, ) q, k, v = [backend.reshape_qkv(x, decode=True) for x in backend.split_qkv(mixed_qkv)] - raw_gate = raw_gate.view(-1, 1, backend.tp_projection_size) - raw_beta = raw_beta.view(-1, 1, backend.tp_num_heads) - output, _ = fused_recurrent_kda( + output, _ = fused_recurrent_gated_delta_rule( 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, + a_raw=raw_gate.view(-1, backend.tp_num_heads, backend.head_dim), + b_raw=raw_beta.view(-1, backend.tp_num_heads), + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, initial_state=ssm_states, lower_bound=backend.lower_bound, inplace_final_state=True, + use_qk_l2norm_in_kernel=True, ssm_state_indices=self.b_ssm_buffer_idx, ) return output diff --git a/lightllm/models/glm5_next/attention.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py similarity index 97% rename from lightllm/models/glm5_next/attention.py rename to lightllm/common/basemodel/attention/nsa/glm5_next.py index 8c334d0fec..c911825119 100644 --- a/lightllm/models/glm5_next/attention.py +++ b/lightllm/common/basemodel/attention/nsa/glm5_next.py @@ -2,7 +2,7 @@ import torch -from lightllm.common.basemodel.attention.nsa.flashmla_sparse import ( +from .flashmla_sparse import ( NsaFlashMlaSparseAttBackend, NsaFlashMlaSparsePrefillAttState, NsaFlashMlaSparseDecodeAttState, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 5dfbd6e4ab..52e7376858 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -20,6 +20,7 @@ { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "IS_SINGLE_TOKEN": lambda args: args["cu_seqlens"] is None and args["T"] == 1, "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, "HAS_SEPARATE_WRITE_INDICES": lambda args: args["ssm_state_write_indices"] is not None, @@ -41,8 +42,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( num_accepted_tokens, # Fused gating parameters (only used when FUSE_GATING=True) A_log, # [HV] per-head log decay - dt_bias, # [HV] per-head dt bias - a_raw, # [B*T, HV] raw alpha values (before softplus) + dt_bias, # [HV] for GDN, [HV, K] for KDA + a_raw, # [B*T, HV] for GDN, [B*T, HV, K] for KDA b_raw, # [B*T, HV] raw beta values (before sigmoid) scale, N: tl.int64, # num of sequences @@ -67,11 +68,13 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( stride_write_indices_tok: tl.constexpr, # NEW: stride for write indices SOFTPLUS_BETA: tl.constexpr, # softplus beta parameter (default 1.0) SOFTPLUS_THRESHOLD: tl.constexpr, # softplus threshold (default 20.0) + LOWER_BOUND: tl.constexpr, # bounded sigmoid gate when provided; otherwise softplus USE_INITIAL_STATE: tl.constexpr, # whether to use initial state INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, IS_VARLEN: tl.constexpr, + IS_SINGLE_TOKEN: tl.constexpr, IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, IS_KDA: tl.constexpr, @@ -81,6 +84,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) + if IS_SINGLE_TOKEN: + T = 1 if IS_VARLEN: bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int64), @@ -98,6 +103,9 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( o_k = i_k * BK + tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] p_q = q + bos * stride_q_tok + i_h * K + o_k p_k = k + bos * stride_k_tok + i_h * K + o_k @@ -105,8 +113,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if FUSE_GATING: # Fused gating: load per-head constants once, compute g/beta inline per token b_A_log = tl.load(A_log + i_hv).to(tl.float32) - b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) - p_a_raw = a_raw + bos * stride_a_tok + i_hv + if IS_KDA: + b_dt_bias = tl.load(dt_bias + i_hv * K + o_k, mask=mask_k, other=0).to(tl.float32) + p_a_raw = a_raw + bos * stride_a_tok + i_hv * K + o_k + else: + b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) + p_a_raw = a_raw + bos * stride_a_tok + i_hv p_b_raw = b_raw + bos * stride_b_tok + i_hv else: if IS_BETA_HEADWISE: @@ -121,10 +133,6 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: if IS_CONTINUOUS_BATCHING: @@ -151,16 +159,24 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_q = b_q * scale # [BK, BV] if FUSE_GATING: - # Compute g = -exp(A_log) * softplus(a_raw + dt_bias) inline - b_a = tl.load(p_a_raw).to(tl.float32) + if IS_KDA: + b_a = tl.load(p_a_raw, mask=mask_k, other=0).to(tl.float32) + else: + b_a = tl.load(p_a_raw).to(tl.float32) x = b_a + b_dt_bias - softplus_x = tl.where( - SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, - (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), - x, - ) - b_g = -tl.exp(b_A_log) * softplus_x - b_h *= exp(b_g) + if LOWER_BOUND is not None: + b_g = LOWER_BOUND * tl.sigmoid(tl.exp(b_A_log) * x) + else: + softplus_x = tl.where( + SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, + (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + if IS_KDA: + b_h *= exp(b_g[:, None]) + else: + b_h *= exp(b_g) # Compute beta = sigmoid(b_raw) inline b_b = tl.load(p_b_raw).to(tl.float32) b_beta = tl.sigmoid(b_b) @@ -169,7 +185,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_g = tl.load(p_g).to(tl.float32) b_h *= exp(b_g) else: - b_gk = tl.load(p_gk).to(tl.float32) + b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) b_h *= exp(b_gk[:, None]) if IS_BETA_HEADWISE: b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) @@ -240,8 +256,8 @@ def _ensure_gate_token_strided(x: torch.Tensor, inner_numel: int): """Return a_raw/b_raw and token stride, copying only when needed.""" if x is None: return None, 0 - # a_raw/b_raw are 2D [tokens, HV]; the tail HV dimension must be packed. - if x.stride(1) != 1: + # Gates use [tokens, HV], or [tokens, HV, K] for KDA's per-channel decay. + if x.stride(-1) != 1 or (x.ndim == 3 and x.stride(-2) != x.shape[-1]): x = x.contiguous() return x, inner_numel return x, x.stride(0) @@ -267,6 +283,7 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + lower_bound: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -274,10 +291,12 @@ def fused_recurrent_gated_delta_rule_fwd( # Qwen3Next MTP verify path passes cu_seqlens for variable-length verify # chunks. Both flow through the per-token strided-view path below. N = B if cu_seqlens is None else len(cu_seqlens) - 1 + fuse_gating = A_log is not None + is_kda = a_raw.ndim == 3 if fuse_gating else g.ndim == 4 q, stride_q_tok = _ensure_qkv_token_strided(q, H * K) k, stride_k_tok = _ensure_qkv_token_strided(k, H * K) v, stride_v_tok = _ensure_qkv_token_strided(v, HV * V) - a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV) + a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV * K if is_kda else HV) b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw, HV) BK = triton.next_power_of_2(K) if T == 1: @@ -294,8 +313,6 @@ def fused_recurrent_gated_delta_rule_fwd( NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" - fuse_gating = A_log is not None - if out is not None: o = out.unsqueeze(0) if out.ndim == v.ndim else out else: @@ -368,10 +385,11 @@ def fused_recurrent_gated_delta_rule_fwd( stride_write_indices_tok=stride_write_indices_tok, SOFTPLUS_BETA=1.0, SOFTPLUS_THRESHOLD=20.0, + LOWER_BOUND=lower_bound, IS_BETA_HEADWISE=False if fuse_gating else (beta.ndim == v.ndim), USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, - IS_KDA=False, + IS_KDA=is_kda, FUSE_GATING=fuse_gating, num_warps=num_warps, num_stages=num_stages, @@ -402,6 +420,7 @@ def forward( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + lower_bound: float | None = None, ): # q/k/v/a_raw/b_raw may be non-contiguous column views of one projection # output; the kernel handles them via per-token strides (no copies). @@ -424,6 +443,7 @@ def forward( a_raw=a_raw, b_raw=b_raw, out=out, + lower_bound=lower_bound, ) return o, final_state @@ -449,6 +469,7 @@ def fused_recurrent_gated_delta_rule( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + lower_bound: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -460,7 +481,7 @@ def fused_recurrent_gated_delta_rule( values of shape `[B, T, HV, V]`. GVA is applied if `HV > H`. g (torch.Tensor): - g (decays) of shape `[B, T, HV]`. + Log decays of shape `[B, T, HV]` for GDN or `[B, T, HV, K]` for KDA. beta (torch.Tensor): betas of shape `[B, T, HV]`. scale (Optional[int]): @@ -481,6 +502,15 @@ def fused_recurrent_gated_delta_rule( Indices to map the input sequences to the initial/final states. num_accepted_tokens (Optional[torch.Tensor]): Number of accepted tokens for each sequence during decoding. + a_raw (Optional[torch.Tensor]): + Raw decay gates of shape `[B*T, HV]` for GDN or `[B*T, HV, K]` + for KDA. With `A_log`, `dt_bias`, and `b_raw`, fuse gate computation. + `A_log` is per-head; `dt_bias` has the same trailing shape as `a_raw`. + b_raw (Optional[torch.Tensor]): + Raw beta gates of shape `[B*T, HV]`, before sigmoid. + lower_bound (Optional[float]): + Use `lower_bound * sigmoid(exp(A_log) * (a_raw + dt_bias))` for + fused log decays. `None` keeps `-exp(A_log) * softplus(a_raw + dt_bias)`. Returns: o (torch.Tensor): @@ -531,5 +561,6 @@ def fused_recurrent_gated_delta_rule( a_raw, b_raw, out, + lower_bound, ) return o, final_state diff --git a/lightllm/models/glm5_next/triton_kernel/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py similarity index 98% rename from lightllm/models/glm5_next/triton_kernel/kda.py rename to lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py index f2d42919a4..f8f4e33f80 100644 --- a/lightllm/models/glm5_next/triton_kernel/kda.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang -"""KDA helpers built on LightLLM's continuous-batching recurrent kernel.""" +"""Chunkwise KDA prefill with per-channel decay gates.""" from __future__ import annotations @@ -10,12 +10,12 @@ import triton import triton.language as tl -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.chunk_delta_h import chunk_gated_delta_rule_fwd_h -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.cumsum import chunk_local_cumsum -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd +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 lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril import solve_tril +from .solve_tril import solve_tril FLA_CHUNK_SIZE = 64 diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index d9a7ff87d7..c7b69ab36e 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -5,12 +5,12 @@ import triton from lightllm.common.build_utils import repair_config +from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend +from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend from lightllm.common.req_manager import ReqManagerForMamba from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel from lightllm.models.registry import ModelRegistry -from .attention import Glm5NextSparseAttBackend from .cache_config import Glm5NextCacheConfig -from .kda_backend import KDALinearAttBackend 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 diff --git a/lightllm/models/glm5_next/triton_kernel/kda_decode.py b/lightllm/models/glm5_next/triton_kernel/kda_decode.py deleted file mode 100644 index 6482114f28..0000000000 --- a/lightllm/models/glm5_next/triton_kernel/kda_decode.py +++ /dev/null @@ -1,91 +0,0 @@ -import torch -import triton -import triton.language as tl - - -@triton.jit -def _kda_decode( - Q, - K, - V, - G, - B, - A, - Bias, - State, - Idx, - O, - SQ: tl.constexpr, - SK: tl.constexpr, - SV: tl.constexpr, - SG: tl.constexpr, - SB: tl.constexpr, - H: tl.constexpr, - D: tl.constexpr, - LOWER: tl.constexpr, - BV: tl.constexpr, -): - row_head = tl.program_id(1) - row, head = row_head // H, row_head % H - ki = tl.arange(0, D) - vi = tl.program_id(0) * BV + tl.arange(0, BV) - q = tl.load(Q + row * SQ + head * D + ki).to(tl.float32) - k = tl.load(K + row * SK + head * D + ki).to(tl.float32) - v = tl.load(V + row * 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 + row * SG + head * D + ki).to(tl.float32) - bias = tl.load(Bias + head * D + ki) - amplitude = tl.exp(tl.load(A + head)) - decay = tl.exp(LOWER * tl.sigmoid(amplitude * (gate + bias))) - beta = tl.sigmoid(tl.load(B + row * SB + head).to(tl.float32)) - req = tl.load(Idx + row) - ptr = State + (req * H + head) * D * D + ki[:, None] * D + vi[None, :] - state = tl.load(ptr).to(tl.float32) * decay[:, None] - delta = (v - tl.sum(state * k[:, None], 0)) * beta - state += k[:, None] * delta[None, :] - tl.store(ptr, state) - out = tl.sum(state * q[:, None], 0) - tl.store(O + row_head * D + vi, out) - - -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, -): - assert inplace_final_state and q.shape[1] == 1 - batch, _, heads, dim = q.shape - assert dim == 128 - 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, - out, - q.stride(0), - k.stride(0), - v.stride(0), - raw_gate.stride(0), - raw_beta.stride(0), - heads, - dim, - lower_bound, - 32, - num_warps=4, - ) - return out, initial_state diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py index 969a3f354d..eff554a2a4 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py @@ -60,6 +60,65 @@ def run(q_, k_, v_, a_, b_, state): assert torch.equal(state_ref, state_strided) +@pytest.mark.parametrize("per_channel,lower_bound", [(False, None), (True, None), (True, -5.0)]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +def test_decode_fused_and_precomputed_gates_match_reference(per_channel, lower_bound, state_dtype): + """Exercise scalar/vector gates, strided projections, and separate state slots.""" + torch.manual_seed(53) + batch, H, HV, K, V = 4, 2, 4, 128, 64 + gate_dim = HV * K if per_channel else HV + widths = [H * K, H * K, HV * V, gate_dim, HV] + mixed = torch.randn(batch, sum(widths), device="cuda", dtype=torch.bfloat16) + q, k, v, a, b = mixed.split(widths, dim=-1) + q, k, v = q.view(batch, 1, H, K), k.view(batch, 1, H, K), v.view(batch, 1, HV, V) + if per_channel: + a = a.view(batch, HV, K) + A_log = torch.randn(HV, device="cuda") * 0.1 + bias = torch.randn(a.shape[1:], device="cuda") * 0.1 + amplitude = A_log.exp().view(HV, 1) if per_channel else A_log.exp() + x = a.float() + bias + if lower_bound is None: + g = -amplitude * torch.nn.functional.softplus(x) + else: + g = lower_bound * torch.sigmoid(amplitude * x) + beta = b.float().sigmoid() + decay = g.exp().unsqueeze(-1) if per_channel else g.exp()[..., None, None] + q_ref = q[:, 0].float() + k_ref = k[:, 0].float() + q_ref = (q_ref * torch.rsqrt(q_ref.square().sum(-1, keepdim=True) + 1e-6) / K ** 0.5).repeat_interleave( + HV // H, dim=1 + ) + k_ref = (k_ref * torch.rsqrt(k_ref.square().sum(-1, keepdim=True) + 1e-6)).repeat_interleave(HV // H, dim=1) + state_ref = torch.randn(12, HV, K, V, device="cuda", dtype=state_dtype) * 0.1 + state_fused, state_precomputed = state_ref.clone(), state_ref.clone() + read_idx = torch.tensor([4, 1, 7, 3], device="cuda", dtype=torch.int32) + write_idx = torch.tensor([2, 9, 0, 6], device="cuda", dtype=torch.int32) + for _ in range(3): + state = state_ref[read_idx].float() * decay + delta = (v[:, 0].float() - torch.einsum("bhkv,bhk->bhv", state, k_ref)) * beta[..., None] + state += k_ref[..., None] * delta[..., None, :] + expected = torch.einsum("bhkv,bhk->bhv", state, q_ref).unsqueeze(1) + state_ref[write_idx] = state.to(state_dtype) + for cache, gates in ( + (state_fused, dict(A_log=A_log, dt_bias=bias, a_raw=a, b_raw=b, lower_bound=lower_bound)), + (state_precomputed, dict(g=g.unsqueeze(1), beta=beta.unsqueeze(1))), + ): + output, _ = fused_recurrent_gated_delta_rule( + q, + k, + v, + initial_state=cache, + ssm_state_indices=read_idx, + ssm_state_write_indices=write_idx, + use_qk_l2norm_in_kernel=True, + **gates, + ) + torch.testing.assert_close(output.float(), expected, atol=5e-4, rtol=1e-2) + # Comparing the entire cache also checks that unrelated slots stay intact. + torch.testing.assert_close(cache, state_ref, atol=5e-4, rtol=1e-2) + read_idx, write_idx = write_idx, read_idx + + # NOTE: the decode-only `cu_seqlens is None` contract from upstream #1349 was # intentionally lifted on this branch so the Qwen3Next MTP verify path can drive # the kernel with variable-length verify chunks (cu_seqlens + 2D SSM index diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index f719c090a0..9c8d08539e 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -7,8 +7,10 @@ from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import set_env_start_args -from lightllm.models.glm5_next.triton_kernel.kda import chunk_kda_with_fused_gate -from lightllm.models.glm5_next.triton_kernel.kda_decode import fused_recurrent_kda +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.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools, gather_pools, expand_topk from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 from lightllm.models.glm5_next.triton_kernel.mhc import hc_pre_norm, hc_pre_reference, hc_post, hc_post_reference @@ -71,16 +73,18 @@ def test_kda_chunk_and_decode_match_recurrence(tokens): states[2] = initial[0] unchanged = states[[0, 1, 3]].clone() for i in range(tokens): - out, _ = fused_recurrent_kda( + out, _ = fused_recurrent_gated_delta_rule( 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), + a_raw=gate[:, i], + b_raw=beta[:, i], + A_log=a, + dt_bias=bias.flatten(), + initial_state=states, + ssm_state_indices=torch.tensor([2], device="cuda", dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + lower_bound=-5.0, ) 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) @@ -216,7 +220,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): def test_nope_attention_with_existing_image_kernels(): from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.models.glm5_next.attention import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState packed = torch.randn(32, 1, 904, dtype=torch.bfloat16, device="cuda") packed[:, :, 512:576] = 0 From d4fa9d7cd9326736c6675f586946f8d9993349ea Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 11 Sep 2026 07:28:12 +0000 Subject: [PATCH 08/24] revert: keep standalone KDA decode with common attention backends --- .../common/basemodel/attention/linear/kda.py | 17 ++-- .../linear_att/fla/ops/fused_recurrent.py | 81 +++++------------ .../linear_att/fla/ops/kda_decode.py | 91 +++++++++++++++++++ .../test_fused_recurrent_strided.py | 59 ------------ unit_tests/models/glm5_next/test_kernels.py | 20 ++-- 5 files changed, 132 insertions(+), 136 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 0b084ad941..34c5b52d33 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -20,9 +20,7 @@ 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.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices if TYPE_CHECKING: @@ -159,18 +157,19 @@ def decode_att( conv_state_indices=self.b_conv_buffer_idx, ) q, k, v = [backend.reshape_qkv(x, decode=True) for x in backend.split_qkv(mixed_qkv)] - output, _ = fused_recurrent_gated_delta_rule( + raw_gate = raw_gate.view(-1, 1, backend.tp_projection_size) + raw_beta = raw_beta.view(-1, 1, backend.tp_num_heads) + output, _ = fused_recurrent_kda( q=q, k=k, v=v, - a_raw=raw_gate.view(-1, backend.tp_num_heads, backend.head_dim), - b_raw=raw_beta.view(-1, backend.tp_num_heads), - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, + 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, - use_qk_l2norm_in_kernel=True, ssm_state_indices=self.b_ssm_buffer_idx, ) return output diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 52e7376858..5dfbd6e4ab 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -20,7 +20,6 @@ { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - "IS_SINGLE_TOKEN": lambda args: args["cu_seqlens"] is None and args["T"] == 1, "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, "HAS_SEPARATE_WRITE_INDICES": lambda args: args["ssm_state_write_indices"] is not None, @@ -42,8 +41,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( num_accepted_tokens, # Fused gating parameters (only used when FUSE_GATING=True) A_log, # [HV] per-head log decay - dt_bias, # [HV] for GDN, [HV, K] for KDA - a_raw, # [B*T, HV] for GDN, [B*T, HV, K] for KDA + dt_bias, # [HV] per-head dt bias + a_raw, # [B*T, HV] raw alpha values (before softplus) b_raw, # [B*T, HV] raw beta values (before sigmoid) scale, N: tl.int64, # num of sequences @@ -68,13 +67,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( stride_write_indices_tok: tl.constexpr, # NEW: stride for write indices SOFTPLUS_BETA: tl.constexpr, # softplus beta parameter (default 1.0) SOFTPLUS_THRESHOLD: tl.constexpr, # softplus threshold (default 20.0) - LOWER_BOUND: tl.constexpr, # bounded sigmoid gate when provided; otherwise softplus USE_INITIAL_STATE: tl.constexpr, # whether to use initial state INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, IS_VARLEN: tl.constexpr, - IS_SINGLE_TOKEN: tl.constexpr, IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, IS_KDA: tl.constexpr, @@ -84,8 +81,6 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) - if IS_SINGLE_TOKEN: - T = 1 if IS_VARLEN: bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int64), @@ -103,9 +98,6 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( o_k = i_k * BK + tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] p_q = q + bos * stride_q_tok + i_h * K + o_k p_k = k + bos * stride_k_tok + i_h * K + o_k @@ -113,12 +105,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if FUSE_GATING: # Fused gating: load per-head constants once, compute g/beta inline per token b_A_log = tl.load(A_log + i_hv).to(tl.float32) - if IS_KDA: - b_dt_bias = tl.load(dt_bias + i_hv * K + o_k, mask=mask_k, other=0).to(tl.float32) - p_a_raw = a_raw + bos * stride_a_tok + i_hv * K + o_k - else: - b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) - p_a_raw = a_raw + bos * stride_a_tok + i_hv + b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) + p_a_raw = a_raw + bos * stride_a_tok + i_hv p_b_raw = b_raw + bos * stride_b_tok + i_hv else: if IS_BETA_HEADWISE: @@ -133,6 +121,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: if IS_CONTINUOUS_BATCHING: @@ -159,24 +151,16 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_q = b_q * scale # [BK, BV] if FUSE_GATING: - if IS_KDA: - b_a = tl.load(p_a_raw, mask=mask_k, other=0).to(tl.float32) - else: - b_a = tl.load(p_a_raw).to(tl.float32) + # Compute g = -exp(A_log) * softplus(a_raw + dt_bias) inline + b_a = tl.load(p_a_raw).to(tl.float32) x = b_a + b_dt_bias - if LOWER_BOUND is not None: - b_g = LOWER_BOUND * tl.sigmoid(tl.exp(b_A_log) * x) - else: - softplus_x = tl.where( - SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, - (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), - x, - ) - b_g = -tl.exp(b_A_log) * softplus_x - if IS_KDA: - b_h *= exp(b_g[:, None]) - else: - b_h *= exp(b_g) + softplus_x = tl.where( + SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, + (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + b_h *= exp(b_g) # Compute beta = sigmoid(b_raw) inline b_b = tl.load(p_b_raw).to(tl.float32) b_beta = tl.sigmoid(b_b) @@ -185,7 +169,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_g = tl.load(p_g).to(tl.float32) b_h *= exp(b_g) else: - b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk).to(tl.float32) b_h *= exp(b_gk[:, None]) if IS_BETA_HEADWISE: b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) @@ -256,8 +240,8 @@ def _ensure_gate_token_strided(x: torch.Tensor, inner_numel: int): """Return a_raw/b_raw and token stride, copying only when needed.""" if x is None: return None, 0 - # Gates use [tokens, HV], or [tokens, HV, K] for KDA's per-channel decay. - if x.stride(-1) != 1 or (x.ndim == 3 and x.stride(-2) != x.shape[-1]): + # a_raw/b_raw are 2D [tokens, HV]; the tail HV dimension must be packed. + if x.stride(1) != 1: x = x.contiguous() return x, inner_numel return x, x.stride(0) @@ -283,7 +267,6 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, - lower_bound: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -291,12 +274,10 @@ def fused_recurrent_gated_delta_rule_fwd( # Qwen3Next MTP verify path passes cu_seqlens for variable-length verify # chunks. Both flow through the per-token strided-view path below. N = B if cu_seqlens is None else len(cu_seqlens) - 1 - fuse_gating = A_log is not None - is_kda = a_raw.ndim == 3 if fuse_gating else g.ndim == 4 q, stride_q_tok = _ensure_qkv_token_strided(q, H * K) k, stride_k_tok = _ensure_qkv_token_strided(k, H * K) v, stride_v_tok = _ensure_qkv_token_strided(v, HV * V) - a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV * K if is_kda else HV) + a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV) b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw, HV) BK = triton.next_power_of_2(K) if T == 1: @@ -313,6 +294,8 @@ def fused_recurrent_gated_delta_rule_fwd( NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" + fuse_gating = A_log is not None + if out is not None: o = out.unsqueeze(0) if out.ndim == v.ndim else out else: @@ -385,11 +368,10 @@ def fused_recurrent_gated_delta_rule_fwd( stride_write_indices_tok=stride_write_indices_tok, SOFTPLUS_BETA=1.0, SOFTPLUS_THRESHOLD=20.0, - LOWER_BOUND=lower_bound, IS_BETA_HEADWISE=False if fuse_gating else (beta.ndim == v.ndim), USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, - IS_KDA=is_kda, + IS_KDA=False, FUSE_GATING=fuse_gating, num_warps=num_warps, num_stages=num_stages, @@ -420,7 +402,6 @@ def forward( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, - lower_bound: float | None = None, ): # q/k/v/a_raw/b_raw may be non-contiguous column views of one projection # output; the kernel handles them via per-token strides (no copies). @@ -443,7 +424,6 @@ def forward( a_raw=a_raw, b_raw=b_raw, out=out, - lower_bound=lower_bound, ) return o, final_state @@ -469,7 +449,6 @@ def fused_recurrent_gated_delta_rule( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, - lower_bound: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -481,7 +460,7 @@ def fused_recurrent_gated_delta_rule( values of shape `[B, T, HV, V]`. GVA is applied if `HV > H`. g (torch.Tensor): - Log decays of shape `[B, T, HV]` for GDN or `[B, T, HV, K]` for KDA. + g (decays) of shape `[B, T, HV]`. beta (torch.Tensor): betas of shape `[B, T, HV]`. scale (Optional[int]): @@ -502,15 +481,6 @@ def fused_recurrent_gated_delta_rule( Indices to map the input sequences to the initial/final states. num_accepted_tokens (Optional[torch.Tensor]): Number of accepted tokens for each sequence during decoding. - a_raw (Optional[torch.Tensor]): - Raw decay gates of shape `[B*T, HV]` for GDN or `[B*T, HV, K]` - for KDA. With `A_log`, `dt_bias`, and `b_raw`, fuse gate computation. - `A_log` is per-head; `dt_bias` has the same trailing shape as `a_raw`. - b_raw (Optional[torch.Tensor]): - Raw beta gates of shape `[B*T, HV]`, before sigmoid. - lower_bound (Optional[float]): - Use `lower_bound * sigmoid(exp(A_log) * (a_raw + dt_bias))` for - fused log decays. `None` keeps `-exp(A_log) * softplus(a_raw + dt_bias)`. Returns: o (torch.Tensor): @@ -561,6 +531,5 @@ def fused_recurrent_gated_delta_rule( a_raw, b_raw, out, - lower_bound, ) return o, final_state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py new file mode 100644 index 0000000000..6482114f28 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py @@ -0,0 +1,91 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _kda_decode( + Q, + K, + V, + G, + B, + A, + Bias, + State, + Idx, + O, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SG: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + LOWER: tl.constexpr, + BV: tl.constexpr, +): + row_head = tl.program_id(1) + row, head = row_head // H, row_head % H + ki = tl.arange(0, D) + vi = tl.program_id(0) * BV + tl.arange(0, BV) + q = tl.load(Q + row * SQ + head * D + ki).to(tl.float32) + k = tl.load(K + row * SK + head * D + ki).to(tl.float32) + v = tl.load(V + row * 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 + row * SG + head * D + ki).to(tl.float32) + bias = tl.load(Bias + head * D + ki) + amplitude = tl.exp(tl.load(A + head)) + decay = tl.exp(LOWER * tl.sigmoid(amplitude * (gate + bias))) + beta = tl.sigmoid(tl.load(B + row * SB + head).to(tl.float32)) + req = tl.load(Idx + row) + ptr = State + (req * H + head) * D * D + ki[:, None] * D + vi[None, :] + state = tl.load(ptr).to(tl.float32) * decay[:, None] + delta = (v - tl.sum(state * k[:, None], 0)) * beta + state += k[:, None] * delta[None, :] + tl.store(ptr, state) + out = tl.sum(state * q[:, None], 0) + tl.store(O + row_head * D + vi, out) + + +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, +): + assert inplace_final_state and q.shape[1] == 1 + batch, _, heads, dim = q.shape + assert dim == 128 + 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, + out, + q.stride(0), + k.stride(0), + v.stride(0), + raw_gate.stride(0), + raw_beta.stride(0), + heads, + dim, + lower_bound, + 32, + num_warps=4, + ) + return out, initial_state diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py index eff554a2a4..969a3f354d 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py @@ -60,65 +60,6 @@ def run(q_, k_, v_, a_, b_, state): assert torch.equal(state_ref, state_strided) -@pytest.mark.parametrize("per_channel,lower_bound", [(False, None), (True, None), (True, -5.0)]) -@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) -def test_decode_fused_and_precomputed_gates_match_reference(per_channel, lower_bound, state_dtype): - """Exercise scalar/vector gates, strided projections, and separate state slots.""" - torch.manual_seed(53) - batch, H, HV, K, V = 4, 2, 4, 128, 64 - gate_dim = HV * K if per_channel else HV - widths = [H * K, H * K, HV * V, gate_dim, HV] - mixed = torch.randn(batch, sum(widths), device="cuda", dtype=torch.bfloat16) - q, k, v, a, b = mixed.split(widths, dim=-1) - q, k, v = q.view(batch, 1, H, K), k.view(batch, 1, H, K), v.view(batch, 1, HV, V) - if per_channel: - a = a.view(batch, HV, K) - A_log = torch.randn(HV, device="cuda") * 0.1 - bias = torch.randn(a.shape[1:], device="cuda") * 0.1 - amplitude = A_log.exp().view(HV, 1) if per_channel else A_log.exp() - x = a.float() + bias - if lower_bound is None: - g = -amplitude * torch.nn.functional.softplus(x) - else: - g = lower_bound * torch.sigmoid(amplitude * x) - beta = b.float().sigmoid() - decay = g.exp().unsqueeze(-1) if per_channel else g.exp()[..., None, None] - q_ref = q[:, 0].float() - k_ref = k[:, 0].float() - q_ref = (q_ref * torch.rsqrt(q_ref.square().sum(-1, keepdim=True) + 1e-6) / K ** 0.5).repeat_interleave( - HV // H, dim=1 - ) - k_ref = (k_ref * torch.rsqrt(k_ref.square().sum(-1, keepdim=True) + 1e-6)).repeat_interleave(HV // H, dim=1) - state_ref = torch.randn(12, HV, K, V, device="cuda", dtype=state_dtype) * 0.1 - state_fused, state_precomputed = state_ref.clone(), state_ref.clone() - read_idx = torch.tensor([4, 1, 7, 3], device="cuda", dtype=torch.int32) - write_idx = torch.tensor([2, 9, 0, 6], device="cuda", dtype=torch.int32) - for _ in range(3): - state = state_ref[read_idx].float() * decay - delta = (v[:, 0].float() - torch.einsum("bhkv,bhk->bhv", state, k_ref)) * beta[..., None] - state += k_ref[..., None] * delta[..., None, :] - expected = torch.einsum("bhkv,bhk->bhv", state, q_ref).unsqueeze(1) - state_ref[write_idx] = state.to(state_dtype) - for cache, gates in ( - (state_fused, dict(A_log=A_log, dt_bias=bias, a_raw=a, b_raw=b, lower_bound=lower_bound)), - (state_precomputed, dict(g=g.unsqueeze(1), beta=beta.unsqueeze(1))), - ): - output, _ = fused_recurrent_gated_delta_rule( - q, - k, - v, - initial_state=cache, - ssm_state_indices=read_idx, - ssm_state_write_indices=write_idx, - use_qk_l2norm_in_kernel=True, - **gates, - ) - torch.testing.assert_close(output.float(), expected, atol=5e-4, rtol=1e-2) - # Comparing the entire cache also checks that unrelated slots stay intact. - torch.testing.assert_close(cache, state_ref, atol=5e-4, rtol=1e-2) - read_idx, write_idx = write_idx, read_idx - - # NOTE: the decode-only `cu_seqlens is None` contract from upstream #1349 was # intentionally lifted on this branch so the Qwen3Next MTP verify path can drive # the kernel with variable-length verify chunks (cu_seqlens + 2D SSM index diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 9c8d08539e..76bf7ef22d 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -8,9 +8,7 @@ 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 -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools, gather_pools, expand_topk from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 from lightllm.models.glm5_next.triton_kernel.mhc import hc_pre_norm, hc_pre_reference, hc_post, hc_post_reference @@ -73,18 +71,16 @@ def test_kda_chunk_and_decode_match_recurrence(tokens): states[2] = initial[0] unchanged = states[[0, 1, 3]].clone() for i in range(tokens): - out, _ = fused_recurrent_gated_delta_rule( + out, _ = fused_recurrent_kda( q[:, i : i + 1], k[:, i : i + 1], v[:, i : i + 1], - a_raw=gate[:, i], - b_raw=beta[:, i], - A_log=a, - dt_bias=bias.flatten(), - initial_state=states, - ssm_state_indices=torch.tensor([2], device="cuda", dtype=torch.int32), - use_qk_l2norm_in_kernel=True, - lower_bound=-5.0, + 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) From ad73d17c0b8ad486d03bb3f0aefbabbb4dd5666f Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:27:12 +0000 Subject: [PATCH 09/24] refactor: simplify KDA gate prefill and use LightLLM autotune --- .../common/basemodel/attention/linear/kda.py | 5 - .../linear_att/fla/ops/chunk_delta_h.py | 4 +- .../triton_kernel/linear_att/fla/ops/kda.py | 205 ++++++++++++------ .../linear_att/fla/ops/solve_tril.py | 4 +- unit_tests/models/glm5_next/test_kernels.py | 52 ++++- 5 files changed, 194 insertions(+), 76 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 34c5b52d33..2a637ed492 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -21,7 +21,6 @@ ) 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.fla.ops.index import prepare_chunk_indices if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -59,13 +58,10 @@ def reshape_qkv(self, value: torch.Tensor, *, decode: bool): class KDAPrefillAttState(BasePrefillAttState): b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None - chunk_indices: 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 - # Build variable-length chunk metadata once for all KDA layers. - self.chunk_indices = prepare_chunk_indices(self.infer_state.b1_cu_q_seq_len, 64) def prefill_att( self, @@ -113,7 +109,6 @@ def prefill_att( output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=self.infer_state.b1_cu_q_seq_len, - chunk_indices=self.chunk_indices, safe_gate=True, lower_bound=backend.lower_bound, ) 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 3ca4d6e39b..7b1495d783 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 @@ -268,7 +268,6 @@ def chunk_gated_delta_rule_fwd_h( save_new_value: bool = True, cu_seqlens: torch.LongTensor | None = None, run_config=None, - chunk_indices: torch.Tensor | None = None, use_exp2: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # This kernel is slightly different from fla to support Q/K with different head numbers. @@ -277,8 +276,7 @@ def chunk_gated_delta_rule_fwd_h( H = u.shape[-2] BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None # N: the actual number of sequences in the batch with either equal or variable lengths if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py index f8f4e33f80..82917c121e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -10,6 +10,8 @@ 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 @@ -685,17 +687,8 @@ def grid(meta): return o -@triton.heuristics( - { - "HAS_BIAS": lambda args: args["g_bias"] is not None, - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - } -) -@triton.autotune( - configs=[triton.Config({"BD": BD}, num_warps=num_warps) for BD in [32, 64] for num_warps in [2, 4, 8]], - key=["H", "D", "BT", "IS_VARLEN"], -) -@triton.jit(do_not_specialize=["T"]) +@triton.heuristics({"HAS_BIAS": lambda args: args["g_bias"] is not None}) +@triton.jit def kda_gate_cumsum_fwd_kernel( g, A, @@ -703,63 +696,67 @@ def kda_gate_cumsum_fwd_kernel( g_bias, cu_seqlens, chunk_indices, + # Element strides for input/output [T, H, D]: token, head, channel. + stride_g_t: tl.constexpr, + stride_g_h: tl.constexpr, + stride_g_d: tl.constexpr, + stride_y_t: tl.constexpr, + stride_y_h: tl.constexpr, + stride_y_d: tl.constexpr, cumsum_scale, beta, threshold, SAFE_GATE: tl.constexpr, LOWER_BOUND: tl.constexpr, - T, H: tl.constexpr, D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, HAS_BIAS: tl.constexpr, - IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = ( - tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), - ) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), - ) - T = eos - bos - else: - bos = i_b * T - + # One program handles one [BT, BD] tile for one request/head. + dim_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 * BT + dim_start = dim_block_id * BD + + # Fix the request/head, then view [T, H, D] as a [seq_len, D] matrix. + # Moving one token/channel advances by stride_*_t/stride_*_d elements. + g_seq_head = g + seq_start * stride_g_t + head_id * stride_g_h + y_seq_head = y + seq_start * stride_y_t + head_id * stride_y_h p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), + base=g_seq_head, + shape=(seq_len, D), + strides=(stride_g_t, stride_g_d), + offsets=(chunk_start, dim_start), + block_shape=(BT, BD), + order=(1, 0), ) p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), + base=y_seq_head, + shape=(seq_len, D), + strides=(stride_y_t, stride_y_d), + offsets=(chunk_start, dim_start), + block_shape=(BT, BD), + order=(1, 0), ) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, boundary_check=(0, 1), padding_option="zero").to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) - b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + dim_indices = dim_start + tl.arange(0, BD) + b_bias = tl.load(g_bias + head_id * D + dim_indices, mask=dim_indices < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] - b_a = tl.load(A + i_h).to(tl.float32) + b_a = tl.load(A + head_id).to(tl.float32) b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) if SAFE_GATE: - # y = lower_bound * sigmoid(exp(A) * (g + g_bias)); bounded to - # (lower_bound, 0). Mirrors the SGlang safe_gate branch used by GLM5-Next - # checkpoints whose linear_attn_config["safe_gate"] is True. + # 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 @@ -780,34 +777,68 @@ def kda_gate_cumsum_fwd_kernel( tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) +def _get_kda_gate_cumsum_configs(): + return [{"BD": BD, "num_warps": num_warps} for BD 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 { + "H": raw_g.shape[1], + "D": raw_g.shape[2], + "BT": chunk_size, + "SAFE_GATE": safe_gate, + "HAS_BIAS": g_bias is not None, + "dtype": str(raw_g.dtype), + "out_dtype": str(output_dtype or raw_g.dtype), + } + + +@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 T. +) 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, - cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, output_dtype: torch.dtype | None = torch.float, safe_gate: bool = False, lower_bound: float = -5.0, + run_config: dict | None = None, ) -> torch.Tensor: - if cu_seqlens is not None: - assert raw_g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - B, T, H, D = raw_g.shape - if chunk_indices is None and cu_seqlens is not None: + """Activate packed decay gates and return chunk-local log2 prefix sums in [T, H, D]. + + raw_g: [T, H, D], packed tokens, local heads, and key channels. + Input/output addressing uses each tensor's strides, measured in elements. + A_log: [H]; g_bias: [H * D] or [H, D], or None to skip the bias. + cu_seqlens: [N + 1], required token boundaries for N packed requests. + run_config: optional LightLLM autotune config with BD and num_warps. + """ + assert raw_g.ndim == 3, "raw_g must have packed shape [T, H, D]" + assert cu_seqlens is not None, "cu_seqlens is required for packed KDA prefill" + H, D = raw_g.shape[1:] + if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + NT = len(chunk_indices) A_log = A_log.reshape(-1) if g_bias is not None: g_bias = g_bias.reshape(-1) y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) - def grid(meta): - return (cdiv(meta["D"], meta["BD"]), NT, B * H) + if run_config is None: + run_config = {"BD": 32, "num_warps": 4} + BD = run_config.get("BD", 32) + num_warps = run_config.get("num_warps", 4) + grid = (cdiv(D, BD), NT, H) kda_gate_cumsum_fwd_kernel[grid]( g=raw_g, A=A_log, @@ -815,6 +846,12 @@ def grid(meta): g_bias=g_bias, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + stride_g_t=raw_g.stride(0), + stride_g_h=raw_g.stride(1), + stride_g_d=raw_g.stride(2), + stride_y_t=y.stride(0), + stride_y_h=y.stride(1), + stride_y_d=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`. @@ -823,10 +860,11 @@ def grid(meta): threshold=threshold, SAFE_GATE=safe_gate, LOWER_BOUND=lower_bound, - T=T, H=H, D=D, BT=chunk_size, + BD=BD, + num_warps=num_warps, ) return y @@ -864,7 +902,6 @@ def _chunk_kda_fwd_with_cumulative_g( A = solve_tril( A=A, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, output_dtype=k.dtype, ) w, u, _, kg = recompute_w_u_fwd( @@ -885,7 +922,6 @@ def _chunk_kda_fwd_with_cumulative_g( initial_state=initial_state, output_final_state=output_final_state, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, chunk_size=chunk_size, use_exp2=True, ) @@ -954,16 +990,19 @@ def chunk_kda_with_fused_gate_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - cu_seqlens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor, chunk_indices: torch.Tensor | None = None, safe_gate: bool = False, lower_bound: float = -5.0, ): + assert raw_g.ndim == 4 and raw_g.shape[0] == 1, "KDA prefill expects packed gates shaped [1, T, H, D]" chunk_size = FLA_CHUNK_SIZE - if chunk_indices is None and cu_seqlens is not None: + if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # The gate kernel uses [T, H, D]; downstream FLA ops use [1, T, H, D]. + # Removing/restoring the leading dimension only creates tensor views. g = fused_kda_gate_chunk_cumsum( - raw_g, + raw_g.squeeze(0), A_log=A_log, g_bias=g_bias, cu_seqlens=cu_seqlens, @@ -971,7 +1010,7 @@ def chunk_kda_with_fused_gate_fwd( 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, @@ -1029,17 +1068,55 @@ def chunk_kda_with_fused_gate( 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, - cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, safe_gate: bool = False, lower_bound: float = -5.0, **kwargs, ): - """Run chunk KDA from raw gate projection using fused gate+cumsum.""" + """Run 64-token chunk KDA with fused decay-gate activation and chunk-local prefix sums. + + Shapes use token count T, local heads H, and key/value dimensions K/V. + The leading dimension is always 1; required cu_seqlens splits the packed tokens + into N=len(cu_seqlens)-1 requests. + + Args: + q, k: [1, T, H, K], query/key projections, optionally L2-normalized by this function. + v: [1, T, H, V], value projections. + raw_g: [1, T, H, K], raw per-token, per-key-channel decay-gate projection. + beta: [1, T, H], per-token/head update strength; sigmoid is applied by the caller. + A_log: [H], learned per-head log gate scale, shared across tokens and key channels. + g_bias: [H * K] or [H, K], learned gate bias per head/key channel; None skips the bias. + initial_state: [N, H, K, V], previous recurrent state; None starts from zeros. + cu_seqlens: [N + 1], required cumulative sequence lengths for packed inputs. + chunk_indices: [num_chunks, 2], optional (request ID, local 64-token chunk ID) pairs. + None prepares the indices from cu_seqlens; provided indices are reused. + + Per-token formulas (one sequence/head, column vectors, state S: [K, V]): + a = exp(A_log), bias = 0 if g_bias is None else g_bias + ell_t = lower_bound * sigmoid(a * (raw_g_t + bias)) # safe_gate=True + ell_t = -a * softplus(raw_g_t + bias) # safe_gate=False + alpha_t = exp(ell_t) + S_decay = diag(alpha_t) @ S_prev + delta_t = beta_t * (v_t - S_decay.T @ k_t) + S_t = S_decay + outer(k_t, delta_t) + o_t = scale * (S_t.T @ q_t) + + When use_qk_l2norm_in_kernel=True, q_t/k_t above are normalized as + x / sqrt(sum(x * x) + 1e-6). The default scale is K ** -0.5. + + The fused gate kernel stores G_t = sum(ell_r, r=chunk_start..t) / ln(2). + This prefix sum resets within each sequence at every 64-token chunk boundary. + For j <= i in the same chunk, exp2(G_i - G_j) = product(alpha_r, r=j+1..i). + + Returns: + Output [1, T, H, V] in v.dtype, and final state [N, H, K, V] in float32. + The final state is None when output_final_state=False. + """ if scale is None: scale = k.shape[-1] ** -0.5 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py index 1053d281a0..b5b6cfc369 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py @@ -412,7 +412,6 @@ def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float, - chunk_indices: torch.Tensor | None = None, ) -> torch.Tensor: """ Compute the inverse of the matrix I + A @@ -434,8 +433,7 @@ def solve_tril( output_dtype = A.dtype if output_dtype is None else output_dtype B, T, H, BT = A.shape - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) Ai = torch.zeros_like(A, dtype=output_dtype) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 76bf7ef22d..63c34bcab8 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -1,4 +1,5 @@ import dataclasses +import math from types import SimpleNamespace import pytest @@ -7,7 +8,10 @@ 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 +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, expand_topk from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 @@ -87,6 +91,52 @@ def test_kda_chunk_and_decode_match_recurrence(tokens): 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) + + def test_kpool_chunk_boundaries_and_fragmented_token_kv(): # Two unaligned requests, with a pool completed after restoring token KV. seqs = [9, 7] From 6a416b862e294bf8f74aa0cbf7044e9282232fbd Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:12:55 +0000 Subject: [PATCH 10/24] feat: finalize GLM5.3 Flash cache, attention and PD support Store compressed indexer pools in packed token KV with per-request tails, and transfer the complete hybrid runtime state across PD stages. Use native 512-wide MLA and vLLM Top512 selection as an indexer method. Share mHC kernels, update packed KDA tuning, and fix GLM tool-call parsing. Validation: 118 targeted tests passed on H200; Python and shell syntax checks and git diff --check passed. --- .../common/basemodel/attention/linear/kda.py | 23 +- .../basemodel/attention/nsa/glm5_next.py | 35 +- .../linear_att/fla/ops/chunk_delta_h.py | 26 +- .../triton_kernel/linear_att/fla/ops/kda.py | 1467 ++++++++++------- .../basemodel/triton_kernel/mhc/__init__.py | 43 + .../basemodel/triton_kernel/mhc/post.py | 132 ++ .../basemodel/triton_kernel/mhc/pre_norm.py | 249 +++ .../basemodel/triton_kernel/mhc/streams.py | 31 + .../common/kv_cache_mem_manager/__init__.py | 2 + .../glm5_next_mem_manager.py | 123 ++ .../operator/linear_att.py | 3 +- .../qwen3next_mem_manager.py | 10 +- lightllm/common/req_manager/__init__.py | 9 +- lightllm/common/req_manager/glm5_next.py | 34 + .../common/state_cache_manager/__init__.py | 10 +- .../state_cache_manager/glm5_next.py} | 19 +- .../common/state_cache_manager/linear_att.py | 4 - lightllm/models/glm5_next/README.md | 127 -- lightllm/models/glm5_next/indexer.py | 57 +- .../glm5_next/layer_infer/pre_layer_infer.py | 20 + .../layer_infer/transformer_layer_infer.py | 61 +- .../layer_weights/transformer_layer_weight.py | 10 +- lightllm/models/glm5_next/mem_manager.py | 44 - lightllm/models/glm5_next/model.py | 20 +- lightllm/models/glm5_next/tokenizer.py | 17 - .../models/glm5_next/triton_kernel/kpool.py | 187 ++- .../models/glm5_next/triton_kernel/mhc.py | 547 ------ lightllm/server/function_call_parser.py | 184 +-- lightllm/server/tokenizer.py | 9 +- requirements.txt | 20 +- .../service/benchmark_glm53_flash.py | 2 +- test/start_scripts/glm53/glm53_pd_1p1d.sh | 77 + .../attention/flashinfer/test_mla_nope.py | 106 ++ .../linear_att/test_kda_autotune.py | 196 +++ .../basemodel/triton_kernel/mhc/test_mhc.py | 100 ++ unit_tests/models/glm5_next/test_cache.py | 61 +- .../models/glm5_next/test_indexer_topk.py | 95 ++ unit_tests/models/glm5_next/test_kernels.py | 367 ++++- unit_tests/models/glm5_next/test_pd_cache.py | 157 ++ unit_tests/models/glm5_next/test_tokenizer.py | 27 - unit_tests/server/test_glm47_tool_parser.py | 127 ++ 41 files changed, 3112 insertions(+), 1726 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/mhc/__init__.py create mode 100644 lightllm/common/basemodel/triton_kernel/mhc/post.py create mode 100644 lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py create mode 100644 lightllm/common/basemodel/triton_kernel/mhc/streams.py create mode 100644 lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py create mode 100644 lightllm/common/req_manager/glm5_next.py rename lightllm/{models/glm5_next/cache_config.py => common/state_cache_manager/glm5_next.py} (80%) delete mode 100644 lightllm/models/glm5_next/README.md create mode 100644 lightllm/models/glm5_next/layer_infer/pre_layer_infer.py delete mode 100644 lightllm/models/glm5_next/mem_manager.py delete mode 100644 lightllm/models/glm5_next/tokenizer.py delete mode 100644 lightllm/models/glm5_next/triton_kernel/mhc.py create mode 100644 test/start_scripts/glm53/glm53_pd_1p1d.sh create mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py create mode 100644 unit_tests/models/glm5_next/test_indexer_topk.py create mode 100644 unit_tests/models/glm5_next/test_pd_cache.py delete mode 100644 unit_tests/models/glm5_next/test_tokenizer.py create mode 100644 unit_tests/server/test_glm47_tool_parser.py diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 2a637ed492..4039ed97da 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -35,7 +35,7 @@ def __init__(self, model: "TpPartBaseModel"): self.head_dim = config["head_dim"] assert self.num_heads % model.tp_world_size_ == 0 self.tp_num_heads = self.num_heads // model.tp_world_size_ - self.tp_projection_size = self.tp_num_heads * self.head_dim + self.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) @@ -46,12 +46,7 @@ def create_att_decode_state(self, infer_state: "InferStateInfo"): return KDADecodeAttState(backend=self, infer_state=infer_state) def split_qkv(self, mixed_qkv: torch.Tensor): - return mixed_qkv.split(self.tp_projection_size, dim=-1) - - def reshape_qkv(self, value: torch.Tensor, *, decode: bool): - if decode: - return value.view(-1, 1, self.tp_num_heads, self.head_dim) - return value.view(1, -1, self.tp_num_heads, self.head_dim) + return mixed_qkv.split(self.tp_hidden_size, dim=-1) @dataclasses.dataclass @@ -92,8 +87,11 @@ def prefill_att( activation="silu", ).transpose(0, 1) - q, k, v = [backend.reshape_qkv(x, decode=False) for x in backend.split_qkv(mixed_qkv)] - raw_gate = raw_gate.view(1, -1, backend.tp_projection_size) + 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() @@ -151,8 +149,11 @@ def decode_att( activation="silu", conv_state_indices=self.b_conv_buffer_idx, ) - q, k, v = [backend.reshape_qkv(x, decode=True) for x in backend.split_qkv(mixed_qkv)] - raw_gate = raw_gate.view(-1, 1, backend.tp_projection_size) + 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) output, _ = fused_recurrent_kda( q=q, diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py index c911825119..5fd631d092 100644 --- a/lightllm/common/basemodel/attention/nsa/glm5_next.py +++ b/lightllm/common/basemodel/attention/nsa/glm5_next.py @@ -19,26 +19,16 @@ def create_att_decode_state(self, infer_state): @dataclasses.dataclass class Glm5NextSparsePrefillState(NsaFlashMlaSparsePrefillAttState): - query_batch: torch.Tensor = None - - def init_state(self): - super().init_state() - state = self.infer_state - self.query_batch = torch.repeat_interleave( - torch.arange(state.batch_size, device=state.b_req_idx.device, dtype=torch.int32), - state.b_q_seq_len, - output_size=state.input_ids.numel(), - ) - def _nsa_prefill_att(self, q, kv, att_control): from sgl_kernel.flash_mla import flash_mla_sparse_fwd tokens, heads, dim = q.shape - # The installed Hopper kernel accepts 576-wide Q/K and 64 heads. - # Zero padding preserves NoPE attention and avoids a runtime fork. + # FlashMLA accepts native 512-wide NoPE Q/K; head counts still use 64-head tiles. padded_heads = ((heads + 63) // 64) * 64 - padded_q = q.new_zeros((tokens, padded_heads, dim + 64)) - padded_q[:, :heads, :dim] = q + 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, @@ -52,12 +42,8 @@ def _nsa_prefill_att(self, q, kv, att_control): @dataclasses.dataclass class Glm5NextSparseDecodeState(NsaFlashMlaSparseDecodeAttState): - query_batch: torch.Tensor = None - def init_state(self): super().init_state() - state = self.infer_state - self.query_batch = torch.arange(state.batch_size, device=state.b_req_idx.device, dtype=torch.int32) pool = self.backend.model.config["index_kpool"] topk = self.backend.model.config["index_topk"] self.nsa_cache_seqlens = ( @@ -69,13 +55,15 @@ def _nsa_decode_att(self, q, kv, att_control): from sgl_kernel.flash_attn import flash_attn_with_kvcache q_nope, _ = q - q_rope = q_nope.new_zeros((*q_nope.shape[:-1], 64)) + kv_nope = kv.view(-1, 1, 1, 512) params = att_control.nsa_decode_dict + # only_qv skips QK entirely. Reuse views for the API's required Q/K tensors + # so the wrapper does not allocate a dummy 64-wide query or KV cache. return flash_attn_with_kvcache( - q=q_rope, + q=q_nope[..., :64], qv=q_nope, - k_cache=kv[:, :, 512:].view(-1, 1, 1, 64), - v_cache=kv[:, :, :512].view(-1, 1, 1, 512), + k_cache=kv_nope[..., :64], + 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, @@ -83,4 +71,5 @@ def _nsa_decode_att(self, q, kv, att_control): 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/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 7b1495d783..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 @@ -60,6 +60,16 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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: @@ -114,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)) @@ -128,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)) @@ -144,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)) @@ -166,6 +177,7 @@ 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, @@ -199,6 +211,8 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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) @@ -270,6 +284,14 @@ def chunk_gated_delta_rule_fwd_h( 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] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py index 82917c121e..6d627078dd 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -2,7 +2,22 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang -"""Chunkwise KDA prefill with per-channel decay gates.""" +"""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 @@ -44,645 +59,817 @@ def kda_safe_gate( ``gate_bias`` is per head/key coordinate. """ - head_count = a_log.numel() - key_dim = gate_bias.numel() // head_count - gate = raw_gate.float().view(*raw_gate.shape[:-1], head_count, key_dim) - amplitude = a_log.float().reshape(*((1,) * (gate.ndim - 2)), head_count, 1).exp() - bias = gate_bias.float().reshape(*((1,) * (gate.ndim - 2)), head_count, key_dim) + 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.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["BC"], -) -@triton.jit(do_not_specialize=["T"]) +@triton.jit def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( q, k, g, beta, - A, + Akk, Aqk, scale, cu_seqlens, chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, + 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, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_i, i_j = i_c // NC, i_c % NC - if IS_VARLEN: - i_n, i_t = ( - tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + """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), ) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + 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), ) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - if i_t * BT + i_i * BC >= T: - return - if i_i <= i_j: - return + 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)) - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - A += (bos * H + i_h) * BT - Aqk += (bos * H + i_h) * BT - - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_kt = tl.make_block_ptr(k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - # [BK,] - b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) - # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) - # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) - # [BC, BC] - b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) - b_A += tl.dot(b_k, b_ktg) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale - b_Aqk += tl.dot(b_qg, b_ktg) - - b_A *= b_b[:, None] - - p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], - key=["BK", "BT"], -) -@triton.jit(do_not_specialize=["T"]) + +@triton.jit def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( q, k, g, beta, - A, + Akk, Aqk, scale, cu_seqlens, chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, + head_num: tl.constexpr, + head_dim: tl.constexpr, + chunk_size: tl.constexpr, + subchunk_size: tl.constexpr, + head_block_size: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = ( - tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), - ) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), - ) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T + """Complete Akk/Aqk inside each diagonal [subchunk_size, subchunk_size] subtile of a chunk. - if i_t * BT + i_i * BC >= T: + 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 - o_i = tl.arange(0, BC) - o_k = tl.arange(0, BK) - m_k = o_k < K - m_A = (i_t * BT + i_i * BC + o_i) < T - o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + 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( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), + 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( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), + 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( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), + 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), ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - - p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h - b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] - - p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k - p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k - - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :]) - b_A = tl.sum(b_k * b_ktg, 1) - b_A = tl.where(o_i > j, b_A, 0.0) - b_Aqk = tl.sum(b_q * b_ktg, 1) - b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) - tl.store(A + o_A + j, b_A, mask=m_A) - tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) - p_kt += H * K - p_gk += H * K def chunk_kda_scaled_dot_kkt_fwd( q: torch.Tensor, k: torch.Tensor, - gk: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.Tensor | None = None, + 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]: - r""" - Compute beta * K * K^T. - - Args: - k (torch.Tensor): - The key tensor of shape `[B, T, H, K]`. - beta (torch.Tensor): - The beta tensor of shape `[B, T, H]`. - gk (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. - cu_seqlens (torch.Tensor): - The cumulative sequence lengths of the input tensor. - Default: None - chunk_size (int): - The chunk size. Default: 64. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float32` - - Returns: - beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """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. """ - B, T, H, K = k.shape - assert K <= 256 - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - BC = min(16, BT) - NC = cdiv(BT, BC) - BK = max(next_power_of_2(K), 16) - A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) - Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) - grid = (NT, NC * NC, B * H) - chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( + _, 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, - g=gk, + gk=gk, beta=beta, - A=A, + Akk=Akk, Aqk=Aqk, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - NC=NC, ) - grid = (NT, NC, B * H) - chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( + _chunk_kda_scaled_dot_kkt_sub_intra( q=q, k=k, - g=gk, + gk=gk, beta=beta, - A=A, + Akk=Akk, Aqk=Aqk, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, ) - return A, Aqk + return Akk, Aqk -@triton.heuristics( - { - "STORE_QG": lambda args: args["qg"] is not None, - "STORE_KG": lambda args: args["kg"] is not None, - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - } -) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], -) -@triton.jit(do_not_specialize=["T"]) +@triton.jit def recompute_w_u_fwd_kernel( - q, k, - qg, kg, v, beta, w, u, - A, + R, gk, cu_seqlens, chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STORE_QG: tl.constexpr, - STORE_KG: tl.constexpr, - IS_VARLEN: tl.constexpr, + 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, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = ( - tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), - ) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), - ) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + """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 i_v in range(tl.cdiv(V, BV)): + for value_block_id in range(tl.cdiv(value_head_dim, value_block_size)): p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), + 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( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), + 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), ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + 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 i_k in range(tl.cdiv(K, BK)): + for key_block_id in range(tl.cdiv(head_dim, head_block_size)): p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + 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( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + 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), ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = b_k * b_b[:, None] + k_tile = tl.load(p_k, boundary_check=(0, 1)) + weighted_k = k_tile * beta_tile[:, None] p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + 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, ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kb *= exp2(b_gk) - if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) - if STORE_KG: - last_idx = min(i_t * BT + BT, T) - 1 - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0) - b_kg = b_k * exp2(b_gn - b_gk) - - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) - - b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + 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, - A: torch.Tensor, - q: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - cu_seqlens: torch.Tensor | None = None, + R: torch.Tensor, + gk: torch.Tensor, + cu_seqlens: torch.Tensor, chunk_indices: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - BK = 64 - BV = 64 + 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 and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + 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) if gk is not None else None - recompute_w_u_fwd_kernel[(NT, B * H)]( - q=q, + kg = torch.empty_like(k) + recompute_w_u_fwd_kernel[(chunk_num, head_num)]( k=k, - qg=None, kg=kg, v=v, beta=beta, w=w, u=u, - A=A, + R=R, gk=gk, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, + 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, None, kg + return w, u, kg -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["BT"], -) -@triton.jit(do_not_specialize=["T"]) +@triton.jit def chunk_gla_fwd_kernel_o( q, v, g, h, o, - A, + Aqk, cu_seqlens, chunk_indices, scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, + 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, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_tg = i_t - i_n, i_t = ( - tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), - ) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), - ) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T + """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 - m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + causal_mask = tl.arange(0, chunk_size)[:, None] >= tl.arange(0, chunk_size)[None, :] - b_o = tl.zeros([BT, BV], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): + 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( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + 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( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + 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( - h + (i_tg * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, i_v * BV), - (BK, BV), - (1, 0), + 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), ) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - # [BT, BK] - b_qg = (b_q * exp2(b_g)).to(b_q.dtype) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - if i_k >= 0: - b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), + # [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( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), + 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), ) - p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) - b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + # 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, - A: torch.Tensor, + Aqk: torch.Tensor, h: torch.Tensor, o: torch.Tensor, scale: float, - cu_seqlens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor, chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, + run_config: dict | None = None, ): - B, T, H, K, V = *q.shape, v.shape[-1] - BT = chunk_size + """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 and cu_seqlens is not None: + if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + chunk_num = len(chunk_indices) - def grid(meta): - return (cdiv(V, meta["BV"]), NT, B * H) + 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, - A=A, + Aqk=Aqk, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, + 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 @@ -691,68 +878,77 @@ def grid(meta): @triton.jit def kda_gate_cumsum_fwd_kernel( g, - A, + A_log, y, g_bias, cu_seqlens, chunk_indices, - # Element strides for input/output [T, H, D]: token, head, channel. - stride_g_t: tl.constexpr, - stride_g_h: tl.constexpr, - stride_g_d: tl.constexpr, - stride_y_t: tl.constexpr, - stride_y_h: tl.constexpr, - stride_y_d: tl.constexpr, + # 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, - H: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, + head_dim: tl.constexpr, + chunk_size: tl.constexpr, + head_block_size: tl.constexpr, HAS_BIAS: tl.constexpr, ): - # One program handles one [BT, BD] tile for one request/head. - dim_block_id, global_chunk_id, head_id = tl.program_id(0), tl.program_id(1), tl.program_id(2) + """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 * BT - dim_start = dim_block_id * BD + chunk_start = chunk_id_in_seq * chunk_size + head_dim_start = head_block_id * head_block_size - # Fix the request/head, then view [T, H, D] as a [seq_len, D] matrix. - # Moving one token/channel advances by stride_*_t/stride_*_d elements. - g_seq_head = g + seq_start * stride_g_t + head_id * stride_g_h - y_seq_head = y + seq_start * stride_y_t + head_id * stride_y_h + # 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, D), - strides=(stride_g_t, stride_g_d), - offsets=(chunk_start, dim_start), - block_shape=(BT, BD), + 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, D), - strides=(stride_y_t, stride_y_d), - offsets=(chunk_start, dim_start), - block_shape=(BT, BD), + 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: - dim_indices = dim_start + tl.arange(0, BD) - b_bias = tl.load(g_bias + head_id * D + dim_indices, mask=dim_indices < D, other=0.0).to(tl.float32) + 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 + head_id).to(tl.float32) + 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)). @@ -771,25 +967,29 @@ def kda_gate_cumsum_fwd_kernel( # b_gate non-zero) participate in the dot product. They only contribute to # out-of-bounds output rows, which are masked away by `boundary_check` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) + 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 [{"BD": BD, "num_warps": num_warps} for BD in [32, 64] for num_warps in [2, 4, 8]] + 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 { - "H": raw_g.shape[1], - "D": raw_g.shape[2], - "BT": chunk_size, + "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), - "out_dtype": str(output_dtype or raw_g.dtype), + "dtype": str(raw_g.dtype).removeprefix("torch."), + "out_dtype": str(output_dtype or raw_g.dtype).removeprefix("torch."), } @@ -797,7 +997,7 @@ def _get_kda_gate_cumsum_static_key(raw_g, g_bias, chunk_size, output_dtype, saf 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 T. + run_key_func=lambda raw_g: raw_g.shape[0], # Total packed token count. ) def fused_kda_gate_chunk_cumsum( raw_g: torch.Tensor, @@ -813,20 +1013,24 @@ def fused_kda_gate_chunk_cumsum( lower_bound: float = -5.0, run_config: dict | None = None, ) -> torch.Tensor: - """Activate packed decay gates and return chunk-local log2 prefix sums in [T, H, D]. + """Activate packed decay gates and return chunk-local log2 prefix sums in [total_tokens, head_num, head_dim]. - raw_g: [T, H, D], packed tokens, local heads, and key channels. + 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: [H]; g_bias: [H * D] or [H, D], or None to skip the bias. - cu_seqlens: [N + 1], required token boundaries for N packed requests. - run_config: optional LightLLM autotune config with BD and num_warps. + 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 [T, H, D]" + 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" - H, D = raw_g.shape[1:] + head_num, head_dim = raw_g.shape[1:] if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = len(chunk_indices) + chunk_num = len(chunk_indices) A_log = A_log.reshape(-1) if g_bias is not None: @@ -834,24 +1038,24 @@ def fused_kda_gate_chunk_cumsum( y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) if run_config is None: - run_config = {"BD": 32, "num_warps": 4} - BD = run_config.get("BD", 32) + 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(D, BD), NT, H) + grid = (cdiv(head_dim, head_block_size), chunk_num, head_num) kda_gate_cumsum_fwd_kernel[grid]( g=raw_g, - A=A_log, + A_log=A_log, y=y, g_bias=g_bias, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - stride_g_t=raw_g.stride(0), - stride_g_h=raw_g.stride(1), - stride_g_d=raw_g.stride(2), - stride_y_t=y.stride(0), - stride_y_h=y.stride(1), - stride_y_d=y.stride(2), + 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`. @@ -860,10 +1064,9 @@ def fused_kda_gate_chunk_cumsum( threshold=threshold, SAFE_GATE=safe_gate, LOWER_BOUND=lower_bound, - H=H, - D=D, - BT=chunk_size, - BD=BD, + head_dim=head_dim, + chunk_size=chunk_size, + head_block_size=head_block_size, num_warps=num_warps, ) return y @@ -878,17 +1081,32 @@ def _chunk_kda_fwd_with_cumulative_g( scale: float, initial_state: torch.Tensor, output_final_state: bool, - cu_seqlens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor, chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, ): - # `g` must already be chunk-local cumulatively-summed AND scaled by - # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). - # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of - # calling this helper directly unless that invariant is upheld. - # the intra Aqk is kept in fp32 - # the computation has very marginal effect on the entire throughput - A, Aqk = chunk_kda_scaled_dot_kkt_fwd( + """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, @@ -899,21 +1117,28 @@ def _chunk_kda_fwd_with_cumulative_g( chunk_size=chunk_size, output_dtype=torch.float32, ) - A = solve_tril( - A=A, + # 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, ) - w, u, _, kg = recompute_w_u_fwd( + 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, - A=A, + R=R, gk=g, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, ) - del A + 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, @@ -926,11 +1151,14 @@ def _chunk_kda_fwd_with_cumulative_g( 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, - A=Aqk, + Aqk=Aqk, h=h, o=v, scale=scale, @@ -951,10 +1179,15 @@ def chunk_kda_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, - cu_seqlens: torch.Tensor | None = None, + 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) if cu_seqlens is not None else None + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) g = chunk_local_cumsum( g, chunk_size=chunk_size, @@ -995,11 +1228,18 @@ def chunk_kda_with_fused_gate_fwd( safe_gate: bool = False, lower_bound: float = -5.0, ): - assert raw_g.ndim == 4 and raw_g.shape[0] == 1, "KDA prefill expects packed gates shaped [1, T, H, D]" + """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 [T, H, D]; downstream FLA ops use [1, T, H, D]. + # 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), @@ -1032,15 +1272,18 @@ def chunk_kda( 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, - cu_seqlens: torch.Tensor | None = None, **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: - scale = k.shape[-1] ** -0.5 + head_dim = k.shape[-1] + scale = head_dim ** -0.5 if use_qk_l2norm_in_kernel: q = l2norm_fwd(q.contiguous()) @@ -1078,47 +1321,81 @@ def chunk_kda_with_fused_gate( lower_bound: float = -5.0, **kwargs, ): - """Run 64-token chunk KDA with fused decay-gate activation and chunk-local prefix sums. - - Shapes use token count T, local heads H, and key/value dimensions K/V. - The leading dimension is always 1; required cu_seqlens splits the packed tokens - into N=len(cu_seqlens)-1 requests. - - Args: - q, k: [1, T, H, K], query/key projections, optionally L2-normalized by this function. - v: [1, T, H, V], value projections. - raw_g: [1, T, H, K], raw per-token, per-key-channel decay-gate projection. - beta: [1, T, H], per-token/head update strength; sigmoid is applied by the caller. - A_log: [H], learned per-head log gate scale, shared across tokens and key channels. - g_bias: [H * K] or [H, K], learned gate bias per head/key channel; None skips the bias. - initial_state: [N, H, K, V], previous recurrent state; None starts from zeros. - cu_seqlens: [N + 1], required cumulative sequence lengths for packed inputs. - chunk_indices: [num_chunks, 2], optional (request ID, local 64-token chunk ID) pairs. - None prepares the indices from cu_seqlens; provided indices are reused. - - Per-token formulas (one sequence/head, column vectors, state S: [K, V]): - a = exp(A_log), bias = 0 if g_bias is None else g_bias - ell_t = lower_bound * sigmoid(a * (raw_g_t + bias)) # safe_gate=True - ell_t = -a * softplus(raw_g_t + bias) # safe_gate=False - alpha_t = exp(ell_t) - S_decay = diag(alpha_t) @ S_prev - delta_t = beta_t * (v_t - S_decay.T @ k_t) - S_t = S_decay + outer(k_t, delta_t) - o_t = scale * (S_t.T @ q_t) - - When use_qk_l2norm_in_kernel=True, q_t/k_t above are normalized as - x / sqrt(sum(x * x) + 1e-6). The default scale is K ** -0.5. - - The fused gate kernel stores G_t = sum(ell_r, r=chunk_start..t) / ln(2). - This prefix sum resets within each sequence at every 64-token chunk boundary. - For j <= i in the same chunk, exp2(G_i - G_j) = product(alpha_r, r=j+1..i). + """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