diff --git a/.agents/skills/kernel-microbenchmark/SKILL.md b/.agents/skills/kernel-microbenchmark/SKILL.md new file mode 100644 index 0000000000..545a79f9fa --- /dev/null +++ b/.agents/skills/kernel-microbenchmark/SKILL.md @@ -0,0 +1,92 @@ +--- +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. +--- + +# Kernel Microbenchmark + +## Workflow + +1. Create an isolated repro or benchmark when the existing harness is noisy. +2. Check correctness before timing. Keep tolerances explicit. +3. Time only the operation under study. Exclude allocation, compilation, random + input generation, logging, and host-device transfers unless those are the + target. +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. +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 + usually enough. + +## Benchmark Defaults + +- Use FlashInfer CUPTI timing by default, with CUDA graph and cold L2 cache: + `from flashinfer.testing import bench_gpu_time_with_cupti`. +- 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 + caveats. TFLOPS and memory bandwidth should be computed from the theoretical + best for the operation, not from a particular kernel implementation. For + example, memory bandwidth should assume all data is read exactly once from + global memory. +- When comparing across shapes, prefer throughput metrics such as TFLOPS or + GB/s as the primary table columns; keep latency for absolute cost. +- If a result exceeds expected peak/SOL, first inspect units, FLOP/byte + formulas, skipped work, sparsity, caching, and whether the baseline is doing + the same operation. +- Force compilation/autotuning before measuring compiled kernels. +- 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. + +## Multi-GPU Benchmarks + +- State whether the run is local or multi-node and report the GPU topology, + world size, GPUs per node, collective backend, and relevant library versions. +- Compare like-for-like TP configurations. Report both per-rank and global + dimensions, and do not compare results from different TP sizes without an + explicit normalization or scaling question. +- Define the timed operation boundary before benchmarking. If the production + wrapper performs input staging, flag resets, generation barriers, padding, + or output copies, keep them in the timed region for an end-to-end comparison. + Use a separate, clearly labeled ablation for kernel-only timing. +- Check distributed correctness before timing. Seed each rank deliberately, + form the reference with the same collective semantics, and synchronize before + reading or comparing outputs. +- Use a device-side barrier immediately before each measured replay. Keep this + common synchronization outside the timed interval, but keep barriers required + by the candidate implementation inside it. +- With CUDA graphs, warm up before capture, coordinate capture across ranks, + rotate pointer-distinct graphs when cache reuse matters, and ensure every + collective is issued in the same order on every rank. +- Measure every rank and reduce each sample with `MAX`; report the median of + those per-sample maxima. A rank-local event time is not a distributed latency. +- Reset reusable symmetric-memory flags before each invocation and establish a + device-side generation barrier before peers may signal them. Otherwise a fast + rank can signal before a slow rank resets its flags, causing a lost arrival + and intermittent deadlock. +- Preserve symmetric-memory handles, CUDA graphs, streams, and graph outputs for + the full measurement lifetime. Allocate and rendezvous symmetric buffers in + identical order and with identical shapes on all ranks. +- Treat hangs and isolated millisecond outliers as synchronization bugs or rank + skew until disproven. Add stage markers, bounded synchronization checks, and + per-rank diagnostics before blaming compilation or kernel performance. +- For multi-node runs, record the scheduler allocation and verify the fabric + supports the required multicast or NVLink-domain assumptions. Do not describe + 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. + +## Included Examples + +- Use [benchmarks/cupti_microbenchmark.py](benchmarks/cupti_microbenchmark.py) + as a minimal single-GPU FlashInfer CUPTI timing pattern. +- Use + [benchmarks/multi_gpu_gemm_rs.py](benchmarks/multi_gpu_gemm_rs.py) as a + minimal distributed CUDA-graph timing pattern. + +Adapt the operation, cases, work formula, and correctness tolerances rather +than copying either example unchanged. diff --git a/.agents/skills/kernel-microbenchmark/agents/openai.yaml b/.agents/skills/kernel-microbenchmark/agents/openai.yaml new file mode 100644 index 0000000000..cbd52cc600 --- /dev/null +++ b/.agents/skills/kernel-microbenchmark/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Kernel Microbenchmark" + short_description: "Build and debug reliable GPU kernel microbenchmarks." + default_prompt: "Use the kernel microbenchmark workflow to build or debug a vLLM GPU kernel benchmark." diff --git a/.agents/skills/kernel-microbenchmark/benchmarks/cupti_microbenchmark.py b/.agents/skills/kernel-microbenchmark/benchmarks/cupti_microbenchmark.py new file mode 100644 index 0000000000..74eb6119b4 --- /dev/null +++ b/.agents/skills/kernel-microbenchmark/benchmarks/cupti_microbenchmark.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""FlashInfer CUPTI microbenchmark template with throughput metrics.""" + +import statistics + +import pandas as pd +import torch +from flashinfer.testing import bench_gpu_time_with_cupti + +WARMUP = 25 +MATMUL_CASES = [ + ("compute-bound", 4096, 4096, 4096), + ("small-M", 16, 16384, 8192), +] + + +def bench_us(fn): + for _ in range(WARMUP): + fn() + torch.accelerator.synchronize() + return statistics.median(bench_gpu_time_with_cupti(fn)) * 1e3 + + +def main() -> None: + if not torch.accelerator.is_available() or torch.version.cuda is None: + raise RuntimeError("CUDA is required for CUPTI kernel timing.") + + 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) + b = torch.randn(k, n, dtype=torch.bfloat16) + out = torch.empty(m, n, dtype=torch.bfloat16) + + def run_matmul(a=a, b=b, out=out): + torch.mm(a, b, out=out) + + run_matmul() + ref = torch.mm(a, b) + torch.accelerator.synchronize() + torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1) + + us = bench_us(run_matmul) + rows.append( + { + "case": name, + "shape": f"{m}x{n}x{k}", + "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() diff --git a/.agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py b/.agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py new file mode 100644 index 0000000000..56b503a297 --- /dev/null +++ b/.agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Minimal multi-GPU GEMM + reduce-scatter microbenchmark. + +Run on one node with, for example: + + torchrun --standalone --nproc-per-node=8 \ + .agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py +""" + +import argparse +import os +import statistics +from collections.abc import Callable + +import pandas as pd +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.distributed.parallel_state import ( + get_tp_group, + init_distributed_environment, + initialize_model_parallel, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--m", type=int, nargs="+", default=[128, 512, 2048]) + parser.add_argument("--n", type=int, default=4096) + parser.add_argument( + "--k", + type=int, + nargs="+", + default=[4096], + help="Per-rank K values", + ) + parser.add_argument("--num-workspaces", type=int, default=10) + parser.add_argument("--warmup-replays", type=int, default=5) + parser.add_argument("--samples", type=int, default=20) + return parser.parse_args() + + +def make_gemm_rs( + x: torch.Tensor, + weight: torch.Tensor, + partial: torch.Tensor, + output: torch.Tensor, + rows: int, + device_group: dist.ProcessGroup, +) -> Callable[[], None]: + def run() -> None: + torch.mm(x, weight.T, out=partial[:rows]) + dist.reduce_scatter_single(output, partial, group=device_group) + + return run + + +def check_correctness( + runs: dict[str, Callable[[], None]], + outputs: dict[str, torch.Tensor], + x: torch.Tensor, + weight: torch.Tensor, + padded_rows: int, + rank: int, + device_group: dist.ProcessGroup, +) -> None: + expected_full = torch.zeros( + (padded_rows, weight.shape[0]), + dtype=x.dtype, + device=x.device, + ) + torch.mm(x, weight.T, out=expected_full[: x.shape[0]]) + dist.all_reduce(expected_full, group=device_group) + expected = expected_full.chunk(dist.get_world_size(device_group))[rank] + for name, run in runs.items(): + run() + torch.accelerator.synchronize() + torch.testing.assert_close( + outputs[name], + expected, + rtol=5e-2, + atol=4.0, + ) + + +def capture_graph( + run: Callable[[], None], + cpu_group: dist.ProcessGroup, +) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + dist.barrier(group=cpu_group) + with torch.cuda.stream(stream): + for _ in range(3): + run() + stream.synchronize() + dist.barrier(group=cpu_group) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run() + torch.cuda.current_stream().wait_stream(stream) + dist.barrier(group=cpu_group) + return graph, stream + + +def benchmark_graphs( + candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]], + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + device_barrier: Callable[[], None], +) -> dict[str, float]: + candidate_names = list(candidate_graphs) + for round_index in range(warmup_replays): + for candidate_index in range(len(candidate_names)): + candidate_id = (round_index + candidate_index) % len(candidate_names) + name = candidate_names[candidate_id] + graphs = candidate_graphs[name] + device_barrier() + graphs[round_index % len(graphs)].replay() + torch.accelerator.synchronize() + + timings: dict[str, list[float]] = {name: [] for name in candidate_names} + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for sample_index in range(samples): + for candidate_index in range(len(candidate_names)): + candidate_id = (sample_index + candidate_index) % len(candidate_names) + name = candidate_names[candidate_id] + graphs = candidate_graphs[name] + device_barrier() + start.record() + graphs[sample_index % len(graphs)].replay() + end.record() + end.synchronize() + + elapsed_us = torch.tensor( + start.elapsed_time(end) * 1000, + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(elapsed_us, op=dist.ReduceOp.MAX, group=device_group) + timings[name].append(elapsed_us.item()) + return {name: statistics.median(values) for name, values in timings.items()} + + +def benchmark_shape( + m: int, + n: int, + k: int, + num_workspaces: int, + warmup_replays: int, + samples: int, + device: torch.device, + rank: int, + world_size: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, + device_barrier: Callable[[], None], +) -> dict[str, float | int]: + padded_m = (m + world_size - 1) // world_size * world_size + local_m = padded_m // world_size + + torch.manual_seed(1000 + rank * 10 + m + k) + inputs = [] + weights = [] + for _ in range(num_workspaces): + inputs.append(torch.randn(m, k, dtype=torch.bfloat16, device=device)) + weights.append(torch.randn(n, k, dtype=torch.bfloat16, device=device)) + + ring_partial = torch.empty(padded_m, n, dtype=torch.bfloat16, device=device) + ldmc_partial = symm_mem.empty( + (padded_m, n), + dtype=torch.bfloat16, + device=device, + ) + ldmc_handle = symm_mem.rendezvous(ldmc_partial, device_group) + ring_output = torch.empty(local_m, n, dtype=torch.bfloat16, device=device) + ldmc_output = torch.empty_like(ring_output) + if padded_m > m: + ring_partial[m:].zero_() + ldmc_partial[m:].zero_() + + candidate_runs = { + "ring_ll_us": [ + make_gemm_rs( + x, + weight, + ring_partial, + ring_output, + m, + device_group, + ) + for x, weight in zip(inputs, weights) + ], + "ldmc_us": [ + make_gemm_rs( + x, + weight, + ldmc_partial, + ldmc_output, + m, + device_group, + ) + for x, weight in zip(inputs, weights) + ], + } + x = inputs[0] + weight = weights[0] + check_correctness( + {name: runs[0] for name, runs in candidate_runs.items()}, + {"ring_ll_us": ring_output, "ldmc_us": ldmc_output}, + x, + weight, + padded_m, + rank, + device_group, + ) + + candidate_graphs = {} + graph_keepalive: list[object] = [ldmc_handle] + for name, runs in candidate_runs.items(): + bundles = [capture_graph(run, cpu_group) for run in runs] + candidate_graphs[name] = [graph for graph, _ in bundles] + graph_keepalive.extend(bundles) + + times = benchmark_graphs( + candidate_graphs, + warmup_replays, + samples, + device_group, + device_barrier, + ) + global_flops = 2 * m * n * k * world_size + return { + "M": m, + "N": n, + "K_per_rank": k, + "K_global": k * world_size, + **times, + "ring_ll_tflops": global_flops / (times["ring_ll_us"] * 1e6), + "ldmc_tflops": global_flops / (times["ldmc_us"] * 1e6), + "ldmc_speedup": times["ring_ll_us"] / times["ldmc_us"], + } + + +def main() -> None: + args = parse_args() + assert args.m and min(args.m) > 0 + assert args.k and min(args.k) > 0 + assert min(args.n, args.num_workspaces, args.samples) > 0 + assert args.warmup_replays >= 0 + local_rank = int(os.environ["LOCAL_RANK"]) + local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + torch.accelerator.set_device_index(local_rank) + init_distributed_environment() + world_size = dist.get_world_size() + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" + symm_mem.set_backend("NCCL") + with set_current_vllm_config(VllmConfig()): + initialize_model_parallel(tensor_model_parallel_size=world_size) + + tp_group = get_tp_group() + device_group = tp_group.device_group + cpu_group = tp_group.cpu_group + rank = tp_group.rank_in_group + device = torch.device("cuda", local_rank) + group_warmup = torch.zeros(1, device=device) + dist.all_reduce(group_warmup, group=device_group) + pynccl_comm = tp_group.device_communicator.pynccl_comm + assert pynccl_comm is not None + sync_input = torch.zeros(1, device=device) + sync_output = torch.empty_like(sync_input) + + def device_barrier() -> None: + # Order the timed launch after a device-side rank rendezvous without + # including the rendezvous itself in the measured event interval. + pynccl_comm.all_reduce(sync_input, sync_output) + + results = [ + benchmark_shape( + m, + args.n, + k, + args.num_workspaces, + args.warmup_replays, + args.samples, + device, + rank, + world_size, + device_group, + cpu_group, + device_barrier, + ) + for k in args.k + for m in args.m + ] + + if rank == 0: + metadata = { + "world_size": world_size, + "local_world_size": local_world_size, + "num_nodes": world_size // local_world_size, + "backend": dist.get_backend(device_group), + "gpu": torch.cuda.get_device_name(local_rank), + "torch": torch.__version__, + "cuda": torch.version.cuda, + } + print(pd.Series(metadata, name="value").to_string()) + df = pd.DataFrame(results) + print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}")) + + dist.barrier(group=cpu_group) + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/.claude/skills/kernel-microbenchmark b/.claude/skills/kernel-microbenchmark new file mode 120000 index 0000000000..c0233605e8 --- /dev/null +++ b/.claude/skills/kernel-microbenchmark @@ -0,0 +1 @@ +../../.agents/skills/kernel-microbenchmark \ No newline at end of file