diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index d6a920902..9d7a8d245 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -185,9 +185,10 @@ 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. 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 +out of the generic dummy run declare a feature-neutral `dummy_run_recipe` +instead; MiniMax-H3 declares a maximum-duration t2va generation, so the +engine's startup warmup runs it and, under DLO, 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 310d8216d..f77f54ca7 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -69,13 +69,13 @@ omni = Omni( ## 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. +Models whose request geometry needs model-specific sampling arguments for a +warmup generation (currently MiniMax-H3: `t2va`, an explicit aspect ratio) +run one at engine startup whatever features are enabled. It runs at the +maximum output duration so it covers every production request; under DLO +this means the component allocator-cache retention policy learns its +footprint during startup and no real request pays the cold-start flushes. +The tradeoff is one maximum-length generation added to boot time. ## Host-weight loading diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 0e22d087a..b2f2d056b 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -70,7 +70,7 @@ 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: +def test_dummy_run_recipe_returns_none_without_declaration(monkeypatch: pytest.MonkeyPatch) -> None: class PlainModel: pass @@ -80,14 +80,14 @@ def test_dlo_dummy_run_recipe_returns_none_without_declaration(monkeypatch: pyte lambda model_class_name: PlainModel, ) - assert io_support.get_dlo_dummy_run_recipe("plain") is None + assert io_support.get_dummy_run_recipe("plain") is None -def test_dlo_dummy_run_recipe_returns_copy_of_declaration(monkeypatch: pytest.MonkeyPatch) -> None: +def test_dummy_run_recipe_returns_copy_of_declaration(monkeypatch: pytest.MonkeyPatch) -> None: declared = {"task": "t2va", "duration": 4.0} class RecipeModel: - dlo_dummy_run_recipe = declared + dummy_run_recipe = declared monkeypatch.setattr( io_support.DiffusionModelRegistry, @@ -95,7 +95,7 @@ def test_dlo_dummy_run_recipe_returns_copy_of_declaration(monkeypatch: pytest.Mo lambda model_class_name: RecipeModel, ) - recipe = io_support.get_dlo_dummy_run_recipe("recipe") + recipe = io_support.get_dummy_run_recipe("recipe") assert recipe == declared assert recipe is not declared @@ -106,7 +106,7 @@ def test_minimax_h3_declares_a_valid_dlo_warmup_recipe() -> None: MiniMaxH3Pipeline, ) - recipe = MiniMaxH3Pipeline.dlo_dummy_run_recipe + recipe = MiniMaxH3Pipeline.dummy_run_recipe assert recipe is not None assert recipe["task"] == "t2va" assert recipe["aspect_ratio"] == "16:9" @@ -142,17 +142,20 @@ def _opt_out_generic_warmup( lambda model_class_name, supports_audio_input: 0, ) monkeypatch.setattr( - "vllm_omni.diffusion.diffusion_engine.get_dlo_dummy_run_recipe", + "vllm_omni.diffusion.diffusion_engine.get_dummy_run_recipe", lambda model_class_name: recipe, ) -def test_make_dummy_request_skips_recipe_without_dlo(monkeypatch: pytest.MonkeyPatch) -> None: - _opt_out_generic_warmup(monkeypatch, {"task": "t2va"}) +def test_make_dummy_request_applies_recipe_without_dlo(monkeypatch: pytest.MonkeyPatch) -> None: + # The recipe is feature-neutral: a model whose geometry needs it warms up + # whatever features are enabled, not only under DLO. + _opt_out_generic_warmup(monkeypatch, {"task": "t2va", "duration": 4.0, "aspect_ratio": "16:9"}) request = _dlo_recipe_engine(enable_dlo=False)._make_dummy_request(height=512, width=512, guidance_scale=0.0) - assert request is None + assert request is not None + assert request.sampling_params.extra_args["task"] == "t2va" def test_make_dummy_request_builds_recipe_request_when_generic_opted_out( diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 144dba961..40510f200 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -34,9 +34,9 @@ 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, + get_dummy_run_recipe, image_color_format, supports_audio_output, supports_multimodal_input, @@ -65,9 +65,9 @@ logger = init_logger(__name__) _ASYNC_OUTPUT_TIMEOUT_ENV = "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT" _ASYNC_OUTPUT_TIMEOUT_DEFAULT = 600.0 # seconds -# DLO warmup-recipe keys applied to the dummy request's sampling params; -# every other recipe entry is forwarded through the request's extra_args. -_DLO_DUMMY_RUN_SAMPLING_KEYS = frozenset({"height", "width", "num_inference_steps", "num_frames", "guidance_scale"}) +# Warmup-recipe keys applied to the dummy request's sampling params; every +# other recipe entry is forwarded through the request's extra_args. +_DUMMY_RUN_RECIPE_SAMPLING_KEYS = frozenset({"height", "width", "num_inference_steps", "num_frames", "guidance_scale"}) def _async_output_timeout() -> float: @@ -1024,19 +1024,18 @@ class DiffusionEngine: supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config) num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) if num_frames <= 0: - # The model opted out of the generic dummy run. Under distributed - # layerwise offload, fall back to its declared warmup recipe so one - # startup generation primes the component allocator-cache - # retention policy ahead of real traffic. The recipe names its own - # task, so the request stays text-only even when the model's other - # tasks accept image or audio conditions. - 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) + # The model opted out of the generic dummy run; fall back to its + # declared warmup recipe, which supplies whatever feature-specific + # sampling arguments the model requires (under distributed + # layerwise offload the generation also primes the component + # allocator-cache retention policy ahead of real traffic). The + # recipe names its own task, so the request stays text-only even + # when the model's other tasks accept image or audio conditions. + recipe = get_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 _DLO_DUMMY_RUN_SAMPLING_KEYS}) + extra_args.update({k: v for k, v in recipe.items() if k not in _DUMMY_RUN_RECIPE_SAMPLING_KEYS}) return OmniDiffusionRequest( prompt={"prompt": "dummy run"}, request_id=DUMMY_DIFFUSION_REQUEST_ID, diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index d1d145c64..cf9d56b44 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -55,18 +55,19 @@ 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. +def get_dummy_run_recipe(model_class_name: str) -> dict[str, Any] | None: + """Get the model-declared 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. + because their request geometry needs model-specific sampling arguments + (task, duration, aspect ratio, ...) declare a recipe instead. The recipe + is feature-neutral: the engine runs it at startup whatever features are + enabled — under distributed layerwise offload the generation also primes + the component allocator-cache retention policy's observed-peak budget. """ 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 + recipe = getattr(model_cls, "dummy_run_recipe", None) if model_cls is not None else None return dict(recipe) if recipe else None 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 e50439f10..72a1eb405 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -704,14 +704,14 @@ 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] = { + # Startup warmup recipe: 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 covers every production request — under DLO this + # makes the learned cache-retention peak dominate, moving all cold-start + # flushes ahead of real traffic. The cost is one max-length generation in + # boot time whatever features are enabled. + dummy_run_recipe: ClassVar[Mapping[str, Any] | None] = { "task": "t2va", "duration": MINIMAX_H3_MAX_OUTPUT_SECONDS, "aspect_ratio": "16:9",