大模型推理的顯存優(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場(chǎng)景）模型參數(shù)量vs.單卡顯存容量單卡算力vs.訪存帶寬發(fā)展顯存管理：從cudaMalloc到CUDA

VMM

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

（基于CUDA

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

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

H4物理地址

H1TensorB1-1B3-2對(duì)齊到

chunk，splitblockB1B3H1H3

H4釋放空閑chunks:H1,

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

cache顯存碎片和attention

kernelvirtualTensor：背景與問題KV

cache顯存碎片問題

vLLM

PagedAttn易用性:attn

kernel改造適配調(diào)優(yōu)的復(fù)雜度、難度；以4~7月計(jì)性能：影響吞吐。業(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)要求高，周期長(zhǎng)業(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目標(biāo)解耦attn

kernel實(shí)現(xiàn)與KV

cache顯存管理3個(gè)小目標(biāo)：碎片問題

+ 性能+

分鐘級(jí)集成-定制各種attnFA3,

Prefix，Sparsity，Quantization

…kernelCUDA

VMMTensorCore(FlashInfer,

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

cache

onvirtualTensorCUDA

VMM

APIsDynamic

MappingAsyncPre-alloc/abs/2407.15309一個(gè)PyTorch

Tensor對(duì)象，專門管理KV

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

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

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

attn

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

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

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

FlashAttn3,

PyTorch-FlexAttn，Double-sparsity

(2024.8)virtualTensor速度對(duì)比效果virtualTensor

+

FA3

2.1x

vs.

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

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

cachePrefill計(jì)算耗時(shí) vs

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

reserve,

alloc,

offloadingPrefill：分層申請(qǐng)KV

cache顯存vs.按prompt長(zhǎng)度一把reservePrefill：分層offloading到CPUKV

cacheDecode：prefetch

KV

cache?

+Scheduler：基于顯存壓力和隊(duì)列請(qǐng)求長(zhǎng)度，動(dòng)態(tài)決定LayerKV效果與對(duì)比: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負(fù)載壓力(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負(fù)載壓力(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)更多模型對(duì)比:7b,34b,70b10X(QPS

-1%)17X(QPS

-2%)6X(QPS

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

(1K)~

3.2X

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

52

120

112

02040801001401601K2K4K8K16K二字時(shí)延（ms）PD PD+

LayerKV1157275

278537

54133783332109724232390050010001500200025003000350040001K2K4K8K16K首字時(shí)延（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）單模型多進(jìn)程推理等1）GMLake:優(yōu)化大模型訓(xùn)練的顯存碎片更大batch、更高吞吐重點(diǎn)：映射元數(shù)據(jù)管理（無(wú)數(shù)據(jù)copy），策略控制GMLake:Efficient

and

Transparent

GPU

Memory

Defragmentation

for

Large-scale

DNN

Training

with

VirtualMemory

Stitching,ASPLOS24,螞蟻集團(tuán)、上海交大等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不同對(duì)比Batch-size,

Strategies,

Distributed

TrainingDevices

and

Different

Training

Engines負(fù)載76

fine-tuneworkloadsAvg

15%Upto

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

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

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

Server、ORT:-

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

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