flashinfer.testing.bench_gpu_time_with_cuda_event¶

flashinfer.testing.bench_gpu_time_with_cuda_event(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, input_args: Tuple = (), input_kwargs: dict | None = None, cold_l2_cache: bool = True)¶

使用 CUDA 事件（不使用 CUDA 图）基准测试内核执行时间。

这是最简单的基准测试方法。最适合内核启动开销与执行时间相比可以忽略不计的情况。

该函数执行：1. 快速估算阶段（5 次迭代）以确定迭代次数 2. 预热迭代（未测量） 3. 通过 CUDA 事件进行逐迭代计时的测量迭代

可以显式指定迭代次数，也可以从目标持续时间派生：- 如果提供了 dry_run_iters/repeat_iters，则直接使用这些计数。- 否则，计数将从 dry_run_time_ms/repeat_time_ms 计算得出。

参数:

fn (Callable) – 要基准测试的内核函数。
dry_run_iters (int, optional) – 热身迭代次数（未计时）。如果为 None，则从 dry_run_time_ms 计算得出。
repeat_iters (int, optional) – 测量迭代的次数。如果为 None，则从 repeat_time_ms 计算得出。
dry_run_time_ms (int) – 目标热身持续时间，以毫秒为单位（默认值：25）。
repeat_time_ms (int) – 目标测量持续时间，以毫秒为单位（默认值：100）。
sleep_after_run (bool) – 如果为 True，则在每次迭代后短暂休眠以减少热限制（默认值：False）。
input_args (tuple) – 传递给 fn 的位置参数。
input_kwargs (dict, optional) – 传递给 fn 的关键字参数。
cold_l2_cache (bool) – 如果为 True，则在每次迭代之前刷新 L2 缓存，以确保冷缓存性能测量（默认值：True）。

返回值:

毫秒级的逐迭代执行时间。

返回值类型:

List[float]

示例

基本用法

>>> def my_kernel(a, b):
...     return torch.matmul(a, b.T)
>>> q = torch.randn(1024, 128, device="cuda")
>>> k = torch.randn(1024, 128, device="cuda")
>>> times = bench_gpu_time_with_cuda_event(
...     fn=my_kernel,
...     input_args=(q, k),
... )
>>> print(f"Median time: {np.median(times):.3f} ms")

注意

此方法不使用 CUDA 图，因此每次迭代都会产生内核启动开销。对于启动延迟很重要的小型基准测试，请考虑使用 bench_gpu_time_with_cudagraph 代替。

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