feature:add head size detect and patch some ops

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

README.md

@@ -1,4 +1,6 @@
-# vLLM Tokenizer 自动修复方案
+# vLLM 自动修复方案
+
+Tokenizer修复 + Head Size补丁 + Ixformer Ops补丁
 
 ## 1. 背景
 
@@ -16,6 +18,10 @@ ValueError: Tokenizer class TokenizersBackend does not exist or is not currently
 - 开启 trust_remote_code=True 时，transformers 会强制加载该 class
 - vLLM 无法通过参数 override tokenizer class
 
+另外，某些模型的 head_size 可能不在 vLLM 默认支持的列表中（64, 80, 96, 112, 120, 128, 192, 256），导致运行时错误。
+
+此外，某些模型需要特定的 ixformer 操作函数（如 `gelu_tanh_and_mul`），但这些函数可能不在默认的 ixformer 库中。
+
 ---
 
 ## 2. 方案目标
@@ -27,6 +33,7 @@ ValueError: Tokenizer class TokenizersBackend does not exist or is not currently
 无需修改模型文件
 无需修改启动命令
 自动修复 tokenizer 并启动 vLLM
+自动检测并 patch 不支持的 head_size
 
 ```
 
@@ -46,6 +53,12 @@ entrypoint.sh
 ↓
 修复 tokenizer_config.json
 ↓
+检测模型 head_size
+↓
+如果 head_size 不在支持列表中，patch vLLM 代码
+↓
+Patch ixformer ops（复制自定义函数并修改 __init__.py）
+↓
 vllm serve --tokenizer /tmp/fixed_tokenizer
 
 ````
@@ -173,12 +186,273 @@ bos_token / eos_token / pad_token 一致
 
 ---
 
-## 9. 总结
+## 9. Head Size 自动补丁
 
-本方案通过在容器启动阶段引入 tokenizer 修复逻辑，实现：
+### 问题
+
+vLLM 默认只支持以下 head_size 值：
 
 ```
