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 6, 12:00 UTC
Adding deterministic lora benchmarking to vLLM Bench (#36057)
Signed-off-by: Ubuntu <ubuntu@ip-172-31-43-201.ap-northeast-1.compute.internal> Signed-off-by: Ronald Xu <ronaldxu@amazon.com>
RonaldBXu committed
Mar 18, 2026 at 09:02 UTC
c9d838fc338db9a5a23cb3906d17c47423c4c9e4
3 files changed
+122
-9
vllm/benchmarks/datasets.py
+84
-4
@@ -183,6 +183,68 @@ class BenchmarkDataset(ABC):
183
)
184
return lora_request
185
186
+ def get_round_robin_lora_request(
187
+ self,
188
+ index: int,
189
+ max_loras: int | None = None,
190
+ lora_path: str | None = None,
191
+ ) -> LoRARequest | None:
192
+ """
193
+ Optionally select a LoRA request using deterministic round-robin.
194
+
195
+ This method cycles through LoRA IDs in order based on the request
196
+ index, providing reproducible LoRA assignment.
197
+
198
+ Args:
199
+ index (int): The request index used for round-robin selection.
200
+ max_loras (Optional[int]): The maximum number of LoRAs available.
201
+ If `None`, LoRA is not used.
202
+ lora_path (Optional[str]): Path to the LoRA parameters on disk.
203
+ If `None`, LoRA is not used.
204
+
205
+ Returns:
206
+ A new [`LoRARequest`][vllm.lora.request.LoRARequest]
207
+ (or `None` if not applicable).
208
+ """
209
+ if max_loras is None or lora_path is None:
210
+ return None
211
+
212
+ # Deterministic round-robin: cycle through [1, max_loras]
213
+ lora_id = index % max_loras + 1
214
+ lora_request = LoRARequest(
215
+ lora_name=str(lora_id),
216
+ lora_int_id=lora_id,
217
+ lora_path=lora_path_on_disk(lora_path),
218
+ )
219
+ return lora_request
220
+
221
+ def get_lora_request(
222
+ self,
223
+ index: int,
224
+ max_loras: int | None = None,
225
+ lora_path: str | None = None,
226
+ lora_assignment: str = "random",
227
+ ) -> LoRARequest | None:
228
+ """
229
+ Select a LoRA request using the specified assignment strategy.
230
+
231
+ Args:
232
+ index (int): The request index (used for round-robin).
233
+ max_loras (Optional[int]): The maximum number of LoRAs available.
234
+ lora_path (Optional[str]): Path to the LoRA parameters on disk.
235
+ lora_assignment (str): Strategy for LoRA selection.
236
+ 'random' (default) or 'round-robin'.
237
+
238
+ Returns:
239
+ A new [`LoRARequest`][vllm.lora.request.LoRARequest]
240
+ (or `None` if not applicable).
241
+ """
242
+ if lora_assignment == "round-robin":
243
+ return self.get_round_robin_lora_request(
244
+ index=index, max_loras=max_loras, lora_path=lora_path
245
+ )
246
+ return self.get_random_lora_request(max_loras=max_loras, lora_path=lora_path)
247
+
248
@abstractmethod
249
def sample(
250
self,
@@ -478,6 +540,9 @@ class RandomDataset(BenchmarkDataset):
540
input_len: int = DEFAULT_INPUT_LEN,
541
output_len: int = DEFAULT_OUTPUT_LEN,
542
batchsize: int = 1,
543
+ max_loras: int | None = None,
544
+ lora_path: str | None = None,
545
+ lora_assignment: str = "random",
546
**kwargs,
547
) -> list[SampleRequest]:
548
# validate total input tokens (prefix + sampled) is at least 1.
@@ -522,11 +587,18 @@ class RandomDataset(BenchmarkDataset):
587
allowed_tokens=allowed_tokens,
588
)
589
token_mismatch_total += token_mismatch
590
+ lora_req = self.get_lora_request(
591
+ index=i,
592
+ max_loras=max_loras,
593
+ lora_path=lora_path,
594
+ lora_assignment=lora_assignment,
595
+ )
596
requests.append(
597
SampleRequest(
598
prompt=prompt,
599
prompt_len=total_input_len,
600
expected_output_len=int(output_lens[i]),
601
+ lora_request=lora_req,
602
request_id=request_id_prefix + str(i),
603
)
604
)
@@ -1263,6 +1335,7 @@ class ShareGPTDataset(BenchmarkDataset):
1335
enable_multimodal_chat: bool = False,
1336
request_id_prefix: str = "",
1337
no_oversample: bool = False,
1338
+ lora_assignment: str = "random",
1339
**kwargs,
1340
) -> list:
1341
samples: list = []
@@ -1275,8 +1348,11 @@ class ShareGPTDataset(BenchmarkDataset):
1348
entry["conversations"][1]["value"],
1349
)
1350
1278
- lora_request = self.get_random_lora_request(
1279
- max_loras=max_loras, lora_path=lora_path
1351
+ lora_request = self.get_lora_request(
1352
+ index=ind,
1353
+ max_loras=max_loras,
1354
+ lora_path=lora_path,
1355
+ lora_assignment=lora_assignment,
1356
)
1357
prompt_ids = tokenizer(prompt).input_ids
1358
completion_ids = tokenizer(completion).input_ids
@@ -2413,6 +2489,7 @@ class BurstGPTDataset(BenchmarkDataset):
2489
lora_path: str | None = None,
2490
request_id_prefix: str = "",
2491
no_oversample: bool = False,
2492
+ lora_assignment: str = "random",
2493
**kwargs,
2494
) -> list[SampleRequest]:
2495
samples = []
@@ -2420,8 +2497,11 @@ class BurstGPTDataset(BenchmarkDataset):
2497
for i in range(num_requests):
2498
input_len = int(data[i][2])
2499
output_len = int(data[i][3])
2423
- lora_req = self.get_random_lora_request(
2424
- max_loras=max_loras, lora_path=lora_path
2500
+ lora_req = self.get_lora_request(
2501
+ index=i,
2502
+ max_loras=max_loras,
2503
+ lora_path=lora_path,
2504
+ lora_assignment=lora_assignment,
2505
)
2506
vocab_size = tokenizer.vocab_size
2507
# Generate a synthetic prompt: a list of token IDs computed as (i +
vllm/benchmarks/serve.py
+28
-5
@@ -624,6 +624,7 @@ async def benchmark(
624
lora_modules: Iterable[str] | None,
625
extra_headers: dict | None,
626
extra_body: dict | None,
627
+ lora_assignment: Literal["random", "round-robin"] = "random",
628
ramp_up_strategy: Literal["linear", "exponential"] | None = None,
629
ramp_up_start_rps: int | None = None,
630
ramp_up_end_rps: int | None = None,
@@ -731,10 +732,20 @@ async def benchmark(
732
print("Starting main benchmark run...")
733
734
if lora_modules:
734
- # For each input request, choose a LoRA module at random.
735
- lora_modules = iter(
736
- [random.choice(lora_modules) for _ in range(len(input_requests))]
737
- )
735
+ lora_modules_list = list(lora_modules)
736
+ if lora_assignment == "round-robin":
737
+ # Deterministic round-robin assignment across requests.
738
+ lora_modules = iter(
739
+ [
740
+ lora_modules_list[i % len(lora_modules_list)]
741
+ for i in range(len(input_requests))
742
+ ]
743
+ )
744
+ else:
745
+ # For each input request, choose a LoRA module at random.
746
+ lora_modules = iter(
747
+ [random.choice(lora_modules_list) for _ in range(len(input_requests))]
748
+ )
749
750
if profile:
751
print("Starting profiler...")
@@ -1523,7 +1534,18 @@ def add_cli_args(parser: argparse.ArgumentParser):
1534
default=None,
1535
help="A subset of LoRA module names passed in when "
1536
"launching the server. For each request, the "
1526
- "script chooses a LoRA module at random.",
1537
+ "script chooses a LoRA module at random by default. "
1538
+ "Use --lora-assignment to control selection strategy.",
1539
+ )
1540
+
1541
+ parser.add_argument(
1542
+ "--lora-assignment",
1543
+ type=str,
1544
+ default="random",
1545
+ choices=["random", "round-robin"],
1546
+ help="Strategy for assigning LoRA modules to requests. "
1547
+ "'random' (default) selects a LoRA at random for each request. "
1548
+ "'round-robin' cycles through LoRA modules deterministically.",
1549
)
1550
1551
parser.add_argument(
@@ -1788,6 +1810,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1810
goodput_config_dict=goodput_config_dict,
1811
max_concurrency=args.max_concurrency,
1812
lora_modules=args.lora_modules,
1813
+ lora_assignment=args.lora_assignment,
1814
extra_headers=headers,
1815
extra_body=extra_body,
1816
ramp_up_strategy=args.ramp_up_strategy,
vllm/benchmarks/throughput.py
+10
@@ -350,6 +350,7 @@ def get_requests(args, tokenizer):
350
"tokenizer": tokenizer,
351
"lora_path": args.lora_path,
352
"max_loras": args.max_loras,
353
+ "lora_assignment": getattr(args, "lora_assignment", "random"),
354
"num_requests": args.num_prompts,
355
}
356
@@ -778,6 +779,15 @@ def add_cli_args(parser: argparse.ArgumentParser):
779
help="Path to the lora adapters to use. This can be an absolute path, "
780
"a relative path, or a Hugging Face model identifier.",
781
)
782
+ parser.add_argument(
783
+ "--lora-assignment",
784
+ type=str,
785
+ default="random",
786
+ choices=["random", "round-robin"],
787
+ help="Strategy for assigning LoRA adapters to requests. "
788
+ "'random' (default) selects a LoRA at random for each request. "
789
+ "'round-robin' cycles through LoRAs deterministically.",
790
+ )
791
parser.add_argument(
792
"--prefix-len",
793
type=int,