A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 14:30 UTC
[Bugfix][LoRA][Spec Decode] Support LoRA with speculative decoding (#21068)
Signed-off-by: Sean Chen <xiaohong_chen1991@hotmail.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Danielle Robinson <dcmaddix@gmail.com> Co-authored-by: Haipeng Li <li2haipeng@gmail.com> Co-authored-by: li2haipeng <44383182+li2haipeng@users.noreply.github.com>
Xiaohong (Sean) Chen committed
Nov 7, 2025 at 20:58 UTC
d0c779200421fed3c8ab010c4e4504a5d3e153e6
7 files changed
+201
-15
tests/v1/e2e/test_lora_with_spec_decode.py
new
+141
@@ -0,0 +1,141 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+This script contains:
5
+1. test lora with speculative decoding for batch inference
6
+"""
7
+
8
+import random
9
+
10
+import numpy as np
11
+import pytest
12
+import torch
13
+
14
+from vllm import LLM, SamplingParams
15
+from vllm.distributed import cleanup_dist_env_and_memory
16
+from vllm.lora.request import LoRARequest
17
+from vllm.platforms import current_platform
18
+
19
+LORA_TEST_PROMPT_MAP: dict[str, str] = {}
20
+
21
+LORA_TEST_PROMPT_MAP["premjatin/qwen-linear-algebra-coder"] = """
22
+### INSTRUCTION:
23
+You are an AI assistant that generates Python code to solve linear
24
+algebra problems.
25
+
26
+### PROBLEM:
27
+Find the eigenvalues and eigenvectors of the following 3x3 matrix:
28
+[[3, 2, 0],
29
+ [2, 3, 0],
30
+ [0, 0, 2]]
31
+
32
+### OUTPUT FORMAT (STRICT):
33
+Numbers should be represented as integers only.
34
+
35
+### PYTHON SOLUTION:
36
+"""
37
+
38
+SEED = 42
39
+
40
+
41
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
42
+@pytest.mark.parametrize(
43
+ "model_setup",
44
+ [
45
+ (
46
+ "eagle3",
47
+ "Qwen/Qwen3-1.7B",
48
+ "AngelSlim/Qwen3-1.7B_eagle3",
49
+ "premjatin/qwen-linear-algebra-coder",
50
+ 1,
51
+ )
52
+ ],
53
+)
54
+def test_batch_inference_correctness(
55
+ monkeypatch: pytest.MonkeyPatch,
56
+ model_setup: tuple[str, str, str, str, int],
57
+):
58
+ """
59
+ Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora.
60
+ Should be the same and no failure when doing batch inference.
61
+ model_setup: (method, model_name, spec_model_name, lora_path, tp_size)
62
+ """
63
+ with monkeypatch.context() as m:
64
+ m.setenv("VLLM_USE_V1", "1")
65
+
66
+ # Disable randomness
67
+ m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
68
+ torch.manual_seed(SEED)
69
+ np.random.seed(SEED)
70
+ random.seed(SEED)
71
+ torch.cuda.manual_seed_all(SEED)
72
+ torch.backends.cudnn.benchmark = False
73
+ torch.backends.cudnn.deterministic = True
74
+
75
+ method, model_name, spec_model_name, lora_path, tp_size = model_setup
76
+
77
+ # without speculative decoding
78
+ ref_llm = LLM(
79
+ model=model_name,
80
+ trust_remote_code=True,
81
+ tensor_parallel_size=tp_size,
82
+ max_model_len=2048,
83
+ max_num_seqs=4,
84
+ enable_lora=True,
85
+ max_loras=1,
86
+ max_cpu_loras=1,
87
+ max_lora_rank=16,
88
+ )
89
+
90
+ prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
91
+ lora_request = LoRARequest("adapter", 1, lora_path)
92
+ sampling_params = SamplingParams(
93
+ temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
94
+ )
95
+
96
+ ref_outputs = ref_llm.generate(
97
+ prompts, sampling_params, lora_request=lora_request
98
+ )
99
+ del ref_llm
100
+ torch.cuda.empty_cache()
101
+ cleanup_dist_env_and_memory()
102
+
103
+ lora_spec_llm = LLM(
104
+ model=model_name,
105
+ trust_remote_code=True,
106
+ tensor_parallel_size=tp_size,
107
+ speculative_config={
108
+ "method": method,
109
+ "model": spec_model_name,
110
+ "num_speculative_tokens": 3,
111
+ "max_model_len": 2048,
112
+ },
113
+ max_model_len=2048,
114
+ max_num_seqs=4,
115
+ enable_lora=True,
116
+ max_loras=1,
117
+ max_cpu_loras=1,
118
+ max_lora_rank=16,
119
+ )
120
+
121
+ lora_spec_outputs = lora_spec_llm.generate(
122
+ prompts, sampling_params, lora_request=lora_request
123
+ )
124
+
125
+ matches = 0
126
+ misses = 0
127
+ for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
128
+ if ref_output.outputs[0].text == spec_output.outputs[0].text:
129
+ matches += 1
130
+ else:
131
+ misses += 1
132
+ print(f"ref_output: {ref_output.outputs[0].text}")
133
+ print(f"spec_output: {spec_output.outputs[0].text}")
134
+
135
+ # Heuristic: expect at least 90% of the prompts to match exactly
136
+ # Upon failure, inspect the outputs to check for inaccuracy.
137
+ print(f"match ratio: {matches}/{len(ref_outputs)}")
138
+ assert matches > int(0.90 * len(ref_outputs))
139
+ del lora_spec_llm
140
+ torch.cuda.empty_cache()
141
+ cleanup_dist_env_and_memory()
vllm/engine/arg_utils.py
+14
@@ -1574,6 +1574,20 @@ class EngineArgs:
1574
else None
1575
)
1576
1577
+ if (
1578
+ lora_config is not None
1579
+ and speculative_config is not None
1580
+ and scheduler_config.max_num_batched_tokens
1581
+ < (
1582
+ scheduler_config.max_num_seqs
1583
+ * (speculative_config.num_speculative_tokens + 1)
1584
+ )
1585
+ ):
1586
+ raise ValueError(
1587
+ "Consider increasing max_num_batched_tokens or "
1588
+ "decreasing num_speculative_tokens"
1589
+ )
1590
+
1591
# bitsandbytes pre-quantized model need a specific model loader
1592
if model_config.quantization == "bitsandbytes":
1593
self.quantization = self.load_format = "bitsandbytes"
vllm/lora/punica_wrapper/punica_gpu.py
+5
-1
@@ -51,8 +51,12 @@ class PunicaWrapperGPU(PunicaWrapperBase):
51
self.max_loras, max_num_batched_tokens, device=device
52
)
53
54
+ # When speculative decoding is enabled, max_num_samples is
55
+ # max_batches * (num_speculative_decoding_tokens + 1).
56
+ # This line can be optimized by replacing max_num_batched_tokens
57
+ # to max_batches * (num_speculative_decoding_tokens + 1).
58
self.prompt_mapping_meta = LoRAKernelMeta.make(
55
- self.max_loras, max_batches, device=device
59
+ self.max_loras, max_num_batched_tokens, device=device
60
)
61
62
def update_metadata(
vllm/v1/worker/gpu_input_batch.py
+6
-4
@@ -859,22 +859,24 @@ class InputBatch:
859
return prompt_token_ids_cpu_tensor.to(device=self.device, non_blocking=True)
860
861
def make_lora_inputs(
862
- self, num_scheduled_tokens: np.ndarray
862
+ self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
863
) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
864
"""
865
Given the num_scheduled_tokens for each request in the batch, return
866
datastructures used to activate the current LoRAs.
867
Returns:
868
- 1. prompt_lora_mapping: A tuple of size self.num_reqs where,
869
- prompt_lora_mapping[i] is the LoRA id to use for the ith prompt.
868
+ 1. prompt_lora_mapping: A tuple of size np.sum(num_sampled_tokens)
869
+ where, prompt_lora_mapping[i] is the LoRA id to use for the ith
870
+ sampled token.
871
2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens)
872
where, token_lora_mapping[i] is the LoRA id to use for ith token.
873
3. lora_requests: Set of relevant LoRA requests.
874
"""
875
876
req_lora_mapping = self.request_lora_mapping[: self.num_reqs]
876
- prompt_lora_mapping = tuple(req_lora_mapping)
877
+ prompt_lora_mapping = tuple(req_lora_mapping.repeat(num_sampled_tokens))
878
token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens))
879
+
880
active_lora_requests: set[LoRARequest] = set(
881
self.lora_id_to_lora_request.values()
882
)
vllm/v1/worker/gpu_model_runner.py
+15
-3
@@ -1268,6 +1268,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1268
logits_indices = query_start_loc[1:] - 1
1269
num_draft_tokens = None
1270
spec_decode_metadata = None
1271
+ num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
1272
else:
1273
# Get the number of draft tokens for each request.
1274
# Iterate over the dictionary rather than all requests since not all
@@ -1294,7 +1295,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1295
num_draft_tokens, cu_num_tokens
1296
)
1297
logits_indices = spec_decode_metadata.logits_indices
1297
-
1298
+ num_sampled_tokens = num_draft_tokens + 1
1299
# For DECODE only cuda graph of some attention backends (e.g., GDN).
1300
self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens
1301
self.num_decode_draft_tokens.np[num_reqs:].fill(-1)
@@ -1445,7 +1446,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1446
1447
# Hot-Swap lora model
1448
if self.lora_config:
1448
- self.set_active_loras(self.input_batch, num_scheduled_tokens)
1449
+ assert (
1450
+ np.sum(num_sampled_tokens)
1451
+ <= self.vllm_config.scheduler_config.max_num_batched_tokens
1452
+ )
1453
+ self.set_active_loras(
1454
+ self.input_batch, num_scheduled_tokens, num_sampled_tokens
1455
+ )
1456
1457
return (
1458
attn_metadata,
@@ -3390,6 +3397,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3397
assert len(num_scheduled_tokens_list) == num_reqs
3398
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
3399
total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
3400
+ num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
3401
3402
# Disable DP padding when running eager
3403
allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
@@ -3485,7 +3493,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3493
attn_metadata[layer_name] = attn_metadata_i
3494
3495
with self.maybe_dummy_run_with_lora(
3488
- self.lora_config, num_scheduled_tokens, activate_lora, remove_lora
3496
+ self.lora_config,
3497
+ num_scheduled_tokens,
3498
+ num_sampled_tokens,
3499
+ activate_lora,
3500
+ remove_lora,
3501
):
3502
# Make sure padding doesn't exceed max_num_tokens
3503
assert num_tokens_after_padding <= self.max_num_tokens
vllm/v1/worker/lora_model_runner_mixin.py
+19
-6
@@ -38,7 +38,6 @@ class LoRAModelRunnerMixin:
38
"Regarding multimodal models, vLLM currently "
39
"only supports adding LoRA to language model."
40
)
41
-
41
# Add LoRA Manager to the Model Runner
42
self.lora_manager = LRUCacheWorkerLoRAManager(
43
vllm_config,
@@ -70,13 +69,19 @@ class LoRAModelRunnerMixin:
69
raise RuntimeError("LoRA is not enabled. Use --enable-lora to enable LoRA.")
70
71
def set_active_loras(
73
- self, input_batch: InputBatch, num_scheduled_tokens: np.ndarray
72
+ self,
73
+ input_batch: InputBatch,
74
+ num_scheduled_tokens: np.ndarray,
75
+ num_sampled_tokens: np.ndarray | None = None,
76
) -> None:
75
- prompt_lora_mapping: tuple[int, ...] # of size input_batch.num_reqs
77
+ if num_sampled_tokens is None:
78
+ num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)
79
+
80
+ prompt_lora_mapping: tuple[int, ...] # of size np.sum(num_sampled_tokens)
81
token_lora_mapping: tuple[int, ...] # of size np.sum(num_scheduled_tokens)
82
lora_requests: set[LoRARequest]
83
prompt_lora_mapping, token_lora_mapping, lora_requests = (
79
- input_batch.make_lora_inputs(num_scheduled_tokens)
84
+ input_batch.make_lora_inputs(num_scheduled_tokens, num_sampled_tokens)
85
)
86
return self._set_active_loras(
87
prompt_lora_mapping, token_lora_mapping, lora_requests
@@ -123,8 +128,12 @@ class LoRAModelRunnerMixin:
128
self,
129
lora_config: LoRAConfig | None,
130
num_scheduled_tokens: np.ndarray,
131
+ num_sampled_tokens: np.ndarray | None = None,
132
activate_lora: bool = True,
133
):
134
+ if num_sampled_tokens is None:
135
+ num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)
136
+
137
if lora_config is None:
138
yield
139
else:
@@ -143,6 +152,9 @@ class LoRAModelRunnerMixin:
152
else:
153
prompt_lora_mapping = np.zeros(num_reqs, dtype=np.int32)
154
155
+ # Make sample lora mapping
156
+ sample_lora_mapping = np.repeat(prompt_lora_mapping, num_sampled_tokens)
157
+
158
# Make token lora mapping
159
token_lora_mapping = np.repeat(prompt_lora_mapping, num_scheduled_tokens)
160
@@ -157,7 +169,7 @@ class LoRAModelRunnerMixin:
169
}
170
171
self._set_active_loras(
160
- tuple(prompt_lora_mapping), tuple(token_lora_mapping), lora_requests
172
+ tuple(sample_lora_mapping), tuple(token_lora_mapping), lora_requests
173
)
174
175
yield
@@ -167,13 +179,14 @@ class LoRAModelRunnerMixin:
179
self,
180
lora_config: LoRAConfig | None,
181
num_scheduled_tokens: np.ndarray,
182
+ num_sampled_tokens: np.ndarray,
183
activate_lora: bool = True,
184
remove_lora: bool = True,
185
):
186
with (
187
self.maybe_setup_dummy_loras(lora_config, remove_lora),
188
self.maybe_select_dummy_loras(
176
- lora_config, num_scheduled_tokens, activate_lora
189
+ lora_config, num_scheduled_tokens, num_sampled_tokens, activate_lora
190
),
191
):
192
yield
vllm/v1/worker/tpu_input_batch.py
+1
-1
@@ -526,7 +526,7 @@ class InputBatch:
526
return prompt_token_ids_cpu_tensor.to(device=self.device, non_blocking=True)
527
528
def make_lora_inputs(
529
- self, num_scheduled_tokens: np.ndarray
529
+ self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
530
) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
531
"""
532
Given the num_scheduled_tokens for each request in the batch, return