智源打造基于Triton的大模型算子库，助力AI芯片软硬件生态建设

2024年大模型进入了新的发展阶段，AI全领域开启了更为迅猛的量变积累。一方面，模型突破了模态的隔离，文本、语音、视觉等各种形式之间产生的丰富的结合，大大增加了模态的多样性；同时，模型参数量从百亿、千亿级膨胀到万亿级，训练数据量从TB级达到PB级，上下文令牌数量也从几千增长到百万级，计算规模空前庞大；此外，算法结构的创新也带来了MoE、模型量化、定制算子等更加复杂的计算需求。

而在硬件设备的层面，由英伟达主导的CUDA生态以SIMT编程模型和CUDA编程语言为核心，从高性能算子库、开发工具链和GPU驱动等各个层次全面协同，建立了一套完整的体系，并长期占据着高性能计算的领先地位。虽然有生态竞争者——诸如OpenCL和ROCm等——在试图挑战和替代CUDA，但是无论在推广程度、使用体验还是计算性能方面都存在一定差距。新兴的各种AI芯片硬件架构不一、指令集不一、各自有自己的AI编译器，算子库也各自实现，整体呈现十分割裂且难以强大的生态。哪怕有部分AI芯片厂商不断模仿、跟进和对齐CUDA，然而受限于芯片架构的差异和底层的封闭属性，厂商的生态适配仍旧面临开发难度大、任务重以及各自为战的难点。

算力需求增长而资源供应紧张的局面下，多元芯片的开源共建无疑是撼动CUDA地位、打造全新格局的机遇。而作为软硬件衔接的关键环节，编译技术这一层次则成为构建统一生态的入口。

为降低大模型新算法的开发门槛、加速芯片架构的创新，智源研究院在今年的智源大会上发布了使用Triton语言实现的算子库FlagGems，以大模型的计算需求为导向，面向多元芯片，借助Triton编译器的开源和轻量级优势，提供了一套易适配、高性能的算子实现，以推动基于Triron的统一、开源的软硬件生态。

1. 算子库技术选型

在编译技术的多条路线中，统一的中间语言、统一的算子接口、统一的开源算子库都经过了一定的探索和实践。

统一的中间语言：提供更自由的表达能力和更灵活的优化空间，但其设计研发显著依赖硬件架构相关的底层信息，要求厂商的深度参与和多元芯片的高度协作
统一的算子接口：对上层框架能够保持良好的一致性，对下层则要求芯片厂商各自开发算子库，无法确保一致的算子特性，也无法配合使用框架的图优化技术
统一的开源算子库：能够在厂商之间做到源码共享，省去重复开发的成本且保障一致的算子实现，并且能够由厂商对编译器的个性化适配来发挥硬件相关的性能优势

FlagGems采用开源算子库的技术路线，基于OpenAI推出的Triton编程语言及编译器，以eager模式接入PyTorch框架。相比使用Triton的Inductor技术，FlagGems不仅具有同等的简便易用的优点，而且能够更加细致深入地提升算子性能，从而提高模型的训练和推理吞吐。

2. 算子库现状

FlagGems以Llama2、Llava2等十余个国内外热门开源大模型为范例，通过profiler工具监测和分析模型的算子调用情况，汇总到一百余个算子。模型包含了语言、视觉、语音、多模态等各种类型，支持前向推理和反向训练的不同场景，基本能够代表主流大模型的应用实践，覆盖大部分的算子调用需求。

FlagGems以PyTorch框架在eager模式下调用的原生CUDA算子为基线，展现出了强势的性能和可观的优化空间。在访存密集型算子上，FlagGems有部分算子性能超越CUDA算子；在计算密集型算子上，FlagGems基本追平CUDA算子；在融合算子上，FlagGems全面优于CUDA算子。

FlagGems目前已经支持Bert和Llama2两个模型的单机单卡推理，支持Aquila模型的多卡分布式训练。

3. 算子库创新点

3.1 运行时优化

在模型运行的过程中，影响性能的不止有GPU上核函数的执行效率，还有CPU端侧的运行时开销。Triton语言在定义算子功能时会包装Autotuner或Heuristics装饰器，其中Autotuner使用自动调优的方法为核函数选择编译参数，而Heuristics则使用启发式的方法计算编译参数。编译器在执行函数时需要调用所有装饰器的执行序列，复杂的包装会为运行时带来较大的时间开销。

为了尽可能压缩运行时的过程，FlagGems设计了LibEntry机制，使用独立维护的缓存来获取和调用核函数，绕过其他运行时装饰器的多层调用。在算子首次运行时，仍然会沿着Triton的原调用路径进行自动调参和启发式计算来取得编译参数，获得的核函数加入缓存；而当算子再次调用到键值相同的核函数时，运行时机制可以直接从缓存取得核函数，快速进行调用运行。

LibEntry优化技术能够将CPU端的运行时开销降低约70%，在运行时开销为主要瓶颈的小规模算子上获得良好的收益，达到与PyTorch ATen持平的CPU端性能。

3.2 代码生成

Triton算子实现以PyTorch框架定义的张量作为输入和输出，依托Python原生特性无缝融入PyTorch生态。然而Triton函数的操作对象是张量指针，PyTorch的输入输出本质是张量视图，二者的差异给FlagGems的代码编写带来了挑战。

Triton函数中的张量指针与张量在存储上的布局密切相关，使用偏移量访问存储和读写数据；而张量视图则是存储的变形，更关心张量的形状尺寸，对于张量布局和连续性感知不强。在Triton函数处理连续张量时，数据的寻址和访存都比较简单，可以无视指针与视图的差异；而当PyTorch向算子输入非连续张量时，Triton就不能以与连续存储相同的方式来获取数据，强制连续化可能增加的拷贝开销，使得张量步长成为函数中必要的输入信息。

对于高维张量输入，步长的数量是不固定的，鉴于Triton无法支持变长元组参数，算子库需要对所有维度分别编写函数代码，这对开发者造成了一定的负担。FlagGems利用pointwise运算的相似性以及偏移量计算代码的模式化特点，使用自动代码生成的方式处理高维张量的非连续访存问题。

自动代码生成机制能够使用模板生成外层包装函数和Triton核心函数，自动地处理张量广播、索引空间计算、并行参数计算、地址偏移计算等任务。对于经过转置、切片、跨步等处理的非连续张量，自动代码生成的GPU性能可以显著超过强制连续化的方法、达到与PyTorch ATen持平的表现。

下面以两个张量相加为例，分别展示CUDA写法、Triton写法和自动代码生成写法，其中自动代码生成写法，通过简单加法逻辑表示，即可生成Triton核心函数。