PaperlayerView exact revision on arXiv

Source-linked AI summary

Affix Cache for Diffusion Large Language Models

Kaihua Liang, An Zhong, Xin Tan, Zafar Ayyub Qazi, Hong Xu, Jian Weng, Marco Canini

arXiv:2608.26140v1cs.CLcs.LG

TL;DR

DLLM bidirectional attention makes shared-cache reuse stale, while full KV-cache recomputation remains expensive. ACache selectively recomputes request-sensitive Anchor Tokens in shared affixes, recovering most accuracy loss with around 20% recomputation and reducing recompute latency by up to 55.7%.

  • Problem

    DLLM bidirectional attention couples shared-token KV states to evolving responses, making direct cache reuse stale and full-sequence recomputation costly.

  • Method

    ACache measures masked-token cross-attention to identify request-specific Anchor Tokens, selectively recomputing them while reusing the remaining shared-affix cache.

  • Results

    Recomputing around 20% of affix tokens recovers most accuracy loss, while reducing recompute latency by up to 55.7% and improving throughput by up to 1.68×.

  • Takeaways & Limitations

    ACache is a practical step toward broader DLLM cache reuse across prefix, infix, and suffix settings.

  • Takeaways & Limitations

    The end-to-end benefits of non-prefix affix reuse remain unvalidated because the inference prototype supports only shared prefixes.

Abstract

from arXiv · show

Diffusion Large Language Models (DLLMs) enable non-autoregressive decoding and bidirectional context modeling, but efficient inference remains challenging. Unlike autoregressive systems, whose key-value (KV) cache can be reused for shared prefixes, DLLMs couple the KV states of shared context tokens with evolving generated tokens through bidirectional attention, making naive cache reuse stale while full recomputation is expensive. We present ACache, an affix-oriented cache reuse mechanism for shared text spans in DLLMs beyond prefixes. ACache identifies a small request-specific subset of critical affix tokens, called Anchor Tokens, by measuring their influence on masked generation tokens, and selectively recomputes the KV states of only these tokens while reusing the remaining affix cache. Built on Fast-dLLM, ACache recovers the accuracy loss caused by direct affix-cache reuse across different settings when recomputing around 20% of affix tokens. We also build a shared-prefix prototype on top of the Nano-vLLM engine, showing that ACache reduces recompute latency by up to 55.7% and improves end-to-end throughput by up to 1.68$\times$.

1 Introduction

ACache enables fine-grained reuse of shared affixes in diffusion large language models by selectively recomputing request-specific Anchor Tokens while reusing the remaining cache. Implemented on Fast-dLLM and Nano-vLLM, it recovers accuracy with around 20% recomputation and improves latency and throughput.

  • Motivation: DLLMs’ bidirectional attention couples shared-context KV states with evolving generated tokens, preventing the naive cache reuse that benefits autoregressive shared prefixes.DLLMs instead perform iterative denoising with non-autoregressive token prediction and bidirectional context modeling.
  • ACache: ACache targets shared contiguous affixes beyond prefixes, enabling reuse when spans occur at the beginning, middle, or end of context.The mechanism is designed specifically for DLLMs and extends cache reuse beyond conventional prefix-centric inference.
  • ACache: Anchor Tokens identify a small request-specific subset of affix tokens whose KV states are selectively recomputed while the remaining affix cache is reused.This selective recomputation replaces recomputing all affix-token states.
  • Evaluation: Around 20% of affix-token recomputation recovers the accuracy loss caused by direct affix-cache reuse across multiple Fast-dLLM benchmarks and settings.ACache replaces Fast-dLLM’s periodic full-cache recomputation with selective Anchor Token recomputation.
  • Evaluation: 55.7% lower recompute latency and 1.68× higher end-to-end throughput are achieved by ACache in a Nano-vLLM shared-prefix prototype under a realistic inference stack.The prototype demonstrates these improvements with a shared-prefix system built on Nano-vLLM.

2 Background

