[Feat] KVShard: PD transfer and e2e test

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

@@ -37,9 +37,11 @@
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
     filter_kv_indices_for_cp_rank,
+    filter_kv_indices_for_shard_rank,
     get_dsv41_spec_layout,
 )
 from sglang.srt.environ import envs
+from sglang.srt.mem_cache.page_interleave import get_kv_shard_group
 from sglang.srt.runtime_context import (
     get_disagg,
     get_parallel,
@@ -115,6 +117,8 @@ class PrefillServerInfo:
     enable_dsa_cache_layer_split: bool = False
     dsv41_spec_layout: Optional[dict] = None
     decode_allocation_policy: str = "early"
+    # Each prefill shard-group rank holds a disjoint stripe of every KV page.
+    enable_kv_cache_sharding: bool = False
 
     # PD true-retraction rebootstrap: the prefill's HTTP API port. The decode
     # already knows the prefill host (the bootstrap_addr host), so it can POST
@@ -280,6 +284,11 @@ class CommonKVManager(BaseKVManager):
 
     dsv41_spec_layout: Optional[dict] = None
 
+    # Logical-page KV sharding (PD prefill only): this rank's slot in the
+    # group that stripes every KV page.
+    kv_shard_rank: int = 0
+    kv_shard_size: int = 1
+
     # Used by decode when the prefill reported Failed without a reason frame.
     DEFAULT_PREFILL_FAILURE_REASON = (
         "Failed to get kvcache from prefill instance, it might be dead"
@@ -337,6 +346,10 @@ def __init__(
         self.pp_size = get_parallel().pp_size
         self.pp_rank = self.kv_args.pp_rank
         self.local_ip = get_local_ip_auto()
+        if get_parallel().enable_kv_cache_sharding:
+            kv_shard_group = get_kv_shard_group(is_mla_backend)
+            self.kv_shard_rank = kv_shard_group.rank_in_group
+            self.kv_shard_size = kv_shard_group.world_size
         cp_sharded_prefill = self.attn_cp_size > 1 and (
             self.is_hybrid_mla_backend or get_parallel().enable_dsa_cache_layer_split
         )
@@ -353,6 +366,8 @@ def __init__(
             envs.SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER.get()
             or cp_sharded_prefill
             or hybrid_decode_pulls_all_ranks
+            # No sharded rank holds a full page, so every CP rank sends.
+            or self.kv_shard_size > 1
         )
 
         # bind zmq socket
@@ -1272,6 +1287,11 @@ def try_ensure_parallel_info(self, bootstrap_addr: str) -> bool:
                     "PD decode DCP currently requires prefill attention CP=1, "
                     f"got {info.attn_cp_size}."
                 )
+            if info.enable_kv_cache_sharding:
+                raise RuntimeError(
+                    "PD decode DCP does not support a prefill with "
+                    "--enable-kv-cache-sharding."
+                )
 
         self._resolve_rank_mapping(info)
         self.prefill_info_table[bootstrap_addr] = info
@@ -1282,7 +1302,23 @@ def _resolve_rank_mapping(self, info: PrefillServerInfo) -> None:
         """Compute TP/CP/PP rank mapping and store on the PrefillServerInfo object.
         Deterministic for a given (bootstrap_addr, decode engine) pair."""
         # TP rank mapping
-        if self.attn_tp_size == info.attn_tp_size:
+        # Only MLA shards across attn-TP, so non-MLA never reads the sharding flag.
+        if (
+            self.is_mla_backend
+            and info.enable_kv_cache_sharding
+            and info.attn_cp_size == 1
+        ):
+            # MLA KV sharded across prefill attn-TP: every decode rank pulls
+            # from all prefill TP ranks (None marks them all as real senders).
+            assert self.attn_tp_size == info.attn_tp_size, (
+                "MLA KV cache sharding requires equal prefill/decode attn TP "
+                f"sizes, got prefill={info.attn_tp_size} decode={self.attn_tp_size}"
+            )
+            target_tp_rank = None
+            target_tp_ranks = list(range(info.attn_tp_size))
+            required_dst_info_num = self.attn_tp_size
+            required_prefill_response_num = info.attn_tp_size
+        elif self.attn_tp_size == info.attn_tp_size:
             target_tp_rank = self.kv_args.engine_rank % self.attn_tp_size
             required_dst_info_num = 1
             required_prefill_response_num = 1
@@ -1336,6 +1372,7 @@ def _resolve_rank_mapping(self, info: PrefillServerInfo) -> None:
             pull_from_all_cp_ranks = (
                 self.enable_all_cp_ranks_for_transfer
                 or info.enable_dsa_cache_layer_split
+                or info.enable_kv_cache_sharding
             )
             if not pull_from_all_cp_ranks:
                 # Only retrieve from prefill CP rank 0 when not using all ranks
@@ -1431,6 +1468,7 @@ def register_to_bootstrap(self):
             "dsv41_spec_layout": self.dsv41_spec_layout,
             "load_balance_method": get_parallel().load_balance_method,
             "enable_dsa_cache_layer_split": get_parallel().enable_dsa_cache_layer_split,
+            "enable_kv_cache_sharding": get_parallel().enable_kv_cache_sharding,
             # Self-register the HTTP API port so the decode can derive the PD
             # retract rebootstrap /generate URL from bootstrap info instead of a
             # router-injected pd_rebootstrap_prefill_url.
@@ -2083,7 +2121,17 @@ def _prepare_send_indices(
         self.curr_idx += len(kv_indices)
         is_last_chunk = self.curr_idx == self.num_kv_indices
 
-        if (
+        if self.kv_mgr.kv_shard_size > 1:
+            kv_indices, index_slice = filter_kv_indices_for_shard_rank(
+                kv_indices=kv_indices,
+                index_slice=index_slice,
+                shard_rank=self.kv_mgr.kv_shard_rank,
+                shard_size=self.kv_mgr.kv_shard_size,
+            )
+            # Like the scheduler, skip a non-final chunk with nothing to send.
+            if len(kv_indices) == 0 and not is_last_chunk:
+                return kv_indices, index_slice, is_last_chunk, True
+        elif (
             self.kv_mgr.enable_all_cp_ranks_for_transfer
             and not get_parallel().enable_dsa_cache_layer_split
         ):
@@ -2648,6 +2696,7 @@ def __init__(self, host: str, port: int):
         self.dsv41_spec_layout: Optional[dict] = None
         self.follow_bootstrap_room: Optional[bool] = None
         self.enable_dsa_cache_layer_split: Optional[bool] = None
+        self.enable_kv_cache_sharding: Optional[bool] = None
         self.prefill_http_port: Optional[int] = None
         self.decode_allocation_policy = "early"
         self.prefill_port_table: Dict[
@@ -2764,6 +2813,11 @@ async def _handle_route_put(self, request: web.Request):
                 data.get("enable_dsa_cache_layer_split", False)
             )
 
+        if self.enable_kv_cache_sharding is None:
+            self.enable_kv_cache_sharding = bool(
+                data.get("enable_kv_cache_sharding", False)
+            )
+
         if system_dp_size == 1:
             dp_group = attn_dp_rank
         else:
@@ -2829,6 +2883,7 @@ async def _handle_route_get(self, request: web.Request):
                     else True
                 ),
                 enable_dsa_cache_layer_split=bool(self.enable_dsa_cache_layer_split),
+                enable_kv_cache_sharding=bool(self.enable_kv_cache_sharding),
                 prefill_http_port=self.prefill_http_port,
                 decode_allocation_policy=self.decode_allocation_policy,
             )

python/sglang/srt/disaggregation/prefill.py

@@ -191,8 +191,10 @@ def __init__(
         self.queue: List[Req] = []
         self.scheduler = scheduler
         self.scheduler_stage_metrics = scheduler_stage_metrics
+        # KV sharding stores 1/N of each page per rank, so capacity widens by N.
         self.max_total_num_tokens = (
             self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens
+            * self.scheduler.kv_shard_widening
         )
         self.transfer_backend = transfer_backend
         self.pp_poll_sync_work_list: List = []

python/sglang/srt/disaggregation/utils.py

@@ -764,6 +764,20 @@ def _get_cp_rank_page_bounds(
     return local_start, local_start + n_pages
 
 
+def filter_kv_indices_for_shard_rank(
+    *,
+    kv_indices: np.ndarray,
+    index_slice: slice,
+    shard_rank: int,
+    shard_size: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Logical page l lives on rank l % N as local page l // N; returns the owned
+    local ids and their send positions, which replace the contiguous index_slice."""
+    owned = kv_indices % shard_size == shard_rank
+    positions = np.arange(index_slice.start, index_slice.stop, dtype=np.int64)
+    return kv_indices[owned] // shard_size, positions[owned]
+
+
 def filter_kv_indices_for_cp_rank(
     kv_mgr: CommonKVManager,
     kv_indices: np.ndarray,

test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py

@@ -0,0 +1,150 @@
+"""PD accuracy and per-rank KV footprint with logical-page KV cache sharding.
+
+Each prefill TP rank stores only its stripe of every KV page and sends only the
+pages it owns; decode pulls from both prefill ranks. A wrong owner, local page
+id, or send position corrupts decode's KV and collapses GSM8K accuracy. A prompt
+whose KV is larger than one prefill rank's whole pool can only be served if no
+rank keeps a full copy.
+"""
+
+import random
+import re
+import unittest
+
+import requests
+
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import DEFAULT_MLA_MODEL_NAME_FOR_TEST
+
+register_cuda_ci(est_time=200, stage="base-c", runner_config="4-gpu-b200")
+
+PAGE_SIZE = 64
+# trtllm_mla is the Blackwell MLA backend that supports KV sharding.
+COMMON_ARGS = ["--attention-backend", "trtllm_mla", "--page-size", str(PAGE_SIZE)]
+# Physical KV rows per prefill rank. 512 logical pages, so GSM8K's prefill
+# radix footprint (~243 pages) never forces eviction.
+PREFILL_POOL_TOKENS = 16384
+# 1.5x one rank's rows, well under the sharded admission bound ~2 * (16384 - 64).
+LONG_PROMPT_TOKENS = 24576
+# Page p of a fresh prompt is owned by prefill rank p % 2, so these needles sit
+# on different ranks. Page 20 ends before the last 16448 tokens; page 341
+# starts after token 16448: a 16448-row full copy could not hold both.
+NEEDLE_PAGES = {"red": 20, "blue": 341}
+FILLER = (
+    "The grass is green. The sky is blue. The sun is yellow. "
+    "Here we go. There and back again. "
+)
+
+
+def build_needle_prompt(tokenizer, keys: dict[str, int], color: str) -> list[int]:
+    """Token ids of LONG_PROMPT_TOKENS; each pass key starts its own page."""
+    filler = tokenizer.encode(FILLER, add_special_tokens=False)
+
+    def pad_to(ids: list[int], length: int) -> list[int]:
+        while len(ids) < length:
+            ids = ids + filler[: length - len(ids)]
+        return ids
+
+    ids = tokenizer.encode(
+        "There is important info hidden inside a lot of irrelevant text. "
+        "Find it and memorize it.\n"
+    )
+    for name, page in sorted(NEEDLE_PAGES.items(), key=lambda item: item[1]):
+        ids = pad_to(ids, page * PAGE_SIZE)
+        ids += tokenizer.encode(
+            f"The pass key for the {name} door is {keys[name]}. Remember it. ",
+            add_special_tokens=False,
+        )
+    # Prefill samples only the space; decode generates every digit from the KV
+    # it received.
+    question = tokenizer.encode(
+        f"\nWhat is the pass key for the {color} door? "
+        f"The pass key for the {color} door is",
+        add_special_tokens=False,
+    )
+    return pad_to(ids, LONG_PROMPT_TOKENS - len(question)) + question
+
+
+class TestDisaggregationKVCacheSharding(PDDisaggregationServerBase, GSM8KMixin):
+    model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+    prefill_tp_size = 2
+    decode_tp_size = 2
+    decode_base_gpu_id = 2
+    # KV sharding with the radix cache on requires the Python tree core.
+    extra_prefill_env = {"SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND": "python"}
+    # 256-token chunks split each 5-shot prompt, so later chunks read a
+    # sharded prefix and every rank sends pages from several chunks.
+    extra_prefill_args = COMMON_ARGS + [
+        "--enable-kv-cache-sharding",
+        "--chunked-prefill-size",
+        "256",
+        "--max-total-tokens",
+        str(PREFILL_POOL_TOKENS),
+    ]
+    extra_decode_args = COMMON_ARGS
+    gsm8k_score_threshold = 0.70
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.launch_all()
+
+    def _server_info(self, url: str) -> dict:
+        response = requests.get(url + "/server_info", timeout=60)
+        response.raise_for_status()
+        return response.json()
+
+    def test_prompt_kv_larger_than_one_rank_pool(self):
+        prefill_info = self._server_info(self.prefill_url)
+        # Fail fast: a prefill-side reject reaches the client only through
+        # decode's 300 s KV-wait timeout.
+        self.assertGreater(prefill_info["max_req_input_len"], LONG_PROMPT_TOKENS)
+        prefill_mem = prefill_info["internal_states"][0]["memory_usage"]
+        decode_mem = self._server_info(self.decode_url)["internal_states"][0][
+            "memory_usage"
+        ]
+        self.assertEqual(prefill_mem["token_capacity"], PREFILL_POOL_TOKENS)
+        # kvcache is summed from the pool tensors. Decode's unsharded pool of
+        # the same model gives the full KV size per token.
+        full_kv_gb_per_token = decode_mem["kvcache"] / decode_mem["token_capacity"]
+        prompt_kv_gb = LONG_PROMPT_TOKENS * full_kv_gb_per_token
+        print(f"{prompt_kv_gb=:.3f} vs one prefill rank's pool {prefill_mem=}")
+        self.assertGreater(prompt_kv_gb, 1.3 * prefill_mem["kvcache"])
+
+        tokenizer = get_tokenizer(self.model)
+        rng = random.Random(0)
+        for color in NEEDLE_PAGES:
+            # New keys per request: a flush does not zero the pools, so a
+            # reused key could be answered from stale rows.
+            keys = {name: rng.randint(10000, 99999) for name in NEEDLE_PAGES}
+            input_ids = build_needle_prompt(tokenizer, keys, color)
+            # Start from an empty prefill cache so every page is allocated now.
+            requests.post(
+                self.prefill_url + "/flush_cache", params={"timeout": 30}, timeout=60
+            ).raise_for_status()
+            response = requests.post(
+                self.lb_url + "/generate",
+                json={
+                    "input_ids": input_ids,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+                },
+                timeout=300,
+            )
+            self.assertEqual(response.status_code, 200, response.text)
+            output = response.json()
+            meta_info = output["meta_info"]
+            print(f"{color}: {output['text']!r} {meta_info['cached_tokens']=}")
+            self.assertEqual(meta_info["prompt_tokens"], LONG_PROMPT_TOKENS)
+            self.assertEqual(meta_info["cached_tokens"], 0)
+            answer = re.search(r"\d+", output["text"])
+            self.assertIsNotNone(answer, output["text"])
+            self.assertEqual(answer.group(), str(keys[color]), output["text"])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_register_to_bootstrap.py

@@ -231,6 +231,8 @@ def test_payload_contains_required_fields(self, mock_put, mock_time):
             "rank_port",
             "page_size",
             "kv_cache_dtype",
+            # Decode pulls from every shard-group rank when this is set.
+            "enable_kv_cache_sharding",
             # Self-registered HTTP API port used to derive the PD retract
             # rebootstrap /generate URL on the decode side.
             "prefill_http_port",