diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md index 8b0b9a18691a..baf4dafe9847 100644 --- a/docs/backend/SYCL.md +++ b/docs/backend/SYCL.md @@ -795,6 +795,9 @@ use 1 SYCL GPUs: [0] with Max compute units:512 | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).| | GGML_SYCL_DISABLE_DNN | 0 (default) or 1 | Disable running computations through oneDNN and always use oneMKL. | | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. | +| GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` | +| GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. | +| GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. | | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.
Recommended to use when --split-mode = layer | | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. | | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. | diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp index 8534bd3581e5..57d1ec974e21 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common.hpp @@ -63,6 +63,7 @@ extern int g_ggml_sycl_disable_optimize; extern int g_ggml_sycl_prioritize_dmmv; extern int g_ggml_sycl_enable_flash_attention; extern int g_ggml_sycl_dev2dev_memcpy; +extern int g_ggml_sycl_fa_onednn; #if defined(__clang__) && __has_builtin(__builtin_expect) diff --git a/ggml/src/ggml-sycl/fattn-hybrid.cpp b/ggml/src/ggml-sycl/fattn-hybrid.cpp new file mode 100644 index 000000000000..d78159cf6221 --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-hybrid.cpp @@ -0,0 +1,410 @@ +// +// MIT license +// Copyright (C) 2025 Intel Corporation +// SPDX-License-Identifier: MIT +// + +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// + +// Hybrid flash attention: dequantize K/V to F16 → oneDNN SDPA (XMX fused kernel). +// Combines the dequant pipeline from PR #25025 with the hardware SDPA from PR #25222. +// Covers quantized, BF16, and F32 KV caches that oneDNN's native F16-only gate rejects. + +#include +#include +#include +#include + +#include "fattn-hybrid.hpp" +#include "fattn.hpp" // for fallback: ggml_sycl_flash_attn_ext_mkl +#include "fattn-tile.hpp" +#include "convert.hpp" +#include "fattn-buffers.hpp" + +#ifdef GGML_SYCL_DNNL +#include "dnnl.hpp" +#include "dnnl_sycl.hpp" +#include "oneapi/dnnl/dnnl_graph.hpp" +#endif + +// --------------------------------------------------------------------------- +// Gate: eligible shapes for dequant + oneDNN SDPA +// --------------------------------------------------------------------------- +bool ggml_sycl_flash_attn_ext_hybrid_supported(const ggml_tensor * dst) { +#if !GGML_SYCL_DNNL + GGML_UNUSED(dst); + return false; +#else + // Respect the global kill switch + if (!g_ggml_sycl_fa_onednn) { + return false; + } + + const ggml_tensor * Q = dst->src[0]; + 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]; + + // HYBRID handles non-F16 types (native F16 goes to ONEDNN when merged). + if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { + return false; + } + + // MKL prefill gate conditions + if (Q->ne[1] < 32 || K->ne[1] < 1024) { + return false; + } + float max_bias = 0.0f, logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + if (!(Q->ne[3] == K->ne[3] || K->ne[3] == 1)) { + return false; + } + + // F16 stride alignment (from MKL gate) + for (const ggml_tensor * t : {K, V}) { + if (t->type == GGML_TYPE_F16 && t->nb[1] % (t->ne[0] * 2) != 0) { + return false; + } + } + + // SDPA conditions + if (!mask || mask->type != GGML_TYPE_F16 || mask->ne[2] != 1 || mask->ne[3] != 1) { + return false; + } + if (sinks) { + return false; + } + const int64_t d = K->ne[0]; + if (V->ne[0] != d || Q->ne[3] != 1) { + return false; + } + if (K->ne[2] == 0 || Q->ne[2] % K->ne[2] != 0) { + return false; + } + + return true; +#endif +} + +// --------------------------------------------------------------------------- +// Implementation +// --------------------------------------------------------------------------- +#if GGML_SYCL_DNNL + +#include "dnnl.hpp" +#include "dnnl_sycl.hpp" +#include "oneapi/dnnl/dnnl_graph.hpp" + +using namespace dnnl; +using namespace dnnl::graph; + +// Build + compile the contiguous-input GQA SDPA graph. +// MatMul → Divide → Add(mask) → SoftMax → MatMul, f16 out. +// Partitions=1, sdp_primitive_kernel_t (the systolic XMX kernel). +// Originally from fattn-onednn.cpp (PR #25222 by @hmscider). +sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) { + using ltype = logical_tensor::layout_type; + using dt = logical_tensor::data_type; + using ldims = logical_tensor::dims; + const dt fi = dt::f32, t = dt::f16; + 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}; + 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 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); + bmm1.add_inputs({query, key}); bmm1.add_outputs({score}); + + auto scale = logical_tensor(id++, t, sc, ltype::strided); + auto scaled = logical_tensor(id++, fi, s_sz, ltype::strided); + auto sdiv = op(id++, op::kind::Divide, "scale_div"); + sdiv.add_inputs({score, scale}); sdiv.add_outputs({scaled}); + + auto mask = logical_tensor(id++, t, msk, ltype::strided); + auto masked = logical_tensor(id++, fi, s_sz, ltype::strided); + auto madd = op(id++, op::kind::Add, "mask_add"); + madd.add_inputs({scaled, mask}); madd.add_outputs({masked}); + + auto probs = logical_tensor(id++, t, s_sz, ltype::strided); + auto smax = op(id++, op::kind::SoftMax, "softmax"); + smax.set_attr(op::attr::axis, -1); + 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); + // f16 output required to hit sdp_primitive_kernel_t. + auto output = logical_tensor(id++, t, o_sz, ltype::strided); + auto bmm2 = op(id++, op::kind::MatMul, "bmm2"); + bmm2.add_inputs({probs, value}); bmm2.add_outputs({output}); + + dnnl::graph::graph g(eng.get_kind()); + g.add_op(bmm1); g.add_op(sdiv); g.add_op(madd); g.add_op(smax); g.add_op(bmm2); + g.finalize(); + + auto parts = g.get_partitions(); + if (parts.size() != 1 || !parts[0].is_supported()) { + return E; + } + E.ins = parts[0].get_input_ports(); + E.out = parts[0].get_output_ports()[0]; + E.cp = parts[0].compile(E.ins, {E.out}, eng); + E.out = E.cp.query_logical_tensor(E.out.get_id()); + E.id_q = query.get_id(); E.id_k = key.get_id(); E.id_v = value.get_id(); + E.id_scale = scale.get_id(); E.id_mask = mask.get_id(); + E.ok = true; + return E; +} + +void ggml_sycl_flash_attn_ext_hybrid(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try { + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + ggml_tensor * KQV = dst; + + const int64_t d = K->ne[0]; + const int64_t seq = K->ne[1]; + const int64_t Hkv = K->ne[2]; + const int64_t H = Q->ne[2]; + const int64_t q = Q->ne[1]; + const int64_t mb = Q->ne[3]; + + float kq_scale = 1.0f; + memcpy(&kq_scale, (const float *) KQV->op_params + 0, sizeof(float)); + + dpct::queue_ptr stream = ctx.stream(); + dnnl::engine eng = ctx.engine_dnnl(stream); + dnnl::stream strm = ctx.stream_dnnl(stream); + + // --- Step 1: Dequant K to dense F16 --- + ggml_sycl_fattn_alloc K_f16(ctx.fattn_buffers().K); + { + const char * K_data = (const char *)K->data; + + // True Gemma interleave packs heads within a row (nb[2] < ne[1]*nb[1]). + // Padded seq-views have nb[2] > ne[1]*nb[1] — use physical strides. + // Both satisfy ne[1]*nb[1] != nb[2] but need different dequant paths. + 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]); + + K_f16.alloc(ggml_nelements(K)); + + if (K->type == GGML_TYPE_F16) { + if (!k_non_dense) { + stream->memcpy(K_f16.ptr, K_data, ggml_nelements(K) * sizeof(sycl::half)); + } else { + const int64_t ne0 = K->ne[0], ne1 = K->ne[1]; + const int64_t ne23 = (int64_t)K->ne[2] * K->ne[3]; + const int64_t src_nb1 = (int64_t)K->nb[1]; + const int64_t src_nb2 = (int64_t)K->nb[2]; + const sycl::half * src = (const sycl::half *)K_data; + sycl::half * dst = K_f16.ptr; + stream->parallel_for( + sycl::range<3>((size_t)ne23, (size_t)ne1, (size_t)ne0), + [=](sycl::item<3> it) { + int64_t hb = it.get_id(0), r = it.get_id(1), c = it.get_id(2); + const sycl::half * src_row = (const sycl::half *)( + (const char *)src + hb * src_nb2 + r * src_nb1); + dst[(hb * ne1 + r) * ne0 + c] = src_row[c]; + }); + } + } else 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_f16.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_f16.ptr, + K->ne[0], K->ne[1], K->ne[2], K->ne[3], + s01, s02, s03, stream); + } + } + + // --- Step 2: Dequant V to dense F16 --- + // K and V share buffers in quantized KV cache (V is a view of K). + // Detect it: same data pointer means reuse K_f16. + ggml_sycl_fattn_alloc V_f16(ctx.fattn_buffers().V); + bool V_is_K_view = (K->type != GGML_TYPE_F32 && V->type != GGML_TYPE_F32 && + K->data == V->data); + if (V_is_K_view) { + // V shares K's buffer — reuse the already-normalized dense K data + V_f16.ptr = K_f16.ptr; // don't allocate, just alias + } else { + const char * V_data = (const char *)V->data; + + // True Gemma interleave packs heads within a row (nb[2] < ne[1]*nb[1]). + // Padded seq-views have nb[2] > ne[1]*nb[1] — use physical strides. + // Both satisfy ne[1]*nb[1] != nb[2] but need different dequant paths. + 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]); + + V_f16.alloc(ggml_nelements(V)); + + if (V->type == GGML_TYPE_F16) { + if (!v_non_dense) { + stream->memcpy(V_f16.ptr, V_data, ggml_nelements(V) * sizeof(sycl::half)); + } else { + const int64_t ne0 = V->ne[0], ne1 = V->ne[1]; + const int64_t ne23 = (int64_t)V->ne[2] * V->ne[3]; + const int64_t src_nb1 = (int64_t)V->nb[1]; + const int64_t src_nb2 = (int64_t)V->nb[2]; + const sycl::half * src = (const sycl::half *)V_data; + sycl::half * dst = V_f16.ptr; + stream->parallel_for( + sycl::range<3>((size_t)ne23, (size_t)ne1, (size_t)ne0), + [=](sycl::item<3> it) { + int64_t hb = it.get_id(0), r = it.get_id(1), c = it.get_id(2); + const sycl::half * src_row = (const sycl::half *)( + (const char *)src + hb * src_nb2 + r * src_nb1); + dst[(hb * ne1 + r) * ne0 + c] = src_row[c]; + }); + } + } else 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_f16.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_f16.ptr, + V->ne[0], V->ne[1], V->ne[2], V->ne[3], + s01, s02, s03, stream); + } + } + + // --- Step 3: Copy Q to dense F16 --- + ggml_sycl_pool_alloc Qf(ctx.pool(), (size_t) H * q * d); + { + const char * Q_data = (const char *)Q->data; + const int64_t n = H * q * d; + sycl::half * Qf_ptr = Qf.get(); + size_t Q_nb1 = Q->nb[1]; + size_t Q_nb2 = Q->nb[2]; + size_t Q_nb3 = Q->nb[3]; + stream->parallel_for(sycl::range<1>(n), [=](sycl::id<1> ix) { + const int64_t gid = ix[0]; + int64_t i = gid; + const int64_t i0 = i % d; i /= d; + const int64_t i1 = i % q; i /= q; + const int64_t i2 = i % H; + const int64_t i3 = i / H; + const float * p = (const float *) (Q_data + i1 * Q_nb1 + i2 * Q_nb2 + i3 * Q_nb3) + i0; + Qf_ptr[gid] = (sycl::half) (*p); + }); + } + + // --- Step 4: Build/run oneDNN SDPA --- + const sycl::half scale_h = (sycl::half) (1.0f / kq_scale); + ggml_sycl_pool_alloc scbuf(ctx.pool(), 1); + stream->memcpy(scbuf.get(), &scale_h, sizeof(sycl::half)); + + ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); + + // Compile once per (device, shape), reuse across layers. + 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); + 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; + } + sdpa_partition & E = it->second; + if (!E.ok) { + // Partition failed — fall back to MKL + ggml_sycl_flash_attn_ext_mkl(ctx, dst); + return; + } + + auto id2ptr = [&](size_t r) -> void * { + if (r == E.id_q) return Qf.get(); + if (r == E.id_k) return K_f16.ptr; + if (r == E.id_v) return V_f16.ptr; + if (r == E.id_scale) return scbuf.get(); + if (r == E.id_mask) return (void *) mask->data; + return nullptr; + }; + std::vector ti; + ti.reserve(E.ins.size()); + for (auto & lt : E.ins) { + ti.emplace_back(lt, eng, id2ptr(lt.get_id())); + } + tensor to(E.out, eng, outf.get()); + E.cp.execute(strm, ti, {to}); + + // --- Step 5: Permute SDPA output → ggml dst --- + const int64_t n = mb * H * q * d; + float * dst_data = (float *) dst->data; + const sycl::half * out_data = outf.get(); + const size_t dst_nb0 = sizeof(float); + const size_t dst_nb1 = dst->nb[1]; + const size_t dst_nb2 = dst->nb[2]; + const size_t dst_nb3 = dst->nb[3]; + stream->parallel_for(sycl::range<1>(n), [=](sycl::id<1> ix) { + const int64_t gid = ix[0]; + int64_t i = gid; + const int64_t e = i % d; i /= d; + const int64_t t = i % q; i /= q; + const int64_t h = i % H; const int64_t b = i / H; + int64_t off = (e * dst_nb0 + h * dst_nb1 + t * dst_nb2 + b * dst_nb3) / sizeof(float); + dst_data[off] = (float) out_data[gid]; + }); + + // Wait for all GPU work to finish before pool allocs go out of scope. + // Without this, the pool frees and reuses memory while the GPU is still + // reading it, corrupting the output. + stream->wait(); +} +catch (const std::exception & e) { + GGML_LOG_WARN("%s: hybrid SDPA failed (%s); falling back to MKL\n", __func__, e.what()); + ggml_sycl_flash_attn_ext_mkl(ctx, dst); +} + +#else // !GGML_SYCL_DNNL + +void ggml_sycl_flash_attn_ext_hybrid(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + // oneDNN not compiled — fall back to MKL + ggml_sycl_flash_attn_ext_mkl(ctx, dst); +} + +#endif // GGML_SYCL_DNNL diff --git a/ggml/src/ggml-sycl/fattn-hybrid.hpp b/ggml/src/ggml-sycl/fattn-hybrid.hpp new file mode 100644 index 000000000000..990cd6cc51b4 --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-hybrid.hpp @@ -0,0 +1,43 @@ +// +// MIT license +// Copyright (C) 2025 Intel Corporation +// SPDX-License-Identifier: MIT +// + +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// + +#ifndef GGML_SYCL_FATTN_HYBRID_HPP +#define GGML_SYCL_FATTN_HYBRID_HPP + +#include "common.hpp" + +// Check whether this FA call qualifies for the hybrid path: +// dequantize K/V to f16, then route through oneDNN's fused SDPA kernel. +bool ggml_sycl_flash_attn_ext_hybrid_supported(const ggml_tensor * dst); + +// Run FA: dequant K/V → oneDNN SDPA → permute to ggml dst. +// Falls back to MKL GEMM on failure. +void ggml_sycl_flash_attn_ext_hybrid(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +#if GGML_SYCL_DNNL +#include "dnnl.hpp" +#include "dnnl_sycl.hpp" +#include "oneapi/dnnl/dnnl_graph.hpp" + +// Compiled SDPA partition — caches the fused MatMul+SoftMax kernel. +struct sdpa_partition { + dnnl::graph::compiled_partition cp; + std::vector ins; + dnnl::graph::logical_tensor out; + size_t id_q = 0, id_k = 0, id_v = 0, id_scale = 0, id_mask = 0; + bool ok = false; +}; + +sdpa_partition build_sdpa(const dnnl::engine & eng, int H, int Hkv, int q, int seq, int d); +#endif + +#endif // GGML_SYCL_FATTN_HYBRID_HPP diff --git a/ggml/src/ggml-sycl/fattn-mkl.cpp b/ggml/src/ggml-sycl/fattn-mkl.cpp new file mode 100644 index 000000000000..6fc719fcddd3 --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-mkl.cpp @@ -0,0 +1,711 @@ +// Flash attention via oneMKL GEMM (XMX-accelerated). +// Uses column_major::gemm for Q*K^T and S*V matmuls +// with an online softmax SYCL kernel. +// +// All GQA query heads sharing a KV head are batched into single +// GEMM calls, amortizing MKL launch overhead across K and V reuse. +// + +#include "common.hpp" +#include "fattn-common.hpp" +#include "fattn-buffers.hpp" +#include "convert.hpp" +#include "fattn.hpp" + +#include +#include +#include + +#define MKL_FA_CHUNK_SIZE_KV 8192 + +using oneapi::mkl::transpose; +using oneapi::mkl::blas::column_major::gemm; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Pack all GQA Q heads for one KV head into fp16, applying q_scale. +// Launches one kernel per GQA group — each kernel copies exactly +// n_queries * DKQ elements using the per-group dst offset and +// per-head source stride. +static void mkl_fa_pack_q_fp16( + dpct::queue_ptr stream, + sycl::half * __restrict dst, + const float * __restrict q_src, + int n_queries, int n_query_rows, int DKQ, + int gqa_ratio, int kvh_base_head, + float q_scale, int64_t q_row_stride, int64_t q_head_stride, + int64_t wg_size) { + + for (int iqg = 0; iqg < gqa_ratio; iqg++) { + int iqh = kvh_base_head + iqg; + sycl::half * dst_g = dst + (int64_t)iqg * n_queries * DKQ; + + const int64_t n_elem = (int64_t)n_queries * DKQ; + const int64_t wg = ((n_elem + wg_size - 1) / wg_size) * wg_size; + + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<1>(wg, wg_size), + [=](sycl::nd_item<1> item) { + int64_t e = item.get_global_id(0); + if (e >= n_elem) return; + + int64_t q = e / DKQ; + int64_t d = e - q * DKQ; + + // Stride-aware source offset: handles permuted, + // sliced, or contiguous Q tensor layouts. + int64_t src_off = d + + q * q_row_stride + + (int64_t)iqh * q_head_stride; + + dst_g[e] = sycl::half( + q_src[src_off] * q_scale); + }); + }); + } +} + +// Zero-initialize the online softmax state arrays. +// KQ_max → -inf, KQ_sum → 0, VKQ_accum → 0. +// Merged into one kernel to avoid per-array launch overhead. +static void mkl_fa_init_softmax_state( + dpct::queue_ptr stream, + float * kmax, float * ksum, float * vacc, + int n_query_rows, int DV, int64_t wg_size) { + + const float neg_inf = -1e30f; + const int64_t n_maxsum = n_query_rows; + const int64_t n_vacc = (int64_t)n_query_rows * DV; + const int64_t total = (n_vacc > n_maxsum) ? n_vacc : n_maxsum; + const int64_t wg = ((total + wg_size - 1) / wg_size) * wg_size; + + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<1>(wg, wg_size), + [=](sycl::nd_item<1> item) { + int64_t i = item.get_global_id(0); + if (i < n_maxsum) { + kmax[i] = neg_inf; + ksum[i] = 0.0f; + } + if (i < n_vacc) { + vacc[i] = 0.0f; + } + }); + }); +} + +// Online softmax over one KV chunk for all GQA query rows. +// For each row: find local max → rescale previous VKQ_accum → +// compute exp(s - max) → write S_f16 → update running max/sum. +static void mkl_fa_online_softmax_chunk( + dpct::queue_ptr stream, + float * __restrict KQ_f32, + sycl::half * __restrict S_f16, + float * __restrict KQ_max, + float * __restrict KQ_sum, + float * __restrict VKQ_accum, + int n_query_rows, int n_queries, int DV, + int chunk_size, int chunk_start, + int kvh_head, int gqa_ratio, + const sycl::half * mask_data, int64_t mask_head_stride, + int64_t mask_row_stride, int mask_n_heads, + float logit_softcap, int64_t wg_size) { + + const int64_t wg = ((n_query_rows + wg_size - 1) / wg_size) * wg_size; + + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<1>(wg, wg_size), + [=](sycl::nd_item<1> item) { + int jc = item.get_global_id(0); + if (jc >= n_query_rows) return; + + const int gqa_group = jc / n_queries; + const int q_row = jc % n_queries; + + const float * __restrict KQ_row = KQ_f32 + + jc * (int64_t)chunk_size; + float * __restrict vkq = VKQ_accum + + jc * (int64_t)DV; + + const sycl::half * mask_h = nullptr; + int64_t m_stride = 0; + if (mask_data) { + int m_head = (mask_n_heads > 1) + ? (kvh_head + gqa_group) : 0; + mask_h = mask_data + (int64_t)m_head * mask_head_stride; + m_stride = mask_row_stride; + } + + // Row-wise local maximum (softcap before mask) + float local_max = -1e30f; + for (int i = 0; i < chunk_size; i++) { + float s = KQ_row[i]; + if (logit_softcap != 0.0f) { + s = logit_softcap * sycl::tanh(s); + } + if (mask_h) { + s += (float)mask_h[q_row * m_stride + + (chunk_start + i)]; + } + if (s > local_max) local_max = s; + } + + // Rescale previous accumulator by exp(old_max - new_max) + float old_max = KQ_max[jc]; + float new_max = (old_max > local_max) ? old_max : local_max; + float rescale = (old_max < -1e29f) ? 1.0f + : sycl::native::exp(old_max - new_max); + + for (int v = 0; v < DV; v++) { + vkq[v] *= rescale; + } + + // Softmax and write S_f16 + float local_sum = 0.0f; + sycl::half * __restrict S_row = S_f16 + + jc * (int64_t)chunk_size; + + for (int i = 0; i < chunk_size; i++) { + float s = KQ_row[i]; + if (logit_softcap != 0.0f) { + s = logit_softcap * sycl::tanh(s); + } + if (mask_h) { + s += (float)mask_h[q_row * m_stride + + (chunk_start + i)]; + } + float val = sycl::native::exp(s - new_max); + S_row[i] = sycl::half(val); + local_sum += val; + } + + KQ_sum[jc] = KQ_sum[jc] * rescale + local_sum; + KQ_max[jc] = new_max; + }); + }); +} + +// Write one GQA group's normalized output to its destination head. +static void mkl_fa_normalize_head( + dpct::queue_ptr stream, + float * __restrict dst_batch, + const float * __restrict VKQ_accum, + const float * __restrict KQ_sum, + int iqh, int n_queries, int DV, int n_q_heads, + int64_t src_offset, int64_t wg_size) { + + const int64_t wg = ((n_queries + wg_size - 1) / wg_size) * wg_size; + + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<1>(wg, wg_size), + [=](sycl::nd_item<1> item) { + int jc = item.get_global_id(0); + if (jc >= n_queries) return; + + int ksum_idx = (int)(src_offset / DV) + jc; + float inv_sum = 1.0f / KQ_sum[ksum_idx]; + const float * __restrict src = VKQ_accum + + src_offset + jc * (int64_t)DV; + // Interleaved dst layout (matching TILE): + // rows alternate between heads, then increment query. + // offset = (query * n_q_heads + head) * DV + float * __restrict dst_row = dst_batch + + ((int64_t)jc * n_q_heads + iqh) * (int64_t)DV; + + for (int v = 0; v < DV; v++) { + dst_row[v] = src[v] * inv_sum; + } + }); + }); +} + +// --------------------------------------------------------------------------- +// MKL Flash Attention orchestrator +// +// Pipeline: dequantize K/V → for each KV head: +// pack GQA Q heads → MKL GEMM KQ → online softmax → +// MKL GEMM VKQ → accumulate → normalize → scatter to dst +// --------------------------------------------------------------------------- +void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + ggml_tensor * KQV = dst; + + GGML_ASSERT(Q->type == GGML_TYPE_F32); + GGML_ASSERT(KQV->type == GGML_TYPE_F32); + + // --- Op params --- + float scale = 1.0f, max_bias = 0.0f, logit_softcap = 0.0f; + memcpy(&scale, (const float *)KQV->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *)KQV->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *)KQV->op_params + 2, sizeof(float)); + + const float q_scale = scale; + + // --- Dimensions --- + const int DKQ = (int)K->ne[0]; + const int DV = (int)V->ne[0]; + const int n_queries = (int)Q->ne[1]; + const int n_q_heads = (int)Q->ne[2]; + const int n_kv_heads = (int)K->ne[2]; + const int n_batch = (int)Q->ne[3]; + const int n_kv = (int)K->ne[1]; + const int gqa_ratio = n_q_heads / n_kv_heads; + const int n_query_rows = n_queries * gqa_ratio; + + GGML_ASSERT(n_q_heads % n_kv_heads == 0); + GGML_ASSERT(max_bias == 0.0f); // ALiBi not supported + GGML_ASSERT(Q->ne[3] == K->ne[3] || K->ne[3] == 1); + + const int chunk_size = std::min(MKL_FA_CHUNK_SIZE_KV, n_kv); + const int64_t wg_size = 256; + + // --- Debug output (gated by GGML_SYCL_MKL_FA_DEBUG=1) --- + static int mkl_call_count = 0; + mkl_call_count++; + static int mkl_debug = -1; + if (mkl_debug < 0) { + mkl_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0); + } + const bool do_print = (mkl_debug == 1); + + const int64_t q_row_stride = Q->nb[1] / sizeof(float); + const int64_t q_head_stride = Q->nb[2] / sizeof(float); + + 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)); + + // Early interleaved detection for debug output. + // True interleaved detection happens after dequant (nb12_fp16 == nb11_fp16), + // but we can pre-detect on the original tensor strides. + const bool k_early_interleaved = + ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]); + const bool v_early_interleaved = + !V_is_K_view && ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]); + + if (do_print) { + GGML_LOG_INFO("[MKL-FA] #%d D=%d DV=%d n_q=%d n_kv=%d " + "n_qh=%d n_kvh=%d gqa=%d batch=%d K=%s V=%s " + "chunk=%d buf=%.1fMB%s%s\n", + mkl_call_count, DKQ, DV, n_queries, n_kv, + n_q_heads, n_kv_heads, gqa_ratio, n_batch, + ggml_type_name(K->type), ggml_type_name(V->type), + chunk_size, + (double)((int64_t)n_query_rows * chunk_size * sizeof(float)) + / (1024.0 * 1024.0), + k_early_interleaved ? " K_ILV" : "", + v_early_interleaved ? " V_ILV" : ""); + GGML_LOG_INFO("[MKL-FA] #%d Q-nb1=%lld Q-nb2=%lld " + "q_rs=%lld q_hs=%lld dst_rs=%lld dst_hs=%lld\n", + mkl_call_count, + (long long)Q->nb[1], (long long)Q->nb[2], + (long long)q_row_stride, (long long)q_head_stride, + (long long)(KQV->nb[1] / sizeof(float)), + (long long)(KQV->nb[2] / sizeof(float))); + } + + // --- Stream and allocators --- + dpct::queue_ptr stream = ctx.stream(); + + ggml_sycl_fattn_kv_buffers & fbuf = ctx.fattn_buffers(); + +#define MKL_TAKE_TIME(t0) auto t0 = std::chrono::steady_clock::now() +#define MKL_ACCUM(acc, t0) do { if (do_print) { \ + acc += (int64_t)std::chrono::duration_cast \ + (std::chrono::steady_clock::now() - (t0)).count(); \ +} } while(0) + + int64_t gemm_kq_time_us = 0; + int64_t gemm_vkq_time_us = 0; + int64_t softmax_time_us = 0; + int64_t dequant_time_us = 0; + + MKL_TAKE_TIME(t_deq); + + // --- Dequantize K, V to fp16 --- + // + // Output is ALWAYS dense row-major (contiguous rows per head). + // For interleaved source layouts (Gemma: nb[2]==nb[1], heads + // interleave at the row level), we reconstruct the PHYSICAL + // source strides so to_fp16_nc writes the correct data into a + // dense buffer. The logical view strides (nb##/ts) are wrong + // for interleaved layouts — e.g. s01=8 aliases Row1/Head0 to + // the same blocks as Row0/Head1. + // + // Detection: ne[1]*nb[1] != nb[2] means heads are interleaved. + ggml_sycl_fattn_alloc K_f16(fbuf.K); + ggml_sycl_fattn_alloc V_f16(fbuf.V); + + const char * K_data = (const char *)K->data; + size_t nb11 = K->nb[1]; + size_t nb12 = K->nb[2]; + size_t nb13 = K->nb[3]; + + const char * V_data = (const char *)V->data; + size_t nb21 = V->nb[1]; + size_t nb22 = V->nb[2]; + size_t nb23 = V->nb[3]; + + const bool k_interleaved = + ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1; + const bool v_interleaved = + ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1; + + { + // Normalize K to dense row-major F16 for GEMM. + // F16 is copied via memcpy (strides guaranteed standard by the gate). + K_f16.alloc(ggml_nelements(K)); + + if (K->type == GGML_TYPE_F16) { + // F16: always copy to dense for GEMM. + // Standard strides guaranteed by gate (nb1 == ne0 * sizeof(half)). + if (!k_interleaved) { + // Dense row-major F16 — flat memcpy + stream->memcpy(K_f16.ptr, K_data, ggml_nelements(K) * sizeof(sycl::half)); + } else { + // Interleaved F16 — strided copy to dense + const int64_t ne0 = K->ne[0]; + const int64_t ne1 = K->ne[1]; + const int64_t ne23 = (int64_t)K->ne[2] * K->ne[3]; + const int64_t src_nb1 = (int64_t)nb11; + const int64_t src_nb2 = (int64_t)nb12; + const sycl::half * src = (const sycl::half *)K_data; + sycl::half * dst = K_f16.ptr; + + stream->parallel_for( + sycl::range<3>((size_t)ne23, (size_t)ne1, (size_t)ne0), + [=](sycl::item<3> it) { + int64_t hb = it.get_id(0); + int64_t r = it.get_id(1); + int64_t c = it.get_id(2); + const sycl::half * src_row = (const sycl::half *)( + (const char *)src + hb * src_nb2 + r * src_nb1); + dst[(hb * ne1 + r) * ne0 + c] = src_row[c]; + }); + + nb11 = K->ne[0] * sizeof(sycl::half); + nb12 = K->ne[1] * nb11; + nb13 = K->ne[2] * nb12; + } + } else if (ggml_is_contiguously_allocated(K) && !k_interleaved) { + const size_t bs = ggml_blck_size(K->type); + const size_t ts = ggml_type_size(K->type); + to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst); + to_fp16(K_data, K_f16.ptr, ggml_nelements(K), stream); + + nb11 = nb11 * bs * sizeof(sycl::half) / ts; + nb12 = nb12 * bs * sizeof(sycl::half) / ts; + nb13 = nb13 * bs * sizeof(sycl::half) / ts; + } 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_interleaved) { + const int64_t blk_per_row = (int64_t)K->ne[0] / bs; + s01 = (int64_t)n_kv_heads * blk_per_row; + s02 = blk_per_row; + s03 = (int64_t)K->ne[1] * s01; + } else { + s01 = nb11 / ts; + s02 = nb12 / ts; + 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, stream); + + nb11 = K->ne[0] * sizeof(sycl::half); + nb12 = K->ne[1] * nb11; + nb13 = K->ne[2] * nb12; + } + K_data = (char *) K_f16.ptr; + } + + { + // Normalize V to dense row-major F16 for GEMM. + if (V_is_K_view) { + // V shares K's buffer — reuse the already-normalized dense K data + V_data = K_data; + nb21 = nb11; + nb22 = nb12; + nb23 = nb13; + } else { + V_f16.alloc(ggml_nelements(V)); + + if (V->type == GGML_TYPE_F16) { + if (!v_interleaved) { + // Dense row-major F16 — flat memcpy + stream->memcpy(V_f16.ptr, V_data, ggml_nelements(V) * sizeof(sycl::half)); + } else { + // Interleaved F16 — strided copy to dense + const int64_t ne0 = V->ne[0]; + const int64_t ne1 = V->ne[1]; + const int64_t ne23 = (int64_t)V->ne[2] * V->ne[3]; + const int64_t src_nb1 = (int64_t)nb21; + const int64_t src_nb2 = (int64_t)nb22; + const sycl::half * src = (const sycl::half *)V_data; + sycl::half * dst = V_f16.ptr; + + stream->parallel_for( + sycl::range<3>((size_t)ne23, (size_t)ne1, (size_t)ne0), + [=](sycl::item<3> it) { + int64_t hb = it.get_id(0); + int64_t r = it.get_id(1); + int64_t c = it.get_id(2); + const sycl::half * src_row = (const sycl::half *)( + (const char *)src + hb * src_nb2 + r * src_nb1); + dst[(hb * ne1 + r) * ne0 + c] = src_row[c]; + }); + + nb21 = V->ne[0] * sizeof(sycl::half); + nb22 = V->ne[1] * nb21; + nb23 = V->ne[2] * nb22; + } + V_data = (char *) V_f16.ptr; + } else if (ggml_is_contiguously_allocated(V) && !v_interleaved) { + const size_t bs = ggml_blck_size(V->type); + const size_t ts = ggml_type_size(V->type); + to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst); + to_fp16(V_data, V_f16.ptr, ggml_nelements(V), stream); + V_data = (char *) V_f16.ptr; + + nb21 = nb21 * bs * sizeof(sycl::half) / ts; + nb22 = nb22 * bs * sizeof(sycl::half) / ts; + nb23 = nb23 * bs * sizeof(sycl::half) / ts; + } 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_interleaved) { + 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 = nb21 / ts; + s02 = nb22 / ts; + 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, stream); + + nb21 = V->ne[0] * sizeof(sycl::half); + nb22 = V->ne[1] * nb21; + nb23 = V->ne[2] * nb22; + + V_data = (char *) V_f16.ptr; + } + } + } + + MKL_ACCUM(dequant_time_us, t_deq); + + // All paths produce dense row-major fp16 — GEMM lda = D. + // Head pointers use dense strides (nb12 = ne[1]*ne[0]*sizeof(half)). + const int64_t nb12_fp16 = nb12 / sizeof(sycl::half); + const int64_t nb22_fp16 = nb22 / sizeof(sycl::half); + + // --- Resolve mask pointers --- + const sycl::half * mask_data = nullptr; + int64_t mask_head_stride = 0; + int64_t mask_row_stride = 0; + int mask_n_heads = 0; + + if (mask) { + // Use actual fp16 device size (2 bytes), NOT sizeof(sycl::half) + // which may be 4 on the host in oneAPI. + mask_head_stride = mask->nb[2] / 2; + mask_row_stride = mask->nb[1] / 2; + mask_n_heads = (int)mask->ne[2]; + } + + // --- Allocate intermediates from pool --- + ggml_sycl_pool & pool = ctx.pool(); + + ggml_sycl_pool_alloc KQ_f32(pool); + ggml_sycl_pool_alloc S_f16(pool); + ggml_sycl_pool_alloc VKQ_chunk(pool); + ggml_sycl_pool_alloc VKQ_accum(pool); + ggml_sycl_pool_alloc KQ_max(pool); + ggml_sycl_pool_alloc KQ_sum(pool); + ggml_sycl_pool_alloc Q_head_f16(pool); + + KQ_f32.alloc((size_t)n_query_rows * chunk_size); + S_f16.alloc((size_t)n_query_rows * chunk_size); + VKQ_chunk.alloc((size_t)n_query_rows * DV); + VKQ_accum.alloc((size_t)n_query_rows * DV); + KQ_max.alloc(n_query_rows); + KQ_sum.alloc(n_query_rows); + Q_head_f16.alloc((size_t)n_query_rows * DKQ); + + sycl::half * Q_head_f16_ptr = Q_head_f16.ptr; + float * KQ_f32_ptr = KQ_f32.ptr; + sycl::half * S_f16_ptr = S_f16.ptr; + float * VKQ_chunk_ptr = VKQ_chunk.ptr; + float * VKQ_accum_ptr = VKQ_accum.ptr; + float * KQ_max_ptr = KQ_max.ptr; + float * KQ_sum_ptr = KQ_sum.ptr; + + const float alpha = 1.0f; + const float beta = 0.0f; + + for (int ib = 0; ib < n_batch; ib++) { + const float * Q_batch = (const float *)Q->data + + ib * (Q->nb[3] / sizeof(float)); + float * dst_batch = (float *)KQV->data + + ib * (KQV->nb[3] / sizeof(float)); + + const sycl::half * mask_batch = nullptr; + if (mask) { + int m_batch = (mask->ne[3] > 1) ? ib : 0; + mask_batch = (const sycl::half *)mask->data + + m_batch * (mask->nb[3] / 2); // 2 = actual fp16 device size + } + + for (int ikvh = 0; ikvh < n_kv_heads; ikvh++) { + int kvh_base_head = ikvh * gqa_ratio; + + // 1. Pack all GQA Q heads into fp16 + mkl_fa_pack_q_fp16(stream, + Q_head_f16_ptr, Q_batch, + n_queries, n_query_rows, DKQ, + gqa_ratio, kvh_base_head, + q_scale, q_row_stride, q_head_stride, wg_size); + + // 2. Initialize softmax state + mkl_fa_init_softmax_state(stream, + KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr, + n_query_rows, DV, wg_size); + + // Sync before MKL GEMM (MKL may use an internal queue) + stream->wait(); + + // K/V base for this KV head. + // Dense row-major: nb12_fp16 = n_kv * DKQ (head stride in half-elements). + const sycl::half * K_head = (const sycl::half *)K_data + + ikvh * nb12_fp16; + const sycl::half * V_head = (const sycl::half *)V_data + + ikvh * nb22_fp16; + + // 3. Chunked KV loop + for (int chunk_start = 0; chunk_start < n_kv; chunk_start += chunk_size) { + int this_chunk = std::min(chunk_size, n_kv - chunk_start); + const sycl::half * K_chunk = K_head + + (int64_t)chunk_start * DKQ; + const sycl::half * V_chunk = V_head + + (int64_t)chunk_start * DV; + + // 3a. GEMM: KQ = Q_batched × K_chunk^T + { + MKL_TAKE_TIME(t0); + sycl::event ev = gemm(*stream, + transpose::trans, transpose::nontrans, + this_chunk, n_query_rows, DKQ, + alpha, + K_chunk, DKQ, + Q_head_f16_ptr, DKQ, + beta, + KQ_f32_ptr, this_chunk); + try { ev.wait_and_throw(); } catch (sycl::exception & e) { + GGML_LOG_INFO("[MKL-FA] GEMM KQ: %s\n", e.what()); + GGML_ABORT("MKL GEMM KQ failed"); + } + MKL_ACCUM(gemm_kq_time_us, t0); + } + // 3b. Online softmax over this chunk + { + MKL_TAKE_TIME(t0); + mkl_fa_online_softmax_chunk(stream, + KQ_f32_ptr, S_f16_ptr, + KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr, + n_query_rows, n_queries, DV, + this_chunk, chunk_start, + kvh_base_head, gqa_ratio, + mask_batch, mask_head_stride, + mask_row_stride, mask_n_heads, + logit_softcap, wg_size); + stream->wait(); // S_f16 must be ready for GEMM + MKL_ACCUM(softmax_time_us, t0); + } + + // 3c. GEMM: VKQ_chunk = S × V_chunk + { + MKL_TAKE_TIME(t0); + sycl::event ev = gemm(*stream, + transpose::nontrans, transpose::nontrans, + DV, n_query_rows, this_chunk, + alpha, + V_chunk, DV, + S_f16_ptr, this_chunk, + beta, + VKQ_chunk_ptr, DV); + try { ev.wait_and_throw(); } catch (sycl::exception & e) { + GGML_LOG_INFO("[MKL-FA] GEMM VKQ: %s\n", e.what()); + GGML_ABORT("MKL GEMM VKQ failed"); + } + MKL_ACCUM(gemm_vkq_time_us, t0); + } + // 3d. VKQ_accum += VKQ_chunk + { + const int64_t n_total = (int64_t)n_query_rows * DV; + const int64_t wg = ((n_total + wg_size - 1) / wg_size) + * wg_size; + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<1>(wg, wg_size), + [=](sycl::nd_item<1> item) { + int64_t i = item.get_global_id(0); + if (i < n_total) { + VKQ_accum_ptr[i] += VKQ_chunk_ptr[i]; + } + }); + }); + } + } + + // 4. Normalize and scatter each GQA head to dst + for (int iqg = 0; iqg < gqa_ratio; iqg++) { + int iqh = kvh_base_head + iqg; + int64_t src_offset = (int64_t)iqg * n_queries * DV; + mkl_fa_normalize_head(stream, + dst_batch, VKQ_accum_ptr, KQ_sum_ptr, + iqh, n_queries, DV, n_q_heads, + src_offset, wg_size); + } + } + } + +#undef MKL_TAKE_TIME +#undef MKL_ACCUM + + if (do_print) { + double total_mb = (double)( + (int64_t)n_query_rows * chunk_size * sizeof(float) + + (int64_t)n_query_rows * chunk_size * sizeof(sycl::half) + + (int64_t)n_query_rows * DV * sizeof(float) + + (int64_t)n_query_rows * DV * sizeof(float) + + (int64_t)n_query_rows * sizeof(float) + + (int64_t)n_query_rows * sizeof(float) + + (int64_t)n_query_rows * DKQ * sizeof(sycl::half) + ) / (1024.0 * 1024.0); + GGML_LOG_INFO("[MKL-FA] #%d n_kv=%d n_q=%d time_us: " + "dequant=%lld GEMM_KQ=%lld softmax=%lld GEMM_VKQ=%lld " + "buf_mb=%.1f\n", + mkl_call_count, n_kv, n_queries, + (long long)dequant_time_us, + (long long)gemm_kq_time_us, + (long long)softmax_time_us, + (long long)gemm_vkq_time_us, + total_mb); + } +} diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index 7c6e6112fdcd..988149c1225f 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -18,6 +18,8 @@ #include "fattn-tile.hpp" #include "fattn-vec.hpp" #include "fattn.hpp" +// #include "fattn-onednn.hpp" // from PR #25222 — add when merged +#include "fattn-hybrid.hpp" #define FATTN_VEC_CASE(D, type_K, type_V) \ @@ -96,9 +98,13 @@ 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, // native F16 SDPA — reserved for PR #25222 BEST_FATTN_KERNEL_TILE = 200, + BEST_FATTN_KERNEL_HYBRID = 250, // dequant K/V → oneDNN SDPA + BEST_FATTN_KERNEL_MKL = 300, }; + static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) { GGML_UNUSED(device); #ifndef SYCL_FLASH_ATTN @@ -113,6 +119,7 @@ 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); @@ -120,7 +127,53 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const 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; + + // MKL path: XMX-accelerated GEMM for prompt processing (all KV cache types). + // The MKL kernel converts non-F16 K/V to F16 via to_fp16_sycl before GEMM, + // so quantized, F16, BF16, and F32 caches all benefit from XMX acceleration. + // Activates automatically when flash-attn is enabled (--flash-attn on or -fa) + // and n_kv >= 1024. Falls through to TILE/VEC for ALiBi, logit softcap, + // and mismatched batch dimensions (unsupported by the MKL kernel). + // Set GGML_SYCL_ENABLE_MKL_FA=0 to force TILE/VEC path for A/B testing. + // Example: GGML_SYCL_ENABLE_MKL_FA=0 llama-cli -m model.gguf -fa -ngl 99 ... + // Note: MKL GEMM calls are incompatible with SYCL graph capture replay. + + // XMX-accelerated paths (prefill only, Q >= 32). + // Dispatch order: HYBRID (dequant→SDPA) → MKL GEMM. + // When PR #25222 merges, add ONEDNN path above HYBRID: + // if (ggml_sycl_flash_attn_ext_onednn_supported(dst)) return BEST_FATTN_KERNEL_ONEDNN; + + // Path 1: Hybrid — dequantize K/V to F16, then oneDNN SDPA. + if (ggml_sycl_flash_attn_ext_hybrid_supported(dst)) { + return BEST_FATTN_KERNEL_HYBRID; + } + + // Path 2: MKL GEMM — handles SDPA-incompatible shapes (no mask, ALiBi, etc.). + static int mkl_enable = ggml_sycl_get_env("GGML_SYCL_ENABLE_MKL_FA", 1); + if (mkl_enable == 1 && Q->ne[0] >= 64 && mask && !sinks && + Q->ne[1] >= 32 && K->ne[1] >= 1024 && + max_bias == 0.0f && logit_softcap == 0.0f && + (Q->ne[3] == K->ne[3] || K->ne[3] == 1)) { + // F16 K/V strides must be a multiple of ne[0]*2 (the natural row size + // in bytes). This passes both dense (nb1 == ne0*2) and interleaved + // (nb1 == H * ne0*2). Only pathological test strides like nb1=32 or + // nb1=75 for ne0=40 fall through to TILE. + bool kv_strides_ok = true; + for (const ggml_tensor * t : {K, V}) { + if (t->type == GGML_TYPE_F16 && t->nb[1] % (t->ne[0] * 2) != 0) { + kv_strides_ok = false; + break; + } + } + if (kv_strides_ok) { + return BEST_FATTN_KERNEL_MKL; + } + } + for (const ggml_tensor * t : {Q, K, V, mask}) { if (t == nullptr || ggml_is_quantized(t->type)) { continue; @@ -168,6 +221,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: @@ -189,8 +243,6 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const // 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; - // Todo: Use the XMX kernel if possible: - // If there are no tensor cores available, use the generic tile kernel: if (can_use_vector_kernel) { if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { @@ -210,7 +262,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_TILE: @@ -219,7 +273,20 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst case BEST_FATTN_KERNEL_VEC: ggml_sycl_flash_attn_ext_vec(ctx, dst); break; +#if GGML_SYCL_DNNL + // ONEDNN dispatch reserved for PR #25222: + // case BEST_FATTN_KERNEL_ONEDNN: + // ggml_sycl_flash_attn_ext_onednn(ctx, dst); + // break; + case BEST_FATTN_KERNEL_HYBRID: + ggml_sycl_flash_attn_ext_hybrid(ctx, dst); + break; +#endif + case BEST_FATTN_KERNEL_MKL: + ggml_sycl_flash_attn_ext_mkl(ctx, dst); + break; } + } bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) { diff --git a/ggml/src/ggml-sycl/fattn.hpp b/ggml/src/ggml-sycl/fattn.hpp index f2a8ffc97dee..c093970a3fed 100644 --- a/ggml/src/ggml-sycl/fattn.hpp +++ b/ggml/src/ggml-sycl/fattn.hpp @@ -19,4 +19,6 @@ 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); +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 41449db665ec..e4070149a381 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -90,6 +90,7 @@ int g_ggml_sycl_use_async_mem_op_requested = 1; int g_ggml_sycl_use_level_zero_api = 0; int g_ggml_sycl_enable_flash_attention = 1; int g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL; +int g_ggml_sycl_fa_onednn = 1; int g_ggml_sycl_usm_system = 0; static ggml_sycl_device_info ggml_sycl_init() { @@ -273,6 +274,7 @@ static void ggml_check_sycl() try { g_ggml_sycl_disable_optimize = ggml_sycl_get_env("GGML_SYCL_DISABLE_OPT", 0); g_ggml_sycl_disable_graph = ggml_sycl_get_env("GGML_SYCL_DISABLE_GRAPH", 1); g_ggml_sycl_disable_dnn = ggml_sycl_get_env("GGML_SYCL_DISABLE_DNN", 0); + g_ggml_sycl_fa_onednn = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN", 1); g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1); g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0); @@ -341,6 +343,7 @@ static void ggml_check_sycl() try { #endif #if GGML_SYCL_DNNL GGML_LOG_INFO(" GGML_SYCL_DISABLE_DNN: %d\n", g_ggml_sycl_disable_dnn); + GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn); #else GGML_LOG_INFO(" GGML_SYCL_DISABLE_DNN: DNN disabled by compile flag\n"); #endif