[AR-Diffusion] Add session-owned streaming VAE decode (#6533)

Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py

@@ -0,0 +1,349 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Stream LingBot-World 2.0 pixels from one stepwise request.
+
+One ``AsyncOmni.generate()`` drives the whole rollout: each AR block leaves
+``post_decode()`` as one video chunk, decoded through that session's own
+temporal VAE cache rather than as an isolated clip. That is what this script
+shows end to end -- chunk *N + 1* continues chunk *N*, so the streamed timeline
+carries the same frame count as an offline whole-clip decode of the same
+latents instead of losing each block's opening frame to a restarted decoder.
+Only the session's first latent frame expands to a single raw frame; on the
+default checkpoint, whose block is three latent frames at a temporal factor of
+four, that is 9 frames for the opening chunk and 12 for every later one.
+
+The decode lives inside the pipeline, keyed by ``request_id``, so nothing here
+owns a VAE: this file is a client of the served path, and the same chunks reach
+a WebSocket client of ``/v1/realtime/video`` unchanged.
+
+Examples::
+
+    python examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py \
+        --image scene.png --prompt "a lit hallway" \
+        --action-script actions.json --output-dir /tmp/lingbot-stream
+
+``--action-script`` is a JSON list with one camera action list per AR block,
+for example ``[[["w"], ["w"], ["w"]], [["a"], [], []]]`` for a checkpoint that
+generates three latent frames per block. The frames per block, the temporal
+compression and the number of blocks that fit the condition horizon are read
+off the checkpoint rather than assumed, so the script can be repointed with
+``--model``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import math
+import os
+import time
+from collections.abc import Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    # Only for type hints -- the real imports stay deferred inside run() so
+    # --help and the offline validation helpers don't require a CUDA-enabled
+    # vLLM install.
+    import numpy as np
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+# Mirrors _MAX_RAW_FRAMES in vllm_omni/diffusion/models/lingbot_world/pipeline.py:
+# the image-condition horizon is the pipeline's constant, not the checkpoint's.
+_MAX_RAW_FRAMES = 117
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Stream decoded LingBot-World 2.0 video chunks from a single stepwise request.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
+    parser.add_argument("--image", required=True, help="Initial RGB image.")
+    parser.add_argument("--prompt", required=True, help="Scene prompt for the rollout.")
+    parser.add_argument(
+        "--action-script",
+        required=True,
+        help=(
+            "JSON file holding one camera action list per AR block. Both the per-block frame count "
+            "and the number of blocks that fit are read off the checkpoint, not assumed."
+        ),
+    )
+    parser.add_argument("--output-dir", required=True, help="Directory for decoded frames and metadata.")
+    parser.add_argument("--request-id", default="lingbot-world-stream", help="Request id; also the session id.")
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=832)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--gpu-memory-fraction", type=float, default=0.6)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--enforce-eager", action="store_true")
+    return parser.parse_args(argv)
+
+
+def _load_action_script(path: Path) -> list[list[list[str]]]:
+    """Parse the per-block camera action script, without assuming its geometry.
+
+    The shape is the one ``sampling_params.extra_args["camera_action_script"]``
+    takes: one entry per AR block, each entry one per-frame action list. How
+    many frames a block holds is a property of the checkpoint, so it is checked
+    in :func:`_validate_against_checkpoint` once the configs have been read.
+    """
+    try:
+        value = json.loads(path.read_text())
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"--action-script is not valid JSON: {exc.msg}.") from None
+    if not isinstance(value, list) or not value:
+        raise ValueError("--action-script must be a non-empty JSON list of per-block action lists.")
+    script: list[list[list[str]]] = []
+    for block_index, block in enumerate(value):
+        if not isinstance(block, list) or not block:
+            raise ValueError(f"--action-script block {block_index} must be a non-empty list of frame action lists.")
+        frames: list[list[str]] = []
+        for frame in block:
+            if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
+                raise ValueError(f"--action-script block {block_index} frames must contain only action strings.")
+            frames.append(list(frame))
+        script.append(frames)
+    return script
+
+
+def _model_config_json(model: str, relative_path: str) -> dict[str, Any]:
+    """Read one config file out of a local checkpoint or a Hub repo."""
+    if os.path.isdir(model):
+        return json.loads((Path(model) / relative_path).read_text())
+    from huggingface_hub import hf_hub_download
+
+    return json.loads(Path(hf_hub_download(repo_id=model, filename=relative_path)).read_text())
+
+
+def _checkpoint_geometry(model: str) -> tuple[int, int]:
+    """Read ``(frames_per_block, temporal_compression)`` off the checkpoint.
+
+    The pipeline takes both from the model --
+    ``transformer.config.num_frames_per_block`` and
+    ``vae.config.scale_factor_temporal`` -- so hardcoding 3 and 4 here would
+    silently build a wrong action script and an inconsistent ``num_frames`` the
+    moment ``--model`` points at a checkpoint with different geometry.
+    """
+    transformer_config = _model_config_json(model, "transformer/config.json")
+    vae_config = _model_config_json(model, "vae/config.json")
+    if "num_frames_per_block" not in transformer_config:
+        raise ValueError(
+            f"{model}'s transformer/config.json declares no num_frames_per_block, "
+            "so this checkpoint does not generate AR blocks and cannot be streamed this way."
+        )
+    frames_per_block = int(transformer_config["num_frames_per_block"])
+    temporal_compression = int(vae_config.get("scale_factor_temporal", 4))
+    if frames_per_block <= 0 or temporal_compression <= 0:
+        raise ValueError(
+            "the checkpoint reports a non-positive block geometry "
+            f"(num_frames_per_block={frames_per_block}, scale_factor_temporal={temporal_compression})."
+        )
+    return frames_per_block, temporal_compression
+
+
+def _validate_against_checkpoint(
+    script: list[list[list[str]]], *, frames_per_block: int, temporal_compression: int
+) -> int:
+    """Check the script against the checkpoint and return the frame count to request."""
+    for block_index, block in enumerate(script):
+        if len(block) != frames_per_block:
+            raise ValueError(
+                f"--action-script block {block_index} holds {len(block)} frame action lists, but this "
+                f"checkpoint generates {frames_per_block} latent frames per AR block."
+            )
+    max_chunks = ((_MAX_RAW_FRAMES - 1) // temporal_compression + 1) // frames_per_block
+    if len(script) > max_chunks:
+        raise ValueError(
+            f"--action-script holds {len(script)} blocks, but this checkpoint's "
+            f"{_MAX_RAW_FRAMES}-frame image-condition horizon fits at most {max_chunks}."
+        )
+    return (len(script) * frames_per_block - 1) * temporal_compression + 1
+
+
+def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
+    image = Path(args.image).expanduser().resolve()
+    action_script = Path(args.action_script).expanduser().resolve()
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    if not image.is_file():
+        raise ValueError("--image must point to an existing file.")
+    if not action_script.is_file():
+        raise ValueError("--action-script must point to an existing JSON file.")
+    if not args.prompt.strip():
+        raise ValueError("--prompt must contain non-whitespace text.")
+    if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
+        raise ValueError("--height and --width must be positive multiples of 16.")
+    if args.tensor_parallel_size <= 0:
+        raise ValueError("--tensor-parallel-size must be positive.")
+    if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
+        raise ValueError("--gpu-memory-fraction must be in (0, 1].")
+    if not args.request_id.strip():
+        raise ValueError("--request-id must contain non-whitespace text.")
+    return image, action_script, output_dir
+
+
+def _as_frames(value: Any) -> np.ndarray:
+    """Normalize one streamed chunk to ``[T, H, W, 3]`` uint8.
+
+    ``output_type="np"`` reaches the client as a float array in ``[0, 1]``,
+    with or without the leading batch axis depending on how the formatter
+    unwrapped it; a nested list is the ``"pil"`` shape. Anything else means
+    the chunk is not pixels -- most likely the request still asked for
+    latents -- and is worth failing on rather than saving.
+    """
+    import numpy as np
+    import PIL.Image
+    import torch
+
+    if isinstance(value, list):
+        if len(value) == 1 and isinstance(value[0], list):
+            value = value[0]
+        if value and all(isinstance(frame, PIL.Image.Image) for frame in value):
+            return np.stack([np.asarray(frame.convert("RGB")) for frame in value])
+    if isinstance(value, torch.Tensor):
+        value = value.detach().float().cpu().numpy()
+    if not isinstance(value, np.ndarray):
+        raise RuntimeError(f"Expected decoded pixels from each AR block, got {type(value).__name__}.")
+    if value.ndim == 5 and value.shape[0] == 1:
+        value = value[0]
+    if value.ndim != 4 or value.shape[-1] != 3:
+        raise RuntimeError(f"Expected one [T, H, W, 3] pixel chunk per AR block, got shape {tuple(value.shape)}.")
+    if value.dtype != np.uint8:
+        value = (np.clip(value, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+    return value
+
+
+def _chunk_metadata(output: Any) -> dict[str, Any]:
+    multimodal = getattr(output, "multimodal_output", None) or {}
+    metadata = multimodal.get("metadata") if isinstance(multimodal, dict) else None
+    ar_diffusion = metadata.get("ar_diffusion") if isinstance(metadata, dict) else None
+    if not isinstance(ar_diffusion, dict):
+        raise RuntimeError("Streamed chunk is missing its ar_diffusion metadata envelope.")
+    return dict(ar_diffusion)
+
+
+async def run(argv: Sequence[str] | None = None) -> Path:
+    args = parse_args(argv)
+    image, action_script_path, output_dir = _validate_args(args)
+    camera_action_script = _load_action_script(action_script_path)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Imports stay below pure input validation so --help and helper tests do not
+    # require a CUDA-enabled vLLM installation.
+    import numpy as np
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    frames_per_block, temporal_compression = _checkpoint_geometry(args.model)
+    num_chunks = len(camera_action_script)
+    num_frames = _validate_against_checkpoint(
+        camera_action_script,
+        frames_per_block=frames_per_block,
+        temporal_compression=temporal_compression,
+    )
+
+    engine = AsyncOmni(
+        model=args.model,
+        engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
+        enforce_eager=args.enforce_eager,
+        tensor_parallel_size=args.tensor_parallel_size,
+        max_num_seqs=1,
+        # One request, many chunks: the stepwise contract from #6844. Streaming
+        # output is what turns each post_decode() into a delivered chunk.
+        step_execution=True,
+        diffusion_streaming_output=True,
+        model_config={
+            "ar_diffusion_height": args.height,
+            "ar_diffusion_width": args.width,
+            "ar_diffusion_kv_config": {
+                "gpu_memory_fraction": args.gpu_memory_fraction,
+                "warmup_cudagraph": True,
+            },
+        },
+    )
+    sampling = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_frames=num_frames,
+        num_inference_steps=4,
+        max_sequence_length=512,
+        seed=args.seed,
+        # Pixels, not latents: this is what routes the chunk through the
+        # session's streaming VAE decode instead of handing back a tensor.
+        output_type="np",
+        extra_args={"flow_shift": 5.0, "camera_action_script": camera_action_script},
+    )
+    prompt = {"prompt": args.prompt, "multi_modal_data": {"image": str(image)}}
+
+    measurements: list[dict[str, Any]] = []
+    t_start = time.perf_counter()
+    first_frame_at: float | None = None
+    chunk_started = t_start
+    try:
+        async for output in engine.generate(prompt, sampling, request_id=args.request_id):
+            images = getattr(output, "images", None)
+            if not images:
+                continue
+            if len(images) != 1:
+                raise RuntimeError("Expected exactly one decoded chunk per AR block.")
+            frames = _as_frames(images[0])
+            now = time.perf_counter()
+            if first_frame_at is None:
+                first_frame_at = now - t_start
+            chunk_index = len(measurements)
+            metadata = _chunk_metadata(output)
+
+            np.save(output_dir / f"chunk_{chunk_index:03d}.npy", frames)
+            (output_dir / f"chunk_{chunk_index:03d}.json").write_text(
+                json.dumps(metadata, indent=2, sort_keys=True) + "\n"
+            )
+            measurements.append(
+                {
+                    "chunk_index": chunk_index,
+                    "chunk_latency_seconds": now - chunk_started,
+                    "frames": int(frames.shape[0]),
+                    "frame_shape": list(frames.shape),
+                    "metadata": metadata,
+                }
+            )
+            chunk_started = now
+            print(json.dumps(measurements[-1], sort_keys=True), flush=True)
+    finally:
+        engine.shutdown()
+
+    delivered_frames = sum(int(chunk["frames"]) for chunk in measurements)
+    summary_path = output_dir / "summary.json"
+    summary_path.write_text(
+        json.dumps(
+            {
+                "request_id": args.request_id,
+                "chunks_delivered": len(measurements),
+                "chunks_expected": num_chunks,
+                "frames_delivered": delivered_frames,
+                "frames_expected": num_frames,
+                # The point of session-owned decode: a streamed rollout carries
+                # the same frame timeline as an offline decode of the same
+                # latents. Per-chunk decode would deliver fewer, because every
+                # block would re-expand its own opening frame.
+                "matches_offline_frame_timeline": delivered_frames == num_frames,
+                "time_to_first_frame_seconds": first_frame_at,
+                "total_seconds": time.perf_counter() - t_start,
+                "chunks": measurements,
+            },
+            indent=2,
+            sort_keys=True,
+        )
+        + "\n"
+    )
+    return summary_path
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+    return asyncio.run(run(argv))
+
+
+if __name__ == "__main__":
+    main()

tests/diffusion/ar_diffusion/test_streaming_decode.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Streaming VAE decode against the real Wan autoencoder, on CPU.
+
+The autoencoder is built from a config rather than a checkpoint -- weights are
+random, only shapes and the temporal cache protocol matter -- so these run
+without a device or a download.
+
+The load-bearing test is equivalence: decoding a session chunk by chunk must
+produce exactly what decoding the whole clip produces. Everything else about
+streaming is worthless if that does not hold.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.streaming_decode import (
+    StreamingDecodeState,
+    SupportsStreamingDecode,
+    WanStreamingDecoder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Small enough to decode on CPU in a test, same architecture as the shipped
+# checkpoint: base_dim 96 / dim_mult [1,2,4,4] would be far too slow, so the
+# widths are scaled down while the causal structure is unchanged.
+VAE_CONFIG = {
+    "base_dim": 8,
+    "z_dim": 4,
+    "dim_mult": [1, 2],
+    "num_res_blocks": 1,
+    "temperal_downsample": [True],
+    "attn_scales": [],
+    "dropout": 0.0,
+}
+LATENT_H = LATENT_W = 4
+
+
+@pytest.fixture(scope="module")
+def vae():
+    diffusers = pytest.importorskip("diffusers")
+    torch.manual_seed(0)
+    model = diffusers.AutoencoderKLWan(**VAE_CONFIG).eval()
+    for parameter in model.parameters():
+        parameter.requires_grad_(False)
+    return model
+
+
+@pytest.fixture(scope="module")
+def decoder(vae):
+    return WanStreamingDecoder(vae)
+
+
+def _latent(num_frames: int, *, seed: int = 0) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    return torch.randn(1, VAE_CONFIG["z_dim"], num_frames, LATENT_H, LATENT_W, generator=generator)
+
+
+@torch.no_grad()
+def _decode_whole(vae, latent: torch.Tensor) -> torch.Tensor:
+    """The non-streaming reference: one call, cache cleared at both ends."""
+    return vae._decode(latent, return_dict=False)[0]
+
+
+# --------------------------------------------------------------------------
+# Equivalence: the property that makes streaming worth doing at all
+# --------------------------------------------------------------------------
+
+
+def _stream(decoder, latent: torch.Tensor, chunk_sizes, *, session_id: str = "s") -> torch.Tensor:
+    state = decoder.new_decode_state(session_id)
+    pieces, offset = [], 0
+    for size in chunk_sizes:
+        pieces.append(decoder.decode_chunk(latent[:, :, offset : offset + size], state))
+        offset += size
+    return torch.cat(pieces, dim=2)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("chunk_sizes", [(3, 3, 3), (1, 1, 1, 1, 1), (2, 4, 3)])
+def test_chunking_is_neutral(vae, decoder, chunk_sizes) -> None:
+    """Splitting a session into chunks must change nothing about its output.
+
+    This is the property streaming actually introduces, so it is the one
+    asserted exactly. The control is the same session decoded in a single
+    call through the same code path -- not ``_decode`` -- because ``_decode``
+    differs in a way that has nothing to do with chunking; see
+    ``test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact``.
+    """
+    total = sum(chunk_sizes)
+    latent = _latent(total)
+    reference = _stream(decoder, latent, (total,), session_id="ref")
+    streamed = _stream(decoder, latent, chunk_sizes)
+
+    assert streamed.shape == reference.shape
+    torch.testing.assert_close(streamed, reference, rtol=0, atol=0)
+
+
+@torch.no_grad()
+def test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact(vae, decoder) -> None:
+    """Explain, and bound, the gap against diffusers' ``_decode``.
+
+    ``_decode`` applies ``post_quant_conv`` to the whole latent at once; the
+    tiled path in this repo, and this decoder, apply it per frame. It is a
+    1x1x1 convolution, so the two are mathematically identical -- but they can
+    select different kernels, and on bf16 CUDA that difference is real and is
+    amplified by the causal convolutions downstream.
+
+    Asserting equality with ``_decode`` would therefore pin a kernel-selection
+    coincidence rather than a property of streaming. Assert instead that the
+    gap is entirely explained by where ``post_quant_conv`` runs.
+    """
+    latent = _latent(6)
+    per_frame_pq = torch.cat([vae.post_quant_conv(latent[:, :, k : k + 1]) for k in range(latent.shape[2])], dim=2)
+    whole_pq = vae.post_quant_conv(latent)
+
+    streamed = _stream(decoder, latent, (3, 3))
+    one_shot = _decode_whole(vae, latent)
+
+    # Same inputs to the decoder => same output, whatever the backend does.
+    if torch.equal(whole_pq, per_frame_pq):
+        torch.testing.assert_close(streamed, one_shot, rtol=0, atol=0)
+    else:  # pragma: no cover - backend dependent
+        pytest.skip("backend batches post_quant_conv differently; the gap is that, not chunking")
+
+
+@torch.no_grad()
+def test_a_restarted_stream_does_not_match_a_continued_one(decoder) -> None:
+    """Negative control: the equivalence test above must be discriminating.
+
+    Clearing state between chunks -- what the non-streaming path does on every
+    call -- has to change the pixels, otherwise the temporal cache carries
+    nothing and the assertion above holds for a write-only cache too.
+
+    Restarting also changes the frame count, because a fresh state expands its
+    opening latent frame to one raw frame. Comparing totals would therefore
+    pass on the frame count alone and never look at a pixel, so this compares
+    the tail both runs agree on: the frames from the chunk's later latents,
+    which are the full temporal factor wide on either side.
+    """
+    latent = _latent(6)
+    state = decoder.new_decode_state("continued")
+    decoder.decode_chunk(latent[:, :, :3], state)
+    continued = decoder.decode_chunk(latent[:, :, 3:], state)
+    restarted = decoder.decode_chunk(latent[:, :, 3:], decoder.new_decode_state("restarted"))
+
+    tail = min(continued.shape[2], restarted.shape[2]) - 1
+    assert tail > 0, "the comparison needs at least one fully expanded frame on both sides"
+    continued_tail, restarted_tail = continued[:, :, -tail:], restarted[:, :, -tail:]
+    assert continued_tail.shape == restarted_tail.shape, "the tails must be comparable for this to mean anything"
+    assert not torch.equal(continued_tail, restarted_tail), (
+        "restarting the decoder produced the same pixels as continuing it, so the carried cache is not being read"
+    )
+
+
+# --------------------------------------------------------------------------
+# Cross-session isolation: the failure that timing metrics cannot see
+# --------------------------------------------------------------------------
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("seed_a, seed_b", [(1, 2), (11, 12)])
+def test_interleaved_sessions_keep_independent_temporal_context(vae, decoder, seed_a: int, seed_b: int) -> None:
+    """Two sessions ticking alternately must each match their solo decode.
+
+    A session that reads another's temporal cache produces a perfectly
+    plausible but wrong video, and every latency metric stays green, so this
+    has to be asserted against a recorded solo run rather than inferred. The
+    control is a solo stream rather than ``_decode``, so it differs from the
+    run under test only in the interleaving.
+    """
+    latent_a = _latent(6, seed=seed_a)
+    latent_b = _latent(6, seed=seed_b)
+    solo_a = _stream(decoder, latent_a, (6,), session_id="solo-a")
+    solo_b = _stream(decoder, latent_b, (6,), session_id="solo-b")
+    assert not torch.equal(solo_a, solo_b), "negative control: the two sessions must differ"
+
+    state_a = decoder.new_decode_state("a")
+    state_b = decoder.new_decode_state("b")
+    out_a, out_b = [], []
+    for start in (0, 3):
+        window = slice(start, start + 3)
+        out_a.append(decoder.decode_chunk(latent_a[:, :, window], state_a))
+        out_b.append(decoder.decode_chunk(latent_b[:, :, window], state_b))
+
+    torch.testing.assert_close(torch.cat(out_a, dim=2), solo_a, rtol=0, atol=0)
+    torch.testing.assert_close(torch.cat(out_b, dim=2), solo_b, rtol=0, atol=0)
+
+
+@torch.no_grad()
+def test_the_module_cache_is_untouched_so_sessions_can_share_one_decoder(vae, decoder) -> None:
+    vae.clear_cache()
+    before = [entry for entry in vae._feat_map]
+    decoder.decode_chunk(_latent(3), decoder.new_decode_state("s"))
+    assert vae._feat_map == before, "streaming decode must not write the module-owned cache"
+
+
+# --------------------------------------------------------------------------
+# Causal geometry and boundedness
+# --------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def test_opening_chunk_is_shorter_and_later_chunks_are_full_length(vae, decoder) -> None:
+    """(n - 1) * factor + 1 for the opening chunk, n * factor after it."""
+    factor = 2 ** sum(VAE_CONFIG["temperal_downsample"])
+    state = decoder.new_decode_state("s")
+    first = decoder.decode_chunk(_latent(3), state)
+    second = decoder.decode_chunk(_latent(3, seed=5), state)
+
+    assert first.shape[2] == (3 - 1) * factor + 1
+    assert second.shape[2] == 3 * factor
+    assert state.chunks_decoded == 2
+    assert state.frames_decoded == 6
+
+
+@torch.no_grad()
+def test_resident_state_does_not_grow_with_session_length(vae, decoder) -> None:
+    """Each causal convolution retains at most CACHE_T temporal slices."""
+    state = decoder.new_decode_state("s")
+    decoder.decode_chunk(_latent(3), state)
+    baseline = state.nbytes()
+    assert baseline > 0
+
+    for step in range(8):
+        decoder.decode_chunk(_latent(3, seed=step + 10), state)
+        assert state.nbytes() == baseline
+
+    assert sum(state.nbytes_by_device().values()) == baseline
+
+
+@torch.no_grad()
+def test_release_returns_the_session_to_its_pre_stream_state(vae, decoder) -> None:
+    state = decoder.new_decode_state("s")
+    decoder.decode_chunk(_latent(3), state)
+    assert state.started and state.nbytes() > 0
+
+    decoder.release(state)
+    assert state.nbytes() == 0
+    assert not state.started
+    assert state.chunks_decoded == 0
+
+    # A released session restarts rather than resumes, so its first chunk is
+    # short again -- the same semantics reset() has for AR KV.
+    factor = 2 ** sum(VAE_CONFIG["temperal_downsample"])
+    assert decoder.decode_chunk(_latent(3), state).shape[2] == (3 - 1) * factor + 1
+
+
+# --------------------------------------------------------------------------
+# Contract
+# --------------------------------------------------------------------------
+
+
+def test_the_vaes_execution_context_is_entered_for_every_frame(vae) -> None:
+    """Decode must run inside the autoencoder's own execution context.
+
+    Every public compute entry point on ``OmniAutoencoderKLWan`` enters
+    ``_execution_context()``, which is where autocast is established for
+    fp16/bf16 parameters -- and on NPU the only place it is established at all.
+    Driving ``post_quant_conv`` and ``decoder`` directly bypasses it unless this
+    decoder enters it itself, which would compute streamed chunks under a
+    different numeric context from the whole-clip decode they are compared to.
+    """
+    entered: list[str] = []
+
+    class _ContextualVAE:
+        """The parts of the contract this decoder uses, plus the hook."""
+
+        config = vae.config
+        dtype = torch.float32
+
+        def __init__(self) -> None:
+            self.decoder = vae.decoder
+            self.post_quant_conv = vae.post_quant_conv
+            self._cached_conv_counts = vae._cached_conv_counts
+
+        @contextmanager
+        def _execution_context(self):
+            entered.append("enter")
+            try:
+                yield
+            finally:
+                entered.append("exit")
+
+    decoder = WanStreamingDecoder(_ContextualVAE())
+    with torch.no_grad():
+        decoder.decode_chunk(_latent(3), decoder.new_decode_state("s"))
+
+    # One context for the call, not one per frame, and it closes.
+    assert entered == ["enter", "exit"]
+
+
+def test_a_plain_diffusers_autoencoder_needs_no_execution_context(decoder) -> None:
+    """The hook is optional: a VAE without one still decodes."""
+    with torch.no_grad():
+        out = decoder.decode_chunk(_latent(1), decoder.new_decode_state("s"))
+    assert out.shape[2] == 1
+
+
+def test_wan_decoder_satisfies_the_protocol(decoder) -> None:
+    assert isinstance(decoder, SupportsStreamingDecode)
+
+
+def test_state_from_another_decoder_is_rejected(decoder) -> None:
+    foreign = StreamingDecodeState(session_id="s", feat_map=[None])
+    with pytest.raises(ValueError, match="does not belong to this decoder"):
+        decoder.decode_chunk(_latent(1), foreign)
+
+
+def test_a_non_wan_module_is_rejected() -> None:
+    with pytest.raises(TypeError, match="Wan-family autoencoder"):
+        WanStreamingDecoder(object())
+
+
+def test_empty_and_malformed_latents_are_rejected(decoder) -> None:
+    state = decoder.new_decode_state("s")
+    with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"):
+        decoder.decode_chunk(torch.zeros(1, 4, 4, 4), state)
+    with pytest.raises(ValueError, match="at least one frame"):
+        decoder.decode_chunk(torch.zeros(1, 4, 0, LATENT_H, LATENT_W), state)
+
+
+def test_a_batch_is_refused_rather_than_threaded_through_one_cache(decoder) -> None:
+    """One state is one session, so a batch has no correct interpretation here.
+
+    ``AutoencoderKLWan.decode`` handles a batch by decoding each sample with
+    its own cleared cache (``use_slicing``). This decoder carries exactly one
+    temporal context, so a batch would mix samples into it and still return
+    correctly shaped pixels.
+    """
+    state = decoder.new_decode_state("s")
+    batched = torch.zeros(2, VAE_CONFIG["z_dim"], 3, LATENT_H, LATENT_W)
+    with pytest.raises(ValueError, match="one sample per session state"):
+        decoder.decode_chunk(batched, state)
+    assert state.frames_decoded == 0, "a refused call must not advance the session"
+
+
+def test_session_id_must_be_meaningful(decoder) -> None:
+    with pytest.raises(ValueError, match="session_id"):
+        decoder.new_decode_state("  ")
+
+
+def test_declared_state_bytes_scales_with_area_and_dtype(vae) -> None:
+    # A local decoder, not the module-scoped fixture: the measured constant
+    # belongs to this test, and leaking it would change what a reordered run
+    # of the default-None case sees. Passing it here is also the constructor
+    # argument's intended use.
+    # Measured for the shipped checkpoint: 37832 KiB of fp32 cache at 64x64.
+    decoder = WanStreamingDecoder(vae, bytes_per_pixel_fp32=37832 * 1024 / (64 * 64))
+    fp32 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.float32)
+    bf16 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.bfloat16)
+    half_area = decoder.declared_state_bytes(height=240, width=832, dtype=torch.float32)
+    assert bf16 == fp32 // 2
+    assert half_area == fp32 // 2
+    assert fp32 / 2**20 == pytest.approx(3602.0, abs=5.0)

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -170,10 +170,49 @@ class _RecordingTransformer(nn.Module):
         return {name for name, _ in self.loaded_weights}
 
 
+class _StubCausalDecoder:
+    """Stand-in for the Wan causal decoder's cache protocol.
+
+    Reproduces only what streaming depends on: each call walks ``feat_idx``
+    across one cache slot per causal convolution, and the session's opening
+    frame (``first_chunk``) expands to a single raw frame while every later
+    frame expands by the temporal factor.
+    """
+
+    NUM_CONVS = 3
+
+    def __init__(self) -> None:
+        self.first_chunk_flags: list[bool] = []
+        # Identity of the cache each call wrote through, so an interleaved run
+        # can assert that a session never advanced another session's context.
+        self.cache_ids: list[int] = []
+
+    def __call__(self, x: torch.Tensor, *, feat_cache, feat_idx, first_chunk: bool = False):
+        self.first_chunk_flags.append(bool(first_chunk))
+        self.cache_ids.append(id(feat_cache))
+        for _ in range(self.NUM_CONVS):
+            index = feat_idx[0]
+            feat_cache[index] = x.detach().clone()
+            feat_idx[0] = index + 1
+        num_frames = 1 if first_chunk else 4
+        return x.new_zeros(x.shape[0], 3, num_frames, x.shape[-2] * 8, x.shape[-1] * 8)
+
+
 class _StubVAE(_FakePretrained):
     dtype = torch.float32
 
-    def __init__(self):
+    def __init__(self, *, streaming: bool = True):
+        if streaming:
+            self.decoder = _StubCausalDecoder()
+            # Recorded, not discarded: on the streaming branch this is the only
+            # place the tensor handed to the decoder can be observed, and it is
+            # what carries the checkpoint's latent rescale.
+            self.post_quant_inputs: list[torch.Tensor] = []
+            self.post_quant_conv = self._record_post_quant
+            self._cached_conv_counts = {"decoder": _StubCausalDecoder.NUM_CONVS, "encoder": 2}
+            # Module-owned cache the shared VAE keeps for whole-clip decode;
+            # streaming must never write through it.
+            self._feat_map = ["module-owned"]
         self.config = SimpleNamespace(
             z_dim=16,
             scale_factor_temporal=4,
@@ -185,6 +224,10 @@ class _StubVAE(_FakePretrained):
         self.decode_inputs: list[torch.Tensor] = []
         self.on_decode = None
 
+    def _record_post_quant(self, latent: torch.Tensor) -> torch.Tensor:
+        self.post_quant_inputs.append(latent.detach().clone())
+        return latent
+
     def encode(self, video: torch.Tensor):
         self.encode_inputs.append(video.detach().clone())
         latent_frames = (video.shape[2] - 1) // 4 + 1
@@ -455,7 +498,41 @@ def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None:
     assert spec.sink_frames == 3
     assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)]
     assert spec.cross_attention_lengths == {"text": 512}
