JustFit、KV状態を圧縮・リアルタイム調整し、24 GiBのMacBookで21万tokenのワークロードを完遂

JustFitは、ローカル大規模モデルのボトルネックを「重みがメモリに収まるか」だけでなく、実行時のlive set全体へと広げて捉える。量子化された重みに加え、KV cache、再構築用workspace、speculative decoder、output head、並列リクエストの状態が、いずれもApple Siliconのユニファイドメモリを奪い合う。著者らはMLXとQwen3.8-27B MXFP4をベースに、相互に連携する3つのコンポーネントを提案した。

KVExecは、TurboQuantから派生した4-bit表現でKVを保存し、256-token pageを使用する。decode時には圧縮状態を直接読み取り、prefill時には各layerでインプレースに展開して逆Hadamard回転を施し、直ちに標準のSDPAへ渡す。これにより、全layerの浮動小数点KVを同時に保持せずに済む。論文の推定では、このモデルにおける位置あたりのKV payloadはFP16の65,536 bytesから16,640 bytesへ減少する。ただし、圧縮だけでは不十分だと強調している。192K contextでは、単一layerの再構築だけでも約768 MiBを必要とする可能性があるためだ。

PhaseSwapは、実行phaseと「owner lease」に基づいてコンポーネントをoffloadまたはrestoreする。たとえばprefillの中間段階では、output headが使用する約644 MiBを一時的に解放する。StateTransは、リクエストの追加時や完了時、あるいは単一リクエストのspeculative decodingからcontinuous batchingへ移行する際に、target modelのcacheを維持し、page referenceを回収してコンポーネントのlifecycleを調整する。これにより、状態遷移のたびにコンテキスト全体を再構築する必要がなくなる。

プロセス上限を21,000 MiBに設定したM4 Pro MacBookでは、3回のテストすべてで196,608-tokenの入力と16,384-tokenの出力、合計212,992位置を処理し終えた。論文で使用されたmlx-vlm baselineが完了できたのは30,720位置だった。別の32K入力テストでは19.11 token/sを記録し、AIME 2026のsingle-seed評価では30問中29問に正解した。

これらの数値を単一の性能評価として扱うことはできない。容量、throughput、数学評価にはそれぞれ異なるワークロードが使用され、データも複数のコードrevisionにまたがっている。著者らはmlx-vlmのbranchと一部の固定snapshotを公開しているが、過去のGPU core数、macOS、Python、MLXのversion、および元のcheckpointのbit-level identityに関する情報は完全ではないと認めている。今後の焦点は、固定imageを用いた第三者による再現、ほかのモデルやhardwareへのgeneralization、そしてKV圧縮下でも長文コンテキストの品質が安定して維持されるかどうかにある。