(2025.12.23更新加速版)AMD RDNA ROCm vllm后端演示，vllm nn.Conv3d 算子triton适配，以及pipeline中DocLayout-YOLO适配分享

[healy-hub]

据AMD反馈的issue，似乎最近RDNA的CK后端就要上线了，期待一手，到时候测试分享。

2025.12.23更新回答分享，删除了第一种慢速方式，仅采用5D矩阵乘法实现。另外增加Triton 后端的flash_attn使用，这里很有意思，LLM部分是flash_attn运算的，但是ViT model部分是Torch SDPA backend，其实是通过环境变量加载AMD预设的AOTriton计算的，最终整体性能提升10%以上了，vllm这里使用的是最新编译的0.14.0rc1.dev3+gde08b8f61.d20251219.rocm711。Torch版本有特定推荐，这个版本的fp16和bf16矩阵乘能到104tflops的结果，相比月早期的94tflops提升明显。

基于Triton的ROCm 不同后端实现优化，基本实现vllm后端正常推理，以及pipeline后端中第一步layout用的DocLayout-YOLO

已有完整python vllm和mineru环境直接跳转第五步！！！ 建议使用推荐版本的vllm和Torch
其他GPU执行问题可以参考，先prof查看定位找到哪个算子问题，然后triton后端实现即可
测试了一下，基本和MinerU官网效果差不多，用AMD的人也不是很多，就在评论区分享给大家了

1.结果介绍

直接gradio网页测试一本531页的第三版 python编程从入门到实践书测试一下速度，可以到1.79it/s：
这个速度我觉得已经有不错的性价比了
Two Step Extraction: 100%|███████████████████████████████████████████████| 531/531 [04:57<00:00, 1.79it/s]

下面为具体的终端执行测试结果：
采用-b vlm-vllm-engine模式举个例子

测试结果为优化为5d矩阵乘代替原来的点积结果：
Adding requests: 100%|██████████████████████████████████████████████| 531/531 [00:07<00:00, 68.05it/s]
Processed prompts: 100%|██████████████| 531/531 [02:25<00:00, 3.64it/s, est. speed input: 5075.22 toks/s, output: 927.59 toks/s]
Warning: line does not match layout format:
Warning: line does not match layout format:
Adding requests: 100%|██████████████████████████████████████████████| 5907/5907 [00:15<00:00, 373.37it/s]
Processed prompts: 100%|███████████| 5907/5907 [01:47<00:00, 54.73it/s, est. speed input: 9060.71 toks/s, output: 2063.93 toks/s]

---

2.原因介绍

AMD RDNA使用vllm后端有严重的性能问题，原因是因为vllm的qwen2_vl.py中有一个算子在rocm kernel上没有对应的实现，导致性能出现严重的卷积计算回退，一次执行花了12s，。。。。。。。。一言难尽。即MIOpen 库中缺少模型中特定 Conv3d(bfloat16) 的优化内核。
DocLayout-YOLO的g2l_crm.py空洞卷积也是这个问题，专业的CDNA MI210也没解决这个问题
正好一起处理了。

---

3.环境介绍

System: Ubuntu 24.04.3 Kernel: Linux 6.14.0-37-generic ROCm version: 7.1.1 CPU 13900K 内存 64G 6800MHz ddr5
python环境：
python 3.13.8
pytorch-triton-rocm 3.6.0+git5261b273
torch 2.10.0.dev20251208+rocm7.1
torchvision 0.25.0.dev20251209+rocm7.1
vllm 0.14.0rc1.dev3+gde08b8f61.d20251219.rocm711
aiter 0.1.7.post3.dev34+ge0fb38eda
flash_attn 2.8.3

不同版本无所谓，处理方法是一样的。

---

4.前置环境安装

uv venv --python python3.13
source .venv/bin/activate
uv pip install --pre torch==2.10.0.dev20251208+rocm7.1 torchvision==0.25.0.dev20251209+rocm7.1 pytorch-triton-rocm==3.6.0+git5261b273 --index-url https://download.pytorch.org/whl/nightly/rocm7.1
# 最近的更新里就这附近的版本最猛
uv pip install pip
# 避免覆盖我们本地的pytorch，改用pip而没有继续使用uv pip
pip install -U "mineru[core]" -i https://pypi.mirrors.ustc.edu.cn/simple/

vllm 安装参考官方手册Vllm

#手动安装aiter，vllm，amd-smi等，自行找一个位置clone，然后进入该目录吧
git clone --recursive https://github.com/ROCm/aiter.git
cd aiter
git submodule sync; git submodule update --init --recursive
python setup.py develop
cd ..
git clone --recursive https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
git checkout main_perf
export FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"    
#至关重要，官方FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" python setup.py install是错的
python setup.py install
cd ..
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

---

5.vllm中关键triton算子添加

手动优化算子到矩阵乘法，7900xtx肯定适用，其他AMD GPU不一定是最佳速度实现，里面的手动部分可能需要微调。

#定位自己vllm位置XXX
pip show vllm

关键更改
XXX/vllm/model_executor/models/qwen2_vl.py文件：
1.qwen2_vl.py文件33行下增加from .qwen2_vl_vision_kernels import triton_conv3d_patchify

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

