Amal-David/mlx-porting-skill

MLXモデルの移植と最適化

ミッション

正確で、再現可能、アーキテクチャを考慮したMLX実装を生成すること。正確性が速度に優先します。速度やメモリに関する主張には、必ずハードウェア、ソフトウェアバージョン、ワークロード、ベースライン、品質ゲートを明記してください。

非交渉ルール

1. 受入時に信頼できないモデルコードを実行しないでください。 JSON、safetensorsヘッダー、ソースファイル、ライセンスを静的に検査してください。auto_map、カスタムモジュール、インストールフック、trust_remote_code はレビューゲートとして扱ってください。
2. ソースを固定してください。 リポジトリ、リビジョン、モデルファイル、トークナイザー/プロセッサーのリビジョン、ライセンス、チェックサムまたはアーティファクトマニフェストを記録してください。
3. 移植前にソースオラクルを構築してください。 決定論的なフィクスチャを固定し、意味のある境界で中間テンソルをキャプチャしてください。
4. 最初に最小限のeagerパスを移植してください。 基本的なパリティが通過するまで、量子化、コンパイル、カスタムカーネル、バッチング、投機的デコードは行わないでください。
5. 最適化は一度に一つの次元のみ変更してください。 測定結果とロールバック記録を保持してください。
6. ネイティブのMLX操作を優先してください。 カスタムMetalカーネルよりも、組み込みの融合操作、レイアウト変更、キャッシュ設計、mx.compile を先に試してください。
7. CUDAの慣用句を機械的に翻訳しないでください。 CUDAのテクニックは、そのMetal/MLXにおけるボトルネックと実装が実証されるまでは、単なる研究候補です。
8. スループットの背後に品質低下を隠さないでください。 音声、言語、ビジョン、生成モデルには、テンソル許容範囲に加えて、タスク固有の品質チェックを使用してください。
9. ライセンスと来歴のチェックなしに変換済み重みを公開しないでください。 元のモデルカード、帰属、生成設定、トークナイザー/プロセッサーファイル、変換レシピを保持してください。
10. 日々の研究自動化はレビューのみとします。 候補を収集しランク付けすることはできますが、実行手順を黙って書き換えたり、推奨事項をマージしたりしてはいけません。

ワークフロー

1. 検査と分類

以下を実行:

python3 scripts/inspect_model.py MODEL_OR_DIRECTORY --output inspection.json
python3 scripts/make_port_plan.py inspection.json --output PORT_PLAN.md
python3 scripts/recommend_optimizations.py inspection.json --markdown OPTIMIZATIONS.md

イントークとルーティング を読んでください。以下を確認:

• タスクとモデルドメイン;
• アーキテクチャファミリーとリカレント/キャッシュ状態;
• ソースフレームワークとカスタム操作;
• パラメータ数、dtype、結合/共有重み、シャード、アダプター;
• 前処理、トークン化、サンプリング、後処理;
• ライセンスとリモートコードリスク;
• ターゲットMac、メモリ予算、レイテンシ/スループット目標、品質目標。

アーキテクチャ、ソースリビジョン、または評価ターゲットが不明確なまま実装を開始しないでください。不明点は推測せずに PORT_PLAN.md に記録してください。

2. 最も近い実績のあるMLXリファレンスを選択

モデルサポートマップ と assets/architectures.yaml を参照してください。

この順序を使用:

1. 公式のMLX / MLX-LM実装;
2. テスト付きのアクティブなMLX-VLMまたはMLX-Audio実装;
3. アクティブなサードパーティMLX実装;
4. 上流ソース実装とアーキテクチャ論文;
5. 新しいMLX実装を必要とする研究プロトタイプ。

アーキテクチャパターンを再利用し、仮定をコピーしないでください。設定のセマンティクスとテンソルレイアウトを固定されたソースに対して検証してください。

3. ソースオラクルを確立

パリティとテスト に従ってください:

• 決定論的シードと推論モードを設定;
• 前処理後の正確な入力を保存;
• 埋め込み/フロントエンド後、すべてのブロックグループ、ボトルネック、キャッシュ更新、ロジット/潜在変数、デコーダー/ボコーダー出力で、形状/dtype/統計情報と選択されたテンソルをキャプチャ;
• ポータブルな .npz、.json、または .wav フィクスチャにソース出力を保存;
• dtypeと操作ごとに許容範囲を記録;
• 少なくとも1つの最小ケース、1つの通常ケース、1つの境界ケースを含めてください。

4. 最小限のeager MLXグラフを実装

コア移植メソッド と対応するアーキテクチャ実行手順書を読んでください:

• 高密度デコーダートランスフォーマー
• 混合専門家トランスフォーマー
• エンコーダートランスフォーマー
• エンコーダー-デコーダートランスフォーマー
• SSM、リカレント、ハイブリッド
• 拡散とフロー
• ビジョン-言語とオムニ
• ニューラルオーディオコーデック
• 自己回帰音声LM / TTS
• フローまたは拡散TTS
• ボコーダー
• ASR
• ストリーミング音声
• 分離と強調

初期実装の制約:

