[Perf]Enable regional compile for Cosmos3 HSDP blocks (#4485)

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

tests/diffusion/test_compile.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch.nn as nn
+
+from vllm_omni.diffusion.compile import regionally_compile
+
+
+class _WrappedBlock(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.compile_called = False
+
+    def compile(self, *args, **kwargs):
+        self.compile_called = True
+        return self
+
+
+class _ModelWithWrappedRepeatedBlocks(nn.Module):
+    _repeated_blocks = ["OriginalBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.transformer_blocks = nn.ModuleList([_WrappedBlock(), _WrappedBlock()])
+        self.other_blocks = nn.ModuleList([_WrappedBlock()])
+
+
+def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr():
+    model = _ModelWithWrappedRepeatedBlocks()
+
+    regionally_compile(model)
+
+    assert all(block.compile_called for block in model.transformer_blocks)
+    assert not model.other_blocks[0].compile_called

vllm_omni/diffusion/compile.py

@@ -8,6 +8,25 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 
 
+def _matches_repeated_block(
+    name: str,
+    module: nn.Module,
+    repeated_blocks: list[str],
+    repeated_block_attrs: list[str],
+) -> bool:
+    class_name = module.__class__.__name__
+    if class_name in repeated_blocks:
+        return True
+
+    for attr in ("_fsdp_wrapped_module", "module", "_orig_mod"):
+        wrapped = getattr(module, attr, None)
+        if wrapped is not None and wrapped.__class__.__name__ in repeated_blocks:
+            return True
+
+    parts = name.split(".")
+    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:
     """
     Apply regional compilation to a PyTorch model.
@@ -27,15 +46,25 @@ def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: A
         logger.warning("Regional compilation skipped because the model does not define `_repeated_blocks`.")
         return model
 
+    repeated_block_attrs = getattr(model, "_layerwise_offload_blocks_attrs", [])
+
     # Check if we have modules with the specified class names
     has_compiled_region = False
-    for submod in model.modules():
-        if submod.__class__.__name__ in repeated_blocks:
+    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)
             has_compiled_region = True
+            compiled_region_count += 1
 
     if not has_compiled_region:
         logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.")
+    else:
+        logger.info(
+            "Regional compilation applied to %d module(s) for repeated blocks %s.",
+            compiled_region_count,
+            repeated_blocks,
+        )
 
     return model