コンパイラファーストの状態空間双対性と推論のためのポータブルな$O(1)$自己回帰キャッシュ

概要: 状態空間モデルのリリースは通常、NVIDIAハードウェアに強く依存するカスタムのCUDAおよびTritonカーネルに結び付けられています。我々は、Mamba-2の状態空間双対性アルゴリズム——対角状態構造、チャンク可能な再帰、および静的制御フローを伴うeinsum支配の計算——がXLAのフュージョンおよびタイルパスが実際に最適化する部分にきれいにマッピングされ、カスタムカーネルが必須ではなく任意になることを示します。ハンドライティングされたカーネルなしで、XLAの下で成形された標準プリミティブとして完全な推論パス（プリフィル、キャッシュされた自己回帰デコード）を実装し、生成時にホスト同期を必要としない、コンパイル済みのオンデバイスキャッシュとして理論的なO(1)状態管理を実現します。実装は単一のJAXソースから、CPU、NVIDIA GPU、およびGoogle Cloud TPU上で修正なしに動作します。TPU v6e上で5つのモデル規模（1億3千万〜27億パラメータ）に渡り、XLA生成コードは単一ストリームプリフィルで約140 TFLOPS（15％のMFU）、デコード時には最大64％の帯域幅利用率に達します。貪欲デコードは64ステップにわたりPyTorch/CUDAの参照トークンと1対1で一致し、隠れ状態はfloat32の丸め誤差範囲内で同意します。このパターンは、同じ構造条件を満たす任意のSSM再帰に移植可能であり、成熟したXLAバックエンドを持つ任意のプラットフォームで利用可能です。実装は https://github.com/CosmoNaught/mamba2-jax にて公開され、Bonsai JAXモデルライブラリに統合されています。