MagiCompiler

KO | EN

gitlite — search

#javascript #hacktoberfest #react #typescript #golang #android #rust #linux #awesome #cli #kubernetes #security #openai #open-source

(★ 330)

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

File Explorer

Download Latest Version (.zip)

- - copyright.hook
- - checkout_pr.py
- - _ci_pipeline.yml
  - check_chinese_chars.py
  - integration_test.yml
  - merge_test.yml
  - pr-description-check.yml
  - sphinx.yaml
- PULL_REQUEST_TEMPLATE.md
- - h100_inference.png
  - magi_compiler_overview.jpeg
  - rtx5090_inference.png
- - conf.py
  - index.md
- AutoCPUOffload.md
- AutoCudaGraphDesign.md
- make.bat
- Makefile
- README.md
- README_zh.md
- requirements.txt
- WhyMagiDepyf.md
- - llama3.py
  - train.py
  - train.sh
- - __init__.py
  - cudart.py
- - __init__.py
  - _aot_compat.py
  - _cache_data_cls.py
  - compile_artifacts.py
  - cuda_graph_mgr.py
  - magi_backend.py
  - magi_compiler_base.py
  - partition_rules.py
  - piecewise_backend.py
  - piecewise_compiler.py
- - __init__.py
- - __init__.py
  - offload_warpper.py
  - profiler.py
  - scheduler.py
- - __init__.py
- - __init__.py
  - materialize_inputs.py
  - runtime_estimator.py
- - __init__.py
  - _utils.py
  - compile_counter.py
  - compile_time_monitor.py
  - device.py
  - envs.py
  - hash.py
  - logger.py
  - nvtx.py
  - ordered_set.py
  - singleton_meta.py
- __init__.py
- _api.py
- _magi_register_custom_op.py
- api.py
- config.py
- - test_dynamic_int_inputs.py
  - test_magi_compile.py
  - test_nested_compile.py
  - test_register_custom_op.py
- - test_qwen35_4b_infer.py
  - test_qwen35_4b_real_e2e.py
- - conftest.py
  - test_autograd_function_cache_flag.py
  - test_build_cleanup.py
  - test_cache_invariant.py
  - test_cache_topology_isolation.py
  - test_capture_scalar.py
  - test_compile_artifacts.py
  - test_conv_channels_last_switch.py
  - test_cpu_offload.py
  - test_cuda_graph.py
  - test_fsdp_overlap_bucket.py
  - test_fsdp_overlap_e2e.py
  - test_fsdp_overlap_lowering.py
  - test_fsdp_overlap_reorder.py
  - test_magi_inductor_pass.py
  - test_matmul_epilogue_fusion.py
  - test_nd_tiling_workaround.py
  - test_piecewise_deferred_assert_scope.py
  - test_profiling_estimator.py
  - test_profiling_registry.py
  - test_python_scalar_constant_folding.py
  - test_restart_analysis_cache.py
  - test_stride_mismatch.py
  - test_symbolic_unification.py
  - test_torch_op_determinism_matrix.py
  - test_transformer_cache_reuse.py
  - test_unbacked_symbol_guard.py
- - __init__.py
- - test_mlp_infer.py
  - test_mlp_training.py
  - test_rms_norm_infer.py
- - __init__.py
  - conftest.py
  - test_conv_channels_last_perf.py
  - test_mlp_perf.py
  - test_nd_tiling_perf_workaround.py
  - test_norm_residual_fusion_perf.py
  - test_pointwise_fusion_perf.py
  - utils.py
- - test_inductor_cache_reuse.py
  - test_symbolic_dim_ability.py
- __init__.py
- conftest.py
- model_definition.py
- utils.py
.dockerignore
.gitignore
.pre-commit-config.yaml
Dockerfile
LICENSE
Makefile
pyproject.toml
README.md
requirements-test.txt
requirements.txt

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# MagiCompiler

**Break the Boundaries of Local Compilation for Large Models**

</div>

---

## 📢 Latest News

