project_6/qwen3_6_scripts/gdn_prefix.py

"""Shared GDN prefix-state cache contracts for the BI100 runtime."""

from __future__ import annotations

import os
from collections import OrderedDict
from dataclasses import dataclass
from typing import Iterable, List, Optional, Sequence, Tuple


GdnPrefixKey = Tuple[int, bytes]
GdnCapturePoint = Tuple[int, GdnPrefixKey]

_VALID_POLICIES = {"fine32", "admission64", "off"}
GDN_KERNEL_CHUNK_TOKENS = 64
GDN_DIRECT_MIN_REPLAY_TOKENS = 2

_VALID_RESTORE_MODES = {"direct", "hybrid64", "chunk64", "aligned"}


def _env_choice(name: str, default: str, choices: set[str]) -> str:
    value = os.getenv(name, default).strip().lower()
    if value not in choices:
        allowed = ", ".join(sorted(choices))
        raise RuntimeError(f"invalid {name}={value!r}; expected one of: {allowed}")
    return value


def gdn_cache_policy_from_env() -> str:
    return _env_choice("BI100_GDN_CACHE_POLICY", "fine32", _VALID_POLICIES)


def gdn_restore_mode_from_env() -> str:
    return _env_choice(
        "BI100_GDN_RESTORE_MODE", "direct", _VALID_RESTORE_MODES)


def gdn_restore_alignment(restore_mode: str, block_size: int,
                          scheduler_chunk_tokens: int) -> int:
    """Return the content boundary required by a restore mode."""
    if block_size <= 0:
        raise ValueError("block_size must be positive")
    if restore_mode == "direct":
        return block_size
    if restore_mode in {"hybrid64", "chunk64"}:
        alignment = GDN_KERNEL_CHUNK_TOKENS
    elif restore_mode == "aligned":
        alignment = scheduler_chunk_tokens
    else:
        raise ValueError(f"unknown GDN restore mode: {restore_mode}")
    if alignment <= 0 or alignment % block_size != 0:
        raise ValueError(
            f"{restore_mode} GDN restore requires a positive alignment "
            f"divisible by block_size={block_size}; got {alignment}")
    return alignment


def make_prefix_key(block_count: int, digest: bytes) -> GdnPrefixKey:
    if block_count <= 0:
        raise ValueError("GDN prefix key requires at least one complete block")
    if not isinstance(digest, bytes) or len(digest) != 32:
        raise ValueError("GDN prefix digest must be exactly 32 bytes")
    return block_count, digest


def keys_from_block_hashes(block_hashes: Sequence[bytes]) -> List[GdnPrefixKey]:
    return [make_prefix_key(i + 1, digest)
            for i, digest in enumerate(block_hashes)]


def strict_prefix_block_count(token_count: int, block_size: int) -> int:
    if block_size <= 0:
        raise ValueError("block_size must be positive")
    if token_count <= 1:
        return 0
    return (token_count - 1) // block_size


def key_at_strict_boundary(block_hashes: Sequence[bytes], token_count: int,
                           block_size: int) -> Optional[GdnPrefixKey]:
    block_count = min(
        len(block_hashes), strict_prefix_block_count(token_count, block_size))
    if block_count <= 0:
        return None
    return make_prefix_key(block_count, block_hashes[block_count - 1])


