triton-ascend-migration

Triton-Ascend Migration

Quick Start

When encountering a migration request, execute in the following order:

1. First identify the input method:
   • File path / specified code snippet
   • User directly pastes code
2. Then identify the input source:
   • GPU/CUDA Triton kernel
   • Python/PyTorch operator implementation
3. Then identify the operator type:
   • elementwise
   • broadcast / mask
   • reduce
   • contains tl.dot
4. First make a minimal runnable version:
   • cuda -> npu
   • add torch_npu
   • Remove GPU-specific device logic
   • Prefer 1D grid
   • For simple tutorial examples, default to providing a "minimal diff migration version"
5. After it runs successfully, then do Ascend-side optimization:
   • Physical core binding
   • BLOCK_SIZE/XBLOCK
   • BLOCK_SIZE_SUB/XBLOCK_SUB
   • Contiguous/aligned memory access
   • coreDim / UB / dtype / mask troubleshooting
6. If there is a clear optimization opportunity, directly output the optimized implementation; do not stop at suggestions only.

How to use this Skill

If the user asks "how do I use this skill", do not immediately go into a lengthy migration analysis; first give a concise usage guide in 3 to 6 lines, then continue executing based on the input provided by the user.

When summarizing briefly, keep only these points:

• The user can provide Triton/CUDA code,PyTorch reference implementation, file path, or error/performance logs.
• The user should preferably also specify the runtime environment: local command line, existing container, CI, or only generate code without executing.
• If the user has preferences, they should also state them:最小 diff 迁移、文档风格、先跑通再优化、直接给优化版。
• You will output according to the scenario:Triton-Ascend 实现、最小验证脚本、执行命令、优化说明。

If the user continues to ask "how exactly should I phrase the question", "how should I write the command", or "how do I run it in the container", then read references/usage.md, and provide local commands, container commands, and example phrasings as needed; do not copy the entire long explanation directly into a regular answer.

Copy this checklist and track progress:

迁移进度
- [ ] 识别输入来源与算子类型
- [ ] 先做最小迁移或语义改写
- [ ] 调整为 Ascend 友好的并行与 grid
- [ ] 重做 block / tiling
- [ ] 审查 stride / block_ptr / 对齐
- [ ] 处理 coreDim / UB / scalar 退化
- [ ] 直接落地可行优化
- [ ] 生成并保存最小 NPU 验证脚本
- [ ] 实际执行验证脚本
- [ ] 输出结果与优化说明

Input identification

First answer three questions:

1. Is the user giving a file path or pasting code directly?
2. Is this a complete script, a local snippet, or a single kernel?
3. Is this a GPU Triton migration or a Python/PyTorch semantic rewrite?

For details of the input method, default handling when information is missing, and priority when a file path conflicts with pasted code, see references/input-modes.md。

Scenario A: GPU Triton -> Triton-Ascend

Check first:

• Whether there is device='cuda'
• Whether there is GPU-specific device acquisition or assertion logic
• Whether a GPU-style multidimensional free grid is retained
• Whether it uses tl.dot
• Whether there is complex shape/stride/block_ptr/order

Scenario B: Python/PyTorch -> Triton-Ascend

First extract the semantics, then write Triton:

• Input/output tensor relationships
• Indexing and broadcasting methods
• mask / reduce logic
• dtype and precision requirements
• Whether the original PyTorch already performs naturally contiguous memory access

If the original operator is only a reference implementation, first write a semantically equivalent Triton-Ascend version, then continue optimizing.

Migration process

1. Collect the minimum necessary information

Prioritize collecting this information; fill in whatever is missing:

• Input code or minimal reproduction
• Input method: file path / specified code section / user directly pastes code
• shape、dtype、stride
• Whether there is mask, broadcast, reduce
• Current error or performance issue
• Whether exactly the same precision is required
• Runtime environment: local command line, inside a container, CI, or only generate code without executing

If the information is incomplete, fill it in in this order:

1. First infer from the existing code
2. Then fill in the verification script with the smallest reasonable assumptions
3. Only at the end ask the user for the necessary information

If what is currently missing is the "execution location" information, infer in the following order:

1. First see whether the user gave a container name,docker exec, container path, image information
2. Then see whether the user gave a local file path, current directory, or terminal command
3. If it still cannot be determined, ask again, "Do you want me to write the verification steps as local command-line commands or in-container commands?"

