作者/机构: Hanchen Ye (University of Illinois Urbana-Champaign), Deming Chen (Inspirit IoT, Inc., University of Illinois Urbana-Champaign)
本文针对数据流架构在深度学习工作负载(特别是大型语言模型LLM)中面临的内存瓶颈和性能优化挑战,提出了一个名为StreamTensor的编译器框架。现有方法在处理内核间相关性、外部存储器访问管理和缓冲区优化方面存在困难。
核心问题与研究目标:
LLM的自回归解码阶段是内存密集型的,对内存高效的架构提出了更高要求。数据流架构通过在计算内核之间流式传输中间结果来提高效率,但其编程和优化非常复杂。当前的设计范式需要大量手动操作来设计数据流调度、布局转换器、DMA和FIFO,这一过程不仅容易出错,而且难以扩展到大型加速器。具体来说,现有方法存在以下缺陷:
1. 内核间相关性:流水线执行的内核需要在延迟和并行化策略上进行平衡,以避免资源浪费,但FIFO的严格顺序访问为切分(Tiling)、置换(Permutation)和向量化(Vectorization)的协同优化带来了挑战。
2. 外部存储器访问:现有编译器大多假设数据能完全放入片上内存,忽略了实际应用中DMA如何与内核执行重叠、数据布局如何匹配流模式等复杂问题。
3. 基于流的内核融合:在生产者和消费者内核具有不兼容流布局时,需要自动检查兼容性、生成最小的流布局转换器并确保其满足片上内存限制,这对于手动设计是不现实的。
4. FIFO大小确定:不合理的FIFO大小会导致溢出、下溢、流水线停顿甚至死锁。现有手动或基于仿真的方法缺乏可扩展性。
本文的创新点与主要贡献:
为了解决上述问题,本文提出StreamTensor框架,旨在将数据流加速器的设计范式从手动设计转变为编译器自动生成和优化(如下图所示)。
其主要贡献如下:
1. 提出StreamTensor框架:这是首个从PyTorch到设备(PyTorch-to-device)的数据流编译器,能够自动生成基于流的数据流加速器及其对应的运行时系统。
2. 提出迭代张量(itensor)类型系统:首次提出一种itensor类型,用于系统性地编码流信息。该类型系统为基于流的内核融合和数据流组件生成奠定了基础,提高了数据流加速器设计的可扩展性和生产力。
3. 提出三个层次化的设计空间:定义了张量切分空间、内核融合空间和资源分配空间,以算法化和层次化的方式覆盖了数据流架构的复杂设计空间。并为每个空间提出了探索算法以减少资源利用并提高延迟和吞吐量。
4. 提出基于分段函数的令牌行为模型和LP解决方法:将数据流加速器的FIFO大小确定问题转化为一个调度问题,并提出一个线性规划(LP)算法来解决此问题,从而在避免死锁的同时减少资源利用。
5. FPGA评估:在FPGA平台上对LLM进行评估,与最先进的FPGA LLM加速器和GPU相比,StreamTensor实现的延迟分别降低了0.76倍和0.64倍,能效比GPU高1.99倍。
数据流架构作为冯·诺依曼架构的替代方案。数据流架构,相较于NVIDIA H100【索引20,Nvidia hopper h100 gpu: Scaling performance,2023,IEEE Micro】和Google TPUv4【索引33,Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings,2023,Proceedings of the 50th Annual International Symposium on Computer Architecture】等冯·诺依曼风格的架构,正越来越多地被采用和研究,以克服新兴AI应用(如大型语言模型LLM)中的“内存墙”问题。由于LLM的自回归特性,其解码阶段高度受内存限制,需要更高效的内存架构。AMD Versal【索引24,Xilinx adaptive compute acceleration platform: Versaltm architecture,2019,Proceedings of the 2019 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】、Sambanova SN40L【索引43,SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts,2024,arXiv preprint arXiv:2405.07518】和IBM AIU【索引12,Meet the ibm artificial intelligence unit,2022,IBM Research.[Online]】是具有可重构数据流架构的商用AI加速器;许多研究【索引14,Understanding the potential of fpgabased spatial acceleration for large language model inference,2024,ACM Transactions on Reconfigurable Technology and Systems】【索引40,Stream-dataflow acceleration,2017,Proceedings of the 44th Annual International Symposium on Computer Architecture】【索引44,Plasticine: A reconfigurable architecture for parallel paterns,2017,ACM SIGARCH Computer Architecture News】也证明了数据流架构在延迟和能效方面的优势。
数据流加速器的典型计算模式。如图1所示,一个数据流加速器包含以下片上组件:(1) 内核(Kernel):使用并行处理器(如脉动阵列)计算一个算子或粗粒度任务,并为输入和输出提供流接口。(2) 令牌(Token):内核间通信的原子元素。(3) 先进先出队列(FIFO):缓存累积的流令牌,以平衡生产者和消费者的不同速率,避免死锁或不必要的内核停顿。(4) 流布局转换器(Stream Layout Converter):通过本地乒乓缓冲,动态转换流布局以适应生产者和消费者内核不同的计算模式。(5) 直接内存访问(DMA):与外部内存通信,并将内存映射接口与流接口相互转换。内核本身的设计可以采用动态调度的数据流电路【索引31,Dynamically scheduled high-level synthesis,2018,Proceedings of the 2018 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】,或不同的数据流策略(如输入静态)以实现高效的片上数据重用【索引17,Eyeriss: An energy-efficient reconfigurable accelerator for deep convolutional neural networks,2016,IEEE journal of solid-state circuits】,这些概念与本文讨论的数据流架构和加速器在概念上是正交的。
数据流架构的核心思想。数据流架构的关键思想是通过片上FIFO在内核之间流式传输中间结果,而不是频繁访问外部内存。例如,在图1(b)中,Kernel0产生的中间结果直接流向Kernel1和Converter0,而不像图1(a)那样经过外部内存。本文将这种内核间流式传输称为基于流的内核融合。此外,如图1(c)所示,数据流加速器的调度允许Kernel1和Converter0在Kernel0完成之前就开始执行。这种重叠执行可以显著提高整体吞命量和延迟。
当前数据流加速器编程的范式。如图2所示,数据流加速器通常分为特定应用加速器和领域特定加速器(DSA)两类,其编程方式也因此有所不同。
特定应用加速器编程。在此类别中,数据流组件和调度是为单个应用量身定制的。因此,编程通常指架构和微架构的设计或生成。传统上,使用硬件描述语言(HDL)、高级综合(HLS)和像Chisel【索引6,Chisel: constructing hardware in a scala embedded language,2012,In Proceedings of the 49th Annual Design Automation Conference】这样的元HDL来完成【索引13,xpilot: A platform-based behavioral synthesis system,2005,SRC TechCon】【索引19,SODA: Stencil with optimized dataflow architecture,2018,2018 IEEE/ACM International Conference on Computer-Aided Design (ICCAD)】【索引49,FlowGNN: A dataflow architecture for real-time workload-agnostic graph neural network inference,2023,2023 IEEE International Symposium on HighPerformance Computer Architecture (HPCA)】【索引63,DNNBuilder: An automated tool for building high-performance DNN hardware accelerators for FPGAs,2018,2018 IEEE/ACM International Conference on Computer-Aided Design (ICCAD)】。最近,加速器设计语言(ADL)【索引15,Allo: A Programming Model for Composable Accelerator Design,2024,Proceedings of the ACM on Programming Languages】【索引23,Type-directed scheduling of streaming accelerators,2020,Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation】【索引53,Fleet: A framework for massively parallel streaming on FPGAs,2020,Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems】的出现提高了生产力,引入了类型系统和原语来描述计算、内存布局和数据流调度。如图2所示,现有解决方案需要手动将应用转换为数据流调度和组件,然后传递给HLS、元HDL转译器或供应商EDA工具进行硬件生成。尽管ADL和HLS框架集成了设计空间探索(DSE)【索引2,An MLIR-based compiler flow for system-level design and hardware acceleration,2022,Proceedings of the 41st IEEE/ACM International Conference on Computer-Aided Design】【索引9,Stateful dataflow multigraphs: A data-centric model for performance portability on heterogeneous architectures,2019,Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】【索引34,Automatic generation of efficient accelerators for reconfigurable hardware,2016,ACM SIGARCH Computer Architecture News】【索引35,Spatial: A language and compiler for application accelerators,2018,Proceedings of the 39th ACM SIGPLAN Conference on Programming Language Design and Implementation】【索引60,Scalehls: A new scalable high-level synthesis framework on multi-level intermediate representation,2022,2022 IEEE international symposium on high-performance computer architecture (HPCA)】【索引62,An Optimizing Framework on MLIR for Efficient FPGA-based Accelerator Generation,2024,2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】,但这些工作主要集中在优化单个内核上。
数据流DSA编程。DSA旨在高效执行特定类别应用或特定领域的计算,而非通用处理器。DSA通常使用类粗粒度可重构架构(CGRA)实现【索引24,Xilinx adaptive compute acceleration platform: Versaltm architecture,2019,Proceedings of the 2019 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】【索引40,Stream-dataflow acceleration,2017,Proceedings of the 44th Annual International Symposium on Computer Architecture】【索引43,SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts,2024,arXiv preprint arXiv:2405.07518】【索引44,Plasticine: A reconfigurable architecture for parallel paterns,2017,ACM SIGARCH Computer Architecture News】,其中片上资源被重新配置以实现不同的数据流设计。现代DSA使用C/C++原语【索引24,Xilinx adaptive compute acceleration platform: Versaltm architecture,2019,Proceedings of the 2019 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】【索引67,CHARM: C omposing H eterogeneous A ccele R ators for M atrix Multiply on Versal ACAP Architecture,2023,Proceedings of the 2023 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】【索引68,CHARM 2.0: Composing Heterogeneous Accelerators for Deep Learning on Versal ACAP Architecture,2024,ACM Transactions on Reconfigurable Technology and Systems】或领域特定语言(DSL)如Spatial【索引35,Spatial: A language and compiler for application accelerators,2018,Proceedings of the 39th ACM SIGPLAN Conference on Programming Language Design and Implementation】、Halide【索引46,Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines,2013,Acm Sigplan Notices】和TVM【索引16,{TVM}: An automated {End-to-End} optimizing compiler for deep learning,2018,13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】进行编程,以生成领域优化的代码。如图2所示,开发人员必须使用这些DSL或API手动将应用转换为逻辑组件。然后,软件编译器将它们映射到物理资源并生成最终的片上执行二进制文件。虽然这些DSL通常提供数据流内核的自动调优功能,但它们主要关注优化单个内核而非整个数据流应用,从而留下了巨大的性能提升空间。
缺陷1:内核间相关性。先前的工作【索引60,Scalehls: A new scalable high-level synthesis framework on multi-level intermediate representation,2022,2022 IEEE international symposium on high-performance computer architecture (HPCA)】【索引61,HIDA: A Hierarchical Dataflow Compiler for High-Level Synthesis,2024,Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1】表明,内核间的相关性会影响加速器性能。由于内核以流水线方式执行,它们的延迟必须平衡以获得最佳吞吐量。此外,通过缓冲区连接的内核需要对齐其并行化策略以避免低效的内存使用。然而,先前的工作只考虑了支持内存映射访问的乒乓缓冲区。FIFO的限制更强,因为数据必须按顺序推入/拉出。这为每个内核带来了以下挑战:(1) 切分(Tiling):选择能够实现流式传输、最小化本地缓冲并保持内存效率的切分大小。(2) 置换(Permutation):重排循环以减少数据流式传输期间的内存利用率。(3) 向量化(Vectorization):选择展开策略以平衡延迟并提高流式传输效率。这些决策在内核之间相互依赖,使得全局优化对于分析模型或手动设计来说具有挑战性。
缺陷2:外部存储器访问。大多数现有编译器【索引2,An MLIR-based compiler flow for system-level design and hardware acceleration,2022,Proceedings of the 41st IEEE/ACM International Conference on Computer-Aided Design】【索引8,Stream-HLS: Towards Automatic Dataflow Acceleration,2025,Proceedings of the 2025 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】【索引60,Scalehls: A new scalable high-level synthesis framework on multi-level intermediate representation,2022,2022 IEEE international symposium on high-performance computer architecture (HPCA)】【索引61,HIDA: A Hierarchical Dataflow Compiler for High-Level Synthesis,2024,Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1】【索引62,An Optimizing Framework on MLIR for Efficient FPGA-based Accelerator Generation,2024,2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】【索引65,Polsca: Polyhedral high-level synthesis with compiler transformations,2022,2022 32nd International Conference on Field-Programmable Logic and Applications (FPL)】假设所有数据都适合片上内存,这对于大型应用来说是不现实的。当涉及片外内存时,每个DMA必须解决以下问题:(1) 如何将内存访问与内核执行重叠?(2) 哪种数据布局最适合流模式?(3) 如何打包/向量化数据以最大化带宽?这些问题需要复杂的模式分析,并且在手动处理时容易出错。DMA的设计也与内核的切分和调度紧密耦合,进一步增加了复杂性。
缺陷3:基于流的内核融合。基于流的内核融合的目标是将所有中间结果在片上流式传输,将外部内存的使用限制在输入和输出上。然而,生产者和消费者内核由于计算模式不同,通常具有不兼容的流布局。这需要:(1) 检查内核间的布局兼容性。(2) 生成最小的即时流布局转换器。(3) 确保转换器能适应可用的片上内存。这些步骤涉及复杂的模式分析,并需要系统的全局视图,使得手动解决方案不切实际。
缺陷4:FIFO大小确定。如图1所示,如果Kernel1比Converter0慢,FIFO可能会溢出或下溢,导致停顿级联并最终死锁。尽管存在动态调度解决方案【索引32,Buffer placement and sizing for high-performance dataflow circuits,2021,ACM Transactions on Reconfigurable Technology and Systems (TRETS)】,粗粒度加速器仍然依赖手动确定大小【索引14,Understanding the potential of fpgabased spatial acceleration for large language model inference,2024,ACM Transactions on Reconfigurable Technology and Systems】【索引15,Allo: A Programming Model for Composable Accelerator Design,2024,Proceedings of the ACM on Programming Languages】,这对于大量的FIFO来说是不可扩展的。最近的一种自动化方法【索引30,Automated Buffer Sizing of Dataflow Applications in a High-Level Synthesis Workflow,2024,ACM Transactions on Reconfigurable Technology and Systems】使用仿真来确定FIFO大小,但这种方法耗时且缺乏可扩展性。
StreamTensor框架概述。StreamTensor是一个编译框架,旨在将PyTorch模型转换为优化的数据流实现。它构建在MLIR【索引36,MLIR: Scaling compiler infrastructure for domain specific computation,2021,2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)】编译框架之上。StreamTensor的整体架构如图4所示。编译过程从Torch-MLIR【索引55,Torch-MLIR: A compiler for the PyTorch ecosystem,2021,https://github.com/llvm/torch-mlir】提供的PyTorch模型开始,经过几个阶段。首先,张量操作被转换为使用MLIR内置的线性代数(Linalg)操作的结构化中间表示(IR)。然后,这个IR通过MLIR 的Linalg pass(如逐元素操作融合)进行优化。接着,StreamTensor应用设计空间探索(DSE)算法,根据计算模式等因素确定最优的切分策略,包括切分大小、展开因子和置换。然后,Linalg IR被转换为数据流IR,其中计算被组织为分层任务。所有数据流组件,包括DMA、流布局转换器和FIFO,都在此阶段生成。关键优化也在这里执行,例如基于流的内核融合以最小化外部内存访问,以及FIFO大小调整以平衡生产者-消费者执行。在最后阶段,StreamTensor生成特定于硬件的代码和主机运行时。StreamTensor处理内存分配、流连接和指令具象化,这使得像HLS这样的供应商编译器能够生成目标数据流架构。同时,它生成主机运行时代码,管理数据传输、内核执行以及主机CPU和数据流加速器之间的同步。
类型系统用于IR验证和优化。StreamTensor引入了一个类型系统,以实现IR的高效验证和优化。通过StreamTensor专用的类型和操作验证器,该类型系统有助于确保在应用任何转换遍后IR的有效性。
传统张量类型的局限性。传统上,张量类型编码了数据类型和表示其形状的整数列表【索引16,{TVM}: An automated {End-to-End} optimizing compiler for deep learning,2018,13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】【索引28,Graphene: An ir for optimized tensor computations on gpus,2023,Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】【索引36,MLIR: Scaling compiler infrastructure for domain specific computation,2021,2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)】【索引46,Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines,2013,Acm Sigplan Notices】。张量可以以内存映射的方式访问,例如,可以根据偏移量和形状提取或插入一个切片。然而,数据流内核通过FIFO进行通信,FIFO强制执行严格的访问顺序并遵循流式访问模式,而非内存映射模式。因此,传统的张量类型可能无法确保数据流通信的正确性。即使生产者和消费者共享相同的张量类型,流访问顺序也可能不明确,导致意外行为。例如,在Graphene【索引28,Graphene: An ir for optimized tensor computations on gpus,2023,Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】中,张量类型仅编码内存映射布局。结果,当生产者以行主序生成流而消费者期望列主序时,即使两者操作于相同的张量类型,也会导致数据解释错误和逻辑损坏。因此,尽管现有解决方案足以进行Linalg级别的优化(如切分),但它们在生成数据流组件和应用数据流优化方面是易错且不可扩展的。
迭代张量(itensor)类型。为了解决这个问题,我们提出了一种新的itensor类型,它明确编码了流布局信息,使得基于类型的验证和优化成为可能且高效。图5展示了从类型为tensor<8x8xf32>的同一个张量转换而来的三个itensor示例。要将张量转换为itensor,我们首先将其划分为相同的张量切片或向量。例如,在图5(b)中,张量被划分为八个形状为4x2的张量切片。这些切片随后在定义的迭代空间(通常是嵌套循环)内被迭代访问。迭代空间由两个列表定义:行程计数(tripcounts)和步长(step sizes)。在图5(b)中,迭代空间为[4,2]*[2,4],它产生迭代索引[0,0], [0,4], [2,0], [2,4]等。从迭代空间到数据空间的映射由一个仿射映射指定——例如,图5(b)中的(d0,d1)->(d1,d0),它转置了迭代索引。因此,数据访问索引变为[0,0], [4,0], [0,2], [4,2]等,反映了图5(b)中所示的这种转置。在itensor中,张量切片可以被多次访问,其模式明确编码在迭代映射中。例如,在图5(c)中,迭代空间为[4,2,2]*[2,1,4],迭代映射为(d0,d1,d2)->(d2,d0),其中维度d1不对应任何数据维度。当d1从0迭代到1时,所有次要维度(如d2)都会被重新迭代。因此,相应的数据维度(例如行维度)也会被重新访问,为形状为4x2的张量切片产生如[0,0], [4,0], [0,0], [4,0], [0,2]等索引。通过在itensor类型中编码元素形状、迭代空间和迭代映射,可以唯一确定数据流内核的流模式。当生产者和消费者的itensor类型匹配时,可以安全地在它们之间建立流式通信(图5的Case1)。否则,必须在它们之间插入一个流布局转换器(图5的Case2),并且用于布局转换的最小乒乓缓冲区大小可以从itensor类型中分析推断出来。布局转换器的生成细节将在5.2.1节中讨论。由于缺乏流信息,现有的基于张量的类型系统不足以支持基于流的内核融合,限制了它们在基于流的数据流优化中的可用性。
流类型(Stream Type)。在传统的张量编译器中,高级张量IR必须被“缓冲化”(bufferized)为低级内存/缓冲区IR,以支持低级优化和代码生成。遵循这一惯例,我们提出了流类型,它是在缓冲化过程中从itensor类型降级而来的。与不可变的itensor对象不同,流对象代表硬件FIFO,并支持通过流读写等操作进行突变。流类型仅编码数据类型和FIFO深度,而流布局信息在缓冲化过程中被剥离。因此,数据流组件的生成和优化必须在itensor级别的IR上完成。缓冲化之后,流IR保留用于较低级别的硬件/运行时优化和代码生成。
itensor和stream操作。基于第3.1节的类型系统,StreamTensor引入了itensor和stream操作来表示不同的数据流行为。此外,还引入了结构操作来表示数据流加速器的多级层次结构,这些结构操作由itensor和stream级别的IR共享。
迭代张量操作。表1列出了itensor级别的完整操作集。总体而言,这些操作的含义不言自明;我们重点介绍那些语义不太明显的操作。itensor_write可以从概念上理解为向FIFO中写入或推送一个元素。它是一个目标携带(destination-carried)操作,其目标是通过dest操作数传递的itensor。例如,迭代地写入图5(b)中的itensor(称为itensor(b))可以表示为:
%empty = itensor_empty() : itensor(b)
%res0 = scf.for 0 to 8 step 2 iter_args={%arg0 = %empty} {
%res1 = scf.for 0 to 8 step 4 iter_args={%arg1 = %arg0} {
%value = ... : tensor<4x2xf32> // %value is defined
%output = itensor_write %value into %arg1 : ...
scf.yield %output : itensor(b)
} : itensor(b)
scf.yield %res1 : itensor(b)
} : itensor(b)
这里,scf是MLIR内置的结构化控制流方言,包括for循环。scf.for也是目标携带的,其中%empty作为参数传递,并通过itensor_write迭代地推送。最终,%res0作为最终结果返回。相比之下,itensor_read表示从FIFO中拉取一个元素。例如,读取itensor(b)可以表示为:
%source = ... : itensor(b) // %source is defined
scf.for 0 to 8 step 2 {
scf.for 0 to 8 step 4 {
%empty = tensor.empty() : tensor<4x2xf32>
%value = itensor_read %source init %empty : ...
... = ... %value ... // %value is used
itensor_converter包含一个本地乒乓缓冲区,用于执行即时的流布局转换。例如,在图5的Case1中,源和目标共享相同的itensor类型,可以通过FIFO连接。在Case2中,它们不同,因此必须插入一个转换器。需要一个最小为8x2的乒乓缓冲区来适应这两种流布局。当源向ping缓冲区写入时,目标从pong缓冲区读取两次,然后它们交换。
流操作。表2列出了流级别的操作。这些操作大多不言自明;我们重点强调与itensor操作的主要区别。如3.1.3节所讨论,流对象是可变的,不再使用目标携带语义。FIFO的推入和拉出可以写成:
%stream = stream() : stream<f32, depth: 32>
scf.for 0 to 8 step 2 {
scf.for 0 to 8 step 4 {
%value = ... : f32 // %value is defined
stream_write %value into %stream : ...
}
}
scf.for 0 to 8 step 2 {
scf.for 0 to 8 step 4 {
%value = stream_read %stream : ...
... = ... %value ... // %value is used
}
}
注意,整个过程中都使用同一个%stream,而没有像itensor的目标携带风格那样创建新的副本。流IR对于代码生成更高效,但使得定义-使用分析变得复杂。因此,itensor更适合高级数据流优化。流操作的正确性是通过构造来保证的,因为它们是从itensor操作降级而来的,而itensor操作由itensor类型系统严格验证。
结构操作。当itensor和stream操作对行为进行建模时,结构操作则对层次结构进行建模。表3列出了StreamTensor中所有的结构操作。kernel操作代表一个数据流内核(如图1所示),包含一个task操作的图。它以张量作为输入/输出,这些张量在边界处与itensor进行转换。这些隐式转换充当DMA的角色。内核内部使用片上流式传输,而内核之间使用外部内存。例如:
%source = ... : tensor<8x8xf32> // %source is defined
%result = kernel(%arg : itensor<b> = %source : tensor<8x8xf32>) {
= ... %arg ... // %arg is used
%output = ... : itensor<c> // %output is defined
yield %output : itensor<c>
} : tensor<8x8xf32>
通过在内核边界进行转换,我们避免了在内核融合期间显式处理DMA,从而提高了转换效率和可分析性。相比之下,task操作是透明的,在其边界不转换类型。它代表内核内的一个数据流任务,并且可以嵌套以实现层次化的数据流设计。在itensor级别,task是目标携带的,其中输出通过inits写入目标,提高了定义-使用分析的效率。例如:
%empty = ... : itensor(b)
%result = task @example inits={%arg = %empty} {
%value = ... : tensor<4x2xf32> // %value is defined
%output = itensor_write %value into %arg : ...
yield %output : itensor(b)
} : itensor(b)
在降级和缓冲化之后,同样的代码变为:
%stream = stream() : stream<f32, depth: 32>
task @example {
%value = ... : f32 // %value is defined
stream_write %value into %stream : ...
}
我们可以观察到,task结合了itensor和stream操作,使其成为一个跨越两种IR的统一结构抽象,服务于不同级别的数据流优化。最终,所有数据流任务都会被降级为MLIR内置的call和func操作以进行代码生成。
编译流程概述。基于类型系统和操作,我们引入了一个编译流程,将Linalg IR编译成硬件实现和相应的运行时。所有编译过程如图4所示。本节我们重点关注Linalg到数据流的转换、数据流内核融合以及对理解编译器至关重要的数据流优化。
转换过程。图6(a)-(c)展示了Linalg到数据流的转换过程。原始的Linalg操作(图6(a))首先被切分(tile)成图6(b)的形式,其中scf.for代表用于切分的循环嵌套。在每次迭代中,extract_slice提取输入张量块以供给切分后的Linalg操作。当操作产生输出块后,insert_slice将它们插回完整的张量中。然后,每个切分后的循环嵌套被原地转换为一个kernel操作,如图6(c)所示。输入和输出张量在kernel的边界处与itensor相互转换。itensor的类型由以下信息推断得出:(1) 嵌套的scf.for循环——迭代的行程计数和步长定义了itensor的迭代空间。(2) extract_slice和insert_slice操作的偏移量和大小——偏移量定义了迭代映射,而大小定义了元素形状。例如,偏移量[%iv2, %iv0]会产生迭代映射(d0,d1,d2)->(d2,d0)。转换后,extract_slice和insert_slice操作分别被itensor_read和itensor_write操作替代。最终的scf.for循环嵌套被包裹在一个task中,形成一个单层的数据流层级:一个包含一个数据流任务的数据流内核。通过将Linalg语义转换为数据流语义,我们为后续面向数据流的转换和优化创造了机会。
融合过程。在所有切分后的Linalg操作都转换为数据流内核后,这些内核最初通过传统的张量进行通信,这些张量最终存储在外部内存中。为了减少这种通信开销,StreamTensor应用了基于流的内核融合。图6(c)-(d)展示了这一过程。要融合Kernel0和Kernel1,我们首先比较Kernel0的输出itensor类型与Kernel1的输入itensor类型。如3.1.2节所述,如果类型匹配,我们可以直接融合内核。如果不匹配,我们插入一个流布局转换器,如图6(d)所示。融合后的内核包含两个任务和一个转换器,它们都通过itensor进行通信,这些itensor将被降级为片上流式FIFO。itensor类型系统使得任何数据流内核在设计上都可以进行融合,代价是可能需要为转换器占用片上内存。在5.2节中,我们将讨论在内存约束下探索内核融合空间。
融合后的优化。融合之后,StreamTensor应用额外的优化遍来提高外部内存访问的效率。特别地,在内核前后插入张量打包(pack)和解包(unpack)操作,以在默认内存布局和切分后内存布局之间进行转换,从而实现突发内存访问。例如,在一个64x64的张量上使用[16,16]的切分大小,打包后的张量形状为4x4x16x16。为了最大化利用外部内存带宽,StreamTensor使用向量来加宽张量。例如,对于512位DDR或HBM和uint8元素,将64个元素分组为vector<64>可以充分利用带宽。在图6中,打包后的张量被加宽为形状4x4x2x2xvector<8x8>。需要注意的是,打包和加宽操作最终会被降级为主机CPU上的运行时操作,这会为加速器准备数据并产生一些延迟和内存开销。然而,对于静态张量(例如,预训练参数),打包和加宽可以直接融合到这些张量中,消除任何运行时成本。对于动态张量(例如,激活值),打包和加宽操作可以通过有效的Linalg切分空间探索与前一层的解包和解宽对应操作折叠。因此,打包和加宽操作仅对模型的输入和输出是必需的,在运行时贡献的内存和延迟开销可以忽略不计。
具象化(Materialization)。图7(a)和(b)展示了转换器和DMA的具象化过程。具象化涉及将一个高级数据流组件转换为其低级实现,通常是包含张量和itensor操作的scf.for循环嵌套。最初,转换器由itensor_converter表示,而DMA则通过内核边界处的张量到itensor或itensor到张量的转换来隐式处理。这种抽象有助于内核融合和转换器优化。例如,为一个生产者的多个消费者生成的冗余转换器可以使用MLIR的公共子表达式消除(CSE)来移除,这在具象化后变得更加困难。相比之下,具象化后,所有数据流组件都表示为嵌套的任务,使得进一步的数据流优化高效且易于实现。对于转换器,如图7(a)所示,Converter0包含两个scf.for循环嵌套,通过一个16x64的乒乓缓冲区连接。这两个循环嵌套被一个共享的父scf.for循环包裹,以迭代原始的完整64x64张量。因此,16x64的乒乓缓冲区被重用四次,有效地将片上内存资源利用率降低了四倍。在5.2节中,我们将讨论如何从itensor类型推断乒乓缓冲区的形状和共享循环。
DMA的具象化。对于DMA,如图7(a)所示,从tensor<4x4x2x2xvector<8x8>>到itensor<16x16...>的输入类型转换表明存在一个DMA,它将:1) 从外部内存加载4x4x2x2次vector<8x8>的数据;2) 将这些数据存储在一个16x16的乒乓缓冲区中以隐藏外部内存访问延迟;3) 将数据推送到一个FIFO中,其布局由itensor类型编码。在图7(b)中,我们观察到DMA0被自动生成,使用scf.for循环嵌套来实现这三种行为。请注意,我们基于itensor的类型系统编码了所有的转换器和DMA信息。这是传统张量类型所不具备的能力,限制了它们在数据流组件生成中的应用。
迭代张量折叠。图7(b)-(c)展示了itensor折叠。假设我们有一个在DMA0中的itensor_write和一个在Kernel0中的itensor_read,它们通过一个FIFO连接。这代表了两个通过流连接的独立本地缓冲区。通过折叠,我们消除了FIFO并将两个缓冲区合并。这种优化可以减少片上内存利用率,同时通过增加内核之间的重叠来提高整体延迟。如图7(c)所示,取出的数据块直接传递给Kernel0中的linalg.generic操作,消除了冗余的缓冲和通信。itensor折叠要求生产者和消费者的内存访问模式完全匹配。这使得它比基于流的内核融合更具限制性,后者可以应用于任何数据流内核之间。因此,我们将itensor折叠实现为在已融合内核之上的一个额外优化。
迭代张量向量化。由于数据流内核通常并行运行,我们必须对数据流FIFO进行向量化以提供足够的带宽。图7(c)-(d)展示了将一个itensor向量化为vector<2x4>的过程。在DMA0+Kernel0一侧,itensor_write变成了一个包含transfer_read(从缓冲区)后跟itensor_write(到FIFO)的循环。在Converter0一侧,也为读取应用了类似的转换。这个过程使FIFO带宽与数据流内核的并行度保持一致。
设计空间划分。为了生成可实现且优化的加速器,我们必须正确配置编译遍的参数。如图4所示,我们将整个设计空间划分为三个子空间:Linalg切分空间、内核融合空间和资源分配空间。
Linalg切分空间定义。Linalg切分空间决定了每个数据流内核的切分因子、展开因子、置换策略以及输入/输出向量化。在StreamTensor中,这个空间由一个Linalg操作图表示,其中每个节点都标注了循环行程计数、步长和循环类型(归约或并行)等属性。探索的结果也会写回这个图,以配置转换遍。
探索算法。对于切分,向用户暴露一个超参数default_tile_size,并应用于所有内核的所有维度。对于展开,我们开发了一种强度感知算法,它通过一个最大堆迭代选择延迟最长的内核,并增加其展开因子,直到达到用户定义的超参数overall_unroll_size。这种方法平衡了内核延迟以提高吞吐量。一旦确定了展开大小,向量化因子就通过分析循环迭代空间和张量形状来推断。置换由一种启发式方法处理,该方法将归约循环向外移动,同时保持并行循环在最内层,从而减少流水线循环的启动间隔(II)。在StreamTensor中,Linalg切分空间的超参数通过黑盒优化器Optuna【索引4,Optuna: A next-generation hyperparameter optimization framework,2019,In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining】进行自动探索,并利用数据流内核融合结果的反馈。
内核融合概述。如4.2节所述,内核融合使得内核之间能够进行流式传输。如果生产者和消费者的itensor类型不同,则必须插入一个转换器。Linalg切分空间的探索决定了所有数据布局和形状,从而固定了所有数据流内核接口处的itensor类型。因此,融合任何一对内核的内存开销也随之确定。由于片上内存有限,融合所有内核通常是不可行的。为了在遵守内存资源约束的同时有效选择要融合的内核对,我们提出了两种算法:算法1用于推断流布局转换器所需的最小乒乓缓冲区形状;算法2用于在片上内存约束下确定一个全局融合计划。
流布局转换器生成。算法1在每个数据维度上比较源和目标itensor(第3-16行)。只有在以下情况下,乒乓缓冲区的大小才能在一个数据维度上被缩减:1) 它们的元素大小相等(第4-5行);2) 它们对应的迭代维度相等,即引用相同的循环嵌套级别(第8-16行)。例如,在图5中,itensor(b)和itensor(c)的第二个数据维度都对应于迭代维度d0,允许这个维度被缩减;我们只需要缓冲单列的块。在具象化过程中,将生成共享循环以沿这个缩减的维度重用缓冲区。相反,它们的第一个数据维度分别对应于迭代维度d1和d2,使得它们不可缩减。因此,我们必须缓冲所有行的块。因此,如图5所示,布局转换器中需要两个块(乒乓缓冲后为四个块)。
可实现性保证。在识别出可缩减的数据维度和相应的共享循环后,算法会过滤掉那些父循环不可共享的循环,以确保缓冲区的可实现性(第17-19行)。例如,如果循环-{0,1,2,4}是可共享的,但循环-3不是,则必须排除循环-4。最后,返回缓冲区的形状和共享循环。这个过程的最坏情况是没有任何维度可缩减,要求将整个数据保留在片上进行融合。这可能导致巨大的内存开销。
内核融合探索。算法2的输入C_max(最大融合成本)代表单个融合内核可以利用的最大片上内存。对于FPGA,这通常设置为总片上内存大小。因此,内核融合过程也可以看作是一个图划分问题。融合后,每个产生的融合内核将占用一个FPGA。如果一个计算图包含多个这样的内核,它们可以跨多个FPGA执行,或在单个FPGA上顺序执行,或采用混合方法。StreamTensor作为编译器支持所有这些方法。然而,将N个内核映射到M个FPGA并管理FPGA间的通信超出了本文的范围。算法2按拓扑顺序遍历所有内核(第3行)。对于每个内核,它首先从前驱节点收集融合候选,并计算融合成本(第4-11行)。如果内核与最近的有效候选融合后不超过资源限制(第15-20行),则进行融合(第13-14行)。融合结果写回图(第22行),并用于配置4.2节中讨论的优化。数据流内核融合总是有可行解,除非单个内核占用的资源超过单个FPGA。在这种情况下,结果会反馈给切分空间进行调整,例如,减少切分和/或展开因子。
资源分配问题。在像FPGA这样的硬件上,由于片上内存和计算资源有限,有效的资源分配极大地影响布线拥塞和时钟频率。在这个空间中,我们需要解决:(1) FIFO大小确定:确定FIFO深度以避免死锁并提高执行重叠。本节将详细介绍。(2) 图划分:在多Die硬件上,我们需要将任务分配给不同的Die。这个分配问题被建模为整数线性规划(ILP)并求解。在我们的ILP模型中,一个二进制列表代表每个任务的分配。一个约束确保这个列表中只有一个元素可以是“1”,其位置表示分配的Die。ILP的目标是最小化Die间的通信和Die间的资源不平衡。由于类似的模型已被研究【索引22,FADO: F loorplan-A ware D irective O ptimization for High-Level Synthesis Designs on Multi-Die FPGAs,2023,Proceedings of the 2023 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】【索引27,AutoBridge: Coupling coarse-grained floorplanning and pipelining for high-frequency HLS design on multi-die FPGAs,2021,The 2021 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】,我们省略进一步的细节。(3) 内存分配:将每个缓冲区放置在FPGA的LUTRAM、BRAM或URAM中,按大小优先。由于这个算法很简单,我们省略进一步的细节。
令牌行为模型。为了解决1.3节中讨论的FIFO大小确定问题,我们首先提出一个基于分段线性函数的令牌生产和消费模型。图8(a)展示了通过InterFIFO融合的Source和Target内核之间的令牌通信。流水线II是两个连续输出令牌之间的时钟周期数,而初始延迟是产生第一个输出令牌所需的时钟周期数。令牌被定义为内核间通信的原子数据元素。在time0,所有五个输入令牌都在InputFIFO中,令牌在time1开始流入Source。在time5,Source将token1推入InterFIFO,而Target消耗token0,使InterFIFO中留下一个令牌。在time6,Target无法消耗token1,因为它需要两个周期来处理token0。同时,token2被推入InterFIFO,使其令牌数增加到两个。在time8,Source处理完所有令牌,此时InterFIFO达到其最大容量三个令牌。Target随后继续消耗和处理剩余的令牌,直到time15,所有令牌被完全处理。
分段线性函数建模。为了用可分析的函数对这些复杂行为进行建模,我们将图8(a)中的令牌状态重组为图8(b),将同一令牌的状态对齐在同一行。我们观察到,Source(蓝色)和InterFIFO(红色)部分之间的边界可以用一个分段线性函数(蓝色曲线)完美建模。这个函数代表Source产生的令牌数。同样,我们可以用橙色曲线对Target消耗的令牌数进行建模。这两条曲线之间的差异代表了InterFIFO中的令牌数。这些曲线可以通过内核的延迟、初始延迟和流水线II来表示。StreamTensor自动调用像HLS这样的供应商工具,在流程中间对每个内核的这些指标进行分析。由于这些指标特定于供应商平台的架构、技术节点和映射策略,它们必须通过这个分析过程获得。由于资源分配是最后一个设计空间,内核设计在后续的StreamTensor流程中保持不变。只要供应商工具使用确定性的调度算法,最终加速器的指标将与早期分析的指标相匹配。这种一致性保证了我们算法的有效性。
最大令牌数计算(源吞吐量低于目标)。如图8(c)所示,我们定义L为Source执行的总延迟;D为从Source执行开始到其产生第一个输出令牌的初始延迟;delay为从Source执行开始到Target执行开始的时间。delay总是大于或等于D,因为Target不能在Source产生第一个令牌之前开始执行。我们定义T为单次加速器执行中从Source传递到Target的精确令牌数。T是一个静态值,可以从StreamTensor中的张量形状分析推断出来。我们将在5.3.5节中讨论如何处理动态张量形状。对于一个静态的T值,InterFIFO中的最大令牌数max_tokens可以从delay分析计算得出:
最大令牌数计算(源吞吐量高于目标)。流水线II决定了曲线的斜率,即内核吞吐量。图8(c)展示了Source吞吐量大于Target吞吐量的情况。相反,当Source吞吐量较低时,数据饥饿可能会限制Target的吞吐量。图8(d)显示,当delay足够大时,Target不受影响,而图8(e)显示,Target最终会饥饿,其吞吐量被均衡为Source的吞吐量。在这两种情况下,max_tokens都可以从delay计算得出:
max_tokens与delay的关系。公式1和2都揭示了max_tokens和delay之间的正相关关系。如图8(c)-(e)所示,将InterFIFO的深度设置为max_tokens可以防止Target对Source产生反压。这确保了在多次加速器执行中,任何一对Source和Target之间的行为都是稳定和周期性的。通过防止反压引起的停顿,max_tokens和delay之间的分析关系得以保持。
均衡策略。5.3.2节中描述的方法被称为“正常”(Normal)均衡策略,它假设内核总是以其原始吞吐量产生令牌。然而,数据流加速器的吞吐量最终由其最慢的内核决定。基于此,我们提出了一种“保守”(Conservative)均衡策略,它将所有内核的流水线II调整为与最慢内核的吞吐量相匹配。由此产生的max_tokens值小于或等于正常策略的值,因为任何一对Source和Target曲线之间的差距都被最小化了。其缺点是较快的内核会频繁地因反压而停顿,可能增加延迟。因此,正常和保守策略在面积和性能之间提供了一个权衡,保守策略以增加总延迟为代价最小化了FIFO缓冲区的大小。保守和正常策略的关键区别在于它们的II最初是如何被调整的。因为这种调整保留了内核曲线的分段线性特性,所以从delay计算max_tokens的公式对于两种策略是相同的。
基于LP的FIFO大小确定。通过引入令牌行为模型,我们将FIFO大小确定问题转化为确定内核间delay值的问题。图8(f)展示了一个数据流图的例子。Kernel0有两个输出;Kernel1依赖于Kernel0;Kernel2有两个操作数,必须等待Kernel0和Kernel1的第一个令牌。鉴于Kernel1在D[0]+D[1]之后产生其第一个令牌,delay[0][2]必须大于或等于这个值。它们的关系如图8(f)所示,绿色曲线代表Kernel1。Kernel0和Kernel2之间FIFO的最大令牌数max_token[0][2],可以根据delay[0][2]计算得出。如果FIFO大小小于这个最大值,Kernel0会因反压而停顿,损害整体性能。这种停顿可能传播到Kernel1和Kernel2,阻止反压解决并可能导致死锁。一个等于max_token[0][2]的FIFO大小足以防止反压和避免死锁;它也是防止因意外内核停顿而导致性能下降所必需的。我们提出了一个LP公式来优化求解delay值。给定图G = (V, E),其中V是内核集合,E是内核间边的集合,LP的目标和约束为:
LP公式解释。e_i,j ∈ E覆盖了图中所有的边;path ∈ P_u,v覆盖了连接任意一对内核u和v的所有完整路径;e_i,j ∈ path覆盖了连接两个内核u和v的路径上的所有边。我们最小化所有边上延迟的总和,由于max_tokens和delay的正相关性,这可以作为优化FIFO大小的代理。threshold(u, v)是连接两个内核u和v的所有路径上累积的D的最大值:
LP公式示例与实现。上述示例的LP公式如图8(f)所示。注意,在此示例中,从Kernel0分叉的两条路径重新汇聚到Kernel2作为两个不同的输入操作数,而不是合并成一个单一输入。我们将在5.3.5节中讨论如何处理像路径合并这样的动态行为。LP问题不需要资源约束,原因有二:首先,如5.2节所讨论,数据流内核融合保证所有融合后的内核都将适应可用的片上资源,通过限制融合成本。其次,流式FIFO的内存利用率与数据流内核和转换器相比可以忽略不计。因此,LP问题可以在多项式时间内得到最优解。值得注意的是,我们不需要强制供应商工具实现这些delay。相反,这些delay通过数据流内核之间的FIFO依赖关系自动实现。在上面的例子中,Kernel2自动等待Kernel1,因为它依赖于Kernel1的输出令牌。
动态行为处理。StreamTensor使用不同的方法来管理数据流加速器内的动态行为:(1) 控制流:StreamTensor利用Torch-MLIR【索引55,Torch-MLIR: A compiler for the PyTorch ecosystem,2021,https://github.com/llvm/torch-mlir】作为其前端。Torch-MLIR可以尽可能从输入中推断出静态张量形状,消除与静态张量形状相关 的if语句并展开for循环。如果控制流依赖于运行时值,相应的子图将回退到在主机上进行朴素的PyTorch执行【索引5,Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation,2024,Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】。(2) 路径合并:这通常出现在存在控制流的情况下,特别是当一个数据流内核被来自不同源的输入重用时。通过消除控制流,Torch-MLIR解决了相应的路径合并问题。(3) 动态张量形状:具有动态形状的张量,如输入令牌和KV缓存,需要形状提示来定义其最大可能的维度大小(例如,最大序列长度)。这些提示决定了可以在任意两个数据流内核之间处理的总令牌数T。从这些最大T值,StreamTensor根据5.3节中讨论的方法推断出max_tokens。(4) FIFO停顿:StreamTensor不为数据流加速器生成静态调度。相反,所有数据流内核通过FIFO互连自动遵循其依赖关系。因此,由运行时事件(例如外部内存流量)引起的意外FIFO停顿不需要特殊处理。一旦导致停顿的事件解决,数据流加速器将从停顿点无缝恢复操作。
实验平台与模型
* 硬件平台:实验主要在AMD U55C FPGA上部署和评估StreamTensor生成的加速器。对比平台包括使用AMD U280 FPGA的Allo和DFX,以及NVIDIA A100和2080Ti GPU。具体硬件参数如表6所示。
* 软件配置:使用Vitis 2024.1将StreamTensor生成的HLS C++代码编译成比特流。前端采用Torch-MLIR处理PyTorch模型。
* 模型与数据集:评估了多个LLM,包括GPT-2【索引45,Language models are unsupervised multitask learners,2019,OpenAI blog】,以及新兴的Qwen【索引7,Qwen technical report,2023,arXiv preprint arXiv:2309.16609】、Llama【索引56,Llama: Open and efficient foundation language models,2023,arXiv preprint arXiv:2302.13971】和Gemma【索引50,Gemma: Open models based on gemini research and technology,2024,arXiv preprint arXiv:2403.08295】。所有模型均从Huggingface模型修改而来,以满足Torch-MLIR前端的要求。模型配置如表7所示。
* 测量方法:所有StreamTensor的实验结果均通过板上实测获得。
与FPGA SOTA工作的比较:在GPT-2模型上,StreamTensor成功地将一整个Transformer块融合到单个FPGA上,所有中间结果都在片上通信。整个模型通过多次触发该FPGA加速器并加载不同权重来顺序执行。如表4所示,与Allo【索引14,Understanding the potential of fpgabased spatial acceleration for large language model inference,2024,ACM Transactions on Reconfigurable Technology and Systems】【索引15,Allo: A Programming Model for Composable Accelerator Design,2024,Proceedings of the ACM on Programming Languages】相比,StreamTensor的总延迟缩短了0.76倍,首个令牌生成时间(TTFT)缩短了0.40倍。与DFX【索引29,Dfx: A low-latency multi-fpga appliance for accelerating transformer-based text generation,2022,2022 55th IEEE/ACM International Symposium on Microarchitecture (MICRO)】相比,提升更为显著,例如TTFT降低了0.19倍。这些性能增益得益于StreamTensor的自动化数据流架构探索,而Allo和DFX需要手动设计所有数据流组件,易出错且可能导致次优设计。
与GPU的比较:如表5所示,与A100和2080Ti GPU相比,StreamTensor的总延迟分别缩短了0.64倍和0.25倍。尽管GPU因其丰富的计算资源在TTFT指标上表现更好,但LLM推理的解码阶段是内存密集型的。StreamTensor生成的数据流加速器通过减少外部内存访问,在解码速度和总延迟上优于GPU。
灵活性与能效:为了评估StreamTensor的灵活性,本文在Qwen、Llama和Gemma等新兴LLM上进行了测试。对于这三个模型,同样成功地将一个Transformer块融合到单个FPGA上。如图9所示,由于FPGA的功耗较低,StreamTensor在Qwen和Gemma模型上的能效比A100分别高出1.99倍和1.59倍。图10a显示,Llama模型产生更多的中间结果,导致StreamTensor采用更保守的FIFO大小调整策略,从而减少了内核间的执行重叠,性能相较于Qwen和Gemma较低。
片上内存减少研究:图10a展示了在所有评估的LLM中,内核融合前后片上内存的使用情况(仅考虑单层内的中间结果)。内核融合将内存使用量减少到原始设计的14.8%–16.8%。若无融合,由于中间缓冲过大,LLM无法以完全数据流的方式部署。
编译时间研究:图10b展示了从PyTorch生成RTL的执行时间分解。HLS过程(从C++生成RTL)占用了大部分时间。下游工具的性能分析也占了很大一部分,因为资源分配决策依赖于准确的分析结果。相比之下,StreamTensor的编译和参数打包时间只占一小部分。如图10c进一步分解,StreamTensor的总编译时间在26.8秒到63.4秒之间。高级阶段(从Linalg优化到资源分配)相对较快,而低级阶段(缓冲化、HLS优化和代码生成)耗时更长,这验证了高级itensor优化的高效性。
相关工作概述。早期的工作【索引10,Parameterized dataflow modeling for DSP systems,2001,IEEE Transactions on Signal Processing】【索引11,Cyclestatic dataflow,1996,IEEE Transactions on signal processing】【索引37,Synchronous data flow,1987,Proc. IEEE】【索引39,Hierarchical reconfiguration of dataflow models,2004,Proceedings. Second ACM and IEEE International Conference on Formal Methods and Models for Co-Design, 2004. MEMOCODE’04】【索引52,StreamIt: A language for streaming applications,2002,Compiler Construction: 11th International Conference, CC 2002 Held as Part of the Joint European Conferences on Theory and Practice of Software, ETAPS 2002 Grenoble, France, April 8–12, 2002 Proceedings 11】奠定了基于流的数据流建模和编译的基础。后续工作【索引26,Minimizing buffer requirements under rate-optimal schedule in regular dataflow networks,2002,Journal of VLSI signal processing systems for signal, image and video technology】【索引38,Slack matching mode-based asynchronous circuits for average-case performance,2013,2013 IEEE/ACM International Conference on Computer-Aided Design (ICCAD)】【索引57,Leveraging protocol knowledge in slack matching,2006,Proceedings of the 2006 IEEE/ACM international conference on Computer-aided design】探索了数据流网络中的缓冲区最小化和松弛匹配问题。【索引18,Synthesis of statically analyzable accelerator networks from sequential programs,2016,Proceedings of the 35th International Conference on Computer-Aided Design】【索引21,Combining computation and communication optimizations in system synthesis for streaming applications,2014,Proceedings of the 2014 ACM/SIGDA international symposium on Fieldprogrammable gate arrays】探索了顺序程序的死锁分析和缓冲区大小确定。需要注意的是,这些论文关注于稳态场景(即5.3.3节中的保守均衡策略),忽略了面积和性能之间的权衡。【索引22,FADO: F loorplan-A ware D irective O ptimization for High-Level Synthesis Designs on Multi-Die FPGAs,2023,Proceedings of the 2023 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】【索引27,AutoBridge: Coupling coarse-grained floorplanning and pipelining for high-frequency HLS design on multi-die FPGAs,2021,The 2021 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】改进了FPGA上流式应用的布局规划和时钟频率。【索引32,Buffer placement and sizing for high-performance dataflow circuits,2021,ACM Transactions on Reconfigurable Technology and Systems (TRETS)】【索引59,Suppressing Spurious Dynamism of Dataflow Circuits via Latency and Occupancy Balancing,2024,Proceedings of the 2024 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】解决了动态调度数据流电路【索引31,Dynamically scheduled high-level synthesis,2018,Proceedings of the 2018 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】中的缓冲区插入和放置问题。
与现有编译器的比较。编译器对于将应用程序映射到DSA和FPGA等空间架构至关重要。SARA【索引64,SARA: Scaling a reconfigurable dataflow accelerator,2021,2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA)】为像Plasticine【索引44,Plasticine: A reconfigurable architecture for parallel paterns,2017,ACM SIGARCH Computer Architecture News】这样的大规模DSA提供了编译器栈。Revet的编译器【索引47,Revet: A Language and Compiler for Dataflow Threads,2023,arXiv preprint arXiv:2302.06124】将其“数据流线程”抽象映射到向量化DSA【索引48,Capstan: A vector RDA for sparsity,2021,MICRO-54: 54th Annual IEEE/ACM International Symposium on Microarchitecture】上。DSAGEN【索引58,DSAGEN: Synthesizing programmable spatial accelerators,2020,2020 ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA)】直接从数据流图描述合成可编程空间加速器。基于约束的调度技术【索引41,A general constraintcentric scheduling framework for spatial architectures,2013,ACM SIGPLAN Notices】通常使用ILP进行优化调度。更高级的编程抽象也很关键,例如Sigma【索引66,Sigma: Compiling Einstein Summations to Locality-Aware Dataflow,2023,Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】。针对FPGA的Stream-HLS【索引8,Stream-HLS: Towards Automatic Dataflow Acceleration,2025,Proceedings of the 2025 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】从C/C++或PyTorch自动生成优化的HLS数据流架构。尽管这些编译器和框架自动化了关键优化,但它们通常只支持部分设计空间探索,并且缺乏一个系统性的类型系统来支持灵活的基于流的内核融合和其他优化。
与Stream-HLS的详细对比。以Stream-HLS【索引8,Stream-HLS: Towards Automatic Dataflow Acceleration,2025,Proceedings of the 2025 ACM/SIGDA International Symposium on Field Programmable Gate Arrays】为例,其与StreamTensor的差异如下:
* DMA生成:由于缺乏系统性的类型系统,Stream-HLS无法自动为外部内存生成DMA,限制了其在实际应用中的可用性和可扩展性。
* FIFO大小确定:Stream-HLS忽略了FIFO大小确定问题,而这对于避免数据流加速器中的死锁和扩展到实际应用至关重要。
* 内核融合条件:Stream-HLS要求满足两个苛刻条件才能在数据流内核间进行流式传输:1) 共享缓冲区的读写次数必须相等;2) 生产者的写顺序必须与消费者的读顺序匹配。在不满足任一条件时,Stream-HLS无法执行内核融合。相比之下,StreamTensor通过基于itensor的类型系统解决了这两个问题,使得任何数据流内核在设计上都是可融合的。
* 内核融合空间探索:由于上述原因,Stream-HLS不支持像StreamTensor那样的内核融合空间探索,限制了其在无法完全部署在片上的大规模工作负载上的应用。例如,Stream-HLS仅分别报告了多头注意力和前馈层的性能,而不是整个Transformer块的性能。
本文介绍了StreamTensor,一个自动化生成和优化基于流的数据流加速器的编译器框架。StreamTensor的主要贡献包括一个基于itensor的类型系统,该系统是整个框架的基础;一个从PyTorch到设备的编译流程;以及一套用于探索关键架构参数的设计空间。通过解决现有框架中的常见缺陷,StreamTensor有效地提高了数据流加速器的效率。随着对高效AI需求的不断增长,StreamTensor为未来可扩展和可延伸的数据流编译工作铺平了道路。
未来工作
StreamTensor的模块化设计和itensor类型系统为未来的工作开辟了有前景的途径,特别是在扩展其与不同数据流架构和专用内核语言的兼容性方面。StreamTensor可以通过重定向其低级编译和代码生成阶段,以适应可编程架构,如AMD Versal【索引24,Xilinx adaptive compute acceleration platform: Versaltm architecture,2019,Proceedings of the 2019 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays】、Sambanova RDU【索引43,SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts,2024,arXiv preprint arXiv:2405.07518】和Groq LPU【索引1,The Groq software-defined scale-out tensor streaming multiprocessor: From chipsto-systems architectural overview,2022,2022 IEEE Hot Chips 34 Symposium (HCS)】。这个过程会将StreamTensor IR中的数据流内核、FIFO和布局转换器映射到特定平台的组件,例如AMD Versal中的AI引擎和路由网络。同样,StreamTensor可以与像Allo【索引15,Allo: A Programming Model for Composable Accelerator Design,2024,Proceedings of the ACM on Programming Languages】这样的内核语言集成,允许开发人员将手动优化的内核作为黑盒组件整合。在这两种情况下,itensor系统都作为一个关键的抽象层,使StreamTensor能够执行高级数据流优化,包括内核融合和数据流组件生成,同时与特定目标的后端和黑盒组件接口。这有望通过利用各种硬件平台和编程语言的独特优势,拓宽StreamTensor的适用性。