Flash normalization: fast normalization for LLMs

2024-07-12

AI Summary

This work addresses the high computational overhead and poor fusion efficiency of RMSNorm, LayerNorm, and emerging normalization methods such as Dynamic Tanh (DyT) in large language models. We propose FlashNorm—the first zero-error, CUDA kernel-level optimized implementation of RMSNorm, enabling seamless fusion with subsequent linear layers. By jointly optimizing operator fusion, memory access reordering, and unified kernel design, FlashNorm achieves the first coordinated optimization across multiple normalization operators (RMSNorm, LayerNorm, and DyT) without precision loss. Experimental evaluation on mainstream Transformer architectures—including Llama and Mistral—demonstrates 15–25% end-to-end inference speedup and 12% reduction in GPU memory footprint. The implementation is open-sourced and integrated into major Transformer libraries, facilitating immediate adoption in production and research settings.

Technology Category

Natural Language Processing: Learning & Optimization for NLPMachine Learning: Large Multimodal Models (LMMs)Computer Vision: Diffusion Models for Vision

Application Category

Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphsSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMsSearch and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

RMSNorm is used by many LLMs such as Llama, Mistral, and OpenELM. This paper details FlashNorm, which is an exact but faster implementation of RMSNorm followed by linear layers. FlashNorm also speeds up Layer Normalization and its recently proposed replacement Dynamic Tanh (DyT) arXiv:2503.10622. See https://github.com/OpenMachine-ai/transformer-tricks for code and more transformer tricks.

Problem

• Faster implementation of RMSNorm for LLMs
• Speeds up Layer Normalization techniques
• Enhances Dynamic Tanh (DyT) performance

Innovation

• FlashNorm: faster exact RMSNorm implementation
• Speeds up Layer Normalization and Dynamic Tanh
• Optimizes normalization followed by linear layers