[osabnis]

Details:

Improve long-context decode performance for grouped-query attention (GQA) in the llama.cpp OpenVINO backend by selecting the key-partitioned single-token SDPA kernel for up to four query rows. Fix two attention-mask indexing bugs exposed when this kernel handles multiple rows.

Performance

Decode throughput as a percentage of the llama.cpp Vulkan backend on Intel Arc B390 (Panther Lake), Windows, using Q4_K_M models and llama-bench with tg64:

Model	Depth 0	Depth 1024	Depth 4096	Depth 8192
Llama-3.2 1B	79% → 80%	75% → 77%	60% → 66%	52% → 59%
Phi-4 mini	76% → 79%	69% → 73%	59% → 62%	50% → 52%
Qwen3.5 4B	64% → 63%	63% → 63%	60% → 61%	57% → 59%

The change narrows the gap that grows with context length. Absolute throughput in tokens/s (before / after / Vulkan):

Model	Depth 0	Depth 1024	Depth 4096	Depth 8192
Llama-3.2 1B	96.76 / 97.33 / 121.72	86.26 / 89.02 / 115.01	61.38 / 67.49 / 102.43	46.49 / 52.46 / 88.72
Phi-4 mini	31.58 / 32.50 / 41.34	27.06 / 28.29 / 39.02	19.98 / 20.89 / 33.74	14.05 / 14.69 / 28.18
Qwen3.5 4B	21.06 / 20.89 / 32.91	20.28 / 20.31 / 32.21	18.77 / 19.08 / 31.18	17.02 / 17.57 / 29.89

At depth 8192 decode is 12.8% faster for Llama-3.2 1B, 4.5% for Phi-4 mini and 3.2% for Qwen3.5 4B. Depth-0 differences are within run-to-run noise (up to 3.5% between rounds). Models that do not stack query rows in attention (Gemma-4 E2B, SmolLM2 1.7B, and Phi-3.5 mini) are unchanged.

Accuracy is unchanged: KL divergence against the llama.cpp CPU backend is 0.041104 → 0.041087 for Llama-3.2 1B, 0.035677 → 0.035695 for Phi-4 mini, and 0.024973 → 0.024896 for Qwen3.5 4B.

On Intel Arc B580, Llama-3.2 1B decode at depth 8192 improved from 119 to 153 tokens/s (+28%) and Phi-4 mini by 10% (measured on an earlier revision of this PR, with the same kernel selection for these models).

Implementation

sdpa_opt.cpp: Previously, sdpa_opt selected multi-token kernels for any query length above one. With GQA, llama.cpp stacks the query heads belonging to each KV head into rows so each KV head is read once. Decode therefore reaches SDPA with 3–4 query rows on common models, but the multi-token path does not split the keys, leaving long contexts with too few work-groups.

Use the key-partitioned single-token kernel for up to max_single_token_rows = 4. This kernel supports multiple query rows and splits the keys into partitions. Keep the micro single-token kernel restricted to one row because it was slower than sdpa_opt for multiple rows.

Query rows stay on the multi-token kernels where the single-token stages are not exact for them (single_token_supports_query_rows()): per-channel compressed KV needs the micro kernel, every key partition adds the sink to its softmax sum, and the finalization stage merges at most get_max_partitions_num() (128) partitions. The finalization kernel now gets MAX_PARTITIONS_NUM from this host constant instead of a kernel-local define, and exp_sums/max_logits are sized by the V head size.

sdpa_utils.hpp: is_prefill_stage() had no other users, so it becomes get_target_seq_len() (query rows, -1 if dynamic), which both checks use.

The four-row limit preserves the existing multi-token path for prompt chunks of two or more tokens while routing typical GQA decode through the faster partitioned kernel.

Why limit it to four rows? With a limit of 16, prompt chunks at depth 4096 regressed as query-row counts increased: +0.3% at 6 rows, -9% at 8, -11% at 9–12, and -20% to -26% at 15–16. At depth 0, all results were within ±2%. The single-token kernel reads keys once per query row, so multi-token kernels win for larger chunks.

Attention-mask indexing fixes

