[LingBot World] Eleven bit-identical removals of repeated work in the realtime path (#7648)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py

@@ -74,6 +74,17 @@ def _install_vllm_stubs() -> None:
         _SeqAllToAll4D,
     )
 
+    def _all_to_all_5D(value, scatter_idx=3, gather_idx=1, group=None, use_sync=False):
+        # Single-rank stub: the fused q/k/v exchange is the identity when N == 1.
+        del scatter_idx, gather_idx, group, use_sync
+        return value
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.comm"],
+        "all_to_all_5D",
+        _all_to_all_5D,
+    )
+
     def get_sp_group():
         return SimpleNamespace(
             ulysses_world_size=1,
@@ -588,26 +599,7 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP
     torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
 
 
-def test_single_token_text_kv_shard_owns_compact_storage(monkeypatch):
-    module = _load_module()
-    monkeypatch.setattr(
-        module,
-        "get_sp_group",
-        lambda: SimpleNamespace(
-            ulysses_world_size=2,
-            ulysses_rank=1,
-            ulysses_group=None,
-        ),
-    )
-    attention = module.LingBotCrossAttention(dim=8, num_heads=4)
-    full = torch.arange(8, dtype=torch.float32).reshape(1, 1, 4, 2)
-    shard = attention.shard_kv_heads(full)
-    torch.testing.assert_close(shard, full[:, :, 2:], rtol=0, atol=0)
-    assert shard.is_contiguous()
-    assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
-
-
-@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention", "LingBotCrossAttention"])
+@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention"])
 def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
     module = _load_module()
     monkeypatch.setattr(
@@ -621,3 +613,35 @@ def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
     )
     with pytest.raises(ValueError, match="heads must be divisible"):
         getattr(module, attention_class)(dim=8, num_heads=4)
+
+
+def test_supplied_camera_modulation_matches_computing_it_inline() -> None:
+    """Passing a prebuilt modulation must be arithmetically identical to letting the block build its own."""
+    module = _load_module()
+    torch.manual_seed(0)
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+    camera = torch.randn(1, 3, 4)
+
+    scale, shift = block.camera_modulation(camera)
+    inline_scale, inline_shift = block.camera_modulation(camera)
+
+    torch.testing.assert_close(scale, inline_scale, rtol=0, atol=0)
+    torch.testing.assert_close(shift, inline_shift, rtol=0, atol=0)
+
+
+def test_the_camera_modulation_reads_only_the_camera_tokens() -> None:
+    """The reuse window is only sound because the modulation depends on nothing else that varies per forward."""
+    module = _load_module()
+    torch.manual_seed(0)
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+    camera = torch.randn(1, 3, 4)
+
+    first = block.camera_modulation(camera)
+    # Same camera, different everything else a forward would carry: the modulation cannot move.
+    second = block.camera_modulation(camera.clone())
+
+    torch.testing.assert_close(first[0], second[0], rtol=0, atol=0)
+    torch.testing.assert_close(first[1], second[1], rtol=0, atol=0)
+    # A different camera must move it, or the cache key would be hiding a real dependency.
+    other = block.camera_modulation(camera + 1)
+    assert not torch.equal(first[0], other[0])

tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py

@@ -837,6 +837,9 @@ def _rollout(model, mode, dtype, batch):
             frames_per_block=_FRAMES,
             max_scratch_tokens_per_branch=_FRAMES * _TOKENS_PER_FRAME,
             cross_attention_lengths={"text": 5},
+            # Cross-attention keeps every local head on every SP rank, so its pool is wider than the
+            # self-attention share above.
+            cross_attention_kv_heads={"text": model.blocks[0].cross_attn.num_local_heads},
             device=device,
         )
         state = ARDiffusionKVState(kv, "numeric", {"main": kv.begin_request("numeric")}, num_layers=_LAYERS)
@@ -896,7 +899,6 @@ def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
     from vllm_omni.diffusion.distributed.parallel_state import (
         destroy_distributed_env,
         destroy_model_parallel,
-        get_sp_group,
         init_distributed_environment,
         initialize_model_parallel,
     )
@@ -948,11 +950,10 @@ def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
                         bound = 1e-5 if dtype == torch.float32 else 1e-2
                         error = (result - expected).double().norm() / expected.double().norm()
                         assert error <= bound, f"rank={rank}, TP{tp} SP{sp}: relative L2 {error.item():.3g} > {bound}"
