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

We analyze the performance bottleneck of skinny MatMuls execution starting from dense MatMul workloads in LLMs. According to our profiling results for OPT-66B (Zhang et al., 2022a), the average utilization of tensor cores is around 5.0%, 10.1%, 19.9%, and 39.7% under typical batch sizes of 8,16,32 and 64 as shown in Fig.4, while the bandwidth of global memory is already fully saturated. The underlining cause for this is that the compute intensity (i.e., FLOP/Byte) of skinny MatMul is very low. For a MatMul described in sec. 2.2, the total operations conducted are $2MNK$ floating-point operations (FLOP), and the corresponding data read is $2(MK+KN)$ bytes with FP16 data type. Thus, the computational intensity ($CI$) is:

According to Equation.1, it is easy to demonstrate that the overall $CI$ of a MatMul can be easily restricted by either a small M or N dimension. For instance, for a skinny MatMul where $N$ is 16, $CI$ will have a strict upper bound of 16 no matter how big the $M$ dimension is. Note that in generative LLM models, the $N$ dimension equals the inference batch size which is usually very small in production environments. Thus, the $CI$ is strongly bounded by the $N$ dimension in real-time LLM inference. According to the roofline model (Williams et al., 2009), the performance of a kernel with low computational intensity will be easily bounded by memory bandwidth.

Given that the bottleneck of skinny MatMul comes from memory access/memory bandwidth rather than arithmetic computation, we propose the basic idea of Load-as-Sparse and Compute-as-Dense here which leverages the performance boost from reduced memory access while enabling the efficient use of tensor cores for unstructured sparsity (refer to Sec.4 for details). Under this basic idea, given the sparsity ratio $\beta$ (the weight matrix $A$ is sparse while the feature map matrix $B$ is dense), the computational intensity can be improved to:

Fig.5 shows the CIs and their corresponding achievable tensor core performance of a typical MatMul ($M$: 48k, $N$: BS, $K$: 12k) in OPT-175B model inference with different batch sizes. According to the figure, MatMuls in generative model inference with different batch sizes all face memory wall issues. As a result, the dense MatMuls kernels can only achieve 5.1%, 10.3%, 20.5%, and 40.1% peak performance of tensor cores bounded by insufficient global memory bandwidth. These theoretical values are consistent with our actual measurements in Fig.4. In theory, with Load-as-Sparse and Compute-as-Dense approach under 40% sparsity, the tensor cores utilization can be improved to 8.5%, 17.1%, 34.2%, and 68.2%.

As described in Sec.4.1, it requires several stages to load the sparse encoding from the global memory to shared memory in dense format for each $A_{Tile}$. Specifically, it requires loading sparse encoding from global memory to the distributed registers (gmem2reg stage), resetting the target shared memory buffer with zero (rst_smem stage), and writing the sparse encoding from registers to the corresponding locations on shared memory buffer (extract stage). As for $B_{Tile}$, which is already in dense format, it only requires loading directly from global memory to the target shared memory buffer (ld_dense stage). Finally, the smem2tc stage loads the shared memory data of $A_{Tile}$ and $B_{Tile}$ and executes tensor core computations.

As shown in Fig.6c, Flash-LLM exploits a two-level overlapping of the above memory and computation stages for efficient execution. On one hand, it leverages the software pipeline through double buffering to overlap off-chip memory loads and Sparse-to-Dense transformation with tensor core computation, called inter-iteration overlapping. On the other hand, it overlaps the stages of off-chip memory load within Sparse-to-Dense transformation for more efficient memory activities, called intra-iteration overlapping. The horizontal axis of Fig.6c represents the execution time while the vertical axis represents the activities with the double buffer. It uses two shared memory buffers for $A_{Tile}$ (corresponds to A1 and A2) and $B_{Tile}$ (corresponds to B1 and B2), and one register buffer reused in different iterations (corresponds to SE). Specifically, SE in Iteration-1/3 (Iteration-2/4) and A1 (A2) correspond to the Sparse-to-Dense transformation process of $A_{Tile}$ on the first (second) set of buffer, and B1 (B2) corresponds to the data movement of $B_{Tile}$ on the first (second) set of buffer. As for inter-iteration overlapping, as shown in Iteration-2 in Fig.6c, while loading data from shared memory and executing tensor core computations for data on the first set of buffers, Flash-LLM loads and extracts data from global memory to shared memory for the second set of buffers. As for intra-iteration overlapping, the activities of A1 and B1 are processed in parallel, and the gmem2reg and rst_smem stages on $A_{Tile}$ are also processed in parallel. In this way, the sparse data loading, dense data loading, and tensor core computations can be overlapped efficiently.

