[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:
co-authored by
maxious
scientist3
hmscider
parent
12127defda
commit
32b741c336
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user