PaperlayerView exact revision on arXiv

Source-linked AI summary

SAKURA: Serving Attention Key-Value State Under Rate Adaptation for Private Image Editing on Phones

Shuncheng Jie, Desen Sun, Ri Hong, Yuyang Tian, Hong Zhang, Sihang Liu

arXiv:2610.29733v1cs.SEcs.AIcs.NI

TL;DR

Reusable template caches can support private, on-device image editing, but their size and transfer demands challenge mobile deployment. SAKURA compresses and streams that state adaptively while overlapping transfer and reconstruction with local execution. It improves latency and storage efficiency over prior streaming and local alternatives while preserving image quality.

  • Problem

    Reusable template caches are large, and compression must preserve quality for future requests whose request-specific token importance is not known in advance.

    Source: §1 · p. 2, §3.1 · p. 4

  • Method

    SAKURA uses offline token-importance-aware cache compression, adaptively selects each shard’s representation online, and pipelines transfer and reconstruction with on-device execution.

    Source: §1 · p. 2

  • Results

    SAKURA is 1.55–1.82× faster and 1.54–2.37× more storage-efficient than prior KV-streaming systems, outperforms fully local execution by 3.27×, and cuts full-FP16 cache storage by 3.06×.

    Source: §8 · p. 13

  • Takeaways & Limitations

    Remote reuse of DiT template caches can make on-device image editing practical without replicating large cache collections on-device or offloading private request data.

    Source: §8 · p. 13

Abstract

from arXiv · show

Diffusion transformers are widely used in image-editing services. To preserve both efficiency and privacy, existing systems commonly rely on cloud-device collaboration. Keeping request-specific inference on the device avoids sending personal images and prompts to the cloud. These workloads often use common templates, such as in face swapping and virtual try-on, enabling KV state to be precomputed and reused across requests. However, this reusable state is large and scales with the number of templates, making it impractical to store all of them on mobile devices. Storing them in the cloud addresses this capacity issue, but cache transfer is placed onto the inference path and can even be slower than on-device computation, especially under fluctuating mobile networks. We introduce SAKURA, a cloud-edge architecture that delivers reusable attention state to mobile devices under a bandwidth budget. It first uses a fine-grained token-importance-aware compression method to preserve image quality under limited transfer budgets. At runtime, SAKURA adaptively selects the cache representation for each upcoming shard based on network bandwidth. To enable efficient execution, we develop a pipelined runtime that overlaps cache transfer and reconstruction with on-device execution. We evaluate SAKURA using three editing tasks and two diffusion models, using a Snapdragon 8 Elite Gen 5 phone. SAKURA reduces latency by 1.82x and 3.27x over CacheGen adapted to DiT and on-device generation, respectively. It also reduces cache storage by 3.06x over uncompressed cache, while preserving image quality.

1 Introduction

SAKURA addresses the tension between keeping private, request-specific editing on-device and storing large reusable template caches across a growing catalog. It compresses and streams reusable state adaptively, overlapping transfer and reconstruction with local execution.

1 Introduction

  • Keeping request-specific inference on-device avoids sending personal images, prompts, and derived intermediate states to a remote execution service.Provider-owned templates can be reused across users, unlike each request’s privacy-sensitive inputs.Source: §1 · pp. 1–2
  • A single reusable FP16 cache can reach 1.46 GiB, making storage across catalogs of hundreds of templates impractical on mobile devices.The cloud stores provider-owned reusable templates while request-specific inference remains on-device.Source: §1 · pp. 1–2
  • At 545 Mbps, full-cache transfer takes approximately 23 s versus about 13 s for local full-model execution, so remote reuse must reduce bytes and hide transfer behind computation.Only 2.3 s of local computation is available to hide transfer at an approximately 10% mask ratio, and even uniform 2-bit quantization requires at least 650 Mbps.Source: §1 · p. 2
  • SAKURA profiles stable token importance offline, adapts each upcoming shard’s representation to bandwidth and execution progress, and pipelines transfer and reconstruction with on-device diffusion computation.Its heterogeneous runtime overlaps CPU cache reception, GPU reconstruction, and NPU model execution.Source: §1 · p. 2
  • SAKURA reduces latency by 1.82× versus CacheGen adapted to DiT and 3.27× versus on-device generation, while cutting uncompressed-cache storage by 3.06× and preserving image quality.Evaluations cover three editing tasks and two diffusion models on a Snapdragon 8 Elite Gen 5 phone.Source: §1 · p. 2

2 Background and Motivation

Masked editing permits reuse of template-token KV state, reducing repeated computation but creating a large per-template cache. The paper places that reusable state in the cloud while keeping request-specific inference at the edge, shifting the bottleneck to transfer.

