mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-27 21:46:57 +02:00
[SYCL] support sparse FA (#28796)
* fix conflict * fix format issue * rm unused code
This commit is contained in:
Neo Zhang
@@ -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.<br>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. |
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
#include "fattn.hpp"
|
||||
#include "fattn-sparse.hpp"
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
|
||||
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<int32_t,
|
||||
sycl::memory_order::relaxed,
|
||||
sycl::memory_scope::device,
|
||||
sycl::access::address_space::global_space> 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<int32_t> idx_alloc(ctx.pool(), (size_t) n_kv_g);
|
||||
ggml_sycl_pool_alloc<int32_t> cnt_alloc(ctx.pool(), 1);
|
||||
ggml_sycl_pool_alloc<uint8_t> k_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * k_row);
|
||||
ggml_sycl_pool_alloc<uint8_t> v_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * v_row);
|
||||
ggml_sycl_pool_alloc<sycl::half> 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;
|
||||
}
|
||||
@@ -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
|
||||
@@ -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);
|
||||
|
||||
@@ -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__);
|
||||
|
||||
Reference in New Issue
Block a user