SketchSSM: Write to the Full State, Read from a Compact Sketch

Omin Kwon; JoongWon Shin; Minseo Kim; Kurt Keutzer; Sehoon Kim; Jae W. Lee

cs.LGSep 27, 2026

SketchSSM: Write to the Full State, Read from a Compact Sketch

Authors: Omin Kwon, JoongWon Shin, Minseo Kim, Kurt Keutzer, Sehoon Kim, Jae W. Lee

Organizations: Seoul National University · UC Berkeley · KAIST

Abstract

Hybrid-attention models replace most softmax attention layers with linear attention, reducing KV-cache growth and enabling larger decode batches where recurrent-state access becomes a major bottleneck. ReplaySSM amortizes state updates by buffering keys and values, but each new query still requires a full-state read even though the state remains unchanged between state updates. We observe that low-rank state-weighted query approximation accurately preserves state-read outputs. Although future queries are unknown, the basis vectors used to approximate them can be fixed offline. Based on this observation, we introduce SketchSSM, which preserves full-state updates while approximating reads. At each state update, SketchSSM reads the full state once to precompute outputs for these basis vectors, storing them in a compact sketch. Each subsequent decode step combines the sketch vectors with query-dependent coefficients to reconstruct the output without a full-state read. Across four Mamba-2-, GDN-, and KDA-based models, SketchSSM at a mean sketch rank of 8 reduces state-access traffic by approximately 10 $\times$ while matching the average accuracy of the FP32 full-state baseline across four decode benchmarks, and preserves recall on four RULER retrieval tasks. At this rank on one NVIDIA B300, linear-attention kernel speedups over the Standard vLLM baseline reach 7.30 $\times$ , 5.02 $\times$ , and 5.24 $\times$ for Mamba-2, GDN, and KDA, respectively, with up to 2.77 $\times$ higher decode throughput on Nemotron 3 Super.

Figures & tables

Fig. 1: Decode latency breakdown and decode throughput at batch sizes

256

512

, and

B_{\max}

for

2

K and

8

K contexts. At

B_{\max}

, linear attention accounts for

42

–

75\%

of decode latency with Standard, and

31

–

53\%

with ReplaySSM.

B_{\max}

denotes the largest tested batch size that fits within GPU memory for each model and context length. Throughput generally increases as batch size grows toward

B_{\max}

. GLM uses only Standard because KDA lacks an official ReplaySSM implementation.

Fig. 2: Retained state-read output fraction

\rho_{G}

versus rank

G

on Nemotron 3 Super (Mamba-2), Qwen3.8 Flash-Next (GDN), and GLM 5.3 Flash (KDA), using a query basis fixed offline and state-dependent coefficients. Lines show the mean across heads; shaded regions denote the interquartile range. At

G=4

, the mean retained output fraction is 88.4–97.5% across the three models.

Fig. 3: (a) ReplaySSM reads the full state to compute

o_{t}^{\mathrm{state}}

, incurring

KV

elements of state-read traffic. (b) SketchSSM reconstructs

\hat{o}_{t}^{\mathrm{state}}

from rank-

G

factors, reducing traffic to

G(K+V)

elements.

Fig. 4: SketchSSM over a window. At a flush step, SketchSSM updates the full state

S_{0}

and refreshes the sketch

U

and coefficient map

C

. At non-flush steps, it combines

U

and

C

with the query to reconstruct the output. Note that

S_{0}

is not accessed between state updates. Algorithm 1 summarizes the procedure.

Fig. 5: Accuracy and verbosity versus state-access traffic reduction relative to Standard as detailed in Appendix D . ReplaySSM and SketchSSM use

W=16

. DSQ uses

\{8,6,4\}

-bit states; GHOST (Nano/Super/GLM) and DRRQR (Qwen) prune

\{50,62.5,75\}\%

of the state. SketchSSM uses

\bar{G}\in\{8,4,2\}

for all four models. At

\bar{G}=8

, which reduces state-access traffic by 9.4–10.7

\times

, SketchSSM preserves average accuracy across the four benchmarks for all four models, whereas pruning and quantization degrade accuracy at smaller reductions. Verbosity is the mean generation length normalized to Standard. Raw accuracy results are provided in Appendix F.2 .

