@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5e379a361e

[Agents] Add kernel microbenchmark skill (#53688)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Codex <codex@openai.com>

Thien Tran committed Aug 25, 2026 at 15:25 UTC 5e379a361e3ea8bb82b7efd768c36f39a0cf32fd
5 files changed +484
.agents/skills/kernel-microbenchmark/SKILL.md new
+92
@@ -0,0 +1,92 @@
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.
4 +---
5 +
6 +# Kernel Microbenchmark
7 +
8 +## Workflow
9 +
10 +1. Create an isolated repro or benchmark when the existing harness is noisy.
11 +2. Check correctness before timing. Keep tolerances explicit.
12 +3. Time only the operation under study. Exclude allocation, compilation, random
13 + input generation, logging, and host-device transfers unless those are the
14 + target.
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.
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
22 + usually enough.
23 +
24 +## Benchmark Defaults
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 +- For compute-heavy kernels, report TFLOPS with the FLOP formula in the
29 + benchmark. For memory-heavy kernels, report estimated bytes moved and GB/s.
30 + For mixed kernels, report the most honest metric available and call out the
31 + caveats. TFLOPS and memory bandwidth should be computed from the theoretical
32 + best for the operation, not from a particular kernel implementation. For
33 + example, memory bandwidth should assume all data is read exactly once from
34 + global memory.
35 +- When comparing across shapes, prefer throughput metrics such as TFLOPS or
36 + GB/s as the primary table columns; keep latency for absolute cost.
37 +- If a result exceeds expected peak/SOL, first inspect units, FLOP/byte
38 + formulas, skipped work, sparsity, caching, and whether the baseline is doing
39 + the same operation.
40 +- Force compilation/autotuning before measuring compiled kernels.
41 +- Seed inputs when correctness comparisons matter.
42 +- Keep metadata setup, plan construction, allocation, random input generation,
43 + and logging outside the timed region unless that overhead is the experiment.
44 +
45 +## Multi-GPU Benchmarks
46 +
47 +- State whether the run is local or multi-node and report the GPU topology,
48 + world size, GPUs per node, collective backend, and relevant library versions.
49 +- Compare like-for-like TP configurations. Report both per-rank and global
50 + dimensions, and do not compare results from different TP sizes without an
51 + explicit normalization or scaling question.
52 +- Define the timed operation boundary before benchmarking. If the production
53 + wrapper performs input staging, flag resets, generation barriers, padding,
54 + or output copies, keep them in the timed region for an end-to-end comparison.
55 + Use a separate, clearly labeled ablation for kernel-only timing.
56 +- Check distributed correctness before timing. Seed each rank deliberately,
57 + form the reference with the same collective semantics, and synchronize before
58 + reading or comparing outputs.
59 +- Use a device-side barrier immediately before each measured replay. Keep this
60 + common synchronization outside the timed interval, but keep barriers required
61 + by the candidate implementation inside it.
62 +- With CUDA graphs, warm up before capture, coordinate capture across ranks,
63 + rotate pointer-distinct graphs when cache reuse matters, and ensure every
64 + collective is issued in the same order on every rank.
65 +- Measure every rank and reduce each sample with `MAX`; report the median of
66 + those per-sample maxima. A rank-local event time is not a distributed latency.
67 +- Reset reusable symmetric-memory flags before each invocation and establish a
68 + device-side generation barrier before peers may signal them. Otherwise a fast
69 + rank can signal before a slow rank resets its flags, causing a lost arrival
70 + and intermittent deadlock.
71 +- Preserve symmetric-memory handles, CUDA graphs, streams, and graph outputs for
72 + the full measurement lifetime. Allocate and rendezvous symmetric buffers in
73 + identical order and with identical shapes on all ranks.
74 +- Treat hangs and isolated millisecond outliers as synchronization bugs or rank
75 + skew until disproven. Add stage markers, bounded synchronization checks, and
76 + per-rank diagnostics before blaming compilation or kernel performance.
77 +- For multi-node runs, record the scheduler allocation and verify the fabric
78 + supports the required multicast or NVLink-domain assumptions. Do not describe
79 + a two-node TP run as equivalent to a local NVLink-domain run without checking.
80 +- Stabilize GPU clocks or run enough untimed work to reach a steady state.
81 + Alternate candidate order so clock, thermal, and rank-skew effects are shared.
82 +
83 +## Included Examples
84 +
85 +- Use [benchmarks/cupti_microbenchmark.py](benchmarks/cupti_microbenchmark.py)
86 + as a minimal single-GPU FlashInfer CUPTI timing pattern.
87 +- Use
88 + [benchmarks/multi_gpu_gemm_rs.py](benchmarks/multi_gpu_gemm_rs.py) as a
89 + minimal distributed CUDA-graph timing pattern.
90 +
91 +Adapt the operation, cases, work formula, and correctness tolerances rather
92 +than copying either example unchanged.
.agents/skills/kernel-microbenchmark/agents/openai.yaml new
+4
@@ -0,0 +1,4 @@
1 +interface:
2 + display_name: "Kernel Microbenchmark"
3 + short_description: "Build and debug reliable GPU kernel microbenchmarks."
4 + default_prompt: "Use the kernel microbenchmark workflow to build or debug a vLLM GPU kernel benchmark."
.agents/skills/kernel-microbenchmark/benchmarks/cupti_microbenchmark.py new
+63
@@ -0,0 +1,63 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""FlashInfer CUPTI microbenchmark template with throughput metrics."""
5 +
6 +import statistics
7 +
8 +import pandas as pd
9 +import torch
10 +from flashinfer.testing import bench_gpu_time_with_cupti
11 +
12 +WARMUP = 25
13 +MATMUL_CASES = [
14 + ("compute-bound", 4096, 4096, 4096),
15 + ("small-M", 16, 16384, 8192),
16 +]
17 +
18 +
19 +def bench_us(fn):
20 + for _ in range(WARMUP):
21 + fn()
22 + torch.accelerator.synchronize()
23 + return statistics.median(bench_gpu_time_with_cupti(fn)) * 1e3
24 +
25 +
26 +def main() -> None:
27 + if not torch.accelerator.is_available() or torch.version.cuda is None:
28 + raise RuntimeError("CUDA is required for CUPTI kernel timing.")
29 +
30 + torch.set_default_device("cuda")
31 + torch.manual_seed(0)
32 +
33 + rows = []
34 + for name, m, n, k in MATMUL_CASES:
35 + a = torch.randn(m, k, dtype=torch.bfloat16)
36 + b = torch.randn(k, n, dtype=torch.bfloat16)
37 + out = torch.empty(m, n, dtype=torch.bfloat16)
38 +
39 + def run_matmul(a=a, b=b, out=out):
40 + torch.mm(a, b, out=out)
41 +
42 + run_matmul()
43 + ref = torch.mm(a, b)
44 + torch.accelerator.synchronize()
45 + torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1)
46 +
47 + us = bench_us(run_matmul)
48 + rows.append(
49 + {
50 + "case": name,
51 + "shape": f"{m}x{n}x{k}",
52 + "us": us,
53 + "tflops": 2 * m * n * k / (us * 1e6),
54 + "gbps": 2 * (m * k + k * n + m * n) / (us * 1e3),
55 + }
56 + )
57 +
58 + df = pd.DataFrame(rows)
59 + print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}"))
60 +
61 +
62 +if __name__ == "__main__":
63 + main()
.agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py new
+324
@@ -0,0 +1,324 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Minimal multi-GPU GEMM + reduce-scatter microbenchmark.
5 +
6 +Run on one node with, for example:
7 +
8 + torchrun --standalone --nproc-per-node=8 \
9 + .agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py
10 +"""
11 +
12 +import argparse
13 +import os
14 +import statistics
15 +from collections.abc import Callable
16 +
17 +import pandas as pd
18 +import torch
19 +import torch.distributed as dist
20 +import torch.distributed._symmetric_memory as symm_mem
21 +
22 +from vllm.config import VllmConfig, set_current_vllm_config
23 +from vllm.distributed import cleanup_dist_env_and_memory
24 +from vllm.distributed.parallel_state import (
25 + get_tp_group,
26 + init_distributed_environment,
27 + initialize_model_parallel,
28 +)
29 +
30 +
31 +def parse_args() -> argparse.Namespace:
32 + parser = argparse.ArgumentParser()
33 + parser.add_argument("--m", type=int, nargs="+", default=[128, 512, 2048])
34 + parser.add_argument("--n", type=int, default=4096)
35 + parser.add_argument(
36 + "--k",
37 + type=int,
38 + nargs="+",
39 + default=[4096],
40 + help="Per-rank K values",
41 + )
42 + parser.add_argument("--num-workspaces", type=int, default=10)
43 + parser.add_argument("--warmup-replays", type=int, default=5)
44 + parser.add_argument("--samples", type=int, default=20)
45 + return parser.parse_args()
46 +
47 +
48 +def make_gemm_rs(
49 + x: torch.Tensor,
50 + weight: torch.Tensor,
51 + partial: torch.Tensor,
52 + output: torch.Tensor,
53 + rows: int,
54 + device_group: dist.ProcessGroup,
55 +) -> Callable[[], None]:
56 + def run() -> None:
57 + torch.mm(x, weight.T, out=partial[:rows])
58 + dist.reduce_scatter_single(output, partial, group=device_group)
59 +
60 + return run
61 +
62 +
63 +def check_correctness(
64 + runs: dict[str, Callable[[], None]],
65 + outputs: dict[str, torch.Tensor],
66 + x: torch.Tensor,
67 + weight: torch.Tensor,
68 + padded_rows: int,
69 + rank: int,
70 + device_group: dist.ProcessGroup,
71 +) -> None:
72 + expected_full = torch.zeros(
73 + (padded_rows, weight.shape[0]),
74 + dtype=x.dtype,
75 + device=x.device,
76 + )
77 + torch.mm(x, weight.T, out=expected_full[: x.shape[0]])
78 + dist.all_reduce(expected_full, group=device_group)
79 + expected = expected_full.chunk(dist.get_world_size(device_group))[rank]
80 + for name, run in runs.items():
81 + run()
82 + torch.accelerator.synchronize()
83 + torch.testing.assert_close(
84 + outputs[name],
85 + expected,
86 + rtol=5e-2,
87 + atol=4.0,
88 + )
89 +
90 +
91 +def capture_graph(
92 + run: Callable[[], None],
93 + cpu_group: dist.ProcessGroup,
94 +) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]:
95 + stream = torch.cuda.Stream()
96 + stream.wait_stream(torch.cuda.current_stream())
97 + dist.barrier(group=cpu_group)
98 + with torch.cuda.stream(stream):
99 + for _ in range(3):
100 + run()
101 + stream.synchronize()
102 + dist.barrier(group=cpu_group)
103 +
104 + graph = torch.cuda.CUDAGraph()
105 + with torch.cuda.graph(graph, stream=stream):
106 + run()
107 + torch.cuda.current_stream().wait_stream(stream)
108 + dist.barrier(group=cpu_group)
109 + return graph, stream
110 +
111 +
112 +def benchmark_graphs(
113 + candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
114 + warmup_replays: int,
115 + samples: int,
116 + device_group: dist.ProcessGroup,
117 + device_barrier: Callable[[], None],
118 +) -> dict[str, float]:
119 + candidate_names = list(candidate_graphs)
120 + for round_index in range(warmup_replays):
121 + for candidate_index in range(len(candidate_names)):
122 + candidate_id = (round_index + candidate_index) % len(candidate_names)
123 + name = candidate_names[candidate_id]
124 + graphs = candidate_graphs[name]
125 + device_barrier()
126 + graphs[round_index % len(graphs)].replay()
127 + torch.accelerator.synchronize()
128 +
129 + timings: dict[str, list[float]] = {name: [] for name in candidate_names}
130 + start = torch.cuda.Event(enable_timing=True)
131 + end = torch.cuda.Event(enable_timing=True)
132 + for sample_index in range(samples):
133 + for candidate_index in range(len(candidate_names)):
134 + candidate_id = (sample_index + candidate_index) % len(candidate_names)
135 + name = candidate_names[candidate_id]
136 + graphs = candidate_graphs[name]
137 + device_barrier()
138 + start.record()
139 + graphs[sample_index % len(graphs)].replay()
140 + end.record()
141 + end.synchronize()
142 +
143 + elapsed_us = torch.tensor(
144 + start.elapsed_time(end) * 1000,
145 + dtype=torch.float64,
146 + device=torch.accelerator.current_device_index(),
147 + )
148 + dist.all_reduce(elapsed_us, op=dist.ReduceOp.MAX, group=device_group)
149 + timings[name].append(elapsed_us.item())
150 + return {name: statistics.median(values) for name, values in timings.items()}
151 +
152 +
153 +def benchmark_shape(
154 + m: int,
155 + n: int,
156 + k: int,
157 + num_workspaces: int,
158 + warmup_replays: int,
159 + samples: int,
160 + device: torch.device,
161 + rank: int,
162 + world_size: int,
163 + device_group: dist.ProcessGroup,
164 + cpu_group: dist.ProcessGroup,
165 + device_barrier: Callable[[], None],
166 +) -> dict[str, float | int]:
167 + padded_m = (m + world_size - 1) // world_size * world_size
168 + local_m = padded_m // world_size
169 +
170 + torch.manual_seed(1000 + rank * 10 + m + k)
171 + inputs = []
172 + weights = []
173 + for _ in range(num_workspaces):
174 + inputs.append(torch.randn(m, k, dtype=torch.bfloat16, device=device))
175 + weights.append(torch.randn(n, k, dtype=torch.bfloat16, device=device))
176 +
177 + ring_partial = torch.empty(padded_m, n, dtype=torch.bfloat16, device=device)
178 + ldmc_partial = symm_mem.empty(
179 + (padded_m, n),
180 + dtype=torch.bfloat16,
181 + device=device,
182 + )
183 + ldmc_handle = symm_mem.rendezvous(ldmc_partial, device_group)
184 + ring_output = torch.empty(local_m, n, dtype=torch.bfloat16, device=device)
185 + ldmc_output = torch.empty_like(ring_output)
186 + if padded_m > m:
187 + ring_partial[m:].zero_()
188 + ldmc_partial[m:].zero_()
189 +
190 + candidate_runs = {
191 + "ring_ll_us": [
192 + make_gemm_rs(
193 + x,
194 + weight,
195 + ring_partial,
196 + ring_output,
197 + m,
198 + device_group,
199 + )
200 + for x, weight in zip(inputs, weights)
201 + ],
202 + "ldmc_us": [
203 + make_gemm_rs(
204 + x,
205 + weight,
206 + ldmc_partial,
207 + ldmc_output,
208 + m,
209 + device_group,
210 + )
211 + for x, weight in zip(inputs, weights)
212 + ],
213 + }
214 + x = inputs[0]
215 + weight = weights[0]
216 + check_correctness(
217 + {name: runs[0] for name, runs in candidate_runs.items()},
218 + {"ring_ll_us": ring_output, "ldmc_us": ldmc_output},
219 + x,
220 + weight,
221 + padded_m,
222 + rank,
223 + device_group,
224 + )
225 +
226 + candidate_graphs = {}
227 + graph_keepalive: list[object] = [ldmc_handle]
228 + for name, runs in candidate_runs.items():
229 + bundles = [capture_graph(run, cpu_group) for run in runs]
230 + candidate_graphs[name] = [graph for graph, _ in bundles]
231 + graph_keepalive.extend(bundles)
232 +
233 + times = benchmark_graphs(
234 + candidate_graphs,
235 + warmup_replays,
236 + samples,
237 + device_group,
238 + device_barrier,
239 + )
240 + global_flops = 2 * m * n * k * world_size
241 + return {
242 + "M": m,
243 + "N": n,
244 + "K_per_rank": k,
245 + "K_global": k * world_size,
246 + **times,
247 + "ring_ll_tflops": global_flops / (times["ring_ll_us"] * 1e6),
248 + "ldmc_tflops": global_flops / (times["ldmc_us"] * 1e6),
249 + "ldmc_speedup": times["ring_ll_us"] / times["ldmc_us"],
250 + }
251 +
252 +
253 +def main() -> None:
254 + args = parse_args()
255 + assert args.m and min(args.m) > 0
256 + assert args.k and min(args.k) > 0
257 + assert min(args.n, args.num_workspaces, args.samples) > 0
258 + assert args.warmup_replays >= 0
259 + local_rank = int(os.environ["LOCAL_RANK"])
260 + local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
261 + torch.accelerator.set_device_index(local_rank)
262 + init_distributed_environment()
263 + world_size = dist.get_world_size()
264 + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
265 + symm_mem.set_backend("NCCL")
266 + with set_current_vllm_config(VllmConfig()):
267 + initialize_model_parallel(tensor_model_parallel_size=world_size)
268 +
269 + tp_group = get_tp_group()
270 + device_group = tp_group.device_group
271 + cpu_group = tp_group.cpu_group
272 + rank = tp_group.rank_in_group
273 + device = torch.device("cuda", local_rank)
274 + group_warmup = torch.zeros(1, device=device)
275 + dist.all_reduce(group_warmup, group=device_group)
276 + pynccl_comm = tp_group.device_communicator.pynccl_comm
277 + assert pynccl_comm is not None
278 + sync_input = torch.zeros(1, device=device)
279 + sync_output = torch.empty_like(sync_input)
280 +
281 + def device_barrier() -> None:
282 + # Order the timed launch after a device-side rank rendezvous without
283 + # including the rendezvous itself in the measured event interval.
284 + pynccl_comm.all_reduce(sync_input, sync_output)
285 +
286 + results = [
287 + benchmark_shape(
288 + m,
289 + args.n,
290 + k,
291 + args.num_workspaces,
292 + args.warmup_replays,
293 + args.samples,
294 + device,
295 + rank,
296 + world_size,
297 + device_group,
298 + cpu_group,
299 + device_barrier,
300 + )
301 + for k in args.k
302 + for m in args.m
303 + ]
304 +
305 + if rank == 0:
306 + metadata = {
307 + "world_size": world_size,
308 + "local_world_size": local_world_size,
309 + "num_nodes": world_size // local_world_size,
310 + "backend": dist.get_backend(device_group),
311 + "gpu": torch.cuda.get_device_name(local_rank),
312 + "torch": torch.__version__,
313 + "cuda": torch.version.cuda,
314 + }
315 + print(pd.Series(metadata, name="value").to_string())
316 + df = pd.DataFrame(results)
317 + print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}"))
318 +
319 + dist.barrier(group=cpu_group)
320 + cleanup_dist_env_and_memory()
321 +
322 +
323 +if __name__ == "__main__":
324 + main()
.claude/skills/kernel-microbenchmark new
+1
@@ -0,0 +1 @@
1 +../../.agents/skills/kernel-microbenchmark
\ No newline at end of file