介绍vLLM硬件插件,昇腾NPU上的最佳实践
自2024年12月以来,通过vLLM社区和vLLM上的昇腾团队的共同努力,我们完成了硬件可插拔RFC。该提案允许以解耦的方式将硬件集成到vLLM中,从而能够快速且模块化地支持不同的硬件平台。
为什么选择vLLM硬件插件?
目前,vLLM已支持多种后端。然而,随着vLLM后端数量的持续增长,出现了一些挑战
- 代码复杂性增加:每个硬件后端都有自己的
Executor
、Worker
、Runner
和Attention
组件。这增加了vLLM代码库的复杂性,非通用的后端特定代码散布在整个项目中。 - 高维护成本:维护后端的成本很高,不仅对后端开发者,也对vLLM社区而言。当后端维护者不在时,社区贡献者资源的稀缺使得有效添加新功能变得困难。
- 缺乏可扩展性:虽然vLLM通过实现
Executor
、Worker
、Runner
和Attention
遵循了结构良好的分层设计,但支持新硬件通常需要侵入性的修改或补丁,而不是动态注册。这使得添加新后端变得麻烦。
认识到需要一种灵活和模块化的方法来集成硬件后端,我们提出了硬件插件作为一个可行的解决方案
- 解耦的代码库:硬件后端插件代码保持独立,使得vLLM核心代码更清晰。
- 减轻维护负担:vLLM开发者可以专注于通用功能,而不会因后端特定实现造成的差异而感到不知所措。
- 更快集成且更独立:新后端可以快速集成,工作量更少,并且可以独立发展。
什么是vLLM硬件插件?
在介绍vLLM硬件插件之前,我们先来看两个前提RFC
- [RFC] vLLM插件系统:此RFC引入了一种基于插件的方法来支持各种定制需求,允许用户定义自定义模型、执行器、调度器等。
- [RFC] 使vLLM与设备无关以支持多种硬件和(vllm-project/vllm#6080):此RFC引入了平台子模块,该子模块集中了与硬件相关的实现,以减少主代码库中的条件逻辑,并为模块化奠定了基础。
基于这些RFC,我们提出了[RFC] 硬件可插拔,它将Platform
模块集成到vLLM中作为插件。此外,我们重构了Executor
、Worker
、ModelRunner
、AttentionBackend
和Communicator
,以更灵活地支持硬件插件。
目前,vLLM社区已成功实现了RFC中介绍的Platform
模块。通过vllm-project/vllm-ascend和vllm-project/vllm-spyre项目验证了其功能。利用此插件机制,我们成功地将vLLM与昇腾NPU和IBM Spyre后端进行了集成。
如何通过vLLM硬件插件机制集成新的后端
本节将从开发者和用户的角度深入探讨如何通过硬件插件集成新的后端。
开发者视角
要使用硬件插件将新的后端集成到vLLM中,请按照以下步骤操作
步骤1:创建新项目并初始化平台
首先为新的后端创建一个Python项目并添加一个platform.py
文件。然后,从vllm.platforms
导入Platform
类并实现所需的属性和方法。
您可以参考vLLM Ascend项目中的platform.py
作为示例。
步骤2:实现自定义Worker、Model Runner、Attention Backend和Communicator模块
根据新后端的需要,实现以下模块
from vllm.worker.worker_base import WorkerBase
from vllm.worker.model_runner_base import ModelRunnerBase
from vllm.attention.backends.abstract import AttentionBackend
from vllm.distributed.device_communicators.base_communicator import CommunicatorBase
这些类在vLLM中都有相应的基类。同样,您可以参考vLLM Ascend的实现作为示例。
步骤3:注册插件
使用Python的入口点机制在setup.py
中注册插件
setup(
entry_points={'vllm.platform_plugins': ["{your_platform_name} = {code_path}:{register_function}"]}
)
{your_platform_name}
:新后端的名称(可以是任意的)。{code_path}
:主Python模块的路径。{register_function}
:注册函数,返回在步骤1中定义的Platform
类的路径。
参考vLLM Ascend中的setup.py
获取实际示例。
用户视角
用户只需在运行前安装vllm和您的插件,以vllm-ascend为例
pip install vllm vllm-ascend
启动时,您将看到以下日志,这意味着后端插件正在正常工作
INFO 02-06 15:49:01 __init__.py:30] Available plugins for group vllm.platform_plugins:
INFO 02-06 15:49:01 __init__.py:32] name=ascend, value=vllm_ascend:register
… …
INFO 02-06 15:49:01 __init__.py:44] plugin ascend loaded.
INFO 02-06 15:49:01 __init__.py:181] Platform plugin ascend is activated
未来展望
未来,我们将继续与vLLM社区的开发者协作,以增强以下方面
- 持续增强V1 Engine和VLMs。
- 扩展对更多模块和功能的插件支持,例如调度器、图模式和自定义算子。
- 更好的用户体验和更高的性能。
- 维护和增强适用于相应硬件平台的稳定插件架构
我们鼓励大家尝试这个新功能!如果您有任何问题,请加入vLLM Slack并在#sig-extensible-hardware频道参与讨论。🚀
致谢
这个灵活的硬件后端插件机制离不开众多vLLM贡献者的努力。因此,我们深切感谢vLLM的维护者,包括Kaichao You、Simon Mo、Cyrus Leung、Robert Shaw、Michael Goin和Jie Li,他们在相关重构、深入讨论和快速评审方面做出了贡献;vLLM上的昇腾团队成员Xiyuan Wang、Shanshan Shen、Chenguang Li和Mengqing Cao,他们进行了机制设计和实现;vLLM上的Spyre团队成员Joe Runde和Yannick Schnider,他们进行了可插拔调度器的设计和实现;以及其他贡献者,包括yancong,他设计和实现了可扩展的量化方法;Aviv Keshet,他贡献了可扩展的SamplingParams
。