chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support spec decoding when LoRA is applied to target model (#12903)

Browse Source
This commit is contained in:
Lifu Huang
2025-11-16 13:20:23 -08:00
committed by GitHub
parent 2b8b9d8496
commit 254f62d879
13 changed files with 363 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/lora/backend/base_backend.py
View File

@@ -1,8 +1,7 @@
from typing import Optional, Tuple, Union
from typing import Tuple, Union
import torch
from sglang.srt.lora.utils import LoRABatchInfo
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -97,8 +96,8 @@ class BaseLoRABackend:
def init_cuda_graph_batch_info(
self,
cuda_graph_batch_info: LoRABatchInfo,
max_bs_in_cuda_graph: int,
num_tokens_per_bs: int,
):
"""Initialize the batch info for CUDA Graph mode.
@@ -108,6 +107,7 @@ class BaseLoRABackend:
Args:
cuda_graph_batch_info: the LoRABatchInfo object created in LoraManager
max_bs_in_cuda_graph: maximum batch size for CUDA Graph mode
num_tokens_per_bs: number of tokens per sequence (1 for decoding, >1 for target_verify)
"""
pass
@@ -117,7 +117,7 @@ class BaseLoRABackend:
weight_indices: list[int],
lora_ranks: list[int],
scalings: list[float],
batch_info: Optional[LoRABatchInfo] = None,
use_cuda_graph: bool,
):
"""Prepare the lora weights and batch info for current forward batch.
@@ -129,7 +129,6 @@ class BaseLoRABackend:
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
batch_info: optional LoRABatchInfo object, if not provided, the backend should use its own
internal batch info (e.g., self.cuda_graph_batch_info for CUDA Graph mode)
use_cuda_graph: whether to use CUDA Graph for this batch
"""
pass

55
python/sglang/srt/lora/backend/chunked_backend.py
View File