Diffusion language models generate text through iterative denoising with bidirectional context, enabling parallel, arbitrary-order decoding and infilling. However, global attention makes standard autoregressive KV-cache reuse stale, leaving full-sequence recomputation as a major inference bottleneck and motivating approximate caching.

  • Diffusion language models: Diffusion language models iteratively denoise partially masked sequences while using bidirectional context instead of a left prefix.This formulation includes continuous-embedding models such as Diffusion-LM and discrete-state models such as D3PM, with later objectives including SEDD and MDLM.
  • Decoding: DLLMs update masked responses through parallel prediction and partial commitment, enabling arbitrary-order decoding and infilling.Different decoding schedules iteratively determine which tokens are predicted and committed.
  • KV-cache challenge: Global bidirectional attention couples prompt-token representations to evolving masked responses, so DLLM KV caches are not append-only and quickly become stale.This prevents direct adoption of the standard autoregressive KV-cache mechanism used in causal decoding.
  • KV-cache challenge: Full-sequence KV-cache recomputation at every decoding step remains a practical DLLM inference bottleneck.The redundancy across iterative denoising steps makes crossstep cache reuse a natural efficiency target.
  • Related work: Approximate caching methods reduce repeated computation by selectively recomputing features or delaying decoded-token updates across decoding steps.dLLM-Cache uses asymmetric prompt and response update intervals with similarity-based selection, while dKV-Cache delays recomputation within fixed-size blocks.

3 ACache Design

ACache accelerates DLLM inference by reusing shared affix KV states while selectively recomputing request-specific Anchor Tokens and evolving positions. It identifies stable, high-impact affix tokens once per request using masked-to-affix attention and reuses that selection throughout decoding.

  • Inference loop: ACache operates within DLLM’s iterative predict-then-commit decoding loop, where masked positions are predicted in parallel and only selected confident predictions are committed.Structured block-wise schedules can define which positions are committed, but ACache builds on this common inference pattern.
  • Affix cache reuse: At recomputation points, ACache avoids recomputing every shared affix token by targeting a request-dependent subset of critical positions.The shared affix occupies a contiguous non-masked span, whose full KV cache would otherwise be recomputed despite identical text across requests.
  • Anchor selection: ACache runs a one-shot probe before decoding, using the precomputed affix cache as fixed past KV states while evaluating the current request’s non-affix positions.The probe keeps affix states fixed and derives masked-to-affix attention signals for the request.
  • Anchor selection: It aggregates masked-to-affix attention across layers, heads, and initial masked positions, then selects the top-K affix tokens as Anchors.The importance mass concentrates on a small subset of affix tokens, whose high-ranked positions remain stable as decoding proceeds.
  • Anchor selection: O(NHmL) aggregation, O(NHPL) probe computation, and O(L log K) extraction define the one-time preprocessing overhead for Anchor selection.Here m is the number of masked positions, P the number of non-affix probe positions, and L the affix length.
  • Cache construction: ACache combines shared cache entries for non-anchor affix tokens with recomputed Anchor and request-specific states, then repeats this selective recomputation at later checkpoints.Non-anchor affix tokens continue reusing the shared cache throughout the request.

4 System Prototype

The Nano-vLLM prototype realizes ACache for shared prefixes using Fast-dLLM, with shared and request-private KV regions plus mappings that preserve a unified logical context. Admission performs Anchor selection and records layout decisions, while runtime recomputes selected states and prioritizes requests needing recomputation.

  • Prototype scope: The prototype builds ACache on Nano-vLLM with Fast-dLLM as its base caching mechanism and scopes support to shared prefixes.Nano-vLLM’s KV cache design motivates the shared-prefix scope; affix support remains future work.
  • KV organization: The KV layout stores each common prefix once in shared blocks, while private regions hold Anchor Tokens and request-specific prompt or response positions.This separates reusable prefix states from request-dependent states.
  • KV mappings: Two mappings preserve ACache semantics: a write-side recomputation map targets refreshed states, while a read-slot map maps every logical position to its physical KV slot.Non-Anchor prefix positions read shared slots; Anchor and request-specific positions read private slots.
  • Admission and execution: At admission, the runtime resolves compatible shared prefixes, selects Anchors using the masked-to-prefix importance criterion, and records recomputation and read mappings before online inference.The remaining prefix positions stay mapped to shared blocks, moving layout decisions out of the online path.
  • Admission and execution: During decoding, requests needing KV recomputation are prioritized, their selected states and current blocks are prepared, and decoding resumes over the resulting active set.This handles requests admitted after departures or advanced to new decoding blocks.

5 Evaluation