-                        local_heads = _HEADS // tp // sp
-                        first = (
-                            get_tensor_model_parallel_rank() * (_HEADS // tp)
-                            + get_sp_group().ulysses_rank * local_heads
-                        )
+                        # Cross-attention K/V is replicated across the Ulysses group, not split over it, so
+                        # every SP rank holds its whole TP shard of heads. Only TP slices the cache now.
+                        local_heads = _HEADS // tp
+                        first = get_tensor_model_parallel_rank() * local_heads
                         for (k, v), (ek, ev) in zip(cross, expected_cross, strict=True):
                             tolerance = 1e-5 if dtype == torch.float32 else 2e-2
                             torch.testing.assert_close(
@@ -1033,3 +1034,56 @@ def test_temporal_rope_extension_keeps_spatial_limit():
         model._rotary_embedding(
             frames=3, height=17, width=1, start_frame=16, dtype=torch.float32, device=torch.device("cpu")
         )
+
+
+@pytest.mark.cpu
+def test_the_camera_injector_runs_once_per_held_cache_and_gives_the_same_output() -> None:
+    """Two forwards sharing one caller-held cache: the injector runs once per block, outputs unchanged."""
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=2).eval()
+    hidden_states = torch.randn(1, 36, 1, 4, 4)
+    timestep = torch.tensor([1.0])
+    encoder_hidden_states = torch.randn(1, 3, 6)
+    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+
+    def run(cache, camera_cache=None):
+        return model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            camera_hidden_states,
+            cache=cache,
+            start_frame=0,
+            update_cache=False,
+            camera_modulation_cache=camera_cache,
+        )
+
+    # Baseline: no cache, so every forward rebuilds the injector in every block.
+    uncached_first = run(_cache(module, model))
+    uncached_second = run(_cache(module, model))
+    injector_calls = [attention_tests._record_outputs(block.cam_injector_layer1) for block in model.blocks]
+    held = module.CameraModulationCache()
+    cached_first = run(_cache(module, model), held)
+    cached_second = run(_cache(module, model), held)
+
+    # One build per block for both forwards, instead of one per block per forward.
+    assert [len(calls) for calls in injector_calls] == [1, 1]
+    torch.testing.assert_close(cached_first, uncached_first, rtol=0, atol=0)
+    torch.testing.assert_close(cached_second, uncached_second, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_timestep_projection_is_staged_into_one_static_buffer_per_shape() -> None:
+    """Two forwards' projections land in the same buffer; the blocks never see the fresh per-forward tensor."""
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1).eval()
+    first = torch.randn(1, 8, 6, model.dim)
+    second = torch.randn(1, 8, 6, model.dim)
+
+    staged_first = model._stage_timestep_projection(first)
+    assert staged_first is not first and torch.equal(staged_first, first)
+    staged_second = model._stage_timestep_projection(second)
+    assert staged_second is staged_first and torch.equal(staged_second, second)
+    # A different shape gets its own buffer; the first one is kept.
+    other = model._stage_timestep_projection(torch.randn(1, 4, 6, model.dim))
+    assert other is not staged_first and len(model._timestep_projection_buffers) == 2

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -4,11 +4,13 @@
 from __future__ import annotations
 
 import gc
+import os
 import weakref
 from pathlib import Path
 from threading import Lock
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
@@ -29,6 +31,7 @@ from vllm_omni.diffusion.models.lingbot_world.camera import (
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.experimental.ar_diffusion.capability import supports_chunk_step_grouping
 from vllm_omni.experimental.ar_diffusion.tick_protocol import (
     ARDiffusionControlInput,
     ARDiffusionTickRequest,
@@ -127,6 +130,10 @@ class _RecordingTransformer(nn.Module):
             sink_size=3,
         )
         self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()])
+        for block in self.blocks:
+            # The spec reads head geometry from both attentions; cross-attention keeps every local head.
+            block.self_attn = SimpleNamespace(num_sp_heads=2)
+            block.cross_attn = SimpleNamespace(num_local_heads=2)
         self.calls: list[dict] = []
         self.cache_allocations: list[dict] = []
         self.raise_on_call = raise_on_call
