diff --git a/tests/diffusion/test_attention_compile.py b/tests/diffusion/test_attention_compile.py new file mode 100644 index 000000000..d39744641 --- /dev/null +++ b/tests/diffusion/test_attention_compile.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.forward_context import set_forward_context + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_attention_uses_compile_boundary_for_hsdp(monkeypatch): + attention = object.__new__(Attention) + calls = [] + + def _boundary(query, key, value, attn_metadata=None): + calls.append("boundary") + return query + + def _impl(query, key, value, attn_metadata=None): + calls.append("impl") + return query + + attention._forward_hsdp_compile_boundary = _boundary + attention._forward_impl = _impl + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True)) + query = torch.empty(1) + with set_forward_context(omni_diffusion_config=config): + assert Attention.forward(attention, query, query, query) is query + + assert calls == ["boundary"] + + +def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch): + attention = object.__new__(Attention) + calls = [] + + def _boundary(query, key, value, attn_metadata=None): + calls.append("boundary") + return query + + def _impl(query, key, value, attn_metadata=None): + calls.append("impl") + return query + + attention._forward_hsdp_compile_boundary = _boundary + attention._forward_impl = _impl + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False)) + query = torch.empty(1) + with set_forward_context(omni_diffusion_config=config): + assert Attention.forward(attention, query, query, query) is query + + assert calls == ["impl"] + + +def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch): + attention = object.__new__(Attention) + calls = [] + + def _boundary(query, key, value, attn_metadata=None): + calls.append("boundary") + return query + + def _impl(query, key, value, attn_metadata=None): + calls.append("impl") + return query + + attention._forward_hsdp_compile_boundary = _boundary + attention._forward_impl = _impl + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + query = torch.empty(1) + with set_forward_context(vllm_config=None): + assert Attention.forward(attention, query, query, query) is query + + assert calls == ["impl"] diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py index 522ef343f..ca898da95 100644 --- a/tests/diffusion/test_compile.py +++ b/tests/diffusion/test_compile.py @@ -1,20 +1,28 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest import torch.nn as nn +import vllm_omni.diffusion.compile as compile_module from vllm_omni.diffusion.compile import regionally_compile +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + class _WrappedBlock(nn.Module): def __init__(self) -> None: super().__init__() self.compile_called = False + self.forward_compiled = False def compile(self, *args, **kwargs): self.compile_called = True return self + def forward(self, x): + return x + class _ModelWithWrappedRepeatedBlocks(nn.Module): _repeated_blocks = ["OriginalBlock"] @@ -26,10 +34,23 @@ class _ModelWithWrappedRepeatedBlocks(nn.Module): self.other_blocks = nn.ModuleList([_WrappedBlock()]) -def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(): +def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(monkeypatch): model = _ModelWithWrappedRepeatedBlocks() + compile_calls = [] - regionally_compile(model) + def _compile(fn, *args, **kwargs): + compile_calls.append((fn, args, kwargs)) - assert all(block.compile_called for block in model.transformer_blocks) + 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, dynamic=True) + + assert len(compile_calls) == 2 + assert all(not block.compile_called for block in model.transformer_blocks) assert not model.other_blocks[0].compile_called + assert model.transformer_blocks[0].forward("ok") == "compiled:ok" diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 95a0d9c27..746377372 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -109,6 +109,37 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar return runner +def _make_compile_runner(*, use_hsdp: bool): + runner = object.__new__(DiffusionModelRunner) + runner.pipeline = SimpleNamespace(transformer=SimpleNamespace()) + runner.od_config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=use_hsdp)) + return runner + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("use_hsdp", [False, True]) +def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp): + runner = _make_compile_runner(use_hsdp=use_hsdp) + compile_calls = [] + + def _regionally_compile(model, *args, **kwargs): + compile_calls.append((model, args, kwargs)) + return model + + monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile) + + DiffusionModelRunner._compile_transformer(runner, "transformer") + + assert compile_calls == [ + ( + runner.pipeline.transformer, + (), + {"dynamic": True}, + ) + ] + + @pytest.mark.core_model @pytest.mark.cpu def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index e0fb8d264..0ea122273 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -230,6 +230,34 @@ class Attention(nn.Module): key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + if torch.compiler.is_compiling() and is_forward_context_available(): + od_config = get_forward_context().omni_diffusion_config + parallel_config = getattr(od_config, "parallel_config", None) + if getattr(parallel_config, "use_hsdp", False): + # Keep HSDP/FSDP2 parameter all-gather outside Inductor's + # attention graph; otherwise scheduler dependency analysis can + # fail on the fused attention region. + return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata) + + return self._forward_impl(query, key, value, attn_metadata) + + @torch.compiler.disable + def _forward_hsdp_compile_boundary( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + return self._forward_impl(query, key, value, attn_metadata) + + def _forward_impl( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: # Get the appropriate parallel strategy based on SP active state strategy = self._get_active_parallel_strategy() diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 0e4a3ff50..45216febc 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any +import torch import torch.nn as nn from vllm.logger import init_logger @@ -27,7 +28,11 @@ def _matches_repeated_block( return len(parts) >= 2 and parts[-2] in repeated_block_attrs and parts[-1].isdigit() -def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: Any) -> nn.Module: +def regionally_compile( + model: nn.Module, + *compile_args: Any, + **compile_kwargs: Any, +) -> nn.Module: """ Apply regional compilation to a PyTorch model. @@ -53,8 +58,9 @@ def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: A compiled_region_count = 0 for name, submod in model.named_modules(): if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs): - # Compile this submodule - submod.compile(*compile_args, **compile_kwargs) + # Compile the block compute while keeping nn.Module.__call__ hooks + # outside the compiled graph. + submod.forward = torch.compile(submod.forward, *compile_args, **compile_kwargs) has_compiled_region = True compiled_region_count += 1 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 9f3598abe..fffea4a40 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -88,6 +88,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): model = getattr(self.pipeline, attr_name, None) if model is None: return + try: setattr(self.pipeline, attr_name, regionally_compile(model, dynamic=True)) logger.info("Model runner: %s compiled with torch.compile.", attr_name)