[Perf][DreamZero]dreamzero supports CUDA graph, torch.compile, and DiT Caching. (#4213)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

examples/offline_inference/dreamzero/client_schedule.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py).
+ACTION_HORIZON = 24
+DEFAULT_ACTION_DIM = 8
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+DEFAULT_NUM_AR_CHUNKS = 15
+
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def build_frame_schedule(
+    total_frames: int,
+    num_chunks: int,
+    *,
+    logger: logging.Logger | None = None,
+) -> list[list[int]]:
+    """Build the frame index schedule for multi-frame chunks.
+
+    Args:
+        total_frames: Number of frames available in each camera video.
+        num_chunks: Number of 4-frame chunks to schedule after the initial frame.
+
+    Returns:
+        A list of frame-index lists. Each inner list has four indices. The
+        returned list may be shorter than ``num_chunks`` when the videos run
+        out of frames (upstream stops early instead of erroring).
+    """
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            if logger is not None:
+                logger.info(
+                    "Frame %s >= %s, stopping at %s chunks",
+                    indices[-1],
+                    total_frames,
+                    len(chunks),
+                )
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
+    prompt: str,
+    session_id: str,
+) -> dict[str, Any]:
+    """Build an observation dict from real video frames.
+
+    For one frame each image key is ``(H, W, 3)``. For four frames each key is
+    ``(4, H, W, 3)``.
+    """
+    obs: dict[str, Any] = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def build_ar_observations(
+    camera_frames: dict[str, np.ndarray],
+    *,
+    prompt: str,
+    session_id: str,
+    num_chunks: int = DEFAULT_NUM_AR_CHUNKS,
+    repeat_chunk_observations: bool = False,
+    logger: logging.Logger | None = None,
+) -> list[dict[str, Any]]:
+    """Build the AR observation sequence used by upstream ``test_client_AR.py``.
+
+    Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk.
+
+    ``num_chunks`` counts only the 4-frame chunks after the initial frame, not
+    the total number of inferences.
+    """
+    if num_chunks < 0:
+        raise ValueError("num_chunks must be non-negative")
+
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    observations = [
+        make_obs_from_video(
+            camera_frames,
+            [0],
+            prompt=prompt,
+            session_id=session_id,
+        )
+    ]
+
+    chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger)
+    if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks:
+        while len(chunk_schedule) < num_chunks:
+            chunk_schedule.append(chunk_schedule[-1])
+
+    for frame_indices in chunk_schedule:
+        observations.append(
+            make_obs_from_video(
+                camera_frames,
+                frame_indices,
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
+    return observations

examples/offline_inference/dreamzero/export_prediction_video.py

@@ -11,6 +11,11 @@ from pathlib import Path
 import cv2
 import numpy as np
 import torch
+from client_schedule import (
+    CAMERA_FILES,
+    DEFAULT_NUM_AR_CHUNKS,
+    build_ar_observations,
+)
 from PIL import Image
 
 from vllm_omni import Omni
@@ -20,13 +25,6 @@ from vllm_omni.outputs import OmniRequestOutput
 WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
 ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
 DEFAULT_SESSION_PREFIX = "dreamzero-export"
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-ACTION_HORIZON = 24
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
 
 
 def _parse_args() -> argparse.Namespace:
@@ -49,6 +47,17 @@ def _parse_args() -> argparse.Namespace:
         default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
     )
     parser.add_argument("--session-id", default=None)
+    parser.add_argument(
+        "--num-chunks",
+        type=int,
+        default=DEFAULT_NUM_AR_CHUNKS,
+        help="Number of 4-frame chunks to send after the initial frame (default: 15).",
+    )
+    parser.add_argument(
+        "--repeat-chunk-observations",
+        action="store_true",
+        help="Repeat the last valid chunk observation when bundled assets run out of frames.",
+    )
     parser.add_argument("--save-input-video", action="store_true")
     parser.add_argument("--save-gif", action="store_true")
     parser.add_argument("--save-actions", action="store_true")
@@ -84,57 +93,25 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
     return camera_frames
 
 
-def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
-    chunks: list[list[int]] = []
-    current_frame = 23
-    for _ in range(num_chunks):
-        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def _make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    *,
+def _build_observations(
+    video_dir: Path,
     prompt: str,
     session_id: str,
-) -> dict:
-    obs: dict = {}
-    for camera_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]
-        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
+    *,
+    num_chunks: int,
+    repeat_chunk_observations: bool,
+) -> tuple[dict[str, np.ndarray], list[dict]]:
     camera_frames = _load_camera_frames(video_dir)
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    chunks = _build_frame_schedule(total_frames, 1)
-    observations = [
-        _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
-    ]
-    if chunks:
-        observations.append(
-            _make_obs_from_video(
-                camera_frames,
-                chunks[0],
-                prompt=prompt,
-                session_id=session_id,
-            )
-        )
+    observations = build_ar_observations(
+        camera_frames,
+        prompt=prompt,
+        session_id=session_id,
+        num_chunks=num_chunks,
+        repeat_chunk_observations=repeat_chunk_observations,
+    )
     if len(observations) < 2:
-        raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
-    return camera_frames, observations[:2]
+        raise RuntimeError("Need at least two DreamZero observations to export a prediction video.")
+    return camera_frames, observations
 
 
 def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
@@ -229,6 +206,26 @@ def _run_generation(
     return omni, outputs
 
 
+def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray:
+    """Decode server-side accumulated latents when the pipeline supports it."""
+    stage_client = omni.engine.stage_clients[0]
+    engine = getattr(stage_client, "_engine", None)
+    if engine is None:
+        raise RuntimeError("DreamZero export requires inline diffusion stage access.")
+
+    decoded = engine.executor.collective_rpc(
+        "decode_accumulated_video_latents_to_uint8",
+        args=(session_id,),
+        unique_reply_rank=0,
+        exec_all_ranks=True,
+    )
+    if isinstance(decoded, torch.Tensor):
+        decoded = decoded.numpy()
+    if not isinstance(decoded, np.ndarray):
+        raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
+    return decoded
+
+
 def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
     stage_client = omni.engine.stage_clients[0]
     engine = getattr(stage_client, "_engine", None)
@@ -248,6 +245,15 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
     return decoded
 
 
+def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray:
+    """Decode full AR rollout by concatenating per-step latents from all outputs."""
+    latent_steps = [_extract_latents(output) for output in outputs]
+    latent_t = [step.shape[2] for step in latent_steps]
+    print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}")
+    full_latents = torch.cat(latent_steps, dim=2)
+    return _decode_with_worker(omni, full_latents)
+
+
 def main() -> None:
     args = _parse_args()
     session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
