PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashDecoding++: Faster Large Language Model Inference on GPUs

Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Yuhan Dong, Yu Wang

arXiv:2311.01282v4cs.LGcs.CL

TL;DR

LLM inference faces synchronization overhead in partial softmax, under-utilized flat GEMMs, and performance loss from static dataflows. FlashDecoding++ combines asynchronous softmax, flat-GEMM optimization, and adaptive dataflow, reporting speedups over Hugging Face and FlashDecoding across GPUs and LLMs.

  • Problem

    LLM inference faces partial-softmax synchronization overhead, under-utilized flat GEMMs, and performance loss when static dataflows encounter varied input shapes and hardware.

    Source: §1 · p. 2, §3 · p. 5

  • Method

    FlashDecoding++ combines unified-maximum asynchronous partial softmax, flat-GEMM optimization with double buffering, and heuristic dataflow adaptation to input and hardware.

    Source: §1 · p. 3

  • Results

    FlashDecoding++ achieves up to 4.86× speedup on NVIDIA GPUs and 3.93× on AMD GPUs over Hugging Face implementations, plus 1.37× average speedup over FlashDecoding across various LLMs.

    Source: §1 · p. 3, §8 · p. 14

  • Takeaways & Limitations

    The engine supports mainstream LLMs on multiple hardware backends and reports gains over Hugging Face and state-of-the-art inference engines.

    Source: §8 · p. 14

    For Llama2-7B, the selected softmax scaling factor is based on the observation that more than 99.99% of input values lie between −16.8 and 6.5.

    Source: §3 · p. 6

Abstract

from arXiv · show

As the Large Language Model (LLM) becomes increasingly important in various domains. However, the following challenges still remain unsolved in accelerating LLM inference: (1) Synchronized partial softmax update. The softmax operation requires a synchronized update operation among each partial softmax result, leading to ~20% overheads for the attention computation in LLMs. (2) Under-utilized computation of flat GEMM. The shape of matrices performing GEMM in LLM inference is flat, leading to under-utilized computation and >50% performance loss after padding zeros in previous designs. (3) Performance loss due to static dataflow. Kernel performance in LLM depends on varied input data features, hardware configurations, etc. A single and static dataflow may lead to a 50.25% performance loss for GEMMs of different shapes in LLM inference. We present FlashDecoding++, a fast LLM inference engine supporting mainstream LLMs and hardware back-ends. To tackle the above challenges, FlashDecoding++ creatively proposes: (1) Asynchronized softmax with unified max value. FlashDecoding++ introduces a unified max value technique for different partial softmax computations to avoid synchronization. (2) Flat GEMM optimization with double buffering. FlashDecoding++ points out that flat GEMMs with different shapes face varied bottlenecks. Then, techniques like double buffering are introduced. (3) Heuristic dataflow with hardware resource adaptation. FlashDecoding++ heuristically optimizes dataflow using different hardware resource considering input dynamics. Due to the versatility of optimizations in FlashDecoding++, FlashDecoding++ can achieve up to 4.86x and 2.18x speedup on both NVIDIA and AMD GPUs compared to Hugging Face implementations. FlashDecoding++ also achieves an average speedup of 1.37x compared to state-of-the-art LLM inference engines on mainstream LLMs.

1 Introduction

FlashDecoding++ targets synchronization overhead, under-utilized flat GEMMs, and static dataflows in LLM inference with three corresponding optimizations. It reports speedups against Hugging Face and FlashDecoding across GPU backends and LLMs.

1 Introduction

  • The prefill phase mainly uses GEMM, while decode uses GEMV or flat GEMM; partial-softmax synchronization, padding, and static dataflows constrain performance.Profiling reports 18.8% attention-computation overhead from synchronized updates, while prior padding designs can under-utilize computation by over 50%.Source: §1 · p. 2, Figure 2 · p. 3
  • A unified maximum lets partial softmax results be processed individually without synchronized updates.This addresses the synchronization required when partial vectors have different maximum values.Source: §1 · p. 3
  • For flat GEMMs, FlashDecoding++ pads matrix size to 8 rather than 64 and uses techniques such as double buffering to improve utilization across varied bottlenecks.The figure compares FlashDecoding++ with state-of-the-art designs using Llama2-7B; its left panel uses batch size 1 and input length 1K, and its right panel covers first-token and per-token latency.Source: Figure 1 · p. 2, §1 · p. 3
  • FlashDecoding++ dynamically applies kernel optimizations according to input dynamics and hardware configurations.The introduction identifies static dataflow as a source of performance loss across GEMM shapes.Source: §1 · pp. 2–3
  • 1.37× average speedup over FlashDecoding on various LLMs, with up to 4.86× on NVIDIA GPUs and 3.93× on AMD GPUs over Hugging Face implementations.The comparison includes models such as Llama2 and ChatGLM2.Source: §1 · p. 3

