sophgo/tpu-mlir

解决的问题

TPU-MLIR 是一个机器学习编译器，允许开发者从 PyTorch、ONNX、TFLite 和 Caffe 等流行框架中获取预训练的神经网络，并将其转换为一种称为 bmodel 的专用格式。这些文件针对 TPU 硬件进行了优化，弥合了高级 AI 模型设计与硬件特定执行之间的差距。

工作原理

基于 MLIR（多级中间表示）框架构建，该编译器使用“方言”和模式重写组成的流水线，将高级模型描述降低为硬件特定指令。它包含一个完整的工具链，支持量化（将模型从 F32 转换为 BF16、F16 或 INT8）和校准，以保持精度。对于大语言模型（LLM），它提供专用的 llm_convert.py 工具，支持 HuggingFace 模型的一次性转换，并具备 KV 缓存管理与分块推理等高级功能。

适用人群

该工具专为需要将优化后的 AI 模型（从 YOLOv5 等视觉模型到 Qwen 等 LLM）部署到基于 TPU 的加速器上的 AI 工程师和硬件开发者设计。

主要亮点

• 多框架支持：直接转换来自 PyTorch、ONNX、TFLite 和 Caffe 的模型。
• LLM 优化：针对 HuggingFace LLM 提供一次性转换和 KV 缓存处理的专用支持。
• 全面的量化支持：支持 F32、BF16、F16 和 INT8（对称/非对称），以及 AWQ、GPTQ 和 AutoRound。
• 生产级工具链：包含模拟器、可视化工具，以及用于精度验证和 bmodel 检查的工具。
• 基于 MLIR 的流水线：采用清晰的下降低级流水线，结合层组内存规划，实现高效的硬件利用。