@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bfb9ebc211

Fork 0

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, 14:45 UTC

[Feature] Add support for timed trace replay in `vllm bench serve` to replay Moonshot and Alibaba workload traces (#39795)

Signed-off-by: Animesh Trivedi <Animesh.Trivedi@ibm.com>

Animesh Trivedi committed May 28, 2026 at 12:31 UTC bfb9ebc211d3b4abaa3fc2c87315ff7801d93d5b

3 files changed +351 -53

docs/benchmarking/cli.md

+35

@@ -918,6 +918,41 @@ vllm bench serve \

918

919 </details>

920

921 +### Replay Timed Traces

922 +

923 +<details class="admonition abstract" markdown="1">

924 +<summary>Show more</summary>

925 +

926 +Example of how to run traces which have timing information

927 +with them.

928 +

929 +#### Running MoonshotAI traces

930 +

931 +Start the server:

932 +

933 +```bash

934 +vllm serve Qwen/Qwen3.5-2B \

935 +--host 127.0.0.1 --port 8000

936 +```

937 +

938 +Run the benchmark:

939 +

940 +```bash

941 +# Download an example trace

942 +# curl -L -o conversation_trace.jsonl \

943 +#https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl

944 +

945 +vllm bench serve --model Qwen/Qwen3.5-2B \

946 +--dataset-name=timed_trace --num-prompts 100 --host 127.0.0.1 \

947 +--port 8000 --dataset-path ./conversation_trace.jsonl \

948 +--ignore-eos --self-timed --timed-trace-chunk-hash-size 512 \

949 +--timed-trace-sec-multiplier 0.001

950 +```

951 +

952 +This will replay the first 100 lines from the trace file `conversation.jsonl`.

953 +

954 +</details>

955 +

956 ### 🧪 Hashing Benchmarks

957

958 <details class="admonition abstract" markdown="1">

vllm/benchmarks/datasets/datasets.py

+219

@@ -83,6 +83,7 @@ class SampleRequest:

83 multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None

84 lora_request: LoRARequest | None = None

85 request_id: str | None = None

86 + timestamp: float | None = None

89 # -----------------------------------------------------------------------------

@@ -1399,6 +1400,168 @@ class ShareGPTDataset(BenchmarkDataset):

1400 return samples

1401

1402

1403 +class TimedTrace(BenchmarkDataset):

1404 + """

1405 + Implements a base class to replay various timed traces.

1406 + Loads data from a JSON file and generates sample requests

1407 + based on the timing information in the traces.

1408 + """

1409 +

1410 + def __init__(self, **kwargs) -> None:

1411 + super().__init__(**kwargs)

1412 + random.seed(self.random_seed)

1413 + np.random.seed(self.random_seed)

1414 +

1415 + # Set parameters with defaults from timed_trace_group arguments

1416 + self.chunk_size = int(kwargs.get("timed_trace_chunk_hash_size", 16))

1417 + self.sec_multiplier = float(kwargs.get("timed_trace_sec_multiplier", 1))

1418 + self.label_ts = str(kwargs.get("timed_trace_label_timestamp"))

1419 + self.label_input_length = str(kwargs.get("timed_trace_label_input_length"))

1420 + self.label_output_length = str(kwargs.get("timed_trace_label_output_length"))

1421 + self.label_hash_ids = str(kwargs.get("timed_trace_label_hash_ids"))

1422 + print(

1423 + f"timed-trace: chunk_size: {self.chunk_size}, "

1424 + f"sec_multiplier: {self.sec_multiplier}, "

1425 + f'label_ts: "{self.label_ts}", '

1426 + f'label_input_length: "{self.label_input_length}", '

1427 + f'label_output_length: "{self.label_output_length}", '

1428 + f'label_hash_ids: "{self.label_hash_ids}"'

1429 + )

1430 + self._expanded_generated_prompts = {}

1431 + self.load_data()

1432 +

1433 + def load_data(self) -> None:

1434 + # check if the file is there

1435 + if self.dataset_path is None:

1436 + raise ValueError("dataset_path must be provided for loading data.")

1437 +

1438 + # load and we will do transformation once we have the Tokenizer available

1439 + # this is jsonl data format

1440 + with open(self.dataset_path) as f:

1441 + self.data = f.readlines()

1442 +

1443 + def _sample_token(

1444 + self, num_tokens: int, tokenizer: TokenizerLike, seed: int | None = None

1445 + ) -> list[int]:

1446 + # Initialize vocab only if it doesn't exist yet

1447 + if not hasattr(self, "vocab"):

1448 + self.vocab = tokenizer.get_vocab()

1449 + # Remove the special tokens.

1450 + self.vocab = {

1451 + k: v

1452 + for k, v in self.vocab.items()

1453 + if v not in tokenizer.all_special_ids

1454 + }

1455 + # Create a sorted list of vocab values for deterministic sampling

1456 + self.vocab_values_sorted = sorted(self.vocab.values())

1457 +

1458 + # Use the provided seed if given, otherwise use global random state

1459 + if seed is not None:

1460 + rng = random.Random(seed)

1461 + sampled_token_ids = rng.choices(self.vocab_values_sorted, k=num_tokens)

1462 + else:

1463 + sampled_token_ids = random.choices(self.vocab_values_sorted, k=num_tokens)

1464 +

1465 + return sampled_token_ids

1466 +

1467 + def _expand_prompt(

1468 + self,

1469 + chunked_hashes: list[int],

1470 + target_input_size: int,

1471 + tokenizer: TokenizerLike,

1472 + ) -> list[int]:

1473 + raw_tokenized_prompt = []

1474 + for h in chunked_hashes:

1475 + # Calculate how many tokens to expand for this chunk

1476 + expanded_size = (

1477 + self.chunk_size

1478 + if target_input_size >= self.chunk_size

1479 + else target_input_size

1480 + )

1481 +

1482 + # Cache key includes size for partial chunks at the end

1483 + key = f"{h}:{expanded_size}"

1484 +

1485 + if key not in self._expanded_generated_prompts:

1486 + # Convert key to a deterministic seed

1487 + key_seed = hash(key) & 0xFFFFFFFF # Convert to 32-bit int

1488 + self._expanded_generated_prompts[key] = self._sample_token(

1489 + expanded_size, tokenizer, seed=key_seed

1490 + )

1491 + # once inserted get the tokenized prompt and append to the list

1492 + raw_tokenized_prompt.extend(self._expanded_generated_prompts[key])

1493 + target_input_size -= expanded_size

1494 +

1495 + if target_input_size <= 0:

1496 + break

1497 +

1498 + return raw_tokenized_prompt

1499 +

1500 + def sample(

1501 + self,

1502 + tokenizer: TokenizerLike,

1503 + num_requests: int,

1504 + request_id_prefix: str = "",

1505 + **kwargs,

1506 + ) -> list:

1507 + samples: list = []

1508 + assert tokenizer is not None, "Tokenizer must be provided, now is Null"

1509 +

1510 + for ind, entry in enumerate(self.data):

1511 + if len(samples) >= num_requests:

1512 + break

1513 +

1514 + # now we create the SampleRequest with timing info

1515 + entry = json.loads(entry.strip())

1516 + input_length = entry.get(self.label_input_length)

1517 + if input_length is None:

1518 + raise ValueError(

1519 + f"Input length field '{self.label_input_length}' "

1520 + f"not found in trace entry. "

1521 + f"Available fields: {list(entry.keys())}. "

1522 + f"Use --label-input-length to specify the correct "

1523 + f"field name."

1524 + )

1525 + new_output_len = entry.get(self.label_output_length)

1526 + if new_output_len is None:

1527 + raise ValueError(

1528 + f"Output length field '{self.label_output_length}' "

1529 + f"not found in trace entry. "

1530 + f"Available fields: {list(entry.keys())}. "

1531 + f"Use --label-output-length to specify the correct "

1532 + f"field name."

1533 + )

1534 + prompt_ids = self._expand_prompt(

1535 + entry.get(self.label_hash_ids, []), input_length, tokenizer

1536 + )

1537 + prompt = tokenizer.decode(prompt_ids)

1538 +

1539 + # Get timestamp with proper error handling

1540 + ts_value = entry.get(self.label_ts)

1541 + if ts_value is None:

1542 + raise ValueError(

1543 + f"Timestamp field '{self.label_ts}' not found in trace entry. "

1544 + f"Available fields: {list(entry.keys())}. "

1545 + f"Use --label-timestamp to specify the correct field name."

1546 + )

1547 + timestamp = float(ts_value) * self.sec_multiplier

1548 +

1549 + prompt_len = len(prompt_ids)

1550 +

1551 + samples.append(

1552 + SampleRequest(

1553 + prompt=prompt,

1554 + prompt_len=prompt_len,

1555 + expected_output_len=new_output_len,

1556 + lora_request=None,

1557 + multi_modal_data=None,

1558 + request_id=request_id_prefix + str(ind),

1559 + timestamp=timestamp,

1560 + )

1561 + )

1562 + return samples

1563 +

1564 +

1565 def add_dataset_parser(parser: FlexibleArgumentParser):

1566 parser.add_argument(

1567 "--trust-remote-code",

@@ -1431,6 +1594,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser):

1594 "prefix_repetition",

1595 "spec_bench",

1596 "speed_bench",

1597 + "timed_trace",

1598 ],