2. First do the minimal migration or semantic rewrite

By default, first aim for "semantic alignment and getting it to run":

• GPU Triton: first change cuda to npu
• import torch_npu
• Remove GPU-specific device logic
• For simple documentation/tutorial-style examples, try to keep the original kernel names, wrapper names,BLOCK_SIZE, grid style, and main code structure unchanged
• Do not proactively add in the first version contiguous(), extra assertions, function renaming, or engineering-style wrappers, unless the user explicitly asks for an "enhanced version/production version," or these changes are necessary to fix deterministic issues on NPU
• Python/PyTorch: first rewrite it into the most direct Triton kernel according to the original computation semantics

Do not over-rewrite in the first step.

If the user explicitly says these signal words:

• official documentation style
• strict minimal migration
• minimal diff
• do not create an engineering-enhanced version
• only refer to official migration examples

then this "minimal migration mode" should override the later, more generalized optimization guidance requirements:

• the code only makes necessary modifications
• 优化说明 can be only 1 to 3 lines, clearly stating "this problem will not expand on optimization for now"
• do not force expansion of TRITON_ALL_BLOCKS_PARALLEL、multibuffer、care_padding=False, physical core binding, etc., just for template completeness
• do not steer the answer style away from a "documentation diff" into an "engineering optimization overview"
• the verification script should also remain "minimally runnable"; do not write it as an engineering test framework by default

For documentation-style minimal migration, single-file example organization, and verification script naming and saving rules, see references/output-and-validation.md。

3. Rewrite the parallel model

On the Ascend side, prioritize following these rules:

• Prefer 1D grid
• Switch from GPU logical grid thinking to Ascend physical core binding thinking
• Vector-only operators should first be thought of along the Vector Core path
• those containing tl.dot operators should first be thought of along the AI Core path

Further judge according to this set of "general convergence rules"; do not mechanically preserve all implementation branches from the GPU:

• If the original implementation has multiple kernels,autotune, environment variable branches, or automatic dispatch across different data paths, first distinguish which are "semantically necessary" and which are merely "performance strategies on GPU"
• For performance branches on Ascend that are clearly no longer necessary, they can be converged into a single kernel or fewer paths; focus on preserving semantics, not retaining all historical branches
• If an operator is essentially still Vector-only, but the original implementation used complex block_ptr, 2D/3D grid, extra tiling, or multiple versions of the kernel, first evaluate whether it can be changed to a more direct 1D grid, fixed configuration, single-path implementation
• If an operator contains tl.dot, do not just think about "compressing the multi-dimensional grid into 1D"; first determine which grid dimensions are merely logical chunk / token / tile dimensions, and whether it is more appropriate to move them into the kernel's internal loop to reduce scheduling dimensions
• Do not mechanically classify only based on "the source code contains tl.dot"; if tl.dot is only used to implement intermediate techniques such as prefix-sum, local scan, or triangular mask aggregation, still first judge based on the operator's main semantics whether it is more like Vector-only a reduction/scan, or whether it should indeed go through the AI Core path
• If an operator naturally has structures such as chunk, tile, window, prefix-sum, or local reduction, do not just continue using the original block-by-block pointer logic; also evaluate whether "first rearrange the layout, then perform vectorized computation" is more suitable for Ascend
• If an auxiliary tensor (such as gate, mask, bias, index, state-gate) is not contiguous in the current access direction, first perform lightweight transpose/contiguous or equivalent layout rearrangement on the wrapper side, then access it inside the kernel using simpler linear ptr or more regular block_ptr access
• If the main loop order has been rearranged, for example from "K first, then T" to "T first, then K", also re-examine the state tensor, cache tensor, and history block tensor's shape/stride/block_ptr/order; do not only change the scheduling order while continuing to use the old view and then rely on trans or extra indexing to patch it up
• If the current project already has common capabilities such as get_vectorcore_num(), device attribute tools, and commonly used layout helpers, prioritize reusing in-project helpers and do not default to writing inline replacement versions by hand
• However, if the current output target is a "standalone runnable script" or a "minimal verification script", continue checking whether these helpers depend on extra initialization; if they depend on project initialization steps, either add the initialization or clearly state the prerequisites in the result
• When you decide to "remove branches / converge the implementation", explain the reason in the result: whether because the branch only serves GPU autotune, only serves shared memory selection, or has no clear benefit on Ascend
• If the migrated Triton-Ascend run log shows Please DO NOT tune args ['num_warps']、['num_stages'] or a similar warning, first go back and check whether GPU-style launch/tuning parameters are still mechanically retained; for a minimal runnable implementation on Ascend, by default do not explicitly retain these parameters unless you can provide clear compilation requirements or measured benefits
• Do not use only one set of generic shapes for the verification script; the test set should be inferred from the operator's characteristics, covering at least one non-divisible block, one case most likely to trigger branch differences, and one case closer to the real workload

