simple-vector-triton-gpu-to-npu

Triton Operator GPU-to-Ascend NPU Migration Skill

This skill helps users migrate simple Vector-type Triton operators from GPU to the Ascend NPU platform, providing a complete migration process, code conversion guide, and accuracy verification method.

When to use this Skill

Use this skill when the following situations occur:

The user needs to migrate Triton operators from GPU to Ascend NPU
The user mentions "GPU-to-NPU migration", "Triton migration", "Ascend adaptation"
The user needs to optimize the performance of Triton operators on NPU
The user encounters NPU-specific compilation errors (such as coreDim exceeding the limit, UB overflow)
The user needs to understand the architectural differences between GPU and NPU

Quick start

# 1. 分析源代码
# 使用 templates/analysis_template.md 生成语义分析报告

2. Minimal migration Only modify the device: device='cuda' -> device='npu' 3. Run tests

python test_your_kernel.py

4. Adjust based on errors Refer to reference/troubleshooting.md to resolve issues


## 适用范围

**✅ 支持迁移的算子类型**:
- 简单Vector类算子（仅使用Vector Core）
- 不涉及复杂控制流的算子
- 1D/2D网格的算子

**❌ 暂不支持自动迁移的情况**:
- 存在编译错误的算子
- 复杂的循环嵌套和条件分支
- 需要特殊内存对齐的CV算子（同时使用Cube和Vector Core）
- 依赖外部库或特殊runtime的算子

## 迁移策略

本技能采用**NPU专用迁移策略**，将GPU实现迁移为NPU实现。

### 策略特点
- **适用场景**：将GPU代码迁移到NPU平台
- **代码特点**：直接替换GPU API为NPU API，代码简洁清晰
- **维护性**：代码量小，易于维护
- **使用场景**：生产部署、性能优化场景

## Instructions

### Step 1: 代码逻辑语义分析（关键步骤）

**⚠️ 在迁移前，必须先分析源代码的语义逻辑，生成分析报告**

使用 [templates/analysis_template.md](templates/analysis_template.md) 生成分析报告。

分析要点：
1. **tl.load 分析**：指针计算、mask语义、other值选择
2. **tl.store 分析**：mask是否只检查输出边界
3. **Mask 逻辑**：每个mask的语义含义、广播维度
4. **数据流**：输入如何变换为输出

详细分析指南请参考 [reference/analysis_guide.md](reference/analysis_guide.md)。

### Step 2: 环境准备