sdpa_opt.cl: Enabling multiple query rows exposed an out-of-bounds attention-mask read. The INPUT3_GET_INDEX_SAFE macro does not parenthesize its arguments, so the modulo in the expression passed to it applied only to seq_idx, producing an incorrect index when the mask was broadcast across query rows. This caused GPU page faults (CL_OUT_OF_RESOURCES).

Parenthesize the arguments at this call site. Single-row decode, the only case that previously used this kernel, was unaffected. This is the only GPU-plugin kernel call site found passing unparenthesized expressions to the macro; fixing the macro itself in fetch_data.cl could cover other kernels in a follow-up.

load_attn_mask: A mask broadcast over the keys (for example a scalar mask) has a single element along X. The _SAFE index wraps the first key to it, but the following keys read past it, so the result depended on the memory after the mask. The keys now step through the mask only when it has more than one element along X.

Tests

New tests in sdpa_gpu_test.cpp: 3, 4 and 5 query rows (static and dynamic shapes, batch 1 and 2), and the cases that stay on the multi-token kernels: per-channel compressed KV, a sink input, and more than 128 key partitions.

• Intel Arc B390: ov_gpu_unit_tests --gtest_filter=*sdpa*:*SDPA*:*scaled_dot* (excluding the pre-existing *990_1_2* crash): 712 passed, 0 failed, in default and shuffled orders.

Measurement setup

