Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity

一、论文概述

项目	内容
标题	Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity
作者	Haojun Xia, Zhen Zheng, Yuchao Li, Donglin Zhuang, Zhongzhu Zhou, Xiafei Qiu, Yong Li, Wei Lin, Shuaiwen Leon Song (*共同第一作者)
机构	阿里巴巴集团 (Alibaba Group) + 悉尼大学 FSA-Lab (University of Sydney)
论文	https://arxiv.org/abs/2309.10285
代码	https://github.com/AlibabaResearch/flash-llm
发布	2023-09-19
会议	VLDB 2024
许可	Apache License 2.0
领域	cs.DC (分布式/并行/集群计算), cs.AR (硬件架构), cs.LG (机器学习)

二、核心思想

问题定义

大语言模型 (LLM) 的参数量快速增长（GPT-2: 1.5B, GPT-3: 175B, Megatron-Turing NLG: 530B），使得模型部署面临严峻挑战：

显存瓶颈：GPT-3 仅参数就需要 350GB（FP16），而 A100 GPU 最大仅有 80GB 显存
计算瓶颈：大模型推理需要大量计算和内存访问
成本问题：多卡并行推理导致 ROI 低下，跨设备通信进一步降低效率
现有稀疏方案低效：非结构化稀疏 SpMM 在 Tensor Core 上难以高效执行，Sputnik/cuSPARSE 等 SIMT 核心方案无法充分利用 Tensor Core 的 16x 峰值性能优势

关键观察：LLM 推理的性能瓶颈在于四个 Skinny MatMul（QKV Projection、Output Projection、MLP1、MLP2），这些矩阵乘法的 N 维度（batch size）远小于 M 和 K 维度（hidden size），导致 Tensor Core 利用率极低（batch_size=8 时仅约 5%），而全局内存带宽已完全饱和。

解决方案概述

Flash-LLM 提出 Load-as-Sparse and Compute-as-Dense (LSCD) 方法论：

核心洞察：Skinny MatMul 的瓶颈是内存带宽而非计算能力，因此可以以稀疏方式加载数据来减少内存访问量，同时容忍冗余的 Tensor Core 计算
稀疏加载：利用稀疏格式减少从全局内存加载的数据量，解决内存带宽瓶颈
密集计算：在 Tensor Core 上执行密集矩阵乘法，充分利用其高吞吐量
软件流水线：设计两级重叠策略，协调稀疏数据提取、密集数据加载和 Tensor Core 计算

三、技术架构

整体框架

Flash-LLM 的架构基于以下核心组件：

┌─────────────────────────────────────────────────────────┐
│                    Flash-LLM Framework                    │
├─────────────────────────────────────────────────────────┤
│  ┌─────────────┐  ┌──────────────┐  ┌───────────────┐  │
│  │  Tiled-CSL   │  │  Sparse-to-  │  │  Two-level    │  │
│  │  Format      │──│  Dense       │──│  Overlapping  │  │
│  │  (稀疏编码)   │  │  Transform   │  │  Pipeline     │  │
│  └─────────────┘  └──────────────┘  └───────────────┘  │
│         │                │                    │          │
│         ▼                ▼                    ▼          │
│  ┌─────────────────────────────────────────────────┐    │
│  │           GPU Hardware Resources                 │    │
│  │  ┌──────────┐ ┌──────────┐ ┌──────────────────┐ │    │
│  │  │ Registers│ │ Shared   │ │ Tensor Cores     │ │    │
│  │  │ (分布式)  │ │ Memory   │ │ (密集计算)        │ │    │
│  │  └──────────┘ └──────────┘ └──────────────────┘ │    │
│  └─────────────────────────────────────────────────┘    │
└─────────────────────────────────────────────────────────┘

Skinny MatMul 在 LLM 推理中的位置（Decoder Layer 的四个关键 MatMul）：

MatMul	权重矩阵 A 形状	特征矩阵 B 形状	说明
QKV Projection	[M, K] = [H, 3H]	[3H, N] = [3H, B]	查询/键/值投影
Output Projection	[M, K] = [H, H]	[H, N] = [H, B]	输出投影
MLP1	[M, K] = [4H, H]	[H, N] = [H, B]	FFN 第一层
MLP2	[M, K] = [H, 4H]	[4H, N] = [4H, B]	FFN 第二层

其中 H 为 hidden dimension（如 GPT-3 的 12K），B 为 inference batch size（通常较小，如 8/16/32/64）。

核心公式

MatMul 形式化定义：

$C = A \times B$