@@ -146,6 +153,7 @@ class _RecordingTransformer(nn.Module):
             "cache_id": id(kwargs["cache"]),
             "start_frame": kwargs["start_frame"],
             "update_cache": kwargs["update_cache"],
+            "camera_cache": kwargs.get("camera_modulation_cache"),
         }
         self.calls.append(call)
         if self.raise_on_call == len(self.calls):
@@ -534,14 +542,15 @@ 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}
-    # One condition block, committed + pending encoder caches, and decoder state.
+    # One condition block, committed + pending encoder caches, the session's
+    # cached prompt embedding (512 tokens x text_dim 8 x fp32), and decoder state.
     # The stub encoder retains 2 RGB frames at 16x16 and 2 feature frames at 2x2.
     encoder_bytes = (2 * 3 * 16 * 16 + 2 * 16 * 2 * 2) * 4
-    assert spec.model_owned_state_bytes_per_session == 960 + 2 * encoder_bytes + 2_421_248
+    assert spec.model_owned_state_bytes_per_session == 960 + 384 + 2 * encoder_bytes + 2_421_248
 
 
 def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
-    """Admission includes encoder state even when the decoder cannot stream."""
+    """Admission includes encoder state and the prompt embedding even when the decoder cannot stream."""
     module = _load_pipeline_module()
 
     def spec_pipeline():
@@ -560,13 +569,29 @@ def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
     # than the image condition alone could account for.
     assert widened - small > 2_000_000
 
-    # Removing decoder support leaves the condition and both encoder histories.
+    # Removing decoder support leaves the condition, both encoder histories and
+    # the prompt-embedding cache.
     bare = spec_pipeline()
     for attribute in ("decoder", "post_quant_conv"):
         if hasattr(bare.vae, attribute):
             delattr(bare.vae, attribute)
     bare.vae._cached_conv_counts.pop("decoder")
