图优化：MLIR量化Pass实现

可视化演示

展示MLIR量化Pass：浮点图转定点图的仿射量化映射q=clamp(round(r/scale)+zp)，PTQ与QAT两种统计注入工作流，及量化沿算子传播的scale推导。

量化原理基础

改造电路之前，工人们必须先理解“电压”这个最基础的概念。一台电器的铭牌上写着额定电压，变压器把 220V 降到 12V 靠的是一个固定的“倍率”。如果倍率选错了，要么电压太低设备罢工，要么电压太高直接烧毁。量化也是一样，它要把模型里连续变化的浮点数，映射到只有 256 个取值的 8bit 整数空间，这个映射的“倍率”就是 scale，“零位对齐”就是 zero point。

从数学上看，量化是一种线性仿射变换。假设真实值是 $r$，量化后的整数是 $q$，步长 $S$ 表示整数每走一格对应的实数大小，零点 $Z$ 表示实数零落在哪个整数刻度上，那么完整公式为：

量化还分对称与非对称两种风格。对称量化令 $Z=0$，整数范围正负对称，适合权重这种正负值都有的张量；非对称量化保留 $Z$，适合 ReLU 之后全为非负的激活值。精度损失主要来自三处：舍入误差（取整时丢掉的小数）、裁剪误差（超出范围被强行截断）、以及溢出误差（累加时整数越界）。理解这些“电力损耗”来源，才能在设计 Pass 时知道该保护什么。

回到电路改造的比喻：变压器倍率定得准不准，决定了所有电器能否安全运行。scale 与 zero point 就是整栋楼变压器的基础参数，后面的所有施工流程，都建立在这两个数值之上。

MLIR核心机制

一栋楼的改造工程不可能让每个施工队各画各的图纸——电工用一套符号，水管工用另一套，最后根本对不上。MLIR 的价值，正是提供了这样一套“统一的施工图纸标准”。它把不同抽象层级的计算表示都收编成统一的 IR，每种表示称为一个“方言”（dialect）：浮点计算用 F32 方言，量化后的计算用 Q 方言，底层指令又属于 LLVM 方言。每支施工队只认自己专业的方言图纸，但所有图纸都遵循同一套规范。

在 MLIR 里，Pass 就是一支支独立施工队。编译器把一个又一个 Pass 串成流水线（pass pipeline），就像总承包商排出施工顺序：先拆旧电路、再装变压器、最后验收测试。Pass 管理机制负责决定每支队伍何时进场、分析哪些信息、输出怎样的新图纸。其中最重要的技术是“模式重写”（pattern rewrite）：在 IR 图上搜索某个子图模式，一旦匹配就整体替换成新模式。例如把“组合插座 + 变压器”两个零件合并成一个“带变压器的组合插座”，这就是一条重写规则。

Pass 遍历 IR 的方式也有讲究。代码结构上，MLIR 的 IR 呈模块→函数→块→操作的树状层级，Pass 可以自上而下或自下而上地扫描。比如“常量折叠”会先处理叶子节点再向上传播结果，而“内联”则要自顶向下找到调用点。遍历顺序直接影响重写规则能否高效触发，好比施工队进场顺序错了，刚装好的墙又要凿开重走管线。

这就像总承包商的调度室：PassManager 拿着施工计划表，逐个调用施工队，每队看完图纸改完图，交给下一队。MLIR 的强大在于这套调度机制高度规范化，谁先谁后、依赖什么分析结果，都在流水线里写得明明白白，量化 Pass 只是其中一支被调度的队伍。

静态量化Pass

静态量化好比施工前把所有电器的额定功率都测量一遍，统一去市场上订购合适规格的变压器。它把最耗时的“测量”工作放在部署之前完成，推理时就按固定规格直接运行。这个“测量”环节在量化里叫校准（calibration）：准备一批有代表性的输入数据，跑一遍前向推理，记录每层激活值的分布，统计出最小值与最大值。有了极值就能算出 scale：

静态量化 Pass 在拿到 scale 之后做三件事：第一，权重离线量化——把每层权重直接转成 INT8，因为权重在推理中恒定不变；第二，激活量化——按校准统计的 scale 把激活范围压到整数，但为了让图中其他算子接口不变，会插入 Q（Quantize）与 DQ（Dequantize）成对节点，Q 负责把浮点转整数，DQ 负责把整数还原成浮点；第三，算子图变换——把 F32 方言的计算图重写为带 Q/DQ 节点的量化方言图，整张图从“220V 标称”变成“12V 标称接线”。

校准数据的选择就像抽样测量电器功率一样，取样的代表性决定了 transformer 是否合适。如果校准数据与真实推理分布差异过大，scale 就会偏，量化误差被放大。静态量化的核心优势是运行时零校准开销：所有“测量”都在施工前完成，推理时只按定好的规格跑整数运算。

动态量化Pass

有些电器的负载波动非常大——比如一台忽开忽关的工业电机，给它装固定变压器，要么在高峰时功率不足，要么在低峰时大马拉小车。动态量化就是“自适应变压器”：不提前测死规格，而是运行时根据实时负载临时算 scale。在模型里，权重仍然离线量化好，只有激活值在推理时“边算边调”。

动态量化最常用的粒度是 per-token：把激活张量按序列位置拆成一行一行，每一行独立算 scale。设某一行（一个 token）为 $x_{i*}$，量化位宽为 $b$，则该行的步长为：

除了 token 粒度，还有 per-channel 粒度：按卷积的输出通道逐通道计算 scale。通道粒度更细，对权重分布不均的模型更友好，但会引入更多元的 scale，增加存储和计算开销。MLIR 中实现动态量化 Pass 时，Q 节点被设计成携带“运行时计算”逻辑的小型子图：每次前向都执行一次极值扫描与除法。这个子图本身有开销，所以 Pass 要控制额外成本，例如尽量让 scale 计算复用已有的归约算子、避免引入不必要的内存拷贝。

动态量化像一位聪明的电工：遇到功率波动大的设备，就装一个能实时调节的变压器，电费省了，设备也稳了。代价是变压器本身要消耗一点能量——运行时算 scale 就是那点额外开销。什么时候值得用？当激活内存带宽成为瓶颈、而校准数据又很难覆盖多变场景时，动态量化是比静态更稳健的选择。

精度优化策略

工地里总有几台精密仪器——比如做化学实验的电子天平，对电压波动极其敏感。硬把它们塞进 12V 低压系统，测出来的数据就全废了。量化也有这样的“敏感层”：某些层的权重分布特别集中，或激活值跨度特别大，一旦量化误差就会被放大，导致整个模型输出漂移。因此需要对每一层做“敏感度体检”。

体检方法很直观：单独把某一层量化，其他层保持浮点，对比量化前后模型在验证集上的损失变化。设原权重为 $W_i$，量化后权重为 $\hat{W}_i$，该层的敏感度定义为：

根据体检结果，编译器制定三套保护措施。第一是混合精度：敏感层保持 F32，普通层用 INT8，整张图形成“高压区”与“低压区”共存的混合电路。第二是算子融合：把 Conv、BatchNorm、ReLU 三个连续算子融合成一个大算子，同时把夹在中间的 Q/DQ 节点消掉——就像把三台独立电器集成一台一体机，省去了中间反复变压的损耗。第三是精度回退：量化完跑一遍验证集，如果整体精度下降超过阈值，就自动把最敏感的那几层回退到浮点，类似备用电源自动切入。

这套策略背后是一套“成本账”：每量化一层省下多少内存与算力，又可能损失多少精度，编译器不断在两者之间权衡。融合算子是把多个环节合并减少插拔损耗，回退则是把高价值设备留在高压专线上。最终交付的模型，既不是全 220V 的奢侈配置，也不是全 12V 的冒险方案，而是一份“该省的省、该保的保”的优化电路图。

工程实践集成

单支施工队再专业，最终也要融入整栋楼的总包工程。MLIR 量化 Pass 的实战价值，体现在它与主流 AI 框架的深度集成——XLA 和 PyTorch 编译链路中，量化 Pass 都扮演着“核心施工队”的角色。

在 XLA 中，模型先被翻译成 HLO 中间表示，再经 MLIR 的桥接层逐步降级到 LLVM。量化 Pass 穿插在这个降级过程中：HLO 层面标记哪些算子可以量化，MLIR 层面的量化 Pass 负责插入 Q/DQ 并完成图重写。在 PyTorch 一侧则是另一条路径：用 torch.fx 捕获模型计算图，转成 torch-mlir 的方言表示，再进入同一个 MLIR 量化流水线。这样，无论模型来自 PyTorch 还是 TensorFlow，最终都能走进同一支“施工队”。

工程落地离不开性能调优与调试排障。性能调优关注吞吐量、延迟、模型体积三个核心指标：量化后 INT8 模型通常只能跑到原来的几分之一延迟，但前提是算子实现了低精度内核，且内存布局对缓存友好。调试排障则靠“留图纸”：MLIR 支持把每个 Pass 执行前后的 IR dump 到文件，开发者可以用 mlir-opt 的 --pass-pipeline 参数逐段回放，检查 Q/DQ 是否成对、scale 是否溢出、融合是否生效。一旦发现某层量化后输出异常，就在 dump 出的 IR 里定位到具体算子。

整栋楼的改造工程最终验收时，电工、水管工、空调工各司其职又协同有序。MLIR 量化 Pass 的价值正在于此：它不是一个孤立的工具，而是一个标准化流程中可插拔的施工队。前端框架负责“翻译需求”，Pass 负责“按图施工”，LLVM 负责“交付验收”。掌握了这支施工队的工作方式，你就能在任意主流框架上，用同一套方法论完成模型的高效量化部署。

AI 应用场景

Python 演示

"""模拟 MLIR 量化 Pass 的核心步骤：校准、量化、反量化、动态量化与误差评估。"""import math  def static_calibrate(values, bits=8):    """根据校准数据求 scale 与 zero point（对应章节3的校准+scale公式）。"""    q_min, q_max = 0, 2 ** bits - 1    v_min, v_max = min(values), max(values)    scale = (v_max - v_min) / (q_max - q_min)    zero = round(q_min - v_min / scale)    return scale, zero  def quantize(values, scale, zero, bits=8):    """把浮点值量化到 [0, 2^bits-1] 的整数范围。"""    q_max = 2 ** bits - 1    return [max(0, min(q_max, round(v / scale + zero))) for v in values]  def dequantize(quants, scale, zero):    """反量化：从整数还原回浮点，用于评估误差。"""    return [scale * (q - zero) for q in quants]  def dynamic_quantize_row(row, bits=8):    """per-token 动态量化：每行独立算 scale，对应章节4公式。"""    q_max = 2 ** (bits - 1) - 1    peak = max(abs(v) for v in row)    scale = peak / q_max    return [round(v / scale) for v in row], scale  def mse(a, b):    """均方误差，衡量量化前后激活值的偏离程度。"""    return sum((x - y) ** 2 for x, y in zip(a, b)) / len(a)  if __name__ == "__main__":    # 模拟某一层激活值的校准集采样    activations = [0.12, 0.55, 1.43, 2.01, 3.77, 5.20, 4.66, 1.99]    scale, zero = static_calibrate(activations)    q = quantize(activations, scale, zero)    r = dequantize(q, scale, zero)    print("静态量化 scale=%.4f zero=%d 误差=%.6f" % (scale, zero, mse(activations, r)))     # 模拟一批 token，每行一个 token，走动态量化    tokens = [[0.1, 0.9, 1.8], [0.4, 1.5, -0.2], [0.9, 0.2, 3.2]]    for idx, row in enumerate(tokens):        q_row, s_row = dynamic_quantize_row(row)        print("第%d个token动态scale=%.4f 量化结果=%s" % (idx + 1, s_row, q_row))

1. def static_calibrate(values, bits=8):对应校准流程，收集激活值的 min/max，算出 scale 与 zero point。
2. scale = (v_max - v_min) / (q_max - q_min)实数值域均匀摊到 255 个整数刻度上，得到步长。
3. quantize(...)` 中的 `round(v / scale + zero)对每个浮点值做尺度变换、零点平移与取整。
4. zero = round(q_min - v_min / scale)用实数最小值反推零点，保证实数零有对应的整数刻度。
5. dynamic_quantize_row(...)` 中的 `peak = max(abs(v) for v in row)找该 token 行的绝对值峰值作为动态范围上限。
6. q_max = 2 ** (bits - 1) - 1对称量化的整数正半轴最大值，8bit 时为 127。
7. mse(a, b)用均方误差度量量化前后激活分布的偏离，作为精度损失指标。

常见误区

面试题