189 lines
10 KiB
Python
189 lines
10 KiB
Python
"""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()
|