Repository: sgl-project/sglang
Pull request #4115: Feat: support cuda graph for LoRA
Author: login: Qiaolin-Yu
Created: 2025-03-06T02:01:18Z | Merged: 2025-04-29T06:30:44Z

<!-- Thank you for your contribution! We appreciate it. The following guidelines will help improve your pull request and facilitate feedback. If anything is unclear, don't hesitate to submit your pull request and ask the maintainers for assistance. -->

## Motivation

<!-- Explain the purpose of this PR and the goals it aims to achieve. -->
Close #3282 

## Modifications
- Pin memory and do in-place updates when CUDA graph is enabled.
- The cuda graph has been enabled for lora by default.

## Benchmarking result
To do benchmarking for lora, run this command to launch server:
```
# Triton backend
python benchmark/lora/launch_server.py --max-loras-per-batch 4 --lora-backend triton
```
Then run this command to request test from client:
```
python benchmark/lora/lora_bench.py
```

Benchmark configurations:

base model: meta-llama/Llama-2-7b-hf
- lora adapter: winddude/wizardLM-LlaMA-LoRA-7B
- GPU: Nvidia H100
- maximum number of serving loras: 4
- number of requests: 50
- input length: uniform random distribution on [1, 1024]
- output length: uniform random distribution on [1, 128]

| Backend    | Total Throughput (tok/s) | Mean E2E Latency (ms) |
|------------|---------------------------|-----------------------|
| Triton (without cuda graph)   |      2188.78             |      7050.41         |
| Triton (with cuda graph)     |        3786.79           |      4173.12         |

## Checklist

