ggml-openvino: Fix SWA layer detection and KV write indices (Gemma-4-E2B past 512 tokens)

[osabnis]

Overview

On dev_backend_openvino (308ccd8), Gemma-4 produces incorrect output once the context exceeds 512 tokens. KLD against the CPU backend is 0.83 at 1024 tokens and 5.4 at 2048 tokens, against 0.36 with this fix, and the CPU device is affected too. Two separate issues in the sliding-window attention (SWA) cache handling cause this. This PR fixes both issues in separate commits.

1. SWA layers are not detected when both caches have the same size (0d3ea01)

The backend identifies SWA layers by comparing the sizes of the full-attention and SWA KV caches. The SWA cache is sized n_swa + n_ubatch, which for Gemma-4 with -c 1024 -ub 512 is equal to the full cache size. As a result, no layers are identified as SWA, causing the SWA layers to slice the KV cache using the full-attention length.

When the cache sizes are equal, the two masks are now distinguished based on the number of layers that use each mask: the larger group is treated as SWA (4 of 5 layers in Gemma-4). If both groups have the same number of layers, there is no way to distinguish them, so the layers remain unclassified as before.

2. SWA and full-attention caches share the same KV write index input (b61f73f)

Since 9655061 (ggml-org#26625), the KV write index tensors for both caches are named attn_inp_k_idxs / attn_inp_v_idxs. The backend merges tensors with the same name into a single OpenVINO input. Once decoding passes the SWA window, this causes the SWA layers to write their keys and values to rows belonging to the full-attention cache.

The SWA index tensors now use a _swa input name, matching the naming already used for the SWA mask inputs.

One file changed: ggml/src/ggml-openvino/ggml-decoder.cpp (+31 / -4).

Additional information

Setup: Panther Lake (Core Ultra X7 358H, Arc B390), Windows 11, OpenVINO 2026.4.0, base 308ccd8. Accuracy is measured as KL divergence against llama.cpp's CPU backend on WikiText-2 (llama-perplexity --kl-divergence); lower is better.

Gemma-4 E2B Q4_K_M:

Check	Before	After
GPU, decode path (-b 1), ctx 1024	0.828	0.364
GPU, decode path (-b 1), ctx 2048	5.370	0.357
GPU, prompt path, ctx 1024	2.554	0.367
CPU device, prompt path, ctx 1024	2.560	0.369
GPU, 4 sequences per batch, ctx 512	0.324	0.324

The fixed results are consistent with the results from an earlier version of these fixes on the previous base (ab07546), based on a sweep of 7 models from 512 to 4096 tokens.

No regressions:

• GPU, 7 models (Gemma-4 E2B, SmolLM2 1.7B, Phi-3.5 mini, Llama-3.2 1B, Phi-4 mini, Gemma-3 4B, Qwen3.5 4B, all Q4_K_M): perplexity with 4 sequences per batch and generated text are unchanged.
• CPU llama-simple with Llama-3.2 and Gemma-4: correct text.
• llama-bench -fa 1, decode (tg128) before -> after:

Device	Llama-3.2 1B	Gemma-4 E2B
CPU	40.0 -> 40.3	27.1 -> 26.9
GPU	66.7 -> 66.8	36.9 -> 36.5
GPU, stateful	83.6 -> 83.8	41.2 -> 41.9

Prefill (pp512) differences are within llama-bench's reported spread.

Requirements

• I have read and agree to the contributing guidelines
• AI usage disclosure: YES. AI ran the unit testing while I did the rest.

[cavusmustafa]

For a model with multiple layer classes but no SWA, wouldn’t this misclassify the larger group as SWA?

[osabnis]

Good catch, thanks!
You're right that two masks don't imply SWA - I found out that llama.cpp's DeepSeek sparse attention builds two full-attention masks with the same name and the fallback would have marked the larger group as SWA.

I pushed a fix. The fallback now classifies the larger group as SWA only once its mask actually keeps fewer keys per row than the other mask, i.e. once the window is visible. Before that the two masks have identical contents, so leaving the layers unclassified gives the same result. The classification is remembered per KV cache buffer, so it doesn't flip back when a new sequence starts, and swa_layers is now part of the graph reuse check. The switch costs one recompile, and only in the equal-extent case (e.g. gemma-4 at -c 1024 -ub 512).

[osabnis]

Follow-up on the SWA classification: I pushed a commit that classifies the layers from the first graph instead of waiting for the window to show in the mask.

Waiting for the window had a cost I missed earlier. The graph inputs change once the sequence passes the window, so the model recompiles partway through the sequence: about 9 s for gemma-4 at -c 1024, when it first passes position 512.

Two masks still don't imply SWA (DeepSeek sparse attention), but the classification doesn't need them to:

• The SWA role only gives that group of layers its own mask input, KV write indices and attention size. That's correct whether or not the second mask is really a sliding window.
• The one place that uses the window itself is the stateful mask rebuild, and it reads the window back from the mask. On a full mask that gives back the plain causal mask.

So the larger of the two mask groups now gets the role from the first graph, and the per-KV-buffer bookkeeping is gone. Groups of equal size are still left unclassified.

[cavusmustafa]

I think we can merge this for now as there seems to be no way to detect swa layers without some model level assumptions. Can you add a comment in the code so we can revisit this layer? Also, probably we can figure out a way to gate this detection to only stateful mode. For stateless, I believe we can ignore this and let llama.cpp handle swa buffer management anyways. All these can be a future work.

[osabnis]

Thanks! Added a TODO above the fallback.

On stateless: llama.cpp manages the SWA buffers, but the graph still names both masks attn_inp_kq_mask and both caches' write indices attn_inp_k_idxs / attn_inp_v_idxs, and we key OV parameters by name.
Without the split, both layer groups read one mask and one set of write indices. Gemma-4 E2B on GPU, ctx 1024, KLD vs CPU on the current fork tip: decode 0.827 / prompt 2.554 without this PR, 0.364 / 0.367 with it.

What stateless really needs is the two groups kept apart, not the SWA role, so we could split them by tensor identity there and keep the "larger group is SWA" guess for stateful only. Happy to do that as a follow-up!