PaperlayerView exact revision on arXiv

Source-linked AI summary

Prefix Keys and Values Matter Late: RIPE Caching for Self-Speculative Decoding in Looped Language Models

Dezhen Lu, Runguo Li, Zixiang Ni

arXiv:2610.34436v1cs.DC

TL;DR

Looped models store keys and values for every recurrent loop, making the KV cache large and raising which planes need to be retained. The paper proposes RIPE, which keeps the deepest plane except for a recent full-depth window; on Huginn, that window removes most of the deepest-plane error tail at a fraction of full-cache memory.

  • Problem

    Each recurrent loop stores its own keys and values, making the KV cache large and leaving which planes must be retained as a direct question.

    Source: §1 · p. 1

  • Method

    RIPE keeps the deepest plane for settled positions and all planes for a recent window, with a proof of block-independent batched verification when the window covers the block.

    Source: §1 · p. 1, §4 · p. 5, §7 · p. 8

  • Results

    At 11.9% of the full cache’s memory on Huginn, RIPE removes most of the deepest plane’s error tail.

    Source: §7 · p. 8

  • Takeaways & Limitations

    When the window covers the verification block, RIPE’s batched verification reads match sequential decoding.

    Source: §4 · p. 5

    RIPE with W > 0 was evaluated only by teacher-forced replay on Huginn, and the deepest-plane tail statistics rest on four prompts.

    Source: §7 · p. 8

Abstract

from arXiv · show

A looped language model runs one block of layers $R$ times per token, so its key-value (KV) cache holds $R$ planes per recurrent layer, and self-speculative decoding (draft with a few loops, verify with all $R$) reads them all. We find that prefix keys and values matter late: at equal perturbation size, loops 25-32 are $e^{15.0}$ (Huginn-0125) and $e^{12.8}$ (Raven-Llama) times more sensitive than loops 2-9, and the planes approach the last one geometrically, at 0.88 and 0.81 per loop. Whether the deepest plane alone suffices then depends on how ripe the late planes are: on Raven, storing only it is indistinguishable from a bf16-to-fp32 change; on Huginn, 1% of positions carry 74% of its error. We turn this into RIPE, which stores the deepest plane everywhere except a window of recent positions kept at full depth. On Huginn, a 64-position window removes 89% of the tail positions at 11.9% of the full cache's memory (1024-token context), and any window at least as long as the draft block makes batched verification block-independent by construction. Nine pre-registered GPU batches back these results. Simple models of acceptance, decoding cost and memory, two of them checked on held-out measurements, account for the measured self-speculative speed-ups ($2.09\times$ with the full cache, $2.58\times$ with the deepest plane) and decompose the $2.24\times$ fixed-memory throughput gain of one shared plane.

1 INTRODUCTION

Because each recurrent loop stores its own keys and values, looped models can have large KV caches; the introduction asks which planes must be retained and presents RIPE as a cache design for that question.

1 INTRODUCTION

  • 2,788 KB per token for Huginn-0125’s full cache contrasts with 169 KB if only one plane is kept.At an 808-token sequence and a 95 GB budget, the full cache holds 36 sequences versus 302 with one plane.Source: §1 · p. 1
  • The paper measures which planes must be stored in nine GPU batches with decision rules and predictions frozen before each run.The measurements are organized with small analytical models and used to develop a cache design.Source: §1 · p. 1
  • RIPE keeps the deepest plane for settled positions and full depth for a recent window, with a proof that batched verification is block-independent when the window covers the block.The paper also reports nine pre-registered GPU batches and analytical models of acceptance, cost, memory, and fixed-memory throughput.Source: §1 · pp. 1–2
  • Sensitivity and ripening models indicate that only late loops need accurate prefix keys and values, and predict when the deepest plane costs no more than bf16 rounding.The introduction also notes that truncating finished tokens to shallow planes was the starting plan, but is backwards.Source: §1 · p. 1
  • The paper’s performance models cover acceptance, decoding cost, memory, and fixed-memory throughput; per-head filtering, sinks, and a second plane do not help.Two models are checked on held-out data, and cuDNN attention makes bf16 decoding nonreproducible on GH200.Source: §1 · p. 2

