hxia7/qwen3-4b-blockdist
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: hxia7/qwen3-4b-blockdist
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-07-27 06:09:10 +08:00
commit b315e39b60
22 changed files with 1799 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
eval/EVAL.md Normal file
View File

@@ -0,0 +1,44 @@
# Evaluation & reproduction — block-distilled Qwen3 + landmark router
All numbers below are produced by the scripts in `scripts/`. `examples_dump.json` contains the actual
per-example inputs (documents + question + gold) and outputs (dense vs sparse) for the accuracy test.
## Headline results (Qwen3-14B, last-8 summary router)
| check | script | result |
|---|---|---|
| **Accuracy — no drop** | `acc_regress.py` | per-example on 80 real LongBench QA: dense 32/80, sparse 33/80, **0 regressions**, +1 gain |
| **Quality (NLL)** | `bench_real.py` | 50 real samples: router k=3 within ~3% NLL of dense, k=2 within ~12% |
| **No data leakage** | `audit_sparse.py` | canary secret in a masked block **cannot** be reproduced (model hallucinates); with block active it can |
| **Router — no train leakage** | `router_leak_check.py` | on unseen docs, coverage 0.92@k=2 ≈ training eval 0.96 |
| **FlexAttention stable** | `flex_attn_test.py` | flex block-attn == eager (max\|Δ\|=0.004), deterministic, T=16384 at 0.2 GB |
| **Decode savings (weights excluded)** | `decode_savings.py` | KV-read HBM & attention-FLOPs each −4× (2k ctx) … ~25× (long ctx); cap = block/summary |
## How to run
```bash
pip install torch transformers peft accelerate
MODEL=hxia7/qwen3-14b-blockdist # or a local path to the merged model
ROUTER=router.pt # ships in the repo (summary_tokens=8)
python scripts/demo_infer.py --model $MODEL --router $ROUTER --k 2 # generation demo
python scripts/audit_sparse.py --model $MODEL --router $ROUTER # leakage + savings
python scripts/acc_regress.py --model $MODEL --router $ROUTER --n 80 --k 2 # accuracy regression
python scripts/bench_real.py --model $MODEL --router $ROUTER --k 2 # NLL + savings
python scripts/router_leak_check.py --model $MODEL --router $ROUTER # router generalization
python scripts/flex_attn_test.py # flex kernel validation
python scripts/decode_savings.py # HBM/FLOP accounting
```
Datasets (LongBench-Seg etc.) load from paths hard-coded in the scripts — adjust to your HF cache.
## Honest caveats
- Absolute LongBench accuracy (~40%) is a base-model + strict-string-match artifact; dense and sparse
are near-identical — the point is **0 regressions**, not the absolute number.
- Block-sparse decode is a **KV-bandwidth (HBM) optimization**; FLOP savings only matter at very long
context. Big wins are in the long-context × batch serving regime.
- Cross-block **aggregation** queries ("list all X") are a limit of block attention itself (even dense
struggles) → serve those via dense fallback.
- Eager `output_attentions` is O(T²) and OOMs past ~2k tokens; use FlexAttention (validated here) to
scale the block mask in a real serving kernel.

89
eval/scripts/acc_regress.py Normal file
View File

@@ -0,0 +1,89 @@
"""Honest accuracy-retention test: per-example regression analysis, not aggregate F1.
For each real LongBench QA: is dense correct? is sparse correct? Then count REGRESSIONS (dense right ->
sparse wrong) vs GAINS (dense wrong -> sparse right). Aggregate F1 can hide regressions; this can't.
Correct = gold answer (normalized) contained in the generated first line.
"""
from __future__ import annotations
import argparse, glob, json, string, torch, torch.nn as nn
import scripts.demo_infer as DM
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"
def norm(s):
return " ".join("".join(c for c in s.lower() if c not in string.punctuation).split())
def load(tok, sink, max_len, n, max_ans=50):
out = []
for line in open(glob.glob(LB)[0]):
if len(out) >= n: 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"]) < 3: continue
ans = r["answers"][0]
if not (1 <= len(ans) <= max_ans): 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][:600], -1, sp=True)
for bi,c in enumerate(r["chunks"][1:11]):
for s in sink: ids.append(s); blk.append(bi); snk.append(True)
add("\n"+c, bi)
q0=len(ids); add(f"\nQuestion: {r['input']}\nAnswer:", -2)
if len(ids) <= max_len and (max(blk)+1) >= 3:
out.append((ids, blk, snk, q0, ans, r.get("dataset","?"), r["input"], list(r["chunks"][1:11])))
return out
@torch.no_grad()
def main():
ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True)
ap.add_argument("--router", required=True); ap.add_argument("--n", type=int, default=80)
ap.add_argument("--k", type=int, default=2); ap.add_argument("--max-len", type=int, default=2000)
ap.add_argument("--max-new", type=int, default=24); ap.add_argument("--dump", default=None)
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(tok, sink, args.max_len, args.n)
print(f"per-example regression test on {len(ex)} short-answer LongBench QA (k={args.k}) ...\n")
dh=sh=reg=gain=both=0; regs=[]; dump=[]
for ids, blk, snk, q0, gold, ds, question, docs in ex:
gd = DM.gen(model,tok,ids,blk,snk,q0,'dense',None,0,args.max_new,0,summ).split("Answer:")[-1].split("\n")[0]
gs = DM.gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,0,summ).split("Answer:")[-1].split("\n")[0]
dc = norm(gold) in norm(gd); sc = norm(gold) in norm(gs)
dh+=dc; sh+=sc; both+=(dc and sc)
if dc and not sc: reg+=1; regs.append((gold, gd.strip()[:50], gs.strip()[:50], ds))
if sc and not dc: gain+=1
dump.append({"dataset": ds, "question": question, "documents": docs, "gold": gold,
"dense_output": gd.strip(), "sparse_output": gs.strip(),
"dense_correct": bool(dc), "sparse_correct": bool(sc)})
n=len(ex)
if args.dump:
import json as J
J.dump({"model": args.model, "k": args.k, "summary_tokens": summ,
"summary": {"n": n, "dense_correct": dh, "sparse_correct": sh,
"regressions": reg, "gains": gain}, "examples": dump},
open(args.dump, "w"), indent=2, ensure_ascii=False)
print(f"dumped {n} examples -> {args.dump}")
print(f"dense correct : {dh}/{n} ({dh/n*100:.0f}%)")
print(f"sparse correct: {sh}/{n} ({sh/n*100:.0f}%)")
print(f"REGRESSIONS (dense right -> sparse wrong): {reg}/{n}")
print(f"GAINS (dense wrong -> sparse right): {gain}/{n}")
print(f"net change: {sh-dh:+d} (agreement both-correct: {both})")
if regs:
print("\nregression cases (gold | dense | sparse | task):")
for g,d,s,ds in regs[:8]:
print(f" gold={g!r} | dense={d!r} | sparse={s!r} | {ds}")
if __name__ == "__main__":
main()

126
eval/scripts/audit_sparse.py Normal file
View File

@@ -0,0 +1,126 @@
"""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()

115
eval/scripts/bench_real.py Normal file
View File

@@ -0,0 +1,115 @@
"""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() # [A, nb]
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())
# savings: attended content fraction per answer pos = (active content + nb*summ)/total
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()

28
eval/scripts/decode_savings.py Normal file
View File

@@ -0,0 +1,28 @@
"""Decode-time savings from block-sparse decode, Qwen3-14B — WEIGHTS EXCLUDED.
Per the agreed accounting we drop the fixed 28 GB weight read (HBM) and the fixed FFN/linear FLOPs, and
report only the subsystem block-sparsity touches: KV-cache reads (HBM) and attention FLOPs. Both scale
with the number of ATTENDED context tokens, so they save the same fraction (and it's batch-independent).
Sparse attends ~k blocks' content + one summary per block: attended ≈ k*block + (S/block)*summ.
As S grows this is ~flat vs dense's O(S), so savings grow — but the per-block resident summary caps it:
attended/S -> summ/block, so max saving ≈ block/summ.
"""
KVtok = 2*40*8*128*2 # Qwen3-14B: KV bytes/token/seq = 163,840
BLOCK, K, SUMM = 200, 2, 8 # measured recipe: ~200-tok blocks, k=2 active, 8-token summary/block
def attended(S):
n = max(1, S // BLOCK)
return K*BLOCK + n*SUMM
print(f"Qwen3-14B decode savings (WEIGHTS EXCLUDED; KV-read HBM & attention-FLOPs; batch-independent %)")
print(f"block={BLOCK} k={K} summary={SUMM} -> savings cap ≈ block/summary = {BLOCK/SUMM:.0f}x\n")
print(f"{'ctx S':>8} | {'attended d→s':>16} | {'KV read/seq d→s':>20} | {'saved':>6} | {'factor':>7}")
for S in [2048, 4096, 8192, 32768, 131072, 524288]:
a = attended(S); kd, ks = KVtok*S/1e6, KVtok*a/1e6
print(f"{S:>8} | {S:>6} → {a:<7} | {kd:>7.1f} → {ks:>6.2f} MB | {(1-a/S)*100:>5.0f}% | {S/a:>5.0f}x")
print("\nBoth KV-read HBM and attention FLOPs drop by this same factor (4x @2k ... ~25x long-context).")
print("Accuracy is retained (sparse ~= dense F1). This is a KV-bandwidth optimization for long-context")
print("batched serving; to raise the cap, use fewer summary tokens (trades off selection quality).")

157
eval/scripts/demo_infer.py Normal file
View File

@@ -0,0 +1,157 @@
"""Qualitative generation demo: full-attn vs dense block-attn vs sparse (router top-k).
Runs several diverse RAG examples through the block-distilled model and prints actual generations so we
can eyeball quality (not just NLL). Blocks = documents (+4 sink tokens each); landmark = last content
token; router selects top-k blocks per decode step.
"""
from __future__ import annotations
import argparse
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
NEG = -1e9
N_SINKS = 4
SYSTEM = "You are a helpful assistant. Answer the question using ONLY the documents. Be concise."
EXAMPLES = [
# (docs, question) — mix of counterfactual (must read), multi-hop, real, numeric, distractor-heavy
([ "The Zorban Reactor at Vale Station reached full output in the year 3471.",
"Mount Kilimanjaro, in Tanzania, is the highest mountain in Africa.",
"The Great Wall of China stretches across northern China.",
"Photosynthesis converts carbon dioxide and water into glucose and oxygen." ],
"In what year did the Zorban Reactor reach full output?"),
([ "The CEO of Vantacorp is Marisa Quen.",
"Marisa Quen was born in the city of Drennholm.",
"Drennholm is famous for its glass bridges.",
"The CEO of Bexil Industries is Tomas Ray." ],
"In which city was the CEO of Vantacorp born?"),
([ "Python is a programming language created by Guido van Rossum, first released in 1991.",
"Rust is a systems language emphasizing memory safety.",
"The mitochondrion is the powerhouse of the cell.",
"Java was developed by James Gosling at Sun Microsystems." ],
"Who created Python and in what year was it first released?"),
([ "Order #4471 shipped on March 3 and contains 2 laptops.",
"Order #4472 shipped on March 5 and contains 1 monitor.",
"Order #4473 is delayed and contains 3 keyboards.",
"Order #4474 shipped on March 6 and contains 1 laptop." ],
"Which orders shipped in March and what did order #4473 contain?"),
([ "The Treaty of Kessel was signed in 1804 between Aldoria and Brenne.",
"Aldoria's capital is Feldspar City.",
"Brenne is known for its copper mines.",
"The Kessel treaty ended the Seven Rivers War.",
"Feldspar City sits on the river Onn.",
"Copper was Brenne's main export in the 1800s." ],
"What war did the Treaty of Kessel end, and in what year was it signed?"),
]
def build(tok, docs, question):
sink = tok("\n", add_special_tokens=False)["input_ids"] * N_SINKS
ids, blk, snk = [], [], []
def add(t, b, special=False):
for x in tok(t, add_special_tokens=special)["input_ids"]:
ids.append(x); blk.append(b); snk.append(False)
add(SYSTEM + "\n", -1, special=True)
for bi, d in enumerate(docs):
for s in sink:
ids.append(s); blk.append(bi); snk.append(True)
add(f"[Document {bi+1}] {d}\n", bi)
q0 = len(ids); add(f"Question: {question}\nAnswer:", -2)
return ids, blk, snk, q0
def landmarks(blk, snk, nb):
lm, isc = {}, [False]*len(blk)
for j,(b,s) in enumerate(zip(blk,snk)):
if b>=0 and not s: isc[j]=True; lm[b]=j
return [lm[b] for b in range(nb)], isc
def mask(blk, snk, q0, nb, active_rows, n, dev, full=False):
b = torch.tensor(blk+[-2]*(n-len(blk)), device=dev)
causal = torch.tril(torch.ones(n,n,dtype=torch.bool,device=dev))
if full:
return torch.where(causal,0.0,NEG).view(1,1,n,n).float()
bi,bj=b.view(n,1),b.view(1,n)
if active_rows is None:
allowed=((bj==-1)|(bi==bj)|(bi==-2))&causal
return torch.where(allowed,0.0,NEG).view(1,1,n,n).float()
lm,isc=landmarks(blk,snk,nb)
is_lm=torch.zeros(n,dtype=torch.bool,device=dev); is_lm[torch.tensor(lm,device=dev)]=True
isc_t=torch.zeros(n,dtype=torch.bool,device=dev); isc_t[:len(isc)]=torch.tensor(isc,device=dev)
static=(b==-1)|(b==-2)|is_lm
allowed=torch.zeros(n,n,dtype=torch.bool,device=dev)
allowed[:q0]=(((bj==-1)|(bi==bj)|(bi==-2))&causal)[:q0]
for i in range(q0,n):
act=active_rows.get(i-q0,set(range(nb))); vis=static.clone()
if act: vis=vis|(isc_t&torch.isin(bj.view(n),torch.tensor(sorted(act),device=dev)))
allowed[i]=vis
allowed&=causal
return torch.where(allowed,0.0,NEG).view(1,1,n,n).float()
def block_content_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
@torch.no_grad()
def gen(model, tok, ids, blk, snk, q0, mode, router=None, k=3, max_new=24, sticky=0, summ=1):
"""mode: full | dense | sparse. summ = block summary size (last-N content tokens) for router feats.
sticky>0 -> keep a block active for `sticky` more steps after last selected."""
dev=model.device; nb=max(blk)+1
seq,bl,sk=list(ids),list(blk),list(snk)
ar={} if mode=="sparse" else None
active=set(range(nb)); last_seen={}
for step in range(max_new):
n=len(seq)
if mode=="sparse": ar[n-q0]=set(range(nb)) if step==0 else set(active)
m=mask(bl,sk,q0,nb,ar,n,dev,full=(mode=="full"))
out=model(input_ids=torch.tensor([seq],device=dev),attention_mask=m,
output_attentions=(mode=="sparse"))
nxt=int(out.logits[0,-1].argmax())
if mode=="sparse":
att=torch.stack(out.attentions,0)[:,0,:,-1,:] # [L,H,n] attn of last pos
bp=block_content_pos(bl,sk,nb)
feat=torch.zeros(nb, att.shape[0]*att.shape[1], device=dev)
for b in range(nb):
cols=torch.tensor(bp[b][-summ:],device=dev)
feat[b]=att[:,:,cols].mean(-1).reshape(-1).float() # attn to block b's last-summ tokens
top=router(feat).squeeze(-1).argsort(descending=True)[:k].tolist()
for b in top: last_seen[b]=step
active=set(b for b,s in last_seen.items() if step-s<=sticky) if sticky else set(top)
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).strip().replace("\n"," ")
def main():
ap=argparse.ArgumentParser()
ap.add_argument("--model",required=True); ap.add_argument("--router",default=None)
ap.add_argument("--k",type=int,default=2); ap.add_argument("--max-new",type=int,default=24)
ap.add_argument("--sticky",type=int,default=0,help="keep a block active this many steps after last selected")
args=ap.parse_args()
tok=AutoTokenizer.from_pretrained(args.model)
model=AutoModelForCausalLM.from_pretrained(args.model,dtype=torch.bfloat16,
attn_implementation="eager",device_map="cuda").eval()
router=None; summ=1
if args.router:
ck=torch.load(args.router,map_location="cuda")
router=nn.Linear(ck["in_dim"],1).to("cuda"); router.load_state_dict(ck["state_dict"]); router.eval()
summ=ck.get("summary_tokens",1)
for docs,q in EXAMPLES:
ids,blk,snk,q0=build(tok,docs,q)
print("\n"+"="*100); print(f"Q: {q} ({max(blk)+1} blocks, {q0} ctx tok)")
print(f" full-attn : {gen(model,tok,ids,blk,snk,q0,'full',max_new=args.max_new)!r}")
print(f" block dense : {gen(model,tok,ids,blk,snk,q0,'dense',max_new=args.max_new)!r}")
if router is not None:
print(f" sparse k={args.k} (summ{summ}): {gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,0,summ)!r}")
print(f" sparse k={args.k} sticky{args.sticky}: {gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,args.sticky,summ)!r}")
if __name__=="__main__":
main()

