MLC-LLM: 具有ML编译的通用LLM部署引擎 原创

介绍

我们正处在大型语言模型和生成式人工智能的时代，其应用场景有可能改变每个人的生活。开放式大型语言模型为提供定制化和领域特定部署带来了重大机会。

我们正处在一个激动人心的开放模型开发年份。一方面，我们看到了在（云）服务器部署方面取得的令人振奋的进展，这些解决方案可以为具有多个GPU的更大模型提供同时服务的能力。同时，我们也开始看到在设备本地部署方面取得了进展，能够将性能强大的量化模型部署到笔记本电脑、浏览器和手机上。未来将会走向何方？我们相信未来是混合的，因此重要的是要使任何人都能在云端和本地环境中运行LLM。

许多LLM推断项目，包括我们过去版本的MLC LLM工作，为服务器和本地使用案例提供了不同的解决方案，具有不同的实现和优化。例如，服务器解决方案通常支持连续批处理和更好的多GPU支持，而本地解决方案则带来了跨平台更好的可移植性。然而，我们认为有必要将所有技术汇集到一起。在一个使用案例的一侧出现的许多技术直接适用于另一侧。虽然像连续批处理这样的技术在当前某些本地使用案例中可能并不实用，但随着LLM成为操作系统的关键组件并支持多个请求以启用代理任务，它们将变得有价值。我们想提出一个问题：是否可能构建一个单一统一的LLM引擎，可在服务器和本地使用案例中运行？

在这篇文章中，我们介绍MLC LLM引擎（简称MLCEngine），这是一个通用的LLM部署引擎。MLCEngine引入了一个单一引擎，用于在服务器上进行高吞吐量、低延迟的服务，同时将小型而性能强大的模型无缝集成到多样化的本地环境中。

图1. MLCEngine：通用LLM部署引擎 实现通用部署的路径具有独特的挑战。首先，它需要支持广泛的GPU编程模型和运行时，以在每个支持的平台上实现加速。这样的过程通常需要大量重复的工程工作。我们需要在可能的情况下利用有效的供应商库，但也需要支持缺乏标准库支持的新兴平台，如Vulkan或WebGPU。此外，每个应用平台都有不同的编程语言环境，这增加了将引擎引入到Swift、Kotlin、Javascript等环境中的复杂性。

为了克服平台支持的挑战，我们利用了通过Apache TVM进行机器学习编译的技术，自动生成便携式GPU库，以适用于各种硬件和平台。此外，我们构建了一个便携式运行时架构，将业界领先的LLM服务优化（如连续批处理、推测解码、前缀缓存等）与对云端和本地平台的最大可移植性相结合。

图2. MLC LLM工作流程 编译的模型库与一个通用运行时共同工作，该运行时在不同的主机应用语言中提供直接的语言绑定，这些语言共享相同的OpenAI风格的聊天完成API。所有本地应用环境共享与云服务器设置相同的核心引擎，但针对每个个体设置进行了专门配置。

通过Chat开始

开始使用MLCEngine的最简单方式是通过聊天CLI界面。我们为此提供了一个命令行聊天界面。下面是一个示例命令，它在终端中启动聊天CLI，运行一个4位量化的Llama3 8B模型。

mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC

图3. Chat CLI 在幕后，我们执行即时模型编译，为每个平台生成GPU代码。因此，它可以跨GPU类型和操作系统工作。

我们支持广泛的模型系列，包括Llama3、Mistral/Mixtral、Phi3、Qwen2、Gemma等等。

通过OpenAI风格API使用MLCEngine

随着我们引入通用LLM引擎，设计一组开发人员熟悉且易于使用的API是非常重要的。我们选择在所有环境中采用OpenAI风格API。本节将逐步介绍这些API。

• 云端的REST API服务器

LLM服务是最典型的LLM应用之一。LLM服务器通过运行内部LLM引擎来处理从网络接收到的所有请求。我们提供了一个具有完整OpenAI API的REST服务器，用于处理LLM生成请求。下面是一个示例命令，它在本地主机上启动REST服务器，以提供4位量化的Llama3 8B模型。

图4. 启动LLM服务器并发送聊天完成请求 为了适应不同的服务场景，我们为服务器提供了三种模式：“服务器”，“本地”和“交互式”。 “服务器”模式最大化了GPU内存的使用，并支持更高的并发性，而“本地”和“交互式”模式则限制了并发请求的数量，并且使用更少的GPU内存。我们建议在专用服务器GPU（如A100/H100）上使用“服务器”模式，利用完整的FP16模型和FP8变体。对于本地环境，建议使用“本地”模式和4位量化模型设置。

• Python API

MLC LLM在Python中提供MLCEngine和AsyncMLCEngine，分别用于同步和异步LLM生成。这两种引擎支持与OpenAI Python软件包相同的API。

图5. AsyncMLCEngine的Python API

• iOS SDK

为了使LLM在iPhone和iPad上运行，我们在Swift SDK中公开了MLCEngine。虽然没有来自OpenAI的官方Swift API，但我们紧密模拟了Python API，因此我们可以找到与Python代码镜像的代码，具有结构化的输入和输出。Swift API还有效地利用了AsyncStream，以启用生成内容的异步流式传输。

图6. MLCEngine的iOS/Swift API

• Android SDK

除了iOS，我们还希望为Android提供相同的体验。我们选择使用基于OpenAI API的Kotlin API。下面是一个在三星S23上运行4位量化Phi-3模型的聊天完成Kotlin API和MLC Chat应用程序截图的示例。