2 BACKGROUND

A looped model embeds each token with a prelude, repeatedly applies a core block, then produces logits with a coda. The background defines its cache policies and self-speculative decoding setup.

2 BACKGROUND

  • A prelude embeds each token, a core block runs R times per token, and a coda produces logits.Each core loop emits keys and values across its core layers, which form that loop’s plane.Source: §2 · p. 2
  • Huginn-0125 uses 32 loops and four core layers, while Raven-Llama uses 32 loops and six core layers with grouped-query attention.The models also differ in their prelude and coda layer counts and KV-head configurations.Source: §2 · p. 2
  • The full cache reads plane r at loop r, whereas the deepest-plane cache reads plane R; Huginn’s compress-sS with S equal to the core-layer count is exactly the deepest-plane cache.The policy assigns which plane of each prefix position a token reads at each loop.Source: §2 · p. 2
  • 2,788 versus 169 KB per token is Huginn’s full-cache versus one-plane storage, while Raven’s corresponding values are 410 versus 28.7 KB.These storage figures compare keeping all recurrent planes with keeping one plane.Source: §2 · p. 2
  • Self-speculative decoding drafts γ tokens with d loops, verifies γ + 1 tokens through all R loops, accepts the longest agreeing prefix, then adds one verifier token.The setup denotes tokens per verification by τ.Source: §2 · p. 2
  • Teacher-forced prefill records keys and values while substituting prefix planes, and generation replays greedy trajectories against an fp32 full-cache reference to report per-position KL.The fp32 full cache sets the precision floor in this measurement setup.Source: §2 · p. 2

3 MODELS

The section fits post hoc analytical models to explain measured behavior and reports which data were used to fit each model and check it.

3.1 Sensitivity: only the deep loops’ prefix KV matters

  • At equal perturbation size, deep loops are e15.0 times more sensitive on Huginn and e12.8 on Raven; the deepest-plane cache matches the full cache in 0.990 of texts.A window of the first four loops reaches 0.845.Source: §3.1 · p. 3

3.2 Ripening: planes converge geometrically

  • Planes converge geometrically at 0.88 per loop on Huginn and 0.81 on Raven; Raven’s deep loops reach the bf16 floor, while Huginn’s remain above it.Across deep loops 25–31, Huginn’s keys remain 0.0116 from plane R in bf16, while Raven’s are at its 0.0047 floor.Source: §3.2 · p. 3

3.3 Acceptance

  • Deepest-plane drafts accept α4 = 0.93 rather than 0.81, and the acceptance model matches simulated tokens per verification within 2.6%.At d = γ = 4, batch 8 measured 3.34 and 4.47 tokens per verification on 24 other GSM8K problems.Source: §3.3 · p. 4

3.4 Cost

  • 2.09× with the full cache and 2.58× with the deepest plane were measured self-speculative speed-ups; the fitted cost model predicts 2.08× and 2.61×.A single η = 1.07 was fitted to all 24 per-problem timings; per-problem errors reached 5.4%.Source: §3.4 · p. 4

3.5 Memory and fixed-memory throughput

  • Under a 95 GB budget, the memory model predicts maximum batch sizes of 36 with the full cache and 302 with one slot.It predicts batch-8 peak reserved memory within 0.15% on cells run later on another node.Source: §3.5 · p. 4
  • A shared slot yields a 2.24× fixed-memory throughput advantage: 1.63× from cache size at the same batch and 1.34× from running a larger batch.The larger batch is B = 64, compared with the full cache’s best batch of B = 32.Source: §3.5 · p. 5

4 RIPE

RIPE keeps the deepest plane for old positions and full depth for a recent window. When the window is at least as long as the verification block, batched verification matches sequential decoding in exact arithmetic.

