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

python/sglang/srt/lora/backend/ascend_backend.py

@@ -211,6 +211,7 @@ class AscendLoRABackend(BaseLoRABackend):
         lora_ranks: list[int],
         scalings: list[float],
         use_cuda_graph: bool,
+        use_prefill_cuda_graph: bool = False,
     ):
         # Use pinned memory to avoid synchronizations during host-to-device transfer
         weight_indices_tensor = torch.tensor(

python/sglang/srt/lora/backend/base_backend.py

@@ -19,11 +19,21 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
         device: the device where the backend runs.
     """
 
+    # Supporting backends implement init_prefill_cuda_graph_batch_info() and
+    # honor use_prefill_cuda_graph in prepare_lora_batch().
+    supports_prefill_cuda_graph: bool = False
+
     def __init__(self, max_loras_per_batch: int, device: torch.device):
         self.max_loras_per_batch = max_loras_per_batch
         self.device = device
         self.init_lm_head_config()
         self._is_moe_lora = False
+        # Static metadata read by prefill-CUDA-graph kernels, refreshed in
+        # place every prefill batch.
+        self.prefill_cuda_graph_batch_info: LoRABatchInfo | None = None
+        # Request/token caps for serving a batch from the static metadata.
+        self.prefill_cuda_graph_max_bs: int | None = None
+        self.prefill_cuda_graph_max_tokens: int | None = None
 
     def run_lora_a_embedding(
         self,
@@ -161,6 +171,13 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
         """
         pass
 
+    def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int):
+        """Allocate static LoRA batch metadata for the prefill CUDA graph,
+        sized for the largest captured token bucket. Called before capture."""
+        raise NotImplementedError(
+            f"LoRA backend {type(self).__name__} does not support the prefill CUDA graph."
+        )
+
     @property
     def is_moe_lora(self) -> bool:
         return self._is_moe_lora
@@ -317,18 +334,12 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
         lora_ranks: list[int],
         scalings: list[float],
         use_cuda_graph: bool,
+        use_prefill_cuda_graph: bool = False,
     ):
         """Prepare the lora weights and batch info for current forward batch.
 
-        This method provides a hook for each backend to conduct its own preparation
-        logic for each forward batch.
-
-        Args:
-            forward_batch: the ForwardBatch object for current forward pass
-            weight_indices: list of indices of lora weights to be applied for current batch
-            lora_ranks: list of lora ranks corresponding to weight_indices
-            scalings: list of scaling factors corresponding to weight_indices
-            use_cuda_graph: whether to use CUDA Graph for this batch
+        use_cuda_graph / use_prefill_cuda_graph select in-place updates of the
+        static decode / prefill CUDA graph batch info respectively.
         """
         pass
 

python/sglang/srt/lora/backend/chunked_backend.py

