[CP] Add CP-v2 support for Kimi-Linear

python/sglang/srt/configs/kimi_linear.py

@@ -7,6 +7,19 @@
 from sglang.srt.runtime_context import get_parallel
 
 
+def _get_kda_head_shard_info() -> tuple[int, int, bool]:
+    """Return the rank and size that own KDA heads.
+
+    CP ranks own KDA heads when attention CP is configured. Without CP, KDA
+    keeps its existing global-TP ownership.
+    """
+
+    parallel = get_parallel()
+    if parallel.attn_cp_size > 1:
+        return parallel.attn_cp_rank, parallel.attn_cp_size, True
+    return parallel.tp_rank, parallel.tp_size, False
+
+
 class KimiLinearConfig(PretrainedConfig):
     model_type = "kimi_linear"
     keys_to_ignore_at_inference = ["past_key_values"]
@@ -152,9 +165,12 @@ def full_attention_layer_ids(self):
 
     @property
     def mamba2_cache_params(self) -> KimiLinearCacheParams:
-
+        parallel = get_parallel()
+        kda_head_shard_size = (
+            parallel.attn_cp_size if parallel.attn_cp_size > 1 else parallel.tp_size
+        )
         shape = KimiLinearStateShape.create(
-            tp_world_size=get_parallel().attn_tp_size,
+            tp_world_size=kda_head_shard_size,
             num_heads=self.linear_attn_config["num_heads"],
             head_dim=self.linear_attn_config["head_dim"],
             conv_kernel_size=self.linear_attn_config["short_conv_kernel_size"],

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -23,6 +23,8 @@
 from sglang.srt.layers.attention.flashinfer_backend import (
     create_flashinfer_kv_indices_triton,
 )
+from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
+from sglang.srt.layers.cp.utils import is_cp_v2_active
 from sglang.srt.layers.dcp import (
     DecodeContextParallelMetadata,
     update_local_kv_lens_for_dcp,
@@ -77,6 +79,13 @@ class PrefillMetadata:
     use_ragged: bool
 
 
+@dataclass
+class CPPrefillMetadata:
+    wrappers: tuple[BatchMLAPagedAttentionWrapper, BatchMLAPagedAttentionWrapper]
+    kv_indptrs: tuple[torch.Tensor, torch.Tensor]
+    kv_indices: tuple[torch.Tensor, torch.Tensor]
+
+
 # Reuse this workspace buffer across all flashinfer wrappers
 
 
@@ -299,6 +308,7 @@ def __init__(
 
         # Other metadata
         self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None
+        self.cp_prefill_metadata: Optional[CPPrefillMetadata] = None
         self.decode_cuda_graph_metadata = {}
         self.prefill_cuda_graph_metadata = {}  # For verify
 
@@ -377,6 +387,12 @@ def init_forward_metadata_out_graph(
             )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
+        self.cp_prefill_metadata = None
+        if is_cp_v2_active(forward_batch):
+            # CP wrappers are planned lazily after the eager runner builds the
+            # per-rank zigzag metadata. The ordinary full-batch plan is unused.
+            self.forward_metadata = None
+            return
         if forward_batch.forward_mode.is_decode_or_idle():
             self.indices_updater_decode.update(
                 forward_batch.req_pool_indices,
@@ -511,6 +527,142 @@ def init_mha_chunk_metadata(
         """Init the metadata for a forward pass."""
         self.mha_chunk_kv_cache.update_wrapper(forward_batch, disable_flashinfer_ragged)
 
+    def _plan_cp_prefill_wrapper(
+        self,
+        forward_batch: ForwardBatch,
+        qo_indptr: torch.Tensor,
+        kv_lens: torch.Tensor,
+        kv_lens_sum: int,
+    ):
+        bs = len(forward_batch.req_pool_indices)
+        kv_indptr = torch.zeros(
+            bs + 1,
+            dtype=torch.int32,
+            device=forward_batch.req_pool_indices.device,
+        )
+        kv_indptr[1:] = torch.cumsum(kv_lens, dim=0)
+        kv_indices = torch.empty(
+            kv_lens_sum,
+            dtype=torch.int32,
+            device=forward_batch.req_pool_indices.device,
+        )
+        req_to_token = self.req_to_token_pool.req_to_token
+        create_flashinfer_kv_indices_triton[(bs,)](
+            req_to_token,
+            forward_batch.req_pool_indices,
+            kv_lens,
+            kv_indptr,
+            None,
+            kv_indices,
+            req_to_token.shape[1],
+        )
+
+        wrapper = BatchMLAPagedAttentionWrapper(
+            self.workspace_buffer,
+            backend="auto",
+        )
+        updater = self.indices_updater_prefill
+        wrapper.plan(
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            kv_lens,
+            updater.num_local_heads,
+            updater.kv_lora_rank,
+            updater.qk_rope_head_dim,
+            1,
+            True,
+            updater.scaling,
+            updater.q_data_type,
+            updater.data_type,
+        )
+        return wrapper, kv_indptr, kv_indices
+
+    def _get_cp_prefill_metadata(
+        self, forward_batch: ForwardBatch
+    ) -> CPPrefillMetadata:
+        if self.cp_prefill_metadata is not None:
+            return self.cp_prefill_metadata
+
+        meta = forward_batch.attn_cp_metadata
+        prev = self._plan_cp_prefill_wrapper(
+            forward_batch,
+            meta.cu_seqlens_q_prev_tensor,
+            meta.kv_len_prev_tensor,
+            sum(meta.kv_len_prev_list),
+        )
+        next_ = self._plan_cp_prefill_wrapper(
+            forward_batch,
+            meta.cu_seqlens_q_next_tensor,
+            meta.kv_len_next_tensor,
+            sum(meta.kv_len_next_list),
+        )
+        self.cp_prefill_metadata = CPPrefillMetadata(
+            wrappers=(prev[0], next_[0]),
+            kv_indptrs=(prev[1], next_[1]),
+            kv_indices=(prev[2], next_[2]),
+        )
+        return self.cp_prefill_metadata
+
+    def _run_cp_paged_attention(
+        self,
+        wrapper: BatchMLAPagedAttentionWrapper,
+        q: torch.Tensor,
+        layer: RadixAttention,
+    ) -> torch.Tensor:
+        q_nope = q[..., : layer.v_head_dim]
+        q_rope = q[..., layer.v_head_dim :]
+        kv_buffer = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype)
+        ckv_cache = kv_buffer[:, :, : layer.v_head_dim]
+        kpe_cache = kv_buffer[:, :, layer.v_head_dim :]
+        output = q_nope.new_empty(q_nope.shape)
+        return wrapper.run(q_nope, q_rope, ckv_cache, kpe_cache, out=output)
+
+    def _forward_extend_cp_v2(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool,
+        q_rope: Optional[torch.Tensor],
+        k_rope: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        strategy = get_cp_strategy()
+        assert strategy is not None
+        assert k_rope is not None
+        if save_kv_cache:
+            strategy.materialize_full_mla_kv(forward_batch, layer, k, k_rope)
+
+        if q_rope is None:
+            q_fused = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+        else:
+            q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+            q_rope = q_rope.view(
+                -1,
+                layer.tp_q_head_num,
+                layer.head_dim - layer.v_head_dim,
+            )
+            q_fused = torch.cat([q_nope, q_rope], dim=-1)
+
+        cp_metadata = self._get_cp_prefill_metadata(forward_batch)
+        wrapper_index = 0
+
+        def _mla_cp_attn(q_chunk, *_):
+            nonlocal wrapper_index
+            wrapper = cp_metadata.wrappers[wrapper_index]
+            wrapper_index += 1
+            return self._run_cp_paged_attention(wrapper, q_chunk, layer)
+
+        output = strategy.run_attention(
+            q_fused,
+            forward_batch,
+            self.device,
+            _mla_cp_attn,
+            attention_backend=CPAttentionBackendKind.FLASH_ATTENTION,
+        )
+        return output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
     def forward_extend(
         self,
         q: torch.Tensor,
@@ -522,6 +674,18 @@ def forward_extend(
         q_rope: Optional[torch.Tensor] = None,
         k_rope: Optional[torch.Tensor] = None,
     ):
+        if is_cp_v2_active(forward_batch):
+            assert k is not None and v is not None
+            return self._forward_extend_cp_v2(
+                q,
+                k,
+                layer,
+                forward_batch,
+                save_kv_cache,
+                q_rope,
+                k_rope,
+            )
+
         if forward_batch.attn_attend_prefix_cache is not None and any(
             forward_batch.extend_prefix_lens_cpu
         ):  # MHA Chunk

python/sglang/srt/layers/cp/kimi_linear.py

@@ -0,0 +1,112 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""CP-v2 token-layout transitions for Kimi-Linear decoder layers."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional, Tuple
+
+import torch
+
+from sglang.srt.layers.cp.utils import get_cp_strategy, is_cp_v2_active
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class KimiLinearCPV2LayerCommunicator:
+    """Convert Kimi-Linear layer inputs between KDA and MLA token layouts."""
+
+    def __init__(
+        self,
+        *,
+        is_kda_layer: bool,
+        previous_is_kda_layer: Optional[bool],
+        is_last_layer: bool = False,
+    ) -> None:
+        self._is_last_layer = is_last_layer
+        is_first_layer = previous_is_kda_layer is None
+        # CP-v2 enters the model sharded. Every MLP exits with a full TP batch.
+        self._gather_before_attn = is_kda_layer and is_first_layer
+        self._shard_before_attn = not is_kda_layer and not is_first_layer
+        self._gather_before_mlp = not is_kda_layer
+
+    def prepare_attn(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        forward_batch: ForwardBatch,
+        stream: Optional[Any] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if not is_cp_v2_active(forward_batch):
+            return hidden_states, residual
+
+        strategy = get_cp_strategy()
+        assert strategy is not None
+        if self._gather_before_attn:
+            if stream is None:
+                stream = torch.cuda.current_stream()
+            hidden_states = strategy.gather_hidden_states(
+                hidden_states, forward_batch, stream
+            )
+            if residual is not None:
+                residual = strategy.gather_hidden_states(
+                    residual, forward_batch, stream
+                )
+        elif self._shard_before_attn:
+            hidden_states = strategy.shard_hidden_states(hidden_states, forward_batch)
+            if residual is not None:
+                residual = strategy.shard_hidden_states(residual, forward_batch)
+        return hidden_states, residual
+
+    def prepare_mlp(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        forward_batch: ForwardBatch,
+        stream: Optional[Any] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Gather MLA outputs so normalization and MLP run with global TP."""
+        if not self._gather_before_mlp or not is_cp_v2_active(forward_batch):
+            return hidden_states, residual
+
+        strategy = get_cp_strategy()
+        assert strategy is not None
+        if stream is None:
+            stream = torch.cuda.current_stream()
+        hidden_states = strategy.gather_hidden_states(
+            hidden_states, forward_batch, stream
+        )
+        if residual is not None:
+            residual = strategy.gather_hidden_states(residual, forward_batch, stream)
+        return hidden_states, residual
+
+    def postprocess_layer(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        forward_batch: ForwardBatch,
+        stream: Optional[Any] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Shard the full TP output for the model-boundary CP gather."""
+        if not self._is_last_layer or not is_cp_v2_active(forward_batch):
+            return hidden_states, residual
+
+        strategy = get_cp_strategy()
+        assert strategy is not None
+        hidden_states = strategy.shard_hidden_states(hidden_states, forward_batch)
+        if residual is not None:
+            residual = strategy.shard_hidden_states(residual, forward_batch)
+        return hidden_states, residual

python/sglang/srt/layers/cp/utils.py

@@ -41,6 +41,7 @@
     {
         "Qwen3MoeForCausalLM",
         "DeepseekV3ForCausalLM",
+        "KimiLinearForCausalLM",
     }
 )