FreeSolo: Linear-time Sparse Attention with Zero-Shot Length Extrapolation

Zekun Shi; Alastair James Arthur Price

cs.LGOct 12, 2026

FreeSolo: Linear-time Sparse Attention with Zero-Shot Length Extrapolation

Authors: Zekun Shi, Alastair James Arthur Price

Organizations: National University of Singapore · Sea AI Lab, Singapore · University of Toronto

Abstract

Linear-time sequence models typically store long-context information in a recurrent state. A fixed-size state limits how much information can be stored, while hybrid architectures that add full attention for better recall incur quadratic sequence cost. Alternatively, sparse attention with continued pretraining (CPT) can extend the effective context length. Generalization beyond the trained length remains challenging, however, and CPT can be costly. We introduce FreeSolo, a native sparse attention mechanism with a short-range sliding window branch and a long-range branch that selects keypoint cells for direct token attention. Its linear-time variant, FS-L, learns to prune the prefix before query-dependent selection through short CPT. FS-B searches balanced trees with $O(N\log^2 N)$ algorithmic cost including radix-sort construction. FreeSolo demonstrates better zero-shot length extrapolation than the baselines without long-context fine-tuning.

Figures & tables

Figure 1: Zero-shot RULER extrapolation at

800

M. Models train at

4

k (dotted line); shading marks extrapolation.

100

examples per task and length.

Panel of Figure 2 — Figure 2: FreeSolo attention. FS scores the full prefix; FS-L prunes per head, independently of the query. Both select $K=4$ cells per query and head. We use $P=64$ high-norm cells, $W=8$ recent cells, cell size $s=32$ , query blocks of $B=128$ tokens, and SR radius $w=512$ .

Panel of Figure 3 — Figure 3: FS-B long-range selection, per head. The query searches two causal forests over the prefix: one over every cell and one over cells with $a_{c}\geq\tau_{\ell h}$ . Each beam keeps $16$ nodes per level down to the leaves, $\lceil\log_{2}M\rceil$ levels for $M$ admitted cells ( $6$ and $5$ in this $64$ -cell example). The union of both beams is reranked with token keys; the top $K=4$ cells form $\mathcal{S}_{i}$ , which is read as in FS.

			S-NIAH-1								S-NIAH-2						S-NIAH-3						Mean
Model	Params	Complexity	2k	4k	8k	16k	32k	64k	131k	256k	2k	4k	8k	16k	32k	64k	2k	4k	8k	16k	32k	64k	2k	4k	8k	16k
TF++	$335$ M	$O(N^{2})$	$95$	$57$	$0$	$0$	–	–	–	–	$\underline{91}$	$24$	$0$	$0$	–	–	$38$	$1$	$0$	$0$	–	–	$26.2$	$13.3$	$0.2$	$0.0$
InfLLM	$335$ M	$O(N^{2})$	$94$	$\underline{58}$	$13$	$9$	–	–	–	–	$89$	$23$	$8$	$2$	–	–	$39$	$1$	$2$	$\underline{1}$	–	–	–	–	–	–
MoBA	$341$ M	$O(N^{2}/b)$	$76$	$12$	$0$	$0$	–	–	–	–	$83$	$23$	$0$	$0$	–	–	$\underline{59}$	$3$	$0$	$0$	–	–	$25.6$	$7.6$	$0.0$	$0.0$
NSA	$513$ M	$O(N^{2}/b)$	$43$	$23$	$0$	$0$	–	–	–	–	$48$	$18$	$0$	$0$	–	–	$1$	$0$	$0$	$0$	–	–	$13.2$	$7.3$	$0.9$	$0.0$
Raven-full	$355$ M	$O(N^{2})$	$\mathbf{100}$	$\mathbf{99}$	$83$	$39$	$44$	$\underline{4}$	$0$	$\underline{0}$	$\mathbf{100}$	$\underline{89}$	$\underline{41}$	$7$	–	–	$38$	$1$	$0$	$0$	–	–	$26.7$	$21.5$	$13.5$	$7.1$

Table 1: RULER accuracy (

0

–

100

). All models train at

4

k; results at

