From eb3da9c1ddf62b1816d880911f20f1bb431372a7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <1109276519@qq.com> Date: Thu, 25 Dec 2025 16:03:15 +0800 Subject: [PATCH] fuse ssm state store into chunk_gated_delta_rule_fwd_h (#15409) --- .../sglang/srt/layers/attention/fla/chunk.py | 33 ++++++++++--------- .../srt/layers/attention/fla/chunk_delta_h.py | 31 ++++++++--------- python/sglang/srt/layers/attention/fla/kda.py | 16 +++++---- .../attention/hybrid_linear_attn_backend.py | 27 ++++++++------- 4 files changed, 58 insertions(+), 49 deletions(-) diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index fc1206135..e7430f4f9 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -31,7 +31,7 @@ def chunk_gated_delta_rule_fwd( beta: torch.Tensor, scale: float, initial_state: torch.Tensor, - output_final_state: bool, + initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, ): g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens) @@ -48,13 +48,13 @@ def chunk_gated_delta_rule_fwd( g_cumsum=g, cu_seqlens=cu_seqlens, ) - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + h, v_new = chunk_gated_delta_rule_fwd_h( k=k, w=w, u=u, g=g, initial_state=initial_state, - output_final_state=output_final_state, + initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, ) o = chunk_fwd_o( @@ -67,9 +67,9 @@ def chunk_gated_delta_rule_fwd( cu_seqlens=cu_seqlens, ) if SUPPRESS_LEVEL < 3: - return g, o, A, final_state, None, h, None + return g, o, A, None, h, None elif SUPPRESS_LEVEL >= 3: - return g, o, A, final_state, w, h, v_new + return g, o, A, w, h, v_new class ChunkGatedDeltaRuleFunction(torch.autograd.Function): @@ -86,7 +86,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): beta: torch.Tensor, scale: float, initial_state: torch.Tensor, - output_final_state: bool, + initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False, ): @@ -97,7 +97,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): q = l2norm_fwd(q) k = l2norm_fwd(k) - g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd( + g, o, A, w, h, v_new = chunk_gated_delta_rule_fwd( q=q, k=k, v=v, @@ -105,10 +105,10 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): beta=beta, scale=scale, initial_state=initial_state, - output_final_state=output_final_state, + initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, ) - return o.to(q.dtype), final_state, h + return o.to(q.dtype), h @torch.compiler.disable @@ -120,7 +120,7 @@ def chunk_gated_delta_rule( beta: torch.Tensor, scale: float = None, initial_state: torch.Tensor = None, - output_final_state: bool = False, + initial_state_indices: torch.Tensor = None, cu_seqlens: Optional[torch.LongTensor] = None, head_first: bool = False, use_qk_l2norm_in_kernel: bool = False, @@ -217,14 +217,17 @@ def chunk_gated_delta_rule( f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." f"Please flatten variable-length inputs before processing." ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + if ( + initial_state_indices is not None + and initial_state_indices.shape[0] != len(cu_seqlens) - 1 + ): raise ValueError( f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}." ) if scale is None: scale = k.shape[-1] ** -0.5 - o, final_state, h = ChunkGatedDeltaRuleFunction.apply( + o, h = ChunkGatedDeltaRuleFunction.apply( q, k, v, @@ -232,10 +235,10 @@ def chunk_gated_delta_rule( beta, scale, initial_state, - output_final_state, + initial_state_indices, cu_seqlens, use_qk_l2norm_in_kernel, ) if head_first: o = rearrange(o, "b t h ... -> b h t ...") - return o, final_state, h + return o, None, h diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index e926fc4d3..38a7c8f29 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -38,8 +38,8 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( g, gk, h, - h0, - ht, + initial_state, + initial_state_indices, cu_seqlens, chunk_offsets, T, @@ -52,7 +52,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( USE_G: tl.constexpr, USE_GK: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, + INPLACE_UPDATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, IS_VARLEN: tl.constexpr, ): @@ -90,10 +90,14 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( stride_h = H * K * V stride_k = Hg * K stride_w = H * K + + index = tl.load(initial_state_indices + i_n).to(tl.int32) + h0 = initial_state + index * stride_h + ht = initial_state + index * stride_h if USE_INITIAL_STATE: - h0 = h0 + i_nh * K * V - if STORE_FINAL_STATE: - ht = ht + i_nh * K * V + h0 = h0 + i_h * K * V + if INPLACE_UPDATE: + ht = ht + i_h * K * V # load initial state if USE_INITIAL_STATE: @@ -247,7 +251,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 += tl.dot(b_k, b_v) # epilogue - if STORE_FINAL_STATE: + if INPLACE_UPDATE: p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 64: @@ -274,7 +278,7 @@ def chunk_gated_delta_rule_fwd_h( g: Optional[torch.Tensor] = None, gk: Optional[torch.Tensor] = None, initial_state: Optional[torch.Tensor] = None, - output_final_state: bool = False, + initial_state_indices: Optional[torch.Tensor] = None, save_new_value: bool = True, cu_seqlens: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -299,9 +303,6 @@ def chunk_gated_delta_rule_fwd_h( assert K <= 256, "current kernel does not support head dimension larger than 256." h = k.new_empty(B, NT, H, K, V) - final_state = ( - k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - ) v_new = torch.empty_like(u) if save_new_value else None @@ -316,8 +317,8 @@ def chunk_gated_delta_rule_fwd_h( g=g, gk=gk, h=h, - h0=initial_state, - ht=final_state, + initial_state=initial_state, + initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, T=T, @@ -330,10 +331,10 @@ def chunk_gated_delta_rule_fwd_h( USE_G=g is not None, USE_GK=gk is not None, USE_INITIAL_STATE=initial_state is not None, - STORE_FINAL_STATE=final_state is not None, + INPLACE_UPDATE=True, SAVE_NEW_VALUE=v_new is not None, IS_VARLEN=cu_seqlens is not None, num_warps=4, num_stages=2, ) - return h, v_new, final_state + return h, v_new diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/srt/layers/attention/fla/kda.py index 9556fa63f..357138993 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/python/sglang/srt/layers/attention/fla/kda.py @@ -1146,6 +1146,7 @@ def chunk_kda_fwd( beta: torch.Tensor, scale: float, initial_state: torch.Tensor, + initial_state_indices: torch.Tensor, output_final_state: bool, cu_seqlens: torch.LongTensor | None = None, ): @@ -1172,12 +1173,13 @@ def chunk_kda_fwd( cu_seqlens=cu_seqlens, ) del A - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + h, v_new = chunk_gated_delta_rule_fwd_h( k=kg, w=w, u=u, gk=g, initial_state=initial_state, + initial_state_indices=initial_state_indices, output_final_state=output_final_state, cu_seqlens=cu_seqlens, ) @@ -1194,7 +1196,7 @@ def chunk_kda_fwd( chunk_size=chunk_size, ) del Aqk, v_new, h - return o, final_state + return o def chunk_kda( @@ -1205,7 +1207,7 @@ def chunk_kda( beta: torch.Tensor, scale: float = None, initial_state: torch.Tensor = None, - output_final_state: bool = False, + initial_state_indices: torch.Tensor = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, **kwargs, @@ -1217,18 +1219,18 @@ def chunk_kda( q = l2norm_fwd(q.contiguous()) k = l2norm_fwd(k.contiguous()) - o, final_state = chunk_kda_fwd( + o = 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, + initial_state=initial_state, + initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, ) - return o, final_state + return o @triton.autotune( 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 20fbb29f8..5d5a64c9d 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -814,22 +814,17 @@ class KimiLinearAttnBackend(MambaAttnBackendBase): beta = beta.unsqueeze(0) g = g.unsqueeze(0) - initial_state = ssm_states[cache_indices].contiguous() - ( - core_attn_out, - last_recurrent_state, - ) = chunk_kda( + core_attn_out = chunk_kda( q=q, k=k, v=v, g=g, beta=beta, - initial_state=initial_state, - output_final_state=True, + initial_state=ssm_states, + initial_state_indices=cache_indices, use_qk_l2norm_in_kernel=True, cu_seqlens=query_start_loc, ) - ssm_states[cache_indices] = last_recurrent_state return core_attn_out @@ -1064,7 +1059,12 @@ class GDNAttnBackend(MambaAttnBackendBase): retrieve_parent_token=retrieve_parent_token, ) else: - recurrent_state = ssm_states[cache_indices] + # Only cuda env uses fuse ssm_states update + recurrent_state = ssm_states + recurrent_state_indices_args = {"initial_state_indices": cache_indices} + if is_npu(): + recurrent_state = ssm_states[cache_indices] + recurrent_state_indices_args = {} core_attn_out, last_recurrent_state, h = chunk_gated_delta_rule( q=query, k=key, @@ -1072,13 +1072,16 @@ class GDNAttnBackend(MambaAttnBackendBase): g=g, beta=beta, initial_state=recurrent_state, - output_final_state=True, cu_seqlens=query_start_loc, head_first=False, use_qk_l2norm_in_kernel=True, + **recurrent_state_indices_args, ) - last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False) - ssm_states[cache_indices] = last_recurrent_state + if is_npu(): + last_recurrent_state = last_recurrent_state.to( + ssm_states.dtype, copy=False + ) + ssm_states[cache_indices] = last_recurrent_state self._track_mamba_state_extend( forward_batch, h, ssm_states, forward_metadata