@@ -256,6 +262,12 @@ def main() -> None:
         video_dir=args.video_dir,
         prompt=args.prompt,
         session_id=session_id,
+        num_chunks=args.num_chunks,
+        repeat_chunk_observations=args.repeat_chunk_observations,
+    )
+    print(
+        f"Running {len(observations)} inferences "
+        f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)"
     )
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
@@ -273,17 +285,19 @@ def main() -> None:
             deploy_config_path=args.deploy_config,
             observations=observations,
         )
-        latent_steps = [_extract_latents(output) for output in outputs]
-        full_latents = torch.cat(latent_steps, dim=2)
-        frames = _decode_with_worker(omni, full_latents)
+        frames = _decode_video(omni, session_id, outputs)
     finally:
         if omni is not None:
             omni.close()
 
+    expected_frames = 8 * len(observations) + 1
     mp4_path = args.output_dir / f"{args.output_stem}.mp4"
 
     _write_mp4(mp4_path, frames, fps=args.fps)
-    print(f"SAVED_MP4={mp4_path}")
+    print(
+        f"SAVED_MP4={mp4_path} frames={len(frames)} "
+        f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s"
+    )
 
     if args.save_gif:
         gif_path = args.output_dir / f"{args.output_stem}.gif"
@@ -294,8 +308,10 @@ def main() -> None:
         npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
         np.savez(
             npz_path,
-            step0=np.asarray(outputs[0].multimodal_output.get("actions")),
-            step1=np.asarray(outputs[1].multimodal_output.get("actions")),
+            **{
+                f"step{index}": np.asarray(output.multimodal_output.get("actions"))
+                for index, output in enumerate(outputs)
+            },
         )
         print(f"SAVED_ACTIONS={npz_path}")
 

tests/diffusion/cache/test_stepcache.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for StepCacheBackend and StepCacheState."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cache.stepcache import (
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheBackend,
+    StepCacheConfig,
+    StepCacheState,
+    is_stepcache_active,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]:
+    v_prev = torch.randn(1, dim)
+    if sim >= 0.99:
+        v_last = v_prev.clone()
+    else:
+        v_last = torch.randn(1, dim)
+    return [(v_prev,), (v_last,)]
+
+
+class TestStepCacheConfig:
+    def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self):
+        enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True))
+        disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False))
+        assert enabled.enabled is True
+        assert disabled.enabled is False
+
+
+class TestStepCacheState:
+    def test_disabled_always_runs(self):
+        state = StepCacheState(StepCacheConfig(enabled=False))
+        assert state.should_run_step(_make_prev_predictions(0.99)) is True
+
+    def test_insufficient_history_always_runs(self):
+        state = StepCacheState(StepCacheConfig(min_history_steps=2))
+        assert state.should_run_step([(torch.randn(1, 4),)]) is True
+
+    def test_high_similarity_skips_with_countdown(self):
+        state = StepCacheState(
+            StepCacheConfig(
+                sim_thresholds=(0.95,),
+                skip_countdowns=(2,),
+            )
+        )
+        preds = _make_prev_predictions(0.99)
+        assert state.should_run_step(preds) is False
+        assert state.skip_countdown == 2
+        assert state.should_run_step(preds) is False
+        assert state.should_run_step(preds) is True
+
+    def test_low_similarity_runs(self):
+        state = StepCacheState(StepCacheConfig())
+        preds = _make_prev_predictions(0.5)
+        assert state.should_run_step(preds) is True
+
+    def test_trim_history(self):
+        state = StepCacheState(StepCacheConfig(max_history=2))
+        history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)]
+        state.trim_history(history)
+        assert len(history) == 2
+
+
+class TestStepCacheBackend:
+    def test_enable_dreamzero_pipeline(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
+
+        backend = StepCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR)
+        assert isinstance(cfg, StepCacheConfig)
+        assert is_stepcache_active(mock_pipeline)
+        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
+        assert isinstance(state, StepCacheState)
+
+    def test_refresh_resets_state(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
+        backend = StepCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
+        state.skip_countdown = 3
+        backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False)
+        assert state.skip_countdown == 0
+
+    def test_enable_unsupported_pipeline_raises(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "QwenImagePipeline"
+        backend = StepCacheBackend(DiffusionCacheConfig())
+
+        with pytest.raises(ValueError, match="step_cache backend does not support"):
+            backend.enable(mock_pipeline)
+
+    def test_selector_returns_backend(self):
+        backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]})
+        assert isinstance(backend, StepCacheBackend)
+        assert backend.config.velocity_sim_thresholds == [0.9]
+
+    def test_legacy_backend_aliases(self):
+        assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend)
+        assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend)

tests/e2e/online_serving/test_dreamzero_expansion.py

@@ -28,8 +28,6 @@ test_params = [
         server_args=[
             "--deploy-config",
             "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml",
-            "--enforce-eager",
-            "--disable-log-stats",
         ],
         env_dict={
             "ATTENTION_BACKEND": "torch",

tests/entrypoints/openai_api/test_openpi_connection.py

@@ -88,6 +88,47 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts():
     np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32))
 
 
