基于MLIR的大语言模型编译方法

大语言模型（LLM）已成为现代人工智能加速器上的主要工作负载，但将其部署在专用硬件上仍然面临两个核心挑战：如何将经过训练的模型导入到编译器友好的中间表示中，以及如何在有限的片上内存下有效地调度自回归推理循环。本文提出了一种基于 MLIR（多级中间表示）的 大语言模型 编译方法，并使用两种运算符方言 TopOp 和 TpuOp 进行了说明。 TopOp作为一种独立于源框架和目标芯片的高级图方言，负责表达模型语义； TpuOp 作为目标硬件方言，承载与芯片相关的决策，例如量化、层组和内存布局。模型首先表示为TopOp，然后逐层降低为TpuOp，最后生成可部署的二进制文件。此外，每个Transformer层都分为三个阶段进行静态编译：预填充、prefill_kv（使用历史键值缓存预填充）和解码，以适应提示并行处理和每个词元生成的不同计算特性。该方法已在TPU-MLIR编译器{https://github.com/sophgo/tpu-mlir}和LLM-TPU部署项目{https://github.com/sophgo/LLM-TPU}中实现，支持包括Qwen、Llama、InternVL、MiniCPM-V系列在内的多种生成模型，以及GPTQ、 AWQ 和自动回合。