- [x] Format your code according to the [Code Formatting with Pre-Commit](https://docs.sglang.ai/references/contribution_guide.html#code-formatting-with-pre-commit).
- [ ] Add unit tests as outlined in the [Running Unit Tests](https://docs.sglang.ai/references/contribution_guide.html#running-unit-tests-adding-to-ci).
- [ ] Update documentation / docstrings / example tutorials as needed, according to [Writing Documentation](https://docs.sglang.ai/references/contribution_guide.html#writing-documentation-running-docs-ci).
- [ ] Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to [Benchmark and Profiling](https://docs.sglang.ai/references/benchmark_and_profiling.html) and [Accuracy Results](https://docs.sglang.ai/references/accuracy_evaluation.html).
- [ ] For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
- [ ] Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.


Linked issue #3282: [Feature] Support compatibility between Cuda Graph and Lora
### Checklist

- [x] 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed.
- [x] 2. Please use English, otherwise it will be closed.

### Motivation

Currently Lora and Cuda Graph cannot be used at the same time. 

Theoretically they should be compatible.

### Related resources

_No response_
Linked issue comments: user: login: zhaochenyang20
created_at: 2025-02-04T07:41:23Z
body: Great!
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-04-27T00:20:11Z
body: This issue has been automatically closed due to inactivity. Please feel free to reopen it if needed.
author_association: CONTRIBUTOR

Diff:
diff --git a/benchmark/lora/launch_server.py b/benchmark/lora/launch_server.py
index 33ad848a6c41..6ed85d9d283c 100644
--- a/benchmark/lora/launch_server.py
+++ b/benchmark/lora/launch_server.py
@@ -19,7 +19,7 @@ def launch_server(args):
         for i in range(NUM_LORAS):
             lora_name = f"lora{i}"
             cmd += f"{lora_name}={lora_path} "
-    cmd += f"--disable-radix --disable-cuda-graph "
+    cmd += f"--disable-radix "
     cmd += f"--max-loras-per-batch {args.max_loras_per_batch} "
     cmd += f"--max-running-requests {args.max_running_requests} "
     cmd += f"--lora-backend {args.lora_backend} "
diff --git a/docs/backend/lora.ipynb b/docs/backend/lora.ipynb
index 9945e60b617e..ffe2f48f65be 100644
--- a/docs/backend/lora.ipynb
+++ b/docs/backend/lora.ipynb
@@ -77,7 +77,7 @@
     "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n",
     "    --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n",
     "    --max-loras-per-batch 1 --lora-backend triton \\\n",
-    "    --disable-cuda-graph --disable-radix-cache\n",
+    "    --disable-radix-cache\n",
     "\"\"\"\n",
     ")\n",
     "\n",
@@ -136,7 +136,7 @@
     "    --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n",
     "    lora1=Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16 \\\n",
     "    --max-loras-per-batch 2 --lora-backend triton \\\n",
-    "    --disable-cuda-graph --disable-radix-cache\n",
+    "    --disable-radix-cache\n",
     "\"\"\"\n",
     ")\n",
     "\n",
@@ -182,7 +182,7 @@
    "source": [
     "## Future Works\n",
     "\n",
-    "The development roadmap for LoRA-related features can be found in this [issue](https://github.com/sgl-project/sglang/issues/2929). Currently Cuda graph and radix attention are not incompatible with LoRA and must be manually disabled. Other features, including Unified Paging, Cutlass backend, and dynamic loading/unloadingm, are still under development."
+    "The development roadmap for LoRA-related features can be found in this [issue](https://github.com/sgl-project/sglang/issues/2929). Currently radix attention is incompatible with LoRA and must be manually disabled. Other features, including Unified Paging, Cutlass backend, and dynamic loading/unloadingm, are still under development."
    ]
   }
  ],
diff --git a/docs/backend/server_arguments.md b/docs/backend/server_arguments.md
index 3d75f180c62a..c6e1e4850d38 100644
--- a/docs/backend/server_arguments.md
+++ b/docs/backend/server_arguments.md
@@ -160,7 +160,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
 
 | Arguments | Description | Defaults |
 |----------|-------------|---------|
-| `lora_paths` | List of adapters to apply to your model. Each batch element uses the proper LoRA adapter. `cuda_graph` and `radix_attention` are not supported with this, so they must be disabled manually. See related [issues](https://github.com/sgl-project/sglang/issues/2929). | None |
+| `lora_paths` | List of adapters to apply to your model. Each batch element uses the proper LoRA adapter. `radix_attention` is not supported with this, so it must be disabled manually. See related [issues](https://github.com/sgl-project/sglang/issues/2929). | None |
 | `max_loras_per_batch` | Maximum number of LoRAs allowed in a running batch, including the base model. | `8` |
 | `lora_backend` | Backend used to run GEMM kernels for LoRA modules. Can be `triton` or `flashinfer`. | `triton` |
 
diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py
index 8da81826f0c4..aa10ef6b7723 100644
--- a/python/sglang/srt/lora/layers.py
+++ b/python/sglang/srt/lora/layers.py
@@ -136,11 +136,19 @@ def set_lora_info(
         self.set_lora = True
         self.A_buffer_gate_up = A_buffer
         if self.lora_backend.fuse_stacked_lora_b:
-            # TODO: avoid using contiguous() in GPU.
             # B_buffer_gate_up: (num_lora, 2 * output_dim, r)
-            self.B_buffer_gate_up = torch.cat(
-                (B_buffer[0], B_buffer[1]), dim=-2
-            ).contiguous()
+            if not hasattr(self, "B_buffer_gate_up") or self.B_buffer_gate_up is None:
+                self.B_buffer_gate_up = torch.empty(
+                    (
+                        B_buffer[0].shape[0],
+                        2 * B_buffer[0].shape[1],
+                        B_buffer[0].shape[2],
+                    ),
+                    dtype=B_buffer[0].dtype,
+                    device=B_buffer[0].device,
+                )
+            self.B_buffer_gate_up[:, : B_buffer[0].shape[1], :].copy_(B_buffer[0])
+            self.B_buffer_gate_up[:, B_buffer[0].shape[1] :, :].copy_(B_buffer[1])
         else:
             self.B_buffer_gate_up = (B_buffer[0], B_buffer[1])
 
@@ -171,7 +179,7 @@ def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
 
 
 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
-    def init__(
+    def __init__(
         self,
         base_layer: QKVParallelLinear,
         lora_backend: BaseLoRABackend,
@@ -194,12 +202,30 @@ def set_lora_info(
             output_dim_q, output_dim_kv = B_buffer_q.shape[-2], B_buffer_kv.shape[-2]
 
             # B_buffer_qkv: (num_lora, output_dim_q + 2 * output_dim_kv, r)
-            self.B_buffer_qkv = torch.cat(
-                (B_buffer_q[0], B_buffer_kv[0], B_buffer_kv[1]), dim=-2
-            ).contiguous()
+            if not hasattr(self, "B_buffer_qkv") or self.B_buffer_qkv is None:
+                self.B_buffer_qkv = torch.empty(
+                    (
+                        B_buffer_q[0].shape[0],
+                        output_dim_q + 2 * output_dim_kv,
+                        B_buffer_q[0].shape[2],
+                    ),
+                    dtype=B_buffer_q[0].dtype,
+                    device=B_buffer_q[0].device,
+                )
+            self.B_buffer_qkv[:, :output_dim_q, :].copy_(B_buffer_q[0])
+            self.B_buffer_qkv[:, output_dim_q : output_dim_q + output_dim_kv, :].copy_(
+                B_buffer_kv[0]
+            )
+            self.B_buffer_qkv[:, output_dim_q + output_dim_kv :, :].copy_(
+                B_buffer_kv[1]
+            )
 
             # Offsets of q/k/v in output dimension
-            self.output_offset = torch.tensor(
+            if not hasattr(self, "output_offset") or self.output_offset is None:
+                self.output_offset = torch.empty(
+                    4, dtype=torch.int32, device=B_buffer_q.device
+                )
+            self.output_offset[:4] = torch.tensor(
                 [
                     0,
                     output_dim_q,
diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index 7b8f1162964a..fb349d11a328 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -72,6 +72,23 @@ def __init__(
         self.init_loras()
         self.init_lora_memory_pool()
 
+    def init_cuda_graph_batch_info(self, max_bs_in_cuda_graph: int):
+        self.max_bs_in_cuda_graph = max_bs_in_cuda_graph
+        with torch.device("cuda"):
+            self.cuda_graph_batch_info = LoRABatchInfo(
+                bs=self.max_bs_in_cuda_graph,
+                seg_lens=torch.zeros(self.max_bs_in_cuda_graph, dtype=torch.int32),
+                seg_indptr=torch.zeros(
+                    self.max_bs_in_cuda_graph + 1, dtype=torch.int32
+                ),
+                max_len=0,
+                weight_indices=torch.zeros(
+                    self.max_bs_in_cuda_graph, 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),
+            )
+
     def init_loras(self):
         # Config of each LoRA adapter
         self.configs: Dict[str, LoRAConfig] = {}
@@ -140,39 +157,73 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled. Note that
+            # if CUDA graph is enabled, the batch whose bs <= max_bs_in_cuda_graph
+            # will also use these preallocated buffers, no matter whether
+            # the batch can use CUDA graph or not.
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
+                self.cuda_graph_batch_info.seg_lens[:bs].copy_(
+                    forward_batch.extend_seq_lens
+                )
+            else:
+                self.cuda_graph_batch_info.seg_lens[:bs].fill_(1)
+            torch.cumsum(
+                self.cuda_graph_batch_info.seg_lens[:bs],
+                dim=0,
+                out=self.cuda_graph_batch_info.seg_indptr[1 : bs + 1],
+            )
+            self.cuda_graph_batch_info.max_len = int(
+                torch.max(self.cuda_graph_batch_info.seg_lens[:bs])
+            )
+
+            for i, lora_path in enumerate(forward_batch.lora_paths):
+                self.cuda_graph_batch_info.weight_indices[i] = (
+                    self.memory_pool.get_buffer_id(lora_path)
+                )
+                lora = self.loras[lora_path]
+                self.cuda_graph_batch_info.lora_ranks[
+                    self.cuda_graph_batch_info.weight_indices[i]
+                ] = lora.config.hf_config["r"]
+                self.cuda_graph_batch_info.scalings[
+                    self.cuda_graph_batch_info.weight_indices[i]
+                ] = lora.scaling
+            batch_info = self.cuda_graph_batch_info
+        else:
+            seg_lens = (
+                forward_batch.extend_seq_lens
+                if forward_batch.forward_mode.is_extend()
+                else torch.ones(bs, device=self.device)
+            )
+            seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
+            seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+            max_len = int(torch.max(seg_lens))
+            weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
+
+            lora_ranks = torch.empty(
+                (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
+            )
+            scalings = torch.empty(
+                (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
+            )
+            for i, lora_path in enumerate(forward_batch.lora_paths):
+                weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
+                lora = self.loras[lora_path]
+                lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
+                scalings[weight_indices[i]] = lora.scaling
+            batch_info = LoRABatchInfo(
+                bs=bs,
+                seg_lens=seg_lens,
+                seg_indptr=seg_indptr,
+                max_len=max_len,
+                weight_indices=weight_indices,
+                lora_ranks=lora_ranks,
+                scalings=scalings,
+            )
         self.lora_backend.set_batch_info(batch_info)
 
         # call set_lora_info for each lora modules
diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py
index ead9beade9a0..4b20071d5ec2 100644
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -220,6 +220,9 @@ def __init__(self, model_runner: ModelRunner):
         if self.enable_torch_compile:
             set_torch_compile_config()
 
+        if self.model_runner.server_args.lora_paths is not None:
+            self.model_runner.lora_manager.init_cuda_graph_batch_info(self.max_bs)
+
         # Graph inputs
         with torch.device("cuda"):
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
@@ -403,6 +406,13 @@ def capture_one_batch_size(self, bs: int, forward: Callable):
             self.capture_hidden_mode = (
                 spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
             )
+        if self.model_runner.server_args.lora_paths is not None:
+            # Currently, if the lora_path in `lora_paths` is None, the lora backend will use a
+            # different logic to handle lora, so we need to set `lora_paths` to a list of non-None
+            # values if lora is enabled.
+            lora_paths = [next(iter(self.model_runner.server_args.lora_paths))] * bs
+        else:
+            lora_paths = None
 
         forward_batch = ForwardBatch(
             forward_mode=self.capture_forward_mode,
@@ -424,8 +434,12 @@ def capture_one_batch_size(self, bs: int, forward: Callable):
             spec_algorithm=self.model_runner.spec_algorithm,
             spec_info=spec_info,
             capture_hidden_mode=self.capture_hidden_mode,
+            lora_paths=lora_paths,
         )
 
+        if lora_paths is not None:
+            self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
+
         # Attention backend
         self.model_runner.attn_backend.init_forward_metadata_capture_cuda_graph(
             bs,
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 7d03e0032248..eace0ade5df2 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -1242,7 +1242,6 @@ def check_server_args(self):
         assert (
             self.max_loras_per_batch > 0
             # FIXME
-            and (self.lora_paths is None or self.disable_cuda_graph)
             and (self.lora_paths is None or self.disable_radix_cache)
         ), "compatibility of lora and cuda graph and radix attention is in progress"
         assert self.base_gpu_id >= 0, "base_gpu_id must be non-negative"
diff --git a/test/srt/models/lora/test_lora_backend.py b/test/srt/models/lora/test_lora_backend.py
index e159f8c6b6b2..2217e8dd6fa5 100644
--- a/test/srt/models/lora/test_lora_backend.py
+++ b/test/srt/models/lora/test_lora_backend.py
@@ -24,7 +24,7 @@
     DEFAULT_PROMPTS,
     TORCH_DTYPES,
     LoRAModelCase,
-    run_batch_lora_test,
+    run_lora_test_one_by_one,
 )
 
 from sglang.test.test_utils import CustomTestCase, is_in_ci
@@ -42,7 +42,7 @@ def _run_backend_on_model_cases(self, model_cases: List[LoRAModelCase]):
             )
             for torch_dtype in TORCH_DTYPES:
                 for backend in BACKENDS:
-                    run_batch_lora_test(
+                    run_lora_test_one_by_one(
                         prompts,
                         model_case,
                         torch_dtype,
diff --git a/test/srt/models/lora/test_lora_cuda_graph.py b/test/srt/models/lora/test_lora_cuda_graph.py
new file mode 100644
index 000000000000..ba68df59a315
--- /dev/null
+++ b/test/srt/models/lora/test_lora_cuda_graph.py
@@ -0,0 +1,110 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+from utils import (
+    ALL_OTHER_LORA_MODELS,
+    CI_LORA_MODELS,
+    DEFAULT_PROMPTS,
+    TORCH_DTYPES,
+    LoRAModelCase,
+    run_lora_test_by_batch,
+    run_lora_test_one_by_one,
+)
+
+from sglang.test.test_utils import CustomTestCase, is_in_ci
+
+TEST_CUDA_GRAPH_PADDING_PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+    "Computer science is the study of",
+]
+
+
+class TestLoRACudaGraph(CustomTestCase):
+
+    def _run_without_cuda_graph_on_model_cases(self, model_cases: List[LoRAModelCase]):
+        # Since we have already enabled CUDA graph by default in other lora tests,
+        # we only need to run lora tests without CUDA graph here.
+        for model_case in model_cases:
+            # If skip_long_prompt is True, filter out prompts longer than 1000 characters
+            prompts = (
+                DEFAULT_PROMPTS
+                if not model_case.skip_long_prompt
+                else [p for p in DEFAULT_PROMPTS if len(p) < 1000]
+            )
+            for torch_dtype in TORCH_DTYPES:
+                run_lora_test_one_by_one(
+                    prompts,
+                    model_case,
+                    torch_dtype,
+                    max_new_tokens=32,
+                    backend="triton",
+                    disable_cuda_graph=True,
+                    test_tag="without_cuda_graph",
+                )
+
+    def _run_cuda_graph_padding_on_model_cases(self, model_cases: List[LoRAModelCase]):
+        for model_case in model_cases:
+            # Run a batch size of 3, which will not be captured by CUDA graph and need padding
+            prompts = TEST_CUDA_GRAPH_PADDING_PROMPTS
+            for torch_dtype in TORCH_DTYPES:
+                run_lora_test_by_batch(
+                    prompts,
+                    model_case,
+                    torch_dtype,
+                    max_new_tokens=32,
+                    backend="triton",
+                    disable_cuda_graph=False,
+                    test_tag="cuda_graph_padding",
+                )
+
+    def test_ci_lora_models(self):
+        self._run_without_cuda_graph_on_model_cases(CI_LORA_MODELS)
+        self._run_cuda_graph_padding_on_model_cases(CI_LORA_MODELS)
+
+    def test_all_lora_models(self):
+        if is_in_ci():
+            return
+
+        # Retain ONLY_RUN check here
+        filtered_models = []
+        for model_case in ALL_OTHER_LORA_MODELS:
+            if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
+                continue
+            filtered_models.append(model_case)
+
+        self._run_without_cuda_graph_on_model_cases(filtered_models)
+        self._run_cuda_graph_padding_on_model_cases(filtered_models)
+
+
+if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
+    unittest.main(warnings="ignore")
diff --git a/test/srt/models/lora/test_lora_tp.py b/test/srt/models/lora/test_lora_tp.py
index 3d3f2ac0b9a1..51a552d78453 100644
--- a/test/srt/models/lora/test_lora_tp.py
+++ b/test/srt/models/lora/test_lora_tp.py
@@ -23,7 +23,7 @@
     DEFAULT_PROMPTS,
     TORCH_DTYPES,
     LoRAModelCase,
-    run_batch_lora_test,
+    run_lora_test_one_by_one,
 )
 
 from sglang.test.test_utils import CustomTestCase, is_in_ci
@@ -43,7 +43,7 @@ def _run_tp_on_model_cases(self, model_cases: List[LoRAModelCase]):
             for tp_size in tp_list:
                 model_case.tp_size = tp_size
                 for torch_dtype in TORCH_DTYPES:
-                    run_batch_lora_test(
+                    run_lora_test_one_by_one(
                         prompts,
                         model_case,
                         torch_dtype,
diff --git a/test/srt/models/lora/test_multi_lora_backend.py b/test/srt/models/lora/test_multi_lora_backend.py
index c40a01d6f192..eb338dfafa4c 100644
--- a/test/srt/models/lora/test_multi_lora_backend.py
+++ b/test/srt/models/lora/test_multi_lora_backend.py
@@ -22,7 +22,7 @@
     TORCH_DTYPES,
     LoRAAdaptor,
     LoRAModelCase,
-    run_batch_lora_test,
+    run_lora_test_one_by_one,
 )
 
 from sglang.test.test_utils import CustomTestCase, is_in_ci
@@ -89,7 +89,7 @@ def _run_multi_lora_test_on_model_cases(self, model_cases: List[LoRAModelCase]):
             )
             for torch_dtype in TORCH_DTYPES:
                 for backend in BACKENDS:
-                    run_batch_lora_test(
+                    run_lora_test_one_by_one(
                         batch_prompts,
                         model_case,
                         torch_dtype,
diff --git a/test/srt/models/lora/utils.py b/test/srt/models/lora/utils.py
index af68709f48f2..69d4cb37e9b7 100644
--- a/test/srt/models/lora/utils.py
+++ b/test/srt/models/lora/utils.py
@@ -94,19 +94,20 @@ def __post_init__(self):
 ]
 
 
-def run_batch_lora_test(
+def run_lora_test_one_by_one(
     prompts: List[str],
     model_case: LoRAModelCase,
     torch_dtype: torch.dtype,
     max_new_tokens: int,
     backend: str,
-    disable_cuda_graph: bool = True,
+    disable_cuda_graph: bool = False,
     disable_radix_cache: bool = True,
     mem_fraction_static: float = 0.88,
     test_tag: str = "",
 ):
     """
-    Run Lora test for a forward batch.
+    Input a batch of prompts, and run lora tests one by one with several generate requests
+    (each request will have bs=1).
     For prompt0, prompt1, ..., promptN,
     we will use adaptor0, adaptor1, ..., adaptorN included in model case,
     We will then compare the outputs of HF and SRT with and without LoRA.
@@ -119,7 +120,7 @@ def run_batch_lora_test(
         torch_dtype (torch.dtype): The torch dtype to use.
         max_new_tokens (int): The maximum number of new tokens to generate.
         backend (str): The lora backend to use.
-        disable_cuda_graph (bool, optional): Whether to disable CUDA graph. Defaults to True.
+        disable_cuda_graph (bool, optional): Whether to disable CUDA graph. Defaults to False.
         disable_radix_cache (bool, optional): Whether to disable radix cache. Defaults to True.
         mem_fraction_static (float, optional): The fraction of memory to use. Defaults to 0.88.
         test_tag (str, optional): The tag to use for the test. Defaults to "".
@@ -237,3 +238,112 @@ def run_batch_lora_test(
                 f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
                 f"for base '{base_path}', adaptor '{adaptor_names}', backend '{backend}', prompt: '{prompts[0][:50]}...'"
             )
+
+
+def run_lora_test_by_batch(
+    prompts: List[str],
+    model_case: LoRAModelCase,
+    torch_dtype: torch.dtype,
+    max_new_tokens: int,
+    backend: str,
+    disable_cuda_graph: bool = False,
+    disable_radix_cache: bool = True,
+    mem_fraction_static: float = 0.88,
+    test_tag: str = "",
+):
+    """
+    Run lora tests as a batch.
+    For prompt0, prompt1, ..., promptN,
+    we will use adaptor0, adaptor1, ..., adaptorN included in model case,
+    We will then compare the outputs of HF and SRT with LoRA.
+    If number of prompts is larger than number of adaptors,
+    the prompt i will use adaptor i % (number of adaptors).
+
+    Args:
+        prompts (List[str]): The batch of prompts to test.
+        model_case (LoRAModelCase): The model case to test.
+        torch_dtype (torch.dtype): The torch dtype to use.
+        max_new_tokens (int): The maximum number of new tokens to generate.
+        backend (str): The lora backend to use.
+        disable_cuda_graph (bool, optional): Whether to disable CUDA graph. Defaults to False.
+        disable_radix_cache (bool, optional): Whether to disable radix cache. Defaults to True.
+        mem_fraction_static (float, optional): The fraction of memory to use. Defaults to 0.88.
+        test_tag (str, optional): The tag to use for the test. Defaults to "".
+    """
+    base_path = model_case.base
+
+    # Create used adaptors for each prompt in batch
+    i, adaptors = 0, []
+    for _ in range(len(prompts)):
+        adaptors.append(model_case.adaptors[i])
+        i = (i + 1) % len(model_case.adaptors)
+    adaptor_names = [adaptor.name for adaptor in adaptors]
+
+    print(
+        f"\n========== Testing {test_tag} on base '{model_case.base}' with backend={backend}, dtype={torch_dtype} --- "
+        f"Using prompts {[p[:50] for p in prompts]} with adaptors: {adaptor_names} ---"
+    )
+    with SRTRunner(
+        base_path,
+        torch_dtype=torch_dtype,
+        model_type="generation",
+        tp_size=model_case.tp_size,
+        lora_paths=[adaptor.name for adaptor in model_case.adaptors],
+        max_loras_per_batch=model_case.max_loras_per_batch,
+        lora_backend=backend,
+        disable_cuda_graph=disable_cuda_graph,
+        disable_radix_cache=disable_radix_cache,
+        mem_fraction_static=mem_fraction_static,
+    ) as srt_runner:
+        srt_outputs = srt_runner.batch_forward(
+            prompts, max_new_tokens=max_new_tokens, lora_paths=adaptor_names
+        )
+
+    with SRTRunner(
+        base_path,
+        torch_dtype=torch_dtype,
+        model_type="generation",
+        tp_size=model_case.tp_size,
+        mem_fraction_static=mem_fraction_static,
+    ) as srt_runner:
+        srt_no_lora_outputs = srt_runner.batch_forward(
+            prompts, max_new_tokens=max_new_tokens
+        )
+
+    with HFRunner(
+        base_path, torch_dtype=torch_dtype, model_type="generation"
+    ) as hf_runner:
+        hf_outputs = hf_runner.forward(
+            prompts, max_new_tokens=max_new_tokens, lora_paths=adaptor_names
+        )
+
+    with HFRunner(
+        base_path, torch_dtype=torch_dtype, model_type="generation"
+    ) as hf_runner:
+        hf_no_lora_outputs = hf_runner.forward(
+            prompts,
+            max_new_tokens=max_new_tokens,
+        )
+
+    for i in range(len(prompts)):
+
+        srt_output_str = srt_outputs.output_strs[i].strip()
+        hf_output_str = hf_outputs.output_strs[i].strip()
+        rouge_score = calculate_rouge_l([srt_output_str], [hf_output_str])[0]
+        print("ROUGE-L score:", rouge_score)
+        print("SRT output:", srt_output_str)
+        print("HF output:", hf_output_str)
+        print("SRT no lora output:", srt_no_lora_outputs.output_strs[i].strip())
+        print("HF no lora output:", hf_no_lora_outputs.output_strs[i].strip())
+        assert srt_outputs.output_strs[i].strip(" ") == hf_outputs.output_strs[i].strip(
+            " "
+        ), (
+            srt_outputs.output_strs[i].strip(" "),
+            hf_outputs.output_strs[i].strip(" "),
+        )
+        assert srt_no_lora_outputs.output_strs[i].strip(
+            " "
+        ) == hf_no_lora_outputs.output_strs[i].strip(" "), (
+            srt_no_lora_outputs.output_strs[i].strip(" "),
+            hf_no_lora_outputs.output_strs[i].strip(" "),
+        )
diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py
index 659dd2814deb..0c0ed04fd1e7 100644
--- a/test/srt/run_suite.py
+++ b/test/srt/run_suite.py
@@ -80,6 +80,7 @@ class TestFile:
         TestFile("test_vlm_accuracy.py", 60),
         TestFile("test_vision_openai_server.py", 637),
         TestFile("test_w8a8_quantization.py", 46),
+        TestFile("models/lora/test_lora_cuda_graph.py", 250),
     ],
     "per-commit-2-gpu": [
         TestFile("models/lora/test_lora_tp.py", 116),


Reviews:
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T04:16:29Z
user: login: Qiaolin-Yu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T19:29:07Z
user: login: Qiaolin-Yu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T19:32:55Z
user: login: Qiaolin-Yu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T19:36:02Z
user: login: Qiaolin-Yu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T19:36:52Z
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T20:02:17Z
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T20:02:44Z
user: login: Qiaolin-Yu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T20:15:26Z
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T20:38:50Z
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-28T01:33:45Z
user: login: Fridge003
body: LGTM! Great Job!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-28T01:34:35Z
user: login: Ying1123
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T06:29:16Z

Review comments:
diff_hunk: @@ -138,9 +138,18 @@ def set_lora_info(
         if self.lora_backend.fuse_stacked_lora_b:
             # TODO: avoid using contiguous() in GPU.
             # B_buffer_gate_up: (num_lora, 2 * output_dim, r)
-            self.B_buffer_gate_up = torch.cat(
-                (B_buffer[0], B_buffer[1]), dim=-2
-            ).contiguous()
+            if not hasattr(self, "B_buffer_gate_up") or self.B_buffer_gate_up is None:
+                self.B_buffer_gate_up = torch.empty(
+                    (
+                        B_buffer[0].shape[0],
+                        2 * B_buffer[0].shape[1],
+                        B_buffer[0].shape[2],
+                    ),
+                    dtype=B_buffer[0].dtype,
+                    device=B_buffer[0].device,
+                ).contiguous()
path: python/sglang/srt/lora/layers.py
user: login: Fridge003
body: no need for calling contiguous here if you use torch.empty
created_at: 2025-04-27T03:56:50Z
author_association: COLLABORATOR
diff_hunk: @@ -194,9 +203,23 @@ def set_lora_info(
             output_dim_q, output_dim_kv = B_buffer_q.shape[-2], B_buffer_kv.shape[-2]
 
             # B_buffer_qkv: (num_lora, output_dim_q + 2 * output_dim_kv, r)
-            self.B_buffer_qkv = torch.cat(
-                (B_buffer_q[0], B_buffer_kv[0], B_buffer_kv[1]), dim=-2
-            ).contiguous()
+            if not hasattr(self, "B_buffer_qkv") or self.B_buffer_qkv is None:
+                self.B_buffer_qkv = torch.empty(
+                    (
+                        B_buffer_q[0].shape[0],
+                        output_dim_q + 2 * output_dim_kv,
+                        B_buffer_q[0].shape[2],
+                    ),
+                    dtype=B_buffer_q[0].dtype,
+                    device=B_buffer_q[0].device,
+                ).contiguous()
path: python/sglang/srt/lora/layers.py
user: login: Fridge003
body: no need for calling contiguous here if you use torch.empty
created_at: 2025-04-27T03:57:28Z
author_association: COLLABORATOR
diff_hunk: @@ -138,9 +138,18 @@ def set_lora_info(
         if self.lora_backend.fuse_stacked_lora_b:
             # TODO: avoid using contiguous() in GPU.
path: python/sglang/srt/lora/layers.py
user: login: Fridge003
body: This TODO can be  removed.
created_at: 2025-04-27T03:58:01Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
path: python/sglang/srt/lora/lora_manager.py
user: login: Fridge003
body: Cuda graph will not be used during extend, so this branch can be removed.
created_at: 2025-04-27T04:07:42Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
+                self.cuda_graph_batch_info.seg_lens[:bs].copy_(
+                    forward_batch.extend_seq_lens
+                )
+            else:
+                self.cuda_graph_batch_info.seg_lens[:bs].fill_(1)
+            self.cuda_graph_batch_info.seg_indptr[0] = 0
path: python/sglang/srt/lora/lora_manager.py
user: login: Fridge003
body: This line can be removed since the first element of seg_indptr is always zero.
created_at: 2025-04-27T04:10:53Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
+                self.cuda_graph_batch_info.seg_lens[:bs].copy_(
+                    forward_batch.extend_seq_lens
+                )
+            else:
+                self.cuda_graph_batch_info.seg_lens[:bs].fill_(1)
+            self.cuda_graph_batch_info.seg_indptr[0] = 0
+            torch.cumsum(
+                self.cuda_graph_batch_info.seg_lens[:bs],
+                dim=0,
+                out=self.cuda_graph_batch_info.seg_indptr[1 : bs + 1],
+            )
+            self.cuda_graph_batch_info.max_len = int(
+                torch.max(self.cuda_graph_batch_info.seg_lens[:bs])
+            )
+
+            for i, lora_path in enumerate(forward_batch.lora_paths):
+                self.cuda_graph_batch_info.weight_indices[i] = (
path: python/sglang/srt/lora/lora_manager.py
user: login: Fridge003
body: Maybe we need a triton kernel here to set cuda_graph_batch_info, so memcpy synchronization can be avoided.
created_at: 2025-04-27T04:15:17Z
author_association: COLLABORATOR
diff_hunk: @@ -138,9 +138,18 @@ def set_lora_info(
         if self.lora_backend.fuse_stacked_lora_b:
             # TODO: avoid using contiguous() in GPU.
             # B_buffer_gate_up: (num_lora, 2 * output_dim, r)
-            self.B_buffer_gate_up = torch.cat(
-                (B_buffer[0], B_buffer[1]), dim=-2
-            ).contiguous()
+            if not hasattr(self, "B_buffer_gate_up") or self.B_buffer_gate_up is None:
+                self.B_buffer_gate_up = torch.empty(
+                    (
+                        B_buffer[0].shape[0],
+                        2 * B_buffer[0].shape[1],
+                        B_buffer[0].shape[2],
+                    ),
+                    dtype=B_buffer[0].dtype,
+                    device=B_buffer[0].device,
+                ).contiguous()
path: python/sglang/srt/lora/layers.py
user: login: Qiaolin-Yu
body: Thanks!
created_at: 2025-04-27T19:29:07Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
path: python/sglang/srt/lora/lora_manager.py
user: login: Qiaolin-Yu
body: Extend will also use this preallocated space. Batches will use this preallocated space when CUDA Graph is enabled and `bs <= self.max_bs_in_cuda_graph`, no matter the result of `cuda_graph_runner.can_run()`. Comments added in the code to make it more clear.
created_at: 2025-04-27T19:32:55Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
path: python/sglang/srt/lora/lora_manager.py
user: login: Qiaolin-Yu
body: But I'm also open to passing the result of `cuda_graph_runner.can_run()` to this function and processing the batch based on `can_run_cuda_graph` if you prefer this way.
created_at: 2025-04-27T19:36:02Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
+                self.cuda_graph_batch_info.seg_lens[:bs].copy_(
+                    forward_batch.extend_seq_lens
+                )
+            else:
+                self.cuda_graph_batch_info.seg_lens[:bs].fill_(1)
+            self.cuda_graph_batch_info.seg_indptr[0] = 0
path: python/sglang/srt/lora/lora_manager.py
user: login: Qiaolin-Yu
body: Thanks for pointing this out!
created_at: 2025-04-27T19:36:52Z
author_association: COLLABORATOR
diff_hunk: @@ -237,3 +237,112 @@ def run_batch_lora_test(
                 f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
                 f"for base '{base_path}', adaptor '{adaptor_names}', backend '{backend}', prompt: '{prompts[0][:50]}...'"
             )
+
+
+def run_lora_test_by_batch(
path: test/srt/models/lora/utils.py
user: login: Fridge003
body: What's the difference between run_lora_test_by_batch and run_lora_test_one_by_one? 
created_at: 2025-04-27T19:57:55Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
path: python/sglang/srt/lora/lora_manager.py
user: login: Fridge003
body: Just make sure this branch is used. 
created_at: 2025-04-27T20:02:44Z
author_association: COLLABORATOR
diff_hunk: @@ -237,3 +237,112 @@ def run_batch_lora_test(
                 f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
                 f"for base '{base_path}', adaptor '{adaptor_names}', backend '{backend}', prompt: '{prompts[0][:50]}...'"
             )
+
+
+def run_lora_test_by_batch(
path: test/srt/models/lora/utils.py
user: login: Qiaolin-Yu
body: Basically, the difference is `run_lora_test_one_by_one ` uses `srt_runner.forward`, but `run_lora_test_by_batch ` uses `srt_runner.batch_forward`. `srt_runner.forward` will [split the batch into several generate requests](https://github.com/sgl-project/sglang/blob/main/python/sglang/test/runners.py#L605-L615).
created_at: 2025-04-27T20:15:24Z
author_association: COLLABORATOR
diff_hunk: @@ -237,3 +237,112 @@ def run_batch_lora_test(
                 f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
                 f"for base '{base_path}', adaptor '{adaptor_names}', backend '{backend}', prompt: '{prompts[0][:50]}...'"
             )
+
+
+def run_lora_test_by_batch(
path: test/srt/models/lora/utils.py
user: login: Fridge003
body: Nice point, I didn't notice the difference before
created_at: 2025-04-27T20:38:50Z
author_association: COLLABORATOR
diff_hunk: @@ -140,39 +157,71 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         if cur_uids == set([None]):
             return
 
-        # set up batch info shared by all lora moruldes
+        # set up batch info shared by all lora modules
         bs = forward_batch.batch_size
-        seg_lens = (
-            forward_batch.extend_seq_lens
-            if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device=self.device)
-        )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
-        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
-        max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
 
-        lora_ranks = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.int64, device="cuda"
-        )
-        scalings = torch.empty(
-            (self.max_loras_per_batch,), dtype=torch.float, device="cuda"
-        )
-        for i, lora_path in enumerate(forward_batch.lora_paths):
-            weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
-            lora = self.loras[lora_path]
-            lora_ranks[weight_indices[i]] = lora.config.hf_config["r"]
-            scalings[weight_indices[i]] = lora.scaling
-
-        batch_info = LoRABatchInfo(
-            bs=bs,
-            seg_lens=seg_lens,
-            seg_indptr=seg_indptr,
-            max_len=max_len,
-            weight_indices=weight_indices,
-            lora_ranks=lora_ranks,
-            scalings=scalings,
-        )
+        if hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph:
+            # Do in-place updates when CUDA graph is enabled
+            self.cuda_graph_batch_info.bs = bs
+            if forward_batch.forward_mode.is_extend():
+                self.cuda_graph_batch_info.seg_lens[:bs].copy_(
+                    forward_batch.extend_seq_lens
+                )
+            else:
+                self.cuda_graph_batch_info.seg_lens[:bs].fill_(1)
+            self.cuda_graph_batch_info.seg_indptr[0] = 0
+            torch.cumsum(
+                self.cuda_graph_batch_info.seg_lens[:bs],
+                dim=0,
+                out=self.cuda_graph_batch_info.seg_indptr[1 : bs + 1],
+            )
+            self.cuda_graph_batch_info.max_len = int(
+                torch.max(self.cuda_graph_batch_info.seg_lens[:bs])
+            )
+
+            for i, lora_path in enumerate(forward_batch.lora_paths):
+                self.cuda_graph_batch_info.weight_indices[i] = (
path: python/sglang/srt/lora/lora_manager.py
user: login: Fridge003
body: Can be delayed to next PR
created_at: 2025-04-28T01:33:45Z
author_association: COLLABORATOR

PR comments:
user: login: Qiaolin-Yu
created_at: 2025-04-27T03:08:39Z
body: PTAL @Fridge003 @Ying1123 
author_association: COLLABORATOR
user: login: Fridge003
created_at: 2025-04-27T03:32:01Z
body: The last paragraph of https://github.com/sgl-project/sglang/blob/main/docs/backend/lora.ipynb can also be updated.
author_association: COLLABORATOR
user: login: Qiaolin-Yu
created_at: 2025-04-27T03:35:28Z
body: > The last paragraph of https://github.com/sgl-project/sglang/blob/main/docs/backend/lora.ipynb can also be updated.

Done
author_association: COLLABORATOR
user: login: Fridge003
created_at: 2025-04-27T03:40:06Z
body: > > The last paragraph of https://github.com/sgl-project/sglang/blob/main/docs/backend/lora.ipynb can also be updated.
> 
> Done

The `--disable-cuda-graph` flags in its examples can also be removed.
author_association: COLLABORATOR
user: login: Qiaolin-Yu
created_at: 2025-04-27T03:55:05Z
body: > > > The last paragraph of https://github.com/sgl-project/sglang/blob/main/docs/backend/lora.ipynb can also be updated.
> > 
> > 
> > Done
> 
> The `--disable-cuda-graph` flags in its examples can also be removed.

Thanks for the reminder! Also fixed now.
author_association: COLLABORATOR
user: login: Fridge003
created_at: 2025-04-27T04:18:25Z
body: Have you tested the accuracy on batch size not captured(which requires padding)?
The cuda_graph_info is only prepared on captured batch sizes.
author_association: COLLABORATOR
user: login: Qiaolin-Yu
created_at: 2025-04-27T19:42:01Z
body: > Have you tested the accuracy on batch size not captured(which requires padding)? The cuda_graph_info is only prepared on captured batch sizes.

I just added tests that requires padding (bs=3) and it could pass. I think it's similar with [`cuda_graph_runner.replay_prepare()`](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/cuda_graph_runner.py#L479-L516). We only update `[:raw_bs]`, but the pre-allcoated space is large enough for `capture_bs`. So it will be fine. 
author_association: COLLABORATOR
user: login: zhaochenyang20
created_at: 2025-04-28T05:19:55Z
body: @Qiaolin-Yu let me rerun all the CI. No need to rebase
author_association: COLLABORATOR