MinerU 在 AMD ROCm 显卡上的实战调优：用 Triton 自定义算子修复 vLLM 与 DocLayout-YOLO 的卷积性能回退

本文为 MinerU 在 AMD GPU（ROCm 生态）上的完整调优指南：定位 vLLM 视觉塔中 Conv3d 算子在 MIOpen 上缺少优化内核导致的严重性能回退，并给出两套可直接落地的 Triton 自定义算子方案（逐点积版与矩阵乘法版），同时覆盖 pipeline 后端中 DocLayout-YOLO 空洞卷积的同源问题。读完本文，你可以在 7900XTX 等 RDNA 显卡上跑通 MinerU 的 vLLM 推理链路，并把单页处理速度从秒级回退恢复到约 1.3 s/it，同时把 layout 识别从约 1.6 it/s 提升到约 15 it/s。

一、问题定位：AMD RDNA 上 vLLM 后端为什么慢

MinerU 的 VLM 链路支持多种推理后端。从源码看，vlm_analyze.py 中 vllm-engine 后端通过 vllm.LLM(**kwargs) 同步引擎加载模型，并支持注入 MinerULogitsProcessor（来自独立的 mineru_vl_utils 包）；后端的可选值集中在 backend_options.py 中维护，包括 pipeline、vlm-engine、hybrid-engine 等本地后端与两个 HTTP client 后端。vLLM 服务化部署入口见 vllm_server.py。

在 AMD RDNA 显卡上使用 vLLM 后端时存在严重的性能问题，根因是：

• vLLM 的 qwen2_vl.py 中有一个算子在 ROCm kernel 上没有对应实现，导致卷积计算发生严重的性能回退——MIOpen 库中缺少该模型中特定 Conv3d(bfloat16) 的优化内核，该算子单次执行耗时一度高达 12 s；
• pipeline 后端中 layout 识别使用的 DocLayout-YOLO 的 g2l_crm.py 中的空洞卷积存在同样的问题，连 CDNA 架构的 MI210 也没有内置解决，因此需要一并处理。

优化思路是：既然 MIOpen 没有现成内核，就用 Triton 后端自己实现对应算子，绕开回退路径。排查方法建议遵循通用流程：先用 prof 查看耗时分布，定位到具体出问题的算子，再写 Triton 后端实现替换它。这一方法适用于其他 AMD GPU 上的类似问题。

二、测试结果速览（7900XTX 实测）

作者在 7900XTX 上的实测数据（2025-10 前后，环境见下节）：

• 用 -b vlm-vllm-engine 模式解析 14 篇学术论文，Processed prompts 输出速度达到 input 5264.66 toks/s、output 2733.31 toks/s 的量级；
• 交互式界面入口测试 Two Step Extraction: 14/14 [00:18, 1.30s/it]；
• 200 页 PDF 编程书测试，Two Step Extraction 达到 1.99 it/s（01:40 处理 200 页）；
• 方案一（Triton 点积实现）下整页速度约 1.6–1.8 s/it；方案二（矩阵乘法实现）进一步提速到约 1.3 s/it。优化后剩余耗时主要集中在 hipBLAS（不可再压）与 vLLM 的 Triton 后端，二者各占约 25%；
• DocLayout-YOLO 的 layout 速度从原来的 1.6 it/s 提高到 15 it/s。注意 Triton kernel 必须缓存输入尺寸（无法动态获取 PDF 尺寸），实现上保留了原模型的输入输出接口，代码改动最小。

整体解析效果与 MinerU 官方在 CUDA 上的效果基本一致。

三、测试环境

项目	版本
系统	Ubuntu 24.04，Kernel Linux 6.14.0-33-generic
ROCm	7.0.1
Python	3.12
pytorch-triton-rocm	3.5.0
torch	2.10.0.dev+rocm7.0（nightly）
torchvision	0.25.0.dev+rocm7.0（nightly）
vllm	0.11.0rc2.dev+rocm701（源码 develop 安装）