• 浮動小数点重みのみ;
• バッチサイズ1（バッチ処理が本質的でない限り）;
• compileデコレーターなし;
• カスタムカーネルなし;
• 明示的な状態とキャッシュオブジェクト;
• すべての非自明な変換の周りに形状アサーション;
• 可逆的な重みマップマニフェスト。

5. 重みを決定論的に変換

ソースキー、ターゲットキー、ソース形状、ターゲット形状、変換、dtype、結合/共有ルールを含む重みマップを作成してください。検証:

python3 scripts/validate_weight_map.py   --source source-manifest.json   --target target-manifest.json   --mapping WEIGHT_MAP.json

ロード時の strict=False に依存して、欠落または余分な重みを隠蔽しないでください。すべての例外を、意図的に無視されたもの、生成されたもの、共有されたもの、またはサポートされていないものとして分類してください。

6. パリティラダーを通過

scripts/compare_tensors.py を使用し、順に以下を通過:

1. 設定と前処理のパリティ;
2. 重みカバレッジと変換形状パリティ;
3. 単一プリミティブ/ブロックパリティ;
4. 段階的な中間パリティ;
5. エンドツーエンドの決定論的パリティ;
6. タスク品質パリティ;
7. キャッシュ/状態およびインクリメンタル生成パリティ;
8. 境界、長入力、ストリーミング、バッチパリティ。

パリティが失敗した場合は、障害アトラス を参照してください。失敗しているグラフを最適化しないでください。

7. 最適化を選択する前にプロファイリング

ベンチマーク を読んでください。以下を分離:

• プリフィル/エンコーダー/フロントエンド時間;
• トークンあたりまたはフレームあたりのデコード時間;
• コーデック/ボコーダー/後処理時間;
• コンパイルと初回実行コスト;
• ピークおよび定常メモリ;
• データ移動、同期、Pythonオーバーヘッド。

使用:

python3 scripts/benchmark_command.py --warmup 2 --runs 8 --output benchmark.json -- COMMAND ...

8. 最適化ラダーを適用

assets/optimization_guidance.yaml、assets/recommendation-taxonomy.yaml、assets/techniques.yaml を参照し、関連するガイドのみをロード:

1. MLXランタイム、コンパイル、カーネル
2. アテンションとKVキャッシュ
3. デコードとサービング
4. 量子化
5. トレーニングとファインチューニング

デフォルトの順序:

1. 意図しない評価とホスト転送を削除;
2. 正しいdtypeとテンソルレイアウト;
3. ネイティブの融合操作と高速SDPAを使用;
4. アロケーションを削減し、状態/キャッシュ更新を明示的に;
5. 安定した領域をコンパイルし、再コンパイルを制御;
6. アーキテクチャが許可する場合、プリフィル/フロントエンドをチャンク化またはストリーム;
7. KV/キャッシュポリシーとバッチングを最適化;
8. 重みを量子化し、正当化されればKV/状態も量子化;
9. 互換性のある自己回帰パスのみに投機的またはマルチトークンデコードを追加;
10. プロファイリングが残りのカーネルボトルネックを証明した後にのみカスタムMetalカーネルを作成。

すべての変更について、仮説、差分、正確性結果、ベンチマーク結果、メモリ結果、品質結果、維持/破棄の決定を記録してください。

9. パッケージ化と公開

パッケージ化と公開 に従ってください。以下を含めてください:

• 固定されたソースと変換コマンド;
• 互換性のあるMLXおよびライブラリバージョン;
• モデル/設定/トークナイザー/プロセッサーファイル;
• 正確な量子化レシピと除外;
• 決定論的スモークテスト;
• ベンチマークプロトコルと生の結果;
• 既知の制限事項;
• 元のライセンスと帰属;
• サポートされていない「より高速」または「ロスレス」という表現は使用しないでください。

10. エンジニアリングレポートを返す

最終応答では以下を要約:

• アーキテクチャと選択された実行手順書;
• ソースリビジョンとエビデンス参照;
• 実装と重みマップのステータス;
• パリティマトリックス;
• ベースラインと最適化後のメトリクス;
• 採用および却下された最適化;
• 残存リスクと未サポートパス;
• 再現可能なコマンドとアーティファクトの場所。

assets/ のテンプレートを使用してください。

停止すべき時

以下の場合、即興ではなく停止して報告してください:

• ソースライセンスが要求された配布を禁止している場合;
• 必要なソース動作が未レビューのリモートコードにのみ存在する場合;
• 段階的チェック後にパリティを特定できない場合;
• 最適化がマイクロベンチマークを改善するが、エンドツーエンドのレイテンシ、メモリ、または品質を悪化させる場合;
• カスタムカーネルにポータブルなフォールバックまたは十分なテストがない場合;
• ハードウェア、MLXバージョン、ワークロード、または品質ターゲットがパフォーマンス主張から欠落している場合。

メンテナンス

ソースレビュー、セキュリティ、および日常の候補パイプラインについては、メンテナンスと来歴 を読んでください。このスキルを配布する前に scripts/audit_skill.py と scripts/validate_sources.py を実行してください。