@@ -1,5 +1,3 @@
from typing import Optional
import torch
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
@@ -52,7 +50,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
output_offset: torch.Tensor,
base_output: torch.Tensor = None,
*args,
**kwargs
**kwargs,
) -> torch.Tensor:
# For simple lora B, we use slice offsets [0, output_dim]
output_dim = weights.shape[-2]
@@ -75,7 +73,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
max_qkv_out_dim: int,
base_output: torch.Tensor = None,
*args,
**kwargs
**kwargs,
) -> torch.Tensor:
# x: (s, input_dim)
@@ -107,7 +105,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
output_offset: torch.Tensor,
base_output: torch.Tensor = None,
*args,
**kwargs
**kwargs,
) -> torch.Tensor:
# x: (s, input_dim)
@@ -160,13 +158,36 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
chunk_size = 16
return min(self.max_chunk_size, chunk_size)
def init_cuda_graph_batch_info(
self,
max_bs_in_cuda_graph: int,
num_tokens_per_bs: int,
):
max_num_segments = (
(num_tokens_per_bs + MIN_CHUNK_SIZE - 1) // MIN_CHUNK_SIZE
) * max_bs_in_cuda_graph
max_num_tokens = max_bs_in_cuda_graph * num_tokens_per_bs
with torch.device("cuda"):
self.cuda_graph_batch_info = LoRABatchInfo(
bs=max_bs_in_cuda_graph,
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, # Not used in CSGMV backend
)
def prepare_lora_batch(
self,
forward_batch: ForwardBatch,
weight_indices: list[int],
lora_ranks: list[int],
scalings: list[float],
batch_info: Optional[LoRABatchInfo] = None,
use_cuda_graph: bool,
):
chunk_size = self._determine_chunk_size(forward_batch)
@@ -188,7 +209,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
scalings, dtype=torch.float, pin_memory=True, device="cpu"
)
if batch_info is None:
if not use_cuda_graph:
batch_info = LoRABatchInfo(
bs=forward_batch.batch_size,
num_segments=num_segments,
@@ -213,6 +234,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
seg_lens=None,
)
else:
batch_info = self.cuda_graph_batch_info
batch_info.bs = forward_batch.batch_size
batch_info.num_segments = num_segments
batch_info.max_len = chunk_size
@@ -262,14 +284,23 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
with torch.device("cpu"):
seq_weight_indices = torch.tensor(seq_weight_indices, dtype=torch.int32)
seg_lens_cpu = (
torch.tensor(
if forward_batch.forward_mode.is_decode():
seg_lens_cpu = torch.ones(forward_batch.batch_size, dtype=torch.int32)
elif forward_batch.forward_mode.is_target_verify():
seg_lens_cpu = torch.full(
size=(forward_batch.batch_size,),
fill_value=forward_batch.spec_info.draft_token_num,
dtype=torch.int32,
)
elif forward_batch.forward_mode.is_extend():
seg_lens_cpu = torch.tensor(
forward_batch.extend_seq_lens_cpu,
dtype=torch.int32,
)
if forward_batch.forward_mode.is_extend()
else torch.ones(forward_batch.batch_size, dtype=torch.int32)
)
else:
raise ValueError(
f"Unsupported forward mode: {forward_batch.forward_mode}"
)
row_weight_indices = torch.repeat_interleave(
seq_weight_indices, seg_lens_cpu

46
python/sglang/srt/lora/backend/triton_backend.py
View File

@@ -1,5 +1,3 @@
from typing import Optional
import torch
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
@@ -97,16 +95,33 @@ class TritonLoRABackend(BaseLoRABackend):
return lora_output
def init_cuda_graph_batch_info(
self, cuda_graph_batch_info: LoRABatchInfo, max_bs_in_cuda_graph: int
self,
max_bs_in_cuda_graph: int,
num_tokens_per_bs: int,
):
# Initialize seg_lens and seg_indptr for CUDA graph as they remain constant
# across batches.
cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph].fill_(1)
torch.cumsum(
cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
dim=0,
out=cuda_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
)
with torch.device("cuda"):
self.cuda_graph_batch_info = LoRABatchInfo(
bs=max_bs_in_cuda_graph,
use_cuda_graph=True,
num_segments=None,
seg_lens=torch.full(
(max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
),
seg_indptr=torch.empty(max_bs_in_cuda_graph + 1, dtype=torch.int32),
max_len=num_tokens_per_bs,
weight_indices=torch.zeros(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),
permutation=None,
)
# Initialize seg_indptr for CUDA graph as they remain constant
# across batches.
torch.cumsum(
self.cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
dim=0,
out=self.cuda_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
)
def prepare_lora_batch(
self,
@@ -114,7 +129,7 @@ class TritonLoRABackend(BaseLoRABackend):
weight_indices: list[int],
lora_ranks: list[int],
scalings: list[float],
batch_info: Optional[LoRABatchInfo] = None,
use_cuda_graph: bool,
):
# Use pinned memory to avoid synchronizations during host-to-device transfer
weight_indices_tensor = torch.tensor(
@@ -129,10 +144,11 @@ class TritonLoRABackend(BaseLoRABackend):
bs = forward_batch.batch_size
if batch_info is not None:
if use_cuda_graph:
assert (
batch_info.use_cuda_graph
), "batch_info.use_cuda_graph must be True when batch_info is provided"
self.cuda_graph_batch_info is not None
), "CUDA Graph batch info is not initialized."
batch_info = self.cuda_graph_batch_info
batch_info.bs = forward_batch.batch_size
batch_info.num_segments = forward_batch.batch_size
else:

23
python/sglang/srt/lora/lora_manager.py
View File

@@ -29,7 +29,6 @@ from sglang.srt.lora.lora_config import LoRAConfig
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.lora.mem_pool import LoRAMemoryPool
from sglang.srt.lora.utils import (
LoRABatchInfo,
LoRAType,
get_layer_id,
get_normalized_target_modules,
@@ -95,25 +94,13 @@ class LoRAManager:
lora_paths=lora_paths,
)
def init_cuda_graph_batch_info(self, max_bs_in_cuda_graph: int):
def init_cuda_graph_batch_info(
self, max_bs_in_cuda_graph: int, num_tokens_per_bs: int
):
self.max_bs_in_cuda_graph = max_bs_in_cuda_graph
with torch.device("cuda"):
self.cuda_graph_batch_info = LoRABatchInfo(
bs=max_bs_in_cuda_graph,
use_cuda_graph=True,
num_segments=None,
seg_lens=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
max_len=1,
weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
permutation=torch.zeros(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),
)
self.lora_backend.init_cuda_graph_batch_info(
cuda_graph_batch_info=self.cuda_graph_batch_info,
max_bs_in_cuda_graph=max_bs_in_cuda_graph,
num_tokens_per_bs=num_tokens_per_bs,
)
def create_lora_update_result(
@@ -297,7 +284,7 @@ class LoRAManager:
weight_indices=weight_indices,
lora_ranks=lora_ranks,
scalings=scalings,
batch_info=self.cuda_graph_batch_info if use_cuda_graph else None,
use_cuda_graph=use_cuda_graph,
)
def update_lora_info(self):

6
python/sglang/srt/lora/utils.py
View File

@@ -19,9 +19,6 @@ class LoRABatchInfo:
# Number of segments. For triton backend, it is equal to batch size.
num_segments: int
# Maximum segment length of current batch
max_len: int
# Indice pointers of each segment in shape (num_segments + 1, )
seg_indptr: torch.Tensor
@@ -34,6 +31,9 @@ class LoRABatchInfo:
# scaling of each lora adapter, in shape (lora_num,)
scalings: torch.Tensor
# Maximum segment length of current batch
max_len: Optional[int]
# Lengths of each segments in shape (num_segments,)
seg_lens: Optional[torch.Tensor]

5
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -308,7 +308,10 @@ class CudaGraphRunner:
set_torch_compile_config()
if self.model_runner.server_args.enable_lora:
self.model_runner.lora_manager.init_cuda_graph_batch_info(self.max_bs)
self.model_runner.lora_manager.init_cuda_graph_batch_info(
max_bs_in_cuda_graph=self.max_bs,
num_tokens_per_bs=self.num_tokens_per_bs,
)
# Graph inputs
with torch.device(self.device):

7
python/sglang/srt/server_args.py
View File

@@ -3874,6 +3874,13 @@ class ServerArgs:
)
if self.enable_lora:
# Validate compatibility with speculative decoding
if self.speculative_algorithm not in ["NGRAM", None]:
raise ValueError(
"Currently LoRA is only compatible with NGRAM speculative decoding."
)
# Parse lora_paths
if isinstance(self.lora_paths, list):
lora_paths = self.lora_paths
self.lora_paths = []

12
python/sglang/test/runners.py
View File

@@ -528,6 +528,8 @@ class SRTRunner:
speculative_num_steps: Optional[int] = None,
speculative_eagle_topk: Optional[int] = None,
speculative_num_draft_tokens: Optional[int] = None,
speculative_ngram_min_match_window_size: Optional[int] = None,
speculative_ngram_max_match_window_size: Optional[int] = None,
disable_overlap_schedule: bool = False,
disable_custom_all_reduce: bool = False,
torchao_config: Optional[str] = None,
@@ -539,6 +541,7 @@ class SRTRunner:
max_loaded_loras: Optional[int] = None,
json_model_override_args: Optional[dict[str, Any]] = None,
lora_eviction_policy: str = "lru",
enable_deterministic_inference: bool = False,
):
self.model_type = model_type
self.is_generation = model_type == "generation"
@@ -554,6 +557,14 @@ class SRTRunner:
spec_kwargs["speculative_num_steps"] = speculative_num_steps
spec_kwargs["speculative_eagle_topk"] = speculative_eagle_topk
spec_kwargs["speculative_num_draft_tokens"] = speculative_num_draft_tokens
elif speculative_algorithm == "NGRAM":
spec_kwargs["speculative_algorithm"] = speculative_algorithm
spec_kwargs["speculative_ngram_min_match_window_size"] = (
speculative_ngram_min_match_window_size
)
spec_kwargs["speculative_ngram_max_match_window_size"] = (
speculative_ngram_max_match_window_size
)
self.engine = Engine(
model_path=model_path,
@@ -594,6 +605,7 @@ class SRTRunner:
else "{}"
),
lora_eviction_policy=lora_eviction_policy,
enable_deterministic_inference=enable_deterministic_inference,
**spec_kwargs,
)