-    assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 2 * 6_656
+    assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 384 + 2 * 6_656
+
+
+@pytest.mark.parametrize("sp_size", [1, 4])
+def test_session_admission_budgets_camera_modulation_per_rank(sp_size: int) -> None:
+    pipeline = _pipeline(_load_pipeline_module())
+    pipeline._ar_height, pipeline._ar_width = 480, 832
+    pipeline.od_config.parallel_config.ulysses_degree = sp_size
+    pipeline.transformer.config.num_layers = 40
+    pipeline.transformer.config.num_attention_heads = 40
+    pipeline.transformer.config.attention_head_dim = 128
+    pipeline.transformer._dtype = torch.bfloat16
+    forty_layers = pipeline.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
+    pipeline.transformer.config.num_layers = 1
+    one_layer = pipeline.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
+    assert forty_layers - one_layer == 2 * 39 * (4680 // sp_size) * 5120 * 2
 
 
 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None:
@@ -1672,7 +1697,6 @@ def test_typed_ticks_generate_one_global_block_and_return_standard_metadata(
         "applied_event_ids": [1],
     }
     assert pipeline._ar_sessions["world-1"].next_chunk_index == 2
-    assert pipeline._ar_sessions["world-1"].prompt == "enter the snowy valley"
     assert pipeline._ar_sessions["world-1"].camera_tail is not None
     assert pipeline._ar_sessions["world-1"].camera_tail.poses.shape == (1, 4, 4)
     expected_generator = torch.Generator(device="cpu").manual_seed(17)
@@ -1846,8 +1870,11 @@ def test_typed_ticks_keep_bounded_encoder_state_beyond_ten_chunks(
         pipeline(_request(sampling=sampling))
 
     assert pipeline.vae.encode_inputs == []
-    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 32
+    # The stub encoder's history settles after its second block; the nine
+    # blocks after that reuse the settled condition without an encode.
+    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 5
     assert all(torch.count_nonzero(video) == 0 for video in pipeline.vae.encoder.inputs[1:])
+    assert pipeline._ar_sessions["world-1"].condition_fixed_point is not None
     encoder_cache = pipeline._ar_sessions["world-1"].encoder_cache
     assert sum(t.numel() * t.element_size() for t in encoder_cache) == pipeline._condition_encoder_cache_bytes()
     assert pipeline.vae._enc_feat_map == ["module-owned"]
@@ -1943,6 +1970,76 @@ def test_condition_continues_real_causal_vae_beyond_initial_horizon(
     assert pipeline._ar_sessions == {}
 
 
+def test_condition_encoder_stops_at_its_fixed_point_and_restarts_with_the_session(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Once a block returns the history it was given, later blocks reuse the condition without encoding.
+
+    The whole-clip encode is the reference for every block, so reusing the
+    stored condition is only allowed to be a speedup, never a change.
+    """
+    pipeline = _tiny_condition_pipeline(monkeypatch)
+    pipeline._ar_diffusion_kv_state = object()
+    blocks = 12
+    inputs = pipeline._parse_request(_request(sampling=_SamplingParams(num_frames=(blocks * 3 - 1) * 4 + 1)))
+    with torch.inference_mode():
+        expected = pipeline._prepare_condition(inputs, dtype=torch.float32)
+    encoder_calls = []
+    original_encoder = pipeline.vae.encoder.forward
+
+    def encoder(*args, **kwargs):
+        encoder_calls.append(len(seen))
+        return original_encoder(*args, **kwargs)
+
+    seen: list[torch.Tensor] = []
+
+    def generate_block(**kwargs):
+        index = len(seen) % blocks
+        torch.testing.assert_close(
+            kwargs["condition"], expected[:, :, index * 3 : (index + 1) * 3], rtol=1e-6, atol=1e-6
+        )
+        seen.append(kwargs["condition"])
+        return torch.zeros_like(kwargs["condition"][:, :16])
+
+    monkeypatch.setattr(pipeline.vae.encoder, "forward", encoder)
+    monkeypatch.setattr(pipeline, "_generate_block", generate_block)
+
+    def tick(index):
+        pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index))))
+
+    settled_at = None
+    for index in range(blocks):
+        tick(index)
+        state = pipeline._ar_sessions["world-1"]
+        if settled_at is None and state.condition_fixed_point is not None:
+            settled_at = index
+        assert state.pending_encoder_cache is None and state.encoder_cache is not None
+    assert settled_at is not None and 0 < settled_at < blocks - 1, settled_at
+    # Three encoder passes per encoded block, none once the fixed point is stored.
+    assert encoder_calls == [index for index in range(settled_at + 1) for _ in range(3)]
+    # The stored condition is the session's own, not shared with the callers.
+    assert all(seen[settled_at] is not later for later in seen[settled_at + 1 :])
+
+    # A session restarted from chunk 0 encodes again from the source image.
+    pipeline.reset_ar_diffusion_session("world-1")
+    tick(0)
+    state = pipeline._ar_sessions["world-1"]
+    assert state.condition_fixed_point is None
+    assert encoder_calls[-3:] == [blocks] * 3
+
+
+def test_condition_encoder_fixed_point_is_released_with_the_session(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = _tiny_condition_pipeline(monkeypatch)
+    pipeline._ar_diffusion_kv_state = object()
+    monkeypatch.setattr(pipeline, "_generate_block", lambda **kwargs: torch.zeros_like(kwargs["condition"][:, :16]))
+    for index in range(8):
+        pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index))))
+    state = pipeline._ar_sessions["world-1"]
+    assert state.condition_fixed_point is not None
+    pipeline.close_ar_diffusion_session("world-1")
+    assert state.condition_fixed_point is None and state.encoder_cache is None
+
+
 def test_condition_encoder_interleaving_retry_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
     pipeline = _tiny_condition_pipeline(monkeypatch)
     pipeline._ar_diffusion_kv_state = object()
@@ -2209,6 +2306,40 @@ def _stepwise_chunks(pipeline, state, ar_state):
             yield output
 
 
+def test_stepwise_block_shares_one_camera_cache_across_its_five_forwards() -> None:
+    """The four probes and the commit of one stepwise block reuse one camera cache; the next block gets its own.
+
+    Request mode creates one cache per loop; the stepwise path's forwards are
+    separate scheduler steps, so the pipeline holds the block's cache and
+    passes it to each one.
+    """
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "latent"
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert len(outputs) == 2 and len(transformer.calls) == 10
+    caches = [call["camera_cache"] for call in transformer.calls]
+    assert all(c is not None for c in caches)
+    assert len({id(c) for c in caches[:5]}) == 1 and len({id(c) for c in caches[5:]}) == 1
+    assert caches[0] is not caches[5]
+    # The block's cache does not outlive its commit.
+    assert "camera_cache" not in state.extra
+
+
+def test_pipeline_declares_chunk_step_grouping() -> None:
+    """A realtime block's probes and commit may run without a scheduler cycle between them."""
+    module = _load_pipeline_module()
+    assert module.LingBotWorldCausalDMDPipeline.supports_chunk_step_grouping is True
+    assert supports_chunk_step_grouping(_pipeline(module))
+
+
 def test_pipeline_declares_step_execution_support() -> None:
     module = _load_pipeline_module()
     pipeline = _pipeline(module)
@@ -2337,7 +2468,10 @@ def test_stepwise_matches_tick_transformer_trace_and_latents(num_chunks) -> None
         torch.testing.assert_close(step_call["camera_hidden_states"], tick_call["camera_hidden_states"])
     assert stepwise_state.extra["condition"].shape[2] == 3
     assert pipeline.vae.encode_inputs == []
-    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (num_chunks * 3 - 1)) * 2
+    # The stub encoder's history settles after its second block, so both
+    # modes encode two blocks and reuse the settled condition after that.
+    encoded_blocks = min(num_chunks, 2)
+    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (encoded_blocks * 3 - 1)) * 2
     assert stepwise_fake.commits == ["main"] * num_chunks
     assert len(stepwise_outputs) == num_chunks and stepwise_outputs[-1].finished
     assert pipeline.vae.decode_inputs == []
@@ -2406,21 +2540,16 @@ def test_stepwise_emits_latents_without_touching_the_vae() -> None:
     assert pipeline.vae.decode_inputs == []
 
 
-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
+def _capture_pixel_conversion(monkeypatch) -> list[tuple[int, ...]]:
+    """Record each decoded block reaching the uint8 conversion, and name its output."""
+    module = _load_pipeline_module()
+    processed: list[tuple[int, ...]] = []
 
-    processed: list[tuple[tuple[int, ...], str]] = []
+    def uint8_frames(video):
+        processed.append(tuple(video.shape))
+        return f"frames-{len(processed)}"
 
-    class VideoProcessor:
-        def __init__(self, *, vae_scale_factor):
-            assert vae_scale_factor == 8
-
-        def postprocess_video(self, video, *, output_type):
-            processed.append((tuple(video.shape), output_type))
-            return f"frames-{len(processed)}"
-
-    monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+    monkeypatch.setattr(module, "_uint8_frames", uint8_frames)
     return processed
 
 
@@ -2434,7 +2563,7 @@ def _streaming_pipeline(module):
 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)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2443,7 +2572,7 @@ 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)
 
