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 10, 10:15 UTC
[Feature][Benchmarks] Custom dataset: read output length from dataset (#31881)
Signed-off-by: Sophie du Couédic <sop@zurich.ibm.com>
Sophie du Couédic committed
Jan 9, 2026 at 13:40 UTC
b474782ad776d4517675a6eae6f424a6d3ddf7ca
1 file changed
+26
-5
vllm/benchmarks/datasets.py
+26
-5
@@ -1376,7 +1376,9 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1376
"--custom-output-len",
1377
type=int,
1378
default=256,
1379
- help="Number of output tokens per request, used only for custom dataset.",
1379
+ help="Number of output tokens per request. Unless it is set to -1, the "
1380
+ "value overrides potential output length loaded from the dataset. It is "
1381
+ "used only for custom dataset.",
1382
)
1383
1384
spec_bench_group = parser.add_argument_group("spec bench dataset options")
@@ -1958,10 +1960,12 @@ class CustomDataset(BenchmarkDataset):
1960
Implements the Custom dataset. Loads data from a JSONL file and generates
1961
sample requests based on conversation turns. E.g.,
1962
```
1961
- {"prompt": "What is the capital of India?"}
1962
- {"prompt": "What is the capital of Iran?"}
1963
- {"prompt": "What is the capital of China?"}
1963
+ {"prompt": "What is the capital of India?", "output_tokens": 10}
1964
+ {"prompt": "What is the capital of Iran?", "output_tokens": 1520}
1965
+ {"prompt": "What is the capital of China?", "output_tokens": 819}
1966
```
1967
+ Note that 'output_tokens' column is optional and has to be provided only if
1968
+ 'custom-output-len' argument is None or -1.
1969
"""
1970
1971
def __init__(self, **kwargs) -> None:
@@ -2031,6 +2035,23 @@ class CustomDataset(BenchmarkDataset):
2035
break
2036
prompt = item["prompt"]
2037
2038
+ new_output_len = output_len
2039
+ if output_len is None or output_len == -1:
2040
+ # check that the request has an 'output_tokens' field
2041
+ if "output_tokens" not in item:
2042
+ raise ValueError(
2043
+ "If no output length is provided the "
2044
+ "custom dataset must contain an 'output_tokens' field."
2045
+ )
2046
+ # Use number of output tokens from the request data
2047
+ try:
2048
+ new_output_len = int(item["output_tokens"])
2049
+ except (ValueError, TypeError) as e:
2050
+ raise ValueError(
2051
+ f"Invalid value for 'output_tokens' in custom dataset: "
2052
+ f"'{item['output_tokens']}'. Must be an integer."
2053
+ ) from e
2054
+
2055
# apply template
2056
if not skip_chat_template:
2057
prompt = tokenizer.apply_chat_template(
@@ -2044,7 +2065,7 @@ class CustomDataset(BenchmarkDataset):
2065
SampleRequest(
2066
prompt=prompt,
2067
prompt_len=prompt_len,
2047
- expected_output_len=output_len,
2068
+ expected_output_len=new_output_len,
2069
request_id=request_id_prefix + str(i),
2070
)
2071
)