mirror of
https://github.com/ggml-org/whisper.cpp.git
synced 2026-10-09 15:55:37 +02:00
webgpu : add fused gdn + cpy (llama/28976)
This commit is contained in:
committed by
Georgi Gerganov
parent
4ef9fd70b4
commit
29e710c4a1
@@ -1383,7 +1383,8 @@ static webgpu_encoded_op ggml_webgpu_gated_delta_net(webgpu_context & ctx,
|
||||
ggml_tensor * src3,
|
||||
ggml_tensor * src4,
|
||||
ggml_tensor * src5,
|
||||
ggml_tensor * dst) {
|
||||
ggml_tensor * dst,
|
||||
ggml_tensor * dst_fuse) {
|
||||
ggml_webgpu_shader_lib_context shader_lib_ctx = {};
|
||||
shader_lib_ctx.src0 = src0;
|
||||
shader_lib_ctx.src1 = src1;
|
||||
@@ -1391,6 +1392,7 @@ static webgpu_encoded_op ggml_webgpu_gated_delta_net(webgpu_context & ctx,
|
||||
shader_lib_ctx.src3 = src3;
|
||||
shader_lib_ctx.src4 = src4;
|
||||
shader_lib_ctx.dst = dst;
|
||||
shader_lib_ctx.dst_fuse = dst_fuse;
|
||||
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
|
||||
|
||||
webgpu_pipeline pipeline = ctx->shader_lib->get_gated_delta_net_pipeline(shader_lib_ctx);
|
||||
@@ -1426,6 +1428,8 @@ static webgpu_encoded_op ggml_webgpu_gated_delta_net(webgpu_context & ctx,
|
||||
(uint32_t) (src2->ne[3] / src0->ne[3]),
|
||||
K,
|
||||
scale_u32,
|
||||
dst_fuse ? (uint32_t) (dst_fuse->nb[2] / ggml_type_size(dst_fuse->type)) : 0,
|
||||
dst_fuse ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst_fuse) / ggml_type_size(dst_fuse->type)) : 0,
|
||||
};
|
||||
|
||||
std::vector<wgpu::BindGroupEntry> entries = {
|
||||
@@ -1435,6 +1439,10 @@ static webgpu_encoded_op ggml_webgpu_gated_delta_net(webgpu_context & ctx,
|
||||
ggml_webgpu_make_tensor_bind_group_entry(ctx, 6, dst),
|
||||
};
|
||||
|
||||
if (dst_fuse) {
|
||||
entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 7, dst_fuse));
|
||||
}
|
||||
|
||||
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, h, n_seqs);
|
||||
}
|
||||
|
||||
@@ -3220,6 +3228,67 @@ static bool ggml_webgpu_can_fuse_rms_norm_mul(const struct ggml_cgraph * cgraph,
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_webgpu_can_fuse_gdn_cache(const struct ggml_cgraph * cgraph, int node_idx, int & num_encoded_ops) {
|
||||
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 false;
|
||||
}
|
||||
|
||||
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 cpy_idx = 0;
|
||||
for (int j = node_idx + 1; j < cgraph->n_nodes && cpy == nullptr; ++j) {
|
||||
const ggml_tensor * n = cgraph->nodes[j];
|
||||
if (ggml_op_is_empty(n->op) || ggml_is_empty(n)) {
|
||||
continue;
|
||||
}
|
||||
if (n->op != GGML_OP_CPY || (n->flags & GGML_TENSOR_FLAG_OUTPUT)) {
|
||||
return false;
|
||||
}
|
||||
cpy = n;
|
||||
cpy_idx = j;
|
||||
}
|
||||
if (cpy == nullptr) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const ggml_tensor * cpy_src = cpy->src[0]; // view of the gdn snapshot tail
|
||||
const ggml_tensor * cpy_dst = cpy->src[1]; // cache view the kernel writes to
|
||||
|
||||
// src must be this gdn's snapshot tail (contiguous, at the tail offset)
|
||||
if (cpy_src->op != GGML_OP_VIEW || cpy_src->view_src != gdn || cpy_src->view_offs != tail_off ||
|
||||
!ggml_is_contiguous(cpy_src)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// 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 (cpy_dst->op != GGML_OP_VIEW || cpy_dst->type != GGML_TYPE_F32 || cpy_dst->data == nullptr ||
|
||||
!std::equal(expected_ne.begin(), expected_ne.end(), cpy_dst->ne) ||
|
||||
cpy_dst->nb[0] != ggml_type_size(GGML_TYPE_F32) || cpy_dst->nb[1] != (size_t) ggml_row_size(GGML_TYPE_F32, D)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
num_encoded_ops = cpy_idx - node_idx + 1;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static webgpu_encoded_op ggml_webgpu_upscale(webgpu_context ctx, ggml_tensor * src, ggml_tensor * dst) {
|
||||
const uint32_t mode_flags = (uint32_t) ggml_get_op_params_i32(dst, 0);
|
||||
std::vector<uint32_t> params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)),
|
||||
@@ -3358,7 +3427,14 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
|
||||
return ggml_webgpu_ssm_scan(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5], node->src[6],
|
||||
node);
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
return ggml_webgpu_gated_delta_net(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5], node);
|
||||
if (ggml_webgpu_can_fuse_gdn_cache(cgraph, node_idx, num_encoded_ops)) {
|
||||
ggml_tensor * dst_fuse = cgraph->nodes[node_idx + num_encoded_ops - 1]->src[1];
|
||||
return ggml_webgpu_gated_delta_net(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5],
|
||||
node, dst_fuse);
|
||||
} else {
|
||||
return ggml_webgpu_gated_delta_net(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5],
|
||||
node, nullptr);
|
||||
}
|
||||
case GGML_OP_PAD:
|
||||
return ggml_webgpu_pad(ctx, src0, node);
|
||||
case GGML_OP_ARGMAX:
|
||||
|
||||
Reference in New Issue
Block a user