```bash
# 安装依赖
pip uninstall triton  # 卸载社区Triton
pip install triton-ascend
pip install torch-npu

# 验证安装
python -c "import torch_npu; print(torch_npu.npu.is_available())"

Step 3: Minimal migration attempt

Core principle: first try according to the GPU source code, and if the attempt reports errors, then make subsequent adjustments

# 第一步：只修改设备指定
# device='cuda' -> device='npu'
x = torch.rand(size, device='npu')

# 第二步：运行测试
try:
    result = kernel_npu(**test_inputs)
    print("✅ 基础运行成功")
except Exception as e:
    print(f"❌ 运行失败: {e}")

Step 3.1: Key API migration

⚠️ GPU-specific APIs that must be migrated：

0. CUDA API migration reference table

GPU API	NPU API	Description
`torch.cuda.is_available()`	`torch.npu.is_available()`	Check whether the device is available
`torch.cuda.empty_cache()`	`torch.npu.empty_cache()`	Clear the cache
`torch.cuda.synchronize()`	`torch.npu.synchronize()`	Synchronize the device
`torch.cuda.mem_get_info()`	`torch.npu.mem_get_info()`	Get memory information
`device="cuda"`	`device="npu"`	Device specification
`@torch.compile`	Delete	NPU does not currently support torch.compile training

1. Device context migration

# 直接使用NPU API
import torch_npu

with torch_npu.npu.device(device_index):
    kernel[grid](...)

2. Device attribute access migration

props = torch_npu.npu.get_device_properties(device)
sm_count = props.vector_core_num  # Ascend910为48

NPU device attribute reference table：

GPU attribute	NPU attribute	Typical value (Ascend910)
multi_processor_count	vector_core_num	48
total_memory	total_memory	62740MB
name	name	'Ascend910_9392'
-	cube_core_num	24
-	L2_cache_size	'192MB'

Step 4: Adjust based on the error type

Based on the type of error encountered, choose the corresponding solution:

Error type	Error message keyword	Solution
Compilation error	compilation failed	Check Triton syntax compatibility
coreDim exceeding the limit	coreDim > UINT16_MAX	Increase BLOCK_SIZE or set an environment variable
UB overflow	ub overflow	Use a sub-block partitioning strategy
Accuracy issue	NaN, Inf, mismatch	Check logical operators and mask usage
Performance issue	Runs slowly	Optimize memory access and use Tiling

For detailed solutions, refer to reference/troubleshooting.md。

Step 5: Accuracy verification

After migration is completemustperform accuracy verification:

def verify_accuracy(result, ref, dtype):
    # 检查NaN/Inf
    assert not torch.isnan(result).any(), "结果包含NaN"
    assert not torch.isinf(result).any(), "结果包含Inf"
    
    # 设置容差
    if dtype in [torch.float16, torch.bfloat16]:
        rtol, atol = 1e-3, 1e-3
    elif dtype == torch.float32:
        rtol, atol = 1e-4, 1e-4
    else:
        rtol, atol = 0, 0
    
    torch.testing.assert_close(result, ref, rtol=rtol, atol=atol)

Examples

Complete migration example: vector addition

Before migration (GPU version):

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

x = torch.rand(98432, device='cuda')
y = torch.rand(98432, device='cuda')

After migration (NPU version):

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

x = torch.rand(98432, device='npu')
y = torch.rand(98432, device='npu')

Note: The example does not addcare_padding=False, following the principle of "ensure functionality correctness first".

For more examples, refer to reference/examples.md。

Best practices

Core principles

Prefer 1D for Grid: 2D grid will be merged into 1D
Memory alignment: VV operators are 32-byte aligned, CV operators are 512-byte aligned
Accuracy verification is required: Accuracy must be verified after migration
Selection of other value for tl.load: When loading by index,otherit should be set to a value outside the valid range (such asN）
BLOCK_SIZE adjustment: When there are accuracy issues, try reducing BLOCK_SIZE
The mask semantics of tl.load and tl.store are different: tl.loadCheck input validity,tl.storecheck output boundaries
Device APIs must be migrated: torch.cuda.* must be replaced with torch.npu.* or torch_npu.npu.*
Device properties are different: NPU uses vector_core_num instead of multi_processor_count
Remove torch.compile: NPU does not currently support@torch.compiledecorators, and they must be removed

CUDA to NPU API conversion rules

GPU API	NPU API	Description
`torch.cuda.is_available()`	`torch.npu.is_available()`	Check whether the device is available
`torch.cuda.empty_cache()`	`torch.npu.empty_cache()`	Clear cache
`torch.cuda.synchronize()`	`torch.npu.synchronize()`	Synchronize device
`torch.cuda.mem_get_info()`	`torch.npu.mem_get_info()`	Get memory information
`device="cuda"`	`device="npu"`	Device specification
`@torch.compile`	Remove	NPU does not currently support torch.compile training

Usage guidelines for the care_padding parameter

Principle: First ensure functional correctness, then consider performance optimization

Recommended process：

Initial migration: Do not addcare_padding=False
Functional verification: Ensure accuracy is correct
Performance optimization: If performance improvement is needed, then try addingcare_padding=False
Regression verification: Accuracy must be re-verified after adding

Example：

# Step 1: 初始迁移（不添加）
x = tl.load(x_ptr + offsets, mask=mask)

# Step 2: 功能验证通过后，可选的性能优化
# x = tl.load(x_ptr + offsets, mask=mask, care_padding=False)

⚠️ Warning: Directly addingcare_padding=Falsemay cause the output to be all 0 or cause accuracy issues.

Guidelines for using mask with tl.load and tl.store

# ✅ 正确示例：load 和 store 使用各自的 mask
out_mask = rows_mask & cols_mask[None, :]  # 输出边界检查
final_mask = out_mask & index_valid_mask[None, :]  # 输入有效性检查
selected = tl.load(inp + inp_off, mask=final_mask, other=0.0)
tl.store(out + out_off, selected, mask=out_mask)  # 正确！

Operation	mask meaning	What should be checked
`tl.load`	Whichinput positionsneed to be read	Index validity, input boundaries
`tl.store`	Whichoutput positionsneed to be written	Output boundaries, row and column ranges

Requirements

pip install triton-ascend torch-npu

Grid optimization patterns (advanced topic)

⚠️ Note: This section applies to the following scenarios:

The original GPU code uses 2D/3D grids
The NPU physical cores need to be fully utilized
Simple 1D grid operators usually do not require this optimization

For simple Vector operators, usually only the following is needed:

Changedevice='cuda'todevice='npu'
Use a fixed number of cores:grid = (num_core,)

Core concepts

GPU uses logical grids (such as 3D grids), while NPU uses physical core grids (1D grids). To optimize NPU performance, GPU-style grids need to be adapted to the NPU physical core architecture.

Grid adaptation patterns

Original GPU style:

grid = (NV, NK, N * H)  # 3D逻辑网格
kernel[grid](...)

Optimized NPU style:

import torch_npu
import triton.runtime.driver as driver

def get_npu_properties():
    device = torch.npu.current_device()
    return driver.active.utils.get_device_properties(device)

num_core = get_npu_properties()["num_vectorcore"]
grid = (num_core,)  # 1D物理核心网格

Task distribution patterns

GPU kernel entry point:

i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64)
i_n, i_h = i_nh // H, i_nh % H

NPU kernel entry point with task distribution:

core_id = tl.program_id(0)
task_num = NV * NK * N * H
knh_step = NK * N * H
nh_step = N * H

for task_id in tl.range(core_id, task_num, num_core):
    i_v = task_id // knh_step
    i_k = task_id % knh_step // nh_step
    i_nh = task_id % knh_step % nh_step
    i_n = task_id % knh_step % nh_step // H
    i_h = task_id % knh_step % nh_step % H
    # ... 原有内核逻辑

Kernel parameter adaptation

Additional NPU parameters:

def kernel(...,    
    knh_step: tl.constexpr,
    nh_step: tl.constexpr,
    N: tl.constexpr,
    task_num: tl.constexpr,
    num_core: tl.constexpr,
    ...):

Grid optimization steps

Step 1: Analyze the original grid structure

Identify grid dimensions:grid = (dim1, dim2, dim3, ...)
Identify the usage of program_id:tl.program_id(0), tl.program_id(1)etc.
Map the program_id index to a logical dimension

Step 2: Calculate task distribution parameters

# 计算总任务数
task_num = dim1 * dim2 * dim3 * ...  # 所有网格维度的乘积

# 计算每个维度的步长
# 3D网格(dim1, dim2, dim3)示例：
step_dim2_dim3 = dim2 * dim3
step_dim3 = dim3

# 在内核中：
# task_id = core_id + i * num_core
# dim1_idx = task_id // step_dim2_dim3
# dim2_idx = (task_id % step_dim2_dim3) // step_dim3
# dim3_idx = task_id % step_dim3

Step 3: Modify the kernel entry point

Replace the direct program_id index with a task distribution loop:

# 之前：
i0 = tl.program_id(0)
i1 = tl.program_id(1)
i2 = tl.program_id(2)

# 之后：
core_id = tl.program_id(0)
for task_id in tl.range(core_id, task_num, num_core):
    i0 = task_id // step_dim2_dim3
    i1 = (task_id % step_dim2_dim3) // step_dim3
    i2 = task_id % step_dim3
    # ... 所有使用这些变量的代码都在循环内部

Step 4: Handle variable scope (critical step)

Important: When migrating from the GPU style to the NPU task distribution mode, the variable scope changes. All code that uses variables inside the loop must be moved inside the loop.

Common variables that need to be checked：

Index variables:pid_b, pid_h, pid_seq, i0, i1, i2etc.
Length variables:seq_len, T, Betc. (especially whenIS_VARLENmay be modified)
Offset variables:seq_offset, bos, eosetc.

Code that needs to be moved：

All computations that use the above variables
Pointer offset calculations
Loop control variables (such asnchunks）
State initialization
Result storage

Incorrect example：

# ❌ 错误：变量在循环内部定义，但在外部使用
core_id = tl.program_id(0)
for task_id in tl.range(core_id, task_num, num_core):
    pid_b = task_id // h_step
    pid_h = task_id % h_step
    # ... 其他变量定义

# 错误：在循环外部使用循环内部定义的变量
nchunks = tl.cdiv(seq_len, CHUNK_SIZE)  # seq_len未定义
ANGLE += pid_b * stride_angle_batch  # pid_b未定义

Correct example：

# ✅ 正确：所有使用循环内部变量的代码都在循环内部
core_id = tl.program_id(0)
for task_id in tl.range(core_id, task_num, num_core):
    pid_b = task_id // h_step
    pid_h = task_id % h_step
    seq_len = ...  # 在循环内部定义
    
    # 所有使用这些变量的代码都在循环内部
    nchunks = tl.cdiv(seq_len, CHUNK_SIZE)
    angle_ptr = ANGLE + pid_b * stride_angle_batch  # 使用局部变量
    # ... 后续所有计算

Step 5: Update the kernel launch configuration

# 之前：
grid = (dim1, dim2, dim3)
kernel[grid](...)

# 之后：
num_core = get_npu_properties()["num_vectorcore"]
grid = (num_core,)
kernel[grid](
    ...,
    knh_step=step_dim2_dim3,
    nh_step=step_dim3,
    N=dim1,  # 或适当的映射
    task_num=task_num,
    num_core=num_core,
)

Grid optimization best practices

Prefer using a 1D grid: The NPU physical cores are a 1D structure, and using a 1D grid can achieve the best performance
Balanced task distribution: Ensure that compute tasks are evenly distributed to all cores
Pay attention to variable scope: All code that uses variables inside the loop must be inside the loop
Use local pointer variables: In the task distribution loop, use local variables to store the offset pointers
Avoid compile-time undefined errors: Check whether all variables are defined and used in the correct scope

Grid optimization example

Original GPU version (2D grid):

@triton.jit
def kernel_gpu(x_ptr, output_ptr, N, M, BLOCK_SIZE: tl.constexpr):
    pid_n = tl.program_id(0)
    pid_m = tl.program_id(1)
    
    # 计算偏移
    x = x_ptr + pid_n * M + pid_m * BLOCK_SIZE
    # ... 计算逻辑

# 启动内核
grid = (N, M // BLOCK_SIZE)
kernel_gpu[grid](x, output, N, M, BLOCK_SIZE=128)

Optimized NPU version (1D grid):

@triton.jit
def kernel_npu(x_ptr, output_ptr, N, M, BLOCK_SIZE: tl.constexpr,
               m_step: tl.constexpr, task_num: tl.constexpr, num_core: tl.constexpr):
    core_id = tl.program_id(0)
    
    for task_id in tl.range(core_id, task_num, num_core):
        pid_n = task_id // m_step
        pid_m = task_id % m_step
        
        # 计算偏移
        x = x_ptr + pid_n * M + pid_m * BLOCK_SIZE
        # ... 计算逻辑（所有代码都在循环内部）

# 启动内核
num_core = get_npu_properties()["num_vectorcore"]
m_step = M // BLOCK_SIZE
task_num = N * m_step
grid = (num_core,)
kernel_npu[grid](x, output, N, M, BLOCK_SIZE=128, 
                m_step=m_step, task_num=task_num, num_core=num_core)

Advanced usage

Detailed architecture differences:reference/architecture.md
Complete troubleshooting:reference/troubleshooting.md
Migration case analysis:reference/examples.md
Analysis template:templates/analysis_template.md

NPU-specific code generation specification

This skill uses an NPU-specific migration strategy to migrate the GPU implementation to an NPU implementation.

Core principles

Output file location: Create a new file in the same directory (such asxxx_optimized.pyorxxx_npu.py）
Keep function names consistent: All function names remain consistent with the original GPU version, without adding_npusuffix
All comments in English: All newly added comments use English

Code conversion rules

1. Function naming conventions

Incorrect example：

# ❌ 不要添加_npu后缀
def _layer_norm_fwd_1pass_kernel_npu(...):
    ...

Correct example：

# ✅ 保持原始函数名
def _layer_norm_fwd_1pass_kernel(...):
    # NPU optimized implementation
    ...

2. Comment conversion conventions

Chinese comments	English comments
`# NPU优化版本`	`# NPU optimized version`
`# NPU支持`	`# NPU support`
`# NPU任务分发参数`	`# NPU task dispatch parameters`
`# 使用1D物理核心网格`	`# Use 1D physical core grid`
`# 从task_id重建原始索引`	`# Reconstruct original indices from task_id`
`# 计算指针偏移`	`# Calculate pointer offsets`
`# 获取设备属性`	`# Get device properties`

