Skip to content

vllm.benchmarks.forward.profile

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

Functions:

  • profile_graph –

    Replay replays times between cudaProfilerStart/Stop, in an NVTX

  • short_name –

    void ns::foo<T, 4>(args) -> ns::foo: return type, template and

  • summarize –

    Per-forward kernel times from the GPU activities of replays replays.

_match_back(text, end, open_ch, close_ch)

Index of the bracket opening the one closing at end, skipping brackets of the other kind nested inside parentheses.

Source code in vllm/benchmarks/forward/profile.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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

profile_graph(graph, name, prefix, skip, replays=3)

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

Source code in vllm/benchmarks/forward/profile.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
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:
        # `cudaProfilerStart`/`Stop` and NVTX ranges are CUDA-only (for
        # Nsight Systems); other accelerators just replay.
        if current_platform.is_xpu():
            for _ in range(replays):
                graph.replay()
                torch.accelerator.synchronize()
        else:
            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 = get_default_torch_profiler_activities()
    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

short_name(name)

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.

Source code in vllm/benchmarks/forward/profile.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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

summarize(events, replays, skip)

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.

Source code in vllm/benchmarks/forward/profile.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
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,
    }