diff --git a/benchmark/disaggregation/bench_cached_prefix.py b/benchmark/disaggregation/bench_cached_prefix.py new file mode 100644 index 000000000..92062c5de --- /dev/null +++ b/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)) diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index fc93da096..11fb4db24 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index b4196f657..526edf2ac 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/common/dcp_pack.py b/python/sglang/srt/disaggregation/common/dcp_pack.py index 779b2b92d..1d5ae6ab1 100644 --- a/python/sglang/srt/disaggregation/common/dcp_pack.py +++ b/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] diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 753dcdb22..d3456c7bd 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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] diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index ddcf502c5..f46db8422 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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)] diff --git a/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py b/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py index 844f0b319..7ee5b1d94 100644 --- a/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py +++ b/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( diff --git a/test/registered/unit/disaggregation/test_dcp_pack.py b/test/registered/unit/disaggregation/test_dcp_pack.py index cef4295a8..dc5190d37 100644 --- a/test/registered/unit/disaggregation/test_dcp_pack.py +++ b/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( - pack_buffer=buf, - kv_data_ptrs=[kv.data_ptr()], - src_token_indices=src, - token_item_lens=[item_len], - pack_offset_bytes=pack_offset, - ) + packed = mgr._pack_dcp_rank_once(buf, dst, src, {}) + dst.dst_dcp_size = 4 + overflow = mgr._pack_dcp_rank_once(buf, dst, src, {}) + self.assertIsNone(overflow) gather_stream.synchronize.assert_called_once_with() self.assertIsNotNone(packed) ptrs, indices = packed