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

[PCG] support piecewise cuda graph for kimi-linear model (#18849)

Browse Source
This commit is contained in:
Minglei Zhu
2026-02-17 07:31:12 -08:00
committed by GitHub
parent bfe34c90ff
commit bf52388354
4 changed files with 157 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,6 +1,6 @@
import logging
import math
from typing import Optional, Tuple, Union
from typing import Optional, Union
import torch
import triton
@@ -664,13 +664,17 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
def forward_decode(
self,
layer: RadixLinearAttention,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
assert isinstance(mixed_qkv, Tuple)
q_proj_states, k_proj_states, v_proj_states = mixed_qkv
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
@@ -734,7 +738,7 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend
@@ -743,8 +747,12 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
causal_conv1d_fn,
)
assert isinstance(mixed_qkv, Tuple)
q_proj_states, k_proj_states, v_proj_states = mixed_qkv
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
@@ -852,7 +860,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend
@@ -862,8 +870,6 @@ class GDNAttnBackend(MambaAttnBackendBase):
ssm_states = layer_cache.temporal
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,
@@ -910,12 +916,11 @@ class GDNAttnBackend(MambaAttnBackendBase):
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
mixed_qkv: 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]
is_target_verify = forward_batch.forward_mode.is_target_verify()
@@ -1582,7 +1587,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[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
mixed_qkv: Optional[torch.Tensor] = None, # For linear attention
a: Optional[torch.Tensor] = None, # For GDN linear attention
b: Optional[torch.Tensor] = None, # For GDN linear attention
**kwargs,
@@ -1614,7 +1619,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[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
mixed_qkv: Optional[torch.Tensor] = None, # For linear attention
a: Optional[torch.Tensor] = None, # For GDN linear attention
b: Optional[torch.Tensor] = None, # For GDN linear attention
**kwargs,
@@ -1646,11 +1651,9 @@ class HybridLinearAttnBackend(AttentionBackend):
layer: RadixAttention = None,
forward_batch: ForwardBatch = None,
save_kv_cache: bool = True,
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
mixed_qkv: Optional[torch.Tensor] = None, # For linear attention
a: Optional[torch.Tensor] = None, # For linear attention
b: Optional[torch.Tensor] = None, # For linear attention
**kwargs,
):
layer_id = layer.layer_id if layer else kwargs["layer_id"]
@@ -1658,15 +1661,9 @@ 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
)
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(

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

@@ -71,7 +71,7 @@ class RadixLinearAttention(nn.Module):
def forward(
self,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
) -> torch.Tensor:

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

@@ -16,12 +16,13 @@ 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.dp_attention import get_attention_tp_rank, get_attention_tp_size
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
ColumnParallelBatchedLinear,
ColumnParallelLinear,
MergedColumnParallelRepeatedLinear,
QKVParallelLinear,
ReplicatedLinear,
RowParallelLinear,
)
@@ -194,7 +195,9 @@ class KimiDeltaAttention(nn.Module):
# TODO: support fusion with quant
self.do_fuse_qkvbfg = quant_config is None
if self.do_fuse_qkvbfg:
# Fuse: q, k, v, beta (column parallel) + f_a, g_a (replicated)
self.qkvb_sizes = [
projection_size,
projection_size,
@@ -202,40 +205,36 @@ class KimiDeltaAttention(nn.Module):
self.num_heads,
]
self.fg_sizes = [self.head_dim, self.head_dim]
self.fused_qkvbfg_proj = MergedColumnParallelRepeatedLinear(
self.fused_qkvbfg_a_proj = MergedColumnParallelRepeatedLinear(
self.hidden_size,
self.qkvb_sizes,
self.fg_sizes,
self.qkvb_sizes, # Column parallel
self.fg_sizes, # Replicated: f_a, g_a
quant_config=quant_config,
prefix=f"{prefix}.fused_qkvbfg_proj",
prefix=f"{prefix}.fused_qkvbfg_a_proj",
)
self.split_sizes = [x // self.tp_size for x in self.qkvb_sizes] + [
2 * self.head_dim
self.split_sizes = [
3 * projection_size // self.tp_size, # qkv
self.num_heads // self.tp_size, # beta
2 * self.head_dim, # f_a, g_a
]
self.fused_fg_b_proj = ColumnParallelBatchedLinear(
2, self.head_dim, projection_size, dtype=config.dtype
)
else:
self.q_proj = ColumnParallelLinear(
# Unfused path: separate QKVParallelLinear
attn_tp_rank = get_attention_tp_rank()
self.qkv_proj = QKVParallelLinear(
self.hidden_size,
projection_size,
self.head_dim,
self.num_heads,
self.num_k_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.q_proj",
)
self.k_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.k_proj",
)
self.v_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.v_proj",
tp_rank=attn_tp_rank,
tp_size=self.attn_tp_size,
v_head_size=self.head_v_dim,
prefix=f"{prefix}.qkv_proj",
)
self.f_a_proj = ReplicatedLinear(
@@ -341,7 +340,7 @@ class KimiDeltaAttention(nn.Module):
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(
self.attn = RadixLinearAttention(
layer_id=self.layer_idx,
num_q_heads=self.num_k_heads // self.attn_tp_size,
num_k_heads=self.num_k_heads // self.attn_tp_size,
@@ -356,30 +355,37 @@ class KimiDeltaAttention(nn.Module):
)
def forward_qkvbfg(self, hidden_states: torch.Tensor):
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]
qkv, _ = self.qkv_proj(hidden_states)
# Compute beta, forget_gate, and g_proj_states
beta = self.b_proj(hidden_states)[0]
forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
return (
(q_proj_states, k_proj_states, v_proj_states),
qkv,
beta,
forget_gate,
g_proj_states,
)
def forward_qkvbfg_fused(self, hidden_states: torch.Tensor):
fused_states = self.fused_qkvbfg_proj(hidden_states)
q_proj_states, k_proj_states, v_proj_states, beta, fg_a_states = torch.split(
fused_states, self.split_sizes, dim=-1
# Single fused projection for all: qkv + beta + f_a + g_a
fused_states = self.fused_qkvbfg_a_proj(hidden_states)
qkv, beta, fg_a_states = torch.split(
fused_states,
self.split_sizes,
dim=-1,
)
# use batch matmul to calculate forget_gate and g_proj_states
forget_gate, g_proj_states = self.fused_fg_b_proj(
fg_a_states.view(-1, 2, self.head_dim).transpose(0, 1)
)
return (
(q_proj_states, k_proj_states, v_proj_states),
qkv,
beta,
forget_gate,
g_proj_states,
@@ -411,7 +417,7 @@ class KimiDeltaAttention(nn.Module):
forget_gate = forget_gate.unsqueeze(0)
beta = beta.unsqueeze(0)
core_attn_out = self.linear_attn(
core_attn_out = self.attn(
forward_batch,
mixed_qkv=mixed_qkv,
a=forget_gate,
@@ -655,6 +661,7 @@ class KimiLinearForCausalLM(nn.Module):
logit_scale = getattr(self.config, "logit_scale", 1.0)
self.logits_processor = LogitsProcessor(config=config, logit_scale=logit_scale)
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
@@ -682,16 +689,20 @@ class KimiLinearForCausalLM(nn.Module):
# (param_name, shard_name, shard_id)
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
(".fused_qkvbfg_proj", ".q_proj", 0),
(".fused_qkvbfg_proj", ".k_proj", 1),
(".fused_qkvbfg_proj", ".v_proj", 2),
(".fused_qkvbfg_proj", ".b_proj", 3),
(".fused_qkvbfg_proj", ".f_a_proj", 4),
(".fused_qkvbfg_proj", ".g_a_proj", 5),
# Fused path
(".fused_qkvbfg_a_proj", ".q_proj", 0),
(".fused_qkvbfg_a_proj", ".k_proj", 1),
(".fused_qkvbfg_a_proj", ".v_proj", 2),
(".fused_qkvbfg_a_proj", ".b_proj", 3),
(".fused_qkvbfg_a_proj", ".f_a_proj", 4),
(".fused_qkvbfg_a_proj", ".g_a_proj", 5),
(".fused_fg_b_proj", ".f_b_proj", 0),
(".fused_fg_b_proj", ".g_b_proj", 1),
# Unfused path: separate qkv_proj (when do_fuse_qkvbfg=False)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
]
fuse_qkvbfg_keys = {x[1] for x in stacked_params_mapping[2:]}
if self.config.is_moe:
# Params for weights, fp8 weight scales, fp8 activation scales
# (param_name, weight_name, expert_id, shard_id)
@@ -726,11 +737,19 @@ class KimiLinearForCausalLM(nn.Module):
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
if ("mlp.experts." in name) and name not in params_dict:
continue
if weight_name in fuse_qkvbfg_keys:
# Check if this mapping targets a fused projection (only apply fusion check to fused params)
if param_name in {".fused_qkvbfg_a_proj", ".fused_fg_b_proj"}:
layer_id = int(name.split(".")[2])
if not self.config.is_kda_layer(layer_id):
continue
layer = self.model.layers[layer_id].self_attn
# Only load to fused projection if fusion is enabled
if not getattr(layer, "do_fuse_qkvbfg", False):
continue
if weight_name in {".q_proj", ".k_proj", ".v_proj"}:
layer_id = int(name.split(".")[2])
if not self.config.is_kda_layer(layer_id):
continue
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict: