基于MLIR的LLM编译新方法：提升大模型在AI加速器上的部署效率

随着大语言模型成为现代AI加速器上的主要负载，如何在专用硬件上高效部署仍面临两大挑战：一是如何将训练好的模型导入编译器友好的中间表示，二是在有限片上内存下如何高效调度自回归推理循环。近日，一篇发布于arXiv的论文提出了一种基于MLIR的编译解决方案。该方法设计了两种操作方言：TopOp作为独立于源框架和目标芯片的高级图方言，负责表达模型语义；TpuOp作为目标硬件方言，承载量化、层分组、内存布局等芯片相关决策。模型首先被表示为TopOp，然后逐层降级为TpuOp，最终生成可部署的二进制文件。为适应提示并行处理和逐令牌生成的不同计算特性，每个Transformer层被拆分为三个阶段进行静态编译：预填充、带历史键值缓存的预填充和解码。该方法已在TPU-MLIR编译器及LLM-TPU部署项目中实现，支持千问、Llama、InternVL、MiniCPM-V等多个生成模型系列，并兼容GPTQ、AWQ、AutoRound等多种量化与部署形式。