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 2, 02:15 UTC
[Benchmark] Enable reasoning-model (thinking) benchmarking via `--chat-template-kwargs` for client-rendered datasets (#44244)
Signed-off-by: Albert Cheng <albertching0112@gmail.com>
Albert Cheng committed
Jun 2, 2026 at 22:49 UTC
e0081ef8cf0e1e36b4363137de430a73979bc1ab
3 files changed
+88
tests/benchmarks/test_custom_dataset_chat_template_kwargs.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import argparse
4
+import json
5
+from pathlib import Path
6
+
7
+import pytest
8
+
9
+from vllm.benchmarks.datasets import get_samples
10
+
11
+
12
+class _RecordingTokenizer:
13
+ """Minimal tokenizer stub that records the kwargs forwarded to
14
+ apply_chat_template, so we can assert chat_template_kwargs propagation
15
+ without loading a real model/template."""
16
+
17
+ def __init__(self) -> None:
18
+ self.captured_kwargs: dict | None = None
19
+ self.chat_template = "dummy-template"
20
+
21
+ def apply_chat_template(
22
+ self,
23
+ conversation,
24
+ add_generation_prompt: bool = True,
25
+ tokenize: bool = False,
26
+ **kwargs,
27
+ ) -> str:
28
+ self.captured_kwargs = kwargs
29
+ return conversation[0]["content"]
30
+
31
+ def __call__(self, text: str):
32
+ return argparse.Namespace(input_ids=list(range(len(text.split()))))
33
+
34
+
35
+def _args(dataset_path: str, chat_template_kwargs) -> argparse.Namespace:
36
+ return argparse.Namespace(
37
+ dataset_name="custom",
38
+ dataset_path=dataset_path,
39
+ disable_shuffle=True,
40
+ num_prompts=1,
41
+ custom_output_len=32,
42
+ skip_chat_template=False,
43
+ chat_template_kwargs=chat_template_kwargs,
44
+ no_oversample=False,
45
+ seed=0,
46
+ request_id_prefix="",
47
+ )
48
+
49
+
50
+def _write_one(path: Path) -> None:
51
+ path.write_text(json.dumps({"prompt": "hello world"}) + "\n")
52
+
53
+
54
+@pytest.mark.benchmark
55
+def test_chat_template_kwargs_forwarded(tmp_path: Path) -> None:
56
+ """--chat-template-kwargs must reach the client-side apply_chat_template."""
57
+ jsonl = tmp_path / "data.jsonl"
58
+ _write_one(jsonl)
59
+
60
+ tok = _RecordingTokenizer()
61
+ get_samples(_args(str(jsonl), {"thinking": True}), tok)
62
+
63
+ assert tok.captured_kwargs == {"thinking": True}
64
+
65
+
66
+@pytest.mark.benchmark
67
+def test_chat_template_kwargs_default_is_noop(tmp_path: Path) -> None:
68
+ """When not provided, no extra kwargs are passed (existing behavior)."""
69
+ jsonl = tmp_path / "data.jsonl"
70
+ _write_one(jsonl)
71
+
72
+ tok = _RecordingTokenizer()
73
+ get_samples(_args(str(jsonl), None), tok)
74
+
75
+ assert tok.captured_kwargs == {}
vllm/benchmarks/datasets/datasets.py
+4
@@ -2080,6 +2080,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2080
tokenizer=tokenizer,
2081
output_len=args.custom_output_len,
2082
skip_chat_template=args.skip_chat_template,
2083
+ chat_template_kwargs=getattr(args, "chat_template_kwargs", None),
2084
request_id_prefix=args.request_id_prefix,
2085
no_oversample=args.no_oversample,
2086
)
@@ -2409,6 +2410,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2410
num_requests=args.num_prompts,
2411
tokenizer=tokenizer,
2412
output_len=args.speed_bench_output_len,
2413
+ chat_template_kwargs=getattr(args, "chat_template_kwargs", None),
2414
enable_multimodal_chat=args.enable_multimodal_chat,
2415
request_id_prefix=args.request_id_prefix,
2416
no_oversample=args.no_oversample,
@@ -2496,6 +2498,7 @@ class CustomDataset(BenchmarkDataset):
2498
output_len: int | None = None,
2499
enable_multimodal_chat: bool = False,
2500
skip_chat_template: bool = False,
2501
+ chat_template_kwargs: dict | None = None,
2502
**kwargs,
2503
) -> list[SampleRequest]:
2504
# load all data if needed
@@ -2543,6 +2546,7 @@ class CustomDataset(BenchmarkDataset):
2546
[{"role": "user", "content": prompt}],
2547
add_generation_prompt=True,
2548
tokenize=False,
2549
+ **(chat_template_kwargs or {}),
2550
)
2551
2552
prompt_len = len(tokenizer(prompt).input_ids)
vllm/benchmarks/serve.py
+9
@@ -1609,6 +1609,15 @@ def add_cli_args(parser: argparse.ArgumentParser):
1609
"in seconds. Ready check will be skipped by default.",
1610
)
1611
1612
+ parser.add_argument(
1613
+ "--chat-template-kwargs",
1614
+ type=json.loads,
1615
+ default=None,
1616
+ help="A JSON string of kwargs forwarded to the tokenizer's "
1617
+ "apply_chat_template when a dataset renders prompts client-side "
1618
+ "(e.g. custom / speed_bench). "
1619
+ "Example: '{\"thinking\": true}' to enable reasoning models.",
1620
+ )
1621
parser.add_argument(
1622
"--extra-body",
1623
help="A JSON string representing extra body parameters to include "