[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>

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(