Method	Setting	Traffic	Nemotron Nano	Nemotron Super	Qwen3.8 Flash-Next	GLM 5.3 Flash
Standard	full-state	$1.00\times$	98.29	100.00	100.00	99.40
ReplaySSM	full-state	$1.88\times$	98.25	100.00	100.00	99.46
SketchSSM	$\bar{G}=8$	$9.4$ – $10.7\times$	98.03	100.00	100.00	99.74
SketchSSM	$\bar{G}=4$	$11.8$ – $12.8\times$	97.92	100.00	100.00	99.70
State pruning	50%	$2.00\times$	93.65	97.36	99.89	95.23
State pruning	62.5%	$2.67\times$	82.40	83.26	98.59	91.14

Table 1: Average recall (%) on four RULER retrieval tasks at 16K context. Traffic is the state-access traffic reduction relative to Standard. State pruning uses GHOST, except DRRQR for Qwen3.8 Flash-Next. Per-task results are in Appendix F.3 .

Ablation	Method	MATH-500	AIME25	GPQA-D	LCB	Average
Full method	SketchSSM	96.73	70.42	58.25	62.33	71.93
Sketch basis $\Omega$	Random orthogonal	94.73 ( $-2.00$ )	53.33 ( $-17.08$ )	53.54 ( $-4.71$ )	53.23 ( $-9.10$ )	63.71 ( $-8.22$ )
Coefficient vector $c_{t}$	Offline coefficient map	90.87 ( $-5.87$ )	42.08 ( $-28.33$ )	53.20 ( $-5.05$ )	27.72 ( $-34.60$ )	53.47 ( $-18.46$ )
Sketch-rank allocation	Uniform allocation	93.13 ( $-3.60$ )	48.33 ( $-22.08$ )	55.22 ( $-3.03$ )	50.90 ( $-11.43$ )	61.90 ( $-10.04$ )
	Error-only score	94.13 ( $-2.60$ )	59.17 ( $-11.25$ )	57.58 ( $-0.67$ )	54.39 ( $-7.94$ )	66.32 ( $-5.61$ )
	Gradient-only score	90.53 ( $-6.20$ )	35.83 ( $-34.58$ )	51.68 ( $-6.57$ )	44.76 ( $-17.57$ )	55.70 ( $-16.23$ )

Table 2: Ablation of internal design choices on Nemotron Nano 9B v2, covering sketch-basis calibration, state-dependent coefficients, and sketch-rank allocation, at

\bar{G}=4

. Pass@1 is averaged over eight seeds for AIME25 and three seeds for the other benchmarks.

Fig. 6: Linear-attention latency on one NVIDIA B300 with

W=16

and batch sizes 128, 256, and 512. Rows show Nemotron 3 Super, Qwen3.8 Flash-Next, and GLM 5.3 Flash, with SketchSSM at

\bar{G}\in\{8,4\}

. (a) Non-flush-step latency, incurred at each of the

W-1

steps between flush steps. (b) Flush-step latency, incurred once per window. Hatching indicates sketch and coefficient-map construction overhead at the flush step. (c) Total linear-attention speedup over Standard across the

W

steps of a window.

Fig. 7: Decoding performance for Nemotron 3 Super on one NVIDIA B300, comparing Standard, ReplaySSM, and SketchSSM at

\bar{G}\in\{8,4\}

. (a) Decode throughput at 2K and 8K up to each method’s maximum batch, averaged over two 128-step runs; labels give the speedup over Standard at batch 512 and at each method’s maximum batch. (b) Decode latency breakdown at batch size 512.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Model	$\bar{G}$	Exact ( $\mu$ s)	$P=4$ ( $\mu$ s)	Speedup
Nemotron 3 Super	8	2,076.21	460.31	$4.51\times$
Qwen3.8 Flash-Next	8	1,705.90	414.02	$4.12\times$
GLM 5.3 Flash	8	2,650.73	675.79	$3.92\times$

Appendix

Table 3: Complete flush-step latency on one NVIDIA B300 at batch size 256 and

W=16

, with FP32 states and BF16 sketch/coefficient map. Kernel measurements use synthetic activations; one CUDA graph holds the flush step of all recurrent layers, and we report the median of five captures divided by the layer count. Exact solves the full regularized Gram system with the same

\lambda=0.003

P=4

; speedup is Exact latency divided by

