flashinfer.testing.bench_gpu_time_with_cupti

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

使用 CUPTI 活动跟踪来精确测量内核时间，从而基准测试 GPU 时间。

CUPTI（CUDA 剖析工具接口）提供硬件级别的剖析，测量实际的 GPU 内核执行时间，不包括 CPU 端的启动开销。这提供了最准确的内核性能测量。

通过迭代之间的 L2 刷新来实现冷 L2 缓存。CUPTI 测量每次迭代，因此 L2 刷新无论 use_cuda_graph 如何，都能正常工作。

行为：- 使用 CUPTI（需要版本 >= 13，即 CUDA 13+）来跟踪内核活动

并从记录的开始/结束时间戳计算每次迭代的 GPU 时间。

  • 可以选择在 CUDA 图中捕获操作（use_cuda_graph=True），以减少测量期间的启动开销。

  • 如果 CUPTI 不可用，则回退到：- bench_gpu_time_with_cudagraph 如果 use_cuda_graph=True（使用旋转缓冲区

    用于冷 L2）

    • bench_gpu_time_with_cuda_event 否则（使用 L2 刷新用于冷 L2）

参数:

  • 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）。

  • use_cuda_graph (bool) – 如果为 True，则捕获并重放 CUDA 图（默认：False）。

  • input_args (tuple) – 要传递给 fn 的位置参数。

  • input_kwargs (dict, optional) – 要传递给 fn 的关键字参数。

  • cold_l2_cache (bool) – 如果为 True，则在每次迭代之前刷新 L2 缓存，以确保冷缓存性能测量（默认：True）。

返回值:

以毫秒为单位的每次迭代的 GPU 内核执行时间。

返回值类型:

List[float]

示例

基本的 CUPTI 基准测试（需要 cupti-python >= 13）

>>> 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_cupti(
...     fn=my_kernel,
...     input_args=(q, k),
... )
>>> print(f"Median GPU time: {np.median(times):.3f} ms")

注意

需要 cupti-python 包版本 >= 13.0.0：pip install -U cupti-python

如果 CUPTI 不可用，将发出警告，并且该函数会自动回退到 CUDA 事件或 CUDA 图计时。

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