| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Graph-replay microbenchmark with rotating weight buffers (CUDA and ROCm). |
| 5 | |
| 6 | Platform-agnostic: HIP implements the torch.cuda.* API, including CUDAGraph. |
| 7 | Pattern: capture one graph (HIP graph on ROCm) containing R*INNER calls |
| 8 | cycling over R rotating weight buffers, then time replays with CUDA/HIP |
| 9 | events. The rotating buffers keep the captured footprint above the GPU's |
| 10 | last-level cache, so every call reads weights cold from HBM. Prefer this over |
| 11 | do_bench on ROCm: do_bench's per-rep launch/event overhead inflates |
| 12 | small/decode-size kernels, and its default flush buffer is only as large as |
| 13 | the LLC — on a 256 MB LLC that evicts marginally and leaves the cache full of |
| 14 | dirty flush-buffer lines. The same pattern extends to multi-GPU benchmarks |
| 15 | (capture per rank, collectives in identical order every replay); see the |
| 16 | Multi-GPU section of SKILL.md. |
| 17 | """ |
| 18 | |
| 19 | import statistics |
| 20 | |
| 21 | import pandas as pd |
| 22 | import torch |
| 23 | |
| 24 | LLC_TARGET_BYTES = 400 * 1024**2 # above 256 MB MI355X LLC and B200 L2 (126 MB) |
| 25 | INNER = 2 # repeat slots within one graph to amortize per-replay overhead |
| 26 | REPLAYS = 10 # replays per round: makes each measurement large vs overhead |
| 27 | ROUNDS = 5 # median across rounds: replay time is not always stable |
| 28 | WARMUP = 5 # warmup replays to reach steady clocks before timing |
| 29 | MATMUL_CASES = [ |
| 30 | ("compute-bound", 4096, 4096, 4096), |
| 31 | ("decode-GEMM", 16, 16384, 8192), |
| 32 | ] |
| 33 | |
| 34 | |
| 35 | def timeit_graph(calls, inner=INNER, replays=REPLAYS, rounds=ROUNDS): |
| 36 | """Time a list of per-slot callables as one graph (zero launch overhead); |
| 37 | weights rotate across slots so they stay HBM-cold. Returns the median over |
| 38 | rounds of per-round means, ms per call.""" |
| 39 | for c in calls: |
| 40 | c() |
| 41 | torch.accelerator.synchronize() |
| 42 | g = torch.cuda.CUDAGraph() |
| 43 | with torch.cuda.graph(g): # capture runs on a side stream |
| 44 | for i in range(len(calls) * inner): |
| 45 | calls[i % len(calls)]() |
| 46 | torch.accelerator.synchronize() |
| 47 | events = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(rounds)] |
| 48 | for _ in range(WARMUP): |
| 49 | g.replay() |
| 50 | # No sync between warmup and the first round: queued replays keep the GPU |
| 51 | # busy, so each round's launch latency stays outside its event window. |
| 52 | for s, e in events: |
| 53 | s.record() |
| 54 | for _ in range(replays): |
| 55 | g.replay() |
| 56 | e.record() |
| 57 | torch.accelerator.synchronize() |
| 58 | samples = [s.elapsed_time(e) / (replays * len(calls) * inner) for s, e in events] |
| 59 | return statistics.median(samples) # ms per call |
| 60 | |
| 61 | |
| 62 | def main() -> None: |
| 63 | if not torch.accelerator.is_available(): |
| 64 | raise RuntimeError("No GPU visible (torch.accelerator covers HIP on ROCm).") |
| 65 | |
| 66 | torch.set_default_device("cuda") |
| 67 | torch.manual_seed(0) |
| 68 | |
| 69 | rows = [] |
| 70 | for name, m, n, k in MATMUL_CASES: |
| 71 | a = torch.randn(m, k, dtype=torch.bfloat16) # shared: producer-warm |
| 72 | out = torch.empty(m, n, dtype=torch.bfloat16) |
| 73 | b_bytes = k * n * 2 |
| 74 | r = max(2, -(-LLC_TARGET_BYTES // b_bytes)) # ceil |
| 75 | bs = [torch.randn(k, n, dtype=torch.bfloat16) for _ in range(r)] |
| 76 | |
| 77 | def make_call(i, a=a, bs=bs, out=out): |
| 78 | def call(): |
| 79 | torch.mm(a, bs[i], out=out) |
| 80 | |
| 81 | return call |
| 82 | |
| 83 | # correctness before timing, per slot (the graph would freeze a wrong |
| 84 | # result into every replay) |
| 85 | for i in range(r): |
| 86 | make_call(i)() |
| 87 | torch.accelerator.synchronize() |
| 88 | torch.testing.assert_close(out, torch.mm(a, bs[i]), atol=1e-1, rtol=1e-1) |
| 89 | |
| 90 | us = timeit_graph([make_call(i) for i in range(r)]) * 1e3 |
| 91 | rows.append( |
| 92 | { |
| 93 | "case": name, |
| 94 | "shape": f"{m}x{n}x{k}", |
| 95 | "slots": r, |
| 96 | "us": us, |
| 97 | "tflops": 2 * m * n * k / (us * 1e6), |
| 98 | "gbps": 2 * (m * k + k * n + m * n) / (us * 1e3), |
| 99 | } |
| 100 | ) |
| 101 | |
| 102 | df = pd.DataFrame(rows) |
| 103 | print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}")) |
| 104 | |
| 105 | |
| 106 | if __name__ == "__main__": |
| 107 | main() |