pto-kernels

A collection of high-performance custom kernels for Ascend NPUs, built on top of pto-isa — the Parallel Tile Operation virtual instruction set architecture designed by Ascend CANN.

PTO focuses on tile-level operations, enabling efficient, composable kernel development targeting Huawei's Ascend AI processors, and ships as ready-to-use PyTorch (torch-npu) operators.

Why pto-kernels?

• Fast — hand-tuned tile-level kernels for Ascend NPUs, benchmarked against torch-npu built-ins.
• Drop-in — kernels are exposed as plain Python functions that operate on torch_npu tensors.
• Broad coverage — everything from elementwise ops (abs, swiglu) to linear-algebra primitives (tri_inv, matmul) to gated-linear-attention building blocks (GDN, KDA chunked recurrence, WY representation, KKT).
• Extensible — built on pto-isa, so new kernels can be written once at the tile level and reused across ops.

Installation

From PyPI (recommended)

Prebuilt wheels are published to PyPI for Python 3.10–3.12 on x86_64 and aarch64:

pip install pto-kernels

Requires a working torch-npu + Ascend CANN runtime environment to run kernels on-device; the package itself installs without one.

From source

Building from source is only needed for unreleased kernels or if you plan to contribute.

Prerequisites:

• A configured torch-npu environment
• Ascend toolkit installed at /usr/local/Ascend/ascend-toolkit

# One-time setup
make setup_once

# Install directly from GitHub
export CMAKE_GENERATOR="Unix Makefiles"
pip install -v git+https://github.com/huawei-csl/pto-kernels.git

Or build a wheel locally:

source /usr/local/Ascend/ascend-toolkit/set_env.sh
pip3 install -r requirements.txt
make wheel                          # produces pto_kernels-X.Y.Z-*.whl
pip install --force-reinstall pto_kernels-*.whl

Quickstart

import torch
import torch_npu  # noqa
from pto_kernels import pto_swiglu

x = torch.randn(4, 2048, device="npu", dtype=torch.float16)
y = pto_swiglu(x)  # fused SwiGLU on Ascend NPU

Available kernels

More end-to-end usage patterns live under examples/ (JIT C++ kernels, AI CPU custom ops) and tests/ (correctness against reference/torch_npu implementations).

Testing

make test

Repository structure

pto-kernels/
├── csrc/                  # C++ kernel source files
├── python/pto_kernels/    # Python bindings and utilities
├── examples/jit_cpp/      # JIT compilation examples
├── examples/aicpu/        # AI CPU custom-op examples
├── tests/                 # Test suite
├── scripts/               # Helper scripts
├── doxygen/               # API documentation config
└── CMakeLists.txt         # CMake build configuration

Contributing

Contributions are welcome! Whether it's a new kernel, a bug fix, or a benchmark, please read CONTRIBUTING.md before opening a pull request.

Release process

See RELEASE.md for how new versions are cut and published.

License

BSD-3-Clause-Clear — see LICENSE for details.