triton_ascend入门学习一

从Triton Ascend文档的快速入门(https://triton-ascend.readthedocs.io/zh-cn/latest/quick_start.html)开始，环境可以选择gitcode的vllm（高性能大语言模型推理框架）那个镜像。

atomgit@676982cc8b2dd94c4aac433a:~$ git clone https://gitcode.com/Ascend/triton-ascend.git
Cloning into 'triton-ascend'...
remote: Enumerating objects: 48718, done.
remote: Counting objects: 100% (4838/4838), done.
remote: Compressing objects: 100% (446/446), done.
remote: Total 48718 (delta 4674), reused 4392 (delta 4392), pack-reused 43880 (from 1)
Receiving objects: 100% (48718/48718), 67.33 MiB | 31.82 MiB/s, done.
Resolving deltas: 100% (34267/34267), done.

atomgit@676982cc8b2dd94c4aac433a:~$ python3 ./triton-ascend/third_party/ascend/tutorials/01-vector-add.py
/usr/local/python3.11.14/lib/python3.11/site-packages/torch_npu/utils/collect_env.py:58: UserWarning: Warning: The /usr/local/Ascend/cann-8.5.0 owner does not match the current owner.
  warnings.warn(f"Warning: The {path} owner does not match the current owner.")

path string is NULLpath string is NULL[W406 20:01:44.991600977 compiler_depend.ts:164] Warning: Device do not support double dtype now, dtype cast replace with float. (function operator())
tensor([0.8329, 1.0024, 1.3639,  ..., 1.0796, 1.0406, 1.5811], device='npu:0')
tensor([0.8329, 1.0024, 1.3639,  ..., 1.0796, 1.0406, 1.5811], device='npu:0')
The maximum difference between torch and triton is 0.0

atomgit@676982cc8b2dd94c4aac433a:~$ pip list
torch                              2.9.0+cpu
torch_npu                          2.9.0
torchvision                        0.24.0
triton-ascend                      3.2.0

atomgit@676982cc8b2dd94c4aac433a:~$ llvm-config-15 --version
15.0.7

atomgit@676982cc8b2dd94c4aac433a:~$ cat /usr/local/Ascend/ascend-toolkit/latest/compiler/version.info 
Version=8.5.0
version_dir=cann
required_package_runtime_version="8.5.0"
required_package_metadef_version="8.5.0"
required_package_opbase_version="8.5.0"
required_package_bisheng-compiler_version="8.5.0"
required_package_ge-executor_version="8.5.0"
required_package_tbe-tik_version="8.5.0"
timestamp=20250725_000000000

这里插入以下基础概念介绍

• LLVM：通用的编译器基础设施框架，提供底层代码优化与机器码生成能力。
• TVM：端到端的深度学习编译器栈，专注于计算图优化与跨硬件的代码生成。参考这个。
• Triton：针对GPU算子开发的编程语言与编译器，专注于简化高性能并行计算内核的编写。参考这个。
• vLLM：大语言模型推理与服务框架，专注于显存管理（PagedAttention）与高吞吐量调度。

技术依赖关系

• TVM与Triton：两者均利用LLVM作为后端编译器，将高层计算逻辑转换为特定硬件（如CPU、GPU）的机器码。
• vLLM：作为应用层框架，主要依赖PyTorch进行计算，不直接依赖LLVM进行核心编译工作。

Triton-Ascend 是适配昇腾芯片的 Triton 版本，提供核函数自动调优、算子编译及部署能力，支持 Atlas A2/A3 等系列产品，兼容 Triton 核心语法，针对昇腾进行了深度优化，包括自动解析核函数参数、优化内存访问逻辑、完善安全部署机制等。（Triton本身是在GPU内核生成这一垂直领域做到极致易用。）架构：

分三块说说在NPU上进行Triton算子开发值得注意的问题：多核任务并行、单核数据搬运、单核数据运算。

在一个Triton算子中，通常使用grid进行分核操作。GPU计算核心SM通常是几十到几百量级，但昇腾 NPU 其计算核心AI Core的数量在几十个的量级。
虽然运行时接口允许下发并发任务数最大为65535，但超过物理核数的部分是通过新一轮的下发来完成的。如果直接将GPU上的Triton算子拿到昇腾平台上运行，这些大量的任务会引入可观的核启动和核初始化时的额外开销。
因此需要修改分核逻辑，最推荐的做法是将分核的数量直接固定为硬件的物理核数：

• 对于纯Vector算子，分核数等于Vector核数量
• 对于CV融合算子，分核数等于Cube核数量（通常为Vector核数量的一半），算子执行时会按1：2的比例调用Vector核

循环内数据分块大小（BLOCK SIZE）。通过修改BLOCK_SIZE可以调整循环内数据分块和计算中间结果占用的大小。如果超过上限则算子编译时会提示预期占用大小并报错。要达到最大计算访存比，BLOCK_SIZE需要在不超出片上空间时尽可能大，可以通过Triton-Ascend的Autotune(https://triton-ascend.readthedocs.io/zh-cn/latest/programming_guide.html#triton-autotune-自动调优)预先设置不同的BLOCK_SIZE，运行时会自动选取最优设置。

对于VV类算子，UB要求Tensor的尾轴大小能被32Bytes整除，而对于CV类算子要求能被512Bytes整除，若尾轴长度不足则会自动补齐。因此，对模型中shape为(2048,3)和(2048,1)Tensor的种种操作，都会因为自动补齐导致性能明显恶化，此时可考虑通过转置操作将对齐轴转到低维，直到store时再转置为原始状态，从而规避自动补齐。同时由于转置操作本身也受自动补齐规则的影响，因此同样需要特殊技巧来规避补齐。

# conv_state = tensor([2048, 3], bfloat16)  尾轴3*2=6bytes 不对齐
conv_state = tl.load(conv_state_ptr + conv_batch_offs * conv_batch_stride + doffs * 3 + tl.arange(0, 2048 * 3)) # 当成1D tensor load，此时由于numel对齐，不会自动补齐。
conv_state_T = conv_state.reshape(128, 16 * 3).trans().reshape(16, 3 * 128).trans().reshape(3 * 2048,) # 长轴(2048)裂出一根对齐轴(16)借给短轴(3)，从而让两个轴都对齐

AI Core进行计算的时候要先将数据搬运至片上内存，其空间通常远小于AI Core要处理的总数据量，以Atlas 800T/I A2产品为例，片上内存容量为192KB，默认开启doublebuffer后容量还会减至原来的一半。因此计算时需要对数据进行分块操作，每次只加载处理其中的一小部分数据。在 Tiling 分块优化中，BLOCK_SIZE、BLOCK_SIZE_SUB等分块参数直接影响算子性能，但手动调试参数组合效率低。triton.autotune是自动调优工具，能遍历预设的参数配置，通过实际运行对比性能，自动选择最优参数组合。自动遍历参数空间：针对BLOCK_SIZE、BLOCK_SIZE_SUB等 constexpr 类型的分块参数，批量测试不同取值的性能。
性能基准对比：以算子的执行耗时为指标，筛选出适配当前硬件的最优参数。
缓存调优结果：调优后的最优配置会被缓存，后续调用算子时直接复用，避免重复调优。UB或者L1 Size存在上限，当出现UB OVERFLOW错误时，需要减少单次搬运的数据量，以for循环的方式处理长序列场景。

单核数据运算暂略。