[SYCL] Flash Attention with XMX engine via oneDNN (#25222)

* [SYCL] F16 (default) Flash Attention with XMX engine via oneDNN graph API; Qwen3.6-27b-Q8_0 prefill speed up x1.21 at p=512 and x4.26 at p=80k

* [SYCL] Address review on FA oneDNN path. Result: llama-bench---pp512; 32% increase with fa1; llama-perplexity---0.11% difference; tested model: mradermacher/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf

* PR-25222 revision v2: addressed audits

* [SYCL] flash-attn oneDNN SDPA KV F16 rev 3.0: add BMG gate + multi-device sync. Narrow the scrope of this PR to Battlemage only (bmg; Xe2). Other archs (e.g., alchemist) fall back to existing FA kernel. When device_count >1, apply stream -> wait_and_throw(), validated working path for multi-gpu sync fix by @maxious.

Co-authored-by: maxious <81432+maxious@users.noreply.github.com>

* updated comment on bmg gate, noted the issue

---------

Co-authored-by: scientist3 <scientist.3@users.noreply.github.com>
Co-authored-by: hmscider <hmscider@users.noreply.github.com>
Co-authored-by: maxious <81432+maxious@users.noreply.github.com>
This commit is contained in:
hmscider
2026-07-15 10:26:53 +03:00
committed by GitHub
co-authored by maxious scientist3 hmscider
parent 12127defda
commit 32b741c336
5 changed files with 298 additions and 1 deletions
+14 -1
View File
@@ -18,6 +18,7 @@
#include "fattn-tile.hpp"
#include "fattn-vec.hpp"
#include "fattn.hpp"
#include "fattn-onednn.hpp"
#define FATTN_VEC_CASE(D, type_K, type_V) \
@@ -96,6 +97,7 @@ 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_TILE = 200,
};
@@ -189,7 +191,11 @@ 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:
// Fused-XMX path: oneDNN Graph SDPA (flash attention). Strictly
// additive -- taken only when statically supported, otherwise falls through to VEC/TILE below.
if (ggml_sycl_flash_attn_ext_onednn_supported(dst)) {
return BEST_FATTN_KERNEL_ONEDNN;
}
// If there are no tensor cores available, use the generic tile kernel:
if (can_use_vector_kernel) {
@@ -213,6 +219,13 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
switch (ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst)) {
case BEST_FATTN_KERNEL_NONE:
GGML_ABORT("Not support Flash-Attention");
case BEST_FATTN_KERNEL_ONEDNN:
// guarded: ggml_sycl_flash_attn_ext_onednn() is only defined under GGML_SYCL_DNNL;
// the reference must be compiled out here or the GGML_SYCL_DNNL=0 build fails to link.
#if GGML_SYCL_DNNL
ggml_sycl_flash_attn_ext_onednn(ctx, dst);
#endif
break;
case BEST_FATTN_KERNEL_TILE:
ggml_sycl_flash_attn_ext_tile(ctx, dst);
break;