1599 help="Name of the dataset to benchmark on.",

1600 )

@@ -1521,6 +1685,53 @@ def add_dataset_parser(parser: FlexibleArgumentParser):

1685 "from the ShareGPT dataset.",

1686 )

1687

1688 + timed_trace_group = parser.add_argument_group("timed-trace dataset options")

1689 + timed_trace_group.add_argument(

1690 + "--timed-trace-chunk-hash-size",

1691 + type=int,

1692 + default=16,

1693 + help=(

1694 + "Each hash tokens, if present, represent how many token "

1695 + "hashes. For example in the Moonshot traces it is 512, while "

1696 + "the Qwen/Alibaba has 16."

1697 + ),

1698 + )

1699 + timed_trace_group.add_argument(

1700 + "--timed-trace-sec-multiplier",

1701 + type=float,

1702 + default=1,

1703 + help=(

1704 + "What multiplier to use when converting timestamps to "

1705 + "seconds. We will multiply timestamps by this. For example"

1706 + "if the timestamps are in milliseconds, then pass 0.001."

1707 + "If they are already in seconds, then the default 1 is sufficient."

1708 + ),

1709 + )

1710 + timed_trace_group.add_argument(

1711 + "--timed-trace-label-timestamp",

1712 + type=str,

1713 + default="timestamp",

1714 + help="What json label to use to index the timestamp in the trace.",

1715 + )

