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

8
python/sglang/srt/configs/mamba_utils.py
View File

@@ -217,11 +217,13 @@ class KimiLinearStateShape:
conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1)
temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
conv_state_shape = conv_state_shape[1], conv_state_shape[0]
conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0]
conv_state_shape = (
conv_state_shape[1],
conv_state_shape[0] + conv_state_k_shape[0] * 2,
)
return KimiLinearStateShape(
conv=[conv_state_shape, conv_state_k_shape, conv_state_k_shape],
conv=[conv_state_shape],
temporal=temporal_state_shape,
num_heads=num_heads,
head_dim=head_dim,

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,

46
python/sglang/srt/models/kimi_linear.py
View File

@@ -22,6 +22,7 @@ from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
ColumnParallelBatchedLinear,
ColumnParallelLinear,
MergedColumnParallelLinear,
MergedColumnParallelRepeatedLinear,
QKVParallelLinear,
ReplicatedLinear,
@@ -283,34 +284,18 @@ class KimiDeltaAttention(nn.Module):
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
self.q_conv1d = ColumnParallelLinear(
self.qkv_conv1d = MergedColumnParallelLinear(
input_size=self.conv_size,
output_size=projection_size,
output_sizes=[projection_size, projection_size, projection_size],
bias=False,
params_dtype=torch.float32,
prefix=f"{prefix}.q_conv1d",
)
self.k_conv1d = ColumnParallelLinear(
input_size=self.conv_size,
output_size=projection_size,
bias=False,
params_dtype=torch.float32,
prefix=f"{prefix}.k_conv1d",
)
self.v_conv1d = ColumnParallelLinear(
input_size=self.conv_size,
output_size=projection_size,
bias=False,
params_dtype=torch.float32,
prefix=f"{prefix}.v_conv1d",
prefix=f"{prefix}.qkv_conv1d",
)
# unsqueeze to fit conv1d weights shape into the linear weights shape.
# Can't do this in `weight_loader` since it already exists in
# `ColumnParallelLinear` and `set_weight_attrs`
# doesn't allow to override it
self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)
self.qkv_conv1d.weight.data = self.qkv_conv1d.weight.data.unsqueeze(1)
self.A_log = nn.Parameter(
torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
@@ -328,18 +313,8 @@ class KimiDeltaAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
self.q_conv_weights = self.q_conv1d.weight.view(
self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
)
self.k_conv_weights = self.k_conv1d.weight.view(
self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2)
)
self.v_conv_weights = self.v_conv1d.weight.view(
self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
)
conv_weights = (self.q_conv_weights, self.k_conv_weights, self.v_conv_weights)
bias = (self.q_conv1d.bias, self.k_conv1d.bias, self.v_conv1d.bias)
conv_weights = self.qkv_conv1d.weight.squeeze(1)
bias = self.qkv_conv1d.bias
self.attn = RadixLinearAttention(
layer_id=self.layer_idx,
@@ -409,12 +384,11 @@ class KimiDeltaAttention(nn.Module):
)
# fused_kda_gate is fused to KimiLinearAttentionBackend with decode
beta = beta.float()
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.sigmoid()
beta = beta.float().sigmoid()
forget_gate = forget_gate.unsqueeze(0)
beta = beta.unsqueeze(0)
@@ -703,6 +677,10 @@ class KimiLinearForCausalLM(nn.Module):
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
# qkv conv fuse
(".qkv_conv1d", ".q_conv1d", 0),
(".qkv_conv1d", ".k_conv1d", 1),
(".qkv_conv1d", ".v_conv1d", 2),
]
if self.config.is_moe:
# Params for weights, fp8 weight scales, fp8 activation scales