Measured at the local console on master cd202900e4 (before) and this PR (after): one OpenVINO build in which only the GPU plugin differs, with the same llama.cpp build (OpenVINO fork 675680b54 plus ravi9/llama.cpp#382, stateless, -fa 1) and the same Q4_K_M models. Decode: llama-bench -p 0 -n 64 -d <depth> -r 3, one depth per call, mean of two rounds with the order reversed in the second round. KL divergence: llama-perplexity -c 1024 -b 1 against stored CPU logits. The PR head before the review fixes measured within ±2% of the final head on all models and depths, so the review fixes do not change these models' kernels.

Tickets:

• CVS-196699

AI Assistance:

• AI assistance was used for analysis, kernel-fault debugging, test execution and measurements.

[p-durandin]

build_jenkins

[osabnis]

The only failing check is Jenkins build_win_vs2022_release, and it’s an infrastructure issue, not this change. The OpenVINO build finished; the job then failed in the GenAI build step because two dependency downloads from github.com returned HTTP 504 (gguf-tools and pybind11). All GitHub Actions checks and the other Jenkins jobs pass, including the Windows debug build and the GPU unit tests on DG2.
@p-durandin could you please re-trigger the Jenkins runs when possible? Thanks!

[maxnick]

@allnes , could you please review?

[allnes]

Q=2–4 now skips micro multi-token execution, while micro single-token is restricted to Q=1. However, validation still accepts per-channel compressed KV based on micro support. The selected regular kernel indexes scales by token Y with X=0, whereas [B,H,1,D] per-channel scales require Y=0 and X=channel. This produces incorrect dequantization and out-of-bounds scale reads. Preserve the micro multi-token route for per-channel KV in both stage construction and execution, or implement channel-indexed dequantization before selecting the regular path.

[osabnis]

Good catch, thanks. Fixed in 3db4909: with per-channel compressed KV, 2-4 query rows now stay on the multi-token route (micro), in stage construction, execution and buffer sizing alike, through one predicate (single_token_supports_query_rows()).
New tests: sdpa_gpu_compressed_kv_per_channel with 3 query rows (int4 and int8).

[allnes]

output_count is B * H * query_length * V_width, but this expression divides it by the Q/K width. These widths may differ, so when V is narrower, both exp_sums and max_logits are undersized even though the kernel indexes B * H * query_length * partitions entries. For example, V width 64 with Q/K width 128 halves the allocation. Derive both buffer sizes directly from the canonicalized batch, head and query dimensions.

[osabnis]

Agreed. Fixed in 3db4909: the two buffers are now sized from the V head size (SDPABase::get_effective_v_head_size(), the value the kernels get as V_HEAD_SIZE), so they hold B * H * query rows * partitions entries. This also fixes the same sizing for a single query row.

[allnes]

Finalization allocates max_logits_u_exp_sum[128], but its load, update and output-reduction loops use the unrestricted partition count. With 256 keys per partition, K=32769 already requires 129 entries and accesses index 128 outside the array. The new Q=2–4 routing exposes this unsupported case. Keep the multi-token fallback when the partition count exceeds 128, consistently in stage construction and execution, or change finalization to reduce arbitrary partition counts without exceeding local storage.

[osabnis]

Fixed for the new route in 3db4909: 2-4 query rows with more than 128 key partitions keep the multi-token kernels. The limit now comes from one host constant (get_max_partitions_num()), passed to the finalization kernel as MAX_PARTITIONS_NUM instead of a kernel-local define.
New test: 3 query rows over 33024 keys (129 partitions).

The single-row path has the same 128-partition limit today (decode above 32K keys when micro SDPA is not used); I'd rather fix the finalization kernel for any partition count in a separate PR than grow this one.

[allnes]

Each single-token partition adds the sink contribution to its local softmax denominator. Finalization then combines all partition denominators, counting the same virtual key once per partition. With P partitions, the global denominator contains P sink contributions instead of one, incorrectly reducing the attention output. Include the sink exactly once in the global softmax normalization, with the correct maximum-logit adjustment, or retain the multi-token route for sink inputs.

[osabnis]

Right, each partition adds the sink to its exp_sum. Fixed for the new route in 3db4909: with a sink input, 2-4 query rows keep the multi-token route.
New test: sdpa_gpu_sink_test with 3 query rows over 300 keys (2 partitions), which fails without this change.

As with the previous comment, the single-row path already behaves this way when it has more than one partition; I can fix the finalization (add the sink once, after the global max) in a follow-up PR.

[allnes]

The Q=2–4 threshold replaces supported micro prefill with the regular single-token path for all eligible inputs. This path reads K/V separately for each query row; with multiple key partitions it also needs scratch buffers and an additional finalization launch. Query length alone therefore discards the faster implementation for some supported shapes. Keep micro prefill eligible and restrict the partitioned route to shapes where its additional work is justified.

[osabnis]

The description's numbers compare against master's multi-token selection on Arc B390: decode with 3-4 query rows gets faster as the context grows (+2.6% to +9.6% at depth 4096-8192 on Llama-3.2 1B, Phi-4 mini and Qwen3.5 4B) and is within run-to-run noise at depth 0. If you prefer, I can keep micro prefill when the keys fit in one partition, since that is where the partitioned kernel has no work-group advantage.

[osabnis]

While adding the 32-row sink case I saw it fail only in some test orders. The cause is in load_attn_mask, not in this
PR's selection: with a mask broadcast over the keys (the sink test's scalar mask), keys after the first read past the
mask's single element, so the result depended on what the allocator left behind it. The second commit (49d071b) fixes the read and makes the test deterministic: the mask sits at the start of a buffer whose other elements are large, so the sink
tests fail without the kernel fix and pass with it in any order.

[allnes]

The new few-row route returns NaN for a fully masked key partition. For an FP32 SDPA primitive with Q=3, KV=512 and D=64, leave the first 256 keys unmasked and set the remaining mask entries to -Inf: all 384 output elements become NaN, while the multi-token base produces finite output. The empty partition divides by zero in sdpa_opt.cl:756, and finalization preserves NaN when multiplying its partial output by zero. Handle zero-sum partitions before normalization/finalization, or exclude this precision from the new route. This is confirmed at primitive level; the standard FP32 frontend currently decomposes SDPA.

[allnes]

V_HEAD_SIZE > K_HEAD_SIZE also needs a fallback here. With Q=3, Q/K width 64 and V width 256, SG_SCALE_FACTOR=1 allocates query_local[64], but the unconditional query loads/stores address indices 0–255. For V width 128, the factor-2 output reduction reuses that array across indices 0–127. These accesses exceed the declared local array, although the numeric repro passes on the tested runtimes. Keep the multi-token route for V>K, or bound every query load and size scratch for both query storage and V-wide reduction; cover asymmetric widths with static and dynamic tests.