本文探讨了如何利用机器学习编译（MLC）技术，使AMD GPU在大型语言模型（LLM）推理方面具备竞争力。

**背景**

随着开源LLM的兴起，出现了许多LLM推理解决方案，其中大多数基于CUDA并针对NVIDIA GPU进行了优化。为了满足日益增长的计算需求，将支持扩展到更广泛的硬件加速器类别，而AMD GPU是一个潜在的候选者。

**硬件和软件讨论**

从规格比较来看，AMD的RX 7900 XTX与NVIDIA的RTX 4090和RTX 3090 Ti性能相当，都拥有24GB内存、相似的内存带宽。虽然4090的FP16性能是7900 XTX的两倍，3090 Ti的FP16性能是7900 XTX的1.3倍，但由于延迟敏感的LLM推理主要受内存限制，因此FP16性能并非瓶颈。此外，RX 7900 XTX的价格比RTX 4090便宜40%。

AMD在过去落后并非完全是硬件原因，主要问题在于缺乏软件支持和针对相关模型的优化。生态系统中存在两个因素开始改变这一现状：AMD正在加大对ROCm栈的投入，而机器学习编译等新兴技术有助于降低跨后端通用软件支持的总体成本。

**ROCm的机器学习编译**

机器学习编译是一种新兴技术，用于编译和自动化机器学习工作负载的优化。与为每个后端（例如ROCm或CUDA）专门编写内核不同，MLC解决方案会自动生成针对不同后端的代码。MLC-LLM基于Apache TVM Unity，这是一个机器学习编译栈，提供高效的Python优先开发和通用部署。MLC-LLM为各种后端（包括CUDA、Metal、ROCm、Vulkan和OpenCL）提供最先进的性能，涵盖从服务器级GPU到移动设备（iPhone和Android）。

**AMD GPU和APU的MLC**

支持AMD GPU有多种途径，包括ROCm、OpenCL、Vulkan和WebGPU。ROCm栈是AMD最近大力推动的，具有与CUDA栈类似的构建模块。Vulkan是最新的图形标准，提供了跨GPU设备最广泛的支持。WebGPU是最新的网络标准，允许计算在网络浏览器上运行。尽管有很多方法，但很少有ML软件解决方案构建除了CUDA以外的其他解决方案，原因是复制新硬件或GPU编程模型的栈的工程成本很高。MLC-LLM支持自动代码生成，无需为每个后端重新设计GPU内核，并支持所有这些方式。然而，性能仍然取决于底层GPU运行时的质量及其在每个平台上的可用性。文中选择ROCm用于Radeon 7900 XTX，选择Vulkan用于SteamDeck的APU。

**MLC Python包的基准测试**

文中对Llama 2 7B和13B进行了4位量化基准测试，并通过设置单个提示标记并生成512个标记来测量解码性能，所有结果均针对单批次推理。

结果表明，单批次推理性能可以达到NVIDIA 4090速度的80%。

**在SteamDeck上使用Vulkan和统一内存进行运行**

文章也探讨了更广泛的AMD设备，例如配备AMD APU的SteamDeck。尽管ROCm中GPU VRAM在BIOS中被限制为4GB，但Mesa Vulkan驱动程序提供了强大的支持，允许缓冲区通过使用统一内存扩展到16GB，足以运行4位量化Llama-7B。

**讨论和未来工作**

在生成式人工智能时代，硬件可用性已成为一个紧迫问题。ML编译可以通过在硬件后端提供高性能通用部署来提供帮助。根据文中提供的证据，认为在合适的价位和可用性下，AMD GPU可以开始用于LLM推理。

未来的工作包括：启用批处理和多GPU支持、与PyTorch生态系统集成、支持更多量化和模型架构、在更多硬件后端引入更多自动优化等。

**总结**

机器学习系统工程是一个持续存在的问题。NVIDIA凭借持续的创新仍然处于领先地位，预计随着H100等新硬件以及软件的演进，格局将会发生变化。因此，关键问题不仅仅是现在构建正确的解决方案，还在于如何持续赶超并将其引入新的平台。机器学习工程中的生产力至关重要，得益于Python优先的ML编译开发流程，文中在几个小时内获得了ROCm优化的支持。随着探索更多想法以实现通用部署并解决硬件可用性问题，预计相关方法将变得更加有用。