# 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.