simple-vector-triton-gpu-to-npu

Triton演算子のGPUから昇腾NPUへの移行スキル

このスキルは、シンプルなVector型のTriton演算子をGPUから昇腾NPUプラットフォームへ移行する際に役立ちます。移行手順、コード変換ガイド、精度検証方法を一通り提供します。

このスキルを使用する場面

以下のような場合にこのスキルを使用してください。

• Triton演算子をGPUから昇腾NPUへ移行する必要がある場合
• 「GPUからNPUへの移行」「Tritonの移行」「昇腾への対応」について言及された場合
• NPU上でのTriton演算子の性能を最適化する必要がある場合
• NPU固有のコンパイルエラー（coreDimの上限超過、UBオーバーフローなど）が発生した場合
• GPUとNPUのアーキテクチャの違いを知る必要がある場合

クイックスタート

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

# 2. 最小化迁移
# 只修改设备：device='cuda' -> device='npu'

# 3. 运行测试
python test_your_kernel.py

# 4. 根据错误调整
# 参考 reference/troubleshooting.md 解决问题

適用範囲

✅ 移行をサポートする演算子の種類:

• シンプルなVector系演算子（Vector Coreのみを使用）
• 複雑な制御フローを含まない演算子
• 1D／2Dグリッドの演算子

❌ 現時点で自動移行をサポートしないケース:

• コンパイルエラーがある演算子
• 複雑なループのネストや条件分岐
• 特殊なメモリアラインメントが必要なCV演算子（CubeとVector Coreを同時に使用）
• 外部ライブラリや特殊なランタイムに依存する演算子

移行方針

このスキルでは、NPU専用の移行方針を採用し、GPUの実装をNPUの実装へ移行します。

方針の特徴

• 適用シナリオ：GPUコードをNPUプラットフォームへ移行します
• コードの特徴：GPU APIをNPU APIに直接置き換えるため、コードが簡潔で分かりやすくなります
• 保守性：コード量が少なく、保守しやすくなります
• 使用シナリオ：本番環境へのデプロイや性能最適化のシナリオ

手順

Step 1: コードのロジックとセマンティクスの分析（重要なステップ）

⚠️ 移行前に、必ずソースコードの意味論的なロジックを分析し、分析レポートを作成してください

templates/analysis_template.md を使用して分析レポートを作成します。

分析のポイント：

1. tl.load の分析：ポインター計算、maskのセマンティクス、other値の選択
2. tl.store の分析：maskが出力境界のみをチェックしているか
3. Mask のロジック：各maskの意味、ブロードキャストの次元
4. データフロー：入力がどのように出力へ変換されるか

詳細な分析ガイドについては、reference/analysis_guide.md を参照してください。

Step 2: 環境の準備

# 安装依赖
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: 最小限の移行を試す

基本原則：まずGPUのソースコードに従って試し、エラーが発生した場合にのみ、以降の調整を行います

# 第一步：只修改设备指定
# 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: 主要APIの移行

⚠️ 必ず移行する必要があるGPU専用API：

0. CUDA API移行対応表

GPU API	NPU API	説明
torch.cuda.is_available()	torch.npu.is_available()	デバイスが利用可能か確認します
torch.cuda.empty_cache()	torch.npu.empty_cache()	キャッシュをクリアします
torch.cuda.synchronize()	torch.npu.synchronize()	デバイスを同期します
torch.cuda.mem_get_info()	torch.npu.mem_get_info()	メモリ情報を取得します
device="cuda"	device="npu"	デバイスを指定します
@torch.compile	削除	NPUは現時点でtorch.compileによる学習をサポートしていません

1. デバイスコンテキストの移行

# 直接使用NPU API
import torch_npu

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

2. デバイス属性アクセスの移行

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

NPUデバイス属性の対応表：

GPU属性	NPU属性	代表値（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: エラーの種類に応じて調整する

発生したエラーの種類に応じて、対応する解決策を選択してください。

