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

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Any

import torch
from torch import nn


@dataclass(frozen=True)
class AxiomKVConfig:
    physical_dim: int = 128
    effective_dim: int = 20_480
    local_window: int = 64
    anchor_slots: int = 32
    anchor_rank: int = 64
    decay: float = 0.97


@dataclass
class AxiomKVState:
    local_exact: torch.Tensor
    anchors: torch.Tensor
    cursor: int
    long_context_kv_tokens_stored: int = 0

    def cached_token_capacity(self) -> int:
        return int(self.local_exact.shape[1] + self.anchors.shape[1])


class AxiomKVLedger(nn.Module):
    """Bounded KV companion for AxiomDim.

    Long context is compressed into fixed anchor slots. Only a small exact local
    window is retained, so KV capacity is constant with respect to sequence
    length and independent of ``effective_dim``.
    """

    def __init__(self, cfg: AxiomKVConfig):
        super().__init__()
        if cfg.physical_dim <= 0 or cfg.effective_dim <= 0 or cfg.local_window <= 0 or cfg.anchor_slots <= 0 or cfg.anchor_rank <= 0:
            raise ValueError("AxiomKV dimensions must be positive")
        self.cfg = cfg
        self.to_anchor = nn.Linear(cfg.physical_dim, cfg.anchor_rank, bias=False)
        self.from_anchor = nn.Linear(cfg.anchor_rank, cfg.physical_dim, bias=False)

    def _empty_state(self, batch: int, device: torch.device, dtype: torch.dtype) -> AxiomKVState:
        return AxiomKVState(
            local_exact=torch.zeros(batch, 0, self.cfg.physical_dim, device=device, dtype=dtype),
            anchors=torch.zeros(batch, self.cfg.anchor_slots, self.cfg.anchor_rank, device=device, dtype=dtype),
            cursor=0,
            long_context_kv_tokens_stored=0,
        )

    def ingest(self, x: torch.Tensor, state: AxiomKVState | None = None) -> AxiomKVState:
        batch, seq_len, dim = x.shape
        if dim != self.cfg.physical_dim:
            raise ValueError(f"expected physical_dim={self.cfg.physical_dim}, got {dim}")
        current = state if state is not None else self._empty_state(batch, x.device, x.dtype)
        local = torch.cat([current.local_exact.to(x.device, x.dtype), x], dim=1)[:, -self.cfg.local_window :]
        anchors = current.anchors.to(x.device, x.dtype)
        cursor = int(current.cursor)
        decay = float(min(max(self.cfg.decay, 0.0), 0.9999))

        compressed = torch.tanh(self.to_anchor(x))
        slot_ids = (torch.arange(seq_len, device=x.device) + cursor) % self.cfg.anchor_slots
        slot_ids = slot_ids.to(torch.long)
        for slot in range(self.cfg.anchor_slots):
            mask = slot_ids == slot
            if bool(mask.any().item()):
                update = compressed[:, mask].mean(dim=1)
                anchors[:, slot] = decay * anchors[:, slot] + (1.0 - decay) * update
        cursor += seq_len

        return AxiomKVState(
            local_exact=local,
            anchors=anchors,
            cursor=cursor,
            long_context_kv_tokens_stored=0,
        )

    def retrieve_summary(self, query: torch.Tensor, state: AxiomKVState) -> torch.Tensor:
        q = torch.tanh(self.to_anchor(query))
        scores = torch.einsum("bd,bsd->bs", q, state.anchors) / max(1.0, self.cfg.anchor_rank ** 0.5)
        weights = torch.softmax(scores, dim=-1)
        anchor = torch.einsum("bs,bsd->bd", weights, state.anchors)
        return self.from_anchor(anchor)

    def bounded_capacity(self) -> int:
        return self.cfg.local_window + self.cfg.anchor_slots


def _measure(cfg: AxiomKVConfig, seq_len: int) -> dict[str, Any]:
    torch.manual_seed(20260527 + int(seq_len))
    ledger = AxiomKVLedger(cfg)
    x = torch.randn(2, seq_len, cfg.physical_dim)
    state = ledger.ingest(x)
    q = torch.randn(2, cfg.physical_dim, requires_grad=True)
    summary = ledger.retrieve_summary(q, state)
    loss = summary.pow(2).mean()
    loss.backward()
    return {
        "seq_len": seq_len,
        "local_exact_shape": list(state.local_exact.shape),
        "anchors_shape": list(state.anchors.shape),
        "cached_token_capacity": state.cached_token_capacity(),
        "long_context_kv_tokens_stored": state.long_context_kv_tokens_stored,
        "summary_shape": list(summary.shape),
        "forward_finite": bool(torch.isfinite(summary).all().item()),
        "backward_finite": q.grad is not None and bool(torch.isfinite(q.grad).all().item()),
    }


def build_axiomkv_report(
    out_dir: str | Path,
    *,
    effective_dim: int = 20_480,
    physical_dim: int = 128,
    seq_lengths: list[int] | None = None,
    local_window: int = 64,
    anchor_slots: int = 32,
    anchor_rank: int = 64,
) -> dict[str, Any]:
    lengths = seq_lengths or [128, 1024, 8192]
    cfg = AxiomKVConfig(
        physical_dim=physical_dim,
        effective_dim=effective_dim,
        local_window=local_window,
        anchor_slots=anchor_slots,
        anchor_rank=anchor_rank,
    )
    measurements = [_measure(cfg, int(length)) for length in lengths]
    capacities = {item["cached_token_capacity"] for item in measurements}
    bounded = len(capacities) == 1 and next(iter(capacities)) == local_window + anchor_slots
    report = {
        "schema": "tinymind.axiomkv.v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "target": {
            "effective_dim": effective_dim,
            "physical_dim": physical_dim,
            "kv_strategy": "exact_local_window_plus_fixed_compressed_anchors",
            "materializes_effective_dim_kv": False,
        },
        "config": {
            "local_window": local_window,
            "anchor_slots": anchor_slots,
            "anchor_rank": anchor_rank,
            "bounded_capacity": local_window + anchor_slots,
        },
        "measurements": measurements,
        "bounded_kv_gate": {
            "passed": bounded and all(item["forward_finite"] and item["backward_finite"] for item in measurements),
            "cached_token_capacity": local_window + anchor_slots,
            "long_context_kv_tokens_stored": 0,
            "capacity_constant_across_lengths": bounded,
        },
        "claim_gate": {
            "kv_growth_blocked": True,
            "full_kv_growth_claim_allowed": False,
            "tier0_claim_allowed": False,
            "world_best_claim_allowed": False,
            "reason": "AxiomKV proves bounded local smoke memory; exact long recall still requires external ledger evidence.",
        },
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    path = out / "axiomkv_report.json"
    report["json_path"] = str(path)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return report