Streaming Tensor Program: A streaming abstraction for dynamic parallelism
Gina Sohn, Genghan Zhang, Konstantin Hossfeld, Jungwoo Kim, Nathan Sobotka, Nathan Zhang, Olivia Hsu, Kunle Olukotun
动态行为在许多张量应用程序中越来越普遍。 例如,在机器学习中,输入张量是动态形状或破烂的,数据依赖的控制流在许多模型中被广泛使用。 然而,空间数据流加速器的先前编程抽象的有限表现力迫使动态行为静态实现,或者缺乏性能关键决策的可见性。 为了应对这些挑战,我们提出了流式张量程序(STeP),这是一个新的流抽象,使动态张量工作负载能够在空间数据流加速器上高效运行。 STeP引入了灵活的路由运算符、显式内存层次和符号形状语义,可以暴露动态数据速率和张量维度。 这些功能解锁了新的优化动态平铺,动态并行和配置时间-多路复用 - 适应动态行为,同时保持数据流效率。 使用具有真实世界痕迹的具有代表性的LLM层上的循环近似模拟器,动态平铺将片上内存需求减少了2.18倍,动态并行化将延迟提高了1.5倍,配置时间复用使计算利用率比先前抽象中可用的实现提高了2.57倍。
Dynamic behaviors are becoming prevalent in many tensor applications. In machine learning, for example, the input tensors are dynamically shaped or ragged, and data-dependent control flow is widely used in many models. However, the limited expressiveness of prior programming abstractions for spatial dataflow accelerators forces the dynamic behaviors to be implemented statically or lacks the visibility for performance-critical decisions. To address these challenges, we present the Streaming Tenso...
