File size: 7,606 Bytes
5f0e4a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """Read-only Full-DiT capture of targets and joint-window Context prefeatures."""
from __future__ import annotations
from typing import Any, Callable
import torch
from .prefeature_schema import CONTEXT_BLOCK_IDS
from .v2_hooks import PredictorV2TeacherCapture, _clone_cpu
class PredictorPrefeatureTeacherCapture(PredictorV2TeacherCapture):
"""Capture ``img_modulated`` before selected Teacher K/V projections."""
def __init__(
self,
transformer: torch.nn.Module,
*,
on_chunk: Callable[[int, dict[str, torch.Tensor], dict[int, dict[str, torch.Tensor]]], None],
num_steps: int = 4,
context_block_ids: tuple[int, ...] = CONTEXT_BLOCK_IDS,
capture_direct_kv: bool = False,
) -> None:
super().__init__(
transformer,
on_chunk=on_chunk,
num_steps=num_steps,
context_block_ids=context_block_ids,
)
self.capture_direct_kv = bool(capture_direct_kv)
self.pending_context: dict[int, dict[str, torch.Tensor]] | None = None
self.pending_direct_kv: dict[int, dict[str, torch.Tensor]] | None = None
self._prefill_active = False
self._prefill_features: dict[int, torch.Tensor] = {}
self._prefill_metadata: dict[str, torch.Tensor] = {}
def __enter__(self) -> "PredictorPrefeatureTeacherCapture":
super().__enter__()
for block_id in self.context_block_ids:
block = self.transformer.double_blocks[block_id]
self._handles.append(
block.img_attn_k.register_forward_pre_hook(
self._make_k_pre_hook(block_id), with_kwargs=True
)
)
return self
def _make_k_pre_hook(self, block_id: int):
def hook(module, args, kwargs) -> None:
if not self._prefill_active:
return
if block_id in self._prefill_features:
raise RuntimeError(f"Context block {block_id} was captured twice")
if not args or not torch.is_tensor(args[0]):
raise RuntimeError(f"Missing img_modulated input for block {block_id}")
self._prefill_features[block_id] = _clone_cpu(args[0])
return hook
def _transformer_pre(self, module, args, kwargs) -> None:
if kwargs.get("ar_vision_inference", False) and kwargs.get("cache_vision", False):
if self.pending_context is not None or self._prefill_active:
raise RuntimeError("A Context prefill was not consumed before the next prefill")
frame_indices = kwargs.get("context_frame_indices")
if frame_indices is None:
raise RuntimeError("Context prefill lacks context_frame_indices metadata")
if not torch.is_tensor(frame_indices):
frame_indices = torch.tensor(frame_indices, dtype=torch.int64)
frame_indices = frame_indices.detach().to(device="cpu", dtype=torch.int64).reshape(-1)
frames = int(kwargs["hidden_states"].shape[2])
if frame_indices.numel() != frames:
raise ValueError("Context frame metadata does not match prefill tensor")
self._prefill_metadata = {
"selected_frame_indices": frame_indices.contiguous(),
"context_viewmats": _clone_cpu(kwargs["viewmats"]),
"context_Ks": _clone_cpu(kwargs["Ks"]),
"rope_temporal_size": torch.tensor(
[int(kwargs["rope_temporal_size"])], dtype=torch.int64
),
"start_rope_start_idx": torch.tensor(
[int(kwargs["start_rope_start_idx"])], dtype=torch.int64
),
}
self._prefill_features = {}
self._prefill_active = True
return
super()._transformer_pre(module, args, kwargs)
if self.active is not None and self.active.step_id == 0:
self.chunk_context = self._consume_context()
def _transformer_post(self, module, args, kwargs, output) -> None:
if self._prefill_active:
self._prefill_active = False
missing = set(self.context_block_ids).difference(self._prefill_features)
if missing:
raise RuntimeError(f"Missing prefeatures for blocks {sorted(missing)}")
token_counts = {int(value.shape[1]) for value in self._prefill_features.values()}
if len(token_counts) != 1:
raise RuntimeError("Selected Context blocks have different window lengths")
tokens = next(iter(token_counts))
mask = torch.ones((1, tokens), dtype=torch.bool)
self.pending_context = {
block_id: {
"img_modulated": feature,
"context_valid_mask": mask.clone(),
**{name: value.clone() for name, value in self._prefill_metadata.items()},
}
for block_id, feature in self._prefill_features.items()
}
if self.capture_direct_kv:
self.pending_direct_kv = {
block_id: {
"k_vision": _clone_cpu(output[block_id]["k_vision"]),
"v_vision": _clone_cpu(output[block_id]["v_vision"]),
}
for block_id in self.context_block_ids
}
self._prefill_features = {}
self._prefill_metadata = {}
return
super()._transformer_post(module, args, kwargs, output)
def _consume_context(self) -> dict[int, dict[str, torch.Tensor]]:
if self.pending_context is not None:
context = self.pending_context
self.pending_context = None
return context
# The first chunk has no history prefill.
empty_metadata = {
"context_valid_mask": torch.ones((1, 0), dtype=torch.bool),
"selected_frame_indices": torch.empty((0,), dtype=torch.int64),
"context_viewmats": torch.empty((1, 0, 4, 4), dtype=torch.bfloat16),
"context_Ks": torch.empty((1, 0, 3, 3), dtype=torch.bfloat16),
"rope_temporal_size": torch.tensor([0], dtype=torch.int64),
"start_rope_start_idx": torch.tensor([0], dtype=torch.int64),
}
return {
block_id: {
"img_modulated": torch.empty((1, 0, 2048), dtype=torch.bfloat16),
**{name: value.clone() for name, value in empty_metadata.items()},
}
for block_id in self.context_block_ids
}
def _capture_context(self, kv_cache):
"""Capture only case-level text KV; Context features come from block hooks."""
if self.text_context is None:
text_context: dict[int, dict[str, torch.Tensor]] = {}
for block_id in self.context_block_ids:
cache = kv_cache[block_id]
k_txt, v_txt = cache.get("k_txt"), cache.get("v_txt")
if k_txt is None or v_txt is None:
raise RuntimeError(f"Block {block_id} text KV is unavailable")
text_context[block_id] = {
"k_txt": _clone_cpu(k_txt),
"v_txt": _clone_cpu(v_txt),
}
self.text_context = text_context
return {block_id: {} for block_id in self.context_block_ids}
def take_direct_kv(self) -> dict[int, dict[str, torch.Tensor]] | None:
value = self.pending_direct_kv
self.pending_direct_kv = None
return value
|