其中 $A \in \mathbb{R}^{M \times K}$ 为权重矩阵（稀疏）， $B \in \mathbb{R}^{K \times N}$ 为特征矩阵（密集）。

计算强度 (Compute Intensity, CI)：

$CI = \frac{2 \times M \times N \times K}{M \times K \times \text{sizeof}(dtype) + K \times N \times \text{sizeof}(dtype) + M \times N \times \text{sizeof}(dtype)}$

对于 Skinny MatMul（N 很小），CI 的上界受 N 限制：

$CI \leq N$

例如当 N=16 时，CI 最大为 16，无论 M 多大。

LSCD 方法下的性能分析：

给定稀疏率 $\beta$ （权重矩阵 A 的稀疏比例）：

内存加载量减少为 $(1-\beta)$ 倍
计算量保持不变（容忍冗余计算）
当 $(1-\beta) \times \text{原始时间} < \text{原始时间}$ 时，即内存瓶颈主导时，LSCD 获得加速

Tensor Core 利用率分析：

在 OPT-175B 的典型 MatMul（M=48K, N=BS, K=12K）中：

Batch size 8/16/32/64 时，Tensor Core 峰值利用率仅 5.1%/10.3%/20.5%/40.1%
全局内存带宽已完全饱和
LSCD 通过减少内存访问量有效提升 Tensor Core 利用率

模型组件

组件	说明	关键参数
Tiled-CSL 格式	按 tile 组织的稀疏编码格式，支持 tile-by-tile SpMM 执行	每 tile 包含非零元素值+位置索引
Sparse-to-Dense Transform	利用分布式寄存器和共享内存将稀疏数据转换为密集格式	使用 SIMT 核心执行
两级重叠流水线	Inter-iteration（双缓冲）+ Intra-iteration（阶段重叠）	gmem2reg, rst_smem, extract, ld_dense, smem2tc
Ahead-of-time 数据重排序	预处理阶段重排稀疏数据以减少共享内存 bank conflict	基于 ldmatrix 布局要求
Tiling 配置	M_TB=128/256, K_TB=64, N_TB=8/16/32/32	根据 batch size 调整

训练/推理流程

推理流程（Pipeline 执行）：

初始化阶段：准备第一个迭代的 A_Tile 和 B_Tile 数据到共享内存
主循环（对每个 tile 迭代）：
- gmem2reg：从全局内存加载稀疏编码到分布式寄存器（LDG 指令）
- rst_smem：重置目标共享内存缓冲区为零
- extract：将寄存器中的稀疏编码提取到共享内存（STS 指令）
- ld_dense：从全局内存异步加载密集特征矩阵 B 到共享内存（cp.async）
- smem2tc：从共享内存加载数据到寄存器供 Tensor Core 使用（ldmatrix）
- TC_compute：Tensor Core 执行密集矩阵乘法
同步与屏障：最小化同步点，利用异步复制原语实现重叠

两级重叠策略：

Inter-iteration 重叠：通过双缓冲，当前迭代的 Tensor Core 计算与下一迭代的内存加载重叠
Intra-iteration 重叠：Sparse-to-Dense 变换内部，gmem2reg 与 rst_smem 阶段重叠

稀疏编码格式：Tiled-CSL

Tiled-CSL (Tiled Compressed Sparse List) 格式的核心设计：

原始稀疏矩阵 A [M x K]:
┌───────────────────────────────┐
│  Tile(0,0)  Tile(0,1)  ...   │
│  Tile(1,0)  Tile(1,1)  ...   │
│  ...         ...        ...   │
└───────────────────────────────┘

Tiled-CSL 编码:
NonZeros[]: [v0,idx0 | v1,idx1 | v2,idx2 | ...]
              ──tile 0──  ──tile 1──  ...
TileOffsets[]: [offset_0, offset_1, offset_2, ...]

每个非零元素 = 32 bits = [16-bit value | 16-bit index]

非零元素按 tile 组织，每个 tile 维护其非零元素及稀疏索引
使用 uint4（128-bit）向量化加载以最大化内存带宽
Tile 偏移量存储在 TileOffsets 数组中，用于运行时定位每个 tile 的稀疏编码

Ahead-of-time 数据重排序算法：


for each tile in A:
    # 收集 tile 内所有非零元素
    nonzeros = collect_nonzeros(tile)

    # 按 32 个一组分组
    groups = split_into_groups(nonzeros, size=32)

    for each group:
        # 按目标共享内存 bank 重排序
        # 确保同一 warp 内的元素映射到不同 bank
        reordered = reorder_by_bank(group)
        # 存入 NonZeros 数组

四、核心创新

