为智能体工作流优化MLX引擎

我们最近在 LM Studio 中发布了 mlx-engine v1.8.5。此次更新通过对 KV cache 进行 checkpointing，大幅提升了重复、长上下文 Agentic 工作流的性能。它还为 VLM 请求添加了连续批处理（continuous batching）。这项工作是开源的，你可以在这里查看 PR。

在这篇帖子中，我将解释这项更新解决的缓存复用问题、为什么当前开源 LLM 模型让回退（rewinding）变得更困难，以及新的磁盘-backed 缓存是如何工作的。我们的基准测试显示，额外 RAM 使用量最多降低 80%，吞吐量最高提升 2 倍，图像请求处理速度最高提升 3.5 倍。

看看这段视频，Adrien 使用新的 mlx-engine 在本地用 codex --oss 审查一个 URL 缩短器应用。

1、什么是 mlx-engine？

MLX Engine（mlx-engine）是一个针对 Apple Silicon 优化的 MIT 许可推理引擎。它由 LM Studio 创建和维护，使用 Apple 的 MLX 机器学习库，并基于 mlx-lm 和 mlx-vlm 等项目构建。MLX Engine 是 LM Studio 用于所有 MLX 推理的后端。

2、当前模型架构以及 mlx-engine 的不足之处

目前最受欢迎的两个开源模型是 Qwen 3.5（和 3.6）以及 Gemma 4。作为每个模型架构的一部分，它们使用了一些巧妙技巧来减少大上下文长度下的 KV cache 大小。Qwen 3.5 使用混合架构（hybrid architecture），Gemma 4 使用滑动窗口架构（sliding window architecture）。这些注意力策略在大上下文长度下减少了内存使用，但它们也使得 KV cache 无法任意回退。

让我们逐步了解 Gemma 4 如何处理推理。这个例子重点关注 Gemma 4 E2B；它交替使用“本地”注意力层（512 token 的滑动窗口）和“全局”注意力层。

步骤 1：Prompt prefill。为系统提示和用户消息计算 KV cache。

步骤 2：Decode。在计算助手推理内容和助手消息的同时，逐步构建 KV cache。

步骤 3：Rewind。将 KV cache 修剪回步骤（1），并附加助手消息（不包含之前的推理内容）。

因此，推理引擎需要解决的一个关键问题是：在为后续响应准备而回退 KV cache 时，避免重复计算。

3、我们如何改进 prompt caching

我们为这些 Agentic 用例设计了一种 KV cache 回退的解决方案。通过将 prompt cache 保存到磁盘并从中恢复，后续请求的 KV cache 就不需要重新计算。

将 KV cache 保存到磁盘

在 256 token 边界处复制并存储这些 KV cache，让我们在对应的 KV cache 张量仍然存在时，能够恢复精确的缓存前缀。如果 prompt 的某部分被编辑、从未计算过，或已从磁盘缓存中逐出，mlx-engine 会回退到重新计算该后缀。256 token 足够小，不会浪费太多重复计算的工作，同时也足够大，能保持磁盘缓存的高效性。

首先，在每个 256 token 边界（sequence len % 256 == 0）处，将本地注意力层的 KV cache 副本流式传输到磁盘写入后端。在模型处理 prompt 或生成新 token 的同时，后台磁盘写入进程一直在运行。每到 256 token 边界，系统就会复制最近 256 token 对应的 KV cache，并将其发送给磁盘写入器，后者再将该块持久化到磁盘。

由于 Apple Silicon 采用统一内存架构，我们将本地注意力的 KV cache 提交到磁盘，并从内存中逐出。这确保 mlx-engine 的内存占用随着活跃序列而扩展，而不是随着所有之前见过的序列而扩展。

从磁盘恢复 KV cache

