PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.33051v2cs.LGcs.DC

TL;DR

Hybrid-attention decoding makes recurrent-state access a bottleneck, and buffered updates do not remove repeated full-state reads. SketchSSM precomputes a compact sketch for read approximation while preserving full-state updates. Across four models, it reduces state-access traffic by approximately 10× while matching baseline average accuracy and improving linear-attention speed.

  • Problem

    ReplaySSM amortizes state writes, but each new query still requires a full-state read even while the state remains unchanged between updates.

    Source: §1 · pp. 1–2

  • Method

    SketchSSM precomputes outputs for fixed query basis vectors with a full-state read at each update, then reconstructs subsequent outputs from the compact sketch while preserving full-state updates.

    Source: §1 · p. 2

  • Results

    At mean sketch rank 8, SketchSSM reduces state-access traffic by approximately 10× while matching the FP32 full-state baseline’s average accuracy across four benchmarks.

    Source: §5.1 · p. 8, §6 · p. 10

  • Takeaways & Limitations

    SketchSSM reduces state-read traffic without propagating state-compression errors through updates.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

Hybrid-attention decoding makes recurrent-state access a bottleneck, and buffering updates leaves repeated full-state reads intact. SketchSSM approximates reads with a compact sketch while preserving full-state updates, reducing traffic and achieving strong accuracy and speed results.

1 INTRODUCTION

  • ReplaySSM buffers keys and values to amortize state writes, but every decode step still reads the full state while it remains unchanged between updates.In the large-batch regime, recurrent-state reads and updates become a major decode bottleneck.Source: §1 · pp. 1–2
  • A small fixed query basis can reconstruct changing-query state-read outputs with low error, even though future queries are unknown.The basis is fixed offline, while reconstruction uses state- and query-dependent coefficients.Source: §1 · p. 2
  • SketchSSM uses one full-state read at each update to precompute basis-vector outputs, then reconstructs each intervening query output from sketch vectors without reading the full state.It approximates reads but preserves exact full-state updates, avoiding state-compression errors propagating through updates.Source: §1 · p. 2
  • At mean sketch rank 8, SketchSSM reduces state-access traffic by approximately 10× while matching the FP32 full-state baseline’s average accuracy across four decode benchmarks.On one NVIDIA B300, linear-attention kernel speedups over Standard reach 7.30× for Mamba-2, 5.02× for GDN, and 5.24× for KDA.Source: §1 · p. 2
  • SketchSSM allocates sketch ranks differently across state heads to improve accuracy.The paper also reports better accuracy–traffic tradeoffs than state pruning and quantization, and optimized kernels across Mamba-2, GDN, and KDA.Source: §1 · p. 2

2 RELATED WORK

Prior work addresses recurrent-state inference through state compression or buffered updates, while matrix sketching offers a framework for compact approximations. SketchSSM combines an offline-calibrated query basis with coefficients chosen to minimize output error.

Efficient Recurrent-State Inference

  • State quantization and pruning reduce traffic by shrinking recurrent states, but can limit information capacity and propagate compression errors through updates.ReplaySSM and KVBuffer buffer keys and values to reduce write traffic while leaving read traffic unchanged.Source: §2 · p. 3
  • SketchSSM complements buffered updates by reducing the read traffic that ReplaySSM and KVBuffer leave unchanged.ReplaySSM and KVBuffer apply accumulated updates at each flush.Source: §2 · p. 3

Matrix Sketching

  • Matrix sketching covers randomized projections, streaming sketches, and learned sketches; prior LLM low-rank methods approximate weights or reduce KV-cache storage.SketchSSM builds on the range-finding framework of Halko et al. (2011).Source: §2 · p. 3
  • SketchSSM sketches the state with an offline-calibrated query basis and reconstructs outputs using coefficients that minimize output error for the current state and query.This differs from the cited prior examples, which approximate weights or reduce KV-cache storage.Source: §2 · p. 3

3 MOTIVATION

In hybrid-attention decoding, linear attention can dominate latency because each step accesses recurrent state. SketchSSM uses a fixed query basis and compact precomputed outputs to approximate reads while retaining exact state updates.