126
test/srt/lora/test_lora.py
View File

@@ -14,138 +14,20 @@
import multiprocessing as mp
import os
import random
import unittest
from typing import List
from utils import (
ALL_OTHER_MULTI_LORA_MODELS,
CI_MULTI_LORA_MODELS,
TORCH_DTYPES,
LoRAModelCase,
ensure_reproducibility,
run_lora_multiple_batch_on_model_cases,
)
from sglang.test.runners import HFRunner, SRTRunner
from sglang.test.test_utils import CustomTestCase, calculate_rouge_l, is_in_ci
TEST_MULTIPLE_BATCH_PROMPTS = [
"""
### 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:
""",
"""
### Instruction:
Write a poem about the transformers Python library.
Mention the word "large language models" in that poem.
### Response:
The Transformers are large language models,
They're used to make predictions on text.
""",
"AI is a field of computer science focused on",
"Computer science is the study of",
"Write a short story.",
"What are the main components of a computer?",
]
from sglang.test.test_utils import CustomTestCase, is_in_ci
class TestLoRA(CustomTestCase):
def _create_test_samples(
self, lora_adapter_paths: List[str], repeated_trials: int = 3
):
random.seed(42) # Ensure reproducibility
patterns = [
[None, lora_adapter_paths[0], lora_adapter_paths[1]],
[lora_adapter_paths[0], None, lora_adapter_paths[1]],
[lora_adapter_paths[0], lora_adapter_paths[1], None],
[None, lora_adapter_paths[1], None],
[None, None, None],
]
batches = [
[random.choice(pattern) for _ in range(3)]
for pattern in patterns
for _ in range(repeated_trials)
]
return batches
def _run_lora_multiple_batch_on_model_cases(self, model_cases: List[LoRAModelCase]):
for model_case in model_cases:
for torch_dtype in TORCH_DTYPES:
max_new_tokens = 32
base_path = model_case.base
lora_adapter_paths = [a.name for a in model_case.adaptors]
assert len(lora_adapter_paths) >= 2
print(
f"\n========== Testing multiple batches on base '{base_path}', dtype={torch_dtype} ---"
)
# Initialize runners
srt_runner = SRTRunner(
base_path,
torch_dtype=torch_dtype,
model_type="generation",
lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
max_loras_per_batch=len(lora_adapter_paths) + 1,
sleep_on_idle=True, # Eliminate non-determinism by forcing all requests to be processed in one batch.
attention_backend="torch_native",
)
hf_runner = HFRunner(
base_path, torch_dtype=torch_dtype, model_type="generation"
)
batches = self._create_test_samples(lora_adapter_paths)
with srt_runner, hf_runner:
for i, lora_paths in enumerate(batches, start=1):
prompts = [
random.choice(TEST_MULTIPLE_BATCH_PROMPTS) for _ in range(3)
]
print(
f"\n--- Running Batch {i} --- prompts: {prompts}, lora_paths: {lora_paths}"
)
ensure_reproducibility()
srt_outputs = srt_runner.batch_forward(
prompts,
max_new_tokens=max_new_tokens,
lora_paths=lora_paths,
)
ensure_reproducibility()
hf_outputs = hf_runner.forward(
prompts,
max_new_tokens=max_new_tokens,
lora_paths=lora_paths,
)
print("SRT outputs:", [s for s in srt_outputs.output_strs])
print("HF outputs:", [s for s in hf_outputs.output_strs])
for srt_out, hf_out in zip(
srt_outputs.output_strs, hf_outputs.output_strs
):
srt_str = srt_out.strip()
hf_str = hf_out.strip()
rouge_tol = model_case.rouge_l_tolerance
rouge_score = calculate_rouge_l([srt_str], [hf_str])[0]
if rouge_score < rouge_tol:
raise AssertionError(
f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
f"for base '{base_path}', adaptor '{lora_paths}', prompt: '{prompts}...'"
)
print(f"--- Batch {i} Comparison Passed --- ")
def test_ci_lora_models(self):
self._run_lora_multiple_batch_on_model_cases(CI_MULTI_LORA_MODELS)
run_lora_multiple_batch_on_model_cases(CI_MULTI_LORA_MODELS)
def test_all_lora_models(self):
if is_in_ci():
@@ -157,7 +39,7 @@ class TestLoRA(CustomTestCase):
continue
filtered_models.append(model_case)
self._run_lora_multiple_batch_on_model_cases(filtered_models)
run_lora_multiple_batch_on_model_cases(filtered_models)
if __name__ == "__main__":

