摘要
arXiv:2504.12984v2 Announce Type: replace-cross
摘要:运行大规模语言模型(LLMs)对于AI驱动的应用至关重要,但同时也需要大量的计算资源,特别是在内存带宽和计算吞吐量方面。低精度计算已经成为了改进效率同时减少资源消耗的关键技术。当前用于生成低精度内核的方法局限于权重位宽为2的幂,并且由于高层次GPU编程抽象而导致性能不佳。这些抽象限制了诸如细粒度寄存器管理和优化的内存访问模式等关键优化,这些都是有效低精度计算所必需的。本文介绍了一种用于通用GPU(GPGPU)计算的虚拟机(VM),它能够支持任意位宽的低精度数据类型,同时保持GPU的编程能力。该提出的VM具有线程块级编程模型、分层内存空间、新颖的代数布局系统,并广泛支持各种低精度数据类型。VM程序编译成高效的GPU程序,具有自动向量化和指令选择。广泛的实验表明,我们的VM能够高效地支持各种低精度数据类型,并在支持的类型上优于最先进的低精度内核。与现有的编译器(如 Triton 和 Ladder)以及手工优化的内核(如 QuantLLM 和 Marlin)相比,我们的VM分别实现了1.75倍、2.61倍、1.29倍和1.03倍的性能提升。