[LoRA] Support LoRA under the breakable/full prefill CUDA graph (#30988)

test/registered/cuda_graph/breakable/test_bcg_with_lora.py

@@ -0,0 +1,305 @@
+"""LoRA under the breakable (BCG) prefill CUDA graph.
+
+For csgmv and triton, launches the same LoRA-enabled server with the prefill
+graph breakable vs disabled and compares per-token prompt logprobs for LoRA,
+base, and mixed batches. Guards against vacuous passes by scraping the logs
+for graph-served prefill batches (including a window where a lone LoRA
+request is the only work in flight) and asserting the adapter moves logprobs.
+"""
+
+import os
+import re
+import shutil
+import tempfile
+import time
+import unittest
+
+import requests
+import torch
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=480, stage="base-b", runner_config="1-gpu-large")
+
+BASE_MODEL = "Qwen/Qwen3-4B"
+LORA_ADAPTER = "nissenj/Qwen3-4B-lora-v2"
+LORA_NAME = "lora0"
+
+# Both servers run with --enable-deterministic-inference (batch-invariant
+# kernels), so logprobs must not depend on how the scheduler composes
+# prefill batches — the graph-vs-eager comparison is exact by construction.
+# Measured diff is 0 on H200 for both csgmv and triton; 1e-2 is headroom.
+GRAPH_VS_EAGER_TOLERANCE = 1e-2
+# The adapter must move at least one prompt logprob by this much.
+LORA_EFFECT_THRESHOLD = 5e-2
+
+# Prompt lengths chosen to land in different captured token buckets.
+PROMPTS = [
+    "What is the capital of France?",
+    "List three benefits of regular exercise. " * 8,
+    "The quick brown fox jumps over the lazy dog. " * 40,
+]
+
+MIXED_BATCH_PROMPTS = [
+    "Summarize the plot of Romeo and Juliet.",
+    "Explain how photosynthesis works. " * 6,
+    "Write a haiku about the sea.",
+    "Describe the water cycle in detail. " * 12,
+]
+MIXED_BATCH_LORA_PATHS = [LORA_NAME, None, LORA_NAME, None]
+
+PREFILL_GRAPH_REPLAY_PATTERN = re.compile(r"Prefill batch.*cuda graph: True")
+
+
+def _generate(text, lora_path):
+    resp = requests.post(
+        DEFAULT_URL_FOR_TEST + "/generate",
+        json={
+            "text": text,
+            "sampling_params": {"max_new_tokens": 0, "temperature": 0.0},
+            "return_logprob": True,
+            "logprob_start_len": 0,
+            "lora_path": lora_path,
+        },
+        timeout=120,
+    )
+    resp.raise_for_status()
+    return resp.json()
+
+
+def _prompt_logprobs(out):
+    # First entry has no logprob (no preceding context).
+    return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:]
+
+
+def _max_abs_diff(a, b):
+    ta, tb = torch.tensor(a, dtype=torch.float64), torch.tensor(b, dtype=torch.float64)
+    assert ta.shape == tb.shape, f"token count mismatch: {ta.shape} vs {tb.shape}"
+    return (ta - tb).abs().max().item()
+
+
+class BCGLoRAServerMixin:
+    """Launches breakable + disabled servers for one LoRA kernel backend and
+    collects prompt logprobs from both. Concrete classes set lora_backend."""
+
+    lora_backend: str
+
+    @classmethod
+    def _server_args(cls, prefill_backend):
+        return [
+            "--enable-lora",
+            "--lora-paths",
+            f"{LORA_NAME}={LORA_ADAPTER}",
+            "--max-loras-per-batch",
+            "2",
+            "--lora-backend",
+            cls.lora_backend,
+            "--cuda-graph-backend-prefill",
+            prefill_backend,
+            "--cuda-graph-max-bs-prefill",
+            "1024",
+            "--cuda-graph-max-bs-decode",
+            "8",
+            "--disable-radix-cache",
+            "--mem-fraction-static",
+            "0.8",
+            "--random-seed",
+            "42",
+            "--enable-deterministic-inference",
+        ]
+
+    @classmethod
+    def _collect(cls):
+        return {
+            "lora": [_prompt_logprobs(_generate(p, LORA_NAME)) for p in PROMPTS],
+            "base": [_prompt_logprobs(_generate(p, None)) for p in PROMPTS],
+            "mixed": [
+                _prompt_logprobs(out)
+                for out in _generate(MIXED_BATCH_PROMPTS, MIXED_BATCH_LORA_PATHS)
+            ],
+        }
+
+    @classmethod
+    def _log_sizes(cls):
+        return [os.path.getsize(path) for path in (cls.stdout_path, cls.stderr_path)]
+
+    @classmethod
+    def _logs_appended_since(cls, sizes):
+        # Binary read: the size snapshot is a byte offset.
+        chunks = []
+        for path, start in zip((cls.stdout_path, cls.stderr_path), sizes):
+            with open(path, "rb") as f:
+                f.seek(start)
+                chunks.append(f.read().decode(errors="replace"))
+        return "\n".join(chunks)
+
+    @classmethod
+    def _probe_lora_replay_window(cls):
+        """Send one adapter-carrying request with nothing else in flight and
+        return the log text appended while it ran; any "Prefill batch" line
+        in the window belongs to that request."""
+        # Baseline once the logs go quiet, so an earlier batch's line cannot
+        # drain into the window (the pipe-dump threads flush asynchronously).
+        sizes = cls._log_sizes()
+        quiet_deadline = time.monotonic() + 10
+        while time.monotonic() < quiet_deadline:
+            time.sleep(1.0)
+            new_sizes = cls._log_sizes()
+            if new_sizes == sizes:
+                break
+            sizes = new_sizes
+
+        _generate(PROMPTS[0], LORA_NAME)
+        deadline = time.monotonic() + 30
+        while True:
+            window = cls._logs_appended_since(sizes)
+            if PREFILL_GRAPH_REPLAY_PATTERN.search(window) or (
+                time.monotonic() > deadline
+            ):
+                return window
+            time.sleep(0.5)
+
+    @classmethod
+    def setUpClass(cls):
+        cls.log_dir = tempfile.mkdtemp(prefix=f"bcg_lora_{cls.lora_backend}_")
+        cls.stdout_path = os.path.join(cls.log_dir, "server.out")
+        cls.stderr_path = os.path.join(cls.log_dir, "server.err")
+        cls.stdout = open(cls.stdout_path, "w")
+        cls.stderr = open(cls.stderr_path, "w")
+        process = popen_launch_server(
+            BASE_MODEL,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=cls._server_args("breakable"),
+            return_stdout_stderr=(cls.stdout, cls.stderr),
+        )
+        try:
+            cls.with_graph = cls._collect()
+            cls.lora_replay_window = cls._probe_lora_replay_window()
+        finally:
+            kill_process_tree(process.pid)
+            cls.stdout.close()
+            cls.stderr.close()
+
+        process = popen_launch_server(
+            BASE_MODEL,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=cls._server_args("disabled"),
+        )
+        try:
+            cls.without_graph = cls._collect()
+        finally:
+            kill_process_tree(process.pid)
+
+    @classmethod
+    def tearDownClass(cls):
+        for attr in ("stdout", "stderr"):
+            f = getattr(cls, attr, None)
+            if f is not None and not f.closed:
+                f.close()
+        log_dir = getattr(cls, "log_dir", None)
+        if log_dir:
+            shutil.rmtree(log_dir, ignore_errors=True)
+
+    def test_prefill_graph_replay_logged(self):
+        # "cuda graph: True" is logged only when PrefillCudaGraphRunner
+        # served the batch; catches an across-the-board eager fallback.
+        logs = ""
+        for path in (self.stdout_path, self.stderr_path):
+            with open(path) as f:
+                logs += f.read()
+        self.assertRegex(
+            logs,
+            PREFILL_GRAPH_REPLAY_PATTERN,
+            "No prefill batch was served by the prefill CUDA graph; the "
+            "breakable-vs-disabled comparison is vacuous.",
+        )
+
+    def test_lora_request_replays_prefill_graph(self):
+        # The probe window held a lone adapter-carrying request, so this
+        # fails if LoRA batches specifically fall back to eager prefill.
+        prefill_lines = [
+            line
+            for line in self.lora_replay_window.splitlines()
+            if "Prefill batch" in line
+        ]
+        self.assertTrue(
+            any(PREFILL_GRAPH_REPLAY_PATTERN.search(line) for line in prefill_lines),
+            "LoRA probe request was not served by the prefill CUDA graph. "
+            f"Prefill lines in the probe window: {prefill_lines}",
+        )
+
+    def test_lora_prefill_logprobs_match_eager(self):
+        for i, prompt in enumerate(PROMPTS):
+            diff = _max_abs_diff(
+                self.with_graph["lora"][i], self.without_graph["lora"][i]
+            )
+            print(f"[lora] prompt {i} ({len(prompt)} chars): max_abs_diff={diff:.5f}")
+            self.assertLess(
+                diff,
+                GRAPH_VS_EAGER_TOLERANCE,
+                f"LoRA prompt logprobs diverge between breakable prefill graph "
+                f"and eager prefill for prompt {i}",
+            )
+
+    def test_base_prefill_logprobs_match_eager(self):
+        for i, prompt in enumerate(PROMPTS):
+            diff = _max_abs_diff(
+                self.with_graph["base"][i], self.without_graph["base"][i]
+            )
+            print(f"[base] prompt {i} ({len(prompt)} chars): max_abs_diff={diff:.5f}")
+            self.assertLess(
+                diff,
+                GRAPH_VS_EAGER_TOLERANCE,
+                f"Base-model prompt logprobs diverge between breakable prefill "
+                f"graph and eager prefill for prompt {i}",
+            )
+
+    def test_mixed_batch_matches_eager(self):
+        for i in range(len(MIXED_BATCH_PROMPTS)):
+            diff = _max_abs_diff(
+                self.with_graph["mixed"][i], self.without_graph["mixed"][i]
+            )
+            print(
+                f"[mixed] req {i} (lora={MIXED_BATCH_LORA_PATHS[i]}): "
+                f"max_abs_diff={diff:.5f}"
+            )
+            self.assertLess(
+                diff,
+                GRAPH_VS_EAGER_TOLERANCE,
+                f"Mixed LoRA/base batch logprobs diverge between breakable "
+                f"prefill graph and eager prefill for request {i}",
+            )
+
+    def test_lora_is_applied_under_graph(self):
+        # If the captured graph dropped the LoRA kernels (or read stale
+        # metadata), LoRA and base logprobs would coincide.
+        for i in range(len(PROMPTS)):
+            diff = _max_abs_diff(self.with_graph["lora"][i], self.with_graph["base"][i])
+            print(f"[effect] prompt {i}: lora-vs-base max_abs_diff={diff:.5f}")
+            self.assertGreater(
+                diff,
+                LORA_EFFECT_THRESHOLD,
+                f"Adapter did not change prompt logprobs under the breakable "
+                f"prefill graph for prompt {i}; LoRA was likely not applied.",
+            )
+
+
+class TestBCGLoRACsgmv(BCGLoRAServerMixin, CustomTestCase):
+    lora_backend = "csgmv"
+
+
+class TestBCGLoRATriton(BCGLoRAServerMixin, CustomTestCase):
+    lora_backend = "triton"
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py

