[Skills] Update kernel-microbenchmark to include ROCm (#58646)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Kimi <noreply@moonshot.cn>
Thien Tran committed
Sep 28, 2026 at 14:18 UTC
469414504538a1bb7f3200f885109dc6adc3c805
2 files changed
+130
-2
.agents/skills/kernel-microbenchmark/SKILL.md
+23
-2
@@ -1,6 +1,6 @@
1
---
2
name: kernel-microbenchmark
3
-description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUDA, Triton, and CuteDSL, including CUPTI timing, correctness checks, generated-code inspection, multi-GPU measurements, and SOL sanity checks.
3
+description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUDA, ROCm/HIP, Triton, and CuteDSL, including CUPTI timing, correctness checks, generated-code inspection, multi-GPU measurements, and SOL sanity checks.
4
---
5
6
# Kernel Microbenchmark
@@ -15,7 +15,7 @@ description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUD
15
4. Compare against a baseline and report enough metadata to reproduce the
16
result: GPU, dtype, shape, command, branch/commit, and relevant env vars.
17
5. Treat explanations as hypotheses until backed by an artifact: ablation,
18
- generated PTX/SASS, profiler output, or controlled benchmark.
18
+ generated PTX/SASS or AMD ISA, profiler output, or controlled benchmark.
19
6. If the result changes the conclusion, preserve the compact lesson in a note,
20
comment, benchmark table, or final summary. For experiments, a short
21
`Question / Change / Correctness / Result / Observation / Next` note is
@@ -25,6 +25,12 @@ description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUD
25
26
- Use FlashInfer CUPTI timing by default, with CUDA graph and cold L2 cache:
27
`from flashinfer.testing import bench_gpu_time_with_cupti`.
28
+- ROCm: CUPTI and FlashInfer do not exist. Use HIP graph replay over rotating
29
+ buffers whose total footprint exceeds the last-level cache (e.g. 256 MB on
30
+ MI355X), so every call reads cold from HBM, timed with HIP events — see
31
+ [benchmarks/graph_replay_benchmark.py](benchmarks/graph_replay_benchmark.py)
32
+ (platform-agnostic). Prefer this over `do_bench`, whose per-rep
33
+ launch/event overhead inflates small/decode-size kernels.
34
- For compute-heavy kernels, report TFLOPS with the FLOP formula in the
35
benchmark. For memory-heavy kernels, report estimated bytes moved and GB/s.
36
For mixed kernels, report the most honest metric available and call out the
@@ -38,6 +44,11 @@ description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUD
44
formulas, skipped work, sparsity, caching, and whether the baseline is doing
45
the same operation.
46
- Force compilation/autotuning before measuring compiled kernels.
47
+- Tune every arm of a comparison before trusting it. A variant's default or
48
+ inherited config can be far from its optimum, so conclusions drawn from
49
+ untuned variants — especially negative ones — are unreliable. A coarse
50
+ sweep of the cheap knobs at a couple of representative shapes is enough;
51
+ report the winning config with the numbers.
52
- Seed inputs when correctness comparisons matter.
53
- Keep metadata setup, plan construction, allocation, random input generation,
54
and logging outside the timed region unless that overhead is the experiment.
@@ -98,11 +109,21 @@ throughput conventions.
109
a two-node TP run as equivalent to a local NVLink-domain run without checking.
110
- Stabilize GPU clocks or run enough untimed work to reach a steady state.
111
Alternate candidate order so clock, thermal, and rank-skew effects are shared.
112
+- ROCm: capturing ProcessGroupNCCL collectives in HIP graphs trips the NCCL
113
+ watchdog (HIP forbids querying an event recorded in a capturing stream;
114
+ the watchdog aborts the process). `TORCH_NCCL_BLOCKING_WAIT=1` before
115
+ `import torch` disables the watchdog — benchmark-only knob. Ask for backend
116
+ `"nccl"`; torch routes it to RCCL automatically (there is no `"rccl"`
117
+ backend string).
118
119
## Included Examples
120
121
- Use [benchmarks/cupti_microbenchmark.py](benchmarks/cupti_microbenchmark.py)
122
as a minimal single-GPU FlashInfer CUPTI timing pattern.
123
+- Use
124
+ [benchmarks/graph_replay_benchmark.py](benchmarks/graph_replay_benchmark.py)
125
+ as a graph-replay timing pattern over rotating buffers sized past the
126
+ last-level cache — the ROCm/HIP default, also usable on CUDA.
127
- Use
128
[benchmarks/multi_gpu_gemm_rs.py](benchmarks/multi_gpu_gemm_rs.py) as a
129
minimal distributed CUDA-graph timing pattern.
.agents/skills/kernel-microbenchmark/benchmarks/graph_replay_benchmark.py
new
+107
@@ -0,0 +1,107 @@
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()