From 469414504538a1bb7f3200f885109dc6adc3c805 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Mon, 28 Sep 2026 14:18:38 +0800 Subject: [PATCH] [Skills] Update kernel-microbenchmark to include ROCm (#58646) Signed-off-by: Thien Tran Co-authored-by: Kimi --- .agents/skills/kernel-microbenchmark/SKILL.md | 25 +++- .../benchmarks/graph_replay_benchmark.py | 107 ++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) create mode 100644 .agents/skills/kernel-microbenchmark/benchmarks/graph_replay_benchmark.py diff --git a/.agents/skills/kernel-microbenchmark/SKILL.md b/.agents/skills/kernel-microbenchmark/SKILL.md index f2dd5be929..98542023ff 100644 --- a/.agents/skills/kernel-microbenchmark/SKILL.md +++ b/.agents/skills/kernel-microbenchmark/SKILL.md @@ -1,6 +1,6 @@ --- name: kernel-microbenchmark -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. +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. --- # Kernel Microbenchmark @@ -15,7 +15,7 @@ description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUD 4. Compare against a baseline and report enough metadata to reproduce the result: GPU, dtype, shape, command, branch/commit, and relevant env vars. 5. Treat explanations as hypotheses until backed by an artifact: ablation, - generated PTX/SASS, profiler output, or controlled benchmark. + generated PTX/SASS or AMD ISA, profiler output, or controlled benchmark. 6. If the result changes the conclusion, preserve the compact lesson in a note, comment, benchmark table, or final summary. For experiments, a short `Question / Change / Correctness / Result / Observation / Next` note is @@ -25,6 +25,12 @@ description: Build, debug, and interpret vLLM GPU kernel microbenchmarks for CUD - Use FlashInfer CUPTI timing by default, with CUDA graph and cold L2 cache: `from flashinfer.testing import bench_gpu_time_with_cupti`. +- ROCm: CUPTI and FlashInfer do not exist. Use HIP graph replay over rotating + buffers whose total footprint exceeds the last-level cache (e.g. 256 MB on + MI355X), so every call reads cold from HBM, timed with HIP events — see + [benchmarks/graph_replay_benchmark.py](benchmarks/graph_replay_benchmark.py) + (platform-agnostic). Prefer this over `do_bench`, whose per-rep + launch/event overhead inflates small/decode-size kernels. - For compute-heavy kernels, report TFLOPS with the FLOP formula in the benchmark. For memory-heavy kernels, report estimated bytes moved and GB/s. 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 formulas, skipped work, sparsity, caching, and whether the baseline is doing the same operation. - Force compilation/autotuning before measuring compiled kernels. +- Tune every arm of a comparison before trusting it. A variant's default or + inherited config can be far from its optimum, so conclusions drawn from + untuned variants — especially negative ones — are unreliable. A coarse + sweep of the cheap knobs at a couple of representative shapes is enough; + report the winning config with the numbers. - Seed inputs when correctness comparisons matter. - Keep metadata setup, plan construction, allocation, random input generation, and logging outside the timed region unless that overhead is the experiment. @@ -98,11 +109,21 @@ throughput conventions. a two-node TP run as equivalent to a local NVLink-domain run without checking. - Stabilize GPU clocks or run enough untimed work to reach a steady state. Alternate candidate order so clock, thermal, and rank-skew effects are shared. +- ROCm: capturing ProcessGroupNCCL collectives in HIP graphs trips the NCCL + watchdog (HIP forbids querying an event recorded in a capturing stream; + the watchdog aborts the process). `TORCH_NCCL_BLOCKING_WAIT=1` before + `import torch` disables the watchdog — benchmark-only knob. Ask for backend + `"nccl"`; torch routes it to RCCL automatically (there is no `"rccl"` + backend string). ## Included Examples - Use [benchmarks/cupti_microbenchmark.py](benchmarks/cupti_microbenchmark.py) as a minimal single-GPU FlashInfer CUPTI timing pattern. +- Use + [benchmarks/graph_replay_benchmark.py](benchmarks/graph_replay_benchmark.py) + as a graph-replay timing pattern over rotating buffers sized past the + last-level cache — the ROCm/HIP default, also usable on CUDA. - Use [benchmarks/multi_gpu_gemm_rs.py](benchmarks/multi_gpu_gemm_rs.py) as a minimal distributed CUDA-graph timing pattern. diff --git a/.agents/skills/kernel-microbenchmark/benchmarks/graph_replay_benchmark.py b/.agents/skills/kernel-microbenchmark/benchmarks/graph_replay_benchmark.py new file mode 100644 index 0000000000..ec051ba789 --- /dev/null +++ b/.agents/skills/kernel-microbenchmark/benchmarks/graph_replay_benchmark.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Graph-replay microbenchmark with rotating weight buffers (CUDA and ROCm). + +Platform-agnostic: HIP implements the torch.cuda.* API, including CUDAGraph. +Pattern: capture one graph (HIP graph on ROCm) containing R*INNER calls +cycling over R rotating weight buffers, then time replays with CUDA/HIP +events. The rotating buffers keep the captured footprint above the GPU's +last-level cache, so every call reads weights cold from HBM. Prefer this over +do_bench on ROCm: do_bench's per-rep launch/event overhead inflates +small/decode-size kernels, and its default flush buffer is only as large as +the LLC — on a 256 MB LLC that evicts marginally and leaves the cache full of +dirty flush-buffer lines. The same pattern extends to multi-GPU benchmarks +(capture per rank, collectives in identical order every replay); see the +Multi-GPU section of SKILL.md. +""" + +import statistics + +import pandas as pd +import torch + +LLC_TARGET_BYTES = 400 * 1024**2 # above 256 MB MI355X LLC and B200 L2 (126 MB) +INNER = 2 # repeat slots within one graph to amortize per-replay overhead +REPLAYS = 10 # replays per round: makes each measurement large vs overhead +ROUNDS = 5 # median across rounds: replay time is not always stable +WARMUP = 5 # warmup replays to reach steady clocks before timing +MATMUL_CASES = [ + ("compute-bound", 4096, 4096, 4096), + ("decode-GEMM", 16, 16384, 8192), +] + + +def timeit_graph(calls, inner=INNER, replays=REPLAYS, rounds=ROUNDS): + """Time a list of per-slot callables as one graph (zero launch overhead); + weights rotate across slots so they stay HBM-cold. Returns the median over + rounds of per-round means, ms per call.""" + for c in calls: + c() + torch.accelerator.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): # capture runs on a side stream + for i in range(len(calls) * inner): + calls[i % len(calls)]() + torch.accelerator.synchronize() + events = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(rounds)] + for _ in range(WARMUP): + g.replay() + # No sync between warmup and the first round: queued replays keep the GPU + # busy, so each round's launch latency stays outside its event window. + for s, e in events: + s.record() + for _ in range(replays): + g.replay() + e.record() + torch.accelerator.synchronize() + samples = [s.elapsed_time(e) / (replays * len(calls) * inner) for s, e in events] + return statistics.median(samples) # ms per call + + +def main() -> None: + if not torch.accelerator.is_available(): + raise RuntimeError("No GPU visible (torch.accelerator covers HIP on ROCm).") + + torch.set_default_device("cuda") + torch.manual_seed(0) + + rows = [] + for name, m, n, k in MATMUL_CASES: + a = torch.randn(m, k, dtype=torch.bfloat16) # shared: producer-warm + out = torch.empty(m, n, dtype=torch.bfloat16) + b_bytes = k * n * 2 + r = max(2, -(-LLC_TARGET_BYTES // b_bytes)) # ceil + bs = [torch.randn(k, n, dtype=torch.bfloat16) for _ in range(r)] + + def make_call(i, a=a, bs=bs, out=out): + def call(): + torch.mm(a, bs[i], out=out) + + return call + + # correctness before timing, per slot (the graph would freeze a wrong + # result into every replay) + for i in range(r): + make_call(i)() + torch.accelerator.synchronize() + torch.testing.assert_close(out, torch.mm(a, bs[i]), atol=1e-1, rtol=1e-1) + + us = timeit_graph([make_call(i) for i in range(r)]) * 1e3 + rows.append( + { + "case": name, + "shape": f"{m}x{n}x{k}", + "slots": r, + "us": us, + "tflops": 2 * m * n * k / (us * 1e6), + "gbps": 2 * (m * k + k * n + m * n) / (us * 1e3), + } + ) + + df = pd.DataFrame(rows) + print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}")) + + +if __name__ == "__main__": + main()