PaperlayerView exact revision on arXiv

Source-linked AI summary

Shallow Prefill, Deep Decoding: Efficient Long-Context Inference via Layer-Asymmetric KV Visibility

Jungsuk Oh, Hyeseo Jeon, Hyunjune Ji, Kyongmin Kong, Jay-Yoon Lee

arXiv:2605.06105v1cs.AI

TL;DR

Long-context inference is expensive because Prefill time, Decode-time KV reads, and active KV memory grow with long prompts and model depth. SPEED makes non-anchor prefill-token KV visibility shallow while keeping Decode tokens full-depth, and SPEED-24+BoS reaches near-baseline benchmark quality while improving 128K-context efficiency. The paper concludes that a minimal BoS anchor can stabilize this regime, while cutoff selection remains task-dependent.

  • Problem

    Long-context inference couples Prefill TTFT, repeated Decode-time KV reads, and active KV memory to prompt length and model depth.

  • Method

    SPEED materializes non-anchor prefill-token KV states only through the first K layers, keeps Decode tokens full-depth, and optionally retains BoS as the sole full-depth prefill anchor.

  • Results

    SPEED-24+BoS reaches 51.2 average score versus 51.4 for Full-IT while improving TTFT by 33%, TPOT by 22%, and active KV memory by 25.0% at 128K context.

  • Takeaways & Limitations

    A single BoS anchor stabilizes shallow Prefill without restoring upper-layer access to the full prefill sequence, supporting a lower-cost KV-visibility regime.

  • Takeaways & Limitations

    The required prefill-visible depth is task-dependent, so no single cutoff is universally optimal.

Abstract

from arXiv · show

Long-context inference in decoder-only language models is costly because long prompts are processed during Prefill, cached at every layer, and repeatedly attended to during autoregressive Decode. We introduce \emph{Shallow Prefill, dEEp Decode} (SPEED), a phase-asymmetric KV-visibility policy that materializes non-anchor prompt-token KV states only in lower layers while keeping Decode-phase tokens full-depth. Unlike previous approaches that make upper-layer prompt KV states cheaper to store or construct, SPEED removes prefill tokens from the upper-layer Decode visibility set altogether. With a minimal BoS anchor, this simple change preserves broad benchmark quality while reducing long-context cost. In a controlled Llama-3.1-8B instruction-tuning study, SPEED using only 75\% of layers for prefill tokens reaches 51.2 average score on OLMES-style benchmarks, compared with 51.4 for the full-depth baseline, while improving TTFT by 33\%, TPOT by 22\%, and reducing active KV memory by 25.0\% at 128K context. Layer-wise diagnostics suggest that this cutoff retains the main prompt-selection and representation-stabilization regions of the full-depth model. These results show that long-context prompt tokens need not always persist as full-depth KV-cache objects when Decode-phase tokens remain full-depth.

1 Introduction

Long-context inference couples Prefill time, Decode memory traffic, and active KV memory to prompt length and model depth. SPEED addresses these costs by making prefill-token KV visibility shallow while retaining full-depth Decode tokens, with BoS anchoring preserving quality near the full-depth baseline.

  • Motivation: Long-context inference incurs coupled costs because Prefill dominates TTFT, Decode repeatedly reads cached KV states, and active KV memory scales with context length and depth.These costs arise in workloads including retrieval-augmented generation, document question answering, summarization, and code assistance.
  • SPEED: SPEED makes non-anchor prefill tokens visible only through the first K layers while Decode tokens traverse all L layers and retain full-depth KV states.This changes KV visibility rather than transformer weights, the language-modeling objective, or positional indices.
  • Evaluation: SPEED-24+BoS uses 75% of layers for prefill tokens and reaches 51.2 average score versus 51.4 for the full-depth instruction-tuned baseline.The controlled sweep uses shared data, formatting, optimizer, and evaluation protocol to isolate KV visibility.
  • Diagnostics: Layer-wise diagnostics guide K by tracking prefill-token selectivity, attention to generated Decode tokens, and representation stabilization in the full-depth model.These diagnostics are intended to reduce reliance on exhaustive cutoff sweeps.

2 Related Work

