[Bugfix] Accept kv_prefetch_jobs in ARDiffusionModelRunner.execute_model (#4941)

Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

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

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)

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