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