From 9d3dfa2886ae0b3b248f563c678ff3616658c5f6 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 14 Oct 2026 10:52:26 -0700 Subject: [PATCH] [Benchmark] Add vllm bench forward: exact-shape forward timing (#61339) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 --- .agents/skills/model-microbenchmark/SKILL.md | 73 ++++++ .claude/skills/model-microbenchmark | 1 + docs/cli/README.md | 15 +- docs/mkdocs/gen_files/generate_argparse.py | 2 + tests/benchmarks/test_forward_cli.py | 127 ++++++++++ vllm/benchmarks/forward/__init__.py | 2 + vllm/benchmarks/forward/batch.py | 148 +++++++++++ vllm/benchmarks/forward/capture.py | 246 ++++++++++++++++++ vllm/benchmarks/forward/kv_init.py | 70 ++++++ vllm/benchmarks/forward/main.py | 250 +++++++++++++++++++ vllm/benchmarks/forward/routing.py | 227 +++++++++++++++++ vllm/benchmarks/forward/shapes.py | 137 ++++++++++ vllm/entrypoints/cli/benchmark/forward.py | 22 ++ vllm/entrypoints/cli/benchmark/main.py | 1 + 14 files changed, 1320 insertions(+), 1 deletion(-) create mode 100644 .agents/skills/model-microbenchmark/SKILL.md create mode 120000 .claude/skills/model-microbenchmark create mode 100644 tests/benchmarks/test_forward_cli.py create mode 100644 vllm/benchmarks/forward/__init__.py create mode 100644 vllm/benchmarks/forward/batch.py create mode 100644 vllm/benchmarks/forward/capture.py create mode 100644 vllm/benchmarks/forward/kv_init.py create mode 100644 vllm/benchmarks/forward/main.py create mode 100644 vllm/benchmarks/forward/routing.py create mode 100644 vllm/benchmarks/forward/shapes.py create mode 100644 vllm/entrypoints/cli/benchmark/forward.py diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md new file mode 100644 index 0000000000..250059cd0a --- /dev/null +++ b/.agents/skills/model-microbenchmark/SKILL.md @@ -0,0 +1,73 @@ +--- +name: model-microbenchmark +description: Time a vLLM model's forward pass on exact batch shapes (decode, prefill, chunked prefill, prefix hit, spec-decode verify, or mixes) with `vllm bench forward`, single-GPU or DP/TP/EP across nodes. Use for per-step model latency, layout/parallelism comparisons, and checking a model-level optimization. Not for single kernels (use kernel-microbenchmark) or end-to-end serving (use `vllm bench serve`). +--- + +# Model Microbenchmark (`vllm bench forward`) + +Builds the batch without the scheduler, runs it through MRV2's real +`execute_model`, and times the forward as a one-shot CUDA graph (prefill +included). Sampling, the speculator, and CPU step overhead are excluded, so +expect results somewhat below a serving step. + +## Usage + +```bash +vllm bench forward --model deepseek-ai/DeepSeek-V4-Flash \ + --case decode=64*1@8192 --case prefill=8192 --case mix=4096,32*1@8192 + +# DP/TP/EP: one process per GPU (multi-node: add --nnodes/--node-rank/--master-addr) +torchrun --nproc-per-node=4 -m vllm.entrypoints.cli.main bench forward \ + --model ... -dp 4 --enable-expert-parallel --case decode=128*1@100000 +``` + +Requests: `[N*]Q[@K][c][:D]`, comma-separated, per DP rank. + +| Form | Meaning | +| --- | --- | +| `8192` | fresh prefill, 8192 tokens | +| `32*1@8192` | 32 decodes, KV length 8192 (KV length **includes** the query) | +| `4@4096:3` | spec-decode verify, 3 draft tokens | +| `2048@10240` | prefill after an 8192-token prefix-cache hit | +| `2048@10240c` | chunked prefill: the prefix came from earlier chunks | + +Hit vs chunk only differs on models with SWA bounded replay (DSV4.1): a hit +also recomputes the last window (128 tokens) of the prefix. + +Or `--shapes file.json|yaml`: +`{"cases": [{"name": ..., "requests": [{"query_len", "kv_len", "num_draft_tokens", "chunked", "count"}]}]}`. + +Key flags (plus all `EngineArgs`): + +- `--kv-len-jitter 0.5` (default): prefixes drawn within ±50%, same on every + DP rank, capped at `max_model_len`. Use `0` for exact lengths. +- `--routing auto|real|balanced`: `auto` = balanced with dummy weights. +- `--iters 20 --samples 5`, `--output-json out.json` (per-rank latency stats, + `samples_us`, graph-vs-eager check). +- Defaults: `--load-format dummy`, prefix caching off, no torch.compile. + +Read `median_us` (each replay is already the max over ranks). + +## Caveats + +- **Dummy weights lie for data-dependent kernels.** Dummy gates route every + token to the same experts (hence balanced routing). Dummy norm weights make + DSV4 indexer logits degenerate, so `topKPerRowPrefill` runs ~2.6x slow + until #61323 lands. Cross-check suspicious kernels with real weights. +- **torch.compile models (e.g. Qwen3/3.5) look slow**: no compile here. Use + them for correctness only; trust perf for models under `vllm/models/`. +- `sample spread > 0.5%` warns of unstable timing. DSV4.1 DEP prefill is + inherently ~0.5–1% (MegaMoE cross-rank sync, host engram reads). +- `graph_vs_eager` may be nonzero for small batches where models branch on + `is_current_stream_capturing()`; that's expected, not a bug. +- KV cache is zero-initialized (`--kv-init zeros`); random KV was within 3%. +- Uniform DP only: every rank gets the same case; no idle ranks. +- CUDA only: it hangs on ROCm (untested cause: likely the in-graph timing + events). +- Large prefills/prefix hits need room: raise `--max-num-batched-tokens` / + `--max-model-len` / `--gpu-memory-utilization` if build fails on KV blocks. +- Startup cost is JIT/autotune, not the bench: install FlashInfer jit-cache + and cubin; DeepGEMM warmup is skipped (each case's eager pass compiles its + own shapes). +- To compare against serving decode: `vllm bench latency` with 65- vs + 1-token outputs, divided by 64, with a FULL decode graph. diff --git a/.claude/skills/model-microbenchmark b/.claude/skills/model-microbenchmark new file mode 120000 index 0000000000..af14d2467d --- /dev/null +++ b/.claude/skills/model-microbenchmark @@ -0,0 +1 @@ +../../.agents/skills/model-microbenchmark \ No newline at end of file diff --git a/docs/cli/README.md b/docs/cli/README.md index a1eff4c8ae..c2b27bed26 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -131,9 +131,22 @@ To use benchmark commands, please install with extra dependencies using `pip ins Available Commands: ```bash -vllm bench {latency, serve, throughput} +vllm bench {forward, latency, serve, throughput} ``` +### forward + +Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; sampling and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP. CUDA only for now. + +```bash +vllm bench forward \ + --model meta-llama/Llama-3.2-1B-Instruct \ + --case decode=32*1@4096 \ + --case mix=2048,32*1@4096 +``` + +See [vllm bench forward](./bench/forward.md) for the full reference of all available arguments. + ### latency Benchmark the latency of a single batch of requests. diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py index ec9511efad..8e71d13ae0 100644 --- a/docs/mkdocs/gen_files/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -145,6 +145,7 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): ) +bench_forward = auto_mock("vllm.benchmarks.forward", "main") bench_latency = auto_mock("vllm.benchmarks", "latency") bench_serve = auto_mock("vllm.benchmarks", "serve") bench_startup = auto_mock("vllm.benchmarks", "startup") @@ -340,6 +341,7 @@ pages = { False, ), "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True), + "cli/bench/forward.md": (create_parser(bench_forward.add_cli_args), True), "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True), # URL kept as `mm_processor` for back-compat; command name is `mm-processor` "cli/bench/mm_processor.md": ( diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py new file mode 100644 index 0000000000..5d1286243a --- /dev/null +++ b/tests/benchmarks/test_forward_cli.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import subprocess +from dataclasses import replace +from types import SimpleNamespace + +import numpy as np +import pytest + +from vllm.benchmarks.forward.batch import build_batch +from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case +from vllm.platforms import current_platform + +MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" + + +def _group(block_size: int, replay: int = 0, dcp_sharded: bool = False): + return SimpleNamespace( + kv_cache_spec=SimpleNamespace( + block_size=block_size, + prefix_replay_tokens=replay, + dcp_sharded=dcp_sharded, + ) + ) + + +def test_build_batch_layout(): + """The runner sees each request kind as the scheduler would schedule it.""" + case = parse_case("c=2*1@101,4@64:3,32,8@48c,16@316") + groups = [_group(16), _group(32, replay=128)] + batch = build_batch(case, groups, 1000, 1000, np.random.default_rng(0)) + + # (scheduled, computed, is_decode) per request: decode, verify, prefill, + # chunked prefill, then a prefix hit replaying the last 128 hit tokens. + assert list(batch.expected.values()) == [ + (1, 100, True), + (1, 100, True), + (4, 60, True), + (32, 0, False), + (8, 40, False), + (16 + 128, 300 - 128, False), + ] + output = batch.scheduler_output + assert output.total_num_scheduled_tokens == 1 + 1 + 4 + 32 + 8 + 144 + assert len(output.scheduled_spec_decode_tokens["bench-2"]) == 3 + assert output.scheduled_new_reqs[-1].replay_start == 300 - 128 + for req_id, (scheduled, _, _) in batch.expected.items(): + assert len(batch.query_token_ids[req_id]) == scheduled + # Groups share one block pool: block IDs must not overlap. + ranges = sorted(r for group in batch.group_block_ranges for r in group) + assert ranges[0][0] == 1 + assert all(a[1] == b[0] for a, b in zip(ranges, ranges[1:])) + + with pytest.raises(ValueError, match="KV blocks"): + build_batch(case, groups, 10, 1000, np.random.default_rng(0)) + + # DCP spreads each block's tokens over the DCP ranks, like the scheduler. + groups = [_group(16, dcp_sharded=True), _group(16)] + batch = build_batch( + parse_case("c=1@640"), groups, 1000, 1000, np.random.default_rng(0), 4 + ) + assert batch.group_block_ranges == [[(1, 11)], [(11, 51)]] + + +def test_jitter_kv_lens(): + """Prefixes vary, capped at the max length; queries and fresh prefills don't. + A case's draw depends only on the seed and its name, so adding or reordering + other cases doesn't change it.""" + case = parse_case("c=1000*1@1001,4@64:3,32") + jittered = jitter_kv_lens(case, 0.5, 1400, 0) + assert jitter_kv_lens(case, 0.5, 1400, 0) == jittered + other = jitter_kv_lens(replace(case, name="d"), 0.5, 1400, 0) + assert other.requests != jittered.requests + prefixes = np.array([r.prefix_len for r in jittered.requests[:1000]]) + assert prefixes.min() >= 500 and max(r.kv_len for r in jittered.requests) == 1400 + assert prefixes.std() > 200 + assert [r.query_len for r in jittered.requests] == [ + r.query_len for r in case.requests + ] + assert jittered.requests[-1] == case.requests[-1] + assert jitter_kv_lens(case, 0, 2048, 0) == case + with pytest.raises(ValueError, match="KV length"): + jitter_kv_lens(case, 0, 1000, 0) + + +@pytest.mark.benchmark +@pytest.mark.skipif( + current_platform.is_rocm(), reason="Hangs on ROCm; the bench is CUDA-only." +) +def test_bench_forward(tmp_path): + output = tmp_path / "forward.json" + command = [ + "vllm", + "bench", + "forward", + "--model", + MODEL_NAME, + "--max-model-len", + "2048", + "--gpu-memory-utilization", + "0.5", + "--case", + "decode=8*1@512", + "--case", + "prefill=256", + "--case", + "mix=128@256,4*1@512", + "--iters", + "3", + "--samples", + "2", + "--output-json", + str(output), + ] + result = subprocess.run(command, capture_output=True, text=True) + print(result.stdout) + print(result.stderr) + assert result.returncode == 0, f"Benchmark failed: {result.stderr}" + + results = {row["case"]: row for row in json.loads(output.read_text())} + assert list(results) == ["decode", "prefill", "mix"] + for row in results.values(): + rank0 = row["ranks"][0] + assert len(rank0["batch_order"]) == len(row["requests"]) + assert rank0["finite"] and rank0["graph_deterministic"] + assert row["median_us"] > 0 diff --git a/vllm/benchmarks/forward/__init__.py b/vllm/benchmarks/forward/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/benchmarks/forward/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/benchmarks/forward/batch.py b/vllm/benchmarks/forward/batch.py new file mode 100644 index 0000000000..61ebf0d8db --- /dev/null +++ b/vllm/benchmarks/forward/batch.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build a `SchedulerOutput` for an exact batch, without the scheduler.""" + +from dataclasses import dataclass + +import numpy as np + +from vllm.benchmarks.forward.shapes import BenchCase +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheSpec, + MambaSpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, +) + + +@dataclass +class BenchBatch: + scheduler_output: SchedulerOutput + # req_id -> token IDs of the query tokens, in order. + query_token_ids: dict[str, np.ndarray] + # Per KV cache group: allocated (begin, end) block ranges. + group_block_ranges: list[list[tuple[int, int]]] + num_blocks: int + # req_id -> the runner's expected (scheduled, computed, is_decode). + expected: dict[str, tuple[int, int, bool]] + + @property + def num_tokens(self) -> int: + return self.scheduler_output.total_num_scheduled_tokens + + @property + def req_ids(self) -> list[str]: + return list(self.scheduler_output.num_scheduled_tokens) + + def finish(self) -> SchedulerOutput: + """A step that only removes this batch's requests.""" + output = SchedulerOutput.make_empty() + output.finished_req_ids = set(self.req_ids) + return output + + +def kv_block_size(spec: KVCacheSpec, dcp_size: int) -> int: + """Tokens per block as the scheduler counts them: DCP spreads each + block's tokens over the DCP ranks.""" + return spec.block_size * (dcp_size if spec.dcp_sharded else 1) + + +def _num_blocks( + spec: KVCacheSpec, computed: int, kv_len: int, dcp_size: int +) -> tuple[int, int]: + """Return (leading null blocks, live blocks) for one request.""" + block_size = kv_block_size(spec, dcp_size) + if isinstance(spec, UniformTypeKVCacheSpecs): + spec = spec.first_spec + if isinstance(spec, (CircularBufferSpec, MambaSpec)): + return 0, 1 + total = -(-kv_len // block_size) + skipped = 0 + if isinstance(spec, SlidingWindowSpec): + # Blocks wholly before the window are freed to the null block. + skipped = max(0, computed - spec.sliding_window + 1) // block_size + return skipped, total - skipped + + +def build_batch( + case: BenchCase, + kv_cache_groups, + num_kv_blocks: int, + vocab_size: int, + rng: np.random.Generator, + dcp_size: int = 1, +) -> BenchBatch: + """Allocate disjoint blocks and build `NewRequestData` for every request. + + Groups share one block pool, so block IDs are unique across groups. Block + 0 is the null block. + + With prefix P = K - Q (KV length K, query length Q): + + | Kind | prompt tokens | num_computed_tokens | spec tokens | + | ------------------ | ------------- | ------------------- | ----------- | + | Prefill | K | P | - | + | Decode | P | P | - | + | Verify (D drafts) | P | P | D | + + Unless chunked, a prefill's prefix is a prefix hit, which replays its last + R = min(window, P) tokens on models with SWA bounded replay (window 0 + otherwise), like the scheduler: it schedules R + Q tokens from + replay_start = num_computed_tokens = P - R. + """ + replay_window = max( + (group.kv_cache_spec.prefix_replay_tokens for group in kv_cache_groups), + default=0, + ) + expected: dict[str, tuple[int, int, bool]] = {} + cursor = 1 + new_reqs = [] + query_token_ids = {} + group_block_ranges: list[list[tuple[int, int]]] = [[] for _ in kv_cache_groups] + output = SchedulerOutput.make_empty() + for i, req in enumerate(case.requests): + req_id = f"bench-{i}" + replay = min(replay_window, req.prefix_len) if req.prefix_hit else 0 + computed = req.prefix_len - replay + block_ids = [] + for group, ranges in zip(kv_cache_groups, group_block_ranges): + skipped, live = _num_blocks( + group.kv_cache_spec, computed, req.kv_len, dcp_size + ) + block_ids.append([0] * skipped + list(range(cursor, cursor + live))) + ranges.append((cursor, cursor + live)) + cursor += live + tokens = rng.integers(0, vocab_size, req.kv_len, dtype=np.int64) + prompt = tokens[: req.prefix_len] if req.is_decode else tokens + query_token_ids[req_id] = tokens[computed:] + expected[req_id] = (req.query_len + replay, computed, req.is_decode) + new_reqs.append( + NewRequestData( + req_id=req_id, + prompt_token_ids=prompt.tolist(), + prefill_token_ids=prompt.tolist(), + mm_features=[], + sampling_params=SamplingParams(max_tokens=16, temperature=0.0), + pooling_params=None, + block_ids=tuple(block_ids), + num_computed_tokens=computed, + lora_request=None, + replay_start=computed if replay else 0, + ) + ) + output.num_scheduled_tokens[req_id] = req.query_len + replay + if req.num_draft_tokens: + output.scheduled_spec_decode_tokens[req_id] = tokens[ + req.prefix_len + 1 : + ].tolist() + if cursor > num_kv_blocks: + raise ValueError( + f"Case {case.name!r} needs {cursor} KV blocks, but only " + f"{num_kv_blocks} exist. Increase the KV cache size." + ) + output.scheduled_new_reqs = new_reqs + output.total_num_scheduled_tokens = sum(output.num_scheduled_tokens.values()) + return BenchBatch(output, query_token_ids, group_block_ranges, cursor, expected) diff --git a/vllm/benchmarks/forward/capture.py b/vllm/benchmarks/forward/capture.py new file mode 100644 index 0000000000..0889026430 --- /dev/null +++ b/vllm/benchmarks/forward/capture.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Capture the forward of one exact batch as a one-shot CUDA graph and time it. + +The graph is captured for a single batch and only replays that batch, so the +host-side decisions baked in at capture (attention plans, kernel choices by +query length, chunk splits) are correct by construction. This deliberately +bypasses vLLM's CUDA graph manager: no capture sizes, padding or dummy +metadata. Each replay first restores the KV blocks the step writes, so state the +forward mutates in place (e.g. compressor state) is the same on every replay. +""" + +import time +from typing import Any + +import numpy as np +import torch + +from vllm.benchmarks.forward.batch import BenchBatch +from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.parallel_state import get_world_group, graph_capture +from vllm.forward_context import get_forward_context +from vllm.v1.worker.gpu.input_batch import InputBatch + + +def _first_tensor(output: Any) -> torch.Tensor: + if isinstance(output, tuple): + output = output[0] + assert isinstance(output, torch.Tensor), type(output) + return output + + +def _compare(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + delta = actual.float() - expected.float() + scale = expected.float().square().mean().sqrt().clamp(min=1e-12) + return { + "max_abs": delta.abs().max().item(), + "relative_rms": (delta.square().mean().sqrt() / scale).item(), + } + + +class ForwardBench: + """Runs one batch through the runner's `execute_model`, timing its forward. + + `execute_model` prepares the inputs, attention metadata and forward context + as in serving; the bench takes over the eager model call. Two steps: the + first ("warmup") runs the model eagerly, which triggers JIT, autotuning and + lazy init. The second ("capture") captures, validates and times the graph. + """ + + def __init__( + self, + runner: Any, + batch: BenchBatch, + iters: int, + samples: int, + state_views: list[torch.Tensor], + ): + self.runner = runner + self.batch = batch + self.device = runner.device + self.iters = iters + self.samples = samples + self.state_views = state_views + self.phase = "warmup" + self.result: dict[str, Any] = {} + + def run(self) -> dict[str, Any]: + runner, model = self.runner, self.runner.model + # The runner's eager branch calls `model(**inputs)`. + assert runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE + assert runner.ubatch_runner is None + prepare_inputs = runner.prepare_inputs + self._model_forward = model.forward + + def prepare_and_check(*args, **kwargs) -> InputBatch: + self._input_batch = prepare_inputs(*args, **kwargs) + self._check_batch(self._input_batch) + return self._input_batch + + # Instance attributes shadow the methods for these two steps only. + runner.prepare_inputs = prepare_and_check + model.forward = self._forward + try: + for self.phase in ("warmup", "capture"): + runner.execute_model(self.batch.scheduler_output) + finally: + del runner.prepare_inputs, model.forward + # No sample_tokens follows. + runner.execute_model_state = None + return self.result + + def _forward(self, **inputs) -> Any: + self._write_input_ids(self._input_batch, inputs["input_ids"]) + # Models may update the forward context in place (e.g. sequence + # parallelism shards is_padding); every call must start from the + # context the runner built. + context = get_forward_context() + saved = dict(vars(context)) + + def run_model() -> Any: + for name, value in saved.items(): + setattr(context, name, value) + return self._model_forward(**inputs) + + self._run_model = run_model + if self.phase == "warmup": + return run_model() + return self._capture_and_time() + + def _check_batch(self, input_batch: InputBatch) -> None: + """Fail unless the runner built exactly the requested batch.""" + expected = self.batch.expected + assert sorted(input_batch.req_ids) == sorted(expected), input_batch.req_ids + for i, req_id in enumerate(input_batch.req_ids): + actual = ( + int(input_batch.num_scheduled_tokens[i]), + int(input_batch.num_computed_tokens_np[i]), + not bool(input_batch.is_prefilling_np[i]), + ) + assert actual == expected[req_id], ( + f"{req_id}: got (scheduled, computed, decode) = {actual}, " + f"wanted {expected[req_id]}" + ) + assert input_batch.num_tokens == self.batch.num_tokens + self.result["num_tokens_after_padding"] = input_batch.num_tokens_after_padding + self.result["batch_order"] = input_batch.req_ids + + def _write_input_ids( + self, input_batch: InputBatch, input_ids: torch.Tensor | None + ) -> None: + """Decode/verify inputs otherwise come from sampler state.""" + if input_ids is None: + return + ids = np.concatenate( + [self.batch.query_token_ids[req_id] for req_id in input_batch.req_ids] + ) + input_ids[: len(ids)].copy_(torch.from_numpy(ids)) + + def _capture_and_time(self) -> Any: + world = get_world_group() + # The forward mutates state it reads (e.g. compressor state), so + # every run starts by restoring the KV blocks the step writes. + snapshots = [(view, view.clone()) for view in self.state_views] + + def restore() -> None: + for view, saved in snapshots: + view.copy_(saved) + + # A rendezvous and the timing events live inside the graph, after the + # restore: neither rank skew, the restore nor host launch gaps reach + # the timed interval. + flag = torch.zeros(1, device=self.device) + self._events = ( + torch.cuda.Event(enable_timing=True, external=True), + torch.cuda.Event(enable_timing=True, external=True), + ) + if world.world_size > 1: + # Initialize the communicator outside capture. + torch.distributed.all_reduce(flag, group=world.device_group) + with graph_capture(self.device) as capture_context: + stream = capture_context.stream + # Some kernels set up per-stream state on first use, outside capture. + eager = _first_tensor(self._run_model()).clone() + restore() + reference = _first_tensor(self._run_model()).clone() + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + restore() + if world.world_size > 1: + torch.distributed.all_reduce(flag, group=world.device_group) + self._events[0].record() + output = self._run_model() + self._events[1].record() + hidden = _first_tensor(output) + + graph.replay() + torch.accelerator.synchronize() + first = hidden.clone() + graph.replay() + torch.accelerator.synchronize() + self.result.update( + # Same step from the same restored state: should match exactly. + eager_repeat=_compare(reference, eager), + graph_vs_eager=_compare(first, reference), + graph_deterministic=bool(torch.equal(first, hidden)), + finite=bool(hidden.isfinite().all()), + state_reset_bytes=sum( + v.numel() * v.element_size() for v in self.state_views + ), + ) + self.result["latency_us"] = self._time(graph) + del graph + return output + + def _time(self, graph: torch.cuda.CUDAGraph) -> dict[str, Any]: + """Per-replay forward time between in-graph events, MAX across ranks.""" + world = get_world_group() + start, end = self._events + self._warmup(graph) + times = [] + for _ in range(self.iters * self.samples): + graph.replay() + end.synchronize() + times.append(start.elapsed_time(end) * 1000) + arr = torch.tensor(times, device=self.device) + if world.world_size > 1: + torch.distributed.all_reduce( + arr, op=torch.distributed.ReduceOp.MAX, group=world.device_group + ) + arr = arr.cpu().numpy() + sample_medians = np.median(arr.reshape(self.samples, -1), axis=1) + median = float(np.median(arr)) + return { + "timer": "graph_events" + ("_rank_max" if world.world_size > 1 else ""), + "median": median, + # Stability: spread of per-sample medians relative to the median. + "sample_spread_pct": float( + 100 * (sample_medians.max() - sample_medians.min()) / median + ), + "mean": float(arr.mean()), + "std": float(arr.std()), + "min": float(arr.min()), + "p90": float(np.percentile(arr, 90)), + "count": len(arr), + "samples_us": [round(float(t), 2) for t in arr], + } + + def _warmup(self, graph: torch.cuda.CUDAGraph, seconds: float = 1.0) -> None: + """Replay until clocks and caches reach a steady state.""" + world = get_world_group() + graph.replay() + torch.accelerator.synchronize() + begin = time.perf_counter() + graph.replay() + torch.accelerator.synchronize() + per_replay = max(time.perf_counter() - begin, 1e-6) + # Same count on every rank: the graph may contain collectives. + count = torch.tensor([max(self.iters, int(seconds / per_replay))]) + torch.distributed.all_reduce( + count, op=torch.distributed.ReduceOp.MAX, group=world.cpu_group + ) + for _ in range(int(count.item())): + graph.replay() + torch.accelerator.synchronize() diff --git a/vllm/benchmarks/forward/kv_init.py b/vllm/benchmarks/forward/kv_init.py new file mode 100644 index 0000000000..03700f9ab3 --- /dev/null +++ b/vllm/benchmarks/forward/kv_init.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KV blocks of a benchmark case: initialization and the blocks a step writes. + +Only the case's blocks are touched, group by group: groups can alias one +allocation with different layouts, and the batch builder keeps their block IDs +disjoint. Zero KV times within a few percent of random KV, even for sparse +attention (DeepSeek-V4.1, ~100K context), so zeros are the only init for now. +""" + +from collections.abc import Iterator + +import torch + +from vllm.benchmarks.forward.batch import BenchBatch, kv_block_size + + +def _tensors(kv_cache) -> Iterator[torch.Tensor]: + if isinstance(kv_cache, torch.Tensor): + if kv_cache.numel(): + yield kv_cache + elif isinstance(kv_cache, (list, tuple)): + for item in kv_cache: + yield from _tensors(item) + + +def _block_dim(tensor: torch.Tensor, num_blocks: int) -> tuple[int, int]: + """Return (block dim, kernel blocks per allocator block).""" + for dim in range(min(2, tensor.dim())): + if tensor.shape[dim] >= num_blocks and tensor.shape[dim] % num_blocks == 0: + return dim, tensor.shape[dim] // num_blocks + raise ValueError(f"No block dim of size {num_blocks} in {tuple(tensor.shape)}") + + +@torch.inference_mode() +def zero_kv_blocks( + static_forward_context: dict, kv_cache_groups, batch: BenchBatch, num_blocks: int +) -> None: + for group, ranges in zip(kv_cache_groups, batch.group_block_ranges): + for name in group.layer_names: + layer = static_forward_context[name] + for tensor in _tensors(getattr(layer, "kv_cache", None)): + dim, factor = _block_dim(tensor, num_blocks) + for begin, end in ranges: + tensor.narrow(dim, begin * factor, (end - begin) * factor).zero_() + + +def touched_kv_views( + static_forward_context: dict, + kv_cache_groups, + batch: BenchBatch, + num_blocks: int, + dcp_size: int, +) -> list[torch.Tensor]: + """Views of the KV blocks a step may write: each request's last + ceil(scheduled / block_size) + 2 blocks, in every group.""" + scheduled_tokens = [batch.expected[req_id][0] for req_id in batch.req_ids] + views = [] + for group, ranges in zip(kv_cache_groups, batch.group_block_ranges): + block_size = kv_block_size(group.kv_cache_spec, dcp_size) + for name in group.layer_names: + layer = static_forward_context[name] + for tensor in _tensors(getattr(layer, "kv_cache", None)): + dim, factor = _block_dim(tensor, num_blocks) + for (begin, end), tokens in zip(ranges, scheduled_tokens): + first = max(begin, end - (-(-tokens // block_size) + 2)) + views.append( + tensor.narrow(dim, first * factor, (end - first) * factor) + ) + return views diff --git a/vllm/benchmarks/forward/main.py b/vllm/benchmarks/forward/main.py new file mode 100644 index 0000000000..3ea40ea6ed --- /dev/null +++ b/vllm/benchmarks/forward/main.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the model forward of exact batches as one-shot CUDA graphs. + +Each case is a list of requests `(query_len, kv_len, num_draft_tokens)` +per DP rank. The bench builds the batch without the scheduler, runs it through +the real `execute_model` (input prep, attention metadata, DP sync), and times +the target-model forward captured as a CUDA graph. Sampling, the speculator +and CPU step overhead aren't measured. + +Single GPU: `vllm bench forward --model M --case decode=32*1@8192`. +Multi-GPU: `torchrun --nproc-per-node=N -m vllm.entrypoints.cli.main bench +forward ...`, one process per GPU. +""" + +import argparse +import json +import os +import socket +from dataclasses import asdict +from typing import Any + +import numpy as np +import torch + +from vllm.benchmarks.forward.shapes import ( + BenchCase, + jitter_kv_lens, + load_cases, + parse_case, +) +from vllm.engine.arg_utils import EngineArgs +from vllm.logger import init_logger +from vllm.utils.argparse_utils import FlexibleArgumentParser + +logger = init_logger(__name__) + + +def add_cli_args(parser: FlexibleArgumentParser) -> None: + parser.add_argument( + "--case", + action="append", + default=[], + help="`NAME=REQS`, where REQS is a comma list of `[N*]Q[@K][c][:D]` " + "(query length Q, KV length K including the query (default Q), D draft " + "tokens, N requests). A prefill's prefix is a prefix-cache hit unless c " + "marks a chunked prefill. E.g. `decode=32*1@8192`, " + "`verify=16*4@4096:3`, `prefill=8192`, `hit=2048@10240`, " + "`mix=8192,32*1@8192`. Repeatable.", + ) + parser.add_argument("--shapes", help="JSON or YAML file with cases.") + parser.add_argument( + "--kv-len-jitter", + type=float, + default=0.5, + help="Draw each request's prefix (KV length minus query length) " + "uniformly within +-this fraction of its value, like the varied lengths " + "of a real batch, capped at the max model length. Seeded by --seed only, " + "so every DP rank gets the same lengths. 0 keeps the exact lengths.", + ) + parser.add_argument( + "--kv-init", + choices=["zeros"], + default="zeros", + help="Initial KV cache contents of each case's blocks.", + ) + parser.add_argument( + "--routing", + choices=["auto", "real", "balanced"], + default="auto", + help="MoE routing. balanced: fixed routing with equal load per EP rank. " + "auto: balanced with dummy weights, else real.", + ) + parser.add_argument("--iters", type=int, default=20, help="Replays per sample.") + parser.add_argument("--samples", type=int, default=5) + parser.add_argument("--output-json", help="Write results to this file.") + EngineArgs.add_cli_args(parser) + parser.set_defaults(load_format="dummy", enable_prefix_caching=False) + + +def _setup_env() -> None: + os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1" + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + # Each case's eager pass compiles exactly the shapes it needs. + os.environ.setdefault("VLLM_DEEP_GEMM_WARMUP", "skip") + if "RANK" in os.environ: + return # Launched by torchrun. + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + os.environ.update( + RANK="0", + LOCAL_RANK="0", + WORLD_SIZE="1", + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + ) + + +def _register_dummy_loader() -> str: + """Dummy loader that also runs load-time model finalization. + + TODO(woosuk): Remove once #60096 lands. DeepSeek-V4.1 fuses its MegaMoE + shared experts at the end of `load_weights`, before the generic linear + post-processing repacks their scales. Dummy loading never calls + `load_weights`, so the fusion would fail and the bench would time unfused + shared experts. + """ + from vllm.model_executor.model_loader import register_model_loader + from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader + + name = "forward_bench_dummy" + + @register_model_loader(name) + class ForwardBenchDummyLoader(DummyModelLoader): + def load_weights(self, model, model_config) -> None: + super().load_weights(model, model_config) + if any(hasattr(m, "finalize_mega_moe_weights") for m in model.modules()): + model.process_weights_after_loading() + + return name + + +def _make_engine_args(args: argparse.Namespace, cases: list[BenchCase]): + from vllm.config.compilation import CompilationMode, CUDAGraphMode + + engine_args = EngineArgs.from_cli_args(args) + if engine_args.load_format == "dummy": + engine_args.load_format = _register_dummy_loader() + engine_args.distributed_executor_backend = "external_launcher" + engine_args.enable_prefix_caching = False + engine_args.compilation_config.mode = CompilationMode.NONE + engine_args.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + if engine_args.max_num_seqs is None: + engine_args.max_num_seqs = max(len(c.requests) for c in cases) + if engine_args.max_num_batched_tokens is None: + # Prefix hits may replay up to a window more per request. + engine_args.max_num_batched_tokens = max( + 8192, + *( + c.num_tokens + 1024 * sum(r.prefix_hit for r in c.requests) + for c in cases + ), + ) + return engine_args + + +def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any]: + from vllm.benchmarks.forward.batch import build_batch + from vllm.benchmarks.forward.capture import ForwardBench + from vllm.benchmarks.forward.kv_init import touched_kv_views, zero_kv_blocks + from vllm.benchmarks.forward.routing import install_routing + + context = runner.vllm_config.compilation_config.static_forward_context + kv_config = runner.kv_cache_config + dcp_size = runner.parallel_config.decode_context_parallel_size + rng = np.random.default_rng([args.seed, runner.dp_rank]) + batch = build_batch( + case, + kv_config.kv_cache_groups, + kv_config.num_blocks, + runner.model_config.get_vocab_size(), + rng, + dcp_size, + ) + zero_kv_blocks(context, kv_config.kv_cache_groups, batch, kv_config.num_blocks) + bench = ForwardBench( + runner, + batch, + args.iters, + args.samples, + touched_kv_views( + context, kv_config.kv_cache_groups, batch, kv_config.num_blocks, dcp_size + ), + ) + routing, balanced = install_routing( + runner, args.routing, args.load_format == "dummy", batch.num_tokens, args.seed + ) + try: + result = bench.run() + if balanced is not None: + balanced.check_applied() + finally: + runner.execute_model(batch.finish()) + if balanced is not None: + balanced.remove() + return { + "num_tokens": batch.num_tokens, + "kv_blocks": batch.num_blocks, + "routing": routing, + **result, + } + + +def main(args: argparse.Namespace) -> None: + cases = [parse_case(c) for c in args.case] + if args.shapes: + cases += load_cases(args.shapes) + if not cases: + raise ValueError("Pass at least one --case or --shapes.") + if not 0 <= args.kv_len_jitter < 1: + raise ValueError("--kv-len-jitter must be in [0, 1).") + _setup_env() + + from vllm import LLM + from vllm.config import set_current_vllm_config + from vllm.distributed.parallel_state import get_world_group + + llm = LLM.from_engine_args(_make_engine_args(args, cases)) + runner = llm.llm_engine.model_executor.driver_worker.worker.model_runner + # Same draw on every rank: shapes are balanced across DP ranks. + max_kv_len = runner.model_config.max_model_len + cases = [ + jitter_kv_lens(c, args.kv_len_jitter, max_kv_len, args.seed) for c in cases + ] + world = get_world_group() + results = [] + with torch.inference_mode(), set_current_vllm_config(runner.vllm_config): + for case in cases: + result = run_case(runner, case, args) + gathered: list[Any] = [None] * world.world_size + torch.distributed.all_gather_object(gathered, result, group=world.cpu_group) + if world.rank != 0: + continue + rank0 = gathered[0] + latency = rank0["latency_us"] # Each replay is the max over ranks. + results.append( + { + "case": case.name, + "requests": [asdict(r) for r in case.requests], + "median_us": latency["median"], + "ranks": gathered, + } + ) + print( + f"[forward] {case.name}: {len(case.requests)} reqs, " + f"{rank0['num_tokens']} tokens, median {latency['median']:.1f} us, " + f"sample spread {latency['sample_spread_pct']:.2f}%, graph vs " + f"eager rel rms {rank0['graph_vs_eager']['relative_rms']:.1e}", + flush=True, + ) + if latency["sample_spread_pct"] > 0.5: + logger.warning( + "%s: unstable timing (sample medians spread %.2f%%).", + case.name, + latency["sample_spread_pct"], + ) + if world.rank == 0 and args.output_json: + with open(args.output_json, "w") as f: + json.dump(results, f, indent=2) diff --git a/vllm/benchmarks/forward/routing.py b/vllm/benchmarks/forward/routing.py new file mode 100644 index 0000000000..6cb9229ce6 --- /dev/null +++ b/vllm/benchmarks/forward/routing.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Balanced MoE routing for benchmarks with dummy weights. + +Dummy gate weights give degenerate routing, which skews MoE and all-to-all +cost. The routing still runs; its result is replaced by the fixture at +whichever boundary the layer exposes: + +- Gate output: fixed logits whose top-k is the fixture. Covers every backend + that routes from logits, including monolithic kernels that run top-k + internally (e.g. FlashInfer TRTLLM). +- Experts input: models that compute top-k themselves and pass + `topk_ids` to their experts (e.g. MegaMoE with a fused gate) get the + fixture's IDs. The routing weights are kept; they don't affect cost. + +Hash-routed layers ignore logits and keep their tables unless their experts +take `topk_ids`. +""" + +import inspect + +import numpy as np +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + +logger = init_logger(__name__) + +# Fixture logits: chosen experts score far above the rest under any monotonic +# scoring function, so top-k picks exactly them. +_CHOSEN_LOGIT = 4.0 +_OTHER_LOGIT = -4.0 + + +def balanced_topk_ids( + num_tokens: int, num_experts: int, top_k: int, ep_size: int, rng +) -> np.ndarray: + """Expert IDs [num_tokens, top_k] giving every EP rank equal load. + + Each token's slots go round-robin over EP ranks, to distinct experts. + """ + assert num_experts % ep_size == 0 + local = num_experts // ep_size + slots = np.arange(num_tokens)[:, None] * top_k + np.arange(top_k) + ranks = slots % ep_size + # Slots of one token that land on the same rank take consecutive experts. + repeat = np.zeros_like(ranks) + for j in range(1, top_k): + repeat[:, j] = (ranks[:, :j] == ranks[:, j : j + 1]).sum(axis=1) + assert repeat.max() < local + base = rng.integers(local, size=(num_tokens, ep_size)) + offset = np.take_along_axis(base, ranks, axis=1) + return ranks * local + (offset + repeat) % local + + +def _moe_layers(model: torch.nn.Module) -> list[tuple[torch.nn.Module, GateLinear]]: + """(MoE module, gate) per layer, in model order.""" + layers: dict[GateLinear, torch.nn.Module] = {} + for module in model.modules(): + gate = getattr(module, "gate", None) + if not isinstance(gate, GateLinear) or gate in layers: + continue + if _top_k(module) is not None: + layers[gate] = module + return [(module, gate) for gate, module in layers.items()] + + +def _top_k(module: torch.nn.Module) -> int | None: + """The MoE layer's top-k, from the module or its experts submodules.""" + for sub in module.modules(): + for name in ("top_k", "n_activated_experts"): + top_k = getattr(sub, name, None) + if isinstance(top_k, int): + return top_k + return None + + +class BalancedRouting: + """Installs balanced routing for one case; `remove()` restores the model.""" + + def __init__( + self, + model: torch.nn.Module, + num_tokens: int, + dp_rank: int, + dp_size: int, + ep_size: int, + seed: int, + device: torch.device, + ): + layers = _moe_layers(model) + if not layers: + raise ValueError("No MoE gates found; use --routing real.") + self._hooks = [] + self._biases: list[tuple[torch.Tensor, torch.Tensor]] = [] + self._fired = [0] * len(layers) + self._layers = layers + # Per-rank rows, with headroom for padding; every rank generates the + # same global fixture and takes its slice. + rows = -(-num_tokens // 256) * 256 + 256 + for layer, (module, gate) in enumerate(layers): + if (getattr(module, "num_expert_group", None) or 1) > 1: + raise NotImplementedError("Balanced routing with grouped top-k.") + num_experts = gate.weight.shape[0] + top_k = _top_k(module) + assert top_k is not None + rng = np.random.default_rng([seed, layer]) + ids = balanced_topk_ids(rows * dp_size, num_experts, top_k, ep_size, rng)[ + dp_rank * rows : (dp_rank + 1) * rows + ] + logits = torch.full((rows, num_experts), _OTHER_LOGIT, dtype=torch.float32) + logits.scatter_(1, torch.from_numpy(ids), _CHOSEN_LOGIT) + # Distinct values: no ties for any top-k implementation. + logits += torch.from_numpy(rng.uniform(0, 0.5, logits.shape)).float() + self._hooks.append( + gate.register_forward_hook( + self._substitute(logits.to(device), self._fired, layer) + ) + ) + bias = getattr(gate, "e_score_correction_bias", None) + if bias is None: + bias = getattr(module, "e_score_correction_bias", None) + if isinstance(bias, torch.Tensor): + self._biases.append((bias, bias.detach().clone())) + bias.data.zero_() + experts = getattr(module, "experts", None) + if isinstance(experts, torch.nn.Module): + params = list(inspect.signature(experts.forward).parameters) + if "topk_ids" in params: + self._hooks.append( + experts.register_forward_pre_hook( + self._substitute_ids( + torch.from_numpy(ids).to(device), + params.index("topk_ids"), + self._fired, + layer, + ), + with_kwargs=True, + ) + ) + + @staticmethod + def _substitute(fixture: torch.Tensor, fired: list[int], layer: int): + cast = [fixture] + + def hook(module, inputs, output): + fired[layer] += 1 + logits = output[0] if isinstance(output, tuple) else output + if cast[0].dtype != logits.dtype: + cast[0] = fixture.to(logits.dtype) # Once, at warmup. + replaced = cast[0][: logits.shape[0]] + return (replaced, *output[1:]) if isinstance(output, tuple) else replaced + + return hook + + @staticmethod + def _substitute_ids(fixture: torch.Tensor, index: int, fired, layer: int): + cast = [fixture] + + def hook(module, args, kwargs): + fired[layer] += 1 + ids = kwargs["topk_ids"] if "topk_ids" in kwargs else args[index] + if cast[0].dtype != ids.dtype: + cast[0] = fixture.to(ids.dtype) # Once, at warmup. + replaced = cast[0][: ids.shape[0]] + if "topk_ids" in kwargs: + return args, {**kwargs, "topk_ids": replaced} + return (*args[:index], replaced, *args[index + 1 :]), kwargs + + return hook + + def check_applied(self) -> None: + """Fail if a layer bypassed both the gate and the experts hooks.""" + missed = [ + type(module).__name__ + for (module, _), fired in zip(self._layers, self._fired) + if not fired + ] + if missed: + raise RuntimeError( + f"Balanced routing didn't reach {len(missed)} MoE layers " + f"({missed[0]}): their gate is fused away. Use --routing real." + ) + + def remove(self) -> None: + for hook in self._hooks: + hook.remove() + for bias, saved in self._biases: + bias.data.copy_(saved) + + +def install_routing( + runner, requested: str, dummy_weights: bool, num_tokens: int, seed: int +) -> tuple[str, BalancedRouting | None]: + """Resolve `--routing` and install balanced routing if selected. + + `auto` is balanced with dummy weights and real otherwise; it falls back to + real routing when the model doesn't support the fixture. + """ + mode = requested + if mode == "auto": + mode = "balanced" if dummy_weights else "real" + if mode == "real": + return mode, None + parallel = runner.parallel_config + ep_size = ( + parallel.data_parallel_size * parallel.tensor_parallel_size + if parallel.enable_expert_parallel + else 1 + ) + try: + routing = BalancedRouting( + runner.model, + num_tokens, + runner.dp_rank, + runner.dp_size, + ep_size, + seed, + runner.device, + ) + except (ValueError, NotImplementedError) as e: + if requested == "balanced": + raise + logger.warning("Using real routing: %s", e) + return "real", None + return mode, routing diff --git a/vllm/benchmarks/forward/shapes.py b/vllm/benchmarks/forward/shapes.py new file mode 100644 index 0000000000..6f5c670a32 --- /dev/null +++ b/vllm/benchmarks/forward/shapes.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exact batch shapes for `vllm bench forward`.""" + +import json +import zlib +from dataclasses import dataclass, replace +from pathlib import Path + +import numpy as np +import regex as re + + +@dataclass(frozen=True) +class BenchRequest: + query_len: int + # KV length after the step: the prefix already in the KV cache plus the + # query tokens. + kv_len: int + # Spec-decode verify: query_len == 1 + num_draft_tokens. + num_draft_tokens: int = 0 + # The prefix was computed by this request's earlier chunks. Otherwise a + # prefill's prefix is a prefix-cache hit, of which models with SWA bounded + # replay (DeepSeek-V4.1) recompute the last window. + chunked: bool = False + + def __post_init__(self) -> None: + if ( + self.query_len < 1 + or self.kv_len < self.query_len + or self.num_draft_tokens < 0 + ): + raise ValueError(f"Invalid request: {self}") + if self.num_draft_tokens and not self.is_decode: + raise ValueError( + f"Verify needs query_len == 1 + num_draft_tokens and a prefix: {self}" + ) + if self.chunked and (self.is_decode or not self.prefix_len): + raise ValueError(f"A chunk needs a prefill with a prefix: {self}") + + @property + def prefix_len(self) -> int: + return self.kv_len - self.query_len + + @property + def is_decode(self) -> bool: + """Decode or verify; otherwise a prefill.""" + return self.prefix_len > 0 and self.query_len == 1 + self.num_draft_tokens + + @property + def prefix_hit(self) -> bool: + return self.prefix_len > 0 and not self.is_decode and not self.chunked + + +@dataclass(frozen=True) +class BenchCase: + name: str + requests: tuple[BenchRequest, ...] + + @property + def num_tokens(self) -> int: + return sum(r.query_len for r in self.requests) + + +def jitter_kv_lens( + case: BenchCase, jitter: float, max_kv_len: int, seed: int +) -> BenchCase: + """Draw each request's prefix uniformly within +-jitter of its value. + + Query lengths are kept, so a fresh prefill stays one and a non-empty + prefix stays non-empty. KV lengths are capped at `max_kv_len`. The draw + depends only on the seed and the case's name, not on the other cases. + """ + rng = np.random.default_rng([seed, zlib.crc32(case.name.encode())]) + requests = [] + for req in case.requests: + if req.kv_len > max_kv_len: + raise ValueError(f"{case.name}: KV length {req.kv_len} > {max_kv_len}") + prefix = req.prefix_len + if prefix and jitter: + prefix = max(1, round(prefix * rng.uniform(1 - jitter, 1 + jitter))) + kv_len = min(req.query_len + prefix, max_kv_len) + requests.append(replace(req, kv_len=kv_len)) + return replace(case, requests=tuple(requests)) + + +_REQ_RE = re.compile(r"^(?:(\d+)\*)?(\d+)(?:@(\d+))?(c)?(?::(\d+))?$") + + +def parse_case(text: str) -> BenchCase: + """Parse `NAME=REQS`, where REQS is a comma list of `[N*]Q[@K][c][:D]`. + + Q is the query length, K the KV length including the query (default Q, a + fresh prefill), D the number of draft tokens (verify) and N the number of + such requests. A prefill's prefix is a prefix-cache hit unless `c` marks + it as a chunked prefill. For example, `decode=32*1@8192`, + `verify=16*4@4096:3`, `prefill=8192`, `hit=8192@40960`, + `chunk=8192@40960c` or `mix=8192,32*1@8192`. + """ + name, sep, body = text.partition("=") + if not sep: + name, body = text, text + requests: list[BenchRequest] = [] + for item in body.split(","): + match = _REQ_RE.match(item.strip()) + if match is None: + raise ValueError(f"Invalid request {item!r} in case {text!r}") + count, query, kv, chunked, drafts = match.groups() + request = BenchRequest( + int(query), int(kv or query), int(drafts or 0), chunked=bool(chunked) + ) + requests.extend([request] * int(count or 1)) + return BenchCase(name, tuple(requests)) + + +def load_cases(path: str) -> list[BenchCase]: + """Load cases from a JSON or YAML file. + + Format: `{"cases": [{"name": ..., "requests": [{"query_len": ..., + "kv_len": ..., "num_draft_tokens": ..., "chunked": ..., "count": ...}]}]}`. + """ + text = Path(path).read_text() + if path.endswith((".yaml", ".yml")): + import yaml + + data = yaml.safe_load(text) + else: + data = json.loads(text) + cases = [] + for case in data["cases"]: + requests: list[BenchRequest] = [] + for req in case["requests"]: + req = dict(req) + count = req.pop("count", 1) + requests.extend([BenchRequest(**req)] * count) + cases.append(BenchCase(case["name"], tuple(requests))) + return cases diff --git a/vllm/entrypoints/cli/benchmark/forward.py b/vllm/entrypoints/cli/benchmark/forward.py new file mode 100644 index 0000000000..5128a1057b --- /dev/null +++ b/vllm/entrypoints/cli/benchmark/forward.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import argparse + +from vllm.benchmarks.forward.main import add_cli_args, main +from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +class BenchmarkForwardSubcommand(BenchmarkSubcommandBase): + """The `forward` subcommand for `vllm bench`.""" + + name = "forward" + help = "Benchmark the model forward of exact batches as CUDA graphs." + + @classmethod + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None: + add_cli_args(parser) + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + main(args) diff --git a/vllm/entrypoints/cli/benchmark/main.py b/vllm/entrypoints/cli/benchmark/main.py index 9ea4998709..d540c9ddde 100644 --- a/vllm/entrypoints/cli/benchmark/main.py +++ b/vllm/entrypoints/cli/benchmark/main.py @@ -38,6 +38,7 @@ def maybe_exec_rust_bench() -> None: def _import_bench_subcommand_modules() -> None: # Imported lazily so `BenchmarkSubcommandBase` subclasses register only # when `vllm bench` is actually invoked. + import vllm.entrypoints.cli.benchmark.forward # noqa: F401 import vllm.entrypoints.cli.benchmark.latency # noqa: F401 import vllm.entrypoints.cli.benchmark.mm_processor # noqa: F401 import vllm.entrypoints.cli.benchmark.serve # noqa: F401