minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fuse ssm state store into chunk_gated_delta_rule_fwd_h (#15409)

Browse Source
This commit is contained in:
Yi Zhang
2025-12-25 16:03:15 +08:00
committed by GitHub
parent 72a980c6d5
commit eb3da9c1dd
4 changed files with 58 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/fla/chunk.py
+18 -15
View File
@@ -31,7 +31,7 @@ def chunk_gated_delta_rule_fwd(
beta: torch.Tensor, beta: torch.Tensor,
scale: float, scale: float,
initial_state: torch.Tensor, initial_state: torch.Tensor,
output_final_state: bool, initial_state_indices: torch.Tensor,
cu_seqlens: Optional[torch.LongTensor] = None, cu_seqlens: Optional[torch.LongTensor] = None,
): ):
g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens) 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, g_cumsum=g,
cu_seqlens=cu_seqlens, 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, k=k,
w=w, w=w,
u=u, u=u,
g=g, g=g,
initial_state=initial_state, initial_state=initial_state,
output_final_state=output_final_state, initial_state_indices=initial_state_indices,
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
) )
o = chunk_fwd_o( o = chunk_fwd_o(
@@ -67,9 +67,9 @@ def chunk_gated_delta_rule_fwd(
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
) )
if SUPPRESS_LEVEL < 3: if SUPPRESS_LEVEL < 3:
return g, o, A, final_state, None, h, None return g, o, A, None, h, None
elif SUPPRESS_LEVEL >= 3: 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): class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
@@ -86,7 +86,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
beta: torch.Tensor, beta: torch.Tensor,
scale: float, scale: float,
initial_state: torch.Tensor, initial_state: torch.Tensor,
output_final_state: bool, initial_state_indices: torch.Tensor,
cu_seqlens: Optional[torch.LongTensor] = None, cu_seqlens: Optional[torch.LongTensor] = None,
use_qk_l2norm_in_kernel: bool = False, use_qk_l2norm_in_kernel: bool = False,
): ):
@@ -97,7 +97,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
q = l2norm_fwd(q) q = l2norm_fwd(q)
k = l2norm_fwd(k) 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, q=q,
k=k, k=k,
v=v, v=v,
@@ -105,10 +105,10 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
beta=beta, beta=beta,
scale=scale, scale=scale,
initial_state=initial_state, initial_state=initial_state,
output_final_state=output_final_state, initial_state_indices=initial_state_indices,
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
) )
return o.to(q.dtype), final_state, h return o.to(q.dtype), h
@torch.compiler.disable @torch.compiler.disable
@@ -120,7 +120,7 @@ def chunk_gated_delta_rule(
beta: torch.Tensor, beta: torch.Tensor,
scale: float = None, scale: float = None,
initial_state: torch.Tensor = None, initial_state: torch.Tensor = None,
output_final_state: bool = False, initial_state_indices: torch.Tensor = None,
cu_seqlens: Optional[torch.LongTensor] = None, cu_seqlens: Optional[torch.LongTensor] = None,
head_first: bool = False, head_first: bool = False,
use_qk_l2norm_in_kernel: 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"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." 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( raise ValueError(
f"The number of initial states is expected to be equal to the number of input sequences, " 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: if scale is None:
scale = k.shape[-1] ** -0.5 scale = k.shape[-1] ** -0.5
o, final_state, h = ChunkGatedDeltaRuleFunction.apply( o, h = ChunkGatedDeltaRuleFunction.apply(
q, q,
k, k,
v, v,
@@ -232,10 +235,10 @@ def chunk_gated_delta_rule(
beta, beta,
scale, scale,
initial_state, initial_state,
output_final_state, initial_state_indices,
cu_seqlens, cu_seqlens,
use_qk_l2norm_in_kernel, use_qk_l2norm_in_kernel,
) )
if head_first: if head_first:
o = rearrange(o, "b t h ... -> b h t ...") o = rearrange(o, "b t h ... -> b h t ...")
return o, final_state, h return o, None, h
python/sglang/srt/layers/attention/fla/chunk_delta_h.py
+16 -15
View File
@@ -38,8 +38,8 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
g, g,
gk, gk,
h, h,
h0, initial_state,
ht, initial_state_indices,
cu_seqlens, cu_seqlens,
chunk_offsets, chunk_offsets,
T, T,
@@ -52,7 +52,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
USE_G: tl.constexpr, USE_G: tl.constexpr,
USE_GK: tl.constexpr, USE_GK: tl.constexpr,
USE_INITIAL_STATE: tl.constexpr, USE_INITIAL_STATE: tl.constexpr,
STORE_FINAL_STATE: tl.constexpr, INPLACE_UPDATE: tl.constexpr,
SAVE_NEW_VALUE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr,
IS_VARLEN: 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_h = H * K * V
stride_k = Hg * K stride_k = Hg * K
stride_w = H * 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: if USE_INITIAL_STATE:
h0 = h0 + i_nh * K * V h0 = h0 + i_h * K * V
if STORE_FINAL_STATE: if INPLACE_UPDATE:
ht = ht + i_nh * K * V ht = ht + i_h * K * V
# load initial state # load initial state
if USE_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) b_h4 += tl.dot(b_k, b_v)
# epilogue # 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)) 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)) tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
if K > 64: if K > 64:
@@ -274,7 +278,7 @@ def chunk_gated_delta_rule_fwd_h(
g: Optional[torch.Tensor] = None, g: Optional[torch.Tensor] = None,
gk: Optional[torch.Tensor] = None, gk: Optional[torch.Tensor] = None,
initial_state: 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, save_new_value: bool = True,
cu_seqlens: Optional[torch.LongTensor] = None, cu_seqlens: Optional[torch.LongTensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ) -> 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." assert K <= 256, "current kernel does not support head dimension larger than 256."
h = k.new_empty(B, NT, H, K, V) 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 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, g=g,
gk=gk, gk=gk,
h=h, h=h,
h0=initial_state, initial_state=initial_state,
ht=final_state, initial_state_indices=initial_state_indices,
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
chunk_offsets=chunk_offsets, chunk_offsets=chunk_offsets,
T=T, T=T,
@@ -330,10 +331,10 @@ def chunk_gated_delta_rule_fwd_h(
USE_G=g is not None, USE_G=g is not None,
USE_GK=gk is not None, USE_GK=gk is not None,
USE_INITIAL_STATE=initial_state 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, SAVE_NEW_VALUE=v_new is not None,
IS_VARLEN=cu_seqlens is not None, IS_VARLEN=cu_seqlens is not None,
num_warps=4, num_warps=4,
num_stages=2, num_stages=2,
) )
return h, v_new, final_state return h, v_new
python/sglang/srt/layers/attention/fla/kda.py
+9 -7
View File
@@ -1146,6 +1146,7 @@ def chunk_kda_fwd(
beta: torch.Tensor, beta: torch.Tensor,
scale: float, scale: float,
initial_state: torch.Tensor, initial_state: torch.Tensor,
initial_state_indices: torch.Tensor,
output_final_state: bool, output_final_state: bool,
cu_seqlens: torch.LongTensor | None = None, cu_seqlens: torch.LongTensor | None = None,
): ):
@@ -1172,12 +1173,13 @@ def chunk_kda_fwd(
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
) )
del A del A
h, v_new, final_state = chunk_gated_delta_rule_fwd_h( h, v_new = chunk_gated_delta_rule_fwd_h(
k=kg, k=kg,
w=w, w=w,
u=u, u=u,
gk=g, gk=g,
initial_state=initial_state, initial_state=initial_state,
initial_state_indices=initial_state_indices,
output_final_state=output_final_state, output_final_state=output_final_state,
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
) )
@@ -1194,7 +1196,7 @@ def chunk_kda_fwd(
chunk_size=chunk_size, chunk_size=chunk_size,
) )
del Aqk, v_new, h del Aqk, v_new, h
return o, final_state return o
def chunk_kda( def chunk_kda(
@@ -1205,7 +1207,7 @@ def chunk_kda(
beta: torch.Tensor, beta: torch.Tensor,
scale: float = None, scale: float = None,
initial_state: torch.Tensor = None, initial_state: torch.Tensor = None,
output_final_state: bool = False, initial_state_indices: torch.Tensor = None,
use_qk_l2norm_in_kernel: bool = False, use_qk_l2norm_in_kernel: bool = False,
cu_seqlens: torch.LongTensor | None = None, cu_seqlens: torch.LongTensor | None = None,
**kwargs, **kwargs,
@@ -1217,18 +1219,18 @@ def chunk_kda(
q = l2norm_fwd(q.contiguous()) q = l2norm_fwd(q.contiguous())
k = l2norm_fwd(k.contiguous()) k = l2norm_fwd(k.contiguous())
o, final_state = chunk_kda_fwd( o = chunk_kda_fwd(
q=q, q=q,
k=k, k=k,
v=v.contiguous(), v=v.contiguous(),
g=g.contiguous(), g=g.contiguous(),
beta=beta.contiguous(), beta=beta.contiguous(),
scale=scale, scale=scale,
initial_state=initial_state.contiguous(), initial_state=initial_state,
output_final_state=output_final_state, initial_state_indices=initial_state_indices,
cu_seqlens=cu_seqlens, cu_seqlens=cu_seqlens,
) )
return o, final_state return o
@triton.autotune( @triton.autotune(
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+15 -12
View File
@@ -814,22 +814,17 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
beta = beta.unsqueeze(0) beta = beta.unsqueeze(0)
g = g.unsqueeze(0) g = g.unsqueeze(0)
initial_state = ssm_states[cache_indices].contiguous() core_attn_out = chunk_kda(
(
core_attn_out,
last_recurrent_state,
) = chunk_kda(
q=q, q=q,
k=k, k=k,
v=v, v=v,
g=g, g=g,
beta=beta, beta=beta,
initial_state=initial_state, initial_state=ssm_states,
output_final_state=True, initial_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True, use_qk_l2norm_in_kernel=True,
cu_seqlens=query_start_loc, cu_seqlens=query_start_loc,
) )
ssm_states[cache_indices] = last_recurrent_state
return core_attn_out return core_attn_out
@@ -1064,7 +1059,12 @@ class GDNAttnBackend(MambaAttnBackendBase):
retrieve_parent_token=retrieve_parent_token, retrieve_parent_token=retrieve_parent_token,
) )
else: 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( core_attn_out, last_recurrent_state, h = chunk_gated_delta_rule(
q=query, q=query,
k=key, k=key,
@@ -1072,13 +1072,16 @@ class GDNAttnBackend(MambaAttnBackendBase):
g=g, g=g,
beta=beta, beta=beta,
initial_state=recurrent_state, initial_state=recurrent_state,
output_final_state=True,
cu_seqlens=query_start_loc, cu_seqlens=query_start_loc,
head_first=False, head_first=False,
use_qk_l2norm_in_kernel=True, use_qk_l2norm_in_kernel=True,
**recurrent_state_indices_args,
) )
last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False) if is_npu():
ssm_states[cache_indices] = last_recurrent_state last_recurrent_state = last_recurrent_state.to(
ssm_states.dtype, copy=False
)
ssm_states[cache_indices] = last_recurrent_state
self._track_mamba_state_extend( self._track_mamba_state_extend(
forward_batch, h, ssm_states, forward_metadata forward_batch, h, ssm_states, forward_metadata