Prior methods reduce or restructure cached prefill computation through selection, compression, quantization, depth-wise sharing, or stage-aware optimization. SPEED instead changes which prefill KV states remain visible during upper-layer Decode, yielding a distinct Decode-time efficiency benefit.

  • KV-cache reduction and serving systems: Token selection, cache compression, quantization, sparse attention, routing, and serving systems reduce stored KV states, bytes per state, or KV traffic.These approaches exploit redundancy or reduce attention and serving costs without the specific visibility policy introduced by SPEED.
  • Depth-wise and phase-aware optimization: Depth-wise KV methods share, merge, condense, or allocate KV budgets across transformer layers, while stage-aware methods restructure later-layer Prefill caches.SwiftKV constructs later-layer caches from earlier representations, and POP removes deep-layer Prefill computation while retaining full-depth Decode.
  • SPEED’s distinction: Table 1 organizes the current Decode token’s visible KV sets by layer, with upper layers excluding prefill-side states in anchor-free SPEED and retaining only BoS under SPEED+BoS.Decode-phase tokens remain full-depth in all SPEED variants.
  • SPEED’s distinction: SPEED differs from related methods by removing non-anchor prefill tokens from upper-layer Decode visibility rather than preserving, sharing, or synthesizing their upper-layer KV states.At comparable K = 24, SPEED-24, POP-24, and SwiftKV-24 have similar TTFT reductions, but SPEED improves TPOT and has the lowest active KV memory.
  • SPEED’s distinction: Unlike early-exit or layer-skipping generation, SPEED keeps Decode tokens traversing all layers and producing full-depth KV states.Its asymmetry applies to prefill-token KV materialization, not to generation depth.

3 SPEED: Shallow Prefill, dEEp Decode

SPEED makes non-anchor prefill-token KV materialization shallow while keeping Decode-phase tokens full-depth, changing which cached states upper layers can access. A BoS anchor stabilizes generation while preserving the main memory and computation savings from truncating prefill depth.

  • KV-visibility policy: SPEED processes non-anchor prefill tokens only through layers 1 through K, while Decode-phase tokens traverse all L layers and produce full-depth KV states.Optional anchors such as BoS remain visible through all layers.
  • Anchoring: Anchor-free SPEED can destabilize generation when early Decode steps have very small upper-layer key sets, motivating SPEED+BoS as the main stabilized variant.The anchor-free setting remains useful diagnostically because it exposes the no-upper-prefill-KV regime.
  • Anchoring: SPEED+BoS retains one full-depth prefill-side KV state, leaving all other prefill tokens lower-layer-only.The BoS token is used as the minimal anchor set A = {s}.
  • Cost model: For long prompts with N ≫ a, T, SPEED reduces dominant prefill-side KV memory from O(LN) to O(KN), with the same layer-token reduction applying to Prefill computation and prefill-token Decode attention.Here N counts non-anchor prefill tokens, a counts full-depth anchors, and T counts cached Decode-phase tokens.
  • Cost model: The scaling expressions are proxies rather than a complete latency model because realized TTFT and TPOT also depend on kernels, bandwidth, cache layout, batching, and serving implementation.This limits direct translation from the cost model to deployment latency.
  • Training and implementation: SPEED-aware training preserves the loss, target tokens, and position indices while controlling KV-cache materialization and layer-wise attention visibility.Prompt positions follow the prefill-token rule, while assistant targets follow the Decode-token rule under teacher forcing.

4 Experimental Setup

