From 31101dcf87c7ad474e530085492420f01d6bd0b4 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Sun, 11 Oct 2026 04:22:58 -0700 Subject: [PATCH 1/6] [kv-shard 4/4] P/D transfer: owned-page send filter, bootstrap, E2E tests A sharded prefill rank holds only its stripe of every KV page, so each rank sends the pages it owns and decode pulls from every rank of the shard group. - disaggregation/utils.py: filter_kv_indices_for_shard_rank. Logical page l lives on rank l % N as local page l // N; the owned send positions replace the contiguous index slice, and decode pairs destinations positionally. - disaggregation/common/conn.py: take the shard identity from get_kv_shard_group; filter every send and skip a non-final chunk the rank owns no page of (NIXL chunk ids stay contiguous); advertise enable_kv_cache_sharding through bootstrap; decode pulls from all CP ranks, or from all prefill TP ranks for MLA; decode DCP rejects a sharded prefill. - disaggregation/prefill.py: widen the bootstrap capacity bound by the shard size, like the scheduler's other capacity gates. Tests: unit tests for the send filter, the empty-chunk skip and the bootstrap payload; a 4-GPU PD E2E (DeepSeek-Coder-V2-Lite, prefill TP2 + decode TP2, trtllm_mla) gating GSM8K accuracy. Co-Authored-By: Claude Opus 5.5 --- .../sglang/srt/disaggregation/common/conn.py | 58 ++++++++++++++++++- python/sglang/srt/disaggregation/prefill.py | 2 + python/sglang/srt/disaggregation/utils.py | 14 +++++ .../test_disaggregation_kv_cache_sharding.py | 47 +++++++++++++++ .../test_register_to_bootstrap.py | 2 + .../mem_cache/test_page_interleave_shard.py | 43 ++++++++++++++ 6 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 45cb4c20bdd3..02006a8e7b91 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/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, @@ -114,6 +116,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 @@ -279,6 +283,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" @@ -331,6 +340,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 ) @@ -343,6 +356,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 @@ -1233,6 +1248,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 @@ -1243,7 +1263,22 @@ 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: + if ( + info.enable_kv_cache_sharding + and self.is_mla_backend + 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 @@ -1297,6 +1332,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 @@ -1387,6 +1423,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. @@ -1929,7 +1966,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 ): @@ -2494,6 +2541,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[ @@ -2610,6 +2658,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: @@ -2675,6 +2728,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, ) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 7ff3603cd7d1..735866b84a46 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -190,8 +190,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 = [] diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 140f82d48f0d..d6c27eccad47 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -766,6 +766,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, diff --git a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py new file mode 100644 index 000000000000..677691cee343 --- /dev/null +++ b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py @@ -0,0 +1,47 @@ +"""PD accuracy with logical-page KV cache sharding on a small MLA model. + +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. +""" + +import unittest + +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") + +# trtllm_mla is the Blackwell MLA backend that supports KV sharding. +COMMON_ARGS = ["--attention-backend", "trtllm_mla", "--page-size", "64"] + + +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", + ] + extra_decode_args = COMMON_ARGS + gsm8k_score_threshold = 0.70 + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.launch_all() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_register_to_bootstrap.py b/test/registered/unit/disaggregation/test_register_to_bootstrap.py index 322bca9e6659..2fdad977d8d6 100644 --- a/test/registered/unit/disaggregation/test_register_to_bootstrap.py +++ b/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", diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index f4675277373a..f50ded6141a2 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -30,6 +30,8 @@ class draws (owners exactly cyclic along a chain), least-full root 5. ``begin_shard_extend`` plan capture (page positions, padded send rows, owner-congruence guard) with the gather stubbed out, following the SimpleNamespace binding pattern of ``test_dsa_layer_shard_utils.py``. +6. The P/D send filter ``filter_kv_indices_for_shard_rank`` and the sender's + skip of a non-final chunk this rank owns no page of. """ import shutil @@ -39,9 +41,12 @@ class draws (owners exactly cyclic along a chain), least-full root from contextlib import contextmanager from types import SimpleNamespace +import numpy as np import torch from parameterized import parameterized_class +from sglang.srt.disaggregation.common.conn import CommonKVSender +from sglang.srt.disaggregation.utils import filter_kv_indices_for_shard_rank from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache import page_interleave @@ -1327,5 +1332,43 @@ def test_partial_tail_page_may_own_nothing(self): self.assertEqual(local_rows.numel(), 0) +class TestShardSendFilter(CustomTestCase): + """P/D send side: rank r sends the pages with l % N == r, as its local page + l // N, at their request-global send positions.""" + + def test_ranks_partition_a_rotated_chunk(self): + # A base-2 rotated chain: owners are not position-congruent. + pages = np.array(_chain_pages(base=2, n_pages=9), dtype=np.int64) + positions = [] + for rank in range(N): + local, pos = filter_kv_indices_for_shard_rank( + kv_indices=pages, + index_slice=slice(20, 29), + shard_rank=rank, + shard_size=N, + ) + np.testing.assert_array_equal(pages[pos - 20] % N, rank) + np.testing.assert_array_equal(local, pages[pos - 20] // N) + positions.append(pos) + np.testing.assert_array_equal( + np.sort(np.concatenate(positions)), np.arange(20, 29) + ) + + def test_empty_chunk_is_skipped_unless_final(self): + # One page per owner 0..3, sent as two 2-page chunks. + pages = np.array(_chain_pages(base=0, n_pages=4), dtype=np.int64) + for rank, expect_skip in ((3, [True, False]), (1, [False, False])): + sender = SimpleNamespace( + kv_mgr=SimpleNamespace(kv_shard_rank=rank, kv_shard_size=N), + curr_idx=0, + num_kv_indices=len(pages), + ) + skips = [ + CommonKVSender._prepare_send_indices(sender, chunk)[3] + for chunk in (pages[:2], pages[2:]) + ] + self.assertEqual(skips, expect_skip) + + if __name__ == "__main__": unittest.main() From 59f3c034adfb6c4d57c19e7a800a789f152855d0 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Sun, 11 Oct 2026 21:43:33 -0700 Subject: [PATCH 2/6] [kv-shard 4/4] Test that sharding cuts per-rank KV memory The accuracy E2E passes even if every rank kept a full KV copy. Prove the copies are gone instead: - E2E: cap the sharded prefill at 16384 physical rows per rank and serve a 24576-token passkey prompt, with needles on pages owned by each rank. The prompt's KV (0.712 GB, measured by decode's unsharded pool) exceeds one prefill rank's whole pool (0.476 GB, from the pool tensors), so it can only fit and answer correctly if no rank holds a full copy. This also covers the widened bootstrap capacity bound, which GSM8K prompts never reach. - Unit (CPU): fill the full logical capacity of N stock-size MLA pools for N in {2, 4, 8}; each rank fills exactly its own rows and the group stores every token exactly once. Co-Authored-By: Claude Opus 5.5 --- .../test_disaggregation_kv_cache_sharding.py | 109 +++++++++++++++++- .../mem_cache/test_page_interleave_shard.py | 96 ++++++++++++++- 2 files changed, 200 insertions(+), 5 deletions(-) diff --git a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py index 677691cee343..1dad15a1a0ce 100644 --- a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py +++ b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py @@ -1,12 +1,19 @@ -"""PD accuracy with logical-page KV cache sharding on a small MLA model. +"""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. +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 ( @@ -16,8 +23,51 @@ 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", "64"] +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): @@ -33,6 +83,8 @@ class TestDisaggregationKVCacheSharding(PDDisaggregationServerBase, GSM8KMixin): "--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 @@ -42,6 +94,57 @@ 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() diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index f50ded6141a2..389671c87b06 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -32,6 +32,8 @@ class draws (owners exactly cyclic along a chain), least-full root SimpleNamespace binding pattern of ``test_dsa_layer_shard_utils.py``. 6. The P/D send filter ``filter_kv_indices_for_shard_rank`` and the sender's skip of a non-final chunk this rank owns no page of. +7. Real sharded MLA pools: N ranks of stock-size pools hold N times one pool's + tokens, each token stored on exactly one rank. """ import shutil @@ -62,7 +64,11 @@ class draws (owners exactly cyclic along a chain), least-full root ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import _evict_until_allocatable -from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool +from sglang.srt.mem_cache.memory_pool import ( + MHATokenToKVPool, + MLATokenToKVPool, + ReqToTokenPool, +) from sglang.srt.mem_cache.page_interleave import ( PageInterleavePlacement, PageShardSpec, @@ -77,10 +83,11 @@ class draws (owners exactly cyclic along a chain), least-full root from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.components import ComponentType from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import ceil_div from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.mem_cache_utils import finish_req -from sglang.test.test_utils import CustomTestCase +from sglang.test.test_utils import CustomTestCase, enter_scope register_cpu_ci(est_time=30, suite="base-a-test-cpu") @@ -1332,6 +1339,91 @@ def test_partial_tail_page_may_own_nothing(self): self.assertEqual(local_rows.numel(), 0) +class TestShardedPoolFootprint(CustomTestCase): + """Real MLA pools on CPU: the shard group stores one copy of the KV.""" + + PAGES_PER_RANK = 8 + POOL_KWARGS = dict( + page_size=PS, + dtype=torch.bfloat16, + kv_lora_rank=16, + qk_rope_head_dim=8, + layer_num=1, + device="cpu", + enable_memory_saver=False, + ) + + def setUp(self): + super().setUp() + # The MLA write path asserts against decode context parallelism. + enter_scope(self, get_parallel().override(dcp_enabled=False)) + + def _sharded_pool(self, shard_size, shard_rank): + spec = PageShardSpec( + shard_rank=shard_rank, + shard_size=shard_size, + page_size=PS, + max_prefix_tokens=shard_size * PS, + chunk_tokens=PS, + ) + group = SimpleNamespace( + world_size=shard_size, + rank_in_group=shard_rank, + cpu_group=None, + device="cpu", + ) + with unittest.mock.patch( + "sglang.srt.mem_cache.page_interleave_pool.PyNcclCommunicator" + ): + return PageInterleaveMLATokenToKVPool( + self.PAGES_PER_RANK * PS, + shard_spec=spec, + shard_group=group, + **self.POOL_KWARGS, + ) + + def test_full_logical_capacity_stores_each_token_once(self): + stock = MLATokenToKVPool(self.PAGES_PER_RANK * PS, **self.POOL_KWARGS) + for shard_size in (2, 4, 8): + with self.subTest(shard_size=shard_size): + pools = [self._sharded_pool(shard_size, r) for r in range(shard_size)] + for pool in pools: + self.assertEqual( + pool.get_kv_size_bytes(), stock.get_kv_size_bytes() + ) + alloc = PageInterleavePoolAllocator( + size=stock.size, + physical_page_size=PS, + shard_size=shard_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=pools[0], + need_sort=False, + ) + num_tokens = shard_size * stock.size + loc = _alloc_extend(alloc, 0, num_tokens, None) + self.assertEqual(alloc.available_size(), 0) + # Token i carries (i % 128 + 1, i // 128 + 1), exact in bf16. + token = torch.arange(num_tokens) + cache_k = torch.zeros(num_tokens, 1, 24, dtype=torch.bfloat16) + cache_k[:, 0, 0] = (token % 128 + 1).to(torch.bfloat16) + cache_k[:, 0, 1] = (token // 128 + 1).to(torch.bfloat16) + stored = [] + for pool in pools: + pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, cache_k, None) + tags = pool.kv_buffer[0][:, 0, :2].float() + rows = torch.nonzero(tags[:, 0]).squeeze(1) + # Every row past the padded page is filled, and no more. + self.assertTrue( + torch.equal(rows, torch.arange(PS, PS + stock.size)) + ) + stored.append( + ((tags[rows, 1] - 1) * 128 + tags[rows, 0] - 1).long() + ) + # One copy across the group: each token on exactly one rank. + self.assertTrue(torch.equal(torch.cat(stored).sort().values, token)) + + class TestShardSendFilter(CustomTestCase): """P/D send side: rank r sends the pages with l % N == r, as its local page l // N, at their request-global send positions.""" From 2069b12621600d1ad04c40fa5dff4057d13af6c3 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Wed, 14 Oct 2026 05:58:28 -0700 Subject: [PATCH 3/6] [kv-shard 4/4] Read the sharding flag only on the MLA rank-mapping path _resolve_rank_mapping read info.enable_kv_cache_sharding before checking is_mla_backend. A non-MLA decode now stops at is_mla_backend, so it reads only the fields main reads. This fixes the hetero-TP SWA wire test, whose SimpleNamespace info has no sharding field. Behavior is unchanged for a real PrefillServerInfo, where the field defaults to False. --- python/sglang/srt/disaggregation/common/conn.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index af0ef851c2b5..99f6217f4898 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -1302,9 +1302,10 @@ 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 + # Only MLA shards across attn-TP, so non-MLA never reads the sharding flag. if ( - info.enable_kv_cache_sharding - and self.is_mla_backend + 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 1ca853ca4d232dfd222997ffcab09beb96900791 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:16:09 -0700 Subject: [PATCH 4/6] [kv-shard 4/4] Reject prompts that sharded prefill admission can never fit The PD bootstrap check rejected only prompts longer than N*T (N shard ranks, T tokens per rank). Sharded PrefillAdder admission needs prompt + 1 new token + one page per shard < N*T, strictly, so prompts in [N*T - N*ps - 1, N*T - 7] passed intake and bootstrap but were refused with NO_TOKEN on every pass, even on an empty pool. Under FCFS such a request blocked the prefill queue until decode's 300 s KV-wait timeout turned it into a 500. With tp=2, page 64 and 16384 tokens per rank, that is 32639..32761. The bootstrap check now rejects prompts longer than N*T - N*ps - 2 (32638 there) with a 400. Unsharded runs keep their limit. Test: a CPU unit test drives the real sharded PrefillAdder on an empty PageInterleavePoolAllocator and the real bootstrap check at the limit and one past it, for N in {2, 4}. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/disaggregation/prefill.py | 21 ++++++- .../test_disaggregation_kv_cache_sharding.py | 2 +- .../unit/managers/test_prefill_adder.py | 58 +++++++++++++++++++ 3 files changed, 78 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 433401b02188..b3ce0f0be286 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -142,6 +142,18 @@ def _transfer_start_layer(*, pool, hf_text_config) -> int: ) +def max_bootstrap_input_len(*, capacity: int, page_size: int, shard_size: int) -> int: + """Longest prompt the bootstrap queue accepts. + + Sharded PrefillAdder admission needs prompt + 1 new token + one page per + shard < capacity, so a longer prompt never fits even an empty pool. The + unsharded limit stays the plain capacity. + """ + if shard_size == 1: + return capacity + return capacity - shard_size * page_size - 2 + + def maybe_release_metadata_buffer( req: Req, allocator: ReqToMetadataIdxAllocator ) -> None: @@ -196,6 +208,11 @@ def __init__( self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens * self.scheduler.kv_shard_widening ) + self.max_input_len = max_bootstrap_input_len( + capacity=self.max_total_num_tokens, + page_size=self.scheduler.token_to_kv_pool_allocator.page_size, + shard_size=self.scheduler.kv_shard_widening, + ) self.transfer_backend = transfer_backend self.pp_poll_sync_work_list: List = [] if envs.SGLANG_DISAGG_STAGING_BUFFER.get(): @@ -446,8 +463,8 @@ def extend(self, reqs: List[Req], num_kv_heads: int) -> None: self.add(req, num_kv_heads) def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: - if len(req.origin_input_ids) > self.max_total_num_tokens: - message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_total_num_tokens}" + if len(req.origin_input_ids) > self.max_input_len: + message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_input_len}" logger.error(message) req.time_stats.trace_ctx.abort(abort_info={"reason": message}) prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) diff --git a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py index 1dad15a1a0ce..23c1bce671aa 100644 --- a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py +++ b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py @@ -29,7 +29,7 @@ # 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). +# 1.5x one rank's rows, well under the sharded admission bound 2 * (16384 - 64) - 2. 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 diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index ae828ee28409..a069599ee523 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -5,6 +5,10 @@ import torch import sglang.srt.managers.schedule_policy as schedule_policy +from sglang.srt.disaggregation.prefill import ( + PrefillBootstrapQueue, + max_bootstrap_input_len, +) from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import ( AddReqResult, @@ -1099,6 +1103,60 @@ def test_sharded_admission_reserves_nothing_when_delayed(self): self.assertEqual(adder.kv_shard_block_bound_pages, 2) self.assertEqual(adder.kv_shard_chunk_pages, 1) + def test_bootstrap_limit_is_the_sharded_admission_boundary(self): + """The P/D bootstrap check rejects exactly the prompts that sharded + admission can never fit, so none of them waits in the queue forever.""" + page_size, tokens_per_rank = 64, 16384 + # tp=2, page 64, 16384 tokens per rank: 32700 used to pass the check. + self.assertEqual( + max_bootstrap_input_len(capacity=32768, page_size=64, shard_size=2), 32638 + ) + self.assertEqual( + max_bootstrap_input_len(capacity=16384, page_size=64, shard_size=1), 16384 + ) + for shard_size in (2, 4): + capacity = shard_size * tokens_per_rank + queue = PrefillBootstrapQueue.__new__(PrefillBootstrapQueue) + queue.scheduler = SimpleNamespace(output_streamer=MagicMock()) + queue.max_input_len = max_bootstrap_input_len( + capacity=capacity, page_size=page_size, shard_size=shard_size + ) + for prompt_len, fits in ( + (queue.max_input_len, True), + (queue.max_input_len + 1, False), + ): + with self.subTest(shard_size=shard_size, prompt_len=prompt_len): + spec = PageShardSpec( + shard_rank=0, + shard_size=shard_size, + page_size=page_size, + max_prefix_tokens=capacity, + chunk_tokens=8192, + ) + adder = self.create_sharded_adder( + shard_spec=spec, + allocator_size=tokens_per_rank, + rem_input_tokens=capacity, + ) + self.assertEqual(adder.rem_total_tokens, capacity) # empty pool + # PD prefill runs every request with max_new_tokens=1. + req = self.create_sharded_req( + "boundary", prefix_len=0, extend_len=prompt_len + ) + req.origin_input_ids = req.full_untruncated_fill_ids + req.return_logprob = False + req.time_stats.trace_ctx = MagicMock() + + result = adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + self.assertEqual(req in adder.can_run_list, fits) + if not fits: + self.assertEqual(result, AddReqResult.NO_TOKEN) + self.assertEqual( + queue._check_if_req_exceed_kv_capacity(req), not fits + ) + def _build_hybrid_swa_chunked_req( self, *, From cbdcfccc5f35e0854f218fa2b25fb0e716ae1494 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:16:40 -0700 Subject: [PATCH 5/6] [kv-shard 4/4] Say each rank owns a subset of logical KV pages Sharding assigns whole logical pages to ranks; it does not split every page into per-rank stripes. Fix the comments and the E2E docstring that said otherwise. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/disaggregation/common/conn.py | 4 ++-- python/sglang/srt/disaggregation/prefill.py | 2 +- .../disaggregation/test_disaggregation_kv_cache_sharding.py | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 99f6217f4898..ae1bf06aa2b8 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -117,7 +117,7 @@ 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. + # Each prefill shard-group rank owns a subset of logical KV pages. enable_kv_cache_sharding: bool = False # PD true-retraction rebootstrap: the prefill's HTTP API port. The decode @@ -285,7 +285,7 @@ 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. + # shard group. Each rank owns a subset of logical KV pages. kv_shard_rank: int = 0 kv_shard_size: int = 1 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b3ce0f0be286..12d36c9124a0 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -203,7 +203,7 @@ 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. + # Each rank owns a subset of logical KV pages, 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 diff --git a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py index 23c1bce671aa..66bb1f0e849d 100644 --- a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py +++ b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py @@ -1,7 +1,7 @@ """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 +Each prefill TP rank owns a subset of logical KV pages and stores and sends only +those pages; 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. From bd98f73aa4982fe30d57c7a97e36a4f1fe441cf5 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:23:19 -0700 Subject: [PATCH 6/6] [kv-shard 4/4] Bound sharded prompts by the stricter ignore_eos admission path With the radix cache disabled, PrefillAdder admits ignore_eos requests through add_one_req_ignore_eos. That path rounds the prompt up to a page and keeps 1 new token plus IGNORE_EOS_RESERVE_TOKENS spare, so on an empty pool it admits only ceil_page(L) <= N*T - N*ps - 3. The previous limit, N*T - N*ps - 2, still let [N*T - (N+1)*ps + 1, N*T - N*ps - 2] wait forever (32577..32638 for tp=2, page 64, 16384 tokens per rank). The limit is now floor_page(N*T - N*ps - 3), 32576 there. A sweep over N in {2, 4, 8} and page sizes {2, 16, 64} matches the largest prompt both real admission paths admit. The boundary test now also drives the ignore_eos path. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/disaggregation/prefill.py | 11 ++++----- .../test_disaggregation_kv_cache_sharding.py | 2 +- .../unit/managers/test_prefill_adder.py | 23 ++++++++++++------- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 12d36c9124a0..82ef832ac845 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -143,15 +143,12 @@ def _transfer_start_layer(*, pool, hf_text_config) -> int: def max_bootstrap_input_len(*, capacity: int, page_size: int, shard_size: int) -> int: - """Longest prompt the bootstrap queue accepts. - - Sharded PrefillAdder admission needs prompt + 1 new token + one page per - shard < capacity, so a longer prompt never fits even an empty pool. The - unsharded limit stays the plain capacity. - """ + """Bootstrap prompt limit; sharded admission never fits a longer prompt.""" if shard_size == 1: return capacity - return capacity - shard_size * page_size - 2 + # Strictest path, add_one_req_ignore_eos: page-rounded prompt, one page per + # shard, 1 new token and 1 spare token must stay below capacity. + return (capacity - shard_size * page_size - 3) // page_size * page_size def maybe_release_metadata_buffer( diff --git a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py index 66bb1f0e849d..11fc1cf5658a 100644 --- a/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py +++ b/test/registered/e2e/disaggregation/test_disaggregation_kv_cache_sharding.py @@ -29,7 +29,7 @@ # 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) - 2. +# 1.5x one rank's rows, well under the sharded bootstrap limit 2 * 16384 - 3 * 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 diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index a069599ee523..f8f38c1e16ec 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -1104,12 +1104,11 @@ def test_sharded_admission_reserves_nothing_when_delayed(self): self.assertEqual(adder.kv_shard_chunk_pages, 1) def test_bootstrap_limit_is_the_sharded_admission_boundary(self): - """The P/D bootstrap check rejects exactly the prompts that sharded - admission can never fit, so none of them waits in the queue forever.""" + """The P/D bootstrap check rejects every prompt that sharded admission + can never fit, so none of them waits in the queue forever.""" page_size, tokens_per_rank = 64, 16384 - # tp=2, page 64, 16384 tokens per rank: 32700 used to pass the check. self.assertEqual( - max_bootstrap_input_len(capacity=32768, page_size=64, shard_size=2), 32638 + max_bootstrap_input_len(capacity=32768, page_size=64, shard_size=2), 32576 ) self.assertEqual( max_bootstrap_input_len(capacity=16384, page_size=64, shard_size=1), 16384 @@ -1121,11 +1120,18 @@ def test_bootstrap_limit_is_the_sharded_admission_boundary(self): queue.max_input_len = max_bootstrap_input_len( capacity=capacity, page_size=page_size, shard_size=shard_size ) - for prompt_len, fits in ( - (queue.max_input_len, True), - (queue.max_input_len + 1, False), + limit = queue.max_input_len + # Both admission paths fit the limit; the stricter ignore_eos path + # (radix cache disabled) refuses one token more. + for prompt_len, ignore_eos, fits in ( + (limit, False, True), + (limit, True, True), + (limit + 1, True, False), ): - with self.subTest(shard_size=shard_size, prompt_len=prompt_len): + with self.subTest( + shard_size=shard_size, prompt_len=prompt_len, ignore_eos=ignore_eos + ): + self.mock_tree_cache.disable = ignore_eos spec = PageShardSpec( shard_rank=0, shard_size=shard_size, @@ -1144,6 +1150,7 @@ def test_bootstrap_limit_is_the_sharded_admission_boundary(self): "boundary", prefix_len=0, extend_len=prompt_len ) req.origin_input_ids = req.full_untruncated_fill_ids + req.sampling_params.ignore_eos = ignore_eos req.return_logprob = False req.time_stats.trace_ctx = MagicMock()