chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

KDA: fuse qkv conv and support stride for fused_sigmoid_gating_delta_rule_update_kernel (#19506)

Browse Source
This commit is contained in:
strgrb
2026-03-04 22:45:53 +08:00
committed by GitHub
parent 6910c1b281
commit 738ebfd330
4 changed files with 63 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py
View File

@@ -4,8 +4,6 @@ import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.fla.utils import input_guard
@triton.jit(do_not_specialize=["T"])
def fused_sigmoid_gating_delta_rule_update_kernel(
@@ -24,6 +22,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
cu_seqlens,
scale,
T,
stride_q,
stride_k,
stride_v,
stride_b,
B: tl.constexpr,
H: tl.constexpr,
HV: tl.constexpr,
@@ -57,10 +59,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
o_k = i_k * BK + tl.arange(0, BK)
o_v = i_v * BV + tl.arange(0, BV)
p_q = q + (bos * H + i_h) * K + o_k
p_k = k + (bos * H + i_h) * K + o_k
p_v = v + (bos * HV + i_hv) * V + o_v
p_b = b + bos * HV + i_hv
p_q = q + bos * stride_q + i_h * K + o_k
p_k = k + bos * stride_k + i_h * K + o_k
p_v = v + bos * stride_v + i_hv * V + o_v
p_b = b + bos * stride_b + i_hv
p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
# Gating computation pointers
@@ -164,7 +166,6 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h)
@input_guard
def fused_sigmoid_gating_delta_rule_update(
A_log: torch.Tensor,
a: torch.Tensor,
@@ -188,6 +189,10 @@ def fused_sigmoid_gating_delta_rule_update(
and the recurrent delta rule update for better performance.
"""
B, T, H, K, V = *k.shape, v.shape[-1]
stride_q = q.stride()[1]
stride_k = k.stride()[1]
stride_v = v.stride()[1]
stride_b = b.stride()[-2]
HV = v.shape[2]
N = B if cu_seqlens is None else len(cu_seqlens) - 1
BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
@@ -220,6 +225,10 @@ def fused_sigmoid_gating_delta_rule_update(
cu_seqlens=cu_seqlens,
scale=scale,
T=T,
stride_q=stride_q,
stride_k=stride_k,
stride_v=stride_v,
stride_b=stride_b,
B=B,
H=H,
HV=HV,

95
python/sglang/srt/layers/attention/linear/kda_backend.py
View File

@@ -136,49 +136,21 @@ class KDAAttnBackend(MambaAttnBackendBase):
b: torch.Tensor,
**kwargs,
):
q_proj_states, k_proj_states, v_proj_states = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights
q_conv_bias, k_conv_bias, v_conv_bias = layer.bias
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
q_conv_state, k_conv_state, v_conv_state = layer_cache.conv
conv_states = layer_cache.conv[0]
ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
q_conv_state = q_conv_state.transpose(-1, -2)
k_conv_state = k_conv_state.transpose(-1, -2)
v_conv_state = v_conv_state.transpose(-1, -2)
q = causal_conv1d_update(
q_proj_states,
q_conv_state,
q_conv_weights,
q_conv_bias,
qkv = causal_conv1d_update(
mixed_qkv,
conv_states.transpose(-1, -2),
layer.conv_weights,
layer.bias,
activation="silu",
conv_state_indices=cache_indices,
)
k = causal_conv1d_update(
k_proj_states,
k_conv_state,
k_conv_weights,
k_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
v = causal_conv1d_update(
v_proj_states,
v_conv_state,
v_conv_weights,
v_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
@@ -205,62 +177,55 @@ class KDAAttnBackend(MambaAttnBackendBase):
b: torch.Tensor,
**kwargs,
):
q_proj_states, k_proj_states, v_proj_states = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights
q_conv_bias, k_conv_bias, v_conv_bias = layer.bias
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_state_q, conv_state_k, conv_state_v = mamba_cache_params.conv
# deal with strides
conv_state_q = conv_state_q.transpose(-1, -2)
conv_state_k = conv_state_k.transpose(-1, -2)
conv_state_v = conv_state_v.transpose(-1, -2)
conv_states = mamba_cache_params.conv[0].transpose(-1, -2)
ssm_states = mamba_cache_params.temporal
has_initial_state = forward_batch.extend_prefix_lens > 0
q_proj_states = q_proj_states.transpose(0, 1)
k_proj_states = k_proj_states.transpose(0, 1)
v_proj_states = v_proj_states.transpose(0, 1)
splits = [layer.q_dim, layer.k_dim, layer.v_dim]
q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0)
q_conv_weight, k_conv_weight, v_conv_weight = layer.conv_weights.split(
splits, dim=0
)
q_conv_state, k_conv_state, v_conv_state = conv_states.split(splits, dim=-2)
if layer.bias is not None:
q_bias, k_bias, v_bias = layer.bias.split(splits, dim=0)
else:
q_bias, k_bias, v_bias = None, None, None
q = causal_conv1d_fn(
q_proj_states,
q_conv_weights,
q_conv_bias,
q,
q_conv_weight,
q_bias,
activation="silu",
conv_states=conv_state_q,
conv_states=q_conv_state,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
k = causal_conv1d_fn(
k_proj_states,
k_conv_weights,
k_conv_bias,
k,
k_conv_weight,
k_bias,
activation="silu",
conv_states=conv_state_k,
conv_states=k_conv_state,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
v = causal_conv1d_fn(
v_proj_states,
v_conv_weights,
v_conv_bias,
v,
v_conv_weight,
v_bias,
activation="silu",
conv_states=conv_state_v,
conv_states=v_conv_state,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,