@@ -32,6 +32,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
     """
 
     name = "csgmv"
+    supports_prefill_cuda_graph = True
 
     def __init__(
         self,
@@ -240,6 +241,34 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
                 req_weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
             )
 
+    def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int):
+        # Worst-case chunk segments for any replay batch: ceil(N / chunk_top)
+        # (bounded by 16 for the small tiers) plus one per adapter group.
+        chunk_top = self._determine_chunk_size_for_tokens(max_num_tokens)
+        max_num_segments = (
+            max((max_num_tokens + chunk_top - 1) // chunk_top, 16)
+            + self.max_loras_per_batch
+        )
+        # Each extend request has >= 1 token, so bs is bounded by the bucket.
+        max_bs = max_num_tokens
+        with torch.device(self.device):
+            self.prefill_cuda_graph_batch_info = LoRABatchInfo(
+                bs=0,  # Set per batch
+                use_cuda_graph=True,
+                seg_lens=torch.zeros(max_num_segments, dtype=torch.int32),
+                seg_indptr=torch.zeros(max_num_segments + 1, dtype=torch.int32),
+                weight_indices=torch.zeros(max_num_segments, dtype=torch.int32),
+                permutation=torch.zeros(max_num_tokens, dtype=torch.int32),
+                lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
+                scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
+                num_segments=None,  # Set per batch
+                max_len=None,  # Set per batch (chunk size)
+                req_seg_indptr=torch.zeros(max_bs + 1, dtype=torch.int32),
+                req_weight_indices=torch.zeros(max_bs, dtype=torch.int32),
+            )
+        self.prefill_cuda_graph_max_bs = max_bs
+        self.prefill_cuda_graph_max_tokens = max_num_tokens
+
     def prepare_lora_batch(
         self,
         forward_batch: ForwardBatch,
@@ -247,6 +276,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
         lora_ranks: list[int],
         scalings: list[float],
         use_cuda_graph: bool,
+        use_prefill_cuda_graph: bool = False,
     ):
         chunk_size = self._determine_chunk_size(forward_batch)
 
@@ -276,7 +306,16 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
         max_num_segments = 0
         has_unused_cuda_graph_segments = False
 
-        if not use_cuda_graph:
+        if use_prefill_cuda_graph:
+            # In-place refresh of the prefill-sized static buffers; unused
+            # segment slots no-op via the padded seg_indptr tail.
+            batch_info = self.prefill_cuda_graph_batch_info
+            batch_info.bs = bs
+            batch_info.num_segments = num_segments
+            batch_info.max_len = chunk_size
+            max_num_segments = batch_info.weight_indices.shape[0]
+            has_unused_cuda_graph_segments = num_segments < max_num_segments
+        elif not use_cuda_graph:
             batch_info = LoRABatchInfo(
                 bs=bs,
                 num_segments=num_segments,
@@ -409,6 +448,8 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
 
         return dataclasses.replace(
             batch_info,
+            # lm_head LoRA runs in the eager tail, outside any captured graph.
+            use_cuda_graph=False,
             num_segments=num_segments,
             max_len=chunk_size,
             seg_indptr=seg_indptr,

python/sglang/srt/lora/backend/torch_backend.py

@@ -197,6 +197,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
         lora_ranks: list[int],
         scalings: list[float],
         use_cuda_graph: bool,
+        use_prefill_cuda_graph: bool = False,
     ):
         # Do not use merge optimization for graph mode
         # Use pinned memory to avoid synchronizations during host-to-device transfer

python/sglang/srt/lora/backend/triton_backend.py

@@ -16,9 +16,14 @@ from sglang.srt.lora.utils import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
+# Fixed segment slots (one per request) baked into the captured prefill LoRA
+# kernel grids; batches with more requests fall back to eager prefill.
+PREFILL_CUDA_GRAPH_LORA_SEGMENTS = 32
+
 
 class TritonLoRABackend(BaseLoRABackend):
     name = "triton"
+    supports_prefill_cuda_graph = True
 
     def __init__(
         self,
@@ -181,6 +186,27 @@ class TritonLoRABackend(BaseLoRABackend):
                 permutation=torch.zeros(max_tokens, dtype=torch.int32),
             )
 
+    def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int):
+        num_slots = PREFILL_CUDA_GRAPH_LORA_SEGMENTS
+        mlpb = self.max_loras_per_batch
+        with torch.device(self.device):
+            # bs pinned at num_slots so the captured grids cover any replay
+            # batch; slots past the live batch keep seg_lens == 0 and no-op.
+            self.prefill_cuda_graph_batch_info = LoRABatchInfo(
+                bs=num_slots,
+                use_cuda_graph=True,
+                num_segments=num_slots,
+                seg_lens=torch.zeros(num_slots, dtype=torch.int32),
+                seg_indptr=torch.zeros(num_slots + 1, dtype=torch.int32),
+                max_len=0,
+                weight_indices=torch.zeros(num_slots, dtype=torch.int32),
+                lora_ranks=torch.zeros(mlpb, dtype=torch.int32),
+                scalings=torch.zeros(mlpb, dtype=torch.float),
+                permutation=None,
+            )
+        self.prefill_cuda_graph_max_bs = num_slots
+        self.prefill_cuda_graph_max_tokens = max_num_tokens
+
     def compute_sgemm_routing(self, use_cuda_graph: bool):
         """Sort tokens by adapter and build merged segments for sgemm LoRA."""
         bi = self.batch_info
@@ -231,6 +257,7 @@ class TritonLoRABackend(BaseLoRABackend):
         lora_ranks: list[int],
         scalings: list[float],
         use_cuda_graph: bool,
+        use_prefill_cuda_graph: bool = False,
     ):
         # Use pinned memory to avoid synchronizations during host-to-device transfer
         weight_indices_tensor = torch.tensor(
@@ -252,6 +279,17 @@ class TritonLoRABackend(BaseLoRABackend):
             batch_info = self.cuda_graph_batch_info
             batch_info.bs = forward_batch.batch_size
             batch_info.num_segments = forward_batch.batch_size
+        elif use_prefill_cuda_graph:
+            batch_info = self.prefill_cuda_graph_batch_info
+            # bs stays pinned at the allocated slot count; slots past the
+            # live batch no-op via seg_lens == 0.
+            batch_info.num_segments = bs
+            batch_info.max_len = max(forward_batch.extend_seq_lens_cpu)
+            batch_info.seg_lens[:bs].copy_(
+                forward_batch.extend_seq_lens, non_blocking=True
+            )
+            batch_info.seg_lens[bs:].zero_()
+            torch.cumsum(batch_info.seg_lens, dim=0, out=batch_info.seg_indptr[1:])
         else:
             max_len = (
                 # Calculate max_len from the CPU copy to avoid D2H transfer.
@@ -364,6 +402,9 @@ class TritonLoRABackend(BaseLoRABackend):
 
         return dataclasses.replace(
             batch_info,
+            # lm_head LoRA runs in the eager tail outside any captured prefill
+            # graph, on freshly allocated pruned metadata.
+            use_cuda_graph=False,
             bs=num_segments,
             num_segments=num_segments,
             max_len=max(seg_lens_cpu),

python/sglang/srt/lora/lora_manager.py

@@ -89,6 +89,7 @@ class LoRAManager:
         self.pending_lora_load_events = {}
 
         self.eviction_policy = server_args.lora_eviction_policy
+        self.enable_dp_attention: bool = server_args.enable_dp_attention
         self._experts_shared_outer_override: Optional[bool] = (
             server_args.experts_shared_outer_loras
         )
@@ -138,6 +139,54 @@ class LoRAManager:
             init_lora_two_stream_resources(self.device)
         # ===== END TO BE REFACTORED ====
 
+    def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int):
+        """Allocate the static prefill-CUDA-graph LoRA metadata, sized by the
+        largest captured token bucket. Called before capture."""
+        self.lora_backend.init_prefill_cuda_graph_batch_info(
+            max_num_tokens=max_num_tokens
+        )
+
+    @property
+    def supports_prefill_cuda_graph(self) -> bool:
+        """Whether LoRA kernels can be captured into the prefill CUDA graph;
+        excludes MoE LoRA and DP attention."""
+        return (
+            self.lora_backend.supports_prefill_cuda_graph
+            and not self.lora_backend.is_moe_lora
+            and not self.enable_dp_attention
+        )
+
+    @property
+    def prefill_cuda_graph_max_bs(self) -> Optional[int]:
+        """Request-count cap for prefill-graph LoRA batches; None until
+        init_prefill_cuda_graph_batch_info() ran."""
+        return self.lora_backend.prefill_cuda_graph_max_bs
+
+    def can_use_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
+        """Whether this batch can use the static prefill-graph LoRA metadata;
+        shared by prepare_lora_batch and can_run_graph so they stay consistent."""
+        max_bs = self.lora_backend.prefill_cuda_graph_max_bs
+        max_tokens = self.lora_backend.prefill_cuda_graph_max_tokens
+        if max_bs is None or max_tokens is None:
+            return False
+        # DP attention: per-rank eligibility could diverge across ranks and
+        # desync collectives; keep LoRA prefill eager.
+        if self.enable_dp_attention:
+            return False
+        # Decode-CUDA-graph extend modes (TARGET_VERIFY, DLLM_EXTEND) are
+        # owned by the decode static batch info path.
+        if (
+            not forward_batch.forward_mode.is_extend()
+            or forward_batch.forward_mode.is_cuda_graph()
+        ):
+            return False
+        if forward_batch.extend_num_tokens is None:
+            return False
+        return (
+            forward_batch.batch_size <= max_bs
+            and forward_batch.extend_num_tokens <= max_tokens
+        )
+
     def init_cuda_graph_moe_buffers(
         self, max_bs: int, max_loras: int, compute_dtype, moe_layer
     ):
@@ -374,6 +423,11 @@ class LoRAManager:
             and bs <= self.max_bs_in_cuda_graph
             and forward_batch.forward_mode.is_cuda_graph()
         )
+        # Eligible extend batches refresh the static prefill batch info in
+        # place so captured kernels read current values at replay.
+        use_prefill_cuda_graph = not use_cuda_graph and self.can_use_prefill_cuda_graph(
+            forward_batch
+        )
 
         weight_indices = [0] * len(forward_batch.lora_ids)
         lora_ranks = [0] * self.max_loras_per_batch
@@ -394,6 +448,7 @@ class LoRAManager:
             lora_ranks=lora_ranks,
             scalings=scalings,
             use_cuda_graph=use_cuda_graph,
+            use_prefill_cuda_graph=use_prefill_cuda_graph,
         )
         self.lora_backend.batch_info.has_active_lora = any(
             lora_ranks[wi] > 0 for wi in weight_indices

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -180,6 +180,17 @@ def capture_prefill_graph(
         )
         return eager_runner
 
+    if (
+        model_runner.server_args.enable_lora
+        and not model_runner.lora_manager.supports_prefill_cuda_graph
+    ):
+        logger.warning(
+            "Disable prefill CUDA graph because the current LoRA "
+            "configuration does not support it (unsupported LoRA backend, "
+            "MoE LoRA, or DP attention)."
+        )
+        return eager_runner
+
     # Resolve the decoder once. Some VLM wrappers (for example Kimi-VL)
     # expose it as ``language_model`` rather than ``model``.
     try:

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -167,6 +167,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # --- model flags ----------------------------------------------
         self.quant_config = getattr(model_runner.model, "quant_config", None)
         self.is_multimodal = model_runner.model_config.is_multimodal
+        self.enable_lora = model_runner.server_args.enable_lora
         # Classification/reward forwards branch on return_pooled_hidden_states;
         # capture must use the same flag value as replay for those models.
         self.capture_return_pooled_hidden_states = not model_runner.is_generation
@@ -267,6 +268,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # default False; the assignment below sets the real value once the
         # backend type is known.
         self._is_full_backend = False
+        # Same ordering requirement: capture_prepare reads this.
+        self._capture_lora = False
         # TcPiecewise does its compile pass during backend construction.
         # Wrap only that path with the prefill CUDA graph failure hint.
         try:
@@ -286,6 +289,30 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 # Auto: scale request slots with the chunked prefill size.
                 max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1)
             self._capture_req_slots = min(max_req, self.max_bs)
+
+        # BCG/Full record LoRA kernels, so the metadata they read must live in
+        # static buffers refreshed in place per batch; unsupported LoRA
+        # configs were already routed to the eager runner.
+        self._capture_lora = self.enable_lora and isinstance(
+            self.backend, (BreakableCudaGraphBackend, FullCudaGraphBackend)
+        )
+        if self._capture_lora:
+            model_runner.lora_manager.init_prefill_cuda_graph_batch_info(
+                max_num_tokens=self.max_num_tokens
+            )
+            # Clamp Full's request slots to the LoRA segment-slot count
+            # rather than fail capture.
+            lora_max_bs = model_runner.lora_manager.prefill_cuda_graph_max_bs
+            if self._capture_req_slots > lora_max_bs:
+                logger.info(
+                    "Clamping full prefill CUDA graph request slots from %d to %d "
+                    "to fit the LoRA backend's static segment slots.",
+                    self._capture_req_slots,
+                    lora_max_bs,
+                )
+                self._capture_req_slots = lora_max_bs
+
+        if self._is_full_backend:
             self._full_cg_seq_lens_cpu = torch.zeros(
                 (self._capture_req_slots,), dtype=torch.int64, device="cpu"
             )
@@ -681,6 +708,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
         if self._is_full_backend and forward_batch.batch_size > self._capture_req_slots:
             return False
+        # LoRA batches may only replay the graph when prepare_lora_batch put
+        # their metadata in the static buffers (same predicate); keyed off
+        # enable_lora, not lora_ids, which is non-None even without LoRA.
+        if self.enable_lora and not (
+            self._capture_lora
+            and self.model_runner.lora_manager.can_use_prefill_cuda_graph(forward_batch)
+        ):
+            return False
         if forward_batch.input_embeds is not None:
             return False
         if forward_batch.replace_embeds is not None:
@@ -850,7 +885,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 num_token_non_padded=self._capture_num_token_non_padded(num_tokens),
                 num_token_non_padded_cpu=num_tokens,
                 global_forward_mode=ForwardMode.EXTEND,
-                lora_ids=None,
+                # All-None ids are safe: kernels no-op at rank 0 and replay
+                # refreshes the static batch info with live values.
+                lora_ids=([None] * bs if self._capture_lora else None),
                 return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
             )
             self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
@@ -895,6 +932,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         """
         num_tokens = size
         forward_batch, attn_backend = self.capture_prepare(num_tokens)
