Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
johnkarlhill wants to merge 2 commits into
base: master
Choose a base branch
from
+153 −26
Draft

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b0b76ae
sycl: extend oneDNN SDPA to Q4_0-Q8_0 and F32 KV caches
johnkarlhill Jul 18, 2026
6e9257c
docs: drop GGML_SYCL_FA_DEBUG from SYCL.md (not shipped in this PR)
johnkarlhill Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 127 additions & 22 deletions ggml/src/ggml-sycl/fattn-onednn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@
#include <cstdio>
#include <cstring>
#include <string>
#include <optional>
#include <unordered_map>
#include <vector>

#include "fattn-onednn.hpp"
#include "fattn-tile.hpp"
#include "convert.hpp"

// set minimum query length to treat as prefill (32)
#define GGML_SYCL_FA_ONEDNN_MIN_Q 32
Expand All @@ -33,10 +35,30 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
const ggml_tensor * mask = dst->src[3];
const ggml_tensor * sinks = dst->src[4];

// gate for f16 KV only for now
// need to implement quantized KV
// F16 KV: native SDPA at any KV length.
// Non-F16: dequant to F16 then SDPA at prefill lengths. Only the
// standard quantized KV cache types (Q4_0-Q8_0) and F32 are accepted
// because their to_fp16_sycl conversion is verified. BF16 and IQ*
// are excluded: BF16 needs a strided conversion kernel that does not
// exist yet; IQ types are model-weight-only quants with no dequant
// registration and are never used as KV caches.
if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) {
return false;
auto kt = K->type, vt = V->type;
bool k_ok = kt == GGML_TYPE_F32 || kt == GGML_TYPE_Q4_0 || kt == GGML_TYPE_Q4_1 ||
kt == GGML_TYPE_Q5_0 || kt == GGML_TYPE_Q5_1 || kt == GGML_TYPE_Q8_0;
bool v_ok = vt == GGML_TYPE_F32 || vt == GGML_TYPE_Q4_0 || vt == GGML_TYPE_Q4_1 ||
vt == GGML_TYPE_Q5_0 || vt == GGML_TYPE_Q5_1 || vt == GGML_TYPE_Q8_0;
if (!k_ok || !v_ok) {
return false;
}
if (Q->ne[1] < 32 || K->ne[1] < 1024) {
return false;
}
for (const ggml_tensor * t : {K, V}) {
if (t->type == GGML_TYPE_F16 && t->nb[1] % (t->ne[0] * 2) != 0) {
return false;
}
}
}
// gate for the following cases
// 1. if the oneDNN graph Add node has no input --> skip
Expand Down Expand Up @@ -199,13 +221,101 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
dnnl::engine eng = ctx.engine_dnnl(stream);
dnnl::stream strm = ctx.stream_dnnl(stream);

// cont/cast inputs to contiguous f16 (head-major) -- the layout the fast systolic path wants.
ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
ggml_sycl_pool_alloc<sycl::half> Kf(ctx.pool(), (size_t) Hkv * seq * d);
ggml_sycl_pool_alloc<sycl::half> Vf(ctx.pool(), (size_t) Hkv * seq * d);
cont_to_f16_sycl<float> ((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf.get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf.get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
// Q: always f32 -- copy to dense f16.
ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
cont_to_f16_sycl<float>((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);

// K/V: use pool-alloc for both F16 and dequant paths.
sycl::half * K_ptr = nullptr;
sycl::half * V_ptr = nullptr;
std::optional<ggml_sycl_pool_alloc<sycl::half>> Kf_pool;
std::optional<ggml_sycl_pool_alloc<sycl::half>> Vf_pool;

if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf_pool->get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
K_ptr = Kf_pool->get();
V_ptr = Vf_pool->get();
} else if (ggml_is_quantized(K->type)) {
// Quantized K/V: dequant to dense F16 using pool, same lifetime as F16 path.
Kf_pool.emplace(ctx.pool(), ggml_nelements(K));
K_ptr = Kf_pool->get();
{
const char * K_data = (const char *)K->data;
const bool k_non_dense = ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1;
const bool k_gemma = k_non_dense &&
((int64_t)K->nb[2] < (int64_t)K->ne[1] * (int64_t)K->nb[1]);
if (ggml_is_contiguously_allocated(K) && !k_non_dense) {
to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst);
to_fp16(K_data, K_ptr, ggml_nelements(K), stream);
} else {
const size_t bs = ggml_blck_size(K->type);
const size_t ts = ggml_type_size(K->type);
to_fp16_nc_sycl_t to_fp16 = ggml_get_to_fp16_nc_sycl(K->type);
int64_t s01, s02, s03;
if (k_gemma) {
const int64_t blk_per_row = (int64_t)K->ne[0] / bs;
s01 = (int64_t)Hkv * blk_per_row;
s02 = blk_per_row;
s03 = (int64_t)K->ne[1] * s01;
} else {
s01 = (int64_t)K->nb[1] / ts;
s02 = (int64_t)K->nb[2] / ts;
s03 = (int64_t)K->nb[3] / ts;
}
to_fp16(K_data, K_ptr,
K->ne[0], K->ne[1], K->ne[2], K->ne[3],
s01, s02, s03, stream);
}
}
// Quantized V: always dequant separately. Even when K and V share
// the same underlying allocation (V is a view of K with the same
// data pointer), their logical values differ because the quantized
// elements at different positions/offsets represent different K/V
// data. Master's F16 path also never aliases K and V.
Vf_pool.emplace(ctx.pool(), ggml_nelements(V));
V_ptr = Vf_pool->get();
{
const char * V_data = (const char *)V->data;
const bool v_non_dense = ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1;
const bool v_gemma = v_non_dense &&
((int64_t)V->nb[2] < (int64_t)V->ne[1] * (int64_t)V->nb[1]);
if (ggml_is_contiguously_allocated(V) && !v_non_dense) {
to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst);
to_fp16(V_data, V_ptr, ggml_nelements(V), stream);
} else {
const size_t bs = ggml_blck_size(V->type);
const size_t ts = ggml_type_size(V->type);
to_fp16_nc_sycl_t to_fp16 = ggml_get_to_fp16_nc_sycl(V->type);
int64_t s01, s02, s03;
if (v_gemma) {
const int64_t blk_per_row = (int64_t)V->ne[0] / bs;
s01 = (int64_t)V->ne[2] * blk_per_row;
s02 = blk_per_row;
s03 = (int64_t)V->ne[1] * s01;
} else {
s01 = (int64_t)V->nb[1] / ts;
s02 = (int64_t)V->nb[2] / ts;
s03 = (int64_t)V->nb[3] / ts;
}
to_fp16(V_data, V_ptr,
V->ne[0], V->ne[1], V->ne[2], V->ne[3],
s01, s02, s03, stream);
}
}
} else {
// F32: strided copy to dense F16 via cont_to_f16_sycl<float>.
Kf_pool.emplace(ctx.pool(), ggml_nelements(K));
K_ptr = Kf_pool->get();
cont_to_f16_sycl<float>((const char *) K->data, K_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3],
K->nb[1], K->nb[2], K->nb[3], stream);
Vf_pool.emplace(ctx.pool(), ggml_nelements(V));
V_ptr = Vf_pool->get();
cont_to_f16_sycl<float>((const char *) V->data, V_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3],
V->nb[1], V->nb[2], V->nb[3], stream);
}

// divide-by-(1/scale) reproduces ggml's score *= kq_scale on the proven probe graph.
const sycl::half scale_h = (sycl::half) (1.0f / kq_scale);
Expand All @@ -230,8 +340,8 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso

auto id2ptr = [&](size_t r) -> void * {
if (r == E.id_q) return Qf.get();
if (r == E.id_k) return Kf.get();
if (r == E.id_v) return Vf.get();
if (r == E.id_k) return K_ptr;
if (r == E.id_v) return V_ptr;
if (r == E.id_scale) return scbuf.get();
if (r == E.id_mask) return (void *) mask->data;
return nullptr;
Expand All @@ -245,16 +355,11 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
E.cp.execute(strm, ti, {to});

permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream);
// Single device: no sync is required, and actually PP perf is ~6% > wait_and_throw() (tested on llama-3.1-8b & qwen3.6-27b, both Q8_0, with Arc B70).
// Any future multi-GPU refactor MUST re-measure this single-device path and keep the best
// single-device PP speed. Otherwise (multiple devices/streams can race the reuse):
if (ggml_sycl_info().device_count > 1) {
// cont_to_f16 -> oneDNN execute -> permute is async on this stream, but the
// pool_alloc*s above free their device buffers at host return. Without this wait the next
// scheduler op re-acquires those bytes while the GPU is still computing the SDPA, turning
// it into garbage and collapsing multi-turn output to a single repeated token ("GGGGG...").
stream->wait_and_throw();
}
// cont_to_f16 -> oneDNN execute -> permute is async on this stream, but the
// pool_alloc*s above free their device buffers at host return. Without this wait the next
// scheduler op re-acquires those bytes while the GPU is still computing the SDPA, turning
// it into garbage and collapsing multi-turn output to a single repeated token ("GGGGG...").
stream->wait_and_throw();
}
catch (const std::exception & e) {
// any oneDNN/SYCL failure is non-fatal: fall back to the existing kernel (strictly additive).
Expand Down
30 changes: 26 additions & 4 deletions ggml/src/ggml-sycl/fattn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,10 +97,11 @@ static void ggml_sycl_flash_attn_ext_vec(ggml_backend_sycl_context & ctx, ggml_t
enum best_fattn_kernel {
BEST_FATTN_KERNEL_NONE = 0,
BEST_FATTN_KERNEL_VEC = 100,
BEST_FATTN_KERNEL_ONEDNN = 150, // added enum for onednn==150
BEST_FATTN_KERNEL_ONEDNN = 150, // oneDNN SDPA: native F16 (PR #25222)
BEST_FATTN_KERNEL_TILE = 200,
};


static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
GGML_UNUSED(device);
#ifndef SYCL_FLASH_ATTN
Expand All @@ -115,14 +116,28 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
const ggml_tensor * K = dst->src[1];
const ggml_tensor * V = dst->src[2];
const ggml_tensor * mask = dst->src[3];
const ggml_tensor * sinks = dst->src[4];

const int gqa_ratio = Q->ne[2] / K->ne[2];
GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);

float max_bias = 0.0f;
memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));

float logit_softcap = 0.0f;
memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));

bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;

// XMX-accelerated path: oneDNN SDPA (native F16 and dequant+non-F16).
// ONEDNN requires min 32 query tokens — short-circuit decode to avoid
// calling _supported() on every decode FA call.
if (Q->ne[1] >= 32
&& ggml_sycl_flash_attn_ext_onednn_supported(dst)) {
return BEST_FATTN_KERNEL_ONEDNN;
}


for (const ggml_tensor * t : {Q, K, V, mask}) {
if (t == nullptr || ggml_is_quantized(t->type)) {
continue;
Expand Down Expand Up @@ -170,6 +185,7 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
switch (K->type) {
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
break;
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
Expand All @@ -188,8 +204,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
return BEST_FATTN_KERNEL_NONE;
}

// For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes:
const bool can_use_vector_kernel = Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0;
// For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes.
// BF16 is excluded: the VEC kernel has no BF16 template (it needs GGML_SYCL_FA_ALL_QUANTS for non-F16/Q4_0/Q8_0).
const bool has_bf16 = (K->type == GGML_TYPE_BF16 || V->type == GGML_TYPE_BF16);
const bool can_use_vector_kernel = Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0
&& !has_bf16;

// Fused-XMX path: oneDNN Graph SDPA (flash attention). Strictly
// additive -- taken only when statically supported, otherwise falls through to VEC/TILE below.
Expand All @@ -216,7 +235,9 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const

void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
ggml_sycl_set_device(ctx.device);
switch (ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst)) {

const best_fattn_kernel fk = ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst);
switch (fk) {
case BEST_FATTN_KERNEL_NONE:
GGML_ABORT("Not support Flash-Attention");
case BEST_FATTN_KERNEL_ONEDNN:
Expand All @@ -233,6 +254,7 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
ggml_sycl_flash_attn_ext_vec(ctx, dst);
break;
}

}

bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) {
Expand Down