enginex-biren-vllm
基于 vLLM 推理框架的壁仞 BR166M 算力卡适配项目。通过 vLLM 的 OOT(Out-of-Tree)平台插件机制,将壁仞 SUPA 设备无缝接入 vLLM 生态,实现大语言模型在壁仞硬件上的高性能推理服务。
项目结构
enginex-biren-vllm/
├── vllm/ # 标准 vLLM 代码(含编译后的二进制)
└── vllm_br/ # 壁仞适配层(核心代码)
├── __init__.py # 插件入口,注册 SUPAPlatform
├── platform.py # SUPA 平台抽象实现
├── envs.py # 壁仞专用环境变量定义
├── utils.py # 工具函数与 monkey-patch
├── forward_context.py # 前向推理上下文扩展
├── attention/ # Attention 层 patch
├── compilation/ # SUPA Graph 图优化与编译
│ ├── supa_graph.py # SUPAGraphWrapper(类似 CUDA Graph)
│ └── supa_piecewise_backend.py # 分段编译后端
├── config/ # 编译配置扩展
├── distributed/ # 分布式通信(SCCL)
│ └── communicator.py # SUPACommunicator
├── executor/ # Ray 执行器 patch
├── model_executor/
│ ├── layers/ # 自定义算子层
│ │ ├── linear.py # NUMA 权重线性层
│ │ ├── activation.py # 激活函数
│ │ ├── layernorm.py # 层归一化
│ │ ├── rotary_embedding.py # 旋转位置编码
│ │ ├── logits_processor.py # Logits 处理
│ │ └── vocab_parallel_embedding.py # 词<><E8AF8D><EFBFBD>并行嵌入
│ └── models/ # 壁仞优化模型实现
│ ├── deepseek_v2.py # DeepSeek V2/V3
│ ├── deepseek_mtp.py # DeepSeek MTP(多 Token 预测)
│ ├── qwen2.py / qwen3.py / qwen3_moe.py # Qwen 系列
│ ├── qwen2_vl.py / qwen2_5_vl.py / qwen3_vl.py # Qwen 视觉语言模型
│ ├── chatglm.py / glm4.py / glm4_moe.py # GLM 系列
│ ├── llama.py # LLaMA
│ ├── internlm2.py # InternLM2
│ ├── intern_vit.py # InternViT
│ ├── roberta.py / bert.py # BERT 系列
│ └── registry.py # 模型注册
├── sample/ # 采样操作
├── v0/ # vLLM V0 引擎适配
│ ├── attention/backends/ # V0 Flash Attention 后端
│ └── worker/ # V0 SUPAWorker
└── v1/ # vLLM V1 引擎适配
├── attention/backends/ # V1 Attention 后端
│ ├── attention_v1.py # SUPAFlashAttentionBackend
│ └── mla/ # MLA 注意力(DeepSeek 专用)
│ ├── flashmla.py # SupaFlashMLABackend
│ └── flashmla_sparse.py # SupaFlashMLASparseBackend
├── core/ # V1 调度器扩展
├── engine/ # V1 引擎 patch
├── executor/ # V1 执行器 patch
├── sample/ # V1 采样
├── spec_decode/ # 投机解码(EAGLE)
└── worker/ # V1 SUPAWorker & ModelRunner
核心架构
插件注册机制
项目通过 vLLM 的 OOT 平台插件系统集成:
vllm_br/__init__.py::register()
→ 返回 "vllm_br.platform.SUPAPlatform"
→ vLLM 加载 SUPAPlatform 作为当前平台
→ 自动选择 SUPAWorker、Attention 后端、通信器等组件
SUPAPlatform 继承自 vllm.platforms.interface.Platform,核心配置:
- 设备名称:
supa(映射到 PyTorch 的PrivateUse1dispatch key) - 分布式后端:
sccl(SUPA Collective Communications Library) - KV Cache Block Size:固定 128
- 支持数据类型:
bfloat16、float32
双引擎支持
同时适配 vLLM V0(旧版)和 V1(新版)两套引擎架构:
| 组件 | V0 | V1 |
|---|---|---|
| Worker | vllm_br.v0.worker.worker.SUPAWorker |
vllm_br.v1.worker.worker.SUPAWorker |
| Attention | SUPAFlashAttentionBackend (V0) |
SUPAFlashAttentionBackend (V1) |
| MLA | 不支持 | SupaFlashMLABackend / SupaFlashMLASparseBackend |
| 投机解码 | 不支持 | 支持(EAGLE) |
推荐使用 V1 引擎(设置 VLLM_USE_V1=1)。
图优化(SUPA Graph)
类似 NVIDIA CUDA Graph 的机制,通过 SUPAGraphWrapper 实现:
- 捕获(Capture)推理计算图并重放(Replay),减少 kernel launch 开销
- 支持多种模式:
FULL、FULL_DECODE_ONLY、PIECEWISE、FULL_AND_PIECEWISE - 图池管理,优化显存复用
分布式通信
SUPACommunicator 基于 SCCL 库实现:
- 自动检测 P2P 直连拓扑,决定 fused allreduce 策略
- 支持 FP32 allreduce(提升数值稳定性)
- CPU allreduce 回退方案
- 张量并行(TP)与数据并行(DP)支持
内存管理
- NUMA 感知的权重分配(默认
VLLM_BR_WEIGHT_TYPE=NUMA) - 禁用不必要的 memset 操作以提升性能
- 自定义内存 profiling(
SUPAMemorySnapshot)
支持的模型
| 模型系列 | 具体模型 |
|---|---|
| DeepSeek | DeepSeek V2/V3、DeepSeek MTP |
| Qwen | Qwen2、Qwen3、Qwen3-MoE |
| Qwen-VL | Qwen2-VL、Qwen2.5-VL、Qwen3-VL、Qwen3-VL-MoE |
| GLM | ChatGLM、GLM-4、GLM-4-MoE、GLM-4-1V |
| LLaMA | LLaMA 系列 |
| InternLM | InternLM2、InternViT |
| BERT | RoBERTa、BERT |
环境变量
| 变量名 | 默认值 | 说明 |
|---|---|---|
VLLM_USE_V1 |
- | 设为 1 启用 V1 引擎(推荐) |
SUPA_VISIBLE_DEVICES |
- | 可见的 SUPA 设备列表 |
VLLM_BR_WEIGHT_TYPE |
NUMA |
权重存储类型 |
VLLM_BR_QUANT_METHOD |
INT8 |
量化方法 |
VLLM_BR_USE_FUSED_ALLREDUCE |
自动检测 | Fused AllReduce 设备数 |
VLLM_BR_USE_FP32_ALL_REDUCE |
0 |
启用 FP32 AllReduce |
VLLM_BR_USE_CPU_ALL_REDUCE |
0 |
启用 CPU AllReduce 回退 |
VLLM_BR_DEVICE_SPC_NUM |
自动检测 | 每设备计算单元数 |
VLLM_BR_STATIC_MOE_DECODER_MAX_LEN |
256 |
MoE 静态解码最大长度 |
VLLM_BR_ENABLE_TP_GROUPS_IN_SUPERNODE |
0 |
超节点内张量并行分组 |
VLLM_SCCL_SO_PATH |
/usr/local/birensupa/base/latest/succl/lib/x86_64-linux-gnu/libsuccl.so |
SCCL 库路径 |
依赖
torch_br— 壁仞 PyTorch 扩展(SUPA 设备支持)pybrml— 壁仞设备管理库(设备信息查询、P2P 拓扑检测)SCCL— SUPA 集合通信库- vLLM 及其标准依赖(torch、transformers 等)
代码模块关系
┌─────────────────────────────────────────────────────┐
│ vLLM 核心框架 │
│ (Engine, Scheduler, TokenizerPool, API Server...) │
└──────────────────────┬──────────────────────────────┘
│ OOT 插件接口
┌──────────────────────▼──────────────────────────────┐
│ vllm_br (壁仞适配层) │
│ │
│ ┌──────────┐ ┌────────────┐ ┌───────────────┐ │
│ │ Platform │──▶│ Worker │──▶│ ModelRunner │ │
│ │ (SUPA) │ │ (V0 / V1) │ │ │ │
│ └────┬─────┘ └─────┬──────┘ └───────┬───────┘ │
│ │ │ │ │
│ ┌────▼─────┐ ┌─────▼──────┐ ┌──────▼────────┐ │
│ │ Attention│ │Distributed │ │Model Executor │ │
│ │ Backend │ │(SCCL Comm) │ │(Layers+Models)│ │
│ └────┬─────┘ └─────┬──────┘ └──────┬────────┘ │
│ │ │ │ │
│ ┌────▼─────┐ ┌─────▼──────┐ ┌──────▼────────┐ │
│ │SUPA Graph│ │ AllReduce │ │ NUMA Weights │ │
│ │(图优化) │ │ (Fused/CPU)│ │ (线性层/MoE) │ │
│ └──────────┘ └────────────┘ └───────────────┘ │
└──────────────────────────────────────────────────────┘
│
┌────────────▼────────────┐
│ 硬件抽象层 │
│ torch_br / pybrml / SCCL│
└────────────┬────────────┘
│
┌────────────▼────────────┐
│ 壁仞 BR166M 算力卡 │
└─────────────────────────┘
数据流:
- vLLM 通过插件机制加载
SUPAPlatform - Platform 根据配置选择对应的 Worker(V0/V1)和 Attention 后端
- Worker 创建 ModelRunner,加载模型权重到 SUPA 设备(NUMA 模式)
- 推理时通过 SUPA Graph 捕获计算图,后续请求直接重放
- 多卡场景下通过 SCCL 进行张量并行通信
License
Apache License 2.0 — Copyright (c) 2020-2025 Shanghai Biren Technology Co., Ltd.
Description
Languages
Python
100%