hxia7/qwen3-14b-blockdist

Qwen3-14B Block-Distilled (+ landmark router)

A block-attention–distilled Qwen3-14B for efficient long-context / RAG serving, plus a light landmark router for decode-time block sparsity. Trained on one H200 with LoRA distillation (merged into the released weights).

What "block attention" means here

The input is laid out as [system] [block_1] … [block_n] [query]. Each context block is encoded independently — it attends only to itself and the shared system prefix, not to other blocks — so a block's KV can be cached once and reused across requests (prefill reuse). The query/answer segment attends to all blocks. Every block is prefixed with 4 sink tokens (\n here) for stability (counters the "lost-in-block-head" pathology).

This model is distilled so that its block-attention predictions track a frozen full-attention teacher (KL + damage-weighted CE), i.e. block attention ≈ full attention in quality.

Decode-time sparsity (optional, via the landmark router)

Keep each block's summary — its last 8 content tokens — resident. A tiny linear router scores blocks from per-(layer, head) attention to those summary tokens and activates only the top-k blocks per decode step (on-demand re-route). In our evaluation this recovers most of dense quality at ~k/n of the KV bandwidth. The router is small (router.pt, summary_tokens=8) and the backbone is frozen for it.

Design note (matters for quality): a single last-token landmark is too weak a summary — on detail queries it mis-ranks the answer block and the sparse decode drops a needed block for one token (e.g. a year "1804"→"1800"). Using the block's last 8 tokens as the summary, with a router trained on real RAG attention, fixes this; sparse top-2 then matches dense on retrieval-style questions. Purely cross-block aggregation questions ("list all X") are a limit of block attention itself (even dense struggles) — route those through dense fallback.

Validation

Real-scale benchmark — 50 real LongBench QA samples (~7 blocks/sample, ~200 tokens/block, ~1.2k context tokens). Gold-answer NLL vs how much block-content the sparse decode actually reads per step:

decode	gold NLL	vs dense	block-content read	context saved
dense (all blocks)	1.019	—	100%	1×
router last-8, k=3	1.052	+0.03 (≈3%)	63%	1.6×
router last-8, k=2	1.145	+0.13 (≈12%)	37%	2.7×
static (freeze k=2)	1.217	+0.20	37%	2.7×

So at k=3 the router is within ~3% NLL of dense while reading 1.6× less; at k=2, ~12% for 2.7× less. It recovers ~36% of the static→dense gap (learned per-head routing beats the untrained last-token landmark). Savings scale with block count — k stays ~2–3 while the number of blocks grows with context, so attended ≈ k/n + (1−k/n)·(summary/blocksize); at tens of blocks this is several-fold.

No data leakage: a masked-out block provably cannot influence the output — a canary secret placed in a masked block cannot be reproduced (the model hallucinates instead); with the block active it can. See scripts/audit_sparse.py.

Stress test (synthetic multi-hop, strong distractors): the router recovers up to 94% of the static→dense gap and beats a full-attention oracle — this harder setup is where the learned last-8 router most clearly beats a single-landmark baseline.

Distillation also lifted block-attention quality (dense NLL 0.78 → 0.48 vs the base model on that set).

Files

• model weights (merged Qwen3-14B) — load with transformers as usual
• router.pt — landmark-router state dict (optional, for sparse decode)
• infer.py — reference inference: block-attention layout, sink/landmark tokens, dense + sparse decode

Usage

python infer.py --model <this-repo> --router router.pt --k 3

infer.py is also the spec for a custom serving kernel (vLLM etc.): block-diagonal context attention, sink + landmark tokens, and router-gated top-k block KV gather at decode.

Caveats

• The block layout / sink-token format must match training (4× \n sinks per block). See infer.py.
• Best in the long-context, batched serving regime (that's where KV-read dominates decode); at batch-1 short-context the win is small (decode is weight-bound).
• Distilled with LoRA on segmented LongBench + SemanticSeg data.