[Benchmark] Add vllm bench forward: exact-shape forward timing (#61339)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

.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.

.claude/skills/model-microbenchmark

@@ -0,0 +1 @@
+../../.agents/skills/model-microbenchmark

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.

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": (

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

vllm/benchmarks/forward/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

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)

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()

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

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)

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

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

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)

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