1716 + timed_trace_group.add_argument(

1717 + "--timed-trace-label-input-length",

1718 + type=str,

1719 + default="input_length",

1720 + help=("What json label to use to index the input length field in the trace."),

1721 + )

1722 + timed_trace_group.add_argument(

1723 + "--timed-trace-label-output-length",

1724 + type=str,

1725 + default="output_length",

1726 + help=("What json label to use to index the output length field in the trace."),

1727 + )

1728 + timed_trace_group.add_argument(

1729 + "--timed-trace-label-hash-ids",

1730 + type=str,

1731 + default="hash_ids",

1732 + help=("What json label to use to index the hash ids for the input prompts."),

1733 + )

1734 +

1735 blazedit_group = parser.add_argument_group("blazedit dataset options")

1736 blazedit_group.add_argument(

1737 "--blazedit-min-distance",

@@ -2051,6 +2262,14 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:

2262 **hf_kwargs,

2263 )

2264

2265 + elif args.dataset_name == "timed_trace":

2266 + dataloader = TimedTrace(**vars(args))

2267 + input_requests = dataloader.sample(

2268 + num_requests=args.num_prompts,

2269 + tokenizer=tokenizer,

2270 + request_id_prefix=args.request_id_prefix,

2271 + )

2272 +

2273 else:

2274 # For datasets that follow a similar structure, use a mapping.