-    assert spec.model_owned_state_bytes_per_session == 9_600
+    # The image condition (9,600 B) plus the streaming decoder's own per-session
+    # temporal cache. The cache dwarfs the condition at any real resolution, and
+    # the runner subtracts this field from the KV budget before sizing its pools,
+    # so leaving the cache out of it over-commits device memory.
+    assert spec.model_owned_state_bytes_per_session == 9_600 + 2_421_248
+
+
+def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
+    """The declaration scales with output area, and is the condition alone when
+    the VAE cannot stream."""
+    module = _load_pipeline_module()
+
+    def spec_pipeline():
+        pipeline = _pipeline(module)
+        # The spec reads the constructed attention's head count; the head count
+        # plays no part in model-owned bytes, so any value will do here.
+        pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1)
+        return pipeline
+
+    small = spec_pipeline().ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
+
+    wide = spec_pipeline()
+    wide._ar_width = wide._ar_width * 2
+    widened = wide.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
+    # Twice the area: the decoder cache doubles, so the total grows by far more
+    # than the image condition alone could account for.
+    assert widened - small > 2_000_000
+
+    # A VAE without the causal-cache contract cannot stream, so there is no
+    # decoder cache to declare and the figure falls back to the condition.
+    bare = spec_pipeline()
+    for attribute in ("decoder", "post_quant_conv", "_cached_conv_counts"):
+        if hasattr(bare.vae, attribute):
+            delattr(bare.vae, attribute)
+    assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 9_600
 
 
 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None:
@@ -1852,6 +1929,27 @@ def _run_stepwise(pipeline, state):
     return outputs
 
 
+def _stepwise_chunks(pipeline, state, ar_state):
+    """Yield one chunk at a time, binding AR state around each invocation.
+
+    The runner binds runner-owned KV for the duration of one stepwise
+    invocation and releases it again, so a generator that suspends at each
+    chunk boundary is what lets two requests be driven the way the scheduler
+    drives them at session_capacity > 1: A-chunk0, B-chunk0, A-chunk1, ...
+    """
+    with pipeline.bind_ar_diffusion_state(state.request_id, ar_state):
+        pipeline.prepare_encode(state)
+    while not state.request_denoise_completed:
+        output = None
+        with pipeline.bind_ar_diffusion_state(state.request_id, ar_state):
+            noise = pipeline.denoise_step(None, states=[state])
+            pipeline.step_scheduler(state, noise)
+            if state.chunk_denoise_completed:
+                output = pipeline.post_decode(state)
+        if output is not None:
+            yield output
+
+
 def test_pipeline_declares_step_execution_support() -> None:
     module = _load_pipeline_module()
     pipeline = _pipeline(module)
@@ -2039,12 +2137,11 @@ def test_stepwise_emits_latents_without_touching_the_vae() -> None:
     assert pipeline.vae.decode_inputs == []
 
 