+def test_unpack_accepts_openpi_client_ndarray_markers():
+    image = np.arange(6, dtype=np.uint8).reshape(2, 3)
+    payload = {
+        "observation/exterior_image_0_left": {
+            b"__ndarray__": True,
+            b"dtype": image.dtype.str,
+            b"shape": image.shape,
+            b"data": image.tobytes(),
+        }
+    }
+
+    decoded = openpi_connection._unpack_numpy(payload)
+
+    np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
+    assert decoded["observation/exterior_image_0_left"].dtype == np.uint8
+
+
+def test_unpack_openpi_client_packed_observation():
+    openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
+    image = np.zeros((180, 320, 3), dtype=np.uint8)
+    obs = {
+        "observation/exterior_image_0_left": image,
+        "observation/joint_position": np.zeros(7, dtype=np.float32),
+    }
+
+    decoded = openpi_connection._unpack(openpi_msgpack.packb(obs))
+
+    np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image)
+    np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"])
+
+
+def test_pack_uses_openpi_client_ndarray_markers():
+    openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy")
+    actions = np.asarray([[1.0, 2.0]], dtype=np.float32)
+
+    packed = openpi_connection._pack(actions)
+    decoded = openpi_msgpack.unpackb(packed)
+
+    np.testing.assert_allclose(decoded, actions)
+
+
 def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged():
     payload = {
         "metadata": {

vllm_omni/deploy/dreamzero.yaml

@@ -12,7 +12,7 @@ stages:
   - stage_id: 0
     devices: "0"
     max_num_seqs: 1
-    enforce_eager: true
+    enforce_eager: false
     model_class_name: DreamZeroPipeline
     model_config:
       default_robot_embodiment: roboarena
@@ -23,3 +23,8 @@ stages:
         needs_stereo_camera: false
         needs_session_id: true
         action_space: joint_position
+    cache_backend: step_cache
+    cache_config:
+      step_cache_dit_enabled: true
+      velocity_sim_thresholds: [0.95, 0.93]
+      velocity_skip_countdowns: [4, 2]

vllm_omni/deploy/dreamzero_tp1_cfg2.yaml

@@ -9,7 +9,7 @@ stages:
   - stage_id: 0
     devices: "0,1"
     max_num_seqs: 1
-    enforce_eager: true
+    enforce_eager: false
     model_class_name: DreamZeroPipeline
     parallel_config:
       tensor_parallel_size: 1
@@ -23,3 +23,8 @@ stages:
         needs_stereo_camera: false
         needs_session_id: true
         action_space: joint_position
+    cache_backend: step_cache
+    cache_config:
+      step_cache_dit_enabled: true
+      velocity_sim_thresholds: [0.95, 0.93]
+      velocity_skip_countdowns: [4, 2]

vllm_omni/diffusion/cache/base.py

@@ -13,6 +13,8 @@ Main cache backend implementations:
    the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit".
 2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from
    CacheBackend. Used via cache_backend="tea_cache".
+3. StepCacheBackend: Velocity cosine step-skipping for DreamZero. Inherits from
+   CacheBackend. Used via cache_backend="step_cache".
 
 All backends implement the same interface:
 - enable(pipeline): Enable cache on the pipeline

vllm_omni/diffusion/cache/selector.py

@@ -3,6 +3,7 @@ from typing import Any
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
+from vllm_omni.diffusion.cache.stepcache import StepCacheBackend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
@@ -14,14 +15,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
     - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface
     - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface
     - mag_cache: Uses MagCacheBackend with enable()/refresh() interface
+    - step_cache: Uses StepCacheBackend for DreamZero velocity step skip
 
     Args:
-        cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None).
+        cache_backend: Cache backend name ("cache_dit", "tea_cache",
+            "mag_cache", "step_cache", or None).
         cache_config: Cache configuration (dict or DiffusionCacheConfig instance).
 
     Returns:
-        Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend)
-        if cache_backend is set, None otherwise.
+        Cache backend instance if cache_backend is set, None otherwise.
 
     Raises:
         ValueError: If cache_backend is unsupported.
@@ -38,7 +40,10 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
         return TeaCacheBackend(cache_config)
     elif cache_backend == "mag_cache":
         return MagCacheBackend(cache_config)
+    elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"):
+        return StepCacheBackend(cache_config)
     else:
         raise ValueError(
-            f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'"
+            f"Unsupported cache backend: {cache_backend}. "
+            "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'"
         )

vllm_omni/diffusion/cache/stepcache/__init__.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+StepCache: velocity-based step skipping for DreamZero-style DiT denoising.
+
+Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the
+pipeline denoise loop when successive velocity predictions are highly aligned.
+
+Usage:
+    from vllm_omni import Omni
+
+    omni = Omni(
+        model="...",
+        cache_backend="step_cache",
+    )
+
+    # Alternative: environment variable
+    # export DIFFUSION_CACHE_BACKEND=step_cache
+"""
+
+from vllm_omni.diffusion.cache.stepcache.backend import (
+    CUSTOM_STEP_CACHE_DIT_ENABLERS,
+    CUSTOM_STEPCACHE_ENABLERS,
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_DIT_CONFIG_ATTR,
+    STEP_CACHE_DIT_STATE_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheBackend,
+    StepCacheDiTBackend,
+    StepCacheDiTConfig,
+    StepCacheDiTState,
+    enable_dreamzero_step_cache_dit,
+    enable_dreamzero_stepcache,
+    get_step_cache_dit_config,
+    get_step_cache_dit_state,
+    get_stepcache_config,
+    get_stepcache_state,
+    is_step_cache_dit_active,
+    is_stepcache_active,
+)
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
+
+__all__ = [
+    "CUSTOM_STEP_CACHE_DIT_ENABLERS",
+    "CUSTOM_STEPCACHE_ENABLERS",
+    "STEP_CACHE_CONFIG_ATTR",
+    "STEP_CACHE_DIT_CONFIG_ATTR",
+    "STEP_CACHE_DIT_STATE_ATTR",
+    "STEP_CACHE_STATE_ATTR",
+    "StepCacheBackend",
+    "StepCacheConfig",
+    "StepCacheDiTBackend",
+    "StepCacheDiTConfig",
+    "StepCacheDiTState",
+    "StepCacheState",
+    "enable_dreamzero_step_cache_dit",
+    "enable_dreamzero_stepcache",
+    "get_step_cache_dit_config",
+    "get_step_cache_dit_state",
+    "get_stepcache_config",
+    "get_stepcache_state",
+    "is_step_cache_dit_active",
+    "is_stepcache_active",
+]

vllm_omni/diffusion/cache/stepcache/backend.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+StepCache backend implementation.
+
+This module provides the stepcache backend that implements the CacheBackend
+interface for DreamZero-style velocity-based step skipping in the denoise loop.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.stepcache.config import (
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheConfig,
+)
+from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+logger = init_logger(__name__)
+
+# Backward-compatible aliases for older imports and pipeline attribute names.
+STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR
+STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR
+StepCacheDiTConfig = StepCacheConfig
+StepCacheDiTState = StepCacheState
+
+
+def is_stepcache_active(pipeline: Any) -> bool:
+    """Return True when stepcache is enabled on *pipeline*."""
+    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
+    return isinstance(cfg, StepCacheConfig) and cfg.enabled
+
+
+def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None:
+    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
+    return cfg if isinstance(cfg, StepCacheConfig) else None
+
+
+def get_stepcache_state(pipeline: Any) -> StepCacheState | None:
+    state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None)
+    return state if isinstance(state, StepCacheState) else None
+
+
+def is_step_cache_dit_active(pipeline: Any) -> bool:
+    return is_stepcache_active(pipeline)
+
+
+def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None:
+    return get_stepcache_config(pipeline)
+
+
+def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None:
+    return get_stepcache_state(pipeline)
+
+
+def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """Enable stepcache for DreamZeroPipeline."""
+    cache_config = StepCacheConfig.from_diffusion_cache_config(config)
+    setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config)
+    setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config))
+
+    logger.info(
+        "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)",
+        cache_config.enabled,
+        cache_config.sim_thresholds,
+        cache_config.skip_countdowns,
+    )
+
+
+enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache
+
+CUSTOM_STEPCACHE_ENABLERS = {
+    "DreamZeroPipeline": enable_dreamzero_stepcache,
+}
+
+CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS
+
+
+class StepCacheBackend(CacheBackend):
+    """
+    Velocity cosine step-skipping cache backend for DreamZero.
+
+    Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported
+    pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to
+    decide whether to run ``predict_noise``.
+
+    Example:
+        >>> from vllm_omni.diffusion.data import DiffusionCacheConfig
+        >>> backend = StepCacheBackend(DiffusionCacheConfig())
+        >>> backend.enable(pipeline)
+        >>> backend.refresh(pipeline, num_inference_steps=16)
+    """
+
+    def enable(self, pipeline: Any) -> None:
+        """Enable stepcache on a supported pipeline."""
+        pipeline_type = pipeline.__class__.__name__
+
+        if pipeline_type in CUSTOM_STEPCACHE_ENABLERS:
+            logger.info("Using custom stepcache enabler for model: %s", pipeline_type)
+            CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config)
+        else:
+            raise ValueError(
+                f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}"
+            )
+
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh stepcache state for a new generation."""
+        state = get_stepcache_state(pipeline)
+        if state is not None:
+            state.reset()
+            if verbose:
+                logger.debug(
+                    "stepcache state refreshed (num_inference_steps=%d)",
+                    num_inference_steps,
+                )
+        elif verbose and is_stepcache_active(pipeline):
+            logger.warning("stepcache config active but state not found on pipeline")
+
+
+StepCacheDiTBackend = StepCacheBackend

