freenode
AI & ML

PyTorch prototypes AOTI registration for DSL kernels via TVM-FFI

A new torch.library.register_aot_impl path aims to compile custom DSL kernels into PyTorch’s ahead-of-time Inductor stack.

PyTorch has proposed a prototype API, torch.library.register_aot_impl, that would let developers compile domain-specific language (DSL) kernels into the framework’s ahead-of-time Inductor (AOTI) path through TVM-FFI.

Today, many custom operators sit outside the compiled export story or force graph breaks. The prototype is meant to give those kernels a first-class registration route so they can ship as AOTI-compatible implementations rather than remaining eager-only. TVM-FFI is the bridge: DSL kernels would be compiled and bound into the same ahead-of-time artifact flow Inductor already uses for exported models.

The work is explicitly labeled prototype and has not been categorized for release notes, so it is still an early integration experiment rather than a committed user-facing feature. If it lands and stabilizes, it would widen how external and DSL-authored kernels plug into PyTorch’s compile-and-export toolchain.