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 8, 08:45 UTC
[Bugfix] Fix Random Dataset Prefix Length Inaccuracy (#33907)
Signed-off-by: frankwang28 <frank.wbb@hotmail.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Frank Wang committed
Feb 12, 2026 at 18:21 UTC
b86bf4417e3172b372ff20cccf4d30289a6db8ae
1 file changed
+29
-10
vllm/benchmarks/datasets.py
+29
-10
@@ -380,7 +380,7 @@ def gen_prompt_decode_to_target_len(
380
max_retry: int = 10,
381
add_special_tokens: bool = False,
382
rng: np.random.Generator | None = None,
383
-) -> tuple[str, list[int]]:
383
+) -> tuple[str, list[int], int]:
384
"""
385
Ensure decoded-then-encoded prompt length matches the target token length.
386
@@ -392,7 +392,9 @@ def gen_prompt_decode_to_target_len(
392
[6880, 6881] -> ['Ġcalls', 'here'] ->
393
[1650, 939, 486] -> ['Ġcall', 'sh', 'ere']
394
395
- Returns a tuple of the final prompt string and the adjusted token sequence.
395
+ Returns a tuple of the final prompt string, the adjusted token sequence,
396
+ and the token mismatch (final_len - target_token_len) if the retry budget
397
+ is exhausted.
398
"""
399
remain_num_try = max_retry
400
token_mismatch = 0
@@ -499,7 +501,7 @@ class RandomDataset(BenchmarkDataset):
501
allowed_tokens = np.array(list(set(all_tokens) - set(prohibited_tokens)))
502
503
# Generate prefix once
502
- prefix_token_ids = self.get_prefix(allowed_tokens, prefix_len)
504
+ prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)
505
506
requests = []
507
token_mismatch_total = 0
@@ -554,19 +556,36 @@ class RandomDataset(BenchmarkDataset):
556
557
def get_prefix(
558
self,
559
+ tokenizer: TokenizerLike,
560
allowed_tokens: np.ndarray,
561
prefix_len: int,
562
) -> list[int]:
563
"""
564
Get the prefix for the dataset.
565
"""
563
- return (
564
- allowed_tokens[
565
- self._rng.integers(0, len(allowed_tokens), size=prefix_len)
566
- ].tolist()
567
- if prefix_len > 0
568
- else []
566
+ if prefix_len <= 0:
567
+ return []
568
+
569
+ prefix_tokens = allowed_tokens[
570
+ self._rng.integers(0, len(allowed_tokens), size=prefix_len)
571
+ ].tolist()
572
+ _, adjusted_tokens, token_mismatch = gen_prompt_decode_to_target_len(
573
+ tokenizer=tokenizer,
574
+ token_sequence=prefix_tokens,
575
+ target_token_len=prefix_len,
576
+ add_special_tokens=False,
577
+ rng=self._rng,
578
)
579
+ if token_mismatch != 0:
580
+ sign = "more" if token_mismatch > 0 else "fewer"
581
+ logger.warning(
582
+ "Prefix tokenization produced %d %s tokens than expected "
583
+ "after decoding and re-encoding. This is expected due to "
584
+ "the imperfect nature of the sampling procedure",
585
+ abs(token_mismatch),
586
+ sign,
587
+ )
588
+ return adjusted_tokens
589
590
def get_sampling_params(
591
self,
@@ -1128,7 +1147,7 @@ class RandomMultiModalDataset(RandomDataset):
1147
"Sampling from %d out of %d (vocab size)", len(allowed_tokens), vocab_size
1148
)
1149
# Generate prefix once
1131
- prefix_token_ids = self.get_prefix(allowed_tokens, prefix_len)
1150
+ prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)
1151
# Add synthetic multimodal items to each request
1152
mm_requests = []
1153
token_mismatch_total = 0