[PD][LoRA] Gate decode admission on adapter slots (#39332)

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:

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:
         """

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()

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

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),