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

Model: hxia7/qwen3-4b-blockdist
Source: Original Platform

eval/scripts/bench_real.py

@@ -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()