Add caching mechanism for completed AV blocks in streaming model. Introduce CompletedBlockPieceCache class to manage reusable blocks during strict-window replay, enhancing efficiency. Update profiling hooks for cache operations and refine README documentation to reflect changes in block handling and replay strategies.

README.md

@@ -477,7 +477,9 @@ two practical changes for this dataset:
 1. the current 0--2 second partial block is evaluated speculatively every
    0.2 seconds and discarded unless it reaches a complete block boundary;
 2. the KV history is bounded and rebuilt from recent complete blocks rather
-   than accumulating for an entire multi-minute session.
+   than accumulating for an entire multi-minute session. Complete-block AV
+   embeddings are retained in a bounded ring, so rolling rebuilds rerun only
+   the Thinker and do not rerun the processor, audio tower, or video tower.
 
 It should therefore be described as **bounded Qwen-style open streaming**, not
 as bit-for-bit reproduction of an unbounded official serving session.
@@ -522,7 +524,10 @@ prediction is emitted every 0.2 seconds.
 budget it rebuilds the open cache from the latest nine completed blocks
 (18 seconds) when a new block starts, then adds the current 0--2 second tail.
 Its effective history therefore varies from 18 to 20 seconds and never drops
-to an empty 0.2-second context at a 20-second boundary.
+to an empty 0.2-second context at a 20-second boundary. The nine blocks'
+prepared AV embeddings and local TMRoPE coordinates are cached once; strict
+window replay still recomputes the Thinker KV so retained blocks no longer
+carry indirect influence from observations that fell out of the window.
 
 For a short native-stream fidelity check, set `--memory-seconds` to at least
 the tested clip duration. For example, the options

profile_streaming_bottleneck_gpu7.py

@@ -231,6 +231,16 @@ def _install_profiling_hooks() -> None:
         "_cache_from_intervals",
         "rolling_cache_rebuild",
     )
+    _patch_function(
+        core.CompletedBlockPieceCache,
+        "put",
+        "completed_block_cache_store",
+    )
+    _patch_function(
+        core.CompletedBlockPieceCache,
+        "get_many",
+        "completed_block_cache_lookup",
+    )
     _patch_function(torch, "save", "output_save")
 
     original_extract = core.OpenStreamReplayExtractor.extract
