Batch Omni AV cache extraction

cache_omni_av_pieces.py

@@ -8,7 +8,9 @@ future samples enter a phase.
 
 The optional verification compares a full-sequence reconstruction from the
 cache with the existing open-stream KV replay.  Run this 24-second pilot before
-caching a full EasyCom session.
+caching a full EasyCom session.  ``--av-batch-size`` batches equal-duration
+phases from different source blocks, increasing frozen AV-tower occupancy while
+preserving the exact per-block cache schema and causal timestamps.
 """
 
 from __future__ import annotations
@@ -16,6 +18,7 @@ from __future__ import annotations
 import argparse
 import math
 import platform
+import time
 from pathlib import Path
 from typing import Any, Sequence
 
@@ -60,6 +63,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
     parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
     parser.add_argument("--frame-cache-size", type=int, default=64)
+    parser.add_argument(
+        "--av-batch-size",
+        type=int,
+        default=1,
+        help=(
+            "Microbatch source blocks that have the same 0.2s phase through "
+            "the processor and frozen AV towers; 1 preserves serial behavior"
+        ),
+    )
     parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
     parser.add_argument("--decoder-gpu-id", type=int, default=0)
     parser.add_argument("--decoder-threads", type=int, default=1)
@@ -118,70 +130,256 @@ def _model_device(model: torch.nn.Module) -> torch.device:
         raise RuntimeError("Model has no parameters") from exc
 
 
+def _one_video_item(value: Any) -> Any:
+    if isinstance(value, (list, tuple)):
+        if len(value) != 1:
+            raise ValueError(f"Expected one prepared video, got {len(value)}")
+        return value[0]
+    return value
+
+
+@torch.inference_mode()
+def _prepare_pieces_batch(
+    *,
+    model: torch.nn.Module,
+    processor: Any,
+    media: Any,
+    config: core.StreamConfig,
+    intervals: list[tuple[float, float]],
+) -> list[core.OpenMediaPieces]:
+    """Encode equal-duration AV intervals as one independent-sample batch."""
+
+    if not intervals:
+        raise ValueError("Cannot prepare an empty AV batch")
+    durations = [end - start for start, end in intervals]
+    if any(
+        not math.isclose(value, durations[0], rel_tol=0.0, abs_tol=1e-6)
+        for value in durations[1:]
+    ):
+        raise ValueError(f"AV microbatch intervals have mixed durations: {durations}")
+
+    texts: list[str] = []
+    audios: list[Any] = []
+    videos: list[Any] = []
+    for interval in intervals:
+        conversation, manual_audios = core.build_chunk_conversation(
+            video_path=None,
+            audio_path=None,
+            intervals=[interval],
+            use_audio_in_video=True,
+            config=config,
+            fixed_grid_av=media,
+        )
+        if manual_audios is None or len(manual_audios) != 1:
+            raise RuntimeError("Direct-media interval did not provide one audio slice")
+        text = core._single_template_text(
+            processor.apply_chat_template(
+                conversation,
+                add_generation_prompt=False,
+                tokenize=False,
+            )
+        )
+        _, images, prepared_videos = core.process_mm_info(
+            conversation,
+            use_audio_in_video=False,
+        )
+        if images is not None and len(images):
+            raise RuntimeError("AV cache batch unexpectedly produced image inputs")
+        if prepared_videos is None:
+            raise RuntimeError("AV cache batch did not produce video inputs")
+        texts.append(text)
+        audios.append(manual_audios[0])
+        videos.append(_one_video_item(prepared_videos))
+
+    batch = processor(
+        text=texts,
+        audio=audios,
+        videos=videos,
+        padding=True,
+        use_audio_in_video=True,
+        fps=config.video_fps,
+    )
+    batch = core._processor_batch_convert_to_tensors(batch)
+    batch = core._processor_batch_to_device(batch, model.device)
+    batch = core._processor_batch_to_dtype(batch, model.dtype)
+    inputs = dict(batch)
+    input_ids = inputs["input_ids"]
+    if input_ids.shape[0] != len(intervals):
+        raise RuntimeError(
+            f"Processor returned batch {input_ids.shape[0]} for {len(intervals)} intervals"
+        )
+
+    embeddings = core._multimodal_embeddings(model, inputs)
+    positions = core._local_position_ids(
+        model,
+        inputs,
+        use_audio_in_video=True,
+    )
+    attention_mask = inputs.get("attention_mask")
+    audio_id = core._config_token_id(model.config, "audio")
+    video_id = core._config_token_id(model.config, "video")
+    positions_per_chunk = core._tmrope_positions_per_chunk(model)
+    pieces_batch: list[core.OpenMediaPieces] = []
+    for row in range(len(intervals)):
+        payload_mask = (input_ids[row] == audio_id) | (input_ids[row] == video_id)
+        indices = payload_mask.nonzero(as_tuple=False).flatten()
+        if indices.numel() == 0:
+            raise ValueError(f"Prepared AV batch row {row} has no payload tokens")
+        first = int(indices[0].item())
+        last = int(indices[-1].item())
+        if not bool(payload_mask[first : last + 1].all().item()):
+            raise RuntimeError(f"AV payload is not contiguous in batch row {row}")
+        if attention_mask is None:
+            valid_end = input_ids.shape[1]
+        else:
+            valid = attention_mask[row].nonzero(as_tuple=False).flatten()
+            if valid.numel() == 0:
+                raise RuntimeError(f"Prepared AV batch row {row} is fully padded")
+            valid_end = int(valid[-1].item()) + 1
+        if valid_end <= last + 1:
+            raise RuntimeError(f"Prepared AV batch row {row} has no closure tokens")
+        pieces = core.OpenMediaPieces(
+            prefix_embeddings=embeddings[row : row + 1, :first, :],
+            prefix_position_ids=positions[:, row : row + 1, :first],
+            payload_embeddings=embeddings[row : row + 1, first : last + 1, :],
+            payload_position_ids=positions[:, row : row + 1, first : last + 1],
+            closure_ids=input_ids[row : row + 1, last + 1 : valid_end],
+            closure_position_ids=positions[
+                :, row : row + 1, last + 1 : valid_end
+            ],
+        )
+        temporal = pieces.payload_position_ids[0]
+        temporal_span = int(temporal.max().item() - temporal.min().item())
+        if temporal_span >= positions_per_chunk:
+            raise RuntimeError(
+                "One batched AV phase spans multiple native TMRoPE chunks: "
+                f"interval={intervals[row]}, span={temporal_span}, "
+                f"chunk={positions_per_chunk}"
+            )
+        pieces_batch.append(pieces)
+    return pieces_batch
+
+
 @torch.inference_mode()
 def _cache_blocks(
     *,
     extractor: core.OpenStreamReplayExtractor,
     model: torch.nn.Module,
+    processor: Any,
+    media: Any,
+    config: core.StreamConfig,
     output_dir: Path,
     num_steps: int,
     steps_per_block: int,
     step_seconds: float,
     block_seconds: float,
+    av_batch_size: int,
     overwrite: bool,
 ) -> list[dict[str, Any]]:
     num_blocks = math.ceil(num_steps / steps_per_block)
     records: list[dict[str, Any]] = []
     input_embeddings = model.get_input_embeddings()
-    for block_index in tqdm(range(num_blocks), desc="cache AV blocks", unit="block"):
-        first_step = block_index * steps_per_block
-        phase_count = min(steps_per_block, num_steps - first_step)
-        block_start = block_index * block_seconds
-        variants: dict[str, dict[str, torch.Tensor]] = {}
-        token_lengths: dict[str, int] = {}
-        for tail_steps in range(1, phase_count + 1):
-            interval_end = block_start + tail_steps * step_seconds
-            pieces = extractor._prepare_pieces(
-                video_path=None,
-                audio_path=None,
-                interval=(block_start, interval_end),
-                use_audio_in_video=True,
-            )
-            closure_embeddings = input_embeddings(pieces.closure_ids)
-            cached = pieces_to_cache(
-                pieces,
-                closure_embeddings=closure_embeddings,
-                floating_dtype=torch.bfloat16,
-            )
-            variants[str(tail_steps)] = cached
-            token_lengths[str(tail_steps)] = int(
-                cached["prefix_embeddings"].shape[1]
-                + cached["payload_embeddings"].shape[1]
-                + cached["closure_embeddings"].shape[1]
-            )
-        path = output_dir / cache_block_filename(block_index)
-        _safe_output_path(path, overwrite)
-        torch.save(
-            {
-                "schema_version": SCHEMA_VERSION,
-                "block_index": block_index,
-                "start_seconds": block_start,
-                "end_seconds": min(
-                    block_start + block_seconds, num_steps * step_seconds
-                ),
-                "variants": variants,
-            },
-            path,
-        )
-        records.append(
-            {
-                "block_index": block_index,
-                "path": path.name,
-                "phase_count": phase_count,
-                "token_lengths_without_query": token_lengths,
-                "size_bytes": path.stat().st_size,
-            }
+    progress = tqdm(
+        total=num_blocks,
+        desc=f"cache AV blocks (batch={av_batch_size})",
+        unit="block",
+    )
+    for tile_start in range(0, num_blocks, av_batch_size):
+        block_indices = list(
+            range(tile_start, min(tile_start + av_batch_size, num_blocks))
         )
+        variants_by_block: dict[int, dict[str, dict[str, torch.Tensor]]] = {
+            index: {} for index in block_indices
+        }
+        lengths_by_block: dict[int, dict[str, int]] = {
+            index: {} for index in block_indices
+        }
+        phase_counts = {
+            index: min(
+                steps_per_block,
+                num_steps - index * steps_per_block,
+            )
+            for index in block_indices
+        }
+        for block_index in block_indices:
+            _safe_output_path(
+                output_dir / cache_block_filename(block_index), overwrite
+            )
+
+        for tail_steps in range(1, steps_per_block + 1):
+            active_indices = [
+                index
+                for index in block_indices
+                if tail_steps <= phase_counts[index]
+            ]
+            if not active_indices:
+                continue
+            intervals = [
+                (
+                    index * block_seconds,
+                    index * block_seconds + tail_steps * step_seconds,
+                )
+                for index in active_indices
+            ]
+            if len(active_indices) == 1 or av_batch_size == 1:
+                pieces_batch = [
+                    extractor._prepare_pieces(
+                        video_path=None,
+                        audio_path=None,
+                        interval=intervals[0],
+                        use_audio_in_video=True,
+                    )
+                ]
+            else:
+                pieces_batch = _prepare_pieces_batch(
+                    model=model,
+                    processor=processor,
+                    media=media,
+                    config=config,
+                    intervals=intervals,
+                )
+            if len(pieces_batch) != len(active_indices):
+                raise RuntimeError("Batched AV piece count changed unexpectedly")
+            for block_index, pieces in zip(active_indices, pieces_batch):
+                closure_embeddings = input_embeddings(pieces.closure_ids)
+                cached = pieces_to_cache(
+                    pieces,
+                    closure_embeddings=closure_embeddings,
+                    floating_dtype=torch.bfloat16,
+                )
+                variants_by_block[block_index][str(tail_steps)] = cached
+                lengths_by_block[block_index][str(tail_steps)] = int(
+                    cached["prefix_embeddings"].shape[1]
+                    + cached["payload_embeddings"].shape[1]
+                    + cached["closure_embeddings"].shape[1]
+                )
+
+        for block_index in block_indices:
+            block_start = block_index * block_seconds
+            path = output_dir / cache_block_filename(block_index)
+            torch.save(
+                {
+                    "schema_version": SCHEMA_VERSION,
+                    "block_index": block_index,
+                    "start_seconds": block_start,
+                    "end_seconds": min(
+                        block_start + block_seconds, num_steps * step_seconds
+                    ),
+                    "variants": variants_by_block[block_index],
+                },
+                path,
+            )
+            records.append(
+                {
+                    "block_index": block_index,
+                    "path": path.name,
+                    "phase_count": phase_counts[block_index],
+                    "token_lengths_without_query": lengths_by_block[block_index],
+                    "size_bytes": path.stat().st_size,
+                }
+            )
+        progress.update(len(block_indices))
+    progress.close()
     return records
 
 
@@ -270,6 +468,8 @@ def main(argv: Sequence[str] | None = None) -> None:
     core._configure_qwen_logging()
     if args.end <= 0:
         raise ValueError("--end must be positive")
+    if args.av_batch_size <= 0:
+        raise ValueError("--av-batch-size must be positive")
     steps_per_block = _integer_ratio(
         args.native_block_seconds, args.step_seconds, "native block / step"
     )
@@ -339,16 +539,22 @@ def main(argv: Sequence[str] | None = None) -> None:
     }
     torch.save(static, static_path)
 
+    cache_started = time.perf_counter()
     block_records = _cache_blocks(
         extractor=extractor,
         model=model,
+        processor=processor,
+        media=media,
+        config=config,
         output_dir=output_dir,
         num_steps=num_steps,
         steps_per_block=steps_per_block,
         step_seconds=args.step_seconds,
         block_seconds=args.native_block_seconds,
+        av_batch_size=args.av_batch_size,
         overwrite=args.overwrite,
     )
+    cache_seconds = time.perf_counter() - cache_started
     anchors = [
         make_anchor(
             step_index=step_index,
@@ -395,6 +601,10 @@ def main(argv: Sequence[str] | None = None) -> None:
         "use_audio_in_video": True,
         "interval_convention": "half_open_[start,end)",
         "phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
+        "av_batch_size": args.av_batch_size,
+        "cache_generation_seconds": cache_seconds,
+        "cache_media_realtime_factor": args.end / cache_seconds,
+        "cache_blocks_per_second": len(block_records) / cache_seconds,
         "label_policy": "target[i+1:i+11]",
         "static_path": static_path.name,
         "blocks": block_records,
@@ -413,6 +623,11 @@ def main(argv: Sequence[str] | None = None) -> None:
     print(
         f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
     )
+    print(
+        f"AV cache generation: {cache_seconds:.3f}s, "
+        f"{len(block_records) / cache_seconds:.3f} block/s, "
+        f"{args.end / cache_seconds:.3f}x media realtime"
+    )
     print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB")
     if checks:
         print(