2275 dataset_mapping = {

vllm/benchmarks/serve.py

+97 -53

@@ -251,6 +251,7 @@ async def get_request(

251 ramp_up_strategy: Literal["linear", "exponential"] | None = None,

252 ramp_up_start_rps: int | None = None,

253 ramp_up_end_rps: int | None = None,

254 + self_timed: bool = False,

255 ) -> AsyncGenerator[tuple[SampleRequest, float], None]:

256 """

257 Asynchronously generates requests at a specified rate

@@ -290,49 +291,59 @@ async def get_request(

291 # Precompute delays among requests to minimize request send laggings

292 request_rates = []

293 delay_ts = []

293 - for request_index, request in enumerate(input_requests):

294 - current_request_rate = _get_current_request_rate(

295 - ramp_up_strategy,

296 - ramp_up_start_rps,

297 - ramp_up_end_rps,

298 - request_index,

299 - total_requests,

300 - request_rate,

301 - )

302 - assert current_request_rate > 0.0, (

303 - f"Obtained non-positive request rate {current_request_rate}."

304 - )

305 - request_rates.append(current_request_rate)

306 - if current_request_rate == float("inf"):

307 - delay_ts.append(0)

308 - elif burstiness == float("inf"):

309 - # when burstiness tends to infinity, the delay time becomes constant

310 - # and tends to the inverse of the request rate

311 - delay_ts.append(1.0 / current_request_rate)

312 - else:

313 - theta = 1.0 / (current_request_rate * burstiness)

314 -

315 - # Sample the request interval from the gamma distribution.

316 - # If burstiness is 1, it follows exponential distribution.

317 - delay_ts.append(np.random.gamma(shape=burstiness, scale=theta))

318 -

319 - # Calculate the cumulative delay time from the first sent out requests.

320 - for i in range(1, len(delay_ts)):

321 - delay_ts[i] += delay_ts[i - 1]

322 - if ramp_up_strategy is None and delay_ts[-1] != 0:

323 - # When ramp_up_strategy is not set, we assume the request rate is fixed

324 - # and all requests should be sent in target_total_delay_s, the following

325 - # logic would re-scale delay time to ensure the final delay_ts

326 - # align with target_total_delay_s.

327 - #

328 - # NOTE: If we simply accumulate the random delta values

329 - # from the gamma distribution, their sum would have 1-2% gap

330 - # from target_total_delay_s. The purpose of the following logic is to

331 - # close the gap for stabilizing the throughput data

332 - # from different random seeds.

333 - target_total_delay_s = total_requests / request_rate

334 - normalize_factor = target_total_delay_s / delay_ts[-1]

335 - delay_ts = [delay * normalize_factor for delay in delay_ts]

294 +

295 + # if the traces have timing info then:

296 + if not self_timed:

297 + for request_index, request in enumerate(input_requests):

298 + current_request_rate = _get_current_request_rate(

299 + ramp_up_strategy,

300 + ramp_up_start_rps,

301 + ramp_up_end_rps,

302 + request_index,

303 + total_requests,

304 + request_rate,

305 + )

306 + assert current_request_rate > 0.0, (

307 + f"Obtained non-positive request rate {current_request_rate}."

308 + )

309 + request_rates.append(current_request_rate)

310 + if current_request_rate == float("inf"):

311 + delay_ts.append(0)

312 + elif burstiness == float("inf"):

313 + # when burstiness tends to infinity, the delay time becomes constant

314 + # and tends to the inverse of the request rate

315 + delay_ts.append(1.0 / current_request_rate)

316 + else:

317 + theta = 1.0 / (current_request_rate * burstiness)

318 +

319 + # Sample the request interval from the gamma distribution.

320 + # If burstiness is 1, it follows exponential distribution.

321 + delay_ts.append(np.random.gamma(shape=burstiness, scale=theta))

322 +

323 + # Calculate the cumulative delay time from the first sent out requests.

324 + for i in range(1, len(delay_ts)):

325 + delay_ts[i] += delay_ts[i - 1]

326 + if ramp_up_strategy is None and delay_ts[-1] != 0:

327 + # When ramp_up_strategy is not set, we assume the request rate is fixed

328 + # and all requests should be sent in target_total_delay_s, the following

329 + # logic would re-scale delay time to ensure the final delay_ts

330 + # align with target_total_delay_s.

331 + #

332 + # NOTE: If we simply accumulate the random delta values

333 + # from the gamma distribution, their sum would have 1-2% gap

334 + # from target_total_delay_s. The purpose of the following logic is to

335 + # close the gap for stabilizing the throughput data

336 + # from different random seeds.

337 + target_total_delay_s = total_requests / request_rate

338 + normalize_factor = target_total_delay_s / delay_ts[-1]

339 + delay_ts = [delay * normalize_factor for delay in delay_ts]

340 + else:

341 + for request_index, request in enumerate(input_requests):

342 + # this is cumulative running ts, from which sleep is calculated later

343 + delay_ts.append(request.timestamp)

344 + # TODO: there is no notion of RPS here, may be we can calculate

345 + # from the trace.

346 + request_rates.append(0.0)

347

348 start_ts = time.time()

349 for request_index, request in enumerate(input_requests):

@@ -634,6 +645,7 @@ async def benchmark(

645 ramp_up_end_rps: int | None = None,

646 ready_check_timeout_sec: int = 600,

647 ssl_context: ssl.SSLContext | bool | None = None,

648 + self_timed: bool = False,

649 ):

650 try:

651 request_func = ASYNC_REQUEST_FUNCS[endpoint_type]

@@ -773,18 +785,20 @@ async def benchmark(

785 print("Profiler started")

786

787 distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"

788 + if not self_timed:

789 + if ramp_up_strategy is not None:

790 + print(f"Traffic ramp-up strategy: {ramp_up_strategy}.")

791 + print(

792 + f"Will increase RPS from {ramp_up_start_rps} to "

793 + f"{ramp_up_end_rps} RPS over the duration of the benchmark."

794 + )

795 + else:

796 + print(f"Traffic request rate: {request_rate}")

797

777 - if ramp_up_strategy is not None:

778 - print(f"Traffic ramp-up strategy: {ramp_up_strategy}.")

779 - print(

780 - f"Will increase RPS from {ramp_up_start_rps} to "

781 - f"{ramp_up_end_rps} RPS over the duration of the benchmark."

782 - )

798 + print(f"Burstiness factor: {burstiness} ({distribution})")

799 + print(f"Maximum request concurrency: {max_concurrency}")

800 else:

784 - print(f"Traffic request rate: {request_rate}")

785 -

786 - print(f"Burstiness factor: {burstiness} ({distribution})")

787 - print(f"Maximum request concurrency: {max_concurrency}")

801 + print("Self timing is set, using the timestamps from the trace file.")

802

803 spec_decode_metrics_before = await fetch_spec_decode_metrics(base_url, session)

804

@@ -823,6 +837,7 @@ async def benchmark(

837 ramp_up_strategy,

838 ramp_up_start_rps,

839 ramp_up_end_rps,

840 + self_timed,

841 ):

842 if ramp_up_strategy is not None:

843 current_int_rps = int(current_request_rate)

@@ -1436,6 +1451,16 @@ def add_cli_args(parser: argparse.ArgumentParser):

1451 help="Set ignore_eos flag when sending the benchmark request."

1452 "Warning: ignore_eos is not supported in deepspeed_mii and tgi.",

1453 )

1454 + parser.add_argument(

1455 + "--self-timed",

1456 + action=argparse.BooleanOptionalAction,

1457 + default=None,

1458 + help="Use timing information from the traces instead of the configuration. "

1459 + "This is useful when replaying traces faithfully based on their timestamps. "

1460 + "When unset, defaults to False, except for --dataset-name=timed_trace where "

1461 + "it defaults to True. Use --no-self-timed to force off. When off, user "

1462 + "defined generation rates are used and in trace timing info is ignored.",

1463 + )

1464 parser.add_argument(

1465 "--percentile-metrics",

1466 type=str,

@@ -1767,6 +1792,24 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:

1792 ):

1793 args.ignore_eos = True

1794

1795 + if args.dataset_name == "timed_trace":

1796 + # timed_trace carries per-request timestamps;

1797 + # ignore EOS so generation runs to the trace's specified output length,

1798 + # and default to using those timestamps for scheduling unless the user

1799 + # opted out.

1800 + args.ignore_eos = True

1801 + if args.self_timed is None:

1802 + args.self_timed = True

1803 + else:

1804 + # if this is set for anything else, it is an error

1805 + if args.self_timed is not None:

1806 + raise ValueError(

1807 + "--self-timed/--no-self-timed is only supported with "

1808 + "--dataset-name=timed_trace"

1809 + )

1810 + # for any non self-timed trace, this is False

1811 + args.self_timed = False

1812 +

1813 # Load the dataset.

1814 input_requests = get_samples(args, tokenizer)

1815 goodput_config_dict = check_goodput_args(args)

@@ -1846,6 +1889,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:

1889 ramp_up_end_rps=args.ramp_up_end_rps,

1890 ready_check_timeout_sec=args.ready_check_timeout_sec,

1891 ssl_context=ssl_context,

1892 + self_timed=args.self_timed,

1893 )

1894

1895 # Save config and results to json