---

2.qwen2_vl.py文件498行class Qwen2VisionPatchEmbed(nn.Module)函数,PS.就是这玩意AMD没有现成的内核算子导致回退，这里我们直接5D张量一步到位，改为矩阵乘法

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)

        return triton_conv3d_patchify(x_reshaped_5d, self.proj.weight)

3.XXX/vllm/model_executor/models/目录下创建qwen2_vl_vision_kernels.py文件，用triton实现

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_base = tl.arange(0, BLOCK_K)

    # 计算指向输入矩阵 A 和 B 的初始指针
    a_ptrs = A + (offs_m[:, None] * stride_am + offs_k_base[None, :] * stride_ak)
    b_ptrs = B + (offs_k_base[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    # 初始化累加器
    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # 在 K 维度上进行循环，执行矩阵乘法的核心计算
    for k in range(0, K, BLOCK_K):
        # 处理 K 维度不是 BLOCK_K 整倍数的情况，创建掩码（mask）
        k_remaining = K - k
        offs_k = tl.arange(0, BLOCK_K)

        # 使用掩码加载数据块，边界外的数据用 0.0 填充
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0)

        # 执行点积运算
        accumulator += tl.dot(a, b)

        # 更新指针到下一个 K 块
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    # 将累加器的结果转换为输出张量 C 的数据类型
    c = accumulator.to(C.dtype.element_ty)

    # 计算指向输出矩阵 C 的指针，并使用掩码存储最终结果
    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:
    """
    使用 Triton 内核执行 3D 卷积的 patchify -> GEMM 流程。
    """
    N_patches, _, _, _, _ = x_5d.shape
    C_out, _, _, _, _ = weight_5d.shape

    # 将输入和权重张量重塑为二维矩阵以进行 GEMM 计算
    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的最优组合可能需要自行寻找，直接代码让AI在这个值附近autotune修改一下吧，不要把范围弄得太大了，会死的很惨，要求就在这个范围附近小点 ---
    best_config = {
        'BLOCK_M': 64,
        'BLOCK_N': 128,
        'BLOCK_K': 32,
        'num_stages': 2,
    }
    num_warps = 4

    # 定义启动网格（grid），即需要多少个程序实例来覆盖整个计算
    grid = (triton.cdiv(M, best_config['BLOCK_M']),
            triton.cdiv(N, best_config['BLOCK_N']))

    # 启动 Triton 内核
    _conv_gemm_kernel[grid](
        A, B, C,
        M, N, K,
        A.stride(0), A.stride(1),
        B.stride(0), B.stride(1),
        C.stride(0), C.stride(1),
        **best_config,
        num_warps=num_warps
    )

    return C

---

4.关闭终端后再次使用mineru-gradio会报一个Lora错误，修改代码跳过它，似乎mineru官方改了，看一下没问题就不用改了

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

最后整三个环境变量后愉快玩耍即可

export MINERU_MODEL_SOURCE=modelscope
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"  #使用时也需要，否则flash_attn不识别

---

6.vllm后端已经没有问题，下面是pipeline 中layout用的doclayout-yolo模型空洞卷积问题

我在 DocLayout-YOLO 下做了一个回答，因此 pipeline 的空洞卷积问题不在这里赘述，直接点击链接查看即可。

查看自己doclayout-yolo安装位置如下，然后进入修改链接中回复介绍的文件即可

pip show doclayout-yolo

[tinafengfun]

非常赞的patch，能详细说说为啥7900 要这样手动调的原因？--- 针对7900xtx的手动调优配置，其他GPU的最优组合可能需要自行寻找，AMD的autotune效果就是没有效果

[healy-hub]

非常赞的patch，能详细说说为啥7900 要这样手动调的原因？--- 针对7900xtx的手动调优配置，其他GPU的最优组合可能需要自行寻找，AMD的autotune效果就是没有效果

我当时开始用autotune来自动调优的时候发现效果比没有算子实现好的多，但是性能还是很低下， 至少是在当时那个triton版本以及自行clone编译的版本这个autotune是失效的，在给定的范围没有找到最优解，可能是我给的范围太大了？所以当时干脆写了一个脚本，直接rocprofv3记录下不同组合的kernel执行时间，粗调一部分到细调一部分,排个序手动找了一下。也许最新的triton仓库已经解决这个问题了，我回头测试一下。
现在脚本测试微调找到的最优大概是：
best_config = {
'BLOCK_M': 64,
'BLOCK_N': 128,
'BLOCK_K': 32,
'num_stages': 2,
}
num_warps = 4

PS. VLLM v11.1.0的性能比最新的v11.1.1，v11.1.2性能要好，感觉最近的更改对基于qwen2_vl的模型并不友好，并发吞吐下降了5%左右，非常稳定的下降。torch.nn.Conv3d在大的batch size（mineru 需要）算子回退到native未优化问题反馈给了amd，不知道rocm 7.2.0能解决不？看他们速度了，相比nvidia cudnn来说，MIOpen的边缘覆盖在RDNA上少的可怜，mineru pipeline中的det_ocr速度也是回退到慢的很，问题似乎出在了一个步长为2的conv2d卷积回退了，有空了再写个triton给实现一下好了。