fix: _patch_cpu_offload_apply handles .to(cuda) in addition to .cuda()

magi_compiler/_api.py

@@ -509,15 +509,26 @@ def _cpu_apply(self, fn):
         id_cpu_lambda = getattr(fn, "__qualname__", "") == "Module.cpu.<locals>.<lambda>"
         is_to_lambda = getattr(fn, "__qualname__", "") == "Module.to.<locals>.convert"
 
+        # Detect .to("cuda:X") by probing the lambda on a small CPU tensor.
+        # .to("cpu") also produces is_to_lambda=True but should not trigger offload.
+        is_to_cuda = False
+        if is_to_lambda and not getattr(self, "_magi_offloaded_once", False):
+            try:
+                is_to_cuda = fn(torch.empty(0, device="cpu")).is_cuda
+            except Exception:
+                pass
+
+        is_moving_to_gpu = is_cuda_lambda or is_to_cuda
+
         # after first time to call _apply(cuda), skip "Module.to" and "Module.cpu" and "Module.cuda"
         if getattr(self, "_magi_offloaded_once", False):
             if is_cuda_lambda or id_cpu_lambda or is_to_lambda:
                 return self
             else:
                 return _orig_apply(self, fn)
         else:
-            # first time to call _apply(cuda), move all parameters/buffers to CPU
-            if not is_cuda_lambda:
+            # first time to call _apply(cuda) or _apply(to_cuda), move all parameters/buffers to CPU
+            if not is_moving_to_gpu:
                 return _orig_apply(self, fn)
 
         # move all parameters/buffers to CPU

magi_compiler/config.py

@@ -92,6 +92,18 @@ class PassConfig(BaseModel):
             "Env var: MAGI_COMPILE_PASS_CONFIG__ENABLE_ND_TILING_WORKAROUND (1/0/true/false)."
         ),
     )
+    nd_tiling_max_tiles: int = Field(
+        2,
+        ge=1,
+        le=3,
+        description=(
+            "max_tiles the ND-tiling workaround sets. 2 (default) is safe: Inductor's Grid2D "
+            "folds a y-grid overflow into z. 3 is experimental -- Grid3D has no z-overflow "
+            "handling, so a conv-heavy dynamic-shape graph (turbo VAE at 1080p) can exceed "
+            "CUDA's 65535 z-grid limit. "
+            "Env var: MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES."
+        ),
+    )
     enable_mm_epilogue_fusion: bool = Field(
         False,
         description=(
@@ -180,6 +192,7 @@ class OffloadConfig(BaseModel):
         OffloadPolicy.COST_EFFECTIVE, description="The policy for offloading the model to CPU."
     )
     bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.")
+    max_prefetch_lookahead: int = Field(2, description="Max layers to prefetch ahead. 0 disables prefetch to save GPU memory.")
 
 
 class FSDPConfig(BaseModel):

magi_compiler/offload/scheduler.py

@@ -154,7 +154,7 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
         except ValueError:
             return
 
-        max_lookahead = 2
+        max_lookahead = self.compile_config.offload_config.max_prefetch_lookahead
         target_node = None
         is_next_iter = False
 

magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py

@@ -18,7 +18,6 @@
 import torch
 
 from ...magi_depyf.timeline import emit_pass_lifecycle
-from ...utils.envs import IS_PT_212
 from ..pass_base import MagiInductorPass
 
 
@@ -40,7 +39,9 @@ def __call__(self, graph: torch.fx.Graph):
         torch._inductor.config.triton.prefer_nd_tiling = True
         torch._inductor.config.triton.tile_reductions = True
 
-        # PT 2.12 Inductor generates invalid 3D-grid reduction kernels with
-        # max_tiles=3 (program_id(2) mapped to a non-existent grid dim).
-        # Cap at 2 on PT >= 2.12 until the upstream fix lands.
-        torch._inductor.config.triton.max_tiles = 2 if IS_PT_212 else 3
+        # Capped at 2 by default: Grid3D has no z-overflow handling, so a
+        # conv-heavy dynamic-shape graph can exceed CUDA's 65535 z-grid limit.
+        # Raise it via MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES to test 3.
+        from ...config import get_compile_config
+
+        torch._inductor.config.triton.max_tiles = get_compile_config().pass_config.nd_tiling_max_tiles