用于张量编译器的配置交叉注意力图神经网络

随着神经网络的日益流行，高效地提供推理工作负载的服务变得尤为重要。神经网络的推理工作负载可以表示为一个计算图，其中节点作为操作符对多维张量进行变换。这些张量可以通过组合数量巨大的方式转置和/或切片，某些配置可以加速推理过程。我们提出了一种名为TGraph的神经图架构，该架构允许筛选目标计算图的快速配置，从而代表了一个基于人工智能（AI）的张量编译器，而传统的编译器则主要依赖于启发式方法。所提出的解决方案在TpuGraphs布局集合上的平均Kendall's $\tau$值从可靠的基线水平29.8%提高到了TGraph的67.4%。我们估计，与我们的工作相关的潜在CO$_2$排放减少量相当于承载AI导向数据中心地区家庭总排放量的50%以上。