File size: 4,047 Bytes
d5e0d8f | 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 | """Cache workspaces for detached Predictor-v4 trajectory rollout."""
from __future__ import annotations
from collections.abc import Sequence
import torch
def reset_main_caches(
kv_cache: Sequence[dict],
crossattn_cache: Sequence[dict],
) -> None:
"""Reset logical cache extents without clearing unused storage."""
for cache in kv_cache:
cache["global_end_index"].zero_()
cache["local_end_index"].zero_()
for cache in crossattn_cache:
cache["is_init"] = False
def build_predictor_workspace(
main_cache: Sequence[dict],
*,
source_block_ids: tuple[int, int],
history_tokens: int,
current_tokens: int,
) -> dict[int, dict[str, torch.Tensor]]:
"""Copy committed clean history into a temporary selected-layer workspace.
The current-chunk region is intentionally empty. Predictor P1/P2/P3
overwrite that region, and the entire workspace is discarded after the
chunk. Only the final Full timestep-zero pass updates persistent history.
"""
history_tokens = int(history_tokens)
current_tokens = int(current_tokens)
if history_tokens < 0 or current_tokens <= 0:
raise ValueError("Invalid history/current token count")
capacity = history_tokens + current_tokens
result: dict[int, dict[str, torch.Tensor]] = {}
for block_id in source_block_ids:
source = main_cache[int(block_id)]
if int(source["global_end_index"].item()) < history_tokens:
raise RuntimeError(
f"Teacher cache block {block_id} ends before committed history: "
f"{int(source['global_end_index'].item())} < {history_tokens}"
)
key = source["k"].new_zeros(
source["k"].shape[0], capacity, *source["k"].shape[2:]
)
value = source["v"].new_zeros(
source["v"].shape[0], capacity, *source["v"].shape[2:]
)
if history_tokens:
key[:, :history_tokens].copy_(source["k"][:, :history_tokens])
value[:, :history_tokens].copy_(source["v"][:, :history_tokens])
result[int(block_id)] = {
"k": key,
"v": value,
"global_end_index": torch.tensor(
[history_tokens], dtype=torch.long, device=key.device
),
"local_end_index": torch.tensor(
[history_tokens], dtype=torch.long, device=key.device
),
}
return result
def reset_predictor_workspace(
workspace: dict[int, dict[str, torch.Tensor]],
*,
history_tokens: int,
) -> None:
"""Discard the previous timestep's differentiable current-cache region."""
history_tokens = int(history_tokens)
for cache in workspace.values():
# Slice assignment from Predictor K/V can attach CopySlices autograd
# history to the workspace tensor. Replace it with a detached tensor
# before the next timestep so P2/P3 cannot backpropagate through cache.
cache["k"] = cache["k"].detach()
cache["v"] = cache["v"].detach()
if history_tokens < cache["k"].shape[1]:
cache["k"][:, history_tokens:].zero_()
cache["v"][:, history_tokens:].zero_()
cache["global_end_index"].fill_(history_tokens)
cache["local_end_index"].fill_(history_tokens)
def assert_clean_history_extent(
kv_cache: Sequence[dict],
*,
expected_tokens: int,
) -> None:
expected_tokens = int(expected_tokens)
for block_id, cache in enumerate(kv_cache):
global_end = int(cache["global_end_index"].item())
local_end = int(cache["local_end_index"].item())
if global_end != expected_tokens or local_end != expected_tokens:
raise RuntimeError(
f"Cache block {block_id} extent is ({global_end}, {local_end}), "
f"expected committed clean history {expected_tokens}"
)
__all__ = [
"assert_clean_history_extent",
"build_predictor_workspace",
"reset_main_caches",
"reset_predictor_workspace",
]
|