fix: keep DLO hooks outside regional compilation

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/test_compile.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 
 import vllm_omni.diffusion.compile as compile_module
 from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -76,6 +77,37 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
     assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
 
 
+def test_regionally_compile_keeps_hook_dispatch_outside_compiled_graph(monkeypatch):
+    model = _ModelWithWrappedRepeatedBlocks()
+    block = model.transformer_blocks[0]
+    registry = HookRegistry.get_or_create(block)
+    hook = ModelHook()
+    registry.register_hook("test", hook)
+
+    wrapped_forward = block.forward
+    original_forward = block._omni_original_forward
+    compile_calls = []
+
+    def _compile(fn, *args, **kwargs):
+        compile_calls.append(fn)
+
+        def _compiled(*fn_args, **fn_kwargs):
+            return f"compiled:{fn(*fn_args, **fn_kwargs)}"
+
+        return _compiled
+
+    monkeypatch.setattr(compile_module.torch, "compile", _compile)
+
+    regionally_compile(model)
+
+    assert compile_calls[0] is original_forward
+    assert compile_calls[0] is not wrapped_forward
+    assert block.forward is wrapped_forward
+    assert block._omni_original_forward is not original_forward
+    assert hook.fn_ref.original_forward is block._omni_original_forward
+    assert block("ok") == "compiled:ok"
+
+
 def test_compiled_block_preserves_forward_signature_for_inspection(monkeypatch):
     """cache-dit matches blocks via inspect.signature(block.forward).
 

vllm_omni/diffusion/compile.py

@@ -56,26 +56,50 @@ def regionally_compile(
     # Build all compiled callables before mutating the model. This keeps setup
     # failures atomic: callers can safely continue with the uncompiled model if
     # torch.compile raises synchronously for any repeated block.
-    compiled_forwards: list[tuple[nn.Module, Any]] = []
+    # Keep the pre-hook callable separately.  DLO and the ordinary layerwise
+    # offloader install a HookRegistry wrapper in ``module.forward``.  That
+    # wrapper performs stream/event synchronization and storage rebinding, so
+    # compiling it pulls offload control flow into the graph.  Compile the
+    # original block compute instead and leave the wrapper outside the graph.
+    compiled_forwards: list[tuple[nn.Module, Any | None, Any]] = []
     for name, submod in model.named_modules():
         if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs):
             # Compile the block compute while keeping nn.Module.__call__ hooks
-            # outside the compiled graph. NOTE: anything that wraps this
-            # callable must stay signature-transparent — cache-dit's
-            # BlockAdapter matches blocks by inspecting forward's parameter
-            # names and return annotation, which torch.compile preserves.
+            # outside the compiled graph. If a HookRegistry is already
+            # installed, ``submod.forward`` is the hook dispatcher and the
+            # original callable is stored in ``_omni_original_forward``.
+            # NOTE: anything that wraps this callable must stay
+            # signature-transparent — cache-dit's BlockAdapter matches blocks
+            # by inspecting forward's parameter names and return annotation,
+            # which torch.compile preserves.
+            original_forward = getattr(submod, "_omni_original_forward", None)
+            forward_to_compile = original_forward if original_forward is not None else submod.forward
             compiled_forwards.append(
                 (
                     submod,
-                    torch.compile(submod.forward, *compile_args, **compile_kwargs),
+                    original_forward,
+                    torch.compile(forward_to_compile, *compile_args, **compile_kwargs),
                 )
             )
 
     if not compiled_forwards:
         logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.")
     else:
-        for submod, compiled_forward in compiled_forwards:
-            submod.forward = compiled_forward
+        for submod, original_forward, compiled_forward in compiled_forwards:
+            if original_forward is None:
+                submod.forward = compiled_forward
+                continue
+
+            # Keep HookRegistry's dispatcher as ``forward`` and replace only
+            # the callable it dispatches to.  Hooks such as MagCache retain a
+            # reference to the original callable as well, so update those
+            # references when they still point at the pre-compile function.
+            submod._omni_original_forward = compiled_forward
+            registry = getattr(submod, "_hook_registry", None)
+            for hook in getattr(registry, "_hooks", {}).values():
+                fn_ref = getattr(hook, "fn_ref", None)
+                if fn_ref is not None and fn_ref.original_forward is original_forward:
+                    fn_ref.original_forward = compiled_forward
         logger.info(
             "Regional compilation applied to %d module(s) for repeated blocks %s.",
             len(compiled_forwards),