咖啡图片
正在将巧克力泡入咖啡
ntainer" style="display: none">
文章

如何研发一块可编程AI加速器

拆解可编程 AI 加速器的研发目标、团队分工、workload 定义、计算系统、内存系统和多卡互联设计

如何研发一块可编程AI加速器

研发目标与定位

可编程 AI 加速器,属于领域专用可编程加速器(domain-specific programmable accelerator)的一种,既不是一块只能执行某个固定算法的电路,也不是像通用处理器那样尽量服务所有通用程序的处理器。“领域专用”即针对某一类 workload 优化,这里的 workload 包括 LLM 推理、LLM 训练、CV、边缘推理等等。“可编程”意思是不把模型直接固化为电路,而是暴露出一套执行模型、指令、内存模型和软件接口,让不同模型、不同 shape、不同 dtype 的计算都能通过 kernel、算子库或编译器落到硬件上。可编程 AI 加速器的研发目标,是围绕 AI 模型中最常见、最耗时的张量计算,把硬件、编译器、运行时和框架后端一起设计出来,让模型可以被映射到这块设备上运行。

研发团队和任务拆解

整个研发任务分成三个大的方面:架构、软件和硬件。

架构的部门划分和相应职责如下:

  • 系统:整个 SoC 怎么组成,设计空间探索、定义 ISA spec(指令集架构规范文档) 作为软硬件之间的接口
  • 性能:开发性能模型(Roofline Model、Cycle Model、Bandwidth Model、Latency Model),性能/功耗面积(PPA)分析、性能瓶颈分析
  • 建模:定义 Simulator 的功能模型、执行模型、内存模型作为软硬件之间的接口

软件的部门划分和相应职责如下(按软件层次划分):

  • 模型/算法:研究算法本身,如多模态和生成算法、推理和训练框架、大语言模型、量化和压缩等等
  • 框架:将算法拆成计算图以及具体的算子
  • 算子库:在定制的硬件平台完成算子的高效实现,算子可细分为通信算子、计算算子、分布式并行计算算子等
  • 编译器:对计算图进行优化,通过 DSL(Domain-Specific Language,领域专用语言)和 builtin library 将源代码实现的算子转化为目标平台硬件能认识的 binary 可执行文件
  • Runtime 和 Profiler(应用软件):为程序员提供调用异构计算设备的运行时 API 和性能分析工具
  • Driver:为 binary 软件提供操作硬件的 Host 侧的软件接口
  • Firmware:根据 ISA spec 开发芯片内部控制处理器和硬件模块之间的 Device 侧软件
  • 测试与 CI:提供软件侧的自动化测试和持续集成工具
  • 工具链开发:负责 SDK 等配套工具链的开发

硬件的部门划分和相应职责如下(按研发阶段划分):

  • SoC 设计:完成芯片的微架构设计
  • IP(Intellectual Property)Core:对各个 Core 模块完成前端 RTL 设计
  • 验证:通过 UVM / Formal / Emulation / FPGA 等方法学对前端设计进行验证和测试
  • HAV(Hardware Assisted Verification):利用 Emulation / FPGA 进行原型验证
  • PTV(Prototype Verification):完成 bring-up 功能验证,自动化测试流程
  • 后端(Physical Design):完成 Place-Route-CTS-Timing 这一套从 RTL 到 Layout 的流程
  • 封装 Package:完成芯片的封装
  • 系统硬件(Board Design):PCIE Card、Power、PCB 设计等等
  • 平台:设计芯片开发平台,包括 EDA Flow、BES(Backend Engineering Support)、DFT(Design For Test)等

定义 Workload

Workload 是芯片实际服务的工作负载,包括模型结构、输入 shape、batch、dtype、服务形态、延迟目标和通信模式。没有 workload,硬件架构目标会变得很空。由于不同模型对计算、访存、通信、延迟和精度的要求差别很大,因此定制加速器的硬件形态一定是基于实际计算负载(Workload)设计的。例如:矩阵阵列多大、片上 SRAM 多大、HBM 带宽多少、是否需要特殊函数单元、是否把 DMA 做得更强、是否支持 FP8、是否需要多卡互联、软件是否必须做自动 tiling 等等,这些硬件规格需要从 Workload 反推出来,凭空设定只会导致无法收敛的后果。

AI 加速器中,定义 Workload 需要收集典型模型和服务形态,统计 OP mix、shape、dtype、batch、seq_len,分析 GEMM、attention、embedding、通信和访存占比,再定义性能、功耗、面积、成本和延迟目标;然后通过 Roofline 模型,用 arithmetic intensity(算术强度,即每字节数据搬运能做多少计算)把 OP 分成 compute-bound 和 memory-bound。Roofline model 是一种性能分析模型,把硬件峰值算力、内存带宽和 OP 算术强度放在同一张图上。它适合快速判断瓶颈在计算还是访存。

LLM 训练更关心大规模 GEMM、反向传播、梯度同步、HBM 容量、多卡 allreduce 和容错。LLM 推理更关心首 token 延迟、decode 阶段的小 batch GEMM、KV cache、PagedAttention、continuous batching、显存容量和显存带宽。CV 模型会带来卷积、resize、feature map layout 等问题。推荐系统里 embedding 和稀疏访存可能比矩阵乘更难。边缘推理则会把功耗、成本、片上存储和 INT8/INT4 量化推到前面。

计算系统

AI 芯片最主要的计算单元是矩阵乘单元,例如 NVIDIA Tensor Core、Google TPU 的 MXU、各种 NPU matrix engine、TMAC 或 systolic array。它们的共同目标是高吞吐地执行矩阵乘和卷积等密集的线性代数运算。

在矩阵乘单元之外,还需要:

  • 向量单元处理 activation、softmax、elementwise、reduce
  • 标量或控制单元处理地址计算、循环、分支和调度
  • 特殊函数单元处理 exp、rsqrt、类型转换、量化和反量化

attention 这种算子尤其能说明问题:它既有矩阵乘,也有 softmax、mask、scale、layout、KV cache 和访存模式。FlashAttention 的论文把 attention 的瓶颈放在 HBM 访问上,而不是只看 FLOPs。

设计计算系统时,需要根据目标模型推导 tile size、array size 和吞吐,决定支持 FP32、BF16、FP16、FP8、INT8 甚至更低精度,模拟不同 shape 下的硬件利用率。峰值算力只是一项指标。更重要的是常见 shape 能否填满阵列,数据能否及时喂进去,计算结果能否以合适 layout 留在片上供下一个 OP 使用。

内存系统

HBM 或 DRAM 提供大容量和高带宽,片上 SRAM、shared memory、scratchpad 或 cache 提供更低延迟和更高能效,寄存器文件和 tile register 支撑最内层计算。DMA 或 tensor copy engine 负责在不同内存层次之间搬运数据。

内存系统要考虑 bank、interleave、address mapping、cache policy、MMU、IOMMU、虚拟地址、多进程隔离和错误恢复。对软件来说,最要紧的是能否表达异步搬运、ping-pong buffer / double buffering 和明确的数据布局。

多卡互联

大模型训练和高吞吐推理还会把问题推到多卡。单卡不够时,需要片内 NoC、chiplet 互联、多卡互联、PCIe、CXL 或自研互联。tensor parallel、pipeline parallel、expert parallel 会产生不同通信模式。allreduce、allgather、reduce-scatter、broadcast、send/recv 的性能会直接影响端到端吞吐。NCCL-like 通信库、PyTorch ProcessGroup 对接和拓扑感知调度,都是软件栈必须补齐的部分。

本文由作者按照 CC BY 4.0 进行授权
/body>