tiny-vllm 是一个用 C++ 和 CUDA 编写的高性能 LLM 推理引擎（vLLM 的轻量版），完整实现了 Llama 3.2 1B 的前向推理、KV cache、continuous batching 与 PagedAttention。本文带你用 ROCm/HIP 兼容层在 AMD GPU 上编译运行 tiny-vllm，全程 CUDA 源码零改动 ⚡

tiny-vllm 是什么？

tiny-vllm 由两部分组成：一套完整的推理服务源码，和一门从零手写推理引擎的课程。引擎能力一览：

• ✅ 从 Safetensors 加载真实 LLM 权重（Llama 3.2 1B Instruct）
• ✅ 全部计算由 GPU kernel 完成（prefill + decode）
• ✅ KV cache、静态批处理、continuous batching
• ✅ FlashAttention 风格的 online softmax 与 PagedAttention

源码高度精简：src/main.cpp 承载推理主流程，src/kernels.cu 集中所有 kernel，唯一外部依赖是单头文件 JSON 解析器 include/json.hpp。

移植原理：一个兼容头搞定 CUDA→HIP

AMD 的 HIP 是 CUDA 的开源对应物：同一套编程模型，编译器为 hipcc，数学库从 cuBLAS 换成 hipBLAS。移植的关键就是一个头文件 src/cuda_to_hip.h——在 AMD 平台上编译时，它用宏把源文件中的"CUDA 拼写"逐一映射到 HIP 符号：

也就是说，所有源码继续用 CUDA 的写法，编译器看到的只是替换后的 HIP 符号。

还有一个容易踩坑的细微差异：CUDA 的 warp shuffle（__shfl_*）使用 32 位掩码，而 HIP 要求 64 位掩码。兼容头为此为两个平台分别定义了WARP_FULL_MASK宏，src/kernels.cu 中 PagedAttention kernel 统一引用该宏，无需手动区分平台。

三步在 AMD GPU 上编译运行

第 1 步：安装 ROCm

在 Linux 上安装与 GPU 匹配的 ROCm 发行版（默认位于/opt/rocm），它提供 hipcc 编译器与 hip、hipBLAS 包。

第 2 步：获取源码

git clone https://gitcode.com/gh_mirrors/ti/tiny-vllm cd tiny-vllm

第 3 步：CMake 配置与构建 🚀

CMakeLists.txt 内置了USE_HIP开关，打开后 CMake 会改用 hipcc 编译main.cpp与kernels.cu（bfloat16 依赖 clang 内建函数，两个文件都必须走 HIP 编译），并链接hip::host与roc::hipblas：

cmake -B build -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1100 -DCMAKE_PREFIX_PATH=/opt/rocm -G Ninja cmake --build build ./build/tiny-vllm

三个关键参数：

• -DUSE_HIP=ON：启用 HIP 构建模式
• -DCMAKE_HIP_ARCHITECTURES：目标 GPU 架构（见下表），不传则自动探测——注意无 GPU 的构建机上自动探测会失败，必须显式指定
• -DCMAKE_PREFIX_PATH=/opt/rocm：让 CMake 找到 hip 与 hipBLAS 包；若/opt/rocm/bin已在PATH中可省略

GPU 架构怎么选？

先用rocminfo查询本机 GPU 的架构名，常见取值：

作者已在gfx90a、gfx1100、gfx1201三种架构上验证通过，架构值不是写死的，按你的卡如实设置即可。

验证移植结果

程序启动时先执行 src/main.cpp 中的checkGPUStatus()，打印 GPU 名称、显存总量与计算单元数——这些调用走兼容头映射到hip*函数，输出格式与 NVIDIA 版完全一致，看到自己的 AMD 卡名即说明运行时正常工作。

随后引擎会加载 Llama 3.2 1B 的model.safetensors权重并生成 token。运行入口可参考仓库根目录的脚本：test.sh（构建+执行）、run.sh（直接运行）、check.sh（用 compute-sanitizer 做内存检查）。模型获取方式见 README.md 的 "Safetensors and your model" 一节。

常见问题 FAQ

Q：CMake 报找不到 hip / hipBLAS 包？加-DCMAKE_PREFIX_PATH=/opt/rocm，或把/opt/rocm/bin加入PATH；ROCm 装在别处则改路径。

Q：构建机上没有 GPU，自动探测架构报错？显式传-DCMAKE_HIP_ARCHITECTURES=gfxXXXX即可，project(... LANGUAGES HIP)会尊重该参数。

Q：编译通过但输出乱码/报错？大概率是架构值与物理 GPU 不匹配（例如给 RDNA3 卡传了gfx90a），用rocminfo核对后重建。

Q：会影响原来的 NVIDIA 构建吗？不会。不传-DUSE_HIP时，CMake 行为与之前完全一致，仍走 nvcc + cuBLAS。

小结

一页 src/cuda_to_hip.h 兼容头 + 一个USE_HIP构建开关 = CUDA 源码零改动跑上 AMD GPU。这种"兼容层头文件 + 编译器开关"的移植模式非常轻量，如果你手里有其他 CUDA 推理项目，同样值得一试 ⚡