不同版本之间差异不影响处理方法本身，核心都是“Triton 替换 MIOpen 缺失的卷积内核”。

四、前置环境安装

已有完整 Python vLLM 与 MinerU 环境的可直接跳到第六节。从零搭建：

uv venv --python python3.12
source .venv/bin/activate
# 安装 ROCm nightly 版 torch（清华镜像 + PyTorch nightly rocm7.0 索引）
uv pip install --pre torch torchvision \
    -i https://pypi.tuna.tsinghua.edu.cn/simple/ \
    --extra-index-url https://download.pytorch.org/whl/nightly/rocm7.0
uv pip install pip
# 注意：为避免覆盖本地 pytorch，改用 pip 安装 mineru
pip install -U "mineru[core]" -i https://pypi.mirrors.ustc.edu.cn/simple/

vLLM 与相关组件需从源码 develop 安装（参考 vLLM 官方 ROCm 安装手册的 AMD/ROCm 部分）：

# 手动安装 aiter
git clone --recursive https://github.com/ROCm/aiter.git
cd aiter
git submodule sync; git submodule update --init --recursive
python setup.py develop
cd ..
# 手动安装 vllm
git clone https://github.com/vllm-project/vllm.git
cd vllm/
cp -r /opt/rocm/share/amd_smi ~/Pytorch/vllm/
pip install amd_smi/
pip install --upgrade numba \
    scipy \
    "huggingface-hub[cli,hf_transfer]" \
    setuptools_scm
pip install -r requirements/rocm.txt
export PYTORCH_ROCM_ARCH="gfx1100"   # 按自己的 GPU 架构设置，可用 rocminfo | grep gfx 查看
python setup.py develop

两个坑位提醒：

• PYTORCH_ROCM_ARCH 必须匹配自己的 GPU 架构（RDNA 3 如 7900XTX 为 gfx1100），先用 rocminfo | grep gfx 确认；
• 如果 pip 把 triton 后端的 flash_attn 依赖卸掉导致反复报错，不要纠缠——按本文方案直接不用该路径即可。

五、vLLM 中关键 Triton 算子的添加（核心）

先用 pip show vllm 定位 vLLM 安装目录，下文记为 XXX。

5.1 公共改动：引入自定义 kernel 模块

在 XXX/vllm/model_executor/models/qwen2_vl.py 第 33 行下方增加导入：

from collections.abc import Iterable, Mapping, Sequence
from functools import partial
from typing import Annotated, Any, Callable, Literal, Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from .qwen2_vl_vision_kernels import triton_conv3d_patchify

问题对象就是 qwen2_vl.py 中第 498 行附近的 class Qwen2VisionPatchEmbed(nn.Module)——它的 nn.Conv3d 在 AMD 上没有现成内核算子，从而发生回退。接下来在方案一（5.2 + 5.3）与方案二（5.4 + 5.5）之间二选一即可：

• 方案一：Triton 逐点积实现，通用性好，实测 1.6–1.8 s/it；
• 方案二：5D 张量直接转 GEMM（矩阵乘法），7900XTX 上约 1.3 s/it；其中手工调优的 block 配置针对 7900XTX，其他 AMD GPU 相对方案一也有提速但不一定是最佳组合，可能需要自行微调（AMD 上的 autotune 基本不起作用）。

5.2 方案一：修改 Qwen2VisionPatchEmbed，调用点积 kernel