def final_capture_key(
        block_hashes: Sequence[bytes], prompt_tokens: int, block_size: int,
        restore_mode: str, replay_alignment: int) -> Optional[GdnPrefixKey]:
    if restore_mode in {"direct", "hybrid64"}:
        block_count = min(
            len(block_hashes), strict_prefix_block_count(
                prompt_tokens, block_size))
        if (block_count > 0
                and prompt_tokens - block_count * block_size
                < GDN_DIRECT_MIN_REPLAY_TOKENS):
            block_count -= 1
        if block_count <= 0:
            return None
        return make_prefix_key(block_count, block_hashes[block_count - 1])
    if restore_mode not in {"chunk64", "aligned"}:
        raise ValueError(f"unknown GDN restore mode: {restore_mode}")
    if (replay_alignment <= 0 or replay_alignment % block_size != 0
            or prompt_tokens <= 1):
        return None
    boundary_tokens = ((prompt_tokens - 1) // replay_alignment
                       * replay_alignment)
    block_count = min(len(block_hashes), boundary_tokens // block_size)
    if block_count <= 0:
        return None
    return make_prefix_key(block_count, block_hashes[block_count - 1])


def restore_key_is_eligible(
        key: GdnPrefixKey, prompt_tokens: int, block_size: int,
        restore_mode: str, replay_alignment: int,
        direct_final_key: Optional[GdnPrefixKey] = None) -> bool:
    """Return whether restoring ``key`` preserves the execution contract."""
    make_prefix_key(*key)
    if block_size <= 0:
        raise ValueError("block_size must be positive")
    boundary_tokens = key[0] * block_size
    remaining_tokens = prompt_tokens - boundary_tokens
    if remaining_tokens <= 0:
        return False
    if restore_mode == "direct":
        return remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
    if restore_mode == "hybrid64":
        if direct_final_key is not None:
            make_prefix_key(*direct_final_key)
        return (remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
                and replay_alignment > 0
                and (boundary_tokens % replay_alignment == 0
                     or key == direct_final_key))
    if restore_mode not in {"chunk64", "aligned"}:
        raise ValueError(f"unknown GDN restore mode: {restore_mode}")
    return (replay_alignment > 0
            and boundary_tokens % replay_alignment == 0)


def capture_points_for_step(
        targets: Iterable[GdnPrefixKey], physical_context_tokens: int,
        logical_end_tokens: int, block_size: int) -> Tuple[GdnCapturePoint, ...]:
    if physical_context_tokens < 0 or logical_end_tokens < 0:
        raise ValueError("token positions must be non-negative")
    if logical_end_tokens <= physical_context_tokens:
        return ()
    selected = {}
    for key in targets:
        make_prefix_key(*key)
        boundary_tokens = key[0] * block_size
        if physical_context_tokens < boundary_tokens <= logical_end_tokens:
            selected[boundary_tokens - physical_context_tokens] = key
    points = tuple(sorted(selected.items()))
    if len(points) > 2:
        raise ValueError("at most two GDN capture points are allowed per step")
    return points


def cap_prefill_end_at_capture_boundary(
        logical_start_tokens: int, logical_end_tokens: int,
        targets: Iterable[GdnPrefixKey], block_size: int) -> int:
    """Stop a physical prefill step at its earliest pending capture boundary."""
    if logical_start_tokens < 0 or logical_end_tokens < 0:
        raise ValueError("token positions must be non-negative")
    if logical_end_tokens < logical_start_tokens:
        raise ValueError("logical end must not precede logical start")
    if block_size <= 0:
        raise ValueError("block_size must be positive")

    capped_end = logical_end_tokens
    for key in targets:
        make_prefix_key(*key)
        boundary_tokens = key[0] * block_size
        if logical_start_tokens < boundary_tokens < capped_end:
            capped_end = boundary_tokens
    return capped_end


def canonical_direct_segment_offsets(
        block_hashes: Sequence[bytes], physical_context_tokens: int,
        logical_end_tokens: int, block_size: int,
        scheduler_chunk_tokens: int) -> Tuple[int, ...]:
    """Reproduce cold fine32/direct segment boundaries after fast-forward."""
    if physical_context_tokens < 0 or logical_end_tokens < 0:
        raise ValueError("token positions must be non-negative")
    if block_size <= 0 or scheduler_chunk_tokens <= 0:
        raise ValueError("block and scheduler chunk sizes must be positive")
    if scheduler_chunk_tokens % block_size != 0:
        raise ValueError("scheduler chunk size must be divisible by block size")
    if logical_end_tokens <= physical_context_tokens:
        return ()

    boundaries = set()
    step_ends = list(range(scheduler_chunk_tokens, logical_end_tokens,
                           scheduler_chunk_tokens))
    for step_end in (*step_ends, logical_end_tokens):
        key = final_capture_key(block_hashes, step_end, block_size,
                                "direct", block_size)
        if key is not None:
            boundaries.add(key[0] * block_size)
    boundaries.update(step_ends)
    return tuple(
        boundary - physical_context_tokens
        for boundary in sorted(boundaries)
        if physical_context_tokens < boundary < logical_end_tokens)


@dataclass(frozen=True)
class GdnCachePlan:
    restore_key: Optional[GdnPrefixKey] = None
    capture_points: Tuple[GdnCapturePoint, ...] = ()
    evict_keys: Tuple[GdnPrefixKey, ...] = ()


class GdnPrefixStatePolicy:
    """Scheduler-owned state index with deterministic worker actions."""

    def __init__(self, policy: str) -> None:
        if policy not in _VALID_POLICIES:
            raise ValueError(f"unknown GDN cache policy: {policy}")
        self.policy = policy
        self.capacity = {"fine32": 32, "admission64": 64, "off": 0}[policy]
        self._resident: OrderedDict[GdnPrefixKey, None] = OrderedDict()

    def __len__(self) -> int:
        return len(self._resident)

    def resident_keys(self) -> Tuple[GdnPrefixKey, ...]:
        return tuple(self._resident)

    def contains(self, key: GdnPrefixKey) -> bool:
        return key in self._resident

    def should_capture_final(self, key: GdnPrefixKey) -> bool:
        """Return whether a final state must be materialized on this request."""
        make_prefix_key(*key)
        if self.policy == "off":
            return False
        if self.policy == "admission64":
            return key not in self._resident
        return True

    def select_restore(
            self, live_prefix_keys: Sequence[GdnPrefixKey],
            max_blocks: int) -> Optional[GdnPrefixKey]:
        if self.capacity == 0 or max_blocks <= 0:
            return None
        best = None
        for key in live_prefix_keys[:max_blocks]:
            if key in self._resident:
                best = key
        if best is not None:
            self._resident.move_to_end(best)
        return best

    def repeated_branch_candidate(
            self, live_prefix_keys: Sequence[GdnPrefixKey],
            max_blocks: int) -> Optional[GdnPrefixKey]:
        """Return a repeated raw-KV branch that lacks recurrent state.

        A live KV hit proves that the content occurred in an earlier request;
        the current request is therefore the second or later occurrence.
        """
        if (self.policy != "admission64" or max_blocks <= 0
                or not live_prefix_keys):
            return None
        candidate = live_prefix_keys[min(len(live_prefix_keys), max_blocks) - 1]
        if candidate in self._resident:
            return None
        return candidate

    def admit(self, keys: Iterable[GdnPrefixKey]) -> Tuple[GdnPrefixKey, ...]:
        evicted: List[GdnPrefixKey] = []
        if self.capacity == 0:
            return ()
        for key in keys:
            make_prefix_key(*key)
            if key in self._resident:
                self._resident.move_to_end(key)
            else:
                self._resident[key] = None
            while len(self._resident) > self.capacity:
                evicted_key, _ = self._resident.popitem(last=False)
                evicted.append(evicted_key)
        return tuple(evicted)

    def forget(self, keys: Iterable[GdnPrefixKey]) -> None:
        for key in keys:
            self._resident.pop(key, None)