- **[03/25/2026]** ⚡️ **[LightX2V-MagiCompiler](https://github.com/SandAI-org/LightX2V-MagiCompiler) is now available!** This fork of [LightX2V](https://github.com/ModelTC/lightx2v) showcases how to seamlessly integrate MagiCompiler into a SOTA framework. With **minimal code changes**, it unlocks even greater acceleration! Try it out, check the [benchmark](https://github.com/SandAI-org/MagiCompiler#-6-benchmark) for details, and stay tuned for more integration demos!

- **[03/23/2026]** 🚀 **MagiCompiler is officially open-sourced!** Delivering whole-graph compilation for multi-modality inference and FSDP-aware whole-layer compilation for large model training.

---

## 📖 About

MagiCompiler is an advanced compiler and runtime augmentation framework built on top of `torch.compile`. Designed specifically for large-scale Transformer-like architectures, it addresses the critical bottlenecks of memory walls and operator overheads.

By stepping beyond traditional local operator optimization, MagiCompiler introduces system-level optimizations, seamlessly accelerating both **training** and **multi-modality inference** workloads with minimal code intrusion.

---

## 💡 Design Philosophy

### Compiler as Manager

> *"Reimagining the compiler: from generating kernels to orchestrating the entire dataflow."*

MagiCompiler's core philosophy is **Compiler as Manager**. We believe a modern deep learning compiler should not be restricted to mere kernel fusion. Instead, it acts as a global manager that owns the full lifecycle of execution. MagiCompiler actively manages subgraph dispatching, dynamically orchestrates dataflow (like offloading and prefetching), and controls memory allocation, ensuring optimal balance between compute efficiency and memory footprint.

### Key Features

#### 🎯 1. Unified Inference & Training
Tailored for Transformer-like architectures with scenario-specific strategies:
- **Inference**: Achieves **full-graph capture** across Transformer boundaries, maximizing kernel fusion scope.
- **Training**: Introduces **FSDP-aware layer-wise compilation**. Unlocks aggressive cross-op fusion while keeping distributed parameter sharding entirely transparent.

#### ⚡️ 2. Easy to Use, Free Gain, Plug and Play
No complex model refactoring needed. Just two decorators deliver up to **20%+ extra speedups** out-of-the-box, seamlessly integrating into SOTA multi-modality frameworks.

#### 🧠 3. Smart Asynchronous Offloading
For memory-constrained setups, our built-in **selective offloading policy** perfectly overlaps H2D transfers with computation, eliminating pipeline bubbles.

#### ♻️ 4. Heuristic Activation Recomputation
Say goodbye to manual `torch.utils.checkpoint`. MagiCompiler automatically saves compute-bound ops (e.g., MatMul, Attention) and recomputes memory-bound ones, slashing peak memory without sacrificing throughput.

#### 🛠 5. Magi Depyf Introspection
Meet **magi_depyf**, MagiCompiler’s native introspection toolkit. Compilation timelines, decompiled bytecode flows, split subgraphs, and backend artifacts are automatically dumped into the cache path as organized, human-readable files for easier debugging.

---

## ⚙️ Installation

**Requirements:**
- Python >= 3.12
- PyTorch 2.9.x / 2.12.x
- CUDA Toolkit

> **Recommended for reproducibility:** start from the prebuilt Docker image first, then run examples inside the container.

```bash
# Option A (recommended) — Use prebuilt image
# Step 1 — Pull the image
docker pull sandai/magi-compiler:latest
# Step 2 - Start the container
docker run --name my-magi-compiler -it -d --privileged --gpus all --network host --ipc host \
  -v /path/on/host:/workspace sandai/magi-compiler:latest /bin/bash
# Step 3 - Attach the container
docker exec -it my-magi-compiler /bin/bash

# Option B — Local source installation
# Step 1 — Clone the repo
git clone https://github.com/SandAI-org/MagiCompiler.git
cd MagiCompiler

# Step 2 — System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
sudo apt update && sudo apt install -y graphviz

# Step 3 — Python dependencies
pip install -r requirements.txt

# Step 4 — Install MagiCompiler (pick one)
pip install .   # End users (recommended)
# pip install -e . --no-build-isolation --config-settings editable_mode=compat  # Developer / editable

# Step 5 (optional) — Install CUTLASS for matmul epilogue fusion
# Required for the CUTLASS-based matmul + epilogue fusion pass (sm_90 / sm_120).
# Without CUTLASS the compiler still works but skips this optimization.
git clone --depth 1 https://github.com/NVIDIA/cutlass.git /usr/local/cutlass
# Or specify a custom path:
#   git clone --depth 1 https://github.com/NVIDIA/cutlass.git /your/path
#   export MAGI_CUTLASS_ROOT=/your/path
export CUDACXX=${CUDA_INSTALL_PATH}/bin/nvcc
mkdir /usr/local/cutlass/build && cd /usr/local/cutlass/build
cmake .. -DCUTLASS_NVCC_ARCHS=90a # compiles for NVIDIA Hopper GPU architecture
# cmake .. -DCUTLASS_NVCC_ARCHS=120a # compiles for NVIDIA consumer Blackwell (RTX 50 series)
```

---

## 🚀 Quick Start

### 🧹 1. One Decorator to Rule Them All (`@magi_compile`)
Remove scattered `torch.compile` or `torch.compiler.disable` calls. Decorate your core Transformer block once for automatic full-graph capture and dynamic shape support (defaulting to dim 0).

```python
import torch
from torch import nn
from magi_compiler import magi_compile

# Decorate your core module once. No more scattered compile tweaks!
@magi_compile
class TransformerBlock(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.attn = Attention(hidden_dim)
        self.mlp = MLP(hidden_dim)

def forward(self, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
        x = x + self.attn(x, mask)
        x = x + self.mlp(x)
        return x

model = TransformerBlock(hidden_dim=1024).cuda()

# Execute normally - whole-graph compilation handles dynamic batches automatically!
out = model(torch.randn(4, 128, 1024, device="cuda"), None)
out = model(torch.randn(8, 128, 1024, device="cuda"), None)
```

### 🛠️ 2. Bridge Custom Kernels (`@magi_register_custom_op`)
Using custom kernels (FlashAttention, MoE routers) that break FX tracing? Don't disable compilation. Wrap them to teach the compiler how to handle them during graph partitioning and recomputation.

```python
from magi_compiler import magi_register_custom_op

@magi_register_custom_op(
    name="athena::flash_attn",
    infer_output_meta_fn=["q"],       # Output shape matches parameter 'q'
    is_subgraph_boundary=True,        # Split graph here for subgraph compilation
    is_compute_sensitive=True,        # Retain this output during recomputation
)
def flash_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    ... # Your custom kernel or C++ extension
```

### 🔧 3. Advanced Configurations
Explore `magi_compiler/config.py` for power-user features like custom backend toggles and fine-grained memory management. *(Comprehensive guides for popular training/inference frameworks are coming soon!)*

---

## 📊 Benchmark

### 🔥 H100 Extreme Acceleration

On a single NVIDIA H100, MagiCompiler outperforms current SOTA solutions (like LightX2V) by 9% to 26% across mainstream open-source video generation models.

### 💻 RTX 5090 Near Real-Time

Thanks to our underlying JIT offloading engine, [daVinci-MagiHuman](https://github.com/GAIR-NLP/daVinci-MagiHuman) achieves near real-time speeds, even on heavily VRAM-constrained consumer GPUs.

---

## 🗺 Roadmap

We are actively developing MagiCompiler. Here is a glimpse into our upcoming milestones:

- [ ] **Ecosystem Integration**: Benchmarks and out-of-the-box integration guides for popular frameworks (e.g., `sglang-diffusion`, `vllm-omni`, and `LLaMA` training).
- [ ] **Official Hub & Tech Blog**: A dedicated website for advanced tutorials, documentation, and frontier engineering insights.
- [ ] **Hardware-Aware Auto-Scheduler**: An adaptive engine that dynamically orchestrates optimal strategies (auto-recomputation boundaries, offloading) based on your hardware constraints.
- [ ] **Next-Gen Custom Backend (v2.0)**: Pushing hardware limits with extreme kernel-level efficiency, native distributed communication and MegaKernels.

---

## 📝 Citation

If you find MagiCompiler useful in your research or production, please consider citing us:

```bibtex
@software{magi_compiler_2026,
  author = {Hongyu Jia and Zhiyao Cen and Taoran Wang and Yunbo Zhang},
  title = {MagiCompiler: Break the Boundaries of Local Compilation for Large Models},
  year = {2026},
  url = {https://github.com/SandAI-org/MagiCompiler}
}
```

---

## 🙏 Acknowledgement

MagiCompiler is deeply inspired by and builds upon the shoulders of giants. We extend our heartfelt gratitude to the [PyTorch](https://pytorch.org/) team for their foundational work on `torch.compile` and `torch.fx`, and to the [vLLM](https://github.com/vllm-project/vllm) community for their pioneering contributions to large model inference.

**We are moving fast, and we want you on board!** MagiCompiler is under rapid development. If you are passionate about pushing the limits of large model compilation, we'd love to have you with us. From opening issues and discussing architectures to submitting core PRs, every contribution matters. Let's engineer the future of AI infrastructure together!

---

## ⭐ Star History

---

## 📜 License

This project is licensed under the [Apache License 2.0](LICENSE).

MagiCompiler

File Explorer

Today's picks

Was this content helpful?

Recent Feedback

MagiCompiler

File Explorer

Today's picks

Was this content helpful?

Recent Feedback

Download README