165
test/srt/lora/test_lora_qwen3.py
View File

@@ -13,15 +13,11 @@
# ==============================================================================
import multiprocessing as mp
import os
import random
import unittest
from typing import List
from utils import TORCH_DTYPES, LoRAAdaptor, LoRAModelCase, ensure_reproducibility
from utils import LoRAAdaptor, LoRAModelCase, run_lora_multiple_batch_on_model_cases
from sglang.test.runners import HFRunner, SRTRunner
from sglang.test.test_utils import CustomTestCase, calculate_rouge_l, is_in_ci
from sglang.test.test_utils import CustomTestCase
LORA_MODELS_QWEN3 = [
LoRAModelCase(
@@ -41,164 +37,9 @@ LORA_MODELS_QWEN3 = [
]
TEST_MULTIPLE_BATCH_PROMPTS = [
"""
### 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:
""",
"""
### Instruction:
Write a poem about the transformers Python library.
Mention the word "large language models" in that poem.
### Response:
The Transformers are large language models,
They're used to make predictions on text.
""",
"AI is a field of computer science focused on",
"Computer science is the study of",
"Write a short story.",
"What are the main components of a computer?",
]
class TestLoRAQwen3(CustomTestCase):
def _run_lora_multiple_batch_on_model_cases(self, model_cases: List[LoRAModelCase]):
for model_case in model_cases:
for torch_dtype in TORCH_DTYPES:
max_new_tokens = 32
base_path = model_case.base
lora_adapter_paths = [a.name for a in model_case.adaptors]
assert len(lora_adapter_paths) >= 2
batches = [
(
[
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
],
[
None,
lora_adapter_paths[0],
lora_adapter_paths[1],
],
),
(
[
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
],
[
lora_adapter_paths[0],
None,
lora_adapter_paths[1],
],
),
(
[
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
],
[lora_adapter_paths[0], lora_adapter_paths[1], None],
),
(
[
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
],
[None, lora_adapter_paths[1], None],
),
(
[
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
],
[None, None, None],
),
]
print(
f"\n========== Testing multiple batches on base '{base_path}', dtype={torch_dtype} ---"
)
# Initialize runners
ensure_reproducibility()
srt_runner = SRTRunner(
base_path,
torch_dtype=torch_dtype,
model_type="generation",
lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
max_loras_per_batch=len(lora_adapter_paths) + 1,
sleep_on_idle=True, # Eliminate non-determinism by forcing all requests to be processed in one batch.
attention_backend="torch_native",
)
ensure_reproducibility()
hf_runner = HFRunner(
base_path,
torch_dtype=torch_dtype,
model_type="generation",
patch_model_do_sample_false=True,
)
with srt_runner, hf_runner:
for i, (prompts, lora_paths) in enumerate(batches):
print(
f"\n--- Running Batch {i+1} --- prompts: {prompts}, lora_paths: {lora_paths}"
)
srt_outputs = srt_runner.batch_forward(
prompts,
max_new_tokens=max_new_tokens,
lora_paths=lora_paths,
)
hf_outputs = hf_runner.forward(
prompts,
max_new_tokens=max_new_tokens,
lora_paths=lora_paths,
)
print("SRT outputs:", [s for s in srt_outputs.output_strs])
print("HF outputs:", [s for s in hf_outputs.output_strs])
for srt_out, hf_out in zip(
srt_outputs.output_strs, hf_outputs.output_strs
):
srt_str = srt_out.strip()
hf_str = hf_out.strip()
rouge_tol = model_case.rouge_l_tolerance
rouge_score = calculate_rouge_l([srt_str], [hf_str])[0]
if rouge_score < rouge_tol:
raise AssertionError(
f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
f"for base '{base_path}', adaptor '{lora_paths}', prompt: '{prompts}...'"
)
print(f"--- Batch {i+1} Comparison Passed --- ")
def test_ci_lora_models(self):
self._run_lora_multiple_batch_on_model_cases(LORA_MODELS_QWEN3)
def test_all_lora_models(self):
if is_in_ci():
return
qwen_filtered_models = []
for model_case in LORA_MODELS_QWEN3:
if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
continue
qwen_filtered_models.append(model_case)
self._run_lora_multiple_batch_on_model_cases(qwen_filtered_models)
run_lora_multiple_batch_on_model_cases(LORA_MODELS_QWEN3)
if __name__ == "__main__":

71
test/srt/lora/test_lora_spec_decoding.py Normal file
View File

@@ -0,0 +1,71 @@
# Copyright 2023-2025 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 unittest
from utils import (
CI_MULTI_LORA_MODELS,
LoRAAdaptor,
LoRAModelCase,
run_lora_multiple_batch_on_model_cases,
)
from sglang.test.test_utils import CustomTestCase
LORA_MODELS_QWEN3 = [
LoRAModelCase(
base="Qwen/Qwen3-4B",
adaptors=[
LoRAAdaptor(
name="nissenj/Qwen3-4B-lora-v2",
prefill_tolerance=3e-1,
),
LoRAAdaptor(
name="y9760210/Qwen3-4B-lora_model",
prefill_tolerance=3e-1,
),
],
max_loras_per_batch=2,
),
]
class TestLoRASpecDecoding(CustomTestCase):
def test_qwen(self):
run_lora_multiple_batch_on_model_cases(
LORA_MODELS_QWEN3,
attention_backend="triton",
use_spec_decoding=True,
disable_cuda_graph=True,
enable_deterministic_inference=True,
)
def test_llama(self):
run_lora_multiple_batch_on_model_cases(
CI_MULTI_LORA_MODELS,
attention_backend="triton",
use_spec_decoding=True,
disable_cuda_graph=True,
enable_deterministic_inference=True,
)
if __name__ == "__main__":
try:
mp.set_start_method("spawn")
except RuntimeError:
pass
unittest.main(warnings="ignore")

174
test/srt/lora/utils.py
View File

@@ -395,3 +395,177 @@ def ensure_reproducibility():
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.use_deterministic_algorithms(True)
TEST_MULTIPLE_BATCH_PROMPTS = [
"""
### 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:
""",
"""
### Instruction:
Write a poem about the transformers Python library.
Mention the word "large language models" in that poem.
### Response:
The Transformers are large language models,
They're used to make predictions on text.
""",
"AI is a field of computer science focused on",
"Computer science is the study of",
"Write a short story.",
"What are the main components of a computer?",
]
def create_multiple_batch_test_samples(
prompts: List[str], lora_adapter_paths: List[str]
):
random.seed(42)
return [
(
[
random.choice(prompts),
random.choice(prompts),
random.choice(prompts),
],
[
None,
lora_adapter_paths[0],
lora_adapter_paths[1],
],
),
(
[
random.choice(prompts),
random.choice(prompts),
random.choice(prompts),
],
[
lora_adapter_paths[0],
None,
lora_adapter_paths[1],
],
),
(
[
random.choice(prompts),
random.choice(prompts),
random.choice(prompts),
],
[lora_adapter_paths[0], lora_adapter_paths[1], None],
),
(
[
random.choice(prompts),
random.choice(prompts),
random.choice(prompts),
],
[None, lora_adapter_paths[1], None],
),
(
[
random.choice(prompts),
random.choice(prompts),
random.choice(prompts),
],
[None, None, None],
),
]
def run_lora_multiple_batch_on_model_cases(
model_cases: List[LoRAModelCase],
use_spec_decoding: bool = False,
attention_backend: str = "torch_native",
disable_cuda_graph: bool = True,
enable_deterministic_inference: bool = False,
):
for model_case in model_cases:
for torch_dtype in TORCH_DTYPES:
max_new_tokens = 32
base_path = model_case.base
lora_adapter_paths = [a.name for a in model_case.adaptors]
assert len(lora_adapter_paths) >= 2
batches = create_multiple_batch_test_samples(
TEST_MULTIPLE_BATCH_PROMPTS, lora_adapter_paths
)
print(
f"\n========== Testing multiple batches on base '{base_path}', dtype={torch_dtype} ---"
)
# Initialize runners
ensure_reproducibility()
spec_args = (
{}
if not use_spec_decoding
else {
"speculative_algorithm": "NGRAM",
"speculative_num_draft_tokens": 5,
"speculative_ngram_min_match_window_size": 2,
"speculative_ngram_max_match_window_size": 15,
}
)
srt_runner = SRTRunner(
base_path,
torch_dtype=torch_dtype,
model_type="generation",
lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
max_loras_per_batch=len(lora_adapter_paths) + 1,
sleep_on_idle=True, # Eliminate non-determinism by forcing all requests to be processed in one batch.
attention_backend=attention_backend,
enable_deterministic_inference=enable_deterministic_inference,
disable_cuda_graph=disable_cuda_graph,
**spec_args,
)
ensure_reproducibility()
hf_runner = HFRunner(
base_path,
torch_dtype=torch_dtype,
model_type="generation",
patch_model_do_sample_false=True,
)
with srt_runner, hf_runner:
for i, (prompts, lora_paths) in enumerate(batches):
print(
f"\n--- Running Batch {i+1} --- prompts: {prompts}, lora_paths: {lora_paths}"
)
srt_outputs = srt_runner.batch_forward(
prompts,
max_new_tokens=max_new_tokens,
lora_paths=lora_paths,
)
hf_outputs = hf_runner.forward(
prompts,
max_new_tokens=max_new_tokens,
lora_paths=lora_paths,
)
print("SRT outputs:", [s for s in srt_outputs.output_strs])
print("HF outputs:", [s for s in hf_outputs.output_strs])
for srt_out, hf_out in zip(
srt_outputs.output_strs, hf_outputs.output_strs
):
srt_str = srt_out.strip()
hf_str = hf_out.strip()
rouge_tol = model_case.rouge_l_tolerance
rouge_score = calculate_rouge_l([srt_str], [hf_str])[0]
if rouge_score < rouge_tol:
raise AssertionError(
f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
f"for base '{base_path}', adaptor '{lora_paths}', prompt: '{prompts}...'"
)
print(f"--- Batch {i+1} Comparison Passed --- ")

1
test/srt/run_suite.py
View File

@@ -26,6 +26,7 @@ suites = {
TestFile("lora/test_lora_eviction.py", 240),
TestFile("lora/test_lora_update.py", 600),
TestFile("lora/test_lora_backend.py", 99),
TestFile("lora/test_lora_spec_decoding.py", 150),
TestFile("lora/test_multi_lora_backend.py", 60),
TestFile("models/test_compressed_tensors_models.py", 42),
TestFile("models/test_cross_encoder_models.py", 100),