《赵军平-大模型推理显存优化.pdf》由会员分享,可在线阅读,更多相关《赵军平-大模型推理显存优化.pdf(41页珍藏版)》请在三个皮匠报告上搜索。
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