[PD] Enable all CP ranks for KVCache transfer (#19765)

Signed-off-by: Shangming Cai <csmthu@gmail.com>

python/sglang/srt/disaggregation/common/conn.py

@@ -110,6 +110,9 @@ class CommonKVManager(BaseKVManager):
         self.pp_size = server_args.pp_size
         self.pp_rank = self.kv_args.pp_rank
         self.local_ip = get_local_ip_auto()
+        self.enable_all_cp_ranks_for_transfer = (
+            envs.SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER.get()
+        )
 
         # bind zmq socket
         context = zmq.Context()
@@ -124,9 +127,12 @@ class CommonKVManager(BaseKVManager):
         self.failure_lock = threading.Lock()
 
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            # TODO(shangming): Fix me when we support MHA/GQA + CP, or when we utilize all cp ranks for KV transfer in CP mode.
+            # When SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER is True, all CP ranks
+            # participate in KV transfer; Otherwise only CP rank 0 sends.
             self.is_dummy_cp_rank = (
-                is_mla_backend and self.attn_cp_size > 1 and self.attn_cp_rank != 0
+                not self.enable_all_cp_ranks_for_transfer
+                and self.attn_cp_size > 1
+                and self.attn_cp_rank != 0
             )
             self.register_to_bootstrap()
             self.transfer_infos = {}
@@ -395,6 +401,9 @@ class CommonKVSender(BaseKVSender):
     def init(self, num_kv_indices: int, aux_index: Optional[int] = None):
         self.num_kv_indices = num_kv_indices
         self.aux_index = aux_index
+        logger.debug(
+            f"CommonKVSender init with num_kv_indices: {num_kv_indices} and aux_index: {aux_index}"
+        )
 
     def send(
         self,
@@ -488,23 +497,22 @@ class CommonKVReceiver(BaseKVReceiver):
                     self.prefill_info.attn_tp_size // self.kv_mgr.attn_tp_size
                 )
 
-        # Decode cp size should be equal to prefill cp size or 1
+        # Decode cp size should be equal to 1
-        assert (
+        assert self.kv_mgr.attn_cp_size == 1, (
-            self.kv_mgr.attn_cp_size == self.prefill_info.attn_cp_size
+            f"Decode cp size ({self.kv_mgr.attn_cp_size}) should be equal to 1",
-            or self.kv_mgr.attn_cp_size == 1
-        ), (
-            f"Decode cp size ({self.kv_mgr.attn_cp_size}) should be equal to prefill cp size ({self.prefill_info.attn_cp_size}) or 1",
         )
         if self.kv_mgr.attn_cp_size == self.prefill_info.attn_cp_size:
+            # This means that the prefill cp size is 1
+            assert self.prefill_info.attn_cp_size == 1, (
+                f"When prefill cp size is 1, attn cp size should be 1, but got {self.kv_mgr.attn_cp_size}",
+            )
             self.target_cp_ranks = [self.kv_mgr.attn_cp_rank]
         else:
             self.target_cp_ranks = [
                 rank for rank in range(self.prefill_info.attn_cp_size)
             ]
-            # TODO(shangming): Support KVCache transfer for multiple prefill cp ranks -> 1 decode cp rank
+            if not self.kv_mgr.enable_all_cp_ranks_for_transfer:
-            # For now, we handle the control plane in advance, but we need to support the data plane in the future.
+                # Only retrieve from prefill CP rank 0 when not using all ranks
-            if self.kv_mgr.is_mla_backend:
-                # For MLA: we only need to retrieve KVCache from the first CP rank now
                 self.target_cp_ranks = self.target_cp_ranks[:1]
                 self.required_prefill_response_num *= 1
             else:

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -28,7 +28,10 @@ from sglang.srt.disaggregation.common.utils import (
 from sglang.srt.disaggregation.mooncake.utils import (
     check_mooncake_custom_mem_pool_enabled,
 )
-from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    filter_kv_indices_for_cp_rank,
+)
 from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
 from sglang.srt.environ import envs
 from sglang.srt.server_args import ServerArgs
@@ -779,7 +782,12 @@ class MooncakeKVManager(CommonKVManager):
                 )
                 polls = []
                 dst_ranks_infos = []
-                local_rank = self.attn_tp_rank * self.pp_size + self.pp_rank
+                # Unique id per prefill sender so decode's response set size matches expected_response_num.
+                prefill_unique_rank = (
+                    self.attn_tp_rank * (self.pp_size * self.attn_cp_size)
+                    + self.pp_rank * self.attn_cp_size
+                    + self.attn_cp_rank
+                )
                 for req in reqs_to_be_processed:
                     if not req.is_dummy:
                         # Early exit if the request has failed
@@ -795,7 +803,7 @@ class MooncakeKVManager(CommonKVManager):
                                     req.dst_port,
                                     req.room,
                                     KVPoll.Failed,
-                                    local_rank,
+                                    prefill_unique_rank,
                                 )
                                 break
 
