[Kimi-Linear] Refactor Kimi-Linear to support RadixLinearAttention (#17506)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Tuple, Union
 
 import torch
 import triton
@@ -624,32 +624,24 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
 
     def forward_decode(
         self,
-        q: torch.Tensor,
+        layer: RadixLinearAttention,
-        k: torch.Tensor,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
-        v: torch.Tensor,
+        a: torch.Tensor,
-        layer: RadixAttention,
+        b: torch.Tensor,
-        forward_batch: ForwardBatch,
-        save_kv_cache: bool = True,
         **kwargs,
     ):
-        q_proj_states = kwargs["q_proj_states"]
+        assert isinstance(mixed_qkv, Tuple)
-        k_proj_states = kwargs["k_proj_states"]
+        (q_proj_states, k_proj_states, v_proj_states) = mixed_qkv
-        v_proj_states = kwargs["v_proj_states"]
+        (q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
-        q_conv_weights = kwargs["q_conv_weights"]
+        (q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
-        k_conv_weights = kwargs["k_conv_weights"]
-        v_conv_weights = kwargs["v_conv_weights"]
 
-        q_conv_bias = kwargs["q_conv_bias"]
+        head_dim = layer.head_qk_dim
-        k_conv_bias = kwargs["k_conv_bias"]
+        layer_id = layer.layer_id
-        v_conv_bias = kwargs["v_conv_bias"]
+        beta = b
+        g = a
 
-        head_dim = kwargs["head_dim"]
+        A_log = layer.A_log
-        layer_id = kwargs["layer_id"]
+        dt_bias = layer.dt_bias
-        beta = kwargs["beta"]
-        g = kwargs["gate"]
-
-        A_log = kwargs["A_log"]
-        dt_bias = kwargs["dt_bias"]
 
         layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id)
         q_conv_state, k_conv_state, v_conv_state = layer_cache.conv
@@ -711,33 +703,29 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
 
     def forward_extend(
         self,
-        q: torch.Tensor,
+        layer: RadixLinearAttention,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        layer: RadixAttention,
         forward_batch: ForwardBatch,
-        save_kv_cache: bool = True,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
-        **kwargs,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,  # Unused, for compatibility with HybridLinearAttnBackend
     ):
         from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
             causal_conv1d_fn,
         )
 
-        q_proj_states = kwargs["q_proj_states"]
+        assert isinstance(mixed_qkv, Tuple)
-        k_proj_states = kwargs["k_proj_states"]
+        (q_proj_states, k_proj_states, v_proj_states) = mixed_qkv
-        v_proj_states = kwargs["v_proj_states"]
+        (q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
-        q_conv_weights = kwargs["q_conv_weights"]
+        (q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
-        k_conv_weights = kwargs["k_conv_weights"]
-        v_conv_weights = kwargs["v_conv_weights"]
 
-        q_conv_bias = kwargs["q_conv_bias"]
+        head_dim = layer.head_qk_dim
-        k_conv_bias = kwargs["k_conv_bias"]
+        layer_id = layer.layer_id
-        v_conv_bias = kwargs["v_conv_bias"]
+        beta = b
+        g = a
 
-        head_dim = kwargs["head_dim"]
+        A_log = layer.A_log
-        layer_id = kwargs["layer_id"]
+        dt_bias = layer.dt_bias
-        beta = kwargs["beta"]
-        g = kwargs["gate"]
 
         query_start_loc = self.forward_metadata.query_start_loc
         cache_indices = self.forward_metadata.mamba_cache_indices
@@ -836,7 +824,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
         self,
         layer: RadixLinearAttention,
         forward_batch: ForwardBatch,
-        mixed_qkv: torch.Tensor,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
         a: torch.Tensor,
         b: torch.Tensor,
         **kwargs,  # Unused, for compatibility with HybridLinearAttnBackend
@@ -856,6 +844,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
         query_start_loc = self.forward_metadata.query_start_loc
         cache_indices = self.forward_metadata.mamba_cache_indices
 
+        assert isinstance(mixed_qkv, torch.Tensor)
         mixed_qkv = causal_conv1d_update(
             mixed_qkv,
             conv_states,
@@ -907,11 +896,12 @@ class GDNAttnBackend(MambaAttnBackendBase):
         self,
         layer: RadixLinearAttention,
         forward_batch: ForwardBatch,
-        mixed_qkv: torch.Tensor,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
         a: torch.Tensor,
         b: torch.Tensor,
         **kwargs,  # Unused, for compatibility with HybridLinearAttnBackend
     ):
+        assert isinstance(mixed_qkv, torch.Tensor)
         seq_len = mixed_qkv.shape[0]
 
         conv_weights = layer.conv_weights
@@ -1234,7 +1224,7 @@ class HybridLinearAttnBackend(AttentionBackend):
         q: Optional[torch.Tensor] = None,  # For full attention
         k: Optional[torch.Tensor] = None,  # For full attention
         v: Optional[torch.Tensor] = None,  # For full attention
-        mixed_qkv: Optional[torch.Tensor] = None,  # For GDN linear attention
+        mixed_qkv: Optional[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
         a: Optional[torch.Tensor] = None,  # For GDN linear attention
         b: Optional[torch.Tensor] = None,  # For GDN linear attention
         **kwargs,
@@ -1266,7 +1256,7 @@ class HybridLinearAttnBackend(AttentionBackend):
         q: Optional[torch.Tensor] = None,  # For full attention
         k: Optional[torch.Tensor] = None,  # For full attention
         v: Optional[torch.Tensor] = None,  # For full attention
-        mixed_qkv: Optional[torch.Tensor] = None,  # For GDN linear attention
+        mixed_qkv: Optional[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
         a: Optional[torch.Tensor] = None,  # For GDN linear attention
         b: Optional[torch.Tensor] = None,  # For GDN linear attention
         **kwargs,
@@ -1298,7 +1288,9 @@ class HybridLinearAttnBackend(AttentionBackend):
         layer: RadixAttention = None,
         forward_batch: ForwardBatch = None,
         save_kv_cache: bool = True,
-        mixed_qkv: Optional[torch.Tensor] = None,  # For GDN linear attention
+        mixed_qkv: Optional[
+            Union[torch.Tensor, Tuple[torch.Tensor, ...]]
+        ] = None,  # For GDN linear attention
         a: Optional[torch.Tensor] = None,  # For GDN linear attention
         b: Optional[torch.Tensor] = None,  # For GDN linear attention
         **kwargs,
@@ -1308,6 +1300,12 @@ class HybridLinearAttnBackend(AttentionBackend):
 
         if forward_batch.forward_mode.is_idle():
             if is_linear_attn:
+                # KDA:
+                if isinstance(mixed_qkv, tuple):
+                    return mixed_qkv[0].new_empty(
+                        mixed_qkv[0].shape[0], layer.num_v_heads, layer.head_v_dim
+                    )
+                else:  # GDN:
                     return mixed_qkv.new_empty(
                         mixed_qkv.shape[0], layer.num_v_heads, layer.head_v_dim
                     )

python/sglang/srt/layers/radix_linear_attention.py

@@ -14,7 +14,7 @@
 """Radix linear attention."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -36,7 +36,8 @@ class RadixLinearAttention(nn.Module):
         head_qk_dim: int,
         head_v_dim: int,
         attention_tp_size: int = 1,
-        conv_weights: Optional[torch.Tensor] = None,
+        # GDN KDA Shared Weights
+        conv_weights: Optional[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
         bias: Optional[torch.Tensor] = None,
         activation: str = "silu",
         A_log: Optional[torch.Tensor] = None,
@@ -64,13 +65,14 @@ class RadixLinearAttention(nn.Module):
         self.conv_weights = conv_weights
         self.bias = bias
         self.activation = activation
+
         self.A_log = A_log
         self.dt_bias = dt_bias
 
     def forward(
         self,
         forward_batch: ForwardBatch,
-        mixed_qkv: torch.Tensor,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
         a: torch.Tensor,
         b: torch.Tensor,
     ) -> torch.Tensor:

python/sglang/srt/models/kimi_linear.py

@@ -16,6 +16,7 @@ from sglang.srt.distributed import (
 )
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.attention.fla.kda import FusedRMSNormGated, fused_kda_gate
+from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -27,6 +28,7 @@ from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.layers.utils import PPMissingLayer
 from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -171,10 +173,15 @@ class KimiDeltaAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.tp_size = get_tensor_model_parallel_world_size()
+        self.attn_tp_size = get_attention_tp_size()
         self.hidden_size = hidden_size
         self.config = config
         self.head_dim = config.linear_attn_config["head_dim"]
         self.num_heads = config.linear_attn_config["num_heads"]
+        self.num_k_heads = config.linear_attn_config["num_heads"]
+        self.num_v_heads = config.linear_attn_config["num_heads"]
+        self.head_k_dim = config.linear_attn_config["head_dim"]
+        self.head_v_dim = config.v_head_dim
         self.layer_idx = layer_idx
         self.prefix = prefix
         assert self.num_heads % self.tp_size == 0
@@ -293,6 +300,32 @@ 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)
+
+        self.linear_attn = RadixLinearAttention(
+            layer_id=self.layer_idx,
+            num_qk_heads=self.num_k_heads // self.attn_tp_size,
+            num_v_heads=self.num_v_heads // self.attn_tp_size,
+            head_qk_dim=self.head_k_dim,
+            head_v_dim=self.head_v_dim,
+            attention_tp_size=self.attn_tp_size,
+            conv_weights=conv_weights,
+            bias=bias,
+            A_log=self.A_log,
+            dt_bias=self.dt_bias,
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -303,54 +336,26 @@ class KimiDeltaAttention(nn.Module):
         q_proj_states = self.q_proj(hidden_states)[0]
         k_proj_states = self.k_proj(hidden_states)[0]
         v_proj_states = self.v_proj(hidden_states)[0]
-
+        mixed_qkv = (q_proj_states, k_proj_states, v_proj_states)
-        q_conv_weights = self.q_conv1d.weight.view(
-            self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
-        )
-        k_conv_weights = self.k_conv1d.weight.view(
-            self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2)
-        )
-        v_conv_weights = self.v_conv1d.weight.view(
-            self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
-        )
 
         forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
 
         beta = self.b_proj(hidden_states)[0].float()
         # fused_kda_gate is fused to KimiLinearAttentionBackend with decode
+        beta = self.b_proj(hidden_states)[0].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.unsqueeze(0)
             forget_gate = forget_gate.unsqueeze(0)
+        beta = beta.unsqueeze(0)
 
-        kwargs = {
+        core_attn_out = self.linear_attn(
-            "q_proj_states": q_proj_states,
+            forward_batch,
-            "k_proj_states": k_proj_states,
+            mixed_qkv=mixed_qkv,
-            "v_proj_states": v_proj_states,
+            a=forget_gate,
-            "q_conv_weights": q_conv_weights,
+            b=beta,
-            "k_conv_weights": k_conv_weights,
-            "v_conv_weights": v_conv_weights,
-            "q_conv_bias": self.q_conv1d.bias,
-            "k_conv_bias": self.k_conv1d.bias,
-            "v_conv_bias": self.v_conv1d.bias,
-            "head_dim": self.head_dim,
-            "layer_id": self.layer_idx,
-            "beta": beta,
-            "gate": forget_gate,
-            "A_log": self.A_log,
-            "dt_bias": self.dt_bias,
-        }
-
-        core_attn_out = forward_batch.attn_backend.forward(
-            q=None,
-            k=None,
-            v=None,
-            layer=None,
-            forward_batch=forward_batch,
-            **kwargs,
         )
 
         g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]