Complete conversion example

Original GPU code (layernorm_gated.py)：

def _layer_norm_fwd_1pass_kernel(...):
    # GPU kernel implementation
    row = tl.program_id(0)
    group = tl.program_id(1)
    ...

def _layer_norm_fwd(...):
    grid = (M, ngroups)
    with torch.cuda.device(x.device.index):
        _layer_norm_fwd_1pass_kernel[grid](...)

Optimized NPU code (layernorm_gated_optimized.py)：

# NPU support
import torch_npu
import triton.runtime.driver as driver

def get_npu_properties():
    """Get NPU device properties, including number of cores"""
    device = torch.npu.current_device()
    return driver.active.utils.get_device_properties(device)

def _layer_norm_fwd_1pass_kernel(..., 
    # NPU task dispatch parameters
    ngroups_step: tl.constexpr,
    task_num: tl.constexpr,
    num_core: tl.constexpr,
):
    # NPU optimization: Use 1D physical core grid with task dispatch
    core_id = tl.program_id(0)
    
    for task_id in tl.range(core_id, task_num, num_core):
        # Reconstruct original 2D indices from task_id
        row = task_id // ngroups_step
        group = task_id % ngroups_step
        
        # Calculate pointer offsets
        X_ptr = X + row * stride_x_row + group * N
        # ... kernel logic