创新点	说明	理论/实验依据
Load-as-Sparse and Compute-as-Dense (LSCD)	以稀疏方式加载权重数据减少内存访问，在 Tensor Core 上执行密集计算	Skinny MatMul 的 CI 上界受 N 限制，内存带宽是瓶颈；A100 上 TC 吞吐量是 SIMT 的 16x
Tiled-CSL 稀疏编码格式	按 tile 组织的稀疏格式，支持 tile-by-tile SpMM 执行和高效稀疏数据提取	支持寄存器级稀疏数据缓冲和向量化加载（uint4）
两级重叠软件流水线	Inter-iteration（双缓冲）+ Intra-iteration（阶段内重叠）的流水线设计	有效隐藏内存延迟，Tensor Core 利用率从 5% 提升至接近饱和
分布式寄存器作为稀疏数据缓冲	使用寄存器而非共享内存作为全局内存到共享内存的中间缓冲	增加指令级并行度（ILP），避免共享内存容量限制
Ahead-of-time 稀疏数据重排序	预处理阶段重排稀疏数据元素，减少共享内存 bank conflict	基于 ldmatrix 布局要求，确保 warp 内元素映射到不同 bank
最小化同步和内存屏障	精确控制同步点位置，使用不同异步复制屏障实现细粒度流水线	rst_smem 和 ld_dense 使用独立屏障，extract 可与 ld_dense 重叠

五、代码实现分析

项目结构

flash-llm/
├── csrc/                           # 核心 CUDA 内核实现
│   ├── SpMM_Kernel.cuh             # SpMM 内核主模板（含 LSCD 流水线）
│   ├── SpMM_API.cu                 # SpMM API 接口（C++ 封装）
│   ├── TilingConfig.h              # Tiling 配置参数
│   ├── MatMulUtilities.cuh         # 矩阵乘法工具函数
│   ├── MMA_PTX.cuh                 # Tensor Core MMA PTX 指令封装
│   ├── AsyncCopy_PTX.cuh           # 异步复制 PTX 指令封装
│   └── Reduction_Kernel.cuh        # 归约内核
├── kernel_benchmark/               # 内核性能基准测试
│   ├── spmm_test.cu                # SpMM 测试主程序
│   ├── spmm_test_utils.h           # 测试工具
│   ├── sputnik_utils.h             # Sputnik 对比工具
│   ├── benchmark.sh                # 基准测试脚本
│   └── profiling.sh                # 性能分析脚本
├── end2end_inference/              # 端到端推理集成
│   ├── ft_tools/                   # FasterTransformer 集成工具
│   └── ds_scripts/                 # DeepSpeed 对比脚本
├── third_party/                    # 第三方依赖
├── docs/                           # 文档
│   ├── 1_Preparations.md           # 环境准备
│   ├── 2_KernelBenchmarking.md     # 内核基准测试指南
│   └── 3_LLMInferenceExample.md    # LLM 推理示例
├── build/                          # 构建目录
├── Init_FlashLLM.sh                # 初始化脚本
├── LICENSE                         # Apache 2.0
└── README.md

关键实现细节

1. SpMM 内核模板（SpMM_Kernel.cuh）

核心内核函数使用模板参数化：

TilingConfig：控制 tile 大小（M_TB=128/256, K_TB=64, N_TB=8/16/32/32）
SparseKernelConfig：控制稀疏内核参数（寄存器数量、向量大小等）

关键函数 SpMM_CopyFromGlobalToReg：

// 从全局内存加载稀疏编码到分布式寄存器
// 使用 uint4 (128-bit) 向量化加载
template<typename TilingConfig, typename SparseKernelConfig>
__device__ __forceinline__ void SpMM_CopyFromGlobalToReg(
    uint32_t* Registers_GlobalToShared1,  // 寄存器缓冲区1
    uint32_t* NNZ_VECTOR_ThreadLocal1,    // 本线程处理的非零向量数
    const uint4* GlobalPTR1,              // 全局内存指针
    int NNZ_VECTOR_ThisTile1,             // 当前 tile 的非零向量总数
    ...);

2. Tiling 配置（TilingConfig.h）

// 根据 batch size (N) 选择配置
// N=8:   M_TB=128/256, K_TB=64, N_TB=8
// N=16:  M_TB=128/256, K_TB=64, N_TB=16
// N=32:  M_TB=128/256, K_TB=64, N_TB=32
// N=64:  M_TB=128/256, K_TB=64, N_TB=32 (使用 2 个 tile)

3. Tensor Core MMA 指令（MMA_PTX.cuh）