-    assert [output_type for _, output_type in processed] == ["np", "np"]
+    assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
     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]
@@ -2457,7 +2586,7 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
     the same latents would: 9 frames, then 12.
     """
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2466,7 +2595,7 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
     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)]
+    assert 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
@@ -2511,7 +2640,7 @@ def test_prepare_next_chunk_rejects_live_camera_with_scripted_paths() -> None:
 def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) -> None:
     """Camera timelines must track the streaming decoder's 9-then-12 media counts."""
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2536,7 +2665,7 @@ def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) ->
 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)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(request_id="req-stream", num_frames=21)
     state.sampling.output_type = "np"
@@ -2560,7 +2689,7 @@ def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
 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)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(request_id="req-reset", num_frames=21)
     state.sampling.output_type = "np"
@@ -2584,7 +2713,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
     interleaving the scheduler actually produces.
     """
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     states = []
     for request_id in ("req-a", "req-b"):
@@ -2617,7 +2746,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
     assert pipeline.vae.decoder.first_chunk_flags == (
         [True, False, False] + [True, False, False] + [False] * 3 + [False] * 3
     )
-    assert [shape for shape, _ in processed] == [
+    assert processed == [
         (1, 3, 9, 16, 16),
         (1, 3, 9, 16, 16),
         (1, 3, 12, 16, 16),
@@ -2628,7 +2757,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
 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)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(request_id="req-boom", num_frames=21)
     state.sampling.output_type = "np"
@@ -2662,7 +2791,7 @@ def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None:
     never reaches, so this asserts on what ``post_quant_conv`` received.
     """
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2695,7 +2824,7 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
     perf regression in the streaming decoder would be too.
     """
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _pipeline(
         module,
         transformer=_RecordingTransformer(),
@@ -2723,21 +2852,18 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
 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)
+    processed = _capture_pixel_conversion(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)
-    outputs = [
-        pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), output_type="np", session_id="tiled")
-        for _ in range(2)
-    ]
+    outputs = [pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), session_id="tiled") for _ in range(2)]
 
     # 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 processed == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)]
     assert outputs == ["frames-1", "frames-2"]
 
 
@@ -2749,7 +2875,7 @@ def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monke
     on a configuration that would never have tiled.
     """
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(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)
@@ -2761,7 +2887,7 @@ def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monke
 
     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)]
