Enable KV cache offloading for a better speed-memory trade-off

@asomoza

We have KV Caching in a bunch of models now, including Flux2 and QwenImage 2.1.

While it helps make inference faster, it also increases memory consumption. For example, in QwenImage 2.1, with KV Caching, peak memory goes from 32.3G -> 36.0G.

To make this trade-off better for our users, I think we could start thinking about KV cache offloading.

The offloading could have different flavors:

Offload to and onload from RAM (sync, will compromise speed)
Offload to and onload from RAM but async (like group offload, should provide the best trade-off)
Same options but with secondary storage

Need to think about a nice abstraction for this. @asomoza do you think this could be helpful? @DN6 let me know your thoughts too.

@sayakpaul

Hi @sayakpaul, I'm really interested in this, is it possible for me to make a contribute to this? Would a focused first task on QwenImage 2.1 be useful?

I'd start by mapping the KV-cache lifecycle and establishing a reproducible memory/latency baseline, then share a short design proposal for opt-in synchronous CPU KV-cache offloading before implementation. I'd keep asynchronous prefetching and disk offloading out of the initial scope.

Validation would compare against KV caching enabled without offloading, covering numerical agreement, cache cleanup between requests, peak GPU memory, and latency.

Is this scope available, or is there an existing design or branch I should coordinate with? Thank you.

Sure, I would love to see the proposal.

@sayakpaul

Thanks @sayakpaul! I propose starting with opt-in synchronous RAM offloading for QwenImage 2.1, with a shared lifecycle that can later support async transfers and disk storage.

Trade-offs across the three flavors

Synchronous RAM: straightforward to implement and validate; removes persistent KV storage from the GPU, but keeps the full cache in host RAM and adds blocking transfer latency.
Asynchronous RAM: can overlap transfers with computation, but moves the same amount of data and requires pinned memory, stream synchronization, and bounded prefetch buffers. Benefits would need profiling.
Secondary storage, sync or async: extends capacity when host RAM is constrained, at the cost of additional I/O and RAM staging.

Async disk support must coordinate storage reads and GPU transfers; filesystem page-cache usage also matters.

Core abstraction

I suggest separating:

Model-specific layout: token selection, layer ordering, and attention assembly.
Persistent storage: ownership of independent K/V payloads, initially GPU-resident or CPU-backed.
Materialization/transfer management: temporary tensors on the consuming device, with future prefetching and bounded staging.

Qwen’s prefix KV is immutable after extraction, so each entry should be stored once and materialized repeatedly, without repeated writeback. The lifecycle would be store → materialize/use → release, followed by request-level cleanup.

Materialization must establish readiness on the consuming stream; temporary buffers cannot be reused before their consumers finish. Requests and CFG branches remain isolated. Future cleanup must finish outstanding work before freeing buffers or deleting files. Prefetch scheduling belongs at block traversal, where execution order is known.

First milestone: synchronous RAM only

Preserve existing behavior by default and validate unsupported option combinations early.
Copy extracted prefixes into independently owned CPU storage, preserving values and dtype.
Retrieve only the current layer’s prefix on the actual key tensor’s device, using blocking transfers.
Keep Qwen’s slicing and attention assembly in _qwenimage21_prepare_qkv.
Retain no temporary GPU copies in persistent cache fields; clear both branches on completion, interruption, or exceptions, before VAE decoding on normal completion.

Async scheduling and disk backends would remain follow-up implementations.

Validation would cover numerical agreement with GPU-resident caching, storage ownership, multiple steps, CFG isolation, and cleanup. Benchmarks would compare caching disabled, GPU-resident caching, and synchronous RAM offloading, reporting full-pipeline latency and GPU/host-memory peaks under identical, reproducible settings.

Hardware questions

Do you prefer SM80, SM90, or testing both?
Are there minimum or recommended GPU-memory and host-RAM capacities?
Is single-GPU validation sufficient for this milestone?
Is there an existing workload, attention backend, or weight-offloading setup you would like reproduced?

