大模型推理的顯存優(yōu)化探索01020304目錄顯存需求與挑戰(zhàn)優(yōu)化1:virtualTensor優(yōu)化KV

cache和attnkernel優(yōu)化2:LayerKV優(yōu)化TTFT圍繞顯存的更多優(yōu)化探索總結(jié)05顯存需求與挑戰(zhàn)LLM推理：顯存需求Llama-65b,

FP16LLM推理：硬件發(fā)展顯存容量、訪存帶寬（特別是推理小batch場景）模型參數(shù)量vs.單卡顯存容量單卡算力vs.訪存帶寬發(fā)展顯存管理：從cudaMalloc到CUDA

VMM

API2層指針與動態(tài)remapping

（基于CUDA

VMM）虛擬地址：對用戶可見，保證連續(xù)物理地址：CHUNK_SIZE（2MB）粒度分配，不要求連續(xù)Remapping：動態(tài)回收chunk與重映射對齊到chunk(2MB*N)虛擬地址物理地址handles顯存碎片問題例子400MB顯存碎片原因分析CUDA不支持直接釋放部分(空閑)顯存區(qū)域訪存特征動態(tài)變化，LLM更加復(fù)雜生命周期不同，大小不同，多stream；強化學(xué)習；。。?；贑UDA

VMM的碎片優(yōu)化思路H3

H4物理地址

H1TensorB1-1B3-2對齊到

chunk，splitblockB1B3H1H3

H4釋放空閑chunks:H1,

3~5新分block(驅(qū)動內(nèi)部對物理顯存進行整理)virtualTensor優(yōu)化KV

cache顯存碎片和attention

kernelvirtualTensor：背景與問題KV

cache顯存碎片問題

vLLM

PagedAttn易用性:attn

kernel改造適配調(diào)優(yōu)的復(fù)雜度、難度；以4~7月計性能：影響吞吐。業(yè)界代表優(yōu)化：PagedAttn(vLLM),

TokenAttnvLLM

PagedAttnOn

CUDA

Core2023.6.202024.1.23/1.31FlashAttn

+

Paged;FlashInferOnTensorCorevLLM

+

(FlashAttn

+

Paged)OnTensorCore2024.6.1120

Days210

DaysPagedAttn不足-1attn

kernel額外適配調(diào)優(yōu)要求高，周期長業(yè)界各種attn仍創(chuàng)新不斷FlashAttn3,

SageAttn1/2SparseAttn,

quant,SeerAttn,

DuoAttn,

reAttnFlexAttn,

FlashMask,…PagedAttn不足-2性能開銷特別早期版本：GQA；不支持TensorCore尚不支持FlashAttn3GPU

HWKV

cachedefragDefrag

releasedateLoC

forDefragHFPyTorchTensorCore---FlashAttnTensorCore-2023.4.12vLLM

PagedAttnCUDA

CorePageable2023.6.20-FlashAttndTensorCorePageable2024.1.23280FlashInferTensorCorePageable2024.1.31-vLLM-FlashAttn-pagedTensorCorePageable2024.6.1365OursTensorCorevirtualTensor2200virtualTensor目標解耦attn

kernel實現(xiàn)與KV

cache顯存管理3個小目標：碎片問題

+ 性能+

分鐘級集成-定制各種attnFA3,

Prefix，Sparsity，Quantization

…kernelCUDA

VMMTensorCore(FlashInfer,

…)PagedAttnvirtualTensor(Ours)NaivevirtualTensor方案設(shè)計KV

cache

onvirtualTensorCUDA

VMM

APIsDynamic

MappingAsyncPre-alloc/abs/2407.15309一個PyTorch

Tensor對象，專門管理KV

cache基于VMM；用戶不需要修改Attention

kernel；內(nèi)置性能優(yōu)化virtualTensor實現(xiàn)virtualTensor接口virtualTensor顯存管理效果KV

cache顯存碎片基本消除；~=PagedAttn且可以按需彈性增長virtualTensor

attn

kernel適配效果3行代碼，1分鐘：完成新attn

kernel到vLLM適配Kernel不需修改，只需要import

package，替換函數(shù)的2個參數(shù)demo:

FlashAttn3,

PyTorch-FlexAttn，Double-sparsity

(2024.8)virtualTensor速度對比效果virtualTensor

+

FA3

2.1x

vs.

FA2LayerKV優(yōu)化顯存不足而導(dǎo)致TTFT耗時激增問題LayerKV：背景問題TTFT=prefill耗時+排隊耗時等顯存，…Llama2-7b,L20*單卡，1QPS原因分析vLLM等按請求prompt長度一次性分配（reserve）KV

