ホームへ戻る

GitHub Repo

Ollama メンテナーが MLX の状態修正を提案、ツール呼び出し後の隠れたメモリ使用量を削減

10月5日に提案された修正は、投機的デコードからのフォールバック後も大きなバッファーを参照し続ける再帰状態に対処し、余分なメモリがプレフィックスキャッシュに保持されるのを防ぐ。作者は、ツールを12回呼び出すテストの最終ターンで使用量が3.52 GiB減ったと報告しているが、提案はまだマージされていない。

Mattruffoni · CC BY-SA 4.0 · Image source
zh-Hant

Ollama のメンテナー dhiltgen は10月5日、投機的デコードからのフォールバック後に MLX 推論パスが過大なバッファーを保持する問題に対処する PR #18805 を提出した。作者によると、ツールを12回呼び出す A/B テストで、修正により最終ターンの MLX 保持メモリが3.52 GiB減少し、論理キャッシュの使用量は同程度だった。確認時点では提案は審査中で、対象ブランチは release_v0.40.0 となっている。修正提案

問題は状態のライフサイクルにある。トークンごとに取得される再帰状態は、フォールバック後にその一つが選択されても、より大きな基底バッファーを参照し続ける場合がある。次のモデル計算の前にリクエストが終了すると、そのバッファーは存続中のキャッシュとともに残る。キャッシュの計上対象は選択された状態だけだが、実際の割り当てはそれより大きい可能性がある。この提案では、バッファーを共有している可能性のある復元状態に印を付け、リクエスト終了時にプレフィックスキャッシュへ保持する前にまとめてコンパクションする。通常のモデルステップで状態が置き換えられる際には印を消去し、復元後にバッファーが分離されることを確認するテストも追加する。仕組みの説明

この修正の発端は、9月24日のコミュニティ報告にさかのぼる。報告者は、32 GB メモリ搭載の M1 Max Mac Studio で Ollama 0.34.2 と 0.34.4 を使い、MTP 投機的デコードを有効にして qwen3.6:27b-mlx を実行し、ツール呼び出しで終わるリクエストのたびに約0.43 GiBが追加されることを確認した。この増加はプレフィックスキャッシュのカウンターには現れなかったが、通常のチャットを対照とした場合には同じ増加は見られなかった。これは特定の構成で再現された事例であり、すべてのモデルやバックエンドに当てはまるとは限らない。元の事例とスクリプト

ローカルエージェントを運用する場合、容量の見積もりではキャッシュ上の計上値と実行時の割り当ての両方を確認する必要がある。MLX の公式ドキュメントによると、get_active_memory() は使用中のメモリをバイト数で返し、キャッシュバッファーは含まない。そのため、OS が表示する使用量とも一致するとは限らない。MLX メモリのドキュメント

エンジニアは今後、マージとリリースを追い、連続したツール呼び出しで修正前後のメモリ使用量の推移、スワップ領域、レイテンシを比較する必要がある。3.52 GiBは作者による単一テストの結果であり、他のモデルでの改善幅、コンパクションのコスト、関連する累積問題を完全に解消できるかどうかは、引き続き検証が必要だ。

出典

  1. mlx: compact restored recurrent state after speculative rollback — PR #18805
  2. MLX runner: tool-call requests retain memory outside the prefix-cache budget — Issue #18620
  3. mlx.core.get_active_memory — MLX documentation