The experiments use controlled Llama-3.1-8B instruction tuning, broad capability and long-context efficiency evaluation, lighter LoRA adaptation, and layer-wise diagnostics. Supporting checks examine transfer, length robustness, repetition loops, SelfOnly variants, and training throughput.

  • Main setup: All main experiments use the 32-layer Llama-3.1-8B architecture and compare prefill-visible cutoffs K ∈ {16, 20, 24, 28}, with K = 32 as full-depth attention.The primary comparison is a controlled instruction-tuning study from the base checkpoint.
  • Main setup: The full-depth baseline and SPEED variants share the same fine-tuning data, formatting, optimizer, schedule, batch construction, and update count to isolate KV-visibility effects.The instruction-tuning mixture contains 178,502 examples and each model is trained for two epochs.
  • Capability evaluation: Quality is evaluated on TULU-3-DEV using an OLMES-style protocol with an unweighted macro-average over 11 benchmark scores and five category aggregates.The categories are Knowledge, Reasoning, Code, Math, and Instruction.
  • Efficiency evaluation: Long-context efficiency measurements vary prompts from 1K to 128K tokens with a fixed 128-token continuation and report TTFT, TPOT, active KV-cache memory, and estimated FLOPs.POP-24 and SwiftKV-24 serve as efficiency-only stage-aware Prefill baselines.
  • LoRA adaptation: A lighter adaptation experiment compares SPEED+BoS LoRA with full-depth LoRA for document-grounded QA transfer and synthetic long-context retrieval.It starts from Llama-3.1-8B-Instruct and uses one epoch of HotpotQA pseudo-labeled adaptation.
  • Diagnostics: Layer-wise diagnostics measure Decode-token attention, conditional prompt entropy, and hidden-trajectory straightening to guide cutoff selection.The diagnostics are run on Full-IT with TULU-3-DEV prompts during greedy Decode.
  • Ablations and checks: The SelfOnly ablation tests whether upper-layer Decode-token attention among generated tokens can also be removed while retaining SPEED’s shallow-Prefill rule.Supporting appendices examine transfer, length robustness, repetition loops, additional SelfOnly variants, and training throughput.

5 Results

SPEED results identify BoS-anchored K = 24 as a strong quality–efficiency operating point, while diagnostics explain why moderate cutoffs work and aggressive visibility restrictions fail.

  • 5.1 BoS anchoring yields a strong quality–efficiency point: 51.2 average score for SPEED-24+BoS is 0.2 points below Full-IT while preserving substantial 128K-context efficiency gains.The same operating point improves TTFT by 33%, improves TPOT by 22%, and reduces active KV memory by 25.0%.
  • 5.1 BoS anchoring yields a strong quality–efficiency point: BoS anchoring recovers most of anchor-free SPEED's quality loss at K = 24 without changing its efficiency profile.Anchor-free SPEED falls from 51.4 to 49.1 average score, whereas SPEED-24+BoS reaches 51.2.
  • 5.1 BoS anchoring yields a strong quality–efficiency point: Task sensitivity varies: code remains relatively robust to shallow prefill visibility, whereas math and instruction are more sensitive to aggressive cutoffs.Knowledge and Reasoning also benefit substantially from BoS anchoring at moderate cutoffs.
  • 5.4 Upper-layer Decode-token attention remains necessary: SPEED improves Decode efficiency by removing non-anchor prefill tokens from upper-layer visibility while preserving full-depth Decode-token computation and interaction.This distinguishes SPEED from efficiency-only approaches that restructure prefill work but do not improve TPOT in the reported implementation.
  • 5.2 Off-the-shelf LoRA compatibility: Moderate SPEED cutoffs remain close to full-depth LoRA across HotpotQA, TriviaQA, and S-NIAH after one epoch of lightweight adaptation.OffShelf-FT-SPEED+BoS-24 reaches 59.5/73.7 on HotpotQA, 81.4/86.5 on TriviaQA, and 99.6 on S-NIAH.
  • 5.3 Layer-wise diagnostics guide cutoff selection: Layer-wise diagnostics place selective prompt access in middle layers and representation stabilization later, supporting cutoffs above those regions rather than raw attention peaks alone.Reasoning and Knowledge show prompt-mass peaks at L1, entropy minima around L13–L14, and straightening peaks around L17–L19; K = 24 covers these regions.
  • 5.3 Layer-wise diagnostics guide cutoff selection: Coding is the main diagnostic exception, with prompt-mass and entropy peaks at L3 but straightening at L19, consistent with relative robustness under aggressive truncation.The results indicate that required prefill-visible depth is task-dependent rather than universally optimal at one cutoff.
  • 5.4 Upper-layer Decode-token attention remains necessary: SelfOnly-24+BoS drops to 47.2 average score versus 51.2 for IT-SPEED-24+BoS, showing that upper-layer Decode-token attention remains necessary.The ablation removes upper-layer attention to other Decode-phase tokens while keeping the shallow-Prefill visibility rule.

6 Limitations

