[Feature] overlap LoRA weight loading with compute

[glenliu21]

Motivation

In #14190, I tried out a method to make LoRA weight loading asynchronous (see #8712). However, one issue is that this method required allocating more GPU memory for LoRA weight storage, which is not ideal. This PR instead makes LoRA weight loading truly free and asynchronous by pipelining the loading of LoRA weights. To illustrate:

In the current implementation, loading in a new batch of LoRA adapters blocks all forward computation:

GPU Compute
R1:  │            idle                                                       │████████ RUN R1..R4 ████████████│
R2:  │            idle                                                       │████████ RUN R1..R4 ████████████│
R3:  │            idle                                                       │████████ RUN R1..R4 ████████████│
R4:  │            idle                                                       │████████ RUN R1..R4 ████████████│

PCIe / LoRA Load
     │[ LOAD R1, R2, R3, R4 (blocking before run) ]                          │
     │─████████████████████──────────────────────────────────────────────────|

With this PR, we pipeline this process so that forward compute can overlap with LoRA weight loading:

GPU Compute
R1:  │      idle       │█████ RUN R1 ████│███ RUN R1+R2 ███│█ RUN R1+R2+R3 ██│RUN R1..R4 █████│
R2:  │      idle       │      idle       │███ RUN R1+R2 ███│█ RUN R1+R2+R3 ██│RUN R1..R4 █████│
R3:  │      idle       │      idle       │      idle       │█ RUN R1+R2+R3 ██│RUN R1..R4 █████│
R4:  │      idle       │      idle       │      idle       │       idle      │RUN R1..R4 █████│

PCIe / LoRA Load
     │ [ LOAD R1 ]     │ [ LOAD R2 ]     │ [ LOAD R3 ]     │ [ LOAD R4 ]     │
     │█████████────────│█████████────────│█████████────────│█████████────────│

Modifications

• Add --enable-lora-overlap-loading as a server argument
• Pin LoRA weights and introduce a separate CUDA stream (load_stream) to ensure adapter loading happens on a separate stream from forward_stream
• Modify existing batch selection logic in Scheduler and add required bookkeeping logic

Accuracy Tests

Implemented an end-to-end test that checks that outputs for workloads with --enable-lora-overlap-loading match workloads without overlap loading enabled.

Benchmarking and Profiling

The hardware I used was a single H200 GPU. I benchmarked with the following scripts:

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --max-loaded-loras 16 \
    --max-loras-per-batch 8 \
    --enable-lora-overlap-loading \ # Disable as needed
    --lora-paths \
        adapter0=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter1=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter2=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter3=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter4=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter5=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter6=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter7=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter8=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter9=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter10=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter11=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter12=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter13=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter14=mkopecki/chess-lora-adapter-llama-3.1-8b \
        adapter15=mkopecki/chess-lora-adapter-llama-3.1-8b

Note that mkopecki/chess-lora-adapter-llama-3.1-8b is a large adapter (>1GB). However, I used it mainly to see how well this implementation can perform in the best case and to establish it as a proof of concept.

python3 -m sglang.bench_serving \
  --backend sglang \
  --base-url http://localhost:30000 \
  --dataset-name random \
  --num-prompts 100 \
  --request-rate 4 \
  --random-input-len 512 \
  --random-output-len 512 \
  --lora-request-distribution distinct \
  --lora-name \
    adapter0 \
    adapter1 \
    adapter2 \
    adapter3 \
    adapter4 \
    adapter5 \
    adapter6 \
    adapter7 \
    adapter8 \
    adapter9 \
    adapter10 \
    adapter11 \
    adapter12 \
    adapter13 \
    adapter14 \
    adapter15

We use distinct for the request distribution to maximize cache misses. Furthermore, we use a relatively small random-input-len and random-output-len so that the bottleneck is our LoRA weight loading, not compute.

Benchmark Comparison: Main vs PR

Metric	main Run 1	PR Run 1	% decrease
Mean E2E Latency	5768.42	4529.93	21.5%
Median E2E Latency	5468.30	4137.39	24.3%
Mean TTFT	2220.95	1533.80	30.9%
Median TTFT	750.58	164.11	78.1%
P99 TTFT	9248.78	5724.13	38.1%
Median TPOT	13.93	7.76	44.3%

Overall, we see meaningful decreases in E2E latency and TTFT.

Running benchmarks with the profiler also shows that weight loading is overlapped with compute. This first image shows GPU compute being blocked by weight loading:

And this image shows weights being loading concurrently with GPU computation:

Benchmarking with Smaller Adapters