-“模型不动，运行时自适应兼容”
+[64, 80, 96, 112, 120, 128, 192, 256]
+```
+
+如果模型的 head_size 不在此列表中，会导致运行时错误。
+
+### 检测逻辑
+
+系统会自动从模型的 `config.json` 中检测 head_size，支持以下多种配置格式：
+
+1. **直接读取 `head_dim` 字段**
+   ```json
+   {
+     "head_dim": 128
+   }
+   ```
+
+2. **从 `hidden_size / num_attention_heads` 计算**
+   ```json
+   {
+     "hidden_size": 2048,
+     "num_attention_heads": 16
+   }
+   // head_size = 2048 / 16 = 128
+   ```
+
+3. **从 `n_embd / n_head` 计算（GPTJ等模型）**
+   ```json
+   {
+     "n_embd": 2048,
+     "n_head": 16
+   }
+   ```
+
+4. **直接读取 `d_kv` 字段（T5等模型）**
+   ```json
+   {
+     "d_kv": 128
+   }
+   ```
+
+### 补丁逻辑
+
+如果检测到的 head_size 不在支持列表中，系统会：
+
+1. **备份原文件**
+   ```
+   /usr/local/corex/lib64/python3/dist-packages/vllam/attention/ops/paged_attn.py.backup
+   ```
+
+2. **修改 get_supported_head_sizes 方法**
+   ```python
+   @staticmethod
+   def get_supported_head_sizes() -> List[int]:
+       return [64, 80, 96, 112, 120, 128, 192, 256, YOUR_NEW_SIZE]
+   ```
+
+3. **保持列表排序**
+   新的 head_size 会被插入到正确的位置，保持列表升序排列。
+
+### 日志示例
+
+**无需补丁的情况**
+```
+[entrypoint] checking model head_size...
+[detect_head_size] Found head_dim in config: 128
+[detect_head_size] Model head_size: 128
+[detect_head_size] head_size 128 is already supported by vLLM, skipping patch
+```
+
+**需要补丁的情况**
+```
+[entrypoint] checking model head_size...
+[detect_head_size] Calculated from hidden_size(4096) / num_attention_heads(32) = 128
+[detect_head_size] Model head_size: 128
+[detect_head_size] head_size 160 is NOT in default supported list: [64, 80, 96, 112, 120, 128, 192, 256]
+[detect_head_size] Attempting to patch vLLM...
+[patch] Backed up original file to /usr/local/.../paged_attn.py.backup
+[patch] Successfully added head_size 160 to supported list: [64, 80, 96, 112, 120, 128, 160, 192, 256]
+[detect_head_size] Successfully patched vLLM to support head_size 160
+```
+
+### 补丁恢复
+
+如需恢复原始文件：
+```bash
+cp /usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py.backup \
+   /usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py
+```
+
+### 容错机制
+
+**重要**：head_size 检测和 patch 功能具有完整的容错机制：
+
+- ✅ **检测失败不影响启动**：如果无法检测 head_size，vLLM 仍会正常启动
+- ✅ **Patch 失败不影响启动**：如果 patch 过程出错，vLLM 仍会正常启动
+- ✅ **代码异常不影响启动**：如果检测脚本本身出现异常，vLLM 仍会正常启动
+
+这确保了：
+```
+本来能跑的模型 → 即使 head_size 检测失败 → 仍然能跑
+```
+
+### 日志示例
+
+**检测失败的情况（仍会启动vLLM）**
+```
+[entrypoint] checking model head_size...
+[detect_head_size] Error during head_size detection/patch: config.json not found
+[detect_head_size] Continuing with vLLM startup anyway...
+[entrypoint] head_size check failed, but continuing with vLLM startup
+[entrypoint] starting vLLM...
+```
+
+---
+
+## 10. Ixformer Ops 自动补丁
+
+### 问题
+
+某些模型需要特定的 ixformer 操作函数，如 `gelu_tanh_and_mul`，但这些函数可能不在默认的 ixformer 库中。
+
+### 通用解决方案
+
+**重要特性**：系统采用**通用扫描机制**，无需每次修改代码：
+
+1. **自动扫描**：自动扫描 `patched_ops` 目录中的所有 `.py` 文件
+2. **批量处理**：批量复制所有文件到目标目录
+3. **自动生成import**：为每个文件自动生成对应的 `from .xxx import *` 语句
+4. **智能去重**：自动检测已存在的import，避免重复添加
+
+### 使用方法
+
+只需将需要补丁的 ops 文件放入 `patched_ops` 目录即可：
+
+```bash
+patched_ops/
+├── gelu_tanh_and_mul.py      # 第一个ops文件
+├── another_op.py              # 第二个ops文件
+├── third_operation.py         # 第三个ops文件
+└── ...
+```
+
+系统会自动：
+- 扫描所有 `.py` 文件（排除 `__init__.py`）
+- 复制到 `/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/`
+- 在 `__init__.py` 中添加对应的 import 语句
+
+### 补丁逻辑
+
+系统会自动执行以下操作：
+
+1. **扫描源目录**
+   ```bash
+   源目录: /opt/patched_ops/
+   自动查找所有 .py 文件（排除 __init__.py）
+   ```
+
+2. **批量复制文件**
+   ```
+   源文件: /opt/patched_ops/*.py
+   目标: /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/
+   ```
+
+3. **自动生成并添加 import 语句**
+   ```
+   目标文件: /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/__init__.py
+   自动生成: from .gelu_tanh_and_mul import *
+            from .another_op import *
+            from .third_operation import *
+            ```
+
+4. **自动备份**
+   在修改前会自动备份原始的 `__init__.py` 文件为 `__init__.py.backup`
+
+### 日志示例
+
+**成功的批量补丁操作**
+```
+[entrypoint] patching ixformer ops...
+[patch_ops] Starting ixformer ops patch...
+[patch_ops] Found 3 ops file(s): gelu_tanh_and_mul.py, another_op.py, third_operation.py
+[patch_ops] Backed up /usr/local/.../__init__.py to /usr/local/.../__init__.py.backup
+[patch_ops] Copied gelu_tanh_and_mul.py to /usr/local/.../ixformer/functions/
+[patch_ops] Copied another_op.py to /usr/local/.../ixformer/functions/
+[patch_ops] Copied third_operation.py to /usr/local/.../ixformer/functions/
+[patch_ops] Added 3 import statement(s) to /usr/local/.../__init__.py
+[patch_ops] Successfully patched ixformer ops
+[patch_ops] Patch completed successfully
+```
+
+**重复运行（已存在import）**
+```
+[entrypoint] patching ixformer ops...
+[patch_ops] Starting ixformer ops patch...
+[patch_ops] Found 3 ops file(s): gelu_tanh_and_mul.py, another_op.py, third_operation.py
+[patch_ops] Backup already exists: /usr/local/.../__init__.py.backup
+[patch_ops] Copied gelu_tanh_and_mul.py to /usr/local/.../ixformer/functions/
+[patch_ops] Copied another_op.py to /usr/local/.../ixformer/functions/
+[patch_ops] Copied third_operation.py to /usr/local/.../ixformer/functions/
+[patch_ops] All imports already exist in /usr/local/.../__init__.py, skipping modification
+[patch_ops] Successfully patched ixformer ops
+```
+
+**目标目录不存在的情况**
+```
+[entrypoint] patching ixformer ops...
+[patch_ops] Starting ixformer ops patch...
+[patch_ops] Target directory not found: /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/
+[patch_ops] Patch failed, but this will not prevent vLLM from starting
+[entrypoint] ixformer ops patch failed, but continuing with vllm startup
+```
+
+### 补丁恢复
+
+如需恢复原始的 `__init__.py` 文件：
+```bash
+cp /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/__init__.py.backup \
+   /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/__init__.py
+```
+
+### 容错机制
+
+与 head_size 检测一样，ixformer ops 补丁也具有完整的容错机制：
+
+- ✅ **源目录不存在不影响启动**：如果源目录不存在，vLLM 仍会正常启动
+- ✅ **目标目录不存在不影响启动**：如果目标目录不存在，vLLM 仍会正常启动
+- ✅ **复制失败不影响启动**：如果复制过程出错，vLLM 仍会正常启动
+- ✅ **修改失败不影响启动**：如果修改 __init__.py 失败，vLLM 仍会正常启动
+- ✅ **部分失败不影响整体**：即使某个文件处理失败，其他文件仍会继续处理
+
+---
+
+## 11. 总结
+
+本方案通过在容器启动阶段引入多种自动修复和补丁逻辑，实现：
 
 ```
+"模型不动，运行时自适应兼容"
 ```
+
+主要功能：
+- ✅ 自动修复不兼容的 tokenizer 配置
+- ✅ 自动检测并补丁不支持的 head_size
+- ✅ 自动补丁 ixformer ops（如 gelu_tanh_and_mul）
+- ✅ 无需修改模型文件，无需修改启动命令
+- ✅ 完全透明，不影响正常模型部署
+- ✅ **完整的容错机制，确保本来能跑的模型不受影响**
+
+### 启动流程总结
+
+```
+容器启动
+↓
+修复 tokenizer（如需要）
+↓
+检查并 patch head_size（如需要）
+↓
+检查并 patch ixformer ops（如需要）
+↓
+启动 vLLM
+```
+
+每个步骤都有完整的容错机制，确保即使某个步骤失败，也不会影响后续步骤的执行。