tvm-ffi — 한 wheel 로 여러 framework 를 지원하는 ABI

강의의 출발점은 kernel 라이브러리 저자의 실제 고통. 한 라이브러리가 PyTorch wheel · JAX plugin · Paddle 확장을 따로 만들고, 거기에 Python 3.9–3.13 까지 곱하면 wheel 매트릭스가 폭발한다. 같은 CUDA kernel 인데.

Tianqi 의 진단 — 현존 FFI 의 두 큰 부류 가 다 이 문제를 풀지 못한다.

tvm-ffi 의 위치는 그 사이 — 최소 C ABI로 framework 사이의 wrapper 비용을 없애면서, tensor 가 first-class인 ML 친화적 디자인.

tvm-ffi 는 새로 만든 게 아니다. Apache TVM 안에서 10년 동안 써온 PackedFunc 메커니즘을 — TVM 본체와 분리해 — 독립 패키지로 추출한 것. “TVM 을 안 써도 PackedFunc 만 쓸 수 있게”가 핵심.

• 2017–18 TVM 의 시작 — Tianqi 가 TVM 을 만들 때부터 “generated kernel 을 어떤 frontend 에서도 부를 수 있어야 한다” 가 디자인 결정. PackedFunc 가 그 contract.
• 2019–22 다중 frontend — TVM 이 PyTorch · TensorFlow · ONNX · MXNet 에 generated kernel 을 export 하면서 PackedFunc 가 공통 layer 가 됨.
• 2023–24 다른 프로젝트의 채택 — FlashInfer 가 같은 메커니즘을 빌려 attention kernel 을 PyTorch/JAX 양쪽에 export. TileLang · cuteDSL 도 비슷.
• 2025 tvm-ffi 분리 — 같은 메커니즘을 “TVM 안의 한 부분” 이 아니라 독립 표준으로. 다른 프로젝트가 TVM 전체 의존 없이 가져다 쓸 수 있게.

tvm-ffi 의 디자인 핵심은 의도적인 작음. C ABI 가 작아야 — 모든 컴파일러, 모든 OS, 모든 Python 버전에서 안정적. 큰 ABI 는 한 군데가 깨지면 전체가 깨진다.

tvm-ffi 의 또 다른 핵심 — language-agnostic. 함수의 진짜 구현은 C++ 또는 Rust 에 있고, Python 은 thin wrapper. 그러나 그 반대도 가능 — Python 함수를 C++ 에서 호출.

# Python 측 — kernel 라이브러리 함수 호출
import tvm_ffi
import torch

flash = tvm_ffi.load_module("libflashinfer.so")

q = torch.randn(1, 8, 1024, 128, device="cuda")
k = torch.randn(1, 8, 1024, 128, device="cuda")
v = torch.randn(1, 8, 1024, 128, device="cuda")

# tvm_ffi 가 torch tensor 를 DLPack 으로 view → kernel 로
out = flash.attention(q, k, v)
# out 도 torch tensor 처럼 보임 (DLPack import)

// C++ 측 — kernel 함수 정의
#include "tvm/ffi/function.h"

void attention(NDArray q, NDArray k, NDArray v,
               NDArray out) {
  // q.data, q.shape, q.strides 직접 접근
  // CUDA kernel 호출
  launch_attention_kernel(q, k, v, out);
}

// 등록 — 한 줄
TVM_FFI_REGISTER_GLOBAL("flashinfer.attention")
    .set_body_typed(attention);

그리고 Rust binding 의 의의 — Rust 로 짠 kernel 라이브러리가 Python 에서 그대로 호출 가능. 새 시스템 언어 (Rust, Mojo, Zig) 가 ML kernel 영역에 들어오는 길이 자연스럽게 열림.

PackedFunc 는 tvm-ffi 의 핵심 추상. 모든 함수가 같은 C 시그니처를 가진다. 인자가 5개든 0개든, 반환값이 tensor 든 string 이든. 호출 시점에 type code 로 dispatch.

// PackedFunc 의 C 시그니처 (단순화)
typedef int (*PackedCFunc)(
    TVMValue* args,
    int* type_codes,
    int num_args,
    TVMValue* ret,
    int* ret_type_code,
    void* resource_handle
);

// 모든 함수가 이 한 형태
// add(a, b), attention(q, k, v), tokenize(s) — 다 같은 시그니처

# Python 측 typed 호출 — 외피는 익숙한 형태
add: tvm_ffi.PackedFunc = registry.get("my.add")
result = add(3, 5)

# 내부적으로는
# [TVMValue(3), TVMValue(5)] + [int_code, int_code]
# → C 함수 호출 → ret 으로 결과 반환

그리고 PackedFunc 와 함께 다니는 추상 — Object. tensor 외 string · list · dict · module 등 다양한 ML primitive 가 같은 refcounted handle 위에 올라간다. PyTorch 의 tensor 든 numpy 의 ndarray 든 — DLPack 으로 같은 PackedFunc 를 통과.

