Post

TVM FFI Deep Dive: Shipping Native Kernels Without Tying Them to One PyTorch Build

A deep dive into why PyTorch native C++ extensions are often tied to a specific Torch build, how TVM FFI provides a stable framework-agnostic ABI, how PyTorch tensors cross the boundary through DLPack and TensorView, and when to choose TVM FFI versus PyTorch's stable ABI.

Posted Updated
By
10 min read
TVM FFI Deep Dive: Shipping Native Kernels Without Tying Them to One PyTorch Build

Why this post

If you ship CUDA/C++ kernels to Python users, sooner or later you run into the native-extension packaging problem:

I built a wheel that works with one PyTorch build. Will it still load and run correctly with the next PyTorch build?

Sometimes the answer is yes. Sometimes the answer is no in a very obvious way: the shared library fails to import. Sometimes it imports, but a symbol is missing only when a specific code path runs. Sometimes it runs, but metadata, allocation, stream, or autograd assumptions subtly differ.

This post is a deep dive into why that happens, what Apache TVM FFI provides, how it integrates with PyTorch tensors, and what tradeoffs it introduces compared with staying inside PyTorch’s extension ecosystem.

All upstream code links to TVM FFI are pinned to commit 88e066e so the references remain stable.

1. The problem: native PyTorch extension ABI instability

The most common way to expose a custom CUDA/C++ kernel to PyTorch is to build a native extension with some combination of:

  • torch::Tensor
  • ATen/C10 APIs
  • torch.utils.cpp_extension
  • pybind11
  • TORCH_LIBRARY custom operator registration

That is convenient. The binding can allocate tensors, call PyTorch helpers, read the current CUDA stream, and integrate with the dispatcher/autograd stack. But it also means the resulting .so is linked to PyTorch’s C++ world.

That can tie the wheel to the exact Torch binary used at build time.

Common failure modes

Native extension ABI problems usually show up in one of four ways.

Failure modeWhat it looks like
Import/load failureimport my_extension fails because a linked libtorch symbol or dependency cannot be resolved.
Runtime symbol failureThe module imports, but a call path fails when it touches a missing ATen/C10/Torch symbol.
Runtime behavior mismatchThe module imports and runs, but tensor metadata, allocation, stream, dispatcher, or autograd behavior differs.
Wheel matrix explosionYou rebuild the same kernel wheel for every Torch/CUDA/Python build combination.

The key issue is not CUDA. It is the Python-facing native boundary depending on PyTorch’s C++ ABI and libtorch symbols. If your CUDA kernel is conceptually a standalone kernel library, you probably do not want the whole wheel to be coupled to a single PyTorch build just because the binding layer uses torch::Tensor.

2. What TVM FFI provides

TVM FFI is an open ABI and FFI layer for ML systems. Its purpose is to expose native functions through a small stable ABI rather than through a framework-specific C++ extension ABI.

At the lowest level, TVM FFI represents every function call with one C ABI:

1
2
3
4
5

int tvm_ffi_c_abi(
    void* handle,
    const TVMFFIAny* args,
    int32_t num_args,
    TVMFFIAny* result);

TVMFFIAny is a tagged value representation that can carry integers, floats, strings, tensors, objects, and other supported values. The layout and calling convention are described in the TVM FFI Stable C ABI guide.

For C++ kernel authors, the API is much nicer than hand-writing the C ABI:

1
2
3
4
5
6
7
8
9
10
11
12

#include <tvm/ffi/tvm_ffi.h>

void add_one(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
  int64_t n = x.size(0);
  float* x_data = static_cast<float*>(x.data_ptr());
  float* y_data = static_cast<float*>(y.data_ptr());
  for (int64_t i = 0; i < n; ++i) {
    y_data[i] = x_data[i] + 1.0f;
  }
}

TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, add_one);

Then Python loads the shared library:

1
2
3
4
5
6
7
8

import torch
import tvm_ffi

mod = tvm_ffi.load_module("my_kernels.so")

