Laya‑MLX – 在 Apple Silicon 上实现快速、类型化决策推理

是什么 – 一个 Python 包，可在 macOS Apple Silicon（M 系列）GPU 上完全本地运行 Laya 决策语言模型系列。该库将原始 Convai Innovations 检查点移植到 MLX（Apple 的 Metal 加速张量库），从而在无需 PyTorch、🤗 Transformers 或云调用的情况下，对单个简短问题实现低于 15ms 的延迟。支持类型化输出（选择、评分、二元“noul”），避免逐 token 解码，实现确定性、低延迟响应。

---

核心理念

概念	Laya‑MLX 的实现方式
类型化决策	不生成自由文本，而是通过一次前向传播返回结构化答案（选择、评分或二元“noul”）。避免逐 token 解码，实现确定性、低延迟输出。
双向编码器	输入（状态 + 问题）使用 ModernBERT-large 或 mmBERT-base 主干网络编码，然后专用头部生成请求类型的概率。
本地运行，无运行时依赖	所有推理均在 MLX 中执行；分词使用 Hugging Face 的 Rust 分词器，已编译进 wheel 包。无需 PyTorch/Transformers 二进制文件。
Apple Silicon 优化	可选 compile=True、前缀缓存和填充技巧在 M3 Max 上可提升约 6% 速度；库还提供预转换的 FP16 检查点以供 GPU 使用。

---

快速开始（30 秒）

pip install laya-mlx            # 核心库
pip install 'laya-mlx[demo]'    # 可选的演示工具

import laya_mlx as laya
agent = laya.load("aac6fef/laya-mlx")   # 首次使用时下载 FP16 检查点
result = agent.predict(
    "I was billed twice. Please refund the duplicate.",
    {
        "department": {
            "type": "choice",
            "instructions": "Who should handle this?",
            "criteria": ["billing", "technical", "sales"]
        }
    },
)
print(result["answers"]["department"])   # → "billing"

支持 macOS 14+、Python 3.11+ 和任意 Apple Silicon GPU（M1–M3）。首次调用会下载模型；后续调用完全离线。

---

可用检查点

模型 ID（加载）	编码器	参数量	上下文长度	语言
aac6fef/laya-mlx	ModernBERT-large	421 M	512	英语
aac6fef/laya-multilingual-mlx	mmBERT-base	322 M	1 024	多语言
aac6fef/laya-typed-decisions-mlx	ModernBERT-large	421 M	1 024	英语（类型化决策工作流）

以上三个均为上游 Convai Innovations 权重的精确 FP16 转换版本，托管于 Hugging Face。您也可将 laya.load 指向原始 Hub ID（如 convaiinnovations/laya 等），库会自动下载并转换。

---

性能（M3 Max，FP16）

指标	英语（Laya 421M）	多语言（Laya-multilingual 322M）
单个简短问题的中位延迟	13.4 ms	7.4 ms
95 百分位延迟	13.9 ms	7.8 ms
50 问题批量（批量大小 64）吞吐量	146 q/s	395 q/s
每请求峰值 GPU 内存	944 MiB	688 MiB
优化版 laya‑snake --optimize --max‑speed	75.4 moves/s（比 eager 快约 6.5%）

数值包含分词、张量准备、推理、校准和结果格式化；不包含模型加载时间。

---

主要 API 接口

agent = laya.load(
    checkpoint,          # HF 仓库 ID 或本地路径
    dtype="float16",   # 或 "float32", "bfloat16"
    batch_size=16,       # 单次前向传播最大问题数
    device="gpu",       # "cpu" 也支持（慢得多）
    compile=False,       # 启用 MLX 编译以提升速度
    cache_prompts=False, # 保留分词后的提示以复用
)

# 预测 – `system_one` 是别名
answers = agent.predict(state, questions)

*state 可为普通字符串、JSON 字典或历史消息列表。questions 为字典，每个条目描述期望的回答类型（choice, score, noul）。返回值包含：

• answers（结构化结果）
• action.act_probability（原始头部概率）
• token 使用统计
• 四舍五入至小数点后四位（与上游模型一致）。

路由器辅助工具

适用于需要自动选择语言特定检查点的应用：

router = laya.Router(dtype="float16", max_loaded=2)
out = router.predict(state, triage_questions())
print(out["routing"])   # 例如："multilingual"

路由器最多可常驻 max_loaded 个模型，也可通过 Router(preload=True) 预加载。

---

命令行工具

命令	用途
laya-mlx predict …	从 JSON 文件或内联字符串运行单次推理。
laya-mlx convert …	将 Hugging Face 检查点转换为 MLX 兼容目录（safetensors + config）。
laya-snake	交互式终端演示，运行经典贪吃蛇游戏，每步调用模型。使用 --optimize 可启用编译后的快速路径。

所有 CLI 均接受与 laya.load 相同的 --model 参数。

---

开发与测试

• 依赖项：项目使用 uv 实现可复现环境。运行 uv sync --extra dev 可拉取测试、基准测试和参考扩展。
• 测试：单元测试将 MLX 实现与原始 Transformers 头部在小型随机模型上进行对比；完整检查点验证涵盖分词、校准概率、确定性重复和内存增长。
• 基准测试：benchmarks/run 测量延迟/吞吐量；结果存储于 benchmarks/results，并在 BENCHMARKS.md 中总结。
• 研究：docs/ 文件夹包含关于性能瓶颈的深度报告及 10 倍加速的设想（数学、工程、实现层面）。脚本位于 experiments/。

---

许可与署名

• 代码 – Apache-2.0（见 LICENSE）。
• 权重 – 原始 Laya 权重版权归 Convai Innovations 所有；通过上述 Hugging Face 仓库以相同许可证重新分发。
• 移植 – MLX 重实现及周边工具由 mizorewww 编写，并借鉴了上游 NandhaKishorM/laya 仓库的部分内容（NOTICE 中有 MIT 风格署名）。

---

谁会使用它？

• 产品团队：需要在设备端实现确定性、低延迟路由或分类（如工单分诊、紧急程度评分、二元策略检查）。
• 开发者：构建仅限 macOS 的 AI 助手或边缘服务，不希望将数据发送至云端。
• 研究人员：希望在 Apple GPU 上对比 MLX 与 PyTorch/Transformers 性能，或探索进一步加速技术。

---

TL;DR – Laya‑MLX 为 Apple Silicon 上的 Laya 决策模型提供即开即用的高性能推理库，拥有简洁的 Python API、终端贪吃蛇演示，以及完整的转换、基准测试和模型卡发布工具链。