PyPTO

导言

• 浦江现场性能优化时，原始 triton 的 GDN性能相对于H200的triton性能很差；
• 接入了 Ascend C的若干GDN算子实现，提速了一倍；
• 接入 mojo_opset 的 casual_conv1d; 或者 huawei-csl的 sglang的高性能 solve_tril 。性能又提升一波，但和竞品还是有差异。
• 2012实验室对 PyPTO 有精细调优，现在想接入

PyPTO（发音为 pai p-t-o），全称是 Parallel Tensor/Tile Operation，是近期（特别是伴随华为昇腾生态 CANN 以及像 DeepSeek 这样的大模型部署时）崭露头角的一个面向 AI 加速器的高性能算子编程框架。

很多开发者在手写 Triton 遇到性能瓶颈或者做跨平台（尤其是向 NPU 迁移）部署时，会听到这个名字。

---

PyPTO 定位

如果在 Triton 的世界里，你是在用 Python 语法写“类似于 CUDA 的底层线程块（Block）逻辑”；那么在 PyPTO 的世界里，你是在写一种大模型时代的“算子 DSL（领域特定语言）”。

PyPTO 的核心设计理念是 Tile-centric（以 Tile 为中心）。现代 AI 芯片（无论是 GPU 还是 NPU）都有复杂的内存层次结构，计算都是以数据块（Tile）为单位送入计算单元（如矩阵乘法引擎、Tensor Core 或华为的 Cube/Vector Core）。PyPTO 将这些复杂的算子逻辑拆解为一系列可组合的 Tile 级指令，并在上层提供非常 Pythonic 的 API。

它不仅仅是一个 Kernel 编译器，而是一个从“模型层级”贯穿到“底层指令”的软垫层，特别擅长处理诸如 Sparse Attention、MoE 动态路由、KV Cache 动态更新等大模型里极其复杂的融合算子。

---

PyPTO vs. OpenAI Triton

这两种框架的目标都是“让开发者用 Python 写出媲美甚至超越手写 C++/CUDA 的高性能算子”，但在架构设计和执行范式上有很大区别：

对比维度	OpenAI Triton	PyPTO
抽象层级与核心	Block 与指针操作：需要开发者手动计算指针偏移（offsets）、掩码（masks）和 block size，非常贴近内存寻址。	多级抽象（Tensor/Tile/Block）：开发者可以用 Tensor 层写算法逻辑，用 Tile 层做性能调优。Tile 是框架的“一等公民”，屏蔽了大量繁琐的指针计算。
调度与执行模型	**SPMD (单程序多数据)**：所有的实例运行相同的 Kernel 代码（类似 CUDA grid/block 模型）。	**MPMD (多程序多数据)**：支持在不同类型的处理器核上调度不同的任务（例如在 NPU 上同时调度标量、向量计算和矩阵计算），这对于异构硬件发挥极致性能至关重要。
编译流水线	Python AST -> Triton IR -> LLVM IR -> PTX/AMDGPU 汇编。	多层计算图转换：Tensor Graph -> Tile Graph -> Block Graph -> 最终生成 PTO 虚拟指令，然后再编译为目标硬件机器码。
硬件倾向	GPU 优先：深度绑定英伟达 CUDA 生态，虽然也在适配 AMD 和其他后端，但设计哲学深受 GPU 影响。	硬件原生（含 NPU 优势）：虽然理论上支持 GPU，但在 NPU（如华为 Ascend/CANN 体系）上适配极深，能通过内置宏和内存编排直接解决“内存墙”问题。

一句话总结他们的区别：

• Triton 给你的感觉是：“我给了你更简单的 Python 语法，但你依然是一个 GPU 内存工程师，你需要去算 offset。”
• PyPTO 给你的感觉是：“你告诉我这些 Tile（数据块）要怎么流动、切分和计算（比如量化、RoPE 旋转、Matmul），我和底层的调度器来帮你榨干硬件算力。”

总结建议

如果你现在正在使用 Triton，并且：

1. 你的目标依然是在 NVIDIA GPU 上开发泛用内核：Triton 目前生态依然最完善。
2. 你的目标是向国产化硬件（如昇腾 NPU）迁移，或者在部署超大模型（如 DeepSeek-V3）时遇到了现有框架无法逾越的性能瓶颈：强烈建议花时间学习 PyPTO。它可以让你把精力集中在“算法的并行切分逻辑”上，而不是跟晦涩的底层汇编或内存边界检查死磕。

---

基本语法与编程范式

PyPTO 提供的是一种分层抽象的设计：对于算法验证，它长得非常像 PyTorch；但对于算子调优，它提供了 Tile 级别的精细控制。

层面一：Tensor 级别的算法构建 (高层)

在这个级别，它与 PyTorch 高度相似，支持动态 Shape 和符号化编程。

import pypto
import pypto.nn.functional as F

# 这是一个使用混合精度的算子/模块示例
class FP16GEMM(pypto.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # 参数定义，可以显式控制数据存储和计算精度
        self.weight = pypto.Parameter(
            pypto.randn(out_features, in_features, dtype=pypto.float32)
        )
        self.bias = pypto.Parameter(
            pypto.zeros(out_features, dtype=pypto.float32)
        )
      
    def forward(self, x):
        # 运行时计算会自动基于 PTO 的 pass 优化，转变为底层的高效 Tile 计算
        # 比如自动转换为 FP16 进行 Matmul 运算
        return F.linear(x, self.weight, self.bias)

层面二：Tile / Block 级别的算子编排 (底层核心)

当你需要像写 Triton 那样去追求极致性能时，PyPTO 不是让你去写 tl.load / tl.store 和一堆掩码，而是通过操作 Tile 对象来完成的。底层 API 常常涉及到明确的 Cast（类型转换）、Matmul、Reshape 和流水线调度。

一个抽象的、伪代码级别的 Tile 计算逻辑大概是这样的：

import pypto.tile as pt

def fused_attention_kernel(Q_tile, K_tile, V_tile, out_tile):
    # PyPTO 的内核逻辑强调 "基于数据块的操作"
  
    # 1. Tile 级别的数据加载（框架处理显存到缓存的搬移，无需手动算指针 mask）
    q_blk = pt.load(Q_tile)
    k_blk = pt.load(K_tile)
  
    # 2. 调用硬件优化的 Tile 级原子操作，如矩阵乘
    # 框架会自动映射到 GPU 的 TensorCore 或 NPU 的 Cube 核
    score_blk = pt.matmul(q_blk, pt.transpose(k_blk))
  
    # 3. 计算 Softmax (同样在局部 Tile 内完成)
    prob_blk = pt.softmax(score_blk)
  
    # 4. 与 V 进行计算并写回
    v_blk = pt.load(V_tile)
    res_blk = pt.matmul(prob_blk, v_blk)
    pt.store(out_tile, res_blk)

*注：具体 API 可能会随版本更新有所迭代，但核心逻辑是声明 Tile 的流水线和计算依赖。*

实践

• https://gitcode.com/cann/pypto/tree/master/models/qwen3_next
• https://github.com/huawei-csl/pto-kernels