diff --git a/ggml/src/ggml-sycl/fattn-common.hpp b/ggml/src/ggml-sycl/fattn-common.hpp index c6cc13cfb0..82813f7a99 100644 --- a/ggml/src/ggml-sycl/fattn-common.hpp +++ b/ggml/src/ggml-sycl/fattn-common.hpp @@ -6,6 +6,7 @@ #include "convert.hpp" #include "vecdotq.hpp" #include "fattn-buffers.hpp" +#include "fattn.hpp" #include "ggml.h" @@ -926,6 +927,7 @@ void launch_fattn( ggml_sycl_fattn_alloc K_f16(fbuf.K); ggml_sycl_fattn_alloc V_f16(fbuf.V); + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); ggml_sycl_pool_alloc KV_max(pool); ggml_sycl_pool_alloc dst_tmp(pool); ggml_sycl_pool_alloc dst_tmp_meta(pool); @@ -944,10 +946,11 @@ void launch_fattn( const size_t bs = ggml_blck_size(K->type); const size_t ts = ggml_type_size(K->type); - K_f16.alloc(ggml_nelements(K)); + sycl::half * K_f16_ptr = extra.K_buffer_ptr ? (sycl::half *) extra.K_buffer_ptr + : K_f16.alloc(ggml_nelements(K)); if (ggml_is_contiguously_allocated(K)) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst); - to_fp16(K_data, K_f16.ptr, ggml_nelements(K), main_stream); + to_fp16(K_data, K_f16_ptr, ggml_nelements(K), main_stream); nb11 = nb11 * bs * sizeof(sycl::half) / ts; nb12 = nb12 * bs * sizeof(sycl::half) / ts; @@ -958,13 +961,13 @@ void launch_fattn( const int64_t s01 = nb11 / ts; const int64_t s02 = nb12 / ts; const int64_t s03 = nb13 / ts; - to_fp16(K_data, K_f16.ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream); + to_fp16(K_data, K_f16_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream); nb11 = K->ne[0] * sizeof(sycl::half); nb12 = K->ne[1] * nb11; nb13 = K->ne[2] * nb12; } - K_data = (char *) K_f16.ptr; + K_data = (char *) K_f16_ptr; } if (need_f16_V && V->type != GGML_TYPE_F16) { @@ -977,11 +980,12 @@ void launch_fattn( const size_t bs = ggml_blck_size(V->type); const size_t ts = ggml_type_size(V->type); - V_f16.alloc(ggml_nelements(V)); + sycl::half * V_f16_ptr = extra.V_buffer_ptr ? (sycl::half *) extra.V_buffer_ptr + : V_f16.alloc(ggml_nelements(V)); if (ggml_is_contiguously_allocated(V)) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst); - to_fp16(V_data, V_f16.ptr, ggml_nelements(V), main_stream); - V_data = (char *) V_f16.ptr; + to_fp16(V_data, V_f16_ptr, ggml_nelements(V), main_stream); + V_data = (char *) V_f16_ptr; nb21 = nb21 * bs * sizeof(sycl::half) / ts; nb22 = nb22 * bs * sizeof(sycl::half) / ts; @@ -992,13 +996,13 @@ void launch_fattn( const int64_t s01 = nb21 / ts; const int64_t s02 = nb22 / ts; const int64_t s03 = nb23 / ts; - to_fp16(V_data, V_f16.ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream); + to_fp16(V_data, V_f16_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream); nb21 = V->ne[0] * sizeof(sycl::half); nb22 = V->ne[1] * nb21; nb23 = V->ne[2] * nb22; } - V_data = (char *) V_f16.ptr; + V_data = (char *) V_f16_ptr; } } diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index d41c2ddce3..4349363a3d 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -14,9 +14,21 @@ // set minimum query length to treat as prefill (32) #define GGML_SYCL_FA_ONEDNN_MIN_Q 32 -bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { +bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V) { + if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) { + return false; + } + 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; + }; + return bindable(K) && bindable(V); +} + +bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit) { #if !GGML_SYCL_DNNL GGML_UNUSED(dst); + GGML_UNUSED(use_shape_limit); return false; #else if (!g_ggml_sycl_fa_onednn) { @@ -44,7 +56,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { if (!k_ok || !v_ok) { return false; } - if (Q->ne[1] < 32 || K->ne[1] < 1024) { + if (use_shape_limit && (Q->ne[1] < 32 || K->ne[1] < 1024)) { return false; } for (const ggml_tensor * t : {K, V}) { @@ -94,7 +106,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { return false; } // Prefill only. - if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) { + if (use_shape_limit && Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) { return false; } return true; @@ -240,9 +252,16 @@ 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); + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); + // Q: always f32 -- copy to dense f16. - 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); + std::optional> Qf_pool; + sycl::half * Qf_ptr = (sycl::half *) extra.Q_buffer_ptr; + if (!Qf_ptr) { + Qf_pool.emplace(ctx.pool(), (size_t) H * q * d); + Qf_ptr = Qf_pool->get(); + } + cont_to_f16_sycl((const char *) Q->data, Qf_ptr, d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream); // 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. @@ -253,11 +272,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso std::array v_str = k_str; std::optional> Kf_pool; std::optional> Vf_pool; + // Helper: hand out reserved space, or fall back to the pool. + auto stage_k = [&](size_t n) { if (extra.K_buffer_ptr) { return (sycl::half *) extra.K_buffer_ptr; } + Kf_pool.emplace(ctx.pool(), n); return Kf_pool->get(); }; + auto stage_v = [&](size_t n) { if (extra.V_buffer_ptr) { return (sycl::half *) extra.V_buffer_ptr; } + Vf_pool.emplace(ctx.pool(), n); return Vf_pool->get(); }; - 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]); @@ -266,22 +286,19 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso return std::array{ s3, s2, s2, s1, 1 }; }; - if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) { + if (ggml_sycl_fattn_onednn_binds_kv(K, 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); - cont_to_f16_sycl((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(); + K_ptr = stage_k((size_t) Hkv * seq * d); + V_ptr = stage_v((size_t) Hkv * seq * d); + cont_to_f16_sycl((const char *) K->data, K_ptr, d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); + cont_to_f16_sycl((const char *) V->data, V_ptr, d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); } 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(); + K_ptr = stage_k((size_t) ggml_nelements(K)); { 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; @@ -315,8 +332,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso // 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(); + V_ptr = stage_v((size_t) ggml_nelements(V)); { 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; @@ -347,12 +363,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso } } else { // F32: strided copy to dense F16 via cont_to_f16_sycl. - Kf_pool.emplace(ctx.pool(), ggml_nelements(K)); - K_ptr = Kf_pool->get(); + K_ptr = stage_k((size_t) ggml_nelements(K)); cont_to_f16_sycl((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(); + V_ptr = stage_v((size_t) ggml_nelements(V)); cont_to_f16_sycl((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); } @@ -366,11 +380,21 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso // instead -- the value is captured into the command, so no host memory has to outlive the // call, and the enqueue stays async. const sycl::half scale_h = (sycl::half) (1.0f / kq_scale); - ggml_sycl_pool_alloc scbuf(ctx.pool(), 1); - sycl::half * const scale_dev = scbuf.get(); + std::optional> scbuf; + sycl::half * scale_dev = (sycl::half *) extra.scale_buffer_ptr; + if (!scale_dev) { + scbuf.emplace(ctx.pool(), 1); + scale_dev = scbuf->get(); + } stream->single_task([=]() { *scale_dev = scale_h; }); - ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] + // f16 contiguous SDPA out [mb,H,q,d] + std::optional> outf_pool; + sycl::half * outf_ptr = (sycl::half *) extra.out_buffer_ptr; + if (!outf_ptr) { + outf_pool.emplace(ctx.pool(), (size_t) H * q * d); + outf_ptr = outf_pool->get(); + } // 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. @@ -392,7 +416,7 @@ 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_q) return Qf_ptr; if (r == E.id_k) return K_ptr; if (r == E.id_v) return V_ptr; if (r == E.id_scale) return scale_dev; @@ -404,10 +428,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso for (auto & lt : E.ins) { ti.emplace_back(lt, eng, id2ptr(lt.get_id())); } - tensor to(E.out, eng, outf.get()); + tensor to(E.out, eng, outf_ptr); E.cp.execute(strm, ti, {to}); - permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream); + permute_sdpa_out_sycl(outf_ptr, (float *) dst->data, mb, H, q, d, stream); // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA // serializes with the staging kernels before it and the permute/pool reuse after it. The // garbage output formerly blamed on the missing sync here was the scale use-after-return diff --git a/ggml/src/ggml-sycl/fattn-onednn.hpp b/ggml/src/ggml-sycl/fattn-onednn.hpp index d3019e8768..9669d1bd27 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.hpp +++ b/ggml/src/ggml-sycl/fattn-onednn.hpp @@ -5,7 +5,11 @@ // Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op // (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.) -bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst); +bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit = true); + +// True when the oneDNN path binds an F16 KV cache in place instead of staging a dense copy of +// it. Depends only on the types and strides of K and V, so the answer holds for every call. +bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V); // Run flash attention through oneDNN's fused xmx SDPA // execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure. diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index a85bca7cb3..b73e6d46ff 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -378,3 +378,76 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) { return ggml_sycl_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; } + +static uintptr_t ggml_sycl_fattn_reserve_halves(ggml_sycl_fattn_extra & extra, size_t n_halves) { + if (n_halves == 0) { + return 0; + } + extra.end = GGML_PAD(extra.end, SYCL_BUFFER_ALIGNMENT); + const uintptr_t block = extra.end; + extra.end += n_halves * sizeof(sycl::half); + return block; +} + +ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst) { + ggml_sycl_fattn_extra extra; + + extra.end = (uintptr_t) dst->data + ggml_nbytes(dst); + + if (dst->op != GGML_OP_FLASH_ATTN_EXT) { + return extra; + } + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + if (!Q || !K || !V) { + return extra; + } + + const int64_t d = K->ne[0]; + const int64_t H = Q->ne[2]; + const int64_t q = Q->ne[1]; + + // calculate the worst-case memory consumption across all kernels + const bool onednn_supported = ggml_sycl_flash_attn_ext_onednn_supported(dst, /* use_shape_limit */ false); + + const bool tile_needs_K = K->type != GGML_TYPE_F16; + const bool tile_needs_V = V->type != GGML_TYPE_F16; + + const bool V_is_K_view = V->view_src && + (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)); + + size_t need_K = 0, need_V = 0, need_Q = 0, need_out = 0, need_scale = 0; + if (onednn_supported) { + need_Q = (size_t) H * q * d; + need_out = (size_t) H * q * d; + need_scale = 1; + // an f16 cache is bound in place, so it needs no staging copy + if (!ggml_sycl_fattn_onednn_binds_kv(K, V)) { + need_K = (size_t) ggml_nelements(K); + need_V = (size_t) ggml_nelements(V); + } + } + if (tile_needs_K) { + need_K = std::max(need_K, (size_t) ggml_nelements(K)); + } + if (tile_needs_V) { + need_V = std::max(need_V, (size_t) ggml_nelements(V)); + } + + extra.Q_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_Q); + extra.K_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_K); + extra.V_buffer_ptr = (V_is_K_view && !onednn_supported && need_V) + ? extra.K_buffer_ptr + : ggml_sycl_fattn_reserve_halves(extra, need_V); + extra.scale_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_scale); + extra.out_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_out); + + return extra; +} + +size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst) { + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); + return (size_t) (extra.end - (uintptr_t) dst->data); +} diff --git a/ggml/src/ggml-sycl/fattn.hpp b/ggml/src/ggml-sycl/fattn.hpp index c093970a3f..f803aa2a80 100644 --- a/ggml/src/ggml-sycl/fattn.hpp +++ b/ggml/src/ggml-sycl/fattn.hpp @@ -19,6 +19,24 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst); +// Scratch that flash attention needs beyond the output tensor +struct ggml_sycl_fattn_extra { + uintptr_t K_buffer_ptr = 0; // F16 copy of the K cache + uintptr_t V_buffer_ptr = 0; // F16 copy of the V cache + uintptr_t Q_buffer_ptr = 0; // dense F16 copy of Q, oneDNN only + uintptr_t scale_buffer_ptr = 0; // the softmax scale as an F16 scalar, oneDNN only + uintptr_t out_buffer_ptr = 0; // F16 SDPA output before conversion to F32, oneDNN only + uintptr_t end = 0; // one past the last reserved byte; sizes the allocation +}; + +// ggml_sycl_fattn_get_extra() is the single source of truth for the layout: it both sizes +// the reservation and hands out the pointers, so the two cannot disagree. +// Each field is the address of one reserved block, or 0 if that block was not reserved, +// in which case the caller allocates from the scratch pool instead. +ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst); + +size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst); + void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst); #endif // GGML_SYCL_FATTN_HPP diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d83..dc8a174432 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -955,7 +955,10 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ } static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - size_t size = ggml_nbytes(tensor); + // Reserve the additional scratch so it's visible to the graph allocator + size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT + ? ggml_sycl_flash_attn_ext_get_alloc_size(tensor) + : ggml_nbytes(tensor); int64_t ne0 = tensor->ne[0]; if (ggml_is_quantized(tensor->type)) {