chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

33
python/sglang/srt/layers/attention/fla/chunk.py
View File

@@ -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

31
python/sglang/srt/layers/attention/fla/chunk_delta_h.py
View File

@@ -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

16
python/sglang/srt/layers/attention/fla/kda.py
View File

@@ -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(

27
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View File

@@ -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