From bcc6d84f93fbfbbb64bf4c86356147acee042750 Mon Sep 17 00:00:00 2001 From: strgrb Date: Wed, 21 Jan 2026 19:24:29 +0800 Subject: [PATCH] Use fused_sigmoid_gating_delta_rule_update_kernel for KDA (#17108) --- .../fla/fused_sigmoid_gating_recurrent.py | 16 ++++++++--- .../attention/hybrid_linear_attn_backend.py | 27 +++++++++++-------- python/sglang/srt/models/kimi_linear.py | 11 +++++--- 3 files changed, 37 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py index 1ed66680d..5d4722904 100644 --- a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py @@ -34,6 +34,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( USE_INITIAL_STATE: tl.constexpr, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, IS_VARLEN: tl.constexpr, + IS_KDA: tl.constexpr, ): """ Fused kernel that combines sigmoid gating computation with recurrent delta rule update. @@ -64,8 +65,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel( # Gating computation pointers p_A_log = A_log + i_hv - p_a = a + bos * HV + i_hv - p_dt_bias = dt_bias + i_hv + if IS_KDA: + p_a = a + (bos * HV + i_hv) * K + o_k + p_dt_bias = dt_bias + i_hv * K + o_k + else: + p_a = a + bos * HV + i_hv + p_dt_bias = dt_bias + i_hv mask_k = o_k < K mask_v = o_v < V @@ -119,7 +124,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_q = b_q * scale # Apply gating to hidden state: h *= exp(g) - b_h *= tl.exp(b_g) + if IS_KDA: + b_h *= tl.exp(b_g[:, None]) + else: + b_h *= tl.exp(b_g) # Delta rule: v -= sum(h * k, dim=0) b_v -= tl.sum(b_h * b_k[:, None], 0) @@ -172,6 +180,7 @@ def fused_sigmoid_gating_delta_rule_update( scale: Optional[float] = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: Optional[torch.Tensor] = None, + is_kda: bool = False, ): """ Fused triton implementation of sigmoid gating delta rule update. @@ -221,6 +230,7 @@ def fused_sigmoid_gating_delta_rule_update( USE_INITIAL_STATE=initial_state_source is not None, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_VARLEN=cu_seqlens is not None, + IS_KDA=is_kda, num_warps=num_warps, num_stages=num_stages, ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index e0ddbcd05..a0fb9d845 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -17,7 +17,7 @@ from sglang.srt.layers.attention.fla.fused_recurrent import ( from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) -from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda +from sglang.srt.layers.attention.fla.kda import chunk_kda from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( PAD_SLOT_ID, causal_conv1d_fn, @@ -647,6 +647,9 @@ class KimiLinearAttnBackend(MambaAttnBackendBase): beta = kwargs["beta"] g = kwargs["gate"] + A_log = kwargs["A_log"] + dt_bias = kwargs["dt_bias"] + layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id) q_conv_state, k_conv_state, v_conv_state = layer_cache.conv ssm_states = layer_cache.temporal @@ -686,21 +689,23 @@ class KimiLinearAttnBackend(MambaAttnBackendBase): lambda x: rearrange(x, "n (h d) -> 1 n h d", d=head_dim), (q, k, v) ) - initial_state = ssm_states[cache_indices].contiguous() - ( - core_attn_out, - last_recurrent_state, - ) = fused_recurrent_kda( + core_attn_out = fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + dt_bias=dt_bias, q=q, k=k, v=v, - g=g, - beta=beta, - initial_state=initial_state, - use_qk_l2norm_in_kernel=True, + a=g, + b=beta, + initial_state_source=ssm_states, + initial_state_indices=cache_indices, cu_seqlens=query_start_loc, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + is_kda=True, ) - ssm_states[cache_indices] = last_recurrent_state + return core_attn_out def forward_extend( diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 8a9029ce4..7efb8c2e9 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -316,9 +316,12 @@ class KimiDeltaAttention(nn.Module): beta = self.b_proj(hidden_states)[0].float().sigmoid() forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0] - forget_gate = fused_kda_gate( - forget_gate, self.A_log, self.head_dim, g_bias=self.dt_bias - ) + + # fused_kda_gate is fused to KimiLinearAttentionBackend with decode + if not forward_batch.forward_mode.is_decode(): + forget_gate = fused_kda_gate( + forget_gate, self.A_log, self.head_dim, g_bias=self.dt_bias + ) beta = beta.unsqueeze(0) forget_gate = forget_gate.unsqueeze(0) @@ -336,6 +339,8 @@ class KimiDeltaAttention(nn.Module): "layer_id": self.layer_idx, "beta": beta, "gate": forget_gate, + "A_log": self.A_log, + "dt_bias": self.dt_bias, } core_attn_out = forward_batch.attn_backend.forward(