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 3, 13:00 UTC
[Misc] Add common random prefix option to structured-output serving benchmark (#41632)
Signed-off-by: Viktor Pus <viktorpus@tenstorrent.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Viktor Pus committed
May 16, 2026 at 02:44 UTC
87a2adcb43513ead1434aff03a535d86f56f768b
1 file changed
+45
benchmarks/benchmark_serving_structured_output.py
+45
@@ -115,6 +115,39 @@ class SampleRequest:
115
def sample_requests(
116
tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace
117
) -> list[SampleRequest]:
118
+ def _apply_random_prefix(
119
+ tokenizer: PreTrainedTokenizerBase,
120
+ requests: list[SampleRequest],
121
+ prefix_len: int,
122
+ seed: int,
123
+ ) -> list[SampleRequest]:
124
+ if prefix_len <= 0:
125
+ return requests
126
+ rng = np.random.default_rng(seed)
127
+ vocab_size = tokenizer.vocab_size
128
+ prohibited = getattr(tokenizer, "all_special_ids", None) or []
129
+ allowed = np.array([i for i in range(vocab_size) if i not in prohibited])
130
+ if len(allowed) == 0:
131
+ return requests
132
+ prefix_ids = rng.integers(0, len(allowed), size=prefix_len)
133
+ prefix_token_ids = allowed[prefix_ids].tolist()
134
+ out = []
135
+ for req in requests:
136
+ prompt_ids = tokenizer(req.prompt, add_special_tokens=False).input_ids
137
+ full_ids = prefix_token_ids + prompt_ids
138
+ full_prompt = tokenizer.decode(full_ids, skip_special_tokens=False)
139
+ out.append(
140
+ SampleRequest(
141
+ prompt=full_prompt,
142
+ prompt_len=len(tokenizer(full_prompt).input_ids),
143
+ expected_output_len=req.expected_output_len,
144
+ schema=req.schema,
145
+ structure_type=req.structure_type,
146
+ completion=req.completion,
147
+ )
148
+ )
149
+ return out
150
+
151
if args.dataset == "json" or args.dataset == "json-unique":
152
if args.json_schema_path is None:
153
dir_path = os.path.dirname(os.path.realpath(__file__))
@@ -261,6 +294,9 @@ def sample_requests(
294
)
295
)
296
297
+ requests = _apply_random_prefix(
298
+ tokenizer, requests, args.random_prefix_len, args.seed
299
+ )
300
return requests
301
302
@@ -945,6 +981,15 @@ def create_argument_parser():
981
"results in a more uniform arrival of requests.",
982
)
983
parser.add_argument("--seed", type=int, default=0)
984
+ parser.add_argument(
985
+ "--random-prefix-len",
986
+ type=int,
987
+ default=0,
988
+ help=(
989
+ "Number of prefix tokens to prepend to every prompt. "
990
+ "The same prefix is used for all prompts to enable prefix caching."
991
+ ),
992
+ )
993
parser.add_argument(
994
"--trust-remote-code",
995
action="store_true",