| """Real-scale benchmark: 50 LongBench RAG samples, real block sizes. Measures (1) quality recovery |
| (gold-answer NLL: dense vs static vs last-8 router-sparse) and (2) active-context savings (fraction of |
| block-content tokens attended per decode step). Uses the trained last-8 router for per-position select. |
| """ |
| from __future__ import annotations |
| import argparse, glob, json, torch, torch.nn as nn |
| import scripts.e2e_nll as N |
|
|
| LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl" |
|
|
|
|
| def load_lb(tok, sink, max_len, max_blocks, n_want): |
| """Real LongBench QA with up to max_blocks real chunks (real block sizes), <= max_len tokens.""" |
| out = [] |
| for line in open(glob.glob(LB)[0]): |
| if len(out) >= n_want: |
| break |
| if len(line) > max_len * 60: |
| continue |
| try: |
| r = json.loads(line) |
| except Exception: |
| continue |
| if not r.get("chunks") or not r.get("answers") or len(r["chunks"]) < 4: |
| continue |
| ids, blk, snk = [], [], [] |
| def add(t, b, sp=False, sk=False): |
| for x in ([t] if sk else tok(t, add_special_tokens=sp)["input_ids"]): |
| ids.append(x); blk.append(b); snk.append(sk) |
| add(r["chunks"][0][:400], -1, sp=True) |
| for bi, c in enumerate(r["chunks"][1:1+max_blocks]): |
| for s in sink: |
| ids.append(s); blk.append(bi); snk.append(True) |
| add("\n" + c, bi) |
| kl0 = len(ids); add(f"\nQuestion: {r['input']}\nAnswer:", -2) |
| ce0 = len(ids); add(" " + r["answers"][0], -2) |
| if len(ids) <= max_len and (max(blk)+1) >= 4 and ce0 < len(ids)-1: |
| out.append((ids, blk, snk, kl0, ce0, "lb")) |
| return out |
|
|
|
|
| def last8_feat(atts, bp, a0, n, nb, summ, dev): |
| """[A, nb, L*H] = per-(layer,head) attn from each answer pos to each block's last-`summ` tokens.""" |
| A = n - a0 |
| feat = torch.zeros(A, nb, atts.shape[0]*atts.shape[1]) |
| for b in range(nb): |
| cols = torch.tensor(bp[b][-summ:], device=dev) |
| feat[:, b, :] = atts[:, :, a0:n, :][:, :, :, cols].mean(-1).permute(2,0,1).reshape(A,-1).cpu() |
| return feat |
|
|
|
|
| @torch.no_grad() |
| 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) |
| ap.add_argument("--n", type=int, default=50); ap.add_argument("--max-len", type=int, default=3500) |
| ap.add_argument("--max-blocks", type=int, default=16) |
| 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() |
| sink = tok("\n", add_special_tokens=False)["input_ids"]*4 |
|
|
| ex = load_lb(tok, sink, args.max_len, args.max_blocks, args.n) |
| print(f"benchmarking {len(ex)} real LongBench QA samples (k={args.k}, summary={summ}, <= {args.max_blocks} blocks) ...") |
| nll = {m: [] for m in ["dense","static","router"]} |
| fracs=[]; nblocks=[]; blk_tok=[]; ctx_tok=[] |
| for ids, blk, snk, kl0, ce0, src in ex: |
| n=len(ids); nb=max(blk)+1 |
| if nb<3 or ce0>=n-1: continue |
| bp = D_block_pos(blk, snk, nb) |
| clen = {b: len(bp[b]) for b in range(nb)}; total=sum(clen.values()) |
| dense_rows={i:set(range(nb)) for i in range(n-kl0)} |
| out = model(input_ids=torch.tensor([ids],device=dev), |
| attention_mask=N.mask_for(blk,snk,kl0,nb,dense_rows,dev), output_attentions=True) |
| atts = torch.stack(out.attentions,0)[:,0].float() |
| A = n-ce0 |
| feat = last8_feat(atts, bp, ce0, n, nb, summ, dev) |
| sc = router(feat.to(dev)).squeeze(-1).cpu() |
| rsel = {t: set(sc[t].argsort(descending=True)[:args.k].tolist()) for t in range(A)} |
| s0 = set(sc[0].argsort(descending=True)[:args.k].tolist()) |
| def rows(sel_fn): |
| r={} |
| for row in range(kl0, n): |
| r[row-kl0] = set(range(nb)) if row<ce0 else sel_fn(row-ce0) |
| return r |
| nll["dense"].append(N.gold_nll(model, ids, N.mask_for(blk,snk,kl0,nb,{i:set(range(nb)) for i in range(n-kl0)},dev), ce0, dev).mean().item()) |
| nll["static"].append(N.gold_nll(model, ids, N.mask_for(blk,snk,kl0,nb,rows(lambda t:s0),dev), ce0, dev).mean().item()) |
| nll["router"].append(N.gold_nll(model, ids, N.mask_for(blk,snk,kl0,nb,rows(lambda t:rsel[t]),dev), ce0, dev).mean().item()) |
| |
| fr = [(sum(clen[b] for b in rsel[t]) + sum(min(summ,clen[b]) for b in range(nb)))/max(total,1) for t in range(A)] |
| fracs.append(sum(fr)/len(fr)); nblocks.append(nb); blk_tok.append(total/nb); ctx_tok.append(total) |
| mean=lambda x: sum(x)/len(x) |
| d,s,r = mean(nll["dense"]), mean(nll["static"]), mean(nll["router"]) |
| print(f"\n=== quality recovery (gold-answer NLL over {len(fracs)} samples) ===") |
| print(f" dense (all blocks) : {d:.3f}") |
| print(f" static (freeze top-{args.k}) : {s:.3f}") |
| print(f" router last-{summ} (top-{args.k}) : {r:.3f} -> recovers {(s-r)/(s-d)*100:.0f}% of the static->dense gap") |
| print(f"\n=== active-context savings ===") |
| print(f" avg blocks/sample: {mean(nblocks):.1f} | avg block size: {mean(blk_tok):.0f} tok | avg context: {mean(ctx_tok):.0f} tok") |
| print(f" avg block-content tokens attended per decode step: {mean(fracs)*100:.0f}% -> ~{1/mean(fracs):.1f}x less context read") |
|
|
|
|
| def D_block_pos(blk, snk, nb): |
| bp={b:[] for b in range(nb)} |
| for j,(b,s) in enumerate(zip(blk,snk)): |
| if b>=0 and not s: bp[b].append(j) |
| return bp |
|
|
|
|
| if __name__=="__main__": |
| main() |
|
|
