[Model] Prime DLO cache retention with a startup warmup

Models whose request geometry opts out of the generic dummy run (H3: t2va
needs an explicit aspect ratio and a bounded duration) can declare a
dlo_dummy_run_recipe. Under distributed layerwise offload the engine's
startup warmup runs it, so the learned cache-retention budget is primed
before real traffic instead of on the first request.

H3 declares a maximum-duration t2va generation: duration is the only
workload variable at the fixed 768 short edge, so the learned peak
dominates every production request and no real request re-learns. Validated
on L20X: all cold-start releases happen during the startup warmup, zero
during requests; first request 15.91 s vs 16.7 s unwarmed, steady 14.9 s;
startup gains one max-length generation (~10-20 s). Non-DLO deployments
keep the opted-out generic warmup unchanged.

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -184,7 +184,10 @@ the weight-reuse payload plus the transient footprint whatever the
 resolution. It starts at zero: a cold start pays one release per observation
 step while the peak climbs, after which boundaries retain steadily, and a
 workload change whose cached footprint grows by more than the margin pays one
-further release while the new peak lands. The budget is capped at a quarter
+further release while the new peak lands. Models whose request geometry opts
+out of the generic dummy run can declare `dlo_dummy_run_recipe`; MiniMax-H3
+declares one minimal t2va generation, so under DLO the engine's startup
+warmup runs it and the cold-start flushes land before real traffic. The budget is capped at a quarter
 of device capacity so a spurious observation cannot monopolize a small
 device, and detaching the last attached stager resets the learned peak so one
 workload never permanently influences another. No user-chosen fraction, no

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -67,6 +67,16 @@ omni = Omni(
 | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
 | `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
 
+## Startup warmup
+
+Under DLO, models that need model-specific sampling arguments for a warmup
+generation (currently MiniMax-H3: `t2va`, an explicit aspect ratio) run one
+at engine startup. It runs at the maximum output duration so the component
+allocator-cache retention policy learns a footprint that covers every
+production request: the cold-start flushes happen during startup and no real
+request pays them. The tradeoff is one maximum-length generation added to
+boot time; on deployments without DLO the behavior is unchanged.
+
 ## Host-weight loading
 
 The diffusion loader chooses host storage before DLO is enabled. It first

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -70,6 +70,122 @@ def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
     assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
 
 
+def test_dlo_dummy_run_recipe_returns_none_without_declaration(monkeypatch: pytest.MonkeyPatch) -> None:
+    class PlainModel:
+        pass
+
+    monkeypatch.setattr(
+        io_support.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda model_class_name: PlainModel,
+    )
+
+    assert io_support.get_dlo_dummy_run_recipe("plain") is None
+
+
+def test_dlo_dummy_run_recipe_returns_copy_of_declaration(monkeypatch: pytest.MonkeyPatch) -> None:
+    declared = {"task": "t2va", "duration": 4.0}
+
+    class RecipeModel:
+        dlo_dummy_run_recipe = declared
+
+    monkeypatch.setattr(
+        io_support.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda model_class_name: RecipeModel,
+    )
+
+    recipe = io_support.get_dlo_dummy_run_recipe("recipe")
+    assert recipe == declared
+    assert recipe is not declared
+
+
+def test_minimax_h3_declares_a_valid_dlo_warmup_recipe() -> None:
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        MINIMAX_H3_MAX_OUTPUT_SECONDS,
+        MiniMaxH3Pipeline,
+    )
+
+    recipe = MiniMaxH3Pipeline.dlo_dummy_run_recipe
+    assert recipe is not None
+    assert recipe["task"] == "t2va"
+    assert recipe["aspect_ratio"] == "16:9"
+    # Maximum duration makes the learned peak dominate every production
+    # request at the fixed short edge, so no real request re-learns.
+    assert float(recipe["duration"]) == MINIMAX_H3_MAX_OUTPUT_SECONDS
+    assert MiniMaxH3Pipeline.dummy_run_num_frames == 0  # generic warmup stays opted out
+
+
+def _dlo_recipe_engine(enable_dlo: bool) -> DiffusionEngine:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        model_class_name="RecipeModel",
+        enable_distributed_layerwise_offload=enable_dlo,
+    )
+    return engine
+
+
+def test_make_dlo_dummy_request_skipped_without_dlo(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_dlo_dummy_run_recipe",
+        lambda model_class_name: {"task": "t2va"},
+    )
+
+    assert _dlo_recipe_engine(enable_dlo=False)._make_dlo_dummy_request() is None
+
+
+def test_make_dlo_dummy_request_builds_recipe_request(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_dlo_dummy_run_recipe",
+        lambda model_class_name: {
+            "task": "t2va",
+            "duration": 4.0,
+            "aspect_ratio": "16:9",
+            "num_inference_steps": 2,
+            "guidance_scale": 1.0,
+        },
+    )
+
+    request = _dlo_recipe_engine(enable_dlo=True)._make_dlo_dummy_request()
+
+    assert request is not None
+    assert request.is_dummy_run()
+    assert request.sampling_params.num_inference_steps == 2
+    assert request.sampling_params.guidance_scale == 1.0
+    extra = request.sampling_params.extra_args
+    assert extra["task"] == "t2va"
+    assert extra["duration"] == 4.0
+    assert extra["aspect_ratio"] == "16:9"
+    assert extra["cfg_text_scale"] == 1.0
+
+
+def test_dummy_run_falls_back_to_dlo_recipe_when_generic_opted_out(monkeypatch: pytest.MonkeyPatch) -> None:
+    engine = _dlo_recipe_engine(enable_dlo=True)
+    recipe_request = OmniDiffusionRequest(
+        prompt="dummy run",
+        request_id="dlo-dummy",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2),
+    )
+    engine._make_dummy_request = Mock(return_value=None)
+    engine._make_dlo_dummy_request = Mock(return_value=recipe_request)
+    engine.add_req_and_wait_for_response = Mock(return_value=SimpleNamespace(error=None, request_id="dlo-dummy"))
+
+    engine._dummy_run()
+
+    engine._make_dlo_dummy_request.assert_called_once()
+    engine.add_req_and_wait_for_response.assert_called_once_with(recipe_request)
+
+
+def test_dummy_run_still_skips_without_recipe(monkeypatch: pytest.MonkeyPatch) -> None:
+    engine = _dlo_recipe_engine(enable_dlo=True)
+    engine._make_dummy_request = Mock(return_value=None)
+    engine._make_dlo_dummy_request = Mock(return_value=None)
+
+    engine._dummy_run()
+
+    engine._make_dlo_dummy_request.assert_called_once()
+
+
 def test_dense_mode_does_not_build_kv_profile_request() -> None:
     engine = object.__new__(DiffusionEngine)
     engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)

vllm_omni/diffusion/diffusion_engine.py

@@ -34,6 +34,7 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_sch
 from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.io_support import (
+    get_dlo_dummy_run_recipe,
     get_dummy_run_num_frames,
     get_dummy_run_num_image_inputs,
     image_color_format,
@@ -1102,6 +1103,10 @@ class DiffusionEngine:
             profile_requests.append(profile_request)
         return profile_requests
 
+    # Sampling-recipe keys consumed by _make_dlo_dummy_request directly;
+    # everything else is forwarded through the request's extra_args.
+    _DLO_DUMMY_RUN_SAMPLING_KEYS = frozenset({"height", "width", "num_inference_steps", "num_frames", "guidance_scale"})
+
     def _dummy_run(self):
         """A dummy run to warm up the model."""
         req = self._make_dummy_request(
@@ -1110,13 +1115,47 @@ class DiffusionEngine:
             guidance_scale=0.0,
         )
         if req is None:
-            logger.info("Skipping dummy warmup run (num_frames=0)")
-            return
+            req = self._make_dlo_dummy_request()
+            if req is None:
+                logger.info("Skipping dummy warmup run (num_frames=0)")
+                return
         logger.info("dummy run to warm up the model")
         output = self.add_req_and_wait_for_response(req)
         if output.error:
             raise RuntimeError(f"Dummy run failed: {output.error}")
 
+    def _make_dlo_dummy_request(self) -> OmniDiffusionRequest | None:
+        """Build the model-declared startup warmup used under distributed offload.
+
+        Models whose request geometry needs model-specific sampling arguments
+        (task, duration, aspect ratio, ...) opt out of the generic dummy run
+        and declare ``dlo_dummy_run_recipe`` instead. Under distributed
+        layerwise offload one startup generation also primes the component
+        allocator-cache retention policy's observed-peak budget, so the
+        cold-start flushes land before real traffic.
+        """
+
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            return None
+        recipe = get_dlo_dummy_run_recipe(self.od_config.model_class_name)
+        if recipe is None:
+            return None
+        extra_args: dict[str, Any] = {"cfg_text_scale": 1.0, "cfg_img_scale": 1.0}
+        extra_args.update({k: v for k, v in recipe.items() if k not in self._DLO_DUMMY_RUN_SAMPLING_KEYS})
+        return OmniDiffusionRequest(
+            prompt={"prompt": "dummy run"},
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
+            sampling_params=OmniDiffusionSamplingParams(
+                height=int(recipe.get("height", 512)),
+                width=int(recipe.get("width", 512)),
+                num_inference_steps=int(recipe.get("num_inference_steps", 1)),
+                num_frames=int(recipe.get("num_frames", 0)),
+                guidance_scale=float(recipe.get("guidance_scale", 0.0)),
+                num_outputs_per_prompt=1,
+                extra_args=extra_args,
+            ),
+        )
+
     def _submit_rpc(
         self,
         method: str,

vllm_omni/diffusion/io_support.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import inspect
+from typing import Any
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
@@ -54,6 +55,21 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool)
     return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
 
 
+def get_dlo_dummy_run_recipe(model_class_name: str) -> dict[str, Any] | None:
+    """Get the model-declared DLO startup-warmup recipe, or None.
+
+    Models that opt out of the generic dummy run (``dummy_run_num_frames = 0``)
+    because their request geometry needs model-specific sampling arguments can
+    still declare a recipe used only when distributed layerwise offload is
+    enabled, where one startup generation primes the component allocator-cache
+    retention policy.
+    """
+
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    recipe = getattr(model_cls, "dlo_dummy_run_recipe", None) if model_cls is not None else None
+    return dict(recipe) if recipe else None
+
+
 def get_dummy_run_num_image_inputs(model_class_name: str) -> int:
     """Return the maximum advertised image-input count for profiling."""
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -704,6 +704,20 @@ class MiniMaxH3Pipeline(
         "post_decode",
     ]
     dummy_run_num_frames: ClassVar[int] = 0
+    # Startup warmup used only when distributed layerwise offload is enabled:
+    # the generic dummy run cannot satisfy t2va's required sampling arguments
+    # (explicit aspect ratio, bounded duration). Duration is the only workload
+    # variable at the fixed 768 short edge, so warming at the maximum makes the
+    # learned cache-retention peak dominate every production request and moves
+    # all cold-start flushes ahead of real traffic, at the cost of one
+    # max-length generation in boot time.
+    dlo_dummy_run_recipe: ClassVar[Mapping[str, Any] | None] = {
+        "task": "t2va",
+        "duration": MINIMAX_H3_MAX_OUTPUT_SECONDS,
+        "aspect_ratio": "16:9",
+        "num_inference_steps": 2,
+        "guidance_scale": 1.0,
+    }
     # Only distilled releases pin a schedule, so the default keeps the legacy
     # uniform path available to partially constructed pipelines.
     _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}