ggml-openvino: enable in-place 4-D stateless KV cache updates on GPU

[osabnis]

Overview

Improve long-context decode performance and reduce GPU memory usage across models on the OpenVINO GPU backend in stateless mode.

The stateless KV cache currently uses a [1, 1, ctx, n_heads_kv * head_size] model input, bound to its first n_kv rows and updated with ScatterUpdate. The GPU plugin's ScatterUpdate copies the entire input to the output before writing the updated rows. Consequently, every decode step copies the full KV cache of every layer. As the context grows, the changing bound shape also forces the plugin to allocate new memory pool buffers that cannot be reused.

This PR uses the 4-D layout already used in stateful mode: [1, ctx, n_heads_kv, head_size]. The input is bound over the entire cache, keeping its shape constant as the context grows. Cache updates use ScatterUpdate along the token axis, while reads use VariadicSplit to select the first attention_size tokens, followed by a Transpose.

This pattern enables the GPU plugin's StatelessKVFusion to replace cache reads and writes with a StatelessKV node that updates the cache in place and exposes it to attention without copying it. In the compiled Llama-3.2 1B graph, all 32 KV caches become StatelessKV nodes, attention uses 16 SDPA nodes, and no ScatterUpdate nodes remain.

The 4-D layout is enabled only on the GPU in stateless mode when the cache holds a single sequence with one token per row, is written exclusively through SET_ROWS, and is read through views that split rows into heads. Other configurations retain the previous layout, including CPU, NPU, stateful mode, static or split graphs, multiple streams, and the transposed V cache used by -fa off.

Set GGML_OPENVINO_DISABLE_KV_4D=1 to restore the previous layout in all cases.

Additional information

Performance

llama-bench -fa 1, Q4_K_M, Arc B390 (Panther Lake) GPU. One console run, two rounds (round 1 / round 2), comparing the same build with and without GGML_OPENVINO_DISABLE_KV_4D=1, on top of #343.

Decode at context depth 4096 (-n 64). Throughput is in tokens/s; peak GPU memory is measured during the same test.

Model	Before (tok/s)	After (tok/s)	Change	Peak GPU memory (MB)
Phi-3.5 mini	12.2 / 12.1	21.5 / 21.5	+77%	14,348 → 4,818
SmolLM2 1.7B	24.3 / 25.1	40.8 / 42.4	+68%	7,210 → 2,450
Phi-4 mini	20.6 / 21.3	26.4 / 26.4	+26%	6,931 → 3,939
Llama-3.2 1B	64.5 / 64.9	75.5 / 76.2	+17%	2,137 → 1,333
Qwen3.5 4B	23.6 / 23.7	25.2 / 25.0	+6%	5,012 → 4,263
Gemma-4 E2B	43.0 / 43.2	44.5 / 43.7	+2%	4,527 → 4,437

At context depth 0, decode changes range from -3% to +4%. Using -ub 256 produces similar decode gains at depth 4096.

Prefill at context depth 4096 (-p 512):

• Phi-3.5 mini: +35% (1,312 / 1,147 → 1,707 / 1,647 tokens/s).
• Other models: changes are within run-to-run variation.

Peak prefill GPU memory also drops substantially:

Model	Before (MB)	After (MB)
SmolLM2 1.7B	7,142	2,362
Phi-3.5 mini	14,665	5,135

Accuracy

KL divergence against the CPU backend at context length 1024 (decode, -b 1) is unchanged from the previous layout:

Model	KL divergence
Llama-3.2 1B	0.0406
SmolLM2 1.7B	0.0728
Gemma-4 E2B	0.3634
Phi-4 mini	0.0357
Qwen3.5 4B	0.0250
Phi-3.5 mini	0.0312

At context length 4096, KL divergence between the new and previous layouts is at most 0.00018 for both decode and prompt evaluation. Perplexity with four sequences per batch (-c 256 -b 1024, 16 chunks) is unchanged: Llama-3.2 1B 21.6656 with both layouts, SmolLM2 1.7B 12.5269 → 12.5265.

test-backend-ops: 8024/8024 passed.

Notes

SDPA takes slightly longer to read the cache through the padded view. In a per-op GPU profile of Llama-3.2 1B decode at context depth 4096, SDPA increases from 8.1 to 9.3 ms. The cache update it replaces costs much more: 3.4 ms with ScatterUpdate against 0.6 ms with StatelessKV.

In prefill, the trade-off is closer. In a per-op profile of one 512-token prompt pass at context depth 4096 (-ub 512), Llama-3.2 1B SDPA increases from 19.3 to 25.5 ms, while the cache update drops from 4.5 to 2.1 ms. That is a net cost of about 4% of the pass. With -ub 2048, SDPA time is about the same in both layouts. For Gemma-4 E2B and Qwen3.5 4B, the difference is about 1 ms or less.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. AI assisted with testing and analysis; I performed the implementation and remaining work.