CuTe-DSL kernels are not captured into torch.cuda.graph, causing silent under-measurement (and stale-output correctness bypass) in flashinfer-bench

[sha7doww]

TL;DR

@cute.kernel .launch() (from nvidia-cutlass-dsl) does not participate in
CUDA graph capture — calling it inside a with torch.cuda.graph(g): block
runs the kernel immediately during the capture, but it is not recorded
into the graph. All subsequent g.replay() calls silently skip it.

Combined with flashinfer-bench's benchmarking protocol (fixed inputs per
workload, ≥100 measured iterations of graph.replay()), any solution that
wraps Triton-anchor + @cute.kernel in torch.cuda.graph ends up with:

• The CuTe kernel runs once (during the with cuda.graph(...) block,
  which does NOT record it).
• All 100+ measured iterations replay only the Triton anchor.
• The output buffer keeps the stale value from the one-time CuTe launch →
  because benchmark inputs don't change across iterations, correctness
  passes against the reference.
• CUPTI reports latency of the Triton anchor only → reported speedup is
  inflated by whatever fraction of the real per-call work was in the
  skipped CuTe kernel(s).

This is not a CUPTI bug. CUPTI faithfully measures what the graph
executes — which simply doesn't include the CuTe kernels. The bug is the
CuTe-DSL ↔ CUDA-graph interaction, with flashinfer-bench silently
benefiting.

Minimal reproduction

Standalone, no flashinfer-bench harness required. Pattern:

import torch, cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack

@cute.kernel
def write_pattern(X: cute.Tensor):
    tid, _, _ = cute.arch.thread_idx()
    X[tid] = cutlass.Float32(3.14)

@cute.jit
def launch(X):
    write_pattern(X).launch(grid=(1,1,1), block=(64,1,1))

buf = torch.zeros(64, dtype=torch.float32, device="cuda")

# Prime (JIT compile)
launch(from_dlpack(buf, assumed_align=16))
torch.cuda.synchronize()

# Capture
buf.zero_()
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    launch(from_dlpack(buf, assumed_align=16))
torch.cuda.synchronize()
print("post-capture  :", buf.abs().max().item())   # 3.14  (ran IMMEDIATELY)

# Replay
buf.zero_()
torch.cuda.synchronize()
g.replay()
torch.cuda.synchronize()
print("after replay  :", buf.abs().max().item())   # 0.00  (NOT in graph)

When the @cute.kernel is captured alone, PyTorch prints UserWarning: The CUDA Graph is empty. When paired with a Triton kernel in the same capture
block, the warning is absent (the graph is non-empty — it contains the
Triton kernel). The CuTe kernel is still silently dropped, but there's no
longer any signal alerting the user.

Three independent tests in flashinfer-bench-style harness

Given any solution that launches _launch(...) which does triton_fwd(q, ..., PO) followed by cute_reduce(PO, ..., OUT) inside torch.cuda.graph:

T1 — Graph-replay populates output?

output.zero_()
with torch.cuda.graph(g):
    _launch(...)                         # runs BOTH kernels immediately
assert output.abs().max() > 0            # ✓ during capture block
output.zero_()
g.replay()
assert output.abs().max() > 0            # ✗ FAILS — reduce wasn't recorded

T2 — Does replay rewrite poisoned cells?

# After graph is cached, poison one cell of the persistent output tensor.
output[0,0,0] = -1e9
run_fn(*inputs)                          # triggers graph.replay()
# If reduce were in the graph, this cell would be rewritten.
# If not, -1e9 survives — stale.

T3 — Different inputs → different outputs?

# Reuse same tensor addresses, change contents in place.
for i in range(5):
    q.normal_(mean=float(i))
    outs.append(run_fn(q, ...).clone())
# If reduce is in the graph, outputs reflect per-iter q; if not, they are
# byte-identical across all 5 iterations.

On Modal B200 / CUDA 13.2, a solution using the above pattern:

• T1 ✗ — replay leaves output all-zeros.
• T2 ✗ — poisoned cell survives across 100+ replays.
• T3 ✗ — 5 very different q tensors produce byte-identical outputs.

Timing consequences (CUPTI vs Event vs Wall-clock)

Comparing two solutions with identical Triton fwd but different reduce
implementations, back-to-back in the same Modal container (no session
drift), same workload:

graph contents	CUPTI	Event	Wall
Triton fwd + @cute.kernel reduce (reduce skipped on replay)	4.5 µs	6.2 µs	6.2 µs
Triton fwd + Triton reduce (both in graph)	9.5 µs	10.9 µs	10.9 µs

The ~5 µs gap is the reduce-kernel GPU time the CuTe row silently
skips. Both rows pass correctness against the same reference. The
all-Triton row's Event/CUPTI ratio is 1.15× (normal instrumentation
overhead); the CuTe row's ratio is 1.37× — slightly higher because
the per-replay graph is shorter and Event's fixed overhead dominates a
larger fraction, not a CUPTI-specific signal.

Environment

• GPU: NVIDIA Blackwell B200 (sm_100) on Modal. Also reproduces on H100.
• Docker: flashinfer/flashinfer-ci-cu132:20260401-2c675fb (CUDA 13.2,
  PyTorch 2.12.0+cu132, Triton 3.6.0).
• nvidia-cutlass-dsl: ≥4.3.4 (pip install --no-deps).
• flashinfer-bench: f7b4d8d185625ab2d609233a1a06e99ee18a0c6b.

Suggested fixes

1. One-shot sanity gate in flashinfer-bench. After warmup but
   before measurement, do one output.zero_() → runner() →
   torch.cuda.synchronize() → assert output.abs().sum() > 0. This
   catches any graph replay that fails to populate its own output
   (this bug, and more generally the Bug: Kernels can cheat the benchmark by caching their outputs. #228 class) with ~1 replay of
   overhead.

2. Upstream fix in nvidia-cutlass-dsl. @cute.kernel.launch()
   should respect the active stream's capture mode — either participate
   in graph capture (preferred) or raise a clear error when called
   inside torch.cuda.graph(). This is the root cause; feat: implement Tracer #1 above is a
   detection guardrail while that lands.

Note: flashinfer.testing.bench_gpu_time_with_cupti itself is not
miscounting — it faithfully reports what the replayed graph actually
executes. A pure-Triton graph on the same workload measures CUPTI and
Event within ~15% of each other (normal instrumentation overhead), so
the CUPTI correlation-ID algorithm is working correctly; there is
simply no CuTe-DSL kernel activity for it to count because the kernel
never got into the graph in the first place.

Related issues

• do_bench: per-iteration torch.cuda.synchronize() inflates runtimes for fast kernels #195 (closed): do_bench: per-iteration torch.cuda.synchronize()
  inflates runtimes — different flashinfer_bench.bench.timing accuracy
  issue.
• Bug: Kernels can cheat the benchmark by caching their outputs. #228 (open): Kernels can cheat the benchmark by caching their
  outputs — related cheat vector (outputs cached across iterations); the
  issue here is that the cache happens silently at the CuTe/graph-capture
  layer, not in user code.