From bf523883545917d225cd5839d0ac8b0d715c5515 Mon Sep 17 00:00:00 2001 From: Minglei Zhu Date: Tue, 17 Feb 2026 07:31:12 -0800 Subject: [PATCH] [PCG] support piecewise cuda graph for kimi-linear model (#18849) --- .../attention/hybrid_linear_attn_backend.py | 53 +++++---- .../srt/layers/radix_linear_attention.py | 2 +- python/sglang/srt/models/kimi_linear.py | 103 +++++++++++------- .../models/test_kimi_linear_models_pcg.py | 70 ++++++++++++ 4 files changed, 157 insertions(+), 71 deletions(-) create mode 100644 test/registered/models/test_kimi_linear_models_pcg.py 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 360d164e5..de3cebabd 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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( diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 571d292c0..89821e774 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -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: diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 2c2520f97..20b1ac3bb 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -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: diff --git a/test/registered/models/test_kimi_linear_models_pcg.py b/test/registered/models/test_kimi_linear_models_pcg.py new file mode 100644 index 000000000..6c0094794 --- /dev/null +++ b/test/registered/models/test_kimi_linear_models_pcg.py @@ -0,0 +1,70 @@ +""" +Kimi-Linear piecewise CUDA graph tests. +""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.few_shot_gsm8k import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci( + est_time=100, + suite="stage-b-test-large-2-gpu", +) + +KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" + +ACC_THRESHOLDS = { + KIMI_LINEAR_MODEL: {"gsm8k": 0.88}, +} + + +class TestKimiLinearPiecewiseCudaGraph(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.model = KIMI_LINEAR_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp", + "2", + "--trust-remote-code", + "--enable-piecewise-cuda-graph", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + +if __name__ == "__main__": + unittest.main()