Feature Request: SWA KV cache compression for Gemma 4 (and other SWA models)

[ormandj]

Feature Request: SWA KV cache compression for Gemma 4 (and other SWA models)

Prerequisites

• Running the latest code (ik_llama.cpp commit 13d7178d, 2026-04-09).
• Followed README.md.
• Searched existing issues — the closest references are Bug: Gemma 4 31B in Graph Mode has insane VRAM usage #1594 (closed) where this feature was suggested and Feature Request: Gemma4 model support #1572 (tracking). No standalone tracking issue exists.
• Reviewed Discussions — no existing SWA compression discussion.

Feature Description

Implement sliding-window-attention (SWA) KV cache compression, so that SWA layers allocate only window-sized KV storage rather than full-context KV. This was suggested by @ikawrakow in #1594:

"You can add a feature request to implement KV compression for SWA layers, along with checkpointing, which becomes necessary in that case."

Filing this as a standalone tracking issue since none exists yet.

Motivation

Gemma 4 31B architecture (from the official config.json)

• num_hidden_layers: 60
• layer_types: 5× sliding_attention followed by 1× full_attention, repeated 10 times → 50 SWA layers, 10 global layers
• sliding_window: 1024 tokens
• max_position_embeddings: 262144

KV cache size with vs without SWA compression

Per-layer KV cost scales linearly with the number of tokens stored. Without SWA compression, ik_llama.cpp allocates the full context length of KV for every layer. With SWA compression, only the 1024-token window is allocated for SWA layers; global layers still get the full context.

Ratio of SWA-compressed KV to uncompressed, as a function of context size N:

compressed_kv    50 * 1024 + 10 * N
-------------- = --------------------
uncompressed_kv       60 * N

context N	compressed / uncompressed	effective KV savings
2048	58%	42%
8192	27%	73%
32768	19%	81%
65536	18%	82%
262144	17%	83%

So at any context past ~8K, SWA compression eliminates ~70-83% of the KV cache memory.

Observed impact on our hardware

3× RTX 3090 (72 GB total VRAM), ik_llama.cpp commit 13d7178d, tuned config with -sm graph -ctk q6_0 -ctv q6_0 -khad -vhad:

model	backend	context tested
Gemma 4 31B dense	llama.cpp (master, q8_0 KV)	262144 (model native max, runs comfortably)
Gemma 4 31B dense	ik_llama.cpp (q6_0 KV + -khad -vhad)	~65536 (higher values OOM on load)

llama.cpp runs at the model's full native context (262144) on this hardware without strain. ik_llama.cpp on the same hardware OOMs during load when I push the context past ~65K, despite using a smaller KV quant and Hadamard transforms. 262144 is the model's trained max_position_embeddings, not a tested ceiling for llama.cpp — I haven't probed beyond it because that's outside the model's valid range. So the observed difference is "runs at model native max" vs "caps around 1/4 of native max". The memory savings table above is consistent with a ~6× ratio between the two KV allocations; I can't prove KV is the only contributor to the observed cap difference, but it is a plausible primary explanation given the architecture.

Throughput numbers

These are included for context only; the request is motivated by the context cap, not by the throughput gap. Tuned -b/-ub via ik-llama-bench and llama-bench on the same hardware. pp512 / tg128 with fa on.

Gemma 4 31B dense (Unsloth UD-Q6_K_XL):

backend	config	PP tok/s	TG tok/s
llama.cpp	-sm layer -b 512 -ub 128 -ctk q8_0 -ctv q8_0	1020	25.0
ik_llama.cpp	-sm graph -muge -b 2048 -ub 512 -ctk q6_0 -ctv q6_0 -khad -vhad	1474	36.6

Gemma 4 26B-A4B MoE (Unsloth UD-Q8_K_XL):

backend	config	PP tok/s	TG tok/s
llama.cpp	-sm layer -b 2048 -ub 512 -ctk q8_0 -ctv q8_0	3154	97.8
ik_llama.cpp	-sm graph -muge -b 2048 -ub 1024 -ctk q6_0 -ctv q6_0 -khad -vhad	3054	73.3

The primary motivation for this request is the context cap, not TG throughput — I can't isolate which of several differences (SWA compression, attention kernel, MoE dispatch, etc.) accounts for the 26B TG gap, so I'm not claiming SWA compression as its cause. The context cap impact, however, is directly explained by the KV size table above.

Quality cost

SWA compression is an architectural optimization, not a quantization. SWA layers attend only to the window by definition, so KV outside the window is not used by the model's math. Discarding it is lossless.

The -khad/-vhad + q6_0 KV combination is also near-lossless per @ikawrakow's own data in #1599:

KV config	Gemma 4 31B PPL
f16	5.2306
q8_0	5.2311
q6_0 + -vhad	5.2323

So this feature request is for a memory optimization with no expected quality cost on its own.

Possible Implementation

Where the gap is in ik_llama.cpp today

ik_llama.cpp already has most of the SWA machinery in place:

• llama_hparams::swa_layers[LLAMA_MAX_LAYERS] — per-layer SWA flag array (src/llama-hparams.h:144)
• n_swa, n_swa_pattern — window size and pattern period (src/llama-hparams.h:30-31)
• n_embd_head_k_swa, n_embd_head_v_swa, n_rot_swa, rope_freq_base_train_swa — per-SWA hparams
• build_inp_KQ_mask_swa() / inp_KQ_mask_swa — SWA-aware attention mask builder (src/llama-build-context.h:130, src/llama-context.h:203)

So the attention math already masks to the window correctly. The missing piece is purely in KV cache sizing.

In llama_kv_cache_init() (src/llama.cpp:854 onwards), the per-layer K/V tensors are allocated as:

// src/llama.cpp:956
k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv * kv_size);
...
// src/llama.cpp:958
int64_t v_ne = int64_t(n_embd_v_row) * kv_size;
// src/llama.cpp:969
v = ggml_new_tensor_1d(ctx, this_type_v, v_ne);

Every layer uses the global kv_size, regardless of whether hparams.swa_layers[i] is set. This is the single point where SWA layers pay the full-context memory cost they don't need.

How llama.cpp solves it

llama.cpp implements this via composition in src/llama-kv-cache-iswa.{h,cpp}. The core design is a wrapper class that holds two existing llama_kv_cache instances, filtered per-layer by SWA-ness:

// llama-kv-cache-iswa.cpp — ctor, edited for brevity
const layer_filter_cb filter_base = [&](int32_t il) { return !hparams.is_swa(il); };
const layer_filter_cb filter_swa  = [&](int32_t il) { return  hparams.is_swa(il); };

const uint32_t size_base = kv_size;
uint32_t size_swa = GGML_PAD(std::min(size_base,
                                      hparams.n_swa * (unified ? n_seq_max : 1) + n_ubatch),
                             256);
if (swa_full) { size_swa = size_base; }

kv_base = std::make_unique<llama_kv_cache>(
    model, type_k, type_v, v_trans, offload, unified,
    size_base, n_seq_max, n_pad,
    0, LLAMA_SWA_TYPE_NONE, filter_base, reuse);

kv_swa = std::make_unique<llama_kv_cache>(
    model, type_k, type_v, v_trans, offload, unified,
    size_swa, n_seq_max, n_pad,
    hparams.n_swa, hparams.swa_type, filter_swa, reuse);

All seq ops (seq_rm, seq_cp, seq_keep, seq_add, seq_div) just delegate to both sub-caches in order (llama-kv-cache-iswa.cpp:80-107). The per-layer routing is done entirely through filter_base / filter_swa, which the existing llama_kv_cache constructor already accepts.

Key supporting pieces in llama.cpp:

• llama_swa_type enum (src/llama-hparams.h:19-24): NONE / STANDARD / CHUNKED / SYMMETRIC.
• llama_hparams::is_swa(il) — per-layer predicate populated from the model's layer_types.
• llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1) (src/llama-hparams.h:316-350) — the window-mask helper used by the attention mask builder.
• Per-arch model files use the ISWA variant when swa_type != NONE: src/models/gemma2-iswa.cpp, gemma3n-iswa.cpp, gemma4-iswa.cpp, cohere2-iswa.cpp, openai-moe-iswa.cpp, mimo2-iswa.cpp, step35-iswa.cpp, llama-iswa.cpp.
• The selection point is in llama-model.cpp: if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) res = new llama_kv_cache_iswa(...) else the plain llama_kv_cache (src/llama-model.cpp:8702-8735).
• A --swa-full cparam exists as an escape hatch: when set, the SWA sub-cache is sized to size_base, giving the legacy "uncompressed" behavior for debugging / correctness comparisons.

Mapping onto ik_llama.cpp

Two broad approaches for the port, listed from biggest to smallest architectural change:

Approach A: Compose two sub-caches (llama.cpp pattern). Introduce an ik_llama_kv_cache_iswa type that owns two llama_kv_cache instances — kv_base (sized kv_size, filtered to non-SWA layers) and kv_swa (sized n_swa + pad, filtered to SWA layers). Seq ops delegate to both. This is the cleanest mapping but requires each sub-cache to independently track its own cells / head / used, which currently live as single fields in struct llama_kv_cache (src/llama-context.h:35-85).

Approach B: Per-layer tensor sizing in the existing cache. Keep the single llama_kv_cache but allow k_l[i] / v_l[i] to have different cell counts depending on hparams.swa_layers[i]. At the SWA layer K/V allocation site above (src/llama.cpp:956-969), replace kv_size with hparams.swa_layers[i] ? swa_kv_size : kv_size, where swa_kv_size = n_swa + pad. This avoids the architectural refactor but needs a way to map logical positions into the SWA layer's smaller slot ring. That's the "context checkpointing" concern from #1594: when a position rolls out of the SWA window, the K/V tensor slot needs to be reusable even though the global cells array still references that position for non-SWA layers.

Approach A mirrors what llama.cpp did and sidesteps the slot-reuse problem by giving SWA layers their own cells. Approach B is a smaller diff but needs new logic in find_slot / seq_rm / position-shift paths to keep per-layer slot accounting consistent.

