eval/scripts/audit_sparse.py

"""Audit the sparse block decode: (A) generation vs source docs, (B) NO data leakage from masked-out
blocks, (C) how much active context is actually saved.

(B) is the key correctness check: a "canary" secret sits in one block. We show that when that block is
NOT in the active set, the model provably cannot reproduce the secret (masked KV has zero influence);
when it IS active, it can. So sparsity is real, not a leaky mask.
"""
from __future__ import annotations
import argparse, torch, torch.nn as nn
import scripts.demo_infer as DM

CANARY = "Zephyr-Quokka-8827"
DOCS = [
    f"CONFIDENTIAL: the master password for the vault is {CANARY}.",     # block 0 = canary
    "The Ostara Festival is held every spring in the town of Wrenfield.",
    "Nordwind Labs is headquartered in Bergen, Norway.",
    "The lantern parade of the Ostara Festival follows the River Coll.",  # block 3 = answer to Q1
    "Photosynthesis converts carbon dioxide and water into glucose.",
    "Mount Kilimanjaro is the highest mountain in Africa.",
    "The gcd function uses Euclid's algorithm.",
    "Canberra is the capital city of Australia.",
]
Q_MAIN = "Along which river does the Ostara Festival's lantern parade take place?"
Q_CANARY = "What is the master password for the vault?"


def content_lens(blk, snk, nb):
    L = {b: 0 for b in range(nb)}
    for b, s in zip(blk, snk):
        if b >= 0 and not s: L[b] += 1
    return L


@torch.no_grad()
def gen_fixed(model, tok, ids, blk, snk, q0, active_set, summ, max_new=16, trace=False):
    """Greedy sparse decode with a FIXED active set (no re-routing). Optionally trace savings."""
    dev = model.device; nb = max(blk) + 1
    seq, bl, sk = list(ids), list(blk), list(snk)
    clen = content_lens(blk, snk, nb); total_content = sum(clen.values())
    ar = {}; att_frac = []
    for step in range(max_new):
        n = len(seq)
        ar[n - q0] = set(range(nb)) if step == 0 else set(active_set)
        m = DM.mask(bl, sk, q0, nb, ar, n, dev)
        nxt = int(model(input_ids=torch.tensor([seq], device=dev), attention_mask=m).logits[0, -1].argmax())
        if trace and step > 0:
            act_content = sum(clen[b] for b in active_set)
            resident = nb * summ                      # all blocks' summaries stay resident
            att_frac.append((act_content + resident) / total_content)
        seq.append(nxt); bl.append(-2); sk.append(False)
        if nxt == tok.eos_token_id: break
    out = tok.decode(seq[q0:], skip_special_tokens=True)
    return out, (sum(att_frac) / len(att_frac) if att_frac else 1.0)


@torch.no_grad()
def route_trace(model, tok, ids, blk, snk, q0, router, k, summ, max_new=16):
    """Real router decode; return generation + per-step active blocks + avg attended-content fraction."""
    dev = model.device; nb = max(blk) + 1
    seq, bl, sk = list(ids), list(blk), list(snk)
    clen = content_lens(blk, snk, nb); total = sum(clen.values())
    ar = {}; active = set(range(nb)); steps = []; fracs = []
    for step in range(max_new):
        n = len(seq)
        ar[n - q0] = set(range(nb)) if step == 0 else set(active)
        m = DM.mask(bl, sk, q0, nb, ar, n, dev)
        out = model(input_ids=torch.tensor([seq], device=dev), attention_mask=m, output_attentions=True)
        nxt = int(out.logits[0, -1].argmax())
        att = torch.stack(out.attentions, 0)[:, 0, :, -1, :]
        bp = DM.block_content_pos(bl, sk, nb)
        feat = torch.zeros(nb, att.shape[0]*att.shape[1], device=dev)
        for b in range(nb):
            feat[b] = att[:, :, torch.tensor(bp[b][-summ:], device=dev)].mean(-1).reshape(-1).float()
        if step > 0:
            steps.append(sorted(active)); fracs.append((sum(clen[b] for b in active) + nb*summ) / total)
        active = set(router(feat).squeeze(-1).argsort(descending=True)[:k].tolist())
        seq.append(nxt); bl.append(-2); sk.append(False)
        if nxt == tok.eos_token_id: break
    return tok.decode(seq[q0:], skip_special_tokens=True), steps, (sum(fracs)/len(fracs) if fracs else 1.0), total, nb


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True)
    ap.add_argument("--router", required=True); ap.add_argument("--k", type=int, default=2)
    args = ap.parse_args(); dev = "cuda"
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16,
              attn_implementation="eager", device_map="cuda").eval()
    ck = torch.load(args.router, map_location=dev); summ = ck["summary_tokens"]
    router = nn.Linear(ck["in_dim"], 1).to(dev); router.load_state_dict(ck["state_dict"]); router.eval()

    print("="*90 + "\nSOURCE DOCUMENTS (blocks):")
    for i, d in enumerate(DOCS): print(f"  [block {i}] {d}")

    # (A) generation vs source
    ids, blk, snk, q0 = DM.build(tok, DOCS, Q_MAIN)
    g, steps, frac, total, nb = route_trace(model, tok, ids, blk, snk, q0, router, args.k, summ, 20)
    print(f"\n(A) Q: {Q_MAIN}")
    print(f"    sparse gen: {g.split('Answer:')[-1].strip()[:120]!r}")
    print(f"    -> grounded in block 3 ('River Coll'). canary block 0 NOT needed.")

    # (C) active-context savings
    print(f"\n(C) active-context savings (k={args.k} of {nb} blocks):")
    print(f"    per-step active blocks (first few): {steps[:6]}")
    print(f"    avg block-content tokens attended per decode step: {frac*100:.0f}%  "
          f"(vs 100% dense) -> ~{1/max(frac,1e-9):.1f}x less context read")
    print(f"    (total block-content tokens = {total}; canary block 0 in active steps: "
          f"{sum(0 in s for s in steps)}/{len(steps)})")

    # (B) leakage: ask the CANARY question under active sets that EXCLUDE vs INCLUDE block 0
    idc, blc, snc, q0c = DM.build(tok, DOCS, Q_CANARY)
    excl, _ = gen_fixed(model, tok, idc, blc, snc, q0c, active_set={3, 5}, summ=summ, max_new=16)  # block 0 masked
    incl, _ = gen_fixed(model, tok, idc, blc, snc, q0c, active_set={0, 3}, summ=summ, max_new=16)  # block 0 active
    print(f"\n(B) DATA-LEAKAGE test — canary secret = {CANARY!r} lives in block 0")
    print(f"    Q: {Q_CANARY}")
    print(f"    active={{3,5}} (block 0 MASKED): {excl.split('Answer:')[-1].strip()[:80]!r}")
    print(f"        -> canary leaked? {CANARY in excl}")
    print(f"    active={{0,3}} (block 0 ACTIVE): {incl.split('Answer:')[-1].strip()[:80]!r}")
    print(f"        -> canary present? {CANARY in incl}")
    print("    PASS: masked block cannot leak" if (CANARY not in excl and CANARY in incl)
          else "    CHECK: unexpected leakage behaviour")


if __name__ == "__main__":
    main()