The reported limitations bound SPEED's conclusions to specific architectures, cutoff and anchor choices, matched-run evaluations, diagnostics, and serving configurations.

  • Scope and assumptions: SPEED's behavior depends on cutoff K, anchor design, adaptation procedure, sequence lengths, task distribution, and model architecture.Aggressive cutoffs can degrade quality, while anchor-free SPEED can destabilize generation; the main variant therefore uses BoS anchoring.
  • Evidence boundaries: Matched-run quality results are not statistical equivalence tests, so small aggregate gaps do not prove that shallow Prefill is lossless.Additional adaptation experiments do not exhaust all long-context task distributions or deployment settings.
  • Evidence boundaries: Layer-wise diagnostics guide cutoff selection but are not causal proofs of layer roles or reliable per-example cutoff predictors.They identify regions associated with prompt selection and representation stabilization in the evaluated full-depth model.
  • Deployment dependence: Realized TTFT and TPOT gains depend on serving-stack factors including kernels, batching, memory bandwidth, and KV-cache management.SPEED should be evaluated under the target deployment configuration, especially with continuous batching, prefix sharing, speculative decoding, or custom serving systems.

7 Conclusion

SPEED makes Prefill shallow while keeping Decode deep, using a minimal BoS anchor to preserve quality without restoring upper-layer access to the full prefill sequence. At 128K context, SPEED+BoS remains close to Full-IT quality while improving latency and reducing active KV memory.

  • Conclusion: SPEED processes and caches prefill tokens only through a lower-layer prefix while Decode-phase tokens traverse all layers and produce full-depth KV states.Upper-layer Decode attention can access the BoS anchor and generated Decode tokens, but not the full prefill sequence.
  • Conclusion: A minimal BoS anchor stabilizes shallow Prefill without restoring upper-layer access to the full prefill sequence.The BoS token is an existing reference, not a learned prompt summary or additional memory module.
  • Conclusion: 33% TTFT speedup, 22% TPOT speedup, and 25.0% active-KV-memory reduction are achieved at 128K context while remaining close to Full-IT quality.Measurements use Llama-3.1-8B, a fixed continuation of 128 generated tokens, and five repeats across prompt lengths up to 128K.
  • Conclusion: SPEED improves efficiency by deciding which prefill-token states persist as full-depth cached memory rather than only compressing or serving an already materialized cache.This design removes upper-layer prefill-token KV states and reduces repeated upper-layer prefill-token attention.

E Layer-wise Diagnostics

Layer-wise diagnostics compare prompt access, conditional prompt selectivity, and hidden-trajectory stabilization across task categories. They support an access-to-stabilization interpretation of the cutoff frontier, while remaining diagnostic rather than causal.

  • Layer-wise Diagnostics: Figure 3 measures Decode-token attention to user-prompt tokens, conditional prompt entropy, and all-token hidden-trajectory straightening across categories.Higher curves in the inverted entropy panel indicate lower entropy and more selective prompt access.
  • Layer-wise Diagnostics: Reasoning and Knowledge have early prompt-mass peaks but later conditional-entropy minima, indicating that broad prompt attention can precede selective access.Their prompt-mass peaks occur at L1, while entropy minima occur at L14 and L13, respectively.
  • Layer-wise Diagnostics: Coding reaches both its prompt peak and conditional-entropy minimum at L3, while its straightening peak occurs at L19.Its entropy–straightening gap is −16, compared with about −4 at the peak level for other categories.
  • Layer-wise Diagnostics: K = 24 retains the observed selection-to-stabilization interval better than K = 16 or K = 20 for most non-coding categories.The diagnostics motivate an access-to-stabilization cutoff rather than a single-peak rule.
  • Layer-wise Diagnostics: The diagnostics are category-level evidence, not causal proofs or per-example cutoff predictors.Straightening indicates geometrically stabilized hidden trajectories, not token independence, and upper-layer Decode-token attention remains important.

F Upper-layer Decode-token Attention Ablation

