Samoyeds: Accelerating MoE Models with Structured Sparsity Leveraging Sparse Tensor Cores

TL;DR

Samoyeds tackles the challenge of efficiently deploying MoE LLMs by exploiting dual-sided structured sparsity (activations and weights) with NVIDIA Sparse Tensor Cores. It introduces a novel sparse data format and a sparse-sparse MMA kernel, paired with tiling, data stationary execution, packing, and layout optimizations to maximize hardware utilization. Empirical results show kernel-level speedups up to 1.99x and end-to-end model speedups up to 2.36x, along with up to 4.41x larger maximum batch sizes and robust portability across modern GPUs. When compared with state-of-the-art sparse libraries and dynamic-sparsity compilers, Samoyeds delivers superior performance, accuracy retention, and hardware compatibility. This work demonstrates practical, scalable improvements for real-world deployment of large MoE LLMs on contemporary accelerators.

Abstract

The escalating size of Mixture-of-Experts (MoE) based Large Language Models (LLMs) presents significant computational and memory challenges, necessitating innovative solutions to enhance efficiency without compromising model accuracy. Structured sparsity emerges as a compelling strategy to address these challenges by leveraging the emerging sparse computing hardware. Prior works mainly focus on the sparsity in model parameters, neglecting the inherent sparse patterns in activations. This oversight can lead to additional computational costs associated with activations, potentially resulting in suboptimal performance. This paper presents Samoyeds, an innovative acceleration system for MoE LLMs utilizing Sparse Tensor Cores (SpTCs). Samoyeds is the first to apply sparsity simultaneously to both activations and model parameters. It introduces a bespoke sparse data format tailored for MoE computation and develops a specialized sparse-sparse matrix multiplication kernel. Furthermore, Samoyeds incorporates systematic optimizations specifically designed for the execution of dual-side structured sparse MoE LLMs on SpTCs, further enhancing system performance. Evaluations show that Samoyeds outperforms SOTA works by up to 1.99$\times$ at the kernel level and 1.58$\times$ at the model level. Moreover, it enhances memory efficiency, increasing maximum supported batch sizes by 4.41$\times$ on average. Additionally, Samoyeds surpasses existing SOTA structured sparse solutions in both model accuracy and hardware portability.

Figures (19)

• Figure 1: MoE LLM Architecture.
• Figure 2: Time Breakdown of MoE Models. Left: Without Flash-Attention; Right: With Flash-Attention.
• Figure 3: Data Patterns in Different Sparse Formats. Blank cells represent sparse elements.
• Figure 4: 2:4 Sparse Encoding and Mapping for SpTC.
• Figure 5: Redundancy in Data Flow of the MoE layer.