3.1 STATE-ACCESS BOTTLENECK IN HYBRID-ATTENTION DECODING

  • At Bmax, linear attention accounts for 42–75% of decode latency with Standard across the three profiled models, and ReplaySSM lowers that share to 31–53% while leaving read traffic unchanged.ReplaySSM buffers keys and values to amortize state-write traffic; reducing state-read traffic is still needed to further reduce the bottleneck.Source: §3.1 · p. 3

3.2 WHY EVERY DECODE STEP READS THE FULL STATE

  • Although the state S0 stays fixed during a buffering window, the effective query changes at each decode step, so directly computing its state-read output requires reading all K × V state entries.The changing query reflects both the new query and the accumulated transition; the reformulation isolates state dependence in the state-read term.Source: §3.2 · pp. 3–4

3.3 REPLACING FULL-STATE READS WITH A COMPACT SKETCH

  • A small fixed query basis can approximate changing effective queries with state- and query-dependent coefficients, preserving the state-read output in a low-dimensional subspace.With G = 4 basis vectors in a K = 128-dimensional query space, the retained output energy is 88.4–97.5% across the three model families.Source: §3.3 · p. 4
  • At each flush, SketchSSM precomputes basis-vector outputs into a sketch; intervening steps reconstruct outputs from sketch vectors and query-dependent coefficients instead of rereading the full state.This reduces read traffic from K × V elements to G(K + V) elements while keeping updates exact from buffered inputs.Source: §3.3 · pp. 4–5

4 SKETCHSSM

SketchSSM keeps full-state updates but approximates state reads using an offline-calibrated query basis, a refreshed sketch, and query-dependent coefficients. It also allocates sketch ranks across heads to balance accuracy and traffic.

4.1 CONSTRUCTING THE SKETCH

  • Because the optimal basis depends on changing S0 and per-state optimization is expensive, SketchSSM fixes Ω during inference using offline calibration on average state geometry.Source: §4.1 · p. 5
  • At each window boundary, buffered updates are applied to full state S0, and one full-state read constructs sketch U before the updated state is written back.Source: §4.1 · p. 6
  • Between flushes, query-dependent coefficients combine precomputed sketch vectors to approximate each state-read output without accessing the full state.Source: §4.1 · p. 6, Figure 4 · p. 6
  • A low-rank-plus-diagonal approximation speeds state-dependent coefficient-map computation 4.51×, while sketch and map construction adds 20.3–27.1% to flush-step latency.Source: §4.1 · p. 6

4.2 ALLOCATING SKETCH RANKS ACROSS LAYERS AND HEADS

  • At mean rank 8, SketchSSM reduces state-access traffic 9.4–10.7× while preserving average accuracy across four benchmarks for all four models.Source: Figure 5 · p. 7, §4.2 · p. 7
  • Rank allocation scores each head’s reconstruction error against loss sensitivity, then uses a knapsack formulation; it improves accuracy over uniform allocation by 10.04 percentage points on average.Source: §4.2 · p. 7

5 EVALUATION

Evaluation across four hybrid-attention models finds that rank-8 SketchSSM preserves average decode-benchmark accuracy while reducing state-access traffic, and delivers kernel and end-to-end decode speedups. Recall results, ablations, and speed measurements characterize its tradeoffs and implementation.

5 EVALUATION

  • The evaluation covers four NVFP4-weight models spanning Mamba-2, GDN, and KDA, with accuracy tests against full-state, quantization, and pruning baselines.SketchSSM calibration uses WikiText-2 only, out of domain for every evaluated benchmark; speed tests run on one NVIDIA B300.Source: §5 · pp. 7–8

5.1 ACCURACY

  • At mean rank 8, SketchSSM reduces state-access traffic by 9.4–10.7× while preserving average accuracy across four decode benchmarks and all four models.At rank 4, traffic falls 11.8–12.8× with at most a 2.4-point accuracy decrease; at approximately 4× traffic reduction, quantization loses 1.6–30.6 points and pruning 74.1–87.6 points.Source: §5.1 · p. 8
  • At 16K context, rank-8 and rank-4 SketchSSM recall stays within 0.4 points of Standard on every model, while 75% pruning yields 43.05–76.70%.SketchSSM approximates reads but updates the full state exactly, so read errors do not accumulate in the state through updates.Source: §5.1 · p. 8, Table 1 · p. 8
  • At the same mean rank, state-weighted basis calibration, state-dependent coefficient maps, and joint output-error/loss-sensitivity rank allocation improve accuracy over their respective alternatives.The reported gains are 2.0–17.1, 5.1–34.6, and 3.0–22.1 percentage points, respectively; the table evaluates these ablations on Nemotron Nano 9B v2 at rank 4.Source: §5.1 · p. 8, Table 2 · p. 9