A critical design for Sparse-to-Dense transformation is explicitly using registers as data buffers between global memory and shared memory. In Flash-LLM, the sparse encoding movement from global memory to shared memory is explicitly split into two stages, i.e. LDG (loading data from global memory to registers) instructions during gmem2reg and STS (storing data to shared memory from registers) instructions during extract as shown in Fig.6c. On one hand, the split two-stage design helps to increase instruction-level-parallelism (ILP) to hide high global memory access latency. Note that each pair of LDG and STS instructions have load-use dependency. If we do not split the gmem2reg and extract into two stages but launch each pair of LDG and STS instructions in the adjacent cycles (e.g. directly storing to shared memory after loading from global memory) , each GPU thread will execute instructions in the order of $LDG_{0},STS_{0},LDG_{1},STS_{1},...$ without ILP of global memory load (i.e., LDGs). By splitting the two instructions into two stages as in Fig.6c, the execution order will be $LDG_{0},LDG_{1},...,STS_{0},STS_{1},...$ and results in a high ILP of global memory load. On the other hand, splitting the data movement into gmem2reg and extract enables the overlapping opportunity between gmem2reg and rst_smem. Note that STS instructions within extract should not be launched before the completion of rst_smem stage, otherwise, the data written by STS instructions might be overwritten by rst_smem incorrectly. The execution of gmem2reg and rst_smem can be overlapped once the gmem2reg contains no shared memory write (all STS instructions are assigned to the extract stage), which further increases ILP.

Given the complex pipeline in Fig.6c, it requires a set of thread synchronizations and memory barriers to ensure correctness. Flash-LLM inserts the minimum range of synchronizations and memory barriers to ensure correctness while keeping the overlapping.

To avoid the data written by extract being overwritten by rst_smem incorrectly, Flash-LLM inserts the explicit thread-block level synchronization between the two stages to ensure that all the threads have finished their work resetting the A1/A2 buffer in shared memory, shown as the first yellow lines in each iteration in Fig.6c. Meanwhile, it also requires another synchronization to ensure that all data movements and tensor core operations of the current iteration are completed before starting the next iteration, shown as the second yellow lines in each iteration in Fig.6c. Take Iteration-1 as an example, we should make sure that all the threads have finished writing the A1 and B1 shared memory buffers before we start Iteration-2, as the data on the shared memory buffers will be loaded to tensor cores as inputs in Iteration-2. Besides, we have to make sure that all the threads have finished reading the data from the register buffer for extract in Iteration-1 before letting them be overwritten by the rst_smem in Iteration-2.

Beside the thread-block synchronizations, it also requires memory barriers after asynchronous copy activities of global-to-shared data movement. Flash-LLM makes use of the asynchronous copy primitives on GPU for the overlapping of data movement and other activities. Note that the asynchronous copy primitive cp.async, starting from NVIDIA Ampere GPU (NVIDIA, 2020), allows moving data from global memory to shared memory in the background asynchronously while executing other computations in the foreground. Specifically, both the rst_smem and ld_dense stages use the cp.async primitives to enable the overlapped execution on the double buffer.

To enable a fine-grained pipeline execution, Flash-LLM uses different async-copy barriers for rst_smem and ld_dense stages. As shown in Fig.6c, the extract stage waits for the completion of only rst_smem, while the final thread-block barrier of each iteration waits for the completion of all previous cp.async operations. In this way, the extract stage could be overlapped with the ld_dense stages.