class Qwen2VisionPatchEmbed(nn.Module):

    def __init__(
        self,
        patch_size: int = 14,
        temporal_patch_size: int = 2,
        in_channels: int = 3,
        embed_dim: int = 1152,
    ) -> None:
        super().__init__()
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        self.embed_dim = embed_dim

        kernel_size = (temporal_patch_size, patch_size, patch_size)
        self.proj = nn.Conv3d(in_channels,
                              embed_dim,
                              kernel_size=kernel_size,
                              stride=kernel_size,
                              bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        L, C = x.shape
        x_reshaped = x.view(L, -1, self.temporal_patch_size, self.patch_size,
                            self.patch_size)

        # 用自定义 Triton kernel 代替 self.proj
        x_out = triton_conv3d_patchify(x_reshaped, self.proj.weight)

        # kernel 输出形状已经是 [L, embed_dim]
        return x_out

5.3 方案一：创建 qwen2_vl_vision_kernels.py（点积实现）

在 XXX/vllm/model_executor/models/ 目录下新建 qwen2_vl_vision_kernels.py，用 Triton 实现“无重叠 3D patch 卷积”（本质是 patch 向量与卷积核向量的点积累加）：

import torch
from vllm.triton_utils import tl, triton

@triton.jit
def _conv3d_patchify_kernel(
    # Pointers to tensors
    X, W, Y,
    # Tensor dimensions
    N, C_in, D_in, H_in, W_in,
    C_out, KD, KH, KW,
    # Stride and padding for memory access
    stride_xn, stride_xc, stride_xd, stride_xh, stride_xw,
    stride_wn, stride_wc, stride_wd, stride_wh, stride_ww,
    stride_yn, stride_yc,
    # Triton-specific metaparameters
    BLOCK_SIZE: tl.constexpr,
):
    """
    Triton kernel for a non-overlapping 3D patching convolution.
    Each kernel instance computes one output value for one patch.
    """
    # Get the program IDs for the N (patch) and C_out (output channel) dimensions
    pid_n = tl.program_id(0)  # The index of the patch we are processing
    pid_cout = tl.program_id(1) # The index of the output channel we are computing

    # --- Calculate memory pointers ---
    # Pointer to the start of the current input patch
    x_ptr = X + (pid_n * stride_xn)
    # Pointer to the start of the current filter (weight)
    w_ptr = W + (pid_cout * stride_wn)
    # Pointer to where the output will be stored
    y_ptr = Y + (pid_n * stride_yn + pid_cout * stride_yc)

    # --- Perform the convolution (element-wise product and sum) ---
    # This is a dot product between the flattened patch and the flattened filter.
    accumulator = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)

    # Iterate over the elements of the patch/filter
    for c_offset in range(0, C_in):
        for d_offset in range(0, KD):
            for h_offset in range(0, KH):
                # Unrolled loop for the innermost dimension (width) for performance
                for w_offset in range(0, KW, BLOCK_SIZE):
                    # Create masks to handle cases where KW is not a multiple of BLOCK_SIZE
                    w_range = w_offset + tl.arange(0, BLOCK_SIZE)
                    w_mask = w_range < KW

                    # Calculate offsets to load data
                    patch_offset = (c_offset * stride_xc + d_offset * stride_xd +
                                    h_offset * stride_xh + w_range * stride_xw)
                    filter_offset = (c_offset * stride_wc + d_offset * stride_wd +
                                     h_offset * stride_wh + w_range * stride_ww)

                    # Load patch and filter data, applying masks
                    patch_vals = tl.load(x_ptr + patch_offset, mask=w_mask, other=0.0)
                    filter_vals = tl.load(w_ptr + filter_offset, mask=w_mask, other=0.0)

                    # Multiply and accumulate
                    accumulator += patch_vals.to(tl.float32) * filter_vals.to(tl.float32)

    # Sum the accumulator block and store the single output value
    output_val = tl.sum(accumulator, axis=0)
    tl.store(y_ptr, output_val)


