From 68205660111304e8d43b10563b1e0b113a77d301 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Tue, 18 Aug 2026 11:48:45 +0800 Subject: [PATCH 1/4] fix: _patch_cpu_offload_apply now handles .to(cuda) in addition to .cuda() When callers use model.to(cuda:X) instead of model.cuda(), the _cpu_apply hook did not recognize the Module.to lambda and fell through to _orig_apply, putting all weights on GPU. This defeats the purpose of model_cpu_offload. Fix: probe the .to() lambda with a small CPU tensor to detect if the target device is CUDA, and treat it the same as .cuda() for offload interception. --- magi_compiler/_api.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index e7c6209..7bf5f21 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -509,6 +509,17 @@ def _cpu_apply(self, fn): id_cpu_lambda = getattr(fn, "__qualname__", "") == "Module.cpu.." is_to_lambda = getattr(fn, "__qualname__", "") == "Module.to..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: @@ -516,8 +527,8 @@ def _cpu_apply(self, fn): 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 From d8cbfd12476e88f97339097288fde55d63fbccfb Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 19 Aug 2026 10:27:31 +0800 Subject: [PATCH 2/4] feat(offload): make the prefetch lookahead configurable The scheduler hard-coded two layers of lookahead. Prefetching costs GPU memory that a small card may not have, so expose it as max_prefetch_lookahead and keep 2 as the default; 0 disables prefetch entirely. --- magi_compiler/config.py | 1 + magi_compiler/offload/scheduler.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index f0206ac..d5592d8 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -180,6 +180,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): diff --git a/magi_compiler/offload/scheduler.py b/magi_compiler/offload/scheduler.py index 566a4ea..c3569b4 100644 --- a/magi_compiler/offload/scheduler.py +++ b/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 From 1054c1f3d5e1eafaa3f356beef150b749c06917e Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Wed, 19 Aug 2026 15:00:55 +0800 Subject: [PATCH 3/4] fix(nd_tiling): cap max_tiles at 2 unconditionally max_tiles=3 can overflow CUDA z-grid limit (65535) on conv-heavy dynamic-shape graphs when coalesce_tiling_analysis collapses high-dim tensors into 3D. Grid3D has no z-overflow handling unlike Grid2DWithYZOverflow. The previous PT-version guard (2 on PT>=2.12, 3 otherwise) left PT<2.12 exposed. Cap at 2 on all versions since max_tiles=3 is documented as experimental. --- .../passes/piecewise_graph/nd_tiling_workaround.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py b/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py index a69c32d..ba044e7 100644 --- a/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py +++ b/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py @@ -40,7 +40,11 @@ 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 + # max_tiles=3 causes two known issues: + # - PT 2.12+: invalid 3D-grid reduction kernels (program_id(2) mapped + # to a non-existent grid dim). + # - All versions: conv-heavy dynamic-shape graphs (e.g. turbo VAE at + # 1080p) can overflow CUDA's z-grid limit (65535) when Inductor's + # coalesce_tiling_analysis collapses high-dim tensors into 3D. + # max_tiles=3 is documented as "experimental and may have bugs". + torch._inductor.config.triton.max_tiles = 2 From 48a9ab00caa31b14c2f1407ec7fa20af3d1eeaf6 Mon Sep 17 00:00:00 2001 From: cenzhiyao <2523403608@qq.com> Date: Thu, 20 Aug 2026 12:08:34 +0800 Subject: [PATCH 4/4] refactor(nd_tiling): make max_tiles configurable instead of hardcoded The pass hardcoded max_tiles=2 after we found Grid3D has no z-overflow handling. That value is right as a default but there was no way to try 3 without editing the pass. Expose it as PassConfig.nd_tiling_max_tiles, reachable via MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES. Turning the whole pass off is not an alternative: it also sets prefer_nd_tiling and tile_reductions, which conv-heavy dynamic-shape graphs need to avoid degrading to Grid1D. --- magi_compiler/config.py | 12 ++++++++++++ .../piecewise_graph/nd_tiling_workaround.py | 15 ++++++--------- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index d5592d8..181c94a 100644 --- a/magi_compiler/config.py +++ b/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=( diff --git a/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py b/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py index ba044e7..21de6d1 100644 --- a/magi_compiler/passes/piecewise_graph/nd_tiling_workaround.py +++ b/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,11 +39,9 @@ def __call__(self, graph: torch.fx.Graph): torch._inductor.config.triton.prefer_nd_tiling = True torch._inductor.config.triton.tile_reductions = True - # max_tiles=3 causes two known issues: - # - PT 2.12+: invalid 3D-grid reduction kernels (program_id(2) mapped - # to a non-existent grid dim). - # - All versions: conv-heavy dynamic-shape graphs (e.g. turbo VAE at - # 1080p) can overflow CUDA's z-grid limit (65535) when Inductor's - # coalesce_tiling_analysis collapses high-dim tensors into 3D. - # max_tiles=3 is documented as "experimental and may have bugs". - torch._inductor.config.triton.max_tiles = 2 + # 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