[fix] compile state with instance-level forward wrapping

[songh11]

🗂️ PR Category

• ✨ New Feature
• 🚀 Optimization (performance, memory, etc.)
• 💥 Breaking Change
• 🐛 Bug Fix
• 🛠️ Development / Refactoring
• 📚 Documentation
• 🧹 Chore (Dependencies, CI/CD, Configuration, etc.)
• 🧪 Testing

📝 Description

• Align MagiCompileState with instance-level wrapping semantics

  • When compiling nn.Module instances, prefer _magi_original_forward (if present) as the target callable, ensuring the compiled graph reflects the original, pre-wrapped forward.
  • In JIT dispatch, temporarily rebind the instance’s forward to the class-level forward bound method while swapping bytecode, then restore the original binding after execution, so hooks/FSDP logic still flow through __call__ while keeping instance-level overrides consistent.

• Refine instance-level compile orchestration path

  • In _magi_compile_instance, store module._magi_original_forward and change the orchestration entry from calling the unbound forward directly to going through module.__class__.__call__, preserving module __call__ behavior and hooks for instance-compiled models.

• Add CUDA timing sanity test for multiple Magi entrypoints

  • Introduce a lightweight CUDA test that compares execution time across class, function, instance, and method @magi_compile entrypoints on the same SimpleModel, asserting their compiled timings stay within a small factor of each other to catch regressions or inconsistent fast-path behavior.

[cennn]

Great catch! — clean fix for the instance-level compilation bypass. On our end, we'll be adding perf regression tests for all compile entry points, and exploring whether instance-level dispatch can avoid touching __class__.forward altogether to keep it fully isolated and multi-instance safe.