4 RIPE

  • RIPE reads the deepest plane for old positions and each position’s own plane for the W most recent positions.Its greedy self-speculative step drafts with d loops, verifies with R loops, and drops non-deepest planes once positions leave the window.Source: §4 · p. 5
  • At a 1024-token context, W = 64 costs 341 MB on Huginn (11.9% of 2,855 MB) and 54 MB on Raven (12.8%).The implementation stores the deepest plane for every position and a ring buffer of W × (R − 1) core planes.Source: §4 · p. 5
  • W ≥ γ makes batched verification compute the same logits and cached keys and values as sequential decoding under RIPE, in exact arithmetic.The proposition’s condition is that the recent full-depth window is at least as long as the block.Source: §4 · p. 5
  • With W = 0, batched verification reads a different plane for earlier block tokens, so outputs depend on where blocks start.The deepest-plane cache does not satisfy the condition for block-independent verification.Source: §4 · p. 5
  • The ripeness test suggests W = γ suffices for exact verification on Raven, while W = 64 removed most tail positions on Huginn; the speculative decoder was not tested with W > 0.Drafts may read any stored plane without affecting correctness, and reading plane R everywhere gives higher acceptance.Source: §4 · p. 5

5 EVALUATION

Across the evaluation, the deepest-plane cache improved real self-speculative decoding speed, while Huginn retained a concentrated error tail that RIPE windows reduced. Results also show that fixed-memory throughput improved with one shared slot, while some proposed cache policies did not help and bf16 reproducibility depended on the attention backend.

5 EVALUATION

  • The evaluation used pre-registered batches with frozen decision rules and predictions, and ran on one GH200 with specified software and model revisions.Prompts were GSM8K 8-shot problems and AlpacaEval instructions or reference texts.Source: §5 · pp. 5–6

5.1 Generation fidelity and the tail

  • On Huginn, 1% of positions carried 74% of deepest-plane KL; on Raven, the deepest plane was indistinguishable from a precision change.Huginn’s 9 positions with KL above 0.1 occurred in four prompts, while Raven had none above 0.1.Source: §5.1 · p. 6

5.2 Self-speculative decoding

  • 2.09× with the full cache and 2.58× with the deepest plane were the measured self-speculative speed-ups over sequential decoding on Huginn.The decoder used depth-4 drafts, four drafts per block, and one 32-loop verification pass per block.Source: §5.2 · p. 6

5.3 Fixed-memory throughput

  • With one shared slot, WaveFront reached 487 tokens/s at B = 64 and 572 tokens/s at B = 256, yielding 2.24× and 2.26× gains over the best no-sharing configurations.The no-sharing configurations reached 218 tokens/s at B = 32 and 253 tokens/s at B = 34, respectively; the batches ran on different nodes and were not pooled.Source: §5.3 · pp. 6–7

5.4 What did not help

  • Per-head position filtering cost 72 times more KL than deepest-plane storage, and neither sinks nor a second plane removed Huginn’s tail positions.Selecting at loop 4 instead of loop 2 cut the filtering cost to 0.33 of its value.Source: §5.4 · p. 7

5.5 Reproducibility on GH200

  • The bf16 full-cache reference matched a same-process rerun bitwise at only 30% of positions, and disabling cuDNN attention made reruns bitwise identical.The authors conclude that a bf16 “lossless” claim needs a stated reference computation.Source: §5.5 · pp. 7–8

6 RELATED WORK

The related work covers recurrent-depth models, self-speculative decoding, and KV compression. It positions this paper’s focus on late-loop prefix K/V alongside prior approaches that reuse early-loop attention support.

Looped models

  • Looped models reuse a core block, and prior work reports that Ouro can keep only the last loop’s KV nearly losslessly on its model.The paper measures when this holds and why it fails on Huginn.Source: §6 · p. 8

Self-speculation in looped models

  • Prior self-speculative methods include zero-shot self-speculation, pipelined branches, cross-recurrence KV sharing, mature values for shallow drafts, and LoopAhead.The passage also identifies layer-skipping Draft&Verify and prior acceptance models.Source: §6 · p. 8