+    assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
 
 
 def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
@@ -2904,3 +3030,193 @@ def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None:
 
     assert encoded == ["move through the room"]
     assert states[0].prompt_embeds is states[1].prompt_embeds
+
+
+def test_session_text_caches_publish_key_only_after_a_successful_build() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    state = module._LingBotARSessionState()
+    invalidations: list[bool] = []
+    fail_next = {"raise": False}
+
+    def fake_text_caches(prompt_embeds, *, invalidate):
+        invalidations.append(invalidate)
+        if fail_next["raise"]:
+            fail_next["raise"] = False
+            raise RuntimeError("injected failure inside _ar_text_caches")
+        return ["caches"]
+
+    pipeline._ar_text_caches = fake_text_caches
+    embeds = torch.ones(1, 512, 8)
+    caches = lambda prompt, length=512, e=embeds: pipeline._session_text_caches(  # noqa: E731
+        e, prompt=prompt, max_sequence_length=length, session_state=state
+    )
+
+    # First build of a session: runner state is empty -> no invalidation, key published.
+    caches("A")
+    assert invalidations == [False] and state.text_cache_key == ("A", 512, str(embeds.dtype))
+    assert state.text_cache_dirty is False
+    # Same inputs: reuse without invalidation.
+    caches("A")
+    assert invalidations == [False, False]
+    # Regression (Codex review): B's build fails after invalidating -> no key is
+    # left behind, so retrying A must invalidate again instead of consuming
+    # partial/mixed K/V.
+    fail_next["raise"] = True
+    with pytest.raises(RuntimeError, match="injected failure"):
+        caches("B")
+    assert state.text_cache_key is None and state.text_cache_dirty is True
+    caches("A")
+    assert invalidations[-2:] == [True, True] and state.text_cache_key == ("A", 512, str(embeds.dtype))
+    assert state.text_cache_dirty is False
+    # The key covers every embedding-defining input: same text, other length or dtype -> rebuild.
+    caches("A", length=256)
+    caches("A", e=embeds.to(torch.bfloat16))
+    assert invalidations[-2:] == [True, True]
+
+
+# ---------------------------------------------------------------------------
+# a source image is decoded once while the file behind it is unchanged
+# ---------------------------------------------------------------------------
+
+
+def _write_png(path, colour=(10, 20, 30)):
+    Image.new("RGB", (64, 64), colour).save(path)
+    return path
+
+
+def _counting_decode(module):
+    """Replace the file decode with a counter, so the tests measure decodes rather than pixels."""
+    decoded: list[str] = []
+    original = module._decode_source_image_file
+
+    def decode(path):
+        decoded.append(os.fspath(path))
+        return Image.new("RGB", (64, 64), (1, 2, 3))
+
+    module._decode_source_image_file = decode
+    return decoded, original
+
+
+def test_the_same_source_image_is_decoded_once(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    path = _write_png(tmp_path / "frame.png")
+    decoded, original = _counting_decode(module)
+    try:
+        first = module._load_source_image(path)
+        second = module._load_source_image(path)
+        assert decoded == [os.fspath(os.path.realpath(path))] or len(decoded) == 1
+        # Every caller owns its copy: mutating one must not reach the next caller or the cache.
+        assert first is not second
+        first.putpixel((0, 0), (255, 255, 255))
+        assert module._load_source_image(path).getpixel((0, 0)) == (1, 2, 3)
+        assert len(decoded) == 1
+    finally:
+        module._decode_source_image_file = original
+        module._decode_source_image_fingerprinted.cache_clear()
+
+
+def test_a_rewritten_source_image_is_decoded_again(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    path = _write_png(tmp_path / "frame.png")
+    decoded, original = _counting_decode(module)
+    try:
+        module._load_source_image(path)
+        module._load_source_image(path)
+        assert len(decoded) == 1
+        # Same path, different bytes: the mtime/size in the key must force a fresh decode.
+        os.utime(path, ns=(0, 0))
+        _write_png(path, colour=(200, 100, 50))
+        module._load_source_image(path)
+        assert len(decoded) == 2
+    finally:
+        module._decode_source_image_file = original
+        module._decode_source_image_fingerprinted.cache_clear()
+
+
+def test_an_image_rewritten_during_the_decode_is_never_cached(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    path = _write_png(tmp_path / "frame.png")
+    decoded: list[str] = []
+    original = module._decode_source_image_file
+
+    def decode_then_rewrite(p):
+        decoded.append(os.fspath(p))
+        _write_png(path, colour=(len(decoded), 0, 0))
+        os.utime(path, ns=(len(decoded) * 1_000_000, len(decoded) * 1_000_000))
+        return Image.new("RGB", (64, 64), (1, 2, 3))
+
+    module._decode_source_image_file = decode_then_rewrite
+    try:
+        module._load_source_image(path)
+        # The entry is keyed by the fingerprint taken before the decode, which the rewrite invalidated: the
+        # next call must decode again rather than serve the image that was decoded mid-rewrite.
+        module._load_source_image(path)
+        assert len(decoded) == 2
+    finally:
+        module._decode_source_image_file = original
+        module._decode_source_image_fingerprinted.cache_clear()
+
+
+def test_an_unstattable_path_keeps_upstream_validation(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    missing = tmp_path / "does-not-exist.png"
+    # The uncached path owns the error contract; caching must not change which exception a caller sees.
+    with pytest.raises(ValueError):
+        module._load_source_image(missing)
+    assert module._decode_source_image_fingerprinted.cache_info().currsize == 0
+
+
+def test_the_text_cache_is_sized_by_cross_attention_heads_not_the_self_attention_share() -> None:
+    """Cross-attention replicates heads per rank, so its pool must not inherit the self-attention head count."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1)
+    pipeline.transformer.blocks[0].cross_attn = SimpleNamespace(num_local_heads=4)
+
+    spec = pipeline.ar_diffusion_kv_cache_spec()
+
+    assert spec.num_kv_heads == 1
+    assert spec.cross_attention_kv_heads == {"text": 4}
+    assert spec.cross_attention_lengths == {"text": 512}
+
+
+@pytest.mark.parametrize("output_type", ["np", "pil"])
+def test_streamed_chunk_is_uint8_frames_whatever_pixel_output_type_was_asked(monkeypatch, output_type) -> None:
+    """A realtime tick delivers (F, H, W, 3) uint8 for every pixel output type; only latent differs."""
+    module = _load_pipeline_module()
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = output_type
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        outputs = _run_stepwise(pipeline, state)
+
+    videos = [output.output["payload"]["video"] for output in outputs]
+    assert [type(video) for video in videos] == [np.ndarray, np.ndarray]
+    assert [video.dtype for video in videos] == [np.uint8, np.uint8]
+    assert [video.shape for video in videos] == [(9, 16, 16, 3), (12, 16, 16, 3)]
+
+
+def test_uint8_frames_match_the_pil_conversion_byte_for_byte() -> None:
+    """The device-side conversion is the PIL path's arithmetic, so the bytes are the same."""
+    from diffusers.video_processor import VideoProcessor
+
+    module = _load_pipeline_module()
+    generator = torch.Generator().manual_seed(7)
+    # Cover the clamp on both sides and values that sit exactly on rounding edges.
+    video = torch.rand(1, 3, 5, 8, 8, generator=generator) * 2.4 - 1.2
+    video[0, :, 0, 0, :4] = torch.tensor([-1.0, 1.0, 0.0, 1 / 255 - 1.0])
+
+    reference = VideoProcessor(vae_scale_factor=8).postprocess_video(video, output_type="pil")[0]
+    expected = np.stack([np.asarray(frame) for frame in reference])
+
+    frames = module._uint8_frames(video)
+
+    assert frames.dtype == np.uint8 and frames.shape == (5, 8, 8, 3)
+    assert frames.flags["C_CONTIGUOUS"]
+    np.testing.assert_array_equal(frames, expected)