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

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

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"