From b315e39b604fc10b7fe7cf405df8da0379625bc4 Mon Sep 17 00:00:00 2001 From: ModelHub XC Date: Mon, 27 Jul 2026 06:09:10 +0800 Subject: [PATCH] =?UTF-8?q?=E5=88=9D=E5=A7=8B=E5=8C=96=E9=A1=B9=E7=9B=AE?= =?UTF-8?q?=EF=BC=8C=E7=94=B1ModelHub=20XC=E7=A4=BE=E5=8C=BA=E6=8F=90?= =?UTF-8?q?=E4=BE=9B=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Model: hxia7/qwen3-4b-blockdist Source: Original Platform --- .gitattributes | 37 +++++ README.md | 90 ++++++++++++ chat_template.jinja | 89 ++++++++++++ config.json | 71 ++++++++++ eval/EVAL.md | 44 ++++++ eval/scripts/acc_regress.py | 89 ++++++++++++ eval/scripts/audit_sparse.py | 126 +++++++++++++++++ eval/scripts/bench_real.py | 115 ++++++++++++++++ eval/scripts/decode_savings.py | 28 ++++ eval/scripts/demo_infer.py | 157 +++++++++++++++++++++ eval/scripts/distill_lora.py | 218 ++++++++++++++++++++++++++++++ eval/scripts/e2e_nll.py | 178 ++++++++++++++++++++++++ eval/scripts/flex_attn_test.py | 92 +++++++++++++ eval/scripts/router_leak_check.py | 85 ++++++++++++ eval/scripts/train_router_real.py | 137 +++++++++++++++++++ generation_config.json | 13 ++ infer.py | 188 ++++++++++++++++++++++++++ model.safetensors | 3 + router.pt | 3 + router.weights | 3 + tokenizer.json | 3 + tokenizer_config.json | 30 ++++ 22 files changed, 1799 insertions(+) create mode 100644 .gitattributes create mode 100644 README.md create mode 100644 chat_template.jinja create mode 100644 config.json create mode 100644 eval/EVAL.md create mode 100644 eval/scripts/acc_regress.py create mode 100644 eval/scripts/audit_sparse.py create mode 100644 eval/scripts/bench_real.py create mode 100644 eval/scripts/decode_savings.py create mode 100644 eval/scripts/demo_infer.py create mode 100644 eval/scripts/distill_lora.py create mode 100644 eval/scripts/e2e_nll.py create mode 100644 eval/scripts/flex_attn_test.py create mode 100644 eval/scripts/router_leak_check.py create mode 100644 eval/scripts/train_router_real.py create mode 100644 generation_config.json create mode 100644 infer.py create mode 100644 model.safetensors create mode 100644 router.pt create mode 100644 router.weights create mode 100644 tokenizer.json create mode 100644 tokenizer_config.json diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..02dc1b2 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,37 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +tokenizer.json filter=lfs diff=lfs merge=lfs -text +router.weights filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000..d291d9b --- /dev/null +++ b/README.md @@ -0,0 +1,90 @@ +--- +license: apache-2.0 +base_model: Qwen/Qwen3-4B +tags: + - block-attention + - long-context + - rag + - kv-cache + - efficient-inference +--- + +# Qwen3-4B Block-Distilled (+ decode-side per-step router) + +A block-attention–distilled **Qwen3-4B** for efficient long-context / RAG serving, plus a light +**decode-side router** that re-selects the active blocks **every generated token**. Distilled on one +H200 with LoRA (merged into the released weights). `infer.py` is a runnable reference and the spec for +a custom serving kernel (e.g. vLLM). + +## Block attention + +Input layout `[system] [block_1] … [block_n] [query]`. Each context **block** attends only to itself +and the system prefix (block-diagonal), **not** to other blocks — so a block's KV is context-independent +and can be **computed once and reused across requests** (the main RAG serving win: no per-request +document re-encoding). Each block is prefixed with **4 sink tokens** (`\n`) and carries a **summary** = +its **last 8 content tokens**, which stay resident. The model is distilled so block-attention tracks a +frozen full-attention teacher (KL + damage-weighted CE) — block ≈ full in quality. + +## Decode paths (`--path`) + +All three share one block-diagonal prefill (query sees all blocks, KV cached once); they differ only at decode: + +| path | behavior | note | +|---|---|---| +| `dense` | decode attends to all blocks | quality reference | +| `reuse-shift` | **select-once**: router picks top-k at step 0, RoPE-shifts them + query to compact positions, decodes reusing that compact KV | == static KV **pruning** | +| `reroute` | **per-step**: all block KV stays resident; every step re-selects top-k from the current token's summary-attention and re-RoPE-shifts to compact positions | follows the block as it **moves** during generation | + +RoPE key-shift is exact (relative RoPE: `R(Δ)·R(p)k = R(p+Δ)k`), so cached post-RoPE keys are retargeted +to compact positions without recompute (verified: offset-invariance 6/6 in fp32). + +## What we validated (honest) + +Measured on this 4B model, real LongBench + a controlled probe: + +- **Machinery is exact.** With `k = all blocks`, `dense == reuse-shift == reroute` token-for-token (0 mismatch). +- **On short-answer QA, static (pruning) already ≈ dense**, and per-step reroute ties it — the answer + comes from a few blocks that don't move, so there's little for per-step to gain. (Real LongBench + single-hop & multi-hop, n=100–150: reroute/static within ~1–2 examples of dense; EM identical.) +- **Per-step reroute becomes *necessary* when the output spans many blocks.** In a controlled + multi-target recall (12 blocks, recite a fact from 8 of them in order, k=3): static is structurally + capped at **0.24** recall (it holds only k=3 blocks), while an **oracle per-step selector recovers + 0.43 ≈ dense 0.46** — a ~2× gap that pruning cannot close because the relevant block moves across the + generation. **This is the regime the per-step router targets** (long-form / multi-doc generation), + and the contribution over static KV pruning. + +## The router (`router.pt`) — work in progress + +Trained by **oracle-distillation**: it predicts, from the serve-time per-step summary-attention feature, +the blocks the full-attention model actually attends to at that step. On the multi-target probe this +lifts learned reroute from static's 0.24 toward the 0.43 oracle ceiling (first cut ~0.32; the +static→oracle gap is the active research target). Small MLP over per-(layer,head) summary attention; +`summary_tokens=8`, backbone frozen. `infer.py:load_router` reads `{in_dim, arch, state_dict, +summary_tokens}`. + +## Usage + +`router.weights` is the current **oracle-distilled per-step router** (load with `torch.load`; same format +as a `.pt`, non-LFS name for reliable upload). The legacy `router.pt` is the older single-landmark linear +router — prefer `router.weights`. + +```bash +python infer.py --model --router router.weights --k 3 --path reroute +python infer.py --model --router router.weights --k 3 --path reuse-shift # static pruning baseline +python infer.py --model --path dense # reference +``` + +## Serving spec (vLLM etc.) + +`infer.py` is the reference: block-diagonal context KV store (reusable per doc chunk across requests), +resident per-block summaries, per-step top-k block gather + RoPE compact-shift at decode. Savings are +**KV-read bandwidth** (~k/n of blocks, plus resident summaries), realized in the **long-context × +batched** regime where decode is KV-bound; at batch-1 short-context the win is small (decode is +weight-bound). The bigger serving win is **prefill KV reuse** of document chunks across requests. + +## Caveats + +- Block layout / 4× `\n` sink format must match training. See `infer.py`. +- The per-step router is WIP: it beats static on span-output but has headroom to the oracle; short-answer + QA does not need per-step at all (static suffices there). +- Distilled with LoRA on segmented LongBench + SemanticSeg. Sibling: `hxia7/qwen3-14b-blockdist`. diff --git a/chat_template.jinja b/chat_template.jinja new file mode 100644 index 0000000..01be9b3 --- /dev/null +++ b/chat_template.jinja @@ -0,0 +1,89 @@ +{%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {{- '\n' }} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file diff --git a/config.json b/config.json new file mode 100644 index 0000000..161afd9 --- /dev/null +++ b/config.json @@ -0,0 +1,71 @@ +{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "dtype": "bfloat16", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 40960, + "max_window_layers": 36, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 36, + "num_key_value_heads": 8, + "pad_token_id": null, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 1000000, + "rope_type": "default" + }, + "sliding_window": null, + "tie_word_embeddings": true, + "transformers_version": "5.14.1", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/eval/EVAL.md b/eval/EVAL.md new file mode 100644 index 0000000..6fc354d --- /dev/null +++ b/eval/EVAL.md @@ -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. diff --git a/eval/scripts/acc_regress.py b/eval/scripts/acc_regress.py new file mode 100644 index 0000000..904e203 --- /dev/null +++ b/eval/scripts/acc_regress.py @@ -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() diff --git a/eval/scripts/audit_sparse.py b/eval/scripts/audit_sparse.py new file mode 100644 index 0000000..a1de6b7 --- /dev/null +++ b/eval/scripts/audit_sparse.py @@ -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() diff --git a/eval/scripts/bench_real.py b/eval/scripts/bench_real.py new file mode 100644 index 0000000..8f030f5 --- /dev/null +++ b/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 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() diff --git a/eval/scripts/decode_savings.py b/eval/scripts/decode_savings.py new file mode 100644 index 0000000..7e952b8 --- /dev/null +++ b/eval/scripts/decode_savings.py @@ -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).") diff --git a/eval/scripts/demo_infer.py b/eval/scripts/demo_infer.py new file mode 100644 index 0000000..19b62bc --- /dev/null +++ b/eval/scripts/demo_infer.py @@ -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() diff --git a/eval/scripts/distill_lora.py b/eval/scripts/distill_lora.py new file mode 100644 index 0000000..a1fb487 --- /dev/null +++ b/eval/scripts/distill_lora.py @@ -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 : 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'', 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() diff --git a/eval/scripts/e2e_nll.py b/eval/scripts/e2e_nll.py new file mode 100644 index 0000000..666346d --- /dev/null +++ b/eval/scripts/e2e_nll.py @@ -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() diff --git a/eval/scripts/flex_attn_test.py b/eval/scripts/flex_attn_test.py new file mode 100644 index 0000000..43fc5a1 --- /dev/null +++ b/eval/scripts/flex_attn_test.py @@ -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() diff --git a/eval/scripts/router_leak_check.py b/eval/scripts/router_leak_check.py new file mode 100644 index 0000000..85d4c1f --- /dev/null +++ b/eval/scripts/router_leak_check.py @@ -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() diff --git a/eval/scripts/train_router_real.py b/eval/scripts/train_router_real.py new file mode 100644 index 0000000..70f733c --- /dev/null +++ b/eval/scripts/train_router_real.py @@ -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() diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000..c37ba92 --- /dev/null +++ b/generation_config.json @@ -0,0 +1,13 @@ +{ + "bos_token_id": 151643, + "do_sample": true, + "eos_token_id": [ + 151645, + 151643 + ], + "pad_token_id": 151643, + "temperature": 0.6, + "top_k": 20, + "top_p": 0.95, + "transformers_version": "5.14.1" +} diff --git a/infer.py b/infer.py new file mode 100644 index 0000000..1ede0a3 --- /dev/null +++ b/infer.py @@ -0,0 +1,188 @@ +"""Reference inference for the block-distilled Qwen3-4B + decode-side per-step router. + +Layout: [system] [block_1 ...] [block_n] [query]. Each context block attends only to itself + the +system prefix (block-diagonal), so its KV is context-independent and REUSABLE across requests. Every +block is prefixed with 4 sink tokens and carries a `summ`-token SUMMARY (its last content tokens). The +query/decode attends to blocks; a light router selects which blocks stay ACTIVE. + +Serving paths (all keep the shared block-diagonal prefill; differ only at decode): + - dense : decode attends to ALL blocks (quality reference). + - reuse-shift : SELECT-ONCE. Router picks top-k active blocks at generation start, RoPE-shifts them + + query to COMPACT positions, decodes reusing that compact KV. == static KV pruning. + - reroute : PER-STEP (the contribution). ALL block KV stays resident; every decode step re-selects + the top-k active blocks from the current token's summary-attention, re-RoPE-shifts the + [system + active-full + all-summaries + query + prior-gen] to compact positions, and + decodes one token. Necessary when the RELEVANT block MOVES across the generated + sequence (long-form / multi-doc), where select-once/pruning structurally cannot follow. + +RoPE key-shift is exact (relative RoPE: R(Δ)·R(p)k = R(p+Δ)k), so cached post-RoPE keys are retargeted +to compact positions WITHOUT recompute. This file is the SPEC for a custom serving kernel (e.g. vLLM): +block-diagonal context KV store + resident summaries + per-step top-k gather + RoPE compact-shift. + +Usage: + python infer.py --model [--router router.pt --k 3] [--path dense|reuse-shift|reroute] +""" +from __future__ import annotations +import argparse, torch, torch.nn as nn +from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache +from transformers.models.qwen3.modeling_qwen3 import rotate_half + +NEG = -1e9; N_SINKS = 4; SINK_TEXT = "\n" +SYS = "You are a helpful assistant. Answer the question using only the documents. Be concise." +DOCS = [ + "The Zorban Reactor at Vale Station reached full output in the year 3471.", + "Mount Kilimanjaro, in Tanzania, is the highest mountain in Africa at 5,895 metres.", + "The Great Wall of China stretches thousands of kilometres across northern China.", + "Photosynthesis converts carbon dioxide and water into glucose and oxygen.", +] +QUESTION = "In what year did the Zorban Reactor reach full output?" + + +def build_prompt(tok, docs, question): + """Qwen3 CHAT layout (matches router training). ids, block id (-1 sys / -2 query / >=0 block), sink flag.""" + sink_ids = tok(SINK_TEXT, add_special_tokens=False)["input_ids"] * N_SINKS + ids, blk, snk = [], [], [] + def add(text, b): + for x in tok(text, add_special_tokens=False)["input_ids"]: + ids.append(x); blk.append(b); snk.append(False) + add(f"<|im_start|>system\n{SYS}<|im_end|>\n<|im_start|>user\n", -1) + for bi, d in enumerate(docs): + for s in sink_ids: ids.append(s); blk.append(bi); snk.append(True) + add(f"[Document {bi+1}] {d}\n", bi) + q0 = len(ids) + add(f"{question}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", -2) + return ids, blk, snk, q0 + + +def block_summaries(blk, snk, nb, S): + cols = {b: [] for b in range(nb)} + for j, (b, s) in enumerate(zip(blk, snk)): + if b >= 0 and not s: cols[b].append(j) + return {b: cols[b][-S:] for b in range(nb)} + + +def prefill_mask(blk, snk, q0, nb, n, dev, dtype): + """block-diagonal context + query-sees-all, causal. [1,1,n,n] additive, cast to model dtype (SDPA-safe).""" + b = torch.tensor(blk + [-2] * (n - len(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) + allowed = ((bj == -1) | (bi == bj) | (bi == -2)) & causal + return torch.where(allowed, 0.0, NEG).view(1, 1, n, n).to(dtype) + + +def rope_shift_keys(keys, delta_pos, rotary, dev): + cos, sin = rotary(keys, delta_pos[None].to(dev).float()) + return keys * cos.unsqueeze(1) + rotate_half(keys) * sin.unsqueeze(1) + + +def block_feat(att, smry, nb): + """att [L,H,T] (decode token's attention). Feature = mean attention to each block's summary cols -> [nb,L*H].""" + return torch.stack([att[:, :, smry[b]].mean(-1) for b in range(nb)], -1).permute(2, 0, 1).reshape(nb, -1) + + +@torch.no_grad() +def generate(model, tok, docs, question, router=None, k=3, max_new=32, summ=8, mode="dense"): + """mode: dense | reuse-shift (select-once) | reroute (per-step). Router required for the sparse modes.""" + dev = model.device + ids, blk, snk, q0 = build_prompt(tok, docs, question) + n = len(ids); nb = max(blk) + 1 + smry = block_summaries(blk, snk, nb, summ); lmset = set(c for b in range(nb) for c in smry[b]) + rotary = model.model.rotary_emb + out = model(input_ids=torch.tensor([ids], device=dev), + attention_mask=prefill_mask(blk, snk, q0, nb, n, dev, model.dtype), + use_cache=True, output_attentions=(mode != "dense")) + pkv = out.past_key_values; nxt = int(out.logits[0, -1].argmax()); gen = [nxt] + + if mode == "dense": + cur = n + for _ in range(max_new - 1): + if nxt == tok.eos_token_id: break + cur += 1 + out = model(input_ids=torch.tensor([[nxt]], device=dev), attention_mask=torch.ones((1, cur), device=dev), + past_key_values=pkv, use_cache=True) + nxt = int(out.logits[0, -1].argmax()); gen.append(nxt) + return tok.decode(gen, skip_special_tokens=True), None + + # router pick at generation start (query-boundary attention to summaries) + att0 = torch.stack(out.attentions, 0)[:, 0, :, -1, :].float() + topk0 = set(router(block_feat(att0, smry, nb)).squeeze(-1).argsort(descending=True)[:k].tolist()) + + if mode == "reuse-shift": # SELECT-ONCE: drop inactive KV, compact-shift, decode + keep = [j for j, b in enumerate(blk) if b < 0 or b in topk0] + old = torch.tensor(keep, device=dev); K = len(keep); delta = torch.arange(K, device=dev) - old + for layer in pkv.layers: + layer.keys = rope_shift_keys(layer.keys[:, :, old, :], delta, rotary, dev).contiguous() + layer.values = layer.values[:, :, old, :].contiguous() + cur = K + for _ in range(max_new - 1): + if nxt == tok.eos_token_id: break + out = model(input_ids=torch.tensor([[nxt]], device=dev), position_ids=torch.tensor([[cur]], device=dev), + cache_position=torch.tensor([K + len(gen) - 1], device=dev), past_key_values=pkv, use_cache=True) + nxt = int(out.logits[0, -1].argmax()); gen.append(nxt); cur += 1 + return tok.decode(gen, skip_special_tokens=True), sorted(topk0) + + # mode == "reroute": PER-STEP. keep ALL KV resident; re-select + re-compact every step. + K0 = [l.keys for l in pkv.layers]; V0 = [l.values for l in pkv.layers]; L = len(K0) + active = set(range(nb)); g = 0; Kg = [None] * L; Vg = [None] * L + def keep_col(j, b): return (b < 0) or (b in active) or (j in lmset) + for step in range(max_new - 1): + kp = [j for j in range(n) if keep_col(j, blk[j])] + pos_in = {j: i for i, j in enumerate(kp)} + sidx = [[pos_in[c] for c in smry[b]] for b in range(nb)] + Kt = len(kp) + g + old = torch.tensor(kp + [n + i for i in range(g)], device=dev); delta = torch.arange(Kt, device=dev) - old + cache = DynamicCache() + for i in range(L): + kk = K0[i][:, :, kp, :]; vv = V0[i][:, :, kp, :] + if g: kk = torch.cat([kk, Kg[i]], 2); vv = torch.cat([vv, Vg[i]], 2) + cache.update(rope_shift_keys(kk, delta, rotary, dev), vv, i) + out = model(input_ids=torch.tensor([[nxt]], device=dev), position_ids=torch.tensor([[Kt]], device=dev), + cache_position=torch.tensor([Kt], device=dev), past_key_values=cache, use_cache=True, + output_attentions=True) + nxt2 = int(out.logits[0, -1].argmax()) + att = torch.stack(out.attentions, 0)[:, 0, :, 0, :].float() + active = set(router(block_feat(att, {b: sidx[b] for b in range(nb)}, nb)).squeeze(-1).argsort(descending=True)[:k].tolist()) + dstore = torch.tensor([(n + g) - Kt], device=dev) + for i in range(L): + Kg[i] = rope_shift_keys(cache.layers[i].keys[:, :, -1:, :], dstore, rotary, dev) if Kg[i] is None \ + else torch.cat([Kg[i], rope_shift_keys(cache.layers[i].keys[:, :, -1:, :], dstore, rotary, dev)], 2) + Vg[i] = cache.layers[i].values[:, :, -1:, :] if Vg[i] is None \ + else torch.cat([Vg[i], cache.layers[i].values[:, :, -1:, :]], 2) + g += 1; gen.append(nxt2); nxt = nxt2 + if nxt == tok.eos_token_id: break + return tok.decode(gen, skip_special_tokens=True), None + + +def load_router(path, dev): + ck = torch.load(path, map_location=dev) + LH = ck["in_dim"]; arch = ck.get("arch", "linear"); summ = ck.get("summary_tokens", 8) + if arch == "mlp": + r = nn.Sequential(nn.Linear(LH, 256), nn.GELU(), nn.Dropout(0.0), nn.Linear(256, 1)) + else: + r = nn.Linear(LH, 1) + r.load_state_dict(ck["state_dict"]); return r.to(dev).eval(), summ + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--router", default=None); ap.add_argument("--k", type=int, default=3) + ap.add_argument("--max-new", type=int, default=32) + ap.add_argument("--path", default="reroute", choices=["dense", "reuse-shift", "reroute"]) + 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 = summ = None + if args.router: router, summ = load_router(args.router, "cuda") + ids, blk, _, q0 = build_prompt(tok, DOCS, QUESTION) + print(f"blocks={max(blk)+1} context_tokens={q0} path={args.path} question={QUESTION!r}") + dense, _ = generate(model, tok, DOCS, QUESTION, max_new=args.max_new, mode="dense") + print("dense :", repr(dense.strip())) + if router is not None and args.path != "dense": + s, act = generate(model, tok, DOCS, QUESTION, router, args.k, args.max_new, summ, args.path) + print(f"{args.path} (top-{args.k}, active={act}) :", repr(s.strip())) + + +if __name__ == "__main__": + main() diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000..8f883f9 --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2d98978bcd4d409d45f3bb2ecf76c187eedeb2f06b76bb15a412556c9e9b57c8 +size 8044982080 diff --git a/router.pt b/router.pt new file mode 100644 index 0000000..90ccd25 --- /dev/null +++ b/router.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c57a3127de4d9d286327e6cd12ca2d5f5ff2d21987f6444b2c03637b562e4094 +size 6797 diff --git a/router.weights b/router.weights new file mode 100644 index 0000000..271bf2a --- /dev/null +++ b/router.weights @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28f151742607d2e8219d8bf5675ab413fe172fe703c4cfbe1292b7930eab89b3 +size 1183955 diff --git a/tokenizer.json b/tokenizer.json new file mode 100644 index 0000000..c7afbed --- /dev/null +++ b/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506 +size 11422650 diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000..5668a4a --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,30 @@ +{ + "add_prefix_space": false, + "backend": "tokenizers", + "bos_token": null, + "clean_up_tokenization_spaces": false, + "eos_token": "<|im_end|>", + "errors": "replace", + "extra_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|object_ref_start|>", + "<|object_ref_end|>", + "<|box_start|>", + "<|box_end|>", + "<|quad_start|>", + "<|quad_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>" + ], + "is_local": true, + "local_files_only": false, + "model_max_length": 131072, + "pad_token": "<|endoftext|>", + "split_special_tokens": false, + "tokenizer_class": "Qwen2Tokenizer", + "unk_token": null +}