I also ran benchmarks using smaller adapters to get a more representative benchmark for the standard LoRA use case. In the below script, the adapter sizes range from 10s to 100s of MB, with most being less than 200MB:

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --max-loaded-loras 16 \
    --max-loras-per-batch 8 \
    --enable-lora-overlap-loading \
    --lora-paths \
        adapter0=faridlazuarda/valadapt-llama-3.1-8B-it-chinese \
        adapter1=LlamaFactoryAI/Llama-3.1-8B-Instruct-cv-job-description-matching \
        adapter2=Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16 \
        adapter3=pbevan11/llama-3.1-8b-ocr-correction \
        adapter4=reissbaker/llama-3.1-8b-abliterated-lora \
        adapter5=Roblox/Llama-3.1-8B-Instruct-RobloxGuard-1.0 \
        adapter6=nvidia/llama-3.1-nemoguard-8b-topic-control \
        adapter7=Kawon/llama3.1-food-finetune_v14_r8 \
        adapter8=faridlazuarda/valadapt-llama-3.1-8B-it-chinese \
        adapter9=LlamaFactoryAI/Llama-3.1-8B-Instruct-cv-job-description-matching \
        adapter10=Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16 \
        adapter11=pbevan11/llama-3.1-8b-ocr-correction \
        adapter12=reissbaker/llama-3.1-8b-abliterated-lora \
        adapter13=Roblox/Llama-3.1-8B-Instruct-RobloxGuard-1.0 \
        adapter14=nvidia/llama-3.1-nemoguard-8b-topic-control \
        adapter15=Kawon/llama3.1-food-finetune_v14_r8

This produces the following results:

Metric	main Run 1	PR Run 1	% decrease
Mean E2E Latency	3744.45	3408.13	8.98%
Median E2E Latency	3569.53	3215.31	9.92%
Mean TTFT	1231.55	1013.58	17.7%
Median TTFT	194.21	126.46	34.88%
P99 TTFT	4819.61	3974.23	17.54%
Median TPOT	9.69	9.29	4.13%

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Fridge003]

/tag-and-rerun-ci

[Fridge003]

There is accuracy issue https://github.com/sgl-project/sglang/actions/runs/20389111310/job/58633965556?pr=15512

[glenliu21]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant performance optimization by overlapping LoRA weight loading with GPU computation. The approach involves a new LoRAPrefetcher to handle asynchronous loading, which is a solid design for separating concerns. The changes are consistent and well-motivated, with impressive benchmark improvements. I have a couple of suggestions to improve maintainability by reducing code duplication and simplifying a method signature. Overall, this is a great contribution.

[glenliu21]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces an excellent optimization by overlapping LoRA weight loading with GPU computation, which significantly improves performance, especially for scenarios with frequent LoRA adapter switching. The implementation is well-structured, using a dedicated LoRAPrefetcher with a ThreadPoolExecutor and a separate CUDA stream to handle asynchronous loading. The changes to the scheduler and memory pool are correct and necessary for this feature.

My review includes a couple of suggestions for improvement:

• Adding a shutdown mechanism for the ThreadPoolExecutor to prevent potential resource leaks in a long-running server.
• Refactoring duplicated code in ScheduleBatch to improve maintainability.

Overall, this is a great contribution that brings substantial performance gains.

[Fridge003]

The accuracy error still exists
https://github.com/sgl-project/sglang/actions/runs/20468910505/job/58819496065?pr=15512
Maybe there is bug on weight synchronization?

[Fridge003]

Can we add a server arg for controlling this feature?

[glenliu21]

The accuracy error still exists https://github.com/sgl-project/sglang/actions/runs/20468910505/job/58819496065?pr=15512 Maybe there is bug on weight synchronization?

Yes, there's an issue with weight synchronization currently - I'm taking a look now.

Can we add a server arg for controlling this feature?

Yes, I can add an enable-lora-prefetch argument. In the future (probably in a follow-up PR), we should add some statistics tracking so that we only load large adapters asynchronously.

[gemini-code-assist]

Code Review

This pull request introduces a significant performance improvement by overlapping LoRA weight loading with GPU computation. The implementation is well-structured, introducing a LoRAPrefetcher class that manages asynchronous loading on a separate CUDA stream. The changes to the scheduler to accommodate prefetching states are logical and correct. The addition of end-to-end tests to verify correctness is also a great addition. My feedback includes a couple of minor suggestions to improve code maintainability.

[Fridge003]

The benchmark used large adaptors as showcase.
But I'm curious how this feature behaves on normal cases (like llama-8b, with some common small adaptors)
On the normal cases it shouldn't cause performance regression

[glenliu21]

The benchmark used large adaptors as showcase. But I'm curious how this feature behaves on normal cases (like llama-8b, with some common small adaptors) On the normal cases it shouldn't cause performance regression

Please see the updated description for a benchmark ran with smaller adapters.

One thing to note is that there are potentially cases in which enabling overlap loading results in lower throughput. This is because distinct LoRA adapters must run their own individual prefill steps, whereas in the baseline case, we wait for all adapters to load then run one single prefill step. So if the time it takes to load an adapter is much less than the time it takes to run a prefill step, then we can see throughput decrease.

[Fridge003]

Can we add an example for lora overlap loading in below section (can be updated in a following PR)

[glenliu21]

Please see #17464.