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