开源 · Apache 2.0 基于 MLIR PyTorch · ONNX · TFLite · Caffe HuggingFace 大模型

把任意模型
编译到 TPU 上的 bmodel。

TPU-MLIR 是一款开源的机器学习编译器，可将来自 PyTorch、ONNX、TFLite、Caffe 以及 HuggingFace LLM 的预训练网络，转换为可在算能（Sophgo）TPU 上高效运行的 .bmodel 文件。两层 MLIR Dialect 设计让从算法到芯片的路径既清晰、可重定向，又易于扩展。

快速开始开发者手册阅读论文

前端框架

120+

已验证模型

F32 → INT4

精度类型

前端 → Top Dialect → TPU Dialect → bmodel → 算能 TPU

编译流程

三条命令，把训练好的模型变成可部署的 bmodel。

围绕三个简洁的 Python 入口构建，整个流程脚本化、可复现。

导入为 MLIR

将原始模型导入与框架无关的 Top Dialect。

$ model_transform.py \
    --model_name yolov5s \
    --model_def  yolov5s.onnx \
    --input_shapes [[1,3,640,640]] \
    --mlir yolov5s.mlir

校准与量化

可选：基于少量数据生成 INT8 校准表。

$ run_calibration.py yolov5s.mlir \
    --dataset   ../COCO2017 \
    --input_num 100 \
    -o yolov5s_cali_table

下沉与部署

下沉到 TPU Dialect，针对目标芯片产出 .bmodel。

$ model_deploy.py \
    --mlir      yolov5s.mlir \
    --quantize  INT8 \
    --processor bm1684x \
    --model     yolov5s_int8.bmodel

为什么选 TPU-MLIR

面向真实部署场景设计的编译器。

支持多框架输入、多精度输出，构建在易于扩展的 MLIR 之上。

多框架前端

开箱即用支持 PyTorch、ONNX、TFLite 与 Caffe；其他框架可经 ONNX 转换接入。

HuggingFace 大模型

原生支持 HuggingFace LLM —— 包含 AWQ、GPTQ、AutoRound 等量化形态 —— 通过统一的 llm_convert.py 一键编译。

完整精度矩阵

F32 / F16 / BF16 / INT8 / W4A16 / W8A16，并提供基于校准的 INT8 训练后量化。

多目标代码生成

同一条 MLIR 流水线即可覆盖 bm1684x、bm1688 等多款算能 TPU 处理器。

Layer-Group 优化

图级与张量级的深度优化 —— Layer Group、前后处理融合 —— 让 TPU 持续繁忙、显存搬运保持低位。

120+ 模型已验证

视觉（YOLO 系列）、语音（Whisper）以及大语言模型（Qwen3.5、Qwen3-VL、MiniCPM-V-4 ……）均在持续回归。

支持矩阵

框架、精度与芯片，一目了然。

选定前端、确定精度，部署到您所使用的目标芯片。

框架

PyTorch ONNX TFLite Caffe ★HuggingFace

精度

F32 F16 BF16 INT8（PTQ + QAT） W4A16 W8A16 ★AWQ ★GPTQ

目标芯片

★BM1684X ★BM1688 …等更多算能 TPU

模型

视觉、语音、大语言模型 — 都跑在 TPU 上。

以下并非完整清单，仅为 TPU-MLIR 回归与示例流程中常见的代表性模型族。

视觉

★YOLOv5 ★YOLOv8 ResNet MobileNet SAM Stable Diffusion

语音

★Whisper Paraformer FunASR

大语言模型

★Qwen3.5 ★Qwen3-VL ★MiniCPM-V-4 …等更多

YOLOv5s 在 F32 / F16 / INT8 bmodel 下的输出对比

YOLOv5s — 不同精度 bmodel 的输出对比。

Qwen3.5-2B INT4 在算能 BM1684X 上经 TPU-MLIR 编译后运行的对话效果

Qwen3.5-2B（INT4 / AutoRound）— 通过 llm_convert.py 编译后在算能 BM1684X 上运行对话。

快速开始

选择一个起点。

两条路径都假设您已使用官方 sophgo/tpuc_dev Docker 镜像并执行了 pip install tpu_mlir。

LLM 编译一个 HuggingFace 大模型（Qwen3.5）

# 1. 从 HuggingFace 拉取量化模型
$ git lfs install
$ git clone https://huggingface.co/Intel/Qwen3.5-2B-int4-AutoRound

# 2. 编译为面向 BM1684X 的 bmodel
$ llm_convert.py \
    -m /workspace/Qwen3.5-2B-int4-AutoRound \
    --max_input_length 1024 -s 2048 \
    -c bm1684x --max_pixels 768,768 \
    -o qwen3.5_2b
# → qwen3.5_2b/*.bmodel  ✔ 即可上 TPU 运行

CNN 编译一个视觉模型（YOLOv5s）

# 1) 前端转换 → Top Dialect MLIR
$ model_transform.py \
    --model_name yolov5s \
    --model_def  ../yolov5s.onnx \
    --input_shapes [[1,3,640,640]] \
    --mlir yolov5s.mlir

# 2) INT8 校准
$ run_calibration.py yolov5s.mlir \
    --dataset ../COCO2017 --input_num 100 \
    -o yolov5s_cali_table

# 3) 下沉为 bmodel
$ model_deploy.py \
    --mlir yolov5s.mlir --quantize INT8 \
    --calibration_table yolov5s_cali_table \
    --processor bm1684x \
    --model yolov5s_int8.bmodel

把任意模型编译到 TPU 上的 bmodel。