介绍 vLLM 硬件插件:来自昇腾(Ascend)NPU 的最佳实践
自 2024 年 12 月以来,通过 vLLM 社区与 vLLM 昇腾团队的共同努力,我们完成了 Hardware Pluggable RFC。该提案允许以解耦的方式将硬件集成到 vLLM 中,从而实现对不同硬件平台的快速且模块化支持。
为什么需要 vLLM 硬件插件?
目前,vLLM 已经支持多个后端。然而,随着 vLLM 后端数量的不断增加,出现了一些挑战:
- 代码复杂度增加:每个硬件后端都有自己的
Executor、Worker、Runner和Attention组件。这增加了 vLLM 代码库的复杂性,非通用的特定后端代码散落在整个项目中。 - 维护成本高:维护后端的成本很高,不仅对后端开发人员如此,对 vLLM 社区也是如此。社区贡献者资源稀缺,当后端维护者不在场时,很难高效地添加新功能。
- 缺乏扩展性:虽然 vLLM 通过实现
Executor、Worker、Runner和Attention遵循了结构良好的分层设计,但支持新硬件通常需要侵入性修改或打补丁,而不是动态注册。这使得添加新后端变得繁琐。
意识到需要一种灵活且模块化的方式来集成硬件后端,我们提出了硬件插件作为一个可行的解决方案:
- 解耦代码库:硬件后端插件代码保持独立,使 vLLM 核心代码更加整洁。
- 减轻维护负担:vLLM 开发人员可以专注于通用功能,而不会被特定后端实现带来的差异所困扰。
- 更快的集成与更强的独立性:新后端可以快速集成,工作量更少,并且可以独立演进。
什么是 vLLM 硬件插件?
在介绍 vLLM 硬件插件之前,让我们先看看两个前提 RFC:
- [RFC] vLLM Plugin System:该 RFC 引入了基于插件的方法来支持各种自定义需求,允许用户定义自定义模型、执行器、调度器等。
- [RFC] Make vLLM Device-Agnostic for Diverse Hardware Support 以及 (vllm-project/vllm#6080):该 RFC 引入了 platform 子模块,它集中了硬件相关的实现,以减少主代码库中的条件逻辑,并为模块化奠定了基础。
基于这些 RFC,我们提出了 [RFC] Hardware Pluggable,将 Platform 模块作为插件集成到 vLLM 中。此外,我们重构了 Executor、Worker、ModelRunner、AttentionBackend 和 Communicator,以更灵活地支持硬件插件。
目前,vLLM 社区已成功实现了 RFC 中引入的 Platform 模块。该功能已通过 vllm-project/vllm-ascend 和 vllm-project/vllm-spyre 项目得到验证。通过这种插件机制,我们成功地将 vLLM 与昇腾(Ascend)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 的 entrypoint 机制在 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 引擎和多模态大模型(VLM)。
- 扩大对更多模块和功能的插件支持,如调度器、图模式和自定义算子。
- 更好的用户体验和更高的性能。
- 维护和增强适用于相应硬件平台的稳定插件架构。
我们鼓励大家尝试这一新功能!如果你有任何问题,请加入 vLLM Slack 并参与 #sig-extensible-hardware 频道进行讨论。🚀
致谢
这种灵活的硬件后端插件机制离不开众多 vLLM 贡献者的努力。因此,我们深表感谢:感谢 vLLM 维护者,包括 尤凯超、Simon Mo、Cyrus Leung、Robert Shaw、Michael Goin 和 Jie Li 进行的相关重构、深入讨论和快速审查;感谢来自 vLLM 昇腾团队的 王希源、沈姗姗、李晨光 和 曹梦清 进行的机制设计和实现;感谢来自 vLLM Spyre 团队的 Joe Runde 和 Yannick Schnider 进行的可插拔调度器设计和实现;以及其他贡献者,包括进行可扩展量化方法设计和实现的 yancong,以及进行可扩展 SamplingParams 实现的 Aviv Keshet。