diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..39a864e --- /dev/null +++ b/Dockerfile @@ -0,0 +1,31 @@ +# 基于壁仞官方 vLLM 推理镜像 +FROM git.modelhub.org.cn:9443/enginex/xc-llm-biren166m:26.01 + +# 镜像元信息 +LABEL maintainer="enginex" +LABEL description="vLLM inference engine for Biren BR166M (SUPA) GPU" +LABEL version="26.01" + +# 设置工作目录 +WORKDIR /workspace + +# 复制壁仞适配代码 +COPY vllm_br/ /workspace/vllm_br/ + +# 如果需要覆盖基础镜像中的 vllm(按需启用) +# COPY vllm/ /workspace/vllm/ + +# 确保 vllm_br 可被 Python 发现 +ENV PYTHONPATH="/workspace:${PYTHONPATH}" + +# 默认使用 V1 引擎 +ENV VLLM_USE_V1=1 + +# SUPA 设备相关默认配置 +ENV VLLM_BR_WEIGHT_TYPE=NUMA +ENV VLLM_BR_QUANT_METHOD=INT8 + +# 暴露 vLLM API 服务端口 +EXPOSE 8000 + +ENTRYPOINT ["python", "-m", "vllm.entrypoints.openai.api_server"] diff --git a/README.md b/README.md index ad14780..71b9655 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,195 @@ # 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 # 词���并行嵌入 + │ └── 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 的 `PrivateUse1` dispatch 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 算力卡 │ + └─────────────────────────┘ +``` + +数据流: +1. vLLM 通过插件机制加载 `SUPAPlatform` +2. Platform 根据配置选择对应的 Worker(V0/V1)和 Attention 后端 +3. Worker 创建 ModelRunner,加载模型权重到 SUPA 设备(NUMA 模式) +4. 推理时通过 SUPA Graph 捕获计算图,后续请求直接重放 +5. 多卡场景下通过 SCCL 进行张量并行通信 + +## License + +Apache License 2.0 — Copyright (c) 2020-2025 Shanghai Biren Technology Co., Ltd.