# NV A100 vLLM Patched v2.0

基于 `harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0` 构建的 NVidia A100 GPU 专用 vLLM Docker 镜像，包含 4 类兼容性补丁。

## 基础镜像

- **镜像**: `harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0`
- **Python**: 3.12
- **vLLM**: 0.11.0

## 补丁概览

此镜像在构建时打了 4 类补丁，启动时再执行运行时修复：

| # | 文件 | 类型 | 说明 |
|---|------|------|------|
| 1 | `patch.py` | 构建时 | transformers `tokenization_utils_base.py` 库级补丁 |
| 2 | `patch_triton.py` | 构建时 | Triton Attention backend `validate_head_size` 绕过 |
| 3 | `detect_tokenizer.py` + `fix_tokenizer.py` | 运行时 | 自动检测并修复 tokenizer 配置 |
| 4 | `detect_head_size.py` | 运行时 | head_size 检测，自动切换 attention backend |

---

## 修复 1: transformers 库级补丁 (`patch.py`)

**问题**: 部分模型的 `tokenizer_config.json` 中 `extra_special_tokens` 字段是 list 而非 dict，导致 `transformers` 加载 tokenizer 时崩溃。

**修复位置**: `/usr/local/lib/python3.12/dist-packages/transformers/tokenization_utils_base.py`

**逻辑**: 在 `SPECIAL_TOKENS_ATTRIBUTES` 赋值前插入类型检查——若 `special_tokens` 是 list，则转为 `{t: t}` 的 dict 形式：

```python
# PATCH: some models have extra_special_tokens as list instead of dict
if isinstance(special_tokens, list):
    special_tokens = {t: t for t in special_tokens}
```

---

## 修复 2: Triton Attention head_size 验证绕过 (`patch_triton.py`)

**问题**: Triton Attention backend 的 `validate_head_size` 方法强制要求 `head_size >= 32`，某些非标准模型的 head_size 不满足此约束会直接报错。

**修复位置**: `/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/triton_attn.py`

**逻辑**: 将整个 `validate_head_size` 方法体替换为直接 `return`，完全绕过验证。Triton 本身在运行时编译，实际支持任意 head_size。

---

## 修复 3: 运行时 tokenizer 配置修复

### `detect_tokenizer.py` — 检测器

根据模型目录中的文件自动判断 tokenizer 类型：

| 特征文件 | 判定类型 |
|----------|----------|
| `tokenizer.json` | `fast` |
| `tokenizer.model` | `sentencepiece` |
| `vocab.json` + `merges.txt` | `bpe` |
| 以上均无 | `unknown` |

同时读取 `tokenizer_config.json` 中的 `tokenizer_class` 字段作为原始类名。

### `fix_tokenizer.py` — 修复器

**运行时机**: 容器启动时（entrypoint 第一步）

**逻辑**:
1. 将模型目录中的 tokenizer 相关文件拷贝到 `/tmp/fixed_tokenizer/`
2. 调用 `detect_tokenizer` 判断 tokenizer 类型
3. 检查原 `tokenizer_class` 是否有效：
   - 有效类：`transformers` 库中所有包含 "Tokenizer" 的类
   - 无效类（黑名单）：`TokenizersBackend`、`TiktokenTokenizer`
   - 缺失或无效时，按类型回退：

| tokenizer 类型 | fallback 类 |
|----------------|-------------|
| `fast` | `PreTrainedTokenizerFast` |
| `sentencepiece` | `LlamaTokenizer` |
| `bpe` | `GPT2TokenizerFast` |
| `unknown` | `PreTrainedTokenizerFast` |

4. 将修复后的配置写回 `/tmp/fixed_tokenizer/tokenizer_config.json`

**环境变量**:
- `MODEL_DIR`：模型目录，默认 `/model`
- `FIX_TOKENIZER_DIR`：修复后 tokenizer 输出目录，默认 `/tmp/fixed_tokenizer`

---

## 修复 4: head_size 检测 (`detect_head_size.py`)

**运行时机**: 容器启动时（entrypoint 第二步）

**逻辑**:
1. 读取 `${MODEL_DIR}/config.json`
2. 获取 `head_dim`；若无此字段，则用 `hidden_size / num_attention_heads` 计算
3. 检查是否在 FlashAttention 支持的 head_size 白名单内：

```
{32, 64, 96, 128, 160, 192, 224, 256}
```

4. **退出码**:
   - `0` — head_size 在白名单内，无需处理
   - `2` — head_size 不在白名单，将数值输出到 stdout，entrypoint 会据此切换 backend

**entrypoint 联动**: 当 `detect_head_size.py` 返回退出码 2 时，entrypoint 自动设置 `VLLM_USE_FLASH_ATTN_PA=0`，切换到 Triton Attention backend。

---

## 入口脚本 (`entrypoint.sh`)

启动流程：

```
1. python3 /opt/fix_tokenizer.py       → 修复 tokenizer 配置
2. python3 /opt/detect_head_size.py    → 检测 head_size 兼容性
3. 若非白名单 head_size → export VLLM_USE_FLASH_ATTN_PA=0（切到 Triton）
4. exec vllm serve ${MODEL_DIR} --tokenizer ${FIX_TOKENIZER_DIR} $@
```

启动时自动以修复后的 tokenizer 目录运行 vLLM，用户只需挂载模型到 `/model`。

---

## 构建

```bash
docker build -t nv-vllm-patched:v2.0 .
```

---

## 与 K100 v2.0 的关系

此版本是 K100 vLLM Patched v2.0 的 NV A100 移植版。与 K100 版本的主要区别：

- 基础镜像相同（`vllm-openai:v0.11.0`）
- K100 版本面向 K100 芯片集群（含 K100 专属 patch），NV A100 版本去掉了 K100 专属补丁
- tokenizer 修复和 head_size 检测逻辑保持兼容

## 参考

- 上游 K100 版本：`k100-vllm-patched-v2.0/K100-vLLM-Patched-v2.0/`
- 第四范式 modelhub 平台：`modelhub.org.cn`
- vLLM 官方文档：https://docs.vllm.ai/