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

python/sglang/srt/configs/mamba_utils.py

@@ -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_shape = (
-        conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0]
+            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,

python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -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_q = q + bos * stride_q + i_h * K + o_k
-    p_k = k + (bos * H + i_h) * K + o_k
+    p_k = k + bos * stride_k + i_h * K + o_k
-    p_v = v + (bos * HV + i_hv) * V + o_v
+    p_v = v + bos * stride_v + i_hv * V + o_v
-    p_b = b + bos * HV + i_hv
+    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,

python/sglang/srt/layers/attention/linear/kda_backend.py

@@ -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)
+        qkv = causal_conv1d_update(
-        k_conv_state = k_conv_state.transpose(-1, -2)
+            mixed_qkv,
-        v_conv_state = v_conv_state.transpose(-1, -2)
+            conv_states.transpose(-1, -2),
-
+            layer.conv_weights,
-        q = causal_conv1d_update(
+            layer.bias,
-            q_proj_states,
-            q_conv_state,
-            q_conv_weights,
-            q_conv_bias,
             activation="silu",
             conv_state_indices=cache_indices,
         )
-        k = causal_conv1d_update(
+        q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
-            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 = 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
+        conv_states = mamba_cache_params.conv[0].transpose(-1, -2)
-        # 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)
 
         ssm_states = mamba_cache_params.temporal
 
         has_initial_state = forward_batch.extend_prefix_lens > 0
 
-        q_proj_states = q_proj_states.transpose(0, 1)
+        splits = [layer.q_dim, layer.k_dim, layer.v_dim]
-        k_proj_states = k_proj_states.transpose(0, 1)
+        q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0)
-        v_proj_states = v_proj_states.transpose(0, 1)
+        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,
-            q_conv_weights,
+            q_conv_weight,
-            q_conv_bias,
+            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,
-            k_conv_weights,
+            k_conv_weight,
-            k_conv_bias,
+            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,
-            v_conv_weights,
+            v_conv_weight,
-            v_conv_bias,
+            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,

python/sglang/srt/models/kimi_linear.py

@@ -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",
+            prefix=f"{prefix}.qkv_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",
         )
         # 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.qkv_conv1d.weight.data = self.qkv_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.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(
+        conv_weights = self.qkv_conv1d.weight.squeeze(1)
-            self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
+        bias = self.qkv_conv1d.bias
-        )
-        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)
 
         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