diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index c148bc67e..d8fc0103f 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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, - 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( diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 2fe1dc749..43d8b582c 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/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: diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 9f6513654..0c5983d7d 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/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] - - 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]