AI芯片架构大揭秘：英伟达GPU与谷歌TPU的设计、扩展及软件全解析

AI芯片架构2018年国际计算机体系结构研讨会上约翰·轩尼诗和大卫·帕特森发表图灵讲座《计算机体系结构的新黄金时代》。20世纪80年代轩尼诗和帕特森开展获图灵奖研究时单线程CPU性能年增长52%到2018年降至3%。特定领域架构应运而生以谷歌TPU v1为例其在神经网络推理方面吞吐量是CPU29倍能效高80倍。他们预测未来十年新型计算机体系结构将迎来寒武纪大爆发如今有数十种架构正大力研发专注AI计算。已实现实际部署的架构有GPU、脉动阵列加速器、Cerebras晶圆级引擎和Groq LPU。英伟达领先AMD紧随OpenAI和Meta向AMD下达6GW订单。TPU用于训练Gemini还将为Anthropic提供芯片Anthropic在Trainium芯片上运行Claude。Cerebras为OpenAI提供推理服务Groq LPU并入英伟达。问题所在AI计算主要由矩阵乘法主导Transformer模型是一系列矩阵乘法运算。训练前沿模型需执行10^25次乘加运算。矩阵乘法形式取决于工作负载训练和预填充是大规模矩阵 - 矩阵乘法计算强度高解码阶段是矩阵 - 向量乘法计算强度下降。推理系统通过批处理将GEMV运算转换为GEMM运算提高计算强度。架构问题是“内存墙”问题即计算能力增长而内存带宽跟不上。每种架构针对数据移动问题有不同策略。英伟达GPU英伟达GPU是大规模并行处理器设计理念是可编程芯片在主机CPU协调下通过CUDA编程适合运行可并行化工作负载。每代产品在流多处理器上添加加速原语不改变编程模型同芯片可用于多种场景。发展历程2006 - 2027年有多代架构和型号如2006年Tesla架构G802027年Rubin Ultra架构等。架构由面向吞吐量的核心、深层内存层次结构和少量调度硅芯片组成。核心是流多处理器每个芯片封装有多个SM。每个SM内部结构相同包含四个子分区共享L1/共享内存块和张量内存加速器。线程组织成线程束同步执行。计算CUDA核心负责除矩阵乘法外的操作绝大部分计算吞吐量来自张量核心。张量核心执行融合矩阵乘加运算不同代产品的张量核心功能不断提升矩阵乘法对发出指令的线程依赖减小提高了Transformer注意力内核执行效率。内存片上内存层次结构由硬件管理缓存结合软件提示片外是HBM。不同层次数据移动逐渐与线程束解耦每代产品让线程束操作减少。线程束专业化Hopper时代编程习惯是线程束专业化通过mbarrier和异步事务屏障实现细粒度生产者/消费者握手成为现代注意力内核参考标准。数值精度FP32曾是默认精度各代产品不断降低精度并采用缩放方案恢复精度同时集成更多硅芯片提高每瓦吞吐量。设计理念包括可编程性、通过大规模多线程隐藏延迟、线程束封装的矩阵乘法、异步内存层次结构、分摊SIMT开销。扩展有纵向扩展和横向扩展两种方式。纵向扩展将多个GPU绑定到一致内存域可直接通过NVLink访问其他GPU的HBM。纵向扩展架构由NVLink和NVSwitch组成定义了HGX 8 - GPU主板等多种形式。高密度连接依赖无源铜不同代产品在纵向扩展上不断升级。横向扩展在机架和集群级别连接内存域数据通过RDMA传输。横向扩展架构源于英伟达对Mellanox收购不同代产品在横向扩展带宽等方面有提升。从铜缆到光纤转换发生在机架边界Rubin推出后光学层集成到交换机ASIC中。NVLink Fusion开放纵向扩展网络。软件CUDA是编程模型开发者编写内核在多线程上启动。每代新产品引入新硬件并在PTX和SASS中暴露。软件栈包括cuBLAS、cuDNN、CUTLASS等。FlashAttention针对英伟达芯片手动优化。软件栈大部分由英伟达之外的人编写英伟达还提供专业技术支持。谷歌TPUTPU是矩阵乘法机器设计理念是专注密集矩阵乘法让XLA编译器提前规划内存使用。每代产品扩大集群规模用于训练和服务谷歌工作负载。发展历程2015年TPU v12017年TPU v2等。