2.1 Cross-Request KV Reuse for Image Editing

  • Masked-region queries combine their K and V with cached K and V from unchanged template tokens, enabling sparse request-time execution.Token-wise operators process only tokens inside the editable mask, and attention is evaluated for masked-region queries.Source: §2.1 · p. 3
  • A 512 × 512 template in a lightweight 4B FLUX2 model requires 1.46 GiB of reusable KV state, which grows with the template library.The cache preserves K and V tensors for every denoising step and DiT block because unchanged-token hidden states vary across them.Source: §2.1 · p. 3

2.2 Cloud–Edge Placement of Reusable State

  • The edge performs request-specific inference while the cloud stores and supplies reusable per-template KV state on demand.Unlike conventional cloud–edge inference, computation does not migrate across the network; only reusable state is transferred.Source: §2.2 · p. 3, Figure 1 · p. 3
  • Full-cache transfer takes approximately 23 s at the 5G trace’s 545 Mb/s median throughput, versus 13 s for local full-model execution without reuse.Remote reuse therefore requires reducing transfer overhead and delivering cache state before execution needs it.Source: §2.2 · p. 3

3 Challenges and High-Level Ideas

SAKURA must fit cache transfer within limited, changing mobile bandwidth while preserving quality and aligning delivery with DiT’s execution. Its approach combines offline quality-efficient representation with online shard-level adaptation.

3 Challenges and High-Level Ideas

  • At a 10% mask ratio, the median aggregate transfer window is 2.3 s, while even 2-bit quantization requires 650 Mbps and leaves quality non-ideal.Compression must preserve the cache information that contributes most to output quality under a fixed transfer budget.Source: §3.1 · p. 3

3.1 Quality-Efficient Cache Representation

  • Quantizing the most-attended 10% of cached tokens accounts for 62% of induced image error, and token importance remains highly ranked across requests, supporting offline importance-guided compression.The offline signal can guide cache construction for future requests even though their masked-region queries differ.Source: §3.1 · p. 4

3.2 Bandwidth-Adaptive DiT Cache Transfer

  • For 2.26 s execution windows, required compression is 9.3× at the median and 53.8× in the slowest 10%, motivating transfer decisions that adapt across successive shards.Block-grained shards align transfer decisions with when DiT consumes each cache state, unlike token-level chunks spanning data needed at different times.Source: §3.2 · pp. 4–5

3.3 Summary of High-Level Ideas

  • SAKURA combines offline construction of a quality-efficient reusable cache with online transfer adaptation as network and execution conditions change.The two mechanisms address complementary needs: preserving quality under transfer limits and adjusting delivery during execution.Source: §3.3 · p. 5

4 Design

SAKURA combines offline, token-importance-aware cache construction with bandwidth-adaptive streaming and a pipelined mobile runtime. The design aims to select cache representations that fit transfer opportunities while overlapping reception and reconstruction with execution.

4.1 Overview

  • Offline, SAKURA builds a compact low-rank cache and profiles rank–precision tradeoffs for later request-time configuration selection.The overview separates per-template cache construction from request-time transfer and execution.Source: §4.1 · p. 5, Figure 4 · p. 5

4.2 Offline Token-Importance-Aware Cache Construction

  • Masked-region attention scores weight reusable-token rows during low-rank factorization, directing limited rank toward tokens that matter more to the final edit.The importance exponent q controls the weighting strength; q = 0 reduces to ordinary SVD, and the weights are used only during offline construction.Source: §4.2.1 · p. 6, §4.2.2 · p. 6, §4.2 · p. 6
  • SAKURA stores ordered low-rank components and independently quantized factors in a progressive layout, so rank and precision prefixes serve multiple transfer budgets without separate cache encodings.The rank components are ordered by weighted singular contribution.Source: §4.2.3 · p. 6, §4.2 · p. 6, Figure 5 · p. 6

4.3 Online Adaptive Cache Streaming

  • SAKURA operates at shard granularity, using sequential shard consumption so later transfer decisions can incorporate earlier measurements.The monitor tracks transfer rate and arrival timing; early arrival becomes carry-over time for the next shard's transfer window.Source: §4.3 · p. 7, §4.3.1 · p. 7, §4.3.2 · p. 7
  • At each shard, SAKURA uses recent transfer rate and execution progress to choose the highest-quality profiled rank–precision configuration that fits the available transfer window.The allocator selects from a nondominated set of configurations, and the server slices the required prefixes from the progressive cache pack.Source: §4.1 · p. 5, §4.3.3 · p. 7, §4.3 · p. 7

4.4 Online Pipelined Mobile Execution

  • SAKURA overlaps CPU cache reception, GPU KV reconstruction, and NPU execution across successive shards, hiding most communication and reconstruction behind computation.After startup, the device buffers the reconstructed cache for the next shard while the NPU executes the current one.Source: §4.4 · p. 7, Figure 9 · p. 8

5 Implementation

SAKURA targets Qualcomm mobile platforms and implements its heterogeneous runtime across CPU, GPU, and NPU stages. The implementation compiles diffusion-model shards for the mobile platform and reconstructs cache units on the GPU.

