FlashDecoding++Next: High Throughput LLM Inference With Latency and Memory Optimization

doi:10.1109/TC.2025.3585339

FlashDecoding++Next: High Throughput LLM Inference With Latency and Memory Optimization

0

PRE

AI

DOI：10.1109/TC.2025.3585339

原文链接

原文求助

摘要

En 中文

随着大型语言模型（LLM）在各个领域的重要性日益凸显，LLM推理的性能对于大规模LLM应用至关重要。然而，围绕计算效率和内存利用率，在实现高吞吐量LLM推理方面仍存在以下未解决的问题：（1）同步部分softmax更新。softmax操作需要在每个部分softmax结果之间进行同步更新操作，导致LLM中注意力计算产生约20%的开销。（2）扁平GEMM计算利用率不足。LLM推理中执行GEMM的矩阵形状趋于扁平，导致计算利用率不足，并在现有设计中（如cuBLAS、CUTLASS等）在填充零后性能损失达50%。（3）激活值导致的内存冗余。推理过程中动态分配激活值导致无用变量的冗余存储，增加了22%的内存消耗。我们提出FlashDecoding++Next，一个支持主流LLM和硬件后端的高吞吐量推理引擎。为解决上述挑战，FlashDecoding++Next创造性地提出：（1）基于统一最大值的异步softmax。FlashDecoding++Next为不同的部分softmax计算引入统一最大值技术以避免同步，并在此基础上提出细粒度流水线，分别使LLM推理的prefill和解码阶段吞吐量提升1.18倍和1.14倍。（2）基于双缓冲的扁平GEMM优化。FlashDecoding++Next指出不同形状的扁平GEMM面临不同的瓶颈，并引入双缓冲等技术，使扁平GEMM操作加速最高达52%。（3）缓冲区复用和统一内存管理。FlashDecoding++Next在整个推理过程中复用预分配的激活缓冲区以消除冗余，并在此基础上统一不同类型存储的管理以进一步利用复用机会。内存优化使能处理最长1.57倍长的序列。FlashDecoding++Next展现出显著的吞吐量提升，相比HuggingFace [1]实现最高提升68.88倍吞吐量。平均而言，FlashDecoding++Next相比vLLM [2]和TensorRT-LLM [3]在主流LLM上分别提升1.25倍和1.46倍吞吐量。

Keyword:

Large language model

inference

computation

memory

efficiency

FlashDecoding++Next: High Throughput LLM Inference With Latency and Memory Optimization

摘要

期刊

机构

引用论文