def triton_conv3d_patchify(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """
    Python wrapper for the 3D patching convolution Triton kernel.
    """
    # Get tensor dimensions
    N, C_in, D_in, H_in, W_in = x.shape
    C_out, _, KD, KH, KW = weight.shape

    # Create the output tensor
    # The output of this specific conv is (N, C_out, 1, 1, 1), which we squeeze
    Y = torch.empty((N, C_out), dtype=x.dtype, device=x.device)

    # Define the grid for launching the Triton kernel
    # Each kernel instance handles one patch (N) for one output channel (C_out)
    grid = (N, C_out)

    # Launch the kernel
    # We pass all strides to make the kernel flexible
    _conv3d_patchify_kernelgrid, x.stride(1), x.stride(2), x.stride(3), x.stride(4),
        weight.stride(0), weight.stride(1), weight.stride(2), weight.stride(3), weight.stride(4),
        Y.stride(0), Y.stride(1),
        BLOCK_SIZE=16, # A reasonable default, can be tuned
    )

    return Y

5.4 方案二：Qwen2VisionPatchEmbed 直接交给 GEMM kernel

class Qwen2VisionPatchEmbed(nn.Module):

    def __init__(
        self,
        patch_size: int = 14,
        temporal_patch_size: int = 2,
        in_channels: int = 3,
        embed_dim: int = 1152,
    ) -> None:
        super().__init__()
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        self.embed_dim = embed_dim

        kernel_size = (temporal_patch_size, patch_size, patch_size)

        self.proj = nn.Conv3d(in_channels,
                              embed_dim,
                              kernel_size=kernel_size,
                              stride=kernel_size,
                              bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        L, C = x.shape
        x_reshaped_5d = x.view(L, -1, self.temporal_patch_size, self.patch_size,
                               self.patch_size)

        # 5D 张量一步到位，转为矩阵乘法
        return triton_conv3d_patchify(x_reshaped_5d, self.proj.weight)

5.5 方案二：创建 qwen2_vl_vision_kernels.py（GEMM 实现）

import torch
from vllm.triton_utils import tl, triton

@triton.jit
def _conv_gemm_kernel(
    A, B, C, M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = A + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
        offs_k += BLOCK_K
    c = accumulator.to(C.dtype.element_ty)
    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

def triton_conv3d_patchify(x_5d: torch.Tensor, weight_5d: torch.Tensor) -> torch.Tensor:
    N_patches, _, _, _, _ = x_5d.shape
    C_out, _, _, _, _ = weight_5d.shape
    A = x_5d.view(N_patches, -1)
    B = weight_5d.view(C_out, -1).transpose(0, 1).contiguous()
    M, K = A.shape
    _K, N = B.shape
    assert K == _K
    C = torch.empty((M, N), device=A.device, dtype=A.dtype)

    # 针对 7900xtx 的手动调优配置，其他 GPU 的最优组合可能需要自行寻找，
    # AMD 上的 autotune 基本没有效果
    best_config = {
        'BLOCK_M': 128,
        'BLOCK_N': 128,
        'BLOCK_K': 32,
    }
    num_stages = 4
    num_warps = 8

    grid = (triton.cdiv(M, best_config['BLOCK_M']),
            triton.cdiv(N, best_config['BLOCK_N']))

    _conv_gemm_kernelgrid, A.stride(1),
        B.stride(0), B.stride(1),
        C.stride(0), C.stride(1),
        **best_config,
        num_stages=num_stages,
        num_warps=num_warps
    )

    return C

原理上，stride 等于 kernel size 的 Conv3d 就是“patchify + 线性投影”：把每个 patch 展平后与展平的卷积核做内积。方案一逐点累加，逻辑直白；方案二利用 tl.dot 把整批 patch 组织成标准 GEMM（A 为 N_patches×K，B 为 K×C_out），从而吃到 tensor core / 矩阵单元，这也是它在 7900XTX 上更快的原因。

六、收尾修复：mineru_vl_utils 的 LoRA 报错与环境变量

修改 vLLM 后，关闭终端再次启动交互界面时会遇到一个 LoRA 相关报错。定位并打补丁：

pip show mineru_vl_utils

打开 XXX/mineru_vl_utils/vlm_client/vllm_async_engine_client.py，将第 58 行的 self.tokenizer = vllm_async_llm.tokenizer.get_lora_tokenizer() 改为带降级逻辑的版本：

try:
    self.tokenizer = vllm_async_llm.tokenizer.get_lora_tokenizer()
except AttributeError:
    # 如果没有 get_lora_tokenizer 方法，直接使用原始 tokenizer
    self.tokenizer = vllm_async_llm.tokenizer

从源码结构看，mineru_vl_utils 是 MinerU 独立出的 VLM 推理工具包（见 changelog 中“vlm 推理相关代码已移至 mineru_vl_utils，降低与主仓库耦合度”的记录），主仓库通过 from mineru_vl_utils import MinerUClient（vlm_analyze.py）与其集成，因此该补丁作用于独立的工具包而非 MinerU 本体。

最后设置两个环境变量即可正常运行：

export MINERU_MODEL_SOURCE=modelscope
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

其中 TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 开启 ROCm 上的 Triton 实验性 AOT 支持，是让上述 Triton kernel 生效的关键开关；MINERU_MODEL_SOURCE=modelscope 控制模型从 ModelScope 下载。

七、pipeline 后端的 DocLayout-YOLO 空洞卷积问题

vLLM 链路修好后，pipeline 后端中 layout 识别使用的 DocLayout-YOLO 模型还存在同样的空洞卷积内核缺失问题（g2l_crm.py）。作者已在 DocLayout-YOLO 项目对应 issue 的回复中给出完整改法，思路与本文一致：用 Triton 重写空洞卷积算子。落地步骤：

# 查看 doclayout-yolo 安装位置
pip show doclayout-yolo

进入该位置后，修改 issue 回复中介绍的文件即可。需要注意该方案要求缓存输入 PDF 的尺寸信息（Triton kernel 无法动态感知尺寸），这也是 layout 速度能从 1.6 it/s 提到 15 it/s 的前提之一。

顺带说明当前仓库的 layout 模型现状：pipeline 后端的 layout 初始化入口在 model_init.py 中，当前仓库内置的是 pp_doclayout_v2 布局模型（实现见 pp_doclayoutv2.py）；而 doclayout_yolo 模型是此前版本引入的自研 layout 方案，据 changelog 记录，它在相近解析效果下比原方案提速 10 倍以上，并升级到了 doclayout_yolo(2501)。如果你当前安装的版本仍使用 doclayout-yolo 作为 pipeline layout 模型，就适用本节修复；如果已经切换到 pp_doclayout_v2，可先实测确认是否还有卷积回退，再决定是否套用同样的 Triton 思路。

八、验证与复现要点小结

1. 确认环境：pip show vllm、pip show mineru_vl_utils、pip show doclayout-yolo 定位三处需要打补丁的位置；rocminfo | grep gfx 确认 PYTORCH_ROCM_ARCH；
2. 选择方案：求稳选方案一（点积 Triton kernel）；求速度且为 7900XTX 选方案二（GEMM Triton kernel，block 配置可自调）；
3. 打补丁顺序：修改 qwen2_vl.py 导入与 Qwen2VisionPatchEmbed → 新建 qwen2_vl_vision_kernels.py → 修复 vllm_async_engine_client.py 的 LoRA 降级；
4. 设置 MINERU_MODEL_SOURCE=modelscope 与 TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1；
5. 用 -b vlm-vllm-engine 或交互入口跑 14 篇论文级别的测试集，对照本文 1.3–1.8 s/it 的量级验证是否生效；layout 环节对照 15 it/s 量级验证空洞卷积修复。

适用前提与限制：本文实测基于 7900XTX（gfx1100）+ ROCm 7.0.1 + vLLM 0.11 开发版，结论对其他 AMD GPU 的迁移属于“同类方法、效果需自测”；方案二中的手工调优参数（BLOCK_M/N/K、num_stages、num_warps）仅对 7900XTX 验证过，其他架构请自行搜索最优组合；由于补丁直接改在 pip/源码安装的 vllm 与 mineru_vl_utils 包内部，升级这两个包时需要重新应用补丁。