実際の Snake ワークロードの最適化

採用したオプトインの経路は、MLX のコンパイル、16 トークン長のバケット、上限付きのトークナイズ済み質問プレフィックスキャッシュを組み合わせます。重みを変えず、モデルを量子化せず、予測をキャッシュせず、質問をまたいで双方向エンコーダの隠れ状態を再利用することも ありません。

ペアの完全ループテストで、同梱の最適化経路は 2,400 手を通して 毎秒 75.40 手を達成し、同じテストの eager 推論の 毎秒 70.82 手と比べられました：1.065×、約 6.5% です。どちらも死亡ゼロ、安全介入 2 回、2,400/2,400 手で実行アクションが同一でした。ウォームアップとキャッシュクリア後のアクティブ MLX メモリ増加は合計 0 バイトでした。

有効にする

laya-snake --optimize
laya-snake --optimize --max-speed

一般 API も同じオプトインの制御を公開します：

import laya_mlx as laya

agent = laya.load(
    "aac6fef/laya-multilingual-mlx",
    compile=True,
    pad_to_multiple=16,
    cache_prompts=True,
)

3 つのオプションは既定で無効で、既存の eager の挙動とベンチマーク構成を保ちます。コンパイルは入力形状に特殊化し、初回利用と新しい形状ではコンパイルコストがかかることがあります。重みやモジュール構造を変えた後は新しい Agent を構築してください。パディングは、設定されたコンテキスト上限を超えずに、系列長を要求された倍数へ切り上げます。マスクはパディングされたトークンを除外します。

cache_prompts=True は、Agent ごとに最大 128 個の不変な PreparedQuestion プレフィックスを、マーカー位置も含めて保持します。キャッシュキーにはトークナイザの同一性、特殊トークン、質問タイプ、順序づけられたレンダリング済み選択肢、instructions、プレフィックス予算が含まれます。状態は prepare 呼び出しごとに 1 回サニタイズしてトークナイズされ、その後、各質問プレフィックスと独立に連結されます。質問が変われば、適切なプレフィックスを生成または選択します。どの質問も依然として完全なモデルのフォワードパスを受けます。

形状と準備のアブレーション

このデモは着手ごとに 3 つの質問を投げます：方向、安全経路の推定、餌への到達可能性の推定です。したがってバッチサイズは 1 ではなく 3 です。サンプリングした 32 個の実際の記録盤面を通して：

• 多言語の系列長は 59、61、63、64 でした。16 トークン倍数では、すべてが 64 トークン のバケットに入ります。
• 英語の系列長は 66、68、69、70 で、80 トークン のバケットに対応します。
• 多言語入力を 64 から 96 にパディングすると、トークン単位の作業が 50% 増えます。別の短文 API フィクスチャが示唆する 93 から 96 への小さな調整ではありません。

候補は、計測対象のすべての形状を一度訪れた後、同一の各状態の中でローテーション順に実行しました。表には、トークナイズと出力変換を含む同期された Agent.predict レイテンシが含まれ、プランナ／UI の作業と初期の形状ウォームアップは除外します。

7 つすべての候補が、チェックポイントごとに 32/32 の盤面で、eager が提案・実行した方向と一致しました。表示される小数 4 桁の確率と推定値の最大差は、これらのサンプルで 0 でした。これは有限サンプルでの丸め済み出力の一致であり、内部の浮動小数点テンソルがビット単位で同一という主張ではありません。

アブレーションは、設計をふるいにかけるために上限付きのプレフィックス準備ラッパーを使いました。下の完全ループテストは、実際に同梱されている compile、pad_to_multiple、cache_prompts の API 実装を使います。その正しさのテストはさらに、状態の切り詰め、変化する criteria、マスクのサニタイズ、キャッシュの立ち退きの下で、準備済み ID とマーカーを比較します。

同梱の最適化経路は、実際のチェックポイントの完全な検証マトリクスにも合格しました：3 つのチェックポイントそれぞれについて、FP32 と FP16 で 63/63 の選択回答一致（合計 378/378）。較正済み確率の誤差は既存の許容範囲内にとどまりました。各構成はさらに、有限で決定的な呼び出しを 10 回追加で行い、計測上のアクティブメモリ増加が 0 バイトでした。最適化後の検証データ。元の eager 経路の構成あたり 100 回反復の結果は元のベンチマークレポートに残っています。

英語では、このサンプルでは実際の系列長でコンパイルする方が、より大きなバケットを強制するより良かったです。デモの既定は多言語です。一般 API の利用者は、コンパイルとプレフィックス再利用を有効にしつつ pad_to_multiple=None のままにできます。

完全ループのペアテスト

4 シード、各 600 手で、候補の順序はシードごとに交互にします。レンダリングにはトゥルーカラーの Rich 構築と ANSI シリアライズが含まれ、ターミナルエミュレータの描画は除外します。各手は新しい予測を行います。結果は 1 回のローカルなペア実行によるものです。

最適化経路は 1 つのシードで遅くなりました。したがって、6.5% はこの計測実行における統合的な改善であり、すべてのエピソードやマシンで保証される改善ではありません。以前の広範な速度スイープとこの後のペアテストは別の実行であり、速度向上を主張するために両者の絶対レートを差し引いてはいけません。完全なループデータ。

レイテンシだけでなくゲームプレイでモデルを選ぶ

両方のチェックポイントは 20 のペアシード × 300 手を実行し、チェックポイントの順序はシードごとに交互にしました。各エピソードは同じ初期状態、餌の RNG シード、簡潔な特徴量記述、サイクルシールドを使いました。範囲は固定です。これらは 300 手後のスコアであり、死亡や盤面の充填で終わる完全なゲームではありません。このモデル比較にターミナルのレンダリングは含まれません。

多言語モデルは餌への進行がより多く、このワークロードではより速かったため、既定のデモ用チェックポイントのままです。この結果が評価するのはこの特徴量支援ポリシーであり、汎用的な推論品質や支援なしの Snake モデルではありません。すべてのエピソードと推論。

再現

uv run --extra demo python -m experiments.snake_runtime \
  --output artifacts/snake/runtime-multilingual.json
uv run --extra demo python -m experiments.snake_runtime \
  --model models/hub/laya-mlx --bucket 80 \
  --output artifacts/snake/runtime-english.json
uv run --extra demo python -m benchmarks.snake_optimized \
  --output artifacts/snake/optimized-paired.json
uv run --extra demo python -m benchmarks.snake_models \
  --episodes 20 --steps 300 --output artifacts/snake/models.json

GPU 計測は順番に実行してください。まず 2 つのローカルモデルディレクトリをダウンロードします。チェックイン済みのソース録画が、サンプリングした正確な盤面状態を提供します。生のアブレーション：多言語、英語、初期の 96 トークンパイロット。

以前の簡潔プロンプト対詳細プロンプトの比較は、64 の状態でプロンプトの順序を交互にし、最初の 8 回のウォームアップ反復を捨てた後に中央値 11.80 → 9.29 ms を得ました。その元の場当たり的な記録は盤面のスナップショットを保存していなかったため、主要な再現可能アブレーションではなく補助的な根拠です。python -m benchmarks.snake_prompt は、状態、シード、完全な判断、手法を保存する再現可能な版を提供します。

実装は MLX の公式コンパイルガイドに従います：長寿命のコンパイル済み呼び出し可能オブジェクトと、通常の形状特殊化を使います。形状に依存する Python のモデルコードに shapeless=True は使いません。Context7 CLI のリクエストがネットワークエラーで失敗した後、現在のドキュメントは公式サイトで確認しました。