SuperAPI AI Directory
Partner with us中文AI newsFavorites
Skillsexternal-cannbot-ops-triton-op-coding

external-cannbot-ops-triton-op-coding

Agents & models

View repository

Updated 10-15 17:19

About this skill

Triton Ascend Code Generation Skill

Your task is to generate optimized kernel code based on the following fixed configuration:

  • Target DSL: triton_ascend
  • Target Framework: torch
  • Target Backend: ascend
  • Target Architecture: {{ arch }}

Core Constraint: PyTorch Degradation Prohibited

⚠️ The generated code must be a pure Triton Ascend implementation; degrading to PyTorch is prohibited.

Operations Prohibited in forward()

Prohibited Operation Example Reason
torch computation functions torch.matmul(x, w), torch.relu(x), torch.sum(x) Must be implemented in the @triton.jit kernel
torch.nn.functional F.softmax(x, dim=-1), F.linear(x, w), F.relu(x) Must be implemented in the @triton.jit kernel
tensor method computations x.sum(), x.mean(), x.softmax(dim=-1), x.relu() Must be implemented in the @triton.jit kernel
tensor operators x @ w, x + y, x * y, x / y Must be implemented in the @triton.jit kernel
nn.Module calls self.conv(x), self.linear(x), self.layer(x) Must be implemented in the @triton.jit kernel

Operations Allowed in forward()

Allowed Operation Example Description
buffer allocation torch.empty(shape), torch.zeros(shape), torch.ones(shape) Used to store kernel output
shape operations x.view(...), x.reshape(...), x.permute(...), x.transpose(...) Does not involve computation
metadata queries x.shape, x.dtype, x.device, x.numel() Used for grid computation
kernel launch kernel[grid](...args) Call a custom @triton.jit kernel

❌ Incorrect Example (degraded to PyTorch)

# ❌ 错误 1：完全无 kernel，纯 PyTorch
def forward(self, x, w):
    return torch.matmul(x, w)

# ❌ 错误 2：有 kernel 但 forward 未调用
@triton.jit
def matmul_kernel(...):
    pass

def forward(self, x, w):
    return torch.matmul(x, w)  # kernel 定义了但没用

# ❌ 错误 3：混合实现（部分 kernel + 部分 torch）
def forward(self, x, w):
    y = self.kernel[grid](x, w)
    return y.sum(dim=-1)  # ← 违规：tensor 方法计算

# ❌ 错误 4：tensor 运算符
def forward(self, x, w):
    y = self.kernel[grid](x, w)
    return y + 1  # ← 违规：+ 是 PyTorch 运算符

✅ Correct Example (pure Triton implementation)

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n, BLOCK_SIZE: tl.constexpr):
    idx = tl.arange(0, BLOCK_SIZE)
    x = tl.load(x_ptr + idx)
    y = tl.load(y_ptr + idx)
    output = x + y  # ← 计算在 kernel 中
    tl.store(output_ptr + idx, output)

class ModelNew(nn.Module):
    def forward(self, x, y):
        output = torch.empty_like(x)  # ✅ 允许：buffer 分配
        add_kernel[(1,)](x, y, output, x.numel(), BLOCK_SIZE=128)  # ✅ 允许：kernel 启动
        return output  # ✅ 允许：直接返回 kernel 输出

Input Information

You will be given the following information:

  1. Task description and specification — operator requirements in operator task format (including the Model class)
  2. Algorithm design sketch（sketch) — algorithm sketch generated by kernel-designer (passed in by the workflow during first-time generation)
  3. Relevant knowledge and examples — Triton Ascend programming knowledge (see the knowledge loading rules below)
  4. Execution history — previous error messages and fix suggestions (during iterative generation)

Knowledge Loading Rules

Required Knowledge (loaded on every generation)

  • Hardware specifications(must be loaded on every generation): @../npu-arch/references/npu-arch-guide-triton.md and @../npu-arch/references/npu-hardware-params.md

  • @references/triton-ascend-fundamentals.md — API reference, programming basics, Grid configuration, memory optimization, performance optimization, debugging checklist

  • @references/triton-ascend-examples.md — complete PyTorch + Triton Ascend example code

Knowledge Selected by Operator Type

Based on the operator type,additionallyload the corresponding reference documents:

Operator Type Identifying Feature Load Document
Element-wise add/mul/relu/sigmoid/tanh/gelu/exp/log/silu and other element-wise operations @references/triton-ascend-elementwise.md
MatMul matmul/bmm/linear/gemm and other matrix multiplications @references/triton-ascend-matmul.md
Reduce sum/mean/max/min/softmax/layernorm/logsoftmax and other reduction operations @references/triton-ascend-reduce.md
Attention self-attention/cross-attention/flash-attention/scaled-dot-product @references/triton-ascend-attention.md
Sort/Select nms and other sort-and-select operations @references/triton-ascend-sort-select.md
Interpolate Interpolation and other interpolation operations @references/triton-ascend-interpolate.md

If the operator involves multiple types (such as a fused operator), load all related documents.

Rules for using the algorithm sketch

When sketch(the algorithm design sketch generated by kernel-designer) is passed in,the code implementation must be based on the sketch , making full use of the algorithmic ideas and optimization strategies in it.