The SelfOnly ablation tests whether upper-layer Decode-token attention remains necessary when prefill visibility is shallow. Removing that attention degrades quality, especially at the lower K = 24 cutoff.

  • Upper-layer Decode-token Attention Ablation: SelfOnly keeps shallow Prefill visibility but removes upper-layer attention from Decode-phase tokens to other Decode-phase tokens.Upper-layer Decode tokens attend only to their current position, optionally with a BoS anchor.
  • Upper-layer Decode-token Attention Ablation: At K = 28, SelfOnly-28+BoS reaches 50.0 average score versus 51.3 for IT-SPEED-28+BoS.The degradation is moderate and spread across categories rather than concentrated primarily in Math.
  • Upper-layer Decode-token Attention Ablation: At K = 24, SelfOnly-24+BoS reaches 47.2 average score versus 51.2 for IT-SPEED-24+BoS.Drops occur across Knowledge, Reasoning, Code, Math, and Instruction.
  • Upper-layer Decode-token Attention Ablation: Upper-layer Decode-token attention becomes more important when fewer lower layers retain direct prefill-token visibility.SelfOnly is used as a diagnostic of Decode-token attention, not as evidence about the optimal anchor design.

G Repetition-loop Analysis

The repetition-loop analysis treats suffix repetition as a separate generation-stability failure mode. Anchor-free SPEED increases these loops, while adding a BoS anchor substantially reduces them without restoring full upper-layer prefill access.

  • Repetition-loop Analysis: The analysis targets suffix repetition loops as a generation-stability failure mode distinct from task accuracy.It specifically examines instability caused by removing all full-depth prefill-side anchors.
  • Repetition-loop Analysis: Loop detection searches the final 256 output tokens for repeated suffix units and flags spans of at least 12 tokens repeated at least three times.The heuristic permits repeated units of length 1 to 20 tokens, with limited trailing tokens and a partial final unit.
  • Repetition-loop Analysis: The loop-rate heuristic does not detect all semantic repetition or non-suffix repetition.It is designed for short exact suffix loops near the end of generation.
  • Repetition-loop Analysis: Anchor-free SPEED increases suffix repetition loops relative to Full-IT, especially on PopQA and GSM.Adding a BoS anchor substantially reduces this failure mode without restoring upper-layer KV states for the full prefill sequence.
  • Repetition-loop Analysis: SPEED+BoS is evaluated across prompt-length buckets on TriviaQA and S-NIAH, but bucket-level S-NIAH scores may reflect instance composition and evaluation variance.S-NIAH contexts extend to approximately 130K tokens.

I.1 Downstream transfer without task-adaptive fine-tuning

Moderate SPEED+BoS cutoffs remain compatible with downstream transfer and task adaptation across document QA, summarization, math, and code, while stronger post-hoc truncation can degrade substantially. Lightweight adaptation from an off-the-shelf instruction model also preserves usability, though transfer results are treated as compatibility evidence.

  • Adaptation requirement: PostHoc-SPEED shows large degradation at K = 24, whereas SPEED-aware BoS-anchored models remain competitive at moderate cutoffs.The contrast supports SPEED-aware adaptation when stronger prefill truncation is used.
  • Task-adaptive transfer: Moderate SPEED+BoS cutoffs remain close to TaskFT-Full on document QA and summarization.The reported comparison covers K = 24 and K = 28.
  • Task-adaptive transfer: TaskFT-SPEED-24+BoS is competitive with or slightly above the full-depth task-adapted baseline on math and code.These results use task-specific adaptation data and are not the primary basis for the quality–efficiency frontier.
  • Off-the-shelf compatibility: Moderate cutoffs remain usable after lightweight adaptation from an off-the-shelf instruction model.Because adaptation uses HotpotQA data, TriviaQA and S-NIAH serve as transfer checks rather than direct fitting measurements.

J Training Efficiency

SPEED improves downstream LoRA fine-tuning efficiency mainly through higher wall-clock throughput rather than substantial peak-memory reduction. The measured gains reflect reduced prefill-token layer computation while other training-memory components remain largely unchanged.

  • Throughput: Effective token throughput increases from 2213.8 to 2863.1 tokens/s/GPU under SPEED.The comparison is reported for downstream LoRA fine-tuning with one GPU.
  • Memory: Peak memory changes only modestly from 63.4 GiB to 61.6 GiB at K = 24.The reported benefit is therefore primarily wall-clock throughput rather than peak-memory reduction.
  • Mechanism and scope: SPEED reduces prefill-token layer computation while optimizer state, LoRA parameters, and much of the training memory footprint remain unchanged.The efficiency measurements use matched data order, effective batch size, optimizer, precision, hardware, checkpointing, and gradient accumulation settings.
Loading 2605.06105v1…