如何通过算子融合减少解量化开销：详解在移动端将 Dequant 与 MatMul 合并的技巧

在移动端部署深度学习模型时，为了追求极致的推理速度和更小的模型体积，INT8 量化几乎是标配。然而，很多开发者在实际部署时发现，虽然权重变成了 INT8，但推理速度提升并不明显，甚至有所下降。这通常是因为大量的 Dequantize（解量化）算子产生了高额的内存拷贝和带宽开销。本文将教你如何通过算子融合（Operator Fusion）技术，将 Dequant 算子并入 MatMul 核心计算中，实现端侧推理的显著加速。

1. 为什么单独的 Dequant 算子是性能杀手？

在未优化的计算图中，量化模型通常表现为：
Input(INT8) -> Dequant -> Input(FP32) -> MatMul(FP32) -> Output(FP32)

这种模式在移动端存在两个致命问题：
1. 内存带宽瓶颈：Dequant 算子会将 INT8 数据扩展为 FP32，数据量瞬间膨胀 4 倍。移动端 SoC 的内存带宽非常有限，频繁的读写会导致 CPU/GPU 长期处于等待状态。
2. 算子启动开销：每个算子都有启动延迟（Kernel Launch Overhead）。如果网络很深，成百上千个 Dequant 算子的累积延迟非常可观。

2. 核心优化方案：算子融合（Fusion）

优化的核心思路是将解量化逻辑“下沉”到矩阵乘法的内部逻辑中。直接读取 INT8 数据进行整数乘累加，仅在最后结果写回内存前进行一次缩放（Scaling）。

融合后的流程：
Input(INT8) + Weight(INT8) -> Fused_Quantized_MatMul -> Output(FP32/INT8)

3. 实战：使用 ONNX GraphSurgeon 实现算子融合逻辑

虽然像 ncnn 或 MNN 这样的推理框架会自动处理这类融合，但理解其底层逻辑并手动干预计算图是资深 AI 工程化的必备技能。以下代码展示了如何使用 onnx-graphsurgeon 识别并清理不必要的解量化节点。

import onnx
import onnx_graphsurgeon as gs
import numpy as np

def fuse_dequant_matmul(model_path, output_path):
    # 加载 ONNX 模型
    graph = gs.import_onnx(onnx.load(model_path))

    # 遍历所有节点寻找 MatMul
    for node in graph.nodes:
        if node.op == "MatMul":
            # 检查输入是否连接着 DequantizeLinear 节点
            for i, input_tensor in enumerate(node.inputs):
                if isinstance(input_tensor, gs.Variable) and len(input_tensor.inputs) > 0:
                    prev_node = input_tensor.inputs[0]
                    if prev_node.op == "DequantizeLinear":
                        print(f"发现可融合的节点: {prev_node.name} -> {node.name}")

                        # 获取 Dequant 的原始 INT8 输入和 scale
                        int8_input = prev_node.inputs[0]
                        scale = prev_node.inputs[1]

                        # 技巧：在移动端通常将 MatMul 替换为专有的 QLinearMatMul 或 Gemm
                        # 这里模拟将 MatMul 的输入直接改为 INT8 变量
                        node.inputs[i] = int8_input

                        # 注意：实际底层 Kernel 需要支持感知 scale，
                        # 通常我们会将 scale 信息存入节点的属性中
                        node.attrs["x_scale"] = scale

    # 清理图中孤立的无用节点（即原先的 Dequant 节点）
    graph.cleanup().toposort()
    onnx.save(gs.export_onnx(graph), output_path)
    print("融合优化完成！")

# 假设有一个名为 base_model.onnx 的量化模型
# fuse_dequant_matmul("base_model.onnx", "fused_model.onnx")

4. 移动端底层实现的技巧

在底层 C++ 实现（如使用 ARM NEON 指令集）时，融合后的算子应当遵循以下原则：

1. 寄存器级解量化：在执行 vmlal.s16 (ARM NEON 整数乘加) 之后，利用累加器中的 32 位整数结果进行浮点转换。此时数据仍在寄存器中，不需要写回内存。
2. 利用 SMLAL 指令：对于 INT8 矩阵乘法，使用 SMLAL (Signed Multiply-Accumulate Long) 可以直接处理乘法并累加到 32 位寄存器中，有效防止精度溢出。
3. 多路并行读写：在融合算子内部，通过 vld1 同时读取多个 INT8 数据块，并在一次 Kernel 执行中完成从“读取 INT8”到“输出 FP32”的全过程。

5. 总结

通过将 Dequant 算子融合进 MatMul，我们不仅减少了内存占用（4:1），还通过减少内存读写次数显著提升了推理速度。在端侧适配国产芯片（如瑞芯微 RK3588 或晶晨 NPU）时，这种“计算与转换合并”的策略是解决性能瓶颈的关键。下次你在优化模型时，请务必检查你的计算图中是否存在零散的解量化节点！