From 7fcbf1f26255d66c20dbd729eff2d31fa6fad596 Mon Sep 17 00:00:00 2001 From: wwwbby Date: Wed, 8 Apr 2026 14:16:09 +0800 Subject: [PATCH 1/9] [skill]add core partition optimizer md --- .../references/core_partition.md | 149 ++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 skills/latency-optimizer/references/core_partition.md diff --git a/skills/latency-optimizer/references/core_partition.md b/skills/latency-optimizer/references/core_partition.md new file mode 100644 index 00000000..1b768375 --- /dev/null +++ b/skills/latency-optimizer/references/core_partition.md @@ -0,0 +1,149 @@ +# Core Partition 分核优化模式 + +## 概述 + +在 Triton NPU kernel 中,**分核策略直接影响硬件利用率和性能**。Ascend 910B4 有 20 个 AI Core(共 40 个 VEC),选择合适的核数、分核维度和任务分配方式是性能优化的关键。 + +## 触发条件 + +**当 Triton 代码中存在以下情况时,应考虑优化分核策略**: + +1. **核数不合理**:grid 大小与数据规模不匹配(过多或过少) +2. **负载不均衡**:某些核空闲或处理数据量差异大 +3. **极端数据形状**:M << N 或 N << M 的场景 +4. **小数据量**:数据量 < 100KB 时使用过多核 +5. **UB 溢出风险**:tile 大小超过 192KB UB 容量 + +## 硬件约束 + +| 硬件资源 | 数量 | 说明 | +| -------- | --------- | --------------------------------------- | +| AI Core | 20 | 每个 AI Core 包含 2 VEC + 1 CUBE + 1 SU | +| VEC | 40 | 向量计算单元,可并行执行 | +| UB | 192KB/VEC | Unified Buffer,每个 VEC 独立 | + +## 优化方法 + +### 1. 核数选择优化 + +#### 原始代码(核数过多) + +```python +# 小数据量使用过多核 +N = 1024 +BLOCK_SIZE = 64 +grid = (triton.cdiv(N, BLOCK_SIZE),) # grid = 16,调度开销大 +``` + +#### 优化后代码(合理核数) + +```python +# 小数据量使用少量核 +N = 1024 +BLOCK_SIZE = 256 +grid = (triton.cdiv(N, BLOCK_SIZE),) # grid = 4,调度开销合理 +``` + +### 2. Reduction 分核维度优化 + +#### 原始代码(非 reduce 轴分核,负载不均) + +```python +# 数据形状:(16, 262144),沿 axis=1 归约 +M, N = 16, 262144 +BLOCK_M = 16 +grid = (triton.cdiv(M, BLOCK_M),) # grid = 1,只用 1 个核! +``` + +#### 优化后代码(reduce 轴分核 + 原子操作) + +```python +# 数据形状:(16, 262144),沿 axis=1 归约 +M, N = 16, 262144 +BLOCK_N = 8192 +SUB_BLOCK_N = 1024 # 二次切分避免 UB 溢出 +grid = (triton.cdiv(N, BLOCK_N),) # grid = 32,充分利用多核 + +@triton.jit +def reduce_min_kernel(x_ptr, y_ptr, M, N, BLOCK_N: tl.constexpr, SUB_BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + n_start = pid * BLOCK_N + + row_min = tl.full([M], float('inf'), dtype=tl.float32) + + # 二次切分 + for sub_start in range(0, BLOCK_N, SUB_BLOCK_N): + n_offsets = n_start + sub_start + tl.arange(0, SUB_BLOCK_N) + mask = n_offsets < N + x = tl.load(x_ptr + tl.arange(0, M)[:, None] * N + n_offsets[None, :], mask=mask) + local_min = tl.min(x, axis=1) + row_min = tl.minimum(row_min, local_min) + + # 原子操作跨核归约 + for m in range(M): + tl.atomic_min(y_ptr + m, row_min[m]) +``` + +### 3. 任务分配优化(非均匀分配) + +#### 原始代码(简单除法,边界处理不当) + +```python +n_cores = 32 +total_tasks = 1000 +tasks_per_core = total_tasks // n_cores # 31,剩余 8 个任务无人处理 + +# Kernel 内 +pid = tl.program_id(0) +start = pid * tasks_per_core +end = start + tasks_per_core # 最后 8 个任务丢失! +``` + +#### 优化后代码(base + rem 模式) + +```python +n_cores = 32 +total_tasks = 1000 +base = total_tasks // n_cores # 31 +rem = total_tasks % n_cores # 8 + +# Kernel 内 +pid = tl.program_id(0) +if pid < rem: + task_start = pid * (base + 1) + task_end = task_start + (base + 1) +else: + task_start = rem * (base + 1) + (pid - rem) * base + task_end = task_start + base +# 总计:8*32 + 24*31 = 256 + 744 = 1000 ✓ +``` + +## 核数选择建议 + +| 数据规模 | 推荐核数 | 理由 | +| ----------- | -------- | -------------- | +| < 1KB | 1 | 调度开销主导 | +| 1KB - 100KB | 1-8 | 平衡调度和计算 | +| 100KB - 1MB | 8-16 | 开始受益于并行 | +| > 1MB | 16-40 | 充分利用硬件 | + +## 分核维度决策 + +| 算子类型 | 数据特征 | 分核维度 | +| ----------- | -------- | --------------------------- | +| Elementwise | 任意 | 按数据量均分 | +| Reduction | M >> N | 按非 reduce 轴分核 | +| Reduction | M << N | 按 reduce 轴分核 + 原子操作 | +| MatMul | 通用 | 按 M-N 维度 2D 分核 | + +## 性能收益 + +- **核数优化**:小数据量场景可提升 2-5x 性能 +- **分核维度优化**:极端形状场景可提升 2-3x 性能 +- **负载均衡**:避免部分核空闲,提升整体吞吐 + +## 注意事项 + +1. **UB 容量约束**:确保 tile_size * dtype_size * buffers <= 192KB +2. **原子操作开销**:原子操作有额外开销,核数过多时可能成为瓶颈 +3. **对齐要求**:数据传输需 256B 对齐 From 48ce789a28731db7c64c4c37b3c5b7867a75e617 Mon Sep 17 00:00:00 2001 From: wwwbby Date: Wed, 8 Apr 2026 14:31:11 +0800 Subject: [PATCH 2/9] register Core Partition skill --- skills/latency-optimizer/SKILL.md | 1 + 1 file changed, 1 insertion(+) diff --git a/skills/latency-optimizer/SKILL.md b/skills/latency-optimizer/SKILL.md index 7c911c0b..b401ca71 100644 --- a/skills/latency-optimizer/SKILL.md +++ b/skills/latency-optimizer/SKILL.md @@ -28,6 +28,7 @@ argument-hint: > | 入参静态化 | `references/constexpr_parameters.md` | | Int32 向量加法 | `references/int32_vector_add.md` | | Load 指令重排序 | `references/load-order.md` | +| 分核策略优化 | `references/core_partition.md` | ### 以下文档通过分析已有代码特征,按需加载 From c91bed3e2e8b1450c616cd36de69f988c09f92ad Mon Sep 17 00:00:00 2001 From: wwwbby <60009003+wwwbby@users.noreply.github.com> Date: Wed, 8 Apr 2026 20:57:48 +0800 Subject: [PATCH 3/9] Update core_partition.md --- .../references/core_partition.md | 545 +++++++++++++++--- 1 file changed, 450 insertions(+), 95 deletions(-) diff --git a/skills/latency-optimizer/references/core_partition.md b/skills/latency-optimizer/references/core_partition.md index 1b768375..657c227b 100644 --- a/skills/latency-optimizer/references/core_partition.md +++ b/skills/latency-optimizer/references/core_partition.md @@ -2,148 +2,503 @@ ## 概述 -在 Triton NPU kernel 中,**分核策略直接影响硬件利用率和性能**。Ascend 910B4 有 20 个 AI Core(共 40 个 VEC),选择合适的核数、分核维度和任务分配方式是性能优化的关键。 +在 Triton NPU kernel 中,**分核策略直接影响硬件利用率和性能**。NPU设备有多个AI Core,选择合适的核数、分核维度和任务分配方式是性能优化的关键。 ## 触发条件 **当 Triton 代码中存在以下情况时,应考虑优化分核策略**: -1. **核数不合理**:grid 大小与数据规模不匹配(过多或过少) -2. **负载不均衡**:某些核空闲或处理数据量差异大 -3. **极端数据形状**:M << N 或 N << M 的场景 -4. **小数据量**:数据量 < 100KB 时使用过多核 -5. **UB 溢出风险**:tile 大小超过 192KB UB 容量 - -## 硬件约束 - -| 硬件资源 | 数量 | 说明 | -| -------- | --------- | --------------------------------------- | -| AI Core | 20 | 每个 AI Core 包含 2 VEC + 1 CUBE + 1 SU | -| VEC | 40 | 向量计算单元,可并行执行 | -| UB | 192KB/VEC | Unified Buffer,每个 VEC 独立 | +1. **发射核数不合理**:grid 大小与数据规模不匹配(过多或过少),npu 设备的物理核数一般为40或48,当grid 大小远大于物理核时或者远小于物理核时将会使得性能极大地弱化。 +2. **tiling大小不合理**:在For循环中调用Vector计算单元时,运算的tile数据量远小于当前设备的UB大小(通常是192KB),导致无法充分使用算力单元。 ## 优化方法 -### 1. 核数选择优化 +### 直接固定发射核数等于设备核数 -#### 原始代码(核数过多) +#### 原始代码(发射核数过多) ```python -# 小数据量使用过多核 -N = 1024 -BLOCK_SIZE = 64 -grid = (triton.cdiv(N, BLOCK_SIZE),) # grid = 16,调度开销大 +@triton.jit +def gather_dim1_kernel( + x_ptr, # *x [B, C] + idx_ptr, # *idx[B, K] + out_ptr, # *out[B, K] + stride_xb, stride_xc, + stride_ib, stride_ik, + stride_ob, stride_ok, + B, K, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_b = tl.program_id(0) # 1 block per batch row + pid_k = tl.program_id(1) # 1 block per K-tile + k_off = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + mask = k_off < K + idx = tl.load(idx_ptr + pid_b * stride_ib + k_off * stride_ik, mask=mask) # [BLOCK_K] + x_val = tl.load(x_ptr + pid_b * stride_xb + idx * stride_xc, mask=mask) + tl.store(out_ptr + pid_b * stride_ob + k_off * stride_ok, x_val, mask=mask) + +# 调用 +B = 128 # batch dim +K = 64 + +BLOCK_B = 4 +BLOCK_K = 128 + +grid = (B, triton.cdiv(K, BLOCK_K)) + +gather_dim1_kernel[grid]( + x, idx, out, + x.stride(0), x.stride(1), + idx.stride(0), idx.stride(1), + out.stride(0), out.stride(1), + B, K, + BLOCK_B=BLOCK_B, + BLOCK_K=BLOCK_K, +) ``` #### 优化后代码(合理核数) ```python -# 小数据量使用少量核 -N = 1024 -BLOCK_SIZE = 256 -grid = (triton.cdiv(N, BLOCK_SIZE),) # grid = 4,调度开销合理 +@triton.jit +def gather_dim1_kernel( + x_ptr, # *x [B, C] + idx_ptr, # *idx[B, K] + out_ptr, # *out[B, K] + stride_xb, stride_xc, + stride_ib, stride_ik, + stride_ob, stride_ok, + B, K, + BLOCK_B: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_b = tl.program_id(0) # 1 block per batch row +- # 原始实现 +- pid_k = tl.program_id(1) # 1 block per K-tile + +- k_off = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) +- mask = k_off < K + +- idx = tl.load(idx_ptr + pid_b * stride_ib + k_off * stride_ik, mask=mask) # [BLOCK_K] + +- x_val = tl.load(x_ptr + pid_b * stride_xb + idx * stride_xc, mask=mask) + +- tl.store(out_ptr + pid_b * stride_ob + k_off * stride_ok, x_val, mask=mask) + ++ # 优化后实现使用向量化处理,一次处理一整个BLOCK_B,因此这里的得到的是一个向量 ++ b_idx = pid_b * BLOCK_B + tl.arange(0, BLOCK_B) ++ b_mask = b_idx < B # 需要判断是否越界 + ++ # 对 K 维进行循环,向量化处理BLOCK_B * BLOCK_K个数据 ++ for k_start in range(0, K, BLOCK_K): ++ ks = tl.arange(0, BLOCK_K) ++ k_mask = ks < K - k_start + ++ idx_off = (b_idx[:, None] * stride_ib + ++ (k_start + ks)[None, :] * stride_ik) ++ col_idx = tl.load(idx_ptr + idx_off, mask=b_mask[:, None] & k_mask) + ++ x_off = (b_idx[:, None] * stride_xb + ++ col_idx * stride_xc) ++ x_val = tl.load(x_ptr + x_off, mask=b_mask[:, None] & k_mask) + ++ out_off = (b_idx[:, None] * stride_ob + ++ (k_start + ks)[None, :] * stride_ok) ++ tl.store(out_ptr + out_off, x_val, mask=b_mask[:, None] & k_mask) + +# 调用 +B = 128 # batch dim +K = 64 + +BLOCK_B = 4 +BLOCK_K = 128 + +— # 原始grid较大,每个核心处理BLOCK_K个数据,分核数=B*K/BLOCK_K +- grid = (B, triton.cdiv(K, BLOCK_K)) ++ # 优化后grid变小,每个核心处理BLOCK_B*K个数据,分核数=B/BLOCK_B,内部展开循环处理BLOCK_K个数据 ++ grid = (triton.cdiv(B, BLOCK_B),) + +gather_dim1_kernel[grid]( + x, idx, out, + x.stride(0), x.stride(1), + idx.stride(0), idx.stride(1), + out.stride(0), out.stride(1), + B, K, + BLOCK_B=BLOCK_B, + BLOCK_K=BLOCK_K, +) ``` -### 2. Reduction 分核维度优化 - -#### 原始代码(非 reduce 轴分核,负载不均) +### 使用Triton-Ascend autotune搜索最佳分核参数 +Triton-Ascend autotune是一个Triton-Ascend提供的tiling超参数性能调优工具,遍历搜索空间,尝试不同参数组合,展示每个组合的运行耗时与最优组合。使用Triton-Ascend autotune需要遵循以下工作流程: +1. 识别出 triton kernel 中哪些 `tl.constexpr` 参数是自由可调的 tiling 参数,包括影响分核(split)和切块(tiling)大小的参数,这里的分核指的是影响 grid 大小,而切块指的是影响 tile 大小,即影响 `tl.load` 或是 `tl.make_block_ptr` 产生的数据大小。 +2. 如果这些参数能从 `tl.program_id`、`tl.arange`、`tl.range/range`、`mask/bounds` 表达式中被唯一识别出来,就尝试使用自动生成tiling `configs=[]` +3. 如果 kernel 语义上适合自动 tiling,但 DSL 写法让 parser 解析不出来,就显式传 `hints` +4. 如果某些 tiling 参数不可自由调整,例如某 kernel dsl 写法要求 grid 第一维必须固定为 `batch_size` 大小,或者根本没有暴露出可调的 tiling 参数,此时建议直接手写 triton.Config。 + +#### 使用方法 +@triton.autotune 入参列表 +| 参数名 | 类型 | 必填 | 说明 | +| --------- | ------------------------------ | ----- | --------------------------------------------------------------- | +| `configs` | `list[Config]` | 否 | 用户自定义的调优配置列表,为空时自动生成 | +| `key` | `dict[str, str]` / `list[str]` | **是** | 缓存键,指定哪些参数变化时需要重新调优。使用 `hints` 时必须用**字典形式**,如 `{"x": "n_rows"}` | +| `hints` | `dict` | 否 | **Ascend 扩展参数**,用于显式指定轴与 tiling 参数的映射关系 | + +hints 字典内部字段 +| 字段名 | 类型 | 必填 | 说明 | +| ----------------- | ---------------- | ----- | ------------------------------------------------------ | +| `split_params` | `dict[str, str]` | **是** | 分核参数映射,如 `{"x": "BLOCK_M"}` 表示沿 `x` 轴切 program | +| `tiling_params` | `dict[str, str]` | **是** | 切块参数映射,如 `{"y": "BLOCK_N"}` 表示沿 `y` 轴切 block | +| `low_dim_axes` | `list[str]` | **是** | 低维轴列表,如 `["y"]`,用于优化 tiling 效果 | +| `reduction_axes` | `list[str]` | **是** | 规约轴列表,如 `[]`,用于优化 tiling 效果 | +| `auto_gen_config` | `bool` | 否 | 是否自动生成 tiling 配置,默认 `True`;当 `configs` 非空时默认变为 `False` | + +有三种使用方法,以下是这三种方法的使用模板 + +##### 自动解析模板 +详情见 Step 2:判断是否能用 configs=[] 自动生成 tiling +```python +@triton.autotune( + configs=[], + key=["n_rows"], +) +@triton.jit +def kernel( + x_ptr, + y_ptr, + n_rows, + BLOCK_M: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) + mask = offs < n_rows + x = tl.load(x_ptr + offs, mask=mask, other=0) + tl.store(y_ptr + offs, x, mask=mask) +``` +##### 显式给 hints +详情见 Step 3:自动解析出错时,显式传 hints +```python +@triton.autotune( + configs=[], + key={"x": "n_rows", "y": "n_cols"}, + hints={ + "split_params": {"x": "BLOCK_M"}, + "tiling_params": {"y": "BLOCK_N"}, + "low_dim_axes": ["y"], + "reduction_axes": [], + }, +) +@triton.jit +def kernel( + x_ptr, + y_ptr, + n_rows, + n_cols, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + ... +``` +##### 完全手写 triton.Config ```python -# 数据形状:(16, 262144),沿 axis=1 归约 -M, N = 16, 262144 -BLOCK_M = 16 -grid = (triton.cdiv(M, BLOCK_M),) # grid = 1,只用 1 个核! +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "multibuffer": True}), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "multibuffer": True}), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "multibuffer": False}), + ], + key=["n_rows", "n_cols"], +) +@triton.jit +def kernel( + x_ptr, + y_ptr, + n_rows, + n_cols, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + ... ``` -#### 优化后代码(reduce 轴分核 + 原子操作) +#### 注意事项 +1. `@triton.autotune` 需要直接包在 `@triton.jit` 外层,示例如下: ```python -# 数据形状:(16, 262144),沿 axis=1 归约 -M, N = 16, 262144 -BLOCK_N = 8192 -SUB_BLOCK_N = 1024 # 二次切分避免 UB 溢出 -grid = (triton.cdiv(N, BLOCK_N),) # grid = 32,充分利用多核 - +@triton.autotune(...) @triton.jit -def reduce_min_kernel(x_ptr, y_ptr, M, N, BLOCK_N: tl.constexpr, SUB_BLOCK_N: tl.constexpr): - pid = tl.program_id(0) - n_start = pid * BLOCK_N - - row_min = tl.full([M], float('inf'), dtype=tl.float32) - - # 二次切分 - for sub_start in range(0, BLOCK_N, SUB_BLOCK_N): - n_offsets = n_start + sub_start + tl.arange(0, SUB_BLOCK_N) - mask = n_offsets < N - x = tl.load(x_ptr + tl.arange(0, M)[:, None] * N + n_offsets[None, :], mask=mask) - local_min = tl.min(x, axis=1) - row_min = tl.minimum(row_min, local_min) - - # 原子操作跨核归约 - for m in range(M): - tl.atomic_min(y_ptr + m, row_min[m]) +def kernel(...): + ... ``` -### 3. 任务分配优化(非均匀分配) +2. 自动生成 tiling 功能主要面向 vector kernel:当前 Triton-Ascend 这套自动解析/自动生成 tiling 面向的是 vector 类 kernel,且该 kernel 中所有分核和分块参数均可调。Cube 类算子目前还不支持自动 tiling 生成。 -#### 原始代码(简单除法,边界处理不当) +#### 详细步骤 +##### Step 1:先识别哪些参数真的是可调参数 +###### 1.1 先看“哪些 tl.constexpr 没有在 launch 时显式传入” +Triton-Ascend autotuner 在自动解析 split/tiling 参数时,首先会看 kernel 调用时 哪些参数没有传入,把这些“缺省的参数”当成候选项。 +简单理解: + +* Tensor 参数不可能是自动解析候选项; +* 普通运行时 shape 参数(如 n_rows、n_cols)通常属于 `key`; +* 真正的候选项通常是没有在 launch 处显式传值的 `tl.constexpr`; +* 如果某个 tl.constexpr 已经在 launch 时手动写死,它就不会再被当成自动解析候选项。 +例如: ```python -n_cores = 32 -total_tasks = 1000 -tasks_per_core = total_tasks // n_cores # 31,剩余 8 个任务无人处理 +@triton.jit +def kernel( + x_ptr, + y_ptr, + n_rows, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, # BLOCK_N是自动解析候选项 +): + ... + +# BLOCK_M 已显式传入,不再是自动解析候选项 +kernel[grid](x, y, n_rows, BLOCK_M=128) +``` + +###### 1.2 如何识别 split 参数 -# Kernel 内 +split 参数控制“一个 program 负责多大的一块数据”,它最常见的写法特征是: +1. 和 `tl.program_id(...)` 有直接关系; +2. 参与构造 block 起始位置; +3. 最后能通过 mask/bounds 表达式对应回某个 shape 轴。 +例如: +```python +# 一维切分 pid = tl.program_id(0) -start = pid * tasks_per_core -end = start + tasks_per_core # 最后 8 个任务丢失! +offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)# 可以知道BLOCK_M 是 split 参数 +mask_m = offs_m < n_rows + +# 二维切分 +pid_m = tl.program_id(0) +pid_n = tl.program_id(1) +offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]# 可以知道BLOCK_M 是 split 参数 +offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]# 可以知道BLOCK_N 是 split 参数 +mask_m = offs_m < n_rows +mask_n = offs_n < n_cols ``` -#### 优化后代码(base + rem 模式) +###### 1.3 如何识别 tiling 参数 +tiling 参数控制“在一个大的 split block 内,再按多大的子块去迭代”,它最常见的写法特征是: + +1. 出现在 `tl.arange(0, PARAM)` 中; +2. 同时还出现在 for 循环的步长或循环次数推导中; +3. 最后能通过 mask/bounds 对应回某个轴长度参数。 +例如: ```python -n_cores = 32 -total_tasks = 1000 -base = total_tasks // n_cores # 31 -rem = total_tasks % n_cores # 8 +# 典型形态 1:步长是 tiling 参数 +for k0 in tl.range(0, BLOCK_K, BLOCK_K_SUB):# 可以知道BLOCK_K_SUB 是 tiling 参数 + offs_k = k0 + tl.arange(0, BLOCK_K_SUB)# 可以知道BLOCK_K_SUB 是 tiling 参数 + mask_k = offs_k < k_size + +# 典型形态 2:先计算循环次数 +num_k_tiles = (k_size + BLOCK_K_SUB - 1) // BLOCK_K_SUB# 可以知道BLOCK_K_SUB 是 tiling 参数 +for tile_id in range(num_k_tiles): + offs_k = tile_id * BLOCK_K_SUB + tl.arange(0, BLOCK_K_SUB)# 可以知道BLOCK_K_SUB 是 tiling 参数 + mask_k = offs_k < k_size +``` -# Kernel 内 -pid = tl.program_id(0) -if pid < rem: - task_start = pid * (base + 1) - task_end = task_start + (base + 1) -else: - task_start = rem * (base + 1) + (pid - rem) * base - task_end = task_start + base -# 总计:8*32 + 24*31 = 256 + 744 = 1000 ✓ +##### Step 2:判断是否能用 configs=[] 自动生成 tiling + +当你已经找到了候选参数后,可以按下面的检查表判断。 +###### 2.1 可以优先尝试 configs=[] 的情况 +一般同时满足下面几条时,`configs=[]` 成功率比较高: + +* split 参数能从 `tl.program_id` 路径判断出来; +* tiling 参数能从 `tl.arange + for(range/tl.range)` 路径判断出来; +* 每个轴都有比较清晰的 mask/bounds 表达式,例如: + * `offs < n` + * `offs < min(block_end, n)` +* `key` 能和运行时 shape 参数一一对应; + +###### 2.2 不适合直接用 configs=[] 的常见信号 + +下面这些情况,直接走自动 tiling 生成可能会出现解析失败的情况: + +* 没有和轴长度直接绑定的 mask/bounds +* 某个参数必须覆盖完整语义维度 + * 例如 `BLOCK_SIZE >= hidden_dim` +* grid 某一维被业务语义固定,不允许自由切块 +* 一个参数同时影响两个轴,或者同时影响“核数 + tile 形状” +* kernel 没暴露出可调 `tl.constexpr` + +###### 2.3 如果出现错误建议打开调试日志 +排查问题时,建议先启动环境变量debug: +``` +export TRITON_PRINT_AUTOTUNING=1 +``` + +日志中可以直接看到: + +* 识别出的 split axes; +* 识别出的 tiling axes; +* 识别出的 low-dimensional axes; +* 识别出的 reduction axes; +* 生成的 config 数量。 +小 shape 算子如果 benchmark 抖动大,也可以按需开启: +``` +export TRITON_BENCH_METHOD=npu ``` +这会测试的时间更准确,但 autotune 时间也会明显变长。 + +##### Step 3:自动解析出错时,显式传 hints + +如果你已经确认 自动 tiling 适用于该 triton kernel,只是因为 DSL 写法不能够被当前 parser 识别,这时可以尝试显式传 `hints`。 +`hints` 是 Triton-Ascend 在 `autotune` 装饰器中新增的一个参数,类型为 `dict`,用于给 Triton-Ascend autotune 提供该 triton kernel 的一些关键信息,帮助 autotune 更好的生成 tiling 配置。 -## 核数选择建议 +###### 3.1 什么情况下可以考虑 hints + +推荐显式传 `hints` 的场景: + +* 你能人工判断哪个参数属于 `split`,哪个属于 `tiling`; +* 你知道每个参数对应轴的长度参数; + +###### 3.2 hints 参数说明 +* 当前 `hints` 参数中可以识别的字段有: + * `split_params`: dict[str, str],分核参数的映射关系,例如 `{"x": "BLOCK_M"}` 表示 `BLOCK_M` 是沿 `x` 轴切 program + * `tiling_params`:dict[str, str],切块参数的映射关系,例如 `{"y": "BLOCK_N"}` 表示 `BLOCK_N` 是沿 `y` 轴切 block + * `low_dim_axes`:list[str],低维轴的列表,例如 `["y"]` 表示 `y` 轴是低维轴 + * `reduction_axes`:list[str],规约轴的列表,例如 `[]` 表示没有规约轴 + * `auto_gen_config`:bool,是否自动生成 tiling 配置,默认值为 `True` +* 注意: + * 通过 `hints` 来显示指定轴关系时,autotune 中原本的参数 `key` 必须改为字典形式传入,因为后续 `split_params`、`tiling_params` 等参数都是按轴名来填写,需要和 `key` 里的轴名对应起来 + * 通过 `hints` 来显示指定轴关系时,`split_params`、`tiling_params`、`low_dim_axes`、`reduction_axes` 必须传入,即使某些参数为空 + * 合法的轴名称是 `x/y/z/w/v/t`,仅仅用做关系映射 + * `split_params` 和 `tiling_params` 为自动生成 tiling 算法必须的输入,`low_dim_axes` 和 `reduction_axes` 为 tiling 算法的可选输入,用于优化 tiling 效果,留空时 tiling 也能够自动生成,但可能会影响生成的候选 tiling 数量和质量 + * 当用户传入的 configs 不为空时,`auto_gen_config` 默认值为 `False`,如果希望此时也希望自动生成 tiling 配置并与用户传入的 configs 合并,需要显式在 `hints` 中传如入 `"auto_gen_config": True` + +使用示例: +```python +import triton +import triton.language as tl +import triton.backends.ascend.runtime + +@triton.autotune( + # configs 为空列表,表示不传入自定义配置 + # 此时 auto_gen_config 默认为 True,会自动生成 tiling 配置 + configs=[], + + # key 使用字典形式,轴名必须与 hints 中的轴名对应 + # "x" 对应 n_rows(行数),"y" 对应 n_cols(列数) + # autotune 会根据这些维度值来缓存和选择最佳配置 + key={"x": "n_rows", "y": "n_cols"}, + + # hints 参数:显式指定轴与 tiling 参数的映射关系 + hints={ + # split_params: 分核参数映射,指定沿哪个轴切分 program(任务) + # "x": "BLOCK_M" 表示 BLOCK_M 沿 x 轴切分,即按行方向分核 + # 每个 program 处理 BLOCK_M 行数据 + "split_params": {"x": "BLOCK_M"}, + + # tiling_params: 切块参数映射,指定沿哪个轴切分 block(数据块) + # "y": "BLOCK_N" 表示 BLOCK_N 沿 y 轴切分,即按列方向切块 + # 每行数据在列方向上被切分为 BLOCK_N 大小的块,通过 for 循环处理 + "tiling_params": {"y": "BLOCK_N"}, + + # low_dim_axes: 低维轴列表,用于优化 tiling 效果 + # ["y"] 表示 y 轴(列方向)是低维轴,访问连续性更好,适合作为内层循环 + "low_dim_axes": ["y"], + + # reduction_axes: 规约轴列表,本 kernel 无规约操作(如 sum/max 等) + # 为空列表表示没有规约轴 + "reduction_axes": [], + + # auto_gen_config: 默认为 True,表示自动生成 tiling 配置 + # 由于 configs 为空,此处使用默认值 True 即可,无需显式传入 + }, +) +@triton.jit +def kernel_with_hints( + x_ptr, + y_ptr, + n_rows, + n_cols, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] + + for n0 in range(0, n_cols, BLOCK_N): + offs_n = n0 + tl.arange(0, BLOCK_N)[None, :] + mask_m = offs_m < n_rows + mask_n = offs_n < n_cols + mask = mask_m & mask_n + + x = tl.load(x_ptr + offs_m * n_cols + offs_n, mask=mask, other=0) + tl.store(y_ptr + offs_m * n_cols + offs_n, x, mask=mask) +``` + +##### Step 4:手写 triton.Config +如果 Step 2 判断该 triton kernel 不适合或者无法使用自动 tiling 生成,那么可以使用社区 triton autotune 的基本功功能:手写一组 `triton.Config` 传入参数 `configs` 中。 +###### 手写 triton.Config 的总体原则 +1. 对于影响 grid 发射核数的参数,一般我们尽量让其能够等于物理核数,如果数据量较小,也可能发射较少核数的时候能获得最优性能;对于影响 tile 块大小的参数,我们尽量在不产生 UB overflow 的情况下让其尽可能大,同时避免尾块的产生 +2. 影响 grid 发射核数的参数:可以按照总长度从高到低设置为 X, X/2, X/4 等等的值,如果输入 shape 较大,可以设置为让 grid 发射核数正好等于物理核数的大小,例如 (X + num_cores - 1) // num_cores;这里是以一个切分轴为例,如果存在多个切分轴,那么就需要按照乘积来计算 +3. 影响 tile 块大小的参数:起始值为切分轴参数(如果存在)或者轴长度,注意当该轴长度特别大的时候,我们可以直接从 16384 这样一个经验值开始取,然后按照 X / 2, X / 4 这样去取值 +4. 上述按 2 的幂次方下降的值采样较为粗粒度,如果用户想要得到极致的最优性能,尤其是在输入大小不规则的情况下,需要在可能的最优区间内细粒度撒点,可以通过粗粒度采样后确认性能最优的大致区间后再进一步细分来实现。 +5. 对于 vector 类算子,在设置了上述 tiling 大小的配置候选集后,可以加上 multibuffer 编译选项的调优。 + +示例: +```python +import triton +import triton.language as tl +import triton.backends.ascend.runtime + + +def get_configs(): + return [ + triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "multibuffer": MB}) + for BM in [256, 128, 64, 32] + for BN in [128, 64, 32, 16] + for MB in [True, False] + ] + + +@triton.autotune( + configs=get_configs(), + key=["n_rows", "n_cols"], +) +@triton.jit +def manual_config_kernel( + x_ptr, + y_ptr, + n_rows, + n_cols, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] + offs_n = tl.arange(0, BLOCK_N)[None, :] + mask = (offs_m < n_rows) & (offs_n < n_cols) + + x = tl.load(x_ptr + offs_m * n_cols + offs_n, mask=mask, other=0) + tl.store(y_ptr + offs_m * n_cols + offs_n, x, mask=mask) +``` -| 数据规模 | 推荐核数 | 理由 | -| ----------- | -------- | -------------- | -| < 1KB | 1 | 调度开销主导 | -| 1KB - 100KB | 1-8 | 平衡调度和计算 | -| 100KB - 1MB | 8-16 | 开始受益于并行 | -| > 1MB | 16-40 | 充分利用硬件 | +##### 常见失败速查 +下面这张表可以直接用来决定你下一步该怎么做。 -## 分核维度决策 +| 现象 | 更可能的原因 | 建议动作 | +| ------------------------------- | ---------------------------- | ------------------------- | +| configs=[] 直接解析失败 | split/tiling 轴没有从 DSL 唯一识别出来 | 先补 hints,再试 | +| parser 能识别一部分,但总差一个参数 | 某个参数没有和轴长度 mask 建立联系 | 改 DSL 写法或改手写 config | +| kernel 完全没有合适的 tl.constexpr 可调项 | DSL 没暴露调参接口 | 先改 kernel dsl,再谈 autotune | +| 自动生成能跑,但候选质量明显差 | 当前算法不适合该 kernel 的参数耦合方式 | 手动构造 config 传入 | -| 算子类型 | 数据特征 | 分核维度 | -| ----------- | -------- | --------------------------- | -| Elementwise | 任意 | 按数据量均分 | -| Reduction | M >> N | 按非 reduce 轴分核 | -| Reduction | M << N | 按 reduce 轴分核 + 原子操作 | -| MatMul | 通用 | 按 M-N 维度 2D 分核 | ## 性能收益 -- **核数优化**:小数据量场景可提升 2-5x 性能 -- **分核维度优化**:极端形状场景可提升 2-3x 性能 -- **负载均衡**:避免部分核空闲,提升整体吞吐 +- **核数优化**:可提升 2-5x 性能 ## 注意事项 1. **UB 容量约束**:确保 tile_size * dtype_size * buffers <= 192KB 2. **原子操作开销**:原子操作有额外开销,核数过多时可能成为瓶颈 -3. **对齐要求**:数据传输需 256B 对齐 From ed3253ba7dd35e2ce1c20a87981cf779defb4fdb Mon Sep 17 00:00:00 2001 From: wwwbby Date: Thu, 9 Apr 2026 20:14:32 +0800 Subject: [PATCH 4/9] add autotune --- skills/latency-optimizer/references/autotune.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 skills/latency-optimizer/references/autotune.md diff --git a/skills/latency-optimizer/references/autotune.md b/skills/latency-optimizer/references/autotune.md new file mode 100644 index 00000000..e69de29b From f6adc6eca9d16069c31655a2c6c412504211dcaa Mon Sep 17 00:00:00 2001 From: wwwbby <60009003+wwwbby@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:19:17 +0800 Subject: [PATCH 5/9] Update autotune.md --- .../latency-optimizer/references/autotune.md | 576 ++++++++++++++++++ 1 file changed, 576 insertions(+) diff --git a/skills/latency-optimizer/references/autotune.md b/skills/latency-optimizer/references/autotune.md index e69de29b..3f6a1838 100644 --- a/skills/latency-optimizer/references/autotune.md +++ b/skills/latency-optimizer/references/autotune.md @@ -0,0 +1,576 @@ +# Autotune 自动调优 + +## 概述 + +Triton autotune 用于自动选择最优的 kernel 配置参数,主要包括影响分核(split)和切块(tiling)大小的参数,主要使用方式如下: + +**三种种使用方式:** + +| 方式 | 说明 | 适用场景 | +|------|------|---------| +| **自动 autotune** | 框架自动解析切分轴、tiling 轴,生成配置 | Vector 类算子,简化使用 | +| **半自动 autotune** | 用户手动传入 `hints` ,框架基于`hints`生成配置 | Vector 类算子,简化使用 | +| **自定义 autotune** | 用户手动传入 `triton.Config` 列表 | 需要精确控制搜索空间 | + +上述三种方法的详细使用场景参考如下: +1. 如果这些参数能从 `tl.program_id`、`tl.arange`、`tl.range/range`、`mask/bounds` 表达式中被唯一识别出来,就尝试使用自动 autotune `configs=[]` +2. 如果 kernel 语义上适合自动 tiling,但 DSL 写法让 parser 解析不出来,就使用半自动 autotune显式传 `hints` +3. 如果某些 tiling 参数不可自由调整,例如某 kernel dsl 写法要求 grid 第一维必须固定为 `batch_size` 大小,或者根本没有暴露出可调的 tiling 参数,此时建议使用自定义 autotune。 + +**说明:** 当前 Triton-Ascend autotune 支持 block size、multibuffer(编译器优化),因硬件架构差异不支持 num_warps、num_stages 参数。 + +--- + +## 一、API 参考 + +### triton.autotune 装饰器 + +```python +@triton.autotune( + configs=[...], # Config 列表 + key=['x_size'], # 触发重新评估的参数名 + hints={...} # 显式指定轴与 tiling 参数的映射关系 + prune_configs_by=None, # 配置剪枝函数 + reset_to_zero=None, # 评估前重置为零的参数 + restore_value=None, # 评估后恢复原值的参数 + pre_hook=None, # 内核调用前的钩子 + post_hook=None, # 内核调用后的钩子 + warmup=None, # 预热时间(已弃用) + rep=None, # 重复时间(已弃用) + do_bench=None, # 自定义基准测试函数 + cache_results=False, # 是否缓存结果到磁盘 +) +@triton.jit +def kernel(...): + ... +``` + +#### 参数说明 + +| 参数 | 类型 | 说明 | +|------|------|------| +| `configs` | `list[triton.Config]` | Config 对象列表,每个代表一种 kernel 配置 | +| `hints` | `dict` | **Ascend 扩展参数**,用于显式指定轴与 tiling 参数的映射关系 | +| `key` | `list[str]` | 参数名列表,这些参数值变化时触发重新评估所有配置 | +| `prune_configs_by` | `dict` | 配置剪枝函数,用于减少评估的配置数量 | +| `reset_to_zero` | `list[str]` | 参数名列表,评估前重置为零(避免累积更新) | +| `restore_value` | `list[str]` | 参数名列表,评估后恢复原值 | +| `pre_hook` | `lambda` | 内核调用前的钩子函数 | +| `post_hook` | `lambda` | 内核调用后的钩子函数 | +| `do_bench` | `lambda` | 自定义基准测试函数 | +| `cache_results` | `bool` | 是否缓存调优结果到磁盘(默认 False) | + +#### 重要提示 + +**避免累积更新:** 当所有配置都评估后,内核会运行多次。如果内核会更新某些值,这些值会被多次更新。使用 `reset_to_zero` 在评估前重置这些张量: + +```python +@triton.autotune( + configs=[...], + key=['n'], + reset_to_zero=['output_ptr'], # 评估前重置 output 为零 +) +@triton.jit +def kernel(output_ptr, ...): + ... +``` + +**调试输出:** 设置环境变量 `TRITON_PRINT_AUTOTUNING=1`,Triton 会打印调优时间和最佳配置: + +```bash +export TRITON_PRINT_AUTOTUNING=1 +``` + +### triton.Config 类 + +```python +triton.Config( + kwargs={'BLOCK_SIZE': 128}, # 传递给内核的元参数 + num_warps=4, # warp 数量(GPU) + num_stages=3, # 流水阶段数(GPU) + num_ctas=1, # 块集群中的块数(SM90+) + maxnreg=None, # 最大寄存器数 + pre_hook=None, # 调用前的钩子 + ir_override=None, # 自定义 IR 文件名 +) +``` + +#### 参数说明 + +| 参数 | 类型 | 说明 | +|------|------|------| +| `kwargs` | `dict` | 传递给内核的元参数字典,如 `{'BLOCK_SIZE': 128}` | +| `num_warps` | `int` | GPU warp 数量,决定并行线程数(8 warps = 256 线程) | +| `num_stages` | `int` | 软件流水线阶段数,用于矩阵乘法优化 | +| `num_ctas` | `int` | 块集群中的块数量(仅 SM90+ GPU) | +| `maxnreg` | `int` | 单线程最大寄存器数量 | +| `pre_hook` | `lambda` | 内核调用前的钩子函数 | +| `ir_override` | `str` | 自定义 IR 文件名(.ttgir/.llir/.ptx/.amdgcn) | + +#### Triton-Ascend 支持情况 + +| 参数 | 社区版 | Triton-Ascend | 说明 | +|------|--------|---------------|------| +| `kwargs` | ✅ | ✅ | 完全支持 | +| `num_warps` | ✅ | ❌ | NPU 架构差异,不支持 | +| `num_stages` | ✅ | ❌ | NPU 架构差异,不支持 | +| `multibuffer` | ❌ | ✅ | NPU 特有,多缓冲优化 | +| `unit_flag` | ❌ | ✅ | NPU 特有,独立计算单元 | + +#### 使用示例 + +```python +# GPU 风格配置 +configs = [ + triton.Config({'BLOCK_SIZE': 128}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_SIZE': 256}, num_warps=8, num_stages=3), +] + +# Triton-Ascend 风格配置 +configs = [ + triton.Config({'BLOCK_SIZE': 128, 'multibuffer': True}), + triton.Config({'BLOCK_SIZE': 256, 'multibuffer': False}), +] +``` + +### prune_configs_by 配置剪枝 + +用于减少需要评估的配置数量,加速 autotune 过程: + +```python +@triton.autotune( + configs=[...], + key=['n'], + prune_configs_by={ + 'perf_model': my_perf_model, # 性能预测模型 + 'top_k': 10, # 只评估 top_k 个配置 + 'early_config_prune': my_prune_fn, # 自定义剪枝函数 + } +) +@triton.jit +def kernel(...): + ... +``` + +**剪枝函数签名:** + +```python +def prune_fn( + configs: List[triton.Config], + named_args: Dict[str, Any], + **kwargs +) -> List[triton.Config]: + # 返回剪枝后的配置列表(至少返回一个) + return pruned_configs +``` + +### pre_hook / post_hook 钩子 + +**pre_hook 签名:** + +```python +def pre_hook(kwargs, reset_only): + # kwargs: 传递给内核的所有参数 + # reset_only: 是否仅用于重置值 + pass +``` + +**post_hook 签名:** + +```python +def post_hook(kwargs, exception): + # kwargs: 传递给内核的所有参数 + # exception: 编译或运行时异常(无异常时为 None) + pass +``` + +--- + +## 二、autotune使用工作流 +使用Triton-Ascend autotune搜索最佳分核参数需要遵循以下工作流: + +### step1.识别出 triton kernel 中哪些 `tl.constexpr` 参数是自由可调的 tiling 参数 + +系统首先识别 kernel 调用时**未传入**的参数作为候选项: +* Tensor 参数不可能是自动解析候选项; +* 普通运行时 shape 参数(如 n_rows、n_cols)通常属于 `key`; +* 真正的候选项通常是没有在 launch 处显式传值的 `tl.constexpr`; +* 如果某个 tl.constexpr 已经在 launch 时手动写死,它就不会再被当成自动解析候选项。 + +例如: +```python +@triton.jit +def kernel_func( + output_ptr, input_ptr, + n_rows, n_cols, + BLOCK_SIZE: tl.constexpr, # 调用时传入,不可调 + XBLOCK: tl.constexpr, # 未传入,可调 + XBLOCK_SUB: tl.constexpr, # 未传入,可调 +): + ... + +# 调用时只传入 BLOCK_SIZE +kernel_func[grid](y, x, n_rows, n_cols, BLOCK_SIZE=block_size) +# 可调参数候选项:XBLOCK, XBLOCK_SUB +``` + +#### step1.1.识别切分参数 + +切分(split)参数控制“一个 program 负责多大的一块数据”,它最常见的写法特征是: +1. 和 `tl.program_id(...)` 有直接关系; +2. 参与构造 block 起始位置; +3. 最后能通过 mask/bounds 表达式对应回某个 shape 轴。 +例如: +```python +# 一维切分 +pid = tl.program_id(0) +offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)# 可以知道BLOCK_M 是 split 参数 +mask_m = offs_m < n_rows + +# 二维切分 +pid_m = tl.program_id(0) +pid_n = tl.program_id(1) +offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]# 可以知道BLOCK_M 是 split 参数 +offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]# 可以知道BLOCK_N 是 split 参数 +mask_m = offs_m < n_rows +mask_n = offs_n < n_cols +``` + +#### step1.2.识别分块参数 + +分块(tiling)参数控制“在一个大的 split block 内,再按多大的子块去迭代”,它最常见的写法特征是: + +1. 出现在 `tl.arange(0, PARAM)` 中; +2. 同时还出现在 for 循环的步长或循环次数推导中; +3. 最后能通过 mask/bounds 对应回某个轴长度参数。 + +例如: +```python +# 典型形态 1:步长是 tiling 参数 +for k0 in tl.range(0, BLOCK_K, BLOCK_K_SUB):# 可以知道BLOCK_K_SUB 是 tiling 参数 + offs_k = k0 + tl.arange(0, BLOCK_K_SUB)# 可以知道BLOCK_K_SUB 是 tiling 参数 + mask_k = offs_k < k_size + +# 典型形态 2:先计算循环次数 +num_k_tiles = (k_size + BLOCK_K_SUB - 1) // BLOCK_K_SUB# 可以知道BLOCK_K_SUB 是 tiling 参数 +for tile_id in range(num_k_tiles): + offs_k = tile_id * BLOCK_K_SUB + tl.arange(0, BLOCK_K_SUB)# 可以知道BLOCK_K_SUB 是 tiling 参数 + mask_k = offs_k < k_size +``` + +#### step1.3.识别低维轴参数 +**依据:** `tl.arange()` 切片操作 + +**识别规则:** +1. 必须通过 `tl.arange()` 计算 +2. 必须进行切片操作 +3. 在**非最低维**进行维度扩充才被识别为低维轴 + +```python +@triton.autotune(configs=[], key=["n_rows", "n_cols"]) +@triton.jit +def kernel(...): + for row_idx in tl.range(0, XBLOCK, XBLOCK_SUB): + # row_offsets:切片在低维扩充,不是低维轴 + row_offsets = row_idx + tl.arange(0, XBLOCK_SUB)[:, None] + # col_offsets:切片在高维扩充,是低维轴 + col_offsets = tl.arange(0, BLOCK_SIZE)[None, :] + + xmask = row_offsets < n_rows + ymask = col_offsets < n_cols + +# 解析结果:low_dim_axes = ["y"] +``` + +#### step1.4.指针参数解析 + +**依据:** 参数是否参与 `tl.load()` 或 `tl.store()` 的第一个参数计算 + +```python +@triton.jit +def kernel(input_ptr, output_ptr, ...): + # 直接参与 + input = tl.load(input_ptr + offsets, mask=mask) + tl.store(output_ptr + offsets, input, mask=mask) + + # 或间接参与 + inputs_ptr = input_ptr + offsets + input = tl.load(inputs_ptr, mask=mask) + +# 解析结果:指针参数 = input_ptr, output_ptr +``` + +### step2.尝试自动autotune +一般同时满足下面几条时,自动autotune的成功率比较高: + +* split 参数能从 `tl.program_id` 路径判断出来; +* tiling 参数能从 `tl.arange + for(range/tl.range)` 路径判断出来; +* 每个轴都有比较清晰的 mask/bounds 表达式,例如: + * `offs < n` + * `offs < min(block_end, n)` +* `key` 能和运行时 shape 参数一一对应; + +下面这些情况,自动autotune可能会出现解析失败的情况: + +* 没有和轴长度直接绑定的 mask/bounds +* 某个参数必须覆盖完整语义维度 + * 例如 `BLOCK_SIZE >= hidden_dim` +* grid 某一维被业务语义固定,不允许自由切块 +* 一个参数同时影响两个轴,或者同时影响“核数 + tile 形状” +* kernel 没暴露出可调 `tl.constexpr` + +自动autotune模板如下: +```python +@triton.autotune( + # configs 为空列表,表示不传入自定义配置 + # 此时 auto_gen_config 默认为 True,会自动生成 tiling 配置 + configs=[], + key=["n_rows"], +) +@triton.jit +def kernel( + x_ptr, + y_ptr, + n_rows, + BLOCK_M: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) + mask = offs < n_rows + x = tl.load(x_ptr + offs, mask=mask, other=0) + tl.store(y_ptr + offs, x, mask=mask) +``` + +### step3.尝试半自动autotune +`hints` 是 Triton-Ascend 在 `autotune` 装饰器中新增的一个参数,类型为 `dict`,用于给 Triton-Ascend autotune 提供该 triton kernel 的一些关键信息,帮助 autotune 更好的生成 tiling 配置。推荐显式传 `hints` 的场景: + +* 能判断哪个参数属于 `split`,哪个属于 `tiling`; +* 知道每个参数对应轴的长度参数; + +hints 参数说明: +* 当前 `hints` 参数中可以识别的字段有: + * `split_params`: dict[str, str],分核参数的映射关系,例如 `{"x": "BLOCK_M"}` 表示 `BLOCK_M` 是沿 `x` 轴切 program + * `tiling_params`:dict[str, str],切块参数的映射关系,例如 `{"y": "BLOCK_N"}` 表示 `BLOCK_N` 是沿 `y` 轴切 block + * `low_dim_axes`:list[str],低维轴的列表,例如 `["y"]` 表示 `y` 轴是低维轴 + * `reduction_axes`:list[str],规约轴的列表,例如 `[]` 表示没有规约轴 + * `auto_gen_config`:bool,是否自动生成 tiling 配置,默认值为 `True` +* 注意: + * 通过 `hints` 来显示指定轴关系时,autotune 中原本的参数 `key` 必须改为字典形式传入,因为后续 `split_params`、`tiling_params` 等参数都是按轴名来填写,需要和 `key` 里的轴名对应起来 + * 通过 `hints` 来显示指定轴关系时,`split_params`、`tiling_params`、`low_dim_axes`、`reduction_axes` 必须传入,即使某些参数为空 + * 合法的轴名称是 `x/y/z/w/v/t`,仅仅用做关系映射 + * `split_params` 和 `tiling_params` 为自动生成 tiling 算法必须的输入,`low_dim_axes` 和 `reduction_axes` 为 tiling 算法的可选输入,用于优化 tiling 效果,留空时 tiling 也能够自动生成,但可能会影响生成的候选 tiling 数量和质量 + * 当用户传入的 configs 不为空时,`auto_gen_config` 默认值为 `False`,如果希望此时也希望自动生成 tiling 配置并与用户传入的 configs 合并,需要显式在 `hints` 中传如入 `"auto_gen_config": True` + +使用示例: +```python +import triton +import triton.language as tl +import triton.backends.ascend.runtime + +@triton.autotune( + # configs 为空列表,表示不传入自定义配置 + # 此时 auto_gen_config 默认为 True,会自动生成 tiling 配置 + configs=[], + + # key 使用字典形式,轴名必须与 hints 中的轴名对应 + # "x" 对应 n_rows(行数),"y" 对应 n_cols(列数) + # autotune 会根据这些维度值来缓存和选择最佳配置 + key={"x": "n_rows", "y": "n_cols"}, + + # hints 参数:显式指定轴与 tiling 参数的映射关系 + hints={ + # split_params: 分核参数映射,指定沿哪个轴切分 program(任务) + # "x": "BLOCK_M" 表示 BLOCK_M 沿 x 轴切分,即按行方向分核 + # 每个 program 处理 BLOCK_M 行数据 + "split_params": {"x": "BLOCK_M"}, + + # tiling_params: 切块参数映射,指定沿哪个轴切分 block(数据块) + # "y": "BLOCK_N" 表示 BLOCK_N 沿 y 轴切分,即按列方向切块 + # 每行数据在列方向上被切分为 BLOCK_N 大小的块,通过 for 循环处理 + "tiling_params": {"y": "BLOCK_N"}, + + # low_dim_axes: 低维轴列表,用于优化 tiling 效果 + # ["y"] 表示 y 轴(列方向)是低维轴,访问连续性更好,适合作为内层循环 + "low_dim_axes": ["y"], + + # reduction_axes: 规约轴列表,本 kernel 无规约操作(如 sum/max 等) + # 为空列表表示没有规约轴 + "reduction_axes": [], + + # auto_gen_config: 默认为 True,表示自动生成 tiling 配置 + # 由于 configs 为空,此处使用默认值 True 即可,无需显式传入 + }, +) +@triton.jit +def kernel_with_hints( + x_ptr, + y_ptr, + n_rows, + n_cols, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] + + for n0 in range(0, n_cols, BLOCK_N): + offs_n = n0 + tl.arange(0, BLOCK_N)[None, :] + mask_m = offs_m < n_rows + mask_n = offs_n < n_cols + mask = mask_m & mask_n + + x = tl.load(x_ptr + offs_m * n_cols + offs_n, mask=mask, other=0) + tl.store(y_ptr + offs_m * n_cols + offs_n, x, mask=mask) +``` + +### step4.尝试自定义autotune +手写一组 `triton.Config` 传入参数 `configs` 中,手写 triton.Config 的总体原则 +1. 对于影响 grid 发射核数的参数,一般我们尽量让其能够等于物理核数,如果数据量较小,也可能发射较少核数的时候能获得最优性能;对于影响 tile 块大小的参数,我们尽量在不产生 UB overflow 的情况下让其尽可能大,同时避免尾块的产生 +2. 影响 grid 发射核数的参数:可以按照总长度从高到低设置为 X, X/2, X/4 等等的值,如果输入 shape 较大,可以设置为让 grid 发射核数正好等于物理核数的大小,例如 (X + num_cores - 1) // num_cores;这里是以一个切分轴为例,如果存在多个切分轴,那么就需要按照乘积来计算 +3. 影响 tile 块大小的参数:起始值为切分轴参数(如果存在)或者轴长度,注意当该轴长度特别大的时候,我们可以直接从 16384 这样一个经验值开始取,然后按照 X / 2, X / 4 这样去取值 +4. 上述按 2 的幂次方下降的值采样较为粗粒度,如果用户想要得到极致的最优性能,尤其是在输入大小不规则的情况下,需要在可能的最优区间内细粒度撒点,可以通过粗粒度采样后确认性能最优的大致区间后再进一步细分来实现。 +5. 对于 vector 类算子,在设置了上述 tiling 大小的配置候选集后,可以加上 multibuffer 编译选项的调优。 + +示例: +```python +import triton +import triton.language as tl +import triton.backends.ascend.runtime + + +def get_configs(): + return [ + triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "multibuffer": MB}) + for BM in [256, 128, 64, 32] + for BN in [128, 64, 32, 16] + for MB in [True, False] + ] + + +@triton.autotune( + configs=get_configs(), + key=["n_rows", "n_cols"], +) +@triton.jit +def manual_config_kernel( + x_ptr, + y_ptr, + n_rows, + n_cols, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] + offs_n = tl.arange(0, BLOCK_N)[None, :] + mask = (offs_m < n_rows) & (offs_n < n_cols) + + x = tl.load(x_ptr + offs_m * n_cols + offs_n, mask=mask, other=0) + tl.store(y_ptr + offs_m * n_cols + offs_n, x, mask=mask) +``` + + + +## 三、性能采集方式 + +### 默认方式:benchmark + +```python +# 默认使用 benchmark 方式获取片上计算时间 +``` + +### Profiling 方式(小 shape 推荐) + +```python +import os +os.environ["TRITON_BENCH_METHOD"] = "npu" # 使用 profiler + +# 对于小 shape 算子,能获取更准确的计算时间 +# 但会显著增加整体 autotune 时间,请谨慎开启 +``` + +--- + +## 四、问题定位 + +排查问题时,建议先启动环境变量debug: +```bash +export TRITON_PRINT_AUTOTUNING=1 +``` +日志中可以直接看到: + +* 识别出的 split axes; +* 识别出的 tiling axes; +* 识别出的 low-dimensional axes; +* 识别出的 reduction axes; +* 生成的 config 数量。 + +### 自动生成 Profiling 结果 + +```python +@triton.autotune( + auto_profile_dir="./profile_result", # 输出目录 + configs=[...], + key=[...], +) +@triton.jit +def kernel(...): + ... +``` + +自动在指定目录生成最优 kernel 配置的 profiling 结果。 + +### 常见问题 1:configs=[] 解析失败 + +**原因:** 切分/tiling 参数无法从 DSL 唯一识别 + +**解决:** +1. 添加 `hints` 显式指定 +2. 或改为手写 `configs` + +### 常见问题 2:自动生成配置质量差 + +**原因:** 参数耦合方式不适合自动算法 + +**解决:** 手写 Config,参考分核优化原则 + +### 常见问题 3:Kernel 没有可调参数 + +**原因:** 所有 `tl.constexpr` 都被显式传入 + +**解决:** 移除部分显式传入,让 autotune 接管 + +### 常见问题 4:性能抖动大 + +**原因:** 小 shape 测试不稳定 + +**解决:** +```bash +export TRITON_BENCH_METHOD=npu +export TRITON_PRINT_AUTOTUNING=1 +``` + +--- + +### 问题快速排查表 + +| 现象 | 可能的原因 | 建议动作 | +| ------------------------------- | ---------------------------- | ------------------------- | +| configs=[] 直接解析失败 | split/tiling 轴没有从 DSL 唯一识别出来 | 先补 hints,再试 | +| parser 能识别一部分,但总差一个参数 | 某个参数没有和轴长度 mask 建立联系 | 改 DSL 写法或改手写 config | +| kernel 完全没有合适的 tl.constexpr 可调项 | DSL 没暴露调参接口 | 先改 kernel dsl,再谈 autotune | +| 自动生成能跑,但候选质量明显差 | 当前算法不适合该 kernel 的参数耦合方式 | 手动构造 config 传入 | + +--- + +## 总结 + +| 方式 | 适用场景 | 复杂度 | +|------|---------|--------| +| 自定义 autotune | 需要精确控制搜索空间 | 中 | +| 自动 autotune (configs=[]) | DSL 规范,参数清晰 | 低 | +| 半自动 autotune (hints) | 自动解析失败,但能人工判断 | 中 | + +**限制:** 进阶用法仅支持 Vector 类算子,不支持 Cube 类算子。 + +**优先级:** 自定义 autotune > 半自动 autotune (hints) > 自定义 autotune From 93dd17f98e2ded6f83f6432492bc962eed88c3cf Mon Sep 17 00:00:00 2001 From: wwwbby <60009003+wwwbby@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:29:50 +0800 Subject: [PATCH 6/9] Update autotune.md --- .../latency-optimizer/references/autotune.md | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/skills/latency-optimizer/references/autotune.md b/skills/latency-optimizer/references/autotune.md index 3f6a1838..221431f5 100644 --- a/skills/latency-optimizer/references/autotune.md +++ b/skills/latency-optimizer/references/autotune.md @@ -300,7 +300,7 @@ def kernel(input_ptr, output_ptr, ...): # 解析结果:指针参数 = input_ptr, output_ptr ``` -### step2.尝试自动autotune +#### step1.5.判断使用什么autotune 一般同时满足下面几条时,自动autotune的成功率比较高: * split 参数能从 `tl.program_id` 路径判断出来; @@ -310,8 +310,7 @@ def kernel(input_ptr, output_ptr, ...): * `offs < min(block_end, n)` * `key` 能和运行时 shape 参数一一对应; -下面这些情况,自动autotune可能会出现解析失败的情况: - +此时应当跳转到step2.尝试自动autotune执行。出现下面的情况,自动autotune可能会出现解析失败的情况: * 没有和轴长度直接绑定的 mask/bounds * 某个参数必须覆盖完整语义维度 * 例如 `BLOCK_SIZE >= hidden_dim` @@ -319,7 +318,17 @@ def kernel(input_ptr, output_ptr, ...): * 一个参数同时影响两个轴,或者同时影响“核数 + tile 形状” * kernel 没暴露出可调 `tl.constexpr` -自动autotune模板如下: +如果step2.尝试自动autotune失败,并且满足以下条件: + +* 能判断哪个参数属于 `split`,哪个属于 `tiling`; +* 知道每个参数对应轴的长度参数; + +此时应当跳转到step3.尝试半自动autotune + +否则,请直接尝试step4.尝试自定义autotune + +### step2.尝试自动autotune +自动autotune完全由编译器决定尝试哪些参数组合,仅需要指定`key`。自动autotune模板如下: ```python @triton.autotune( # configs 为空列表,表示不传入自定义配置 @@ -342,10 +351,7 @@ def kernel( ``` ### step3.尝试半自动autotune -`hints` 是 Triton-Ascend 在 `autotune` 装饰器中新增的一个参数,类型为 `dict`,用于给 Triton-Ascend autotune 提供该 triton kernel 的一些关键信息,帮助 autotune 更好的生成 tiling 配置。推荐显式传 `hints` 的场景: - -* 能判断哪个参数属于 `split`,哪个属于 `tiling`; -* 知道每个参数对应轴的长度参数; +`hints` 是 Triton-Ascend 在 `autotune` 装饰器中新增的一个参数,类型为 `dict`,用于给 Triton-Ascend autotune 提供该 triton kernel 的一些关键信息,帮助 autotune 更好的生成 tiling 配置。 hints 参数说明: * 当前 `hints` 参数中可以识别的字段有: From 917a1ca6aa7b056ba5d33cd36274b43a79b85132 Mon Sep 17 00:00:00 2001 From: wwwbby <60009003+wwwbby@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:31:22 +0800 Subject: [PATCH 7/9] Update core_partition.md --- .../references/core_partition.md | 830 +++++++++--------- 1 file changed, 430 insertions(+), 400 deletions(-) diff --git a/skills/latency-optimizer/references/core_partition.md b/skills/latency-optimizer/references/core_partition.md index 657c227b..e244a511 100644 --- a/skills/latency-optimizer/references/core_partition.md +++ b/skills/latency-optimizer/references/core_partition.md @@ -1,504 +1,534 @@ -# Core Partition 分核优化模式 +# 分核优化 -## 概述 +## 核心原则 -在 Triton NPU kernel 中,**分核策略直接影响硬件利用率和性能**。NPU设备有多个AI Core,选择合适的核数、分核维度和任务分配方式是性能优化的关键。 +**NPU 设备有多个 AI Core(通常 40 或 48 个),选择合适的核数是性能优化的关键。** -## 触发条件 +| 问题 | 影响 | +|------|------| +| Grid 远大于物理核数 | Kernel launch 开销大,调度开销大 | +| Grid 远小于物理核数 | 硬件利用率低,算力浪费 | +| Grid ≈ 物理核数 | **最优** | -**当 Triton 代码中存在以下情况时,应考虑优化分核策略**: +## 一、Grid 大小优化 -1. **发射核数不合理**:grid 大小与数据规模不匹配(过多或过少),npu 设备的物理核数一般为40或48,当grid 大小远大于物理核时或者远小于物理核时将会使得性能极大地弱化。 -2. **tiling大小不合理**:在For循环中调用Vector计算单元时,运算的tile数据量远小于当前设备的UB大小(通常是192KB),导致无法充分使用算力单元。 +### 问题:发射核数不合理 -## 优化方法 +```python +# 问题:Grid = (128,) 核数过多 +grid = (batch_size,) # 如果 batch_size=128,远超 48 核 +kernel[grid](...) + +# 问题:Grid = (2,) 核数过少 +grid = (batch_size // 64,) # 如果 batch_size=128,只有 2 核 +kernel[grid](...) +``` + +### 优化:匹配物理核数 + +```python +# NPU 通常有 40 或 48 个物理核 +num_cores = 48 + +# 方案1:固定核数 +grid = (num_cores,) -### 直接固定发射核数等于设备核数 +# 方案2:根据数据量调整 +grid = (triton.cdiv(total_work, work_per_core),) +# 确保 grid ≤ num_cores +``` + +### 案例:Softmax 算子 -#### 原始代码(发射核数过多) +**原始实现:Grid 过大** ```python +M = 112 # 行数 +N = 256 # 列数 + +# Grid = (M,) = 112 个核,远超 48 核 +grid = (M,) + @triton.jit -def gather_dim1_kernel( - x_ptr, # *x [B, C] - idx_ptr, # *idx[B, K] - out_ptr, # *out[B, K] - stride_xb, stride_xc, - stride_ib, stride_ik, - stride_ob, stride_ok, - B, K, - BLOCK_B: tl.constexpr, - BLOCK_K: tl.constexpr, -): - pid_b = tl.program_id(0) # 1 block per batch row - pid_k = tl.program_id(1) # 1 block per K-tile - k_off = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) - mask = k_off < K - idx = tl.load(idx_ptr + pid_b * stride_ib + k_off * stride_ik, mask=mask) # [BLOCK_K] - x_val = tl.load(x_ptr + pid_b * stride_xb + idx * stride_xc, mask=mask) - tl.store(out_ptr + pid_b * stride_ob + k_off * stride_ok, x_val, mask=mask) - -# 调用 -B = 128 # batch dim -K = 64 - -BLOCK_B = 4 -BLOCK_K = 128 - -grid = (B, triton.cdiv(K, BLOCK_K)) - -gather_dim1_kernel[grid]( - x, idx, out, - x.stride(0), x.stride(1), - idx.stride(0), idx.stride(1), - out.stride(0), out.stride(1), - B, K, - BLOCK_B=BLOCK_B, - BLOCK_K=BLOCK_K, -) +def softmax_kernel_naive(...): + row_idx = tl.program_id(0) + # 每个 program 只处理 1 行 + row_data = tl.load(ptr + row_idx * stride + col_offs, mask=mask) + # ... 计算 softmax ``` -#### 优化后代码(合理核数) +**优化后:合理核数** ```python +M = 112 +N = 256 +ROWS_PER_BLOCK = 4 # 每个 program 处理 4 行 + +# Grid = (M / ROWS_PER_BLOCK,) = 28 个核,接近最优 +grid = (triton.cdiv(M, ROWS_PER_BLOCK),) + @triton.jit -def gather_dim1_kernel( - x_ptr, # *x [B, C] - idx_ptr, # *idx[B, K] - out_ptr, # *out[B, K] - stride_xb, stride_xc, - stride_ib, stride_ik, - stride_ob, stride_ok, - B, K, - BLOCK_B: tl.constexpr, - BLOCK_K: tl.constexpr, -): - pid_b = tl.program_id(0) # 1 block per batch row -- # 原始实现 -- pid_k = tl.program_id(1) # 1 block per K-tile +def softmax_kernel_optimized(...): + pid = tl.program_id(0) + # 每个 program 处理 ROWS_PER_BLOCK 行 + row_offs = pid * ROWS_PER_BLOCK + tl.arange(0, ROWS_PER_BLOCK) + row_mask = row_offs < M + + # 2D 加载:[ROWS_PER_BLOCK, N] + row_data = tl.load( + ptr + row_offs[:, None] * stride + col_offs[None, :], + mask=row_mask[:, None] & col_mask[None, :] + ) + # ... 计算 softmax(向量化处理多行) +``` -- k_off = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) -- mask = k_off < K +**收益对比:** -- idx = tl.load(idx_ptr + pid_b * stride_ib + k_off * stride_ik, mask=mask) # [BLOCK_K] +| 指标 | 原始 | 优化后 | +|------|------|--------| +| Grid 大小 | 112 | 28 | +| 每个 program 处理 | 1 行 = 256 元素 | 4 行 = 1024 元素 | +| 核数利用率 | 过饱和 | **接近最优** | -- x_val = tl.load(x_ptr + pid_b * stride_xb + idx * stride_xc, mask=mask) +## 二、UB 大小约束 -- tl.store(out_ptr + pid_b * stride_ob + k_off * stride_ok, x_val, mask=mask) +**NPU UB (Unified Buffer) 通常为 192KB,tile 大小必须满足:** -+ # 优化后实现使用向量化处理,一次处理一整个BLOCK_B,因此这里的得到的是一个向量 -+ b_idx = pid_b * BLOCK_B + tl.arange(0, BLOCK_B) -+ b_mask = b_idx < B # 需要判断是否越界 +``` +tile_size × dtype_size × buffers ≤ 192KB +``` -+ # 对 K 维进行循环,向量化处理BLOCK_B * BLOCK_K个数据 -+ for k_start in range(0, K, BLOCK_K): -+ ks = tl.arange(0, BLOCK_K) -+ k_mask = ks < K - k_start +### 计算示例 -+ idx_off = (b_idx[:, None] * stride_ib + -+ (k_start + ks)[None, :] * stride_ik) -+ col_idx = tl.load(idx_ptr + idx_off, mask=b_mask[:, None] & k_mask) +```python +# float32 (4 bytes) +# 单 buffer 最大元素数 +max_elements = 192 * 1024 / 4 = 49152 + +# 考虑多 buffer (如 load + store) +# 单 buffer 最大元素数 +max_elements_per_buffer = 192 * 1024 / 4 / 2 = 24576 + +# 常见的 BLOCK_SIZE 选择 +BLOCK_SIZE = 8192 # 32KB,安全 +BLOCK_SIZE = 16384 # 64KB,安全 +BLOCK_SIZE = 32768 # 128KB,接近上限 +``` -+ x_off = (b_idx[:, None] * stride_xb + -+ col_idx * stride_xc) -+ x_val = tl.load(x_ptr + x_off, mask=b_mask[:, None] & k_mask) +### Tiling 大小选择 -+ out_off = (b_idx[:, None] * stride_ob + -+ (k_start + ks)[None, :] * stride_ok) -+ tl.store(out_ptr + out_off, x_val, mask=b_mask[:, None] & k_mask) +| BLOCK_SIZE | 内存占用 (float32) | 安全性 | +|-----------|------------------|--------| +| 4096 | 16 KB | ✅ 非常安全 | +| 8192 | 32 KB | ✅ 安全 | +| 16384 | 64 KB | ✅ 安全 | +| 32768 | 128 KB | ⚠️ 接近上限 | +| 49152 | 192 KB | ❌ 可能溢出 | -# 调用 -B = 128 # batch dim -K = 64 +**原则:在不溢出的前提下,尽量使用大的 BLOCK_SIZE** -BLOCK_B = 4 -BLOCK_K = 128 +## 三、多行并行优化 (1D → 2D Tiling) -— # 原始grid较大,每个核心处理BLOCK_K个数据,分核数=B*K/BLOCK_K -- grid = (B, triton.cdiv(K, BLOCK_K)) -+ # 优化后grid变小,每个核心处理BLOCK_B*K个数据,分核数=B/BLOCK_B,内部展开循环处理BLOCK_K个数据 -+ grid = (triton.cdiv(B, BLOCK_B),) +### 问题描述 -gather_dim1_kernel[grid]( - x, idx, out, - x.stride(0), x.stride(1), - idx.stride(0), idx.stride(1), - out.stride(0), out.stride(1), - B, K, - BLOCK_B=BLOCK_B, - BLOCK_K=BLOCK_K, -) +**问题:** 每个 program 只处理 1 行数据,导致 kernel launch 开销大,向量化效率低。 + +```python +# 问题代码:每个 program 处理 1 行 +row_idx = tl.program_id(0) +x = tl.load(ptr + row_idx * stride + cols, mask=mask) +# ... 处理 1 行数据 ``` -### 使用Triton-Ascend autotune搜索最佳分核参数 -Triton-Ascend autotune是一个Triton-Ascend提供的tiling超参数性能调优工具,遍历搜索空间,尝试不同参数组合,展示每个组合的运行耗时与最优组合。使用Triton-Ascend autotune需要遵循以下工作流程: -1. 识别出 triton kernel 中哪些 `tl.constexpr` 参数是自由可调的 tiling 参数,包括影响分核(split)和切块(tiling)大小的参数,这里的分核指的是影响 grid 大小,而切块指的是影响 tile 大小,即影响 `tl.load` 或是 `tl.make_block_ptr` 产生的数据大小。 -2. 如果这些参数能从 `tl.program_id`、`tl.arange`、`tl.range/range`、`mask/bounds` 表达式中被唯一识别出来,就尝试使用自动生成tiling `configs=[]` -3. 如果 kernel 语义上适合自动 tiling,但 DSL 写法让 parser 解析不出来,就显式传 `hints` -4. 如果某些 tiling 参数不可自由调整,例如某 kernel dsl 写法要求 grid 第一维必须固定为 `batch_size` 大小,或者根本没有暴露出可调的 tiling 参数,此时建议直接手写 triton.Config。 - -#### 使用方法 -@triton.autotune 入参列表 -| 参数名 | 类型 | 必填 | 说明 | -| --------- | ------------------------------ | ----- | --------------------------------------------------------------- | -| `configs` | `list[Config]` | 否 | 用户自定义的调优配置列表,为空时自动生成 | -| `key` | `dict[str, str]` / `list[str]` | **是** | 缓存键,指定哪些参数变化时需要重新调优。使用 `hints` 时必须用**字典形式**,如 `{"x": "n_rows"}` | -| `hints` | `dict` | 否 | **Ascend 扩展参数**,用于显式指定轴与 tiling 参数的映射关系 | - -hints 字典内部字段 -| 字段名 | 类型 | 必填 | 说明 | -| ----------------- | ---------------- | ----- | ------------------------------------------------------ | -| `split_params` | `dict[str, str]` | **是** | 分核参数映射,如 `{"x": "BLOCK_M"}` 表示沿 `x` 轴切 program | -| `tiling_params` | `dict[str, str]` | **是** | 切块参数映射,如 `{"y": "BLOCK_N"}` 表示沿 `y` 轴切 block | -| `low_dim_axes` | `list[str]` | **是** | 低维轴列表,如 `["y"]`,用于优化 tiling 效果 | -| `reduction_axes` | `list[str]` | **是** | 规约轴列表,如 `[]`,用于优化 tiling 效果 | -| `auto_gen_config` | `bool` | 否 | 是否自动生成 tiling 配置,默认 `True`;当 `configs` 非空时默认变为 `False` | - -有三种使用方法,以下是这三种方法的使用模板 - -##### 自动解析模板 -详情见 Step 2:判断是否能用 configs=[] 自动生成 tiling +### 优化方案 + +**方案:** 每个 program 处理 `ROWS_PER_BLOCK` 行,使用 2D 向量化加载。 + ```python -@triton.autotune( - configs=[], - key=["n_rows"], +# 优化代码:每个 program 处理多行 +pid = tl.program_id(0) +row_offs = pid * ROWS_PER_BLOCK + tl.arange(0, ROWS_PER_BLOCK) # 多行索引 +row_mask = row_offs < n_rows + +# 2D 加载:[ROWS_PER_BLOCK, BLOCK_SIZE] +x = tl.load( + ptr + row_offs[:, None] * stride + col_offs[None, :], + mask=row_mask[:, None] & col_mask[None, :], ) -@triton.jit -def kernel( - x_ptr, - y_ptr, - n_rows, - BLOCK_M: tl.constexpr, -): - pid = tl.program_id(0) - offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) - mask = offs < n_rows - x = tl.load(x_ptr + offs, mask=mask, other=0) - tl.store(y_ptr + offs, x, mask=mask) ``` -##### 显式给 hints -详情见 Step 3:自动解析出错时,显式传 hints + +### 性能收益 + +| 指标 | 原始 | 优化后 | +|-----|------|--------| +| Grid 大小 | `(M,)` | `(M / ROWS_PER_BLOCK,)` | +| Kernel launch 开销 | 高 | 降低 ROWS_PER_BLOCK 倍 | +| 向量化效率 | 1D | 2D 向量化 | +| 典型加速 | 1x | **10-30x** | + +### ROWS_PER_BLOCK 选择 + +| N (列数) | 推荐 ROWS_PER_BLOCK | 原因 | +|---------|-------------------|------| +| < 256 | 16-32 | 列向量化已足够 | +| 256-1024 | 8-16 | 平衡寄存器压力 | +| > 1024 | 4-8 | 避免寄存器溢出 | + +### 变换规则 + +| 单行版本 | 多行版本 | +|---------|---------| +| `row_id = tl.program_id(0)` | `row_offs = pid * ROWS_PER_BLOCK + tl.arange(0, ROWS_PER_BLOCK)` | +| `grid = (M,)` | `grid = (triton.cdiv(M, ROWS_PER_BLOCK),)` | +| 1D 索引 `ptr + row_id * stride` | 2D 索引 `ptr + row_offs[:, None] * stride` | + +### 模板代码 + +#### 模式 1: Row-wise Reduction (softmax, row-max, row-sum) + ```python -@triton.autotune( - configs=[], - key={"x": "n_rows", "y": "n_cols"}, - hints={ - "split_params": {"x": "BLOCK_M"}, - "tiling_params": {"y": "BLOCK_N"}, - "low_dim_axes": ["y"], - "reduction_axes": [], - }, -) @triton.jit -def kernel( - x_ptr, - y_ptr, - n_rows, - n_cols, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, +def row_reduce_kernel( + input_ptr, output_ptr, + stride_in, stride_out, + n_rows, n_cols, + BLOCK: tl.constexpr, + ROWS_PER_BLOCK: tl.constexpr, ): - ... + pid = tl.program_id(0) + row_offs = pid * ROWS_PER_BLOCK + tl.arange(0, ROWS_PER_BLOCK) + row_mask = row_offs < n_rows + col_offs = tl.arange(0, BLOCK) + col_mask = col_offs < n_cols + + # 2D 加载 + x = tl.load( + input_ptr + row_offs[:, None] * stride_in + col_offs[None, :], + mask=row_mask[:, None] & col_mask[None, :], + other=0.0 + ) + + # 沿列方向 reduce (每行一个值) + row_result = tl.max(x, axis=1, keep_dims=True) # 或 tl.sum + + # 存储结果 + tl.store(output_ptr + row_offs[:, None], row_result, mask=row_mask[:, None]) ``` -##### 完全手写 triton.Config +#### 模式 2: Element-wise Operation (activation, copy) + ```python -@triton.autotune( - configs=[ - triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "multibuffer": True}), - triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "multibuffer": True}), - triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "multibuffer": False}), - ], - key=["n_rows", "n_cols"], -) @triton.jit -def kernel( - x_ptr, - y_ptr, - n_rows, - n_cols, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, +def elementwise_kernel( + input_ptr, output_ptr, + stride_in, stride_out, + n_rows, n_cols, + BLOCK: tl.constexpr, + ROWS_PER_BLOCK: tl.constexpr, ): - ... + pid = tl.program_id(0) + row_offs = pid * ROWS_PER_BLOCK + tl.arange(0, ROWS_PER_BLOCK) + row_mask = row_offs < n_rows + col_offs = tl.arange(0, BLOCK) + col_mask = col_offs < n_cols + + # 2D 加载 + x = tl.load( + input_ptr + row_offs[:, None] * stride_in + col_offs[None, :], + mask=row_mask[:, None] & col_mask[None, :], + ) + + # 逐元素操作 (自动向量化) + y = tl.exp(x) # 或其他 elementwise 操作 + + # 2D 存储 + tl.store( + output_ptr + row_offs[:, None] * stride_out + col_offs[None, :], + y, + mask=row_mask[:, None] & col_mask[None, :], + ) ``` -#### 注意事项 +#### 模式 3: Reduction + Element-wise (softmax, layer-norm) -1. `@triton.autotune` 需要直接包在 `@triton.jit` 外层,示例如下: ```python -@triton.autotune(...) @triton.jit -def kernel(...): - ... +def softmax_kernel( + input_ptr, output_ptr, + stride_in, stride_out, + n_rows, n_cols, + BLOCK: tl.constexpr, + ROWS_PER_BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + row_offs = pid * ROWS_PER_BLOCK + tl.arange(0, ROWS_PER_BLOCK) + row_mask = row_offs < n_rows + col_offs = tl.arange(0, BLOCK) + col_mask = col_offs < n_cols + mask_2d = row_mask[:, None] & col_mask[None, :] + + # 加载数据 + x = tl.load( + input_ptr + row_offs[:, None] * stride_in + col_offs[None, :], + mask=mask_2d, other=-float('inf') + ) + + # Phase 1: 计算 max + row_max = tl.max(x, axis=1, keep_dims=True) + + # Phase 2: 计算 sum(exp(x - max)) + x_shifted = x - row_max + exp_x = tl.exp(x_shifted) + row_sum = tl.sum(tl.where(mask_2d, exp_x, 0.0), axis=1, keep_dims=True) + + # Phase 3: 计算 output + output = x_shifted - tl.log(row_sum) + + # 存储 + tl.store( + output_ptr + row_offs[:, None] * stride_out + col_offs[None, :], + output, mask=mask_2d + ) ``` -2. 自动生成 tiling 功能主要面向 vector kernel:当前 Triton-Ascend 这套自动解析/自动生成 tiling 面向的是 vector 类 kernel,且该 kernel 中所有分核和分块参数均可调。Cube 类算子目前还不支持自动 tiling 生成。 +## 四、分核策略选择 + +### 策略对比 -#### 详细步骤 -##### Step 1:先识别哪些参数真的是可调参数 -###### 1.1 先看“哪些 tl.constexpr 没有在 launch 时显式传入” -Triton-Ascend autotuner 在自动解析 split/tiling 参数时,首先会看 kernel 调用时 哪些参数没有传入,把这些“缺省的参数”当成候选项。 +| 策略 | 适用场景 | Grid 大小 | +|------|---------|---------| +| **一维分核** | 单维度处理(如逐行) | (N / BLOCK,) | +| **二维分核** | 矩阵运算(如 matmul) | (M / BM, N / BN) | +| **多行并行** | 行级 reduce(如 softmax) | (M / ROWS_PER_BLOCK,) | -简单理解: +### 选择依据 + +1. **计算数据量与核数匹配** + - 总数据量 / 每个 program 处理量 ≈ 物理核数 + +2. **避免过度细粒度** + - 每个 program 处理足够数据(至少几千元素) + +3. **考虑内存访问模式** + - 连续访问优于随机访问 + - 2D 加载优于多次 1D 加载 + +## 常见错误 + +### 错误 1:Grid 远超物理核数 -* Tensor 参数不可能是自动解析候选项; -* 普通运行时 shape 参数(如 n_rows、n_cols)通常属于 `key`; -* 真正的候选项通常是没有在 launch 处显式传值的 `tl.constexpr`; -* 如果某个 tl.constexpr 已经在 launch 时手动写死,它就不会再被当成自动解析候选项。 -例如: ```python -@triton.jit -def kernel( - x_ptr, - y_ptr, - n_rows, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, # BLOCK_N是自动解析候选项 -): - ... +# ❌ 错误:Grid = (1024,) 远超 48 核 +grid = (batch_size * height * width // BLOCK,) -# BLOCK_M 已显式传入,不再是自动解析候选项 -kernel[grid](x, y, n_rows, BLOCK_M=128) +# ✅ 正确:Grid ≈ 48 +grid = (num_cores,) ``` -###### 1.2 如何识别 split 参数 +### 错误 2:Tile 过小 -split 参数控制“一个 program 负责多大的一块数据”,它最常见的写法特征是: -1. 和 `tl.program_id(...)` 有直接关系; -2. 参与构造 block 起始位置; -3. 最后能通过 mask/bounds 表达式对应回某个 shape 轴。 -例如: ```python -# 一维切分 -pid = tl.program_id(0) -offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)# 可以知道BLOCK_M 是 split 参数 -mask_m = offs_m < n_rows - -# 二维切分 -pid_m = tl.program_id(0) -pid_n = tl.program_id(1) -offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]# 可以知道BLOCK_M 是 split 参数 -offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]# 可以知道BLOCK_N 是 split 参数 -mask_m = offs_m < n_rows -mask_n = offs_n < n_cols -``` +# ❌ 错误:BLOCK_SIZE = 64,UB 利用率低 +BLOCK_SIZE = 64 # 只用 256 bytes -###### 1.3 如何识别 tiling 参数 -tiling 参数控制“在一个大的 split block 内,再按多大的子块去迭代”,它最常见的写法特征是: +# ✅ 正确:BLOCK_SIZE = 8192,充分利用 UB +BLOCK_SIZE = 8192 # 用 32KB +``` -1. 出现在 `tl.arange(0, PARAM)` 中; -2. 同时还出现在 for 循环的步长或循环次数推导中; -3. 最后能通过 mask/bounds 对应回某个轴长度参数。 +### 错误 3:忽略 UB 上限 -例如: ```python -# 典型形态 1:步长是 tiling 参数 -for k0 in tl.range(0, BLOCK_K, BLOCK_K_SUB):# 可以知道BLOCK_K_SUB 是 tiling 参数 - offs_k = k0 + tl.arange(0, BLOCK_K_SUB)# 可以知道BLOCK_K_SUB 是 tiling 参数 - mask_k = offs_k < k_size - -# 典型形态 2:先计算循环次数 -num_k_tiles = (k_size + BLOCK_K_SUB - 1) // BLOCK_K_SUB# 可以知道BLOCK_K_SUB 是 tiling 参数 -for tile_id in range(num_k_tiles): - offs_k = tile_id * BLOCK_K_SUB + tl.arange(0, BLOCK_K_SUB)# 可以知道BLOCK_K_SUB 是 tiling 参数 - mask_k = offs_k < k_size +# ❌ 错误:BLOCK_SIZE 过大导致 UB 溢出 +BLOCK_SIZE = 65536 # 256KB > 192KB UB + +# ✅ 正确:确保不溢出 +BLOCK_SIZE = 32768 # 128KB < 192KB UB ``` -##### Step 2:判断是否能用 configs=[] 自动生成 tiling +## 五、NPU vs GPU 分核语义差异 -当你已经找到了候选参数后,可以按下面的检查表判断。 -###### 2.1 可以优先尝试 configs=[] 的情况 -一般同时满足下面几条时,`configs=[]` 成功率比较高: +### 关键差异 -* split 参数能从 `tl.program_id` 路径判断出来; -* tiling 参数能从 `tl.arange + for(range/tl.range)` 路径判断出来; -* 每个轴都有比较清晰的 mask/bounds 表达式,例如: - * `offs < n` - * `offs < min(block_end, n)` -* `key` 能和运行时 shape 参数一一对应; +| 特性 | GPU | NPU | +|------|-----|-----| +| Grid 含义 | 逻辑并行实例数 | 直接映射到物理核 | +| 超额订阅 | SM 会自动调度 | AI Core 按顺序执行 | +| 最优 Grid | 可远大于 SM 数 | 应接近 AI Core 数 | +| 多核利用率 | SM 动态调度 | 静态绑定 | -###### 2.2 不适合直接用 configs=[] 的常见信号 +**GPU 行为:** Grid 可以远大于 SM 数,GPU 会自动调度,多余的 block 在 SM 上排队等待。 -下面这些情况,直接走自动 tiling 生成可能会出现解析失败的情况: +**NPU 行为:** Grid 直接映射到 AI Core,超额部分串行执行,导致调度开销累积。 -* 没有和轴长度直接绑定的 mask/bounds -* 某个参数必须覆盖完整语义维度 - * 例如 `BLOCK_SIZE >= hidden_dim` -* grid 某一维被业务语义固定,不允许自由切块 -* 一个参数同时影响两个轴,或者同时影响“核数 + tile 形状” -* kernel 没暴露出可调 `tl.constexpr` +### 案例对比 -###### 2.3 如果出现错误建议打开调试日志 -排查问题时,建议先启动环境变量debug: -``` -export TRITON_PRINT_AUTOTUNING=1 -``` +```python +# 场景:处理 128 个 batch -日志中可以直接看到: +# GPU 风格(不适用于 NPU) +grid = (128,) # GPU: 128 block 在 SM 上动态调度,效率高 + # NPU: 128 个核串行执行,调度开销大 -* 识别出的 split axes; -* 识别出的 tiling axes; -* 识别出的 low-dimensional axes; -* 识别出的 reduction axes; -* 生成的 config 数量。 -小 shape 算子如果 benchmark 抖动大,也可以按需开启: -``` -export TRITON_BENCH_METHOD=npu +# NPU 优化风格 +grid = (48,) # NPU: 48 核并行,每个处理 128/48 ≈ 3 个 batch ``` -这会测试的时间更准确,但 autotune 时间也会明显变长。 -##### Step 3:自动解析出错时,显式传 hints +## 六、Tiling 策略选择 -如果你已经确认 自动 tiling 适用于该 triton kernel,只是因为 DSL 写法不能够被当前 parser 识别,这时可以尝试显式传 `hints`。 -`hints` 是 Triton-Ascend 在 `autotune` 装饰器中新增的一个参数,类型为 `dict`,用于给 Triton-Ascend autotune 提供该 triton kernel 的一些关键信息,帮助 autotune 更好的生成 tiling 配置。 +### 案例:GELU 算子 -###### 3.1 什么情况下可以考虑 hints +**Easy Kernel(简单但效率低):** -推荐显式传 `hints` 的场景: +```python +@triton.jit +def gelu_kernel_easy( + x_ptr, y_ptr, n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < n_elements -* 你能人工判断哪个参数属于 `split`,哪个属于 `tiling`; -* 你知道每个参数对应轴的长度参数; + x = tl.load(x_ptr + offs, mask=mask) + # GELU(x) = x * 0.5 * (1 + erf(x / sqrt(2))) + y = x * 0.5 * (1.0 + tl.erf(x * 0.7071067811865475)) + tl.store(y_ptr + offs, y, mask=mask) -###### 3.2 hints 参数说明 -* 当前 `hints` 参数中可以识别的字段有: - * `split_params`: dict[str, str],分核参数的映射关系,例如 `{"x": "BLOCK_M"}` 表示 `BLOCK_M` 是沿 `x` 轴切 program - * `tiling_params`:dict[str, str],切块参数的映射关系,例如 `{"y": "BLOCK_N"}` 表示 `BLOCK_N` 是沿 `y` 轴切 block - * `low_dim_axes`:list[str],低维轴的列表,例如 `["y"]` 表示 `y` 轴是低维轴 - * `reduction_axes`:list[str],规约轴的列表,例如 `[]` 表示没有规约轴 - * `auto_gen_config`:bool,是否自动生成 tiling 配置,默认值为 `True` -* 注意: - * 通过 `hints` 来显示指定轴关系时,autotune 中原本的参数 `key` 必须改为字典形式传入,因为后续 `split_params`、`tiling_params` 等参数都是按轴名来填写,需要和 `key` 里的轴名对应起来 - * 通过 `hints` 来显示指定轴关系时,`split_params`、`tiling_params`、`low_dim_axes`、`reduction_axes` 必须传入,即使某些参数为空 - * 合法的轴名称是 `x/y/z/w/v/t`,仅仅用做关系映射 - * `split_params` 和 `tiling_params` 为自动生成 tiling 算法必须的输入,`low_dim_axes` 和 `reduction_axes` 为 tiling 算法的可选输入,用于优化 tiling 效果,留空时 tiling 也能够自动生成,但可能会影响生成的候选 tiling 数量和质量 - * 当用户传入的 configs 不为空时,`auto_gen_config` 默认值为 `False`,如果希望此时也希望自动生成 tiling 配置并与用户传入的 configs 合并,需要显式在 `hints` 中传如入 `"auto_gen_config": True` +# Grid = (n_elements // BLOCK_SIZE,) +# 如果 n_elements = 128 * 1024 * 1024, BLOCK_SIZE = 1024 +# Grid = 131072,远超 48 核 +``` + +**Better Kernel(优化的 Tiling 策略):** -使用示例: ```python -import triton -import triton.language as tl -import triton.backends.ascend.runtime - -@triton.autotune( - # configs 为空列表,表示不传入自定义配置 - # 此时 auto_gen_config 默认为 True,会自动生成 tiling 配置 - configs=[], - - # key 使用字典形式,轴名必须与 hints 中的轴名对应 - # "x" 对应 n_rows(行数),"y" 对应 n_cols(列数) - # autotune 会根据这些维度值来缓存和选择最佳配置 - key={"x": "n_rows", "y": "n_cols"}, - - # hints 参数:显式指定轴与 tiling 参数的映射关系 - hints={ - # split_params: 分核参数映射,指定沿哪个轴切分 program(任务) - # "x": "BLOCK_M" 表示 BLOCK_M 沿 x 轴切分,即按行方向分核 - # 每个 program 处理 BLOCK_M 行数据 - "split_params": {"x": "BLOCK_M"}, - - # tiling_params: 切块参数映射,指定沿哪个轴切分 block(数据块) - # "y": "BLOCK_N" 表示 BLOCK_N 沿 y 轴切分,即按列方向切块 - # 每行数据在列方向上被切分为 BLOCK_N 大小的块,通过 for 循环处理 - "tiling_params": {"y": "BLOCK_N"}, - - # low_dim_axes: 低维轴列表,用于优化 tiling 效果 - # ["y"] 表示 y 轴(列方向)是低维轴,访问连续性更好,适合作为内层循环 - "low_dim_axes": ["y"], - - # reduction_axes: 规约轴列表,本 kernel 无规约操作(如 sum/max 等) - # 为空列表表示没有规约轴 - "reduction_axes": [], - - # auto_gen_config: 默认为 True,表示自动生成 tiling 配置 - # 由于 configs 为空,此处使用默认值 True 即可,无需显式传入 - }, -) @triton.jit -def kernel_with_hints( - x_ptr, - y_ptr, - n_rows, - n_cols, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, +def gelu_kernel_better( + x_ptr, y_ptr, n_elements, + BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) - offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] + start = pid * BLOCK_SIZE + end = tl.minimum(start + BLOCK_SIZE, n_elements) + + # 循环处理这个 block 负责的所有数据 + for i in range(start, end, 256): # 内部 tiling + offs = i + tl.arange(0, 256) + mask = offs < n_elements - for n0 in range(0, n_cols, BLOCK_N): - offs_n = n0 + tl.arange(0, BLOCK_N)[None, :] - mask_m = offs_m < n_rows - mask_n = offs_n < n_cols - mask = mask_m & mask_n + x = tl.load(x_ptr + offs, mask=mask) + y = x * 0.5 * (1.0 + tl.erf(x * 0.7071067811865475)) + tl.store(y_ptr + offs, y, mask=mask) - x = tl.load(x_ptr + offs_m * n_cols + offs_n, mask=mask, other=0) - tl.store(y_ptr + offs_m * n_cols + offs_n, x, mask=mask) +# Grid = (48,),直接使用物理核数 +# 每个核处理 n_elements / 48 个数据 ``` -##### Step 4:手写 triton.Config -如果 Step 2 判断该 triton kernel 不适合或者无法使用自动 tiling 生成,那么可以使用社区 triton autotune 的基本功功能:手写一组 `triton.Config` 传入参数 `configs` 中。 -###### 手写 triton.Config 的总体原则 -1. 对于影响 grid 发射核数的参数,一般我们尽量让其能够等于物理核数,如果数据量较小,也可能发射较少核数的时候能获得最优性能;对于影响 tile 块大小的参数,我们尽量在不产生 UB overflow 的情况下让其尽可能大,同时避免尾块的产生 -2. 影响 grid 发射核数的参数:可以按照总长度从高到低设置为 X, X/2, X/4 等等的值,如果输入 shape 较大,可以设置为让 grid 发射核数正好等于物理核数的大小,例如 (X + num_cores - 1) // num_cores;这里是以一个切分轴为例,如果存在多个切分轴,那么就需要按照乘积来计算 -3. 影响 tile 块大小的参数:起始值为切分轴参数(如果存在)或者轴长度,注意当该轴长度特别大的时候,我们可以直接从 16384 这样一个经验值开始取,然后按照 X / 2, X / 4 这样去取值 -4. 上述按 2 的幂次方下降的值采样较为粗粒度,如果用户想要得到极致的最优性能,尤其是在输入大小不规则的情况下,需要在可能的最优区间内细粒度撒点,可以通过粗粒度采样后确认性能最优的大致区间后再进一步细分来实现。 -5. 对于 vector 类算子,在设置了上述 tiling 大小的配置候选集后,可以加上 multibuffer 编译选项的调优。 +**收益对比:** + +| 指标 | Easy Kernel | Better Kernel | +|------|-------------|---------------| +| Grid 大小 | 131072 | 48 | +| 核数利用率 | 过饱和,调度开销大 | 最优 | +| 调度开销 | 131072 次 kernel launch | 48 次内部循环 | + +### Tiling 策略选择原则 -示例: ```python -import triton -import triton.language as tl -import triton.backends.ascend.runtime +# 策略 1:外部 tiling(适合小数据量) +# Grid = (n_elements // BLOCK_SIZE,) +# 每个 program 处理 BLOCK_SIZE 数据 +# 适用:n_elements < 48 * BLOCK_SIZE + +# 策略 2:内部 tiling(适合大数据量) +# Grid = (num_cores,) +# 每个 program 内部循环处理 +# 适用:n_elements >> 48 * BLOCK_SIZE +``` +## 七、编译优化选项 -def get_configs(): - return [ - triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "multibuffer": MB}) - for BM in [256, 128, 64, 32] - for BN in [128, 64, 32, 16] - for MB in [True, False] - ] +### 常用编译选项 +| 选项 | 作用 | 适用场景 | +|------|------|---------| +| `multibuffer=True` | 启用多缓冲,隐藏内存延迟 | Vector 算子,内存密集型 | +| `multibuffer=False` | 禁用多缓冲 | 计算密集型,寄存器压力大 | +| `unit_flag=True` | 生成独立的计算单元 | 简单算子,无复杂控制流 | +| `unit_flag=False` | 不生成独立单元 | 复杂算子,有分支 | -@triton.autotune( - configs=get_configs(), - key=["n_rows", "n_cols"], -) +### 使用方式 + +```python @triton.jit -def manual_config_kernel( - x_ptr, - y_ptr, - n_rows, - n_cols, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - pid = tl.program_id(0) - offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] - offs_n = tl.arange(0, BLOCK_N)[None, :] - mask = (offs_m < n_rows) & (offs_n < n_cols) +def kernel(...): + ... + +# 在调用时指定 +kernel[grid]( + ..., + multibuffer=True, # 编译选项 + unit_flag=True, +) +``` + +### 选择建议 + +| 算子类型 | multibuffer | unit_flag | 原因 | +|---------|-------------|-----------|------| +| Element-wise (add, mul, gelu) | True | True | 简单、内存密集 | +| Reduction (sum, mean) | True | False | 有归约操作 | +| 复杂控制流 | False | False | 寄存器压力大 | + +## 常见错误 - x = tl.load(x_ptr + offs_m * n_cols + offs_n, mask=mask, other=0) - tl.store(y_ptr + offs_m * n_cols + offs_n, x, mask=mask) +### 错误 1:Grid 远超物理核数 + +```python +# ❌ 错误:Grid = (1024,) 远超 48 核 +grid = (batch_size * height * width // BLOCK,) + +# ✅ 正确:Grid ≈ 48 +grid = (num_cores,) ``` -##### 常见失败速查 -下面这张表可以直接用来决定你下一步该怎么做。 +### 错误 2:Tile 过小 -| 现象 | 更可能的原因 | 建议动作 | -| ------------------------------- | ---------------------------- | ------------------------- | -| configs=[] 直接解析失败 | split/tiling 轴没有从 DSL 唯一识别出来 | 先补 hints,再试 | -| parser 能识别一部分,但总差一个参数 | 某个参数没有和轴长度 mask 建立联系 | 改 DSL 写法或改手写 config | -| kernel 完全没有合适的 tl.constexpr 可调项 | DSL 没暴露调参接口 | 先改 kernel dsl,再谈 autotune | -| 自动生成能跑,但候选质量明显差 | 当前算法不适合该 kernel 的参数耦合方式 | 手动构造 config 传入 | +```python +# ❌ 错误:BLOCK_SIZE = 64,UB 利用率低 +BLOCK_SIZE = 64 # 只用 256 bytes +# ✅ 正确:BLOCK_SIZE = 8192,充分利用 UB +BLOCK_SIZE = 8192 # 用 32KB +``` + +### 错误 3:忽略 UB 上限 -## 性能收益 +```python +# ❌ 错误:BLOCK_SIZE 过大导致 UB 溢出 +BLOCK_SIZE = 65536 # 256KB > 192KB UB -- **核数优化**:可提升 2-5x 性能 +# ✅ 正确:确保不溢出 +BLOCK_SIZE = 32768 # 128KB < 192KB UB +``` + +### 错误 4:GPU 风格分核 + +```python +# ❌ 错误:GPU 风格,Grid 远大于核数 +grid = (n_elements // 1024,) # 可能是 131072 + +# ✅ 正确:NPU 风格,Grid 匹配物理核数 +grid = (48,) +# 每个 program 内部循环处理大数据 +``` -## 注意事项 +## 总结 -1. **UB 容量约束**:确保 tile_size * dtype_size * buffers <= 192KB -2. **原子操作开销**:原子操作有额外开销,核数过多时可能成为瓶颈 +| 优化点 | 原则 | 方法 | +|-------|------|------| +| Grid 大小 | ≈ 物理核数 (40-48) | 调整每个 program 处理的数据量 | +| Tile 大小 | 接近但不超过 UB | BLOCK_SIZE ≤ 49152 (float32) | +| 多行并行 | 减少 Grid 大小 | 2D 向量化加载 | +| Tiling 策略 | 大数据用内部 tiling | Grid=48,内部循环 | +| 编译选项 | 内存密集用 multibuffer | multibuffer=True | From d0808456468dc46987813e12400b0749d5ccb527 Mon Sep 17 00:00:00 2001 From: wwwbby <60009003+wwwbby@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:34:23 +0800 Subject: [PATCH 8/9] Update SKILL.md --- skills/latency-optimizer/SKILL.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/skills/latency-optimizer/SKILL.md b/skills/latency-optimizer/SKILL.md index b401ca71..fbe7a83e 100644 --- a/skills/latency-optimizer/SKILL.md +++ b/skills/latency-optimizer/SKILL.md @@ -28,7 +28,8 @@ argument-hint: > | 入参静态化 | `references/constexpr_parameters.md` | | Int32 向量加法 | `references/int32_vector_add.md` | | Load 指令重排序 | `references/load-order.md` | -| 分核策略优化 | `references/core_partition.md` | +| 分核策略优化 | `references/vector_core_partition.md` | +| 分核超参数自动搜索调优 | `references/autotune.md` | ### 以下文档通过分析已有代码特征,按需加载 From 47ef5b2d4131af98d3fffa148fdb5b401bb71a0f Mon Sep 17 00:00:00 2001 From: wwwbby Date: Thu, 9 Apr 2026 20:35:22 +0800 Subject: [PATCH 9/9] rename skill --- .../references/{core_partition.md => vector_core_partition.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename skills/latency-optimizer/references/{core_partition.md => vector_core_partition.md} (100%) diff --git a/skills/latency-optimizer/references/core_partition.md b/skills/latency-optimizer/references/vector_core_partition.md similarity index 100% rename from skills/latency-optimizer/references/core_partition.md rename to skills/latency-optimizer/references/vector_core_partition.md