x = torch.randn(1024, device="cuda", dtype=torch.float32)
y = torch.empty_like(x)
mod.add_one(x, y)

At a high level, the Python binding flow is:

  1. TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, add_one) generates an exported C symbol named __tvm_ffi_add_one. The generated wrapper has the stable TVM FFI C ABI: (void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result).
  2. tvm_ffi.load_module("my_kernels.so") calls into TVM FFI’s module loader, which dlopens the shared library and exposes exported __tvm_ffi_* functions as Python-callable tvm_ffi.Function objects.
  3. mod.add_one(x, y) calls Function.__call__. The binding layer packs Python arguments into TVMFFIAny values. For torch.Tensor, the packer uses DLPack to create a TVM FFI tensor view without copying storage.
  4. The exported C ABI wrapper unpacks those TVMFFIAny arguments into the typed C++ signature, so the C++ function receives tvm::ffi::TensorView x and tvm::ffi::TensorView y.
  5. The kernel reads metadata and pointers from TensorView and writes into the output tensor allocated by the caller.

The key source links for that binding path are:

The native .so can be built ahead of time and packaged in a wheel. That avoids runtime JIT build latency. The TVM FFI quickstart shows both raw compiler commands and CMake integration:

3. How PyTorch tensors cross the TVM FFI boundary

The important part is that PyTorch tensors do not need to be converted into a custom Python object manually. In normal code, you pass torch.Tensor objects directly to the TVM FFI function:

1
2
3
4
5
6
7
8

import torch
import tvm_ffi

mod = tvm_ffi.load_module("my_kernel.so")

x = torch.randn(1024, device="cuda", dtype=torch.float32)
y = torch.empty_like(x)
mod.my_kernel(x, y)

There is no explicit from_dlpack(...) call here. The conversion happens inside TVM FFI’s Python binding layer when it packs the arguments for mod.my_kernel(...). PyTorch tensors implement the DLPack protocol, and TVM FFI uses that protocol to create a zero-copy tensor view for the native call.

The relevant upstream code path is:

At a high level:

1
2
3
4
5

PyTorch Tensor
  -> implicit DLPack-compatible view while packing Python call arguments
  -> TVM FFI tensor representation
  -> tvm::ffi::TensorView in C++
  -> raw data pointer + metadata in the kernel

The memory is shared. No GPU copy is required just to call the native function.

You can make the same conversion explicit for debugging or explanation:

1
2
3
4
5
6

import torch
import tvm_ffi

x_torch = torch.randn(1024, device="cuda")
x_tvm = tvm_ffi.from_dlpack(x_torch, require_contiguous=True)
x_torch_again = torch.from_dlpack(x_tvm)

That explicit round trip is not required for mod.my_kernel(x, y); it just shows the mechanism. The memory is shared either way.

The TVM FFI tensor docs describe this pattern in more detail:

4. What tvm::ffi::TensorView contains

On the C++ side, a TVM FFI kernel usually accepts tvm::ffi::TensorView arguments:

1
2
3
4
5
6
7
8
9
10
11
12

void kernel(tvm::ffi::TensorView input, tvm::ffi::TensorView output) {
  if (input.dtype() != DLDataType{kDLFloat, 32, 1}) {
    TVM_FFI_THROW(TypeError) << "expected float32";
  }

  if (!input.IsContiguous()) {
    TVM_FFI_THROW(ValueError) << "expected contiguous input";
  }

  float* input_data = static_cast<float*>(input.data_ptr());
  float* output_data = static_cast<float*>(output.data_ptr());
}

The view carries:

MetadataAPI
data pointerdata_ptr()
dtypedtype()
devicedevice()
rankndim()
shapesize(i)
contiguityIsContiguous()

This is enough for most kernels that only need raw pointers, shapes, dtypes, and devices.

Stream handling

For CUDA kernels, you also need to launch on the framework’s current stream. TVM FFI provides:

1
2
3
4

#include <tvm/ffi/extra/c_env_api.h>