8

k and above test zero-shot extrapolation. Means cover available tasks; cost is defined in Section 4.1 . FS-B’s algorithmic cost includes radix-sort construction (Appendix E.2 ). Bold/underline mark the best/second-best scores within each block (super-linear or linear cost); dashes denote unavailable results.

Model	Params	HellaS.	PIQA	ARC-e	ARC-c	WinoG.	OBQA	SciQ	BoolQ	Avg
GDN	$335$ M	$\underline{39.2}$	$\underline{66.8}$	$55.8$	$\mathbf{28.2}$	$\underline{51.9}$	$\underline{33.0}$	$\mathbf{83.2}$	$59.8$	$\mathbf{52.2}$
FS	$341$ M	$38.9$	$\underline{66.8}$	$\mathbf{57.6}$	$25.9$	$50.9$	$32.4$	$81.2$	$\mathbf{61.1}$	$\underline{51.9}$
TF++	$335$ M	$\underline{39.2}$	$66.0$	$56.2$	$\underline{28.1}$	$50.6$	$32.8$	$\underline{81.5}$	$59.6$	$51.8$
Raven	$357$ M	$37.6$	$65.2$	$55.8$	$27.6$	$51.2$	$32.2$	$80.6$	$\underline{60.8}$	$51.4$
SWA	$341$ M	$38.8$	$66.2$	$\underline{56.4}$	$26.7$	$50.6$	$30.8$	$80.4$	$\mathbf{61.1}$	$51.4$
FS-L	$341$ M	$38.4$	$66.3$	$55.0$	$25.3$	$51.0$	$31.4$	$79.8$	$60.4$	$51.0$

Table 2: Zero-shot downstream accuracy (

0

–

100

). Rows are ordered by average within each scale; bold and underline mark the best and second-best scores.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

	$340$ M		$800$ M
Model	Steps	Tokens	Steps	Tokens
FS	$28{,}000$	$14.68$ B	$60{,}000$	$31.46$ B
TF++	$28{,}610$	$15.00$ B	$60{,}000$	$31.46$ B
SWA	$28{,}000$	$14.68$ B	$60{,}000$	$31.46$ B
MoBA	$28{,}000$	$14.68$ B	–	–
NSA	$28{,}000$	$14.68$ B	$57{,}500$	$30.15$ B

Appendix

Table 3: Evaluated pretraining checkpoints. Steps and training tokens at

524{,}288

tokens per step. MoBA is trained only at the smaller scale. FS-L and FS-B continue from the FS checkpoints: FS-L adds

0.73

B (

340

M) and

1.57

B (

800

M) CPT tokens; FS-B uses the FS-L weights at

340

M and

0.42

B retrieval-weighted CPT tokens at

800

M (Appendix E.4 ).

$\approx$ 340M scale
	Super-linear												Linear
	Raven-full				FS				FS-B				FS-L				NAtS-L				Raven
Task	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k
S-NIAH-1	$\mathbf{100}$	$\mathbf{99}$	$83$	$39$	$\underline{99}$	$\mathbf{99}$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{99}$	$\underline{99}$	$\underline{99}$	$\underline{99}$	$88$	$\underline{90}$	$\mathbf{87}$	$\mathbf{100}$	$\underline{94}$	$37$	$19$	$95$	$\mathbf{96}$	$\mathbf{93}$	$\underline{81}$
S-NIAH-2	$\mathbf{100}$	$\underline{89}$	$\underline{41}$	$7$	$\mathbf{100}$	$77$	$\underline{41}$	$\underline{17}$	$\mathbf{100}$	$\mathbf{99}$	$\mathbf{89}$	$\mathbf{60}$	$\mathbf{99}$	$\mathbf{70}$	$\underline{34}$	$\mathbf{20}$	$89$	$\underline{22}$	$20$	$\underline{3}$	$\underline{92}$	$\mathbf{70}$	$\mathbf{36}$	$\underline{3}$
S-NIAH-3	$38$	$1$	$0$	$\underline{0}$	$\underline{58}$	$\underline{14}$	$\underline{7}$	$\underline{0}$	$\mathbf{91}$	$\mathbf{78}$	$\mathbf{64}$	$\mathbf{27}$	$\mathbf{70}$	$\mathbf{17}$	$\mathbf{9}$	$\mathbf{1}$	$\underline{36}$	$\underline{6}$	$\underline{1}$	$\underline{0}$	$3$	$0$	$0$	$\underline{0}$

