From 43f47e8f3a2e6bafdce1f4bf2299e0d7cee45800 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Tue, 18 Aug 2026 11:02:42 +0800 Subject: [PATCH] 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 | 9 +- profile_streaming_bottleneck_gpu7.py | 12 +- streaming_turn_model.py | 216 ++++++++++++++++++++++++--- streaming_turn_model_directly.py | 5 + 4 files changed, 220 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index e5e6e99..8cf57e1 100644 --- a/README.md +++ b/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 diff --git a/profile_streaming_bottleneck_gpu7.py b/profile_streaming_bottleneck_gpu7.py index 39c6bf1..9711844 100644 --- a/profile_streaming_bottleneck_gpu7.py +++ b/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", diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 72018a5..7d6dff7 100644 --- a/streaming_turn_model.py +++ b/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, diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 6463a53..a879e6e 100644 --- a/streaming_turn_model_directly.py +++ b/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