cudaStream_t stream = reinterpret_cast<cudaStream_t>(
    TVMFFIEnvGetStream(input.device().device_type, input.device().device_id));

In a PyTorch process, that maps to PyTorch’s current stream for the tensor’s device. See:

5. Pros and cons of TVM FFI

Pros

  • Decouples from PyTorch C++ ABI. The Python-facing native module does not need to link against libtorch just to accept PyTorch tensors.
  • Framework-agnostic path. The same ABI can work with PyTorch, JAX, Paddle, NumPy, CuPy, or any framework that supports DLPack.
  • Smaller wheel matrix. A kernel library can target CUDA/platform/Python ABI instead of every PyTorch build.
  • Good generated-code story. DSLs and codegen systems can emit a stable C ABI or a TVM FFI C++ wrapper.
  • Stream interop without framework C++ linkage. GPU kernels can use TVMFFIEnvGetStream(...).

Cons

  • Another runtime dependency. Users need apache-tvm-ffi.
  • Less PyTorch C++ convenience. You should not call ATen ops or use torch::Tensor unless you intentionally link PyTorch again.
  • More explicit validation. You must validate dtype, shape, device, and contiguity yourself.
  • Output allocation discipline. It pushes you toward preallocated outputs, which may require API changes.
  • Potential extra shim layer. Existing libraries may still need a conversion from tvm::ffi::TensorView to their local descriptors.

6. Alternative: PyTorch stable ABI

The alternative is to stay PyTorch-specific and use PyTorch’s stable ABI support where possible.

That may be the right choice when:

  • the kernel is only ever meant for PyTorch;
  • the binding needs deep ATen or dispatcher integration;
  • cross-framework reuse is not important;
  • the extension naturally lives as a PyTorch custom op.

But PyTorch stable ABI is still a PyTorch solution. It does not make the same kernel callable from JAX, Paddle, NumPy, or CuPy. TVM FFI is a better fit when the native code is fundamentally a framework-neutral kernel library.

OptionBest forLimitation
TVM FFIFramework-neutral kernels and DSL-generated librariesRequires TVM FFI runtime and explicit metadata handling
PyTorch stable ABIPyTorch-only extensionsDoes not generalize beyond PyTorch

7. Practical checklist

If I were moving a native kernel to TVM FFI, I would check:

  • Can the core implementation run with raw pointers, shapes, dtypes, and device IDs?
  • Can Python allocate outputs before calling the kernel?
  • Can the C++ side avoid torch::Tensor, ATen, pybind11, and TORCH_LIBRARY?
  • Does CUDA code launch on TVMFFIEnvGetStream(...)?
  • Are all tensor inputs validated for dtype, shape, device, and contiguity?
  • Is the TVM FFI .so built during wheel build rather than JIT-built at import?
  • Does the wheel contain all dependent .so files needed at load time?

If the answer is yes, TVM FFI is a strong way to make the kernel wheel more portable across framework and PyTorch versions.

Appendix: what DLPack is and where zero-copy happens

DLPack is a small C data structure convention for sharing tensor metadata and data pointers across array frameworks. The important object is a managed tensor capsule containing a data pointer, dtype, shape/strides, device type/device id, and a deleter/lifetime callback.

DLPack does not allocate memory or copy data by itself. It describes an existing tensor allocation so another framework can view the same memory safely. That is why TVM FFI can accept a PyTorch tensor without copying it: PyTorch exports the tensor as DLPack, and TVM FFI imports that DLPack capsule as a TVM FFI tensor.

In PyTorch, the Python-level tensor type exposes the DLPack protocol:

On the TVM FFI side, the binding layer consumes that DLPack object:

The zero-copy property comes from preserving the original data pointer through this chain. TVM FFI constructs a view/handle over the same allocation; it does not allocate a second GPU buffer just to call the kernel.

References

Machine Learning Performance, Compiler
tvm-ffi pytorch cuda dlpack abi native-extension kernel-library python-wheels cpp gpu
This post is licensed under CC BY 4.0 by the author.