5 Implementation

  • The implementation focuses on Qualcomm platforms but is described as generally applicable to other mobile platforms.Source: §5 · p. 8
  • The server compiles diffusion models with the Qualcomm ExecuTorch backend and HTP v81 library, using W8A16 quantization and one compiled graph per transformer-block shard.The same block groups define the cache units.Source: §5 · p. 8
  • The C++ device runtime runs network transfer on the CPU, cache dequantization and reconstruction on the GPU, and compiled diffusion shards on the NPU.These stages operate concurrently on separate shards.Source: §5 · p. 8

6 Evaluation

The evaluation tests SAKURA across three editing tasks, two diffusion models, and multiple network, mask, and resolution settings. It reports lower latency and storage use alongside higher image quality, with benefits varying by bandwidth and workload conditions.

6.1 Setup

  • The setup evaluates face swap, virtual try-on, and MagicBrush with FLUX2, plus MagicBrush with SD3.5, on a Snapdragon 8 Elite Gen 5 phone.The main experiments replay 5G throughput traces and sample 100 starting points per workload.Source: §6.1 · pp. 8–9, Figure 10 · p. 9

6.2 Performance Evaluation

  • 44% lower average p90 latency than INTQ and 42% than CacheGen; median latency is 2.1–4.1× lower than local generation.SAKURA achieves the lowest median and tail latency on all four workloads.Source: Figure 12 · p. 9, §6.2 · p. 9

6.3 Quality Evaluation

  • 30–56% less data than INTQ and 23–42% less than CacheGen while maintaining higher image quality; user-study similarity scores range from 85.1% to 90.0%.SAKURA has the highest mean PSNR and SSIM across all four 5G workloads.Source: §6.3 · p. 10, Table 2 · p. 10, Table 3 · p. 10

6.4 Sensitivity Studies

  • As masks grow, latency rises but cache-transfer demands fall, narrowing the quality gap; smaller masks give SAKURA its largest quality advantage over INTQ and CacheGen.The smaller-mask setting combines a larger reusable cache with a shorter transfer window.Source: §6.4 · pp. 10–11, Figure 16 · p. 11
  • Under mixed 4G+5G traces, SAKURA reduces p90 latency by 55–76% over INTQ and 47–62% over CacheGen; under Wi-Fi, its advantage narrows.On the three FLUX workloads under Wi-Fi, p90 latency falls by up to 15% over INTQ and 38–44% over CacheGen.Source: Figure 17 · p. 11, §6.4 · p. 11
  • Across tested resolutions, SAKURA has the lowest latency and highest PSNR/SSIM; a 4× increase in pixel count raises latency and transferred cache by only 2.3–2.5×.The tested resolutions are 384×384, 512×512, and 768×768.Source: Figure 18 · p. 11, §6.4 · p. 11

6.5 Ablation Studies

  • q=0.5 gives the best overall quality at 5×, 15×, and 50× compression; shard adaptation reduces deadline violations, while adaptive bitwidth raises p90 PSNR by 2.5–3.9 dB and p90 SSIM by 0.12–0.17.Using fewer bits per value leaves room for more low-rank components under tight transfer budgets.Source: Figure 19 · p. 12, Figure 20 · p. 12, Figure 21 · p. 12

7 Related Work

Prior work studies diffusion-state reuse, KV-cache compression, and cache-transfer optimization, largely in settings distinct from SAKURA’s on-device diffusion inference. SAKURA combines offline template-cache factorization with token-importance weighting and transfer decisions tied to network bandwidth and denoising progress.

Diffusion cache reuse

  • Diffusion reuse methods exploit within-request denoising redundancy, while NIRVANA and FlashPS [25] reuse state across requests under server-side inference.Examples include U-Net feature or block reuse and DiT attention, residual, token, or timestep reuse [12] [51].Source: §7 · p. 12

KV cache compression and memory optimization

  • KV-cache work compresses through quantization or factorization, selects retained tokens, or selectively prefetches entries.Diffusion-model quantization work targets weights and activations, whereas SAKURA factorizes template KV matrices and weights approximation error using shared token importance.Source: §7 · p. 12

Cache transfer optimization

  • CacheGen, LMCache [13], KVServe, KV-Codec, and LLM.265 address cache transfer, while SAKURA streams segments to meet upcoming DiT execution deadlines.Its transfer decisions jointly depend on bandwidth and denoising execution progress.Source: §7 · p. 12

8 Conclusion

SAKURA keeps request-specific image editing on the device while storing reusable DiT template caches in the cloud and streaming them to the device. Its evaluation reports improvements in latency, storage efficiency, and output quality.

8 Conclusion

  • SAKURA is 1.55–1.82× faster and 1.54–2.37× more storage-efficient than prior KV-streaming systems, is 3.27× faster than fully local execution, and uses 3.06× less storage than full-FP16 cache.The reported conclusion attributes these outcomes to token-importance-aware offline compression, bandwidth-adaptive streaming, and a pipelined mobile runtime.Source: §8 · p. 13

Related papers

Loading 2610.29733v1…