vllm_omni/diffusion/cache/stepcache/config.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Configuration for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+# Pipeline attributes set by StepCacheBackend.enable().
+STEP_CACHE_CONFIG_ATTR = "_stepcache_config"
+STEP_CACHE_STATE_ATTR = "_stepcache_state"
+
+
+@dataclass(frozen=True)
+class StepCacheConfig:
+    """Runtime config for velocity-based step skipping.
+
+    Skips entire DiT forwards when successive video velocity predictions are
+    highly aligned (cosine similarity above configured thresholds).
+
+    Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``.
+    """
+
+    enabled: bool = True
+    min_history_steps: int = 2
+    max_history: int = 2
+    sim_thresholds: tuple[float, ...] = (0.95, 0.93)
+    skip_countdowns: tuple[int, ...] = (4, 2)
+
+    @classmethod
+    def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig:
+        enabled = config.step_cache_dit_enabled
+
+        thresholds = tuple(config.velocity_sim_thresholds)
+        countdowns = tuple(config.velocity_skip_countdowns)
+        if len(thresholds) != len(countdowns):
+            raise ValueError(
+                "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; "
+                f"got {len(thresholds)} and {len(countdowns)}."
+            )
+
+        return cls(
+            enabled=enabled,
+            min_history_steps=config.step_cache_dit_min_history,
+            max_history=config.step_cache_dit_max_history,
+            sim_thresholds=thresholds,
+            skip_countdowns=countdowns,
+        )

vllm_omni/diffusion/cache/stepcache/state.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Mutable state for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+
+
+class StepCacheState:
+    """Per-generation mutable state (skip countdown)."""
+
+    def __init__(self, config: StepCacheConfig) -> None:
+        self.config = config
+        self.skip_countdown = 0
+
+    def reset(self) -> None:
+        self.skip_countdown = 0
+
+    def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool:
+        """Return True when the DiT forward should execute at this scheduler step."""
+        if not self.config.enabled:
+            return True
+
+        if len(prev_predictions) < self.config.min_history_steps:
+            return True
+
+        if self.skip_countdown > 1:
+            self.skip_countdown -= 1
+            return False
+        if self.skip_countdown == 1:
+            self.skip_countdown = 0
+            return True
+
+        v_last = prev_predictions[-1][0].flatten(1).float()
+        v_prev = prev_predictions[-2][0].flatten(1).float()
+        sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean()
+
+        for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns):
+            if sim > threshold:
+                self.skip_countdown = countdown
+                return False
+
+        return True
+
+    def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None:
+        while len(prev_predictions) > self.config.max_history:
+            prev_predictions.pop(0)

vllm_omni/diffusion/data.py

@@ -302,6 +302,8 @@ class DiffusionCacheConfig:
                     scm_steps_mask_policy, scm_steps_policy
         - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
                     mag_ratios, mag_calibrate
+        - step_cache: step_cache_dit_enabled, velocity_sim_thresholds,
+                          velocity_skip_countdowns, step_cache_dit_min_history
 
     Example:
         >>> # From dict (user-facing API) - partial config uses defaults for missing keys
@@ -375,6 +377,13 @@ class DiffusionCacheConfig:
     # "repeat" refreshes every force_refresh_step_hint steps.
     force_refresh_step_policy: str = "once"
 
+    # step_cache parameters [step_cache only] — DreamZero velocity schedule
+    step_cache_dit_enabled: bool = True
+    velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93])
+    velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2])
+    step_cache_dit_min_history: int = 2
+    step_cache_dit_max_history: int = 2
+
     # Additional parameters that may be passed but not explicitly defined
     _extra_params: dict[str, Any] = field(default_factory=dict, repr=False)
 

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -23,6 +23,10 @@ from huggingface_hub import hf_hub_download
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.cache.stepcache import (
+    get_stepcache_state,
+    is_stepcache_active,
+)
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
@@ -55,10 +59,6 @@ logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
 
 
-# ---------------------------------------------------------------------------
-# ---------------------------------------------------------------------------
-
-
 class VideoActionScheduler:
     """Wraps video + action schedulers into single .step() interface."""
 
@@ -205,6 +205,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         self._max_session_states = MAX_DREAMZERO_SESSIONS
         self.state = self._get_or_create_state("default")
 
