[Benchmark] Add per-case kernel profiles to vllm bench forward (#61692)

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

.agents/skills/model-microbenchmark/SKILL.md

@@ -51,6 +51,46 @@ Key flags (plus all `EngineArgs`):
 
 Read `median_us` (each replay is already the max over ranks).
 
+## Kernel breakdown
+
+Add `--profile-dir DIR` to see where the time goes. After timing, each case is
+replayed 3 more times under `torch.profiler`; every rank writes
+`DIR/<case>.rank<r>.kernels.json` (GPU time per kernel per forward, sorted)
+and `DIR/<case>.rank<r>.trace.json.gz` (open in Perfetto). Rank 0's top
+kernels are printed and included in `--output-json`.
+
+```bash
+python -m vllm.benchmarks.forward.profile DIR/decode.rank0.kernels.json
+# Before/after a change, same case: per-kernel deltas, largest first.
+python -m vllm.benchmarks.forward.profile base/decode.rank0.kernels.json new/decode.rank0.kernels.json
+```
+
+- Kernels are grouped by name without template or function arguments; the
+  JSON keeps one full name per group in `example`.
+- Only the timed forward is included: the in-graph KV restore and rank
+  rendezvous are excluded (`pre_forward_us`). `span_us` (first to last
+  forward kernel) should be close to `median_us`, a bit higher from profiler
+  overhead; `busy_us` can exceed it when streams overlap.
+- Compare the same rank across runs: under EP, ranks differ (e.g. collective
+  wait time shows up inside communication kernels).
+
+For Nsight Systems (NVTX, NCCL, CPU side, GPU metrics), skip `--profile-dir`:
+the two can't share CUPTI. Each case's 3 profiled replays run between
+`cudaProfilerStart`/`Stop` in an NVTX range `forward/<case>`, so a capture range
+records only them, one report per case (and rank):
+
+```bash
+nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat \
+  --cuda-graph-trace=node -o prof vllm bench forward ...       # prof.1, prof.2, ...
+torchrun --nproc-per-node=4 --no-python nsys profile ... -o prof_rank%q{RANK} \
+  python -m vllm.entrypoints.cli.main bench forward ...
+nsys stats --report cuda_gpu_kern_sum --format csv -o kern prof.1.nsys-rep
+```
+
+Keep `--capture-range-end=repeat`: nsys's default stops after the first case
+and leaves the bench hanging. `nsys stats` also counts the KV restore copies
+that the JSON summary excludes.
+
 ## Caveats
 
 - **Dummy weights lie for data-dependent kernels.** Dummy gates route every

docs/cli/README.md

@@ -136,7 +136,7 @@ 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.
+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. Add `--profile-dir DIR` to write a Perfetto trace and a per-kernel GPU time summary of each case; `python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]` prints a summary or compares two. Under Nsight Systems, `nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat --cuda-graph-trace=node` records just those replays, one report per case.
 
 ```bash
 vllm bench forward \

tests/benchmarks/test_forward_cli.py

@@ -3,6 +3,7 @@
 import json
 import subprocess
 from dataclasses import replace
+from pathlib import Path
 from types import SimpleNamespace
 
 import numpy as np
@@ -10,6 +11,7 @@ import pytest
 
 from vllm.benchmarks.forward.batch import build_batch
 from vllm.benchmarks.forward.main import run_case
+from vllm.benchmarks.forward.profile import summarize
 from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case
 from vllm.platforms import current_platform
 
@@ -97,6 +99,35 @@ def test_bench_forward_rejects_oversized_verify(capacity):
         run_case(runner, case, SimpleNamespace())
 
 
+def test_profile_summary():
+    """Each replay's leading restore/rendezvous activities are left out, and
+    kernels are grouped by name without template or function arguments."""
+
+    def act(name, ts, dur, cat="kernel"):
+        return {"cat": cat, "name": name, "ts": ts, "dur": dur}
+
+    replay = [
+        ("Memcpy DtoD (Device -> Device)", "gpu_memcpy", 5),
+        ("void ns::gemm<128>(Params)", "kernel", 10),
+        ("void ns::gemm<64>(Params)", "kernel", 20),
+        ("std::enable_if<(1>0), void>::type ns::norm<float>(float*)", "kernel", 4),
+    ]
+    events = [{"cat": "cpu_op", "name": "x", "ts": 0, "dur": 1}]
+    for r in range(2):
+        for i, (name, cat, dur) in enumerate(replay):
+            events.append(act(name, 1000 * r + 100 * i, dur, cat))
+    summary = summarize(events, replays=2, skip=1)
+    assert summary["pre_forward_us"] == 5
+    assert summary["busy_us"] == 34
+    assert summary["span_us"] == 204
+    assert [(k["name"], k["us"], k["calls"]) for k in summary["kernels"]] == [
+        ("ns::gemm", 30, 2),
+        ("ns::norm", 4, 1),
+    ]
+    with pytest.raises(ValueError, match="replays"):
+        summarize(events[:-1], replays=2, skip=1)
+
+
 @pytest.mark.benchmark
 # TODO(woosuk): Support ROCm and re-enable this test there.
 @pytest.mark.skipif(
@@ -126,6 +157,8 @@ def test_bench_forward(tmp_path):
         "2",
         "--output-json",
         str(output),
+        "--profile-dir",
+        str(tmp_path / "profile"),
     ]
     result = subprocess.run(command, capture_output=True, text=True)
     print(result.stdout)
@@ -139,3 +172,7 @@ def test_bench_forward(tmp_path):
         assert len(rank0["batch_order"]) == len(row["requests"])
         assert rank0["finite"] and rank0["graph_deterministic"]
         assert row["median_us"] > 0
+        profile = rank0["profile"]
+        assert profile["top_kernels"] and profile["span_us"] > 0
+        summary = json.loads(Path(profile["summary"]).read_text())
+        assert Path(summary["trace"]).exists()

vllm/benchmarks/forward/capture.py

@@ -55,6 +55,8 @@ class ForwardBench:
         iters: int,
         samples: int,
         state_views: list[torch.Tensor],
+        name: str,
+        profile_prefix: str | None = None,
     ):
         self.runner = runner
         self.batch = batch
@@ -62,6 +64,8 @@ class ForwardBench:
         self.iters = iters
         self.samples = samples
         self.state_views = state_views
+        self.name = name
+        self.profile_prefix = profile_prefix
         self.phase = "warmup"
         self.result: dict[str, Any] = {}
 
@@ -193,6 +197,20 @@ class ForwardBench:
             ),
         )
         self.result["latency_us"] = self._time(graph)
+        from vllm.benchmarks.forward.profile import profile_graph
+
+        # Each replay restores the state views, then (multi-rank) rendezvouses,
+        # before the timed forward.
+        skip = len(self.state_views) + (world.world_size > 1)
+        summary = profile_graph(graph, self.name, self.profile_prefix, skip)
+        if summary is not None:
+            self.result["profile"] = {
+                "summary": f"{self.profile_prefix}.kernels.json",
+                "span_us": summary["span_us"],
+                "busy_us": summary["busy_us"],
+                "activities": summary["activities"],
+                "top_kernels": summary["kernels"][:10],
+            }
         del graph
         return output
 

vllm/benchmarks/forward/main.py

@@ -74,6 +74,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None:
     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.")
+    parser.add_argument(
+        "--profile-dir",
+        help="After timing, trace a few replays of each case and write "
+        "`<case>.rank<r>.trace.json.gz` (Perfetto) and `<case>.rank<r>.kernels.json` "
+        "(GPU time per kernel per forward) here. Summarize or compare with "
+        "`python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]`.",
+    )
     EngineArgs.add_cli_args(parser)
     parser.set_defaults(load_format="dummy", enable_prefix_caching=False)
 
@@ -146,6 +153,13 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any
         dcp_size,
     )
     zero_kv_blocks(context, kv_config.kv_cache_groups, batch, kv_config.num_blocks)
+    profile_prefix = None
+    if args.profile_dir:
+        from vllm.distributed.parallel_state import get_world_group
+
+        os.makedirs(args.profile_dir, exist_ok=True)
+        rank = get_world_group().rank
+        profile_prefix = os.path.join(args.profile_dir, f"{case.name}.rank{rank}")
     bench = ForwardBench(
         runner,
         batch,
@@ -154,6 +168,8 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any
         touched_kv_views(
             context, kv_config.kv_cache_groups, batch, kv_config.num_blocks, dcp_size
         ),
+        case.name,
+        profile_prefix,
     )
     routing, balanced = install_routing(
         runner, args.routing, args.load_format == "dummy", batch.num_tokens, args.seed
@@ -221,6 +237,14 @@ def main(args: argparse.Namespace) -> None:
                 f"eager rel rms {rank0['graph_vs_eager']['relative_rms']:.1e}",
                 flush=True,
             )
+            if profile := rank0.get("profile"):
+                from vllm.benchmarks.forward.profile import format_summary
+
+                print(
+                    f"[forward] {case.name} rank 0 kernels ({profile['summary']}):\n"
+                    + format_summary({**profile, "kernels": profile["top_kernels"]}),
+                    flush=True,
+                )
             if latency["sample_spread_pct"] > 0.5:
                 logger.warning(
                     "%s: unstable timing (sample medians spread %.2f%%).",

vllm/benchmarks/forward/profile.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GPU kernel profiles of `vllm bench forward` cases.
+
+With `--profile-dir`, each case's graph is replayed a few more times under
+`torch.profiler` after timing. Each rank writes a Perfetto trace and a kernel
+summary (`<case>.rank<r>.kernels.json`): per-kernel time per forward, sorted.
+
+Summarize one file or compare two (e.g. before/after a change):
+`python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]`.
+
+The same replays run between `cudaProfilerStart`/`Stop` in an NVTX range per
+case (with or without `--profile-dir`), so Nsight Systems can record just them:
+`nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat
+--cuda-graph-trace=node vllm bench forward ...`.
+"""
+
+import argparse
+import gzip
+import json
+import statistics
+from collections import defaultdict
+from typing import Any
+
+import regex as re
+import torch
+
+_GPU_CATEGORIES = ("kernel", "gpu_memcpy", "gpu_memset")
+# Kernels named after a generic launcher are identified by their first
+# template argument instead.
+_GENERIC = re.compile(r"(^|::)(device_kernel|Kernel|kernel|KernelWrapper)$")
+# CuTe DSL kernel names embed per-shape layouts after these markers.
+_CUTE_DSL_SUFFIX = re.compile(r"_object_at_|_tensorptr")
+
+
+def _match_back(text: str, end: int, open_ch: str, close_ch: str) -> int:
+    """Index of the bracket opening the one closing at `end`, skipping
+    brackets of the other kind nested inside parentheses."""
+    depth, parens = 0, 0
+    for i in range(end, -1, -1):
+        ch = text[i]
+        if open_ch == "<" and ch in "()":
+            parens += 1 if ch == ")" else -1
+            continue
+        if parens:
+            continue
+        if ch == close_ch:
+            depth += 1
+        elif ch == open_ch:
+            depth -= 1
+            if depth == 0:
+                return i
+    return -1
+
+
+def short_name(name: str) -> str:
+    """`void ns::foo<T, 4>(args)` -> `ns::foo`: return type, template and
+    function arguments dropped. Generic launchers (e.g.
+    `cutlass::device_kernel<Op<...>>`) keep their first template argument."""
+    if name.startswith("kernel_cutlass_"):
+        return _CUTE_DSL_SUFFIX.split(name, maxsplit=1)[0]
+    head = name.strip()
+    if head.endswith(")"):
+        head = head[: max(_match_back(head, len(head) - 1, "(", ")"), 0)].rstrip()
+    template = ""
+    if head.endswith(">"):
+        start = _match_back(head, len(head) - 1, "<", ">")
+        if start > 0:
+            head, template = head[:start], head[start + 1 : -1]
+    base = head.split()[-1] if head.split() else name
+    if _GENERIC.search(base) and template:
+        arg = re.split(r"[<>(,]", template, maxsplit=1)[0].strip()
+        if arg:
+            return f"{base}<{arg.removeprefix('void ')}>"
+    return base
+
+
+def summarize(events: list[dict], replays: int, skip: int) -> dict[str, Any]:
+    """Per-forward kernel times from the GPU activities of `replays` replays.
+
+    Each replay first runs `skip` activities that aren't part of the timed
+    forward (the KV state restore and the rank rendezvous); they are reported
+    separately as `pre_forward_us`.
+    """
+    acts = sorted(
+        (e for e in events if e.get("cat") in _GPU_CATEGORIES and "dur" in e),
+        key=lambda e: e["ts"],
+    )
+    if not acts:
+        raise ValueError(
+            "No GPU activity recorded. Another profiler (e.g. nsys) may hold "
+            "CUPTI; don't combine --profile-dir with it."
+        )
+    if len(acts) % replays:
+        raise ValueError(
+            f"{len(acts)} GPU activities don't split into {replays} replays"
+        )
+    per_replay = len(acts) // replays
+    if per_replay <= skip:
+        raise ValueError(f"{per_replay} activities per replay, {skip} to skip")
+    kernels: dict[str, dict[str, Any]] = defaultdict(
+        lambda: {"total_us": 0.0, "calls": 0, "example": ""}
+    )
+    spans, busy, pre = [], [], []
+    for r in range(replays):
+        chunk = acts[r * per_replay : (r + 1) * per_replay]
+        pre.append(sum(e["dur"] for e in chunk[:skip]))
+        forward = chunk[skip:]
+        spans.append(
+            max(e["ts"] + e["dur"] for e in forward) - min(e["ts"] for e in forward)
+        )
+        busy.append(sum(e["dur"] for e in forward))
+        for e in forward:
+            name = short_name(e["name"]) if e["cat"] == "kernel" else e["name"]
+            k = kernels[name]
+            k["total_us"] += e["dur"]
+            k["calls"] += 1
+            k["example"] = k["example"] or e["name"][:300]
+    total = sum(k["total_us"] for k in kernels.values())
+    rows = [
+        {
+            "name": name,
+            "us": round(k["total_us"] / replays, 2),
+            "pct": round(100 * k["total_us"] / total, 2),
+            "calls": k["calls"] // replays,
+            "example": k["example"],
+        }
+        for name, k in kernels.items()
+    ]
+    rows.sort(key=lambda row: -row["us"])
+    return {
+        "replays": replays,
+        # First to last forward activity; busy can exceed it when streams overlap.
+        "span_us": round(statistics.median(spans), 2),
+        "busy_us": round(statistics.median(busy), 2),
+        "pre_forward_us": round(statistics.median(pre), 2),
+        "activities": per_replay - skip,
+        "kernels": rows,
+    }
+
+
+def profile_graph(
+    graph: torch.cuda.CUDAGraph,
+    name: str,
+    prefix: str | None,
+    skip: int,
+    replays: int = 3,
+) -> dict[str, Any] | None:
+    """Replay `replays` times between `cudaProfilerStart`/`Stop`, in an NVTX
+    range named after the case, for `nsys --capture-range=cudaProfilerApi`.
+
+    With `prefix`, also trace them with `torch.profiler` and write
+    `prefix.trace.json.gz` and `prefix.kernels.json`. Every rank must call
+    this (replays may contain collectives).
+    """
+
+    def run() -> None:
+        torch.cuda.profiler.start()
+        with torch.cuda.nvtx.range(f"forward/{name}"):
+            for _ in range(replays):
+                graph.replay()
+                torch.accelerator.synchronize()
+        torch.cuda.profiler.stop()
+
+    if prefix is None:
+        run()
+        return None
+    activities = [torch.profiler.ProfilerActivity.CUDA]
+    with torch.profiler.profile(activities=activities) as prof:
+        run()
+    trace = f"{prefix}.trace.json.gz"
+    prof.export_chrome_trace(trace)
+    with gzip.open(trace, "rt") as f:
+        events = json.load(f)["traceEvents"]
+    summary = {"trace": trace, **summarize(events, replays, skip)}
+    with open(f"{prefix}.kernels.json", "w") as f:
+        json.dump(summary, f, indent=1)
+    return summary
+
+
+def format_summary(summary: dict[str, Any], top: int = 20) -> str:
+    lines = [
+        f"span {summary['span_us']:.1f} us, busy {summary['busy_us']:.1f} us, "
+        f"{summary['activities']} activities/forward",
+        f"{'us':>10} {'%':>6} {'calls':>6}  kernel",
+    ]
+    for row in summary["kernels"][:top]:
+        lines.append(
+            f"{row['us']:10.1f} {row['pct']:6.2f} {row['calls']:6d}  {row['name']}"
+        )
+    return "\n".join(lines)
+
+
+def format_diff(a: dict[str, Any], b: dict[str, Any], top: int = 20) -> str:
+    rows_a = {row["name"]: row for row in a["kernels"]}
+    rows_b = {row["name"]: row for row in b["kernels"]}
+    deltas = []
+    for name in rows_a.keys() | rows_b.keys():
+        ra, rb = rows_a.get(name), rows_b.get(name)
+        us_a, us_b = ra["us"] if ra else 0.0, rb["us"] if rb else 0.0
+        calls_a, calls_b = ra["calls"] if ra else 0, rb["calls"] if rb else 0
+        deltas.append((us_b - us_a, us_a, us_b, calls_a, calls_b, name))
+    deltas.sort(key=lambda d: -abs(d[0]))
+    lines = [
+        f"span {a['span_us']:.1f} -> {b['span_us']:.1f} us "
+        f"({b['span_us'] - a['span_us']:+.1f}), busy {a['busy_us']:.1f} -> "
+        f"{b['busy_us']:.1f} us ({b['busy_us'] - a['busy_us']:+.1f})",
+        f"{'delta us':>10} {'A us':>10} {'B us':>10} {'A n':>5} {'B n':>5}  kernel",
+    ]
+    for delta, us_a, us_b, calls_a, calls_b, name in deltas[:top]:
+        lines.append(
+            f"{delta:+10.1f} {us_a:10.1f} {us_b:10.1f} {calls_a:5d} {calls_b:5d}  "
+            f"{name}"
+        )
+    return "\n".join(lines)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("summaries", nargs="+", help="One or two *.kernels.json.")
+    parser.add_argument("--top", type=int, default=20)
+    args = parser.parse_args()
+    if len(args.summaries) > 2:
+        parser.error("Pass one summary, or two to compare.")
+    loaded = []
+    for path in args.summaries:
+        with open(path) as f:
+            loaded.append(json.load(f))
+    if len(loaded) == 1:
+        print(format_summary(loaded[0], args.top))
+    else:
+        print(format_diff(loaded[0], loaded[1], args.top))
+
+
+if __name__ == "__main__":
+    main()