autotune

[Truth-Ke]

No description provided.

[Truth-Ke]

Ascend 自动调优使用说明

本文档面向算子使用者，说明如何用 Ascend autotune 自动搜索 tile block 参数和编译参数。

1. 最简单配置

先在定义 @triton.autotune kernel 之前导入 hook：

import triton
import triton.language as tl
import triton.backends.dicp_triton.ascend_autotune_hooks  # noqa: F401

然后声明需要搜索的 tile block 参数，例如 BLOCK_M / BLOCK_N：

ATTENTION_SEARCH_HINTS = {
    "search_params": {
        "params": ["BLOCK_M", "BLOCK_N"],
    },
}

在 kernel 上使用：

@triton.autotune(
    configs=[],
    key=["N_CTX", "HEAD_DIM"],
    hints=ATTENTION_SEARCH_HINTS,
)
@triton.jit
def attention_kernel(
    Q, K, V, Out,
    N_CTX: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    ...

这里推荐写 configs=[]，表示不手工枚举候选 config，而是交给 Ascend autotune 自动生成并搜索 BLOCK_M / BLOCK_N 以及编译参数。

key=["N_CTX", "HEAD_DIM"] 是可选的，但推荐写。key 用来区分不同输入 shape 的调优结果；相同 key 会复用 cache，不同 key 会重新调优。如果不写 key，调优器仍然可以运行，但不同 shape 可能复用同一个调优结果，性能不一定最优。

这就是最小推荐用法。它会自动做两类搜索：

• tile block 参数：例如 BLOCK_M、BLOCK_N
• Ascend 编译参数：由 operator classifier 推断 mixcv 或 vector profile；FA/matmul 这类 tl.dot regular kernel 会搜索 workspace / tile mix / multibuffer 相关 profile，vector kernel 只搜索 vector-safe profile

search_params 会自动推断 compile-profile family，不需要再写 compile_options="mixcv"。

2. 调用方式

调用 kernel 时不需要传 BLOCK_M / BLOCK_N，调优器会选择：

grid = lambda meta: (
    triton.cdiv(N_CTX, meta["BLOCK_M"]),
)

attention_kernel[grid](
    Q, K, V, Out,
    N_CTX,
    HEAD_DIM,
)

第一次遇到新的 key=["N_CTX", "HEAD_DIM"] 会运行调优；同一进程内后续相同 key 会直接复用已选出的 best config。

3. 可选：控制 block 搜索范围

默认搜索值为：

[16, 32, 64, 128, 256, 512, 1024, 2048]

如果想缩小范围，可以指定 values：

ATTENTION_SEARCH_HINTS = {
    "search_params": {
        "params": ["BLOCK_M", "BLOCK_N"],
        "values": [32, 64, 128, 256],
    },
}

也可以调整最终进入编译参数搜索的 shape 数量：

ATTENTION_SEARCH_HINTS = {
    "search_params": {
        "params": ["BLOCK_M", "BLOCK_N"],
        "shape_final_top_k": 3,
    },
}

4. 可选：打开调试输出

Python 中打开 search debug：

ATTENTION_SEARCH_HINTS = {
    "search_params": {
        "params": ["BLOCK_M", "BLOCK_N"],
        "debug": True,
    },
}

环境变量打开 autotune 结果和 timing 输出：

export TRITON_PRINT_AUTOTUNING=1
export TRITON_PRINT_AUTOTUNING_TIMINGS=1

5. 可选：手写候选 config

如果不想自动搜索 block 参数，也可以手写候选 triton.Config：

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}),
    ],
    key=["N_CTX", "HEAD_DIM"],
)
@triton.jit
def attention_kernel(..., BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    ...

这种模式只在给定候选里选最快；推荐优先使用 search_params 自动搜索。
如果 configs 中的每个 triton.Config 都显式提供了全部
search_params.params，则手写 configs 优先，search_params 的 block
搜索会被忽略。

6. 注意事项

• key 必须包含会影响最优配置的 shape 或 meta 参数，例如 N_CTX、HEAD_DIM。
• search_params.params 中的名字必须是 kernel 的 tl.constexpr 参数。
• 自动搜索模式不需要手动传 BLOCK_M / BLOCK_N；手写 configs 模式由
  triton.Config 提供这些参数。
• 不要把 unit_flag 放进 search_params；调优器会按内部策略处理。
• search_params.params 必须是性能 tile，不要放覆盖语义、atomic 粒度、状态切分或手工 UB-safe 常量。典型反例包括 ROW_BLOCK、SEQ_BLOCK、state/chunk 分块、只验证过固定值的 CBLOCK/BK/BV/LARGE_BLOCK_T。这些参数要么保持固定，要么提供 reference_fn 做 Stage 3 correctness gate。
• 如果运行时固定传入 set_workspace_multibuffer、tile_mix_*、num_stages 等 compile options，仍然走 search_params 新路径；固定项会覆盖调优器生成的 compile profile，并参与 autotune cache key。
• autotune cache 是进程内 cache；同一进程内相同 key 会复用结果。