エラーの種類	エラーメッセージのキーワード	解決策
コンパイルエラー	compilation failed	Triton構文の互換性を確認します
coreDimの上限超過	coreDim > UINT16_MAX	BLOCK_SIZEを大きくするか、環境変数を設定します
UBオーバーフロー	ub overflow	サブブロック分割方式を使用します
精度の問題	NaN, Inf, 不一致	論理演算子とmaskの使用を確認します
性能の問題	実行が遅い	メモリアクセスを最適化し、Tilingを使用します

詳細な解決策については、reference/troubleshooting.md を参照してください。

Step 5: 精度検証

移行完了後は、必ず精度検証を行ってください。

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)

例

移行の完全な例：ベクトル加算

移行前（GPU版）:

@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')

移行後（NPU版）:

@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')

説明：この例ではcare_padding=Falseを追加していません。「まず機能の正しさを確保する」という原則に従っています。

その他の例については、reference/examples.md を参照してください。

ベストプラクティス

基本原則

1. まず1DのGridを使用する: 2D gridは1Dに統合されます
2. メモリアラインメント: VV演算子は32バイト、CV演算子は512バイトにアラインメントします
3. 精度検証は必須: 移行後に必ず精度を検証します
4. tl.loadのother値の選択: インデックスを読み込む場合、otherには有効範囲を超える値（例：N）を設定します
5. BLOCK_SIZEの調整: 精度に問題がある場合は、BLOCK_SIZEを小さくしてみます
6. tl.loadとtl.storeではmaskのセマンティクスが異なる: tl.loadは入力の有効性を確認し、tl.storeは出力境界を確認します
7. デバイスAPIは必ず移行する: torch.cuda.*をtorch.npu.*またはtorch_npu.npu.*に置き換えます
8. デバイス属性は異なる: NPUではmulti_processor_countではなくvector_core_numを使用します
9. torch.compileを削除する: NPUは現時点で@torch.compileデコレーターをサポートしていないため、必ず削除します

CUDAからNPUへのAPI変換ルール

GPU API	NPU API	説明
torch.cuda.is_available()	torch.npu.is_available()	デバイスが利用可能か確認します
torch.cuda.empty_cache()	torch.npu.empty_cache()	キャッシュをクリアします
torch.cuda.synchronize()	torch.npu.synchronize()	デバイスを同期します
torch.cuda.mem_get_info()	torch.npu.mem_get_info()	メモリ情報を取得します
device="cuda"	device="npu"	デバイスを指定します
@torch.compile	削除	NPUは現時点でtorch.compileによる学習をサポートしていません

care_paddingパラメーターの使用規則

原則：まず機能の正しさを確保し、その後に性能を最適化します

推奨手順：

1. 初期移行：care_padding=Falseを追加しません
2. 機能検証：精度が正しいことを確認します
3. 性能最適化：性能を向上させる必要がある場合は、care_padding=Falseの追加を試します
4. 回帰検証：追加後は必ず精度を再検証します

例：

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

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

⚠️ 警告：care_padding=Falseを直接追加すると、出力がすべて0になったり、精度に問題が生じたりする可能性があります。

tl.load と tl.store における mask の使用規則

# ✅ 正确示例：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)  # 正确！

操作	maskの意味	確認する内容
tl.load	読み込む必要がある入力位置	インデックスの有効性、入力境界
tl.store	書き込む必要がある出力位置	出力境界、行と列の範囲

Requirements

pip install triton-ascend torch-npu

Grid优化模式（高级主题）

⚠️ 注意：本节内容适用于以下场景：

• 原始GPU代码使用2D/3D网格
• 需要充分利用NPU物理核心
• 简单1D网格算子通常不需要此优化

对于简单Vector算子，通常只需要：

1. 将device='cuda'改为device='npu'
2. 使用固定核数：grid = (num_core,)

核心概念