Alg.1 shows the implementations of the pipelined computation in Flash-LLM. In line 3, the software pipeline will be initialized, preparing the data of $A_{Tile}$ and $B_{Tile}$ on shared memory for the tensor core computations of the first iteration in the main loop. The following iterations described in Fig.6c is implemented in lines 7-28. For each iteration, it issues the instructions for the asynchronous data loading for the next iteration and does the tensor core computation of the current iteration in a double buffer manner. Specifically, one $A_{Tile}$ for the next iteration will be loaded and extracted from global memory to shared memory (rst_smem, gmem2reg, and extract), and one dense $B_{Tile}$ will be loaded directly from global memory (ld_dense). The rst_smem stage is in line 17, where each thread issues cp.async operation to set buffer A to zeros. In line 19, gmem2reg is accomplished, where sparse encoding is loaded from global memory to the distributed registers. The ld_dense stage is in line 20, where the data for $B_{Tile}$ is loaded from global memory to shared memory buffer with cp.async operations. After launching these asynchronous memory operations, tensor core operations are launched in line 23. Note that we load dense matrices from shared memory to registers using ldmatrix.sync instruction and utilize tensor cores for the core computations by explicitly launching mma.sync instruction in the function Pipelined_Shared2Reg_TensorCoreOps().

The first async-copy barrier in Fig.6c is in line 25, guaranteeing that all asynchronous operations launched in line 17 are completed while the operations launched in line 20 can still be in progress. The extract stage is in line 26, where the data on the distributed registers are extracted to the shared memory buffer of $A_{Tile}$. Finally, the async-copy and thread-block barriers are called in line 28 to make sure that all threads have completed their work in this iteration.

Different from dense MatMul where the data size to be loaded from global memory can be inferred by the tile sizes, the size of sparse encoding is determined by the number of non-zeros (nnz) within $A_{Tile}$, which is unpredictable. Before loading the sparse encoding for each tile, Flash-LLM identifies its start offset and the length in global memory buffers. Such information is maintained in TileOffsets array, which is stored in global memory (more details in Sec.4.3). To avoid instruction stalls caused by long latency global memory access, this metadata should be pre-fetched. In lines 5-6, the start offset and the size of sparse encoding for the first iteration are pre-fetched. At the beginning of each iteration, the start offset and the size of the current $A_{Tile}$ are updated using the value in registers pre-fetched in advance. In lines 11-12, it pre-fetches the start offset and the length for the next iteration.

The sparse encoding format of $A$ matrix is essential for efficient sparse data storage and Sparse-to-Dense Transformation. We propose a tile-by-tile sparse encoding format to support the optimizations in Sec.4.2 effectively. The non-zero elements are organized tile-by-tile, where each tile maintains its non-zero elements accompanied by the sparse index. As shown in Fig.7, the data of each tile within the sparse matrix are encoded into a small array, and combining all tiles will form the overall array (NonZeros Array). The TileOffsets Array maintains the starting offset of each tile in NonZeros Array. The number of non-zero elements of each Tiled-CSL tile is the difference between two corresponding elements in TileOffsets Array. For each tile in NonZeros Array, each element is stored along with its location within the tile. As shown in Fig.7, each non-zero weight is in 16-bit half-precision and each location is encoded into a 16-bit short integer.

As described in Alg.2, each thread extracts $nnz\_thread$ non-zero values from its sparse encoding buffer Reg[] to the shared memory buffer A with a loop. The v() and idx() functions are used to extract the value (high 16 bits) and its location (low 16 bits). There are some special considerations when using registers as intermediate buffers for sparse encoding. Different from shared memory and global memory, GPU registers are not addressable, which means we can not access an array of registers using a variable offset. As a result, forcing an array defined in CUDA into registers requires that, all the indices used to access the array can be determined statically at compile-time. Otherwise, the array will be stored in global memory instead. In line 1 in Alg.2, #pragma unrol is used to notify the GPU compiler to fully unroll the main loop, so that all the indices used to access the $Reg[]$ can be determined statically. Note that adding such compiling directive alone is not enough as a loop with a variable number of iterations can not be fully unrolled. Thus, we use a constant value $\#REG$ (the upper bound of the number of iterations, typically 32/64 in our implementation) in line 2 instead of using the variable value $nnz\_thread$.

There are two types of shared memory access in the computation pipeline for the sparse weight matrix $A$. The one is shared memory load for tensor core computation in smem2tc stage. The other one is shared memory store in extract stage. It is essential to avoid bank conflict for good performance. However, the random sparsity makes it challenging to avoid bank conflict of both shared memory load and store.

