Merge branch 'main' into fix/pr4709-review-followup

vllm_omni/diffusion/models/dreamzero/causal_wan_model.py

@@ -37,6 +37,20 @@ from vllm_omni.diffusion.models.dreamzero.action_encoder import (
     MultiEmbodimentActionEncoder,
 )
 
+# AR-Diffusion paged self-attention (in-tree experimental engine). Import at
+# module level so the isinstance check + custom-op call trace cleanly inside
+# the fullgraph-compiled DiT block (an import inside the traced region would
+# graph-break). The model still works without the engine: the payload type is
+# only ever constructed by the AR-Diffusion runner.
+try:
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+        ARDiffusionPagedLayerInputs,
+        paged_write_attn,
+    )
+except ImportError:  # pragma: no cover - experimental package always ships in-tree
+    ARDiffusionPagedLayerInputs = None
+    paged_write_attn = None
+
 # ── RoPE utilities ──────────────────────────────────────────────────
 
 
@@ -510,7 +524,7 @@ class CausalWanSelfAttention(nn.Module):
         freqs_action: torch.Tensor,
         freqs_state: torch.Tensor,
         action_register_length: int | None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -567,24 +581,44 @@ class CausalWanSelfAttention(nn.Module):
             action_v = v[:, -action_register_length:]
             v = v[:, :-action_register_length]
 
-        updated_k = kv_cache[0]
-        updated_v = kv_cache[1]
-        new_k = torch.cat([updated_k, roped_key], dim=1)
-        new_v = torch.cat([updated_v, v], dim=1)
-        new_k = new_k[:, -self.max_attention_size :]
-        new_v = new_v[:, -self.max_attention_size :]
-
-        if action_register_length is not None:
-            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
-            k_cat = torch.cat([new_k, roped_action_key], dim=1)
-            v_cat = torch.cat([new_v, action_v], dim=1)
+        if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs):
+            # Fused write+attend custom op: one opaque node in the compiled
+            # graph (slot writes + FlashAttention block-table kernel inside).
+            # Metadata tensors were prepared once per forward in _forward_blocks.
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_act, v_act = roped_action_key[0], action_v[0]
+            else:
+                q_cat = roped_query
+                k_act = v_act = None
+            x = paged_write_attn(
+                kv_cache,
+                q_cat[0],
+                roped_key[0],
+                v[0],
+                k_act,
+                v_act,
+                self.head_dim**-0.5,
+            ).unsqueeze(0)
         else:
-            q_cat = roped_query
-            k_cat = new_k
-            v_cat = new_v
+            updated_k = kv_cache[0]
+            updated_v = kv_cache[1]
+            new_k = torch.cat([updated_k, roped_key], dim=1)
+            new_v = torch.cat([updated_v, v], dim=1)
+            new_k = new_k[:, -self.max_attention_size :]
+            new_v = new_v[:, -self.max_attention_size :]
 
-        x = self.attn(q_cat, k_cat, v_cat)
-        updated_kv_cache = torch.stack([new_k, new_v], dim=0)
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_cat = torch.cat([new_k, roped_action_key], dim=1)
+                v_cat = torch.cat([new_v, action_v], dim=1)
+            else:
+                q_cat = roped_query
+                k_cat = new_k
+                v_cat = new_v
+
+            x = self.attn(q_cat, k_cat, v_cat)
+            updated_kv_cache = torch.stack([new_k, new_v], dim=0)
 
         x = x.flatten(2)
         x = self.o(x)
@@ -646,7 +680,7 @@ class CausalWanAttentionBlock(nn.Module):
         freqs_state: torch.Tensor,
         context: torch.Tensor,
         action_register_length: int | None = None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         crossattn_cache: dict | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
@@ -925,10 +959,10 @@ class CausalWanModel(nn.Module):
         action: torch.Tensor | None,
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[dict] | None,
+        kv_cache: list[Any],
         current_start_frame: int,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        crossattn_cache: list[dict] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
         B = x.shape[0]
         F_t = timestep.shape[1]
@@ -966,7 +1000,28 @@ class CausalWanModel(nn.Module):
             clip_embedding = self.img_emb(clip_feature)
             context = torch.cat([clip_embedding, context], dim=1)
 
