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

大模型推理的显存优化探索演讲人：赵军平蚂蚁集团 / 异构计算与推理负责人0102030405目录显存需求与挑战围绕显存的更多优化探索优化1: virtualTensor优化KV cache和attn kernel优化2: LayerKV优化TTFT总结自我介绍•赵军平，蚂蚁异构计算与推理负责人•中国计算机协会CCF HPC、存储专委委员，~200 中/美技术专利•异构加速，虚拟化，K8S，推理优化，文件系统，企业级存储-保护等•“数据密集型应用系统设计” 译者显存需求与挑战LLM推理：显存需求Llama-65b, FP16LLM推理：硬件发展•显存容量、访存带宽（特别是推理小batch场景）模型参数量 vs. 单卡显存容量单卡算力 vs. 访存带宽发展显存管理：从cudaMalloc 到 CUDA VMM API2层指针与动态remapping （基于CUDA VMM）-虚拟地址：对用户可见，保证连续-物理地址：CHUNK_SIZE（2MB）粒度分配，不要求连续-Remapping：动态回收chunk与重映射对齐到chunk (2MB*N)虚拟地址物理地址handles显存碎片问题•例子400MB显存碎片原因分析CUDA不支持直接释放部分 (空闲)显存区域访存特征动态变化，LLM更加复杂生命周期不同，大小不同，多stream；强化学习；。。。基于CUDA VMM 的碎片优化思路H1H4Tensor H3B1-1B3-2对齐到chunk，split blockB1B3H1H4H3释放空闲chunks: H1, 3~5新分block(驱动内部对物理显存进行整理)物理地址virtualTensor优化KV cache显存碎片 和 attention kernelvirtualTensor：背景与问题•KV cache显存碎片问题  vLLM PagedAttn•易用性: attn kernel改造适配调优的复杂度、难度；以4~7月计•性能：影响吞吐。业界代表优化：PagedAttn(vLLM), TokenAttnvLLM PagedAttnOn CUDA Core2023.6.202024.1.23/1.31FlashAttn + Paged;FlashInferOn TensorCorevLLM + (FlashAttn + Paged)On TensorCo