[Bugfix][CI]Fix HSDP regional compile boundary (#4668)

Signed-off-by: bjf-frz <frz123db@gmail.com>

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"]

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"

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):

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()

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
 

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)