-        updated_kv_caches: list[torch.Tensor] = []
+        # AR-Diffusion paged path: host-side prep ONCE per branch forward, outside
+        # the compiled blocks — allocate current video/action slots, build the
+        # padded block-table metadata all 40 layers share, and hand the compiled
+        # region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python
+        # context object. Lazy-allocation contract preserved: only the branch this
+        # CFG-parallel rank executes reaches its _forward_blocks.
+        if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False):
+            if B != 1:
+                raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
+            fctx = kv_cache[0].forward_ctx
+            if seq_len != fctx.seq_len:
+                raise RuntimeError(
+                    f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens"
+                )
+            fctx.prepare(
+                device=x.device,
+                action_len=int(action_register_length or 0),
+                query_len=int(x.shape[1]),
+            )
+            kv_cache = [c.to_layer_inputs() for c in kv_cache]
+
+        updated_kv_caches: list[torch.Tensor | None] = []
         for block_index, block in enumerate(self.blocks):
             x, updated_kv_cache = block(
                 x=x,
@@ -1000,9 +1055,9 @@ class CausalWanModel(nn.Module):
         timestep: torch.Tensor,
         context: torch.Tensor,
         seq_len: int,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[torch.Tensor],
+        kv_cache: list[Any],
         current_start_frame: int,
+        crossattn_cache: list[dict] | None = None,
         y: torch.Tensor | None = None,
         clip_feature: torch.Tensor | None = None,
         action: torch.Tensor | None = None,

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -3,7 +3,7 @@
 
 """DreamZero pipeline for vllm-omni.
 
-Entry point for DiffusionEngine.step() → pipeline.forward(req)
+Entry point for DiffusionEngine.step() -> pipeline.forward(req)
 """
 
 from __future__ import annotations
@@ -94,17 +94,102 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     Multi-output: predict_noise() returns (video_pred, action_pred).
     CFG: video gets standard CFG, action takes positive branch only.
-    State: DreamZeroState manages KV cache + frame buffer across forward() calls.
+
+    KV is managed by the AR-Diffusion engine: ``self._ar_diffusion_kv_state`` is set by the
+    runner before ``forward()`` and the pipeline routes all KV access (get / update /
+    commit / reset) through the pool-backed state. Purely duck-typed (no engine import).
     """
 
+    _ar_diffusion_kv_state = None  # set by the runner before each forward
+
+    def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False):
+        return self._ar_diffusion_kv_state.get_kv_caches(
+            is_negative,
+            seq_len=seq_len,
+            commit_current=update_kv_cache,
+        )
+
+    def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
+        # The engine owns all KV allocation: self-attn is allocated lazily from
+        # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross.
+        # Nothing is created model-side.
+        return
+
+    def _kv_commit(self, is_negative: bool):
+        self._ar_diffusion_kv_state.commit_paged_context(is_negative)
+
+    def _kv_get_cross(self, state, is_negative):
+        """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img)."""
+        return self._ar_diffusion_kv_state.get_cross_kv_caches(is_negative)
+
+    def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
+        """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool.
+
+        Caches the session-invariant cross-attn projections once, per half: the text
+        ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets
+        (prompt unchanged within a session — only session resets clear them), while
+        the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the
+        257 image tokens the forward splits off, cached model-side by #4154) are
+        re-projected on every window restart from the fresh CLIP features. Must run
+        after the image is encoded so ``clip_feature`` is available.
+        """
+        s = self._ar_diffusion_kv_state
+        need_text = not s._cross_text_populated.get(is_negative, False)
+        need_img = (
+            clip_feature is not None
+            and getattr(self.transformer, "model_type", "t2v") == "i2v"
+            and not s._cross_img_populated.get(is_negative, False)
+        )
+        if not need_text and not need_img:
+            return
+        projected = self.transformer.text_embedding(context) if need_text else None
+        img_ctx = self.transformer.img_emb(clip_feature) if need_img else None
+        for i, block in enumerate(self.transformer.blocks):
+            ca = block.cross_attn
+            n, d = ca.tp_num_heads, ca.head_dim
+            k = v = None
+            if projected is not None:
+                k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+                v = ca.v(projected).unflatten(2, (n, d))
+            k_img = v_img = None
+            if img_ctx is not None:
+                k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d))
+                v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
+            s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
+        if need_text:
+            s._cross_text_populated[is_negative] = True
+        if need_img:
+            s._cross_img_populated[is_negative] = True
+        logger.info(
+            "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s",
+            "neg" if is_negative else "pos",
+            len(self.transformer.blocks),
+            "kept" if projected is None else tuple(context.shape),
+            None if img_ctx is None else tuple(img_ctx.shape),
+        )
+
+    def _kv_reset(self, state, *, clear_video_latents: bool = True):
+        """Reset the engine's pooled session window plus the model's non-KV state.
+
+        DreamZero resets at the attention-window boundary; the engine pool drops the
+        same window so the next forward starts fresh. ``clear_video_latents=False``
+        keeps the accumulated video latents for export.
+
+        ``clear_video_latents=False`` also marks a window ("inference") reset: the
+        prompt is unchanged, so the pool keeps the text cross-attn K/V and only the
+        image half repopulates on the restart forward.
+        """
+        state.reset(clear_video_latents=clear_video_latents)
+        self._ar_diffusion_kv_state.reset(keep_cross_text=not clear_video_latents)
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
 
         DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID):
-          config.json                     — root config (action_head_cfg, architectures, etc.)
-          model-*.safetensors             — all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
-          experiment_cfg/metadata.json    — per-embodiment action normalization stats
-          vae/                            — symlink to Wan2.1 VAE (diffusers-compatible)
+          config.json                     -- root config (action_head_cfg, architectures, etc.)
+          model-*.safetensors             -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
+          experiment_cfg/metadata.json    -- per-embodiment action normalization stats
+          vae/                            -- symlink to Wan2.1 VAE (diffusers-compatible)
 
         Components are instantiated from config (not from_pretrained), then filled
         by load_weights() which reads root safetensors and remaps key prefixes.
@@ -119,6 +204,18 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         """
         super().__init__()
 
