[PD] Bound cached-prefix DCP transfers by pack capacity (#40376)

benchmark/disaggregation/bench_cached_prefix.py

@@ -0,0 +1,113 @@
+"""Measure PD TTFT with a shared cached prefix and unique request suffixes.
+
+Start a PD router and workers before running this client. Both worker caches are
+flushed after warmup; the shared prefix is then warmed on prefill. Decode radix
+caching should be disabled to exercise transfer of the entire missing KV range.
+Saves per-request responses and client-observed TTFT to the requested JSON file.
+"""
+
+import argparse
+import asyncio
+import json
+import random
+import time
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+
+
+async def main(args):
+    rng = random.Random(35762)
+    prefix = [rng.randrange(1000, 30000) for _ in range(args.prefix)]
+    prompts = [
+        prefix + [rng.randrange(1000, 30000) for _ in range(args.unique)]
+        for _ in range(args.requests)
+    ]
+    timeout = aiohttp.ClientTimeout(total=1800)
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+
+        async def generate(ids, stream=True):
+            start = time.perf_counter()
+            first = None
+            result = None
+            body = {
+                "input_ids": ids,
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": args.output,
+                    "ignore_eos": True,
+                },
+                "stream": stream,
+            }
+            async with session.post(args.url + "/generate", json=body) as response:
+                response.raise_for_status()
+                if stream:
+                    async for line in response.content:
+                        if not line.startswith(b"data:"):
+                            continue
+                        payload = line[5:].strip()
+                        if payload == b"[DONE]":
+                            continue
+                        result = json.loads(payload)
+                        if first is None:
+                            first = time.perf_counter() - start
+                else:
+                    result = await response.json()
+            elapsed = time.perf_counter() - start
+            if result is None or "error" in result:
+                raise RuntimeError(result)
+            return {"ttft_s": first, "elapsed_s": elapsed, "response": result}
+
+        # Warm both cold-prefill and cached-prefix/concurrent shapes. Clear both
+        # sides afterwards so measured requests never reuse a unique suffix.
+        await generate(prompts[0])
+        await asyncio.gather(*(generate(p) for p in prompts[1 : args.concurrency + 1]))
+        for url in args.workers:
+            async with session.post(url + "/flush_cache", params={"timeout": 30}) as r:
+                r.raise_for_status()
+        if prefix:
+            await generate(prefix)
+        sem = asyncio.Semaphore(args.concurrency)
+
+        async def run(ids):
+            async with sem:
+                return await generate(ids)
+
+        start = time.perf_counter()
+        records = await asyncio.gather(*(run(ids) for ids in prompts))
+        elapsed = time.perf_counter() - start
+    ttfts = [r["ttft_s"] for r in records]
+    summary = {
+        "args": vars(args),
+        "elapsed_s": elapsed,
+        "requests_per_s": len(records) / elapsed,
+        "ttft_mean_ms": float(np.mean(ttfts) * 1000),
+        "ttft_p50_ms": float(np.percentile(ttfts, 50) * 1000),
+        "ttft_p99_ms": float(np.percentile(ttfts, 99) * 1000),
+        "records": records,
+    }
+    Path(args.result).write_text(json.dumps(summary, indent=2))
+    print(json.dumps({k: v for k, v in summary.items() if k != "records"}), flush=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", default="http://127.0.0.1:30000")
+    parser.add_argument(
+        "--workers",
+        nargs="+",
+        default=["http://127.0.0.1:30001", "http://127.0.0.1:30002"],
+    )
+    parser.add_argument("--prefix", type=int, default=83264)
+    parser.add_argument("--unique", type=int, default=6720)
+    parser.add_argument("--output", type=int, default=1)
+    parser.add_argument("--requests", type=int, default=16)
+    parser.add_argument("--concurrency", type=int, default=1)
+    parser.add_argument("--result", required=True)
+    args = parser.parse_args()
+    if args.prefix < 0 or args.unique <= 0:
+        parser.error("--prefix must be nonnegative and --unique must be positive")
+    if min(args.output, args.requests, args.concurrency) <= 0:
+        parser.error("--output, --requests and --concurrency must be positive")
+    asyncio.run(main(args))

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

@@ -171,6 +171,10 @@ class BaseKVSender(ABC):
     def pop_decode_prefix_len(self) -> int:
         return 0
 
+    def get_max_transfer_tokens(self) -> Optional[int]:
+        """Optional page-aligned limit for one scheduler KV send."""
+        return None
+
     def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
         return num_pages > 0
 

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

@@ -35,9 +35,12 @@ from sglang.srt.environ import envs
 from sglang.srt.runtime_context import (
     get_disagg,
     get_parallel,
+    get_schedule,
     get_serving,
+    max_prefill_buffer_tokens,
 )
 from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.common import ceil_align
 from sglang.srt.utils.network import (
     NetworkAddress,
     get_local_ip_auto,
@@ -181,6 +184,7 @@ class CommonKVManager(BaseKVManager):
             envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get()
         )
         self._dcp_pack_buffers = None
+        self._dcp_pack_max_tokens: Optional[int] = None
         # for p/d multi node infer
         self.bootstrap_host = get_serving().host
         self.bootstrap_port = get_disagg().disaggregation_bootstrap_port
@@ -385,12 +389,16 @@ class CommonKVManager(BaseKVManager):
             return
         from sglang.srt.disaggregation.common.dcp_pack import init_dcp_pack_buffers
 
+        max_tokens = max_prefill_buffer_tokens() or get_schedule().max_prefill_tokens
+        max_tokens = ceil_align(max_tokens, self.kv_args.page_size)
         self._dcp_pack_buffers = init_dcp_pack_buffers(
             self._register_staging_memory,
             self.kv_args,
             len(self.transfer_queues),
             dcp_size,
+            max_tokens,
         )
+        self._dcp_pack_max_tokens = max_tokens
 
     def check_status(self, bootstrap_room: int) -> KVPoll:
         return self.request_status[bootstrap_room]
@@ -1541,6 +1549,19 @@ class CommonKVSender(BaseKVSender):
     def pop_decode_prefix_len(self) -> int:
         return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
 
+    def get_max_transfer_tokens(self) -> Optional[int]:
+        if self.kv_mgr._dcp_pack_max_tokens is None:
+            return None
+        for peer, info in self.kv_mgr.transfer_infos.get(
+            self.bootstrap_room, {}
+        ).items():
+            if (
+                not info.is_dummy
+                and self.kv_mgr.decode_kv_args_table[peer].requires_dcp_relayout
+            ):
+                return self.kv_mgr._dcp_pack_max_tokens
+        return None
+
     def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
         return num_pages > 0 or last_chunk
 

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

@@ -9,7 +9,6 @@ import torch
 
 from sglang.kernels.ops.kvcache.pd_dcp_gather import copy_mla_rows_into_pack
 from sglang.srt.disaggregation.common.staging_buffer import StagingBuffer
-from sglang.srt.runtime_context import get_schedule, max_prefill_buffer_tokens
 
 logger = logging.getLogger(__name__)
 
@@ -43,6 +42,7 @@ def try_pack_dcp_src(
     src_token_indices: npt.NDArray[np.integer],
     token_item_lens: Sequence[int],
     pack_offset_bytes: int = 0,
+    pack_capacity_bytes: Optional[int] = None,
 ) -> Optional[Tuple[List[int], npt.NDArray[np.int64]]]:
     if pack_offset_bytes < 0:
         raise ValueError(
@@ -54,13 +54,17 @@ def try_pack_dcp_src(
         return [], empty
     required = n * sum(int(item_len) for item_len in token_item_lens)
     required_end = pack_offset_bytes + required
-    if not pack_buffer.fits(required_end):
+    if (
+        pack_capacity_bytes is not None and required > pack_capacity_bytes
+    ) or not pack_buffer.fits(required_end):
         logger.warning(
-            "PD DCP pack buffer too small for byte range [%s, %s) (have %s); "
+            "PD DCP pack buffer too small for byte range [%s, %s) "
+            "(have %s, region capacity %s); "
             "falling back to per-token RDMA",
             pack_offset_bytes,
             required_end,
             pack_buffer.get_size(),
+            pack_capacity_bytes,
         )
         return None
 
@@ -88,14 +92,12 @@ def init_dcp_pack_buffers(
     kv_args,
     count: int,
     dcp_size: int,
+    max_tokens: int,
 ) -> List[StagingBuffer]:
     from sglang.srt.disaggregation.common.staging_handler import (
         _get_custom_mem_pool,
     )
 
-    max_tokens = max_prefill_buffer_tokens()
-    if max_tokens <= 0:
-        max_tokens = get_schedule().max_prefill_tokens
     kv_item_lens = kv_args.kv_item_lens
     if kv_args.num_draft_entries > 0:
         kv_item_lens = kv_item_lens[: len(kv_item_lens) - kv_args.num_draft_entries]

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

@@ -1867,6 +1867,7 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
             src_token_indices=src_token_indices,
             token_item_lens=token_item_lens[:num_target],
             pack_offset_bytes=rank * rank_stride,
+            pack_capacity_bytes=rank_stride,
         )
         return packed_source_by_dcp_rank[rank]
 

python/sglang/srt/disaggregation/prefill.py

@@ -1449,15 +1449,21 @@ class SchedulerDisaggregationPrefillMixin:
                 payloads[st]() if st in payloads else None for st in state_types
             ]
 
+        transfer_chunk_tokens = req.disagg_kv_sender.get_max_transfer_tokens()
         if self.enable_staging:
             # One sender.send per grid slot; the sender's cumulative page
             # counter marks only the final sub-send of the final chunk as
             # is_last, routing aux/state correctly.
+            transfer_chunk_tokens = staging_grid_tokens(
+                get_schedule().chunked_prefill_size, page_size
+            )
+        if transfer_chunk_tokens is not None:
+            # Prefill cache hits can leave more KV to transfer than the DCP pack buffer holds.
             segments = compute_grid_segments(
                 start_idx,
                 end_idx,
                 req.disagg_decode_prefix_len,
-                staging_grid_tokens(get_schedule().chunked_prefill_size, page_size),
+                transfer_chunk_tokens,
             )
         else:
             segments = [(start_idx, end_idx)]

test/registered/disaggregation/test_kimi_linear_pd_dcp4.py

@@ -464,6 +464,29 @@ class TestKimiLinearPDDCP4(GSM8KMixin, PDDisaggregationServerBase):
                         f"niah prompt_tokens={LONG_CONTEXT_TOKENS} depth={needle_depth}"
                     ),
                 )
+                # Prefill now holds the long prefix. Decode must receive it
+                # again, even though prefill computes almost no new tokens.
+                # OSL > 1 checks that decode actually reads the transferred KV.
+                self._flush_cache(self.decode_url)
+                cached_actual = self._generate(
+                    self.base_url,
+                    prompt,
+                    max_new_tokens=16,
+                    ignore_eos=False,
+                )
+                self.assertGreater(
+                    cached_actual["meta_info"]["cached_tokens"],
+                    CHUNKED_PREFILL_SIZE,
+                )
+                self.assertIn(NIAH_KEY, cached_actual["text"])
+                self._assert_output_parity(
+                    reference,
+                    cached_actual,
+                    label=(
+                        f"cached niah prompt_tokens={LONG_CONTEXT_TOKENS} "
+                        f"depth={needle_depth}"
+                    ),
+                )
 
     def _assert_batch_completes(self, batch_size: int):
         response = requests.post(

test/registered/unit/disaggregation/test_dcp_pack.py

@@ -6,15 +6,18 @@ from unittest.mock import Mock, patch
 import numpy as np
 import torch
 
+from sglang.srt.disaggregation.base.conn import StateType
 from sglang.srt.disaggregation.common.conn import CommonKVManager
 from sglang.srt.disaggregation.common.dcp_pack import (
     dcp_pack_buffer_bytes,
-    try_pack_dcp_src,
 )
 from sglang.srt.disaggregation.common.utils import (
     build_dcp_token_transfer_plan,
     group_concurrent_contiguous,
 )
+from sglang.srt.disaggregation.nixl.conn import NixlKVManager, NixlKVSender
+from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
+from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -187,6 +190,111 @@ class TestPrepareDcpTokenItemLens(CustomTestCase):
             )
 
 