P=4

latency.

Method	Read	Write	Total	Reduction
Standard	$D$	$D$	$2D$	$1$
ReplaySSM	$D$	$D/W$	$D+D/W$	$\frac{2W}{W+1}$
Quant./Pruning	$D/C_{\mathrm{rate}}$	$D/C_{\mathrm{rate}}$	$2D/C_{\mathrm{rate}}$	$C_{\mathrm{rate}}$
SketchSSM	$B_{G}$	$D/W$	$B_{G}+D/W$	$\frac{2D}{B_{G}+D/W}$

Appendix

Table 4: Average state-memory traffic per state head per decode step, including reads and writes, and reduction relative to Standard.

D=4KV

B_{G}=((W-1)\mathcal{T}_{G}+D)/W

, and

\mathcal{T}_{G}=2G(V+K+eW)

. Full states use FP32; sketches, coefficient maps, projected erase vectors, and the state copies that dense-fallback heads read at non-flush steps use BF16.

Fig. 8: Head-level rank allocation at $\bar{G}=8$ for each model (rows). (a) Loss sensitivity $s_{l,h}=\mathbb{E}_{t}[\|g_{t,(l,h)}\|_{2}^{2}]$ of head $h$ in layer $l$ . (b) Rank-8 state-read error $\varepsilon_{l,h}(8)$ . (c) Allocated rank $G_{l,h}$ ; heads that retain the dense state are shown at the maximum color intensity. Nemotron Nano: ranks 1–41, 108 dense heads of 3,456 ( $9.94\times$ traffic reduction); Nemotron Super: 1–41, 159 of 5,120 ( $9.39\times$ ); Qwen3.8 Flash-Next: 1–47, 1 of 1,728 ( $10.68\times$ ); GLM 5.3 Flash: 1–44, 3 of 2,176 ( $10.68\times$ ). Sensitivity and error are normalized by their mean over heads.

Fig. 9: Head-level rank allocation at $\bar{G}=4$ , as in Fig. 8 . (b) shows the rank-4 state-read error $\varepsilon_{l,h}(4)$ . Nemotron Nano: ranks 1–39, 12 dense heads of 3,456 ( $12.26\times$ traffic reduction); Nemotron Super: 1–38, 5 of 5,120 ( $11.84\times$ ); Qwen3.8 Flash-Next: 1–46, no dense heads ( $12.81\times$ ); GLM 5.3 Flash: 1–44, no dense heads ( $12.81\times$ ).

			MATH-500		AIME25		GPQA-D		LiveCodeBench		Average
Method	Setting	Traffic	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.
Standard	full-state	$1.00\times$	$97.33_{\pm 0.42}$	5,257	$71.67_{\pm 3.56}$	18,868	$59.43_{\pm 1.17}$	13,499	$68.99_{\pm 1.20}$	14,073	$74.36_{\pm 0.99}$
ReplaySSM	full-state	$1.88\times$	$97.33_{\pm 0.61}$	5,142	$73.75_{\pm 5.18}$	18,344	$60.10_{\pm 1.01}$	13,288	$69.42_{\pm 2.40}$	14,174	$75.15_{\pm 1.46}$
SketchSSM	$\bar{G}=8$	$9.94\times$	$97.07_{\pm 0.58}$	5,548	$71.67_{\pm 3.09}$	19,341	$60.77_{\pm 0.77}$	13,969	$68.15_{\pm 1.63}$	14,915	$74.41_{\pm 0.91}$
SketchSSM	$\bar{G}=4$	$12.26\times$	$96.73_{\pm 0.31}$	5,801	$70.42_{\pm 5.18}$	20,336	$58.25_{\pm 0.58}$	14,257	$62.33_{\pm 1.50}$	16,705	$71.93_{\pm 1.36}$
SketchSSM	$\bar{G}=2$	$13.88\times$	$94.87_{\pm 0.64}$	6,510	$60.83_{\pm 4.27}$	21,692	$59.60_{\pm 1.01}$	14,673	$57.67_{\pm 1.81}$	18,963	$68.24_{\pm 1.20}$

Appendix