+        # DreamZero is engine-only: every KV access in forward() routes through
+        # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale
+        # or programmatic config that leaves engine_backend="default" would
+        # otherwise only crash mid-forward on the first KV access.
+        engine_backend = str(getattr(od_config, "engine_backend", "") or "")
+        if "ar_diffusion" not in engine_backend.lower().replace("-", "_"):
+            raise ValueError(
+                "DreamZeroPipeline requires the AR-Diffusion engine; set "
+                "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine "
+                f"in the deploy config (got engine_backend={engine_backend!r})."
+            )
+
         model_path = od_config.model
         model_config = od_config.model_config
         local_files_only = os.path.exists(model_path)
@@ -229,8 +326,11 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         self.max_action_dim: int = ah_config["max_action_dim"]
 
         self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT)
+        # The negative prompt is a model constant: encode it once, lazily, and
+        # reuse across every forward/session (UMT5 encode is deterministic).
+        self._negative_prompt_embeds_cache: torch.Tensor | None = None
 
-        # Embodiment name → numeric ID mapping (model knowledge)
+        # Embodiment name -> numeric ID mapping (model knowledge)
         self.embodiment_name_to_id: dict[str, int] = model_config.get(
             "embodiment_name_to_id",
             DEFAULT_EMBODIMENT_NAME_TO_ID,
@@ -345,11 +445,17 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state=kwargs.get("state_features"),
             embodiment_id=kwargs.get("embodiment_id"),
         )
-        if kwargs.get("update_kv_cache", False) and updated_kv_caches:
-            state = kwargs.get("dreamzero_state", self.state)
+        if kwargs.get("update_kv_cache", False):
             is_neg = kwargs.get("is_negative", False)
-            for i, kv in enumerate(updated_kv_caches):
-                state.update_kv_cache(i, kv, is_negative=is_neg)
+            logger.debug(
+                "AR-Diffusion pipeline predict_noise -> commit paged context: "
+                "is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                is_neg,
+                kwargs.get("seq_len"),
+                kwargs.get("current_start_frame"),
+                len(updated_kv_caches),
+            )
+            self._kv_commit(is_neg)
 
         return video_pred, action_pred
 