Appendix

Table 4: RULER results by task (

0

–

100

) at

340

M and

800

M. Bold and underline mark the best and second-best scores within each group (super-linear or linear cost) and scale. Dashes denote unavailable results.

Model	Train tok/s	Peak mem (GiB)
TF++ (full attention)	$\mathbf{303}$ K	$\mathbf{15.0}$
SWA (sliding- $512$ )	$\underline{228}$ K	$\underline{15.5}$
GDN	$158$ K	$21.6$
MoBA	$140$ K	$27.0$
NSA	$28$ K	$31.9$
Raven	$140$ K	–

Appendix

Table 5: Training throughput and peak memory at

340

M scale on eight A100-40GB GPUs. Raven and NAtS-L use the fla/flame FSDP harness; see the text.

	Prefill (Ktok/s)				Decode (tok/s)
Model	$4$ k	$16$ k	$64$ k	$131$ k	$4$ k	$16$ k	$64$ k	$131$ k
TF++ (full attention)	$144.4$	$123.8$	$48.8$	$26.2$	$34.7$	$34.7$	$34.7$	$20.5$
SWA (sliding- $512$ )	$143.5$	$218.8$	$248.6$	$245.3$	$33.1$	$33.2$	$33.2$	$33.3$
MoBA	$34.3$	$48.6$	$33.3$	$20.2$	$27.4$	$26.9$	$26.1$	OOM
NSA	$87.0$	$121.4$	$89.8$	$65.2$	$15.8$	$15.7$	$15.8$	$15.8$
GDN	$65.7$	$141.8$	$146.8$	$146.2$	$25.5$	$25.7$	$25.6$	$25.4$

Appendix

Table 6: Prefill and eager decode throughput at

340

M scale on one A100-40GB (bf16, batch

1

). OOM denotes the

40

GB memory limit.

$340$ M scale
Length	Attention buffers (MiB)	Peak (GiB)	Prefill (Ktok/s)	Decode (tok/s)
$4$ k	$766.2$	$1.71$	$2.20$	$89.3$
$64$ k	$766.2$	$1.74$	$2.21$	$89.4$
$256$ k	$766.2$	$1.77$	$2.17$	$88.3$
$800$ M scale
$4$ k	$1148.8$	$2.96$	$2.09$	$75.2$

Appendix

Table 7: Bounded FS-L inference. Throughput uses CUDA-graph decode; peak memory is the maximum across three trials.

	S-NIAH-1							S-NIAH-2				S-NIAH-3
Model	2k	4k	8k	16k	32k	64k	131k	2k	4k	8k	16k	2k	4k	8k	16k
FS	$\underline{99}$	$\underline{99}$	$\mathbf{100}$	$\mathbf{100}$	$\underline{98}$	$\mathbf{100}$	$\underline{98}$	$\mathbf{100}$	$\underline{77}$	$\underline{41}$	$\underline{17}$	$\underline{58}$	$14$	$7$	$\underline{0}$
SR-NoPE	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{100}$	$\underline{98}$	$\mathbf{90}$	$\mathbf{71}$	$\mathbf{29}$	$47$	$\underline{23}$	$\mathbf{13}$	$\underline{0}$
LR-content-only	$95$	$97$	$\underline{96}$	$\underline{97}$	$95$	$\underline{95}$	$96$	$80$	$57$	$38$	$14$	$\mathbf{71}$	$\mathbf{25}$	$\underline{11}$	$\mathbf{1}$

Appendix

Table 8: RULER ablations at

340

M (

0

–

100

). Bold and underline mark the best and second-best scores, including ties.