Does this abstraction and first-milestone scope match what you had in mind?

One data point for the RAM flavors: on Apple Silicon the CPU and the GPU share one pool of memory, so offloading from mps to RAM moves tensors within the same memory rather than freeing it.

I measured model offload, not KV cache offload, but the reason is the same. On a 48 GB M5 Pro, Qwen-Image-2.1 at 1024 × 1024 with enable_model_cpu_offload(device="mps") kept the process footprint at 18.5 GB, yet swap grew by 9.6 GB and my memory guard stopped the run in the VAE decode. Loading one model at a time and releasing it peaked at 19.0 GB with no swap growth. The trace: https://github.com/nefayran/stageload/tree/main/bench/results/2026-10-08-qwen-image-offload

So on MPS the RAM flavors would likely add transfers without lowering peak memory, and only the secondary-storage flavor would reduce it. Gating the RAM flavors on the execution device would avoid that cost there.

kept the process footprint at 18.5 GB, yet swap grew by 9.6 GB and my memory guard stopped the run in the VAE decode.

The VAE doesn't have any KV caching. So, why would it get affected?

Loading one model at a time and releasing it peaked at 19.0 GB with no swap growth.

That should have been done by enable_model_cpu_offload().

@zzf2024

@zzf2024 can you simplify your comments in your own words? Let's only focus on the most important details.

The VAE has no KV cache, agreed. That run measured model offload, not KV cache offload. The decode is only where my guard's limit (8 GB of swap growth) was crossed: swap had already grown by 3.5 GB during denoising and reached 9.6 GB in the decode. The footprint never went above 18.5 GB because it does not count pages that are in swap.

On the second point: enable_model_cpu_offload() moves each model to the CPU after its stage, so its weights stay in the process. On Apple Silicon the CPU and mps use the same physical memory, so that move frees nothing for the system. The staged run loads each model from disk when its stage starts and afterwards replaces its tensors with meta tensors, which gives the memory back. I didn't measure which pages went to swap; my guess is the CPU copies of the text encoder and the transformer.

On the second point: enable_model_cpu_offload() moves each model to the CPU after its stage, so its weights stay in the process. On Apple Silicon the CPU and mps use the same physical memory, so that move frees nothing for the system. The staged run loads each model from disk when its stage starts and afterwards replaces its tensors with meta tensors, which gives the memory back. I didn't measure which pages went to swap; my guess is the CPU copies of the text encoder and the transformer.

Do you have a minimal snippet to establish the behaviour?

@sayakpaul

@sayakpaul, I tested the RAM and disk variants on an A100 80GB with Qwen-Image 2.1: one condition image, 1024×1024, 40 steps, BF16. The evidence can be found at: https://github.com/zzf2024/diffusers-kv-offload-experiments/tree/experiments/qwen-image21-a100

Async RAM supports your expectation: 17.94 s versus 17.93 s for GPU caching, with overlapping timing ranges. Sync RAM took 24.57 s. Async RAM did increase process host-memory usage by about 4 GiB.
Disk performance depended heavily on page caching. Async disk took 18.35 s with warm reads served from page cache, but 105.06 s when storage reads were accounted for. The warm result should not be interpreted as disk throughput.
Cache lifetime matters for peak VRAM. Offloading saved about 2 GiB at the end of denoising. However, when every variant cleared its cache before VAE decoding, the overall peak stayed around 36.81 GiB. A separate control retaining caches through VAE showed a roughly 2 GiB lower peak with async RAM.
Tested offload outputs matched GPU caching exactly. These results cover one CUDA setup with model weights kept on the GPU; weight offloading and other models remain untested.
For the abstraction, I would suggest keeping cache storage and request cleanup separate from prefetch scheduling. Async RAM looks promising, but host-memory cost and cache lifetime need explicit handling.

Here is a minimal one, with a 4 GiB tensor standing in for a model's weights (torch 2.14.0, macOS 26.5.2, M5 Pro 48 GB):

import ctypes, os, time
import torch

