[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",