Apache Software Foundation が管理するオープンソースの機械学習コンパイラフレームワークで、多様なハードウェアバックエンドに対応した深層学習モデルの最適化デプロイを実現する。
Apache TVM は、深層学習モデルを CPU・GPU・FPGA・ASIC・マイクロコントローラなど多様なハードウェア上で効率的に実行するためのオープンソースコンパイラフレームワークである。ワシントン大学の陳天奇(Tianqi Chen)らが 2018 年に発表し、Apache Software Foundation のトップレベルプロジェクトとして開発が進められている。
TVM の最大の特徴は、単一のモデル定義から異なるハードウェア向けの最適化コードを自動生成できる汎用性にある。クラウド GPU での推論からエッジデバイスでの組み込み推論まで、一貫したツールチェーンで対応する。
TVM のフロントエンド IR である Relay は、型付き関数型言語として設計されている。PyTorch・TensorFlow・ONNX などのフレームワークからインポートしたモデルは Relay IR に変換され、グラフレベルの最適化が適用される。
Relay の各演算ノードは TIR(Tensor IR)に展開される。TIR はループネスト・メモリスコープ・ベクトル化幅などの低水準情報を表現し、ハードウェア固有の最適化がこのレベルで行われる。
コンパイル済みモデルは TVM ランタイム上で実行される。ランタイムは C++ で実装されており、Python・Java・Rust・Go・JavaScript からバインディング経由で呼び出せる。モバイル(Android/iOS)やブラウザ(WebAssembly)上での実行も可能。
| レイヤー | 役割 | 主要最適化 |
|---|---|---|
| フロントエンド | モデルインポート | ONNX, PyTorch, TF 変換 |
| Relay IR | グラフ最適化 | 演算融合, 定数畳み込み, 量子化 |
| TIR | ループ最適化 | タイリング, ベクトル化, アンローリング |
| AutoScheduler | 自動チューニング | Ansor による探索 |
| コード生成 | デバイス固有コード |
| CUDA, LLVM, Metal, Vulkan, WASM |
TIR のスケジュール空間(タイルサイズ・アンローリング係数・メモリ配置など)をパラメータ化し、実機上でのベンチマーク結果に基づいて最適な組み合わせを探索する。XGBoost をコストモデルとして使用し、探索効率を向上させる。
AutoTVM の後継として開発された自動スケジューラ。人手のスケジュールテンプレートを必要とせず、TIR のループ構造から自動的に探索空間を生成する。探索にはスケッチベースのアルゴリズムが使用され、AutoTVM と比較して 2〜3 倍の性能向上を達成する場合がある。
TVM の最新の自動チューニングフレームワーク。AutoTVM と Ansor の機能を統合し、よりモジュール化された設計となっている。探索空間の定義・コストモデル・探索アルゴリズムを個別にカスタマイズ可能。
MLC-LLM(Machine Learning Compilation for LLM)は、TVM ベースの LLM 推論エンジンである。以下の特徴を持つ。
単一のモデル定義から、CUDA GPU・Metal(macOS/iOS)・Vulkan(Android/Windows)・WebGPU(ブラウザ)・OpenCL 向けの最適化推論エンジンを生成する。
INT4/INT3 量子化をコンパイラレベルで最適化し、量子化カーネルの自動生成を行う。GroupQuantize 方式により、4-bit 量子化でも品質劣化を最小限に抑える。
Llama 2/3、Mistral、Phi、Gemma などの主要 LLM アーキテクチャに対応。Hugging Face からのモデル変換パイプラインが整備されている。
TVM の汎用性は対応バックエンドの広さに現れている。
| バックエンド | ターゲット | 用途 |
|---|---|---|
| CUDA | NVIDIA GPU | クラウド推論・学習 |
| ROCm | AMD GPU | クラウド推論 |
| Metal | Apple GPU | macOS/iOS 推論 |
| Vulkan | Android/Windows GPU | モバイル・デスクトップ |
| OpenCL | 汎用 GPU/FPGA | 組み込み |
| LLVM | x86/ARM CPU | CPU 推論 |
| WebAssembly | ブラウザ | Web 推論 |
| microTVM | マイクロコントローラ | エッジ AI |
| Vitis AI | Xilinx FPGA | FPGA 推論 |
Triton は NVIDIA GPU に特化した高性能カーネル記述言語であり、TVM は多様なハードウェアに対応する汎用コンパイラフレームワークである。NVIDIA GPU のみをターゲットにする場合は Triton(+ torch.compile)の方がエコシステムとの統合が深い。複数のハードウェアをサポートする必要がある場合は TVM が優位。
TensorRT は NVIDIA GPU 上での推論に特化した最適化エンジンであり、TVM は学習・推論の両方をカバーする。TensorRT は NVIDIA が公式にサポートする製品であり安定性が高い一方、TVM はオープンソースコミュニティ駆動で柔軟性が高い。
スマートフォンや IoT デバイスへのモデルデプロイでは、TVM の自動チューニング + クロスコンパイルが威力を発揮する。ターゲットデバイス上でのプロファイリング結果をフィードバックし、デバイス固有の最適化を自動適用する。
WebAssembly + WebGPU バックエンドにより、サーバーレスの完全クライアントサイド推論を実現する。プライバシー要件が厳しいアプリケーションやオフライン動作が求められる場面で有用。
モデルサイズとターゲットハードウェアに依存しますが、ResNet-50 クラスのモデルで数時間、LLM(7B パラメータ)で数十時間〜数日かかる場合があります。ただし、一度チューニングした結果はキャッシュされ、同じモデル・同じハードウェアでは再利用可能です。
TVM は量子化演算(INT8/INT4 GEMM)のカーネルも自動チューニングの対象とするため、ハードウェア固有の量子化命令(NVIDIA の INT8 Tensor Core、ARM の SDOT/UDOT)を効率的に活用できます。TensorRT のような専用エンジンと同等の性能を、より多様なハードウェア上で実現できます。
MLC-LLM プロジェクトが最も簡単な方法です。mlc_llm パッケージをインストールし、Hugging Face からモデルをダウンロード → mlc_llm compile で CUDA/Vulkan 向けにコンパイル → mlc_llm serve で推論サーバーを起動できます。NVIDIA GPU であれば CUDA バックエンド、AMD GPU であれば Vulkan バックエンドを選択します。