repo README 가 자랑스럽게 보여주는 채택 사례. tvm-ffi 가 풀려는 문제가 가짜가 아닌, 실제 라이브러리 저자들이 마주친 문제임을 증명.

• FlashInfer — LLM attention kernel 라이브러리. PyTorch 가 1차 frontend 였지만, JAX 사용자도 같은 attention 을 원함. tvm-ffi 채택 후 — 같은 wheel 이 양쪽에서.
• TileLang — kernel DSL. 생성된 kernel 을 PyTorch/JAX/PaddlePaddle 에 동시에 export.
• cuteDSL — NVIDIA 의 CUTE DSL Python wrapper. 같은 패턴.
• framework extension — PyTorch 의 새 op 를 framework 의 ABI 변경에 영향 없이 배포. PyTorch 2.x → 2.y 의 silent break 문제 회피.
• multi-language — Python / C++ / Rust 의 cross-language ML 인프라.

각자가 자기 자리가 있다. 한 framework 에서 한 op 만 추가는 PyTorch ext 가 빠르고 쉬움. generic Python ↔ C++ 는 pybind. 여러 framework 에 같은 kernel 을 배포는 tvm-ffi. 강의의 메시지는 “하나가 다른 둘을 대체한다” 가 아니라 “각자의 자리가 있다”.

repo README 가 명시 — v0.1.0 이후 RFC stage. C ABI 안정성을 우선시하지만 일부는 여전히 변경 가능. 0.X.Y 의 버전 표기 자체가 “X 가 breaking change” 라고 선언.

• ABI freeze 는 1.0 부터 — 0.X 단계에서는 minor 버전 변화에 binary 호환이 깨질 수 있음.
• autograd 통합 미흡 — kernel 을 부르는 건 OK 인데, 자동 미분 그래프에 등록하는 건 framework 별로 wrap 필요.
• schema 시스템 미성숙 — type code 가 단순한 dispatch 까지만. 복잡한 ML schema (tensor of nested structure) 는 커스텀.
• error reporting — C ABI 의 한계로 stack trace 가 제한적. Python 레이어에서 보강 중.
• build system 통합 — CMake / setup.py 와의 매끄러운 통합이 framework-specific FFI 보다 한 단계 거칠다.

강의에서 Tianqi 가 짚는 야망. CPython 의 limited API 가 “한 wheel 이 여러 Python 버전을 지원” 을 가능케 한 것처럼 — tvm-ffi 가 ML 영역의 limited API가 되는 길.

• 1.0 의 ABI freeze — 그 이후 모든 wheel 이 binary forward-compatible.
• frontend 표준화 — DLPack 이 이미 거의 universal. PackedFunc 가 그 옆에 자리 잡으면 ML kernel 의 lingua franca.
• 새 frontend 의 진입 비용 절감 — Mojo, Zig, Rust 가 ML kernel 영역에 진입할 때 “tvm-ffi binding 만 만들면 모든 기존 라이브러리 사용 가능”.
• edge / mobile 까지 — TVM 의 본래 강점인 edge 배포에 자연스러운 확장. 같은 PackedFunc contract 가 desktop · server · mobile 위에서.

실습 시퀀스

1. tvm-ffi hello world — repo example 의 첫 “add” 함수. C++ 로 짜고 등록, Python 에서 호출. 빌드 매트릭스 확인.
2. tensor 통과 — torch.Tensor 를 인자로 받는 C++ 함수. DLPack import/export 가 어떻게 일어나는지 코드에서.
3. 같은 wheel JAX 에서 호출 — 같은 빌드된 .so 를 JAX 에서 import. 둘 다 동작 확인.
4. FlashInfer attention 호출 — 실제 채택 사례 코드 읽기. tvm-ffi 가 어떻게 wrap 되는지 패턴 학습.
5. Rust binding 시도 — Rust crate 로 같은 함수 호출. cross-language 가 실제로 어떻게 도는지 손에 잡아 보기.

• transcript 가 실패해서 본 노트의 호출 비용 비교는 일반적 추정. 실제 benchmark 는 강의 영상 또는 repo 의 perf 자료에서 확인 필요.
• ABI 호환 범위의 정확한 명세 — “PyTorch 2.0–2.6” 같은 구체 호환은 추정. 실제 지원 매트릭스는 repo 의 CI 결과에서.
• autograd 통합 로드맵 — 1.0 에 들어갈지, 별도 layer 로 남을지 강의에서 다뤘을 가능성.
• JAX custom_call 과의 깊은 비교 — 강의에서 JAX 진영의 FFI 와 어떻게 다른지 정량적 비교.
• TVM 본가와의 의존 관계 — tvm-ffi 가 정말 standalone 인지, TVM runtime 의 일부 의존이 남아있는지 확인.
• Mojo / Zig 의 실제 binding — “새 언어가 진입할 수 있다” 는 미래형. 실제 구현이 어디까지 진척됐는지.