triton.autotune

triton.autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, do_bench=None, cache_results=False)

用于自动调优 triton.jit 函数的装饰器。

@triton.autotune(configs=[
    triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
    triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
  ],
  key=['x_size'] # the two above configs will be evaluated anytime
                 # the value of x_size changes
)
@triton.jit
def kernel(x_ptr, x_size, BLOCK_SIZE: tl.constexpr):
    ...
注意:

当评估所有配置时，内核会多次运行。这意味着内核更新的任何值都会被更新多次。为了避免这种不期望的行为，您可以使用 reset_to_zero 参数，它会在运行任何配置之前将所提供张量的值重置为 zero。

如果环境变量 TRITON_PRINT_AUTOTUNING 设置为 "1"，Triton 将在自动调优每个内核后向 stdout 打印一条消息，包括自动调优所花费的时间和最佳配置。

参数:

  • configs (list[triton.Config]) – triton.Config 对象的列表。

  • key (list[str]) – 参数名称列表，这些参数的值发生变化将触发所有提供配置的评估。

  • prune_configs_by –

    一个包含用于剪枝配置的函数的字典，字段包括：‘perf_model’：用于根据不同配置预测运行时间的性能模型，返回运行时间；‘top_k’：用于基准测试的配置数量；‘early_config_prune’：用于剪枝配置的函数。它应具有以下签名：

    prune_configs_by( configs: List[triton.Config], named_args: Dict[str, Any], **kwargs: Dict[str, Any]) -> List[triton.Config]: 并返回剪枝后的配置。它至少应该返回一个配置。

  • reset_to_zero (list[str]) – 参数名称列表，这些参数的值在评估任何配置之前将被重置为零。

  • restore_value (list[str]) – 参数名称列表，这些参数的值在评估任何配置后将被恢复。

  • pre_hook (lambda args, reset_only) – 在调用内核之前将调用的函数。这将覆盖用于 ‘reset_to_zero’ 和 ‘restore_value’ 的默认 pre_hook。‘kwargs’：传递给内核的所有参数的字典。‘reset_only’：一个布尔值，指示是否仅为了重置值（且没有对应的 post_hook）而调用 pre_hook。

  • post_hook (lambda args, exception) – 在调用内核之后将调用的函数。这将覆盖用于 ‘restore_value’ 的默认 post_hook。‘kwargs’：传递给内核的所有参数的字典。‘exception’：如果发生编译或运行时错误，内核引发的异常。

  • warmup (int) – 传递给基准测试的预热时间（以毫秒为单位）（已弃用）。

  • rep (int) – 传递给基准测试的重复时间（以毫秒为单位）（已弃用）。

  • do_bench (lambda fn, quantiles) – 用于测量每次运行时间的基准测试函数。

  • cache_results – 是否将自动调优时间结果缓存到磁盘。默认为 False。

“type cache_results: bool