As for the smem2tc stage, it makes use of ldmatrix intrinsic for efficient data loading from shared memory to registers for tensor core computation. Fig.8a shows the behavior of ldmatrix where eight threads collectively load an $8\times 8$ matrix in FP16 from shared memory. This $8\times 8$ matrix reading can be served by a single shared memory wavefront⁶⁶6A wavefront is the maximum unit of work that can pass through the GPU hardware pipeline per cycle. At most 1,024 bits can be loaded per wavefront for shared memory. if there is no bank conflict. The bank-conflict-free memory load of ldmatrix requires that all scalars within the $8\times 8$ matrix can be read from disjoint memory banks. Fig.8a shows an example shared memory data layout demonstrating the bank assignment (bank ID ranging from 1 to 32) to achieve bank-conflict-free ldmatrix. Thus, each scalar can be assigned a bank ID according to its location within $A_{Tile}$, given a specific data layout.

However, this requirement will easily cause bank conflict of shared memory store during extract stage due to the random position of the sparse elements in matrix $A$. In other words, we can guarantee the bank-conflict-free load according to the layout requirement of ldmatrix, but cannot avoid bank conflict store in this way during extract stage. Fig.8b gives an example, where each non-zero (NonZero) value should be stored in a target shared memory bank (SMemBank) according to their relative position within $A_{Tile}$ to meet the requirement of bank-conflict-free ldmatrix. As the distribution of NonZeros is random, the target SMemBank for each NonZero value is also random. As a result in Fig.8b, all CUDA WARPs (the NonZeros with the same color are processed by the same WARP) suffer from bank conflict and lead to multiple shared memory wavefront (SMem WF).

To reduce the bank conflict, we propose the ahead of time sparse data reordering approach. The basic insight is that the bank-conflict-free ldmatrix already determines the target bank of each data element, we can thus reorder the data elements in each Tiled-CSL tile so that elements correspond to different banks could be organized into the same WARP for extract execution. Specifically, it iteratively selects the sparse element that corresponds to different memory banks when generating the NonZeros sub-array for each Tiled-CSL tile. Fig.8c shows an ideal case where only one shared memory wavefront is needed to serve one WARP executing extract stage after data reordering. Note that the data reordering is within Tiled-CSL tile scope, which only changes the data distribution on global memory but does not change that on shared memory buffers. In other words, it is a way to change the data placement on global memory to achieve more efficient shared memory access.

Alg.3 shows the algorithm to generate Tiled-CSL format from the original sparse matrix according to ahead of time sparse data reordering approach. The input is the sparse matrix A with $M$ rows and $K$ columns in dense format, where some elements are already set to 0 through model pruning. The outputs are NonZeros and TileOffsets, the key components of Tiled-CSL format. NonZeros are split into groups each of which contains 32 non-zeros in line 2. Each group of non-zeros will be written to shared memory during extract stage within one shared memory request. At lines 7-24, the Tiled-CSL format of one tile ($128\times 64$) will be generated. At lines 11-14, each non-zero within matrix A will be encoded into a 32-bit word containing $(val,loc)$. Besides, non-zeros are distributed to 32 different buckets ($NZ\_Bucket[32]$ at line 4) according to their target shared memory bank ID ranging from 0 to 31 as calculated in line 13. In line 16, the total number of non-zeros (NNZ) is counted. At line 19-22, one group of non-zeros are formed by iteratively picking non-zeros from NZ_Bucket[id] where $NZ\_Bucket[id]$ is the bucket with the most non-zeros not processed at that time.

Model Pruning. To show that pruned model has comparable performance with the original model, we evaluate the accuracy of the pruned model with OPT-30B (Zhang et al., 2022a) and GPT-NEOX-20B (Black et al., 2022; EleutherAI, 2022) on Recognizing Textual Entailment task in SuperGLUE (Wang et al., 2019). To achieve better accuracy, we adopt the popular pruning method Taylor Pruning (Molchanov et al., 2019) to prune these models and keep the front quarter and the last quarter feedforward input layers dense. Based on that, we achieve 80% sparsity on OPT-30B and GPT-NEOX-20B with only 1.44% and 0.72% accuracy decrease, respectively. Specifically, accuracy decreases from 85.55% to 84.11% on OPT-30B, and from 83.03% to 82.31% on GPT-NEOX-20B.