+class TestDcpCachedPrefixSend(CustomTestCase):
+    def test_cached_prefix_fits_pack_capacity_and_preserves_pages_and_state(self):
+        """Bound DCP sends by the allocation without splitting TP sends on the same worker."""
+        total, page_size, token_bytes = 2055, 64, 16
+        mgr = object.__new__(NixlKVManager)
+        mgr.kv_args = SimpleNamespace(
+            kv_item_lens=[page_size * token_bytes],
+            num_draft_entries=0,
+            page_size=page_size,
+            gpu_id=0,
+            state_types=[StateType.MAMBA],
+        )
+        mgr._dcp_pack_buffers = None
+        mgr._dcp_pack_max_tokens = None
+        mgr.transfer_queues = [None]
+        mgr._register_staging_memory = Mock()
+        mgr.request_status = {}
+        mgr.is_dummy_cp_rank = False
+        mgr.enable_all_cp_ranks_for_transfer = False
+        mgr.decode_kv_args_table = {
+            peer: SimpleNamespace(requires_dcp_relayout=relayout)
+            for peer, relayout in (("dcp", True), ("tp", False))
+        }
+
+        def allocate(size, *args, **kwargs):
+            return SimpleNamespace(get_ptr=lambda: 0x1000, get_size=lambda: size)
+
+        with (
+            get_context().override_server_args(chunked_prefill_size=250),
+            patch(
+                "sglang.srt.disaggregation.common.staging_handler._get_custom_mem_pool",
+                return_value=(None, None),
+            ),
+            patch(
+                "sglang.srt.disaggregation.common.dcp_pack.StagingBuffer",
+                side_effect=allocate,
+            ),
+        ):
+            mgr._init_dcp_pack_buffers_once(dcp_size=4)
+        limit = mgr._dcp_pack_buffers[0].get_size() // token_bytes
+
+        for peer, prefix in (("dcp", 0), ("dcp", 256), ("tp", 0), ("tp", 256)):
+            with self.subTest(peer=peer, decode_prefix=prefix):
+                mgr.transfer_infos = {
+                    1: {
+                        "dummy": SimpleNamespace(is_dummy=True),
+                        peer: SimpleNamespace(is_dummy=False),
+                    }
+                }
+                mgr.add_transfer_request = Mock()
+                with get_context().override_server_args(dp_size=1):
+                    sender = NixlKVSender(mgr, "unused", 1, [0], 0)
+                sender.init((total - prefix + page_size - 1) // page_size, 3)
+                req = SimpleNamespace(
+                    rid="cached-prefix",
+                    kv=SimpleNamespace(req_pool_idx=0),
+                    origin_input_ids=[0] * total,
+                    extend_range=SimpleNamespace(end=total),
+                    start_send_idx=prefix,
+                    disagg_decode_prefix_len=prefix,
+                    disagg_kv_sender=sender,
+                )
+                scheduler = SimpleNamespace(
+                    enable_staging=False,
+                    token_to_kv_pool_allocator=SimpleNamespace(
+                        page_size=page_size,
+                        translate_kv_indices_for_transfer=lambda x: x,
+                    ),
+                    req_to_token_pool=SimpleNamespace(
+                        req_to_token=torch.arange(total).reshape(1, -1),
+                        req_index_to_mamba_index_mapping=torch.tensor([17]),
+                        translate_mamba_indices=lambda x: x,
+                    ),
+                    disagg_metadata_buffers=Mock(),
+                    disagg_prefill_bootstrap_queue=SimpleNamespace(kv_manager=mgr),
+                    disagg_prefill_pending_chunk_rids=set(),
+                )
+                SchedulerDisaggregationPrefillMixin._send_kv_chunk(
+                    scheduler, req, last_chunk=True
+                )
+                calls = mgr.add_transfer_request.call_args_list
+                token_counts = [c.args[7] for c in calls]
+                if peer == "dcp":
+                    self.assertLessEqual(max(token_counts), limit)
+                else:
+                    self.assertEqual(len(calls), 1)
+                self.assertEqual(sum(token_counts), total - prefix)
+                np.testing.assert_array_equal(
+                    np.concatenate([c.args[1] for c in calls]),
+                    np.arange(prefix // 64, 33),
+                )
+                page_offset = 0
+                for call in calls:
+                    pages = len(call.args[1])
+                    self.assertEqual(
+                        call.args[2], slice(page_offset, page_offset + pages)
+                    )
+                    page_offset += pages
+                self.assertEqual(
+                    [c.args[3] for c in calls], [False] * (len(calls) - 1) + [True]
+                )
+                self.assertTrue(all(c.args[6] is None for c in calls[:-1]))
+                self.assertEqual(int(calls[-1].args[6][0][0]), 17)
+
+
 class TestDcpPackBufferBytes(CustomTestCase):
     def test_sizes_fixed_regions_for_each_dcp_rank(self):
         self.assertEqual(
@@ -208,6 +316,7 @@ class TestDcpPackBufferBytes(CustomTestCase):
 
 class TestTryDcpPack(CustomTestCase):
     def test_try_pack_uses_requested_region_and_dense_indices(self):
+        """A gather must fit its rank region even when the total buffer has space."""
         dim = 4
         kv = torch.arange(16 * dim, dtype=torch.float32).view(16, 1, dim)
         item_len = int(kv[0].nbytes)
@@ -225,7 +334,12 @@ class TestTryDcpPack(CustomTestCase):
             },
         )()
         src = np.array([1, 5, 9, 13], dtype=np.int64)
-        pack_offset = 2 * item_len
+        pack_offset = 4 * item_len
+        mgr = object.__new__(NixlKVManager)
+        mgr.kv_args = SimpleNamespace(kv_data_ptrs=[kv.data_ptr()], num_draft_entries=0)
+        dst = SimpleNamespace(
+            dst_dcp_rank=1, dst_dcp_size=2, dcp_token_item_lens=[item_len]
+        )
         with (
             patch(
                 "sglang.srt.disaggregation.common.dcp_pack.torch.cuda.default_stream"
@@ -238,14 +352,11 @@ class TestTryDcpPack(CustomTestCase):
                 "sglang.srt.disaggregation.common.dcp_pack.copy_mla_rows_into_pack"
             ) as copy_mock,
         ):
-            packed = try_pack_dcp_src(
+            packed = mgr._pack_dcp_rank_once(buf, dst, src, {})
-                pack_buffer=buf,
+            dst.dst_dcp_size = 4
-                kv_data_ptrs=[kv.data_ptr()],
+            overflow = mgr._pack_dcp_rank_once(buf, dst, src, {})
-                src_token_indices=src,
-                token_item_lens=[item_len],
-                pack_offset_bytes=pack_offset,
-            )
 
+        self.assertIsNone(overflow)
         gather_stream.synchronize.assert_called_once_with()
         self.assertIsNotNone(packed)
         ptrs, indices = packed