mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-15 12:08:31 +02:00
sycl: fuse the gated-delta-net state writeback cpy (#26643)
Port of https://github.com/ggml-org/llama.cpp/pull/23940. Arc Pro B70, Qwen 3.6 27B Q4_K - Medium (48 of its 64 blocks run gated_delta_net), -ngl 99 -fa 1 -ctk f16 -ctv f16 -b 2048 -ub 2048, interleaved A/B passes of r=3: tg128 23.91 / 23.90 / 23.90 -> 24.19 / 24.17 / 24.20 +1.2% tg128 (rebuild) 23.81 / 23.81 -> 24.09 / 24.10 +1.2% pp2048 1050.8 / 1053.9 -> 1053.8 / 1054.5 flat 2 seqs, tg128 32.73 / 32.75 -> 33.11 / 33.10 +1.1%
This commit is contained in:
@@ -14,9 +14,9 @@ void gated_delta_net_sycl(const float * q,
|
||||
const float * beta,
|
||||
const float * curr_state,
|
||||
float * dst,
|
||||
float * state,
|
||||
int64_t H,
|
||||
int64_t n_tokens,
|
||||
int64_t n_seqs,
|
||||
int64_t sq1,
|
||||
int64_t sq2,
|
||||
int64_t sq3,
|
||||
@@ -29,6 +29,7 @@ void gated_delta_net_sycl(const float * q,
|
||||
const sycl::uint3 neqk1_magic,
|
||||
const sycl::uint3 rq3_magic,
|
||||
float scale,
|
||||
int64_t state_slot_stride,
|
||||
int K) {
|
||||
auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
|
||||
const uint32_t h_idx = item_ct1.get_group(2);
|
||||
@@ -40,15 +41,12 @@ void gated_delta_net_sycl(const float * q,
|
||||
const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
|
||||
const uint32_t iq3 = fastdiv(sequence, rq3_magic);
|
||||
|
||||
const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
|
||||
float * attn_data = dst;
|
||||
float * state = dst + attn_score_elems;
|
||||
|
||||
// input state holds s0 only [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v.
|
||||
// output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
|
||||
const int64_t state_in_offset = sequence * H * S_v * S_v + h_idx * S_v * S_v;
|
||||
const int64_t state_out_offset = (sequence * H + h_idx) * S_v * S_v;
|
||||
const int64_t state_size_per_token = S_v * S_v * H * n_seqs; // per-slot stride in output
|
||||
state += state_out_offset;
|
||||
curr_state += state_in_offset + col * S_v;
|
||||
attn_data += (sequence * n_tokens * H + h_idx) * S_v;
|
||||
@@ -145,7 +143,7 @@ void gated_delta_net_sycl(const float * q,
|
||||
if constexpr (keep_rs_t) {
|
||||
const int target_slot = (int) n_tokens - 1 - t;
|
||||
if (target_slot >= 0 && target_slot < K) {
|
||||
float * curr_state = (dst + attn_score_elems) + target_slot * state_size_per_token + state_out_offset;
|
||||
float * curr_state = state + target_slot * state_slot_stride;
|
||||
#pragma unroll
|
||||
for (int r = 0; r < rows_per_lane; r++) {
|
||||
const int i = r * warp_size + lane;
|
||||
@@ -172,6 +170,7 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
const float * b_d,
|
||||
const float * s_d,
|
||||
float * dst_d,
|
||||
float * state_d,
|
||||
int64_t S_v,
|
||||
int64_t H,
|
||||
int64_t n_tokens,
|
||||
@@ -188,6 +187,7 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
int64_t neqk1,
|
||||
int64_t rq3,
|
||||
float scale,
|
||||
int64_t state_slot_stride,
|
||||
int K,
|
||||
dpct::queue_ptr stream) {
|
||||
//TODO: Add chunked kernel for even faster pre-fill
|
||||
@@ -206,9 +206,9 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
constexpr int sv = 16;
|
||||
stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
|
||||
[=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
|
||||
gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens,
|
||||
n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
|
||||
sb3, neqk1_magic, rq3_magic, scale, K);
|
||||
gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens,
|
||||
sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
|
||||
sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
|
||||
});
|
||||
}
|
||||
break;
|
||||
@@ -217,9 +217,9 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
constexpr int sv = 32;
|
||||
stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
|
||||
[=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
|
||||
gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens,
|
||||
n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
|
||||
sb3, neqk1_magic, rq3_magic, scale, K);
|
||||
gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens,
|
||||
sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
|
||||
sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
|
||||
});
|
||||
}
|
||||
break;
|
||||
@@ -229,8 +229,8 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
|
||||
[=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
|
||||
gated_delta_net_sycl<sv, KDA, keep_rs_t>(
|
||||
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, n_seqs, sq1, sq2,
|
||||
sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
|
||||
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, sq1, sq2,
|
||||
sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
|
||||
});
|
||||
}
|
||||
break;
|
||||
@@ -241,8 +241,8 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
|
||||
[=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
|
||||
gated_delta_net_sycl<sv, KDA, keep_rs_t>(
|
||||
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, n_seqs, sq1, sq2,
|
||||
sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
|
||||
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, sq1, sq2,
|
||||
sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
|
||||
});
|
||||
}
|
||||
break;
|
||||
@@ -253,7 +253,8 @@ static void launch_gated_delta_net(const float * q_d,
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
|
||||
static void ggml_sycl_op_gated_delta_net_impl(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
|
||||
const ggml_sycl_gated_delta_net_fused_cache * cache) {
|
||||
ggml_tensor * src_q = dst->src[0];
|
||||
ggml_tensor * src_k = dst->src[1];
|
||||
ggml_tensor * src_v = dst->src[2];
|
||||
@@ -318,30 +319,48 @@ void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor *
|
||||
const int K = ggml_get_op_params_i32(dst, 0);
|
||||
const bool keep_rs = K > 1;
|
||||
|
||||
// recurrent state -> dst tail (after attention scores), or the cache when fusing
|
||||
float * state_d = dst_d + S_v * H * n_tokens * n_seqs;
|
||||
int64_t state_slot_stride = S_v * S_v * H * n_seqs;
|
||||
if (cache != nullptr) {
|
||||
state_d = cache->data;
|
||||
state_slot_stride = cache->slot_stride;
|
||||
}
|
||||
|
||||
if (kda) {
|
||||
if (keep_rs) {
|
||||
launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
|
||||
launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
|
||||
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
|
||||
} else {
|
||||
launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
|
||||
launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
|
||||
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
|
||||
}
|
||||
} else {
|
||||
if (keep_rs) {
|
||||
launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
|
||||
launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
|
||||
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
|
||||
} else {
|
||||
launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
|
||||
launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
|
||||
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
|
||||
sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
|
||||
ggml_sycl_op_gated_delta_net_impl(ctx, dst, nullptr);
|
||||
}
|
||||
|
||||
void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
|
||||
scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6);
|
||||
ggml_sycl_op_gated_delta_net(ctx, dst);
|
||||
}
|
||||
|
||||
void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
|
||||
ggml_sycl_gated_delta_net_fused_cache cache) {
|
||||
scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6);
|
||||
ggml_sycl_op_gated_delta_net_impl(ctx, dst, &cache);
|
||||
}
|
||||
|
||||
@@ -5,5 +5,15 @@
|
||||
#include "common.hpp"
|
||||
#include "ggml.h"
|
||||
|
||||
// fused-kernel recurrent-state output; strides in elements (per-seq stride is always D, set in-kernel)
|
||||
struct ggml_sycl_gated_delta_net_fused_cache {
|
||||
float * data; // rollback slot 0
|
||||
int64_t slot_stride; // between rollback slots (0 when K==1)
|
||||
};
|
||||
|
||||
void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
|
||||
void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
|
||||
|
||||
// same op, but writes the snapshot(s) into the cache instead of dst (see ggml_sycl_try_gdn_cache_fusion)
|
||||
void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
|
||||
ggml_sycl_gated_delta_net_fused_cache cache);
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
//
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <assert.h>
|
||||
#include <atomic>
|
||||
#include <cinttypes>
|
||||
@@ -5464,12 +5465,90 @@ catch (sycl::exception const &exc) {
|
||||
std::exit(1);
|
||||
}
|
||||
|
||||
static bool ggml_sycl_is_view_or_noop(const ggml_tensor * t) {
|
||||
return ggml_is_empty(t) || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_TRANSPOSE ||
|
||||
t->op == GGML_OP_VIEW || t->op == GGML_OP_PERMUTE || t->op == GGML_OP_NONE;
|
||||
}
|
||||
|
||||
// match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
|
||||
// (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
|
||||
// returns the number of following nodes to skip (0 = no fusion)
|
||||
// ported from ggml_cuda_try_gdn_cache_fusion - pure graph inspection, backend-agnostic
|
||||
static int ggml_sycl_try_gdn_cache_fusion(const ggml_cgraph * cgraph, int node_idx,
|
||||
ggml_sycl_gated_delta_net_fused_cache & fused_state_cpy) {
|
||||
if (!g_ggml_sycl_enable_fusion) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const ggml_tensor * gdn = cgraph->nodes[node_idx];
|
||||
// the kernel skips the snapshot tail, so the gdn output must not be a graph output, and the cpy
|
||||
// found below is taken to be its only reader, as it is in every graph that builds this op
|
||||
if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->type != GGML_TYPE_F32 ||
|
||||
(gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const ggml_tensor * src_v = gdn->src[2];
|
||||
const int64_t S_v = src_v->ne[0];
|
||||
const int64_t H = src_v->ne[1];
|
||||
const int64_t n_tokens = src_v->ne[2];
|
||||
const int64_t n_seqs = src_v->ne[3];
|
||||
const int64_t D = S_v * S_v * H;
|
||||
const int64_t K = ggml_get_op_params_i32(gdn, 0); // snapshot slot count
|
||||
const int64_t n_written = std::min<int64_t>(n_tokens, K); // newest n_written slots are written
|
||||
|
||||
// snapshot tail starts right after the attention scores
|
||||
const size_t tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs);
|
||||
|
||||
// the cpy must be the first node the compute loop below runs, so nothing can read the cache first.
|
||||
// skip exactly what that loop skips: views, no-ops, and nodes the graph does not compute.
|
||||
const ggml_tensor * cpy = nullptr;
|
||||
int skip = 0;
|
||||
for (int j = node_idx + 1; j < cgraph->n_nodes && cpy == nullptr; ++j) {
|
||||
const ggml_tensor * n = cgraph->nodes[j];
|
||||
if (ggml_sycl_is_view_or_noop(n) || (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
|
||||
continue;
|
||||
}
|
||||
if (n->op != GGML_OP_CPY || (n->flags & GGML_TENSOR_FLAG_OUTPUT)) {
|
||||
return 0;
|
||||
}
|
||||
cpy = n;
|
||||
skip = j - node_idx;
|
||||
}
|
||||
if (cpy == nullptr) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const ggml_tensor * src = cpy->src[0]; // view of the gdn snapshot tail
|
||||
const ggml_tensor * dst = cpy->src[1]; // cache view the kernel writes to
|
||||
|
||||
// src must be this gdn's snapshot tail (contiguous, at the tail offset)
|
||||
if (src->op != GGML_OP_VIEW || src->view_src != gdn || src->view_offs != tail_off ||
|
||||
!ggml_is_contiguous(src)) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// dst is the [D, n_seqs, n_written] cache view, with the per-seq stride D that the kernel assumes.
|
||||
// ggml_cpy pins src to the same element count, so src needs no shape check of its own.
|
||||
const std::array<int64_t, GGML_MAX_DIMS> expected_ne = { D, n_seqs, n_written, 1 };
|
||||
if (dst->op != GGML_OP_VIEW || dst->type != GGML_TYPE_F32 || dst->data == nullptr ||
|
||||
!std::equal(expected_ne.begin(), expected_ne.end(), dst->ne) ||
|
||||
dst->nb[0] != ggml_type_size(GGML_TYPE_F32) ||
|
||||
dst->nb[1] != (size_t) ggml_row_size(GGML_TYPE_F32, D)) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
fused_state_cpy.data = (float *) dst->data; // rollback group 0 (newest)
|
||||
fused_state_cpy.slot_stride = K > 1 ? (int64_t) (dst->nb[2] / sizeof(float)) : 0;
|
||||
return skip;
|
||||
}
|
||||
|
||||
static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * sycl_ctx, ggml_cgraph * cgraph) {
|
||||
ggml_sycl_set_main_device(sycl_ctx->device);
|
||||
|
||||
for (int i = 0; i < cgraph->n_nodes; i++) {
|
||||
ggml_tensor * node = cgraph->nodes[i];
|
||||
if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) {
|
||||
if (ggml_sycl_is_view_or_noop(node)) {
|
||||
continue;
|
||||
}
|
||||
if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
|
||||
@@ -5489,6 +5568,16 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
|
||||
}
|
||||
}
|
||||
#endif
|
||||
// gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
|
||||
if (node->op == GGML_OP_GATED_DELTA_NET) {
|
||||
ggml_sycl_gated_delta_net_fused_cache fused_state_cpy;
|
||||
const int gdn_nodes_to_skip = ggml_sycl_try_gdn_cache_fusion(cgraph, i, fused_state_cpy);
|
||||
if (gdn_nodes_to_skip > 0) {
|
||||
ggml_sycl_op_gated_delta_net_fused_cache(*sycl_ctx, node, fused_state_cpy);
|
||||
i += gdn_nodes_to_skip;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
if (node->op == GGML_OP_RMS_NORM &&
|
||||
ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) {
|
||||
ggml_sycl_op_rms_norm_fused(*sycl_ctx, node, cgraph->nodes[i + 1]);
|
||||
|
||||
Reference in New Issue
Block a user