MLX Engine Python API¶

MLX Engine Python API 是 LM Studio 开发的一套用于大语言模型（LLM）推理的高层 Python 接口。它基于 Apple 的 MLX 框架构建，专门针对 Apple Silicon (M系列芯片) 硬件进行了优化^[001-TODO__mlx-engine.md]。

该 API 提供了模型加载、流式文本生成、结构化输出以及视觉-语言模型（VLM）推理等功能，并集成了推测解码（Speculative Decoding）、KV 缓存量化等高级优化技术^[001-TODO__mlx-engine.md]。

核心架构¶

MLX Engine 采用模块化设计，根据模型类型的不同，主要有两条处理路径^[001-TODO__mlx-engine.md]：

双路径设计¶

系统根据模型的 config.json 中的 model_type 自动选择初始化路径^[001-TODO__mlx-engine.md]：

1. ModelKit：
   • 用于处理标准的文本模型或拥有专用视觉插件（Vision Add-on）的视觉模型。
   • 全功能支持：支持 KV 缓存量化、跨提示词缓存和推测解码^[001-TODO__mlx-engine.md]。
2. VisionModelKit：
   • 用于处理通用的视觉模型（基于 mlx-vlm）。
   • 兼容性广：适用于大多数视觉模型，但不支持量化、缓存复用和推测解码等高级特性^[001-TODO__mlx-engine.md]。

关键组件¶

• CacheWrapper：负责 KV 缓存的管理与增量更新，实现了跨提示词缓存¹，允许在不同会话间复用计算结果^[001-TODO__mlx-engine.md]。
• Vision Add-ons：位于 mlx_engine/model_kit/vision_add_ons/，提供针对特定模型架构（如 Gemma3, Pixtral, Mistral3）的专用处理逻辑^[001-TODO__mlx-engine.md]。

核心 API 用法¶

1. 加载模型¶

使用 load_model 函数加载模型。对于支持高级优化的模型，可以配置 KV 缓存参数^[001-TODO__mlx-engine.md]。

from mlx_engine import load_model

model_kit = load_model(
    "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
    max_kv_size=4096,        # 最大 KV cache 大小
    kv_bits=None,            # KV 缓存量化位数 (None=不量化)
    kv_group_size=64,        # KV 缓存量化组大小
)

2. 文本生成 (流式)¶

使用 create_generator 进行流式推理。它支持采样参数（如 temp, top_p）和结构化输出约束^[001-TODO__mlx-engine.md]。

from mlx_engine import create_generator, tokenize

prompt_tokens = tokenize(model_kit, "你好，世界")

for result in create_generator(
    model_kit, 
    prompt_tokens,
    temp=0.7, 
    top_p=0.9,
    stop_strings=["结束"],
    json_schema='{"type":"object",...}',  # 结构化输出 (基于 Outlines)
):
    print(result.text, end="", flush=True)

3. 视觉模型推理¶

视觉模型的生成需要传入 Base64 编码的图像数据^[001-TODO__mlx-engine.md]。

model_kit = load_model("mlx-community/pixtral-12b-4bit")

for result in create_generator(
    model_kit, 
    prompt_tokens,
    images_b64=["base64_encoded_image_string..."],
):
    print(result.text)

高级功能¶

推测解码¶

通过加载一个小型的 Draft Model（草稿模型）来加速主模型的推理速度^[001-TODO__mlx-engine.md]。

from mlx_engine import load_draft_model, is_draft_model_compatible

# 检查兼容性并加载
if is_draft_model_compatible(model_kit, "path/to/draft-model"):
    load_draft_model(model_kit, "path/to/draft-model")
    # 后续生成将自动利用草稿模型加速

KV 缓存量化¶

通过减少 KV Cache 的位宽（例如设置为 4bit 或 8bit）来显著降低显存占用，从而支持更长的上下文或更大的批处理大小^[001-TODO__mlx-engine.md]。

结构化输出¶

集成 Outlines 库，允许通过 JSON Schema 强制模型输出符合特定结构的 JSON 数据^[001-TODO__mlx-engine.md]。

限制与注意事项¶

• 视觉模型限制：使用通用 VisionModelKit 加载的模型不支持 KV 缓存量化、跨提示词缓存和推测解码^[001-TODO__mlx-engine.md]。
• 已知 Bug：Qwen VL 系列模型（如 qwen2_vl）的视觉适配器当前存在 Bug，已被注释掉^[001-TODO__mlx-engine.md]。
• Python 版本：项目针对 Python 3.11 编译，其他版本可能存在兼容性问题^[001-TODO__mlx-engine.md]。

支持的模型¶

除了标准的文本模型外，API 明确支持以下视觉模型系列^[001-TODO__mlx-engine.md]：

• Llama-3.2-Vision (11B)
• Pixtral (12B)
• Qwen2-VL (7B)
• Llava-v1.6 (Mistral 7B)

Sources¶

• 001-TODO__mlx-engine.md