赵军平-大模型推理显存优化.pdf

1、大模型推理的显存优化探索演讲人：赵军平蚂蚁集团/异构计算与推理负责人01010202030304040505目录显存需求与挑战围绕显存的更多优化探索优化1:virtualTensor优化KV cache和attn kernel优化2:LayerKV优化TTFT总结自我介绍 赵军平，赵军平，蚂蚁异构计算与推理异构计算与推理负责人 中国计算机协会CCF HPC、存储专委委员，200 中/美技术专利 异构加速，虚拟化，K8S，推理优化，文件系统，企业级存储-保护等“数据密集型应用系统设计”译者显存需求与挑战LLM推理：显存需求Llama-65b,FP16LLM推理：硬件发展 显存容量容量、访存带宽带

2、宽（特别是推理小batch场景）模型参数量 vs.单卡显存容量单卡算力 vs.访存带宽发展显存管理：从cudaMalloc 到 CUDA VMM API2层指针与动态remapping（基于CUDA VMM）-虚拟地址虚拟地址：对用户可见，保证连续-物理地址物理地址：CHUNK_SIZE（2MB）粒度分配，不要求连续-RemappingRemapping：动态回收chunk与重映射对齐到chunk(2MB*N)虚拟地址物理地址handles显存碎片问题 例子400MB显存碎片原因分析CUDA不支持直接释放部分(空闲)显存区域访存特征动态变化，LLM更加复杂 生命周期不同，大小不同，多strea

3、m；强化学习；。基于CUDA VMM 的碎片优化思路H1H4Tensor H3B1-1B3-2对齐到chunk，split blockB1B3H1H4H3释放空闲chunks:H1,35新分block(驱动内部对物理显存进行整理)物理地址virtualTensor优化KV cache显存碎片 和 attention kernelvirtualTensor：背景与问题 KV cacheKV cache显存碎片显存碎片问题 vLLM PagedAttn 易用性:attn kernelattn kernel改造适配调优改造适配调优的复杂度、难度；以47月计 性能：影响吞吐。业界代表优化：PagedA

4、ttn(vLLM),TokenAttnvLLM PagedAttnOn CUDA Core2023.6.202024.1.23/1.31FlashAttn+Paged;FlashInferOn TensorCorevLLM+(FlashAttn+Paged)On TensorCore2024.6.1210 Days120 DaysPagedAttn不足-1 attn kernel额外适配调优 要求高，周期长 业界各种attn仍创新不断FlashAttn3,SageAttn1/2SparseAttn,quant,SeerAttn,DuoAttn,reAttnFlexAttn,FlashMask,

5、PagedAttn不足-2 性能开销 特别早期版本：GQA；不支持TensorCore 尚不支持FlashAttn3 GPU HWKV cache defragDefrag release dateLoC for DefragHF PyTorchTensorCore-FlashAttnTensorCore-2023.4.12vLLM PagedAttnCUDA CorePageable2023.6.20-FlashAttn-pagedTensorCorePageable2024.1.23280FlashInferTensorCorePageable2024.1.31-vLLM-FlashAtt

6、n-pagedTensorCorePageable2024.6.1365OursTensorCorevirtualTensor2200virtualTensor目标 解耦attn kernel实现与KV cache显存管理 3 3个小目标：碎片问题 +性能 +分钟级集成分钟级集成-定制各种定制各种attnattn kernelCUDA VMMTensorCore(FlashInfer,)FA3,Prefix，Sparsity，Quantization PagedAttnvirtualTensor(Ours)NaivevirtualTensor 方案设计KV cache on virtualTe