4 个月前

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

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

摘要

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

代码仓库

基准测试

基准方法指标
runtime-ranking-on-tpugraphs-layout-meanTGraph
Kendall's Tau: 0.674

用 AI 构建 AI

从想法到上线——通过免费 AI 协同编程、开箱即用的环境和市场最优价格的 GPU 加速您的 AI 开发

AI 协同编程
即用型 GPU
最优价格
立即开始

Hyper Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供
用于张量编译器的配置交叉注意力图神经网络 | 论文 | HyperAI超神经