218
eval/scripts/distill_lora.py Normal file
View File

@@ -0,0 +1,218 @@
"""LoRA block distillation for Qwen3-14B on one H200.
ONE frozen base model plays both roles via adapter toggling (memory: ~1 model, not 2):
teacher = LoRA disabled + FULL attention (no_grad)
student = LoRA enabled + BLOCK attention (grad, only LoRA params train)
Layout per example: [system][block_1..n-1 isolated (+4 sink tokens each)][last segment attends to all].
Loss on the last segment: KL(teacher_full || student_block) + damage-weighted CE(student_block, gold),
w = max(CE(teacher_block) - CE(teacher_full), 0) * alpha + beta.
Data (real, segmented):
LongbenchSeg / LoCoMoSeg : chunks=blocks, last segment = "Question: {q}\nAnswer: {a}" (CE on answer)
SemanticSeg : <cut N> markers -> blocks, last segment = final block (CE on whole block)
Run: source env.sh && HF_HOME=/projects/bdjx/hxia3/hf_cache_proj \
python scripts/distill_lora.py --model Qwen/Qwen3-14B --steps 3000 --out checkpoints/qwen3-14b-blockdist-lora
"""
from __future__ import annotations
import argparse
import glob
import json
import random
import re
import torch
import torch.nn.functional as F
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer
NEG = -1e9
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"
LOCOMO = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LoCoMoSeg/snapshots/*/locomo_segmented.json"
SEMSEG = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--SemanticSeg/snapshots/*/cut_*.jsonl"
def _tok(tok, text, special=False):
return tok(text, add_special_tokens=special)["input_ids"]
def emit(tok, sink_ids, system, blocks, last_seg, ce_from_ratio=0.0):
"""Return ids, blk, sink, kl_start, ce_start for one training example."""
ids, blk, sink = [], [], []
def add(text, b, is_sink=False, special=False):
t = [text] if is_sink else _tok(tok, text, special)
for x in t:
ids.append(x); blk.append(b); sink.append(is_sink)
if system:
for x in _tok(tok, system + "\n", special=True):
ids.append(x); blk.append(-1); sink.append(False)
for bi, bt in enumerate(blocks):
for s in sink_ids:
ids.append(s); blk.append(bi); sink.append(True)
add("\n" + bt.strip() + "\n", bi)
kl_start = len(ids)
seg_ids = _tok(tok, last_seg, special=False)
ce_start = kl_start + int(len(seg_ids) * ce_from_ratio)
for x in seg_ids:
ids.append(x); blk.append(-2); sink.append(False)
return ids, blk, sink, kl_start, ce_start
def group_blocks(segs, target_chars=1000, max_blocks=8):
"""Merge tiny segments into ~target_chars blocks (char proxy for ~256 tok); cap count."""
out, cur, cur_n = [], [], 0
for c in segs:
cur.append(c); cur_n += len(c)
if cur_n >= target_chars:
out.append(" ".join(cur)); cur, cur_n = [], 0
if len(out) >= max_blocks:
break
if cur and len(out) < max_blocks:
out.append(" ".join(cur))
return out
def load_data(tok, sink_ids, max_len, max_per_source, seed):
rng = random.Random(seed)
data = []
CHAR = max_len * 4 # ~4 chars/token budget for char prefilters
# LongbenchSeg (QA) — most examples are long; cheap char prefilter before json.loads
lb = 0
for line in open(glob.glob(LB)[0]):
if lb >= max_per_source:
break
if len(line) > CHAR * 40: # skip only pathologically long docs; blocks are capped below
continue
try:
r = json.loads(line)
except Exception:
continue
if not r.get("chunks") or not r.get("answers") or len(r["chunks"]) < 3:
continue
blocks = [c[:CHAR // 4] for c in r["chunks"][1:9]] # cap each block's chars
ex = emit(tok, sink_ids, r["chunks"][0][:600], blocks,
f"Question: {r['input']}\nAnswer: {r['answers'][0]}")
q_ids = _tok(tok, f"Question: {r['input']}\nAnswer:")
ex = (ex[0], ex[1], ex[2], ex[3], ex[3] + len(q_ids))
if len(ex[0]) <= max_len and ex[4] < len(ex[0]):
data.append((*ex, "lb")); lb += 1
print(f" LongbenchSeg: {lb} examples", flush=True)
# SemanticSeg (text) — big diverse corpus, bounded scan per file, char-based blocks
files = sorted(glob.glob(SEMSEG)); rng.shuffle(files)
per_file = max(1, max_per_source // max(1, len(files)))
for f in files:
got = 0; scanned = 0
for line in open(f):
if got >= per_file or scanned > per_file * 20:
break
scanned += 1
if len(line) > CHAR * 2:
continue
try:
txt = json.loads(line).get("cut_item", [{}])[0].get("txt_marker", "")
except Exception:
continue
segs = [s for s in re.split(r'<cut \d+>', txt) if s.strip()]
if len(segs) < 4:
continue
blocks = group_blocks(segs, target_chars=1000, max_blocks=8)
if len(blocks) < 3:
continue
ex = emit(tok, sink_ids, "", blocks[:-1], blocks[-1])
if len(ex[0]) <= max_len and ex[3] < len(ex[0]) - 1:
data.append((*ex, "ss")); got += 1
print(f" {f.split('cut_')[-1].split('.')[0]}: {got}", flush=True)
rng.shuffle(data)
return data
def masks(blk, dev):
n = len(blk); b = torch.tensor(blk, device=dev)
causal = torch.tril(torch.ones(n, n, dtype=torch.bool, device=dev))
bi, bj = b.view(n, 1), b.view(1, n)
full = torch.where(causal, 0.0, NEG).view(1, 1, n, n).float()
allowed = ((bj == -1) | (bi == -2) | (bi == bj)) & causal
block = torch.where(allowed, 0.0, NEG).view(1, 1, n, n).float()
return full, block
def ce_span(logits, ids, start):
lp = F.log_softmax(logits[0, start - 1:-1].float(), -1)
return -lp.gather(-1, ids[0, start:].view(-1, 1)).squeeze(-1)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen3-14B")
ap.add_argument("--steps", type=int, default=3000)
ap.add_argument("--lr", type=float, default=1e-4)
ap.add_argument("--accum", type=int, default=8)
ap.add_argument("--rank", type=int, default=16)
ap.add_argument("--n-sinks", type=int, default=4)
ap.add_argument("--alpha", type=float, default=0.3)
ap.add_argument("--beta", type=float, default=0.1)
ap.add_argument("--max-len", type=int, default=1536)
ap.add_argument("--max-per-source", type=int, default=6000)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--out", default="/projects/bdjx/hxia3/lazy2/checkpoints/qwen3-14b-blockdist-lora")
ap.add_argument("--save-every", type=int, default=500)
args = ap.parse_args()
dev = "cuda"; rng = random.Random(args.seed)
print(f"loading {args.model} (bf16, eager, grad-checkpoint) ...")
tok = AutoTokenizer.from_pretrained(args.model)
base = AutoModelForCausalLM.from_pretrained(
args.model, dtype=torch.bfloat16, attn_implementation="eager").to(dev)
base.gradient_checkpointing_enable()
base.enable_input_require_grads()
lora = LoraConfig(r=args.rank, lora_alpha=args.rank * 2, lora_dropout=0.0,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"])
model = get_peft_model(base, lora)
model.print_trainable_parameters()
sink_ids = tok("\n", add_special_tokens=False)["input_ids"] * args.n_sinks
print("loading data ...")
data = load_data(tok, sink_ids, args.max_len, args.max_per_source, args.seed)
from collections import Counter
print(f"examples: {len(data)} | sources: {dict(Counter(d[5] for d in data))}")
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr, betas=(0.9, 0.95))
opt.zero_grad()
model.train()
run = 0.0
for step in range(1, args.steps + 1):
ids, blk, sink, kl0, ce0, src = data[rng.randrange(len(data))]
t = torch.tensor([ids], device=dev)
fm, bm = masks(blk, dev)
with torch.no_grad(), model.disable_adapter(): # TEACHER (full attn, no LoRA)
tf = model(input_ids=t, attention_mask=fm).logits
tf_ce = ce_span(tf, t, ce0)
tb_ce = ce_span(model(input_ids=t, attention_mask=bm).logits, t, ce0)
w = torch.clamp(tb_ce - tf_ce, min=0) * args.alpha + args.beta
sb = model(input_ids=t, attention_mask=bm).logits # STUDENT (block attn, LoRA)
kl = F.kl_div(F.log_softmax(sb[0, kl0:].float(), -1),
F.log_softmax(tf[0, kl0:].float(), -1),
reduction="none", log_target=True).sum(-1).mean()
wce = (w * ce_span(sb, t, ce0)).mean()
loss = (kl + wce) / args.accum
loss.backward(); run += loss.item() * args.accum
if step % args.accum == 0:
torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
opt.step(); opt.zero_grad()
if step % 10 == 0:
print(f"step {step:5d}/{args.steps} | loss {run/10:.3f} | KL {kl.item():.3f} | wCE {wce.item():.3f}", flush=True)
run = 0.0
if step % args.save_every == 0:
model.save_pretrained(args.out); print(f" saved adapter -> {args.out}", flush=True)
model.save_pretrained(args.out)
tok.save_pretrained(args.out)
print(f"done. LoRA adapter at {args.out}")
if __name__ == "__main__":
main()

178
eval/scripts/e2e_nll.py Normal file
View File

@@ -0,0 +1,178 @@
"""Step 3d (clean metric) — end-to-end re-route validated by NLL, not generation.
Teacher-force the gold answer; at each answer position choose the ACTIVE block set, build the
attention mask, and measure the gold token's NLL. No generation loop -> immune to the rambly-output /
memory-fallback noise that broke the accuracy metric. Modes (gold NLL, lower = better):
dense : all blocks active (ceiling)
static : top-k blocks frozen from the first answer position's real attention
landmark-reroute : top-k by attention to each block's resident LANDMARK token, re-picked per position
oracle-reroute : top-k by full-block attention mass, re-picked per position (best-case selection)
If landmark-reroute NLL ~= dense ~= oracle-reroute, and static NLL is high, on-demand re-route works.
Counterfactual facts (fictional capitals) force reading the docs; gold blocks are randomly placed.
Run: source env.sh && python scripts/e2e_nll.py --model Qwen/Qwen3-8B --n-hops 3
"""
from __future__ import annotations
import argparse
import random
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
NEG = -1e9
SYS = "You are a helpful assistant. Use the documents to answer. Give only the capitals, comma-separated, in order."
PAIRS = [("France", "Paris"), ("Japan", "Tokyo"), ("Egypt", "Cairo"), ("Brazil", "Brasilia"),
("Canada", "Ottawa"), ("Kenya", "Nairobi"), ("Norway", "Oslo"), ("Peru", "Lima"),
("India", "Delhi"), ("Spain", "Madrid"), ("Greece", "Athens"), ("Cuba", "Havana")]
FICT = ["Zandar", "Qeropolis", "Vunbry", "Marnex", "Trellin", "Osquith", "Balmoor", "Kestol",
"Yarrow", "Drennik", "Fenwick", "Lorvath", "Sundeep", "Ashgard", "Perrin", "Wexley"]
def build(tok, rng, n_blocks, n_hops, sink_ids):
docs = rng.sample(range(len(PAIRS)), n_blocks) # which countries are documents
asked = rng.sample(docs, n_hops) # asked subset (gold blocks, random placement)
slot = {d: bi for bi, d in enumerate(docs)}
fic = FICT[:]; rng.shuffle(fic)
cap = {d: fic[i] for i, d in enumerate(docs)} # fictional capital per country
ids, blk, sink, gold = [], [], [], [] # gold[t] = active gold block per answer token
def add(text, b, is_sink=False, special=False, gblock=None, is_ans=False):
t = tok(text, add_special_tokens=special)["input_ids"]
ids.extend(t); blk.extend([b] * len(t)); sink.extend([is_sink] * len(t))
if is_ans:
gold.extend([gblock] * len(t))
add(SYS + "\n", -1, special=True)
for d in docs:
for s in sink_ids:
ids.append(s); blk.append(slot[d]); sink.append(True)
add(f"[Document {slot[d]+1}] The capital of {PAIRS[d][0]} is {cap[d]}.\n", slot[d])
add(f"List the capitals of: {', '.join(PAIRS[d][0] for d in asked)}.\nAnswer:", -2)
a0 = len(ids)
for j, d in enumerate(asked):
prefix = " " + cap[d] if j == 0 else ", " + cap[d]
add(prefix, -2, gblock=slot[d], is_ans=True)
return ids, blk, sink, a0, gold
def landmarks_content(blk, sink, n_blocks):
lm, isc = {}, [False] * len(blk)
for j, (b, s) in enumerate(zip(blk, sink)):
if b >= 0 and not s:
isc[j] = True; lm[b] = j
return [lm[b] for b in range(n_blocks)], isc
def mask_for(blk, sink, a0, n_blocks, active_rows, dev):
n = len(blk)
b = torch.tensor(blk, device=dev)
lm, isc = landmarks_content(blk, sink, n_blocks)
causal = torch.tril(torch.ones(n, n, dtype=torch.bool, device=dev))
is_static = (b == -1) | (b == -2) # system + question/answer visible to answer rows
is_lm = torch.zeros(n, dtype=torch.bool, device=dev); is_lm[torch.tensor(lm, device=dev)] = True
static_key = is_static | is_lm
isc_t = torch.tensor(isc, device=dev)
allowed = torch.zeros(n, n, dtype=torch.bool, device=dev)
bi, bj = b.view(n, 1), b.view(1, n)
allowed[:a0] = ((bj == -1) | (bi == bj) | (bi == -2))[:a0] # prompt: block attention
for i in range(a0, n):
act = active_rows[i - a0]
vis = static_key.clone()
if act:
vis = vis | (isc_t & torch.isin(bj.view(n)[:], torch.tensor(sorted(act), device=dev)))
allowed[i] = vis
allowed &= causal
return torch.where(allowed, 0.0, NEG).view(1, 1, n, n).float()
@torch.no_grad()
def gold_nll(model, ids, mask, a0, dev):
t = torch.tensor([ids], device=dev)
lp = F.log_softmax(model(input_ids=t, attention_mask=mask).logits[0].float(), -1)
tgt = t[0, a0:]
return -lp[a0 - 1:-1].gather(-1, tgt.view(-1, 1)).squeeze(-1) # per-answer-token NLL
@torch.no_grad()
def signals(model, ids, blk, sink, a0, n_blocks, dev):
"""Dense pass -> per-answer-position full-block mass and landmark attention."""
n = len(ids)
dense_rows = {i: set(range(n_blocks)) for i in range(n - a0)}
m = mask_for(blk, sink, a0, n_blocks, dense_rows, dev)
out = model(input_ids=torch.tensor([ids], device=dev), attention_mask=m, output_attentions=True)
a = torch.stack(out.attentions, 0)[:, 0].mean(dim=(0, 1)).float() # [n,n]
lm, isc = landmarks_content(blk, sink, n_blocks)
content = torch.zeros(n_blocks, n, device=dev)
for j, c in enumerate(isc):
if c:
content[blk[j], j] = 1.0
full_mass = a[a0:n] @ content.T # [A, nb]
lm_mass = a[a0:n][:, torch.tensor(lm, device=dev)] # [A, nb]
return full_mass.cpu(), lm_mass.cpu(), m
def topk_rows(scores, k, A, static=False):
if static:
sel = set(scores[0].argsort(descending=True)[:k].tolist())
return {t: set(sel) for t in range(A)}
return {t: set(scores[t].argsort(descending=True)[:k].tolist()) for t in range(A)}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen3-8B")
ap.add_argument("--n-blocks", type=int, default=8)
ap.add_argument("--n-hops", type=int, default=3)
ap.add_argument("--k", type=int, default=1)
ap.add_argument("--n-examples", type=int, default=30)
ap.add_argument("--n-sinks", type=int, default=4)
ap.add_argument("--seed", type=int, default=1)
args = ap.parse_args()
dev = "cuda"; rng = random.Random(args.seed)
print(f"loading {args.model} ...")
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(
args.model, dtype=torch.bfloat16, attn_implementation="eager").to(dev).eval()
sink_ids = tok("\n", add_special_tokens=False)["input_ids"] * args.n_sinks
agg = {m: [] for m in ["dense", "static", "landmark", "oracle"]}
cov = {m: [] for m in ["static", "landmark", "oracle"]}
loads = {m: [] for m in ["dense", "static", "landmark", "oracle"]}
for _ in range(args.n_examples):
ids, blk, sink, a0, gold = build(tok, rng, args.n_blocks, args.n_hops, sink_ids)
A = len(ids) - a0
full_mass, lm_mass, dense_mask = signals(model, ids, blk, sink, a0, args.n_blocks, dev)
sel = {
"dense": {t: set(range(args.n_blocks)) for t in range(A)},
"static": topk_rows(full_mass, args.k, A, static=True),
"landmark": topk_rows(lm_mass, args.k, A),
"oracle": topk_rows(full_mass, args.k, A),
}
for mode, rows in sel.items():
mask = dense_mask if mode == "dense" else mask_for(blk, sink, a0, args.n_blocks, rows, dev)
nll = gold_nll(model, ids, mask, a0, dev)
agg[mode].append(nll.mean().item())
loads[mode].append(sum(len(rows[t]) for t in range(A)) / A)
if mode != "dense":
cov[mode].append(sum(gold[t] in rows[t] for t in range(A)) / A)
def mean(x):
return sum(x) / len(x)
print(f"\nmodel={args.model} n_blocks={args.n_blocks} n_hops={args.n_hops} k={args.k} n={args.n_examples}")
print(f"{'mode':>10} | {'gold NLL':>9} | {'avg active':>10} | {'gold-block coverage':>19}")
for mode in ["dense", "static", "landmark", "oracle"]:
c = f"{mean(cov[mode]):.2f}" if mode in cov else "-"
print(f"{mode:>10} | {mean(agg[mode]):>9.3f} | {mean(loads[mode]):>10.2f} | {c:>19}")
d, s, l = mean(agg['dense']), mean(agg['static']), mean(agg['landmark'])
print(f"\nlandmark-reroute closes {(s-l)/(s-d)*100 if s>d else 0:.0f}% of the static->dense NLL gap "
f"at {mean(loads['landmark']):.1f}/{args.n_blocks} blocks active.")
if __name__ == "__main__":
main()

92
eval/scripts/flex_attn_test.py Normal file
View File

@@ -0,0 +1,92 @@
"""Operator-level validation of block attention under FlexAttention.
1) CORRECTNESS: flex block-attention output == reference SDPA-with-additive-block-mask (small T).
2) STABILITY: flex is deterministic across repeated runs.
3) SCALE: flex runs at long T (O(T) memory) where the eager T x T mask would need terabytes.
This validates the kernel/mask logic your serving stack (vLLM) needs, independent of the HF model.
"""
from __future__ import annotations
import argparse, torch
import torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
NEG = torch.finfo(torch.float32).min
def make_block_ids(T, n_blocks, dev):
"""token 0..s = system(-1); then n_blocks equal blocks; last quarter = query(-2)."""
ids = torch.empty(T, dtype=torch.long, device=dev)
s = max(1, T // 20)
q = T - max(1, T // 4)
ids[:s] = -1
ids[q:] = -2
body = q - s
per = max(1, body // n_blocks)
for b in range(n_blocks):
ids[s + b*per : s + (b+1)*per] = b
ids[s + n_blocks*per : q] = n_blocks - 1
return ids
def block_mask_mod(bids):
def mm(b, h, qi, ki):
causal = ki <= qi
bq = bids[qi]; bk = bids[ki]
return causal & ((bk == -1) | (bq == bk) | (bq == -2))
return mm
def ref_sdpa(q, k, v, bids):
T = q.shape[-2]; dev = q.device
causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=dev))
bq = bids.view(T, 1); bk = bids.view(1, T)
allowed = (((bk == -1) | (bq == bk) | (bq == -2)) & causal).view(1, 1, T, T) # boolean mask (no NaN)
return F.scaled_dot_product_attention(q, k, v, attn_mask=allowed)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--H", type=int, default=8); ap.add_argument("--D", type=int, default=128)
ap.add_argument("--T", type=int, default=1024); ap.add_argument("--blocks", type=int, default=8)
ap.add_argument("--big-T", type=int, default=16384)
args = ap.parse_args()
dev = "cuda"; torch.manual_seed(0)
# (1) correctness at moderate T
T = args.T
q = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16)
k = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16)
v = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16)
bids = make_block_ids(T, args.blocks, dev)
bm = create_block_mask(block_mask_mod(bids), B=None, H=None, Q_LEN=T, KV_LEN=T, device=dev)
flex = torch.compile(flex_attention)
o_flex = flex(q, k, v, block_mask=bm)
o_ref = ref_sdpa(q, k, v, bids)
d = (o_flex.float() - o_ref.float()).abs()
print(f"(1) CORRECTNESS T={T} H={args.H} D={args.D} blocks={args.blocks}")
print(f" max|Δ|={d.max():.4f} mean|Δ|={d.mean():.6f} (bf16 noise floor ~1e-2) -> "
f"{'MATCH' if d.max()<0.05 else 'MISMATCH'}")
# (2) determinism
o2 = flex(q, k, v, block_mask=bm)
print(f"(2) STABILITY flex run twice: max|Δ|={(o_flex.float()-o2.float()).abs().max():.6f} "
f"-> {'deterministic' if (o_flex.float()-o2.float()).abs().max()<1e-3 else 'nondeterministic'}")
# (3) scale — flex at big T (eager T x T mask would be H*T*T*2 bytes)
bt = args.big_T
qb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16)
kb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16)
vb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16)
bidsb = make_block_ids(bt, args.blocks * 8, dev)
bmb = create_block_mask(block_mask_mod(bidsb), B=None, H=None, Q_LEN=bt, KV_LEN=bt, device=dev)
torch.cuda.reset_peak_memory_stats()
ob = flex(qb, kb, vb, block_mask=bmb); torch.cuda.synchronize()
peak = torch.cuda.max_memory_allocated() / 1e9
eager_mask_gb = args.H * bt * bt * 2 / 1e9
print(f"(3) SCALE flex at T={bt}: OK, peak mem {peak:.1f} GB, out finite={bool(torch.isfinite(ob).all())}")
print(f" (an eager [1,{args.H},{bt},{bt}] score/mask alone would be ~{eager_mask_gb:.0f} GB -> impossible)")
if __name__ == "__main__":
main()