Model	HellaS.	PIQA	ARC-e	ARC-c	WinoG.	OBQA	SciQ	BoolQ	Avg
FS	$\underline{38.9}$	$\mathbf{66.8}$	$\mathbf{57.6}$	$\underline{25.9}$	$\underline{50.9}$	$\mathbf{32.4}$	$\underline{81.2}$	$\mathbf{61.1}$	$\underline{51.9}$
SR-NoPE	$28.4$	$55.7$	$36.9$	$23.7$	$49.2$	$25.6$	$48.2$	$38.7$	$38.3$
LR-content-only	$\mathbf{39.3}$	$\underline{66.1}$	$\underline{57.1}$	$\mathbf{26.9}$	$\mathbf{52.8}$	$\underline{31.8}$	$\mathbf{83.2}$	$\underline{60.0}$	$\mathbf{52.1}$

Appendix

Table 9: Downstream accuracy of the ablations at

340

M (

0

–

100

), using the metrics in Table 2 .

Figure 4: Retrieving a number at

256

k (

340

M). Selected cells, routing scores, prompt text, and attention during generation. Both answers are correct.

Figure 5: Long-range read distance at $800$ M on one $256$ k S-NIAH-1 prompt. Orange guides mark the needle position and its distance from each query.

Figure 6: Retrieving a number at

64

k (

800

M). Selected cells, routing scores, prompt text, and attention to the S-NIAH-2 value. Both answers are correct.

Figure 7: Retrieving a UUID at

32

k (

800

M). Selected cells, routing scores, prompt text, and attention to the S-NIAH-3 value. Both answers are correct.

$340$ M scale
		Source retained (%)		Source selected (%)
Task	Prompts	FS (offline)	FS-L	FS	FS-L
S-NIAH-1	$47$	$22.31$	$27.98$	$38.15$	$17.55$
S-NIAH-2	$50$	$18.34$	$19.87$	$1.27$	$1.30$
S-NIAH-3	$47$	$24.64$	$25.76$	$2.49$	$2.65$
$800$ M scale

Appendix

Table 10: Source-cell availability and selection at

16

k. The source is the cell containing the answer value. Rates average over all heads and the final

16

prompt positions within each prompt, then over prompts; FS (offline) applies the FS-L retention rule to FS keys.

Model	S-NIAH-1	S-NIAH-2	S-NIAH-3
FS base	$100$	$74$	$14$
Immediate pruning	$62$	$72$	$14$
FS-L	$90$	$64$	$14$
FS, seed 42 (run 1)	$100$	$50$	$16$
FS, seed 42 (run 2)	$100$	$52$	$20$
FS, seed 43	$100$	$74$	$6$

Appendix

Table 11: Pruning recovery and seed variation at

340

M. Accuracy (%) on

50

shared prompts per task at

4

k, the training context. FS-L is the

1400

-step checkpoint of Table 1 ; seeded rows use

1500

CPT steps without (FS) or with (FS-L) pruning.

Continuation arm	$\Delta\,2$ k	$\Delta\,4$ k	$\Delta\,8$ k	$\Delta\,16$ k
SWA, $743$ M, plain	$-1.4$	$+0.4$	$+0.2$	$-0.2$
Raven-full, $723$ M, plain	$-7.2$	$-5.6$	$-1.4$	$+0.8$
Raven, $726$ M, plain	$-2.2$	$-1.0$	$-2.8$	$-1.9$
NSA, $1128$ M, plain	$-1.2$	$-2.2$	$+0.3$	$+0.1$
FS, $743$ M, plain	$-1.6$	$-1.3$	$-1.0$	$0.0$
FS, $743$ M $\to$ FS-L (pruning)	$-2.2$	$-2.0$	$-2.0$	$-0.8$

Appendix

Table 12: Only FS-B’s weighted CPT improves the RULER mean at all four lengths. Suite-mean changes at

800

M scale (

+

CPT minus base, points, on the same prompts and protocols as Table 1 ; the base means are the Table 1 rows). Each “

+

CPT” arm continues the Table 1 checkpoint of the same architecture with plain next-token loss at the FS-L budget (

1.57

B tokens). The FS-L arm instead keeps the pruning loop in training at the same budget. The single-forest ablation uses FS-B’s recall-weighted CPT, spending

