verify_anchor_memory_and_kv_buffer_agillm4.py

#!/usr/bin/env python3
"""Verifier for AGILLM-4 anchor-memory wiring and KV-buffer harvest.

Two harvested/integrated features share this file because both touch
nB300_agillm4 and both need a regression-safe smoke test:

  1. Anchor memory: AnchorMemoryLayer wired into Encoder at a configurable
     mid-stack position. Default-off; enabled with ``--anchor_memory`` and
     parameterised by ``--anchor_stride / --anchor_max / --anchor_position``.

  2. KV buffer: ``KVBuffer`` replaces ``torch.cat``-based decode-time cache
     growth. Default-off (legacy tuple cache still works); enabled per
     call site by passing a pre-sized ``KVBuffer`` as ``kv_cache``.

Run:

    python /workspace/agillm-4/verify_anchor_memory_and_kv_buffer_agillm4.py
"""
from __future__ import annotations

import pathlib
import sys

HERE = pathlib.Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

import torch  # noqa: E402

import nB300_agillm4 as m  # noqa: E402
from anchor_memory import AnchorMemoryLayer  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"
FAILS = 0


def check(name: str, cond: bool, detail: str = ""):
    global FAILS
    tag = "OK " if cond else "FAIL"
    if not cond:
        FAILS += 1
    print(f"  [{tag}] {name}" + (f" — {detail}" if detail else ""))


def section(title: str):
    print(f"\n=== {title} ===")


# ----- 1. KVBuffer -----
section("KVBuffer fill / view / overflow")
buf = m.KVBuffer(batch=2, heads=4, capacity=32, d_k=8, device=DEV, dtype=torch.float32)
check("initial length is 0", buf.length == 0)
check("k/v allocated", buf.k.shape == (2, 4, 32, 8) and buf.v.shape == (2, 4, 32, 8))

k1 = torch.randn(2, 4, 5, 8, device=DEV)
v1 = torch.randn(2, 4, 5, 8, device=DEV)
buf.append(k1, v1)
check("after append(5) length is 5", buf.length == 5)
k_view, v_view = buf.view()
check("view shape", k_view.shape == (2, 4, 5, 8))
check("k content correct", torch.allclose(k_view, k1))
check("v content correct", torch.allclose(v_view, v1))

buf.append(torch.randn(2, 4, 27, 8, device=DEV), torch.randn(2, 4, 27, 8, device=DEV))
check("filled to capacity (32)", buf.length == 32)
try:
    buf.append(torch.zeros(2, 4, 1, 8, device=DEV), torch.zeros(2, 4, 1, 8, device=DEV))
    check("overflow raises", False, "no exception")
except RuntimeError:
    check("overflow raises", True)

# ----- 2. Encoder without anchor memory (regression) -----
section("Encoder regression (anchor disabled)")
cfg = {"d": 64, "layers": 4, "heads": 4, "rank": 8}
enc_plain = m.Encoder(cfg).to(DEV)
check("anchor module is None", enc_plain.anchor is None)
check("anchor_memory_enabled is False", enc_plain.anchor_memory_enabled is False)
ids = torch.randint(0, m.VOCAB, (1, 16), device=DEV)
with torch.no_grad():
    out = enc_plain(ids, mask=None)
check("plain forward shape", out.shape == (1, 16, 64), str(out.shape))

# ----- 3. Encoder with anchor memory -----
section("Encoder with anchor memory enabled")
enc = m.Encoder(
    cfg,
    anchor_memory=True,
    anchor_stride=8,
    anchor_max=16,
    anchor_position=-1,
).to(DEV)
check("anchor_memory_enabled", enc.anchor_memory_enabled is True)
check("anchor is AnchorMemoryLayer", isinstance(enc.anchor, AnchorMemoryLayer))
check("position resolved to mid-stack (4//2=2)", enc.anchor_position == 2, f"got {enc.anchor_position}")

ids2 = torch.randint(0, m.VOCAB, (1, 32), device=DEV)
out2 = enc(ids2, mask=None)
check("anchor-on forward shape", out2.shape == (1, 32, 64), str(out2.shape))

loss = out2.sum()
loss.backward()
grads = [p.grad for p in enc.anchor.parameters() if p.grad is not None]
check("anchor params got gradients", len(grads) > 0, f"{len(grads)} grads")
nonzero = sum(int(g.abs().sum() > 0) for g in grads)
check("anchor gradients non-zero", nonzero > 0, f"{nonzero}/{len(grads)}")

# ----- 4. Anchor position resolution -----
section("Anchor position resolution")
check("explicit position=1", m.Encoder(cfg, anchor_memory=True, anchor_position=1).anchor_position == 1)
check("oversized position clamped", m.Encoder(cfg, anchor_memory=True, anchor_position=99).anchor_position == cfg["layers"] - 1)

# ----- 5. Module-level defaults exposed -----
section("Module-level defaults present")
for name in ("DEFAULT_ANCHOR_MEMORY", "DEFAULT_ANCHOR_STRIDE", "DEFAULT_ANCHOR_MAX", "DEFAULT_ANCHOR_POSITION", "DEFAULT_KV_BUFFER"):
    check(name, hasattr(m, name))

print()
if FAILS:
    print(f"=== {FAILS} CHECK(S) FAILED ===")
    sys.exit(1)
print("=== ALL CHECKS PASSED ===")