diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py index fc59e54ac..fb55e01e9 100644 --- a/tests/ar_diffusion/test_engine_routing.py +++ b/tests/ar_diffusion/test_engine_routing.py @@ -134,3 +134,17 @@ def test_runner_rejects_multi_sequence_config(): fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4)) with pytest.raises(ValueError, match="max_num_seqs=1"): ARDiffusionModelRunner._preallocate_kv_cache(fake) + + +def test_execute_model_accepts_kv_prefetch_jobs_kwarg(): + """Regression: the worker passes ``kv_prefetch_jobs`` to every runner + (added with the KV prefetch path, #4448); the AR-Diffusion override must + accept and forward it or DreamZero crashes on the warm-up dummy run.""" + import inspect + + from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner + + for cls in (DiffusionModelRunner, ARDiffusionModelRunner): + params = inspect.signature(cls.execute_model).parameters + assert "kv_prefetch_jobs" in params, cls diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py index 83c710efa..28c3c959d 100644 --- a/tests/ar_diffusion/test_pipeline_bridge.py +++ b/tests/ar_diffusion/test_pipeline_bridge.py @@ -203,7 +203,7 @@ def test_failed_forward_tears_down_session(monkeypatch): runner.device = None runner._perf_e2e_times = [] - def boom(self, req): + def boom(self, req, kv_prefetch_jobs=None): raise RuntimeError("layer 17 exploded") monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom) diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py index c8b8e93e1..862b3d9e8 100644 --- a/vllm_omni/experimental/ar_diffusion/runner.py +++ b/vllm_omni/experimental/ar_diffusion/runner.py @@ -234,10 +234,10 @@ class ARDiffusionModelRunner(DiffusionModelRunner): num_frame_per_block=num_frame_per_block, ) - def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput: # KV disabled -> base behavior, unchanged. if self.kv_cache is None: - return super().execute_model(req) + return super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs) kv = self.kv_cache # DreamZero KV is session-scoped (the model state persists across a @@ -273,7 +273,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner): # sync below makes the stop reflect completed GPU work, not just dispatch. _e2e_t0 = time.perf_counter() try: - out = super().execute_model(req) + out = super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs) except Exception: # Transactional containment: a forward that died partway may have # written some layers' K/V into allocated-but-uncommitted blocks