0.42

B tokens (

27\%

; Appendix E.4 ), and compares the full-prefix width-

32

read before and after CPT (base means

40.9/36.0/30.4/28.2

). FS-B itself denotes the two-forest design of Table 1 .

	S-NIAH-1				S-NIAH-2				S-NIAH-3				RULER mean				Down.
Recipe	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	avg
From scratch	$\mathbf{100}$	$\mathbf{97}$	$\mathbf{98}$	$\mathbf{93}$	$78$	$49$	$8$	$3$	$53$	$11$	$3$	$0$	$27.8$	$19.2$	$11.7$	$9.6$	$\mathbf{52.2}$
CPT	$99$	$88$	$90$	$87$	$\mathbf{99}$	$\mathbf{70}$	$\mathbf{34}$	$\mathbf{20}$	$\mathbf{70}$	$\mathbf{17}$	$\mathbf{9}$	$\mathbf{1}$	$\mathbf{32.4}$	$\mathbf{21.6}$	$\mathbf{15.5}$	$\mathbf{10.7}$	$51.0$

Appendix

Table 13: FS-L training recipes at

340

M (

0

–

100

). Both use pruning during evaluation, with

100

examples per RULER task and length. Bold marks the better recipe.

	S-NIAH-1				S-NIAH-2				S-NIAH-3
Model	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k
$P=16$	$84$	$64$	$68$	$68$	$\mathbf{100}$	$\mathbf{56}$	$\underline{24}$	$\mathbf{14}$	$56$	$14$	$4$	$\mathbf{2}$
$P=32$	$\underline{98}$	$\underline{78}$	$\underline{72}$	$66$	$96$	$\underline{50}$	$\mathbf{32}$	$\underline{10}$	$\mathbf{74}$	$\mathbf{20}$	$\underline{6}$	$\mathbf{2}$
$P=64$ , run 1	$\mathbf{100}$	$\mathbf{82}$	$\mathbf{82}$	$\mathbf{80}$	$\underline{98}$	$44$	$22$	$2$	$\underline{66}$	$12$	$\mathbf{8}$	$\mathbf{2}$
$P=64$ , run 2	$\mathbf{100}$	$76$	$70$	$\underline{74}$	$96$	$46$	$20$	$2$	$58$	$\underline{18}$	$\mathbf{8}$	$\mathbf{2}$

Appendix

Table 14: Adapted retention budgets at

340

M. Each model trains and evaluates with its stated

P

Figure 8: Accuracy versus attention storage at

340

M. The

P=64

points show the mean and range of two seed-

42

continuations.

$P$	Buffers (MiB)	Peak (GiB)	Prefill (Ktok/s)	Decode (tok/s)
$16$	$\mathbf{320.0}$	$\mathbf{1.34}$	$2.13$	$\underline{88.4}$
$32$	$\underline{468.8}$	$\underline{1.49}$	$\mathbf{2.19}$	$\mathbf{88.8}$
$64$	$766.2$	$1.77$	$\underline{2.17}$	$88.3$

Appendix

Table 15: Retention budget and bounded inference at

340

256

$340$ M, $4253$ – $4277$ prefix tokens
Facts per prefix	$1$	$2$	$4$	$8$	$16$
Questions	$50$	$100$	$200$	$400$	$400$
FS	$\underline{94}$	$41$	$19$	$7.75$	$2.75$
FS-L	$90$	$\underline{48}$	$\underline{21}$	$\underline{8.00}$	$\underline{3.25}$
FS-B	$\mathbf{100}$	$\mathbf{56}$	$\mathbf{28.5}$	$\mathbf{13.0}$	$\mathbf{4.25}$
$800$ M, strictly within $4096$ tokens

Appendix

Table 16: Multiple-question retrieval. Accuracy (%) on shared prefixes: an answer is correct if the target value appears in

16

greedy tokens;

50

prefixes per condition. Bold and underline mark the best and second-best score within each scale, including ties. FS-B prefills each question afresh; at

340

M this raises FS and FS-L by about

2

points relative to the shared-prefix rows.