If sketchis not passed in, then based on task_desc design the implementation plan yourself.

Code generation modes

Mode 1: First generation (no historical information)

When there are only op_name、task_desc and other basic parameters:

  1. Carefully read task_desc in Model.forward() the reference implementation of
  2. Understand the operator's mathematical logic and computation pattern
  3. Determine the operator type and load the corresponding knowledge document
  4. Choose an appropriate parallelization strategy and memory access pattern
  5. Generate the kernel function and the ModelNew class

Mode 2: Code modification (with previous_code + user_requirements)

When the user requests modification of existing code:

  1. Only modify the parts requested by the user, do not refactor unrelated code
  2. Keep the code structure and interface unchanged(unless the user requests modification)
  3. Ensure that the modified code is still complete and runnable
  4. Output the complete modified code

Mode 3: Iterative repair (with verifier_error / conductor_suggestion)

When the previous round of verification fails:

  1. Analyze the error: carefully read verifier_error, understand the specific reason for the failure
  2. Refer to the suggestions: strictly follow the repair direction in conductor_suggestion to make modifications
  3. Preserve the strengths: preserve the correct parts of the previous round's code, and only modify the problematic parts
  4. Targeted repair: do not perform unnecessary large-scale refactoring
  5. Avoid repetition: if the suggestions mention historical lessons, ensure that the same mistake is not made

Output requirements

The generated codemustbe a complete Python file containing the following structure:

import torch
import torch.nn as nn
import triton
import triton.language as tl
# 其他必要的 import（如 torch_npu）

# Kernel 函数（一个或多个）
@triton.jit
def {op_name}_kernel(...):
    # 高性能内核实现
    ...

# 新 Model 类
class ModelNew(nn.Module):
    def __init__(self, <与原 Model 完全相同的参数>):
        super().__init__()
        # 与原 Model 相同的初始化逻辑
        # 在此获取核心数（如需要）

    def forward(self, <与原 Model 完全相同的输入>):
        # 调用自定义 kernel
        ...
        return output

Key constraints

Constraint Description
Class name ModelNew must use ModelNew，cannotbe Model
Interface consistency __init__ The signatures of forward and Model must be exactly the same as the original
Output consistency The output shape and data type must be the same as the original Model consistent
Self-contained All kernel functions and auxiliary functions must be defined in the same file
Executable The code must be directly importable and runnable
No test code No need to generate test code
Weight consistency Operators with random weights (Conv2d/Linear, etc.) must ensure weight consistency through a fixed seed
PyTorch degradation prohibited All core computations in forward() must be implemented in the @triton.jit kernel; using torch./F./tensor methods/tensor operators is prohibited

Weight consistency for operators with random weights (critical!)

When the Model class in the task description contains nn.Conv2d、nn.Linear、nn.ConvTranspose2d and other modules with learnable parameters, or uses torch.randn / nn.Parameter(torch.randn(...)) to create random parameters,you mustin ModelNew.__init__ ensure complete weight consistency with the original Model by fixing the random seed.

Rationale: the validation framework will call Model before creating torch.manual_seed(0), then call ModelNew again before creating torch.manual_seed(0). As long as both create parameters in the same order inside __init__ , completely identical weights will be obtained.

Standard pattern：

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, ...):
        super().__init__()
        # 1. 固定种子 — 必须与验证框架中的种子一致 (0)
        torch.manual_seed(0)

        # 2. 按照原 Model 的**完全相同的顺序**创建模块并提取权重
        #    确保随机数消耗顺序一致
        conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.weight = nn.Parameter(conv.weight.clone())
        self.bias = nn.Parameter(conv.bias.clone()) if conv.bias is not None else None

        # 如果原 Model 还有其他随机参数（如 nn.Parameter(torch.randn(...))），
        # 也必须在此按相同顺序创建
        self.extra_bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        # 使用提取的权重调用自定义 kernel
        return custom_conv_kernel(x, self.weight, self.bias, ...)

Key points：

  1. ModelNew.__init__ ofthe first linemust call torch.manual_seed(0)
  2. The order of parameter creationordermust be completely identical to the original Model.__init__ (because each torch.randn call advances the random state)
  3. Obtain weights by creating the same nn.Module(such as nn.Conv2d) rather than manually using torch.randn — this ensures the shape and initialization method of the internal parameters are consistent
  4. If the original Model has multiple modules containing weights, you must create and extract them one by one inthe original orderCreate and extract one by one

Thinking requirements

Important: During the thinking process, only perform framework-level analysis and decisions, for example:

  • Determining the operator type (elementwise / reduce / matmul, etc.)
  • What optimization strategy to choose (loop unrolling, vectorization, etc.)
  • How to handle data types
  • The rough skeleton of the code structure

Do not write complete code during the thinking process; the complete code is only given in the final output.

Generation principles

  • Generatecomplete, compilablecode
  • Follow Triton Ascend best practices
  • Optimize for the Ascend NPU architecture
  • Correctly handle boundary cases and exceptional conditions
  • Include the necessary imports and wrapper functions
  • Numerical correctness takes priority over performance
  • Strictly comply with the constraint prohibiting PyTorch fallback — all core computations must be implemented in the @triton.jit kernel