predictor_training/rollout_cache.py

"""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",
]