diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 34abbfec8..d6a920902 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/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 diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 2c6bbacb4..310d8216d 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/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 diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 247fee029..6c648dc93 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/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) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index f3294c5fa..380b45939 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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, diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 8fecc50df..d1d145c64 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/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.""" diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index ff3990b32..e50439f10 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/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]] = {}