封装了 NVIDIA PTX 级别的 mma.sync 指令，支持 FP16 输入和 FP32 累加：

m16n8k16 指令格式
使用 ldmatrix 从共享内存高效加载到寄存器

4. 异步复制（AsyncCopy_PTX.cuh）

封装了 cp.async 指令（Ampere 架构+），支持：

全局内存到共享内存的异步复制
不同粒度的异步屏障（cp.async.wait_group / cp.async.wait_all）

5. FasterTransformer 集成

扩展 DenseWeight 类支持 Tiled-CSL 格式
扩展 cuBlasMMWrapper 类支持动态选择密集/稀疏 MatMul
支持分布式推理（模型并行）

六、实验结果

实验设置

硬件平台：NVIDIA A100-SMX8-80GB（80GB 显存），128 核 Intel Xeon Platinum 8369B CPU @2.90GHz
软件环境：Ubuntu 18.04, CUDA 11.8, Auto-Mixed Precision (AMP)
评估模型：OPT-30B, OPT-66B, OPT-175B
稀疏率：70%, 80%, 90%（随机非结构化稀疏）
Batch size：8, 16, 32, 64
评估指标：Kernel TFLOPs, tokens per GPU-second

基准测试：Kernel 级性能

Flash-LLM vs Sputnik/SparTA（平均加速比）：

稀疏率	vs Sputnik	vs SparTA	vs cuBLAS (dense)
70%	3.6x	1.4x	1.4x
80%	3.0x	1.4x	1.7x
90%	2.0x	1.6x	2.1x

关键发现：

Flash-LLM 在所有稀疏率和 MatMul 形状下均优于所有基线
即使与密集内核 cuBLAS 相比，Flash-LLM 也能通过减少内存访问获得 1.4x-2.1x 加速
稀疏率越高，相对于 cuBLAS 的加速越明显（内存节省更多）

Kernel 分析

GPU 硬件利用率（N=16, 90% sparsity）：

硬件单元	Flash-LLM	cuBLAS (dense)
Tensor Core (TC)	显著提升	基线
L1 + Shared Memory	高利用率	基线
L2 Cache	适中	基线
HBM (全局内存)	降低（稀疏加载）	基线

流水线平衡分析：

Flash-LLM 的三类主要操作：

全局内存访问（Gmem）- 最耗时
共享内存访问（Smem）
Tensor Core 计算（TC）

通过 LSCD 方法，全局内存访问时间显著减少，三类操作更好地重叠，整体性能由 Gmem 决定（Bottleneck effect）。

N 维度影响：

N ≤ 256 时：Flash-LLM 优于 cuBLAS（内存瓶颈主导）
N > 256 时：Flash-LLM 开始慢于 cuBLAS（计算瓶颈出现，cuBLAS 更高效）
这验证了 Flash-LLM 专门针对 Skinny MatMul 的设计目标

端到端模型评估

OPT-30B（单 GPU）：

方法	最大 Batch Size	峰值 tokens/GPU-sec	vs Flash-LLM
Flash-LLM	64	1187	-
DeepSpeed-1GPU	16	348	3.4x slower
FasterTransformer-1GPU	16	359	3.3x slower

Flash-LLM 支持更大 batch size（因稀疏格式减少权重存储）
单 GPU 即可运行 OPT-30B，无需模型并行

OPT-66B（2 GPU）：

方法	GPU 数	最大 Batch Size	峰值 tokens/GPU-sec	vs Flash-LLM
Flash-LLM	2	64	522	-
DeepSpeed-2GPU	2	16	139	3.8x slower
FasterTransformer-2GPU	2	16	144	3.6x slower
DeepSpeed-4GPU	4	64	~280	1.85x slower
FasterTransformer-4GPU	4	64	~310	1.68x slower

Flash-LLM 用 2 GPU 超越 DS/FT 用 4 GPU 的性能
MatMul 时间和跨 GPU 通信时间均显著减少

OPT-175B（4 GPU）：

方法	GPU 数	tokens/GPU-sec (BS=8)	tokens/GPU-sec (BS=64)
Flash-LLM	4	最高	最高
FasterTransformer	8	较低	较低

Flash-LLM 用 4 GPU 成功运行 OPT-175B（传统方案需要 8 GPU）
相比 FT-8GPU，Flash-LLM-4GPU 在 BS=8/16/32/64 分别获得 2.0x/1.9x/1.7x/1.5x 加速
MatMul 时间和跨 GPU 通信时间均显著减少

模型精度评估