# rusage_info_v2 from <libproc.h>; phys_footprint is the Memory column of Activity Monitor
FIELDS = ("user_time", "system_time", "pkg_idle_wkups", "interrupt_wkups", "pageins",
          "wired_size", "resident_size", "phys_footprint")
class RUsageInfoV2(ctypes.Structure):
    _fields_ = [("uuid", ctypes.c_uint8 * 16)] + [(f, ctypes.c_uint64) for f in FIELDS] \
        + [(f"rest{i}", ctypes.c_uint64) for i in range(10)]

libproc = ctypes.CDLL("/usr/lib/libproc.dylib")

def show(step):
    torch.mps.synchronize()
    info = RUsageInfoV2()
    libproc.proc_pid_rusage(os.getpid(), 2, ctypes.byref(info))
    print(f"{step:<14} footprint {info.phys_footprint / 2**30:4.1f} GiB, "
          f"mps driver {torch.mps.driver_allocated_memory() / 2**30:4.1f} GiB")

show("start")
w = torch.randn(2**30, device="mps")      # 4 GiB of float32, standing in for a model's weights
show("on mps")
w = w.to("cpu")                           # what model CPU offload does when a model's turn ends
torch.mps.empty_cache(); time.sleep(2)
show("moved to cpu")
float(w.sum())                            # read every page of the cpu copy
show("cpu copy read")

start          footprint  0.1 GiB, mps driver  0.0 GiB
on mps         footprint  4.2 GiB, mps driver  4.0 GiB
moved to cpu   footprint  0.1 GiB, mps driver  0.0 GiB
cpu copy read  footprint  4.2 GiB, mps driver  0.0 GiB

Right after the move, both the MPS driver count and the process footprint are back near zero, yet the copy is all there: once the CPU reads it, the footprint is 4.2 GiB again. So the moved weights still take the same memory, and a footprint sampled between stages, as in my trace, doesn't show them. That fits the offload run's 18.5 GB footprint next to 9.6 GB of swap growth. I got the same result in two runs. I couldn't get a clean system-wide number on this machine, because other jobs move its free memory by several GiB between runs.

@zzf2024

@zzf2024 could you maybe present a table containing the latency numbers as well as the peak GPU memory consumption from a full pipeline run?

@nefayran thanks for your inputs. Since this is a bit specific to shared memory now, let's tackle the unified architectures like MPS in a separate thread.

Sure, here is the table.

KV-cache configuration	Full-pipeline latency (s)	Peak allocated GPU memory (GiB)
GPU, upstream cache lifetime	18.33	38.83
GPU, pre-VAE cleanup	18.34	36.81
No cache	32.72	36.81
RAM, synchronous	24.86	36.81
RAM, asynchronous	18.35	36.81
File-backed, synchronous, warm cache	55.20	36.81
File-backed, asynchronous, warm cache	18.75	36.81
File-backed, synchronous, eviction hint	128.41	36.81
File-backed, asynchronous, eviction hint	103.71	36.81

Latency includes the full pipeline and VAE decoding, excluding model loading; values are medians of 5 runs after 1 warmup. Peak GPU memory is the maximum torch.cuda.max_memory_allocated() across those runs, resetting peak statistics before each request.

All offload modes clear the KV cache before decoding the image.

Navigation Menu

Repository navigation

Enable KV cache offloading for a better speed-memory trade-off #14833

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Enable KV cache offloading for a better speed-memory trade-off #14833

Description

Activity

zzf2024 commented on Oct 9, 2026

sayakpaul commented on Oct 9, 2026

zzf2024 commented on Oct 10, 2026

nefayran commented on Oct 11, 2026

sayakpaul commented on Oct 12, 2026

sayakpaul commented on Oct 12, 2026

nefayran commented on Oct 12, 2026

sayakpaul commented on Oct 12, 2026

zzf2024 commented on Oct 12, 2026

nefayran commented on Oct 12, 2026

sayakpaul commented on Oct 13, 2026

zzf2024 commented on Oct 13, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions