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

37
.gitattributes vendored Normal file
View File

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

90
README.md Normal file
View File

@@ -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 <this-repo> --router router.weights --k 3 --path reroute
python infer.py --model <this-repo> --router router.weights --k 3 --path reuse-shift # static pruning baseline
python infer.py --model <this-repo> --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`.

89
chat_template.jinja Normal file
View File

@@ -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 <tools></tools> XML tags:\n<tools>" }}
{%- for tool in tools %}
{{- "\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|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('<tool_response>') and message.content.endswith('</tool_response>')) %}
{%- 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 '</think>' in content %}
{%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
{%- set content = content.split('</think>')[-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<think>\n' + reasoning_content.strip('\n') + '\n</think>\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 %}
{{- '<tool_call>\n{"name": "' }}
{{- tool_call.name }}
{{- '", "arguments": ' }}
{%- if tool_call.arguments is string %}
{{- tool_call.arguments }}
{%- else %}
{{- tool_call.arguments | tojson }}
{%- endif %}
{{- '}\n</tool_call>' }}
{%- endfor %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
{{- '<|im_start|>user' }}
{%- endif %}
{{- '\n<tool_response>\n' }}
{{- content }}
{{- '\n</tool_response>' }}
{%- 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 %}
{{- '<think>\n\n</think>\n\n' }}
{%- endif %}
{%- endif %}

71
config.json Normal file
View File

@@ -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
}

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

13
generation_config.json Normal file
View File

@@ -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"
}

188
infer.py Normal file
View File

@@ -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 <path> [--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<think>\n\n</think>\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()

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:2d98978bcd4d409d45f3bb2ecf76c187eedeb2f06b76bb15a412556c9e9b57c8
size 8044982080

3
router.pt Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:c57a3127de4d9d286327e6cd12ca2d5f5ff2d21987f6444b2c03637b562e4094
size 6797

3
router.weights Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:28f151742607d2e8219d8bf5675ab413fe172fe703c4cfbe1292b7930eab89b3
size 1183955

3
tokenizer.json Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
size 11422650

30
tokenizer_config.json Normal file
View File

@@ -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
}