const float * beta,
const float * curr_state,
float * dst,
+ float * state,
int64_t H,
int64_t n_tokens,
int64_t n_seqs,
const uint3 neqk1_magic,
const uint3 rq3_magic,
float scale,
+ int64_t state_slot_stride,
int K) {
const uint32_t h_idx = blockIdx.x;
const uint32_t sequence = blockIdx.y;
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.
if constexpr (keep_rs_t) {
// snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
// When n_tokens < K only slots 0..n_tokens-1 are written; older slots are caller-owned.
- const int64_t state_size_per_token = S_v * S_v * H * n_seqs; // per-slot stride in output
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;
static void launch_gated_delta_net(
const float * q_d, const float * k_d, const float * v_d,
const float * g_d, const float * b_d, const float * s_d,
- float * dst_d,
+ float * dst_d, float * state_d,
int64_t S_v, int64_t H, int64_t n_tokens, int64_t n_seqs,
int64_t sq1, int64_t sq2, int64_t sq3,
int64_t sv1, int64_t sv2, int64_t sv3,
int64_t sb1, int64_t sb2, int64_t sb3,
int64_t neqk1, int64_t rq3,
- float scale, int K, cudaStream_t stream) {
+ float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
//TODO: Add chunked kernel for even faster pre-fill
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
const int num_warps = 4;
const uint3 neqk1_magic = init_fastdiv_values(neqk1);
const uint3 rq3_magic = init_fastdiv_values(rq3);
- int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
-
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream);
switch (S_v) {
case 16:
ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t>, launch_params,
- q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+ q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
break;
case 32:
ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t>, launch_params,
- q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+ q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
break;
case 64: {
ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t>, launch_params,
- q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+ q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
break;
}
case 128: {
ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t>, launch_params,
- q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+ q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
break;
}
default:
}
}
-void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+static void ggml_cuda_op_gated_delta_net_impl(
+ ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_cuda_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];
const int K = ggml_get_op_params_i32(dst, 0);
const bool keep_rs = K > 1;
+ // recurrent state -> gdn_out 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_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+ ggml_cuda_op_gated_delta_net_impl(ctx, dst, nullptr);
+}
+
+void ggml_cuda_op_gated_delta_net_fused_cache(
+ ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_cuda_gated_delta_net_fused_cache cache) {
+ ggml_cuda_op_gated_delta_net_impl(ctx, dst, &cache);
+}
GGML_UNUSED(backend);
}
+static bool ggml_cuda_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;
+}
+
#ifdef USE_CUDA_GRAPH
static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
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_cuda_is_view_or_noop(node)) {
continue;
}
return true;
}
+// 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.
+static int ggml_cuda_try_gdn_cache_fusion(
+ const ggml_cgraph * cgraph, int node_idx, ggml_cuda_gated_delta_net_fused_cache & fused_state_cpy) {
+ const ggml_tensor * gdn = cgraph->nodes[node_idx];
+ // the kernel skips the snapshot tail, so the gdn output must not be a graph output
+ 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);
+
+ // snapshot cpy is the first real node after the gdn (skip views/no-ops)
+ 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_cuda_is_view_or_noop(n)) {
+ 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; require nb[1] == D (the per-seq stride the kernel
+ // assumes). ggml_cpy pins src to the same element count.
+ 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 bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
args.sigmoid = false;
args.softmax = false;
ggml_tensor * node = cgraph->nodes[i];
+ // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
+ if (node->op == GGML_OP_GATED_DELTA_NET) {
+ ggml_cuda_gated_delta_net_fused_cache fused_state_cpy;
+ const int nodes_to_skip = ggml_cuda_try_gdn_cache_fusion(cgraph, i, fused_state_cpy);
+ if (nodes_to_skip > 0) {
+#ifdef GGML_CUDA_DEBUG
+ GGML_LOG_INFO("%s: fused gated_delta_net snapshot copies for %s (skipped %d nodes)\n",
+ __func__, node->name, nodes_to_skip);
+#endif
+ ggml_cuda_op_gated_delta_net_fused_cache(*cuda_ctx, node, fused_state_cpy);
+ return nodes_to_skip;
+ }
+ }
+
//topk-moe
if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX ||
cgraph->nodes[i]->op == GGML_OP_ARGSORT) {
#endif
prev_i = 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_cuda_is_view_or_noop(node)) {
continue;
}
overview / pkg / ggml / sources / llama.cpp / commitdiff