[Diffusion] Tighten runtime cache compatibility guards

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -260,7 +260,8 @@ must use `--dlo-no-use-allgather` or disable online quantization until their
 runtime layouts are validated.
 
 Runtime-cache v1 is no-AllGather only and rejects all quantization
-configurations, HSDP/DTensor, CFG parallelism, PP, non-contiguous tensors, and
+configurations, HSDP/DTensor, expert parallelism, CFG parallelism, PP,
+non-contiguous tensors, and
 tied/shared storage.
 
 ## Validation coverage

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -197,7 +197,7 @@ must enter each collective.
   online quantization methods require no-AllGather until their runtime layouts
   are validated.
 - Runtime-cache v1 rejects all quantized, non-contiguous, aliased/tied,
-  device-only, HSDP, CFG-parallel, and PP layouts.
+  device-only, HSDP, expert-parallel, CFG-parallel, and PP layouts.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -12,6 +12,7 @@ import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
@@ -436,6 +437,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     model = nn.Module()
     model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnquantizedLinearMethod()
     events: list[str] = []
     runtime_plan = HostWeightPlan(
         backing_kind="runtime_cache",
@@ -470,6 +472,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
         assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
         assert kwargs["tensor_parallel_rank"] == 0
         assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
+        assert kwargs["quantization_config"] is None
         return HostWeightPlanResult(runtime_plan)
 
     loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
@@ -513,6 +516,7 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     model = nn.Module()
     model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object()
     loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
     loader.load_weights = lambda _model: None  # type: ignore[method-assign]
     loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
@@ -522,11 +526,12 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat
         "build_checkpoint_mmap_plan",
         lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
     )
-    monkeypatch.setattr(
-        cache_mod,
-        "build_runtime_weight_cache_plan",
-        lambda *_args, **_kwargs: HostWeightPlanResult(None, "writer timed out", "lock_timeout"),
-    )
+
+    def build_cache(_pipeline, **kwargs):
+        assert kwargs["quantization_config"] == "module_quant_method"
+        return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
+
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
 
     original_weight = model.transformer.weight.detach().clone()
     assert loader.load_model(load_device="cpu") is model

tests/diffusion/model_loader/test_runtime_weight_cache.py

@@ -61,6 +61,7 @@ def _build(pipeline: nn.Module, cache_root: Path, **overrides):
         "tensor_parallel_rank": 0,
         "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
         "use_hsdp": False,
+        "enable_expert_parallel": False,
         "quantization_config": None,
         "cfg_parallel_size": 1,
         "pipeline_parallel_size": 1,
@@ -125,6 +126,14 @@ def test_final_runtime_content_is_authoritative(tmp_path):
     assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
 
 
+def test_expert_parallel_ownership_fails_closed(tmp_path):
+    result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
+
+    assert result.plan is None
+    assert result.fallback_code == "unsupported_expert_parallel"
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
 def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
     first = _build(_Pipeline(), tmp_path)
     assert first.plan is not None

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -14,6 +14,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
@@ -556,6 +557,7 @@ class DiffusersPipelineLoader:
 
         modules = ModuleDiscovery.discover(model)
         from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+            DEFAULT_LOCK_TIMEOUT_SECONDS,
             DEFAULT_SHARD_SIZE_BYTES,
             build_runtime_weight_cache_plan,
         )
@@ -577,12 +579,21 @@ class DiffusersPipelineLoader:
             "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"),
         }
         model_identity = getattr(self.od_config, "model", None)
+        runtime_quantization = self.quant_config
+        if runtime_quantization is None and any(
+            (quant_method := getattr(module, "quant_method", None)) is not None
+            and not isinstance(quant_method, UnquantizedLinearMethod)
+            for module in model.modules()
+        ):
+            runtime_quantization = "module_quant_method"
         result = build_runtime_weight_cache_plan(
             model,
             dit_modules=tuple(zip(modules.dit_names, modules.dits)),
             loader_type=type(self),
             cache_root=getattr(self.od_config, "dlo_runtime_cache_dir", None),
-            lock_timeout_seconds=float(getattr(self.od_config, "dlo_runtime_cache_lock_timeout", 600.0)),
+            lock_timeout_seconds=float(
+                getattr(self.od_config, "dlo_runtime_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS)
+            ),
             max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES,
             model_identity=str(model_identity) if model_identity is not None else None,
             revision=getattr(self.od_config, "revision", None),
@@ -593,7 +604,8 @@ class DiffusersPipelineLoader:
             tensor_parallel_rank=tp_rank,
             sequence_parallel_guard=sp_guard,
             use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)),
-            quantization_config=self.quant_config,
+            enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)),
+            quantization_config=runtime_quantization,
             cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)),
             pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)),
         )

vllm_omni/diffusion/model_loader/runtime_weight_cache.py

@@ -572,6 +572,7 @@ def build_runtime_weight_cache_plan(
     tensor_parallel_rank: int,
     sequence_parallel_guard: Mapping[str, Any],
     use_hsdp: bool,
+    enable_expert_parallel: bool,
     quantization_config: Any,
     cfg_parallel_size: int,
     pipeline_parallel_size: int,
@@ -585,6 +586,12 @@ def build_runtime_weight_cache_plan(
         )
     if use_hsdp:
         return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp")
+    if enable_expert_parallel:
+        return HostWeightPlanResult(
+            None,
+            "expert-parallel weight ownership is outside runtime-cache v1",
+            "unsupported_expert_parallel",
+        )
     if quantization_config is not None:
         return HostWeightPlanResult(
             None,