Amal-David/mlx-porting-skill

MLX 模型移植与优化

使命

生成一个正确、可复现、且考虑架构特点的 MLX 实现。正确性重于速度。任何有关速度或内存的声明都必须注明硬件、软件版本、工作负载、基线和质量门限。

不可协商的规则

1. 在初始阶段不执行不可信的模型代码。 以静态方式检查 JSON、safetensors 头部、源文件和许可证。将 auto_map、自定义模块、安装钩子和 trust_remote_code 视为审查门限。
2. 锁定来源。 记录仓库、修订版本、模型文件、分词器/处理器修订版本、许可证和校验和或工件清单。
3. 在进行移植前建立源 Oracle。 冻结确定性测试夹具，并在有意义的边界处捕获中间张量。
4. 首先移植最小的即时执行路径。 在基础一致性测试通过之前，不进行量化、编译、自定义内核、批处理或推测解码。
5. 一次只改变一个优化维度。 保留测量和回滚记录。
6. 优先使用原生 MLX 操作。 在编写自定义 Metal 内核之前，尝试内置融合操作、布局更改、缓存设计和 mx.compile。
7. 不要机械地翻译 CUDA 的经验法则。 在证明其 Metal/MLX 瓶颈和实现之前，CUDA 技术仅作为研究候选。
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

阅读 接受与路由。确认:

• 任务和模型领域；
• 架构系列和循环/缓存状态；
• 源框架和自定义操作；
• 参数数量、数据类型、绑定/共享权重、分片和适配器；
• 预处理、分词、采样和后处理；
• 许可证和远程代码风险；
• 目标 Mac、内存预算、延迟/吞吐量目标和质量目标。

如果架构、源修订版本或评估目标仍然不明确，则不要开始实现。将不确定性记录在 PORT_PLAN.md 中，而不是猜测。

2. 选择最接近的经过验证的 MLX 参考

查阅 模型支持映射 和 assets/architectures.yaml。

按以下顺序使用:

1. 官方的 MLX / MLX-LM 实现；
2. 活跃的 MLX-VLM 或 MLX-Audio 实现（附带测试）；
3. 活跃的第三方 MLX 实现；
4. 上游源实现加架构论文；
5. 需要全新 MLX 实现的研究原型。

重用架构模式，而不是复制假设。对照锁定的源验证配置语义和张量布局。

3. 建立源 Oracle

遵循 一致性与测试：

• 设置确定性种子和推理模式；
• 预处理后保存确切的输入；
• 在嵌入/前端之后、每个块组、瓶颈、缓存更新、logits/潜在向量和解码器/声码器输出处捕获形状/数据类型/统计信息和选定的张量；
• 以可移植的 .npz、.json 或 .wav 夹具保存源输出；
• 按数据类型和操作记录容差；
• 至少包含一个最小、一个普通和一个边界情况。

4. 实现最小的即时 MLX 图

阅读 核心移植方法 和匹配的架构操作手册:

• 密集解码器 Transformer
• 混合专家 Transformer
• 编码器 Transformer
• 编码器-解码器 Transformer
• SSM、循环和混合模型
• 扩散和流模型
• 视觉语言和全能模型
• 神经音频编解码器
• 自回归音频语言模型 / TTS
• 流或扩散 TTS
• 声码器
• ASR
• 流式语音
• 分离与增强

初始实现约束:

• 仅使用浮点权重；
• 批量大小为 1，除非批处理是内在的；
• 不使用编译装饰器；
• 不使用自定义内核；
• 显式的状态和缓存对象；
• 在每个非平凡变换周围放置形状断言；
• 可逆的权重映射清单。

5. 确定性地转换权重

创建一个权重映射，包含源键、目标键、源形状、目标形状、变换、数据类型和绑定/共享规则。验证它:

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. 在选择优化前进行性能分析

阅读 基准测试。区分:

• 预填充/编码器/前端时间；
• 每 token 或每帧解码时间；
• 编解码器/声码器/后处理时间；
• 编译和首次运行成本；
• 峰值和稳态内存；
• 数据移动、同步和 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. 纠正数据类型和张量布局；
3. 使用原生融合操作和快速 SDPA；
4. 减少分配并使状态/缓存更新显式化；
5. 编译稳定区域并控制重新编译；
6. 在架构允许的情况下，对预填充/前端进行分块或流式处理；
7. 优化 KV/缓存策略和批处理；
8. 量化权重，然后如理由充分则量化 KV/状态；
9. 仅为兼容的自回归路径添加推测解码或多 token 解码；
10. 仅在性能分析证明存在剩余的内核瓶颈后，才编写自定义 Metal 内核。

对于每次更改，记录假设、差异、正确性结果、基准测试结果、内存结果、质量结果以及保留/回退决策。

9. 打包和发布

遵循 打包与发布。包括:

• 锁定的源和转换命令；
• 兼容的 MLX 和库版本；
• 模型/配置/分词器/处理器文件；
• 精确的量化配方和排除项；
• 确定性冒烟测试；
• 基准测试协议和原始结果；
• 已知限制；
• 原始许可证和署名；
• 不使用未经支持的“更快”或“无损”措辞。

10. 返回工程报告

最终响应必须总结:

• 架构和选定的操作手册；
• 源修订版本和证据参考；
• 实现和权重映射状态；
• 一致性矩阵；
• 基线和优化后的指标；
• 接受和拒绝的优化；
• 剩余风险和不受支持的路径；
• 可复现的命令和工件位置。

使用 assets/ 中的模板。

何时停止

在以下情况下停止并报告，而不是临时拼凑:

• 源许可证禁止所请求的分发；
• 所需的源行为仅存在于未经审查的远程代码中；
• 经过分阶段检查后无法定位一致性失败原因；
• 某项优化改进了微基准测试，但恶化了端到端延迟、内存或质量；
• 自定义内核缺少可移植的回退或充分的测试；
• 性能声明中缺少硬件、MLX 版本、工作负载或质量目标。

维护

有关源审查、安全性和每日候选管道，请阅读 维护与来源追溯。在分发此技能之前，运行 scripts/audit_skill.py 和 scripts/validate_sources.py。