-def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
-    """A non-latent request must stream pixels under the primary "video" key."""
+def _capture_video_processor(monkeypatch) -> list[tuple[tuple[int, ...], str]]:
+    """Record what reaches postprocess_video, and name each chunk's output."""
     import diffusers.video_processor as video_processor_module
 
-    module = _load_pipeline_module()
-    processed = []
+    processed: list[tuple[tuple[int, ...], str]] = []
 
     class VideoProcessor:
         def __init__(self, *, vae_scale_factor):
@@ -2055,9 +2152,21 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non
             return f"frames-{len(processed)}"
 
     monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+    return processed
+
+
+def _streaming_pipeline(module):
     pipeline = _pipeline(module, transformer=_RecordingTransformer())
     pipeline._ar_height = 16
     pipeline._ar_width = 16
+    return pipeline
+
+
+def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
+    """A non-latent request must stream pixels under the primary "video" key."""
+    module = _load_pipeline_module()
+    processed = _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
     fake = _FakeARState(state.request_id)
@@ -2065,14 +2174,268 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non
     with pipeline.bind_ar_diffusion_state(state.request_id, fake):
         outputs = _run_stepwise(pipeline, state)
 
-    # One decode per AR block, each seeing only that block's three latent frames.
-    assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
     assert [output_type for _, output_type in processed] == ["np", "np"]
     assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
     # Identity metadata survives the switch to pixel output.
     assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
 
 
