AI 推論與開發工具
Axon、強い型付けのDSLから5種類のLLMバックエンドを生成、MLXの中央値レイテンシをTransformersの48.3％に低減
Axonは、モデルアーキテクチャをシンボリックなテンソル形状を持つ関数型仕様として記述し、PyTorch、Triton、JAX、MLX、ネイティブvLLM向けにコンパイルする。467件の推論テストでは多くの生成モデルが高速化した一方、一部のBF16モデル、エンコーダーモデル、シーケンス・ツー・シーケンスモデルでは、精度または性能の低下も確認された。
新たに提案されたAxonは、LLMアーキテクチャを特定のフレームワークに依存するPythonクラスから、複数のバックエンド向けにコンパイル可能な強い型付けのドメイン固有言語（DSL）へと抽象化することを目指している。開発者はHaskellに似た構文を使い、Attention、MoE、状態空間モジュール、パラメーターパス、KV cacheを記述する。`Tensor[B,S,D]`のような型により、コンパイラーはコード生成前にシンボリック形状を検査できる。コンパイル処理では、まず関数とライブラリーを展開し、型付きのGraph IRへ正規化した後、PyTorch、Triton、JAX、MLX、またはvLLM向けにloweringする。デプロイ環境にAxon runtimeを残す必要はない。
[論文](https://arxiv.org/abs/2608.19889)では、47種類のモデル定義と、パラメーター数135M～32Bのモデルを使用し、467件の推論比較を実施した。Hugging Face Transformersとの比較では、PyTorch、Triton、JAX、MLXの公式な中央値ベースの高速化率は、それぞれ7％、12％、91％、107％だった。MLXの126件のテストでは、実行時間の中央値がベースラインの48.3％となった。PagedAttentionとKV cacheをサポートするネイティブvLLMアーキテクチャとしてコンパイルした場合の高速化率の中央値は58％で、より低速なTransformers互換パスによるロードとは異なる。この区別は、[vLLMのモデル実装ドキュメント](https://docs.vllm.ai/en/latest/models/supported_models/)に記載されているネイティブbackendとTransformers backendの違いとも一致する。
技術的な価値は、監査可能な単一の仕様を、トレーニング、デスクトップ、サーバー推論で共通して利用できる点にある。これにより、新しいアーキテクチャごとに5種類の実装を書き直すコストを削減できる。ただし、現時点の結果は著者によるテストのみに基づいている。T5やmT5などの短いforward workloadでは、顕著に低速な外れ値が見られ、5つのモデルではMLX BF16の出力が数値的一致の基準を満たさなかった。また、論文には公開コードリポジトリーが付属していない。エンジニアリングチームは今後、コンパイラーと完全なベンチマークが公開されるか、動的制御フローと量子化モデルがサポートされるか、さらに継続的バッチ処理やマルチユーザーservingでも高速化を維持できるかに注目すべきだ。