tilelang-vector-ascend-ops-migration

TileLang GPU to NPU Operator Migration Skill

Skill Overview

This skill is used to guide the migration of TileLang operators from the GPU (CUDA) platform to the Huawei Ascend NPU platform. By analyzing the GPU implementation, the corresponding NPU implementation code is automatically generated.

Applicable Scenarios

Migrate the TileLang operators under thegpu/directory to thenpu/directory (or, in other words, migrate tilelang operators adapted for GPU to tilelang operators adapted for NPU. If the user does not specify where the migrated operators should be placed, create a new npu directory, place them under the npu directory, and inform the user of the storage location of the migrated operators)
Automatically adapt to NPU hardware constraints and API differences
Generate NPU operator code that can be run directly

📚 Important Reference Documents

During the migration process, there are significant differences between GPU and NPU API interfaces. It is strongly recommended to refer to the following documents:

Core Documents

docs/GPU-To-NPU-Migration-Methods.md - Detailed migration experience and cases

Supplementary Documents

references/debugging-guide.md - NPU Operator Debugging Guide
- Debugging methods for precision issues, compilation failures, and runtime errors
- Debugging techniques at the TVM IR and MLIR layers
- T.print print debugging, GDB debugging
references/hardware-optimization.md - NPU Hardware Features and Performance Optimization
- On-chip cache capacity limits (UB 192KB, L1 512KB)
- Detailed explanation of data alignment requirements
- Compute-storage parallelism and pipeline optimization
- Physical core count limits and core partitioning strategies

Usage Recommendations:

When encountering uncertain APIs, prioritize consulting TileLang-NPU-API-Reference.md
Some GPU APIs may not exist on NPU or may have different names (such as T.vconv → T.vcast）
The parameter order and type requirements of vectorized operations may differ
When encountering debugging issues, refer to debugging-guide.md
When performing performance optimization, refer to hardware-optimization.md

⚠️ Key Precautions

Most Common Mistake 1: Using element-wise processing instead of block processing

This is in NPU migrationthe most fatal and most easily made mistake. The NPU's vectorized architecture requires block processing, not element-wise processing.

# ❌ 错误写法：逐元素处理（效率极低且容易出错）
@T.prim_func
def main(X: T.Tensor((M, N), "float16"), Y: T.Tensor((M, N), "float16")):
    with T.Kernel(M * N, is_npu=True) as (cid, _):
        idx = cid
        bx = idx // N
        by = idx % N

        # 为每个元素分配独立的 buffer（错误！）
        x_ub = T.alloc_shared((1,), "float16")
        y_ub = T.alloc_shared((1,), "float32")

        # 单个元素访问（错误！）
        T.copy(X[bx, by], x_ub)
        T.vsigmoid(x_ub, y_ub)
        T.copy(y_ub, Y[bx, by])

# ✅ 正确写法：分块处理
@T.prim_func
def main(X: T.Tensor((M, N), "float16"), Y: T.Tensor((M, N), "float16")):
    with T.Kernel(T.ceildiv(M, block_M) * T.ceildiv(N, block_N), is_npu=True) as (cid, _):
        grid_y = T.ceildiv(N, block_N)
        bx = cid // grid_y
        by = cid % grid_y

        # 为整个 block 分配 buffer
        x_ub = T.alloc_shared((block_M, block_N), "float16")
        y_ub = T.alloc_shared((block_M, block_N), "float32")

        # 使用切片访问整个 block
        T.copy(X[bx * block_M : (bx + 1) * block_M,
                 by * block_N : (by + 1) * block_N], x_ub)
        T.vsigmoid(x_ub, y_ub)
        T.copy(y_ub, Y[bx * block_M : (bx + 1) * block_M,
                       by * block_N : (by + 1) * block_N])

Why block processing must be used:

NPU vectorized instructions are designed to process relatively large data blocks (usually 64x64)
Element-wise processing cannot fully utilize the NPU's parallel computing capability
The memory access efficiency of single-element buffers is extremely low
It may cause runtime errors or incorrect results

Recommended block size:

2D tensor：block_M=64, block_N=64
1D tensor：block_size=256 or block_size=512

Most Common Mistake 2: Wrong location for calculating grid size

During the migration process, anothermost easily made and most frequent mistakeis defining local variables inside @T.prim_func to calculate the grid size.

# ❌ 错误写法（会导致编译失败）
@T.prim_func
def main(...):
    num_blocks = M * N  # 局部变量
    with T.Kernel(num_blocks, is_npu=True):  # 编译时找不到 num_blocks！
        ...

# ✅ 正确写法
@T.prim_func
def main(...):
    with T.Kernel(M * N, is_npu=True):  # 直接使用参数计算
        ...

Reason: T.Kernel The grid parameters of this are evaluated at compile time, and at that point the local variables inside the function do not yet exist.

Solution: Always use function parameters and T.Kernel(...) directly in the parameters for calculation. Do not use any local variables. T.ceildiv for calculation. Do not use any local variables.

Most Common Mistake 3: Using non-existent scalar operation APIs

NPU APIs that support direct scalar-vector operations are not supported(such as T.vmuls、T.vadds , etc.). Scalars must first be broadcast to vectors.

# ❌ 错误写法：使用不存在的 T.vmuls
alpha_scalar = T.cast(0.01, "float32")
T.vmuls(x_fp32, alpha_scalar, result)  # API 不存在！

# ✅ 正确写法：先广播再向量运算
alpha_scalar = T.cast(0.01, "float32")
alpha_ub = T.alloc_shared((block_M, block_N), "float32")
T.vbrc(alpha_scalar, alpha_ub)  # 广播标量到向量
T.vmul(x_fp32, alpha_ub, result)  # 向量乘法

Key point:

NPU only has vector operation APIs:T.vadd、T.vmul、T.vsub、T.vdiv
There is no scalar operation API:T.vadds、T.vmuls etc. do not exist
You must use T.vbrc(scalar, tensor) to broadcast a scalar into a vector
Note:T.vbrc The first argument of must be a variable, not a literal

Most common mistake 4: Wrong argument order for T.vcmp

The NPU's T.vcmp API requiresexplicitly specifying the output buffer, and the argument order is different from GPU.

# ❌ 错误写法：返回值方式（GPU 风格）
mask = T.vcmp(x_fp32, zero_ub, "GT")  # NPU 不支持！

# ❌ 错误写法：参数顺序错误
T.vcmp(x_fp32, zero_ub, "gt", mask_ub)  # 参数顺序错误！

# ✅ 正确写法：显式输出 buffer
mask_ub = T.alloc_shared((block_M, block_N), "bool")
T.vcmp(x_fp32, zero_ub, mask_ub, "gt")  # 正确顺序
T.vselect(mask_ub, true_val, false_val, result)

Key point:

T.vcmp The correct signature:T.vcmp(src0, src1, dst, cmp_op)
You must first allocate "bool" an output buffer of type
The comparison operator usesa lowercase string："gt", "lt", "ge", "le", "eq", "ne"
The return value is None, and the result is stored in the dst parameter

Most common mistake 5: Using T.copy for broadcast operations

The NPU's T.copy does not support copying with mismatched shapes, and cannot be used for broadcast operations. You must use T.vbrc for scalar broadcasting.

# ❌ 错误写法：使用 T.copy 广播
max_ub = T.alloc_shared((1, 1), "float32")  # 归约结果
T.reduce_max(x_fp32, max_ub, dim=-1)

max_brc = T.alloc_shared((1, N), "float32")
T.copy(max_ub, max_brc)  # ❌ 错误：形状不匹配 (1,1) → (1,N)

# ✅ 正确写法：提取标量后使用 T.vbrc 广播
max_ub = T.alloc_shared((1, 1), "float32")
T.reduce_max(x_fp32, max_ub, dim=-1)

max_brc = T.alloc_shared((1, N), "float32")
max_scalar = max_ub[0, 0]  # 提取标量值
T.vbrc(max_scalar, max_brc)  # 广播到向量

Key point:

T.copy requires the source and destination shapes to becompletely identical, and does not support automatic broadcasting
The result of reduction operations (T.reduce_max、T.reduce_sum) usually needs to be broadcast back to the original dimensions
You must first extract the scalar value through indexing (such as [0, 0])
Use T.vbrc(scalar, tensor) to broadcast the scalar to the entire tensor
This is a for reduction-type operators such as softmax and layernormrequired operation pattern

Most common mistake 6: Ignoring boundary checks during tiled processing

This is in NPU tiled processinga problem that is very easy to overlook. When the tensor dimension is not an integer multiple of the block size, the last block will access out-of-bounds memory, causing a runtime crash.

# ❌ 错误写法：忽略边界检查
@T.prim_func
def main(X: T.Tensor((M, N), "float16"), Y: T.Tensor((M, N), "float16")):
    with T.Kernel(T.ceildiv(M, block_M) * T.ceildiv(N, block_N), is_npu=True) as (cid, _):
        grid_y = T.ceildiv(N, block_N)
        bx = cid // grid_y
        by = cid % grid_y

        x_ub = T.alloc_shared((block_M, block_N), "float16")
        y_ub = T.alloc_shared((block_M, block_N), "float16")

        # 当 M 或 N 不是 64 的倍数时，最后一个 block 会越界！
        T.copy(X[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], x_ub)
        T.vrelu(x_ub, y_ub)
        T.copy(y_ub, Y[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N])

# ✅ 正确写法：使用 size 参数处理边界
@T.prim_func
def main(X: T.Tensor((M, N), "float16"), Y: T.Tensor((M, N), "float16")):
    with T.Kernel(T.ceildiv(M, block_M) * T.ceildiv(N, block_N), is_npu=True) as (cid, _):
        grid_y = T.ceildiv(N, block_N)
        bx = cid // grid_y
        by = cid % grid_y

        row_start = bx * block_M
        col_start = by * block_N

        # 计算当前 block 的实际大小
        actual_M = T.min(block_M, M - row_start)
        actual_N = T.min(block_N, N - col_start)

        x_ub = T.alloc_shared((block_M, block_N), "float16")
        y_ub = T.alloc_shared((block_M, block_N), "float16")

        # 使用 size 参数指定实际拷贝大小
        T.copy(X[row_start, col_start], x_ub, size=[actual_M, actual_N])
        T.vrelu(x_ub, y_ub)
        T.copy(y_ub, Y[row_start, col_start], size=[actual_M, actual_N])

Key point:

You must calculate the actual size：actual_M = T.min(block_M, M - row_start)
Use the size parameter：T.copy(src, dst, size=[actual_M, actual_N])
You cannot use slice syntax together with size：T.copy(X[start:end, ...], dst, size=[...]) is wrong
Correct way to write it：T.copy(X[row_start, col_start], dst, size=[actual_M, actual_N])

Error symptom:

The write address of the MTE instruction is out of range
vector core exception

Most common mistake 7: Conflict between T.copy slice syntax and the size parameter

TileLang does not allow using slice syntax together with the size parameter, which will cause a compilation error.

# ❌ 错误写法：切片语法 + size 参数
T.copy(X[row_start:row_end, col_start:col_end], x_ub, size=[actual_M, actual_N])
# 错误信息：T.copy: cannot use both slice syntax and the size parameter.

# ✅ 正确写法：起始位置索引 + size 参数
T.copy(X[row_start, col_start], x_ub, size=[actual_M, actual_N])

Mnemonic:

When using the size parameter, the source tensor can only be indexed by the starting position (X[row, col]), and cannot use slice syntax (X[start:end]）。

Most common mistake 8: Test data exceeds the actual model range

This is amistake that is extremely easy to overlook but has a major impact. The NPU's low-level APIs are specifically optimized (such as using approximate function fitting), and only guarantee precision within a specific numerical range.

# ❌ 错误写法：生成超出实际范围的测试数据
def test_sigmoid_npu():
    # sigmoid在实际模型中输入通常在[-10, 10]范围
    X = torch.randn(M, N) * 100  # 生成[-300, 300]，超出实际范围！
    Y = sigmoid_npu(X)
    # 可能导致精度测试失败，但实际模型中不会出现这种输入

# ✅ 正确写法：符合实际模型场景的测试数据
def test_sigmoid_npu():
    # 分析实际模型中sigmoid的输入分布
    X = torch.randn(M, N) * 3  # 生成[-9, 9]范围，符合实际
    Y = sigmoid_npu(X)

Key principle:

Analyze the actual usage scenario of the operator: check the input range of the operator in the model
Typical ranges of common operators：
- sigmoid/tanh: input is usually in [-10, 10]
- softmax: input is normalized, usually in [-5, 5]
- LayerNorm: after normalization the output is in [-3, 3]
- Probability values: in the [0, 1] range
Do not arbitrarily expand the test range: even if the test passes, it may mask real problems
Refer to the PyTorch documentation: Understand the numerical characteristics of each operator

Why this matters:

The NPU uses optimization techniques such as lookup tables and polynomial fitting
These optimizations are designed for the numerical distributions of real models
Out-of-range tests may fail, but this does not affect actual deployment

Must-do checklist

Must use block-based processing (Block-based Processing): NPU operators must use block-based processing and cannot process element by element. Each kernel processes the data of one block (such as 64x64), using slices to access memory.
Must handle boundary cases (Boundary Handling): When the tensor dimension is not an integer multiple of the block size, you must use the size parameter to handle boundaries. Compute the actual size:actual_M = T.min(block_M, M - row_start), and in T.copy use size=[actual_M, actual_N]。
Check the NPU API documentation to confirm that all APIs exist: Before migration, you must confirm in ./docs/TileLang-NPU-API-Reference.md that every API used exists. Pay special attention to:
- The NPU has no scalar operation APIs (T.vmuls、T.vadds etc. do not exist)
- When scalar operations are needed, you must use T.vbrc broadcast first and then use vector APIs
- For type conversion use T.vcast instead of T.vconv
- T.vcmp You must explicitly allocate an output buffer of type "bool" , and the parameter order is T.vcmp(src0, src1, dst, cmp_op)
- T.copy Copying with mismatched shapes is not supported; broadcasting reduction results must use T.vbrc
Test data must match real model scenarios: When generating test cases, the range of input data must match situations that may occur in real models. The NPU's low-level APIs are specifically optimized (such as approximate function fitting) and only guarantee precision within specific numerical ranges.

Key principles:
- Analyze the input range of the operator in real models (for example, sigmoid input is usually in [-10, 10], and softmax input is normalized)
- Do not generate test data beyond the actual range (for example, if the input to some activation functions will not be >1, do not generate test values >1)
- Refer to PyTorch/model documentation to understand the typical input distribution of each operator
Incorrect example:
```
# ❌ 错误：sigmoid在模型中输入通常有界，不会出现极大值
X = torch.randn(M, N) * 100  # 生成[-300, 300]范围，不符合实际

# ✅ 正确：符合实际模型中sigmoid的输入范围
X = torch.randn(M, N) * 3    # 生成[-9, 9]范围，符合实际
```
Refuse to migrate operators that need to use APIs for which precision issues have already been discovered

If these APIs need to be used, either find a way around them and use other alternatives, or if they cannot be replaced, directly report that migration is impossible
- APIs refused for use
  
  Temporarily empty

Migration workflow

Step 1: Analyze the GPU operator

When the user requests migration of an operator, first read the GPU implementation file:

gpu/<算子名称>.py

Analyze the following key information:

The operator's input/output tensor shapes and data types
Core computation logic (GEMM, reduction, activation functions, etc.)
Memory allocation patterns (shared memory, fragment)
Parallel strategy (T.Parallel, T.serial, T.Pipelined)
Dynamic shape definition

Step 2: Apply migration rules

Refer todocs/TileLang-NPU-API-Reference.mdanddocs/GPU-To-NPU-Migration-Methods.md, and apply the following core migration rules:

2.1 Environment and decorators

GPU version:

import tilelang as tl
import tilelang.language as T

@tl.jit(pass_configs={...})
def kernel_impl(...):

NPU version:

import os
import tilelang
import tilelang.language as T

os.environ['TILELANG_ASCEND_MODE'] = 'Developer'

@tilelang.jit(target="npuir")
def kernel_impl(...):

2.2 Data type conversion

GPU	NPU
`T.bfloat16`	`"float16"`
`T.float32`	`"float32"`
`T.int32`	`"int32"`
`T.dynamic("name")`	`T.symbolic("name")`

2.3 Kernel declaration

GPU version:

with T.Kernel(T.ceildiv(seq_len, block_Q), threads=512) as bx:
    tx = T.thread_binding(0, 512, thread="threadIdx.x")

NPU Version:

with T.Kernel(num_blocks, is_npu=True) as (cid, _):
    # cid是core ID，NPU自动管理线程

2.4 Memory Operations

Operation	GPU	NPU
Zeroing	`T.fill(tensor, 0)`	`T.clear(tensor)`
GEMM Initialization	`clear_accum=True`	`initC=True`
GEMM Transpose	`transpose_A=True`	`a_transpose=True`
Atomic Operations	`T.atomic_add(dst, src)`	`T.atomic_add(dst, src, [size])`

2.5 Parallel Loop Constraints

WithinT.Parallel:

Must operate on UB buffer, cannot directly access global memory
UseT.vmax/T.vmininstead ofT.max/T.min
Avoid index operations (e.g.A[i+1]）
UseT.vselectinstead of conditional checks

2.6 Vectorization First

GPU Version (Loop):

for i, j in T.Parallel(M, N):
    C[i, j] = T.max(A[i, j], 0) * B[i, j]

NPU Version (Vectorized):

T.copy(A, A_ub)
T.vrelu(A_ub, A_ub)
T.vmul(A_ub, B_ub, C_ub)

Step 3: NPU Hardware Constraint Check

💡 Detailed Optimization Guide: This section provides basic hardware constraint check points. For more in-depth performance optimization strategies (including tiling strategies, compute-storage parallelism, core partitioning optimization, etc.), please refer to references/hardware-optimization.md

3.1 UB Memory Limit (Key Constraint)

Hard Constraint: The sum of allT.alloc_shared()andT.alloc_fragment()within a single Kernel must be < 96KB Recommended Value: < 85KB (reserve space for temporary variables)

Memory Calculation Formula:

# FP16: 2 bytes, FP32: 4 bytes, INT32: 4 bytes
total_bytes = sum(shape[0] * shape[1] * dtype_bytes for each allocation)

Example:

# ✅ 安全（约48KB）
q_shared = T.alloc_shared([64, 128], "float16")    # 16KB
k_shared = T.alloc_shared([64, 128], "float16")    # 16KB
acc_ub = T.alloc_shared([64, 64], "float32")       # 16KB

# ❌ 超限（256KB）
q_shared = T.alloc_shared([512, 128], "float16")   # 128KB
k_shared = T.alloc_shared([512, 128], "float16")   # 128KB

If memory exceeds the limit, you need to:

reduce block_size
process in blocks
reuse buffer

3.2 Data Alignment Requirements

FP16/BF16: dimensions must be multiples of 16 (32-byte alignment)
FP32/INT32: dimensions must be multiples of 8 (32-byte alignment)
All input tensors must call.contiguous()

3.3 Known BUG Workarounds

Cube followed by Vector operation hangs：T.gemm()immediately after usingT.vcmp+T.vselectwill hang
- Solution: useT.vreluinstead of
GEMM crashes after Gather: directly performing GEMM after gathering data within the loop will cause compilation failure
- Solution: complete the Gather operation in advance on the CPU side
CV core partitioning Bug: in some cases, meaningless operations need to be added to trigger optimization
- Solution:T.vadd(k_shared, T.cast(0.0, dtype), k_shared)

Step 4: Generate NPU Code

The generated NPU code should contain the following structure:

import os
import torch
import tilelang
import tilelang.language as T

os.environ['TILELANG_ASCEND_MODE'] = 'Developer'

FP16 = "float16"
FP32 = "float32"
INT32 = "int32"

@tilelang.jit(target="npuir")
def _kernel_impl(...):
    dtype = FP16
    accum_dtype = FP32

    # 使用T.symbolic定义动态shape
    seq_len = T.symbolic("seq_len")

    @T.prim_func
    def kernel(...):
        with T.Kernel(num_blocks, is_npu=True) as (cid, _):
            # 1. 内存分配
            # 2. 数据加载
            # 3. 计算逻辑
            # 4. 结果写回

    return kernel

def wrapper_function(...):
    # PyTorch接口封装
    kernel = _kernel_impl(...)
    kernel(...)
    return output

def run_test():
    # 测试代码，对比参考实现
    pass

if __name__ == "__main__":
    run_test()

Step 5: Verification and Testing

💡 Debugging Guide: If the test fails or you encounter compilation/runtime errors, please refer to references/debugging-guide.md to get detailed debugging methods.

After generating the code, run the following command to test:

:warning:Before debugging, please first ask the user whether the tilelang-ascend related environment is set up

python3 <迁移的算子名称、路径>.py"

Test Code Requirements

Must include the following:

Compute error: use the PyTorch reference implementation for comparison
Output error value: print the maximum absolute error
Pass condition: error < 1e-3
Clearly output the result: print ✅ or ❌

Standard test code template:

def test_xxx_npu():
    # 1. 准备输入数据
    X = torch.randn(..., dtype=torch.float16, device="npu").contiguous()
    Y_npu = torch.empty(..., dtype=torch.float16, device="npu")

    # 2. 运行 NPU kernel
    kernel = xxx_npu(...)
    kernel(X, Y_npu)

    # 3. 运行 PyTorch 参考实现
    Y_ref = pytorch_reference(X)

    # 4. 计算误差
    max_diff = torch.max(torch.abs(Y_npu - Y_ref)).item()
    print(f"Max difference: {max_diff}")

    # 5. 判断并输出结果（必须）
    if max_diff < 1e-3:
        print("✅ Validation Passed!")
    else:
        print("❌ Validation Failed.")

if __name__ == "__main__":
    test_xxx_npu()

Key Requirements:

✅ Must output Max difference: xxx
✅ Must output ✅ Validation Passed! or ❌ Validation Failed.
✅ Condition:max_diff < 1e-3
✅ All input tensors must call .contiguous()

Migration Checklist

When performing migration, the following items must be checked:

Core Architecture (Most Important)

Use chunked processing: Must use the block_M, block_N parameter, and must not process element by element.
Grid size calculation: Use T.ceildiv(M, block_M) * T.ceildiv(N, block_N)
Boundary check handling: Calculate the actual size actual_M = T.min(block_M, M - row_start), and use the size parameter to handle non-integer-multiple boundaries.
T.copy syntax is correct: Use T.copy(X[row_start, col_start], dst, size=[...]) instead of slice syntax with size.
Buffer size: Allocate a buffer for the entire block, such as (block_M, block_N) instead of (1,)

Environment configuration

Addos.environ['TILELANG_ASCEND_MODE'] = 'Developer'
Change the decorator to@tilelang.jit(target="npuir")
Add to the Kernel declarationis_npu=True

Data types

T.bfloat16 → "float16"
T.float32 → "float32"
T.int32 → "int32"
T.dynamic → T.symbolic

Memory constraints

Calculate total UB memory < 96KB (recommended < 85KB)
Check dimension alignment (FP16 requires a multiple of 16, FP32 requires a multiple of 8)
Call for all input tensors.contiguous()

API mapping

T.fill(x, 0) → T.clear(x)
clear_accum → initC
transpose_A/B → a_transpose/b_transpose
T.max/T.min → T.vmax/T.vmin(inside T.Parallel)
T.atomic_add(dst, src) → T.atomic_add(dst, src, [size])
T.vcmp Use the correct parameter order:T.vcmp(src0, src1, dst, cmp_op), and first allocate a dst buffer of type "bool" .

BUG avoidance

Avoid immediate Vector operations after Cube (use T.vrelu)
Avoid GEMM after Gather inside a loop
Add the CV core-splitting bug fix when necessary

Test verification

Generate test code to compare against the reference implementation
Use docker commands to execute tests
Verify numerical precision (rtol=1e-3, atol=1e-3)

Reference documentation

During migration, please refer to:

docs/TileLang-NPU-API-Reference.md - NPU API brief documentation (contains all APIs but with brief content)
docs/GPU-To-NPU-Migration-Methods.md - Detailed migration experience
docs/examplesBelow, each python file is an example of migrating a simple vector tilelang operator from GPU to NPU.

Common migration pitfalls and solutions

Pitfall 1: T.Buffer is deprecated

Error symptom:

DeprecationWarning: T.Buffer(...) is deprecated, use T.Tensor(...) instead

Solution:

# ❌ 错误
X: T.Buffer((M, N), "float16")

# ✅ 正确
X: T.Tensor((M, N), "float16")

Pitfall 2: Vector operation types must match

Error symptom:

error: 'hivm.hir.vmul' op requires the same element type for all operands

Cause: NPU vector operations (T.vmul, T.vadd, etc.) require all operands to have exactly the same type, and FP16 and FP32 cannot be mixed.

Solution:

# ❌ 错误：混用FP16和FP32
x1_ub = T.alloc_shared((M, N), "float16")
y_ub = T.alloc_shared((M, N), "float32")
T.vmul(x1_ub, y_ub, y_ub)  # 类型不匹配！

# ✅ 正确：先转换类型
x1_ub = T.alloc_shared((M, N), "float16")
x1_fp32 = T.alloc_shared((M, N), "float32")
y_ub = T.alloc_shared((M, N), "float32")
T.vcast(x1_ub, x1_fp32)  # FP16 → FP32
T.vmul(x1_fp32, y_ub, y_ub)  # 类型一致

Pitfall 3: Type conversion API error

Error symptom:

error: module 'tilelang.language' has no attribute 'vconv'

Cause: NPU does not supportT.vconv; useT.vcastfor type conversion.

Solution:

# ❌ 错误
T.vconv(x_fp16, "none", x_fp32, 1.0)

# ✅ 正确
T.vcast(x_fp16, x_fp32)

Pitfall 4: Grid size calculation is in the wrong location (high-frequency error)

Error symptom:

ValueError: Failed to evaluate grid expression 'num_blocks': name 'num_blocks' is not defined
ValueError: Failed to evaluate grid expression 'total_tokens': name 'total_tokens' is not defined

Root cause: When the TileLang compiler processes T.Kernel(grid_size, ...) , it will, during thecompilation phase, try to parse the grid_size expression. At this time, @T.prim_func local variables defined internally have not yet been executed and therefore cannot be accessed.

Technical details:

T.Kernel The grid parameter of must bean expression evaluable at compile time
Only function parameters (such as M, N, batch_size) and TileLang built-in functions (such as T.ceildiv）
cannot use @T.prim_func any local variables inside

Solution:

Option 1: 2D Grid flattening

# ❌ 错误示例1
@T.prim_func
def main(...):
    grid_x = (M + block_M - 1) // block_M
    grid_y = (N + block_N - 1) // block_N
    num_blocks = grid_x * grid_y
    with T.Kernel(num_blocks, is_npu=True) as (cid, _):  # 编译时 num_blocks 不存在！
        ...

# ✅ 正确：直接在 Kernel 参数中计算
@T.prim_func
def main(...):
    with T.Kernel(T.ceildiv(M, block_M) * T.ceildiv(N, block_N), is_npu=True) as (cid, _):
        grid_y = T.ceildiv(N, block_N)
        bx = cid // grid_y
        by = cid % grid_y
        ...

Option 2: Simple multiplication

# ❌ 错误示例2
@T.prim_func
def main(...):
    total_tokens = batch_size * seq_len  # 局部变量
    with T.Kernel(total_tokens, is_npu=True) as (cid, _):  # 编译时 total_tokens 不存在！
        ...

# ✅ 正确：直接使用参数计算
@T.prim_func
def main(...):
    with T.Kernel(batch_size * seq_len, is_npu=True) as (cid, _):
        b = cid // seq_len
        s = cid % seq_len
        ...

Mnemonic:

In T.Kernel's grid parameter, only function parameters and T.ceildiv can be used directly; no local variables can be used.

Pitfall 5: Activation function computation order

Problem: When computing composite activation functions such as Swish, attention must be paid to the reuse of intermediate results.

Best practice:

# Swish(x1) * x2 = (x1 * sigmoid(x1)) * x2
T.vcast(x1_ub, x1_fp32)  # 转FP32
T.vcast(x2_ub, x2_fp32)
T.vsigmoid(x1_fp32, y_ub)  # sigmoid(x1) → y_ub
T.vmul(x1_fp32, y_ub, y_ub)  # x1 * sigmoid(x1) → y_ub
T.vmul(y_ub, x2_fp32, y_ub)  # swish(x1) * x2 → y_ub
T.vcast(y_ub, y_fp16)  # 转回FP16

Pitfall 6: Block boundary out-of-bounds (high-frequency error)

Error symptom:

The write address of the MTE instruction is out of range
vector core exception

Cause: When the tensor dimension (M or N) is not an integer multiple of the block size (e.g., 64), the slice access of the last block will exceed the tensor boundary, causing a memory out-of-bounds error.

Incorrect example:

# ❌ 错误：当 M=100, block_M=64 时，第二个 block 会访问 [64:128]，超出边界
T.copy(X[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], x_ub)

Solution:

# ✅ 正确：计算实际大小并使用 size 参数
row_start = bx * block_M
col_start = by * block_N
actual_M = T.min(block_M, M - row_start)  # 第二个 block: min(64, 100-64) = 36
actual_N = T.min(block_N, N - col_start)

T.copy(X[row_start, col_start], x_ub, size=[actual_M, actual_N])

Key points:

GPU element-wise processing naturally avoids out-of-bounds, but NPU block processing requires explicit boundary checks
Must use the T.copy parameter in size to limit the actual access range
Note: Slice syntax and the size parameter cannot be used at the same time

Pitfall 7: Conflict between T.copy slice syntax and the size parameter

Error symptom:

error: T.copy: cannot use both slice syntax and the size parameter.

Cause: The TileLang compiler does not allow T.copy to use slice syntax and the size parameter at the same time.

Incorrect example:

# ❌ 错误：切片语法 + size 参数
T.copy(X[row_start:row_end, col_start:col_end], x_ub, size=[actual_M, actual_N])

Solution:

# ✅ 正确：使用起始位置索引 + size 参数
T.copy(X[row_start, col_start], x_ub, size=[actual_M, actual_N])

Mnemonic:

When boundary handling is needed,T.copy use the starting position index together with the size parameter, and do not use slice syntax.

Usage instructions

When the user requests operator migration:

Read<gpu算子名称>.pyfile
Refer to the API documentation and migration experience to apply the conversion rules
Generate<npu算子名称>.pyfile
Prompt the user to run the test command for verification (the environment must be ensured to work properly in advance)

About this skill

Skill Overview

Applicable Scenarios

📚 Important Reference Documents

Core Documents

Supplementary Documents

⚠️ Key Precautions

Must-do checklist

Migration workflow

Step 1: Analyze the GPU operator

Step 2: Apply migration rules

2.1 Environment and decorators

2.2 Data type conversion

2.3 Kernel declaration

2.4 Memory Operations

2.5 Parallel Loop Constraints

2.6 Vectorization First

Step 3: NPU Hardware Constraint Check

3.1 UB Memory Limit (Key Constraint)

3.2 Data Alignment Requirements

3.3 Known BUG Workarounds

Step 4: Generate NPU Code

Step 5: Verification and Testing

Test Code Requirements

Migration Checklist

Core Architecture (Most Important)

Environment configuration

Data types

Memory constraints

API mapping

BUG avoidance

Test verification

Reference documentation

Common migration pitfalls and solutions

Pitfall 1: T.Buffer is deprecated

Pitfall 2: Vector operation types must match

Pitfall 3: Type conversion API error

Pitfall 4: Grid size calculation is in the wrong location (high-frequency error)

Pitfall 5: Activation function computation order

Pitfall 6: Block boundary out-of-bounds (high-frequency error)

Pitfall 7: Conflict between T.copy slice syntax and the size parameter

Usage instructions