+        # DiT step cache is configured by StepCacheBackend
+        # (cache_backend="step_cache") via pipeline._stepcache_config.
+
         # Keep runtime inference settings separate from the training-time config.
         self.num_inference_steps: int = model_config.get(
             "num_inference_steps",
@@ -306,6 +309,27 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
         """Call CausalWanModel, return (video_pred, action_pred)."""
+        video_pred, action_pred = self._predict_noise_eager(kwargs)
+
+        if is_stepcache_active(self):
+            video_pred = video_pred.clone()
+            if action_pred is not None:
+                action_pred = action_pred.clone()
+
+        if action_pred is None:
+            batch_size = kwargs["hidden_states"].shape[0]
+            action_pred = torch.empty(
+                batch_size,
+                0,
+                self.transformer.action_dim,
+                device=video_pred.device,
+                dtype=video_pred.dtype,
+            )
+        return (video_pred, action_pred)
+
+    def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Eager DiT forward; also handles KV-cache write-back on prefill."""
+        self._cudagraph_mark_step_begin()
         video_pred, action_pred, updated_kv_caches = self.transformer(
             x=kwargs["hidden_states"],
             timestep=kwargs["timestep_video"],
@@ -327,19 +351,118 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             for i, kv in enumerate(updated_kv_caches):
                 state.update_kv_cache(i, kv, is_negative=is_neg)
 
-        video_pred = video_pred.clone()
-        if action_pred is not None:
-            action_pred = action_pred.clone()
-        else:
-            batch_size = kwargs["hidden_states"].shape[0]
-            action_pred = torch.empty(
-                batch_size,
-                0,
-                self.transformer.action_dim,
-                device=video_pred.device,
-                dtype=video_pred.dtype,
-            )  # CFG-parallel-safe dummy action pred
-        return (video_pred, action_pred)
+        return video_pred, action_pred
+
+    # -----------------------------------------------------------------------
+    # torch.compile setup (paper D.2 extended: encoders + VAE + DiT blocks)
+    # -----------------------------------------------------------------------
+
+    def setup_compile(self) -> None:
+        """Compile DreamZero encoders, VAE, and per-block DiT for inference.
+
+        Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
+        on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
+        compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
+        DiT blocks use per-block ``fullgraph=True``.
+        """
+        if not torch.cuda.is_available():
+            logger.info("DreamZero setup_compile skipped: CUDA not available.")
+            return
+
+        from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import (
+            apply_wan_vae_feat_cache_tensor_patch,
+        )
+
+        apply_wan_vae_feat_cache_tensor_patch()
+
+        compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
+        # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
+
+        logger.info(
+            "DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)."
+        )
+
+        try:
+            self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro)
+        except Exception as exc:
+            logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc)
+
+        try:
+            self.image_encoder.model.visual.forward = torch.compile(
+                self.image_encoder.model.visual.forward,
+                **compile_ro,
+            )
+        except Exception as exc:
+            logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
+
+        try:
+            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+        except Exception as exc:
+            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+
+        compiled_blocks = 0
+        for block in self.transformer.blocks:
+            try:
+                block.forward = torch.compile(block.forward, **dit_compile)
+                compiled_blocks += 1
+            except Exception as exc:
+                logger.warning(
+                    "DreamZero: transformer block %d compile failed (%s); leaving remaining eager.",
+                    compiled_blocks,
+                    exc,
+                )
+                break
+        if compiled_blocks:
+            logger.info(
+                "DreamZero: compiled %d/%d transformer blocks.",
+                compiled_blocks,
+                len(self.transformer.blocks),
+            )
+
+        self.warmup_compile()
+
+    def warmup_compile(self) -> None:
+        """Warm up compiled text/image/VAE paths before timed inference."""
+        if not torch.cuda.is_available():
+            return
+
+        device = next(self.text_encoder.parameters()).device
+        with torch.inference_mode():
+            try:
+                text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device)
+                attention_mask = torch.ones_like(text_tokens)
+                self._encode_text(text_tokens, attention_mask)
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc)
+
+            try:
+                image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
+                self._encode_image(image, self.num_frames, 180, 320)
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
+
+            try:
+                latent_h, latent_w = 180 // 8, 320 // 8
+                dummy_latent = torch.zeros(
+                    1,
+                    16,
+                    self.num_frame_per_block,
+                    latent_h * 2,
+                    latent_w * 2,
+                    dtype=self.vae.dtype,
+                    device=device,
+                )
+                mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype)
+                inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype)
+                dummy_latent_denorm = dummy_latent / inv_std + mean
+                self._cudagraph_mark_step_begin()
+                self.vae.decode(dummy_latent_denorm, return_dict=False)
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc)
+
+        torch.accelerator.synchronize(device)
+        logger.info("DreamZero compile warmup finished (text / image / vae decode).")
 
     def combine_cfg_noise(
         self,
@@ -388,12 +511,20 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4)
         return videos.to(dtype=torch.bfloat16)
 
+    @staticmethod
+    def _cudagraph_mark_step_begin() -> None:
+        try:
+            torch.compiler.cudagraph_mark_step_begin()
+        except Exception:
+            pass
+
     # -----------------------------------------------------------------------
     # Text encoding
     # -----------------------------------------------------------------------
 
     def _encode_text(self, text_tokens: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
         """Encode text prompt via UMT5."""
+        self._cudagraph_mark_step_begin()
         seq_lens = attention_mask.gt(0).sum(dim=1).long()
         prompt_emb = self.text_encoder(
             text_tokens,
@@ -422,6 +553,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         batch_size = image.shape[0]
 
         with torch.amp.autocast(dtype=torch.bfloat16, device_type=device.type):
+            self._cudagraph_mark_step_begin()
             clip_context = self.image_encoder.encode_image(image)
 
             msk = torch.ones(batch_size, num_frames, height // 8, width // 8, device=device)
@@ -469,8 +601,23 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         inv_std = self.vae_latents_inv_std.to(device=vae_device, dtype=vae_dtype)
         latents = latents / inv_std + mean
         with torch.no_grad():
+            self._cudagraph_mark_step_begin()
             return self.vae.decode(latents, return_dict=False)[0]
 
+    def decode_accumulated_video_latents(self, session_id: str | None = None) -> torch.Tensor:
+        """Decode all AR-chunk latents accumulated for ``session_id``."""
+        state = self._get_or_create_state(session_id)
+        latents = state.get_concatenated_video_latents()
+        if latents is None:
+            session_key = str(session_id or "default")
+            raise RuntimeError(f"No accumulated video latents for session {session_key!r}.")
+        return self.decode_video_latents(latents)
+
+    def clear_accumulated_video_latents(self, session_id: str | None = None) -> None:
+        """Clear accumulated video latents for ``session_id`` without resetting KV state."""
+        state = self._get_or_create_state(session_id)
+        state.clear_video_latents()
+
     # -----------------------------------------------------------------------
     # KV cache prefill
     # -----------------------------------------------------------------------
@@ -635,6 +782,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         noisy_input = video_latents
         noisy_input_action = action_latents
+        # ---- step_cache (StepCacheBackend / dreamzero.git) ----
+        _cached_flow_pred: torch.Tensor | None = None
+        _cached_flow_pred_action: torch.Tensor | None = None
+        _prev_predictions: list[tuple[torch.Tensor]] = []
+        _step_cache = get_stepcache_state(self) if is_stepcache_active(self) else None
+
         for index in range(len(timesteps_video)):
             video_timestep = timesteps_video[index]
             action_timestep = timesteps_action[index]
@@ -663,46 +816,58 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             else:
                 y = state.ys[:, :, -self.num_frame_per_block :]
 
-            positive_kwargs = dict(
-                hidden_states=noisy_input.transpose(1, 2),
-                timestep_video=timestep,
-                encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
-                y=y,
-                clip_feature=state.clip_feas,
-                action=noisy_input_action,
-                timestep_action=timestep_action,
-                is_negative=False,
-                **common_kwargs,
-            )
-
-            # Negative (uncond) kwargs
-            if do_true_cfg and negative_prompt_embeds is not None:
-                negative_kwargs = dict(
+            run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions)
+            if run_dit:
+                positive_kwargs = dict(
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
-                    encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    encoder_hidden_states=prompt_embeds,
+                    kv_cache=state.get_kv_caches(False),
+                    crossattn_cache=state.get_crossattn_caches(False),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
                     timestep_action=timestep_action,
-                    is_negative=True,
+                    is_negative=False,
                     **common_kwargs,
                 )
-            else:
-                negative_kwargs = None
 
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                do_true_cfg=do_true_cfg,
-                true_cfg_scale=self.cfg_scale,
-                cfg_normalize=False,
-            )
-            flow_pred, flow_pred_action = noise_pred
+                if do_true_cfg and negative_prompt_embeds is not None:
+                    negative_kwargs = dict(
+                        hidden_states=noisy_input.transpose(1, 2),
+                        timestep_video=timestep,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        kv_cache=state.get_kv_caches(True),
+                        crossattn_cache=state.get_crossattn_caches(True),
+                        y=y,
+                        clip_feature=state.clip_feas,
+                        action=noisy_input_action,
+                        timestep_action=timestep_action,
+                        is_negative=True,
+                        **common_kwargs,
+                    )
+                else:
+                    negative_kwargs = None
+
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=self.cfg_scale,
+                    cfg_normalize=False,
+                )
+                flow_pred, flow_pred_action = noise_pred
+                _cached_flow_pred = flow_pred
+                _cached_flow_pred_action = flow_pred_action
+
+                _prev_predictions.append((flow_pred,))
+                if _step_cache is not None:
+                    _step_cache.trim_history(_prev_predictions)
+            else:
+                # Reuse previous prediction (DiT step-skipping cache).
+                assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip."
+                flow_pred = _cached_flow_pred
+                flow_pred_action = _cached_flow_pred_action
 
             latents = (noisy_input, noisy_input_action)
             t = (video_timestep, action_timestep)
@@ -811,9 +976,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
             state.reset()
-        # Auto-reset based on model state (before accumulation)
-        if state.should_reset(text_tokens, 0, self.transformer.local_attn_size):
-            state.reset()
+        else:
+            reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
+            if reset_reason == "session":
+                state.reset()
+            elif reset_reason == "inference":
+                state.reset_inference_state()
         state.language = text_tokens
 
         # Frame accumulation: stitched single frame → multi-frame video
@@ -959,6 +1127,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             video_out = torch.cat([image, video_out], dim=1)
         state.current_start_frame += self.num_frame_per_block
 
+        state.append_video_latents(video_out)
+
         # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -23,7 +23,8 @@ class DreamZeroState:
     Lifecycle:
         - Created once in DreamZeroPipeline.__init__()
         - Mutated every forward() call (frame append, KV cache grow)
-        - reset() on new session / language change / local_attn_size exceeded
+        - reset() on new session / language change
+        - reset_inference_state() on local_attn_size exceeded (KV only)
     """
 
     def __init__(self) -> None:
@@ -65,50 +66,106 @@ class DreamZeroState:
         self.call_count += 1
         return np.stack(frames, axis=0)  # (T, H, W, C)
 
+    def append_video_latents(self, video_out: torch.Tensor) -> None:
+        """Append one AR chunk of normalized video latents for later decode.
+
+        Args:
+            video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop
+                (before the ``transpose(1, 2)`` stored in ``DiffusionOutput``).
+        """
+        if video_out.dim() != 5:
+            raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}")
+        # Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W).
+        chunk = video_out.transpose(1, 2).detach().cpu()
+        self.video_latents_across_time.append(chunk)
+        logger.info(
+            "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d",
+            tuple(chunk.shape),
+            len(self.video_latents_across_time),
+            int(sum(c.shape[2] for c in self.video_latents_across_time)),
+        )
+
+    def get_concatenated_video_latents(self) -> torch.Tensor | None:
+        """Return all accumulated chunks concatenated along the time dimension."""
+        if not self.video_latents_across_time:
+            return None
+        if len(self.video_latents_across_time) == 1:
+            return self.video_latents_across_time[0]
+        return torch.cat(self.video_latents_across_time, dim=2)
+
+    def clear_video_latents(self) -> None:
+        """Drop accumulated video latents without resetting KV/frame state."""
+        self.video_latents_across_time = []
+
     # ------------------------------------------------------------------
     # Reset / should_reset
     # ------------------------------------------------------------------
 
