From ebd37705e47476d76a1dbf11985e62a72b83f37c Mon Sep 17 00:00:00 2001 From: Yanbin Jiang Date: Mon, 14 Sep 2026 22:11:04 -0700 Subject: [PATCH] [PD][LoRA] Gate decode admission on adapter slots (#39332) --- python/sglang/srt/disaggregation/decode.py | 20 ++++ python/sglang/srt/managers/scheduler.py | 4 +- .../test_disaggregation_lora.py | 107 ++++++++++++++++++ .../test_decode_queue_cleanup.py | 3 + ...test_priority_scheduling_disaggregation.py | 52 +++++++++ 5 files changed, 184 insertions(+), 2 deletions(-) create mode 100644 test/registered/e2e/disaggregation/test_disaggregation_lora.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index c1cf16029..6d806fda3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1188,6 +1188,19 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): - len(self.transfer_queue.queue), ) + if self.scheduler.enable_lora: + running_batches = ( + self.scheduler.running_mbs + if is_pp_mode + else (self.scheduler.running_batch,) + ) + # Include finished requests; GPU work may still use their adapters. + running_loras = { + req.lora_id for batch in running_batches for req in batch.reqs + } + running_loras.update(r.req.lora_id for r in self.transfer_queue.queue) + running_loras.update(req.lora_id for req in self.scheduler.waiting_queue) + # Then, preallocate the remaining requests if possible for i, decode_req in enumerate(self.queue): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: @@ -1208,6 +1221,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): if hisparse_req_budget <= 0: break + if self.scheduler.enable_lora and not self.scheduler.can_schedule_lora_req( + decode_req.req, running_loras + ): + continue + # Memory estimation: don't add if the projected memory cannot be met # TODO: add new_token ratio origin_input_len = self._rebootstrap_prefill_len(decode_req.req) @@ -1559,6 +1577,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): self._num_published_destinations += 1 preallocated_reqs.append(decode_req) indices_to_remove.add(i) + if self.scheduler.enable_lora: + running_loras.add(decode_req.req.lora_id) decode_req.req.time_stats.set_decode_transfer_queue_entry_time() if failed_reqs: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e9cba392f..2e3d1bcb2 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3825,7 +3825,7 @@ class Scheduler( mamba_allocator.alloc_group_begin(len(self.waiting_queue)) # Get requests from the waiting queue to a new prefill batch for req in self.waiting_queue: - if self.enable_lora and not self._can_schedule_lora_req(req, running_loras): + if self.enable_lora and not self.can_schedule_lora_req(req, running_loras): continue running_bs = len(running_batch.reqs) @@ -4035,7 +4035,7 @@ class Scheduler( return new_batch, running_batch - def _can_schedule_lora_req( + def can_schedule_lora_req( self, req: Req, running_loras: set[Optional[str]] ) -> bool: """ diff --git a/test/registered/e2e/disaggregation/test_disaggregation_lora.py b/test/registered/e2e/disaggregation/test_disaggregation_lora.py new file mode 100644 index 000000000..18f792539 --- /dev/null +++ b/test/registered/e2e/disaggregation/test_disaggregation_lora.py @@ -0,0 +1,107 @@ +"""PD LoRA requests remain correct under adapter-slot pressure.""" + +import unittest +from concurrent.futures import ThreadPoolExecutor + +import requests + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, + assert_process_healthy, +) +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + popen_launch_server, + terminate_and_kill_process_tree, +) + +register_cuda_ci(est_time=400, stage="base-b", runner_config="2-gpu-large") + +ADAPTERS = { + "fact": "algoprog/fact-generation-llama-3.1-8b-instruct-lora", + "guard": "nvidia/llama-3.1-nemoguard-8b-topic-control", + "sql": "philschmid/code-llama-3-1-8b-text-to-sql-lora", +} +LORA_ARGS = [ + "--enable-deterministic-inference", + "--cuda-graph-max-bs-prefill", + "1024", + "--enable-lora", + "--lora-paths", + *[f"{name}={path}" for name, path in ADAPTERS.items()], + "--max-loras-per-batch", + "2", +] +PROMPTS = [ + "Give three facts about the planet Mars.", + "Translate this SQL request into a query: list the names of all customers.", + "Is the following on topic for a cooking assistant? 'How do I fix my car?'", + "Write one sentence about the sea.", +] + + +def generate(url: str, prompt: str, lora_path=None) -> str: + payload = { + "text": prompt, + "sampling_params": {"temperature": 0, "max_new_tokens": 128}, + } + if lora_path is not None: + payload["lora_path"] = lora_path + response = requests.post(f"{url}/generate", json=payload, timeout=600) + assert response.status_code == 200, response.text + result = response.json() + assert result["meta_info"]["finish_reason"]["type"] != "abort", result + assert result["text"].strip(), result + return result["text"] + + +class TestDisaggregationLoRA(PDDisaggregationServerBase): + extra_prefill_args = LORA_ARGS + extra_decode_args = LORA_ARGS + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.model = DEFAULT_MODEL_NAME_FOR_TEST + + def test_more_adapters_than_slots_match_single_server(self): + """Excess adapters must wait without aborting or changing outputs.""" + names = [None, *ADAPTERS] + process = popen_launch_server( + self.model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=LORA_ARGS, + ) + try: + expected = { + (prompt, name): generate(self.base_url, prompt, name) + for prompt in PROMPTS + for name in names + } + finally: + terminate_and_kill_process_tree(process) + + for name in ADAPTERS: + self.assertTrue( + any( + expected[prompt, name] != expected[prompt, None] + for prompt in PROMPTS + ), + name, + ) + + self.launch_all() + jobs = list(zip(PROMPTS, names)) * 2 + with ThreadPoolExecutor(max_workers=len(jobs)) as pool: + texts = list(pool.map(lambda job: generate(self.lb_url, *job), jobs)) + for job, text in zip(jobs, texts): + self.assertEqual(text, expected[job], job) + assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url) + assert_process_healthy(self, "decode", self.process_decode, self.decode_url) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index b203b4491..2e8005c26 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -127,6 +127,7 @@ class TestDecodeQueueCleanup(CustomTestCase): scheduler.running_batch.reqs = [] scheduler.enable_priority_scheduling = False scheduler.enable_hisparse = False + scheduler.enable_lora = False scheduler.metrics_reporter.enable_metrics = False scheduler.output_streamer = MagicMock() queue.scheduler = scheduler @@ -181,6 +182,7 @@ class TestDecodeQueueCleanup(CustomTestCase): scheduler.running_batch.reqs = [] scheduler.enable_priority_scheduling = False scheduler.enable_hisparse = False + scheduler.enable_lora = False scheduler.output_streamer = MagicMock() queue.scheduler = scheduler @@ -241,6 +243,7 @@ class TestDecodeQueueCleanup(CustomTestCase): scheduler.running_batch.reqs = [] scheduler.enable_priority_scheduling = False scheduler.enable_hisparse = False + scheduler.enable_lora = False scheduler.server_args.disaggregation_decode_enable_radix_cache = False scheduler.output_streamer = MagicMock() queue.scheduler = scheduler diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index bc710275e..800729990 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -13,6 +13,7 @@ from sglang.srt.disaggregation.decode import ( # noqa: E402 SchedulerDisaggregationDecodeMixin, ) from sglang.srt.disaggregation.utils import DisaggregationMode # noqa: E402 +from sglang.srt.lora.lora_manager import LoRAManager # noqa: E402 from sglang.srt.managers.schedule_batch import ( # noqa: E402 FINISH_ABORT, Req, @@ -169,12 +170,63 @@ class TestDecodePreallocQueuePriority(unittest.TestCase): scheduler.running_batch.reqs = [] scheduler.server_args.disaggregation_decode_enable_radix_cache = False scheduler.enable_hisparse = False + scheduler.enable_lora = False scheduler.waiting_queue = [] scheduler.last_batch = None scheduler.output_streamer = MagicMock() queue.scheduler = scheduler return queue + def test_prealloc_lora_slots_cover_inflight_microbatches_and_queues(self): + """In-flight requests retain their adapter slots.""" + running, finished, transferring, waiting = [ + self._new_decode_req(rid, 0) + for rid in ("running", "finished", "transferring", "waiting") + ] + reqs = [ + self._new_decode_req("new", 3), + self._new_decode_req("blocked", 2), + self._new_decode_req("same-adapter", 1), + ] + for entry in [running, finished, transferring, waiting, *reqs]: + entry.req.lora_id = entry.req.rid + entry.req.finished = MagicMock(return_value=entry is finished) + reqs[2].req.lora_id = "running" + + queue = self._new_queue(reqs) + queue.pp_size = 2 + queue.transfer_queue.queue = [transferring] + scheduler = queue.scheduler + scheduler.running_batch.reqs = [finished.req] + scheduler.running_mbs = [ + scheduler.running_batch, + SimpleNamespace(reqs=[running.req]), + ] + scheduler.waiting_queue = [waiting.req] + scheduler.enable_decode_hicache = False + scheduler.enable_lora = True + scheduler.enable_lora_overlap_loading = False + scheduler.lora_drainer = None + scheduler.can_schedule_lora_req = Scheduler.can_schedule_lora_req.__get__( + scheduler + ) + manager = LoRAManager.__new__(LoRAManager) + manager.max_loras_per_batch = 5 + manager.num_pinned_loras = 0 + scheduler.tp_worker.model_runner.lora_manager = manager + + preallocated, failed = queue.pop_preallocated( + pp_good_rids=[entry.req.rid for entry in reqs], pp_bad_rids=[] + ) + + self.assertEqual(preallocated, [reqs[0], reqs[2]]) + self.assertEqual(failed, []) + self.assertEqual(queue.queue, [reqs[1]]) + self.assertEqual( + [call.args[0] for call in queue._pre_alloc.call_args_list], + [reqs[0].req, reqs[2].req], + ) + def test_prealloc_queue_schedules_higher_priority_values_first_by_default(self): reqs = [ self._new_decode_req("low", 1),