diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py index 665aa6825..3cf8f1428 100644 --- a/tests/diffusion/test_compile.py +++ b/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). diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 4a050ab38..11fc2c113 100644 --- a/vllm_omni/diffusion/compile.py +++ b/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),