[MIRROR][Feature Request] Use rotating buffers for in benchmark functions

[yzh119]

Copied from flashinfer-ai#2187

flashinfer.testing.bench_gpu_time_with_cupti today supports flushing the L2 cache before each run, but bench_gpu_time_with_cuda_event does not.

This means that only in environments where cupti-python is installed we are able to microbenchmark performance with a cold L2, which is not widespread today (requires CUDA 13).

It would be helpful to implement a rotating buffer using input[round % N] and output[round % N] where N could be set according to L2 cache size to use a cold L2.

[yzh119]

@claude

[claude]

Claude finished @yzh119's task —— View job • claude/issue-8-20251208-2226 • Create PR ➔

---

Implementation Complete ✓

I've successfully implemented rotating buffer support for the benchmark functions as requested in issue #8.

Tasks

• Read repository setup instructions (CLAUDE.md)
• Understand current benchmark implementation
• Design rotating buffer implementation approach
• Implement rotating buffer support for bench_gpu_time_with_cuda_event
• Implement rotating buffer support for bench_gpu_time_with_cudagraph
• Update bench_gpu_time wrapper function
• Commit and push changes
• Provide PR link

Changes Made

Added rotating_buffer_size parameter to the benchmark functions to enable cold L2 cache benchmarking without requiring CUPTI (CUDA 13+):

1. bench_gpu_time_with_cuda_event (flashinfer/testing/utils.py:551)

   • Added rotating_buffer_size parameter (default 0 = disabled)
   • When enabled, passes iteration index to the benchmark function
   • Users can implement fn(iter_idx) and use input[iter_idx % N] pattern

2. bench_gpu_time_with_cudagraph (flashinfer/testing/utils.py:898)

   • Added rotating_buffer_size parameter
   • Includes warning that CUDA graphs have limited compatibility with rotating buffers
   • CUDA graphs capture specific tensor addresses, so rotation within graphs is constrained

3. bench_gpu_time (flashinfer/testing/utils.py:1038)

   • Updated wrapper to pass through rotating_buffer_size parameter

How to Use

# Example: Create multiple input/output buffers
N = 10  # Number of buffer copies (size based on L2 cache size)
inputs = [torch.randn(...) for _ in range(N)]
outputs = [torch.empty(...) for _ in range(N)]

def benchmark_fn(iter_idx):
    # Use rotating buffers to ensure cold L2 cache
    result = my_kernel(inputs[iter_idx], outputs[iter_idx])
    return result

# Benchmark with rotating buffers
times = bench_gpu_time_with_cuda_event(
    benchmark_fn,
    rotating_buffer_size=N,
    l2_flush=False  # Can disable buffer.zero_() if using rotating buffers
)

Design Notes

• Backward compatible: When rotating_buffer_size=0 (default), functions work as before
• Flexible: Works alongside existing l2_flush mechanism (can use both or either)
• CUDA graph limitation: Rotating buffers have limited support with CUDA graphs since graphs capture specific operations. A warning is issued when combining these features.

---

• Branch: claude/issue-8-20251208-2226