SparkAttention: high-performance multi-head attention for large models on Volta GPU architecture

2025-02-12

CCF Transactions on High Performance Computing

AI Summary

To address low computational efficiency and high memory pressure in Multi-Head Attention (MHA) during large Transformer model training on NVIDIA Volta architecture (e.g., V100), this work proposes a fine-grained kernel fusion and dynamic shared-memory scheduling strategy tailored to Volta Tensor Cores. Our method integrates low-level CUDA optimizations, FP16/INT8 mixed-precision arithmetic, attention computation graph rewriting, and TCU-customized scheduling—achieving performance gains without accuracy degradation. Experimental results demonstrate a 3.2× improvement in MHA throughput and a 67% reduction in end-to-end inference latency. Notably, we achieve real-time inference for a 13B-parameter model on a single V100 GPU—the first such result on this hardware—while significantly enhancing GPU memory utilization efficiency and training scalability.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Large Vision Models
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Problem

• Accelerate Multi-Head Attention training
• Optimize for Volta GPU architecture
• Reduce memory access and overhead

Innovation

• Optimizes Multi-Head Attention on Volta GPU
• Utilizes Tensor Core Units efficiently
• Reduces memory access with kernel fusion