If the user provides a 2D/3D grid, first evaluate whether it can be folded into a 1D grid and then recover the indices inside the kernel.coreDim、UB、shape/stride/block_ptr/order、care_padding=False、TRITON_ALL_BLOCKS_PARALLEL、multibuffer For details such as references/reference.md。

Optimization and Troubleshooting

Default rules for direct optimization

If any of the following conditions is met, directly provide the optimized implementation:

• coreDim clearly exceeds the limit
• UB usage is clearly too large
• Memory access is scattered and can be restructured into contiguous access
• mask load/store has a better form
• dtype clearly causes vector operations to degrade to scalar

If these conditions are not met, especially for examples like simple vector addition, do not default to outputting an enhanced wrapped version just to "look more complete." First provide the minimal migration version, then put the enhancements under "optional optimizations."

Optimization action priority

1. Adjust grid and number of cores
2. Adjust main block size
3. Introduce or refactor sub-block loops
4. Fix shape/stride/block_ptr/order
5. Evaluate care_padding=False
6. Evaluate TRITON_ALL_BLOCKS_PARALLEL
7. Evaluate multibuffer and related compilation optimization items
8. Adjust the dtype path without breaking semantics

Key points that must be covered

The output must cover the following:

• cuda -> npu
• torch_npu
• 1D grid
• Physical core binding
• Vector-only and containing tl.dot distinction
• coreDim <= 65535
• UB limit
• Contiguous / aligned memory access
• shape/stride/block_ptr/order re-review
• TRITON_ALL_BLOCKS_PARALLEL
• multibuffer
• care_padding=False
• scalar degradation caused by dtype

Fixed output template

Always output according to this structure:

## 迁移结论
- 输入来源：
- 算子类型：
- 主要迁移动作：

## Triton-Ascend 实现
- 给出最终 kernel 和调用包装代码
- 如果当前场景只是基础迁移，先给“最小 diff 迁移版”
- 只有在用户要求增强版，或确有明确优化空间时，再额外给“工程增强版/优化版”
- 如果存在明确优化空间，直接给出优化后的版本
- 说明生成文件的保存路径和命名

## 验证脚本
- 给出最小可执行验证脚本
- 使用 PyTorch reference 对比
- 至少包含 `allclose` 或最大误差输出
- 说明验证脚本保存路径
- 明确是否已实际执行，以及执行命令与结果

## 优化说明
- 说明 grid / 核数 / block / 子块 的调整原因
- 说明是否处理了 `coreDim`、UB、访存、dtype、mask 性能问题
- 说明是否使用 `TRITON_ALL_BLOCKS_PARALLEL`、`multibuffer`、`care_padding=False`

如果当前题目是“文档风格最小迁移”，这里可以极简：
- 只说明当前先做最小迁移
- 一句话说明本题未展开 `coreDim` / UB / `multibuffer` 等优化
- 不要为了套模板而展开长篇优化分析

## 风险与限制
- 列出仍未验证的边界条件
- 列出需要用户补充的信息
- 如果脚本未跑通，明确卡在哪一步

If the user's question itself is "how to use this skill," first add a minimal "How to use" subsection before the formal template, controlled to 3 to 6 lines, explaining:

• What input the user should provide
• Whether it is currently handled as a local or container scenario
• What you will produce next

Then proceed to the normal migration output.

If the user continues to ask about command lines, containers, directory switching, or validation command templates, then read references/usage.md, and do not stuff these details into every migration answer by default.

Additional Resources

When rule details are needed, continue reading:

• Usage, local commands, and container scenarios
• Input methods and context completion
• Output, naming, and minimal validation scripts
• Migration and optimization reference
• Typical examples and output samples
• Manual review test checklist