文件大小：1.48M
摘要：针对迈创3000（MT-3000）异构多核处理器在带宽不足场景下多头注意力（Multi-Head Attention， MHA）计算效率低的问题，本文提出一套涵盖算子优化、访存优化与调度优化的综合方案，以加速PyTorch框架中的MHA推理。通过融合MT-3000的VLIW指令集、片上存储层次和DMA传输特性，设计了三方面的优化策略：在算子方面，对Linear和Softmax等算子进行内核级优化与算子融合，减少计算与访存开销；在访存方面，利用广播机制与全局共享内存（GSM）优化数据流，降低主存带宽依赖；在调度方面，以行为粒度分块并行，隐藏数据传输延迟。实验表明，优化后的Linear算子单簇峰值性能达1.53 TFLOPS，占理论峰值的37.7%，较NVIDIA V100 GPU加速比最高达5.34倍；在典型大语言模型配置下（嵌入维度4096/8192，头数32/64），MHA机制相较NVIDIA V100 GPU实现最高23.53倍加速，且在单节点多簇环境中扩展性良好。研究为MT-3000在长序列推理任务中的高效部署提供了解决方案，并为天河超算支持大语言模型等复杂AI任务奠定了技术基础。
文章目录
1. 引言
2. 相关工作
2.1 MT-3000异构多核处理器
2.2 Multi-Head Attention
2.3 PyTorch
3. 动机
4. 优化方案
4.1 算子优化
4.1.1 Linear算子优化
4.1.2 softmax算子优化
4.1.3 算子融合
4.2 算子调度优化
4.3 访存优化
5. 实验结果与分析
5.1 Linear算子的各部分开销测试
5.2 Linear算子的加速比测试
5.3 单节点上的MHA加速测试
5.4 多簇上的MHA扩展性测试
6. 总结与未来工作