+def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> None:
+    """Chunk N + 1 must continue chunk N instead of restarting the decoder.
+
+    Only the session's opening latent frame collapses to a single raw frame,
+    so the streamed timeline carries the same frame count an offline decode of
+    the same latents would: 9 frames, then 12.
+    """
+    module = _load_pipeline_module()
+    processed = _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        _run_stepwise(pipeline, state)
+
+    assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
+    # One opening frame for the session, not one per chunk.
+    assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5
+    # The whole-clip decode path is not involved, and the module-owned cache
+    # the shared VAE keeps for it is untouched.
+    assert pipeline.vae.decode_inputs == []
+    assert pipeline.vae._feat_map == ["module-owned"]
+
+
+def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
+    """Decoder state is keyed by request id and released with the AR session."""
+    module = _load_pipeline_module()
+    _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(request_id="req-stream", num_frames=21)
+    state.sampling.output_type = "np"
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        _run_stepwise(pipeline, state)
+
+    decode_state = pipeline._streaming_decode_states["req-stream"]
+    assert decode_state.session_id == "req-stream"
+    assert decode_state.frames_decoded == 6
+    assert decode_state.nbytes() > 0
+
+    pipeline.close_ar_diffusion_session("req-stream")
+    assert "req-stream" not in pipeline._streaming_decode_states
+    # Release drops the cache itself, not just the pipeline's reference to it.
+    assert decode_state.nbytes() == 0
+    assert decode_state.frames_decoded == 0
+
+
+def test_streaming_decode_state_is_dropped_on_session_reset(monkeypatch) -> None:
+    """A reset session must not resume the temporal context it just abandoned."""
+    module = _load_pipeline_module()
+    _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(request_id="req-reset", num_frames=21)
+    state.sampling.output_type = "np"
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        _run_stepwise(pipeline, state)
+
+    pipeline.reset_ar_diffusion_session("req-reset")
+    assert pipeline._streaming_decode_states == {}
+    # Releasing a session that never decoded is a no-op, not an error.
+    pipeline.reset_ar_diffusion_session("req-reset")
+
+
+def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> None:
+    """Two sessions ticking alternately must not share one temporal cache.
+
+    Running one rollout to completion before the other starts would only show
+    that the dict is keyed by request id. The failure worth catching is chunk N
+    of one session advancing the other's cache mid-rollout, which needs the
+    interleaving the scheduler actually produces.
+    """
+    module = _load_pipeline_module()
+    processed = _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    states = []
+    for request_id in ("req-a", "req-b"):
+        state = _stepwise_state(request_id=request_id, num_frames=21)
+        state.sampling.output_type = "np"
+        states.append(state)
+    runs = [_stepwise_chunks(pipeline, state, _FakeARState(state.request_id)) for state in states]
+
+    chunks: list[tuple[str, int]] = []
+    while runs:
+        for run in list(runs):
+            try:
+                output = next(run)
+            except StopIteration:
+                runs.remove(run)
+                continue
+            metadata = output.output["metadata"]["ar_diffusion"]
+            chunks.append((metadata["session_id"], metadata["chunk_index"]))
+
+    # A-chunk0, B-chunk0, A-chunk1, B-chunk1: the sessions really are interleaved.
+    assert chunks == [("req-a", 0), ("req-b", 0), ("req-a", 1), ("req-b", 1)]
+
+    caches = pipeline._streaming_decode_states
+    assert set(caches) == {"req-a", "req-b"}
+    assert caches["req-a"].feat_map is not caches["req-b"].feat_map
+    # Every frame of a session was decoded through that session's own cache,
+    # three frames per chunk, and each session opened exactly one rollout.
+    cache_a, cache_b = (id(caches["req-a"].feat_map), id(caches["req-b"].feat_map))
+    assert pipeline.vae.decoder.cache_ids == [cache_a] * 3 + [cache_b] * 3 + [cache_a] * 3 + [cache_b] * 3
+    assert pipeline.vae.decoder.first_chunk_flags == (
+        [True, False, False] + [True, False, False] + [False] * 3 + [False] * 3
+    )
+    assert [shape for shape, _ in processed] == [
+        (1, 3, 9, 16, 16),
+        (1, 3, 9, 16, 16),
+        (1, 3, 12, 16, 16),
+        (1, 3, 12, 16, 16),
+    ]
+
+
+def test_streaming_decode_failure_drops_the_half_advanced_cache(monkeypatch) -> None:
+    """A failed chunk must not leave a cache a later chunk would continue from."""
+    module = _load_pipeline_module()
+    _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(request_id="req-boom", num_frames=21)
+    state.sampling.output_type = "np"
+
+    def explode(*args, **kwargs):
+        raise RuntimeError("decoder blew up")
+
+    pipeline.vae.decoder = explode
+
+    with pytest.raises(RuntimeError, match="decoder blew up"):
+        with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+            _run_stepwise(pipeline, state)
+
+    assert pipeline._streaming_decode_states == {}
+
+
+def _enable_vae_tiling(pipeline, *, tile_sample_min: int) -> None:
+    """Configure the stub VAE the way registry.py configures the real one."""
+    pipeline.vae.use_tiling = True
+    pipeline.vae.spatial_compression_ratio = 8
+    pipeline.vae.tile_sample_min_height = tile_sample_min
+    pipeline.vae.tile_sample_min_width = tile_sample_min
+
+
+def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None:
+    """The streaming branch must invert the checkpoint's latent statistics.
+
+    Nothing downstream can catch a miss here: the pixels stay finite and
+    correctly shaped, only wrongly scaled. Request mode and the fallback path
+    observe the rescale through ``vae.decode``'s recorded input, which streaming
+    never reaches, so this asserts on what ``post_quant_conv`` received.
+    """
+    module = _load_pipeline_module()
+    _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        pipeline.prepare_encode(state)
+        while not state.chunk_denoise_completed:
+            pipeline.step_scheduler(state, pipeline.denoise_step(None, states=[state]))
+        model_space = state.latents.clone()
+        pipeline.post_decode(state)
+
+    shape = (1, -1, 1, 1, 1)
+    latent_mean = torch.as_tensor(pipeline.vae.config.latents_mean, dtype=model_space.dtype).view(*shape)
+    latent_std = torch.as_tensor(pipeline.vae.config.latents_std, dtype=model_space.dtype).view(*shape)
+    expected = (model_space * latent_std + latent_mean).to(dtype=pipeline.vae.dtype)
+
+    recorded = torch.cat(pipeline.vae.post_quant_inputs, dim=2)
+    assert recorded.shape == expected.shape, "one recorded frame per latent frame in the chunk"
+    torch.testing.assert_close(recorded, expected, rtol=0, atol=0)
+    # And the rescale is not a no-op on this fixture, so the assertion bites.
+    assert not torch.equal(expected, model_space.to(dtype=pipeline.vae.dtype))
+
+
+def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
+    """A streamed rollout must report decode time, not zero.
+
+    The profiler wraps ``vae.decode``, which streaming never calls, so without
+    its own stage the one stage this path changes would be invisible -- and a
+    perf regression in the streaming decoder would be too.
+    """
+    module = _load_pipeline_module()
+    _capture_video_processor(monkeypatch)
+    pipeline = _pipeline(
+        module,
+        transformer=_RecordingTransformer(),
+        od_config=_od_config(
+            enable_diffusion_pipeline_profiler=True,
+        ),
+    )
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        outputs = _run_stepwise(pipeline, state)
+
+    stage = f"{type(pipeline).__name__}._streaming_decode_chunk"
+    durations = pipeline.stage_durations
+    assert stage in durations, f"streamed decode is unprofiled; stages seen: {sorted(durations)}"
+    assert durations[stage] > 0.0
+    # The stage reaches the consumer on every streamed chunk, not just at the end.
+    assert all(stage in (output.stage_durations or {}) for output in outputs)
+
+
+def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> None:
+    """Tiled decode drives its own cache, so a tiling shape keeps per-chunk decode."""
+    module = _load_pipeline_module()
+    processed = _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    # A latent of 2x2 exceeds a one-latent-cell tile, so vae.decode would tile.
+    _enable_vae_tiling(pipeline, tile_sample_min=8)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        outputs = _run_stepwise(pipeline, state)
+
+    # One whole-clip decode per AR block, each seeing only that block's frames.
+    assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
+    assert pipeline.vae.decoder.first_chunk_flags == []
+    assert pipeline._streaming_decode_states == {}
+    # Nine frames per chunk instead of 9 then 12: the frame loss being removed.
+    assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)]
+    assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
+
+
+def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monkeypatch) -> None:
+    """``use_tiling`` alone does not tile, so such a shape must still stream.
+
+    ``AutoencoderKLWan._decode`` tiles only when the latent exceeds a tile. A
+    gate on the flag alone would give up streaming -- and a frame per block --
+    on a configuration that would never have tiled.
+    """
+    module = _load_pipeline_module()
+    processed = _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    # A 2x2 latent is within a tile this size, so vae.decode would not tile.
+    _enable_vae_tiling(pipeline, tile_sample_min=256)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        _run_stepwise(pipeline, state)
+
+    assert pipeline.vae.decode_inputs == []
+    assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5
+    assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
+
+
 def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
     module = _load_pipeline_module()
     resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(

vllm_omni/diffusion/models/lingbot_world/pipeline.py

@@ -18,6 +18,7 @@ import torch.nn.functional as F
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
+from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -62,11 +63,26 @@ from vllm_omni.experimental.ar_diffusion.capability import (
     ARDiffusionKVBranchSpec,
     ARDiffusionKVCacheSpec,
 )
+from vllm_omni.experimental.ar_diffusion.streaming_decode import (
+    StreamingDecodeState,
+    WanStreamingDecoder,
+)
 from vllm_omni.experimental.ar_diffusion.tick_protocol import (
     ARDiffusionChunkMetadata,
     ARDiffusionTickRequest,
 )
 
+logger = init_logger(__name__)
+
+# Resident streaming-decoder cache per session, as bytes per output pixel in
+# fp32: measured once on the reference checkpoint at 37832 KiB for a 64x64
+# output. Every cached tensor is [B, C, <= CACHE_T, H/8, W/8] per causal
+# convolution, so the total scales with height * width and never with session
+# length; tests/diffusion/ar_diffusion/test_streaming_decode.py pins that
+# scaling, and at 832x480 in bf16 this yields the 1801 MiB the decoder was
+# measured to hold.
+_STREAMING_DECODE_BYTES_PER_PIXEL_FP32 = 37832 * 1024 / (64 * 64)
+
 if TYPE_CHECKING:
     from diffusers.video_processor import VideoProcessor
     from tqdm.std import tqdm as TqdmProgressBar
@@ -552,10 +568,20 @@ class LingBotWorldCausalDMDPipeline(
         self._ar_width = int(model_config.get("ar_diffusion_width", 832))
         self._ar_diffusion_kv_state: ARDiffusionKVState | None = None
         self._ar_sessions: dict[str, _LingBotARSessionState] = {}
+        # One temporal decoder cache per stepwise session, keyed the way the
+        # runner keys AR sessions (session_id == request_id) so both are
+        # released together rather than through two independent lifecycles.
+        self._streaming_decode_states: dict[str, StreamingDecodeState] = {}
+        self._cached_streaming_decoder: WanStreamingDecoder | None = None
+        self._streaming_decode_unsupported = False
         self.setup_diffusion_pipeline_profiler(
             profiler_targets=[
                 "vae.encode",
                 "vae.decode",
+                # Streaming decode never reaches vae.decode, so it needs its
+                # own stage or a streamed rollout reports no decode time at
+                # all -- the one stage this path changes.
+                "_streaming_decode_chunk",
                 "_generate_block",
                 "text_encoder.forward",
                 "tokenizer.forward",
@@ -608,7 +634,32 @@ class LingBotWorldCausalDMDPipeline(
                     _MAX_SEQUENCE_LENGTH,
                 ),
             ),
-            model_owned_state_bytes_per_session=condition_bytes_per_session,
+            model_owned_state_bytes_per_session=(
+                condition_bytes_per_session + self._streaming_decode_bytes_per_session()
+            ),
+        )
+
+    def _streaming_decode_bytes_per_session(self) -> int:
+        """Resident streaming-decoder bytes one session can hold, for admission.
+
+        ``model_owned_state_bytes_per_session`` is what the runner subtracts
+        from the KV budget before sizing the block pools, and this cache is
+        three orders of magnitude larger than the image condition that used to
+        be the only entry: leaving it out lets the profiler report a
+        configuration as fitting and then OOM partway through the second
+        session's rollout.
+
+        Declared for every session rather than only for pixel-output ones: the
+        field is documented as a conservative upper bound, it is fixed at load
+        time, and whether a request asks for pixels is not known then.
+        """
+        decoder = self._streaming_decoder()
+        if decoder is None:
+            return 0
+        return decoder.declared_state_bytes(
+            height=self._ar_height,
+            width=self._ar_width,
+            dtype=self.vae.dtype,
         )
 
     @contextmanager
@@ -629,9 +680,11 @@ class LingBotWorldCausalDMDPipeline(
 
     def reset_ar_diffusion_session(self, session_id: str) -> None:
         self._ar_sessions.pop(session_id, None)
+        self._release_streaming_decode_state(session_id)
 
     def close_ar_diffusion_session(self, session_id: str) -> None:
         self._ar_sessions.pop(session_id, None)
+        self._release_streaming_decode_state(session_id)
 
     def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs:
         if req.num_reqs != 1 or len(req.prompts) != 1:
@@ -1279,15 +1332,90 @@ class LingBotWorldCausalDMDPipeline(
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
+    def _vae_would_tile(self, latents: torch.Tensor) -> bool:
+        """Whether ``vae.decode`` would take its tiled path for this shape.
+
+        Tiled decode splits a frame spatially and drives the module's own
+        ``feat_cache`` across those tiles, so a per-session cache threaded
+        through it is a separate design and is not attempted here. The test
+        mirrors ``AutoencoderKLWan._decode``: the flag alone does not tile, the
+        latent also has to exceed a tile, so a sub-threshold resolution with
+        tiling enabled still streams. A VAE that sets the flag but reports no
+        tile size is treated as tiling, because the safe answer when the
+        threshold is unknown is the one that keeps ``vae.decode`` in charge.
+        """
+        if not bool(getattr(self.vae, "use_tiling", False)):
+            return False
+        compression = int(getattr(self.vae, "spatial_compression_ratio", None) or self.vae_scale_factor_spatial)
+        tile_latent_min_height = int(getattr(self.vae, "tile_sample_min_height", None) or 0) // compression
+        tile_latent_min_width = int(getattr(self.vae, "tile_sample_min_width", None) or 0) // compression
+        return bool(latents.shape[-2] > tile_latent_min_height or latents.shape[-1] > tile_latent_min_width)
+
+    def _streaming_decoder(self) -> WanStreamingDecoder | None:
+        """The session-owned streaming decoder, or ``None`` on a VAE that cannot stream.
+
+        Only the decoder's own capability is decided here, and the decision is
+        cached including the negative. Whether a particular *shape* would have
+        been tiled is a per-call question and lives in :meth:`_vae_would_tile`.
+        """
+        if self._streaming_decode_unsupported:
+            return None
+        if self._cached_streaming_decoder is not None:
+            return self._cached_streaming_decoder
+        try:
+            decoder = WanStreamingDecoder(
+                self.vae,
+                bytes_per_pixel_fp32=_STREAMING_DECODE_BYTES_PER_PIXEL_FP32,
+            )
+            # Materialize the cache geometry now: a decoder that cannot report
+            # it would otherwise fail on the first chunk instead.
+            _ = decoder.num_causal_convs
+        except TypeError as exc:
+            logger.warning(
+                "LingBot streaming decode is unavailable (%s); AR blocks will be decoded independently, "
+                "which drops each block's opening frame and restarts temporal context at every boundary.",
+                exc,
+            )
+            self._streaming_decode_unsupported = True
+            return None
+        self._cached_streaming_decoder = decoder
+        return decoder
+
+    def _streaming_decode_state(self, decoder: WanStreamingDecoder, session_id: str) -> StreamingDecodeState:
+        state = self._streaming_decode_states.get(session_id)
+        if state is None:
+            state = decoder.new_decode_state(session_id)
+            self._streaming_decode_states[session_id] = state
+        return state
+
+    def _release_streaming_decode_state(self, session_id: str) -> None:
+        """Drop one session's decoder cache. Safe for a session that never decoded.
+
+        The state is released directly rather than through
+        ``WanStreamingDecoder.release``, which is the same one-line call: the
+        decoder's method exists for a caller holding only the
+        ``SupportsStreamingDecode`` protocol, and this pipeline holds the state.
+        """
+        state = self._streaming_decode_states.pop(session_id, None)
+        if state is not None:
+            state.release()
+
     def _decode_chunk_to_pixels(
-        self, latents: torch.Tensor, *, output_type: str
+        self, latents: torch.Tensor, *, output_type: str, session_id: str | None = None
     ) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]:
         """Decode one AR block so streaming consumers receive pixels, not latents.
 
-        Blocks are decoded independently, which keeps the shared VAE stateless
-        and lets a rollout be served without a per-session temporal cache. That
-        is the same shape Helios streams today; cross-block ``feat_cache``
-        continuity is a separate decision and is not made here.
+        With a ``session_id`` the block is decoded through that session's own
+        temporal cache, so chunk *N + 1* continues chunk *N*: the causal decoder
+        expands the session's opening latent frame to one raw frame and every
+        later one to the full temporal factor, which is the frame timeline an
+        offline whole-clip decode of the same latents produces. The cache lives
+        on the session rather than on the shared VAE module, so concurrent
+        sessions cannot overwrite one another's temporal context.
+
+        Without a ``session_id`` -- or on a VAE that cannot stream -- blocks are
+        decoded independently: every block re-expands its own opening frame, so
+        the rollout loses frames and continuity at each block boundary.
 
         The registered post-process hook returns a ``{"payload", "metadata"}``
         envelope untouched, so the pixel conversion it would do for a bare
@@ -1295,9 +1423,34 @@ class LingBotWorldCausalDMDPipeline(
         """
         latent_mean, latent_std = self._vae_latent_stats(latents)
         vae_latents = (latents * latent_std + latent_mean).to(dtype=self.vae.dtype)
-        video = self.vae.decode(vae_latents, return_dict=False)[0]
+        decoder = None
+        if session_id is not None and not self._vae_would_tile(vae_latents):
+            decoder = self._streaming_decoder()
+        if decoder is None:
+            video = self.vae.decode(vae_latents, return_dict=False)[0]
+        else:
+            video = self._streaming_decode_chunk(decoder, vae_latents, cast(str, session_id))
         return self._video_processor().postprocess_video(video, output_type=output_type)
 
+    def _streaming_decode_chunk(
+        self, decoder: WanStreamingDecoder, vae_latents: torch.Tensor, session_id: str
+    ) -> torch.Tensor:
+        """Decode one chunk through ``session_id``'s temporal cache.
+
+        A separate method so the profiler has a stage to wrap: streaming never
+        calls ``vae.decode``, and without this a profiled rollout would report
+        no decode time for the one stage this path changes.
+        """
+        state = self._streaming_decode_state(decoder, session_id)
+        try:
+            return decoder.decode_chunk(vae_latents, state)
+        except Exception:
+            # A half-advanced cache cannot be resumed: the session's next chunk
+            # would continue from a temporal context that was never completed,
+            # so drop it here rather than at close.
+            self._release_streaming_decode_state(session_id)
+            raise
+
     def _video_processor(self) -> VideoProcessor:
         processor = getattr(self, "_cached_video_processor", None)
         if processor is None:
@@ -1566,7 +1719,13 @@ class LingBotWorldCausalDMDPipeline(
         if inputs.output_type == "latent":
             payload: dict[str, Any] = {"latents": latents}
         else:
-            payload = {"video": self._decode_chunk_to_pixels(latents, output_type=inputs.output_type)}
+            payload = {
+                "video": self._decode_chunk_to_pixels(
+                    latents,
+                    output_type=inputs.output_type,
+                    session_id=state.request_id,
+                )
+            }
         output = {
             "payload": payload,
             "metadata": {

vllm_omni/experimental/ar_diffusion/streaming_decode.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Session-owned streaming VAE decode for realtime AR-Diffusion.
+
+A causal video decoder already decodes one latent frame at a time, carrying a
+bounded temporal cache between frames. What it does not do today is carry that
+cache *between calls*: every decode call clears it, so a chunk can only be
+decoded as part of a whole clip.
+
+This module moves the cache from the decoder module to the session, which is
+the entire difference between "decode a clip" and "stream a session":
+
+* the cache becomes :class:`StreamingDecodeState`, threaded in and out by the
+  caller, so two concurrent sessions cannot overwrite each other's temporal
+  context -- a failure that produces a plausible but wrong video while every
+  timing metric stays green;
+* ``first_chunk`` becomes a property of the session's first frame rather than
+  of the call, which is what makes chunk ``N + 1`` continue chunk ``N``.
+
+Resident state is bounded by construction: each causal convolution retains at
+most ``CACHE_T`` temporal slices, so the cache is a function of resolution and
+never of session length. It is not free -- see :meth:`StreamingDecodeState.nbytes`,
+which is the quantity session admission has to account for.
+
+Measured on the reference checkpoint at 832x480 in bf16, on one RTX PRO 6000
+with torch 2.11: chunking is neutral -- a session decoded in four chunks is
+bit-identical to the same session decoded in one call -- and resident decoder
+state holds at 1801 MiB across every chunk. What a session gains is structural
+rather than a speedup: a chunk becomes decodable on its own, so pixels leave on
+the block that produced them instead of after the last one, and the session's
+opening frame is expanded once rather than once per block.
+
+Note what "neutral" is measured against. ``AutoencoderKLWan._decode`` applies
+``post_quant_conv`` to the whole latent at once, while the tiled path in this
+repo -- and this decoder -- apply it per frame. It is a 1x1x1 convolution, so
+the two are mathematically identical, but they can select different kernels;
+on bf16 CUDA that seeds a 0.004 difference which the causal convolutions
+downstream amplify to 0.047. Comparing against ``_decode`` therefore measures
+where ``post_quant_conv`` runs, not anything about streaming, and the result
+moves with the torch version. The control that isolates chunking is the same
+session decoded in a single call through this same path.
+
+The consumer is the stepwise execution path: one ``generate()`` is one
+session, ``post_decode`` emits one chunk per AR block, and the pipeline holds
+one state per ``request_id``, releasing it when the runner closes or resets
+that AR session -- so decoder state and KV state share a single lifetime.
+
+Only ``torch`` is imported here so the contract can be exercised without a
+device, a checkpoint, or the distributed VAE stack.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+import torch
+
+
+@dataclass
+class StreamingDecodeState:
+    """One session's temporal decoder cache.
+
+    ``feat_map`` holds one entry per causal convolution. Entries are tensors,
+    ``None`` before that convolution has run, or a non-tensor sentinel the
+    decoder uses for its first upsample step, so byte accounting must skip
+    anything that is not a tensor.
+
+    The cache has **no recompute source**: it is the compressed tail of every
+    frame decoded so far and cannot be rebuilt from the prompt or the latents.
+    Releasing it ends the stream; there is no partial eviction.
+    """
+
+    session_id: str
+    feat_map: list[Any]
+    conv_idx: list[int] = field(default_factory=lambda: [0])
+    frames_decoded: int = 0
+    chunks_decoded: int = 0
+
+    def nbytes(self) -> int:
+        """Resident bytes of the cache, skipping non-tensor entries."""
+        return sum(t.numel() * t.element_size() for t in self.feat_map if torch.is_tensor(t))
+
+    def nbytes_by_device(self) -> dict[str, int]:
+        """Resident bytes grouped by device, for admission accounting."""
+        totals: dict[str, int] = {}
+        for tensor in self.feat_map:
+            if torch.is_tensor(tensor):
+                key = str(tensor.device)
+                totals[key] = totals.get(key, 0) + tensor.numel() * tensor.element_size()
+        return totals
+
+    @property
+    def started(self) -> bool:
+        return self.frames_decoded > 0
+
+    def release(self) -> None:
+        """Drop the cache and return the session to its pre-stream state."""
+        self.feat_map = [None] * len(self.feat_map)
+        self.conv_idx = [0]
+        self.frames_decoded = 0
+        self.chunks_decoded = 0
+
+
+@runtime_checkable
+class SupportsStreamingDecode(Protocol):
+    """A decoder that can emit a session's chunks incrementally.
+
+    The lifecycle mirrors the AR session's: ``new_decode_state`` at session
+    creation, ``decode_chunk`` per committed chunk, and ``release`` on reset or
+    close, so decoder state and KV state are released together rather than
+    through two independent lifecycles.
+    """
+
+    def new_decode_state(self, session_id: str) -> StreamingDecodeState:
+        """Create empty decoder state for a new session."""
+        ...
+
+    def decode_chunk(self, latent: torch.Tensor, state: StreamingDecodeState) -> torch.Tensor:
+        """Decode one committed chunk, advancing ``state`` in place."""
+        ...
+
+    def declared_state_bytes(self, *, height: int, width: int, dtype: torch.dtype) -> int:
+        """Upper bound on resident decoder bytes per session at this shape."""
+        ...
+
+    def release(self, state: StreamingDecodeState) -> None:
+        """Drop ``state``'s cache and return it to its pre-stream state.
+
+        Declared here, not just implemented on :class:`WanStreamingDecoder`,
+        so a caller programming against this protocol -- a session runtime
+        releasing decoder state alongside KV state on close or failure -- can
+        rely on it being present without depending on the concrete class.
+        """
+        ...
+
+
+class WanStreamingDecoder:
+    """Streaming decode over a Wan-family causal autoencoder.
+
+    Wraps any module exposing ``decoder``, ``post_quant_conv`` and the cached
+    causal-convolution counts. It never touches the module's own ``_feat_map``,
+    so a module shared by several sessions stays stateless between calls.
+    """
+
+    def __init__(self, vae: Any, *, bytes_per_pixel_fp32: float | None = None) -> None:
+        for attribute in ("decoder", "post_quant_conv"):
+            if not hasattr(vae, attribute):
+                raise TypeError(f"Streaming decode requires a Wan-family autoencoder exposing {attribute!r}.")
+        self._vae = vae
+        # A one-off measured constant for this checkpoint; see declared_state_bytes.
+        # Optional at construction because most callers never need admission
+        # accounting, but a constructor argument (rather than requiring an
+        # external attribute assignment after the fact) means the value can't
+        # be silently unset for a caller that does.
+        self._bytes_per_pixel_fp32 = bytes_per_pixel_fp32
+
+    def _execution_context(self) -> AbstractContextManager[None]:
+        """The autoencoder's own execution context, when it exposes one.
+
+        Every public compute entry point on ``OmniAutoencoderKLWan`` --
+        ``encode``, ``decode``, ``decode_with_chunks``, ``tile_exec`` -- runs
+        inside ``_execution_context()``, which is where autocast is entered for
+        fp16/bf16 parameters (and on NPU is the only place
+        ``torch.npu.amp.autocast`` is entered at all). This decoder drives
+        ``post_quant_conv`` and ``decoder`` directly, so it has to enter the
+        same context or streamed chunks compute under a different numeric and
+        dispatch context from the whole-clip decode they are compared against.
+
+        A plain ``diffusers`` autoencoder has no such hook, and then there is
+        nothing to enter.
+        """
+        context = getattr(self._vae, "_execution_context", None)
+        if callable(context):
+            return context()
+        return nullcontext()
+
+    @property
+    def num_causal_convs(self) -> int:
+        counts = getattr(self._vae, "_cached_conv_counts", None)
+        if not isinstance(counts, dict) or "decoder" not in counts:
+            raise TypeError("The autoencoder does not expose cached decoder convolution counts.")
+        return int(counts["decoder"])
+
+    def new_decode_state(self, session_id: str) -> StreamingDecodeState:
+        if not isinstance(session_id, str) or not session_id.strip():
+            raise ValueError("session_id must be a non-empty string.")
+        return StreamingDecodeState(session_id=session_id, feat_map=[None] * self.num_causal_convs)
+
+    def decode_chunk(self, latent: torch.Tensor, state: StreamingDecodeState) -> torch.Tensor:
+        """Decode ``latent`` as the continuation of ``state``'s session.
+
+        ``latent`` is ``[1, C, T, H, W]`` in latent space, already rescaled by
+        the pipeline's latent statistics. Frames are decoded one at a time, the
+        same loop the non-streaming path runs, with the cache carried across
+        calls instead of cleared. The returned tensor is ``[1, 3, T', H', W']``
+        where ``T'`` is smaller for a session's first chunk: the opening latent
+        frame expands to a single raw frame and every later one to the full
+        temporal factor.
+        """
+        if latent.ndim != 5:
+            raise ValueError(f"latent must be [B, C, T, H, W]; got shape {tuple(latent.shape)}.")
+        if latent.shape[0] != 1:
+            # One state is one session's temporal context, so a batch would
+            # thread several samples through a single cache. ``decode()``
+            # handles a batch by decoding each sample with its own cleared
+            # cache (``use_slicing``); there is no equivalent here, and the
+            # streaming contract is one session per state, so refuse rather
+            # than silently mix samples.
+            raise ValueError(
+                "streaming decode takes one sample per session state; "
+                f"got batch size {latent.shape[0]}. Decode each sample through its own state."
+            )
+        if len(state.feat_map) != self.num_causal_convs:
+            raise ValueError(
+                "Decoder state does not belong to this decoder: "
+                f"{len(state.feat_map)} cache slots against {self.num_causal_convs} causal convolutions."
+            )
+        num_frames = int(latent.shape[2])
+        if num_frames == 0:
+            raise ValueError("latent must carry at least one frame.")
+
+        decoded_frames = []
+        with self._execution_context():
+            for index in range(num_frames):
+                state.conv_idx[0] = 0
+                frame = self._vae.post_quant_conv(latent[:, :, index : index + 1])
+                # first_chunk marks the session's opening frame, not the call's:
+                # that is what makes chunk N + 1 continue chunk N rather than
+                # restart the causal expansion.
+                decoded_frames.append(
+                    self._vae.decoder(
+                        frame,
+                        feat_cache=state.feat_map,
+                        feat_idx=state.conv_idx,
+                        first_chunk=(state.frames_decoded == 0),
+                    )
+                )
+                state.frames_decoded += 1
+
+        out = torch.cat(decoded_frames, dim=2)
+        patch_size = getattr(getattr(self._vae, "config", None), "patch_size", None)
+        if patch_size is not None:
+            from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
+
+            out = unpatchify(out, patch_size=patch_size)
+        out = torch.clamp(out, min=-1.0, max=1.0)
+        state.chunks_decoded += 1
+        return out
+
+    def release(self, state: StreamingDecodeState) -> None:
+        state.release()
+
+    def declared_state_bytes(self, *, height: int, width: int, dtype: torch.dtype) -> int:
+        """Resident decoder bytes per session at this output shape.
+
+        Every cached tensor is ``[B, C_l, <= CACHE_T, H_l, W_l]`` with ``H_l``
+        and ``W_l`` fixed fractions of the output, so the total scales exactly
+        with ``height * width``. Callers that need the constant should measure
+        it once for their checkpoint and pass it as ``bytes_per_pixel_fp32`` to
+        the constructor; this helper scales that measured constant so admission
+        does not have to re-measure per resolution.
+        """
+        constant = self._bytes_per_pixel_fp32
+        if constant is None:
+            raise RuntimeError(
+                "declared_state_bytes needs a measured bytes-per-pixel constant for this checkpoint; "
+                "pass bytes_per_pixel_fp32= to WanStreamingDecoder() from a one-off measurement."
+            )
+        element_ratio = torch.empty((), dtype=dtype).element_size() / 4
+        return int(constant * height * width * element_ratio)