@@ -363,6 +469,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         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).
+        Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because
+        Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture.
         DiT blocks use per-block ``fullgraph=True``.
         """
         if not torch.cuda.is_available():
@@ -377,6 +485,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
         # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        # The AR-Diffusion paged self-attention is a registered custom op, so the
+        # block stays fullgraph even on that path.
         dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
 
         logger.info(
@@ -397,9 +507,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
 
         try:
-            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+            self.vae._encode = torch.compile(self.vae._encode, **compile_ro)
         except Exception as exc:
-            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+            logger.warning("DreamZero: vae._encode compile failed (%s); skipping.", exc)
 
         compiled_blocks = 0
         for block in self.transformer.blocks:
@@ -438,10 +548,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
             try:
                 image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
-                self._encode_image(image, self.num_frames, 180, 320)
+                self._encode_image(image, self.num_frames, 180, 320, state=self.state)
             except Exception as exc:
                 logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
 
+            try:
+                self.state.reset_vae_encoder_stream()
+                dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device)
+                self._vae_stream_seed(self.state, dummy_video[:, :, :1])
+                for _ in range(4):
+                    self._vae_stream_append_frame(self.state, dummy_video[:, :, :1])
+                self._vae_stream_get_observation_latents(
+                    self.state,
+                    self.num_frame_per_block,
+                    dtype=torch.bfloat16,
+                )
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae encode stream) skipped: %s", exc)
+
             try:
                 latent_h, latent_w = 180 // 8, 320 // 8
                 dummy_latent = torch.zeros(
@@ -500,7 +624,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     # -----------------------------------------------------------------------
 
     def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
-        """uint8 [B,T,H,W,C] → bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
+        """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
         videos = videos.permute(0, 4, 1, 2, 3)
         if videos.dtype == torch.uint8:
             videos = videos.float() / 255.0
@@ -547,6 +671,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         num_frames: int,
         height: int,
         width: int,
+        *,
+        state: DreamZeroState | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Encode first frame via CLIP + VAE.
         Returns: (clip_feas, ys, image_latent)
@@ -582,11 +708,147 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             new_image = y[:, :, 0:1]
             y = torch.concat([msk, y], dim=1)
 
+            if state is not None:
+                if not state.vae_stream_initialized:
+                    # Seed the AR streaming encoder after the compiled full encode above;
+                    # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
+                    self._vae_stream_seed(state, image_input[:, :, :1])
+                else:
+                    # Window ("inference") restart with a live stream: keep the real
+                    # frame history in the Wan feat_cache and append the restart
+                    # observation instead of reseeding from scratch.
+                    self._cudagraph_mark_step_begin()
+                    self._vae_stream_append_frame(state, image_input[:, :, :1])
+
         return clip_context, y, new_image
 
+    def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor:
+        if self.vae.config.patch_size is not None:
+            from diffusers.models.autoencoders.autoencoder_kl_wan import patchify
+
+            return patchify(videos, patch_size=self.vae.config.patch_size)
+        return videos
+
+    @staticmethod
+    def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]:
+        return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map]
+
+    def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]:
+        self.vae.clear_cache()
+        return self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_encode_encoder_chunk(
+        self,
+        chunk: torch.Tensor,
+        feat_map: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """Run one Wan encoder chunk while mutating ``feat_map`` in place.
+
+        Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates
+        conflict with ``reduce-overhead`` CUDAGraph capture.
+        """
+        self.vae._enc_feat_map = feat_map
+        self.vae._enc_conv_idx = [0]
+        out = self.vae.encoder(
+            chunk,
+            feat_cache=self.vae._enc_feat_map,
+            feat_idx=self.vae._enc_conv_idx,
+        )
+        return out, self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor:
+        enc = self.vae.quant_conv(encoder_out)
+        mu, _ = enc.chunk(2, dim=1)
+        mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
+        inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype)
+        return (mu - mean) * inv_std
+
+    def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None:
+        """Seed incremental VAE encode with the first observation frame."""
+        state.reset_vae_encoder_stream()
+        feat_map = self._vae_init_enc_feat_map()
+        chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype))
+        self._cudagraph_mark_step_begin()
+        encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map)
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = None
+        state.vae_stream_initialized = True
+
+    def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None:
+        """Append one pixel frame and encode a 4-frame body chunk when ready."""
+        if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream is not initialized.")
+
+        frame = new_frame.to(dtype=self.vae.dtype)
+        if state.vae_pending_body_frames is None:
+            state.vae_pending_body_frames = frame
+        else:
+            state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2)
+
+        while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4:
+            body = state.vae_pending_body_frames[:, :, :4]
+            if state.vae_pending_body_frames.shape[2] > 4:
+                state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:]
+            else:
+                state.vae_pending_body_frames = None
+            chunk = self._vae_patchify(body)
+            encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map)
+            state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2)
+            state.vae_enc_feat_map = feat_map
+
+    def _vae_stream_get_observation_latents(
+        self,
+        state: DreamZeroState,
+        num_latent_frames: int,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream has no accumulated encoder output.")
+        latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype)
+        if latents.shape[2] >= num_latent_frames:
+            return latents[:, :, -num_latent_frames:]
+        pad_count = num_latent_frames - latents.shape[2]
+        pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1)
+        return torch.cat([pad, latents], dim=2)
+
+    def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames_raw, _, _ = videos.shape
+        if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
+            return videos
+        if num_frames_raw // 4 != self.num_frame_per_block:
+            repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
+            videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
+            first_frame = videos[:, :, 0:1]
+            return torch.cat([first_frame, videos], dim=2)
+        first_frame = videos[:, :, 0:1]
+        return torch.cat([first_frame, videos], dim=2)
+
+    def _encode_observation_latents(
+        self,
+        state: DreamZeroState,
+        videos: torch.Tensor,
+        *,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode current robot observation into normalized VAE latents."""
+        if state.vae_stream_initialized:
+            self._cudagraph_mark_step_begin()
+            self._vae_stream_append_frame(state, videos[:, :, -1:])
+            return self._vae_stream_get_observation_latents(
+                state,
+                self.num_frame_per_block,
+                dtype=latent_dtype,
+            )
+
+        videos = self._preprocess_vae_observation_window(videos)
+        return self._encode_vae_latents(videos).to(dtype=latent_dtype)
+
     def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor:
         """Encode videos into normalized VAE latents."""
         input_dtype = videos.dtype