Table 5: Nemotron Nano 9B v2: numerical results underlying Fig. 5 . ReplaySSM is additionally reported. Accuracies are in percent; token counts are mean generated tokens. Standard, ReplaySSM, and SketchSSM report pass@1 averaged over eight seeds for AIME25 and three seeds for the other benchmarks; subscripts are the standard deviation across seeds (for Average, the per-benchmark deviations combined as

\sqrt{\sum_{b}\sigma_{b}^{2}}/4

). GHOST and DSQ are single runs. Traffic is state-read plus state-write traffic reduction relative to Standard. Average is the unweighted mean accuracy across the four benchmarks.

			MATH-500		AIME25		GPQA-D		LiveCodeBench		Average
Method	Setting	Traffic	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.
Standard	full-state	$1.00\times$	$98.27_{\pm 0.46}$	3,439	$82.50_{\pm 2.95}$	22,635	$77.27_{\pm 0.87}$	18,868	$80.21_{\pm 0.80}$	19,445	$84.56_{\pm 0.80}$
ReplaySSM	full-state	$1.88\times$	$98.27_{\pm 0.23}$	3,309	$84.17_{\pm 4.96}$	23,683	$76.09_{\pm 1.05}$	18,650	$79.47_{\pm 2.11}$	19,572	$84.50_{\pm 1.37}$
SketchSSM	$\bar{G}=8$	$9.39\times$	$98.40_{\pm 0.35}$	3,148	$82.08_{\pm 3.96}$	23,456	$77.44_{\pm 2.49}$	17,142	$80.74_{\pm 0.80}$	18,249	$84.67_{\pm 1.19}$
SketchSSM	$\bar{G}=4$	$11.84\times$	$98.20_{\pm 0.53}$	3,512	$85.00_{\pm 2.52}$	23,134	$78.45_{\pm 0.29}$	17,578	$76.08_{\pm 2.25}$	18,634	$84.43_{\pm 0.86}$
SketchSSM	$\bar{G}=2$	$13.61\times$	$97.87_{\pm 0.42}$	4,111	$80.83_{\pm 2.36}$	25,879	$75.42_{\pm 0.77}$	15,677	$72.91_{\pm 1.75}$	18,807	$81.76_{\pm 0.77}$

Appendix

Table 6: Nemotron 3 Super 120B-A12B: numerical results underlying Fig. 5 . ReplaySSM is additionally reported. Accuracies are in percent; token counts are mean generated tokens. Standard, ReplaySSM, and SketchSSM report pass@1 averaged over eight seeds for AIME25 and three seeds for the other benchmarks; subscripts are the standard deviation across seeds (for Average, the per-benchmark deviations combined as

\sqrt{\sum_{b}\sigma_{b}^{2}}/4

). GHOST and DSQ are single runs. Traffic is state-read plus state-write traffic reduction relative to Standard. Average is the unweighted mean accuracy across the four benchmarks.

			MATH-500		AIME25		GPQA-D		LiveCodeBench		Average
Method	Setting	Traffic	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.
Standard	full-state	$1.00\times$	$99.07_{\pm 0.23}$	1,434	$94.17_{\pm 1.54}$	19,455	$87.88_{\pm 1.82}$	14,736	$81.48_{\pm 0.92}$	22,039	$90.65_{\pm 0.64}$
ReplaySSM	full-state	$1.88\times$	$99.07_{\pm 0.12}$	1,456	$94.17_{\pm 2.36}$	19,714	$91.25_{\pm 0.77}$	14,921	$79.58_{\pm 2.16}$	21,819	$91.01_{\pm 0.82}$
SketchSSM	$\bar{G}=8$	$10.68\times$	$99.00_{\pm 0.20}$	1,474	$94.58_{\pm 2.48}$	19,560	$89.90_{\pm 0.87}$	14,232	$79.05_{\pm 1.38}$	21,665	$90.63_{\pm 0.74}$
SketchSSM	$\bar{G}=4$	$12.81\times$	$99.07_{\pm 0.31}$	1,786	$89.17_{\pm 2.95}$	22,457	$88.89_{\pm 1.52}$	15,947	$77.25_{\pm 0.18}$	25,061	$88.59_{\pm 0.83}$
SketchSSM	$\bar{G}=2$	$14.23\times$	$98.47_{\pm 0.58}$	2,819	$85.00_{\pm 5.63}$	29,609	$82.66_{\pm 1.05}$	21,334	$66.46_{\pm 0.97}$	30,263	$83.15_{\pm 1.46}$

