Reproduce “FlashQuant: Sparse-Dense Fusion for Memory-Efficient Outlier-Aware LLM Inference” on OpenResearch

About this paper

Junqing Lin, Jingwei Sun, Zhengding Hu, Guangzhong Sun

FlashQuant introduces a framework for memory-efficient outlier-aware W4A16 Large Language Model (LLM) inference by fusing low-bit dense GEMM and high-precision sparse SpMM computations into a single GPU kernel. This approach reduces redundant global memory accesses and maximizes on-chip data reuse, resulting in up to 2.01x end-to-end speedup for LLaMA3-8B and enabling execution of models like Qwen2.5-14B that previously caused Out-Of-Memory errors.

Key insights

• Fusing heterogeneous computational paths (dense GEMM and sparse SpMM) derived from the same mathematical operation can significantly reduce memory traffic and enhance data locality for mixed-precision LLM inference.
• Co-designing the tiling hierarchy, data encoding, and scheduling mechanisms is crucial to effectively integrate irregular sparse computations with regular dense operations within a unified GPU kernel.
• Optimizations focused on minimizing global memory accesses and maximizing on-chip data reuse yield substantial performance gains for memory-bound LLM decoding workloads, especially on modern GPU architectures.

Summary by alphaXiv.