SandAI-org/MagiCompiler

A plug-and-play compiler that delivers free-lunch optimizations for both inference and training.

What it solves

MagiCompiler 解决了限制大规模类 Transformer 架构性能的内存瓶颈和算子开销。它超越了简单的本地算子优化，提供系统层面的优化，支持训练和多模态推理。

How it works

基于 torch.compile，MagiCompiler 充当执行生命周期的全局管理器。它负责子图调度、协同数据流（包括卸载和预取），并控制内存分配。关键技术机制包括：

  • Full-graph capture：在推理时最大化跨 Transformer 边界的 kernel 融合范围。
  • FSDP-aware layer-wise compilation：在训练时实现激进的跨算子融合，同时保持分布式参数分片的透明性。
  • Selective offloading：将主机到设备 (H2D) 的传输与计算重叠，以消除内存受限环境中的流水线气泡。
  • Heuristic activation recomputation：自动判断哪些操作需要保存、哪些需要重算，根据计算密集或内存密集的特性降低峰值内存使用。

Who it’s for

该框架面向使用大规模 Transformer 模型的开发者和研究者，特别是实现多模态框架或使用 FSDP 等分布式策略训练大模型的用户。

Highlights

  • Plug-and-play integration：使用简单的装饰器（@magi_compile 与 @magi_register_custom_op）即可获得加速，无需复杂的模型重构。
  • Significant acceleration：在 NVIDIA H100 GPU 上对视频生成模型展示 9%~26% 的性能提升，超越现有最先进方案。
  • Introspection toolkit：包含 magi_depyf 用于调试编译时间线、字节码流和子图拆分。
  • Hardware-aware：支持高端 GPU（H100）和消费级 GPU（RTX 5090），通过 JIT 卸载实现硬件感知。