Appendix

Table 7: Qwen3.8 Flash-Next: numerical results underlying Fig. 5 . ReplaySSM is additionally reported. Accuracies are in percent; token counts are mean generated tokens. Standard, ReplaySSM, and SketchSSM report pass@1 averaged over eight seeds for AIME25 and three seeds for the other benchmarks; subscripts are the standard deviation across seeds (for Average, the per-benchmark deviations combined as

\sqrt{\sum_{b}\sigma_{b}^{2}}/4

). DRRQR and DSQ are single runs. Traffic is state-read plus state-write traffic reduction relative to Standard. Average is the unweighted mean accuracy across the four benchmarks.

			MATH-500		AIME25		GPQA-D		LiveCodeBench		Average
Method	Setting	Traffic	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.	Tokens	Acc.
Standard	full-state	$1.00\times$	$97.73_{\pm 0.61}$	1,239	$81.67_{\pm 4.36}$	22,038	$90.24_{\pm 0.77}$	10,790	$74.50_{\pm 2.64}$	24,524	$86.03_{\pm 1.30}$
ReplaySSM	full-state	$1.88\times$	$97.40_{\pm 0.40}$	1,165	$86.25_{\pm 5.47}$	18,556	$89.56_{\pm 0.58}$	10,797	$74.18_{\pm 0.73}$	25,005	$86.85_{\pm 1.39}$
SketchSSM	$\bar{G}=8$	$10.68\times$	$97.60_{\pm 0.20}$	1,119	$85.42_{\pm 4.34}$	18,984	$88.89_{\pm 1.52}$	10,014	$82.01_{\pm 1.28}$	18,647	$88.48_{\pm 1.19}$
SketchSSM	$\bar{G}=4$	$12.81\times$	$97.53_{\pm 0.64}$	1,170	$86.67_{\pm 4.71}$	17,298	$85.69_{\pm 1.05}$	10,720	$79.58_{\pm 1.94}$	18,153	$87.37_{\pm 1.31}$
SketchSSM	$\bar{G}=2$	$14.23\times$	$95.40_{\pm 0.69}$	1,717	$80.00_{\pm 6.90}$	19,880	$81.14_{\pm 2.96}$	12,944	$67.09_{\pm 1.63}$	18,740	$80.91_{\pm 1.93}$

Appendix

Table 8: GLM 5.3 Flash: numerical results underlying Fig. 5 . ReplaySSM is additionally reported. Accuracies are in percent; token counts are mean generated tokens. Standard, ReplaySSM, and SketchSSM report pass@1 averaged over eight seeds for AIME25 and three seeds for the other benchmarks; subscripts are the standard deviation across seeds (for Average, the per-benchmark deviations combined as

\sqrt{\sum_{b}\sigma_{b}^{2}}/4

). GHOST and DSQ are single runs. Traffic is state-read plus state-write traffic reduction relative to Standard. Average is the unweighted mean accuracy across the four benchmarks.

Method	Setting	Traffic	Single UUID	Multi-key	Multi-value	Multi-query	Average
Standard	full-state	$1.00\times$	100.00	99.80	97.45	95.90	98.29
ReplaySSM	full-state	$1.88\times$	100.00	99.80	97.20	96.00	98.25
SketchSSM	$\bar{G}=8$	$9.94\times$	100.00	99.80	96.75	95.55	98.03
SketchSSM	$\bar{G}=4$	$12.26\times$	100.00	99.80	94.90	97.00	97.92
GHOST	50%	$2.00\times$	100.00	85.60	93.45	95.55	93.65
GHOST	62.5%	$2.67\times$	99.80	39.40	96.10	94.30	82.40

Appendix

Table 9: Recall scores (%) on four RULER retrieval tasks at 16K context for Nemotron Nano. Traffic is the state-read plus state-write traffic reduction relative to Standard, using the same accounting as the accuracy tables. Average is the unweighted mean of the four task scores.

