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

# 基于MLIR的大语言模型编译方法
来源：https://arxiv.org/html/2607.15865
彭超 胡  辛志彬  陈一帆  周阳阳  王亮 \{pengchao\.hu,zhibin\.xin,yifan\.chen,yangyang\.zhou,liang\.wang01\}@sophgo\.com 算能科技

###### 摘要

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

## 1 引言

基于Transformer的大语言模型[5 (https://arxiv.org/html/2607.15865#bib.bib3)]已从研究原型演变为生产级服务。PyTorch和HuggingFace Transformers[6 (https://arxiv.org/html/2607.15865#bib.bib4)]等框架极大地简化了训练和实验，但生成的模型权重通常无法直接在专用AI加速器上运行。传统的手动为每种硬件编写算子库的方法成本高昂，且难以跟上模型演进的步伐。因此，行业越来越依赖特定领域的编译器来自动完成从模型描述到硬件指令的转换。

MLIR[4 (https://arxiv.org/html/2607.15865#bib.bib1)]提供了可重用且可扩展的编译器基础设施，允许开发者通过自定义方言在不同抽象级别表达计算。本文提出了一种面向大语言模型的MLIR编译方法，并以TopOp和TpuOp为例具体说明：TopOp首先捕获与框架无关的模型语义，然后通过标准的降级流水线转换为目标特定的TpuOp。为了解决自回归推理中提示处理与令牌生成之间的差异，每个Transformer层被静态编译为三种变体：prefill、prefill\_kv和decode。

本文聚焦于两个主题：如何使用TopOp导入大语言模型并通过TpuOp部署到加速器；以及为何将LLM拆分为prefill、prefill\_kv和decode这三个阶段，以及每个阶段如何在编译流程中表示。我们有意将讨论保持在方法论层面；具体实现细节（如芯片级算子融合、量化调优和层组分片）可根据目标硬件灵活调整。

## 2 背景

### 2.1 MLIR与多级方言

MLIR[4 (https://arxiv.org/html/2607.15865#bib.bib1)]是一种新颖的编译器基础设施，强调可重用且可扩展的中间表示。其核心抽象包括操作、值、类型、属性和方言。一个方言逻辑上将一组相关的操作、类型和属性组织在一起，允许同一程序在不同抽象级别共存，并逐步降级到目标表示。

当面向深度学习编译时，通常定义两个级别的方言：

- •高层图方言：独立于源框架和目标芯片，表达神经网络图的语义，例如MatMul、RMSNorm、Rope、FAttention（全注意力）、Reshape、Concat、MLP等。
- •目标方言：在确定量化模式、数据类型、内存布局和硬件指令后，表达相同的计算。该方言直接面向代码生成。

一个标准的pass流水线负责将高层图方言算子降级为目标方言算子，在需要时执行校准，运行层组和内存规划，最后为目标硬件生成可执行的二进制文件。

### 2.2 TopOp和TpuOp示例

为便于说明，本文使用TopOp和TpuOp作为两个算子方言的具体示例。

TopOp是高层图方言中的算子集合，用于编码深度学习图的语义。它独立于PyTorch和TensorFlow等源框架，也独立于任何特定加速器。典型的TopOp包括：

- •top.MatMul：矩阵乘法；
- •top.RMSNorm：RMS归一化；
- •top.Gather：根据索引获取数据；
- •top.Rope：LLM的RoPE操作；
- •top.Reshape：张量重塑；
- •top.Concat：沿指定轴拼接张量；
- •top.FAttention：LLM的全注意力（FAttention），封装了缩放点积注意力和因果掩码；
- •top.MLP：LLM的MLP算子；
- •top.Weight：引用权重数据。

TopOp通常操作于RankedTensorType类型的张量值，保持高抽象级别，便于进行硬件无关的图变换和等价重写。其中，top.MLP和top.FAttention被保留为融合的高层算子，而非展开为MatMul+激活+残差等基本算子，因为它们在目标芯片上的实现通常需要专用的融合指令序列（例如注意力融合核[1 (https://arxiv.org/html/2607.15865#bib.bib6)]和MLP融合核）。在高层保留融合边界，允许降级阶段为不同芯片选择最优的融合策略，而非被动地在基本算子粒度上进行匹配。

TpuOp是目标硬件方言中的算子集合，由降级pass根据目标芯片特性在TopOp表达语义后生成。与TopOp相比，TpuOp额外携带芯片相关的属性，例如：

- •数据类型（F32、BF16、F16、INT8等）；
- •量化格式支持，包括对称/非对称INT8，以及对已量化权重模型（如AWQ/GPTQ/AutoRound）的直接通过编译；
- •层组信息（group_info），用于片上内存调度；
- •算子拆分信息（例如按核数拆分）；
- •地址信息（所需内存大小和存储偏移）；
- •目标芯片指令封装。

TopOp和TpuOp之间的关系如图1 (https://arxiv.org/html/2607.15865#S2.F1)所示。

参见标题图1：TopOp与TpuOp的关系。
### 2.3 LLM推理概述

LLM的大致执行过程如下：

1. 1. 将输入文本转换为输入令牌；
2. 2. 通过词嵌入获得隐藏状态；
3. 3. 经过num_layers个块生成隐藏状态；
4. 4. 通过LM头获得logits；
5. 5. 通过采样算法从logits中输出一个令牌；
6. 6. 对该令牌重复步骤2–5，直到遇到序列结束符号。

块计算主要包括RoPE、全注意力、MLP等，每个块为每个令牌生成对应的K缓存和V缓存，统称为KV缓存。大模型的整体执行流程如图2 (https://arxiv.org/html/2607.15865#S2.F2)所示。

为便于理解，图中假设：num_layers为32，输入令牌数为100，hidden_size为4096，kv_head为2，head_dim为128；o0、o1、o2表示顺序输出的令牌。

参见标题图2：LLM执行流程。

## 3 使用TopOp构建大型模型

本文的第一个主题是从训练好的检查点到目标硬件二进制的完整路径。以TPU-MLIR[2 (https://arxiv.org/html/2607.15865#bib.bib2)]中的llm_convert.py为例，该脚本加载HuggingFace模型配置和权重分片，根据模型系列（例如，对于纯解码器Transformer、分组查询注意力变体以及多模态语言模型）分派到相应的转换器，然后依次完成TopOp MLIR生成、每个模块降级到TpuOp、编译和结果合并。

### 3.1 模型导入与TopOp模块生成

下面是一个块的TopOp伪代码示例（为简洁起见，省略了部分参数和形状注释）。

1

2input=top.InputOp\("input\_states",0\)

3pos\_ids=top.InputOp\("position\_ids",1\)

4

5x=top.RMSNorm\(input\)

6q=top.MatMul\(x,w\_q\)

7k=top.MatMul\(x,w\_k\)

8v=top.MatMul\(x,w\_v\)

9cos=top.Gather\(rotary\_cos\_w,pos\_ids\)

10sin=top.Gather\(rotary\_sin\_w,pos\_ids\)

11q=top.Rope\(q,cos,sin\)

12k=top.Rope\(k,cos,sin\)

13attn=top.FAttention\(q,k,v\)

14h=top.Add\(input,top.MatMul\(attn,w\_o\)\)

15o=top.Add\(h,top.MLP\(top.RMSNorm\(h\)\)\)

16returntop.Return\(o,k,v\)

清单1：一个块（无历史KV）。1

2input=top.InputOp\("input\_states",0\)

3pos\_ids=top.InputOp\("position\_ids",1\)

4k\_cache=top.InputOp\("history\_k",2\)

5v\_cache=top.InputOp\("history\_v",3\)

6

7x=top.RMSNorm\(input\)

8q=top.MatMul\(x,w\_q\)

9k=top.MatMul\(x,w\_k\)

10v=top.MatMul\(x,w\_v\)

11cos=top.Gather\(rotary\_cos\_w,pos\_ids\)

12sin=top.Gather\(rotary\_sin\_w,pos\_ids\)

13q=top.Rope\(q,cos,sin\)

14k=top.Rope\(k,cos,sin\)

15k\_all=top.Concat\(k\_cache,k\)

16v\_all=top.Concat\(v\_cache,v\)

17attn=top.FAttention\(q,k\_all,v\_all\)

18h=top.Add\(input,top.MatMul\(attn,w\_o\)\)

19o=top.Add\(h,top.MLP\(top.RMSNorm\(h\)\)\)

20returntop.Return\(o,k,v\)

清单2：一个块（带历史KV）。这里，top.Gather用于根据position_ids获取预计算的cos/sin值，而不是在运行时直接计算RoPE；原因在4.6节（性能优化）中说明。

前端构建后，生成的MLIR文件格式大致如下：

1func.func@block\_0\(%arg0:tensor<1x2048x4096xf32\>,\.\.\.\)\-\>\.\.\.\{

2%0="top.RMSNorm"\(%arg0,%weight0\)\.\.\.

3%1="top.MatMul"\(%0,%weight\_q\)\.\.\.

4%2="top.MatMul"\(%0,%weight\_k\)\.\.\.

5%3="top.MatMul"\(%0,%weight\_v\)\.\.\.

6\.\.\.

7%i="top.FAttention"\(%q,%k\_cache,%v\_cache,\.\.\.\)\.\.\.

8\.\.\.

9%n="top.MLP"\(%m\)\.\.\.

10\.\.\.

11return%hidden,%k\_cache,%v\_cache:\.\.\.

12\}

清单3：一个块生成的TopOp MLIR。
### 3.2 从TopOp到TpuOp的降级

生成TopOp模块后，编译器通过降级流水线将其转换为TpuOp。在降级过程中，每个TopOp根据目标芯片的量化模式、数据类型和内存约束被替换为对应的TpuOp，并附加芯片相关属性。例如：

- •top.MatMul降级为tpu.MatMul后，在W4A16/INT8量化模式下附加权重量化参数和缩放因子；对于已量化的权重（如GPTQ/AWQ），编译器可以直接解析其压缩格式并生成对应的TpuOp；
- •top.FAttention降级为对应芯片的融合注意力TpuOp后，后端生成具体的TPU指令序列；
- •所有TpuOp还可能携带group_info（层组信息）、物理地址和算子拆分信息，供后续的层组调度、内存分配和代码生成使用。

TPU-MLIR支持F32、BF16、F16、INT8（对称/非对称）以及AWQ/GPTQ/AutoRound等量化模型的直接通过编译。对于INT8模式，会插入一个校准pass，使用少量样本数据确定每个张量的缩放因子和零点。

降级完成后，编译器继续运行标准优化和代码生成pass：层组分片、片上/片外内存分配和codegen，最终输出每个模块的目标二进制文件。由于大语言模型无法作为一个编译单元处理，前端将模型拆分为多个小的TopOp模块——每个Transformer层变体一个模块，加上嵌入层、语言模型头和辅助头——并独立编译每个模块。最终的二进制文件通过合并所有模块的结果获得。这种模块化方法有两个优点：编译可以跨层并行化；运行时可以按自回归生成所需的顺序调用函数。

从训练好的检查点到部署二进制的整体编译流水线如图3 (https://arxiv.org/html/2607.15865#S3.F3)所示。

参见标题图3：整体编译流水线（检查点→TopOp→TpuOp→二进制）。模块化拆分、并行编译和合并的过程如图4 (https://arxiv.org/html/2607.15865#S3.F4)所示。前端按层和按阶段将模型拆分为多个小的TopOp模块；每个模块独立完成TopOp→TpuOp→codegen，结果合并为最终的部署二进制文件。

参见标题图4：模块化拆分、并行编译和合并。

## 4 三阶段拆分：Prefill、Prefill_kv和Decode

本文的第二个主题是执行模型。自回归LLM推理并不仅仅是一个前向传递，而是在处理（可能很长的）提示和逐个生成令牌之间交替。因此，我们将每个Transformer层拆分为三个TopOp模块变体进行静态编译：prefill、prefill_kv和decode，然后分别降级为对应的TpuOp变体。

### 4.1 为什么拆分为三个阶段？

- •Prefill：并行处理输入提示中的所有令牌。查询长度等于提示长度，并为每个令牌生成Key/Value张量。
- •Decode：仅处理最新生成的令牌，但必须关注所有先前计算的Key和Value。查询长度为1，而Key/Value长度随生成序列增长。
- •Prefill_kv：当运行时要继续对话并需要将历史Key/Value缓存与新提示合并时使用。其行为类似于prefill，但首先将历史缓存连接到新令牌之前。

拆分为三个阶段的主要原因是这三种情况的张量形状和内存布局差异很大；如果编译为单个函数，将需要大量填充或动态重新编译。大多数专用AI加速器是面向张量操作的并行架构，高度依赖静态编译；动态编译不仅增加了运行时编译开销，还降低了内存分配和算子调度的确定性。为了确保尽可能多的部分能够