cache顯存如果顯存不足，等待……直到decode釋放足夠多顯存；高負載時導(dǎo)致TTFT急劇增長LayerKV思路Prefill分層實際消耗KV

cachePrefill計算耗時 vs

PCIe傳輸分析LayerKV設(shè)計Layer-wise

reserve,

alloc,

offloadingPrefill：分層申請KV

cache顯存vs.按prompt長度一把reservePrefill：分層offloading到CPUKV

cacheDecode：prefetch

KV

cache?

+Scheduler：基于顯存壓力和隊列請求長度，動態(tài)決定LayerKV效果與對比:Llama2-7b，1024->300，L2034243027首字延遲<1sec

&&生成>7tokens/sec下QPS提升2+X

(1.5

->

3+)2398177512243624556871897186919582109182052092172172192222242530500100015002000250030003500400011.1

1.2

1.3

1.4

1.45

1.5

1.6

1.7

1.8

1.92負載壓力(Req/sec)首字延遲平均(msec):llama2-7b,1024-300,L20基線-vLLM

LayerKV27593770474568538246900510780102924078374644553

4518470

472

477

491

501

490

50368702000400060008000100001200011.11.21.31.4

1.45

1.51.61.71.81.92負載壓力(Req/sec)首字延遲P99(msec):llama2-7b,1024-300,L20基線-vLLM

LayerKV生成速度不影響下(QPS:1.4)平均首字延遲降低1.8XP99首字延遲降低5.9X首字延遲<1sec

&&生成>10tokens/sec下：QPS提升1.3X

(1.5->2)更多模型對比:7b,34b,70b10X(QPS

-1%)17X(QPS

-2%)6X(QPS

-3%)LayerKV+PD分離:效果對比二字時延相比vLLM優(yōu)化1.4X

(1K)~

3.2X

(16K時)，TTFT維持不變61791383839404360

52

120

112

02040801001401601K2K4K8K16K二字時延（ms）PD PD+

LayerKV1157275

278537

54133783332109724232390050010001500200025003000350040001K2K4K8K16K首字時延（ms）PD PD+

LayerKVvLLM

0.5.5(backport

vLLM

0.6PD分離),QWen2.5-14b,L20

*2,

1P1D3.2X2.8X2X1.6X13.7

4X圍繞顯存的更多優(yōu)化1）優(yōu)化大模型訓(xùn)練，2）單模型多進程推理等1）GMLake:優(yōu)化大模型訓(xùn)練的顯存碎片更大batch、更高吞吐重點：映射元數(shù)據(jù)管理（無數(shù)據(jù)copy），策略控制GMLake:Efficient

and

Transparent

GPU

Memory

Defragmentation

for

Large-scale

DNN

Training

with

VirtualMemory

Stitching,ASPLOS24,螞蟻集團、上海交大等GMLake:大模型訓(xùn)練顯存優(yōu)化效果硬件環(huán)境16

×

NVIDIA

A100

SXM

（80GB）訓(xùn)練框架FSDP,

DeepSpeed

ZeRO,

Colossal-AI

Gemini模型8

LLMs

includingOPT,

Vicuna,

GPT-2,

GPT-NeoX

and

so

on不同對比Batch-size,

Strategies,

Distributed

TrainingDevices

and

Different

Training

Engines負載76

fine-tuneworkloadsAvg

15%Upto

33%2）多進程推理顯存優(yōu)化單模型多進程實例（python）顯存占用大跨進程的權(quán)重、中間結(jié)果顯存、CUDA

context等權(quán)重(RO)：全局、細粒度、透明共享，類似存儲

“重刪”全局：支持跨線程、跨進程、跨容器共享透明：無需用戶參與、無需配置細粒度：層或者block級共享，基于hash和字節(jié)比對中間結(jié)果：多模型（DAG）、跨進程中間結(jié)果的時分復(fù)用跨進程、跨模型透明共享核心：虛擬-顯存指針的動態(tài)管理與映射跨進程export、引用計數(shù)等不同于Triton

Server、ORT:-

粒度：模型所有W完全相同vs.顯存塊級（通常比tensor粒度更?。┓秶簡芜M程使用：手工配置vs.跨進程、跨容器vs.自動發(fā)現(xiàn)，透明使用顯存優(yōu)化效果6個進程總顯存13.2->4.6GB

(-65%)權(quán)重:激活: