From be876204aa23bc0d8f890981fde2853555b132f9 Mon Sep 17 00:00:00 2001 From: Titaniumtown Date: Fri, 28 Aug 2026 01:53:31 -0700 Subject: [PATCH] sycl: bind the f16 KV cache in place for the oneDNN SDPA path (#27468) Measured at a live KV length of 34816 (32768 depth plus one 2048 ubatch), on Qwen3.8 27B Q4_K_S: per tensor 4 * 34816 * 256 * 2 B = 71.3 MB staged per call K and V, so 2x = 142.6 MB traffic per call read once, write once = 285.2 MB traffic per ubatch 285.2 MB * 16 calls = 4.56 GB One ubatch is one ggml_cgraph submission (llama_context::process_ubatch -> graph_compute), so that 4.56 GB is the cost of a single 2048-token prefill chunk, and it scales with the live KV length: the first ubatch of the same run, at seq = 2048, moves 0.27 GB. Reproduce the two measured inputs with: GGML_SCHED_DEBUG=2 llama-bench -m MODEL -p 8 -n 0 -r 1 -ngl 0 \ -fa on -ctk f16 -ctv f16 -v > nd.txt 2>&1 grep -E 'n_layer|n_head_kv|n_embd_head_k' nd.txt awk '/node # 0 /{g++} g==1 && /\(FLASH_ATTN\)/{n++} END{print n+0}' nd.txt --- ggml/src/ggml-sycl/fattn-onednn.cpp | 62 +++++++++++++++++++++++------ 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index a50129519..d41c2ddce 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -150,7 +151,8 @@ struct sdpa_partition { // Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out. // Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t). -static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) { +static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d, + const std::array & k_str, const std::array & v_str) try { using ltype = logical_tensor::layout_type; using dt = logical_tensor::data_type; using ldims = logical_tensor::dims; @@ -158,11 +160,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int const int rep = H / Hkv; const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq}, sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d}; + const ldims k_st(k_str.begin(), k_str.end()), v_st(v_str.begin(), v_str.end()); int64_t id = 0; sdpa_partition E; auto query = logical_tensor(id++, t, q_sz, ltype::strided); - auto key = logical_tensor(id++, t, kv_sz, ltype::strided); + auto key = logical_tensor(id++, t, kv_sz, k_st); auto score = logical_tensor(id++, fi, s_sz, ltype::strided); auto bmm1 = op(id++, op::kind::MatMul, "bmm1"); bmm1.set_attr(op::attr::transpose_b, true); // key is [.., seq, d] @@ -184,7 +187,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int smax.set_attr(op::attr::mode, "inf_as_zero"); smax.add_inputs({masked}); smax.add_outputs({probs}); - auto value = logical_tensor(id++, t, kv_sz, ltype::strided); + auto value = logical_tensor(id++, t, kv_sz, v_st); // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp). // converted to the f32 ggml dst in the permute below. @@ -198,6 +201,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int auto parts = g.get_partitions(); if (parts.size() != 1 || !parts[0].is_supported()) { + GGML_LOG_WARN("%s: oneDNN did not fuse the SDPA graph; falling back to TILE kernel\n", __func__); return E; // ok stays false -> caller falls back to TILE } E.ins = parts[0].get_input_ports(); @@ -209,6 +213,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int E.ok = true; return E; } +catch (const std::exception & e) { + // compile() can reject a stride set the partitioner never inspects; memoise the failure so the + // fallback costs one build rather than one per call. + GGML_LOG_WARN("%s: oneDNN SDPA partition build failed (%s); falling back to TILE kernel\n", __func__, e.what()); + return {}; +} void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try { const ggml_tensor * Q = dst->src[0]; @@ -234,13 +244,34 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso ggml_sycl_pool_alloc Qf(ctx.pool(), (size_t) H * q * d); cont_to_f16_sycl((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. + // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head + // plane is strided rather than dense, which is what an explicit stride vector expresses. + // Quantized and f32 KV still stage a dense copy -- the layout the k_str/v_str defaults describe. sycl::half * K_ptr = nullptr; sycl::half * V_ptr = nullptr; + std::array k_str{ Hkv * seq * d, seq * d, seq * d, d, 1 }; + std::array v_str = k_str; std::optional> Kf_pool; std::optional> Vf_pool; - if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { + auto bindable = [](const ggml_tensor * t) { + return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 && + t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0; + }; + auto elem_strides = [](const ggml_tensor * t) { + const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]); + const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]); + const int64_t s3 = (int64_t) (t->nb[3] / t->nb[0]); + // dims are {mb=1, Hkv, rep=1, seq, d}; the size-1 dims at 0 and 2 never advance an address. + return std::array{ s3, s2, s2, s1, 1 }; + }; + + if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) { + K_ptr = (sycl::half *) K->data; + V_ptr = (sycl::half *) V->data; + k_str = elem_strides(K); + v_str = elem_strides(V); + } else 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((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); @@ -341,19 +372,24 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] - // compile once per (device, shape), reuse across layers/calls. + // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always + // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ. static std::unordered_map cache; - char keyb[96]; - snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(), - (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d); + char keyb[256]; + snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(), + (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d, + (long long) k_str[0], (long long) k_str[1], (long long) k_str[3], + (long long) v_str[0], (long long) v_str[1], (long long) v_str[3]); auto it = cache.find(keyb); if (it == cache.end()) { - it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first; + it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d, k_str, v_str)).first; } sdpa_partition & E = it->second; - // _supported() is authoritative: if it accepted this op the partition must build. - // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback. - GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape"); + if (!E.ok) { + // oneDNN can decline a shape or a stride set that _supported() never sees; build_sdpa warns per key. + ggml_sycl_flash_attn_ext_tile(ctx, dst); + return; + } auto id2ptr = [&](size_t r) -> void * { if (r == E.id_q) return Qf.get();