首先，为每个 256 token 块计算一个 key。然后确定需要检索哪些全局和本地 KV cache 块。使用 prompt 的 key 列表和缓存类型，尽可能多地从磁盘加载 KV cache。对于从未计算过 KV cache（或其 KV cache 已从磁盘逐出）的 prompt 部分，则安排这些部分进行 prompt prefill。磁盘缓存是一个 LRU 存储，因此每当我们向磁盘存储保存或从中加载时，存储都会逐出最久未使用的 KV cache 张量。

这确保我们的磁盘存储能针对使用模式进行优化。如果引擎收到使用相同系统提示的短 prompt，那么系统提示的本地注意力 KV cache 就不会被逐出，但陈旧对话的 KV cache 会被逐出。如果引擎只接收一个不断增长的对话请求，那么较早的本地注意力 KV cache 将被逐出，以便为更长的全局注意力 KV cache 腾出空间。

磁盘缓存设计

我们将磁盘缓存设计为在模型卸载后自动清理。换句话说，缓存是临时的，不会留下持久化文件。

磁盘缓存是一个 scratch 文件，而不是一文件夹的独立缓存文件。我们将许多缓存记录打包到这一个文件中。每个 KV cache 条目都是一个序列化的 safetensors blob，引擎在内存中维护一张表，记录：“条目 X 从字节偏移 Y 开始，长度为 Z 字节。”当 KV cache 条目被逐出时，它们的字节范围会返回到空闲列表，并被后续记录重用；如果空闲空间到达文件末尾，文件就会被缩小。

我们通过使用操作系统在 /tmp 中的临时文件机制，并将所有查找元数据视为仅在模型生命周期内有效的状态，使磁盘缓存成为临时的。在模型卸载时，缓存存储会清除其内存索引并关闭 scratch 文件。如果模型进程退出，操作系统会关闭文件句柄并释放存储空间。

以及连续批处理

我们还为视觉模型运行器添加了连续批处理。关于连续批处理的实现和好处已经有很多讨论；Hugging Face 有一个很好的解释器。

连续批处理允许用户使用同一个模型进行并发请求处理。结合前面描述的 KV cache 改进，mlx-engine 现在可以用于严肃的 Agentic 工作负载。

4、基准测试

为了让性能改进更加具体，我们在配备 36 GB RAM 的 M3 Max MacBook Pro 上，使用 lmstudio-community/Qwen3.6-27B-MLX-4bit 运行了几项端到端的 LM Studio API 基准测试。

这些基准测试重点关注此次更新旨在改进的工作负载：并行聊天、长 prompt 处理以及重复的高分辨率图像 prompt。

4.1 基准测试：并行聊天吞吐量

设置：模型以 parallel=4 加载，然后通过 LM Studio API 并发发送四个短聊天请求。每个响应都允许自然停止。

结果：对于这个四路并行聊天工作负载，mlx-engine v1.8.5 端到端完成运行的速度大约快了 2.2 倍，输出 token 数量几乎相同。

4.2 基准测试：并行长 prompt 下的内存使用

设置：模型以 parallel=4 加载，然后通过 LM Studio API 并发发送四个大型 prompt。在模型加载后和运行完成后分别测量 RAM 使用量。

结果：对于这个并行长 prompt 工作负载，mlx-engine v1.8.5 在运行结束后额外 RAM 使用量减少了约 82%，同时保持了相似的挂钟时间和略高的总 token 吞吐量。这正是将非活跃 prompt-cache 记录移出统一内存所带来的预期收益。活跃序列仍需驻留在内存中，但陈旧的缓存记录不再需要持续累积在 RAM 中。

4.3 基准测试：重复高分辨率图像 prompt

设置：发送相同的图像 prompt 两次，每个请求生成一个 token。这隔离了处理图像扩展 prompt 和恢复 prompt cache 的成本。

结果：对于这个重复高分辨率图像 prompt，mlx-engine v1.8.5 完成第二个请求的速度大约快了 3.5 倍。胜利来自于恢复大部分图像扩展 prompt 的缓存。

原文链接：Improving LM Studio's MLX Engine for Agentic Workflows

汇智网翻译整理，转载请标明出处