Method	Setting	Traffic	Single UUID	Multi-key	Multi-value	Multi-query	Average
Standard	full-state	$1.00\times$	100.00	100.00	100.00	100.00	100.00
ReplaySSM	full-state	$1.88\times$	100.00	100.00	100.00	100.00	100.00
SketchSSM	$\bar{G}=8$	$9.39\times$	100.00	100.00	100.00	100.00	100.00
SketchSSM	$\bar{G}=4$	$11.84\times$	100.00	100.00	100.00	100.00	100.00
GHOST	50%	$2.00\times$	100.00	90.20	99.60	99.65	97.36
GHOST	62.5%	$2.67\times$	99.40	44.80	93.65	95.20	83.26

Appendix

Table 10: Recall scores (%) on four RULER retrieval tasks at 16K context for Nemotron Super. Traffic is the state-read plus state-write traffic reduction relative to Standard, using the same accounting as the accuracy tables. Average is the unweighted mean of the four task scores.

Method	Setting	Traffic	Single UUID	Multi-key	Multi-value	Multi-query	Average
Standard	full-state	$1.00\times$	100.00	100.00	100.00	100.00	100.00
ReplaySSM	full-state	$1.88\times$	100.00	100.00	100.00	100.00	100.00
SketchSSM	$\bar{G}=8$	$10.68\times$	100.00	100.00	100.00	100.00	100.00
SketchSSM	$\bar{G}=4$	$12.81\times$	100.00	100.00	100.00	100.00	100.00
DRRQR	50%	$2.00\times$	99.80	100.00	99.75	100.00	99.89
DRRQR	62.5%	$2.67\times$	99.40	98.60	97.10	99.25	98.59

Appendix

Table 11: Recall scores (%) on four RULER retrieval tasks at 16K context for Qwen3.8 Flash-Next. Traffic is the state-read plus state-write traffic reduction relative to Standard, using the same accounting as the accuracy tables. Average is the unweighted mean of the four task scores.

Method	Setting	Traffic	Single UUID	Multi-key	Multi-value	Multi-query	Average
Standard	full-state	$1.00\times$	99.80	100.00	97.80	100.00	99.40
ReplaySSM	full-state	$1.88\times$	99.80	100.00	98.05	100.00	99.46
SketchSSM	$\bar{G}=8$	$10.68\times$	100.00	100.00	98.95	100.00	99.74
SketchSSM	$\bar{G}=4$	$12.81\times$	100.00	100.00	98.80	100.00	99.70
GHOST	50%	$2.00\times$	100.00	98.80	82.10	100.00	95.23
GHOST	62.5%	$2.67\times$	99.80	95.00	73.20	96.55	91.14

Appendix

Table 12: Recall scores (%) on four RULER retrieval tasks at 16K context for GLM 5.3 Flash. Traffic is the state-read plus state-write traffic reduction relative to Standard, using the same accounting as the accuracy tables. Average is the unweighted mean of the four task scores.

Method	Update	Read	Traffic	MATH-500	AIME25	GPQA-D	LCB	Average
Baseline	full state	full state	$1.00\times$	97.33	71.67	59.43	68.99	74.36
Baseline	6-bit state	6-bit state	$5.06\times$	70.60	8.30	27.80	12.70	29.85
Ours	full state	4-bit state copy	$7.95\times$	97.40	68.33	55.05	63.49	71.07
Ours	full state	low-rank state, $\bar{G}=8$	$9.94\times$	94.93	57.92	49.83	45.61	62.07
Ours (SketchSSM)	full state	sketch, $\bar{G}=8$	$9.94\times$	97.07	71.67	60.77	68.15	74.41
Baseline	4-bit state	4-bit state	$7.40\times$	9.40	0.00	0.00	5.40	3.70

Appendix

Table 13: Accuracy ablation for maintaining full states and approximating reads on Nemotron Nano 9B v2 (

W=16

). Ours decouples full-state updates from approximate reads and varies only the read representation. Traffic is the state-read plus state-write traffic reduction relative to Standard; sketch reads store the sketch rows and coefficient map in BF16. Low-rank traffic includes BF16 reads of both the compressed state and fixed projection map, using the same per-head ranks as the sketch.

Explore similar work

May 7, 2026cs.LG

Echo: KV-Cache-Free Associative Recall with Spectral Koopman Operators

