[Agents] Add kernel microbenchmark skill (#53688)

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

.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.

.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."

.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()

.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()