85
eval/scripts/router_leak_check.py Normal file
View File

@@ -0,0 +1,85 @@
"""Router leakage / generalization check.
The router was trained on LongBench examples [0:N] (disjoint train/eval slices). Here we load the SAVED
router and evaluate its gold-block coverage on a FRESH set of LongBench examples the training never saw
(offset far past N). If coverage on unseen docs ~ training-time eval coverage, there is no train/eval
leakage or overfit. Also checks the serving condition: features come from the resident summary attention,
which IS available at sparse serving (no train-only signal).
"""
from __future__ import annotations
import argparse, glob, json, torch, torch.nn as nn
import scripts.e2e_nll as N
import scripts.train_router_real as TR
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"
def load_offset(tok, sink, max_len, skip, take):
"""Load LongBench QA, SKIPPING the first `skip` valid examples (training set), take next `take`."""
out, seen = [], 0
for line in open(glob.glob(LB)[0]):
if len(out) >= take:
break
if len(line) > max_len * 40:
continue
try:
r = json.loads(line)
except Exception:
continue
if not r.get("chunks") or not r.get("answers") or len(r["chunks"]) < 3:
continue
seen += 1
if seen <= skip: # skip the examples used during router training
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][:600], -1, sp=True)
for bi, c in enumerate(r["chunks"][1:9]):
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) >= 3 and ce0 < len(ids)-1:
out.append((ids, blk, snk, kl0, ce0, "lb"))
return out
@torch.no_grad()
def main():
ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True)
ap.add_argument("--router", required=True); ap.add_argument("--skip", type=int, default=400)
ap.add_argument("--take", type=int, default=60); ap.add_argument("--max-len", type=int, default=1200)
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_offset(tok, sink, args.max_len, args.skip, args.take)
print(f"FRESH examples (skipped first {args.skip} training docs): {len(ex)}")
data = TR.collect(model, ex, summ, dev, args.max_len) # (feat, oracle_top1) via dense forward
ks = [1, 2, 3]
cov = {k: 0.0 for k in ks}; base = {k: 0.0 for k in ks}; tot = 0
for feat, tgt in data:
sc = router(feat.to(dev)).squeeze(-1).cpu()
rank = sc.argsort(-1, descending=True)
mrank = feat.mean(-1).argsort(-1, descending=True) # untrained baseline
for k in ks:
cov[k] += (rank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
base[k] += (mrank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
tot += feat.shape[0]
print(f"\ncoverage of the model's oracle top-1 block on UNSEEN docs ({tot} answer positions):")
print(f"{'k':>3} | {'untrained mean-summary':>22} | {'trained router':>14}")
for k in ks:
print(f"{k:>3} | {base[k]/tot:>22.3f} | {cov[k]/tot:>14.3f}")
print("\nIf trained >> untrained AND ~matches training-time eval (0.96@k=2), the router generalizes")
print("to unseen documents -> no train/eval leakage, no overfit.")
if __name__ == "__main__":
main()

137
eval/scripts/train_router_real.py Normal file
View File

@@ -0,0 +1,137 @@
"""Train the landmark router on REAL RAG data (LongbenchSeg), distilling to the model's own attention.
Fixes the synthetic-capitals distribution gap. No gold-block labels needed: the target at each answer
position is the ORACLE top-1 block = argmax of full-block attention mass there (the model's own
retrieval signal). Router = linear over per-(layer,head) landmark attention -> predict that block.
Run: source env.sh && HF_HOME=/projects/bdjx/hxia3/hf_cache_proj PYTHONPATH=. \
python scripts/train_router_real.py --model hf_release/model --save-router checkpoints/qwen3-14b-router-real.pt
"""
from __future__ import annotations
import argparse
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
import scripts.distill_lora as D # load_data (real LongbenchSeg/SemanticSeg), masks
import scripts.e2e_nll as N # landmarks_content, mask_for
@torch.no_grad()
def collect(model, examples, summary_n, dev, max_tokens=1100):
"""Per example -> (feat [A,nb,L*H], oracle_top1 [A]). feat = per-(layer,head) attention to each
block's SUMMARY = its last `summary_n` content tokens (mean). summary_n=1 -> old single-landmark."""
data = []
for ids, blk, sink, kl0, ce0, src in examples:
if len(ids) > max_tokens:
continue
nb = max(blk) + 1
if nb < 3:
continue
n = len(ids)
rows = {i: set(range(nb)) for i in range(n - kl0)}
mask = N.mask_for(blk, sink, kl0, nb, rows, dev)
out = model(input_ids=torch.tensor([ids], device=dev), attention_mask=mask,
output_attentions=True)
atts = torch.stack(out.attentions, 0)[:, 0].float() # [L,H,n,n]
_, isc = N.landmarks_content(blk, sink, nb)
bpos = {b: [] for b in range(nb)}
content = torch.zeros(nb, n, device=dev)
for j, c in enumerate(isc):
if c:
content[blk[j], j] = 1.0; bpos[blk[j]].append(j)
a0 = ce0
if a0 >= n - 1:
continue
A = n - a0
# feat[t, b] = mean over block b's last-summary_n content tokens of attn, per (layer,head)
feat = torch.zeros(A, nb, atts.shape[0] * atts.shape[1])
for b in range(nb):
cols = torch.tensor(bpos[b][-summary_n:], device=dev)
v = atts[:, :, a0:n, :][:, :, :, cols].mean(-1) # [L,H,A]
feat[:, b, :] = v.permute(2, 0, 1).reshape(A, -1).cpu()
full_mass = (atts.mean(dim=(0, 1))[a0:n] @ content.T)
data.append((feat, full_mass.argmax(-1).cpu()))
return data
def coverage(scorer, data, ks, nb_cap):
cov = {k: 0.0 for k in ks}; tot = 0
for feat, tgt in data:
sc = scorer(feat)
rank = sc.argsort(-1, descending=True)
for k in ks:
cov[k] += (rank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
tot += feat.shape[0]
return {k: cov[k] / tot for k in ks}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--n-examples", type=int, default=350)
ap.add_argument("--max-len", type=int, default=1100)
ap.add_argument("--epochs", type=int, default=400)
ap.add_argument("--wd", type=float, default=1e-3)
ap.add_argument("--summary-tokens", type=int, default=8, help="block summary = last N content tokens")
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--save-router", default=None)
args = ap.parse_args()
dev = "cuda"
print(f"loading {args.model} ...")
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(
args.model, dtype=torch.bfloat16, attn_implementation="eager").to(dev).eval()
sink_ids = tok("\n", add_special_tokens=False)["input_ids"] * 4
print("loading real RAG data (LongbenchSeg) ...")
ex = D.load_data(tok, sink_ids, args.max_len, args.n_examples, args.seed)
ex = [e for e in ex if e[5] == "lb"] # QA (query+answer) examples
random.Random(1).shuffle(ex)
print(f" {len(ex)} QA examples")
print(f"collecting features (summary = last {args.summary_tokens} tokens) + oracle targets ...")
data = collect(model, ex, args.summary_tokens, dev, args.max_len)
print(f" usable: {len(data)} examples, {sum(f.shape[0] for f,_ in data)} answer positions")
n_eval = max(20, len(data) // 5)
ev, tr = data[:n_eval], data[n_eval:]
LH = tr[0][0].shape[-1]; nb_max = max(f.shape[1] for f, _ in data)
print(f" feat dim L*H={LH} | train {len(tr)} / eval {len(ev)}")
router = nn.Linear(LH, 1).to(dev)
opt = torch.optim.Adam(router.parameters(), lr=1e-3, weight_decay=args.wd)
Xtr = [(f.to(dev), t.to(dev)) for f, t in tr]
for epq in range(1, args.epochs + 1):
opt.zero_grad(); tl = 0.0
for feat, tgt in Xtr:
sc = router(feat).squeeze(-1)
loss = F.cross_entropy(sc, tgt); loss.backward(); tl += loss.item()
opt.step()
if epq % 100 == 0 or epq == 1:
print(f"epoch {epq:4d} | CE {tl/len(Xtr):.4f}")
ks = [1, 2, 3]
with torch.no_grad():
rt = lambda f: router(f.to(dev)).squeeze(-1).cpu()
mean_lm = lambda f: f.mean(-1)
ev_r = coverage(rt, ev, ks, nb_max)
ev_m = coverage(mean_lm, ev, ks, nb_max)
print(f"\nheld-out coverage of the model's own oracle top-1 block (real LongBench QA):")
print(f"{'k':>3} | {'untrained mean-lm':>17} | {'trained router (eval)':>21}")
for k in ks:
print(f"{k:>3} | {ev_m[k]:>17.3f} | {ev_r[k]:>21.3f}")
if args.save_router:
torch.save({"state_dict": router.state_dict(), "in_dim": LH,
"summary_tokens": args.summary_tokens,
"note": f"linear over per-(layer,head) attn to block's last-{args.summary_tokens} tokens; real LongBench QA"},
args.save_router)
print(f"saved router -> {args.save_router} (summary_tokens={args.summary_tokens})")
if __name__ == "__main__":
main()