2 Background

LLM inference generates tokens from an input sequence through prefill and autoregressive decode phases. Its operations include linear projections, attention, and feedforward layers, while partial softmax reduces memory use but introduces synchronization.

2.1 LLM Inference Dataflow Overview

  • Inference uses a prefill phase to process the prompt and generate the first token, followed by autoregressive decode steps that generate subsequent tokens.The decode phase’s output token becomes the input to the next decode step.Source: §2.1 · p. 4

2.2 Operations in LLM Inference

  • Transformer inference includes linear projections, attention, and feedforward operations; decode inputs are flat-shaped matrices or vectors because few tokens are processed at once.The feedforward network expands feature dimensions and then restores them.Source: §2.2 · pp. 4–5

2.3 Attention Optimization

  • Partial softmax divides a vector into smaller vectors to reduce memory cost and improve parallelism, but updating their results requires synchronization.For Llama2-7B on an NVIDIA Tesla A100 with input length 1024, synchronized updates add 18.8% overhead to attention computation.Source: §2.3 · p. 5

3 Asynchronized Softmax with Unified Maximum Value

FlashDecoding++ avoids synchronizing partial softmax computations by sharing a scaling value across them. It recomputes with the synchronized scheme when values fall outside a safe range.

3 Asynchronized Softmax with Unified Maximum Value

  • Synchronization among partial vectors in partial softmax causes approximately 20% overhead in the attention operation, because each partial maximum is computed before previous partial-softmax results are updated.Partial maxima differ and are used to avoid exponent overflow; the exponents are summed as the softmax denominator.Source: §3 · pp. 5–6
  • The shared scaling value can be arbitrary mathematically, but exponent overflow or precision loss constrains its safe range.For Llama2-7B, more than 99.99% of xi fall between −16.8 and 6.5; the technique is not applied to OPT-6.7B because its range is large.Source: §3 · p. 6
  • A unified scaling value lets partial softmax computations proceed independently, avoiding synchronization among partial results.Their inner numerator and denominator accumulations use only their own partial vectors; the outer accumulation waits until all partials finish.Source: §3 · p. 6
  • When an input falls outside the selected range, the method stops asynchronous processing for its partial vector and recomputes synchronously.This fallback is described as avoiding overflow with negligible overhead based on the statistical data.Source: §3 · p. 7

4 Flat GEMM Optimization with Double Buffering

FlashDecoding++ targets decode-phase GEMV and flat GEMM, where padding small M dimensions to 64 wastes computation. It pads to 8 and adapts N tiling and buffering to workload size.

4 Flat GEMM Optimization with Double Buffering

  • Padding decode-phase GEMV or flat GEMM from a small M dimension to 64 underutilizes computation; padding to 8 targets this waste.The paper identifies GEMV and flat GEMM as decode workloads with M much smaller than 64.Source: §4 · p. 7
  • Small-N flat GEMMs are parallelism-bounded, while large-N cases are memory-bounded and can benefit from hiding memory latency.The N-tile size trades computation-to-memory ratio against parallelism; the paper relates the latter to GPU hardware parallelism.Source: §4 · p. 8
  • For large N, double buffering overlaps GEMM computation on one shared-memory buffer with loading the next tile into another.The M dimension is padded to 8 and is not tiled in the illustrated configuration.Source: §4 · p. 8, Figure 8 · p. 9

5 Heuristic Dataflow with Hardware Resource Adaption

Because optimal GEMM implementations vary with workload and GPU characteristics, FlashDecoding++ profiles implementation crossovers and selects a dataflow at runtime.

5 Heuristic Dataflow with Hardware Resource Adaption

  • Tensor Core and CUDA Core implementations favor different workloads: on A100, cuBLAS Tensor Core reaches 82.15% of FastGEMV performance for Llama2-7B decode at one batch, while CUDA Core reaches 49.75% of Tensor Core performance at batch size four.This motivates selecting among implementations rather than using one static dataflow.Source: §5 · p. 9
  • Input dynamics, model diversity, GPU capacity, and engineering effects make mapping workloads to optimal implementations a large search problem.The paper cites batch size, sequence length, model structure, memory bandwidth, cache size, and compute capacity as factors.Source: §5 · p. 9
  • For each LLM, only four [K, N] shapes occur, enabling offline profiling to find M thresholds between FastGEMV, the flat-GEMM kernel, and CUTLASS.Each [N, K] shape gets individual crossover points M1 and M2.Source: §5 · pp. 9–10
  • At runtime, FlashDecoding++ selects CUDA Core below M1, the flat-GEMM implementation between M1 and M2, and CUTLASS at or above M2.The offline decision flow forms a lookup table and does not affect runtime inference performance.Source: §5 · p. 10, Figure 9 · p. 10

6 Evaluation

Across three LLMs, FlashDecoding++ improves decode and prefill performance against Hugging Face and reports average gains against several inference engines. Evaluation covers NVIDIA and AMD GPUs, with broader baseline coverage on NVIDIA.

6.1 Experiments Setup

  • The evaluation compares performance across different GPUs and LLMs against state-of-the-art inference engines.Source: §6.1 · p. 10

6.1.1 Hardware Platforms

  • Tests span Tesla A100 and RTX3090 GPUs from NVIDIA and MI210 and RX7900XTX GPUs from AMD.Source: §6.1.1 · p. 11

6.1.2 LLM Engine Baselines

  • The evaluation compares FlashDecoding++ with Hugging Face, vLLM, DeepSpeed, TensorRT-LLM, OpenPPL, and FlashAttention2/FlashDecoding.Comparisons cover both prefill and decode phases.Source: §6.1.2 · p. 11

6.1.3 Models

  • The models are Llama2, OPT, and ChatGLM2, including model variants with different configurations.Llama2 ranges from 7B to 70B parameters; OPT ranges from 125M to 175B, and ChatGLM2 supports bilingual Chinese-English chat.Source: §6.1.3 · p. 11

6.2 Comparison with State-of-the-art

  • Decode speedup reaches 4.86× over Hugging Face, while average speedups are 1.24×, 1.44×, 1.13×, 1.24×, and 1.21× over vLLM, DeepSpeed, TensorRT-LLM, OpenPPL, and FlashDecoding, respectively.Prefill reaches up to 1.40× over Hugging Face, with average speedups of 1.05×, 1.06×, 1.08×, 1.09×, and 1.08× over DeepSpeed, TensorRT-LLM, OpenPPL, FlashAttention2, and FlashDecoding.Source: §6.2 · p. 11

(a) Llama2-7B@A100

The figure lists input-length settings.

(b) OPT-6.7B@A100

The figure lists input-length settings.

(c) ChatGLM2-6B@A100

The figure lists input-length settings.

(d) Llama2-7B@3090

The supplied material identifies GPU speedup figures and summarizes several earlier LLM inference approaches.

(d) Llama2-7B@3090

  • Figure 10 reports decode-phase speedups on NVIDIA GPUs; blank bars indicate models that cannot be executed.Examples include OpenPPL not supporting OPT-6.7B/ChatGLM2-6B and TensorRT-LLM failing to compile models with input lengths over 8K.Source: Figure 10 · p. 12

7 Related Works

  • DeepSpeed combines kernel fusion with GPU-memory management, while vLLM uses PageAttention and memory management to increase maximum batch sizes.DeepSpeed is described as supporting both training and inference, with a particular focus on KV-cache memory use.Source: §7 · pp. 12–13
  • The supplied captions identify NVIDIA prefill speedups and decode speedups on AMD RX7900XTX and MI210 GPUs.Source: Figure 11 · p. 13, Figure 12 · p. 14, Figure 13 · p. 14
  • FlashAttention optimizes prefill-phase self-attention, and FlashDecoding extends it with split K and V for decode-phase parallelism on long sequences.Source: §7 · p. 13
  • FasterTransformer and OpenPPL use C++ inference engines to reduce kernel-scheduling overhead; TensorRT-LLM builds on TensorRT and FasterTransformer and incorporates open-source technologies including FlashAttention.Source: §7 · p. 13

8 Conclusion

FlashDecoding++ is a fast LLM inference engine with multiple hardware backends, combining three kernel and dataflow designs. It reports speedups over Hugging Face implementations and state-of-the-art inference engines.

8 Conclusion

  • Up to 4.86× and 3.93× speedup on NVIDIA and AMD GPUs, respectively, compared with Hugging Face implementations.Source: §8 · p. 14
  • FlashDecoding++ achieves an average 1.37× speedup over the state-of-the-art inference engine FlashDecoding on various LLMs.Source: §8 · p. 14
  • The engine combines asynchronized softmax with a unified max value, flat GEMM optimization with double buffering, and heuristic dataflow with hardware resource adaptation.Source: §8 · p. 14

Related papers

Loading 2311.01282v4…