flashinfer.testing.bench_gpu_time_with_cudagraph

flashinfer.testing.bench_gpu_time_with_cudagraph(fn, dry_run_iters: int = None, repeat_iters: int = None, dry_run_time_ms: int = 25, repeat_time_ms: int = 100, num_iters_within_graph: int = 10, 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 图和摊销的内核启动开销来基准测试 GPU 时间。

CUDA 图捕获一系列 GPU 操作并以最少的 CPU 开销重放它们。通过在单个图中运行多个迭代，内核启动延迟被摊销，从而得到更接近纯 GPU 时间的测量结果。

冷 L2 基准测试:

当 cold_l2_cache=True 时，该函数使用 **旋转缓冲区** 来确保图中每个内核调用的冷 L2 缓存。在 input_args / input_kwargs 中 GPU 张量的多个副本在捕获图时被创建和轮换，确保每个内核调用都在不同的内存区域上操作。缓冲区副本的数量是根据设备 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）。

  • num_iters_within_graph (int) – 图中捕获的内核调用次数（默认值：10）。较高的值可以更好地摊销启动开销，但在轮换缓冲区时会使用更多的内存。

  • sleep_after_run (bool) – 如果为 True，则在每次迭代后短暂休眠（默认值：False）。

  • input_args (tuple) – 传递给 fn 的位置参数。当 cold_l2_cache=True 时，此结构中的 GPU 张量将被克隆。

  • input_kwargs (dict, optional) – 传递给 fn 的关键字参数。当 cold_l2_cache=True 时，此结构中的 GPU 张量将被克隆。

  • cold_l2_cache (bool) – 如果为 True，则使用旋转缓冲区来确保图中每个内核调用的冷 L2 缓存（默认值：True）。

返回值:

毫秒级的每迭代执行时间。每个时间是图重放持续时间除以 num_iters_within_graph。

返回值类型:

List[float]

示例

冷 L2 基准测试（默认，用于受内存限制的内核）

>>> def run_attention(q, k, v, o):
...     flashinfer.single_prefill_with_kv_cache(q, k, v, o)
...
>>> q = torch.randn(batch, heads, seq_len, head_dim, device="cuda")
>>> k = torch.randn(batch, heads, seq_len, head_dim, device="cuda")
>>> v = torch.randn(batch, heads, seq_len, head_dim, device="cuda")
>>> o = torch.empty_like(q)
>>> times = bench_gpu_time_with_cudagraph(
...     fn=run_attention,
...     input_args=(q, k, v, o),
... )
>>> print(f"Cold-L2 median time: {np.median(times):.3f} ms")

示例

热 L2 基准测试（用于受计算限制的内核）

>>> times = bench_gpu_time_with_cudagraph(
...     fn=lambda: torch.matmul(q, k.T),
...     cold_l2_cache=False,
... )

注意

  • 当使用 input_args / input_kwargs 时，该函数必须接受张量作为参数（而不是从闭包中捕获它们）。

  • GPU 张量会自动检测并克隆。非张量参数（标量、布尔值等）在所有副本中都会被保留。

  • 内存使用量随着需要超过 L2 缓存的旋转次数而增加。

参见

  • calculate_rotation_count：计算冷 L2 所需的缓冲区副本。

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