@@ -30,6 +30,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
     def _make_prefill_runner(self, backend):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner._is_full_backend = False
+        runner.enable_lora = False
         runner.prefill_backend_name = backend
         runner.has_mha_companion_layers = backend == Backend.BREAKABLE
         runner.capture_hidden_mode = CaptureHiddenMode.NULL

test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py

@@ -19,6 +19,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
     def _make_runner(self):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner._is_full_backend = False
+        runner.enable_lora = False
         runner.prefill_backend_name = Backend.TC_PIECEWISE
         runner.has_mha_companion_layers = False
         runner.capture_hidden_mode = CaptureHiddenMode.NULL

test/registered/unit/server_args/test_server_args.py

@@ -1296,6 +1296,60 @@ class TestCudaGraphDisaggregationRoles(CustomTestCase):
         self.assertIn((Phase.DECODE, "backend"), args._cuda_graph_config_locked)
 
 
+class TestPrefillCudaGraphLoRACompatibility(CustomTestCase):
+    """LoRA no longer auto-disables the breakable prefill CUDA graph; guards
+    test_bcg_with_lora.py against a rule re-disabling it (vacuous pass)."""
+
+    def _handled_args(self, **overrides):
+        args = ServerArgs(model_path="dummy", **overrides)
+        args.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]),
+            is_piecewise_cuda_graph_disabled_model=False,
+            is_multimodal=False,
+            is_multimodal_piecewise_cuda_graph_supported=False,
+        )
+        with (
+            patch("sglang.srt.utils.is_cuda", return_value=True),
+            patch.object(ServerArgs, "use_mla_backend", return_value=False),
+        ):
+            args._handle_cuda_graph_config()
+        return args
+
+    def test_enable_lora_keeps_breakable_prefill_graph(self):
+        args = self._handled_args(enable_lora=True)
+
+        self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE)
+
+    def test_lora_paths_keep_breakable_prefill_graph(self):
+        args = self._handled_args(lora_paths=["dummy/lora-adapter"])
+
+        self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE)
+
+    def test_lora_still_disables_tc_piecewise_prefill_graph(self):
+        # Pin the tc_piecewise LoRA rule itself, with the hardware rule
+        # neutralized so this runs on CPU-only CI.
+        args = ServerArgs(model_path="dummy", enable_lora=True)
+        args.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]),
+            is_piecewise_cuda_graph_disabled_model=False,
+            is_multimodal=False,
+            is_multimodal_piecewise_cuda_graph_supported=False,
+        )
+        args.cuda_graph_config = CudaGraphConfig(
+            prefill=PhaseConfig(backend=Backend.TC_PIECEWISE)
+        )
+        with (
+            patch("sglang.srt.server_args.is_hip", return_value=False),
+            patch("sglang.srt.server_args.is_npu", return_value=False),
+            patch("sglang.srt.server_args.is_cpu", return_value=False),
+            patch("sglang.srt.server_args.is_mps", return_value=False),
+            patch("sglang.srt.server_args.is_xpu", return_value=False),
+        ):
+            args._disable_tc_piecewise_cudagraph_if_incompatible()
+
+        self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED)
+
+
 class TestBreakableCudaGraphMultimodalAllowlist(CustomTestCase):
     """The BCG "multimodal model" rule exempts archs on the BCG multimodal
     opt-in allowlist (multimodal_breakable_cuda_graph_supported_model_archs)."""