5.2 SPEEDUPS

  • At rank 8, total linear-attention speedups over Standard reach 7.30× for Nemotron 3 Super, 5.02× for Qwen3.8 Flash-Next, and 5.24× for GLM 5.3 Flash.The Fig. 6 measurements use one NVIDIA B300, W = 16, and batch sizes 128, 256, and 512; construction adds 20.3–27.1% to Nemotron 3 Super flush latency.Source: §5.2 · p. 9, Figure 6 · p. 9
  • With rank 8, maximum decode throughput over Standard rises 2.77× at 2K and 2.33× at 8K; at batch 512, linear attention's decode-time share falls to 22.3% and 21.1%, respectively.Fig. 7 compares Standard, ReplaySSM, and SketchSSM throughput and shows the batch-512 latency breakdown.Source: §5.2 · pp. 9–10, Figure 7 · p. 10

6 CONCLUSION

SketchSSM keeps exact full-state updates while reconstructing outputs from compact sketches, and at mean rank 8 matches average full-state accuracy with about 10× lower state-access traffic. Its derivation separates calibration of a shared sketching matrix from computation of current-state, current-query coefficients.

6 CONCLUSION

  • At mean rank 8, SketchSSM reduces state-access traffic by approximately 10× while matching average FP32 full-state accuracy across four hybrid-attention models.On one NVIDIA B300, it achieves 5.02–7.30× linear-attention speedups and up to 2.77× higher maximum throughput over the full-state baseline.Source: §6 · p. 10
  • SketchSSM preserves exact full-state updates and uses compact sketches to reconstruct outputs.The sketch is constructed at runtime from the current state using an offline-calibrated matrix; sketch ranks are allocated offline using each state head's output error and impact on model loss.Source: §6 · p. 10

A DERIVATION OF SKETCHSSM

  • The derivation treats shared sketch-matrix calibration from state–query samples separately from computing coefficients for the current state and query.Source: A DERIVATION OF SKETCHSSM · p. 14

A.1 DERIVATION OF THE OPTIMAL OFFLINE SKETCHING MATRIX (Ω)

SketchSSM derives its offline sketch basis through PCA and computes query-dependent coefficients from the current state. A low-rank-plus-diagonal approximation reduces coefficient-map construction cost while leaving the recurrent state and sketch exact.

A.1 DERIVATION OF THE OPTIMAL OFFLINE SKETCHING MATRIX (Ω)

  • The offline basis minimizes calibration-averaged reconstruction error by selecting the G eigen-directions with the largest captured energy.The derivation converts weighted error to projection error, then uses the PCA variational principle; the residual is energy in the remaining K − G directions.Source: §A.1 · pp. 14–15

A.2.1 DERIVATION OF THE OPTIMAL COEFFICIENT VECTOR

  • Within each window, SketchSSM reuses a state-dependent coefficient map to approximate the exact query output as a linear combination of sketch vectors.The coefficient map is constructed at the flush step; subsequent steps use its matrix–vector product with the changing query.Source: §A.2.1 · pp. 15–16

A.2.2 LOW-RANK-PLUS-DIAGONAL APPROXIMATION FOR THE COEFFICIENT MAP

  • The coefficient solve approximates the dense state metric with a low-rank-plus-diagonal form, preserving selected output-space correlations and remaining coordinate-wise energy.Only the metric used to compute coefficients is approximated; the recurrent state and sketch remain exact.Source: §A.2.2 · pp. 16–17
  • 3.92–4.51× lower complete flush-step latency across Mamba-2, GDN, and KDA with P = 4 than with the exact regularized coefficient map.Both paths use the same sketch bases and rank allocations and fuse coefficient construction into the state update without an additional full-state read.Source: §A.2.2 · p. 17, Table 3 · p. 17