Panel of Figure 9 — Figure 9: Streaming construction of one FS-B forest, per head. Solid edges form balanced trees within power-of-two chunks; dashed edges join their roots, smallest first. Leaf numbers denote arrival order, while their arrangement follows mean-key splits. The $13=8+4+1$ example has height $\lceil\log_{2}13\rceil=4$ . Right: the two construction steps; node summaries are defined in Eqs. ( 20 )–( 21 ).

$\approx$ 340M scale
	FS-B (two forests)				full, $W^{\mathrm{b}}{=}32$				filtered, $W^{\mathrm{b}}{=}16$				full, $W^{\mathrm{b}}{=}16$
Task	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k	$2$ k	$4$ k	$8$ k	$16$ k
S-NIAH-1	$\mathbf{100}$	$\mathbf{99}$	$\mathbf{99}$	$\mathbf{99}$	$\mathbf{100}$	$\mathbf{99}$	$\mathbf{99}$	$\mathbf{99}$	$99$	$\mathbf{99}$	$\mathbf{99}$	$\mathbf{99}$	$\mathbf{100}$	$\mathbf{99}$	$\mathbf{99}$	$\mathbf{99}$
S-NIAH-2	$\mathbf{100}$	$99$	$89$	$60$	$\mathbf{100}$	$\mathbf{100}$	$\mathbf{92}$	$\mathbf{67}$	$97$	$99$	$68$	$49$	$\mathbf{100}$	$99$	$77$	$44$
S-NIAH-3	$91$	$78$	$64$	$\mathbf{27}$	$\mathbf{95}$	$\mathbf{84}$	$\mathbf{71}$	$24$	$75$	$61$	$51$	$22$	$90$	$79$	$70$	$26$
MK-NIAH-1	$27$	$25$	$19$	$\mathbf{24}$	$27$	$\mathbf{28}$	$\mathbf{23}$	$21$	$29$	$24$	$22$	$19$	$\mathbf{30}$	$20$	$\mathbf{23}$	$19$

Appendix

Table 17: Forest ablation of FS-B (

0

–

100

, RULER protocol of Table 1 ). All arms run on the same checkpoints with the same four-cell token read; only the inference-time candidate generator differs: FS-B’s two forests as rendered in Table 1 , the full-prefix forest alone with a width-

32

or width-

16

beam, or the norm-filtered forest alone with a width-

16

beam. Bold marks the best of the four arms per column.

Read path	Weights	$2$ k	$4$ k	$8$ k	$16$ k	Mean
Full-prefix forest (W32)	base FS	$40.9$	$36.0$	$30.4$	$28.2$	$33.9$
	$+$ weighted CPT	$41.3$	$37.2$	$31.6$	$30.4$	$35.1$
Norm-filtered forest (W16)	base FS	$39.6$	$35.1$	$29.8$	$27.5$	$33.0$
	$+$ weighted CPT	$40.7$	$36.2$	$31.6$	$29.6$	$34.5$

Appendix

Table 18: Recall-weighted CPT under single-forest ablations (

0

–

100

; RULER available-task means over

12

tasks at

2

4

k and

13

8

16

k, protocol of Table 1 ). All rows start from the base

743

M FS checkpoint; CPT rows are the Table 1 FS-B checkpoint evaluated with a single-forest ablation. FS-B denotes the two-forest design, whose results appear in Table 1 .

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunLong-Context Language Model Inference LLM Inference Acceleration
KAIST
Simplified Sparse Attention via Gist Tokens
Apr 22, 2026Yuzhen Mao, Michael Y. Li, Emily B. FoxSelf-Attention Long-Context Language Model Inference
Stanford University
Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling
Jul 3, 2026Xiang Hu, Xinyu Wei, Hao Gu +10Long-Context Language Model Inference Long-Context Language Modeling
Tencent HY Team · ShanghaiTech University · The Hong Kong University of Science and Technology +1

Jun 22, 2026J/K move · Enter open · S save

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Huzama Ahmad, Se-Young Yun

KAIST

Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Block-Sparse Attention

Open paper AskPDF arXiv

Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.

Abstract

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Simplified Sparse Attention via Gist Tokens

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling