feature:add head size detect and patch some ops

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

Dockerfile

@@ -13,5 +13,8 @@ ENV PATH=/root/apps/apache-jmeter-5.6.3/bin:/root/apps/jdk1.8.0_411/bin:/usr/loc
 
 COPY fix_tokenizer.py /opt/
 COPY detect_tokenizer.py /opt/
+COPY detect_head_size.py /opt/
+COPY patch_ops.py /opt/
+COPY patched_ops /opt/patched_ops/
 COPY entrypoint.sh /opt/
 RUN chmod +x /opt/entrypoint.sh

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
+```
+
+每个步骤都有完整的容错机制，确保即使某个步骤失败，也不会影响后续步骤的执行。

detect_head_size.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+"""
+检测模型的head_size，如果不在vLLM支持列表中，则patch到vllm代码中
+"""
+import os
+import json
+import re
+import shutil
+from typing import Optional, List
+
+# vLLM默认支持的head_size列表
+DEFAULT_SUPPORTED_HEAD_SIZES = [64, 80, 96, 112, 120, 128, 192, 256]
+
+# vLLM需要patch的文件路径
+VLLM_PAGED_ATTN_PATH = "/usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py"
+
+
+def get_model_head_size(model_dir: str) -> Optional[int]:
+    """
+    从模型config.json中获取head_size
+
+    Args:
+        model_dir: 模型目录路径
+
+    Returns:
+        head_size，如果无法获取则返回None
+    """
+    config_path = os.path.join(model_dir, "config.json")
+
+    if not os.path.exists(config_path):
+        print(f"[detect_head_size] config.json not found in {model_dir}")
+        return None
+
+    try:
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+
+        # 尝试多种方式获取head_size
+        # 1. 直接从head_dim字段获取
+        if 'head_dim' in config:
+            head_size = config['head_dim']
+            print(f"[detect_head_size] Found head_dim in config: {head_size}")
+            return head_size
+
+        # 2. 从hidden_size和num_attention_heads计算
+        if 'hidden_size' in config and 'num_attention_heads' in config:
+            hidden_size = config['hidden_size']
+            num_heads = config['num_attention_heads']
+            head_size = hidden_size // num_heads
+            print(f"[detect_head_size] Calculated from hidden_size({hidden_size}) / num_attention_heads({num_heads}) = {head_size}")
+            return head_size
+
+        # 3. 对于GPTJ等模型，可能使用n_embd和n_head
+        if 'n_embd' in config and 'n_head' in config:
+            n_embd = config['n_embd']
+            n_head = config['n_head']
+            head_size = n_embd // n_head
+            print(f"[detect_head_size] Calculated from n_embd({n_embd}) / n_head({n_head}) = {head_size}")
+            return head_size
+
+        # 4. 对于T5等模型，可能使用d_kv和num_heads
+        if 'd_kv' in config:
+            head_size = config['d_kv']
+            print(f"[detect_head_size] Found d_kv in config: {head_size}")
+            return head_size
+
+        print(f"[detect_head_size] Cannot determine head_size from config.json")
+        return None
+
+    except Exception as e:
+        print(f"[detect_head_size] Error reading config.json: {e}")
+        return None
+
+
+def patch_vllm_head_sizes(new_head_size: int) -> bool:
+    """
+    将新的head_size patch到vLLM的get_supported_head_sizes方法中
+
+    Args:
+        new_head_size: 需要添加的新head_size
+
+    Returns:
+        是否成功patch
+    """
+    if not os.path.exists(VLLM_PAGED_ATTN_PATH):
+        print(f"[patch] vLLM file not found: {VLLM_PAGED_ATTN_PATH}")
+        return False
+
+    try:
+        # 备份原文件
+        backup_path = VLLM_PAGED_ATTN_PATH + ".backup"
+        if not os.path.exists(backup_path):
+            shutil.copy2(VLLM_PAGED_ATTN_PATH, backup_path)
+            print(f"[patch] Backed up original file to {backup_path}")
+
+        # 读取文件内容
+        with open(VLLM_PAGED_ATTN_PATH, 'r') as f:
+            content = f.read()
+
+        # 找到get_supported_head_sizes方法并patch
+        pattern = r'(@staticmethod\s+def get_supported_head_sizes\(\)\s*->\s*List\[int\]:\s*return\s*\[)([^\]]+)(\])'
+        match = re.search(pattern, content, re.MULTILINE | re.DOTALL)
+
+        if not match:
+            print(f"[patch] Could not find get_supported_head_sizes method in {VLLM_PAGED_ATTN_PATH}")
+            return False
+
+        prefix = match.group(1)
+        current_list = match.group(2)
+        suffix = match.group(3)
+
+        # 解析当前列表
+        current_sizes = [int(x.strip()) for x in current_list.split(',') if x.strip()]
+
+        # 添加新的head_size并排序
+        if new_head_size not in current_sizes:
+            current_sizes.append(new_head_size)
+            current_sizes.sort()
+
+            # 生成新的列表字符串
+            new_list_str = ', '.join(str(size) for size in current_sizes)
+
+            # 替换原方法
+            new_method = f"{prefix}{new_list_str}{suffix}"
+            new_content = re.sub(pattern, new_method, content, flags=re.MULTILINE | re.DOTALL)
+
+            # 写回文件
+            with open(VLLM_PAGED_ATTN_PATH, 'w') as f:
+                f.write(new_content)
+
+            print(f"[patch] Successfully added head_size {new_head_size} to supported list: {current_sizes}")
+            return True
+        else:
+            print(f"[patch] head_size {new_head_size} already in list, skipping")
+            return True
+
+    except Exception as e:
+        print(f"[patch] Error patching vLLM file: {e}")
+        return False
+
+
+def main():
+    """主函数"""
+    try:
+        model_dir = os.environ.get("MODEL_DIR", "/model")
+
+        print(f"[detect_head_size] Checking model in: {model_dir}")
+
+        # 获取head_size
+        head_size = get_model_head_size(model_dir)
+
+        if head_size is None:
+            print(f"[detect_head_size] Could not determine head_size, skipping")
+            return
+
+        print(f"[detect_head_size] Model head_size: {head_size}")
+
+        # 检查是否需要patch
+        if head_size in DEFAULT_SUPPORTED_HEAD_SIZES:
+            print(f"[detect_head_size] head_size {head_size} is already supported by vLLM, skipping patch")
+            return
+
+        print(f"[detect_head_size] head_size {head_size} is NOT in default supported list: {DEFAULT_SUPPORTED_HEAD_SIZES}")
+        print(f"[detect_head_size] Attempting to patch vLLM...")
+
+        # 执行patch
+        success = patch_vllm_head_sizes(head_size)
+
+        if success:
+            print(f"[detect_head_size] Successfully patched vLLM to support head_size {head_size}")
+        else:
+            print(f"[detect_head_size] Failed to patch vLLM for head_size {head_size}")
+
+    except Exception as e:
+        print(f"[detect_head_size] Error during head_size detection/patch: {e}")
+        print(f"[detect_head_size] Continuing with vLLM startup anyway...")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        print(f"[detect_head_size] Unexpected error: {e}")
+        print(f"[detect_head_size] This will not prevent vLLM from starting")
+        import sys
+        sys.exit(0)  # 总是成功退出，不阻止vLLM启动

entrypoint.sh

@@ -34,6 +34,14 @@ else
     TOKENIZER_ARG=""
 fi
 
+# 检查并patch head_size支持（即使失败也不影响启动）
+echo "[entrypoint] checking model head_size..."
+python3 /opt/detect_head_size.py || echo "[entrypoint] head_size check failed, but continuing with vllm startup"
+
+# Patch ixformer ops（即使失败也不影响启动）
+echo "[entrypoint] patching ixformer ops..."
+python3 /opt/patch_ops.py || echo "[entrypoint] ixformer ops patch failed, but continuing with vllm startup"
+
 echo "[entrypoint] starting vllm..."
 
 exec vllm serve "$MODEL_DIR" $TOKENIZER_ARG "$@"

patch_ops.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""
+Patch ixformer functions - 自动扫描并复制自定义的ops文件并修改__init__.py
+
+支持批量处理patched_ops目录中的所有.py文件，无需每次修改代码。
+只需将新的ops文件放入patched_ops目录即可。
+"""
+import os
+import shutil
+from typing import List
+
+# 目标路径
+TARGET_DIR = "/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/"
+TARGET_INIT_FILE = os.path.join(TARGET_DIR, "__init__.py")
+SOURCE_DIR = "/opt/patched_ops/"
+
+
+def get_ops_files(source_dir: str) -> List[str]:
+    """
+    扫描源目录中的所有.py文件（排除__init__.py）
+
+    Args:
+        source_dir: 源目录路径
+
+    Returns:
+        .py文件名列表
+    """
+    if not os.path.exists(source_dir):
+        print(f"[patch_ops] Source directory not found: {source_dir}")
+        return []
+
+    ops_files = []
+    for filename in os.listdir(source_dir):
+        if filename.endswith('.py') and filename != '__init__.py':
+            ops_files.append(filename)
+
+    return ops_files
+
+
+def generate_import_statement(filename: str) -> str:
+    """
+    根据文件名生成import语句
+
+    Args:
+        filename: .py文件名（不含扩展名）
+
+    Returns:
+        import语句
+    """
+    module_name = filename.replace('.py', '')
+    return f"from .{module_name} import *"
+
+
+def patch_ixformer_ops() -> bool:
+    """
+    Patch ixformer functions - 批量复制ops文件并修改__init__.py
+
+    Returns:
+        是否成功patch
+    """
+    try:
+        print(f"[patch_ops] Starting ixformer ops patch...")
+
+        # 检查目标目录是否存在
+        if not os.path.exists(TARGET_DIR):
+            print(f"[patch_ops] Target directory not found: {TARGET_DIR}")
+            return False
+
+        # 扫描源目录中的所有ops文件
+        ops_files = get_ops_files(SOURCE_DIR)
+
+        if not ops_files:
+            print(f"[patch_ops] No ops files found in {SOURCE_DIR}")
+            return False
+
+        print(f"[patch_ops] Found {len(ops_files)} ops file(s): {', '.join(ops_files)}")
+
+        # 备份原__init__.py文件
+        if os.path.exists(TARGET_INIT_FILE):
+            backup_file = TARGET_INIT_FILE + ".backup"
+            if not os.path.exists(backup_file):
+                shutil.copy2(TARGET_INIT_FILE, backup_file)
+                print(f"[patch_ops] Backed up {TARGET_INIT_FILE} to {backup_file}")
+            else:
+                print(f"[patch_ops] Backup already exists: {backup_file}")
+
+        # 读取现有的__init__.py内容
+        existing_imports = set()
+        if os.path.exists(TARGET_INIT_FILE):
+            with open(TARGET_INIT_FILE, 'r') as f:
+                init_content = f.read()
+            # 提取现有的from .xxx import *语句
+            for line in init_content.split('\n'):
+                line = line.strip()
+                if line.startswith('from .') and 'import *' in line:
+                    existing_imports.add(line)
+
+        # 批量复制文件并收集需要添加的import语句
+        new_imports = []
+        for filename in ops_files:
+            source_file = os.path.join(SOURCE_DIR, filename)
+            target_file = os.path.join(TARGET_DIR, filename)
+
+            # 复制文件
+            shutil.copy2(source_file, target_file)
+            print(f"[patch_ops] Copied {filename} to {TARGET_DIR}")
+
+            # 生成import语句
+            import_stmt = generate_import_statement(filename)
+            if import_stmt not in existing_imports:
+                new_imports.append(import_stmt)
+
+        # 修改__init__.py文件，添加新的import语句
+        if new_imports:
+            with open(TARGET_INIT_FILE, 'a') as f:
+                f.write(f"\n# Auto-patched ixformer ops\n")
+                for import_stmt in new_imports:
+                    f.write(f"{import_stmt}\n")
+            print(f"[patch_ops] Added {len(new_imports)} import statement(s) to {TARGET_INIT_FILE}")
+        else:
+            print(f"[patch_ops] All imports already exist in {TARGET_INIT_FILE}, skipping modification")
+
+        print(f"[patch_ops] Successfully patched ixformer ops")
+        return True
+
+    except Exception as e:
+        print(f"[patch_ops] Error patching ixformer ops: {e}")
+        return False
+
+
+def main():
+    """主函数"""
+    try:
+        success = patch_ixformer_ops()
+        if success:
+            print(f"[patch_ops] Patch completed successfully")
+        else:
+            print(f"[patch_ops] Patch failed, but this will not prevent vLLM from starting")
+    except Exception as e:
+        print(f"[patch_ops] Unexpected error: {e}")
+        print(f"[patch_ops] This will not prevent vLLM from starting")
+
+
+if __name__ == "__main__":
+    import sys
+    try:
+        main()
+    except Exception as e:
+        print(f"[patch_ops] Unexpected error: {e}")
+        print(f"[patch_ops] This will not prevent vLLM from starting")
+        sys.exit(0)  # 总是成功退出，不阻止vLLM启动

patched_ops/gelu_tanh_and_mul.py

@@ -0,0 +1,29 @@
+import torch
+
+__all__ = ["gelu_tanh_and_mul"]
+
+
+def gelu_tanh_and_mul(
+    input: "torch.Tensor",
+    output: "torch.Tensor" = None
+):
+    assert isinstance(input, torch.Tensor)
+
+    if output is None:
+        output_shape = list(input.shape)
+        output_shape[-1] = output_shape[-1] // 2
+        output = input.new_empty(output_shape)
+
+    hidden_size = input.shape[-1] // 2
+
+    x = input[..., :hidden_size]
+    gate = input[..., hidden_size:]
+
+    result = x * torch.nn.functional.gelu(
+        gate,
+        approximate="tanh"
+    )
+
+    output.copy_(result)
+
+    return output