def _layer_norm_fwd(...):
    # NPU optimization: Use 1D physical core grid
    npu_props = get_npu_properties()
    num_core = npu_props["num_vectorcore"]
    grid = (num_core,)
    
    _layer_norm_fwd_1pass_kernel[grid](...)

Conversion checklist

When generating NPU-specific code, ensure:

The output file is newly created in the same directory
All function names are consistent with the original GPU version (no_npusuffix)
Removedwith torch.cuda.device(...)context
All Chinese comments have been converted to English
Device attributes usenum_vectorcoreinstead ofmulti_processor_count
Grid uses 1D physics core grid(num_core,)
Added task dispatch parameters (ngroups_step, task_num, num_core）
torch.cuda.is_available() replaced with torch.npu.is_available()
torch.cuda.empty_cache() replaced with torch.npu.empty_cache()
torch.cuda.synchronize() replaced with torch.npu.synchronize()
torch.cuda.mem_get_info() replaced with torch.npu.mem_get_info()
device="cuda" replaced with device="npu"
Removed all @torch.compile decorators

About this skill

When to use this Skill

Quick start

Step 3: Minimal migration attempt

Step 3.1: Key API migration

0. CUDA API migration reference table

1. Device context migration

2. Device attribute access migration

Step 4: Adjust based on the error type

Step 5: Accuracy verification

Examples

Complete migration example: vector addition

Best practices

Core principles

CUDA to NPU API conversion rules

Usage guidelines for the care_padding parameter

Guidelines for using mask with tl.load and tl.store

Requirements

Grid optimization patterns (advanced topic)

Core concepts

Grid adaptation patterns

Task distribution patterns

Kernel parameter adaptation

Grid optimization steps

Step 1: Analyze the original grid structure

Step 2: Calculate task distribution parameters

Step 3: Modify the kernel entry point

Step 4: Handle variable scope (critical step)

Step 5: Update the kernel launch configuration

Grid optimization best practices

Grid optimization example

Advanced usage

NPU-specific code generation specification

Core principles

Code conversion rules

1. Function naming conventions

2. Comment conversion conventions

Complete conversion example

Conversion checklist