Either way, the attention mask path (build_inp_KQ_mask_swa) already exists and does not need to change. Only the allocator and the slot-tracking paths are affected.

Scope: Gemma 2, Gemma 3, Gemma 4 all use SWA; other architectures with ISWA implementations in llama.cpp master include Cohere2, OpenAI-MoE, Mimo2, Step-3.5. A single ISWA implementation would cover all of them.

Hardware / build

• EPYC 7663, 3× RTX 3090, 512 GB RAM
• Talos Linux, NVIDIA driver 595.58.03, CUDA 13.1.1
• ik_llama.cpp commit 13d7178d
• llama.cpp commit d6f303004 (for the comparison numbers)

Happy to rerun benchmarks against any prototype branch.

[ikawrakow]

Thank you for the issue, I'll look into it if Gemma4 remains popular and people don't forget about it the moment a new model comes out.

A few corrections to the estimated compression:

• In Gemma4 non-SWA layers have 2x smaller key and value dimensions than SWA layers
• One cannot keep just the n_swa tokens in the cache. Or, more precisely, one could but that would result in way too many copies and bad performance. llama.cpp keeps n_ubatch + n_swa tokens in the cache, where n_ubatch is the micro-batch size. For split mode graph and a dense model, performance tends to be best for n_ubatch = 1024, so it would be 2048 values per SWA layer, not 1024.

With that, for Gemma4-31B the compression becomes

compressed_kv    50 * 2048  + 5 * N
-------------- = -------------------- 
uncompressed_kv        55 * N 

and the table looks like this

context N	compressed / uncompressed	effective KV savings
2048	100%	0%
8192	32%	68%
32768	15%	85%
65536	12%	88%
262144	10%	90%

[FNsi]

Imo, comparing to the 31b or 26b, E4B could be more popular in time, I already see some people claim their finetunes surpass 26A4b in many benchmarks, indicating that could be a good start point for running in edge or some auto RL agent.

[ikawrakow]

So, what is the sentiment here? With MiniMax-2.7 and GLM-5.1 out, Qwen-3.5 a very worthy contender, and Qwen-3.6 on the horizon, is Gemma4 SWA KV cache compression still worth the effort?

[FNsi]

Personally I think we better have it, in some case people need different model from different companies, argue each other opinions or take a brainstorm? Current selection is limited in that case. In this case only Stepfun/minimax + Qwen 122b / 27b for most people?

[Cyrus-Hei]

I personally would love to see this being implemented as someone who really likes the Gemma 4 models, the current implementation is very vram-heavy without quanting the kv cache(especially for the 31b dense model). At higher context this swa kv cache implementation can easily pose a huge amount of vram savings as seen from the tables in this discussion, this would be very nice for people like me with less vram to play with.

[mindkrypted]

I just wanted to chime in and say that it would be very useful to have, particularly if Google decides to release the 124B mentioned in this post on X: https://xcancel.com/JeffDean/status/2039736943693668800
Discussion on HF: https://huggingface.co/google/gemma-4-31B-it/discussions/10

[jagmarques]

Tested Gemma 4 E2B (28 SWA + 7 global layers). The SWA layers have fixed memory cost independent of context length, so only the 7 global layers are worth compressing. Detection via sliding_window_pattern in the config works cleanly, but the E2B only has 1 KV head which kills retrieval quality regardless of compression, so the 31B dense model would be the real test.

[ormandj]

So, what is the sentiment here? With MiniMax-2.7 and GLM-5.1 out, Qwen-3.5 a very worthy contender, and Qwen-3.6 on the horizon, is Gemma4 SWA KV cache compression still worth the effort?

I would still like to see it, and it should still be useful if future gemma (or other) models come out with this architecture. Thank you!

[jagmarques]

For the 31B dense model, the 1-KV-head constraint does not apply (16 KV heads on SWA layers vs 1 on E2B). The SWA layers also stay at fixed memory regardless of context length, so compression on the global layers gives the most leverage per byte. Worth implementing for the 31B case specifically.

[AQJustice1]

i ran into this same issue, sadly with my hardware (dual 3090s) no backend is working correctly with the 31b SWA and therefore all have poor context window, except for mainline llama.cpp, which has suboptimal performance for production use.
I think gemma 4 is exceptional in some use cases and deserves to be used in full context.

[RBAncient]

I am using Gemma 4, and imo it's a better model than Qwen for general use. SWA KV cache compression would really help,
Please add this feature, Thanks!

[quasar-of-mikus]

The lack of an SWA implementation is the one thing preventing me from using ik_ for all my models.
With Gemma 4 31b q8_0, an unusable (for my use cases) max 14k context on ik_ but over 131k on mainline.

[mpetruc]

+1.
Spent the whole weekend trying to figure out why -sm layer is so slow (vs. vanilla llama.cpp, which is already slow) and -sm graph doesn't work beyond -c 2048 on 2xRTX 5070ti + 128Gb RAM. Should have searched first.
Gemma-4-31B is a fantastic model, very versatile, i hope this gets prioritized. Thank you.