大模型推理显存优化-蚂蚁集团.docx

大模型推理的显存优化探索

演讲人：赵军平

蚂蚁集团/异构计算与推理负责人

.

.显存需求与挑战

优化

优化1:virtualTensor优化KVcache和attnkernel

目

目录

.

.优化2:LayerKV优化TTFT

.

.围绕显存的更多优化探索

总结

05

自我介绍

?赵军平，蚂蚁异构计算与推理负责人

?中国计算机协会CCFHPC、存储专委委员，~200中/美技术专利

?异构加速，虚拟化，K8S，推理优化，文件系统，企业级存储-保护等

?“数据密集型应用系统设计”译者

显存需求与挑战

LLM推理：显存需求

Llama-65b,FP16

LLM推理：硬件发展

?显存容量、访存带宽（特别是推理小batch场景）

单卡算力vs.访存带宽发展模型参数量

单卡算力vs.访存带宽发展

显存管理：从cudaMalloc到CUDAVMMAPI

2层指针与动态remapping（基于CUDAVMM）

-虚拟地址：对用户可见，保证连续

-物理地址：CHUNK_SIZE（2MB）粒度分配，不要求连续

-Remapping：动态回收chunk与重映射

对齐到chunk(2MB*N)

虚拟地址厂TTT

物理地址handles

显存碎片问题

?例子

400MB

400MB

显存碎片原因分析

CUDA不支持直接释放部分(空闲)显存区域

访存特征动态变化，LLM更加复杂

n生命周期不同，大小不同，多stream；强化学习；。。。

基于CUDAVMM的碎片优化思路

Tensor物理地址

H1

H3

H4

B1

B1

B3

B1-1

B1-1

新分block

(驱动内部对物理显存进行整理)

B3-2对齐到chunk，splitblock

B3-2

对齐到

H3H4H1释放闲chunks:H1,35

H3

H4

H1

virtualTensor

优化KVcache显存碎片和attentionkernel

virtualTensor：背景与问题

?KVcache显存碎片问题vLLMPagedAttn

?易用性:attnkernel改造适配调优的复杂度、难度；以4~7月计

?性能：影响吞吐。业界代表优化：PagedAttn(vLLM),TokenAttn

vLLMPagedAttnOnCUDACore

2023.6.20

FlashAttn+Paged;FlashInfer

OnTensorCore

2024.1.23/1.31

vLLM+(FlashAttn+Paged)

OnTensorCore

2024.6.1

210Days

120Days

PagedAttn不足-1

?attnkernel额外适配调优

?要求高，周期长

?业界各种attn仍创新不断

?FlashAttn3,SageAttn1/2

?SparseAttn,quant,

?SeerAttn,DuoAttn,reAttn

?FlexAttn,FlashMask,

?…

PagedAttn不足-2

?性能开销

?特别早期版本：GQA；不支持TensorCore

?尚不支持FlashAttn3

GPUHW

KVcachedefrag

Defragreleasedate

LoCforDefrag

HFPyTorch

TensorCore

-

-

-

FlashAttn

TensorCore

-

2023.4.12

vLLMPagedAttn

CUDACore

Pageable

2023.6.20

-

FlashAttnd

TensorCore

Pageable

2024.1.23

280

FlashInfer

TensorCore

Pageable

2024.1.31

-

vLLM-FlashAttnd

TensorCore

Pageable

2024.6.1

365

Ours

TensorCore

virtualTensor

2200

virtualTensor目标

?解耦attnkernel实现与KVcache显存管理

?3个小目标：碎片问题+性能+分钟级集成-定制各种attnkernelCUDAVMM