Extended SYCL oneDNN SDPA to non-FP16 KV caches (Q4_0–Q8_0 and FP32)

[johnkarlhill]

Overview

Extends the SYCL oneDNN SDPA path (#25222, now merged) to handle non-FP16 KV caches. The existing implementation gives great XMX SDPA speedup during prefill, but only for native FP16 K/V. If you run a quantized KV cache (q8_0, q4_0, etc.) to save VRAM, you were stuck on the TILE path and paying the performance cost.

Simple approach: before passing K/V into the oneDNN SDPA graph, dequantize or convert them to dense FP16 on-device. From there, the SDPA executes the same fused systolic kernel as the standard FP16 pipeline.

Benchmarks

All runs: B70 or B50, 32K prefill, quantized KV cache (q8_0 or q4_0), MTP off. ONEDNN=1 vs ONEDNN=0 (forces TILE fallback). Token gen is identical across all runs since the ONEDNN path does not fire during decode (Q=1).

B70 (Arc Pro B70, 32 GB)

Model	KV Cache	ONEDNN=1 (t/s)	ONEDNN=0 (t/s)	Speedup
Qwen 3.6 27B Q5_K_XL	q8_0	823	330	2.49x
Gemma 4 31B Q4_K_XL	q4_0	445	177	2.51x
Qwen 3.6 35B A3B MoE Q4_K_XL	q8_0	1184	834	1.42x

B50 (Arc Pro B50, 16 GB)

Model	KV Cache	ONEDNN=1 (t/s)	ONEDNN=0 (t/s)	Speedup
Qwen 3.5 9B Q6_K_XL	q8_0	893	403	2.21x
Qwen 3.5 4B Q8_K_XL	q8_0	1197	465	2.57x
Gemma 4 12B Q4_K_XL	q8_0	463	195	2.37x
Gemma 4 E4B Q4_K_XL	q8_0	1098	342	3.21x

test-backend-ops passes 3961/3961. Multi-turn coherence tested on Gemma 26B and Qwen 27B with assorted prompt sequences: clean output, draft acceptance around 90%.

Supported KV types and gate

F16 KV runs native SDPA at any length, same as #25222.

Non-FP16 types gate at K >= 1024 and Q >= 32 (prefill only):

• Q4_0, Q4_1, Q5_0, Q5_1, Q8_0 dequantized via to_fp16_sycl / to_fp16_nc_sycl, covering both contiguous and non-contiguous layouts like Gemma's interleave
• F32 copied to FP16 via cont_to_f16_sycl

Inherits #25222's BMG-only architecture gate for the underlying oneDNN kernel. Includes an unconditional stream->wait_and_throw() after SDPA compute, without this the pool allocator reuses device buffers while the GPU is still executing, corrupting the KV cache on subsequent turns (#25741 by @malsbat).

Requirements

• I have read and agree with the contributing guidelines.
• AI usage disclosure: YES. Claude Code + DeepSeek V4 Pro assisted with code review, debugging the stream sync issue. All code changes were reviewed and tested by me.

[ggml-gh-bot]

Hi @johnkarlhill, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 3 open PRs.

• AI-generated content: This project does not accept PRs, descriptions or commit messages that are fully or predominantly AI-generated. If you have used AI to assist you in writing code, please make sure to disclose that explicitly.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[mndodd]

Any reason not to use mkl_fa_dequant_chunk() from your #25025? I'm testing something similar with TILE and it seems to be worth doing. This bounds transient memory allocation pressures that may not be apparent until higher context utilization.

[johnkarlhill]

Any reason not to use mkl_fa_dequant_chunk() from your #25025? I'm testing something similar with TILE and it seems to be worth doing. This bounds transient memory allocation pressures that may not be apparent until higher context utilization.

Short answer: we can't use chunked dequant here because oneDNN's SDPA graph (sdp_primitive_kernel_t) takes the full K, V, Q, and mask as monolithic inputs. It's a single fused kernel call, one graph execution. There is no API to feed K or V in pieces and rejoin with online softmax rescaling the way MKL does in #25025.

The MKL kernel can chunk because it owns the outer loop: it dequants one 8192-token chunk of K/V, runs GEMM, does the softmax rescale, and moves to the next chunk. The oneDNN graph compiles all of that (MatMul, Scale, Mask, SoftMax, MatMul) into one systolic kernel.

That said, the memory pressure is real. At typical prefill contexts (32K tokens, single head), the dequant buffer is about 64 MB for K+V combined. At 256K it grows to roughly 512 MB transient. This is pool-allocated and released at function return, so it does not creep, but it is a spike. Thinking out loud, once PR25025 is merged, I could see a flag that would shift processing from oneDNN SDPA to to MKL at a certain context, thus minimizing the spike. That way high-context users get bounded memory and still stay accelerated, just via MKL GEMM rather than the fused SDPA kernel. Not sure if the maintainers would let that go through...

[mndodd]

One of the problems here (and elsewhere) is that we're providing a user with levers to buy VRAM by spending compute or quality, and then turning around and charging them VRAM behind their backs.

The mechanism you're describing about shifting processing does work though.

[johnkarlhill]

One of the problems here (and elsewhere) is that we're providing a user with levers to buy VRAM by spending compute or quality, and then turning around and charging them VRAM behind their backs.

The mechanism you're describing about shifting processing does work though.

I'm not quite sure what the expectation is. The llama.cpp pipeline is FP16. To push anything non-FP16 through, you have to convert it first. The quantized cache has to be decompressed before it can be processed, and that decompression has to happen somewhere. With #25025 I was able to decompress in smaller chunks because I control the outer loop. With this PR the kernel is fused, so I'm not in control of the chunks. And to be clear, this is a very, very temporary spike... the memory is given back immediately at function return.

[arthw]

Please provide a detail example to show the benefit of this PR.
I can't get the perf increase on Qwen3.5-4B-Q4_0.gguf B60.

[Titaniumtown]

@johnkarlhill can you follow the PR template? and do you have some benchmarks or figures that support your changes?

[johnkarlhill]

Please provide a detail example to show the benefit of this PR. I can't get the perf increase on Qwen3.5-4B-Q4_0.gguf B60.

Did you use quant the cache (for example, cache-type-k = q8_0 / cache-type-v = q8_0)? This PR only extends PR25222's kernel to certain quanted kv caches.
I was able to show PP perf gains using 4B and E4B.

[johnkarlhill]

@johnkarlhill can you follow the PR template? and do you have some benchmarks or figures that support your changes?

Updated!