-    def reset(self) -> None:
-        """Clear all state."""
-        self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK)
-        self.call_count: int = 0
+    def reset(self, *, clear_video_latents: bool = True) -> None:
+        """Clear session state.
+
+        Args:
+            clear_video_latents: When ``False``, keep ``video_latents_across_time``
+                so offline/online export can decode the full AR rollout.
+        """
+        saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time)
+        self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK)
+        self.call_count = 0
+        # Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU.
+        # Concatenated along T before decode (matches upstream video_across_time).
+        self.video_latents_across_time = saved_video_latents
 
         # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache: list[torch.Tensor] | None = None
-        self.kv_cache_neg: list[torch.Tensor] | None = None
-        self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None
-        self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None
-        self.current_start_frame: int = 0
+        self.kv_cache = None
+        self.kv_cache_neg = None
+        self.crossattn_cache = None
+        self.crossattn_cache_neg = None
+        self.current_start_frame = 0
 
-        self.clip_feas: torch.Tensor | None = None
-        self.ys: torch.Tensor | None = None
-        self.language: torch.Tensor | None = None
+        self.clip_feas = None
+        self.ys = None
+        if clear_video_latents:
+            self.language = None
 
-    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
-        """Determine if state should be reset before this forward()."""
+    def reset_inference_state(self) -> None:
+        """Reset KV/frame state after local attention rolls without dropping video latents."""
+        self.reset(clear_video_latents=False)
+
+    def reset_reason(
+        self,
+        text_tokens: torch.Tensor | None,
+        num_video_frames: int,
+        local_attn_size: int,
+    ) -> str | None:
+        """Return why state should reset before the next forward(), if any."""
         if self.language is None:
             logger.info("language is None, resetting")