B SKETCHSSM ALGORITHM

  • Each window uses W − 1 non-flush steps that read the sketch and coefficient map, followed by a flush that applies buffered updates to the full state exactly.The flush computes the output from the updated state and refreshes the sketch and coefficient map.Source: B SKETCHSSM ALGORITHM · p. 18

C WY-BASED EFFECTIVE-QUERY AND COEFFICIENT COMPUTATION

  • SketchSSM adapts a WY representation to decay-and-erase transitions, deriving exact effective-query evaluation before expressing it in coefficient space.The adaptation is described for the buffered transitions in the WY-based computation.Source: C WY-BASED EFFECTIVE-QUERY AND… · p. 19

C.1 PARALLEL EFFECTIVE-QUERY EVALUATION

A WY-style factorization replaces serial application of buffered transitions with parallel accumulation of their contributions to the current query. Projecting transition factors into coefficient space preserves exact coefficient computation for a fixed map.

C.1 PARALLEL EFFECTIVE-QUERY EVALUATION

  • The WY-style factorization evaluates buffered transition contributions in parallel without reading the full state, while retaining causal dependence between arriving tokens.For Mamba-2, which has no erase terms, the effective query reduces to d[t] ⊙ qt.Source: §C.1 · p. 19

C.2 COEFFICIENT COMPUTATION WITH PROJECTED FACTORS

  • Projecting transition factors through the fixed coefficient map computes Cq̃t exactly, so each appended slot needs only its G-dimensional projected factor rather than a K-dimensional one.For Mamba-2, the projected erase vectors and associated sums are unnecessary.Source: §C.2 · pp. 19–20

D STATE MEMORY TRAFFIC REDUCTION

SketchSSM reduces recurrent-state read traffic by storing and reading compact representations between flushes while retaining full-state reads and writes at flushes. Its traffic accounting includes state reads and writes but excludes smaller auxiliary traffic.

D STATE MEMORY TRAFFIC REDUCTION

  • Standard and ReplaySSM read D bytes per step, whereas ReplaySSM and SketchSSM average D/W bytes of state writes by flushing buffered updates once every W steps.Quantization and pruning reduce both read and write traffic to D/Crate per step.Source: D STATE MEMORY TRAFFIC REDUCTION · p. 21
  • At non-flush steps, SketchSSM reads the BF16 sketch and coefficient map; each flush reads the full FP32 state, yielding average read cost BG = ((W − 1)TG + D)/W.For GDN and KDA, projected erase vectors add to non-flush read cost; dense-fallback heads instead read a BF16 state copy.Source: D STATE MEMORY TRAFFIC REDUCTION · p. 21, Table 4 · p. 21
  • The traffic metric excludes ring-buffer traffic and writes of the sketch and coefficient map, although measured speedups reflect those costs.At W = 16, ReplaySSM's reported traffic reduction is 32/17 ≈ 1.88×.Source: D STATE MEMORY TRAFFIC REDUCTION · p. 21

E OFFLINE SKETCH-RANK ALLOCATION

Sketch ranks are allocated offline using head-level reconstruction error and loss sensitivity under a traffic budget. The evaluation also reports rank-dependent traffic reduction and generated-token verbosity across four models.

E OFFLINE SKETCH-RANK ALLOCATION

  • For each head, calibration evaluates candidate prefix ranks using reconstruction error paired with output-gradient sensitivity, rather than multiplying separately averaged error and gradient statistics.Basis fitting and rank allocation use disjoint calibration examples.Source: E OFFLINE SKETCH-RANK ALLOCATION · p. 22
  • The rank objective uses a traffic price to make head choices independent, then adjusts the price by bisection until total traffic meets the budget.The procedure retains the best feasible allocation and uses remaining discrete slack for improving single-head substitutions.Source: E OFFLINE SKETCH-RANK ALLOCATION · p. 22
  • Allocated ranks are larger for heads with both higher loss sensitivity and higher state-read error, which vary widely across layers and heads.The separable first-order surrogate omits cross-head terms.Source: E OFFLINE SKETCH-RANK ALLOCATION · p. 22
  • 9.39–10.68× traffic reduction at mean rank 8 and 11.84–12.81× at mean rank 4 across the four models.At rank 8, dense fallback is used by 108/3,456 Nemotron Nano heads, 159/5,120 Nemotron Super heads, 1/1,728 Qwen3.8 Flash-Next heads, and 3/2,176 GLM 5.3 Flash heads; at rank 4, the corresponding counts are 12, 5, 0, and 0.Source: Figure 8 · p. 23, Figure 9 · p. 24

