【vLLM】源码解读:vllm如何识别到第三方自定义设备的
🔌 第三方平台插件识别机制
1️⃣ Python Entry Points 机制
第三方平台插件使用 Python 标准的 entry_points 机制进行注册和发现:
# 在第三方包的 setup.py 或 pyproject.toml 中
[project.entry-points."vllm.platform_plugins"]
gcu = "vllm_gcu.platform:gcu_platform_plugin"
2️⃣ GCU 平台插件示例
假设你要开发一个 GCU(燧原科技)平台插件,需要创建以下结构:
步骤 1: 创建插件包
# vllm_gcu/platform.pyfrom typing import Optionaldef gcu_platform_plugin() -> Optional[str]:"""GCU 平台检测函数"""is_gcu = Falsetry:# 检测 GCU 设备的存在import torch_gcu # GCU 的 Python 库# 检查是否有可用的 GCU 设备if hasattr(tops, 'device') and tops.device.is_available():device_count = torch.device.device_count()if device_count > 0:is_gcu = Trueprint(f"Detected {device_count} GCU device(s)")except Exception as e:print(f"GCU platform is not available: {e}")return None# 返回 GCU Platform 类的完全限定名return "vllm_gcu.gcu_platform.GCUPlatform" if is_gcu else None
步骤 2: 实现 Platform 类
# vllm_gcu/gcu_platform.pyfrom vllm.platforms.interface import Platform, PlatformEnumclass GCUPlatform(Platform):_enum = PlatformEnum.OOT # Out-Of-Treedevice_name = "gcu"device_type = "gcu"dispatch_key = "CPU" # 或者 GCU 专用的 dispatch key@classmethoddef get_device_name(cls, device_id: int = 0) -> str:import topsreturn tops.device.get_device_name(device_id)@classmethoddef get_device_capability(cls, device_id: int = 0):# 实现设备能力查询pass# 实现其他必要的方法...
步骤 3: 注册插件
# setup.py
from setuptools import setupsetup(name='vllm-gcu',version='0.1.0',packages=['vllm_gcu'],entry_points={'vllm.platform_plugins': ['gcu = vllm_gcu.platform:gcu_platform_plugin']}
)
或者使用 pyproject.toml
:
[project.entry-points."vllm.platform_plugins"]
gcu = "vllm_gcu.platform:gcu_platform_plugin"
3️⃣ 插件发现和加载流程
启动 vLLM↓
resolve_current_platform_cls_qualname()↓
load_plugins_by_group('vllm.platform_plugins') ← 181行↓
使用 importlib.metadata.entry_points() 扫描所有已安装包↓
发现所有注册到 'vllm.platform_plugins' 的插件↓
┌─────────────────────────────────────────┐
│ 发现的插件: │
│ - gcu (来自 vllm-gcu 包) │
│ - npu (来自 vllm-npu 包) │
│ - ... (其他第三方插件) │
└─────────────────────────────────────────┘↓
遍历所有插件,调用检测函数↓
gcu_platform_plugin() 返回:- "vllm_gcu.gcu_platform.GCUPlatform" (如果检测到 GCU)- None (如果没有 GCU 设备)↓
只允许一个平台插件被激活↓
动态加载并实例化平台类
4️⃣ 环境变量控制
可以通过 VLLM_PLUGINS
环境变量控制加载哪些插件:
# 加载所有插件(默认)
python your_script.py# 只加载特定插件
VLLM_PLUGINS=gcu python your_script.py# 不加载任何插件
VLLM_PLUGINS="" python your_script.py# 加载多个插件(用逗号分隔)
VLLM_PLUGINS=gcu,custom_plugin python your_script.py
5️⃣ 关键代码分析
在 vllm/plugins/__init__.py
的 load_plugins_by_group
函数:
def load_plugins_by_group(group: str) -> dict[str, Callable[[], Any]]:# 1. 使用 entry_points 发现插件discovered_plugins = entry_points(group=group)# 2. 过滤插件(根据 VLLM_PLUGINS 环境变量)allowed_plugins = envs.VLLM_PLUGINSplugins = {}for plugin in discovered_plugins:if allowed_plugins is None or plugin.name in allowed_plugins:# 3. 动态加载插件函数func = plugin.load()plugins[plugin.name] = funcreturn plugins
6️⃣ 优先级规则
在 vllm/platforms/__init__.py
的 200-220 行:
# 外部插件优先级 > 内置插件
if len(activated_oot_plugins) == 1:platform_cls_qualname = platform_plugins[activated_oot_plugins[0]]()logger.info("Platform plugin %s is activated", activated_oot_plugins[0])
elif len(activated_builtin_plugins) == 1:platform_cls_qualname = builtin_platform_plugins[activated_builtin_plugins[0]]()logger.info("Automatically detected platform %s.", activated_builtin_plugins[0])
📝 总结
GCU 设备识别流程:
- ✅ 安装
vllm-gcu
包(包含 entry_points 注册) - ✅ vLLM 启动时自动扫描所有
vllm.platform_plugins
组的插件 - ✅ 调用
gcu_platform_plugin()
检测函数 - ✅ 检测函数尝试导入 GCU 库(如
tops
)并查询设备 - ✅ 如果检测成功,返回平台类的完全限定名
- ✅ vLLM 动态加载并实例化 GCUPlatform 类
- ✅ 使用 GCU 平台进行推理
这就是第三方平台插件的完整识别机制!🎯