MBZUAI团队:CASS如何实现跨平台代码转译?

B站影视 内地电影 2025-06-09 15:23 2

摘要:2025年5月29日,由MBZUAI(穆罕默德·本·扎耶德人工智能大学)和澳大利亚国立大学的研究团队共同发表了一项开创性研究成果。Ahmed Heakl、Sarim Hashmi、Gustavo Bertolo Stahl、Seung Hun Eddie Ha

2025年5月29日,由MBZUAI(穆罕默德·本·扎耶德人工智能大学)和澳大利亚国立大学的研究团队共同发表了一项开创性研究成果。Ahmed Heakl、Sarim Hashmi、Gustavo Bertolo Stahl、Seung Hun Eddie Han、Salman Khan和Abdulrahman Mahmoud带来了名为"CASS: Nvidia to AMD Transpilation with Data, Models, and Benchmark"的研究论文,该论文已上传至arXiv预印本平台(arXiv:2505.16968v3)。感兴趣的读者可以通过GitHub(https://github.com/GustavoStahl/CASS)和HuggingFace(https://huggingface.co/datasets/MBZUAI/cass)访问相关代码和数据集。

想象一下,你有一辆性能超强的跑车,但它只能在特定品牌的道路上行驶,换到其他品牌的道路就寸步难行。在计算机世界中,GPU(图形处理单元)就面临着类似的困境。Nvidia的GPU因其强大的计算能力成为了机器学习和科学计算的首选,但它们使用的编程语言CUDA只能在Nvidia自家的硬件上运行,就像那辆只能在特定道路上行驶的跑车。而AMD生产的GPU虽然性价比更高,但无法直接运行为Nvidia开发的程序。

这种"供应商锁定"问题严重限制了科研机构和企业的灵活性。想象一下,如果你的组织已经投入大量精力开发了基于CUDA的代码库,但现在想切换到价格更实惠的AMD硬件,你将面临一个艰巨的挑战:要么完全重写所有代码,要么继续被锁定在Nvidia的生态系统中。

为了解决这个问题,MBZUAI的研究团队开发了CASS,这是首个大规模的跨架构GPU代码转译系统,能够在源代码级别(CUDA HIP)和汇编级别(Nvidia SASS AMD RDNA3)之间进行转换。如果将代码比作菜谱,那么CASS就像一位精通多国烹饪的主厨,能够将一份意大利菜谱无缝翻译成法式烹饪指南,同时保留原始菜肴的口感和风味。

CASS数据集包含7万对经过验证的代码对,涵盖了主机(CPU)和设备(GPU)层面的代码。这就像是收集了成千上万对完美对应的双语句子,每一对都表达相同的意思,但使用不同的"语言"(Nvidia和AMD的编程语言)。利用这一丰富的"双语词典",研究团队训练了CASS系列的特定领域语言模型,在源代码转换方面达到了95%的准确率,在汇编代码转换方面达到了37.5%的准确率,大幅超越了商业基准,如GPT-4o、Claude和Hipify。

想象你有一份中文食谱,需要翻译成英文给外国朋友。普通翻译软件(类似GPT-4o或Claude)可能会做出语法正确但烹饪步骤错误的翻译,而专业烹饪翻译师(类似CASS)则能确保翻译后的食谱做出来的菜肴口感完全相同。事实上,CASS生成的代码在超过85%的测试案例中,其运行性能与原生代码相匹配,保持了原有的运行时间和内存使用特性。

为了进行严格评估,研究团队还引入了CASS-Bench,这是一个精心策划的基准测试,涵盖16个GPU领域的应用,并提供真实执行结果验证。这就像是建立了一个跨语言翻译比赛,评委会确保每个翻译不仅语法正确,还能准确传达原意。

所有数据、模型和评估工具都已作为开源资源发布,目的是推动GPU编译工具、二进制兼容性和基于LLM的硬件转译技术的进步。这相当于向全世界分享了一本珍贵的"双语字典"和翻译指南,让任何人都能在Nvidia和AMD平台之间自由切换。

一、现状分析:GPU世界中的语言障碍

图形处理单元(GPU)已经成为现代机器学习和科学计算工作负载的基础,这要归功于它们的高吞吐并行计算能力。想象GPU就像一支高效的工厂流水线团队,能同时处理成千上万个简单任务,而传统CPU则像一名能力超强但一次只能专注一项复杂任务的工匠。

Nvidia的计算统一设备架构(CUDA)已经成为GPU加速的主导编程模型,但它与专有硬件的紧密耦合引入了严重的供应商锁定问题:CUDA代码无法在非Nvidia的GPU上运行,因为它们使用不同的指令集架构。这就好比你学会了一门小众语言,但发现只有一个国家的人能听懂,想要与其他国家交流,你就必须完全重新学习新语言。

这种锁定带来了实际问题:拥有大型CUDA代码库的组织在迁移到替代平台时面临着高昂的工程成本。同时,AMD的GPU因其潜在的高性价比,正在数据中心和消费设备中获得越来越多的采用,这创造了一个日益增长的需求:如何在不完全重写软件的情况下在AMD硬件上执行旧的CUDA程序?

为了应对这一挑战,AMD推出了异构计算接口(HIP),这是内置于ROCm堆栈中的C++ GPU API,旨在镜像CUDA的功能,同时支持跨平台开发。HIP使得为Nvidia和AMD GPU维护统一代码库成为可能。像HIPIFY这样的工具可以通过将CUDA特定构造转换为HIP等效项来辅助迁移。然而,HIPIFY只在源代码级别运行,无法执行预编译的CUDA二进制文件。此外,在转换CUDA程序时,它表现出高失败率,凸显了对更可靠、更低级别的转译方法的需求。

跨供应商转译GPU汇编代码面临着指令集和编译管道的巨大差异。这就像尝试在两种完全不同的音乐记谱系统之间转换一首复杂的交响乐作品,不仅需要理解每个音符,还要保持整体和谐。Nvidia采用以nvcc为中心的专有工具链,生成PTX和低级SASS代码,而AMD则使用基于开源ROCm堆栈的GCN/RDNA架构,通过hipcc进行编译。

在汇编级别弥合这一差距对于民主化硬件计算格局、跨供应商转移特定于硬件的优化、以及启用超越源代码重写的自动化至关重要,特别是对于富含低级调优的传统CUDA代码库。CASS模型引入了首个Nvidia到AMD汇编和源代码转换的基础,侧重于正确性和对齐。虽然它尚未考虑优化,但它为未来跨GPU后端保存和适应性能关键模式的系统铺平了道路。

二、CASS:打破GPU编程的语言壁垒

为了解决缺乏跨架构GPU转换数据集的问题,研究团队引入了CASS(CUDA-AMD汇编和源代码映射),这是一个包含7万对语义对齐的CUDA-HIP源代码对及其对应的主机(CPU - x86 ISA)和设备(GPU)汇编代码的大型语料库,涵盖Nvidia(SASS)和AMD(RDNA3)平台。

想象CASS就像一本包含两种语言的对照词典,但不是将法语单词翻译成英语单词,而是将Nvidia的"方言"翻译成AMD的"方言"。每个样本包含跨供应商的功能等效低级代码,通过成功的编译和执行进行验证,实现了跨执行边界的指令级分析。

与像The Stack这样缺乏GPU对齐和可编译内容的通用代码语料库不同,CASS提供了跨两个GPU计算堆栈的完整源代码和二进制表示。为了构建CASS,研究团队开发了一个完全开源的管道,用于抓取、合成、转译(通过HIPIFY)、编译和对齐GPU代码。

这个管道就像一条复杂的工厂生产线,原材料(CUDA代码)被输入后,经过多个精确控制的处理步骤:首先清洗和准备原始代码,然后通过自动翻译工具进行初步转换,接着编译检查代码的有效性,最后将结果与原始代码精确对齐,以确保它们表达相同的功能。

团队从两个主要来源收集CUDA代码:一是从公共代码库抓取真实世界的CUDA程序,二是使用大型语言模型(Qwen2.5-Coder32B)生成各种GPU算法的新代码示例。这就像同时从专业菜谱书和创新厨师那里收集食谱,确保数据集的多样性和全面性。

使用变量增强的提示策略,研究团队能够生成丰富多样的CUDA内核代码。这个过程首先定义带有变量占位符的模板,如"为{size}X{size}网格生成布料模拟的CUDA内核,针对{optimization}优化"。然后,用预定义的值列表填充这些模板。例如,{size}可能是32、64或128,而{optimization}可能是"内存带宽"、"寄存器使用"或"多GPU扩展"。

这种方法使他们能够系统地生成各种提示,每个提示都为模板中的占位符指定不同的值,从而创建出多样化的代码示例。虽然这种方法引入了一些功能不一致,需要大量的后期处理,但它使团队能够创建丰富多样的CUDA样本。在生成的8.5万个CUDA样本中,约49.1%成功编译,最终提供了4.63万个有效文件。

收集CUDA文件后,团队进行了去重处理,以确保数据集中的所有样本都是唯一的。然后,他们使用AMD的Hipify工具将CUDA源文件转换为HIP,替换CUDA特定的API调用。无法转换的文件(约43.9%)被丢弃。一旦CUDA-HIP对可用,它们就被编译成主机和设备汇编代码,使用-Os编译标志来减少代码大小,与O3相比平均减少了9.3%的标记。

Nvidia和AMD的编译管道在如何管理主机和设备汇编分离方面存在重大差异。在ROCm中,设备二进制通常在BitCode到汇编转换期间嵌入到主机二进制中。研究团队修改了这一行为,将插入推迟到主机汇编转换为目标代码之后,从而实现了纯主机和设备汇编的独立提取,以及用于受控转换和评估的选择性重组合。

而Nvidia则不提供对其二进制注入过程的访问,设备和主机汇编保持交织,没有官方的提取或重新集成方法。为了支持主机到主机和设备到设备的转译,研究团队开发了一个基于正则表达式的过滤管道,在CUDA编译期间解开主机和设备汇编部分。

编译完成后,团队只保留了在Nvidia和AMD管道上都成功编译的样本,考虑到不对称失败。最终数据集包括匹配的CUDA-HIP源代码对、SASS-RDNA3设备汇编和主机汇编,总计6.4万个样本。

除了CUDA/HIP管道外,研究团队还利用OpenCL作为生成Nvidia到AMD映射数据集的独立管道。OpenCL作为GPU代码开发的单一"源代码",不经过上述堆栈,直接编译到汇编级别。通过这种方式,团队收集了约6千个OpenCL代码片段并将其编译为设备汇编。

这些管道总共产生了7万个对齐的汇编样本,涵盖了广泛的GPU计算领域,如机器学习、图形学和高性能计算。所有编译都在Nvidia A100 PCIe机器(SASS sm85 ISA)和AMD Radeon RX 7900 XT GPU(RDNA3 ISA)上执行,确保了硬件特定优化的捕获。

三、深入分析:CASS数据集的结构与特点

最终的指令训练数据集(CASS-Instruct)包含70,694个样本,涵盖了广泛的领域,主要集中在GPU计算和GPU相关数据结构上。该数据集还包括相应的CUDA和HIP源代码以及它们编译后的汇编表示。所有样本都经过验证,可以成功编译,并具有成对的源代码/汇编对齐。

CASS揭示了CUDA和HIP在源代码和汇编级别都存在显著的结构差异,凸显了跨架构GPU转译的固有复杂性。研究团队通过分析汇编文件的长度、语法相似性和操作码多样性来研究这些差异。

在汇编文件长度方面,AMD设备汇编在合成和Stack子集中平均比Nvidia的长两倍,而在OpenCL集中,Nvidia的设备汇编比HIP设备汇编多50%。这就像两种不同的语言描述同一事物,但一种语言需要更多的词汇来表达相同的意思。研究团队发现源代码复杂性和汇编大小之间存在指数关系,CUDA生成的输出比HIP更冗长。这凸显了随着代码复杂性增加,汇编级别转译难度也随之增加。

在代码效率和分析方面,汇编准确性在不同领域差异很大:在数学、数据结构和图形任务中为0%,在线性代数和内存操作中为25-50%,在物理模拟中高达100%,凸显了保持低级语义的挑战。尽管如此,转译的代码在执行方面与原始代码非常接近:内存使用偏差小于±0.3%,执行时间在±11.8%以内,超过85%的样本在这两个指标上都在±5.6%范围内。

语法相似性分析显示,HIP和CUDA汇编在设备代码方面表现出较低的语法相似性,在主机代码方面表现出中等相似性,尤其是在OpenCL和Stackv2子集中。相比之下,源代码转译,特别是在合成子集中,显示出高度重叠,凸显了表面级语法在源代码中比在编译后的汇编表示中保留得更好。

在操作码多样性方面,研究团队注意到张量操作在CUDA和HIP汇编中都占主导地位,特别是在设备代码中,内存相关指令如mov和call出现频率最高。此外,HIP操作码如s_mov_b32和v_add_co_u32被广泛使用,反映了AMD ISA独特的低级向量和内存操作,而Nvidia则以其自己的常见指令变体如movq、call和jmp为主,具有更大的主机端集成。

两个堆栈共享常见的控制和内存操作(如mov、test),但HIP提供了对GPU内部更细粒度的访问,揭示了对并行性的更深入可见性。合成子集强调内存导向指令,与LLM驱动的模板优化一致。进一步分析显示,尽管后端差异,Nvidia和AMD在设备和主机级别上共享语义对齐的操作码分布。

为了进行严格评估,研究团队还引入了CASS-Bench,这是一个涵盖16个以GPU为中心的领域的40个样本评估套件,每个领域由1-5个精心策划的提示代表。对于每个提示,团队首先使用Claude-3.7生成CUDA实现,然后在Nvidia硬件上编译和执行以获取参考输出,然后提示Claude-3.7生成相应的AMD代码。如果由于编译错误、格式差异或随机生成器差异导致输出不匹配,则重新生成AMD代码。只有具有手动验证的输出等效性的样本才被包括在内。所有最终的Nvidia-AMD对都通过管道处理,以提取对齐的主机和设备汇编。

四、CASS模型:跨架构代码转译的新标准

研究团队通过对不同参数规模的Qwen2.5-Coder模型进行指令监督微调,对CASS数据集进行了评估。他们开发了两种变体:一种用于汇编转译,另一种用于源代码转译,并将这些模型与专有和开源基线进行了比较,包括更大规模的系统。

为了确保输入样本适合LLM的16K标记上下文窗口,研究团队通过删除冗余空白和注释来规范化CUDA汇编代码,这减少了约15%的标记计数。由于HIP汇编代码对空白更改很敏感,因此没有对其进行预处理。

团队在4个A100 GPU上以1.5B、3B和7B参数规模微调了Qwen2.5-Coder模型,使用批量大小为4,梯度累积为32(有效批量大小为512),学习率为1×10^-5。相对激进的学习率是因为数据集与模型预训练语料库的分布差异。

训练采用了DeepSpeed与优化器状态分片来最大化硬件效率,实现了98%的GPU利用率。此外,团队还结合了Liger Kernel和Paged Adam优化器,以加速训练并更有效地管理内存,使用LLaMA-Factory实现了所有这些优化。

所有模型都使用16K标记上下文窗口进行训练。在推理时,团队应用了RoPE外推法,支持高达32.7K标记。推理效率很高,处理一个16K标记样本约需56秒。

对于源代码和汇编转译,LLM生成的代码(HIP源代码或主机/设备汇编)被编译和执行。然后,将结果与CASS-Bench的基准进行比较,以验证功能正确性。

在汇编到汇编性能方面,所有基线,包括专有和大型开放模型,都失败,准确率为0%,只有Qwen2.5-Coder-32B达到了25%。ZLUDA,一个运行时级系统,尽管直接在编译的二进制文件上操作,但只达到了2.5%的汇编准确率,这可能归因于其与RNDA1的兼容性。相比之下,CASS模型达到了高达37.5%的准确率,突显了数据集在现有工具和模型中缺失的关键汇编级知识。

在源代码到源代码性能方面,为了进一步验证数据集的实用性,评估了源代码转译性能。这项任务与许多专有模型的预训练目标更一致,这反映在它们相对较强的性能中(从80%到90%)。尽管如此,即使是最小的CASS模型(1.5B)也显著优于所有基线,达到了90%的准确率。7B变体显示了出色的最先进性能,达到了95%的准确率。尽管CUDA数据集完全由Hipify转译,并且只保留了语义对齐的样本,但模型的表现超过了Hipify 7.5%。

消融研究显示,仅使用Stack数据产生了17.5%的汇编准确率。添加合成数据将其提高了+12.5%,凸显了其在学习低级模式中的作用。OpenCL增加了+2.5%,提供了互补覆盖,而RoPE外推将准确率提高到37.5%,通过扩展上下文容量。

在领域分析方面,研究显示汇编准确率在不同领域差异很大:在数学、数据结构和图形任务中为0%,在线性代数和内存操作中为25-50%,在物理模拟中高达100%,反映了保持低级语义的挑战。尽管如此,转译的代码在执行方面与原始代码非常接近:内存使用偏差小于±0.3%,执行时间在±11.8%以内,超过85%的样本在这两个指标上都在±5.6%范围内。

与其他大型语言模型相比,CASS-7B在忠实转译CUDA到HIP方面表现优异。例如,在一个案例中,CASS-7B正确转译了CUDA代码,同时保留了原始程序中的确切字符串常量,包括输出格式字符串中的CUDA标签。在另一个例子中,CASS-7B保留了经典的CUDA风格内核启动语法,同时确保生成的代码通过正确包含所需的HIP头文件保持可编译。这展示了对源代码的高度结构保真度,对熟悉标准CUDA约定的开发人员特别重要。

五、局限性与未来展望

尽管在汇编转译方面取得了最先进的成果,但当前性能由于复杂或代表性不足的领域中的有限准确性而不足以用于生产。扩展类别多样性对解决这一问题至关重要。数据集目前只涵盖每个供应商的一个主机/设备对(RTX 4090和RX7900),限制了跨GPU架构的泛化能力。需要更广泛的架构表示来支持实际部署。

此外,数据集大小被最小化以适应16K标记的上下文窗口,排除了许多特定于供应商的低级优化。纳入这些将需要未来具有更大上下文容量或更高级分块和注意力策略的模型。

未来的工作方向包括扩展数据集以涵盖更多GPU架构和优化模式,提高汇编级转译的准确性,以及开发能够在保持功能的同时跨架构转移性能关键优化的技术。此外,团队计划研究如何结合符号推理和神经方法来改进低级代码生成和验证。

六、结论:打破GPU计算的围墙

CASS代表了跨架构GPU代码转译领域的重大进步,为源代码到源代码(CUDA到HIP)和汇编到汇编(SASS到RDNA3)映射提供了首个大规模数据集和模型套件。通过7万对对齐的源代码和汇编对,研究团队创建了一个独特的资源,使研究人员和开发人员能够探索GPU代码可移植性的新前沿。

CASS模型家族在源代码转译方面达到了95%的准确率,在汇编转译方面达到了37.5%的准确率,大幅超越了现有的工具和大型语言模型。重要的是,转译的代码保留了功能行为,在内存使用和运行时方面与原生执行相匹配。

通过开源CASS数据集、模型和基准,研究团队为未来在编译器工具、硬件互操作性和性能感知代码生成方面的研究奠定了基础。这项工作不仅推进了技术边界,还有可能彻底改变GPU计算格局,使组织能够在供应商之间更自由地移动,最终促进更具竞争力和创新的高性能计算生态系统。

随着人工智能和科学计算对GPU资源的需求持续增长,CASS这样的工具将在确保计算灵活性和降低依赖特定供应商的风险方面发挥关键作用。虽然仍有挑战需要克服,特别是在汇编级别的转译方面,但这项研究代表了朝着更开放、更互操作的GPU计算未来迈出的重要一步。

来源:至顶网一点号

相关推荐