flashinfer.testing.bench_gpu_time¶

flashinfer.testing.bench_gpu_time(fn, dry_run_iters： int = None, repeat_iters： int = None, dry_run_time_ms： int = 25, repeat_time_ms： int = 100, l2_flush： bool | None = None, l2_flush_size_mb： int | None = None, l2_flush_device： str | None = None, sleep_after_run： bool = False, enable_cupti： bool = False, use_cuda_graph： bool = False, num_iters_within_graph： int = 10, input_args： Tuple = (), input_kwargs： dict | None = None, cold_l2_cache： bool = True)¶

具有可配置计时后端的统一 GPU 基准测试接口。

这是 GPU 内核基准测试的推荐入口点。它提供了一个单一的接口，根据配置标志分派到适当的计时实现。

计时后端（按优先级顺序）

CUPTI (enable_cupti=True)：最准确，通过硬件分析测量纯 GPU 内核时间。需要 cupti-python >= 13。
CUDA 图形 (use_cuda_graph=True)：通过捕获和重放多个内核调用来摊销启动开销。在准确性和可用性之间取得了良好的平衡。
CUDA 事件（默认）：最简单的方法，测量启动 + 执行。在所有地方都可用，但包含 CPU 开销。

冷 L2 策略（根据计时后端自动选择）

返回值:: 毫秒级的逐次迭代执行时间。
返回值类型:: List[float]

示例

使用 CUDA 事件（默认）进行简单的基准测试

>>> times = bench_gpu_time(fn=lambda: my_kernel())
>>> print(f"Median: {np.median(times):.3f} ms")

示例

使用 CUDA 图形进行基准测试，以减少启动开销

>>> def run_kernel(x, y, out):
...     my_memory_bound_kernel(x, y, out)
>>> times = bench_gpu_time(
...     fn=run_kernel,
...     input_args=(x, y, out),
...     use_cuda_graph=True,
... )

示例

使用 CUPTI 进行基准测试，以获得最准确的 GPU 内核时间

>>> times = bench_gpu_time(
...     fn=run_kernel,
...     input_args=(x, y, out),
...     enable_cupti=True,
... )

参见

bench_gpu_time_with_cuda_event：直接 CUDA 事件计时。
bench_gpu_time_with_cudagraph：直接 CUDA 图形计时。
bench_gpu_time_with_cupti：直接 CUPTI 计时。

自版本以来已弃用：l2_flush、l2_flush_size_mb 和 l2_flush_device 参数已弃用。请改用 cold_l2_cache。