diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 1af1f61af1eb..ae1bf06aa2b8 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, @@ -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 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 # 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 + # shard group. Each rank owns a subset of logical KV pages. + 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, ) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 369a016a8a61..82ef832ac845 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -142,6 +142,15 @@ def _transfer_start_layer(*, pool, hf_text_config) -> int: ) +def max_bootstrap_input_len(*, capacity: int, page_size: int, shard_size: int) -> int: + """Bootstrap prompt limit; sharded admission never fits a longer prompt.""" + if shard_size == 1: + return capacity + # 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( req: Req, allocator: ReqToMetadataIdxAllocator ) -> None: @@ -191,8 +200,15 @@ def __init__( self.queue: List[Req] = [] self.scheduler = scheduler self.scheduler_stage_metrics = scheduler_stage_metrics + # 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 + ) + 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 = [] @@ -444,8 +460,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/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index f428081b492b..2ba665ef8ace 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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, 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..11fc1cf5658a --- /dev/null +++ b/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 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. +""" + +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 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 +# 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() diff --git a/test/registered/unit/disaggregation/test_register_to_bootstrap.py b/test/registered/unit/disaggregation/test_register_to_bootstrap.py index 4e01acfbbf96..aad8a143c411 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/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index ae828ee28409..f8f38c1e16ec 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,67 @@ 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 every prompt that sharded admission + can never fit, so none of them waits in the queue forever.""" + page_size, tokens_per_rank = 64, 16384 + self.assertEqual( + 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 + ) + 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 + ) + 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, ignore_eos=ignore_eos + ): + self.mock_tree_cache.disable = ignore_eos + 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.sampling_params.ignore_eos = ignore_eos + 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, *, 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 a6d9809de2b0..3849583f310b 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -31,6 +31,10 @@ 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. +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. The GPU tests exercise real pools and NCCL collectives: MLA and MHA gathers on 2 ranks (``TestPageInterleaveGatherMultiGpu``), and byte-exact DSA indexer @@ -49,10 +53,13 @@ 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 import torch.multiprocessing as mp 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.distributed import ( init_distributed_environment, initialize_model_parallel, @@ -72,7 +79,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, @@ -93,7 +104,11 @@ class draws (owners exactly cyclic along a chain), least-full root from sglang.srt.utils import ceil_div from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci from sglang.test.mem_cache_utils import finish_req -from sglang.test.test_utils import CustomTestCase, publish_build_topology +from sglang.test.test_utils import ( + CustomTestCase, + enter_scope, + publish_build_topology, +) register_cpu_ci(est_time=30, suite="base-a-test-cpu") register_cuda_ci(est_time=120, stage="base-c", runner_config="4-gpu-b200") @@ -1346,6 +1361,129 @@ 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.""" + + 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) + + # ============================================================================= # Multi-GPU: the real NCCL layer-ahead gather (2 GPUs). #