Evaluation shows that ACache preserves DLLM accuracy across affix reuse settings while reducing recomputation cost. Its gains hold across models and affix positions, and retaining non-anchor cache entries remains important for accuracy.

  • Experimental Setup: Experiments evaluate LLaDA-8B-Instruct and Dreamv0-Instruct-7B using Transformers for quality tests and Nano-vLLM for efficiency tests on NVIDIA A100-SXM4-40GB GPUs.LLaDA is trained from scratch, while Dream is initialized from an autoregressive LLM.
  • Accuracy Evaluation: ACache restores accuracy across models, tasks, and prefix, infix, and suffix reuse settings, with an Anchor ratio of 0.2 recovering most lost accuracy.Direct reuse is especially inaccurate for infix and suffix affixes, while infix reuse is harder for retrieval-style BABILong prompts.
  • Efficiency Evaluation: 15.3% to 55.7%: ACache reduces LLaDA recompute latency across all tested shared-prefix and few-shot settings.System measurements use an Anchor ratio of 0.2; Dream results are reported separately in the appendix.
  • Efficiency Evaluation: ACache improves end-to-end throughput under batched inference when measured across shared-prefix lengths, batch sizes, and few-shot settings.Throughput counts the full generation length and is reported relative to the Fast-dLLM-based Nano-vLLM baseline.
  • Cache Ablation: KeepNA outperforms DropNA on LLaDA 1-shot, showing that non-anchor affix cache entries still provide reusable context after Anchor recomputation.DropNA discards all non-anchor cache entries, whereas KeepNA retains them.

6 Related Work

ACache is positioned as complementary to existing DLLM acceleration methods, including cache-reuse techniques and other approaches that reduce decoding cost through selective computation or attention-derived signals.

  • DLLM acceleration: ACache is orthogonal to most DLLM acceleration techniques.
  • DLLM acceleration: dKV-Cache, dLLM-Cache, Fast-dLLM, and FlashDLM reduce intra-request cost through cache reuse and recomputation.
  • DLLM acceleration: Other optimizations reduce decoding cost through compression, eviction, adaptive recomputation, scheduling, pruning, token-selective compute, or attention-derived signals.

7 Conclusion · A Prompt Construction

ACache selectively recomputes request-sensitive Anchor Tokens while retaining the remaining shared affix cache, recovering accuracy with limited recomputation and improving shared-prefix inference efficiency. The evaluated prompt constructions serialize few-shot affixes and task queries differently across prefix, infix, and suffix layouts, while the current runtime remains limited to shared prefixes.

  • 7 Conclusion: ACache recomputes request-sensitive Anchor Tokens while retaining the remaining shared cache to preserve context consistency for DLLMs.It treats shared text spans as cross-request cache objects.
  • 7 Conclusion: 20%: recomputing around 20% of affix tokens recovers most accuracy lost by direct affix-cache reuse across prefix, infix, and suffix settings.
  • 7 Conclusion: 55.7%: ACache reduces recompute latency by up to 55.7% in the shared-prefix Nano-vLLM prototype.
  • 7 Conclusion: 1.68×: ACache improves end-to-end throughput by up to 1.68× in the shared-prefix Nano-vLLM prototype.
  • Limitations: The prototype supports only shared prefixes, while infix and suffix reuse require position-aware span registration and physical-to-logical indirection.End-to-end benefits for non-prefix affixes remain unvalidated in a full runtime.
  • Limitations: Shared spans must currently be declared before inference, unlike AR prefix caching, which can discover common left prefixes online.Future systems could track recurring spans and precompute shared KV caches when reuse frequency justifies it.
  • A Prompt Construction: Prompt construction serializes few-shot examples, task queries, final-answer prompts, and masked generation spans according to prefix, infix, or suffix placement.The notation defines Ex as examples, Qry as the query, Ans as the final prompt, and Mask as the masked generation span.
  • A Prompt Construction: Task-specific templates cover general question answering, MBPP code generation, and BABILong location answering, with normalized few-shot answers and infix prompts redirecting the model to the initial task.Few-shot affixes are serialized as user queries followed by normalized assistant answers in prefix and suffix modes, while infix mode uses a plaintext block.

B Additional Evaluation

The additional evaluation uses lm-eval with model-specific wrappers and evaluates GSM8K, MBPP, and BABILong under specified affix serialization and task settings.

  • Quality evaluations use lm-eval with model-specific wrappers for LLaDA and Dream.The setup applies the Language Model Evaluation Harness to both models.
  • GSM8K and MBPP use standard task data with the affix serialization described in Appendix A.
  • BABILong evaluation uses a custom task YAML for BABILong-0k/qa1, the 0k/qa1 split of RMT-team/babilong-1k-samples, and exact-match scoring.

B.1 Dream System Efficiency