+        if forward_batch.lora_ids is not None:
+            # Fill the static prefill LoRA batch info the captured kernels
+            # will read (all-None ids: ranks stay 0, kernels no-op).
+            lora_manager = self.model_runner.lora_manager
+            assert lora_manager.can_use_prefill_cuda_graph(forward_batch), (
+                f"Capture batch (req slots {self._capture_req_slots}, bucket "
+                f"{num_tokens}) exceeds the LoRA backend's prefill CUDA graph "
+                "limits; the graph would read stale LoRA metadata at replay."
+            )
+            lora_manager.prepare_lora_batch(forward_batch)
         if self._is_full_backend:
             attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
         else:

python/sglang/srt/server_args.py

@@ -4280,6 +4280,8 @@ class ServerArgs:
                 "MoE A2A backend",
                 lambda: _resolved_view(self).moe_a2a_backend != "none",
             ),
+            # Dynamo blocks LoRA under tc_piecewise (per-batch LoRABatchInfo
+            # rebinds break guards); breakable/full support LoRA.
             ("LoRA", lambda: bool(self.lora_paths) or self.enable_lora),
             (
                 "multimodal model",
@@ -4355,8 +4357,6 @@ class ServerArgs:
                 "decode context parallel (dcp_size > 1)",
                 lambda: self.dcp_size > 1,
             ),
-            # BCG capture + LoRA adapter weights exceed host RAM headroom.
-            ("LoRA", lambda: bool(self.lora_paths) or bool(self.enable_lora)),
             # BCG bucket sizes exceed FlashInfer MoE A2A's dispatch cap.
             (
                 "MoE A2A backend",

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)."""