hxia7/qwen3-4b-blockdist

Qwen3-4B Block-Distilled (+ decode-side per-step router)

Block-attention–distilled Qwen3-4B for efficient long-context / RAG serving, plus a decode-side router that re-selects active blocks every generated token. infer.py is a runnable reference and the spec for a custom serving kernel. Sibling: hxia7/qwen3-32b-blockdist.

Block attention

Input [system] [block_1] … [block_n] [query]. Each context block attends only to itself + the system prefix (block-diagonal) → its KV is context-independent and reusable across requests (the RAG prefill win: no per-request document re-encoding). Each block has 4 sink tokens + an 8-token summary that stays resident. Distilled so block-attention tracks a full-attention teacher.

Decode paths (--path)

path	behavior
full	true full causal attention (no block mask) — quality ceiling
dense	block attention, all blocks active — block ceiling
reuse-shift	select top-k once, RoPE-compact — static KV pruning
reroute	per-step re-select top-k; RoPE-compact active context right-aligned to the query start (query + generated tokens keep fixed positions; only active context packs in before them)

Verified: k=all → dense == reuse-shift == reroute token-for-token. RoPE key-shift is exact (relative RoPE).

Key findings (measured on this recipe)

1. Block attention ≈ full — when measured correctly. Apparent large "block costs" on multi-hop QA were measurement artifacts: token-F1 penalizes the block model's verbosity, and merging a LoRA inflates full-attention. On LongBench subspan-EM with a no-merge eval (base + adapter; full = adapter-off teacher, block = adapter-on student), block attention matches or beats full.

2. Per-step reroute is necessary for span-output, free on QA. On short-answer QA, reroute ≈ static (sparse serving is lossless). On long-form / multi-doc generation the relevant block moves during decoding, and per-step reroute > static.

3. The block-attention deficit on multi-hop is training-distribution-limited, not structural. Standard distillation on single-answer QA only teaches "isolate blocks + retrieve one," never cross-block assembly. Distilling on multi-hop data closes it: on held-out subspan-EM, 4B block attention reaches hotpotqa +0.067, 2wikimqa +0.039, musique +0.000 vs the full teacher (block token-F1 beats full on all three). The musique gap closed monotonically as unique multi-hop data grew (−0.10 → 0.00). There is no query-bottleneck floor — more cross-block data suffices.

4. Scale. On Qwen3-32B the block-vs-full gap is smaller still; the block-distillation + sparse-serving pipeline retains full-attention quality, with block-sparse decode cutting KV-read bandwidth ~k/n.

Measured results (LongBench subspan-EM, held-out)

Block cost is a metric artifact, and the multi-hop fix works. "block-cost" = block − full (positive = block wins). Baseline is the general-distilled release; "multi-hop-distilled" is the block student trained on hotpotqa/2wikimqa/musique, evaluated no-merge (full = adapter-off teacher, block = adapter-on student).

dataset	hops	baseline full	baseline block	baseline cost	multi-hop-distilled cost
triviaqa	1	0.290	0.300	+0.010	— (already ≥)
multifieldqa_en	1	0.184	0.158	−0.026	—
hotpotqa	2	0.490	0.370	−0.120	+0.067 (block 0.600 vs full 0.533)
2wikimqa	2	0.330	0.290	−0.040	+0.039 (block 0.392 vs full 0.353)
musique	2–4	0.200	0.120	−0.080	+0.000 (block 0.150 = full 0.150)

For comparison, the token-F1 "cost" looked far worse (hotpotqa −32%, 2wikimqa −46%, musique −51%) — almost entirely because token-F1 penalizes the block model's verbosity, not real accuracy. On subspan-EM the baseline gaps are small, and after multi-hop distillation block ≥ full on every multi-hop set (block also beats full on token-F1 there). The deepest task, musique, closed monotonically as unique multi-hop training data grew:

musique block−full (subspan-EM):  −0.100  →  −0.075  →  −0.050  →  0.000
unique musique training examples:    16    →    16    →   110    →   330

→ No structural floor: the multi-hop deficit is fully training-distribution-limited and closes with cross-block data. (Absolute musique ~0.15 is 4B's own 3-hop ceiling — full attention scores the same — so block-sparse serving is lossless vs full even there.)

The router (router.weights)

Oracle-distilled: predicts, from the serve-time per-(layer,head) summary-attention feature, the blocks full attention actually attends to per step. infer.py:load_router reads {in_dim, arch, state_dict, summary_tokens}. Legacy router.pt is the older single-landmark linear router — prefer router.weights.

Usage

python infer.py --model <this-repo> --router router.weights --k 3 --path reroute
python infer.py --model <this-repo> --path full      # full-attention ceiling
python infer.py --model <this-repo> --path dense     # block ceiling

Reproduce the multi-hop fix

Distill the block student on multi-hop LongBench (hotpotqa/2wikimqa/musique), long enough context to include deep-hop examples (16k → ~7× more musique), LoRA rank 32; evaluate with subspan-EM using the no-merge adapter-off/on protocol. (Adapter weights for the improved multi-hop variant available on request.)

Caveats

• Block layout / 4× \n sink format must match training (see infer.py).
• Savings are KV-read bandwidth (~k/n blocks + resident summaries), realized in long-context × batched serving; at batch-1 short-context the win is small (weight-bound). The bigger win is prefill KV reuse of document chunks across requests.
• Report subspan-EM (adapter-off/on), not merged-model token-F1 — the latter understates block quality.