Results. As shown in Fig.13a, Flash-LLM achieves $3.4\times$ and $3.3\times$ higher performance than DeepSpeed (DS) and FasterTransformer(FT) with a single GPU. DS and FT can at most achieve 348 and 359 tokens per GPU-second with a single A100 GPU. If further increase the inference batch size, DS/FT will run out of memory as inference tasks with larger batch sizes need more GPU memory to store the cached-KV and activations. In contrast, Flash-LLM achieves up to 1187 tokens per GPU-second on batch size 64. It is because the memory used for storing model weights is reduced with the Tiled-CSL format, and thus more Cached-KV and activations can be accommodated. We also compare the performance of Flash-LLM to DS and FT with two-way model-parallelism(Shoeybi et al., 2019), with which DS and FT can support batch size 64. As shown in Fig.13b, FT and DS achieve similar performance in terms of tokens per GPU-second. Compared to DS/FT, Flash-LLM achieves $1.91\times$/$1.75\times$, $1.87\times$/$1.70\times$, $1.67\times$/$1.55\times$, and $1.54\times$/$1.41\times$ higher performance at batch sizes 8, 16, 32, and 64 respectively. The detailed GPU memory usage of Flash-LLM, FT, and DS are shown in Table.1.

Breakdown. In order to figure out why Flash-LLM can achieve better performance, we conduct the time breakdown of end-to-end inference shown in Fig.14a. Note that we conduct all the end-to-end breakdowns in this paper leveraging the NSight System (NVIDIA, 2023f). Compared to FT-2GPU (FasterTransformer with 2 GPU used), Flash-LLM with 1 GPU can achieve lower normalized inference latency ⁹⁹9To also take inference cost into consideration and compare the inference efficiency with different system configs (e.g. different numbers of GPUs may be used), we sum the execution time on all used GPUs as the normalized latency. mainly because of 1) the more efficient MatMul and 2) the elimination of cross-GPU communication overhead.

Result As shown in Fig.15a, Flash-LLM achieves $3.8\times$ and $3.6\times$ higher token generation throughput than DS and FT with two GPUs. DS and FT can at most achieve 139 and 144 tokens per GPU-second with batch size 16 as they will run out of memory if further increasing the batch size. In contrast, Flash-LLM achieves up to 522 tokens per GPU-second at batch size 64. We also compare the performance of Flash-LLM to DS-4GPU/FT-4GPU where Flash-LLM still uses two GPUs, while DS-4GPU/FT-4GPU uses four GPUs to enable bigger batch sizes for the baselines. Compared to DS-4GPU/FT-4GPU, Flash-LLM achieves $1.85\times$/$1.68\times$, $1.78\times$/$1.61\times$, $1.7\times$/$1.58\times$, and $1.55\times$/$1.45\times$ higher performance of tokens per GPU-second at batch sizes 8, 16, 32, and 64 respectively.

Breakdown. We conduct the time breakdown of end-to-end inference with FT and Flash-LLM for OPT-66B as shown in Fig.14b. Compared to FT-4GPU (FasterTransformer with 4 GPU used), Flash-LLM with two GPUs can achieve lower normalized inference latency than FT-4GPU mainly because of 1) the reduction of MatMul time and 2) the reduction of cross-GPU communication overhead.

Results & Breakdown We successfully run the inference of OPT-175B models with Flash-LLM using 4 A100 GPUs. In contrast, the weight of OPT-175B can not fit into 4 A100 GPUs with traditional solutions. Thus, we do not show the performance of FT/DS using 4 GPUs here as they all run out of GPU memory. In addition, we failed to run OPT-175B with DS using 8 GPUs. Fig.16a compares the performance of Flash-LLM and FT where we only use 4 GPUs while FT uses 8 GPUs. Compared to FT-8GPU, Flash-LLM achieves $2.0\times$, $1.9\times$, $1.7\times$, and $1.5\times$ higher performance at batch sizes 8, 16, 32, and 64 respectively. As shown in fig.16b, both the MatMul time and the cross-GPU communication time are significantly reduced using Flash-LLM.