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

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

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-01-24 21:27:13 +08:00
committed by GitHub
parent bf19d20d89
commit 0c8165ffbd
3 changed files with 94 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache: bool = True,
layer: RadixLinearAttention,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
q_proj_states = kwargs["q_proj_states"]
k_proj_states = kwargs["k_proj_states"]
v_proj_states = kwargs["v_proj_states"]
q_conv_weights = kwargs["q_conv_weights"]
k_conv_weights = kwargs["k_conv_weights"]
v_conv_weights = kwargs["v_conv_weights"]
assert isinstance(mixed_qkv, Tuple)
(q_proj_states, k_proj_states, v_proj_states) = mixed_qkv
(q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
(q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
q_conv_bias = kwargs["q_conv_bias"]
k_conv_bias = kwargs["k_conv_bias"]
v_conv_bias = kwargs["v_conv_bias"]
head_dim = layer.head_qk_dim
layer_id = layer.layer_id
beta = b
g = a
head_dim = kwargs["head_dim"]
layer_id = kwargs["layer_id"]
beta = kwargs["beta"]
g = kwargs["gate"]
A_log = kwargs["A_log"]
dt_bias = kwargs["dt_bias"]
A_log = layer.A_log
dt_bias = layer.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,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
save_kv_cache: bool = True,
**kwargs,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
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"]
k_proj_states = kwargs["k_proj_states"]
v_proj_states = kwargs["v_proj_states"]
q_conv_weights = kwargs["q_conv_weights"]
k_conv_weights = kwargs["k_conv_weights"]
v_conv_weights = kwargs["v_conv_weights"]
assert isinstance(mixed_qkv, Tuple)
(q_proj_states, k_proj_states, v_proj_states) = mixed_qkv
(q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
(q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
q_conv_bias = kwargs["q_conv_bias"]
k_conv_bias = kwargs["k_conv_bias"]
v_conv_bias = kwargs["v_conv_bias"]
head_dim = layer.head_qk_dim
layer_id = layer.layer_id
beta = b
g = a
head_dim = kwargs["head_dim"]
layer_id = kwargs["layer_id"]
beta = kwargs["beta"]
g = kwargs["gate"]
A_log = layer.A_log
dt_bias = layer.dt_bias
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,9 +1300,15 @@ class HybridLinearAttnBackend(AttentionBackend):
if forward_batch.forward_mode.is_idle():
if is_linear_attn:
return mixed_qkv.new_empty(
mixed_qkv.shape[0], layer.num_v_heads, layer.head_v_dim
)
# 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
)
return q.new_empty(q.shape[0], layer.tp_q_head_num * layer.v_head_dim)
elif forward_batch.forward_mode.is_decode():
return self.forward_decode(

8
python/sglang/srt/layers/radix_linear_attention.py
View File

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

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

@@ -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]
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)
)
mixed_qkv = (q_proj_states, k_proj_states, v_proj_states)
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()
forget_gate = forget_gate.unsqueeze(0)
beta = beta.unsqueeze(0)
forget_gate = forget_gate.unsqueeze(0)
kwargs = {
"q_proj_states": q_proj_states,
"k_proj_states": k_proj_states,
"v_proj_states": v_proj_states,
"q_conv_weights": q_conv_weights,
"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,
core_attn_out = self.linear_attn(
forward_batch,
mixed_qkv=mixed_qkv,
a=forget_gate,
b=beta,
)
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]