@@ -857,7 +865,7 @@ class MooncakeKVManager(CommonKVManager):
                                 req.dst_port,
                                 req.room,
                                 KVPoll.Failed,
-                                local_rank,
+                                prefill_unique_rank,
                             )
                             break
 
@@ -888,7 +896,11 @@ class MooncakeKVManager(CommonKVManager):
                                 self.update_status(req.room, status)
                                 for endpoint, dst_port, room in dst_ranks_infos:
                                     self.sync_status_to_decode_endpoint(
-                                        endpoint, dst_port, room, status, local_rank
+                                        endpoint,
+                                        dst_port,
+                                        room,
+                                        status,
+                                        prefill_unique_rank,
                                     )
                     else:
                         # Dummy request means the decode instance is not used, so its status can be marked as success directly
@@ -1136,7 +1148,14 @@ class MooncakeKVSender(CommonKVSender):
         self.curr_idx += len(kv_indices)
         is_last_chunk = self.curr_idx == self.num_kv_indices
 
-        if self.kv_mgr.is_dummy_cp_rank:
+        # Special handling for cp
+        if self.kv_mgr.enable_all_cp_ranks_for_transfer:
+            kv_indices, index_slice = filter_kv_indices_for_cp_rank(
+                self.kv_mgr,
+                kv_indices,
+                index_slice,
+            )
+        elif self.kv_mgr.is_dummy_cp_rank:
             if not is_last_chunk:
                 return
             else:

python/sglang/srt/disaggregation/mori/conn.py

@@ -32,7 +32,10 @@ from sglang.srt.disaggregation.common.conn import (
     CommonKVSender,
 )
 from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