ACache reduces Dream recompute latency in every tested setting, with the largest gains occurring for the longest shared affix. At 4-shot, reductions reach 56.3% on GSM8K and 52.5% on MBPP across batch sizes.

  • Dream System Efficiency: 15.3%–56.3%: ACache reduces Dream recompute latency in every tested setting.These results complement the LLaDA recomputation results in Table 1.
  • Dream System Efficiency: 50.8%–56.3%: At 4-shot, ACache reduces GSM8K recompute latency across batch sizes.The largest gains occur with the longest shared affix.
  • Dream System Efficiency: 49.5%–52.5%: At 4-shot, ACache reduces MBPP recompute latency across batch sizes.These reductions also occur with the longest shared affix.

B.2 Peak KV Cache Usage · B.3 Anchor Selection Overhead · LLaDA

ACache lowers peak KV-cache usage in larger shared-prefix batches and selects Anchor Tokens with a mostly one-time attention-probe cost. On LLaDA, its selector costs 48.1–74.0 ms per request, while Figure 7 compares its accuracy with a CacheBlend-style variant across Anchor ratios.

  • B.2 Peak KV Cache Usage: 43.3%: ACache reduces LLaDA GSM8K peak KV usage from 11.12 GB to 6.31 GB at batch size 16 and 4-shot.The reduction grows with batching and shared-prefix length.
  • B.2 Peak KV Cache Usage: At batch size 1, ACache can use the same or slightly more KV memory because shared and request-private blocks are allocated separately.The fixed shared-cache cost can outweigh savings at block granularity.
  • B.2 Peak KV Cache Usage: Within-model Baseline-to-ACache reductions are the meaningful comparison because Dream and LLaDA use different KV-cache architectures.Dream has smaller effective blocks from grouped-query KV, whereas LLaDA stores full multi-head KV states.
  • B.3 Anchor Selection Overhead: Anchor selection is profiled as a one-time cost split across prompt preparation, the masked-to-affix attention probe, and topk selection.The profiling separates the selector into these three components.
  • B.3 Anchor Selection Overhead: 48.1–74.0 ms: Anchor selection costs per request across LLaDA and Dream on GSM8K/MBPP.The attention probe accounts for nearly all selector time.
  • B.3 Anchor Selection Overhead: 98.3–99.0%: The attention probe’s selector-time share exceeds preparation costs of at most 0.65 ms and top-k costs of at most 0.41 ms.The probe itself takes 47.5–73.1 ms.

B.4 Dream KeepNA/DropNA … C.1 Kernel Integration

ACache’s non-anchor affix cache retains useful context, while its masked-to-affix Anchor Token selector outperforms a CacheBlend-style alternative. The prototype integrates request-specific cache indirection into a customized Triton attention kernel for recomputation and decoding.

  • B.4 Dream KeepNA/DropNA: At Anchor ratio 0.2, Dream 1-shot KeepNA averages 62.21% accuracy versus DropNA’s 55.64%, a 6.57-point gap.The gap disappears at ratio 1.0, when no non-anchor tokens remain.
  • B.4 Dream KeepNA/DropNA: The Dream results confirm that non-anchor affix-cache tokens carry useful context for generation.This conclusion follows from KeepNA outperforming DropNA when only a subset of affix tokens is retained.
  • B.5 CacheBlend-Style Anchor Selection: CacheBlend is a natural analogue because it reuses KV caches beyond prefixes for autoregressive retrieval-augmented-generation serving.ACache uses this comparison to evaluate an alternative affix-level anchor-selection strategy.
  • B.5 CacheBlend-Style Anchor Selection: At Anchor ratio 0.2 across 1-shot settings, ACache averages 52.61% accuracy versus 41.57% for the CacheBlend-style variant, an 11.04-point gap.The variant preserves ACache’s generation path and cache layout but replaces masked-to-affix attention selection with high-KV-deviation selection.
  • B.5 CacheBlend-Style Anchor Selection: At Anchor ratio 0.3, ACache and the CacheBlend-style variant still differ by 11.65 points, becoming identical only at ratio 1.0.At ratio 1.0, both methods fully recompute the affix.
  • C.1 Kernel Integration: The prototype implements a customized Triton attention kernel supporting paged KV caches, ragged query/KV lengths, and a flattened read_slot_map.The read_slot_map supplies physical slots for each KV tile before keys and values are gathered.
  • C.1 Kernel Integration: Logical KV positions can reference shared prefix or request-private blocks, while query tiles remain standard dense tiles.This enables per-request read-slot mapping within the paged-cache attention path.
  • C.1 Kernel Integration: ACache preserves a store-then-attend interface, writing new KV states normally before reading the effective KV view through the read-slot kernel.Recomputation uses Anchor Tokens and request-specific positions, whereas decoding uses the current block; storage indirection remains in the attention read path.
Loading 2608.26140v1…