+        self._cudagraph_mark_step_begin()
         hidden = self.vae._encode(videos.to(dtype=self.vae.dtype))
         mu, _ = hidden.chunk(2, dim=1)
         mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
@@ -636,7 +898,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     ) -> None:
         """Prefill KV cache with first frame and/or current observation.
 
-        Uses predict_noise_maybe_with_cfg() for CFG parallel — same path as
+        Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as
         the denoise loop. The mixin handles rank dispatch automatically.
         KV cache update happens as a side effect inside predict_noise().
         """
@@ -647,7 +909,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         head_dim = self.transformer.dim // self.transformer.num_heads
 
         if state.current_start_frame == 0:
-            state.create_kv_caches(
+            self._kv_create(
+                state,
                 batch_size,
                 dtype,
                 device,
@@ -672,16 +935,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -722,16 +985,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -764,8 +1027,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         For each timestep:
           1. Build positive_kwargs / negative_kwargs
-          2. predict_noise_maybe_with_cfg()    → (video_pred, action_pred)
-          3. scheduler_step_maybe_with_cfg()   → VideoActionScheduler
+          2. predict_noise_maybe_with_cfg()    -> (video_pred, action_pred)
+          3. scheduler_step_maybe_with_cfg()   -> VideoActionScheduler
           4. _synchronize_cfg_parallel_step_output()
         """
         seq_len = kwargs["seq_len"]
@@ -824,8 +1087,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                     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),
+                    kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False),
+                    crossattn_cache=self._kv_get_cross(state, False),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
@@ -839,8 +1102,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                         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),
+                        kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False),
+                        crossattn_cache=self._kv_get_cross(state, True),
                         y=y,
                         clip_feature=state.clip_feas,
                         action=noisy_input_action,
@@ -937,7 +1200,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             device=device,
         )
 
-        # State: raw from transform → pad to (B, state_horizon=1, max_state_dim)
+        # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim)
         raw_state = unified_obs["state"]
         state_for_postprocess = None
         if raw_state is not None:
@@ -977,38 +1240,66 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # Explicit reset from OpenPI serving is carried by `extra_args["reset"]`
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
-            state.reset()
+            self._kv_reset(state)
         else:
+            # Auto-reset based on model state (before accumulation). Both a "session"
+            # reset and an "inference" (window-boundary) reset clear the model-local
+            # KV window, so route both through _kv_reset to drop the AR-Diffusion pool window
+            # too and stay parity-exact; "inference" keeps the accumulated video
+            # latents (reset_inference_state == reset(clear_video_latents=False)).
             reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
             if reset_reason == "session":
-                state.reset()
+                self._kv_reset(state)
             elif reset_reason == "inference":
-                state.reset_inference_state()
+                self._kv_reset(state, clear_video_latents=False)
         state.language = text_tokens
 
-        # Frame accumulation: stitched single frame → multi-frame video
+        # Frame accumulation: stitched single frame -> multi-frame video
         video_frames = state.accumulate_frames(stitched)  # (T, H, W, C)
         videos = torch.from_numpy(video_frames).unsqueeze(0).to(device)  # (B=1, T, H, W, C)
 
-        videos = self._preprocess_video(videos)  # → [B,C,T,H,W] bf16
+        videos = self._preprocess_video(videos)  # -> [B,C,T,H,W] bf16
         _, _, num_frames_raw, height, width = videos.shape
 
-        prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs
+        # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark
+        # synchronizes CUDA, so leave it off for timed benchmark runs.
+        _pt = None
+        if os.environ.get("DZ_PHASE_TIMING"):
+            import time as _time
+
+            torch.accelerator.synchronize()
+            _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []}
+
+        def _pt_mark(name: str) -> None:
+            if _pt is not None:
+                torch.accelerator.synchronize()
+                _pt["marks"].append((name, _pt["time"].perf_counter()))
+
+        # Prompt embeds are constant within a session (a prompt change triggers a
+        # "session" reset above, which clears this cache alongside state.language).
+        if state.prompt_embeds is None:
+            state.prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        prompt_embeds = state.prompt_embeds
         # Negative prompt for CFG uncond branch (model constant)
         negative_prompt_embeds = None
         if self.cfg_scale > 1.0:
-            neg_inputs = self.tokenizer(
-                self.negative_prompt,
-                max_length=512,
-                padding="max_length",
-                truncation=True,
-                return_tensors="pt",
-                add_special_tokens=True,
-            )
-            negative_prompt_embeds = self._encode_text(
-                neg_inputs["input_ids"].to(device),
-                neg_inputs["attention_mask"].to(device),
-            )
+            if self._negative_prompt_embeds_cache is None:
+                neg_inputs = self.tokenizer(
+                    self.negative_prompt,
+                    max_length=512,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="pt",
+                    add_special_tokens=True,
+                )
+                self._negative_prompt_embeds_cache = self._encode_text(
+                    neg_inputs["input_ids"].to(device),
+                    neg_inputs["attention_mask"].to(device),
+                )
+            negative_prompt_embeds = self._negative_prompt_embeds_cache
+
+        _pt_mark("text_encode")
 
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
@@ -1022,27 +1313,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 self.num_frames,
                 height,
                 width,
+                state=state,
             )
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
 
-        if state.current_start_frame != 0:
-            # Subsequent calls: encode current observation via VAE
-            if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
-                pass
-            elif num_frames_raw // 4 != self.num_frame_per_block:
-                repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
-                videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
-            else:
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
+            # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V
+            # into the pool now that the image is encoded (clip_feas available).
+            # Runs on the first forward of a session and after each window-boundary
+            # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry.
+            self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False)
+            if negative_prompt_embeds is not None:
+                self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
 
+        if state.current_start_frame != 0:
             latent_dtype = videos.dtype
             with torch.no_grad():
-                image = self._encode_vae_latents(videos)
-            image = image.to(dtype=latent_dtype)
+                image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype)
+        _pt_mark("obs_vae_encode")
 
         batch_size = image.shape[0]
         generator = torch.Generator(device=device).manual_seed(self.seed)
@@ -1083,6 +1371,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             do_true_cfg,
             state,
         )
+        _pt_mark("prefill_kv")
 
         sample_scheduler = copy.deepcopy(self.scheduler)
         sample_scheduler_action = copy.deepcopy(self.scheduler)
@@ -1124,6 +1413,19 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state_features=state_features,
             embodiment_id=embodiment_id,
         )
+        if _pt is not None:
+            _pt_mark("diffuse")
+            prev_t = _pt["t0"]
+            parts = []
+            for name, t in _pt["marks"]:
+                parts.append(f"{name}={1000 * (t - prev_t):.1f}ms")
+                prev_t = t
+            logger.info(
+                "DZ_PHASE_TIMING csf=%s total=%.1fms %s",
+                state.current_start_frame,
+                1000 * (prev_t - _pt["t0"]),
+                " ".join(parts),
+            )
 
         if state.current_start_frame == 1:
             video_out = torch.cat([image, video_out], dim=1)
@@ -1134,7 +1436,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
-        # Relative → absolute: only for relative_action_keys (joint_position only)
+        # Relative -> absolute: only for relative_action_keys (joint_position only)
         # gripper_position is NOT relative, so don't add state back to it
         if self.relative_action and state_for_postprocess is not None:
             n_relative = self.relative_action_dim  # 7 for DROID (joint only)
@@ -1146,7 +1448,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 action_out[..., :n_relative] + last_state.unsqueeze(1)  # broadcast over horizon
             )
 
-        # Squeeze batch dim for output: (B, horizon, dim) → (horizon, dim)
+        # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim)
         actions_np = action_out.squeeze(0).float().cpu().numpy()  # (horizon, max_action_dim)
         actions_np = transform.transform_action_output(actions_np)
 
@@ -1246,7 +1548,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         q01 = stats["q01"].to(device=action.device, dtype=action.dtype)
         q99 = stats["q99"].to(device=action.device, dtype=action.dtype)
         # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim)
-        # q01/q99 shape: (actual_action_dim,) — only denorm actual dims
+        # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims
         actual_dim = q01.shape[0]
         action_real = action.clone()
         action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -115,17 +115,24 @@ class DreamZeroState:
         # 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 = None
-        self.kv_cache_neg = None
-        self.crossattn_cache = None
-        self.crossattn_cache_neg = None
         self.current_start_frame = 0
 
         self.clip_feas = None
         self.ys = None
         if clear_video_latents:
+            # Session reset: drop the prompt-embed cache and VAE encoder stream.
             self.language = None
+            self.prompt_embeds = None
+            self.reset_vae_encoder_stream()
+        # Window ("inference") resets keep both: the prompt is unchanged, and the
+        # Wan feat_cache history is independent of the DiT attention window.
+
+    def reset_vae_encoder_stream(self) -> None:
+        """Clear incremental Wan VAE encoder state used across AR steps."""
+        self.vae_stream_initialized = False
+        self.vae_enc_feat_map: list[torch.Tensor | None] | None = None
+        self.vae_encoder_out: torch.Tensor | None = None
+        self.vae_pending_body_frames: torch.Tensor | None = None
 
     def reset_inference_state(self) -> None:
         """Reset KV/frame state after local attention rolls without dropping video latents."""
@@ -166,57 +173,3 @@ class DreamZeroState:
     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
-    # ------------------------------------------------------------------
-
-    def create_kv_caches(
-        self,
-        batch_size: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        num_layers: int,
-        num_heads: int,
-        head_dim: int,
-    ) -> None:
-        """Initialize empty KV caches and cross-attention caches."""
-        self.kv_cache = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-        self.kv_cache_neg = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-
-        self.crossattn_cache = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-        self.crossattn_cache_neg = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-
-    def update_kv_cache(
-        self,
-        layer_index: int,
-        updated_kv: torch.Tensor,
-        is_negative: bool = False,
-    ) -> None:
-        """Update a single layer's KV cache after prefill."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
-        cache[layer_index] = updated_kv.clone()
-
-    def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
-        """Get KV caches for the specified branch."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized.")
-        return cache
-
-    def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
-        """Get cross-attention caches for the specified branch."""
-        cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
-        if cache is None:
-            raise RuntimeError("Cross-attn caches not initialized.")
-        return cache

vllm_omni/diffusion/models/dreamzero/video_export_worker.py

@@ -9,6 +9,35 @@ import torch
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    def gpu_mem_stats(self) -> dict:
+        """Peak GPU memory (GiB) for this worker's device, for profiling.
+
+        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
+        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
+        tensor high-water mark. Both are monotonic from process start, so the
+        end-of-run values capture the whole rollout's peak.
+        """
+        dev = torch.accelerator.current_device_index()
+        return {
+            "device": int(dev),
+            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
+        }
+
+    def ar_diffusion_perf_stats(self, reset: bool = True) -> dict:
+        """Per-request server-side forward E2E times (s) recorded by the AR-Diffusion runner.
+
+        The worker-side compute time per ``execute_model`` — the analog of the
+        upstream server's per-request E2E — which excludes the engine<->worker IPC
+        that the client's ``omni.generate`` wall time includes. Returns ``[]`` when
+        the runner is the base (KV-disabled) runner that does not record timings.
+        """
+        runner = self.model_runner
+        times = list(getattr(runner, "_perf_e2e_times", []) or [])
+        if reset and hasattr(runner, "_perf_e2e_times"):
+            runner._perf_e2e_times = []
+        return {"server_e2e_s": times}
+
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()