近年来，大语言模型和生成式 AI 丰富和改变了我们几乎每个人的生活。在 ChatGPT 这类模型之外，开源模型在这年内同样飞速发展，日益涌现的开源大模型在拥有更好质量的同时，也为个性化的和领域特定的模型部署带来了充分的机遇和挑战：一方面，在云服务器的大模型部署上，能更高效地在多卡上跑更大模型、支持更多并发请求的解决方案取得了新突破，另一方面，大模型在本地设备上的部署也开始取得可喜的进展，这些进展让大家在笔记本、浏览器、手机等设备上运行高质量的量化大模型变成可能。（文档/代码链接见文末）
大模型部署的未来会去向何方？我们相信这两种部署场景会一直共存下去，于是，让每个人都有能力在云服务器和本地部署大模型就变得极其重要。
在过往和现在，包括我们自己 MLC LLM 项目的上一个版本在内，许多 LLM 推理项目都在服务器部署和本地部署上有两套不同的解决方案，两侧各自有各自的实现和优化。比如，服务器上的 LLM 部署方案通常有更好的 continuous batching 和多卡支持，而本地的 LLM 部署方案在跨平台支持上更加便捷。然而，我们觉得整个领域对于能够将所有这些技术结合在一起的大模型部署方案是有需求的。事实上，很多应用在服务器部署或者本地部署的技术往往能直接被应用在另一侧。虽然站在现在这个时间点，诸如 continuous batching 之类的服务器部署技术对于本地的部署场景没有十分的使用价值，我们觉得随着 LLM 的应用发展，这些技术有能够在未来大显身手的潜力。比如，如果 LLM 今后成为操作系统的一部分，那么在本地支持处理多请求和 agent 任务的能力就会尤为重要。因此，我们关心的问题是，我们是否可以打造一个统一的、高效的、能够兼具服务器部署和本地部署能力的 LLM 引擎？
在这篇文章里，我们想为大家介绍 MLC LLM Engine (MLCEngine)，一个能够在不同平台上 universally 部署的 LLM 引擎。MLCEngine 既在服务器上具备 high-throughput, low-latency 的 LLM serving 能力，同时又支持在各种本地环境下无缝部署当今高质量的大语言模型。
图1. MLCEngine: Universal LLM Deployment Engine
要支持在各种平台上做 universal deployment 有其独特的挑战。首先，我们需要支持一系列 GPU 编程模型和 runtime，从而能够在不同平台上加速 LLM 推理。一般而言，这样的支持过程往往伴随着大量且重复的 engineering effort，在 CUDA 这样的平台上我们想尽可能利用成熟的 library，而在 Vulkan、WebGPU 这类标准 library 还有待完善的平台上也要能够有相应支持。此外，不同平台上的应用开发通常也使用不同的编程语言和环境，这让在诸如 Swift, Kotlin, Javascript 等平台上引入 LLM 引擎变得愈加复杂。
机器学习编译 (Machine Learning Compilation) 这一技术是帮助我们克服这些挑战、并使 MLCEngine 能够在不同平台上运行的核心 —— 我们使用 Apache TVM 在各种硬件和平台上生成 GPU 代码。同时，我们搭建了一个支持工业级 state-of-the-art LLM serving 优化（包含 continuous batching, speculative decoding, prefix caching 等）的 portable 的 runtime 架构，这大大提升了我们在云服务器和本地环境中同时部署 MLCEngine 的能力。
图2. MLC LLM Workflow
我们将通过 TVM 编译器编译好的模型 library 和这一 universal runtime 结合在一起，在不同的应用开发语言中提供一套直接的 OpenAI-style 的 chat completion 接口，使得我们在不同平台、不同环境下的都能够以同一套标准的 OpenAI 接口与 MLCEngine 进行交互。同时，虽然 MLCEngine 在不同语言里有各自的 binding，所有这些本地应用背后都和 MLCEngine 在服务器上的应用共享同一个引擎、同一套代码。
如何使用 MLCEngine？让我们从 Chat 开始
接下来我们展示如何使用 MLCEngine。最便捷的方式是直接启动 chat CLI，这是一个与 LLM 做交互式对话的界面。下面这个命令能够在终端里通过 chat CLI 和 4-bit 量化的 Llama3 8B 模型进行交互对话：
mlcllm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f161-MLC
图3. Chat CLI
在这行命令背后，我们首先会以 just-in-time 的方式编译 Llama3 模型，生成其在 GPU 上的代码。这种本地、即时的模型编译方式让我们在不同的 GPU 和操作系统上都能够直接运行大模型。
MLCEngine 支持一系列大语言模型，包括 Llama3, Mistral/Mixtral, Phi3, Qwen2, Gemma 等 LLM 模型。
通过 OpenAI-style API 使用 MLCEngine
MLCEngine 作为一个 universal LLM 引擎，仅仅能够在不同平台上运行大模型还不足以达成我们的目标。实际上，拥有一套标准的、开发者熟悉的、易用的 API 非常关键。我们在所有的平台上都使用 OpenAI-style API 作为与 MLCEngine 交互的接口。现在我们来看看在不同平台上的这套 API。
云服务器上的 REST API
LLM serving 作为最经典的 LLM 应用之一，其 server 内部运行着一个 LLM 引擎以处理从网络接受的到所有请求。MLCEngine 提供了一个支持完整 OpenAI API 的 REST server。下面这个命令可以在 localhost 里启动一个REST server，用 4-bit 量化的 Llama3 模型的处理用户发来的 generation 请求。
mlcllm serve HF://mlc-ai/Llama-3-8B-Instruct-q4f161-MLC
图4. Launching LLM Server and Sending Chat Completion Requests
为了满足不同 serving 场景的需求，我们为 server 提供了三种不同的模式：“server”, “local” 和 “interactive”。
对于 “server” 模式，我们在最大限度下使用 GPU 显存，从而提高并发度，使 MLCEngine 能够同时并行处理更多的用户请求。搭配上完整的 FP16 模型和 FP8 variant 模型。“server” 模式在 A100/H100 这类服务器级 GPU 上更为我们所推荐。
相反，在 “local” 和 “interactive” 模式下 MLCEngine 会限制同时处理的请求数量、使用更少的 GPU memory，从而减少 MLCEngine 在整个系统里的压力，也与较为轻量的量化后的模型更加搭配。
Python API
在 Python 中，MLC LLM 提供了 MLCEngine 和 AsyncMLCEngine 两个选项，分别服务于同步和异步两种编程模型下的 LLM generation。这两个选项都支持与 OpenAI Python package 相同的 API。
图5. Python API of AsyncMLCEngine
iOS SDK
为了在 iPhone 和 iPad 上运行 LLM，我们将 MLCEngine 包装为了一个 Swift SDK。尽管 OpenAI 并没有官方的 Swift API，我们的 Swift SDK 模仿了 OpenAI 的 Python API，所以你可以看到下面的代码在接口上和 Python API 比较相像，都具有相同结构的输入和输出。与此同时，MLCEngine 的 Swift API 利用了 AsyncStream 从而能够以异步 streaming 的方式处理 LLM 生成出来的内容。
图6. iOS/Swift API of MLCEngine
Android SDK
对于 Android，我们也想提供同样的体验，因此我们在 Kotlin 中提供了 MLCEngine 的 Android SDK（一样具有和 OpenAI Python API 相同的接口）。下面这张图展示了 chat completion 的 Kotlin API 和在 Samsung S23 上运行 MLC Chat app 的例子。
图7. Android/Kotlin API of MLCEngine
WebLLM SDK
WebGPU 的整个生态系统在过去几年内逐渐成熟。WebGPU 的工作模式是在运行时将 WSGL (WebGPU Shading Language) shader 翻译为原生的 GPU shader，这一特性让我们在浏览器环境内做大模型的 GPU 计算变得可能。在浏览器环境内通过 WebGPU 做 GPU 计算本身具有高效（原生 GPU 加速）、便捷（无需额外环境设置）、高隐私性（100% 本地浏览器内部的计算）的优点。
图8. JavaScript/TypeScript API of MLCEngine
利用 WebGPU，我们的 MLCEngine 可以通过 JavaScript 在浏览器里运行，这个可以跳转到我们 WebLLM project，有直接在浏览器里跑大模型的 demo 和详细说明。
https://github.com/mlc-ai/web-llm
webllm.mlc.ai
一些讨论
在支持这些平台和环境的过程中，为了让各个平台上的开发都有一致的体验，我们特别注意让所有的接口都与 OpenAI API engine.chat.completions.create 保持对齐。
此外，对于不同的平台上的应用编程语言，我们也尽可能地利用它们各自的语言特性来设计结构化的输入和输出，同时利用 async streaming 从而使得 MLCEngine 在后端的运行不会阻塞应用的前端 UI。
高效的结构化输出生成
我们想要在支持基本的 chat 功能外更进一步。这段时间来，LLM 的结构化输出生成极大扩展了这些大模型的能力，使我们对于大模型的使用能够跳出基本的 chat 和 plain text 生成的范围。基于可控可自定义的结构化输出生成，LLM 能够作为标准工具被投入和整合到生活生产中。在所有的 structured format 里，JSON 是在许多场景下都最被广泛使用的标准格式。所以，让 LLM 引擎具备生成 JSON strings 的能力具有非凡意义。
MLCEngine 提供了 state-of-the-art 的 JSON 结构化输出生成模式。对于每一个收到的 request，MLCEngine 都会构造和维护一个高性能的 GrammarBNF 状态机（一种基于巴科斯-诺尔范式，Backus-Naur Form，简称 BNF 的语法规则。BNF是一种用于描述编程语言或数据格式语法的正式语法，它使用产生式规则来定义语言的合法结构。在这里，GrammarBNF可能是用来定义和限制输出内容格式的一组规则）。这个状态机在 LLM auto-regressive 的生成过程中限制输出的格式，从而确保生成的内容一定符合所要求的 JSON 格式。。
MLCEngine 支持两种 JSON 模式。第一种模式支持普通的 JSON 输出，保证 LLM 的输出内容符合 JSON 语法。我们可以通过在调用 chat completion 时提供参数 responseformat={"type": "jsonobject"} 来启用这一 JSON 模式：
图9. JSON mode in MLCEngine
更为高级的 JSON 模式是为用户请求自定义输出的 JSON schema。在看到一个请求的 JSON schema 被指定时，MLCEngine 会生成严格符合这个 schema 的输出：
图10. JSON mode with specified JSON schema in MLCEngine
更多平台上的 LLM 支持
我们团队的工作目标之一一直是让大家在各个平台都能够使用 LLM。在前面我们已经分享了 MLCEngine 在服务器 GPU、Android 和 iOS 上的使用。除此之外，我们验证了 MLCEngine 在更多的硬件和平台上都能够部署大模型， 包括但不限于：
NVIDIA RTX 4090,
NVIDIA Jetson Orin,
NVIDIA T4（in Google Colab),
AMD 7900 XTX,
Steam Deck,
Orange Pi.
这个列表里比较有趣的一个平台是 Steam Deck，因为 Steam Deck 只有有限的 GPU 驱动支持。为了把 LLM 带到 Steam Deck 上，我们选择生成模型的 Vulkan 代码。下面这张照片里，我们在 Steam Deck 上运行了 4-bit 量化的 Llama3 8B 模型。
图11. MLCEngine running Llama 3 8B on SteamDeck
我们在支持各个平台的过程中深切感受到了机器学习编译这项技术所带来的便利。相比起在各个平台上都投入大量的精力，机器学习编译让我们能够对各个平台都复用同一套模型优化 pipeline，再分别做各自的代码生成。这极大地减少了我们的工作量，同时让我们能够广泛支持这么多的硬件，让更多人能够使用这些大模型。
MLCEngine 性能优化
尽管这是 MLCEngine 的第一个 release，我们在性能优化上已经投入了不少努力。整个 MLCEngine 包含了很多关键的系统优化。比如 continuous batching, speculative decoding, paged KV management, common prefix caching, cascade inference 等。在 CUDA 上，我们使用 FlashInfer 这一高性能的 attention 计算库，并通过机器学习编译的代码生成将 FlashInfer 的技术拓展到更多的平台上，从而获得各个平台上的 attention 计算加速。
MLCEngine 有开箱即用的 multi-GPU 多卡支持。下面这个命令会在两张 GPU 上搭起我们的 REST server:
mlcllm serve HF://mlc-ai/Qwen2-72B-Instruct-q0f16-MLC --overrides "tensorparallel_shards=2"
图12. REST Server on 2x NVIDIA RTX 4090
图13. Python API 2x AMD 7900 XTX
MLCEngine 在服务器场景下具备有竞争力的性能，尤其是对于 high-throughput, low-latency 的场景，MLCEngine 的性能能够 scale 到多卡部署上。此外，我们在其它比如 Apple GPU 的平台上保有 state-of-the-art 的性能，还能在诸如 WebGPU 这样逐渐发展的平台提供高性能的 LLM 推理。
值得一提的是，我们发现搭建跨平台 LLM 推理解决方案有很多协同作用。比如，我们能够较为轻松地把比如 attention management, prefix caching, speculative execution 和结构化输出生成这样的技术直接部署在所有这些平台上。相比之下，我们项目过去的解决方案由于工程资源有限，没有支持这些优化。同时，编译器的代码生成让我们能够在不同的平台上应用同样的优化，WebGPU 是其中的一个例子。我们在 WebGPU 上的性能很大程度上来源于我们在本地 GPU runtime 上的优化效果，而这些优化能够通过机器学习编译被直接带到 WebGPU 环境里。
小结
在这篇文章里我们介绍了 MLCEngine，这是我们打造的一个统一的、高效的、同时兼容从云服务器到手机等 edge device 在内各个硬件和平台的 LLM 引擎。
MLCEngine 标志着 MLC LLM 这一项目翻开了崭新一页。在开发 MLCEngine 的过程中，我们非常享受同时也感谢和社区小伙伴一起工作，大家共同为 MLCEngine 贡献包括模型 coverage、系统优化和更多机器学习编译技术在内的很多内容，让 MLCEngine 能够具有更优秀、更 universal 的 LLM 部署。最后我们想感谢来自于 CMU Catalyst group, OctoAI, UW SAMPL group, SJTU 的合作者和开源社区里的众多贡献者。今后我们也会继续和整个开源社区一起，将开源模型带到更多人手中。
欢迎大家尝试 MLCEngine！
快速上手：https://llm.mlc.ai/docs/get_started/quick_start.html#quick-start
文档：https://llm.mlc.ai/docs/index.html
Github repo: https://github.com/mlc-ai/mlc-llm
相关文章
- 查看往期，【阅读原文】 -
