Reproduce “SOLE: Hardware-Software Co-design of Softmax and LayerNorm for Efficient Transformer Inference” on OpenResearch

About this paper

Wenxun Wang, Shuchang Zhou, Wenyu Sun, Peiqin Sun, Yongpan Liu

The SOLE research introduces a hardware-software co-design to optimize Softmax and Layer Normalization operations for transformer inference. It achieves an average end-to-end speedup of 1.50x to 2.09x over FP32 models on GPUs, with negligible accuracy loss, by employing specialized algorithms and custom hardware units for low-precision computation and storage.

Key insights

• Addressing Softmax and LayerNorm with a co-designed hardware-software approach effectively resolves an emerging bottleneck in transformer inference as other operations become more efficient.
• Log2 quantization and dynamic compression, combined with hardware-friendly approximations, enable significant reductions in computational complexity and intermediate data bit-width (4-bit for Softmax, 8-bit for LayerNorm).
• Careful algorithmic design allows for maintaining inference accuracy with negligible drops (average 0.38%), completely eliminating the need for expensive model retraining, which is a common drawback of prior approximation methods.

Summary by alphaXiv.