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

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 = torch.split(
-        q_proj_states, k_proj_states, v_proj_states = mixed_qkv
+            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 = torch.split(
-        q_proj_states, k_proj_states, v_proj_states = mixed_qkv
+            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[
+        mixed_qkv: Optional[torch.Tensor] = None,  # For linear attention
-            Union[torch.Tensor, Tuple[torch.Tensor, ...]]
+        a: Optional[torch.Tensor] = None,  # For linear attention
-        ] = None,  # For GDN linear attention
+        b: 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,
     ):
         layer_id = layer.layer_id if layer else kwargs["layer_id"]
@@ -1658,12 +1661,6 @@ 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

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

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.qkvb_sizes,  # Column parallel
-                self.fg_sizes,
+                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] + [
+            self.split_sizes = [
-                2 * self.head_dim
+                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",
+                tp_rank=attn_tp_rank,
-            )
+                tp_size=self.attn_tp_size,
-            self.k_proj = ColumnParallelLinear(
+                v_head_size=self.head_v_dim,
-                self.hidden_size,
+                prefix=f"{prefix}.qkv_proj",
-                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",
             )
 
             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]
+        qkv, _ = self.qkv_proj(hidden_states)
-        k_proj_states = self.k_proj(hidden_states)[0]
+
-        v_proj_states = self.v_proj(hidden_states)[0]
+        # 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)
+        # Single fused projection for all: qkv + beta + f_a + g_a
-        q_proj_states, k_proj_states, v_proj_states, beta, fg_a_states = torch.split(
+        fused_states = self.fused_qkvbfg_a_proj(hidden_states)
-            fused_states, self.split_sizes, dim=-1
+
+        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 path
-            (".fused_qkvbfg_proj", ".k_proj", 1),
+            (".fused_qkvbfg_a_proj", ".q_proj", 0),
-            (".fused_qkvbfg_proj", ".v_proj", 2),
+            (".fused_qkvbfg_a_proj", ".k_proj", 1),
-            (".fused_qkvbfg_proj", ".b_proj", 3),
+            (".fused_qkvbfg_a_proj", ".v_proj", 2),
-            (".fused_qkvbfg_proj", ".f_a_proj", 4),
+            (".fused_qkvbfg_a_proj", ".b_proj", 3),
-            (".fused_qkvbfg_proj", ".g_a_proj", 5),
+            (".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:

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()