wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Batch Omni AV cache extraction

Browse Source
This commit is contained in:
Xin Wang
2026-08-23 16:03:07 +08:00
parent c2162e1170
commit 85ed8a670c
1 changed files with 264 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cache_omni_av_pieces.py
+232 -17
View File
@@ -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,49 +130,233 @@ 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(
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=(block_start, interval_end),
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[str(tail_steps)] = cached
token_lengths[str(tail_steps)] = int(
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)
_safe_output_path(path, overwrite)
torch.save(
{
"schema_version": SCHEMA_VERSION,
@@ -169,7 +365,7 @@ def _cache_blocks(
"end_seconds": min(
block_start + block_seconds, num_steps * step_seconds
),
"variants": variants,
"variants": variants_by_block[block_index],
},
path,
)
@@ -177,11 +373,13 @@ def _cache_blocks(
{
"block_index": block_index,
"path": path.name,
"phase_count": phase_count,
"token_lengths_without_query": token_lengths,
"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(