图7. MLCEngine的Android/Kotlin API

• WebLLM SDK

近年来，我们看到了WebGPU生态系统的不断成熟。WebGPU通过将WGSL（WebGPU着色语言）着色器转换为本地GPU着色器来工作。这使得在浏览器中执行GPU计算成为可能，这是快速的（本地GPU加速），方便的（无需任何环境设置）和私密的（100％客户端计算）。MLCEngine也可以通过WebLLM项目在JavaScript中访问。

图8. MLCEngine的JavaScript/TypeScript API

• 讨论

我们可能意识到的一件事是，这些API遵循完全相同的风格engine.chat.completions.create，因此我们可以在这些平台上拥有相同的开发体验。我们尽可能多地利用语言本地功能来创建结构化的输入和输出，并启用异步流式传输，以便引擎活动不会阻塞UI。

高效的结构化生成

除了构建聊天应用程序外，我们还将使应用程序超越聊天。结构化生成大大提高了LLM的能力，超越了基本的聊天或纯文本生成。通过可控的结构化生成，LLM能够成为标准工具，并且可以更好地集成到生产中的其他应用程序中。在各种场景中，JSON是最常用的结构化格式之一，并且对于LLM引擎支持JSON字符串的结构化生成具有重要意义。

MLCEngine提供了最先进的JSON模式结构化生成。对于每个请求，MLCEngine运行和维护一个高性能的GrammarBNF状态机，该状态机在自回归生成期间约束响应格式。

MLCEngine支持两种级别的JSON模式：通用JSON响应和JSON模式定制。通用JSON模式将响应约束为符合JSON语法。要使用通用JSON模式，请将参数response_format={"type": "json_object"}传递给聊天完成。以下是使用JSON模式的请求示例：

图9. MLCEngine中的JSON模式 

此外，MLCEngine允许为每个单独的请求定制响应的JSON模式。当提供了JSON模式时，MLCEngine将生成严格符合该模式的响应。以下是带有自定义JSON模式的请求示例：

图10. MLCEngine中指定JSON模式的JSON模式 重要的是，结构化生成支持内置于引擎中，这意味着它可以在MLCEngine支持的所有API平台上使用。

使LLM在各种平台上可访问

我们的使命之一是使LLM在各种平台上可访问。我们已经讨论了服务器GPU、Android和iOS上的用例。我们验证了MLCEngine可以在各种平台上运行。包括但不限于：

• NVIDIA RTX 4090，
• NVIDIA Jetson Orin，
• NVIDIA T4（在Google Colab中），
• AMD 7900 XTX，
• Steam Deck，
• Orange Pi。

其中，Steam Deck是一个有趣的平台，因为它只提供有限的GPU驱动程序支持。我们生成了Vulkan代码，在其上运行了一个Llama3-8B模型。

图11. MLCEngine在Steam Deck上运行Llama 3 8B

我们发现机器学习编译帮助我们减少了在构建MLCEngine上的工程工作量，并帮助我们在更广泛的硬件范围内提供高性能解决方案，使这些模型对更多人可用。

优化引擎性能

虽然这是MLCEngine的第一个发布里程碑，但我们已经付出了大量努力来优化其性能。引擎结合了几个关键的系统改进，如连续批处理、推测解码、分页KV管理、常见前缀缓存和级联推断。我们利用FlashInfer库在CUDA平台上计算快速的注意力，并通过编译器代码生成扩展其支持到更广泛的平台。

MLCEngine支持多GPU开箱即用。以下命令在两个GPU上启动REST服务器：

mlc_llm serve HF://mlc-ai/Qwen2-72B-Instruct-q0f16-MLC --overrides "tensor_parallel_shards=2"

图12. 2个NVIDIA RTX 4090上的REST服务器

图13. Python API 2x AMD 7900 XTX

MLCEngine在服务器使用案例中可以实现高竞争性的性能，特别是在高吞吐低延迟设置下，对多个GPU的强大扩展。我们还在维护着在苹果GPU等平台上的最新性能，并支持新兴平台，如WebGPU。

重要的是，我们发现在各平台上构建统一的解决方案存在很多协同作用。例如，拥有相同的引擎使我们能够将高级优化（如注意力管理、前缀缓存、推测执行和结构化生成）带到这些平台上，在过去的解决方案中由于有限的工程带宽而缺乏支持。我们还可以通过利用代码生成在平台之间转移优化知识。我们在WebGPU上的性能受益于本地GPU运行时的优化，然后直接编译到目标WebGPU环境。

总结

我们介绍了MLCEngine，一个统一的高效LLM引擎，可以在不同的硬件和平台上运行LLM，从云服务器到边缘设备。这标志着MLC LLM项目的新篇章。我们乐于与社区合作，带来更多改进，包括提供更好的模型覆盖、更多的系统优化以及先进的机器学习编译，以实现更加高效的通用部署。

本项目得以实现，多亏了开源社区的许多贡献，包括CMU Catalyst、OctoAI、UW SAMPL、SJTU以及更广泛的MLC社区的贡献。我们计划继续与开源社区合作，为所有人带来开放的基础模型。

译自（有删改）：https://blog.mlc.ai/2024/06/07/universal-LLM-deployment-engine-with-ML-compilation

本文转载自公众号AIGC最前线   

原文链接：​​https://mp.weixin.qq.com/s/GAVn6vm_rWUBOACnm8_EHA​​