tinymind-native-8b-remote-handoff/bundle/model/regen_kv.py

"""ReGenesis-KV: regenerate compact K/V from exact retrieved ledger chunks."""

from __future__ import annotations

from collections.abc import Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import OmegaConfig
from .layers import RMSNorm


class ReGenesisKVBlock(nn.Module):
    """Fuse local hidden state, bounded PureField memory, and retrieved exact chunks.

    The block never stores full historical KV. Retrieved exact tokens are encoded
    on demand into compact regenerated K/V tensors, then fused back into hidden
    states through a small gated residual.
    """

    def __init__(self, cfg: OmegaConfig):
        super().__init__()
        self.cfg = cfg
        self.dim = int(cfg.dim)
        self.rank = int(cfg.regen_kv_rank)
        self.top_k = int(cfg.retrieval_top_k)
        self.token_embed = nn.Embedding(int(cfg.vocab_size), self.dim)
        self.query_norm = RMSNorm(self.dim)
        self.chunk_norm = RMSNorm(self.dim)
        self.memory_proj = nn.Linear(self.dim, self.dim, bias=False)
        self.k_proj = nn.Linear(self.dim, self.rank * self.dim, bias=False)
        self.v_proj = nn.Linear(self.dim, self.rank * self.dim, bias=False)
        self.fuse = nn.Linear(self.dim * 3, self.dim, bias=False)
        self.gate = nn.Linear(self.dim, self.dim, bias=True)

    def _encode_chunks(self, retrieved_tokens: torch.Tensor) -> torch.Tensor:
        safe_ids = retrieved_tokens.clamp(min=0, max=self.cfg.vocab_size - 1)
        embedded = self.token_embed(safe_ids)
        return self.chunk_norm(embedded.mean(dim=2))

    def _chunks_to_tensor(
        self,
        retrieved_chunks: Sequence[Sequence[int]],
        *,
        batch: int,
        device: torch.device,
    ) -> torch.Tensor:
        if not retrieved_chunks:
            return torch.zeros(batch, self.top_k, 1, device=device, dtype=torch.long)
        width = max(1, max(len(chunk) for chunk in retrieved_chunks))
        rows: list[list[int]] = []
        for chunk in list(retrieved_chunks)[: self.top_k]:
            row = [int(token) for token in chunk[:width]]
            row.extend([0] * (width - len(row)))
            rows.append(row)
        while len(rows) < self.top_k:
            rows.append([0] * width)
        tensor = torch.tensor(rows, device=device, dtype=torch.long).unsqueeze(0)
        return tensor.expand(batch, self.top_k, width).contiguous()

    def forward(
        self,
        hidden: torch.Tensor,
        memory_state: torch.Tensor | None = None,
        retrieved_tokens: torch.Tensor | None = None,
        retrieved_chunks: Sequence[Sequence[int]] | None = None,
        return_aux: bool = False,
    ) -> tuple[torch.Tensor, dict] | tuple[torch.Tensor, dict, dict[str, torch.Tensor]]:
        batch, seq_len, _ = hidden.shape
        source = "retrieved_tokens"
        if retrieved_chunks is not None:
            retrieved_tokens = self._chunks_to_tensor(retrieved_chunks, batch=batch, device=hidden.device)
            source = "retrieved_chunks"
        elif retrieved_tokens is None:
            retrieved_tokens = torch.zeros(
                batch,
                self.top_k,
                1,
                device=hidden.device,
                dtype=torch.long,
            )
        retrieved_tokens = retrieved_tokens.to(device=hidden.device, dtype=torch.long)
        if retrieved_tokens.shape[1] != self.top_k:
            if retrieved_tokens.shape[1] > self.top_k:
                retrieved_tokens = retrieved_tokens[:, : self.top_k]
            else:
                pad = torch.zeros(
                    batch,
                    self.top_k - retrieved_tokens.shape[1],
                    retrieved_tokens.shape[2],
                    device=hidden.device,
                    dtype=torch.long,
                )
                retrieved_tokens = torch.cat([retrieved_tokens, pad], dim=1)

        chunk_repr = self._encode_chunks(retrieved_tokens)
        regen_k = self.k_proj(chunk_repr).view(batch, self.top_k, self.rank, self.dim)
        regen_v = self.v_proj(chunk_repr).view(batch, self.top_k, self.rank, self.dim)

        query = self.query_norm(hidden)
        query_summary = query.mean(dim=1)
        chunk_key = F.normalize(regen_k.mean(dim=2), dim=-1, eps=1e-6)
        query_key = F.normalize(query_summary, dim=-1, eps=1e-6)
        weights = torch.softmax(torch.einsum("bd,bkd->bk", query_key, chunk_key), dim=-1)
        retrieved_value = torch.einsum("bk,bkrd->brd", weights, regen_v).mean(dim=1)
        retrieved_value = retrieved_value.unsqueeze(1).expand(batch, seq_len, self.dim)

        if memory_state is None:
            memory_summary = torch.zeros_like(retrieved_value)
        else:
            memory_summary = memory_state.mean(dim=tuple(range(1, memory_state.dim() - 1)))
            memory_summary = self.memory_proj(memory_summary).unsqueeze(1).expand(batch, seq_len, self.dim)

        fused = self.fuse(torch.cat([query, memory_summary, retrieved_value], dim=-1))
        gate = torch.sigmoid(self.gate(query))
        out = hidden + gate * torch.tanh(fused)
        hash_consistency_loss = retrieved_tokens.float().remainder(997).mean() * 0.0
        cache = {
            "regen_k": regen_k,
            "regen_v": regen_v,
            "retrieved_tokens": retrieved_tokens,
            "retrieval_weights": weights,
            "kv_tokens_stored": 0,
            "source": source,
        }
        if return_aux:
            return out, cache, {"hash_consistency_loss": hash_consistency_loss}
        return out, cache