-            return True
+            return "session"
 
         if text_tokens is not None and not torch.equal(self.language, text_tokens):
             logger.info("language changed, resetting")
-            return True
+            return "session"
 
         # NOTE: after accumulate_frames, num_video_frames is the accumulated T
         # (1 for first call, 4 for subsequent). Only reset on true single-frame
         # which happens when the stitched_buffer was cleared externally.
         if num_video_frames == 1 and self.call_count > 1:
             logger.info("single frame input after first call, resetting")
-            return True
+            return "session"
 
         if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
             logger.info(
-                "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size
+                "current_start_frame %d >= local_attn_size %d, resetting inference state",
+                self.current_start_frame,
+                local_attn_size,
             )
-            return True
+            return "inference"
 
-        return False
+        return None
+
+    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
+        """Determine if state should be reset before this forward()."""
+        return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
 
     # ------------------------------------------------------------------
     # KV cache management

vllm_omni/diffusion/models/dreamzero/video_export_worker.py

@@ -9,13 +9,30 @@ import torch
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    @staticmethod
+    def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
+        decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
+        decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
+        return (decoded * 255.0).round().to(torch.uint8).cpu()
+
     def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor:
         if self.model_runner is None or self.model_runner.pipeline is None:
             raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
 
         with torch.inference_mode():
             decoded = self.model_runner.pipeline.decode_video_latents(video_latents)
-            decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
-            decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
-            decoded = (decoded * 255.0).round().to(torch.uint8).cpu()
-        return decoded
+            return self._latents_to_uint8_frames(decoded)
+
+    def decode_accumulated_video_latents_to_uint8(self, session_id: str = "default") -> torch.Tensor:
+        """Decode all AR-chunk latents accumulated on the server for one session."""
+        if self.model_runner is None or self.model_runner.pipeline is None:
+            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
+
+        with torch.inference_mode():
+            decoded = self.model_runner.pipeline.decode_accumulated_video_latents(session_id)
+            return self._latents_to_uint8_frames(decoded)
+
+    def clear_accumulated_video_latents(self, session_id: str = "default") -> None:
+        if self.model_runner is None or self.model_runner.pipeline is None:
+            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
+        self.model_runner.pipeline.clear_accumulated_video_latents(session_id)

vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Make diffusers Wan VAE feat_cache compile-friendly.
+
+diffusers ``WanResample`` stores the string sentinel ``"Rep"`` in feat_cache slots.
+``tensor == "Rep"`` returns a non-Tensor and breaks ``torch.compile(fullgraph=True)``.
+We replace it with a scalar int8 tensor marker (value 0).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+
+
+def _is_feat_rep(entry: Any) -> bool:
+    if not isinstance(entry, torch.Tensor):
+        return False
+    return entry.numel() == 1 and entry.dtype == torch.int8
+
+
+def _make_feat_rep(x: torch.Tensor) -> torch.Tensor:
+    return torch.zeros(1, device=x.device, dtype=torch.int8)
+
+
+def _wan_resample_forward_tensorized(self, x, feat_cache=None, feat_idx=[0]):
+    """``WanResample.forward`` with tensor feat_cache sentinel (no ``"Rep"`` string)."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T
+
+    b, c, t, h, w = x.size()
+    if self.mode == "upsample3d":
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = _make_feat_rep(x)
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not _is_feat_rep(feat_cache[idx]):
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                        dim=2,
+                    )
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and _is_feat_rep(feat_cache[idx]):
+                    cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                if _is_feat_rep(feat_cache[idx]):
+                    x = self.time_conv(x)
+                else:
+                    x = self.time_conv(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+
+                x = x.reshape(b, 2, c, t, h, w)
+                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                x = x.reshape(b, c, t * 2, h, w)
+    t = x.shape[2]
+    x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+    x = self.resample(x)
+    x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
+
+    if self.mode == "downsample3d":
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = x.clone()
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -1:, :, :].clone()
+                x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+    return x
+
+
+def apply_wan_vae_feat_cache_tensor_patch() -> None:
+    """Monkeypatch diffusers WanResample once per process."""
+    global _PATCHED
+    if _PATCHED:
+        return
+
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample
+
+    WanResample.forward = _wan_resample_forward_tensorized
+    _PATCHED = True
+    logger.info("DreamZero: applied Wan VAE feat_cache tensor sentinel patch.")

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -168,8 +168,17 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:
             if current_omni_platform.supports_torch_inductor():
-                self._compile_transformer("transformer")
-                self._compile_transformer("transformer_2")
+                if hasattr(self.pipeline, "setup_compile"):
+                    try:
+                        self.pipeline.setup_compile()
+                    except Exception as exc:
+                        logger.warning(
+                            "Model runner: setup_compile() failed (%s); running without compile.",
+                            exc,
+                        )
+                else:
+                    self._compile_transformer("transformer")
+                    self._compile_transformer("transformer_2")
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",
@@ -306,8 +315,15 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 # num_inference_steps at all (i.e., just resets state and clears
                 # stale residuals).
                 num_inference_steps = req.sampling_params.num_inference_steps
-                if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None:
-                    num_inference_steps = 0
+                if num_inference_steps is None and self.od_config.cache_backend in (
+                    "tea_cache",
+                    "step_cache",
+                ):
+                    # TeaCache refresh ignores the value; step_cache refresh is a
+                    # no-op (per-chunk state resets in the denoise loop). DreamZero often
+                    # leaves sampling_params.num_inference_steps unset and uses the
+                    # pipeline default instead.
+                    num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0
 
                 if num_inference_steps is not None:
                     self.cache_backend.refresh(self.pipeline, num_inference_steps)

vllm_omni/engine/async_omni_engine.py

@@ -520,7 +520,7 @@ class AsyncOmniEngine:
             return
 
         mm_data = prompt.get("multi_modal_data")
-        if not isinstance(mm_data, Mapping) or not mm_data:
+        if not isinstance(mm_data, dict) or not mm_data:
             return
 
         from vllm.multimodal.hasher import MultiModalHasher
@@ -842,6 +842,14 @@ class AsyncOmniEngine:
                 "mag_max_skip_steps": 5,
                 "mag_retention_ratio": 0.1,
             }
+        if cache_backend in ("step_cache"):
+            return {
+                "step_cache_dit_enabled": True,
+                "velocity_sim_thresholds": [0.95, 0.93],
+                "velocity_skip_countdowns": [4, 2],
+                "step_cache_dit_min_history": 2,
+                "step_cache_dit_max_history": 2,
+            }
         return None
 
     @staticmethod

vllm_omni/entrypoints/cli/serve.py

@@ -512,7 +512,7 @@ class OmniServeCommand(CLISubcommand):
             "--cache-backend",
             type=str,
             default="none",
-            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'",
+            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'",
         )
         omni_config_group.add_argument(
             "--cache-config",

vllm_omni/entrypoints/openpi/connection.py

@@ -8,7 +8,13 @@ Protocol (compatible with OpenPI policy clients):
     Infer    -> client sends msgpack(obs), server sends msgpack(ndarray)
     Reset    -> client sends msgpack({endpoint:reset}), server sends msgpack(status)
 
-NumPy values use the msgpack-numpy marker mapping:
+NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the
+openpi-client format so official OpenPI clients can decode action responses:
+    ndarray -> {__ndarray__: true, dtype, shape, data}
+    scalar  -> {__npgeneric__: true, dtype, data}
+
+Inbound payloads accept both the openpi-client markers above and the legacy
+vLLM-native markers:
     ndarray -> {nd: true, type, kind, shape, data}
     scalar  -> {nd: false, type, kind, data}
 """
