面向张量编译器的配置交叉注意力图神经网络
机器学习
2024-11-26 v2 硬件体系结构
性能
摘要
随着神经网络近年来的普及，高效提供推理工作负载的需求日益增长。神经网络推理工作负载可以表示为一个计算图，其中节点是转换多维张量的算子。张量可以以组合数量庞大的方式进行转置和/或分块，某些配置能加速推理。我们提出了TGraph，一种神经图架构，能够筛选目标计算图的快速配置，从而代表了一种人工智能（AI）张量编译器，与传统的基于启发式的编译器形成对比。所提出的方案将TpuGraphs布局集合上的平均Kendall's 从可靠基线的29.8%提升至TGraph的67.4%。我们估计，与我们的工作相关的潜在CO减排量相当于托管AI数据中心地区总家庭排放量的50%以上。
引用
@article{arxiv.2405.16623,
title = {Graph neural networks with configuration cross-attention for tensor compilers},
author = {Dmitrii Khizbullin and Eduardo Rocha de Andrade and Thanh Hau Nguyen and Matheus Pedroza Ferreira and David R. Pugh},
journal= {arXiv preprint arXiv:2405.16623},
year = {2024}
}