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