首个端到端框架,让神经网络高效运行在AMD下一代AI引擎上
AIE4ML: An End-to-End Framework for Compiling Neural Networks for the Next Generation of AMD AI Engines
- 自动编译模型为优化固件,支持2D阵列并行与片上数据流动
- 单核性能逼近理论峰值,98.6%效率,97.4%芯片算力利用率
- 兼容量化模型,适合粒子物理等超低延迟场景
由于VLIW执行、显式数据通路和本地内存管理的紧密耦合,AMD Versal AI Engine(AIE)上的高效人工智能推理极具挑战性。以往工作仅聚焦于第一代AIE内核优化,未解决跨二维阵列的完整神经网络执行问题。本文提出AIE4ML,首个针对AIE-ML生成设备的全自动神经网络编译框架,同时具备对新AIE-MLv2架构的前向兼容性。在单核层面,性能接近架构峰值;在图与系统层面,提供结构化并行方法,可扩展至二维AIE-ML硬件,并利用专用内存单元实现全程片上数据流动。我们设计了一种通用高效的线性层实现,原生支持融合偏置加法与ReLU激活。由于框架需生成多层实现,我们通过一种新型图布局与搜索算法,系统性地推导出确定性、紧凑且拓扑优化的物理网格布局。框架无缝接收来自hls4ml或PyTorch的量化模型,保持位级精确。在层缩放基准测试中,相对单核基线达到98.6%效率,使用了304个AIE单元中的296个(97.4%),全程片上数据移动。在真实模型拓扑评估中,AIE4ML实现在微秒级延迟约束下的类GPU吞吐量,是粒子物理实验触发系统等超低延迟环境的理想工具。
原文摘要 · Abstract (English)
Efficient AI inference on AMD's Versal AI Engine (AIE) is challenging due to tightly coupled VLIW execution, explicit datapaths, and local memory management. Prior work focused on first-generation AIE kernel optimizations, without tackling full neural network execution across the 2D array. In this work, we present AIE4ML, the first comprehensive framework for converting AI models automatically into optimized firmware targeting the AIE-ML generation devices, also with forward compatibility for the newer AIE-MLv2 architecture. At the single-kernel level, we attain performance close to the architectural peak. At the graph and system levels, we provide a structured parallelization method that can scale across the 2D AIE-ML fabric and exploit its dedicated memory tiles to stay entirely on-chip throughout the model execution. As a demonstration, we designed a generalized and highly efficient linear-layer implementation with intrinsic support for fused bias addition and ReLU activation. Also, as our framework necessitates the generation of multi-layer implementations, our approach systematically derives deterministic, compact, and topology-optimized placements tailored to the physical 2D grid of the device through a novel graph placement and search algorithm. Finally, the framework seamlessly accepts quantized models imported from high-level tools such as hls4ml or PyTorch while preserving bit-exactness. In layer scaling benchmarks, we achieve up to 98.6% efficiency relative to the single-kernel baseline, utilizing 296 of 304 AIE tiles (97.4%) of the device with entirely on-chip data movement. With evaluations across real-world model topologies, we demonstrate that AIE4ML delivers GPU-class throughput under microsecond latency constraints, making it a practical companion for ultra-low-latency environments such as trigger systems in particle physics experiments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。