@@ -423,7 +433,7 @@ def _print_summary(
     bottleneck = max(totals, key=totals.get)
     advice = {
         "media_frames": "bulk-decode/cache the selected 2-FPS resized frames",
-        "processor": "cache processor-ready AV tensors for every causal prefix",
+        "processor": "prefetch or microbatch the ten causal prefixes per 2-s block",
         "audio_tower": "microbatch the ten 0.2-s prefixes in each 2-s block",
         "video_tower": "cache unique 2-FPS visual inputs and batch equal shapes",
         "tmrope": "cache position IDs for the ten repeated prefix shapes",

streaming_turn_model.py

@@ -33,6 +33,7 @@ import json
 import logging
 import math
 import subprocess
+from collections import OrderedDict
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Literal
@@ -176,6 +177,102 @@ class OpenMediaPieces:
     closure_position_ids: torch.LongTensor
 
 
+@dataclass(frozen=True)
+class CachedOpenMediaBlock:
+    """Prepared model inputs for one complete source-aligned AV block."""
+
+    start: float
+    end: float
+    pieces: OpenMediaPieces
+
+
+class CompletedBlockPieceCache:
+    """Bounded ring of complete blocks reusable during strict-window replay.
+
+    The modality encoders operate on each native block independently of the
+    Thinker history.  Their embeddings and local TMRoPE coordinates can
+    therefore be reused when the strict rolling window rebuilds the Thinker
+    KV.  Only the Thinker replay must be recomputed.
+    """
+
+    def __init__(
+        self,
+        *,
+        origin: float,
+        block_seconds: float,
+        max_blocks: int,
+    ) -> None:
+        if block_seconds <= 0:
+            raise ValueError("block_seconds must be positive")
+        if max_blocks < 0:
+            raise ValueError("max_blocks must be >= 0")
+        self.origin = origin
+        self.block_seconds = block_seconds
+        self.max_blocks = max_blocks
+        self._blocks: OrderedDict[int, CachedOpenMediaBlock] = OrderedDict()
+
+    def _index(self, start: float) -> int:
+        ratio = (start - self.origin) / self.block_seconds
+        index = round(ratio)
+        if not math.isclose(ratio, index, rel_tol=0.0, abs_tol=1e-6):
+            raise ValueError(
+                f"Block start {start:g} is not aligned to origin "
+                f"{self.origin:g} at {self.block_seconds:g}s"
+            )
+        return index
+
+    def put(self, block: CachedOpenMediaBlock) -> None:
+        duration = block.end - block.start
+        if not math.isclose(
+            duration,
+            self.block_seconds,
+            rel_tol=0.0,
+            abs_tol=1e-6,
+        ):
+            raise ValueError(
+                f"Cached block duration is {duration:g}s, expected "
+                f"{self.block_seconds:g}s"
+            )
+        if self.max_blocks == 0:
+            return
+        index = self._index(block.start)
+        self._blocks[index] = block
+        self._blocks.move_to_end(index)
+        while len(self._blocks) > self.max_blocks:
+            self._blocks.popitem(last=False)
+
+    def get_many(
+        self,
+        intervals: list[tuple[float, float]],
+    ) -> list[CachedOpenMediaBlock]:
+        blocks = []
+        for start, end in intervals:
+            index = self._index(start)
+            try:
+                block = self._blocks[index]
+            except KeyError as exc:
+                available = list(self._blocks)
+                raise RuntimeError(
+                    "Completed-block embedding cache miss for interval "
+                    f"[{start:g}, {end:g}); available block indices="
+                    f"{available}"
+                ) from exc
+            if not (
+                math.isclose(block.start, start, rel_tol=0.0, abs_tol=1e-6)
+                and math.isclose(block.end, end, rel_tol=0.0, abs_tol=1e-6)
+            ):
+                raise RuntimeError(
+                    "Completed-block embedding cache interval mismatch: "
+                    f"requested [{start:g}, {end:g}), cached "
+                    f"[{block.start:g}, {block.end:g})"
+                )
+            blocks.append(block)
+        return blocks
+
+    def __len__(self) -> int:
+        return len(self._blocks)
+
+
 class FixedGridAV:
     """Decode only frames on a prepared MP4's fixed global grid.
 
@@ -1182,11 +1279,25 @@ class OpenStreamReplayExtractor:
         audio_path: str | Path | None,
         intervals: list[tuple[float, float]],
         use_audio_in_video: bool,
+        prepared_blocks: list[CachedOpenMediaBlock] | None = None,
     ) -> tuple[Any, int, int, torch.LongTensor, torch.LongTensor]:
-        """Build an open payload cache and expose its local TMRoPE base."""
+        """Build an open payload cache and expose its local TMRoPE base.
+
+        ``prepared_blocks`` is used only by the production rolling-window
+        path.  Correctness/reference calls omit it and independently rerun the
+        processor and modality encoders from the source media.
+        """
 
         if not intervals:
             raise ValueError("At least one interval is required")
+        if (
+            prepared_blocks is not None
+            and len(prepared_blocks) != len(intervals)
+        ):
+            raise ValueError(
+                "prepared block count does not match replay intervals: "
+                f"{len(prepared_blocks)} vs {len(intervals)}"
+            )
         media_cache = None
         next_offset = 0
         closure_ids = None
@@ -1194,12 +1305,35 @@ class OpenStreamReplayExtractor:
         payload_base = None
         positions_per_chunk = _tmrope_positions_per_chunk(self.model)
         for index, interval in enumerate(intervals):
+            if prepared_blocks is None:
                 pieces = self._prepare_pieces(
                     video_path=video_path,
                     audio_path=audio_path,
                     interval=interval,
                     use_audio_in_video=use_audio_in_video,
                 )
+            else:
+                prepared = prepared_blocks[index]
+                if not (
+                    math.isclose(
+                        prepared.start,
+                        interval[0],
+                        rel_tol=0.0,
+                        abs_tol=1e-6,
+                    )
+                    and math.isclose(
+                        prepared.end,
+                        interval[1],
+                        rel_tol=0.0,
+                        abs_tol=1e-6,
+                    )
+                ):
+                    raise ValueError(
+                        "prepared block/replay interval mismatch: "
+                        f"prepared [{prepared.start:g}, {prepared.end:g}), "
+                        f"requested [{interval[0]:g}, {interval[1]:g})"
+                    )
+                pieces = prepared.pieces
             if payload_base is None:
                 payload_base = int(pieces.payload_position_ids.min().item())
             payload_offset = payload_base + index * positions_per_chunk
@@ -1250,6 +1384,24 @@ class OpenStreamReplayExtractor:
         positions_per_chunk = _tmrope_positions_per_chunk(self.model)
         states: list[torch.Tensor] = []
         checks: list[dict[str, Any]] = []
+        replay_block_count = max(
+            0,
+            _require_integer_ratio(
+                self.config.memory_seconds,
+                self.config.native_block_seconds,
+                "memory / native block",
+            )
+            - 1,
+        )
+        needs_rolling_rebuild = (
+            self.config.memory_seconds
+            < segment_end - segment_start - 1e-6
+        )
+        completed_blocks = CompletedBlockPieceCache(
+            origin=segment_start,
+            block_seconds=self.config.native_block_seconds,
+            max_blocks=replay_block_count if needs_rolling_rebuild else 0,
+        )
 
         for local_step, current_time in enumerate(times, start=1):
             while current_time > block_start + self.config.native_block_seconds + 1e-6:
@@ -1272,6 +1424,10 @@ class OpenStreamReplayExtractor:
                     block_start,
                     self.config.native_block_seconds,
                 )
+                if history_intervals:
+                    prepared_history = completed_blocks.get_many(
+                        history_intervals
+                    )
                     (
                         committed_cache,
                         _,
@@ -1283,7 +1439,13 @@ class OpenStreamReplayExtractor:
                         audio_path=audio_path,
                         intervals=history_intervals,
                         use_audio_in_video=use_audio_in_video,
+                        prepared_blocks=prepared_history,
                     )
+                else:
+                    # A one-block memory budget intentionally retains no
+                    # completed history before the current partial block.
+                    committed_cache = None
+                    payload_base = None
 
             pieces = self._prepare_pieces(
                 video_path=video_path,
@@ -1361,6 +1523,19 @@ class OpenStreamReplayExtractor:
             if math.isclose(current_time, block_end, rel_tol=0.0, abs_tol=1e-6):
                 # The cache ends at AV payload, before media/chat closure.
                 committed_cache = media_cache
+                if math.isclose(
+                    block_end - block_start,
+                    self.config.native_block_seconds,
+                    rel_tol=0.0,
+                    abs_tol=1e-6,
+                ):
+                    completed_blocks.put(
+                        CachedOpenMediaBlock(
+                            start=block_start,
+                            end=block_end,
+                            pieces=pieces,
+                        )
+                    )
                 block_start = block_end
 
             progress.set_postfix_str(
@@ -1739,6 +1914,9 @@ def main() -> None:
                 "temporary": "partial_payload_plus_media_closure_plus_readout",
                 "tmrope_positions_per_chunk": _tmrope_positions_per_chunk(model),
                 "memory": "rolling_complete_blocks_plus_current_partial_block",
+                "rolling_replay_inputs": (
+                    "reuse_cached_complete_block_av_embeddings"
+                ),
                 "memory_seconds_range": [
                     config.memory_seconds - config.native_block_seconds,
                     config.memory_seconds,

streaming_turn_model_directly.py

@@ -1130,6 +1130,11 @@ def main() -> None:
                     if persistent_memory
                     else "rolling_complete_blocks_plus_current_partial_block"
                 ),
+                "rolling_replay_inputs": (
+                    None
+                    if persistent_memory
+                    else "reuse_cached_complete_block_av_embeddings"
+                ),
             }
             if args.mode == "open-stream-replay"
             else None