F.1 DETAILED EVALUATION SETUP

  • Offline calibration uses WikiText-2 rather than benchmark data, with separate training-split prompts for basis fitting and validation-split sequences for rank allocation.The resulting sketching matrices and rank allocations remain fixed during inference; the coefficient approximation uses P = 4 pivots and W = 16.Source: §F.1 · p. 25

F.2 ACCURACY RESULTS ON DECODE BENCHMARKS

  • At mean rank 8, normalized verbosity is 1.04, 0.95, 1.00, and 0.86 for Nemotron Nano, Nemotron Super, Qwen3.8 Flash-Next, and GLM 5.3 Flash, respectively.The reported results describe a tradeoff between state-access traffic reduction and generated response length.Source: §F.2 · p. 26

F.3 RECALL EVALUATION

SketchSSM maintains strong recall on four RULER retrieval tasks across the evaluated models, while state pruning produces larger average-score losses. The evaluation uses fixed offline-calibrated bases and rank allocations, with no calibration on recall examples.

Evaluation setup

  • The evaluation covers four RULER tasks at a 16K context budget, using 500 examples per task and greedy generation capped at 128 tokens.Scores use answer coverage: reference answers found in generated text by case-insensitive substring matching, averaged over examples; the four-task average is unweighted.Source: §F.3 · p. 28
  • SketchSSM uses offline-calibrated bases and rank allocations without calibration on recall examples; its full state is FP32, while its sketch and coefficient map are BF16.Methods are compared with Standard execution, ReplaySSM at W = 16, SketchSSM at W = 16 and P = 4, and pruning at 50%, 62.5%, and 75% sparsity.Source: §F.3 · p. 28

Results

  • On RULER, SketchSSM averages 97.92–98.03% vs 98.29% for Standard on Nano, 100% on Super and Qwen, and 99.70–99.74% vs 99.40% for Standard on GLM.On Super and Qwen, SketchSSM matches both full-state methods at 100% on all four tasks for each reported rank.Source: §F.3 · p. 28, Table 9 · p. 28, Table 10 · p. 29
  • At 75% sparsity, pruning averages 76.70% on Nano, 53.03% on Super, 43.05% on Qwen, and 46.78% on GLM.The largest task-level SketchSSM decrease is Multi-value on Nano at rank 4: 94.90% versus 97.45% for Standard.Source: §F.3 · p. 28

G ABLATION OF PRESERVING FULL-STATE UPDATES WITH APPROXIMATE READS

This ablation holds full-state updates fixed while varying the compact representation used for reads. Among the tested representations, sketches improve average accuracy over quantized state copies at comparable state-access traffic reductions.

Ablation design

  • The low-rank baseline uses an offline basis from leading eigenvectors of the average state Gram matrix and fixed rather than state-dependent query coefficients.It builds on truncated SVD and directly projects queries for coefficients.Source: G ABLATION OF PRESERVING FULL-STATE… · p. 30

Motivation

  • Keeping the full state for updates avoids propagating state-compression errors, and full-state update traffic is amortized because updates occur only at flush steps.The approximation is applied only to the read path.Source: G ABLATION OF PRESERVING FULL-STATE… · p. 30

Ablation design

  • All ablation variants build a compact read representation from the updated state at a flush step and use it between flushes, while retaining the full state for later updates.The variants differ in read representation: a quantized state copy, a low-rank approximation, or a sketch.Source: G ABLATION OF PRESERVING FULL-STATE… · p. 30

Results

  • Within the full-state-update design, sketches improve average accuracy over quantized read-only copies at comparable state-access traffic reductions.Table 13 compares quantized read-only copies, low-rank state representations, and sketches, with state quantization as a reference; it reports Nemotron Nano 9B v2 at W = 16.Source: G ABLATION OF PRESERVING FULL-STATE… · p. 30, Table 13 · p. 30

Related papers

Loading 2609.33051v2…