diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md index 7f18c2bb87..f630f698a1 100644 --- a/.agents/skills/model-microbenchmark/SKILL.md +++ b/.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/.rank.kernels.json` (GPU time per kernel per forward, sorted) +and `DIR/.rank.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/`, 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 diff --git a/docs/cli/README.md b/docs/cli/README.md index c2b27bed26..3cb6ae68d8 100644 --- a/docs/cli/README.md +++ b/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 \ diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py index d44c6b9466..07b216e341 100644 --- a/tests/benchmarks/test_forward_cli.py +++ b/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*)", "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() diff --git a/vllm/benchmarks/forward/capture.py b/vllm/benchmarks/forward/capture.py index 8b378bb18c..8f4f3396ae 100644 --- a/vllm/benchmarks/forward/capture.py +++ b/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 diff --git a/vllm/benchmarks/forward/main.py b/vllm/benchmarks/forward/main.py index d7841c641e..7525c93d9f 100644 --- a/vllm/benchmarks/forward/main.py +++ b/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 " + "`.rank.trace.json.gz` (Perfetto) and `.rank.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%%).", diff --git a/vllm/benchmarks/forward/profile.py b/vllm/benchmarks/forward/profile.py new file mode 100644 index 0000000000..d06b50aba0 --- /dev/null +++ b/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 (`.rank.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(args)` -> `ns::foo`: return type, template and + function arguments dropped. Generic launchers (e.g. + `cutlass::device_kernel>`) 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()