模型	任务	原始精度	80% 稀疏精度	精度下降
OPT-30B	SuperGLUE RTE	85.55%	84.11%	1.44%
GPT-NEOX-20B	SuperGLUE RTE	83.03%	82.31%	0.72%

使用 Taylor Pruning 方法
保留前 1/4 和后 1/4 FFN 输入层为密集
80% 稀疏率下精度损失极小

与现有方法对比

方法	核心技术	稀疏类型	Tensor Core 支持	性能特点
Flash-LLM	LSCD + Tiled-CSL	非结构化	是	最优，1.4x-3.8x 加速
Sputnik	SIMT 优化	非结构化	否	98%+ 稀疏率才优于 cuBLAS
SparTA	结构化+SIMT	结构化	部分	受限于结构化约束
cuSPARSE	CUDA 库	结构化/非结构化	部分	为科学计算设计，99%+ 稀疏
cuBLAS	密集 MatMul	无	是	密集基线
FasterTransformer	模型并行	无	是	多 GPU 通信开销大
DeepSpeed	ZeRO + 推理优化	无	是	内存优化但通信开销大

七、总结

核心贡献

首个在 Tensor Core 上高效支持非结构化稀疏 SpMM 的框架：开辟了在高性能 Tensor Core 上探索非结构化稀疏性的新方向
Load-as-Sparse and Compute-as-Dense 方法论：基于 Skinny MatMul 内存瓶颈的洞察，通过稀疏加载减少内存访问、容忍冗余计算来充分利用 Tensor Core
高效软件流水线设计：包括 Tiled-CSL 稀疏格式、稀疏到密集转换、两级重叠策略，有效协调稀疏数据提取、密集数据加载和 Tensor Core 计算
显著的性能提升：
- Kernel 级：比 Sputnik/SparTA 平均快 2.9x/1.5x
- 端到端：比 DeepSpeed/FasterTransformer 最高快 3.8x/3.6x
- 更低的推理成本（更少 GPU 即可运行相同模型）

技术影响

降低 LLM 部署成本：通过稀疏化减少 GPU 需求（如 OPT-175B 从 8 GPU 降至 4 GPU）
推动非结构化稀疏研究：证明了非结构化稀疏在 Tensor Core 上的可行性，此前该方向被忽视
启发后续工作：LSCD 方法论可推广到其他受限于内存带宽的计算场景
开源贡献：代码集成到 FasterTransformer，便于社区使用

局限性

稀疏率范围限制：主要针对 60%-90% 的中等稀疏率优化，极高稀疏率（95%+）或低稀疏率下优势不明显
N 维度限制：当 batch size 较大（N > 256）时，Flash-LLM 性能开始低于 cuBLAS
剪枝成本：基于重训练的剪枝（retraining-based pruning）需要较高的微调成本
硬件依赖：主要针对 NVIDIA A100 GPU 优化，对其他架构（TPU、Intel CPU）需要适配
精度-稀疏率权衡：非结构化稀疏虽然精度损失小，但仍需在精度和效率之间权衡
与量化的正交性：本文聚焦于剪枝，未探索与量化技术的结合

八、参考资源

论文与代码

arXiv 论文: https://arxiv.org/abs/2309.10285
PDF: https://arxiv.org/pdf/2309.10285
GitHub 仓库: https://github.com/AlibabaResearch/flash-llm
VLDB 2024 发表

关键参考文献

Sputnik: Gale et al., “Sparse GPU Kernels for Deep Learning”, SC 2020
SparTA: Zheng et al., “SparTA: Deep-Learning Model Sparsification via Tensor with Compiler-based N:M Mask Support”, OSDI 2022
FasterTransformer: NVIDIA, https://github.com/NVIDIA/FasterTransformer
DeepSpeed: Aminabadi et al., “DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale”, SC 2022
OPT 模型: Zhang et al., “OPT: Open Pre-trained Transformer Language Models”, 2022
cuBLAS/cuSPARSE: NVIDIA CUDA Toolkit 库
Taylor Pruning: Molchanov et al., “Importance Estimation for Neural Network Pruning”, CVPR 2019

文章目录

Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity

Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity

一、论文概述

二、核心思想

问题定义

解决方案概述

三、技术架构

整体框架

核心公式

模型组件

训练/推理流程

稀疏编码格式：Tiled-CSL

四、核心创新

五、代码实现分析

项目结构

关键实现细节

六、实验结果

实验设置

基准测试：Kernel 级性能

Kernel 分析

端到端模型评估

模型精度评估

与现有方法对比

七、总结

核心贡献

技术影响

局限性

八、参考资源

论文与代码

关键参考文献

相关工作

文档