From cd74ef6274e55b1591373637cc35324af000a5ee Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Fri, 25 Sep 2026 15:17:51 +0800 Subject: [PATCH] [SYCL] support sparse FA (#28796) * fix conflict * fix format issue * rm unused code --- docs/backend/SYCL.md | 3 + ggml/src/ggml-sycl/fattn-sparse.cpp | 267 ++++++++++++++++++++++++++++ ggml/src/ggml-sycl/fattn-sparse.hpp | 10 ++ ggml/src/ggml-sycl/fattn.cpp | 7 +- ggml/src/ggml-sycl/ggml-sycl.cpp | 12 ++ 5 files changed, 298 insertions(+), 1 deletion(-) create mode 100644 ggml/src/ggml-sycl/fattn-sparse.cpp create mode 100644 ggml/src/ggml-sycl/fattn-sparse.hpp diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md index 91b2097419..f9da1ab3ed 100644 --- a/docs/backend/SYCL.md +++ b/docs/backend/SYCL.md @@ -811,6 +811,9 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm | 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. | | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. | | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. | +| GGML_SYCL_SPARSE_FA | 0 (default) or 1 | Enable Sparse Flash-attention.| +| GGML_SYCL_SPARSE_FA_DEBUG | 0 (default) or 1 | Enable to debug for Sparse Flash-attention.| +| GGML_SYCL_SPARSE_FA_MARGIN | [0,..] default:256 | Set the margin value for Sparse Flash-attention.| | 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/fattn-sparse.cpp b/ggml/src/ggml-sycl/fattn-sparse.cpp new file mode 100644 index 0000000000..df28370946 --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-sparse.cpp @@ -0,0 +1,267 @@ +#include "fattn.hpp" +#include "fattn-sparse.hpp" + +#include +#include +#include + +static constexpr int64_t SPARSE_FA_PAD = 256; +static constexpr int64_t SPARSE_FA_MIN_RATIO = 2; + +extern int g_ggml_sycl_enable_sparse_fa; +extern int g_ggml_sycl_debug_sparse_fa; +extern int g_ggml_sycl_sparse_fa_margin; + +static int sparse_fa_enabled(void) { + return g_ggml_sycl_enable_sparse_fa; +} + +static int sparse_fa_debug(void) { + return g_ggml_sycl_debug_sparse_fa; +} + +// slack above n_kv_max; callers may exceed the hint by a few always-attended positions +static int sparse_fa_margin(void) { + return g_ggml_sycl_sparse_fa_margin; +} + +// Unordered output is fine: softmax over the selected set is permutation invariant. +static void sparse_fa_compact_mask(sycl::queue * stream, + const sycl::half * __restrict__ mask, + int32_t * __restrict__ indices, + int32_t * __restrict__ count, + const int64_t n_kv, + const int64_t n_kv_g) { + constexpr size_t WG = 256; + const size_t global = (size_t) GGML_PAD(n_kv, (int64_t) WG); + + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(global), sycl::range<1>(WG)), + [=](sycl::nd_item<1> item) { + const int64_t i = (int64_t) item.get_global_id(0); + if (i >= n_kv || !sycl::isfinite((float) mask[i])) { + return; + } + + sycl::atomic_ref ctr(*count); + + const int32_t pos = ctr.fetch_add(1); + if (pos < (int32_t) n_kv_g) { + indices[pos] = (int32_t) i; + } + }); +} + +// Rows along ne[0] are contiguous for every type used as a KV cache, so this is +// a plain byte copy and needs no per-type code. Padding slots are zeroed. +static void sparse_fa_gather_rows(sycl::queue * stream, + const uint8_t * __restrict__ src, + uint8_t * __restrict__ dst, + const int32_t * __restrict__ indices, + const int32_t * __restrict__ count, + const size_t row_size, + const size_t src_nb1, + const size_t src_nb2, + const int64_t n_kv_g, + const int64_t n_head) { + GGML_ASSERT(row_size % sizeof(uint32_t) == 0); + const size_t words = row_size / sizeof(uint32_t); + + stream->parallel_for( + sycl::range<3>((size_t) n_head, (size_t) n_kv_g, words), + [=](sycl::id<3> id) { + const int64_t h = (int64_t) id[0]; + const int64_t slot = (int64_t) id[1]; + const size_t w = id[2]; + + uint32_t * dst_row = + (uint32_t *) (dst + ((size_t) (h * n_kv_g + slot)) * row_size); + + if (slot >= (int64_t) *count) { + dst_row[w] = 0; + return; + } + + const uint32_t * src_row = + (const uint32_t *) (src + (size_t) indices[slot] * src_nb1 + + (size_t) h * src_nb2); + dst_row[w] = src_row[w]; + }); +} + +static void sparse_fa_gather_mask(sycl::queue * stream, + const sycl::half * __restrict__ mask, + sycl::half * __restrict__ mask_g, + const int32_t * __restrict__ indices, + const int32_t * __restrict__ count, + const int64_t n_kv_g, + const int64_t n_rows, + const size_t mask_s1) { + stream->parallel_for( + sycl::range<2>((size_t) n_rows, (size_t) n_kv_g), + [=](sycl::id<2> id) { + const int64_t r = (int64_t) id[0]; + const int64_t slot = (int64_t) id[1]; + + sycl::half v = sycl::half(-INFINITY); + if (slot < (int64_t) *count) { + v = mask[(size_t) r * mask_s1 + (size_t) indices[slot]]; + } + mask_g[(size_t) r * n_kv_g + slot] = v; + }); +} + +static bool sparse_fa_applicable(const ggml_tensor * dst, int64_t & n_kv_g_out) { + 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]; + + if (!Q || !K || !V || !mask) { + return false; + } + + const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4); + if (n_kv_max <= 0) { + return false; + } + + float max_bias = 0.0f; + float 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; + } + + // single-token decode only; prefill amortises the scan already + if (Q->ne[1] != 1) { + return false; + } + if (K->ne[3] != 1 || V->ne[3] != 1 || mask->ne[2] != 1 || mask->ne[3] != 1) { + return false; + } + if (mask->type != GGML_TYPE_F16 || mask->ne[0] < K->ne[1]) { + return false; + } + if (K->ne[2] != V->ne[2]) { + return false; + } + + // nb[1] may stride over heads (interleaved cache); only ne[0] must be contiguous + if (K->nb[0] != ggml_type_size(K->type) || V->nb[0] != ggml_type_size(V->type)) { + return false; + } + + const size_t k_row = ggml_row_size(K->type, K->ne[0]); + const size_t v_row = ggml_row_size(V->type, V->ne[0]); + if (k_row % sizeof(uint32_t) || v_row % sizeof(uint32_t)) { + return false; + } + + const int64_t n_kv_g = GGML_PAD((int64_t) n_kv_max + sparse_fa_margin(), SPARSE_FA_PAD); + if (n_kv_g * SPARSE_FA_MIN_RATIO > K->ne[1]) { + return false; + } + + n_kv_g_out = n_kv_g; + return true; +} + +bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + int64_t n_kv_g = 0; + if (!sparse_fa_enabled() || !sparse_fa_applicable(dst, n_kv_g)) { + return false; + } + + ggml_tensor * K = dst->src[1]; + ggml_tensor * V = dst->src[2]; + ggml_tensor * mask = dst->src[3]; + + const int64_t n_kv = K->ne[1]; + const int64_t n_head_k = K->ne[2]; + const int64_t n_rows_m = mask->ne[1]; + + const size_t k_row = ggml_row_size(K->type, K->ne[0]); + const size_t v_row = ggml_row_size(V->type, V->ne[0]); + + dpct::queue_ptr stream = ctx.stream(); + + ggml_sycl_pool_alloc idx_alloc(ctx.pool(), (size_t) n_kv_g); + ggml_sycl_pool_alloc cnt_alloc(ctx.pool(), 1); + ggml_sycl_pool_alloc k_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * k_row); + ggml_sycl_pool_alloc v_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * v_row); + ggml_sycl_pool_alloc m_alloc(ctx.pool(), (size_t) n_rows_m * n_kv_g); + + int32_t * d_idx = idx_alloc.get(); + int32_t * d_cnt = cnt_alloc.get(); + uint8_t * d_K = k_alloc.get(); + uint8_t * d_V = v_alloc.get(); + sycl::half * d_mask = m_alloc.get(); + + SYCL_CHECK(CHECK_TRY_ERROR(stream->memset(d_cnt, 0, sizeof(int32_t)))); + + sparse_fa_compact_mask(stream, (const sycl::half *) mask->data, + d_idx, d_cnt, n_kv, n_kv_g); + + sparse_fa_gather_rows(stream, (const uint8_t *) K->data, d_K, d_idx, d_cnt, + k_row, K->nb[1], K->nb[2], n_kv_g, n_head_k); + + sparse_fa_gather_rows(stream, (const uint8_t *) V->data, d_V, d_idx, d_cnt, + v_row, V->nb[1], V->nb[2], n_kv_g, n_head_k); + + sparse_fa_gather_mask(stream, (const sycl::half *) mask->data, d_mask, + d_idx, d_cnt, n_kv_g, n_rows_m, + mask->nb[1] / sizeof(sycl::half)); + + if (sparse_fa_debug()) { + int32_t h_cnt = 0; + SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(&h_cnt, d_cnt, sizeof(int32_t)))); + SYCL_CHECK(CHECK_TRY_ERROR(stream->wait())); + fprintf(stderr, "[FA-SPARSE] n_kv=%lld n_kv_max=%d n_kv_g=%lld finite=%d%s\n", + (long long) n_kv, ggml_get_op_params_i32(dst, 4), + (long long) n_kv_g, (int) h_cnt, + h_cnt > (int32_t) n_kv_g ? " OVERFLOW" : ""); + } + + // shallow copies retargeted at the gathered buffers; kernels are unchanged + ggml_tensor K_g = *K; + K_g.data = d_K; + K_g.ne[1] = n_kv_g; + K_g.nb[1] = k_row; + K_g.nb[2] = (size_t) n_kv_g * k_row; + K_g.nb[3] = (size_t) n_head_k * n_kv_g * k_row; + K_g.view_src = nullptr; + K_g.view_offs = 0; + + ggml_tensor V_g = *V; + V_g.data = d_V; + V_g.ne[1] = n_kv_g; + V_g.nb[1] = v_row; + V_g.nb[2] = (size_t) n_kv_g * v_row; + V_g.nb[3] = (size_t) V->ne[2] * n_kv_g * v_row; + V_g.view_src = nullptr; + V_g.view_offs = 0; + + ggml_tensor M_g = *mask; + M_g.data = d_mask; + M_g.ne[0] = n_kv_g; + M_g.nb[1] = (size_t) n_kv_g * sizeof(sycl::half); + M_g.nb[2] = M_g.nb[1] * mask->ne[1]; + M_g.nb[3] = M_g.nb[2]; + M_g.view_src = nullptr; + M_g.view_offs = 0; + + ggml_tensor dst_g = *dst; + dst_g.src[1] = &K_g; + dst_g.src[2] = &V_g; + dst_g.src[3] = &M_g; + dst_g.op_params[4] = 0; // avoid re-entering this path + + ggml_sycl_flash_attn_ext(ctx, &dst_g); + + return true; +} diff --git a/ggml/src/ggml-sycl/fattn-sparse.hpp b/ggml/src/ggml-sycl/fattn-sparse.hpp new file mode 100644 index 0000000000..98b06bddc4 --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-sparse.hpp @@ -0,0 +1,10 @@ +#ifndef GGML_SYCL_FATTN_SPARSE_HPP +#define GGML_SYCL_FATTN_SPARSE_HPP + +#include "common.hpp" + +// Gather the K/V rows selected by a sparse mask and re-dispatch the dense +// kernels onto them. Returns false if the caller should use the dense path. +bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +#endif // GGML_SYCL_FATTN_SPARSE_HPP diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index 394cda593f..541ae8a82f 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -19,7 +19,7 @@ #include "fattn-vec.hpp" #include "fattn.hpp" #include "fattn-onednn.hpp" - +#include "fattn-sparse.hpp" #define FATTN_VEC_CASE(D, type_K, type_V) \ { \ @@ -276,6 +276,11 @@ 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); + // sparse nodes are gathered down to n_kv_max rows and re-dispatched here + if (ggml_sycl_flash_attn_ext_sparse(ctx, dst)) { + return; + } + // n_kv watchdog: log when n_kv differs from the last FA call with // the same D — helps detect cache-truncation issues. static int nkv_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index e13ec85251..99ebdb3e15 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -113,6 +113,9 @@ int g_ggml_sycl_usm_system = 0; int g_ggml_sycl_enable_host_pinned_mem = 1; int g_ggml_sycl_host_pinned_mem_2g = 0; int g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_LEVEL_ZERO; +int g_ggml_sycl_enable_sparse_fa = 0; +int g_ggml_sycl_debug_sparse_fa = 0; +int g_ggml_sycl_sparse_fa_margin = 256; static ggml_sycl_device_info ggml_sycl_init() { GGML_SYCL_DEBUG("[SYCL] call ggml_sycl_init\n"); @@ -384,6 +387,10 @@ static void ggml_check_sycl() try { g_ggml_sycl_host_pinned_mem_2g = ggml_sycl_get_env("GGML_SYCL_HOST_PINNED_MEM_2G", 0) & g_ggml_sycl_enable_host_pinned_mem; + g_ggml_sycl_enable_sparse_fa = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA", 0); + g_ggml_sycl_debug_sparse_fa = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA_DEBUG", 0); + g_ggml_sycl_sparse_fa_margin = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA_MARGIN", 256); + GGML_SYCL_DEBUG("[SYCL] call ggml_check_sycl\n"); GGML_LOG_INFO("Build with Macros:\n"); @@ -422,6 +429,7 @@ static void ggml_check_sycl() try { GGML_LOG_INFO(" GGML_SYCL_SUPPORT_VMM: no\n"); #endif + //Print the running environment variables for SYCL backend GGML_LOG_INFO("Running with Environment Variables:\n"); GGML_LOG_INFO(" GGML_SYCL_DEBUG: %d\n", g_ggml_sycl_debug); GGML_LOG_INFO(" GGML_SYCL_DEV_DEBUG: %d\n", g_ggml_sycl_dev_debug); @@ -491,6 +499,10 @@ static void ggml_check_sycl() try { GGML_LOG_INFO(" GGML_SYCL_ENABLE_HOST_PINNED_MEM: %d\n", g_ggml_sycl_enable_host_pinned_mem); GGML_LOG_INFO(" GGML_SYCL_HOST_PINNED_MEM_2G: %d\n", g_ggml_sycl_host_pinned_mem_2g); + GGML_LOG_INFO(" GGML_SYCL_SPARSE_FA: %d\n", g_ggml_sycl_enable_sparse_fa); + GGML_LOG_INFO(" GGML_SYCL_SPARSE_FA_DEBUG: %d\n", g_ggml_sycl_debug_sparse_fa); + GGML_LOG_INFO(" GGML_SYCL_SPARSE_FA_MARGIN: %d\n", g_ggml_sycl_sparse_fa_margin); + /* NOT REMOVE, keep it for next optimize for XMX. #if defined(SYCL_USE_XMX) fprintf(stderr, "%s: SYCL_USE_XMX: yes\n", __func__);