obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-07-31 19:08:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Animesh Trivedi <[email protected]>
This commit is contained in:
Animesh Trivedi
2026-05-28 03:31:34 -07:00
committed by GitHub
parent a9bc0ad8e4
commit bfb9ebc211
3 changed files with 349 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/benchmarking/cli.md
+35
View File
@@ -918,6 +918,41 @@ vllm bench serve \
</details>
### Replay Timed Traces
<details class="admonition abstract" markdown="1">
<summary>Show more</summary>
Example of how to run traces which have timing information
with them.
#### Running MoonshotAI traces
Start the server:
```bash
vllm serve Qwen/Qwen3.5-2B \
--host 127.0.0.1 --port 8000
```
Run the benchmark:
```bash
# Download an example trace
# curl -L -o conversation_trace.jsonl \
#https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl
vllm bench serve --model Qwen/Qwen3.5-2B \
--dataset-name=timed_trace --num-prompts 100 --host 127.0.0.1 \
--port 8000 --dataset-path ./conversation_trace.jsonl \
--ignore-eos --self-timed --timed-trace-chunk-hash-size 512 \
--timed-trace-sec-multiplier 0.001
```
This will replay the first 100 lines from the trace file `conversation.jsonl`.
</details>
### 🧪 Hashing Benchmarks
<details class="admonition abstract" markdown="1">
vllm/benchmarks/datasets/datasets.py
+219
View File
@@ -83,6 +83,7 @@ class SampleRequest:
multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
lora_request: LoRARequest | None = None
request_id: str | None = None
timestamp: float | None = None
# -----------------------------------------------------------------------------
@@ -1399,6 +1400,168 @@ class ShareGPTDataset(BenchmarkDataset):
return samples
class TimedTrace(BenchmarkDataset):
"""
Implements a base class to replay various timed traces.
Loads data from a JSON file and generates sample requests
based on the timing information in the traces.
"""
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
random.seed(self.random_seed)
np.random.seed(self.random_seed)
# Set parameters with defaults from timed_trace_group arguments
self.chunk_size = int(kwargs.get("timed_trace_chunk_hash_size", 16))
self.sec_multiplier = float(kwargs.get("timed_trace_sec_multiplier", 1))
self.label_ts = str(kwargs.get("timed_trace_label_timestamp"))
self.label_input_length = str(kwargs.get("timed_trace_label_input_length"))
self.label_output_length = str(kwargs.get("timed_trace_label_output_length"))
self.label_hash_ids = str(kwargs.get("timed_trace_label_hash_ids"))
print(
f"timed-trace: chunk_size: {self.chunk_size}, "
f"sec_multiplier: {self.sec_multiplier}, "
f'label_ts: "{self.label_ts}", '
f'label_input_length: "{self.label_input_length}", '
f'label_output_length: "{self.label_output_length}", '
f'label_hash_ids: "{self.label_hash_ids}"'
)
self._expanded_generated_prompts = {}
self.load_data()
def load_data(self) -> None:
# check if the file is there
if self.dataset_path is None:
raise ValueError("dataset_path must be provided for loading data.")
# load and we will do transformation once we have the Tokenizer available
# this is jsonl data format
with open(self.dataset_path) as f:
self.data = f.readlines()
def _sample_token(
self, num_tokens: int, tokenizer: TokenizerLike, seed: int | None = None
) -> list[int]:
# Initialize vocab only if it doesn't exist yet
if not hasattr(self, "vocab"):
self.vocab = tokenizer.get_vocab()
# Remove the special tokens.
self.vocab = {
k: v
for k, v in self.vocab.items()
if v not in tokenizer.all_special_ids
}
# Create a sorted list of vocab values for deterministic sampling
self.vocab_values_sorted = sorted(self.vocab.values())
# Use the provided seed if given, otherwise use global random state
if seed is not None:
rng = random.Random(seed)
sampled_token_ids = rng.choices(self.vocab_values_sorted, k=num_tokens)
else:
sampled_token_ids = random.choices(self.vocab_values_sorted, k=num_tokens)
return sampled_token_ids
def _expand_prompt(
self,
chunked_hashes: list[int],
target_input_size: int,
tokenizer: TokenizerLike,
) -> list[int]:
raw_tokenized_prompt = []
for h in chunked_hashes:
# Calculate how many tokens to expand for this chunk
expanded_size = (
self.chunk_size
if target_input_size >= self.chunk_size
else target_input_size
)
# Cache key includes size for partial chunks at the end
key = f"{h}:{expanded_size}"
if key not in self._expanded_generated_prompts:
# Convert key to a deterministic seed
key_seed = hash(key) & 0xFFFFFFFF # Convert to 32-bit int
self._expanded_generated_prompts[key] = self._sample_token(
expanded_size, tokenizer, seed=key_seed
)
# once inserted get the tokenized prompt and append to the list
raw_tokenized_prompt.extend(self._expanded_generated_prompts[key])
target_input_size -= expanded_size
if target_input_size <= 0:
break
return raw_tokenized_prompt
def sample(
self,
tokenizer: TokenizerLike,
num_requests: int,
request_id_prefix: str = "",
**kwargs,
) -> list:
samples: list = []
assert tokenizer is not None, "Tokenizer must be provided, now is Null"
for ind, entry in enumerate(self.data):
if len(samples) >= num_requests:
break
# now we create the SampleRequest with timing info
entry = json.loads(entry.strip())
input_length = entry.get(self.label_input_length)
if input_length is None:
raise ValueError(
f"Input length field '{self.label_input_length}' "
f"not found in trace entry. "
f"Available fields: {list(entry.keys())}. "
f"Use --label-input-length to specify the correct "
f"field name."
)
new_output_len = entry.get(self.label_output_length)
if new_output_len is None:
raise ValueError(
f"Output length field '{self.label_output_length}' "
f"not found in trace entry. "
f"Available fields: {list(entry.keys())}. "
f"Use --label-output-length to specify the correct "
f"field name."
)
prompt_ids = self._expand_prompt(
entry.get(self.label_hash_ids, []), input_length, tokenizer
)
prompt = tokenizer.decode(prompt_ids)
# Get timestamp with proper error handling
ts_value = entry.get(self.label_ts)
if ts_value is None:
raise ValueError(
f"Timestamp field '{self.label_ts}' not found in trace entry. "
f"Available fields: {list(entry.keys())}. "
f"Use --label-timestamp to specify the correct field name."
)
timestamp = float(ts_value) * self.sec_multiplier
prompt_len = len(prompt_ids)
samples.append(
SampleRequest(
prompt=prompt,
prompt_len=prompt_len,
expected_output_len=new_output_len,
lora_request=None,
multi_modal_data=None,
request_id=request_id_prefix + str(ind),
timestamp=timestamp,
)
)
return samples
def add_dataset_parser(parser: FlexibleArgumentParser):
parser.add_argument(
"--trust-remote-code",
@@ -1431,6 +1594,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
"prefix_repetition",
"spec_bench",
"speed_bench",
"timed_trace",
],
help="Name of the dataset to benchmark on.",
)
@@ -1521,6 +1685,53 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
"from the ShareGPT dataset.",
)
timed_trace_group = parser.add_argument_group("timed-trace dataset options")
timed_trace_group.add_argument(
"--timed-trace-chunk-hash-size",
type=int,
default=16,
help=(
"Each hash tokens, if present, represent how many token "
"hashes. For example in the Moonshot traces it is 512, while "
"the Qwen/Alibaba has 16."
),
)
timed_trace_group.add_argument(
"--timed-trace-sec-multiplier",
type=float,
default=1,
help=(
"What multiplier to use when converting timestamps to "
"seconds. We will multiply timestamps by this. For example"
"if the timestamps are in milliseconds, then pass 0.001."
"If they are already in seconds, then the default 1 is sufficient."
),
)
timed_trace_group.add_argument(
"--timed-trace-label-timestamp",
type=str,
default="timestamp",
help="What json label to use to index the timestamp in the trace.",
)
timed_trace_group.add_argument(
"--timed-trace-label-input-length",
type=str,
default="input_length",
help=("What json label to use to index the input length field in the trace."),
)
timed_trace_group.add_argument(
"--timed-trace-label-output-length",
type=str,
default="output_length",
help=("What json label to use to index the output length field in the trace."),
)
timed_trace_group.add_argument(
"--timed-trace-label-hash-ids",
type=str,
default="hash_ids",
help=("What json label to use to index the hash ids for the input prompts."),
)
blazedit_group = parser.add_argument_group("blazedit dataset options")
blazedit_group.add_argument(
"--blazedit-min-distance",
@@ -2051,6 +2262,14 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
**hf_kwargs,
)
elif args.dataset_name == "timed_trace":
dataloader = TimedTrace(**vars(args))
input_requests = dataloader.sample(
num_requests=args.num_prompts,
tokenizer=tokenizer,
request_id_prefix=args.request_id_prefix,
)
else:
# For datasets that follow a similar structure, use a mapping.
dataset_mapping = {
vllm/benchmarks/serve.py
+95 -51
View File
@@ -251,6 +251,7 @@ async def get_request(
ramp_up_strategy: Literal["linear", "exponential"] | None = None,
ramp_up_start_rps: int | None = None,
ramp_up_end_rps: int | None = None,
self_timed: bool = False,
) -> AsyncGenerator[tuple[SampleRequest, float], None]:
"""
Asynchronously generates requests at a specified rate
@@ -290,49 +291,59 @@ async def get_request(
# Precompute delays among requests to minimize request send laggings
request_rates = []
delay_ts = []
for request_index, request in enumerate(input_requests):
current_request_rate = _get_current_request_rate(
ramp_up_strategy,
ramp_up_start_rps,
ramp_up_end_rps,
request_index,
total_requests,
request_rate,
)
assert current_request_rate > 0.0, (
f"Obtained non-positive request rate {current_request_rate}."
)
request_rates.append(current_request_rate)
if current_request_rate == float("inf"):
delay_ts.append(0)
elif burstiness == float("inf"):
# when burstiness tends to infinity, the delay time becomes constant
# and tends to the inverse of the request rate
delay_ts.append(1.0 / current_request_rate)
else:
theta = 1.0 / (current_request_rate * burstiness)
# Sample the request interval from the gamma distribution.
# If burstiness is 1, it follows exponential distribution.
delay_ts.append(np.random.gamma(shape=burstiness, scale=theta))
# if the traces have timing info then:
if not self_timed:
for request_index, request in enumerate(input_requests):
current_request_rate = _get_current_request_rate(
ramp_up_strategy,
ramp_up_start_rps,
ramp_up_end_rps,
request_index,
total_requests,
request_rate,
)
assert current_request_rate > 0.0, (
f"Obtained non-positive request rate {current_request_rate}."
)
request_rates.append(current_request_rate)
if current_request_rate == float("inf"):
delay_ts.append(0)
elif burstiness == float("inf"):
# when burstiness tends to infinity, the delay time becomes constant
# and tends to the inverse of the request rate
delay_ts.append(1.0 / current_request_rate)
else:
theta = 1.0 / (current_request_rate * burstiness)
# Calculate the cumulative delay time from the first sent out requests.
for i in range(1, len(delay_ts)):
delay_ts[i] += delay_ts[i - 1]
if ramp_up_strategy is None and delay_ts[-1] != 0:
# When ramp_up_strategy is not set, we assume the request rate is fixed
# and all requests should be sent in target_total_delay_s, the following
# logic would re-scale delay time to ensure the final delay_ts
# align with target_total_delay_s.
#
# NOTE: If we simply accumulate the random delta values
# from the gamma distribution, their sum would have 1-2% gap
# from target_total_delay_s. The purpose of the following logic is to
# close the gap for stabilizing the throughput data
# from different random seeds.
target_total_delay_s = total_requests / request_rate
normalize_factor = target_total_delay_s / delay_ts[-1]
delay_ts = [delay * normalize_factor for delay in delay_ts]
# Sample the request interval from the gamma distribution.
# If burstiness is 1, it follows exponential distribution.
delay_ts.append(np.random.gamma(shape=burstiness, scale=theta))
# Calculate the cumulative delay time from the first sent out requests.
for i in range(1, len(delay_ts)):
delay_ts[i] += delay_ts[i - 1]
if ramp_up_strategy is None and delay_ts[-1] != 0:
# When ramp_up_strategy is not set, we assume the request rate is fixed
# and all requests should be sent in target_total_delay_s, the following
# logic would re-scale delay time to ensure the final delay_ts
# align with target_total_delay_s.
#
# NOTE: If we simply accumulate the random delta values
# from the gamma distribution, their sum would have 1-2% gap
# from target_total_delay_s. The purpose of the following logic is to
# close the gap for stabilizing the throughput data
# from different random seeds.
target_total_delay_s = total_requests / request_rate
normalize_factor = target_total_delay_s / delay_ts[-1]
delay_ts = [delay * normalize_factor for delay in delay_ts]
else:
for request_index, request in enumerate(input_requests):
# this is cumulative running ts, from which sleep is calculated later
delay_ts.append(request.timestamp)
# TODO: there is no notion of RPS here, may be we can calculate
# from the trace.
request_rates.append(0.0)
start_ts = time.time()
for request_index, request in enumerate(input_requests):
@@ -634,6 +645,7 @@ async def benchmark(
ramp_up_end_rps: int | None = None,
ready_check_timeout_sec: int = 600,
ssl_context: ssl.SSLContext | bool | None = None,
self_timed: bool = False,
):
try:
request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
@@ -773,18 +785,20 @@ async def benchmark(
print("Profiler started")
distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"
if not self_timed:
if ramp_up_strategy is not None:
print(f"Traffic ramp-up strategy: {ramp_up_strategy}.")
print(
f"Will increase RPS from {ramp_up_start_rps} to "
f"{ramp_up_end_rps} RPS over the duration of the benchmark."
)
else:
print(f"Traffic request rate: {request_rate}")
if ramp_up_strategy is not None:
print(f"Traffic ramp-up strategy: {ramp_up_strategy}.")
print(
f"Will increase RPS from {ramp_up_start_rps} to "
f"{ramp_up_end_rps} RPS over the duration of the benchmark."
)
print(f"Burstiness factor: {burstiness} ({distribution})")
print(f"Maximum request concurrency: {max_concurrency}")
else:
print(f"Traffic request rate: {request_rate}")
print(f"Burstiness factor: {burstiness} ({distribution})")
print(f"Maximum request concurrency: {max_concurrency}")
print("Self timing is set, using the timestamps from the trace file.")
spec_decode_metrics_before = await fetch_spec_decode_metrics(base_url, session)
@@ -823,6 +837,7 @@ async def benchmark(
ramp_up_strategy,
ramp_up_start_rps,
ramp_up_end_rps,
self_timed,
):
if ramp_up_strategy is not None:
current_int_rps = int(current_request_rate)
@@ -1436,6 +1451,16 @@ def add_cli_args(parser: argparse.ArgumentParser):
help="Set ignore_eos flag when sending the benchmark request."
"Warning: ignore_eos is not supported in deepspeed_mii and tgi.",
)
parser.add_argument(
"--self-timed",
action=argparse.BooleanOptionalAction,
default=None,
help="Use timing information from the traces instead of the configuration. "
"This is useful when replaying traces faithfully based on their timestamps. "
"When unset, defaults to False, except for --dataset-name=timed_trace where "
"it defaults to True. Use --no-self-timed to force off. When off, user "
"defined generation rates are used and in trace timing info is ignored.",
)
parser.add_argument(
"--percentile-metrics",
type=str,
@@ -1767,6 +1792,24 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
):
args.ignore_eos = True
if args.dataset_name == "timed_trace":
# timed_trace carries per-request timestamps;
# ignore EOS so generation runs to the trace's specified output length,
# and default to using those timestamps for scheduling unless the user
# opted out.
args.ignore_eos = True
if args.self_timed is None:
args.self_timed = True
else:
# if this is set for anything else, it is an error
if args.self_timed is not None:
raise ValueError(
"--self-timed/--no-self-timed is only supported with "
"--dataset-name=timed_trace"
)
# for any non self-timed trace, this is False
args.self_timed = False
# Load the dataset.
input_requests = get_samples(args, tokenizer)
goodput_config_dict = check_goodput_args(args)
@@ -1846,6 +1889,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
ramp_up_end_rps=args.ramp_up_end_rps,
ready_check_timeout_sec=args.ready_check_timeout_sec,
ssl_context=ssl_context,
self_timed=args.self_timed,
)
# Save config and results to json