Long chain-of-thought reasoning and agentic tool-calling produce traces spanning tens of thousands of tokens, yet Transformer KV caches grow linearly with sequence length, creating a memory bottleneck on commodity hardware. State-space models offer constant-memory recurrence but suffer a memory cliff: retrieval accuracy collapses once the gap between a stored fact and its query exceeds the effective horizon of the recurrent state. We introduce Echo, a KV-cache-free associative recall architecture built around Spectral Koopman Attention (SKA); a drop-in replacement for attention layers that augments SSM blocks with a closed-form dynamical operator whose sufficient statistics are accumulated in constant memory with no KV cache. Echo fits a spectral linear system to the key and value history via kernel ridge regression and retrieves through a learned power-iterated filter, all from

O(r^{2})

streaming state where

r

is a small projection rank. On the Multi-Query Associative Recall benchmark, a pure Mamba-2 SSM fails to exceed chance accuracy (

{\sim}3\%

) across all gap lengths and KV-pair counts, while at the 50M parameter scale SKA-augmented models achieve

100\%

retrieval accuracy on every configuration tested, including distractor gaps of

4{,}096

tokens with

32

KV pairs. Across five additional transfer benchmarks including needle-in-a-haystack, tool-trace, and multi-hop retrieval, SKA consistently outperforms both pure SSM and SSM+Attention hybrids while maintaining constant inference memory. Ablations confirm that the spectral operator, not the prefix masking strategy, drives the retrieval gain.

Anupama Sridhar, Alexander Johansen

California, USA · Stanford University Computer Science California, USA

Koopman Operator

Ask PDF

Aug 3, 2026cs.LG

DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling

Modern language models are built primarily from Transformers, recurrent models, and their hybrid architectures. Transformers rely on token-level attention memories, while recurrent models such as state space models (SSMs) and linear attention maintain compact recurrent states. These architectures are typically instantiated separately or interleaved at the layer level, leaving open whether a shared memory representation can support both recurrent compression and attention-style retrieval. We study this question through the state space duality (SSD) view of Mamba-2, where the SSM state can be interpreted as a compressed associative key--value (KV) cache. We observe that Mamba-2 decodes token-conditioned values from this state but does not decode token-conditioned keys. Based on this observation, we propose DART (Decoded Attention over Recurrent sTates), which retains the chunk state contributions produced by the Mamba-2 chunked scan as chunk state memories, decodes token-conditioned keys and values from these memories, and performs state-memory attention (SMA) over the resulting KV pairs. The retrieved output is then combined with the native Mamba-2 output through a gated residual connection. DART supports practical training by reusing the Mamba-2 chunked scan and implementing SMA as a FlashAttention-style computation. Our analysis and experiments show that DART substantially reduces the length-dependent inference cache compared with a matched attention baseline (e.g.,

75\%

savings when the chunk size is

S=256

and the state size is

N=128

). Compared with Mamba-2, DART substantially improves associative recall and retrieval while preserving general language-modeling quality.

Yixiao Qian, Song Chen, Pengkai Wang +3

College of Control Science and Engineering, Zhejiang University · Department of Mathematics, National University of Singapore · Hong Kong Polytechnic University +1

Recurrent State Kimi Delta Attention Efficient Long-Context Inference Mamba

Ask PDF

Jul 18, 2026cs.CL

SpecLA: Efficient Speculative Decoding for Linear-Attention Models

Linear-attention models replace the growing KV cache with recurrent states, but autoregressive decoding still reads, updates, and writes these states one token at a time. Speculative decoding can reduce this cost by verifying several draft tokens in one target pass, yet existing speculative systems are designed for Transformer KV caches. For stateful linear-attention targets, verification must follow recurrent dependencies across chains and branches, acceptance must update only the accepted state trajectory, and the drafter must avoid submitting candidates that waste stateful verification work. This paper presents SpecLA, a speculative decoding runtime for stateful linear-attention models. SpecLA verifies chains and trees with topology-aware kernels, stores compact factors produced during verification to recover accepted states, and uses confidence pruning plus a target-aligned EAGLE-style drafter to feed useful candidates to the verifier. On an NVIDIA H100 with a public GDN-1.3B target, SpecLA achieves up to 1.70x end-to-end speedup over autoregressive decoding.

Zhibin Wang, Xuying Han, Zhaohua Yang +5

State Key Laboratory of Novel Software Technology Nanjing University Nanjing, China · Alibaba Group Hangzhou, China

Speculative Decoding Kimi Delta Attention Self-Speculative Recurrent State

Ask PDF