sophgo/tpu-mlir

Machine learning compiler based on MLIR for Sophgo TPU.

解决的问题

TPU-MLIR 解决了将来自各种主流框架的预训练神经网络部署到 TPU 硬件上的问题。它弥合了高层模型定义（如 PyTorch 或 ONNX）与 TPU 芯片高效执行所需的特定二进制格式（bmodel）之间的鸿沟。

工作原理

该项目使用基于 MLIR（多级中间表示）的编译流水线，通过以下几个阶段转换模型：

  1. 前端导入：从 PyTorch、ONNX、TFLite、Caffe 或 HuggingFace 导入模型并转换为统一的 MLIR 表示。
  2. 降低（Lowering）：编译器通过模式重写和内存规划，将高层表示转换为 TPU 特有的方言（Top 和 Tpu）。
  3. 量化：应用各种量化技术（F32、BF16、F16、INT8）和校准，以优化硬件效率。
  4. 部署：最终输出是可以在目标 TPU 处理器上运行的 bmodel 文件。

适用对象

专为需要将深度学习模型（包括大语言模型和视觉模型）部署到 Sophgo TPU 硬件的 AI 工程师和开发人员设计。

亮点

  • 广泛的框架支持：兼容 PyTorch、ONNX、TFLite、Caffe 和 HuggingFace。
  • LLM 优化：针对 HuggingFace LLM 的专用单次转换，支持历史 KV 缓存和分块推理。
  • 全面的量化：支持对称/非对称 INT8、F16、BF16，以及 AWQ、GPTQ 和 AutoRound 的透传。
  • 生产级工具链：包括通用推理运行器 (model_runner)、bmodel 检查工具 (model_tool) 和精度验证实用程序。