| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark hidden state extraction throughput. |
| 4 | |
| 5 | Measures two modes: |
| 6 | 1. Baseline: bulk inference with max_tokens=1, no extraction. |
| 7 | 2. Extract: async hidden state extraction via ExampleHiddenStatesConnector |
| 8 | with N concurrent clients, each consuming hidden states as |
| 9 | soon as their request finishes (overlapping I/O with generation). |
| 10 | |
| 11 | Reports tokens/s and prompts/s for each mode. |
| 12 | |
| 13 | Usage: |
| 14 | python benchmarks/benchmark_hidden_state_extraction.py \ |
| 15 | --model Qwen/Qwen3-0.6B \ |
| 16 | --num-prompts 64 \ |
| 17 | --num-clients 8 \ |
| 18 | --prompt-len 8192 \ |
| 19 | --layers 1 2 3 4 |
| 20 | """ |
| 21 | |
| 22 | import argparse |
| 23 | import asyncio |
| 24 | import time |
| 25 | from concurrent.futures import ThreadPoolExecutor |
| 26 | |
| 27 | import torch |
| 28 | from transformers import AutoConfig |
| 29 | |
| 30 | from vllm import LLM, SamplingParams |
| 31 | from vllm.config.kv_transfer import KVTransferConfig |
| 32 | from vllm.distributed.kv_transfer.kv_connector.v1 import ( |
| 33 | example_hidden_states_connector, |
| 34 | ) |
| 35 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 36 | from vllm.sampling_params import RequestOutputKind |
| 37 | from vllm.v1.engine.async_llm import AsyncLLM |
| 38 | |
| 39 | |
| 40 | def _make_profiler_config(profile_dir: str) -> dict: |
| 41 | """Build a profiler_config dict for torch profiling.""" |
| 42 | return { |
| 43 | "profiler": "torch", |
| 44 | "torch_profiler_dir": profile_dir, |
| 45 | "torch_profiler_with_stack": True, |
| 46 | } |
| 47 | |
| 48 | |
| 49 | def make_random_prompts( |
| 50 | num_prompts: int, prompt_len: int, vocab_size: int, seed: int = 42 |
| 51 | ) -> list[list[int]]: |
| 52 | """Generate lists of random token IDs.""" |
| 53 | # Set seed for reproducibility |
| 54 | torch.manual_seed(seed) |
| 55 | return [ |
| 56 | torch.randint(0, vocab_size, (prompt_len,)).tolist() for _ in range(num_prompts) |
| 57 | ] |
| 58 | |
| 59 | |
| 60 | def consume_hidden_states(path: str) -> float: |
| 61 | """Load hidden states from disk and compute per-position mean. |
| 62 | |
| 63 | Returns a single float: the grand mean of all hidden state values. |
| 64 | This forces the benchmark to actually read and reduce the data. |
| 65 | |
| 66 | Uses :func:`load_hidden_states` which acquires a shared flock, |
| 67 | blocking (without polling) until the async writer releases its |
| 68 | exclusive lock. |
| 69 | """ |
| 70 | obj = example_hidden_states_connector.load_hidden_states(path) |
| 71 | hs = obj["hidden_states"] |
| 72 | total = hs.mean().item() |
| 73 | |
| 74 | example_hidden_states_connector.cleanup_hidden_states(path) |
| 75 | |
| 76 | return total |
| 77 | |
| 78 | |
| 79 | def run_baseline( |
| 80 | model: str, |
| 81 | prompts: list[list[int]], |
| 82 | extra_args: dict, |
| 83 | profile_dir: str | None = None, |
| 84 | ) -> dict: |
| 85 | """Baseline: bulk inference, no hidden state extraction.""" |
| 86 | if profile_dir: |
| 87 | extra_args = { |
| 88 | **extra_args, |
| 89 | "profiler_config": _make_profiler_config(profile_dir), |
| 90 | } |
| 91 | llm = LLM( |
| 92 | model=model, |
| 93 | enable_prefix_caching=False, |
| 94 | **extra_args, |
| 95 | ) |
| 96 | sampling_params = SamplingParams(max_tokens=1) |
| 97 | prompt_inputs = [{"prompt_token_ids": p} for p in prompts] |
| 98 | |
| 99 | # Warmup |
| 100 | llm.generate(prompt_inputs[:4], sampling_params, use_tqdm=False) |
| 101 | |
| 102 | if profile_dir: |
| 103 | llm.start_profile() |
| 104 | |
| 105 | t0 = time.perf_counter() |
| 106 | outputs = llm.generate(prompt_inputs, sampling_params, use_tqdm=True) |
| 107 | elapsed = time.perf_counter() - t0 |
| 108 | |
| 109 | if profile_dir: |
| 110 | llm.stop_profile() |
| 111 | |
| 112 | total_prompt_tokens = sum(len(o.prompt_token_ids) for o in outputs) |
| 113 | num_prompts = len(outputs) |
| 114 | |
| 115 | del llm |
| 116 | torch.accelerator.empty_cache() |
| 117 | |
| 118 | return { |
| 119 | "mode": "baseline", |
| 120 | "elapsed_s": elapsed, |
| 121 | "num_prompts": num_prompts, |
| 122 | "total_prompt_tokens": total_prompt_tokens, |
| 123 | "tokens_per_s": total_prompt_tokens / elapsed, |
| 124 | "prompts_per_s": num_prompts / elapsed, |
| 125 | } |
| 126 | |
| 127 | |
| 128 | # ---- Async extraction benchmark ---- |
| 129 | |
| 130 | |
| 131 | async def _client_loop( |
| 132 | engine: AsyncLLM, |
| 133 | prompt_queue: asyncio.Queue, |
| 134 | consume_pool: ThreadPoolExecutor, |
| 135 | results: list[dict], |
| 136 | client_id: int, |
| 137 | ): |
| 138 | """A single async client: pulls prompts, submits to engine, consumes |
| 139 | hidden states as soon as each request finishes.""" |
| 140 | loop = asyncio.get_event_loop() |
| 141 | while True: |
| 142 | item = await prompt_queue.get() |
| 143 | if item is None: |
| 144 | prompt_queue.task_done() |
| 145 | break |
| 146 | idx, token_ids = item |
| 147 | |
| 148 | request_id = f"req-{idx}" |
| 149 | sampling_params = SamplingParams( |
| 150 | max_tokens=1, |
| 151 | output_kind=RequestOutputKind.FINAL_ONLY, |
| 152 | ) |
| 153 | |
| 154 | final_output = None |
| 155 | async for output in engine.generate( |
| 156 | request_id=request_id, |
| 157 | prompt={"prompt_token_ids": token_ids}, |
| 158 | sampling_params=sampling_params, |
| 159 | ): |
| 160 | if output.finished: |
| 161 | final_output = output |
| 162 | |
| 163 | # Consume hidden states on a thread (disk I/O) |
| 164 | path = final_output.kv_transfer_params["hidden_states_path"] |
| 165 | mean_val = await loop.run_in_executor(consume_pool, consume_hidden_states, path) |
| 166 | num_tokens = len(final_output.prompt_token_ids) |
| 167 | |
| 168 | results.append( |
| 169 | { |
| 170 | "request_id": request_id, |
| 171 | "num_prompt_tokens": num_tokens, |
| 172 | "mean_hidden_value": mean_val, |
| 173 | } |
| 174 | ) |
| 175 | prompt_queue.task_done() |
| 176 | |
| 177 | |
| 178 | async def _run_extraction_async( |
| 179 | model: str, |
| 180 | prompts: list[list[int]], |
| 181 | num_clients: int, |
| 182 | layers: list[int], |
| 183 | tmpdir: str, |
| 184 | extra_args: dict, |
| 185 | profile_dir: str | None = None, |
| 186 | ) -> dict: |
| 187 | if profile_dir: |
| 188 | extra_args = { |
| 189 | **extra_args, |
| 190 | "profiler_config": _make_profiler_config(profile_dir), |
| 191 | } |
| 192 | engine_args = AsyncEngineArgs( |
| 193 | model=model, |
| 194 | enable_prefix_caching=False, |
| 195 | max_num_batched_tokens=40960, |
| 196 | max_model_len=40960, |
| 197 | speculative_config={ |
| 198 | "method": "extract_hidden_states", |
| 199 | "num_speculative_tokens": 1, |
| 200 | "draft_model_config": { |
| 201 | "hf_config": { |
| 202 | "eagle_aux_hidden_state_layer_ids": layers, |
| 203 | }, |
| 204 | }, |
| 205 | }, |
| 206 | kv_transfer_config=KVTransferConfig( |
| 207 | kv_connector="ExampleHiddenStatesConnector", |
| 208 | kv_role="kv_producer", |
| 209 | kv_connector_extra_config={ |
| 210 | "shared_storage_path": tmpdir, |
| 211 | }, |
| 212 | ), |
| 213 | **extra_args, |
| 214 | ) |
| 215 | engine = AsyncLLM.from_engine_args(engine_args) |
| 216 | |
| 217 | try: |
| 218 | # Warmup: run a few prompts sequentially, cleaning up generated files |
| 219 | for i in range(min(4, len(prompts))): |
| 220 | sp = SamplingParams(max_tokens=1, output_kind=RequestOutputKind.FINAL_ONLY) |
| 221 | final_output = None |
| 222 | async for output in engine.generate( |
| 223 | request_id=f"warmup-{i}", |
| 224 | prompt={"prompt_token_ids": prompts[i]}, |
| 225 | sampling_params=sp, |
| 226 | ): |
| 227 | if output.finished: |
| 228 | final_output = output |
| 229 | if final_output and final_output.kv_transfer_params: |
| 230 | path = final_output.kv_transfer_params.get("hidden_states_path") |
| 231 | if path: |
| 232 | example_hidden_states_connector.cleanup_hidden_states(path) |
| 233 | |
| 234 | if profile_dir: |
| 235 | await engine.start_profile() |
| 236 | |
| 237 | # Fill prompt queue |
| 238 | prompt_queue: asyncio.Queue = asyncio.Queue() |
| 239 | for idx, token_ids in enumerate(prompts): |
| 240 | prompt_queue.put_nowait((idx, token_ids)) |
| 241 | # Sentinel per client |
| 242 | for _ in range(num_clients): |
| 243 | prompt_queue.put_nowait(None) |
| 244 | |
| 245 | results: list[dict] = [] |
| 246 | consume_pool = ThreadPoolExecutor(max_workers=num_clients) |
| 247 | |
| 248 | t0 = time.perf_counter() |
| 249 | tasks = [ |
| 250 | asyncio.create_task( |
| 251 | _client_loop(engine, prompt_queue, consume_pool, results, i) |
| 252 | ) |
| 253 | for i in range(num_clients) |
| 254 | ] |
| 255 | await asyncio.gather(*tasks) |
| 256 | elapsed = time.perf_counter() - t0 |
| 257 | |
| 258 | consume_pool.shutdown(wait=True) |
| 259 | |
| 260 | if profile_dir: |
| 261 | await engine.stop_profile() |
| 262 | |
| 263 | total_prompt_tokens = sum(r["num_prompt_tokens"] for r in results) |
| 264 | num_prompts = len(results) |
| 265 | mean_hidden = sum(r["mean_hidden_value"] for r in results) / max( |
| 266 | len(results), 1 |
| 267 | ) |
| 268 | |
| 269 | return { |
| 270 | "mode": "extract", |
| 271 | "elapsed_s": elapsed, |
| 272 | "num_prompts": num_prompts, |
| 273 | "total_prompt_tokens": total_prompt_tokens, |
| 274 | "tokens_per_s": total_prompt_tokens / elapsed, |
| 275 | "prompts_per_s": num_prompts / elapsed, |
| 276 | "mean_hidden_value": mean_hidden, |
| 277 | } |
| 278 | finally: |
| 279 | engine.shutdown() |
| 280 | |
| 281 | |
| 282 | def run_extraction( |
| 283 | model: str, |
| 284 | prompts: list[list[int]], |
| 285 | num_clients: int, |
| 286 | layers: list[int], |
| 287 | extra_args: dict, |
| 288 | profile_dir: str | None = None, |
| 289 | ) -> dict: |
| 290 | return asyncio.run( |
| 291 | _run_extraction_async( |
| 292 | model, |
| 293 | prompts, |
| 294 | num_clients, |
| 295 | layers, |
| 296 | "/dev/shm", |
| 297 | extra_args, |
| 298 | profile_dir=profile_dir, |
| 299 | ) |
| 300 | ) |
| 301 | |
| 302 | |
| 303 | def print_results(results: dict): |
| 304 | mode = results["mode"] |
| 305 | print(f"\n{'=' * 60}") |
| 306 | print(f" {mode.upper()} RESULTS") |
| 307 | print(f"{'=' * 60}") |
| 308 | print(f" Prompts: {results['num_prompts']}") |
| 309 | print(f" Total prompt tokens: {results['total_prompt_tokens']:,}") |
| 310 | print(f" Wall time: {results['elapsed_s']:.2f}s") |
| 311 | print(f" Tokens/s: {results['tokens_per_s']:,.0f}") |
| 312 | print(f" Prompts/s: {results['prompts_per_s']:.2f}") |
| 313 | if mode == "extract": |
| 314 | print(f" Mean hidden value: {results['mean_hidden_value']:.6f}") |
| 315 | print(f"{'=' * 60}\n") |
| 316 | |
| 317 | |
| 318 | def main(): |
| 319 | parser = argparse.ArgumentParser( |
| 320 | description="Benchmark hidden state extraction throughput" |
| 321 | ) |
| 322 | parser.add_argument("--model", type=str, required=True) |
| 323 | parser.add_argument("--num-prompts", type=int, default=64) |
| 324 | parser.add_argument("--num-clients", type=int, default=8) |
| 325 | parser.add_argument("--prompt-len", type=int, default=8192) |
| 326 | parser.add_argument("--layers", type=int, nargs="+", default=[1, 2, 3, 4]) |
| 327 | parser.add_argument("--skip-baseline", action="store_true") |
| 328 | parser.add_argument("--skip-extract", action="store_true") |
| 329 | parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) |
| 330 | parser.add_argument("--max-num-batched-tokens", type=int, default=None) |
| 331 | parser.add_argument("--max-cudagraph-capture-size", type=int, default=None) |
| 332 | parser.add_argument("--max-model-len", type=int, default=None) |
| 333 | parser.add_argument("--enforce-eager", action="store_true") |
| 334 | parser.add_argument("--load-format", type=str, default=None) |
| 335 | parser.add_argument( |
| 336 | "--profile", |
| 337 | action="store_true", |
| 338 | help="Enable torch profiler for both baseline and extraction runs.", |
| 339 | ) |
| 340 | parser.add_argument( |
| 341 | "--torch-profiler-dir", |
| 342 | type=str, |
| 343 | default="./vllm_profile", |
| 344 | help="Directory to save torch profiler traces (default: ./vllm_profile).", |
| 345 | ) |
| 346 | parser.add_argument( |
| 347 | "--enable-flashinfer-autotune", |
| 348 | action="store_true", |
| 349 | default=False, |
| 350 | help="Enable FlashInfer autotuning (can be slow).", |
| 351 | ) |
| 352 | args = parser.parse_args() |
| 353 | |
| 354 | extra_args = { |
| 355 | "gpu_memory_utilization": args.gpu_memory_utilization, |
| 356 | } |
| 357 | if args.max_model_len is not None: |
| 358 | extra_args["max_model_len"] = args.max_model_len |
| 359 | if args.max_num_batched_tokens is not None: |
| 360 | extra_args["max_num_batched_tokens"] = args.max_num_batched_tokens |
| 361 | if args.max_model_len and args.max_num_batched_tokens < args.max_model_len: |
| 362 | raise ValueError( |
| 363 | "max_num_batched_tokens must be >= max_model_len since chunked prefill" |
| 364 | " is not supported by hidden state extraction." |
| 365 | ) |
| 366 | if args.enforce_eager: |
| 367 | extra_args["enforce_eager"] = True |
| 368 | if args.load_format is not None: |
| 369 | extra_args["load_format"] = args.load_format |
| 370 | if args.max_cudagraph_capture_size is not None: |
| 371 | extra_args["max_cudagraph_capture_size"] = args.max_cudagraph_capture_size |
| 372 | extra_args["enable_flashinfer_autotune"] = args.enable_flashinfer_autotune |
| 373 | |
| 374 | # Get vocab size from HF config without loading the full model |
| 375 | hf_config = AutoConfig.from_pretrained(args.model, trust_remote_code=True) |
| 376 | vocab_size = hf_config.vocab_size |
| 377 | prompts = make_random_prompts(args.num_prompts, args.prompt_len, vocab_size) |
| 378 | print( |
| 379 | f"Generated {args.num_prompts} prompts, " |
| 380 | f"{args.prompt_len} tokens each (vocab {vocab_size})" |
| 381 | ) |
| 382 | |
| 383 | profile_dir = args.torch_profiler_dir if args.profile else None |
| 384 | if profile_dir: |
| 385 | print(f"Torch profiler enabled, traces will be saved to {profile_dir}/") |
| 386 | |
| 387 | if not args.skip_baseline: |
| 388 | baseline_profile_dir = f"{profile_dir}/baseline" if profile_dir else None |
| 389 | baseline = run_baseline( |
| 390 | args.model, prompts, extra_args, profile_dir=baseline_profile_dir |
| 391 | ) |
| 392 | print_results(baseline) |
| 393 | |
| 394 | if not args.skip_extract: |
| 395 | extract_profile_dir = f"{profile_dir}/extract" if profile_dir else None |
| 396 | extract = run_extraction( |
| 397 | args.model, |
| 398 | prompts, |
| 399 | args.num_clients, |
| 400 | args.layers, |
| 401 | extra_args, |
| 402 | profile_dir=extract_profile_dir, |
| 403 | ) |
| 404 | print_results(extract) |
| 405 | |
| 406 | if not args.skip_baseline and not args.skip_extract: |
| 407 | slowdown = baseline["tokens_per_s"] / extract["tokens_per_s"] |
| 408 | print("Extraction slowdown factor: {:.2f}x".format(slowdown)) |
| 409 | |
| 410 | |
| 411 | if __name__ == "__main__": |
| 412 | main() |