-from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    filter_kv_indices_for_cp_rank,
+)
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils.common import (
     format_tcp_address,
@@ -865,6 +868,19 @@ class MoriKVSender(CommonKVSender):
         self.curr_idx += len(kv_indices)
         is_last = self.curr_idx == self.num_kv_indices
 
+        # Special handling for cp
+        if self.kv_mgr.enable_all_cp_ranks_for_transfer:
+            kv_indices, index_slice = filter_kv_indices_for_cp_rank(
+                self.kv_mgr,
+                kv_indices,
+                index_slice,
+            )
+        elif self.kv_mgr.is_dummy_cp_rank:
+            if not is_last:
+                return
+            else:
+                self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
+                return
         statuses, infos = self.kv_mgr.add_transfer_request(
             self.bootstrap_room,
             kv_indices,

python/sglang/srt/disaggregation/nixl/conn.py

@@ -20,7 +20,10 @@ from sglang.srt.disaggregation.common.conn import (
     CommonKVSender,
 )
 from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
-from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    filter_kv_indices_for_cp_rank,
+)
 from sglang.srt.environ import envs
 from sglang.srt.server_args import ServerArgs
 
@@ -905,6 +908,20 @@ class NixlKVSender(CommonKVSender):
         self.curr_idx += len(kv_indices)
         is_last = self.curr_idx == self.num_kv_indices
 
+        # Special handling for cp
+        if self.kv_mgr.enable_all_cp_ranks_for_transfer:
+            kv_indices, index_slice = filter_kv_indices_for_cp_rank(
+                self.kv_mgr,
+                kv_indices,
+                index_slice,
+            )
+        elif self.kv_mgr.is_dummy_cp_rank:
+            if not is_last:
+                return
+            else:
+                self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
+                return
+
         new_xfer_handles = self.kv_mgr.add_transfer_request(
             self.bootstrap_room,
             kv_indices,

python/sglang/srt/disaggregation/utils.py

@@ -5,7 +5,7 @@ import random
 from collections import deque
 from contextlib import nullcontext
 from enum import Enum
-from typing import TYPE_CHECKING, Literal, Optional, Type, overload
+from typing import TYPE_CHECKING, Literal, Optional, Tuple, Type, overload
 
 import numpy as np
 import torch
@@ -424,6 +424,87 @@ def kv_to_page_num(num_kv_indices: int, page_size: int):
     return (num_kv_indices + page_size - 1) // page_size
 
 
+def page_indices_to_cp_rank_page_indices(
+    page_indices: np.ndarray,
+    total_pages: int,
+    cp_rank: int,
+    cp_size: int,
+) -> np.ndarray:
+    """
+    Filter page_indices (which are *global* page ids in the KV pool) to those
+    belonging to the given CP rank for this request.
+
+    For a single request, its pages occupy a contiguous global range
+    [first_page, first_page + total_pages). We first compute the local
+    split [0, total_pages) across cp_size ranks, then shift that local
+    range by first_page back into the global page id space and take
+    the intersection with page_indices.
+
+    Returns:
+        Subset of page_indices that fall in this rank's global
+        [start_page, end_page) slice for the given CP rank.
+    """
+    if cp_size <= 1:
+        return page_indices
+
+    if page_indices.size == 0:
+        return np.asarray(page_indices)
+
+    first_page = int(page_indices.min())
+    base = total_pages // cp_size
+    rem = total_pages % cp_size
+
+    if rem == 0:
+        local_start = cp_rank * base
+        local_end = local_start + base
+    else:
+        local_start = cp_rank * base + min(cp_rank, rem)
+        n_pages = base + (1 if cp_rank < rem else 0)
+        local_end = local_start + n_pages
+
+    # Map back to global page ids.
+    start_page = first_page + local_start
+    end_page = first_page + local_end
+
+    mask = (page_indices >= start_page) & (page_indices < end_page)
+    return np.asarray(page_indices)[mask]
+
+
+def filter_kv_indices_for_cp_rank(
+    kv_mgr: CommonKVManager, kv_indices: np.ndarray, index_slice: slice
+) -> Tuple[np.ndarray, slice]:
+    """Filters kv_indices and index_slice for the current CP rank."""
+    total_pages = len(kv_indices)
+    cp_rank = kv_mgr.attn_cp_rank
+    cp_size = kv_mgr.attn_cp_size
+
+    rank_page_indices = page_indices_to_cp_rank_page_indices(
+        page_indices=kv_indices,
+        total_pages=total_pages,
+        cp_rank=cp_rank,
+        cp_size=cp_size,
+    )
+
+    if rank_page_indices.size == 0:
+        new_kv_indices = kv_indices[:0]
+        new_index_slice = slice(index_slice.start, index_slice.start)
+    else:
+        mask = np.isin(kv_indices, rank_page_indices)
+        if not mask.any():
+            new_kv_indices = kv_indices[:0]
+            new_index_slice = slice(index_slice.start, index_slice.start)
+        else:
+            first_pos = int(mask.argmax())
+            last_pos = len(mask) - int(mask[::-1].argmax())
+
+            new_kv_indices = kv_indices[first_pos:last_pos]
+            new_index_slice = slice(
+                index_slice.start + first_pos,
+                index_slice.start + last_pos,
+            )
+    return new_kv_indices, new_index_slice
+
+
 #########################
 # Misc
 #########################

python/sglang/srt/environ.py

@@ -242,6 +242,7 @@ class Envs:
     SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2)
     SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
     SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX")
+    SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False)
 
     # Scheduler: others:
     SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1)  # in seconds. Set if you observe high memory accumulation over a long serving period.