KV compression

  • Prior KV-compression work spans cross-loop compression and quantization, while RIPE’s recent window is compared to a streaming window along depth rather than sequence.The related work also notes standard-transformer methods including sinks, quantization, and paging.Source: §6 · p. 8

7 LIMITATIONS AND CONCLUSION

The paper’s conclusion is that late-loop prefix keys and values matter, while plane ripening determines whether the deepest plane alone suffices. RIPE uses a recent full-depth window with the deepest plane elsewhere, but its evaluation and supporting evidence have defined scope limits.

7 LIMITATIONS AND CONCLUSION

  • RIPE with W > 0 was evaluated only by teacher-forced replay on Huginn, not inside the speculative decoder.The tail statistics rest on four prompts, and evaluation used greedy decoding, GSM8K, and AlpacaEval.Source: §7 · p. 8
  • The conclusion identifies late-loop prefix K/V as the important planes and plane ripening as the determinant of whether the deepest plane alone suffices.Simple sensitivity, ripening, acceptance, cost, and memory models account for measurements from nine pre-registered batches; memory and acceptance models predict held-out measurements.Source: §7 · p. 8
  • At 11.9% of the full cache’s memory on Huginn, RIPE removes most of the deepest-plane error tail and keeps batched verification consistent with sequential decoding.RIPE retains the deepest plane except for a short recent window.Source: §7 · p. 8

A BATCHES

The batches used a pre-registered protocol with frozen predictions and documented checks. Table 5 identifies the pre-registered batches run on one GH200 each.

A BATCHES

  • Each batch had frozen decision rules and predictions, a frozen SHA-256 file table, on-node CPU self-tests, and an unscored Nsight Compute pass.The archive was verified at both ends.Source: A BATCHES · p. 10
  • The archived batch protocol included checks before and after each run.These included on-node CPU self-tests and verification of the archive at both ends.Source: A BATCHES · p. 10
  • Table 5 lists the pre-registered batches, one GH200 each.Source: Table 5 · p. 10

B PREDICTION LEDGER

Table 6 records the pre-registered predictions and their outcomes, including predictions that failed.

B PREDICTION LEDGER

C FITTED PARAMETERS

This section reports fitted parameters for ripening, throughput, and acceptance models, alongside fit visualizations. Its numbers are generated from batch summaries and per-position records by a shared script.

C FITTED PARAMETERS

  • All reported numbers are generated from batch summaries and raw per-position records by one script, rather than typed by hand.The account identifies Tables 8 and 9 and Figure 6 among the outputs.Source: C FITTED PARAMETERS · p. 11
  • Table 7 reports fp32 ripening fits for keys and values over loops 8–31, with ρ and c fitted and ∆∞ fitted to bf16 at fixed ρ and c.Source: Table 7 · p. 11
  • Table 8 fits T(B) = B/(a + bB) per WaveFront batched-emulator configuration and reports saturation throughput, half-saturation batch, and largest relative residual.Batch-8 fits cover a narrow batch range, ran on another node, and are descriptive only.Source: Table 8 · p. 11
  • Table 9 compares Eq. (7)'s predicted tokens per verification with simulation on the batch-3 simulation at γ = 4.It also reports mean conditional acceptance, its spread across four draft positions, and the number of blocks; rejection is approximated by 1 − αd ≈ kρd.Source: Table 9 · p. 12
  • The reported fits give k = 0.62 and ρa = 0.76 for the full cache, versus k = 0.42 and ρa = 0.66 for the deepest plane.Source: C FITTED PARAMETERS · p. 12
  • Figure 6 plots prefix-value ripening with fp32 and bf16 series and a geometric fit to fp32 over loops 8–31.The caption distinguishes filled fp32 points from open bf16 points.Source: Figure 6 · p. 12, C FITTED PARAMETERS · p. 12

Related papers

Loading 2610.34436v1…