GPU使用逻辑网格（如3D网格），而NPU使用物理核心网格（1D网格）。为了优化NPU性能，需要将GPU风格的网格适配到NPU的物理核心架构。

网格适配模式

原始GPU风格：

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

优化后的NPU风格：

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物理核心网格

任务分发模式

GPU内核入口：

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内核入口：

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
    # ... 原有内核逻辑

内核参数适配

额外的NPU参数：

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

Grid优化步骤

步骤1：分析原始网格结构

1. 识别网格维度：grid = (dim1, dim2, dim3, ...)
2. 识别program_id用法：tl.program_id(0), tl.program_id(1)等
3. 将program_id索引映射到逻辑维度

步骤2：计算任务分发参数

# 计算总任务数
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

步骤3：修改内核入口点

用任务分发循环替换直接的program_id索引：

# 之前：
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
    # ... 所有使用这些变量的代码都在循环内部

步骤4：处理变量作用域（关键步骤）

重要：从GPU风格迁移到NPU任务分发模式时，变量作用域发生变化。必须将所有使用循环内部变量的代码移到循环内部。

需要检查的常见变量：

• 索引变量：pid_b, pid_h, pid_seq, i0, i1, i2等
• 长度变量：seq_len, T, B等（特别是当IS_VARLEN时可能被修改）
• 偏移变量：seq_offset, bos, eos等

需要移动的代码：

1. 所有使用上述变量的计算
2. 指针偏移计算
3. 循环控制变量（如nchunks）
4. 状态初始化
5. 结果存储

错误示例：

# ❌ 错误：变量在循环内部定义，但在外部使用
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未定义

正确示例：

# ✅ 正确：所有使用循环内部变量的代码都在循环内部
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  # 使用局部变量
    # ... 后续所有计算

步骤5：更新内核启动配置

# 之前：
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优化最佳实践

1. 优先使用1D网格：NPU物理核心是1D结构，使用1D网格可以获得最佳性能
2. 任务均衡分发：确保计算任务均匀分配到所有核心
3. 注意变量作用域：所有使用循环内部变量的代码必须在循环内部
4. 使用局部指针变量：在任务分发循环中，使用局部变量存储偏移后的指针
5. 避免编译时未定义错误：检查所有变量是否在正确的作用域内定义和使用

Grid优化示例

原始GPU版本（2D网格）：

@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)

优化后的NPU版本（1D网格）：

@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

• 详细架构差异：reference/architecture.md
• 完整故障排查：reference/troubleshooting.md
• 迁移案例分析：reference/examples.md
• 分析模板：templates/analysis_template.md

NPU专用代码生成规范

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

核心原则

1. 输出文件位置：在同目录下新建文件（如xxx_optimized.py或xxx_npu.py）
2. 函数名保持一致：所有函数名与原始GPU版本保持一致，不添加_npu后缀
3. 注释全英文：所有新增注释使用英文

代码转换规则

1. 函数命名规范

错误示例：

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

正确示例：

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

2. 注释转换规范

中文注释	英文注释
# 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

完整转换示例

原始GPU代码（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](...)

优化后NPU代码（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](...)

转换检查清单

在生成NPU专用代码时，确保：

• 输出文件在同目录下新建
• 所有函数名与原始GPU版本一致（无_npu后缀）
• 移除了with torch.cuda.device(...)上下文
• 所有中文注释已转换为英文
• 设备属性使用num_vectorcore而非multi_processor_count
• Grid使用1D物理核心网格(num_core,)
• 添加了任务分发参数（ngroups_step, task_num, num_core）
• torch.cuda.is_available() 替换为 torch.npu.is_available()
• torch.cuda.empty_cache() 替换为 torch.npu.empty_cache()
• torch.cuda.synchronize() 替换为 torch.npu.synchronize()
• torch.cuda.mem_get_info() 替换为 torch.npu.mem_get_info()
• device="cuda" 替换为 device="npu"
• 删除所有 @torch.compile 装饰器