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.

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):
-            pieces = self._prepare_pieces(
-                video_path=video_path,
-                audio_path=audio_path,
-                interval=interval,
-                use_audio_in_video=use_audio_in_video,
-            )
+            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,18 +1424,28 @@ class OpenStreamReplayExtractor:
                     block_start,
                     self.config.native_block_seconds,
                 )
-                (
-                    committed_cache,
-                    _,
-                    payload_base,
-                    _,
-                    _,
-                ) = self._cache_from_intervals(
-                    video_path=video_path,
-                    audio_path=audio_path,
-                    intervals=history_intervals,
-                    use_audio_in_video=use_audio_in_video,
-                )
+                if history_intervals:
+                    prepared_history = completed_blocks.get_many(
+                        history_intervals
+                    )
+                    (
+                        committed_cache,
+                        _,
+                        payload_base,
+                        _,
+                        _,
+                    ) = self._cache_from_intervals(
+                        video_path=video_path,
+                        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,