From 3aece59716ddc9ffd1a0e6dc3ee3e7bc70739504 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:28:32 +0000 Subject: [PATCH 01/10] Fix PD retraction checkpoint boundaries with overlap decoding --- python/sglang/srt/managers/scheduler.py | 25 ++- .../test_pd_retraction_overlap_boundary.py | 154 ++++++++++++++++++ 2 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 test/registered/unit/managers/test_pd_retraction_overlap_boundary.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c134c098ebd5..b8b408e45582 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -4173,10 +4173,29 @@ def update_running_batch(self, batch: ScheduleBatch) -> Optional[ScheduleBatch]: batch.batch_is_full = False return batch - # Check if decode out of memory - if (kv_full_retract_flag := not batch.check_decode_mem()) or ( - TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0 + # PD offload must checkpoint CPU history at the GPU's boundary; + # settle the in-flight result first. + kv_full_retract_flag = not batch.check_decode_mem() + test_retract = TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0 + if ( + (kv_full_retract_flag or test_retract) + and self.disaggregation_mode == DisaggregationMode.DECODE + and self.enable_overlap + and self.result_queue ): + while self.result_queue: + pending_batch, pending_result = self.result_queue.popleft() + self.process_batch_result(pending_batch, pending_result) + # The overlap loop must not process this result a second time. + self.last_batch = None + batch.filter_batch() + if batch.is_empty(): + batch.batch_is_full = False + return batch + # Completed requests may have released enough memory already. + kv_full_retract_flag = not batch.check_decode_mem() + + if kv_full_retract_flag or test_retract: old_available_tokens = self.token_to_kv_pool_allocator.available_size() old_ratio = self.new_token_ratio_tracker.current mamba_allocator = getattr( diff --git a/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py b/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py new file mode 100644 index 000000000000..5784df2fe8a7 --- /dev/null +++ b/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py @@ -0,0 +1,154 @@ +"""PD offload must save the same boundary in CPU metadata and GPU state.""" + +import unittest +from collections import deque +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +import sglang.srt.managers.scheduler as scheduler_module +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class Batch: + def __init__(self, events, enough): + self.events = events + self.enough = enough + self.reqs = [ + SimpleNamespace( + output_ids=list(range(100)), is_retracted=False, finished=False + ) + for _ in range(2) + ] + self.batch_is_full = True + self.snapshots = [] + + def batch_size(self): + return len(self.reqs) + + def filter_batch(self): + self.reqs = [req for req in self.reqs if not req.finished] + + def is_empty(self): + return not self.reqs + + def check_decode_mem(self): + return self.enough + + def retract_decode(self): + req = self.reqs.pop() + self.snapshots.append(list(req.output_ids)) + self.events.append("offload") + req.is_retracted = True + return [req], 0.5, [] + + def prepare_for_decode(self): + self.events.append("prepare") + + +class RetractionOverlapBoundaryTest(CustomTestCase): + def setup_case( + self, + *, + forced=True, + enough=True, + mode=DisaggregationMode.DECODE, + overlap=True, + completed=0, + free_on_result=False, + ): + self.enterContext(patch.object(scheduler_module, "TEST_RETRACT", forced)) + self.enterContext(patch.object(scheduler_module, "TEST_RETRACT_INTERVAL", 16)) + events = [] + batch = Batch(events, enough) + pending = list(batch.reqs) + queue = deque([(pending, [100, 101, 102, 103])]) + + def process_batch_result(reqs, accepted): + events.append("settle") + for index, req in enumerate(reqs): + # Matches output processing's intentional retracted-request guard. + if not req.is_retracted: + req.output_ids.extend(accepted) + req.finished = index < completed + if free_on_result: + batch.enough = True + + # A real instance, so helper methods update_running_batch calls resolve. + scheduler = Scheduler.__new__(Scheduler) + vars(scheduler).update( + enable_hierarchical_cache=False, + enable_overlap=overlap, + disaggregation_mode=mode, + forward_ct=16, + result_queue=queue, + last_batch=pending, + process_batch_result=process_batch_result, + token_to_kv_pool_allocator=SimpleNamespace(available_size=lambda: 1024), + new_token_ratio_tracker=SimpleNamespace( + current=0.5, decay_step=lambda: events.append("decay") + ), + tree_cache=SimpleNamespace(req_to_token_pool=SimpleNamespace()), + metrics_reporter=SimpleNamespace(enable_metrics=False), + server_args=SimpleNamespace(), + _add_request_to_queue=lambda req, is_retracted: events.append("enqueue"), + ) + return scheduler, batch, events + + def test_pending_accepted_block_is_settled_before_pd_snapshot(self): + for forced, enough in ((True, True), (False, False)): + with self.subTest(forced=forced): + scheduler, batch, events = self.setup_case(forced=forced, enough=enough) + Scheduler.update_running_batch(scheduler, batch) + self.assertEqual(batch.snapshots, [list(range(104))]) + self.assertLess(events.index("settle"), events.index("offload")) + self.assertFalse(scheduler.result_queue) + self.assertIsNone(scheduler.last_batch) + + def test_settlement_filters_completions_and_rechecks_capacity(self): + for completed in (0, 1, 2): + with self.subTest(completed=completed): + scheduler, batch, events = self.setup_case( + forced=False, enough=False, completed=completed, free_on_result=True + ) + self.assertIs(Scheduler.update_running_batch(scheduler, batch), batch) + self.assertEqual(batch.batch_size(), 2 - completed) + self.assertEqual( + [len(req.output_ids) for req in batch.reqs], [104] * (2 - completed) + ) + self.assertEqual(batch.snapshots, []) + self.assertEqual( + events, + ["settle"] if completed == 2 else ["settle", "decay", "prepare"], + ) + self.assertIsNone(scheduler.last_batch) + if completed: + self.assertFalse(batch.batch_is_full) + + def test_other_steps_preserve_pending_results(self): + for mode, overlap, forced in ( + (DisaggregationMode.DECODE, True, False), + (DisaggregationMode.NULL, True, True), + (DisaggregationMode.DECODE, False, True), + ): + with self.subTest(mode=mode, overlap=overlap, forced=forced): + scheduler, batch, events = self.setup_case( + mode=mode, overlap=overlap, forced=forced + ) + Scheduler.update_running_batch(scheduler, batch) + self.assertNotIn("settle", events) + self.assertEqual(len(scheduler.result_queue), 1) + self.assertIsNotNone(scheduler.last_batch) + if not forced: + self.assertEqual(events, ["decay", "prepare"]) + + +if __name__ == "__main__": + unittest.main() From 746cf63cd60cd28749232b24eca47d27eeeca3dd Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:31:32 +0000 Subject: [PATCH 02/10] Use the 3.10-safe enter_scope helper in the PD retraction test Co-Authored-By: Claude Opus 5.5 --- .../unit/managers/test_pd_retraction_overlap_boundary.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py b/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py index 5784df2fe8a7..20cf6ca033e7 100644 --- a/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py +++ b/test/registered/unit/managers/test_pd_retraction_overlap_boundary.py @@ -6,7 +6,7 @@ from unittest.mock import patch from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel +from sglang.test.test_utils import CustomTestCase, enter_scope, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() @@ -64,8 +64,8 @@ def setup_case( completed=0, free_on_result=False, ): - self.enterContext(patch.object(scheduler_module, "TEST_RETRACT", forced)) - self.enterContext(patch.object(scheduler_module, "TEST_RETRACT_INTERVAL", 16)) + enter_scope(self, patch.object(scheduler_module, "TEST_RETRACT", forced)) + enter_scope(self, patch.object(scheduler_module, "TEST_RETRACT_INTERVAL", 16)) events = [] batch = Batch(events, enough) pending = list(batch.reqs) From f58a7d1a8956a27893e0eefed244a6e3c0fd3593 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 23:03:45 +0000 Subject: [PATCH 03/10] [Spec] Support DFLASH for Kimi K3 K3 exposed only set_dspark_layers_to_capture, so configure_aux_hidden_state_capture rejected --speculative-algorithm DFLASH at startup. K3's DSPARK taps already capture layer outputs, which is what DFLASH target_layer_ids name, so the new set_dflash_layers_to_capture hooks on the LM and the multimodal wrapper reuse them without the +1 shift that layer-input targets apply. The Kimi K3 cookbook enables the DFLASH option with the public modal-labs/Kimi-K3-DFlash draft at block size 8, gated off where the recipe is pipelined, runs DP attention, or is an NPU recipe, and on Hopper and AMD until it has been validated there. Co-Authored-By: Claude Opus 5.5 --- .../autoregressive/Moonshotai/Kimi-K3.mdx | 2 +- .../snippets/configs/moonshotai/kimi-k3.jsx | 32 +++++--- python/sglang/srt/models/kimi_k3.py | 12 +++ .../test_attention_backend_setup.py | 81 ++++++++++++++++++- 4 files changed, 115 insertions(+), 12 deletions(-) diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index 6025fa67f8b8..2401f4f625d4 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -70,7 +70,7 @@ Pick your hardware, then the deployment shape and operating point. Node count fo `Long-Context` appears only under the `Prefill` PD mode; for long-context unified serving on B200, start from High-Throughput and raise `--context-length`. -**Spec Decode** — layers onto the strategy without changing it, on every platform except B200. DSPARK proposes 7 draft tokens per step (tune in the Playground) and requires `pp_size == 1`, so on B200 it also drops the pipeline and re-lays the same 16 GPUs flat: PP2 × TP8 → TP16, PP2 × DCPEP8 → DCPEP16. DFLASH has no published draft checkpoint. The win is largest on short interactive traffic and fades as the prompt grows. +**Spec Decode** — layers onto the strategy without changing it, on every platform except B200. DSPARK proposes 7 draft tokens per step (tune in the Playground) and requires `pp_size == 1`, so on B200 it also drops the pipeline and re-lays the same 16 GPUs flat: PP2 × TP8 → TP16, PP2 × DCPEP8 → DCPEP16. DFLASH uses the [Kimi-K3-DFlash](https://huggingface.co/modal-labs/Kimi-K3-DFlash) draft at block size 8; it also requires `pp_size == 1` and no DP attention, but does not re-lay pipelined recipes, so it is unavailable on the pipelined B200 Unified recipes. It has only been validated on Blackwell, so the Hopper and AMD recipes don't offer it yet. The win is largest on short interactive traffic and fades as the prompt grows. `--mamba-full-memory-ratio` is the one sizing flag, computed live: set your average request length in the [Mamba ratio calculator](#mamba-ratio-calculator); everything else follows the panels, and the result is pinned into the command. (The Ascend NPU recipes don't use the calculator and set no ratio at all — the Ascend path sizes both pools itself: the A3 Series pins `--max-mamba-cache-size` instead, and the A5 recipe serves with the radix cache off.) diff --git a/docs/src/snippets/configs/moonshotai/kimi-k3.jsx b/docs/src/snippets/configs/moonshotai/kimi-k3.jsx index 00daa4ce3b6f..a322724a4d75 100644 --- a/docs/src/snippets/configs/moonshotai/kimi-k3.jsx +++ b/docs/src/snippets/configs/moonshotai/kimi-k3.jsx @@ -334,16 +334,28 @@ export const config = { { id: "dflash", label: "DFLASH", - // Listed so the axis is complete, but not selectable: no K3 DFLASH draft - // checkpoint has been published, so there is nothing to point - // --speculative-draft-model-path at. DFLASH is also CUDA-only, rejects DP - // attention, and requires pp_size == 1. - disabled: true, - disableReason: - "No K3 DFLASH draft checkpoint published yet — DSPARK is the available speculative path.", - flags: [ + // DFLASH requires pp_size == 1 and rejects DP attention off NPU, and it + // does not re-lay a pipelined cell the way DSPARK does, so those + // recipes are unavailable. The NPU recipes ship DSPARK only. K3 DFLASH + // has only been validated on Blackwell, so Hopper and AMD stay off. + disabled: (s) => + config.isNpuHw(s) || + ["h100", "h200", "mi350x", "mi355x"].includes(s.hw) || + config.isPipelined(s) || + !!config.flagOf(config.cellFor(s), "--enable-dp-attention"), + disableReason: (s) => + config.isNpuHw(s) + ? "Only DSPARK is supported on this recipe." + : ["h100", "h200", "mi350x", "mi355x"].includes(s.hw) + ? "K3 DFLASH has not been validated on this hardware yet; use DSPARK." + : "DFLASH requires pp_size == 1 and no DP attention. Pick a recipe that runs a single pipeline stage without DP attention, or use DSPARK.", + flags: (s) => [ "--speculative-algorithm DFLASH", - "--speculative-draft-model-path ", + "--speculative-draft-model-path modal-labs/Kimi-K3-DFlash", + // 8 is the recommended default block size for this draft. + "--speculative-dflash-block-size 8", + // Same ReplaySSM rule as DSPARK: the PD prefill role rejects the flag. + ...(s.pdMode !== "prefill" ? ["--enable-linear-replayssm-spec"] : []), ], }, ], @@ -785,7 +797,7 @@ export const config = { // DSPARK --speculative-dspark-block-size N (gamma, == proposed) // DFLASH --speculative-dflash-block-size N+1 (verify window) // EAGLE --speculative-num-steps N (chain; topk>1 is a tree) - // Only DSPARK is selectable today, so only its form is emitted. + // Only DSPARK's form is emitted; the DFLASH option pins its block size. id: "proposedDraftTokens", title: "Proposed Draft Tokens", // The NPU recipes pin the shipped block size (7). showWhen: (b) => b.spec === "dspark" && !config.isNpuHw(b), diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 383a6f47812b..5341f7394a77 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -3150,6 +3150,11 @@ def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: self.capture_aux_hidden_states = True self.model.dspark_layers_to_capture = list(layer_ids) + def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None: + # DFLASH target_layer_ids name layer outputs, which is what the DSPARK + # taps already capture here, so reuse them without the usual +1 shift. + self.set_dspark_layers_to_capture(layer_ids) + @torch.no_grad() def forward( self, @@ -3606,6 +3611,13 @@ def get_pp_proxy_dspark_hidden_size(self) -> int: return 0 return self.language_model.get_pp_proxy_dspark_hidden_size() + def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None: + if self.language_model is None: + raise AttributeError( + "DFLASH layer capture is not available in encoder-only mode" + ) + self.language_model.set_dflash_layers_to_capture(layer_ids) + def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: if self.language_model is None: raise AttributeError( diff --git a/test/registered/unit/model_executor/model_runner_components/test_attention_backend_setup.py b/test/registered/unit/model_executor/model_runner_components/test_attention_backend_setup.py index 426145ea36c1..b2f2c4d7be64 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_attention_backend_setup.py +++ b/test/registered/unit/model_executor/model_runner_components/test_attention_backend_setup.py @@ -1,5 +1,5 @@ import sys -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import patch import pytest @@ -118,5 +118,84 @@ def test_equal_resolved_backends_ignore_stale_global_backend(): assert result.name == "resolved" +def test_kimi_k3_dflash_captures_the_dspark_post_layer_taps(): + """DFLASH target_layer_ids name layer outputs. K3 taps already capture layer + outputs, so DFLASH must reach them unshifted (targets that tap layer inputs + add 1) and land on the same states as DSPARK. K3 used to expose only the + DSPARK hook, so DFLASH was rejected here at startup.""" + import torch + + from sglang.srt.layers.aux_hidden_states import pack_aux_hidden_states + from sglang.srt.models import kimi_k3 + from sglang.srt.models.kimi_k3 import ( + KimiK3ForConditionalGeneration, + KimiK3LinearForCausalLM, + KimiK3LinearModel, + ) + + num_tokens, hidden_size, layer_ids = 3, 4, [1, 3] + pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True, world_size=1) + + def make_layer(idx): + def layer(*, hidden_states, residual, **_): + return hidden_states + (idx + 1), residual, False + + return layer + + captured = {} + for is_dspark in (False, True): + model = SimpleNamespace( + config=SimpleNamespace(attn_res_block_size=None), + pp_group=pp_group, + start_layer=0, + end_layer=4, + layers=[make_layer(i) for i in range(4)], + norm=lambda hidden_states: hidden_states, + dspark_layers_to_capture=None, + _trim_padded_attn=False, + ) + model._dspark_capture_stream = MethodType( + KimiK3LinearModel._dspark_capture_stream, model + ) + lm = KimiK3LinearForCausalLM.__new__(KimiK3LinearForCausalLM) + torch.nn.Module.__init__(lm) + lm.model = model + lm.capture_aux_hidden_states = False + # The checkpoint architecture is the multimodal wrapper around the LM. + target = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration) + torch.nn.Module.__init__(target) + target.language_model = lm + + attention_backend_setup.configure_aux_hidden_state_capture( + model=target, + eagle_use_aux_hidden_state=False, + eagle_aux_hidden_state_layer_ids=None, + dflash_use_aux_hidden_state=True, + dflash_target_layer_ids=list(layer_ids), + is_dspark=is_dspark, + ) + assert lm.capture_aux_hidden_states + + with patch.object( + kimi_k3, "get_parallel", return_value=SimpleNamespace(pp_group=pp_group) + ): + _, aux_hidden_states = KimiK3LinearModel.forward( + model, + None, + torch.arange(num_tokens), + SimpleNamespace(), + inputs_embeds=torch.zeros(num_tokens, hidden_size), + ) + captured[is_dspark] = pack_aux_hidden_states(aux_hidden_states) + + # Layer i adds i + 1, so the outputs of layers 1 and 3 hold 3 and 10; + # the draft reads them packed in target_layer_ids order. + expected = torch.cat( + [torch.full((num_tokens, hidden_size), v) for v in (3.0, 10.0)], dim=-1 + ) + for packed in captured.values(): + torch.testing.assert_close(packed, expected) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) From bc3a82a3e4d9888e98f518976a23041555afaa1c Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:28:34 +0000 Subject: [PATCH 04/10] [DCP] Use logical token capacity for PD admission and load reporting --- python/sglang/srt/disaggregation/prefill.py | 2 +- python/sglang/srt/managers/scheduler.py | 5 +- python/sglang/srt/managers/tp_worker.py | 12 +- .../sglang/srt/model_executor/model_runner.py | 15 ++ .../managers/test_dcp_logical_capacity.py | 240 ++++++++++++++++++ 5 files changed, 261 insertions(+), 13 deletions(-) create mode 100644 test/registered/unit/managers/test_dcp_logical_capacity.py diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 0067c984899f..9522b420d75b 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -183,7 +183,7 @@ def __init__( self.scheduler = scheduler self.scheduler_stage_metrics = scheduler_stage_metrics self.max_total_num_tokens = ( - self.scheduler.tp_worker.model_runner.effective_max_total_num_tokens + self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens ) self.transfer_backend = transfer_backend if envs.SGLANG_DISAGG_STAGING_BUFFER.get(): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b8b408e45582..de5ec437ed25 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2360,8 +2360,7 @@ def init_pool_stats_observer(self) -> None: enable_hisparse=self.enable_hisparse, full_tokens_per_layer=self.full_tokens_per_layer, swa_tokens_per_layer=self.swa_tokens_per_layer, - max_total_num_tokens=self.max_total_num_tokens - * get_parallel().attn_dcp_size, + max_total_num_tokens=self.max_total_num_tokens, get_last_batch=lambda: self.last_batch, get_running_batch=lambda: self.running_batch, ) @@ -2546,7 +2545,7 @@ def init_req_max_new_tokens(self, req): max_new_tokens = self.token_to_kv_pool_allocator.max_new_tokens_for_memory( input_len, max_new_tokens, - token_capacity=self.max_total_num_tokens * get_parallel().attn_dcp_size, + token_capacity=self.max_total_num_tokens, sliding_window_size=self.sliding_window_size, chunk_size=self.chunked_prefill_size, ) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index f00aa167d471..643a7409c6c8 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -449,9 +449,7 @@ def alloc_memory_pool( assert self.model_runner.max_running_requests > 0, "max_running_request is zero" max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens - * get_parallel().attn_dcp_size - - 1, + self.model_runner.effective_logical_max_total_num_tokens - 1, ) assert max_req_len > 0, "Memory pool size is too small" @@ -575,14 +573,10 @@ def register_hisparse_coordinator(self, coordinator): def get_worker_info(self): max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens - * get_parallel().attn_dcp_size - - 1, + self.model_runner.effective_logical_max_total_num_tokens - 1, ) return ( - self.model_runner.req_to_token_pool.schedulable_token_capacity( - self.model_runner.max_total_num_tokens - ), + self.model_runner.logical_max_total_num_tokens, get_schedule().max_prefill_tokens, self.model_runner.max_running_requests, get_schedule().max_queued_requests, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9f7a69c7b87f..ecbc85d16f12 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1385,6 +1385,21 @@ def unload_lora_adapter(self, lora_ref: LoRARef): """Unload a lora adapter that was previously loaded during initialization or dynamic loading.""" return self.lora_manager.unload_lora_adapter(lora_ref) + @property + def logical_max_total_num_tokens(self): + """Request-token capacity; the DCP allocator already accounts for sharding.""" + capacity = self.max_total_num_tokens + if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: + capacity = self.token_to_kv_pool_allocator.size + return self.req_to_token_pool.schedulable_token_capacity(capacity) + + @property + def effective_logical_max_total_num_tokens(self): + """Logical request limit, preserving hybrid SWA's separate pool bounds.""" + if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: + return self.logical_max_total_num_tokens + return self.effective_max_total_num_tokens + @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py new file mode 100644 index 000000000000..5218ddb39eee --- /dev/null +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -0,0 +1,240 @@ +"""Logical DCP capacities must agree across validation, admission and telemetry.""" + +import unittest +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.observability.metrics_collector import SchedulerStats +from sglang.srt.runtime_context import get_parallel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, published_topology + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +PHYSICAL = 690240 +CONTEXT = 1048576 + + +def make_worker(dcp_size, *, allocator_size=None): + logical = PHYSICAL * dcp_size if allocator_size is None else allocator_size + kv = NS(size=PHYSICAL, mem_usage=8.89) + allocator = PagedTokenToKVPoolAllocator( + logical, 64 * dcp_size, torch.uint8, "cpu", kv, False + ) + runner = ModelRunner.__new__(ModelRunner) + runner.server_args = NS(dcp_size=dcp_size) + runner.kv_cache_configurator = NS( + hybrid_swa_token_capacity=lambda **kw: kw["full_capacity"] or kw["swa_capacity"] + ) + runner.is_hybrid_swa = False + runner.max_total_num_tokens = PHYSICAL + runner.max_running_requests = 64 + runner.token_to_kv_pool_allocator = allocator + runner.token_to_kv_pool = kv + runner.req_to_token_pool = ReqToTokenPool.__new__(ReqToTokenPool) + runner.req_to_token_pool.size = 96 + runner.req_to_token_pool.max_context_len = CONTEXT + runner.req_to_token_pool._aux_cache = None + runner.forward_stream = None + runner.weight_load_mem_usage = 0 + return NS( + model_runner=runner, + model_config=NS(context_len=CONTEXT), + server_args=NS(max_prefill_tokens=16384, max_queued_requests=None), + random_seed=0, + device="cpu", + graph_memory_usage={}, + ) + + +def make_scheduler(worker): + info = TpModelWorker.get_worker_info(worker) + runner = worker.model_runner + scheduler = NS( + tp_worker=worker, + token_to_kv_pool_allocator=runner.token_to_kv_pool_allocator, + req_to_token_pool=runner.req_to_token_pool, + tree_cache=NS(evictable_size=lambda: 0), + session_controller=None, + hisparse_coordinator=None, + is_hybrid_swa=False, + is_hybrid_ssm=False, + enable_hisparse=False, + sliding_window_size=None, + chunked_prefill_size=16384, + full_tokens_per_layer=None, + swa_tokens_per_layer=None, + max_total_num_tokens=info[0], + max_running_requests=info[2], + max_req_len=info[4], + max_req_input_len=info[5], + max_new_tokens_limit=None, + page_size=64 * runner.server_args.dcp_size, + server_args=NS(dcp_size=runner.server_args.dcp_size, enable_lora=False), + running_batch=NS(reqs=[]), + last_batch=None, + waiting_queue=[], + chunked_req=None, + disaggregation_mode=DisaggregationMode.DECODE, + disagg_decode_prealloc_queue=NS(queue=[], retracted_queue=[]), + disagg_decode_transfer_queue=NS(queue=[]), + spec_algorithm=NS(is_none=lambda: True), + metrics_reporter=NS(stats=SchedulerStats()), + ) + Scheduler.init_pool_stats_observer(scheduler) + Scheduler.init_load_inquirer(scheduler) + return scheduler + + +class TestDcpLogicalCapacity(CustomTestCase): + def make_worker(self, dcp_size, **kwargs): + self.enterContext(published_topology(tp_size=dcp_size, dcp_size=dcp_size)) + return make_worker(dcp_size, **kwargs) + + def setUp(self): + for config in ( + patch( + "sglang.srt.managers.tp_worker.get_schedule", + return_value=NS(max_prefill_tokens=16384, max_queued_requests=None), + ), + patch( + "sglang.srt.managers.scheduler_components.load_inquirer.get_lora", + return_value=NS(enable_lora=False), + ), + ): + config.start() + self.addCleanup(config.stop) + + def test_auxiliary_dense_capacity_applies_after_dcp_translation(self): + worker = self.make_worker(8) + worker.model_config.context_len = 4_000_000 + worker.model_runner.req_to_token_pool._aux_cache = NS(dense_capacity=2_000_000) + info = TpModelWorker.get_worker_info(worker) + self.assertEqual(info[0], 2_000_000) + self.assertEqual(info[4], 1_999_999) + + def test_worker_capacity_uses_effective_dcp(self): + for dcp_size in (1, 2, 8): + with self.subTest(dcp_size=dcp_size): + worker = self.make_worker(dcp_size) + # The runner's configured DCP is not the published attention group. + worker.model_runner.server_args.dcp_size = 8 + info = TpModelWorker.get_worker_info(worker) + capacity = PHYSICAL * dcp_size + self.assertEqual(info[0], capacity) + self.assertEqual(info[4], min(CONTEXT, capacity) - 1) + self.assertEqual(info[5], info[4] - 5) + + def test_allocator_is_authoritative_not_another_dcp_multiplier(self): + worker = self.make_worker(8, allocator_size=PHYSICAL * 4) + self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) + # The size is already logical even when a draft runner's local buffer + # configuration also covers the full logical token domain. + worker.model_runner.max_total_num_tokens = PHYSICAL * 8 + self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) + + def test_non_dcp_and_hybrid_swa_bounds(self): + runner = self.make_worker(1, allocator_size=PHYSICAL * 2).model_runner + self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) + runner.is_hybrid_swa = True + runner.swa_max_total_num_tokens = PHYSICAL // 4 + for dcp_size, full_capacity, expected in ( + (1, PHYSICAL // 2, PHYSICAL // 2), + (1, 0, PHYSICAL // 4), + (8, PHYSICAL // 2, PHYSICAL // 2), + ): + with self.subTest(dcp_size=dcp_size, full_capacity=full_capacity): + self.enterContext(get_parallel().override(attn_dcp_size=dcp_size)) + runner.full_max_total_num_tokens = full_capacity + self.assertEqual( + runner.effective_logical_max_total_num_tokens, expected + ) + self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) + + def test_output_budget_does_not_multiply_logical_capacity_again(self): + worker = self.make_worker(8) + worker.model_config.context_len = PHYSICAL * 16 + scheduler = make_scheduler(worker) + req = NS( + rid="near-capacity", + origin_input_ids=range(scheduler.max_total_num_tokens - 1024), + sampling_params=NS(max_new_tokens=600, min_new_tokens=0), + ) + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 511) + + def test_one_million_token_context_is_not_clipped_to_per_rank_rows(self): + scheduler = make_scheduler(self.make_worker(8)) + self.assertEqual(scheduler.max_req_input_len, CONTEXT - 6) + req = NS( + rid="dcp-long", + origin_input_ids=range(1_000_000), + output_ids=[], + sampling_params=NS(max_new_tokens=64, min_new_tokens=0), + ) + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 64) + self.assertFalse( + PrefillBootstrapQueue._check_if_req_exceed_kv_capacity( + NS( + max_total_num_tokens=scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens + ), + req, + ) + ) + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.max_total_num_tokens = scheduler.max_total_num_tokens + queue.token_to_kv_pool_allocator = scheduler.token_to_kv_pool_allocator + queue.token_to_kv_pool = scheduler.tp_worker.model_runner.token_to_kv_pool + queue.num_reserved_decode_tokens = 0 + queue.scheduler = scheduler + scheduler.output_streamer = Mock() + self.assertFalse( + DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) + ) + # Context safety still applies independently of the larger pool. + req.origin_input_ids = range(CONTEXT - 32) + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 30) + + req.origin_input_ids = range(scheduler.max_total_num_tokens + 1) + req.return_logprob = False + with patch("sglang.srt.disaggregation.decode.prepare_abort") as abort: + self.assertTrue( + DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) + ) + abort.assert_called_once() + queue.scheduler.output_streamer.stream_output.assert_called_once() + + def test_load_usage_uses_logical_capacity_exactly_once(self): + for dcp_size in (1, 8): + with self.subTest(dcp_size=dcp_size): + scheduler = make_scheduler(self.make_worker(dcp_size)) + allocator = scheduler.token_to_kv_pool_allocator + # Keep half the real allocator's pages available; no CUDA calls. + allocator.free_pages = allocator.free_pages[ + : len(allocator.free_pages) // 2 + ] + loads = scheduler.load_inquirer.get_loads() + self.assertEqual(loads.max_total_num_tokens, allocator.size) + self.assertEqual(loads.memory.token_capacity, allocator.size) + used = allocator.size - (allocator.num_pages // 2) * allocator.page_size + self.assertEqual(loads.num_used_tokens, used) + self.assertEqual(loads.token_usage, 0.5) + self.assertEqual( + scheduler.pool_stats_observer.max_total_num_tokens, allocator.size + ) + + +if __name__ == "__main__": + unittest.main() From e6eca5b9a324ccb0ac3b965de2717e3004a9a66c Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:44:20 +0000 Subject: [PATCH 05/10] Derive DCP logical capacity from configured rows in the KV configurator Under DCP with the unified Mamba pool, the allocator size also counts Mamba state bytes as tokens. Compute logical capacity as the configured per-rank rows times attn_dcp_size in KVCacheConfigurator (draft sizes are already widened; SWA allocators never widen), divide sglang:num_pages by the widened allocator page size, and use the 3.10-safe enter_scope helper in the test. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/managers/scheduler.py | 3 +- .../srt/mem_cache/kv_cache_configurator.py | 8 ++ .../sglang/srt/model_executor/model_runner.py | 17 +-- .../managers/test_dcp_logical_capacity.py | 128 +++++++++++++----- 4 files changed, 115 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index de5ec437ed25..72b81ace1517 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1255,7 +1255,8 @@ def emit_metrics_constants(self) -> None: # TODO: max_running_requests_under_SLO has no setter — dead chain. max_running_requests_under_SLO=None, page_size=self.page_size, - num_pages=self.max_total_num_tokens // self.page_size, + num_pages=self.max_total_num_tokens + // self.token_to_kv_pool_allocator.page_size, context_len=self.model_config.context_len, startup_available_gpu_memory_gb=self.startup_available_gpu_memory_gb, ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 2157a39fccdb..3ab13adb1adc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -322,6 +322,14 @@ def hybrid_swa_token_capacity( ) return full_capacity or swa_capacity + def logical_token_capacity(self, *, max_total_num_tokens: int) -> int: + """Request tokens a pool of `max_total_num_tokens` per-rank rows holds.""" + # The paged target allocator widens each row into attn_dcp_size ids; draft + # sizes are already widened (loc_space_scale) and SWA allocators never widen. + if self.is_hybrid_swa or self.is_draft_worker: + return max_total_num_tokens + return max_total_num_tokens * get_parallel().attn_dcp_size + def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): return None diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ecbc85d16f12..72a92b9696a8 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1387,18 +1387,19 @@ def unload_lora_adapter(self, lora_ref: LoRARef): @property def logical_max_total_num_tokens(self): - """Request-token capacity; the DCP allocator already accounts for sharding.""" - capacity = self.max_total_num_tokens - if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: - capacity = self.token_to_kv_pool_allocator.size - return self.req_to_token_pool.schedulable_token_capacity(capacity) + """Request-token capacity in logical tokens, not per-rank DCP rows.""" + return self.req_to_token_pool.schedulable_token_capacity( + self.kv_cache_configurator.logical_token_capacity( + max_total_num_tokens=self.max_total_num_tokens + ) + ) @property def effective_logical_max_total_num_tokens(self): """Logical request limit, preserving hybrid SWA's separate pool bounds.""" - if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: - return self.logical_max_total_num_tokens - return self.effective_max_total_num_tokens + if self.is_hybrid_swa: + return self.effective_max_total_num_tokens + return self.logical_max_total_num_tokens @property def effective_max_total_num_tokens(self): diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py index 5218ddb39eee..a0a129e391ce 100644 --- a/test/registered/unit/managers/test_dcp_logical_capacity.py +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -12,12 +12,21 @@ from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.unified_mamba import ( + UnifiedMambaTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MLASubPoolSpec, + UnifiedKVPool, +) from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.observability.metrics_collector import SchedulerStats from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, published_topology +from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -25,17 +34,56 @@ CONTEXT = 1048576 -def make_worker(dcp_size, *, allocator_size=None): - logical = PHYSICAL * dcp_size if allocator_size is None else allocator_size +def make_configurator(*, is_hybrid_swa=False, is_draft_worker=False): + configurator = KVCacheConfigurator.__new__(KVCacheConfigurator) + configurator.is_hybrid_swa = is_hybrid_swa + configurator.is_draft_worker = is_draft_worker + return configurator + + +def make_unified_mamba_allocator(*, n_full_tokens, n_mamba_slots): + full = MLASubPoolSpec( + name="full", + layer_num=3, + kv_lora_rank=6, + qk_rope_head_dim=2, + store_dtype=torch.bfloat16, + grow_direction="down", + ) + mamba = MambaSubPoolSpec( + name="mamba", + layer_num=2, + conv_state_shapes=((4, 3),), + conv_dtype=torch.float32, + temporal_state_shape=(2, 2, 2), + temporal_dtype=torch.float32, + grow_direction="up", + ) + pool = UnifiedKVPool( + total_bytes=full.entry_bytes() * n_full_tokens + + mamba.entry_bytes() * n_mamba_slots, + sub_pool_specs=[full, mamba], + device="cpu", + enable_memory_saver=False, + page_size=1, + ) + kvcache = NS( + full_kv_pool=NS(buf=torch.empty(pool.max_slots("full"))), + mamba_pool=NS(buf=torch.empty(pool.max_slots("mamba"))), + ) + return UnifiedMambaTokenToKVPoolAllocator( + unified_buffer=pool, kvcache=kvcache, device="cpu", page_size=1 + ) + + +def make_worker(dcp_size): kv = NS(size=PHYSICAL, mem_usage=8.89) allocator = PagedTokenToKVPoolAllocator( - logical, 64 * dcp_size, torch.uint8, "cpu", kv, False + PHYSICAL * dcp_size, 64 * dcp_size, torch.uint8, "cpu", kv, False ) runner = ModelRunner.__new__(ModelRunner) runner.server_args = NS(dcp_size=dcp_size) - runner.kv_cache_configurator = NS( - hybrid_swa_token_capacity=lambda **kw: kw["full_capacity"] or kw["swa_capacity"] - ) + runner.kv_cache_configurator = make_configurator() runner.is_hybrid_swa = False runner.max_total_num_tokens = PHYSICAL runner.max_running_requests = 64 @@ -79,7 +127,7 @@ def make_scheduler(worker): max_req_len=info[4], max_req_input_len=info[5], max_new_tokens_limit=None, - page_size=64 * runner.server_args.dcp_size, + page_size=64, server_args=NS(dcp_size=runner.server_args.dcp_size, enable_lora=False), running_batch=NS(reqs=[]), last_batch=None, @@ -90,6 +138,10 @@ def make_scheduler(worker): disagg_decode_transfer_queue=NS(queue=[]), spec_algorithm=NS(is_none=lambda: True), metrics_reporter=NS(stats=SchedulerStats()), + metrics_collector=Mock(), + draft_worker=None, + startup_available_gpu_memory_gb=None, + model_config=worker.model_config, ) Scheduler.init_pool_stats_observer(scheduler) Scheduler.init_load_inquirer(scheduler) @@ -97,9 +149,14 @@ def make_scheduler(worker): class TestDcpLogicalCapacity(CustomTestCase): - def make_worker(self, dcp_size, **kwargs): - self.enterContext(published_topology(tp_size=dcp_size, dcp_size=dcp_size)) - return make_worker(dcp_size, **kwargs) + def make_worker(self, dcp_size, **server_args_fields): + enter_scope( + self, + published_topology( + tp_size=dcp_size, dcp_size=dcp_size, **server_args_fields + ), + ) + return make_worker(dcp_size) def setUp(self): for config in ( @@ -115,14 +172,6 @@ def setUp(self): config.start() self.addCleanup(config.stop) - def test_auxiliary_dense_capacity_applies_after_dcp_translation(self): - worker = self.make_worker(8) - worker.model_config.context_len = 4_000_000 - worker.model_runner.req_to_token_pool._aux_cache = NS(dense_capacity=2_000_000) - info = TpModelWorker.get_worker_info(worker) - self.assertEqual(info[0], 2_000_000) - self.assertEqual(info[4], 1_999_999) - def test_worker_capacity_uses_effective_dcp(self): for dcp_size in (1, 2, 8): with self.subTest(dcp_size=dcp_size): @@ -135,18 +184,28 @@ def test_worker_capacity_uses_effective_dcp(self): self.assertEqual(info[4], min(CONTEXT, capacity) - 1) self.assertEqual(info[5], info[4] - 5) - def test_allocator_is_authoritative_not_another_dcp_multiplier(self): - worker = self.make_worker(8, allocator_size=PHYSICAL * 4) - self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) - # The size is already logical even when a draft runner's local buffer - # configuration also covers the full logical token domain. - worker.model_runner.max_total_num_tokens = PHYSICAL * 8 - self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) - - def test_non_dcp_and_hybrid_swa_bounds(self): - runner = self.make_worker(1, allocator_size=PHYSICAL * 2).model_runner - self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) + def test_capacity_is_rows_times_dcp_not_allocator_size(self): + """The unified Mamba allocator's size also counts Mamba state bytes as + tokens, so capacity must come from the configured rows, not the allocator.""" + runner = self.make_worker(8).model_runner + runner.max_total_num_tokens = 64 + runner.token_to_kv_pool_allocator = make_unified_mamba_allocator( + n_full_tokens=64, n_mamba_slots=8 + ) + self.assertGreater(runner.token_to_kv_pool_allocator.size, 64 * 8) + self.assertEqual(runner.logical_max_total_num_tokens, 64 * 8) + # Draft sizes are already widened by loc_space_scale; SWA never widens. + for configurator in ( + make_configurator(is_draft_worker=True), + make_configurator(is_hybrid_swa=True), + ): + runner.kv_cache_configurator = configurator + self.assertEqual(runner.logical_max_total_num_tokens, 64) + + def test_hybrid_swa_bounds_do_not_widen_under_dcp(self): + runner = self.make_worker(1).model_runner runner.is_hybrid_swa = True + runner.kv_cache_configurator = make_configurator(is_hybrid_swa=True) runner.swa_max_total_num_tokens = PHYSICAL // 4 for dcp_size, full_capacity, expected in ( (1, PHYSICAL // 2, PHYSICAL // 2), @@ -154,7 +213,7 @@ def test_non_dcp_and_hybrid_swa_bounds(self): (8, PHYSICAL // 2, PHYSICAL // 2), ): with self.subTest(dcp_size=dcp_size, full_capacity=full_capacity): - self.enterContext(get_parallel().override(attn_dcp_size=dcp_size)) + enter_scope(self, get_parallel().override(attn_dcp_size=dcp_size)) runner.full_max_total_num_tokens = full_capacity self.assertEqual( runner.effective_logical_max_total_num_tokens, expected @@ -219,7 +278,9 @@ def test_one_million_token_context_is_not_clipped_to_per_rank_rows(self): def test_load_usage_uses_logical_capacity_exactly_once(self): for dcp_size in (1, 8): with self.subTest(dcp_size=dcp_size): - scheduler = make_scheduler(self.make_worker(dcp_size)) + scheduler = make_scheduler( + self.make_worker(dcp_size, enable_metrics=True) + ) allocator = scheduler.token_to_kv_pool_allocator # Keep half the real allocator's pages available; no CUDA calls. allocator.free_pages = allocator.free_pages[ @@ -234,6 +295,9 @@ def test_load_usage_uses_logical_capacity_exactly_once(self): self.assertEqual( scheduler.pool_stats_observer.max_total_num_tokens, allocator.size ) + Scheduler.emit_metrics_constants(scheduler) + constants = scheduler.metrics_collector.emit_constants.call_args + self.assertEqual(constants.kwargs["num_pages"], PHYSICAL // 64) if __name__ == "__main__": From 38e92b3a98bf7fc5b130d9d4573e8a1166f07e47 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:28:36 +0000 Subject: [PATCH 06/10] Preserve full DFlash draft KV during PD CPU retraction --- python/sglang/srt/managers/schedule_batch.py | 16 ++ python/sglang/srt/mem_cache/allocator/base.py | 5 + python/sglang/srt/mem_cache/common.py | 1 + .../srt/speculative/dflash_worker_v2.py | 6 + .../test_pd_full_draft_cpu_retraction.py | 202 ++++++++++++++++++ .../mem_cache/test_retraction_mamba_backup.py | 1 + 6 files changed, 231 insertions(+) create mode 100644 test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index bb0616c717d7..bf400c3cf629 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2004,6 +2004,13 @@ def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): if mamba_pool is not None and self.kv.holds_mamba else None ), + draft_cpu=( + token_to_kv_pool_allocator.cpu_retraction_draft_pool.get_cpu_copy( + token_indices, req_pool_index=self.kv.req_pool_idx + ) + if token_to_kv_pool_allocator.cpu_retraction_draft_pool is not None + else None + ), ) def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): @@ -2023,6 +2030,15 @@ def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): mamba_indices=self.kv.mamba_pool_idx, req_pool_index=self.kv.req_pool_idx, ) + draft_cpu = self.kv.retraction_backup.draft_cpu + if draft_cpu is not None: + draft_pool = token_to_kv_pool_allocator.cpu_retraction_draft_pool + assert draft_pool is not None + # These remain logical token ids, even when the target MLA pool + # localizes them to this DCP rank's physical rows. + draft_pool.load_cpu_copy( + draft_cpu, token_indices, req_pool_index=self.kv.req_pool_idx + ) self.kv.retraction_backup = None def build_rebootstrap_payload(self) -> dict: diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 870013ed0354..490e18ddb6c3 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -39,6 +39,11 @@ def prepare_mamba_allocation(self, target_size: int) -> None: class BaseTokenToKVPoolAllocator(abc.ABC): + # Optional independent draft storage sharing this allocator's token ids. + # PD CPU retraction must preserve it before those ids can be reused. + # Class-level: SWA and HiSparse allocators do not call this __init__. + cpu_retraction_draft_pool: KVCache | None = None + @abc.abstractmethod def __init__( self, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 9640071cc6d7..70ac3fc4eeb1 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -37,6 +37,7 @@ class RetractionBackup(NamedTuple): pool_transfers: Optional[list[PoolTransfer]] = None # Set when the KV pool leaves the recurrent state to the caller. mamba_cpu: Any = None + draft_cpu: Any = None def kv_to_page_indices(kv_indices: torch.Tensor, page_size: int) -> np.ndarray: diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index c55007f3df14..d0a37ec0f4e4 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -597,6 +597,12 @@ def alloc_memory_pool( ), token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) + if not self.use_compact_draft_cache: + # Full draft KV shares the target's logical token ids but has + # independent tensors, so CPU retraction must save it too. + token_to_kv_pool_allocator.cpu_retraction_draft_pool = ( + self.draft_model_runner.token_to_kv_pool + ) def init_attention_backends(self): with ( diff --git a/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py b/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py new file mode 100644 index 000000000000..7141681c530e --- /dev/null +++ b/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py @@ -0,0 +1,202 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.common import retraction_backup, retraction_restore +from sglang.srt.mem_cache.memory_pool import ( + HybridLinearKVPool, + MambaPool, + MHATokenToKVPool, + MLATokenToKVPool, + ReqToTokenPool, +) +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.runtime_context import get_parallel +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2 +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, published_topology + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +def make_pools(dcp_size): + # Use real copy/restore methods with small CPU tensors. Physical target + # rows are sharded; full draft rows retain the logical token-id domain. + target = object.__new__(MLATokenToKVPool) + target.layer_num = 2 + target.cpu_offloading_chunk_size = 7 + target.kv_buffer = [ + (torch.arange(128, dtype=torch.float32) + 1000 * layer).view(128, 1, 1) + for layer in range(2) + ] + draft = object.__new__(MHATokenToKVPool) + draft.use_hnd = False + draft.layer_num = 3 + draft.cpu_offloading_chunk_size = 11 + draft.k_buffer = [ + (torch.arange(128 * dcp_size, dtype=torch.float32) + 3000 * layer).view( + -1, 1, 1 + ) + for layer in range(3) + ] + draft.v_buffer = [buffer + 1500 for buffer in draft.k_buffer] + mamba = object.__new__(MambaPool) + mamba._slot_siblings = [] + mamba.mamba_cache = SimpleNamespace( + conv=[torch.arange(24, dtype=torch.float32).view(2, 4, 3)], + temporal=torch.arange(40, dtype=torch.float32).view(2, 4, 5), + ) + hybrid = object.__new__(HybridLinearKVPool) + hybrid.full_kv_pool = target + hybrid.mamba_pool = mamba + hybrid._mamba_translate = lambda indices: indices + return target, draft, mamba, hybrid + + +class TestFullDraftCpuRetraction(CustomTestCase): + def setUp(self): + self.enterContext(published_topology(speculative_algorithm="DFLASH")) + + def test_full_draft_registers_but_compact_mapping_does_not(self): + for compact in (False, True): + with self.subTest(compact=compact): + _, draft, _, hybrid = make_pools(1) + allocator = TokenToKVPoolAllocator( + 128, torch.float32, "cpu", hybrid, need_sort=False + ) + worker = object.__new__(DFlashWorkerV2) + worker.use_compact_draft_cache = compact + worker._draft_worker = SimpleNamespace(alloc_memory_pool=Mock()) + worker.draft_model_runner = SimpleNamespace(token_to_kv_pool=draft) + req_pool = object() + worker.alloc_memory_pool( + req_to_token_pool=req_pool, + token_to_kv_pool_allocator=allocator, + ) + self.assertIs( + allocator.cpu_retraction_draft_pool, + None if compact else draft, + ) + + def test_poison_and_relocate_preserves_target_mamba_and_full_draft(self): + for dcp_size in (1, 2, 8): + for rank in range(dcp_size): + for length in (0, 1, 7, 17, 63): + for with_draft in (False, True): + with self.subTest( + dcp_size=dcp_size, + rank=rank, + length=length, + with_draft=with_draft, + ): + self._round_trip(dcp_size, rank, length, with_draft) + + def _round_trip(self, dcp_size, rank, length, with_draft): + target, draft, mamba, hybrid = make_pools(dcp_size) + allocator = TokenToKVPoolAllocator( + 128 * dcp_size, torch.float32, "cpu", hybrid, need_sort=False + ) + if with_draft: + allocator.cpu_retraction_draft_pool = draft + req = Req("test", "", [1] * (length + 1), SamplingParams(max_new_tokens=16)) + req.kv.req_pool_idx = 0 + req.kv.mamba_pool_idx = torch.tensor([1]) + # Two fragments; destination relocation preserves each logical offset. + old = torch.cat((torch.arange(32), torch.arange(64, 96)))[:length] + new = old + 32 * dcp_size + req_pool = SimpleNamespace(req_to_token=old.unsqueeze(0).clone()) + target_expected = [buf.clone() for buf in target.kv_buffer] + draft_expected = [ + buf[old].clone() for buf in (*draft.k_buffer, *draft.v_buffer) + ] + mamba_expected = [ + buf[:, 1].clone() + for buf in (*mamba.mamba_cache.conv, mamba.mamba_cache.temporal) + ] + with get_parallel().override( + dcp_enabled=dcp_size > 1, attn_dcp_size=dcp_size, attn_dcp_rank=rank + ): + self.assertTrue( + retraction_backup(req, None, req_pool, allocator, "cpu_tensor") + ) + for buf in (*target.kv_buffer, *draft.k_buffer, *draft.v_buffer): + buf.fill_(-100) + for buf in (*mamba.mamba_cache.conv, mamba.mamba_cache.temporal): + buf.fill_(-200) + req_pool.req_to_token = new.unsqueeze(0) + req.kv.mamba_pool_idx = torch.tensor([2]) + retraction_restore(req, None, req_pool, allocator, "cpu_tensor") + self.assertIsNone(req.kv.retraction_backup) + mask = old.remainder(dcp_size) == rank + source_rows = old[mask] // dcp_size + destination_rows = new[mask] // dcp_size + for actual, expected in zip(target.kv_buffer, target_expected, strict=True): + torch.testing.assert_close(actual[destination_rows], expected[source_rows]) + untouched = torch.ones(actual.shape[0], dtype=torch.bool) + untouched[destination_rows] = False + self.assertTrue(torch.all(actual[untouched] == -100)) + for actual, expected in zip( + (*draft.k_buffer, *draft.v_buffer), draft_expected, strict=True + ): + if with_draft: + torch.testing.assert_close(actual[new], expected) + else: + self.assertTrue(torch.all(actual == -100)) + for actual, expected in zip( + (*mamba.mamba_cache.conv, mamba.mamba_cache.temporal), + mamba_expected, + strict=True, + ): + torch.testing.assert_close(actual[:, 2], expected) + self.assertTrue(torch.all(actual[:, 1] == -200)) + + def test_swa_allocator_round_trips_without_draft_pool(self): + """SWA allocators skip the base __init__; retraction must not need it.""" + req_pool = ReqToTokenPool( + size=1, max_context_len=32, device="cpu", enable_memory_saver=False + ) + kv_pool = SWAKVPool( + size=64, + size_swa=64, + page_size=1, + dtype=torch.float32, + head_num=1, + head_dim=1, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + device="cpu", + ) + kv_pool.swa_req_ring_size = None + allocator = SWATokenToKVPoolAllocator( + size=64, + size_swa=64, + page_size=1, + dtype=torch.float32, + device="cpu", + kvcache=kv_pool, + need_sort=False, + req_to_token_pool=req_pool, + ) + old = allocator.alloc(9) + full_k = kv_pool.full_kv_pool.k_buffer[0] + full_k[old] = torch.arange(9, dtype=torch.float32).view(9, 1, 1) + 1 + expected = full_k[old[:8]].clone() + req = Req("test", "", [1] * 9, SamplingParams(max_new_tokens=16)) + req.kv.req_pool_idx = 0 + req_pool.req_to_token[0, :9] = old + self.assertTrue(retraction_backup(req, None, req_pool, allocator, "cpu_tensor")) + full_k.fill_(-100) + new = allocator.alloc(9) + req_pool.req_to_token[0, :9] = new + retraction_restore(req, None, req_pool, allocator, "cpu_tensor") + torch.testing.assert_close(full_k[new[:8]], expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py index adc138b6de73..24d083be56ae 100644 --- a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py +++ b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py @@ -26,6 +26,7 @@ class _Allocator: def __init__(self, carries_mamba: bool): self._kv = type("_KV", (), {"cpu_copy_carries_mamba": carries_mamba})() self.loaded_kv = None + self.cpu_retraction_draft_pool = None def get_kvcache(self): return self._kv From 6a1cd3ad8fd3fb5281c1adade2eb008164deace6 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:51:37 +0000 Subject: [PATCH 07/10] Register full DFlash draft KV through one declared allocator slot Replace the externally assigned cpu_retraction_draft_pool with DFlashWorkerV2.full_draft_kv_pool, which accepts only a non-compact MHATokenToKVPool (MLA pools localize token ids to DCP rows), and BaseTokenToKVPoolAllocator.register_full_draft_kv_pool. Trim the round-trip matrix to the DCP1 and DCP8 cases that fail without the fix and use the 3.10-safe enter_scope helper. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/mem_cache/allocator/base.py | 9 ++- .../srt/speculative/dflash_worker_v2.py | 19 ++++-- .../test_pd_full_draft_cpu_retraction.py | 67 +++++++++---------- .../mem_cache/test_retraction_mamba_backup.py | 2 +- 5 files changed, 55 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index bf400c3cf629..8ceba42eccd8 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2005,10 +2005,10 @@ def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): else None ), draft_cpu=( - token_to_kv_pool_allocator.cpu_retraction_draft_pool.get_cpu_copy( + token_to_kv_pool_allocator.full_draft_kv_pool.get_cpu_copy( token_indices, req_pool_index=self.kv.req_pool_idx ) - if token_to_kv_pool_allocator.cpu_retraction_draft_pool is not None + if token_to_kv_pool_allocator.full_draft_kv_pool is not None else None ), ) @@ -2032,7 +2032,7 @@ def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): ) draft_cpu = self.kv.retraction_backup.draft_cpu if draft_cpu is not None: - draft_pool = token_to_kv_pool_allocator.cpu_retraction_draft_pool + draft_pool = token_to_kv_pool_allocator.full_draft_kv_pool assert draft_pool is not None # These remain logical token ids, even when the target MLA pool # localizes them to this DCP rank's physical rows. diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 490e18ddb6c3..679bf743289b 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -39,10 +39,8 @@ def prepare_mamba_allocation(self, target_size: int) -> None: class BaseTokenToKVPoolAllocator(abc.ABC): - # Optional independent draft storage sharing this allocator's token ids. - # PD CPU retraction must preserve it before those ids can be reused. # Class-level: SWA and HiSparse allocators do not call this __init__. - cpu_retraction_draft_pool: KVCache | None = None + full_draft_kv_pool: KVCache | None = None @abc.abstractmethod def __init__( @@ -70,6 +68,11 @@ def __init__( def size_full(self): return self.size + def register_full_draft_kv_pool(self, pool: KVCache) -> None: + """Declare a draft KV pool with its own tensors indexed by this allocator's + token ids; PD CPU retraction and PD transfer carry it with the target KV.""" + self.full_draft_kv_pool = pool + # -- scheduler-facing capacity hooks -- # The scheduler calls these unconditionally, with no allocator-type branches # on its side; byte-accounted composites override the token-count defaults. diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index d0a37ec0f4e4..5b36e7a41ba6 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -25,6 +25,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -597,12 +598,18 @@ def alloc_memory_pool( ), token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) - if not self.use_compact_draft_cache: - # Full draft KV shares the target's logical token ids but has - # independent tensors, so CPU retraction must save it too. - token_to_kv_pool_allocator.cpu_retraction_draft_pool = ( - self.draft_model_runner.token_to_kv_pool - ) + if (draft_pool := self.full_draft_kv_pool) is not None: + token_to_kv_pool_allocator.register_full_draft_kv_pool(draft_pool) + + @property + def full_draft_kv_pool(self) -> Optional[MHATokenToKVPool]: + """The draft pool when it holds full draft KV: MHA tensors of its own, + indexed by the target's logical token ids rather than a compact table.""" + # MLA pools localize token ids to DCP rows; only MHA rows stay logical. + draft_pool = self.draft_model_runner.token_to_kv_pool + if self.use_compact_draft_cache or not isinstance(draft_pool, MHATokenToKVPool): + return None + return draft_pool def init_attention_backends(self): with ( diff --git a/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py b/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py index 7141681c530e..b989f339c35a 100644 --- a/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py +++ b/test/registered/unit/mem_cache/test_pd_full_draft_cpu_retraction.py @@ -20,7 +20,7 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2 from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, published_topology +from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -59,51 +59,51 @@ def make_pools(dcp_size): return target, draft, mamba, hybrid +def alloc_dflash_memory_pool(allocator, draft_pool, *, compact=False): + worker = object.__new__(DFlashWorkerV2) + worker.use_compact_draft_cache = compact + worker._draft_worker = SimpleNamespace(alloc_memory_pool=Mock()) + worker.draft_model_runner = SimpleNamespace(token_to_kv_pool=draft_pool) + worker.alloc_memory_pool( + req_to_token_pool=object(), token_to_kv_pool_allocator=allocator + ) + + class TestFullDraftCpuRetraction(CustomTestCase): def setUp(self): - self.enterContext(published_topology(speculative_algorithm="DFLASH")) + enter_scope(self, published_topology(speculative_algorithm="DFLASH")) - def test_full_draft_registers_but_compact_mapping_does_not(self): - for compact in (False, True): - with self.subTest(compact=compact): - _, draft, _, hybrid = make_pools(1) + def test_only_full_mha_draft_pool_registers(self): + _, draft, _, hybrid = make_pools(1) + mla_draft, _, _, _ = make_pools(1) + for compact, pool, registered in ( + (False, draft, True), + (True, draft, False), + # MLA pools localize token ids to DCP rows, so they are not full drafts. + (False, mla_draft, False), + ): + with self.subTest(compact=compact, pool=type(pool).__name__): allocator = TokenToKVPoolAllocator( 128, torch.float32, "cpu", hybrid, need_sort=False ) - worker = object.__new__(DFlashWorkerV2) - worker.use_compact_draft_cache = compact - worker._draft_worker = SimpleNamespace(alloc_memory_pool=Mock()) - worker.draft_model_runner = SimpleNamespace(token_to_kv_pool=draft) - req_pool = object() - worker.alloc_memory_pool( - req_to_token_pool=req_pool, - token_to_kv_pool_allocator=allocator, - ) + alloc_dflash_memory_pool(allocator, pool, compact=compact) self.assertIs( - allocator.cpu_retraction_draft_pool, - None if compact else draft, + allocator.full_draft_kv_pool, pool if registered else None ) def test_poison_and_relocate_preserves_target_mamba_and_full_draft(self): - for dcp_size in (1, 2, 8): - for rank in range(dcp_size): - for length in (0, 1, 7, 17, 63): - for with_draft in (False, True): - with self.subTest( - dcp_size=dcp_size, - rank=rank, - length=length, - with_draft=with_draft, - ): - self._round_trip(dcp_size, rank, length, with_draft) + """Reused slots must not leak into a resumed request's draft KV.""" + # A 17-token single fragment, and two fragments across DCP8 on rank 3. + for dcp_size, rank, length in ((1, 0, 17), (8, 3, 63)): + with self.subTest(dcp_size=dcp_size, rank=rank, length=length): + self._round_trip(dcp_size, rank, length) - def _round_trip(self, dcp_size, rank, length, with_draft): + def _round_trip(self, dcp_size, rank, length): target, draft, mamba, hybrid = make_pools(dcp_size) allocator = TokenToKVPoolAllocator( 128 * dcp_size, torch.float32, "cpu", hybrid, need_sort=False ) - if with_draft: - allocator.cpu_retraction_draft_pool = draft + alloc_dflash_memory_pool(allocator, draft) req = Req("test", "", [1] * (length + 1), SamplingParams(max_new_tokens=16)) req.kv.req_pool_idx = 0 req.kv.mamba_pool_idx = torch.tensor([1]) @@ -144,10 +144,7 @@ def _round_trip(self, dcp_size, rank, length, with_draft): for actual, expected in zip( (*draft.k_buffer, *draft.v_buffer), draft_expected, strict=True ): - if with_draft: - torch.testing.assert_close(actual[new], expected) - else: - self.assertTrue(torch.all(actual == -100)) + torch.testing.assert_close(actual[new], expected) for actual, expected in zip( (*mamba.mamba_cache.conv, mamba.mamba_cache.temporal), mamba_expected, diff --git a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py index 24d083be56ae..f42de43a3db1 100644 --- a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py +++ b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py @@ -26,7 +26,7 @@ class _Allocator: def __init__(self, carries_mamba: bool): self._kv = type("_KV", (), {"cpu_copy_carries_mamba": carries_mamba})() self.loaded_kv = None - self.cpu_retraction_draft_pool = None + self.full_draft_kv_pool = None def get_kvcache(self): return self._kv From 135b1575fbe6e5873ca2af232975c8edd565bc36 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:28:37 +0000 Subject: [PATCH 08/10] [PD] Transfer DFlash draft KV separately from DCP target indices --- python/sglang/srt/disaggregation/base/conn.py | 1 + python/sglang/srt/disaggregation/decode.py | 19 ++- python/sglang/srt/disaggregation/dflash_kv.py | 19 +++ python/sglang/srt/disaggregation/nixl/conn.py | 19 ++- python/sglang/srt/disaggregation/prefill.py | 19 ++- python/sglang/srt/disaggregation/utils.py | 5 + .../sglang/srt/mem_cache/kv_cache_builder.py | 14 ++- .../test_dflash_logical_transfer.py | 114 ++++++++++++++++++ 8 files changed, 203 insertions(+), 7 deletions(-) create mode 100644 python/sglang/srt/disaggregation/dflash_kv.py create mode 100644 test/registered/unit/disaggregation/test_dflash_logical_transfer.py diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 0cc63380b991..7226e9ebd54a 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -16,6 +16,7 @@ class StateType(str, enum.Enum): MAMBA = "mamba" + DFLASH_KV = "dflash_kv" QSA_PENDING = "qsa_pending" QSA_COMPRESSED = "qsa_compressed" SWA = "swa" diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ecb39d3e7635..caf069216397 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -596,7 +596,8 @@ def _init_kv_manager(self) -> CommonKVManager: kv_item_lens += device_kv_item_lens[c4_layer_num:] kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:]) num_draft_entries = 0 - if self.draft_token_to_kv_pool is not None: + draft_full = getattr(self.draft_token_to_kv_pool, "_pd_dflash_full_kv", False) + if self.draft_token_to_kv_pool is not None and not draft_full: # Draft KV shares target virtual ids. Unified target KV is transferred # with physical ids, so it needs a separate draft index vector. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( @@ -614,7 +615,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_args.num_draft_entries = num_draft_entries kv_args.kv_layer_ids = build_kv_layer_ids( token_to_kv_pool=self.token_to_kv_pool, - draft_token_to_kv_pool=self.draft_token_to_kv_pool, + draft_token_to_kv_pool=None if draft_full else self.draft_token_to_kv_pool, num_draft_entries=num_draft_entries, num_hidden_layers=self.scheduler.model_config.num_hidden_layers, ) @@ -1504,6 +1505,19 @@ def pop_preallocated( seq_len = origin_input_len + def _draft_payload(): + from sglang.srt.disaggregation.dflash_kv import draft_transfer_start + + draft_pool = self.draft_token_to_kv_pool + wire_page = self.token_to_kv_pool.page_size + start = draft_transfer_start( + seq_len, getattr(draft_pool, "_pd_dflash_window", None), wire_page + ) + logical_indices = self.req_to_token_pool.req_to_token[ + decode_req.req.kv.req_pool_idx, start:seq_len + ] + return kv_to_page_indices(logical_indices, wire_page) + def _mamba_payload(): return [ self.req_to_token_pool.translate_mamba_indices( @@ -1574,6 +1588,7 @@ def _request_state_payload(): clear_request_state(int(decode_req.req.kv.req_pool_idx)) payloads = { StateType.MAMBA: _mamba_payload, + StateType.DFLASH_KV: _draft_payload, StateType.QSA_PENDING: _qsa_pending_payload, StateType.QSA_COMPRESSED: _full_kv_pages_payload, StateType.SWA: _swa_payload, diff --git a/python/sglang/srt/disaggregation/dflash_kv.py b/python/sglang/srt/disaggregation/dflash_kv.py new file mode 100644 index 000000000000..1bd3392d70d8 --- /dev/null +++ b/python/sglang/srt/disaggregation/dflash_kv.py @@ -0,0 +1,19 @@ +"""Logical draft KV transfer independent of target DCP-local addressing.""" + + +def draft_transfer_buffers(pool, wire_page_size): + """Use target wire-page units even when draft allocation pages are wider.""" + if wire_page_size <= 0 or pool.page_size % wire_page_size: + raise ValueError("Draft allocation pages must contain whole transfer pages") + ptrs, lengths, item_lengths = pool.get_contiguous_buf_infos() + ratio = pool.page_size // wire_page_size + if any(length % ratio for length in item_lengths): + raise ValueError("Draft buffer page bytes must divide into transfer pages") + return ptrs, lengths, [length // ratio for length in item_lengths] + + +def draft_transfer_start(seq_len, window_size, page_size): + if seq_len < 0 or page_size <= 0 or (window_size is not None and window_size <= 0): + raise ValueError("Invalid draft transfer extent") + start = max(0, seq_len - window_size) if window_size is not None else 0 + return start // page_size * page_size diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 580572c4f3de..47c34c74b598 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -2590,6 +2590,13 @@ def maybe_send_extra( dst_lids = dst_state_layer_ids[i] if i < len(dst_state_layer_ids) else [] comp_notif = f"{notif}_{i}" + if ( + st == StateType.DFLASH_KV + and self.pp_size is not None + and self.pp_size > 1 + ): + raise RuntimeError("Logical DFlash KV transfer requires PP=1") + if st == StateType.MAMBA: if self.attn_tp_size != decode_tp_size: if 0 in src_dims: @@ -2660,13 +2667,16 @@ def maybe_send_extra( force_flat=True, ) elif st in ( + StateType.DFLASH_KV, StateType.SWA, StateType.QSA_PENDING, StateType.QSA_COMPRESSED, StateType.SWA_RING, StateType.DSV4_REQUEST_STATE, ): - if not self.is_mla_backend and self.attn_tp_size != decode_tp_size: + if ( + not self.is_mla_backend or st == StateType.DFLASH_KV + ) and self.attn_tp_size != decode_tp_size: raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet." ) @@ -2691,7 +2701,12 @@ def maybe_send_extra( dst_gpu_id=dst_gpu_id, notif=comp_notif, state_type=st, - force_flat=st in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED), + force_flat=st + in ( + StateType.DFLASH_KV, + StateType.QSA_PENDING, + StateType.QSA_COMPRESSED, + ), src_layer_ids=src_lids, dst_layer_ids=dst_lids, dst_item_lens=dst_lens, diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9522b420d75b..1765f0d5f6c7 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -273,7 +273,8 @@ def _init_kv_manager(self) -> CommonKVManager: else None ) num_draft_entries = 0 - if draft_kv_pool is not None: + draft_full = getattr(draft_kv_pool, "_pd_dflash_full_kv", False) + if draft_kv_pool is not None and not draft_full: # Draft KV shares target virtual ids. Unified target KV is transferred # with physical ids, so it needs a separate draft index vector. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( @@ -290,7 +291,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_args.num_draft_entries = num_draft_entries kv_args.kv_layer_ids = build_kv_layer_ids( token_to_kv_pool=self.token_to_kv_pool, - draft_token_to_kv_pool=draft_kv_pool, + draft_token_to_kv_pool=None if draft_full else draft_kv_pool, num_draft_entries=num_draft_entries, num_hidden_layers=self.scheduler.model_config.num_hidden_layers, ) @@ -1352,6 +1353,19 @@ def _send_kv_chunk( seq_len = min(req.extend_range.end, transfer_input_len) c128_seq_len = transfer_input_len + def _draft_payload(): + from sglang.srt.disaggregation.dflash_kv import draft_transfer_start + + draft_pool = self.disagg_prefill_bootstrap_queue.draft_token_to_kv_pool + wire_page = self.token_to_kv_pool.page_size + start = draft_transfer_start( + seq_len, getattr(draft_pool, "_pd_dflash_window", None), wire_page + ) + logical_indices = self.req_to_token_pool.req_to_token[ + req.kv.req_pool_idx, start:seq_len + ] + return kv_to_page_indices(logical_indices, wire_page) + def _mamba_payload(): return [ self.req_to_token_pool.translate_mamba_indices( @@ -1419,6 +1433,7 @@ def _request_state_payload(): ) payloads = { StateType.MAMBA: _mamba_payload, + StateType.DFLASH_KV: _draft_payload, StateType.QSA_PENDING: _qsa_pending_payload, StateType.QSA_COMPRESSED: _full_kv_pages_payload, StateType.SWA: _swa_payload, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 7f2b8f45b9fd..92b15208873a 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -1349,6 +1349,11 @@ def setup_state_kv_args( kv_args.state_layer_ids = [] kv_args.is_hybrid_mla_backend = False kv_args.state_conv_shard_groups = [] + if getattr(draft_token_to_kv_pool, "_pd_dflash_full_kv", False): + from sglang.srt.disaggregation.dflash_kv import draft_transfer_buffers + + dp, dl, il = draft_transfer_buffers(draft_token_to_kv_pool, kv_args.page_size) + append_state_component(kv_args, StateType.DFLASH_KV, dp, dl, il) # V4's KVCache is organized by compression-ratio buckets rather than by layer. kv_args.mla_compression_ratios = ( list(token_to_kv_pool.compression_ratios) diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index f403e7cd43b6..9dc5ab7f738b 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -138,7 +138,19 @@ def get_draft_kv_pool( draft_runner = draft_worker.draft_worker.draft_runner_list[0] else: draft_runner = draft_worker.draft_worker.draft_runner - return draft_runner.token_to_kv_pool + pool = draft_runner.token_to_kv_pool + if ( + spec_algorithm.is_dflash() + and get_disagg().disaggregation_transfer_backend == "nixl" + and not draft_worker.use_compact_draft_cache + and isinstance(pool, MHATokenToKVPool) + ): + # Only full draft pools share the target's logical token indices. + pool._pd_dflash_full_kv = True + pool._pd_dflash_window = getattr( + draft_runner.model_config.hf_text_config, "sliding_window", None + ) + return pool def maybe_register_hicache_draft( diff --git a/test/registered/unit/disaggregation/test_dflash_logical_transfer.py b/test/registered/unit/disaggregation/test_dflash_logical_transfer.py new file mode 100644 index 000000000000..5dea9a5ccf01 --- /dev/null +++ b/test/registered/unit/disaggregation/test_dflash_logical_transfer.py @@ -0,0 +1,114 @@ +import sys +from types import SimpleNamespace + +import pytest + +from sglang.srt.disaggregation.dflash_kv import ( + draft_transfer_buffers, + draft_transfer_start, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, published_topology + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestDraftTransferRegistration(CustomTestCase): + def test_only_full_nixl_dflash_pools_use_logical_transfer(self): + from sglang.srt.disaggregation.base.conn import StateType + from sglang.srt.disaggregation.utils import setup_state_kv_args + from sglang.srt.mem_cache.kv_cache_builder import get_draft_kv_pool + from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + for backend, compact in (("nixl", False), ("nixl", True), ("mooncake", False)): + with ( + self.subTest(backend=backend, compact=compact), + published_topology(disaggregation_transfer_backend=backend) as args, + ): + pool = MHATokenToKVPool.__new__(MHATokenToKVPool) + pool.page_size = 128 + pool.get_contiguous_buf_infos = lambda: ([1000], [8192], [1024]) + runner = SimpleNamespace( + token_to_kv_pool=pool, + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(sliding_window=4096) + ), + ) + worker = SimpleNamespace( + use_compact_draft_cache=compact, + draft_worker=SimpleNamespace(draft_runner=runner), + ) + selected = get_draft_kv_pool( + draft_worker=worker, + spec_algorithm=SpeculativeAlgorithm.DFLASH, + server_args=args, + ) + kv_args = SimpleNamespace(page_size=64) + setup_state_kv_args(kv_args, SimpleNamespace(), selected) + enabled = backend == "nixl" and not compact + self.assertEqual( + kv_args.state_types, [StateType.DFLASH_KV] if enabled else [] + ) + if enabled: + self.assertEqual(kv_args.state_item_lens, [[512]]) + self.assertEqual(selected._pd_dflash_window, 4096) + + +@pytest.mark.parametrize("dcp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("wire_page", [16, 64]) +def test_wire_page_geometry_is_independent_of_dcp(dcp_size, wire_page): + bytes_per_token = 512 + allocation_page = wire_page * dcp_size + pool = SimpleNamespace( + page_size=allocation_page, + get_contiguous_buf_infos=lambda: ( + [1000, 2000], + [1048576, 1048576], + [allocation_page * bytes_per_token] * 2, + ), + ) + ptrs, lengths, items = draft_transfer_buffers(pool, wire_page) + assert ptrs == [1000, 2000] + assert lengths == [1048576, 1048576] + assert items == [wire_page * bytes_per_token] * 2 + # Raw logical locations, unlike DCP-local target slots, address every token. + logical_page = 3 + assert ptrs[0] + logical_page * items[0] == 1000 + 3 * wire_page * bytes_per_token + + +@pytest.mark.parametrize( + "seq_len,window,page,expected", + [ + (10, None, 64, 0), + (8192, None, 64, 0), + (8192, 4096, 64, 4096), + (8193, 4096, 64, 4096), + (2048, 4096, 64, 0), + (0, 4096, 64, 0), + ], +) +def test_transfer_covers_full_prefix_or_page_aligned_draft_window( + seq_len, window, page, expected +): + assert draft_transfer_start(seq_len, window, page) == expected + + +@pytest.mark.parametrize("seq,window,page", [(-1, None, 64), (1, 0, 64), (1, None, 0)]) +def test_invalid_extent_rejected(seq, window, page): + with pytest.raises(ValueError): + draft_transfer_start(seq, window, page) + + +def test_incompatible_page_geometry_rejected(): + pool = SimpleNamespace( + page_size=64, get_contiguous_buf_infos=lambda: ([1], [512], [512]) + ) + with pytest.raises(ValueError): + draft_transfer_buffers(pool, 128) + with pytest.raises(ValueError): + draft_transfer_buffers(pool, 0) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) From 55d0cd79cc61c4f5264ef9371d542f476c8cd5dd Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 21:05:18 +0000 Subject: [PATCH 09/10] Gate logical DFlash draft transfer on DCP and drop pool monkey-patching get_draft_kv_pool, which tagged the draft pool with private attributes, has no caller: the scheduler hands the PD queues primary_draft_kv_pool. Resolve a DFlashDraftTransfer once per PD queue from the allocator's registered full draft pool (#39749), and only under DCP on NIXL, so non-DCP deployments keep the draft in the target KV entries. Take the tail window only when every draft layer slides, start the payload at the decode prefix, append the state component last, share one payload helper between prefill and decode, check TP equality at peer registration (DFLASH already requires pp_size 1 at startup), and cover the payload through NixlKVManager.maybe_send_extra. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/disaggregation/decode.py | 32 +- python/sglang/srt/disaggregation/dflash_kv.py | 102 ++++++- python/sglang/srt/disaggregation/nixl/conn.py | 21 +- python/sglang/srt/disaggregation/prefill.py | 34 ++- python/sglang/srt/disaggregation/utils.py | 16 +- .../sglang/srt/mem_cache/kv_cache_builder.py | 14 +- .../test_dflash_logical_transfer.py | 277 ++++++++++++------ 7 files changed, 334 insertions(+), 162 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index caf069216397..f59fc612d6af 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -51,6 +51,7 @@ HiCacheRestoreGatedKVReceiver, HiCacheRestoreResult, ) +from sglang.srt.disaggregation.dflash_kv import resolve_dflash_draft_transfer from sglang.srt.disaggregation.utils import ( DisaggregationMode, KVClassType, @@ -398,6 +399,11 @@ def __init__( self.pp_size = get_parallel().pp_size self.num_reserved_decode_tokens = num_reserved_decode_tokens self.transfer_backend = transfer_backend + self.dflash_draft_transfer = resolve_dflash_draft_transfer( + allocator=token_to_kv_pool_allocator, + draft_worker=scheduler.draft_worker, + transfer_backend=transfer_backend, + ) # Queue for requests pending pre-allocation self.queue: List[DecodeRequest] = [] self.retracted_queue: List[Req] = [] @@ -596,8 +602,10 @@ def _init_kv_manager(self) -> CommonKVManager: kv_item_lens += device_kv_item_lens[c4_layer_num:] kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:]) num_draft_entries = 0 - draft_full = getattr(self.draft_token_to_kv_pool, "_pd_dflash_full_kv", False) - if self.draft_token_to_kv_pool is not None and not draft_full: + draft_kv_pool = ( + self.draft_token_to_kv_pool if self.dflash_draft_transfer is None else None + ) + if draft_kv_pool is not None: # Draft KV shares target virtual ids. Unified target KV is transferred # with physical ids, so it needs a separate draft index vector. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( @@ -615,7 +623,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_args.num_draft_entries = num_draft_entries kv_args.kv_layer_ids = build_kv_layer_ids( token_to_kv_pool=self.token_to_kv_pool, - draft_token_to_kv_pool=None if draft_full else self.draft_token_to_kv_pool, + draft_token_to_kv_pool=draft_kv_pool, num_draft_entries=num_draft_entries, num_hidden_layers=self.scheduler.model_config.num_hidden_layers, ) @@ -633,6 +641,7 @@ def _init_kv_manager(self) -> CommonKVManager: self.draft_token_to_kv_pool, total_kv_layers=self.scheduler.model_config.num_hidden_layers, req_to_token_pool=getattr(self, "req_to_token_pool", None), + dflash_draft_transfer=self.dflash_draft_transfer, ) kv_args.ib_device = get_disagg().disaggregation_ib_device @@ -1506,17 +1515,14 @@ def pop_preallocated( seq_len = origin_input_len def _draft_payload(): - from sglang.srt.disaggregation.dflash_kv import draft_transfer_start - - draft_pool = self.draft_token_to_kv_pool - wire_page = self.token_to_kv_pool.page_size - start = draft_transfer_start( - seq_len, getattr(draft_pool, "_pd_dflash_window", None), wire_page + return self.dflash_draft_transfer.page_indices( + req_to_token=self.req_to_token_pool.req_to_token[ + decode_req.req.kv.req_pool_idx + ], + prefix_len=total_prefix_len, + seq_len=seq_len, + wire_page_size=self.token_to_kv_pool.page_size, ) - logical_indices = self.req_to_token_pool.req_to_token[ - decode_req.req.kv.req_pool_idx, start:seq_len - ] - return kv_to_page_indices(logical_indices, wire_page) def _mamba_payload(): return [ diff --git a/python/sglang/srt/disaggregation/dflash_kv.py b/python/sglang/srt/disaggregation/dflash_kv.py index 1bd3392d70d8..453bc9f348ab 100644 --- a/python/sglang/srt/disaggregation/dflash_kv.py +++ b/python/sglang/srt/disaggregation/dflash_kv.py @@ -1,19 +1,91 @@ -"""Logical draft KV transfer independent of target DCP-local addressing.""" +"""Full DFlash draft KV transfer in the draft pool's logical token-id domain.""" +from __future__ import annotations -def draft_transfer_buffers(pool, wire_page_size): - """Use target wire-page units even when draft allocation pages are wider.""" - if wire_page_size <= 0 or pool.page_size % wire_page_size: - raise ValueError("Draft allocation pages must contain whole transfer pages") - ptrs, lengths, item_lengths = pool.get_contiguous_buf_infos() - ratio = pool.page_size // wire_page_size - if any(length % ratio for length in item_lengths): - raise ValueError("Draft buffer page bytes must divide into transfer pages") - return ptrs, lengths, [length // ratio for length in item_lengths] +from typing import TYPE_CHECKING, List, Optional, Tuple +import msgspec +import numpy as np +import torch -def draft_transfer_start(seq_len, window_size, page_size): - if seq_len < 0 or page_size <= 0 or (window_size is not None and window_size <= 0): - raise ValueError("Invalid draft transfer extent") - start = max(0, seq_len - window_size) if window_size is not None else 0 - return start // page_size * page_size +from sglang.srt.disaggregation.utils import TransferBackend +from sglang.srt.mem_cache.common import kv_to_page_indices +from sglang.srt.runtime_context import get_parallel +from sglang.srt.speculative.dflash_utils import ( + get_dflash_attention_sliding_window_size, + get_dflash_layer_types, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool + from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2 + + +class DFlashDraftTransfer(msgspec.Struct, frozen=True): + """Full DFlash draft KV sent as its own PD state component. + + Under DCP the target KV travels as rank-local rows, while the full draft + pool is replicated and indexed by the target allocator's logical token ids. + """ + + pool: MHATokenToKVPool + # Draft tokens decode attends to; None transfers the whole prompt. + window: Optional[int] + + def buffer_infos(self, wire_page_size: int) -> Tuple[List, List, List]: + """Draft buffers in target wire-page units; draft allocation pages are wider.""" + if wire_page_size <= 0 or self.pool.page_size % wire_page_size: + raise ValueError("Draft allocation pages must contain whole transfer pages") + ptrs, lengths, item_lengths = self.pool.get_contiguous_buf_infos() + ratio = self.pool.page_size // wire_page_size + if any(length % ratio for length in item_lengths): + raise ValueError("Draft buffer page bytes must divide into transfer pages") + return ptrs, lengths, [length // ratio for length in item_lengths] + + def page_indices( + self, + *, + req_to_token: torch.Tensor, + prefix_len: int, + seq_len: int, + wire_page_size: int, + ) -> np.ndarray: + """Wire pages of one request's draft KV past the decode prefix.""" + # The decode prefix already holds its draft KV, like the SWA payload. + start = prefix_len + if self.window is not None: + start = max(prefix_len, seq_len - self.window) + start = start // wire_page_size * wire_page_size + return kv_to_page_indices(req_to_token[start:seq_len], wire_page_size) + + +def resolve_dflash_draft_transfer( + *, + allocator: BaseTokenToKVPoolAllocator, + draft_worker: Optional[DFlashWorkerV2], + transfer_backend: TransferBackend, +) -> Optional[DFlashDraftTransfer]: + # Without DCP the target's KV indices are logical too, so the draft rides + # the target KV entries (num_draft_entries) instead. + if ( + allocator.full_draft_kv_pool is None + or transfer_backend != TransferBackend.NIXL + or get_parallel().attn_dcp_size == 1 + ): + return None + return DFlashDraftTransfer( + pool=allocator.full_draft_kv_pool, + window=draft_transfer_window( + draft_worker.draft_model_runner.model_config.hf_config + ), + ) + + +def draft_transfer_window(draft_hf_config) -> Optional[int]: + """Draft KV tokens decode needs, or None unless every draft layer slides.""" + layer_types = get_dflash_layer_types(draft_hf_config) + if not layer_types or any(t != "sliding_attention" for t in layer_types): + return None + # The helper returns window_left; the HF window also counts the current token. + return get_dflash_attention_sliding_window_size(draft_hf_config) + 1 diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 47c34c74b598..02d4cd517e96 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -1606,6 +1606,16 @@ def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo): decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout( decode_kv_args.dst_dcp_size, decode_kv_args.dst_dcp_rank ) + # Draft MHA heads shard by attention TP even when the target is MLA. + if ( + StateType.DFLASH_KV in self.kv_args.state_types + and decode_kv_args.decode_tp_size != self.attn_tp_size + ): + raise RuntimeError( + "PD transfer of full DFlash draft KV requires equal prefill and " + f"decode attention TP sizes, got prefill={self.attn_tp_size}, " + f"decode={decode_kv_args.decode_tp_size}." + ) if decode_kv_args.requires_dcp_relayout: self._init_dcp_pack_buffers_once(decode_kv_args.dst_dcp_size) self.decode_kv_args_table[agent_name] = decode_kv_args @@ -2590,13 +2600,6 @@ def maybe_send_extra( dst_lids = dst_state_layer_ids[i] if i < len(dst_state_layer_ids) else [] comp_notif = f"{notif}_{i}" - if ( - st == StateType.DFLASH_KV - and self.pp_size is not None - and self.pp_size > 1 - ): - raise RuntimeError("Logical DFlash KV transfer requires PP=1") - if st == StateType.MAMBA: if self.attn_tp_size != decode_tp_size: if 0 in src_dims: @@ -2674,9 +2677,7 @@ def maybe_send_extra( StateType.SWA_RING, StateType.DSV4_REQUEST_STATE, ): - if ( - not self.is_mla_backend or st == StateType.DFLASH_KV - ) and self.attn_tp_size != decode_tp_size: + if not self.is_mla_backend and self.attn_tp_size != decode_tp_size: raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet." ) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1765f0d5f6c7..2d773c7168c2 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -42,6 +42,7 @@ compute_grid_segments, staging_grid_tokens, ) +from sglang.srt.disaggregation.dflash_kv import resolve_dflash_draft_transfer from sglang.srt.disaggregation.utils import ( FAKE_BOOTSTRAP_HOST, DisaggregationMode, @@ -186,6 +187,11 @@ def __init__( self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens ) self.transfer_backend = transfer_backend + self.dflash_draft_transfer = resolve_dflash_draft_transfer( + allocator=self.scheduler.token_to_kv_pool_allocator, + draft_worker=self.scheduler.draft_worker, + transfer_backend=transfer_backend, + ) if envs.SGLANG_DISAGG_STAGING_BUFFER.get(): if self.is_mla_backend: raise RuntimeError( @@ -273,8 +279,12 @@ def _init_kv_manager(self) -> CommonKVManager: else None ) num_draft_entries = 0 - draft_full = getattr(draft_kv_pool, "_pd_dflash_full_kv", False) - if draft_kv_pool is not None and not draft_full: + dflash_draft_transfer = ( + self.dflash_draft_transfer if transfer_draft_cache else None + ) + if dflash_draft_transfer is not None: + draft_kv_pool = None + if draft_kv_pool is not None: # Draft KV shares target virtual ids. Unified target KV is transferred # with physical ids, so it needs a separate draft index vector. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( @@ -291,7 +301,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_args.num_draft_entries = num_draft_entries kv_args.kv_layer_ids = build_kv_layer_ids( token_to_kv_pool=self.token_to_kv_pool, - draft_token_to_kv_pool=None if draft_full else draft_kv_pool, + draft_token_to_kv_pool=draft_kv_pool, num_draft_entries=num_draft_entries, num_hidden_layers=self.scheduler.model_config.num_hidden_layers, ) @@ -315,6 +325,7 @@ def _init_kv_manager(self) -> CommonKVManager: self.draft_token_to_kv_pool if transfer_draft_cache else None, self.scheduler.model_config.num_hidden_layers, req_to_token_pool=req_to_token_pool, + dflash_draft_transfer=dflash_draft_transfer, ) kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER) @@ -1354,17 +1365,14 @@ def _send_kv_chunk( c128_seq_len = transfer_input_len def _draft_payload(): - from sglang.srt.disaggregation.dflash_kv import draft_transfer_start - - draft_pool = self.disagg_prefill_bootstrap_queue.draft_token_to_kv_pool - wire_page = self.token_to_kv_pool.page_size - start = draft_transfer_start( - seq_len, getattr(draft_pool, "_pd_dflash_window", None), wire_page + return self.disagg_prefill_bootstrap_queue.dflash_draft_transfer.page_indices( + req_to_token=self.req_to_token_pool.req_to_token[ + req.kv.req_pool_idx + ], + prefix_len=req.disagg_decode_prefix_len, + seq_len=seq_len, + wire_page_size=self.token_to_kv_pool.page_size, ) - logical_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, start:seq_len - ] - return kv_to_page_indices(logical_indices, wire_page) def _mamba_payload(): return [ diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 92b15208873a..8ec3f1529349 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -36,6 +36,7 @@ CommonKVReceiver, CommonKVSender, ) + from sglang.srt.disaggregation.dflash_kv import DFlashDraftTransfer from sglang.srt.managers.schedule_batch import Req if is_npu(): @@ -1326,6 +1327,7 @@ def setup_state_kv_args( draft_token_to_kv_pool=None, total_kv_layers: int = None, req_to_token_pool=None, + dflash_draft_transfer: Optional[DFlashDraftTransfer] = None, ) -> None: from sglang.srt.disaggregation.base.conn import StateType from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool @@ -1349,11 +1351,6 @@ def setup_state_kv_args( kv_args.state_layer_ids = [] kv_args.is_hybrid_mla_backend = False kv_args.state_conv_shard_groups = [] - if getattr(draft_token_to_kv_pool, "_pd_dflash_full_kv", False): - from sglang.srt.disaggregation.dflash_kv import draft_transfer_buffers - - dp, dl, il = draft_transfer_buffers(draft_token_to_kv_pool, kv_args.page_size) - append_state_component(kv_args, StateType.DFLASH_KV, dp, dl, il) # V4's KVCache is organized by compression-ratio buckets rather than by layer. kv_args.mla_compression_ratios = ( list(token_to_kv_pool.compression_ratios) @@ -1731,6 +1728,15 @@ def append_dsa_tail(pool) -> None: slice_outer_counts, ) + # Last, so the other components keep their positions when only one side + # registers the draft (decode-only speculative decoding). + if dflash_draft_transfer is not None: + append_state_component( + kv_args, + StateType.DFLASH_KV, + *dflash_draft_transfer.buffer_infos(kv_args.page_size), + ) + def get_dsv41_spec_layout(kv_args: KVArgs) -> Optional[dict]: """Describe the positional transfer layout without pool capacities or pointers.""" diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 9dc5ab7f738b..f403e7cd43b6 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -138,19 +138,7 @@ def get_draft_kv_pool( draft_runner = draft_worker.draft_worker.draft_runner_list[0] else: draft_runner = draft_worker.draft_worker.draft_runner - pool = draft_runner.token_to_kv_pool - if ( - spec_algorithm.is_dflash() - and get_disagg().disaggregation_transfer_backend == "nixl" - and not draft_worker.use_compact_draft_cache - and isinstance(pool, MHATokenToKVPool) - ): - # Only full draft pools share the target's logical token indices. - pool._pd_dflash_full_kv = True - pool._pd_dflash_window = getattr( - draft_runner.model_config.hf_text_config, "sliding_window", None - ) - return pool + return draft_runner.token_to_kv_pool def maybe_register_hicache_draft( diff --git a/test/registered/unit/disaggregation/test_dflash_logical_transfer.py b/test/registered/unit/disaggregation/test_dflash_logical_transfer.py index 5dea9a5ccf01..859ae79eccb8 100644 --- a/test/registered/unit/disaggregation/test_dflash_logical_transfer.py +++ b/test/registered/unit/disaggregation/test_dflash_logical_transfer.py @@ -1,114 +1,205 @@ -import sys +"""PD transfer of full DFlash draft KV in its logical token-id domain.""" + +import unittest from types import SimpleNamespace -import pytest +import numpy as np +import torch +from sglang.srt.disaggregation.base.conn import StateType from sglang.srt.disaggregation.dflash_kv import ( - draft_transfer_buffers, - draft_transfer_start, + DFlashDraftTransfer, + draft_transfer_window, + resolve_dflash_draft_transfer, ) +from sglang.srt.disaggregation.nixl.conn import NixlKVManager +from sglang.srt.disaggregation.utils import TransferBackend, setup_state_kv_args +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, published_topology - -register_cpu_ci(est_time=2, suite="base-a-test-cpu") - +from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology -class TestDraftTransferRegistration(CustomTestCase): - def test_only_full_nixl_dflash_pools_use_logical_transfer(self): - from sglang.srt.disaggregation.base.conn import StateType - from sglang.srt.disaggregation.utils import setup_state_kv_args - from sglang.srt.mem_cache.kv_cache_builder import get_draft_kv_pool - from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool - from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +register_cpu_ci(est_time=5, suite="base-a-test-cpu") - for backend, compact in (("nixl", False), ("nixl", True), ("mooncake", False)): - with ( - self.subTest(backend=backend, compact=compact), - published_topology(disaggregation_transfer_backend=backend) as args, - ): - pool = MHATokenToKVPool.__new__(MHATokenToKVPool) - pool.page_size = 128 - pool.get_contiguous_buf_infos = lambda: ([1000], [8192], [1024]) - runner = SimpleNamespace( - token_to_kv_pool=pool, - model_config=SimpleNamespace( - hf_text_config=SimpleNamespace(sliding_window=4096) - ), - ) - worker = SimpleNamespace( - use_compact_draft_cache=compact, - draft_worker=SimpleNamespace(draft_runner=runner), - ) - selected = get_draft_kv_pool( - draft_worker=worker, - spec_algorithm=SpeculativeAlgorithm.DFLASH, - server_args=args, - ) - kv_args = SimpleNamespace(page_size=64) - setup_state_kv_args(kv_args, SimpleNamespace(), selected) - enabled = backend == "nixl" and not compact - self.assertEqual( - kv_args.state_types, [StateType.DFLASH_KV] if enabled else [] - ) - if enabled: - self.assertEqual(kv_args.state_item_lens, [[512]]) - self.assertEqual(selected._pd_dflash_window, 4096) +WIRE_PAGE = 64 +TOKEN_BYTES = 512 -@pytest.mark.parametrize("dcp_size", [1, 2, 4, 8]) -@pytest.mark.parametrize("wire_page", [16, 64]) -def test_wire_page_geometry_is_independent_of_dcp(dcp_size, wire_page): - bytes_per_token = 512 - allocation_page = wire_page * dcp_size - pool = SimpleNamespace( - page_size=allocation_page, +def make_draft_pool(dcp_size): + # A replicated draft pool pages the widened (logical) id space. + page_size = WIRE_PAGE * dcp_size + return SimpleNamespace( + page_size=page_size, get_contiguous_buf_infos=lambda: ( [1000, 2000], - [1048576, 1048576], - [allocation_page * bytes_per_token] * 2, + [1 << 20, 1 << 20], + [page_size * TOKEN_BYTES] * 2, ), ) - ptrs, lengths, items = draft_transfer_buffers(pool, wire_page) - assert ptrs == [1000, 2000] - assert lengths == [1048576, 1048576] - assert items == [wire_page * bytes_per_token] * 2 - # Raw logical locations, unlike DCP-local target slots, address every token. - logical_page = 3 - assert ptrs[0] + logical_page * items[0] == 1000 + 3 * wire_page * bytes_per_token - - -@pytest.mark.parametrize( - "seq_len,window,page,expected", - [ - (10, None, 64, 0), - (8192, None, 64, 0), - (8192, 4096, 64, 4096), - (8193, 4096, 64, 4096), - (2048, 4096, 64, 0), - (0, 4096, 64, 0), - ], -) -def test_transfer_covers_full_prefix_or_page_aligned_draft_window( - seq_len, window, page, expected -): - assert draft_transfer_start(seq_len, window, page) == expected -@pytest.mark.parametrize("seq,window,page", [(-1, None, 64), (1, 0, 64), (1, None, 0)]) -def test_invalid_extent_rejected(seq, window, page): - with pytest.raises(ValueError): - draft_transfer_start(seq, window, page) +def make_draft_worker(layer_types, sliding_window=4096): + hf_config = SimpleNamespace(layer_types=layer_types, sliding_window=sliding_window) + return SimpleNamespace( + draft_model_runner=SimpleNamespace( + model_config=SimpleNamespace(hf_config=hf_config) + ) + ) -def test_incompatible_page_geometry_rejected(): - pool = SimpleNamespace( - page_size=64, get_contiguous_buf_infos=lambda: ([1], [512], [512]) - ) - with pytest.raises(ValueError): - draft_transfer_buffers(pool, 128) - with pytest.raises(ValueError): - draft_transfer_buffers(pool, 0) +class RecordingAgent: + def __init__(self): + self.descs = [] + + def get_xfer_descs(self, reqs, mem_kind): + self.descs.append(np.asarray(reqs)) + return len(self.descs) - 1 + + def initialize_xfer(self, op, src, dst, peer, notif): + return (src, dst, notif) + + def transfer(self, handle): + return "PROC" + + +class TestDFlashDraftTransfer(CustomTestCase): + def setUp(self): + enter_scope(self, published_topology(speculative_algorithm="DFLASH")) + + def resolve(self, *, dcp_size, backend=TransferBackend.NIXL, registered=True): + enter_scope(self, get_parallel().override(attn_dcp_size=dcp_size)) + allocator = SimpleNamespace( + full_draft_kv_pool=make_draft_pool(dcp_size) if registered else None + ) + return resolve_dflash_draft_transfer( + allocator=allocator, + draft_worker=make_draft_worker(["full_attention"]), + transfer_backend=backend, + ) + + def test_logical_component_only_under_dcp_on_nixl(self): + """Without DCP the draft must keep riding the target KV entries.""" + for dcp_size, backend, registered, enabled in ( + (8, TransferBackend.NIXL, True, True), + (1, TransferBackend.NIXL, True, False), + (8, TransferBackend.MOONCAKE, True, False), + (8, TransferBackend.NIXL, False, False), + ): + with self.subTest( + dcp_size=dcp_size, backend=backend, registered=registered + ): + transfer = self.resolve( + dcp_size=dcp_size, backend=backend, registered=registered + ) + kv_args = SimpleNamespace(page_size=WIRE_PAGE) + setup_state_kv_args( + kv_args, SimpleNamespace(), dflash_draft_transfer=transfer + ) + self.assertEqual( + kv_args.state_types, [StateType.DFLASH_KV] if enabled else [] + ) + if enabled: + # Item lengths are per target wire page, not per draft page. + self.assertEqual( + kv_args.state_item_lens, [[WIRE_PAGE * TOKEN_BYTES] * 2] + ) + + def test_window_tail_only_when_every_draft_layer_slides(self): + """A full-attention draft layer reads the whole prompt; HF windows + include the current token.""" + sliding = ["sliding_attention"] * 2 + for layer_types, expected in ( + (sliding, 4096), + (["sliding_attention", "full_attention"], None), + (["full_attention"] * 2, None), + (None, None), + ): + with self.subTest(layer_types=layer_types): + hf_config = SimpleNamespace( + layer_types=layer_types, sliding_window=4096 + ) + self.assertEqual(draft_transfer_window(hf_config), expected) + + def test_payload_skips_decode_prefix_and_nixl_writes_logical_rows(self): + """Draft KV is addressed by logical token ids even when the target's + DCP rows are rank-local, and the decode prefix already holds it.""" + dcp_size = 8 + transfer = DFlashDraftTransfer(pool=make_draft_pool(dcp_size), window=None) + # Logical ids of a 512-token request on prefill and on decode. + prefill_row = torch.arange(4096, 4096 + 512) + decode_row = torch.arange(8192, 8192 + 512) + src_pages = transfer.page_indices( + req_to_token=prefill_row, + prefix_len=128, + seq_len=512, + wire_page_size=WIRE_PAGE, + ) + dst_pages = transfer.page_indices( + req_to_token=decode_row, + prefix_len=128, + seq_len=512, + wire_page_size=WIRE_PAGE, + ) + np.testing.assert_array_equal(src_pages, np.arange(66, 72)) + np.testing.assert_array_equal(dst_pages, np.arange(130, 136)) + + ptrs, _, item_lens = transfer.buffer_infos(WIRE_PAGE) + manager = NixlKVManager.__new__(NixlKVManager) + manager.agent = RecordingAgent() + manager.is_mla_backend = True + manager.attn_tp_size = 8 + manager.pp_size = 1 + manager.kv_args = SimpleNamespace( + gpu_id=0, + kv_data_ptrs=[], + state_types=[StateType.DFLASH_KV], + state_data_ptrs=[ptrs], + state_item_lens=[item_lens], + state_dim_per_tensor=[[]], + state_conv_shard_groups=[[]], + state_slice_outer_counts=[[]], + state_layer_ids=[[]], + mla_compression_ratios=None, + ) + handles = manager.maybe_send_extra( + "decode", + prefill_state_indices=[src_pages.tolist()], + dst_state_data_ptrs=[[5000, 6000]], + dst_state_indices=[dst_pages.tolist()], + dst_gpu_id=1, + notif="room_state_0", + decode_tp_size=8, + dst_state_item_lens=[item_lens], + ) + self.assertEqual(len(handles), 1) + src_descs, dst_descs = manager.agent.descs + page_bytes = WIRE_PAGE * TOKEN_BYTES + # One contiguous run of six pages per draft buffer. + np.testing.assert_array_equal( + src_descs[:, :2], + [ + [1000 + 66 * page_bytes, 6 * page_bytes], + [2000 + 66 * page_bytes, 6 * page_bytes], + ], + ) + np.testing.assert_array_equal( + dst_descs[:, :2], + [ + [5000 + 130 * page_bytes, 6 * page_bytes], + [6000 + 130 * page_bytes, 6 * page_bytes], + ], + ) + + def test_window_tail_starts_at_page_boundary(self): + transfer = DFlashDraftTransfer(pool=make_draft_pool(1), window=100) + pages = transfer.page_indices( + req_to_token=torch.arange(1024), + prefix_len=0, + seq_len=1000, + wire_page_size=WIRE_PAGE, + ) + # 1000 - 100 = 900 floors to 896, the first page decode still reads. + np.testing.assert_array_equal(pages, np.arange(14, 16)) if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"])) + unittest.main() From 537482328b81054422a9a0520829ba969596c1df Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 21:28:16 +0000 Subject: [PATCH 10/10] Keep DFlash draft KV entries on decode for DCP relayout prefills Neither PD side knows the peer's DCP size when it registers: prefill learns it at the handshake and decode never does. A DCP>1 decode therefore keeps the draft among its KV entries, for a DCP=1 prefill's relayout plan, and also registers the trailing DFLASH_KV component that a same-DCP prefill sends instead. NIXL registers each buffer once when a state component reuses a KV region. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/disaggregation/decode.py | 12 ++- python/sglang/srt/disaggregation/dflash_kv.py | 16 +++- python/sglang/srt/disaggregation/nixl/conn.py | 5 ++ python/sglang/srt/disaggregation/prefill.py | 10 ++- .../test_dflash_logical_transfer.py | 84 ++++++++++++++++++- 5 files changed, 121 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f59fc612d6af..eab1a0e1ef64 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -51,7 +51,10 @@ HiCacheRestoreGatedKVReceiver, HiCacheRestoreResult, ) -from sglang.srt.disaggregation.dflash_kv import resolve_dflash_draft_transfer +from sglang.srt.disaggregation.dflash_kv import ( + lists_draft_as_kv_entries, + resolve_dflash_draft_transfer, +) from sglang.srt.disaggregation.utils import ( DisaggregationMode, KVClassType, @@ -603,7 +606,12 @@ def _init_kv_manager(self) -> CommonKVManager: kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:]) num_draft_entries = 0 draft_kv_pool = ( - self.draft_token_to_kv_pool if self.dflash_draft_transfer is None else None + self.draft_token_to_kv_pool + if lists_draft_as_kv_entries( + mode=DisaggregationMode.DECODE, + dflash_draft_transfer=self.dflash_draft_transfer, + ) + else None ) if draft_kv_pool is not None: # Draft KV shares target virtual ids. Unified target KV is transferred diff --git a/python/sglang/srt/disaggregation/dflash_kv.py b/python/sglang/srt/disaggregation/dflash_kv.py index 453bc9f348ab..ec9c53087cb9 100644 --- a/python/sglang/srt/disaggregation/dflash_kv.py +++ b/python/sglang/srt/disaggregation/dflash_kv.py @@ -8,7 +8,7 @@ import numpy as np import torch -from sglang.srt.disaggregation.utils import TransferBackend +from sglang.srt.disaggregation.utils import DisaggregationMode, TransferBackend from sglang.srt.mem_cache.common import kv_to_page_indices from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.dflash_utils import ( @@ -82,6 +82,20 @@ def resolve_dflash_draft_transfer( ) +def lists_draft_as_kv_entries( + *, + mode: DisaggregationMode, + dflash_draft_transfer: Optional[DFlashDraftTransfer], +) -> bool: + """Whether this side registers the draft pool among its KV entries. + + Neither side knows the peer's DCP size when it registers, so decode keeps + the entries for a DCP=1 prefill's relayout path and also registers + DFLASH_KV for a same-DCP prefill, which sends only the latter. + """ + return mode == DisaggregationMode.DECODE or dflash_draft_transfer is None + + def draft_transfer_window(draft_hf_config) -> Optional[int]: """Draft KV tokens decode needs, or None unless every draft layer slides.""" layer_types = get_dflash_layer_types(draft_hf_config) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 02d4cd517e96..40712e4436f3 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -1579,6 +1579,9 @@ def register_buffer_to_engine(self): if not self.aux_descs: raise Exception("NIXL memory registration failed for aux tensors") + # Decode lists DFlash draft buffers both as KV entries and as DFLASH_KV + # state; register each buffer once. + kv_regions = set(zip(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens)) state_addrs = [] for comp_ptrs, comp_lens in zip( self.kv_args.state_data_ptrs or [], @@ -1587,6 +1590,8 @@ def register_buffer_to_engine(self): for state_data_ptr, state_data_len in zip(comp_ptrs, comp_lens): if state_data_ptr == 0 or state_data_len == 0: continue + if (state_data_ptr, state_data_len) in kv_regions: + continue state_addrs.append( (state_data_ptr, state_data_len, self.kv_args.gpu_id, "") ) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 2d773c7168c2..04714924b676 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -42,7 +42,10 @@ compute_grid_segments, staging_grid_tokens, ) -from sglang.srt.disaggregation.dflash_kv import resolve_dflash_draft_transfer +from sglang.srt.disaggregation.dflash_kv import ( + lists_draft_as_kv_entries, + resolve_dflash_draft_transfer, +) from sglang.srt.disaggregation.utils import ( FAKE_BOOTSTRAP_HOST, DisaggregationMode, @@ -282,7 +285,10 @@ def _init_kv_manager(self) -> CommonKVManager: dflash_draft_transfer = ( self.dflash_draft_transfer if transfer_draft_cache else None ) - if dflash_draft_transfer is not None: + if not lists_draft_as_kv_entries( + mode=DisaggregationMode.PREFILL, + dflash_draft_transfer=dflash_draft_transfer, + ): draft_kv_pool = None if draft_kv_pool is not None: # Draft KV shares target virtual ids. Unified target KV is transferred diff --git a/test/registered/unit/disaggregation/test_dflash_logical_transfer.py b/test/registered/unit/disaggregation/test_dflash_logical_transfer.py index 859ae79eccb8..e40d41ea83b0 100644 --- a/test/registered/unit/disaggregation/test_dflash_logical_transfer.py +++ b/test/registered/unit/disaggregation/test_dflash_logical_transfer.py @@ -10,10 +10,16 @@ from sglang.srt.disaggregation.dflash_kv import ( DFlashDraftTransfer, draft_transfer_window, + lists_draft_as_kv_entries, resolve_dflash_draft_transfer, ) from sglang.srt.disaggregation.nixl.conn import NixlKVManager -from sglang.srt.disaggregation.utils import TransferBackend, setup_state_kv_args +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + TransferBackend, + resolve_dcp_dst_entry_indices, + setup_state_kv_args, +) from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology @@ -22,6 +28,7 @@ WIRE_PAGE = 64 TOKEN_BYTES = 512 +TARGET_TOKEN_BYTES = 576 def make_draft_pool(dcp_size): @@ -49,6 +56,11 @@ def make_draft_worker(layer_types, sliding_window=4096): class RecordingAgent: def __init__(self): self.descs = [] + self.registered = [] + + def register_memory(self, addrs, mem_kind): + self.registered += [addr[0] for addr in addrs] + return True def get_xfer_descs(self, reqs, mem_kind): self.descs.append(np.asarray(reqs)) @@ -103,6 +115,76 @@ def test_logical_component_only_under_dcp_on_nixl(self): kv_args.state_item_lens, [[WIRE_PAGE * TOKEN_BYTES] * 2] ) + def side_kv_args(self, *, mode, dcp_size): + """KV entries and state components one PD side registers.""" + transfer = self.resolve(dcp_size=dcp_size) + draft = make_draft_pool(dcp_size) + kv_ptrs, kv_lens, kv_items = [100], [1 << 20], [WIRE_PAGE * TARGET_TOKEN_BYTES] + num_draft = 0 + if lists_draft_as_kv_entries(mode=mode, dflash_draft_transfer=transfer): + ptrs, lens, items = draft.get_contiguous_buf_infos() + kv_ptrs, kv_lens, kv_items = ( + kv_ptrs + ptrs, + kv_lens + lens, + kv_items + items, + ) + num_draft = len(ptrs) + kv_args = SimpleNamespace( + page_size=WIRE_PAGE, + kv_data_ptrs=kv_ptrs, + kv_data_lens=kv_lens, + kv_data_mem_kinds=["VRAM"] * len(kv_ptrs), + kv_item_lens=kv_items, + num_draft_entries=num_draft, + aux_data_ptrs=[], + aux_data_lens=[], + gpu_id=0, + ) + setup_state_kv_args(kv_args, SimpleNamespace(), dflash_draft_transfer=transfer) + return kv_args + + def test_prefill_and_decode_layouts_agree_across_dcp_topologies(self): + """Neither side knows the peer's DCP size when it registers, so each + layout must serve every topology the handshake admits.""" + for prefill_dcp, decode_dcp in ((1, 8), (8, 8), (1, 1)): + with self.subTest(prefill_dcp=prefill_dcp, decode_dcp=decode_dcp): + decode = self.side_kv_args( + mode=DisaggregationMode.DECODE, dcp_size=decode_dcp + ) + prefill = self.side_kv_args( + mode=DisaggregationMode.PREFILL, dcp_size=prefill_dcp + ) + manager = NixlKVManager.__new__(NixlKVManager) + manager.kv_args = prefill + manager.dcp_size, manager.dcp_rank = prefill_dcp, 0 + manager.is_mla_backend, manager.is_hybrid_mla_backend = True, False + n_src, n_dst = len(prefill.kv_item_lens), len(decode.kv_item_lens) + # The handshake rejects a decode with fewer KV regions. + self.assertGreaterEqual(n_dst, n_src) + # Prefill sends its own state list; decode may hold a trailing extra. + self.assertEqual( + decode.state_types[: len(prefill.state_types)], + prefill.state_types, + ) + if manager.requires_dcp_relayout(decode_dcp, 0): + # P1 -> D>1: the draft travels in the relayout plan. + dst = resolve_dcp_dst_entry_indices([], [], n_src, n_dst) + manager.prepare_dcp_token_item_lens( + [decode.kv_item_lens[j] for j in dst], decode_dcp + ) + self.assertEqual(prefill.num_draft_entries, 2) + else: + uses_state = StateType.DFLASH_KV in prefill.state_types + self.assertEqual(uses_state, prefill_dcp > 1) + self.assertEqual(prefill.num_draft_entries, 0 if uses_state else 2) + + decode_manager = NixlKVManager.__new__(NixlKVManager) + decode_manager.kv_args = decode + decode_manager.agent = RecordingAgent() + decode_manager.register_buffer_to_engine() + registered = decode_manager.agent.registered + self.assertEqual(len(registered), len(set(registered))) + def test_window_tail_only_when_every_draft_layer_slides(self): """A full-attention draft layer reads the whole prompt; HF windows include the current token."""