Backends & kernels
How one codebase covers CUDA, ROCm, Metal, Vulkan, WebGPU, and TPUs.
Pie shares one runtime, batch scheduler, paged KV cache, and inferlet host across all hardware targets. Each backend adds only a thin engine that launches compiled passes and its target kernels.
Engines
| Engine crate | Target devices |
|---|---|
engine-cuda | NVIDIA GPUs (CUDA), and AMD GPUs when built for ROCm |
engine-metal | Apple silicon, with eligible projections offloaded to the Neural Engine |
engine-vulkan | Vulkan GPUs on Linux, Windows, and Android |
engine-wgpu | WebGPU in browsers, and native desktop via wgpu |
engine-xla | Google TPUs via PJRT |
engine-cpu | Host RAM page pool and CPU fallback used for offloading |
Kernels
As the figure above shows, Pie targets five kernel backends:
- CUDA, ROCm, and Metal use hand-tuned native kernel sets.
- Vulkan and WebGPU share a single Slang shader source tree, compiled to SPIR-V for Vulkan and WGSL for WebGPU.
- TPUs compile graphs to StableHLO.
The Poem compiler fuses a model's forward pass into regions and lowers each region to these kernels. The eta compiler compiles custom epilogues into CUDA, HIP, Metal, or WGSL source at runtime.
Because Pie compiles kernels for the local device rather than bundling precompiled fat binaries for every GPU architecture, the entire binary is about 60 MB.
Device profiles
Kernel performance depends on the exact GPU chip. A profile records measured crossover sizes and the fastest kernel variant for each shape on one device. Pie ships profiles for common GPUs, and pie tune measures your hardware directly:
pie tune # quick run (a few minutes)
pie tune --full # full grid sweep (about an hour; use this when contributing a profile)
pie profile export apple-m2-max > apple-m2-max.tomlA profile file contains only hardware metadata and timing numbers, so contributing a profile for a new chip is a small pull request.
Build from source
cargo build --release -p pie --bin pie --features cuda # or: rocm, metal, vulkan, wgpu, xlaSee CONTRIBUTING.md for required toolchains and tests.