Forward-port the tkv/turbo-attn vLLM seam onto upstream v0.31.0

.agents/skills/debug-ima/SKILL.md

@@ -0,0 +1,9 @@
+---
+name: debug-ima
+description: Debug CUDA illegal memory access (IMA) errors in vLLM with CUDA core dumps and cuda-gdb.
+---
+
+# Debug IMA
+
+Follow the vLLM guide
+[CUDA Core Dump: An Effective Tool to Debug Memory Access Issues and Beyond](https://vllm.ai/blog/2025-08-11-cuda-debugging).

.agents/skills/debug-ima/agents/openai.yaml

@@ -0,0 +1,4 @@
+interface:
+  display_name: "Debug IMA"
+  short_description: "Debug CUDA illegal memory access errors in vLLM."
+  default_prompt: "Use the vLLM CUDA core-dump guide to debug this illegal memory access error."

.agents/skills/kernel-microbenchmark/SKILL.md

@@ -0,0 +1,132 @@
+---
+name: kernel-microbenchmark
+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
+
+## 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 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
+   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`.
+- 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
+  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.
+- 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.
+
+## Sanity-Check Reference Numbers
+
+Use these as rough reference points for large, well-shaped workloads, not as
+gold standards, guaranteed peaks, or hard limits. Hardware SKU, clocks, shape,
+precision conventions, and the FLOP/byte accounting can move the result. A
+large gap is a prompt to investigate, not proof that a kernel is poor.
+
+| Kernel regime | Hardware | Rough reference |
+| --- | --- | ---: |
+| Memory-bound, large batch | Blackwell | 6 TB/s |
+| BF16 GEMM | Blackwell | 2 PFLOP/s |
+| BF16 attention | B200 | 1.6 PFLOP/s |
+| FP8 GEMM | Blackwell | 4 PFLOP/s |
+| FP8 attention | B300 | 2.8 PFLOP/s |
+
+The BF16 attention reference is approximately the 1613 TFLOP/s result reported
+by the FlashAttention-4 paper. Compare kernels only with matching workload and
+throughput conventions.
+
+## 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.
+- 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.
+
+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/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()