@@ -41,18 +47,16 @@ def _pack_numpy(obj: Any) -> Any:
         if not obj.flags.c_contiguous:
             obj = np.ascontiguousarray(obj)
         return {
-            b"nd": True,
+            b"__ndarray__": True,
             b"data": obj.tobytes(),
-            b"type": obj.dtype.str,
-            b"kind": obj.dtype.kind,
+            b"dtype": obj.dtype.str,
             b"shape": obj.shape,
         }
     if isinstance(obj, np.generic):
         return {
-            b"nd": False,
-            b"data": obj.tobytes(),
-            b"type": obj.dtype.str,
-            b"kind": obj.dtype.kind,
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
         }
     raise TypeError(f"Unsupported type: {type(obj)!r}")
 
@@ -67,27 +71,67 @@ def _decode_marker_text(value: Any) -> str:
     return str(value)
 
 
+def _is_truthy_marker(value: Any) -> bool:
+    return value is True or value == 1
+
+
+def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any:
+    nd = _mapping_get(obj, "nd", _MISSING)
+    dtype = _mapping_get(obj, "type", _MISSING)
+    kind = _mapping_get(obj, "kind", _MISSING)
+    data = _mapping_get(obj, "data", _MISSING)
+    if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING:
+        return _MISSING
+
+    dtype_obj = np.dtype(_decode_marker_text(dtype))
+    kind_text = _decode_marker_text(kind)
+    if dtype_obj.kind != kind_text:
+        raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
+    if dtype_obj.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {dtype_obj}")
+
+    array = np.frombuffer(data, dtype=dtype_obj).copy()
+    if nd:
+        shape = _mapping_get(obj, "shape", _MISSING)
+        if shape is _MISSING:
+            raise ValueError("NumPy ndarray marker is missing shape")
+        return array.reshape(tuple(shape))
+    return array[0]
+
+
+def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any:
+    if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)):
+        data = _mapping_get(obj, "data", _MISSING)
+        dtype = _mapping_get(obj, "dtype", _MISSING)
+        shape = _mapping_get(obj, "shape", _MISSING)
+        if data is _MISSING or dtype is _MISSING or shape is _MISSING:
+            raise ValueError("OpenPI ndarray marker is missing required fields")
+
+        dtype_obj = np.dtype(_decode_marker_text(dtype))
+        if dtype_obj.kind in ("V", "O", "c"):
+            raise ValueError(f"Unsupported dtype: {dtype_obj}")
+        array = np.frombuffer(data, dtype=dtype_obj).copy()
+        return array.reshape(tuple(shape))
+
+    if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)):
+        data = _mapping_get(obj, "data", _MISSING)
+        dtype = _mapping_get(obj, "dtype", _MISSING)
+        if data is _MISSING or dtype is _MISSING:
+            raise ValueError("OpenPI scalar marker is missing required fields")
+        dtype_obj = np.dtype(_decode_marker_text(dtype))
+        if dtype_obj.kind in ("V", "O", "c"):
+            raise ValueError(f"Unsupported dtype: {dtype_obj}")
+        return dtype_obj.type(data)
+
+    return _MISSING
+
+
 def _unpack_numpy(obj: Any) -> Any:
     if isinstance(obj, dict):
-        nd = _mapping_get(obj, "nd", _MISSING)
-        dtype = _mapping_get(obj, "type", _MISSING)
-        kind = _mapping_get(obj, "kind", _MISSING)
-        data = _mapping_get(obj, "data", _MISSING)
-        if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING:
-            dtype_obj = np.dtype(_decode_marker_text(dtype))
-            kind_text = _decode_marker_text(kind)
-            if dtype_obj.kind != kind_text:
-                raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
-            if dtype_obj.kind in ("V", "O", "c"):
-                raise ValueError(f"Unsupported dtype: {dtype_obj}")
-
-            array = np.frombuffer(data, dtype=dtype_obj).copy()
-            if nd:
-                shape = _mapping_get(obj, "shape", _MISSING)
-                if shape is _MISSING:
-                    raise ValueError("NumPy ndarray marker is missing shape")
-                return array.reshape(tuple(shape))
-            return array[0]
+        for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker):
+            decoded = decoder(obj)
+            if decoded is not _MISSING:
+                return decoded
         return {key: _unpack_numpy(value) for key, value in obj.items()}
     if isinstance(obj, list):
         return [_unpack_numpy(value) for value in obj]